#!/usr/bin/env python3
"""Task 7016bd94 Part 1: ResGCN vs GRAFT, 20 seeds, paired t-tests."""

import torch
import numpy as np
import json
import os
from scipy import stats as scipy_stats
from src.data import load_dataset
from src.trainers import BPTrainer, GraphGrAPETrainer
from run_deep_baselines import ResGCNTrainer
from run_dblp_depth import load_dblp

device = 'cuda:0'
SEEDS = list(range(20))
EPOCHS = 200
OUT_DIR = 'results/resgcn_20seeds'

grape_extra = dict(diffusion_alpha=0.5, diffusion_iters=10,
                   lr_feedback=0.5, num_probes=64, topo_mode='fixed_A')


def train_one(cls, common, extra, seed):
    torch.manual_seed(seed); np.random.seed(seed); torch.cuda.manual_seed_all(seed)
    t = cls(**common, **extra)
    if hasattr(t, 'align_mode'):
        t.align_mode = 'chain_norm'
    bv, bt = 0, 0
    for ep in range(EPOCHS):
        t.train_step()
        if ep % 5 == 0:
            v = t.evaluate('val_mask')
            te = t.evaluate('test_mask')
            if v > bv: bv, bt = v, te
    del t; torch.cuda.empty_cache()
    return bt


def main():
    os.makedirs(OUT_DIR, exist_ok=True)
    per_seed_file = os.path.join(OUT_DIR, 'per_seed_data.json')
    if os.path.exists(per_seed_file):
        with open(per_seed_file) as f:
            per_seed_data = json.load(f)
    else:
        per_seed_data = {}

    METHODS = {
        'BP': (BPTrainer, {}),
        'ResGCN': (ResGCNTrainer, {}),
        'GRAFT': (GraphGrAPETrainer, grape_extra),
    }

    datasets_cfg = {
        'Cora': lambda: load_dataset('Cora', device=device),
        'CiteSeer': lambda: load_dataset('CiteSeer', device=device),
        'DBLP': lambda: load_dblp(),
    }

    results = {}

    for ds_name, loader in datasets_cfg.items():
        data = loader()
        common = dict(data=data, hidden_dim=64, lr=0.01, weight_decay=5e-4,
                      num_layers=6, residual_alpha=0.0, backbone='gcn')

        for mname, (cls, extra) in METHODS.items():
            key = f"{ds_name}_{mname}"
            print(f"\n=== {key} (20 seeds) ===", flush=True)

            if key not in per_seed_data:
                per_seed_data[key] = {}

            for seed in SEEDS:
                sk = str(seed)
                if sk in per_seed_data[key]:
                    print(f"  seed {seed}: cached", flush=True)
                    continue
                acc = train_one(cls, common, extra, seed)
                per_seed_data[key][sk] = acc
                print(f"  seed {seed}: {acc*100:.1f}%", flush=True)

                with open(per_seed_file, 'w') as f:
                    json.dump(per_seed_data, f, indent=2)

            accs = np.array([per_seed_data[key][str(s)] for s in SEEDS]) * 100
            results[key] = {
                'mean': float(accs.mean()), 'std': float(accs.std()),
                'accs': accs.tolist(),
            }
            print(f"  {mname}: {accs.mean():.1f} ± {accs.std():.1f}%")

        del data; torch.cuda.empty_cache()

    # Paired t-tests: GRAFT vs ResGCN
    print("\n" + "=" * 70)
    print("Paired t-tests: GRAFT vs ResGCN (20 seeds)")
    print("-" * 70)

    for ds in ['Cora', 'CiteSeer', 'DBLP']:
        bp_accs = np.array(results[f"{ds}_BP"]['accs'])
        res_accs = np.array(results[f"{ds}_ResGCN"]['accs'])
        gr_accs = np.array(results[f"{ds}_GRAFT"]['accs'])

        # GRAFT vs ResGCN
        t_stat, p_val = scipy_stats.ttest_rel(gr_accs, res_accs)
        delta = gr_accs.mean() - res_accs.mean()
        sig = '***' if p_val < 0.001 else ('**' if p_val < 0.01 else ('*' if p_val < 0.05 else 'ns'))

        results[f"{ds}_GRAFT_vs_ResGCN"] = {
            'delta': float(delta), 't_stat': float(t_stat),
            'p_value': float(p_val), 'significant': bool(p_val < 0.05),
        }

        # GRAFT vs BP
        t2, p2 = scipy_stats.ttest_rel(gr_accs, bp_accs)
        d2 = gr_accs.mean() - bp_accs.mean()
        sig2 = '***' if p2 < 0.001 else ('**' if p2 < 0.01 else ('*' if p2 < 0.05 else 'ns'))

        results[f"{ds}_GRAFT_vs_BP"] = {
            'delta': float(d2), 't_stat': float(t2),
            'p_value': float(p2), 'significant': bool(p2 < 0.05),
        }

        # ResGCN vs BP
        t3, p3 = scipy_stats.ttest_rel(res_accs, bp_accs)
        d3 = res_accs.mean() - bp_accs.mean()

        results[f"{ds}_ResGCN_vs_BP"] = {
            'delta': float(d3), 't_stat': float(t3),
            'p_value': float(p3), 'significant': bool(p3 < 0.05),
        }

        print(f"\n{ds}:")
        print(f"  BP:     {bp_accs.mean():.1f} ± {bp_accs.std():.1f}")
        print(f"  ResGCN: {res_accs.mean():.1f} ± {res_accs.std():.1f}")
        print(f"  GRAFT:  {gr_accs.mean():.1f} ± {gr_accs.std():.1f}")
        print(f"  GRAFT vs ResGCN: Δ{delta:+.1f}% p={p_val:.6f} {sig}")
        print(f"  GRAFT vs BP:     Δ{d2:+.1f}% p={p2:.6f} {sig2}")

    with open(os.path.join(OUT_DIR, 'results.json'), 'w') as f:
        json.dump(results, f, indent=2)
    print(f"\nSaved to {OUT_DIR}/results.json")


if __name__ == '__main__':
    main()