from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
from scipy import sparse
from scipy.sparse.linalg import eigsh
import torch

from .common import normalized, write_json
from .vg_diagnose import (
    aligned_text_indices,
    arbitrary_target_retrieval,
    bundle_signature,
    rank_normalize_columns,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--vision", default="artifacts/vg/vision_features.pt"
    )
    parser.add_argument("--text", default="artifacts/vg/text_features.pt")
    parser.add_argument(
        "--truth", default="artifacts/vg/ground_truth.private.jsonl"
    )
    parser.add_argument("--max-nodes", type=int, default=5_000)
    parser.add_argument("--knn", type=int, default=32)
    parser.add_argument("--eigenpairs", type=int, default=96)
    parser.add_argument("--heat-scales", type=int, default=24)
    parser.add_argument("--message-passing-steps", type=int, default=3)
    parser.add_argument("--output", default="artifacts/vg/graph_diagnostics.json")
    return parser.parse_args()


def scene_features(state: dict, indices: torch.Tensor | None = None) -> torch.Tensor:
    if "global_features" in state:
        features = state["global_features"]
    else:
        features = state["region_features"].mean(1)
    return features if indices is None else features[indices]


def exact_rank_knn_graph(
    features: torch.Tensor, neighbors: int
) -> tuple[sparse.csr_matrix, torch.Tensor]:
    """Build a symmetrized rank-weighted graph and retain density statistics."""
    x = normalized(features.float())
    similarity = x @ x.T
    similarity.fill_diagonal_(-torch.inf)
    k = min(neighbors, len(x) - 1)
    values, indices = similarity.topk(k, dim=1)

    # Rank weights avoid assuming that cosine scales agree across modalities.
    rank_weight = np.exp(
        -np.arange(k, dtype=np.float64) / max(k / 4.0, 1.0)
    )
    rows = np.repeat(np.arange(len(x), dtype=np.int64), k)
    cols = indices.cpu().numpy().reshape(-1)
    data = np.tile(rank_weight, len(x))
    directed = sparse.csr_matrix(
        (data, (rows, cols)), shape=(len(x), len(x))
    )
    adjacency = directed.maximum(directed.T)
    adjacency.setdiag(0)
    adjacency.eliminate_zeros()

    density_quantiles = torch.quantile(
        values,
        torch.tensor([0.0, 0.25, 0.5, 0.75, 1.0]),
        dim=1,
    ).T
    return adjacency, density_quantiles


def heat_kernel_signature(
    adjacency: sparse.csr_matrix, eigenpairs: int, scales: int
) -> tuple[torch.Tensor, np.ndarray]:
    """Return basis-invariant diagonal heat-kernel signatures."""
    degree = np.asarray(adjacency.sum(axis=1)).reshape(-1)
    inv_sqrt_degree = 1.0 / np.sqrt(np.maximum(degree, 1e-12))
    normalizer = sparse.diags(inv_sqrt_degree)
    normalized_adjacency = normalizer @ adjacency @ normalizer
    k = min(eigenpairs, adjacency.shape[0] - 2)
    eigenvalues, eigenvectors = eigsh(
        normalized_adjacency, k=k, which="LA", tol=1e-4
    )
    order = np.argsort(eigenvalues)[::-1]
    eigenvalues = eigenvalues[order]
    eigenvectors = eigenvectors[:, order]
    laplacian_eigenvalues = np.clip(1.0 - eigenvalues, 0.0, None)

    positive = laplacian_eigenvalues[laplacian_eigenvalues > 1e-8]
    lower = 0.25 / max(float(positive.max(initial=1.0)), 1e-8)
    upper = 8.0 / max(float(positive.min(initial=1.0)), 1e-8)
    times = np.geomspace(lower, upper, scales)
    heat = (eigenvectors**2) @ np.exp(
        -laplacian_eigenvalues[:, None] * times[None, :]
    )
    return torch.from_numpy(heat).float(), degree


def graph_signature(
    features: torch.Tensor,
    neighbors: int,
    eigenpairs: int,
    heat_scales: int,
    message_passing_steps: int,
) -> torch.Tensor:
    adjacency, density = exact_rank_knn_graph(features, neighbors)
    heat, degree = heat_kernel_signature(
        adjacency, eigenpairs=eigenpairs, scales=heat_scales
    )
    base = rank_normalize_columns(
        torch.cat(
            [
                density,
                torch.from_numpy(np.log1p(degree)).float()[:, None],
                heat,
            ],
            dim=1,
        )
    )
    transition = sparse.diags(
        1.0 / np.maximum(np.asarray(adjacency.sum(axis=1)).reshape(-1), 1e-12)
    ) @ adjacency
    current = base
    for _ in range(message_passing_steps):
        current_np = current.numpy()
        neighbor_mean = torch.from_numpy(transition @ current_np).float()
        neighbor_square_mean = torch.from_numpy(
            transition @ np.square(current_np)
        ).float()
        neighbor_std = torch.sqrt(
            torch.clamp(neighbor_square_mean - neighbor_mean.square(), min=0)
        )
        current = rank_normalize_columns(
            torch.cat([base, neighbor_mean, neighbor_std], dim=1)
        )
    return current


def main() -> None:
    args = parse_args()
    vision = torch.load(args.vision, map_location="cpu", weights_only=False)
    text = torch.load(args.text, map_location="cpu", weights_only=False)
    n = min(args.max_nodes, len(vision["node_ids"]), len(text["node_ids"]))
    vision_ids = vision["node_ids"][:n]
    target_full = aligned_text_indices(
        vision_ids, text["node_ids"], args.truth
    )
    candidate_indices = torch.unique(target_full, sorted=False)
    if len(candidate_indices) != n:
        raise ValueError("Ground truth is not a one-to-one permutation")
    candidate_lookup = {
        int(old): new for new, old in enumerate(candidate_indices.tolist())
    }
    targets = torch.tensor(
        [candidate_lookup[int(old)] for old in target_full.tolist()]
    )

    vision_scene = scene_features(vision)[:n]
    text_scene = scene_features(text, candidate_indices)
    vision_graph = graph_signature(
        vision_scene,
        args.knn,
        args.eigenpairs,
        args.heat_scales,
        args.message_passing_steps,
    )
    text_graph = graph_signature(
        text_scene,
        args.knn,
        args.eigenpairs,
        args.heat_scales,
        args.message_passing_steps,
    )

    vision_bundle = rank_normalize_columns(
        bundle_signature(vision["region_features"][:n])
    )
    text_bundle = rank_normalize_columns(
        bundle_signature(text["region_features"][candidate_indices])
    )
    combined_vision = torch.cat(
        [normalized(vision_graph), normalized(vision_bundle)], dim=1
    )
    combined_text = torch.cat(
        [normalized(text_graph), normalized(text_bundle)], dim=1
    )

    result = {
        "nodes": n,
        "knn": min(args.knn, n - 1),
        "eigenpairs": min(args.eigenpairs, n - 2),
        "heat_scales": args.heat_scales,
        "message_passing_steps": args.message_passing_steps,
        "graph_signature_retrieval": arbitrary_target_retrieval(
            vision_graph, text_graph, targets
        ),
        "bundle_signature_retrieval": arbitrary_target_retrieval(
            vision_bundle, text_bundle, targets
        ),
        "combined_signature_retrieval": arbitrary_target_retrieval(
            combined_vision, combined_text, targets
        ),
        "evaluation_note": (
            "Each signature is computed independently from one modality. "
            "Ground truth is loaded only after graph construction to score "
            "the hidden permutation."
        ),
    }
    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
    write_json(args.output, result)
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()