#!/usr/bin/env python3
"""Validate thumbnail K-means on the same Task06 test patients as siamese."""

import os, sys, json, argparse
import numpy as np
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
from PIL import Image as PILImage

os.environ["OMP_NUM_THREADS"] = "1"
os.environ["OPENBLAS_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
from classes import PatientIdentifier
from sklearn.metrics import adjusted_rand_score, normalized_mutual_info_score

ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
PNG_DIR = os.path.join(ROOT, "features", "task06_pngs", "test")
RESULTS_DIR = os.path.join(ROOT, "results", "clustering_validation")
PLOTS_DIR = os.path.join(ROOT, "plots")
SEED = 42

def cluster_purity_stats(y_true, y_pred):
    purities = []; overall_correct = 0
    for c in np.unique(y_pred):
        mask = y_pred == c
        _, counts = np.unique(y_true[mask], return_counts=True)
        purities.append(counts.max() / mask.sum())
        overall_correct += counts.max()
    purities = np.array(purities)
    return {"overall": float(overall_correct / len(y_true)),
            "median": float(np.median(purities)), "mean": float(np.mean(purities)),
            "frac_gt_70": float((purities > 0.7).mean()),
            "frac_gt_90": float((purities > 0.9).mean())}

def patient_capture_stats(y_true, y_pred):
    captures = []
    for p in np.unique(y_true):
        mask = y_true == p
        p_clusters = y_pred[mask]
        _, counts = np.unique(p_clusters, return_counts=True)
        captures.append(counts.max() / mask.sum())
    captures = np.array(captures)
    return {"median": float(np.median(captures)), "mean": float(np.mean(captures)),
            "frac_gt_50": float((captures > 0.5).mean())}

def plot_assignment_matrix(y_true, y_pred, out_path, title):
    true_patients = sorted(np.unique(y_true))
    pred_clusters = sorted(np.unique(y_pred))
    matrix = np.zeros((len(true_patients), len(pred_clusters)))
    for i, p in enumerate(true_patients):
        for j, c in enumerate(pred_clusters):
            matrix[i, j] = ((y_true == p) & (y_pred == c)).sum()
    matrix_norm = matrix / (matrix.sum(axis=1, keepdims=True) + 1e-8)
    fig, ax = plt.subplots(figsize=(max(14, len(pred_clusters)*0.22), max(10, len(true_patients)*0.18)))
    cmap = plt.cm.YlOrRd.copy()
    cmap.set_under('white')
    ax.imshow(matrix_norm, aspect="auto", cmap=cmap, vmin=1e-6, vmax=1)
    ax.set_xticks(range(len(pred_clusters)))
    ax.set_xticklabels([f"c{c}" for c in pred_clusters], fontsize=6, rotation=90)
    ax.set_yticks(range(len(true_patients)))
    ax.set_yticklabels(true_patients, fontsize=7)
    ax.set_xlabel("Predicted cluster"); ax.set_ylabel("True patient")
    ax.set_title(title, fontsize=11)
    plt.colorbar(ax.images[0], ax=ax, label="Fraction of patient's slices")
    plt.tight_layout()
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    plt.savefig(out_path, dpi=150); plt.close()
    print(f"  Heatmap → {out_path}")

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--tag", default="")
    ap.add_argument("--full", action="store_true",
                    help="Run on full NIfTI dataset (63 patients, not just test).")
    args = ap.parse_args()
    tag = f"_{args.tag}" if args.tag else ""
    os.makedirs(RESULTS_DIR, exist_ok=True); os.makedirs(PLOTS_DIR, exist_ok=True)

    if args.full:
        tag = (tag or "") + "_full"
        from classes import NiftiSliceDataset
        VOL = os.path.join(os.path.dirname(ROOT), "Task06_Lung", "imagesTr")
        ds = NiftiSliceDataset(VOL, random_slices=False, seed=SEED, rotate_deg=90)
        ds.load_all_slices(stride=1)
        test_paths = None
        test_pids = ds.patient_labels
        test_fnames = ds.filenames
        y_true_map = {f: pid for f, pid in zip(test_fnames, test_pids)}
        k = ds.n_patients
        print("=" * 60)
        print("Thumbnail K-means validation (manuscript method) — FULL")
        print("=" * 60)
        print(f"  FULL dataset: {len(test_fnames)} slices, {k} patients")
        # Build thumbnails from loaded slices
        thumbs = []
        for sl in ds.slices_rgb:
            img = PILImage.fromarray(sl).convert("L").resize((64, 64))
            thumbs.append(np.asarray(img, dtype=np.float32).flatten() / 255.0)
        thumbs = np.array(thumbs, dtype=np.float32)
    else:
        manifest_path = os.path.join(PNG_DIR, "manifest.json")
        if not os.path.exists(manifest_path):
            print(f"Test PNGs not found at {PNG_DIR}. Run train_siamese.py first.")
            sys.exit(1)
        with open(manifest_path) as f: png = json.load(f)
        test_paths = png["paths"]
        test_pids = np.array(png["patient_ids"])
        test_fnames = [p.split("/")[-1].replace(".png", "") for p in test_paths]
        y_true_map = {f: pid for f, pid in zip(test_fnames, test_pids)}
        k = len(np.unique(test_pids))
        print("=" * 60)
        print("Thumbnail K-means validation (manuscript method)")
        print("=" * 60)
        print(f"  {len(test_fnames)} slices, {k} patients (same test set as siamese)")
        # Build thumbnails from PNGs
        print("  Building 64x64 thumbnails ...", flush=True)
        thumbs = []
        for p in test_paths:
            img = PILImage.open(p).convert("L").resize((64, 64))
            thumbs.append(np.asarray(img, dtype=np.float32).flatten() / 255.0)
        thumbs = np.array(thumbs, dtype=np.float32)

    print(f"  Clustering into k={k} ...", flush=True)
    ident = PatientIdentifier(patient_estimates={"lung": k}, random_state=SEED)
    groups = ident.identify_from_features(thumbs, np.full(len(thumbs), "lung"))
    y_pred = np.array([str(g) for g in groups])
    y_true = np.array([y_true_map[f] for f in test_fnames])

    purity = cluster_purity_stats(y_true, y_pred)
    capture = patient_capture_stats(y_true, y_pred)
    ari = adjusted_rand_score(y_true, y_pred)
    nmi = normalized_mutual_info_score(y_true, y_pred)

    print(f"\n{'='*50}")
    print("RESULTS — Thumbnail K-means (manuscript)")
    print(f"{'='*50}")
    print(f"  ARI:              {ari:.4f}")
    print(f"  NMI:              {nmi:.4f}")
    print(f"  Overall purity:   {purity['overall']:.3f} ({purity['overall']*100:.1f}%)")
    print(f"  Mean capture:     {capture['mean']:.3f} ({capture['mean']*100:.1f}%)")

    results = {"method": "thumbnail_64x64", "n_patients": int(k),
               "n_slices": len(test_fnames), "ARI": ari, "NMI": nmi,
               "cluster_purity": purity, "patient_capture": capture}
    out_json = os.path.join(RESULTS_DIR, f"validate_thumbnail{tag}.json")
    with open(out_json, "w") as f: json.dump(results, f, indent=2)
    print(f"  Metrics → {out_json}")

    plot_assignment_matrix(y_true, y_pred,
        os.path.join(PLOTS_DIR, "clustering_validation", "thumbnail", f"assignment_matrix{tag}.png"),
        f"Thumbnail K-means (ARI={ari:.3f}, purity={purity['overall']:.1%}, capture mean={capture['mean']:.1%})")
    print("DONE")

if __name__ == "__main__": main()