"""Dataset and graph-operator utilities for the MVP."""

from __future__ import annotations

from pathlib import Path
from typing import Any
from urllib.request import urlopen

import torch
import torch_geometric.transforms as transforms
from torch_geometric.datasets import Planetoid


PLANETOID_FILES = (
    "x",
    "tx",
    "allx",
    "y",
    "ty",
    "ally",
    "graph",
    "test.index",
)


def _ensure_planetoid_raw(root: str | Path, name: str) -> None:
    """Pre-fetch from raw.githubusercontent.com to avoid a flaky redirect."""
    raw_dir = Path(root) / name / "raw"
    raw_dir.mkdir(parents=True, exist_ok=True)
    lower_name = name.lower()
    base_url = (
        "https://raw.githubusercontent.com/kimiyoung/"
        "planetoid/master/data"
    )
    for suffix in PLANETOID_FILES:
        filename = f"ind.{lower_name}.{suffix}"
        destination = raw_dir / filename
        if destination.exists() and destination.stat().st_size > 0:
            continue
        temporary = destination.with_suffix(destination.suffix + ".part")
        with urlopen(f"{base_url}/{filename}", timeout=60) as response:
            temporary.write_bytes(response.read())
        temporary.replace(destination)


def spmm(matrix: torch.Tensor, dense: torch.Tensor) -> torch.Tensor:
    """Multiply a native PyTorch sparse matrix by a dense matrix."""
    return torch.sparse.mm(matrix, dense)


def normalized_adjacency(
    edge_index: torch.Tensor,
    num_nodes: int,
) -> torch.Tensor:
    """Return A-hat = D^-1/2 (A + I) D^-1/2 as a coalesced COO tensor."""
    row, col = edge_index
    loops = torch.arange(num_nodes, device=edge_index.device)
    row = torch.cat((row, loops))
    col = torch.cat((col, loops))

    degree = torch.zeros(num_nodes, device=edge_index.device)
    degree.scatter_add_(0, row, torch.ones_like(row, dtype=torch.float32))
    inverse_sqrt = degree.pow(-0.5)
    inverse_sqrt.masked_fill_(torch.isinf(inverse_sqrt), 0.0)
    values = inverse_sqrt[row] * inverse_sqrt[col]

    return torch.sparse_coo_tensor(
        torch.stack((row, col)),
        values,
        (num_nodes, num_nodes),
    ).coalesce()


def load_cora(
    root: str | Path = "data",
    device: str | torch.device = "cpu",
) -> dict[str, Any]:
    """Download Cora and return the tensors used by both training rules."""
    _ensure_planetoid_raw(root, "Cora")
    dataset = Planetoid(
        root=str(root),
        name="Cora",
        transform=transforms.NormalizeFeatures(),
    )
    graph = dataset[0]
    adjacency = normalized_adjacency(
        graph.edge_index,
        graph.num_nodes,
    ).to(device)
    return {
        "x": graph.x.to(device),
        "y": graph.y.to(device),
        "adjacency": adjacency,
        "train_mask": graph.train_mask.to(device),
        "val_mask": graph.val_mask.to(device),
        "test_mask": graph.test_mask.to(device),
        "num_nodes": graph.num_nodes,
        "num_features": dataset.num_features,
        "num_classes": dataset.num_classes,
    }