#!/usr/bin/env python3
"""compare_patient_groupings.py

Compare the three estimated patient groupings of the IQ-OTH/NCCD
("lung_effnet") dataset against each other:

    siamese    -> results/siamese_manifest.csv
    pca50      -> results/simple_patient_manifest.csv   (PCA-50 CNN-feature K-means)
    thumbnail  -> results/thumbnail_patient_manifest.csv (64x64 grayscale K-means)

Downstream classification accuracy is affected similarly by all three, so the
question this answers is: do the three methods actually partition the images
differently? If they agree closely, the method choice is cosmetic; if they
diverge, the groupings are genuinely different partitions — which, combined with
the Task06 ground-truth validation (where siamese scored higher purity/capture),
is what makes the siamese work worthwhile.

Metrics are all label-invariant (cluster-id names don't matter):
    ARI  - Adjusted Rand Index (chance-corrected pair agreement)
    NMI  - Normalized Mutual Information
    V    - V-measure (harmonic mean of homogeneity & completeness)

Outputs:
    results/compare_patient_groupings.json
    plots/analysis/grouping_agreement_heatmap.png

Usage:
    conda activate fundus_imaging
    python scripts/analysis/compare_patient_groupings.py
"""

import os
import csv
import json
from itertools import combinations

import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from sklearn.metrics import (
    adjusted_rand_score,
    normalized_mutual_info_score,
    homogeneity_completeness_v_measure,
)

ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
RESULTS_DIR = os.path.join(ROOT, "results")
PLOTS_DIR = os.path.join(ROOT, "plots", "analysis")

METHODS = {
    "siamese": "siamese_manifest.csv",
    "pca50": "simple_patient_manifest.csv",
    "thumbnail": "thumbnail_patient_manifest.csv",
}
METHOD_ORDER = ["siamese", "pca50", "thumbnail"]
CLASS_OF = {"B": "Benign", "M": "Malignant", "N": "Normal"}


def load_manifest(path):
    """Return {image_short_name: group_id} from a manifest CSV."""
    mapping = {}
    with open(path, newline="") as f:
        for row in csv.DictReader(f):
            for img in row["images"].split(";"):
                if img:
                    mapping[img] = row["patient_id"]
    return mapping


def size_stats(labels):
    """Group-size distribution for one method's label array."""
    _, counts = np.unique(labels, return_counts=True)
    return {
        "n_groups": int(len(counts)),
        "min": int(counts.min()),
        "median": float(np.median(counts)),
        "mean": float(counts.mean()),
        "max": int(counts.max()),
        "singletons": int((counts == 1).sum()),
    }


def agreement(a, b):
    h, c, v = homogeneity_completeness_v_measure(a, b)
    return {"ARI": float(adjusted_rand_score(a, b)),
            "NMI": float(normalized_mutual_info_score(a, b)),
            "V": float(v)}


def main():
    maps = {m: load_manifest(os.path.join(RESULTS_DIR, f))
            for m, f in METHODS.items()}

    # Align on images present in all three (should be the full 1,097).
    common = sorted(set.intersection(*[set(mp) for mp in maps.values()]))
    classes = np.array([CLASS_OF.get(img.split("_")[0], "?") for img in common])
    labels = {m: np.array([maps[m][img] for img in common]) for m in METHOD_ORDER}
    print(f"Aligned on {len(common)} images common to all three methods.\n")

    # --- Structural summary --------------------------------------------------
    print("=" * 72)
    print("GROUP STRUCTURE PER METHOD")
    print("=" * 72)
    print(f"{'method':<12s} {'groups':>7s} {'min':>5s} {'median':>7s} "
          f"{'mean':>6s} {'max':>5s} {'singletons':>11s}")
    structure = {}
    for m in METHOD_ORDER:
        s = size_stats(labels[m])
        structure[m] = s
        print(f"{m:<12s} {s['n_groups']:>7d} {s['min']:>5d} {s['median']:>7.1f} "
              f"{s['mean']:>6.1f} {s['max']:>5d} {s['singletons']:>11d}")

    # --- Pairwise agreement --------------------------------------------------
    print("\n" + "=" * 72)
    print("PAIRWISE AGREEMENT (how similarly the methods partition the images)")
    print("=" * 72)
    print(f"{'pair':<24s} {'ARI':>8s} {'NMI':>8s} {'V':>8s}")
    print("-" * 52)
    pairwise = {}
    ari_matrix = np.eye(len(METHOD_ORDER))
    for i, j in combinations(range(len(METHOD_ORDER)), 2):
        a, b = METHOD_ORDER[i], METHOD_ORDER[j]
        g = agreement(labels[a], labels[b])
        pairwise[f"{a}_vs_{b}"] = g
        ari_matrix[i, j] = ari_matrix[j, i] = g["ARI"]
        print(f"{a+' vs '+b:<24s} {g['ARI']:>8.3f} {g['NMI']:>8.3f} {g['V']:>8.3f}")

    # --- Per-class ARI -------------------------------------------------------
    print("\n" + "=" * 72)
    print("PER-CLASS ARI (agreement within each diagnostic class)")
    print("=" * 72)
    print(f"{'pair':<24s} " + " ".join(f"{c:>10s}" for c in ["Benign", "Malignant", "Normal"]))
    print("-" * 60)
    per_class = {}
    for i, j in combinations(range(len(METHOD_ORDER)), 2):
        a, b = METHOD_ORDER[i], METHOD_ORDER[j]
        row = {}
        cells = []
        for c in ["Benign", "Malignant", "Normal"]:
            mask = classes == c
            ari = float(adjusted_rand_score(labels[a][mask], labels[b][mask]))
            row[c] = ari
            cells.append(f"{ari:>10.3f}")
        per_class[f"{a}_vs_{b}"] = row
        print(f"{a+' vs '+b:<24s} " + " ".join(cells))

    # --- Odd-one-out: mean ARI of each method vs the other two ---------------
    print("\n" + "=" * 72)
    print("MEAN ARI OF EACH METHOD VS THE OTHER TWO (lower = most distinct)")
    print("=" * 72)
    mean_ari = {}
    for i, m in enumerate(METHOD_ORDER):
        others = [ari_matrix[i, j] for j in range(len(METHOD_ORDER)) if j != i]
        mean_ari[m] = float(np.mean(others))
        print(f"  {m:<12s} {mean_ari[m]:.3f}")
    odd = min(mean_ari, key=mean_ari.get)
    print(f"\n  Most distinct grouping: {odd}")

    # --- Save + heatmap ------------------------------------------------------
    out = {
        "n_images": len(common),
        "structure": structure,
        "pairwise": pairwise,
        "per_class_ARI": per_class,
        "mean_ari_vs_others": mean_ari,
        "most_distinct": odd,
    }
    os.makedirs(RESULTS_DIR, exist_ok=True)
    out_path = os.path.join(RESULTS_DIR, "compare_patient_groupings.json")
    with open(out_path, "w") as f:
        json.dump(out, f, indent=2)
    print(f"\nSaved metrics → {out_path}")

    fig, ax = plt.subplots(figsize=(5.5, 4.5))
    im = ax.imshow(ari_matrix, vmin=0, vmax=1, cmap="viridis")
    ax.set_xticks(range(len(METHOD_ORDER)))
    ax.set_yticks(range(len(METHOD_ORDER)))
    ax.set_xticklabels(METHOD_ORDER)
    ax.set_yticklabels(METHOD_ORDER)
    for i in range(len(METHOD_ORDER)):
        for j in range(len(METHOD_ORDER)):
            ax.text(j, i, f"{ari_matrix[i, j]:.2f}", ha="center", va="center",
                    color="white" if ari_matrix[i, j] < 0.6 else "black", fontsize=11)
    ax.set_title("Patient-grouping agreement (ARI)\nIQ-OTH/NCCD, 1,097 images")
    fig.colorbar(im, ax=ax, label="Adjusted Rand Index")
    plt.tight_layout()
    os.makedirs(PLOTS_DIR, exist_ok=True)
    plot_path = os.path.join(PLOTS_DIR, "grouping_agreement_heatmap.png")
    plt.savefig(plot_path, dpi=150)
    plt.close()
    print(f"Saved heatmap → {plot_path}")


if __name__ == "__main__":
    main()