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

End-to-end pipeline for all five CNNs from the manuscript:
  VGG16, DenseNet121, EfficientNetB1, MobileNetV2, ResNet50

For each model:
  1. Extract deep features and save to features/{Model}_features.npz
  2. Generate PCA / t-SNE plot → plots/{Model}_pca_tsne.png
  3. Estimate patient groups via K-means → features/{Model}_patient_groups.npy

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

import os
import sys
import numpy as np

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

from classes import FeatureExtractor, PatientIdentifier
from classes.visualizations import plot_pca_tsne

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

BASE_PATH = os.path.join(
    os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))),
    "The IQ-OTHNCCD lung cancer dataset"
)

PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
FEATURES_DIR = os.path.join(PROJECT_ROOT, "features")
PLOTS_DIR = os.path.join(PROJECT_ROOT, "plots")

os.makedirs(FEATURES_DIR, exist_ok=True)
os.makedirs(PLOTS_DIR, exist_ok=True)

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

PATIENT_ESTIMATES = {
    "Bengin cases": 15,
    "Malignant cases": 40,
    "Normal cases": 55,
}

LEGEND_NAMES = {
    "Bengin cases": "Benign",
    "Malignant cases": "Malignant",
    "Normal cases": "Normal",
}

# ---------------------------------------------------------------------------
# Run pipeline for each model
# ---------------------------------------------------------------------------

for model_name in MODELS:
    print("\n" + "=" * 60)
    print(f"MODEL: {model_name}")
    print("=" * 60)

    # --- Step 1: Extract features ---
    print("\n  [1/3] Feature extraction ...")
    extractor = FeatureExtractor(model_name=model_name)
    features_path = os.path.join(FEATURES_DIR, f"{model_name}_features.npz")

    if os.path.exists(features_path):
        print(f"    Loading cached features from {features_path}")
        X, Y, filenames = FeatureExtractor.load_features(features_path)
    else:
        X, Y, filenames = extractor.extract(BASE_PATH)
        extractor.save_features(X, Y, filenames, FEATURES_DIR)

    print(f"    {model_name}: X shape = {X.shape}")

    # --- Step 2: PCA / t-SNE ---
    print(f"\n  [2/3] PCA / t-SNE visualization ...")
    plot_path = os.path.join(PLOTS_DIR, f"{model_name}_pca_tsne.png")
    plot_pca_tsne(X, Y, legend_names=LEGEND_NAMES, output_path=plot_path)

    # --- Step 3: Patient identification ---
    print(f"\n  [3/3] Patient identification (K-means) ...")
    identifier = PatientIdentifier(patient_estimates=PATIENT_ESTIMATES)
    groups = identifier.identify(BASE_PATH, filenames, Y)

    groups_path = os.path.join(FEATURES_DIR, f"{model_name}_patient_groups.npy")
    np.save(groups_path, groups)
    print(f"    Saved patient groups → {groups_path}")

# ---------------------------------------------------------------------------
# Summary
# ---------------------------------------------------------------------------

print("\n" + "=" * 60)
print("ALL MODELS COMPLETE")
print("=" * 60)
for model_name in MODELS:
    fp = os.path.join(FEATURES_DIR, f"{model_name}_features.npz")
    gp = os.path.join(FEATURES_DIR, f"{model_name}_patient_groups.npy")
    pp = os.path.join(PLOTS_DIR, f"{model_name}_pca_tsne.png")
    print(f"  {model_name:18s}  features: {os.path.basename(fp):30s}  groups: {os.path.basename(gp)}")