from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import torch
from scipy.stats import spearmanr

from .common import linear_cka, normalized, read_json, write_json
from .io import load_feature_pair, select_rows


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser()
    p.add_argument("--manifest", default="artifacts/manifest.json")
    p.add_argument("--vision", default="artifacts/vision.pt")
    p.add_argument("--text", default="artifacts/text.pt")
    p.add_argument("--split", choices=["val", "test"], default="val")
    p.add_argument("--max-samples", type=int, default=1_000)
    p.add_argument("--permutations", type=int, default=20)
    p.add_argument("--seed", type=int, default=20260728)
    p.add_argument("--output", default="artifacts/diagnostics.json")
    return p.parse_args()


def upper_triangle(x: torch.Tensor) -> np.ndarray:
    n = x.shape[0]
    i, j = torch.triu_indices(n, n, offset=1)
    return x[i, j].cpu().numpy()


def main() -> None:
    args = parse_args()
    manifest = read_json(args.manifest)
    vision, text, vlookup, tlookup = load_feature_pair(args.vision, args.text)
    rows = manifest[args.split][: args.max_samples]
    x = select_rows(vision["features"], vlookup, rows)
    y = select_rows(text["features"], tlookup, rows)

    gx = normalized(x) @ normalized(x).T
    gy = normalized(y) @ normalized(y).T
    gx_upper = upper_triangle(gx)
    gy_upper = upper_triangle(gy)
    rho = spearmanr(gx_upper, gy_upper).statistic
    generator = torch.Generator().manual_seed(args.seed)
    shuffled_rhos = []
    for _ in range(args.permutations):
        permutation = torch.randperm(len(y), generator=generator)
        shuffled_gy = gy[permutation][:, permutation]
        shuffled_rhos.append(
            float(spearmanr(gx_upper, upper_triangle(shuffled_gy)).statistic)
        )
    shuffled_mean = float(np.mean(shuffled_rhos))
    shuffled_std = float(np.std(shuffled_rhos))
    result = {
        "split": args.split,
        "samples": len(rows),
        "linear_cka": linear_cka(x, y),
        "pairwise_cosine_spearman": float(rho),
        "shuffled_spearman_mean": shuffled_mean,
        "shuffled_spearman_std": shuffled_std,
        "spearman_shuffle_z": float(
            (rho - shuffled_mean) / max(shuffled_std, 1e-12)
        ),
        "permutations": args.permutations,
        "vision_dim": x.shape[-1],
        "text_dim": y.shape[-1],
    }
    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
    write_json(args.output, result)
    print(result)


if __name__ == "__main__":
    main()