#!/usr/bin/env python3
"""
classification_thumbnail.py — Run the classification pipeline using the
manuscript's thumbnail-based K-means patient manifest.

Saves results alongside the feature-based results for comparison.

Usage:
    conda activate fundus_imaging
    python scripts/classification_thumbnail.py
"""

import os, sys, csv, json, re
import numpy as np

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(
    os.path.abspath(__file__)))))

from classes import PatientIdentifier, PatientLeakageClassifier

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

ROOT = os.path.dirname(os.path.dirname(os.path.dirname(
    os.path.abspath(__file__))))
DATASET_PATH = os.path.join(os.path.dirname(ROOT),
                            "The IQ-OTHNCCD lung cancer dataset")
FEATURES_DIR = os.path.join(ROOT, "features")
RESULTS_DIR = os.path.join(ROOT, "results")

SEED = 20
MODELS = ["VGG16", "DenseNet121", "EfficientNetB1", "MobileNetV2", "ResNet50"]

# Known patient counts per class
PATIENT_COUNTS = {
    "Bengin cases": 15,
    "Malignant cases": 40,
    "Normal cases": 55,
}

# ---------------------------------------------------------------------------
# Step 1: Build thumbnail-based patient manifest
# ---------------------------------------------------------------------------

print("=" * 60)
print("Building thumbnail-based patient manifest")
print("=" * 60)

# Collect all image filenames and class labels (same order as feature extraction)
image_paths, all_labels, all_fnames = [], [], []
for class_name in sorted(os.listdir(DATASET_PATH)):
    class_path = os.path.join(DATASET_PATH, class_name)
    if not os.path.isdir(class_path):
        continue
    for file in sorted(os.listdir(class_path)):
        if file.lower().endswith((".png", ".jpg", ".jpeg")):
            image_paths.append(os.path.join(class_path, file))
            all_labels.append(class_name)
            all_fnames.append(file)

print(f"Found {len(image_paths)} images across "
      f"{len(set(all_labels))} classes")

# Run thumbnail K-means
labels_arr = np.array(all_labels)
fnames_arr = np.array(all_fnames)

identifier = PatientIdentifier(
    patient_estimates=PATIENT_COUNTS, random_state=42)
groups = identifier.identify_from_thumbnails(
    DATASET_PATH, fnames_arr, labels_arr)

# Build manifest
manifest = identifier.build_assignment_dict(groups, fnames_arr, labels_arr)

# Save manifest
manifest_path = os.path.join(RESULTS_DIR, "thumbnail_patient_manifest.csv")


def f2n(fname):
    m = re.search(r"\((\d+)\)", fname)
    num = int(m.group(1)) if m else None
    for cls_key, prefix in [
        ("Bengin cases", "B"), ("Malignant cases", "M"), ("Normal cases", "N"),
    ]:
        if fname.startswith(cls_key.rstrip("s")):
            return f"{prefix}_{num:03d}" if num else fname
    return fname


with open(manifest_path, "w", newline="") as f:
    writer = csv.writer(f)
    writer.writerow(["patient_id", "class", "n_images", "images"])
    for pid in sorted(manifest.keys()):
        imgs = manifest[pid]
        short_names = [f2n(img) for img in imgs]
        cls = pid.split("_")[0]  # "benign_0" → "benign"
        cls = {"benign": "Benign", "malig": "Malignant",
               "normal": "Normal"}.get(cls, cls)
        writer.writerow([pid, cls, len(imgs), ";".join(short_names)])

print(f"Manifest saved → {manifest_path}")
print(f"  {len(manifest)} estimated patients")

# Per-class stats
for cls in ["Bengin cases", "Malignant cases", "Normal cases"]:
    short_cls = {"Bengin cases": "benign", "Malignant cases": "malig",
                 "Normal cases": "normal"}[cls]
    cls_patients = {k: v for k, v in manifest.items()
                    if k.startswith(short_cls)}
    n_pat = len(cls_patients)
    n_img = sum(len(v) for v in cls_patients.values())
    print(f"  {cls}: {n_pat} patients, {n_img} images")

# ---------------------------------------------------------------------------
# Step 2: Run classification with thumbnail manifest
# ---------------------------------------------------------------------------

print(f"\n{'=' * 60}")
print("Running classification (thumbnail manifest, seed=20)")
print("=" * 60)

clf = PatientLeakageClassifier(manifest_path, FEATURES_DIR, n_jobs=6)

results = []
for model_name in MODELS:
    print(f"\n  {model_name} ...", flush=True)
    img = clf.run(model_name, SEED, "image")
    pat = clf.run(model_name, SEED, "patient")
    results.append({
        "model": model_name,
        "manifest": "thumbnail",
        "image_cv": img["cv"],
        "image_test": img["test"],
        "patient_cv": pat["cv"],
        "patient_test": pat["test"],
        "drop": img["test"] - pat["test"],
    })
    print(f"    Image:  CV={img['cv']:.4f}  Test={img['test']:.4f}")
    print(f"    Patient: CV={pat['cv']:.4f}  Test={pat['test']:.4f}  "
          f"Drop={img['test'] - pat['test']:.4f}")

# ---------------------------------------------------------------------------
# Save results
# ---------------------------------------------------------------------------

out_path = os.path.join(RESULTS_DIR, "classification_thumbnail.json")
with open(out_path, "w") as f:
    json.dump(results, f, indent=2)

# Also load feature-based results for side-by-side comparison
feat_path = os.path.join(RESULTS_DIR, "classification_results.json")
if os.path.exists(feat_path):
    with open(feat_path) as f:
        feat_results = json.load(f)

    print(f"\n{'=' * 80}")
    print("COMPARISON: Feature-based (PCA-50) vs Thumbnail K-means")
    print(f"{'=' * 80}")
    print(f"{'Model':<18s} {'Feat Image':>11s} {'Thumb Image':>12s} "
          f"{'Feat Pat':>9s} {'Thumb Pat':>10s} {'Feat Drop':>10s} "
          f"{'Thumb Drop':>11s}")
    print("-" * 78)
    for tr, fr in zip(results, feat_results):
        assert tr["model"] == fr["model"]
        print(f"{tr['model']:<18s} {fr['image_test']:>11.4f} "
              f"{tr['image_test']:>12.4f} {fr['patient_test']:>9.4f} "
              f"{tr['patient_test']:>10.4f} {fr['drop']:>10.4f} "
              f"{tr['drop']:>11.4f}")

print(f"\nSaved → {out_path}")
print("DONE")