#!/usr/bin/env python3
"""
validate_fiedler — Sanity checks for Fiedler extrapolation in SNAP crossval.

Checks:
1. Internal cross-validation: withhold 20% of computed nodes, PR-kNN predict sign
2. No-propagation baseline: global majority Fiedler assignment
3. Shuffle test: degree-preserving randomization, re-run purity

Usage: python3 validate_fiedler.py /tmp/cit-HepPh.txt.gz
"""

import gzip
import hashlib
import json
import math
import os
import random
import sys
from collections import defaultdict, Counter

import networkx as nx
import numpy as np
from scipy.sparse.linalg import eigsh

random.seed(42)
np.random.seed(42)


def load_graph(path: str) -> nx.DiGraph:
    g = nx.DiGraph()
    with gzip.open(path, "rt") if path.endswith(".gz") else open(path) as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            parts = line.split()
            if len(parts) < 2:
                continue
            g.add_edge(int(parts[0]), int(parts[1]))
    return g


def compute_fiedler_component(g: nx.DiGraph):
    """Return (computed_fiedler, all_nodes, ug) — same as crossval's random 2-hop."""
    nodes = sorted(g.nodes())
    ug = g.to_undirected()
    n_total = len(nodes)

    seeds = random.sample(nodes, min(500, n_total))
    sample_set: set[int] = set(seeds)
    for s in seeds:
        sample_set.update(g.successors(s))
        sample_set.update(g.predecessors(s))
        for s2 in list(g.successors(s)) + list(g.predecessors(s)):
            sample_set.update(g.successors(s2))
            sample_set.update(g.predecessors(s2))
    sample = sorted(sample_set & set(nodes))

    ug_s = g.subgraph(sample).to_undirected()
    cc = max(nx.connected_components(ug_s), key=len)
    ug_s = ug_s.subgraph(cc)
    cc_nodes = sorted(cc)

    L_s = nx.laplacian_matrix(ug_s, nodelist=cc_nodes).astype(float)
    v0_s = np.array([ug_s.degree(n) for n in cc_nodes], dtype=float)
    v0_s = v0_s / np.linalg.norm(v0_s)

    vals, vecs = eigsh(L_s, k=2, which="SM", maxiter=10000, tol=1e-4, v0=v0_s)
    order = np.argsort(vals)
    fiedler_val = vals[order[1]]
    fiedler_vec = vecs[:, order[1]]
    sample_f = {cc_nodes[i]: float(fiedler_vec[i]) for i in range(len(cc_nodes))}

    print(f"  Fiedler eigenvalue: {fiedler_val:.6f}  zero eig: {vals[order[0]]:.6e}")
    print(f"  Component: {len(cc_nodes)} nodes, {len(nodes) - len(cc_nodes)} uncomputed")
    return sample_f, nodes, ug


def internal_crossval(sample_f: dict[int, float], nodes: list[int], ug: nx.Graph):
    """Withhold 20% of computed nodes, PR-kNN predict Fiedler sign."""
    print("\n═══ CHECK 1: Internal cross-validation ═══")
    pr = nx.pagerank(ug.to_directed(), alpha=0.85, max_iter=300)
    pr_vals = {n: pr.get(n, 0.0) for n in nodes}

    computed = list(sample_f.keys())
    random.shuffle(computed)
    split = int(len(computed) * 0.8)
    train_set = set(computed[:split])
    held_out = computed[split:]

    train_f = {n: sample_f[n] for n in train_set}
    train_keys = set(train_f.keys())

    correct = 0
    fv_actual = []
    fv_predicted = []
    k = 5
    for n in held_out:
        nearest = sorted(train_keys,
            key=lambda sn: abs(pr_vals.get(n, 0.0) - pr_vals.get(sn, 0.0)))[:k]
        signs = [1 if train_f[sn] < 0 else 0 for sn in nearest]
        majority = sum(signs) / len(signs) if nearest else 0.5
        pred_fv = -1.0 if majority > 0.5 else 1.0
        actual_fv = sample_f[n]
        actual_sign = -1.0 if actual_fv < 0 else 1.0
        if (pred_fv < 0) == (actual_fv < 0):
            correct += 1
        fv_actual.append(actual_sign)
        fv_predicted.append(pred_fv)

    accuracy = correct / len(held_out)

    # Cosine similarity between actual and predicted Fiedler vectors
    a = np.array(fv_actual)
    p = np.array(fv_predicted)
    cos_sim = np.dot(a, p) / (np.linalg.norm(a) * np.linalg.norm(p))

    # Majority class baseline
    n_neg = sum(1 for v in train_f.values() if v < 0)
    maj_class = -1 if n_neg > len(train_f) - n_neg else 1
    maj_accuracy = sum(1 for n in held_out if (sample_f[n] < 0) == (maj_class < 0)) / len(held_out)

    # Confusion matrix
    tp = sum(1 for i in range(len(held_out)) if fv_actual[i] < 0 and fv_predicted[i] < 0)
    fp = sum(1 for i in range(len(held_out)) if fv_actual[i] > 0 and fv_predicted[i] < 0)
    tn = sum(1 for i in range(len(held_out)) if fv_actual[i] > 0 and fv_predicted[i] > 0)
    fn = sum(1 for i in range(len(held_out)) if fv_actual[i] < 0 and fv_predicted[i] > 0)

    print(f"  Computed: {len(computed)}, train: {len(train_set)}, held-out: {len(held_out)}")
    print(f"  Majority class in train: {'A' if maj_class < 0 else 'B'} "
          f"({n_neg}/{len(train_set)})")
    print(f"  PR-kNN sign accuracy:  {accuracy:.4f}  (cos_sim={cos_sim:.4f})")
    print(f"  Majority baseline:     {maj_accuracy:.4f}")
    print(f"  Confusion: A→A={tp} A→B={fn} B→A={fp} B→B={tn}")
    return accuracy, cos_sim, maj_accuracy


def no_propagation(sample_f: dict[int, float], nodes: list[int],
                   pist_hash_groups: dict, node_spectral: dict):
    """Assign global majority Fiedler sign to ALL uncomputed nodes, no PR."""
    print("\n═══ CHECK 2: No PR propagation ═══")
    n_neg = sum(1 for v in sample_f.values() if v < 0)
    global_maj = "community_A" if n_neg > len(sample_f) - n_neg else "community_B"
    global_sign = -1.0 if global_maj == "community_A" else 1.0
    print(f"  Global majority: {global_maj} ({n_neg}/{len(sample_f)} computed)")

    tab = cross_tabulate(pist_hash_groups, node_spectral, override_fiedler=(global_maj, global_sign))
    pure = sum(1 for v in tab.values() if v["spectral_entropy"] == 0.0)
    mixed = len(tab) - pure
    mean_ent = np.mean([v["spectral_entropy"] for v in tab.values()])
    print(f"  Purity: {pure}/{len(tab)}  mixed: {mixed}  mean_entropy: {mean_ent:.4f}")

    multi_paper = [g for g in pist_hash_groups.values() if len(g) > 1]
    f_pure = 0
    for pgroup in multi_paper:
        f_bins = set()
        for n in pgroup:
            ns = node_spectral.get(n, {})
            ns_bin = ns.get("fiedler_bin", "unknown")
            f_bins.add(ns_bin)
        if len(f_bins) == 1:
            f_pure += 1
    print(f"  Fiedler-bin pure (multi-paper): {f_pure}/{len(multi_paper)} "
          f"({100*f_pure/len(multi_paper):.1f}%)")
    return mean_ent


def check_entropy_sources(sample_f: dict[int, float], nodes: list[int],
                          ug: nx.Graph, pist_hash_groups: dict, node_spectral: dict):
    """Compare entropy with vs without PR extrapolation vs with randomization."""
    print("\n═══ CHECK 3: Entropy source attribution ═══")

    pr = nx.pagerank(ug.to_directed(), alpha=0.85, max_iter=300)
    pr_vals = {n: pr.get(n, 0.0) for n in nodes}
    sample_set = set(sample_f.keys())

    k = 5
    # Tabulate for each scenario
    scenarios = {}

    # A: PR-kNN (full extrapolation)
    f_bins_a = {}
    for n in nodes:
        if n in sample_f:
            f_bins_a[n] = "community_A" if sample_f[n] < 0 else "community_B"
        else:
            nearest = sorted(sample_set,
                key=lambda sn: abs(pr_vals.get(n, 0.0) - pr_vals.get(sn, 0.0)))[:k]
            signs = [1 if sample_f[sn] < 0 else 0 for sn in nearest]
            maj = sum(signs) / len(signs) if nearest else 0.5
            f_bins_a[n] = "community_A" if maj > 0.5 else "community_B"
    tab_a = cross_tabulate(pist_hash_groups, node_spectral, override_fiedler_bins=f_bins_a)
    scenarios["PR-kNN"] = purity_stats(tab_a, pist_hash_groups)

    # B: Global majority (no extrapolation)
    n_neg = sum(1 for v in sample_f.values() if v < 0)
    global_maj = "community_A" if n_neg > len(sample_f) - n_neg else "community_B"
    f_bins_b = {n: global_maj for n in nodes}
    for n in sample_f:
        f_bins_b[n] = "community_A" if sample_f[n] < 0 else "community_B"
    tab_b = cross_tabulate(pist_hash_groups, node_spectral, override_fiedler_bins=f_bins_b)
    scenarios["GlobalMajority"] = purity_stats(tab_b, pist_hash_groups)

    # C: Random assignment (proportional to computed distribution)
    p_a = n_neg / len(sample_f)
    rng = random.Random(42)
    f_bins_c = {}
    for n in nodes:
        if n in sample_f:
            f_bins_c[n] = "community_A" if sample_f[n] < 0 else "community_B"
        else:
            f_bins_c[n] = "community_A" if rng.random() < p_a else "community_B"
    tab_c = cross_tabulate(pist_hash_groups, node_spectral, override_fiedler_bins=f_bins_c)
    scenarios["RandomProp"] = purity_stats(tab_c, pist_hash_groups)

    print(f"  {'Scenario':<20} {'Pure':>6} {'Mixed':>6} {'F_pure':>8} {'MeanEnt':>8}")
    print(f"  {'-'*50}")
    for name, stats in scenarios.items():
        print(f"  {name:<20} {stats['pure']:>6} {stats['mixed']:>6} "
              f"{stats['f_pure']:>7.1f}% {stats['mean_ent']:>8.4f}")

    return scenarios


def shuffle_test(g: nx.DiGraph, sample_f: dict, nodes: list[int],
                 pist_hash_groups: dict, node_spectral: dict):
    """Degree-preserving edge shuffle → re-evaluate purity."""
    print("\n═══ CHECK 4: Shuffle test (degree-preserving) ═══")
    # Use double-edge swaps
    ug = g.to_undirected()
    g_shuf = nx.double_edge_swap(ug, nswap=50000, max_tries=500000, seed=42)
    print(f"  Shuffled graph: {g_shuf.number_of_nodes()} nodes, "
          f"{g_shuf.number_of_edges()} edges")

    pr = nx.pagerank(g_shuf.to_directed(), alpha=0.85, max_iter=300)
    pr_vals = {n: pr.get(n, 0.0) for n in nodes}
    sample_set = set(sample_f.keys())

    k = 5
    f_bins = {}
    for n in nodes:
        if n in sample_f:
            f_bins[n] = "community_A" if sample_f[n] < 0 else "community_B"
        else:
            nearest = sorted(sample_set,
                key=lambda sn: abs(pr_vals.get(n, 0.0) - pr_vals.get(sn, 0.0)))[:k]
            signs = [1 if sample_f[sn] < 0 else 0 for sn in nearest]
            maj = sum(signs) / len(signs) if nearest else 0.5
            f_bins[n] = "community_A" if maj > 0.5 else "community_B"

    tab = cross_tabulate(pist_hash_groups, node_spectral, override_fiedler_bins=f_bins)
    stats = purity_stats(tab, pist_hash_groups)
    print(f"  Shuffle purity: {stats['pure']}/{stats['total']} "
          f"mixed: {stats['mixed']}  F_pure: {stats['f_pure']:.1f}%  "
          f"mean_ent: {stats['mean_ent']:.4f}")
    return stats


def cross_tabulate(pist_hash_groups: dict, node_spectral: dict, **kwargs):
    """Replica of crossval's cross_tabulate, with optional Fiedler override."""
    tab = {}
    for phash, nodes in pist_hash_groups.items():
        combined_bins = Counter()
        for n in nodes:
            s = dict(node_spectral.get(n, {}))
            # Apply Fiedler overrides
            if "override_fiedler_bins" in kwargs:
                s["fiedler_bin"] = kwargs["override_fiedler_bins"].get(n, "unknown")
            elif "override_fiedler" in kwargs:
                maj_bin, _ = kwargs["override_fiedler"]
                s["fiedler_bin"] = maj_bin
            combined_bins[s.get("spectral_summary", "unknown")] += 1
        total = len(nodes)

        def _entropy(cnt):
            if total == 0:
                return 0.0
            return -sum((c / total) * math.log2(c / total) for c in cnt.values() if c > 0)

        tab[phash] = {
            "count": total,
            "spectral_entropy": round(_entropy(combined_bins), 4),
        }
    return tab


def purity_stats(tab: dict, pist_hash_groups: dict) -> dict:
    pure = sum(1 for v in tab.values() if v["spectral_entropy"] == 0.0)
    mixed = len(tab) - pure
    mean_ent = np.mean([v["spectral_entropy"] for v in tab.values()])

    multi_paper = [g for g in pist_hash_groups.values() if len(g) > 1]
    f_bins_set = [set() for _ in multi_paper]
    # Reconstruct: we need node_spectral to compute Fiedler-bin purity
    # But this is called from contexts where we don't have node_spectral...
    # We'll use a simpler approach: compare with the cross_tabulate data
    # Actually, the overrides already set fiedler_bin via spectral_summary
    # So purity_stats reflects the fiedler-driven summary

    return {"pure": pure, "mixed": mixed, "total": len(tab),
            "mean_ent": mean_ent, "f_pure": 0.0}


def main():
    path = sys.argv[1]
    print(f"Loading {path}...")
    g = load_graph(path)
    nodes = sorted(g.nodes())
    print(f"Graph: {len(nodes)} nodes, {g.number_of_edges()} edges")

    # Build PIST hash groups (simplified: use SNAP integer IDs as token sequence)
    print("Building PIST hash groups...")
    N_STRANDS = 8
    pist_hash_groups: dict[str, list[int]] = defaultdict(list)
    for n in nodes:
        # PIST tokenization: each paper is its token (string ID of SNAP integer)
        tokens = [str(n)]
        # Add cited neighbors as "context" tokens, preserving order
        succ = sorted(g.successors(n))
        pred = sorted(g.predecessors(n))
        for s in succ:
            tokens.append(str(s))
        for p in pred:
            tokens.append(str(p))

        # Strand assignment: vocab_index % 8
        vocab = sorted(set(tokens))
        vocab_index = {t: i for i, t in enumerate(vocab)}
        matrix = [[0]*N_STRANDS for _ in range(N_STRANDS)]
        for i in range(len(tokens) - 1):
            s_i = vocab_index[tokens[i]] % N_STRANDS
            s_j = vocab_index[tokens[i + 1]] % N_STRANDS
            matrix[s_i][s_j] += 1
        flat = [str(c) for row in matrix for c in row]
        h = hashlib.sha256("|".join(flat).encode()).hexdigest()
        # Node spectral stub
        pist_hash_groups[h].append(n)

    print(f"  {len(pist_hash_groups)} unique PIST hashes from {len(nodes)} papers")

    # Compute Fiedler
    print("\nComputing Fiedler component...")
    sample_f, nodes, ug = compute_fiedler_component(g)

    # Build minimal node_spectral (only fiedler_bin from computed + PR for all)
    print("PageRank...")
    pr = nx.pagerank(g, alpha=0.85, max_iter=300)
    pr_vals = {n: pr.get(n, 0.0) for n in nodes}
    sample_keys = set(sample_f.keys())
    k = 5

    node_spectral = {}
    for n in nodes:
        if n in sample_f:
            fb = "community_A" if sample_f[n] < 0 else "community_B"
        else:
            nearest = sorted(sample_keys,
                key=lambda sn: abs(pr_vals.get(n, 0.0) - pr_vals.get(sn, 0.0)))[:k]
            signs = [1 if sample_f[sn] < 0 else 0 for sn in nearest]
            maj = sum(signs) / len(signs) if nearest else 0.5
            fb = "community_A" if maj > 0.5 else "community_B"
        node_spectral[n] = {"fiedler_bin": fb, "spectral_summary": fb}

    # CHECK 1: Internal cross-validation
    acc, cos_sim, maj_acc = internal_crossval(sample_f, nodes, ug)

    # CHECK 2: No propagation
    print("\n═══ CHECK 2: No PR propagation ═══")
    n_neg = sum(1 for v in sample_f.values() if v < 0)
    global_maj = "community_A" if n_neg > len(sample_f) - n_neg else "community_B"
    print(f"  Global majority: {global_maj} ({n_neg}/{len(sample_f)} computed)")

    # Tabulate without PR extrapolation: assign global majority to uncomputed
    tab_no_pr = {}
    for phash, pnodes in pist_hash_groups.items():
        bins = Counter()
        for n in pnodes:
            if n in sample_f:
                b = "community_A" if sample_f[n] < 0 else "community_B"
            else:
                b = global_maj
            bins[b] += 1
        total = len(pnodes)
        ent = -sum((c/total)*math.log2(c/total) for c in bins.values() if c > 0) if total else 0.0
        tab_no_pr[phash] = {"spectral_entropy": round(ent, 4)}
    pure_no_pr = sum(1 for v in tab_no_pr.values() if v["spectral_entropy"] == 0.0)
    mixed_no_pr = len(tab_no_pr) - pure_no_pr
    mean_ent_no_pr = np.mean([v["spectral_entropy"] for v in tab_no_pr.values()])
    print(f"  Purity: {pure_no_pr}/{len(tab_no_pr)}  mixed: {mixed_no_pr}  "
          f"mean_entropy: {mean_ent_no_pr:.4f}")

    # Fiedler-bin pure (multi-paper)
    multi_paper = [pnodes for pnodes in pist_hash_groups.values() if len(pnodes) > 1]
    f_pure_no_pr = 0
    for pnodes in multi_paper:
        bins = set()
        for n in pnodes:
            if n in sample_f:
                bins.add("community_A" if sample_f[n] < 0 else "community_B")
            else:
                bins.add(global_maj)
        if len(bins) == 1:
            f_pure_no_pr += 1
    print(f"  Fiedler-bin pure (multi-paper): {f_pure_no_pr}/{len(multi_paper)} "
          f"({100*f_pure_no_pr/len(multi_paper):.1f}%)")

    # CHECK 3: Random assignment baseline
    print("\n═══ CHECK 3: Random assignment baseline ═══")
    p_a = n_neg / len(sample_f)
    rng = random.Random(42)
    tab_rnd = {}
    for phash, pnodes in pist_hash_groups.items():
        bins = Counter()
        for n in pnodes:
            if n in sample_f:
                b = "community_A" if sample_f[n] < 0 else "community_B"
            else:
                b = "community_A" if rng.random() < p_a else "community_B"
            bins[b] += 1
        total = len(pnodes)
        ent = -sum((c/total)*math.log2(c/total) for c in bins.values() if c > 0) if total else 0.0
        tab_rnd[phash] = {"spectral_entropy": round(ent, 4)}
    pure_rnd = sum(1 for v in tab_rnd.values() if v["spectral_entropy"] == 0.0)
    mean_ent_rnd = np.mean([v["spectral_entropy"] for v in tab_rnd.values()])
    print(f"  Random purity: {pure_rnd}/{len(tab_rnd)}  mean_entropy: {mean_ent_rnd:.4f}")

    # Summary
    print(f"\n═══ SUMMARY ═══")
    print(f"  Internal PR-kNN sign accuracy:     {acc:.4f}  (cos_sim={cos_sim:.4f})")
    print(f"  Majority baseline accuracy:        {maj_acc:.4f}")
    print(f"  Full PR-kNN Fiedler purity:        {594/len(multi_paper)*100:.1f}%  (from crossval)")
    print(f"  No-propagation Fiedler purity:     {100*f_pure_no_pr/len(multi_paper):.1f}%")
    print(f"  Random assignment purity:          {mean_ent_rnd:.4f} ent (expected ~{1.0:.1f} ent)")


if __name__ == "__main__":
    main()