"""
EEL4759 Final Project — Diabetic Retinopathy Classification
Compares ResNet-50, EfficientNet-B0, and ViT-B/16 with and without CLAHE preprocessing.

Usage:
    python main.py                          # run all 6 experiments (default)
    python main.py --models resnet50 --clahe 0 --epochs 2   # quick sanity check
"""

import argparse
import os

import torch
import torch.nn as nn
from torch.optim.lr_scheduler import CosineAnnealingLR
import kagglehub

from utils import set_seed, compute_class_weights, get_device, CLASS_NAMES
from dataset import split_dataset, build_dataloaders
from models import create_model
from train import train_model
from evaluate import evaluate_model, plot_confusion_matrix, plot_training_curves, generate_summary

RESULTS_DIR = 'results'


def parse_args():
    parser = argparse.ArgumentParser(description='Diabetic Retinopathy Classification')
    parser.add_argument(
        '--models', nargs='+',
        default=['resnet50', 'efficientnet_b0', 'vit_b_16'],
        help='Models to run'
    )
    parser.add_argument(
        '--clahe', nargs='+', type=int, default=[0, 1],
        help='CLAHE settings to run (0=off, 1=on)'
    )
    parser.add_argument('--epochs', type=int, default=30, help='Max training epochs')
    parser.add_argument('--patience', type=int, default=5, help='Early stopping patience')
    parser.add_argument('--batch-size', type=int, default=32, help='Training batch size')
    parser.add_argument('--lr', type=float, default=1e-4, help='Learning rate')
    parser.add_argument('--weight-decay', type=float, default=1e-4, help='AdamW weight decay')
    parser.add_argument('--seed', type=int, default=42, help='Random seed')
    parser.add_argument('--data-root', type=str, default=None,
                        help='Path to dataset root (auto-downloaded if not set)')
    parser.add_argument('--num-workers', type=int, default=4, help='DataLoader workers')
    return parser.parse_args()


def main():
    args = parse_args()
    set_seed(args.seed)
    device = get_device()
    print(f"Using device: {device}")

    # ── Dataset ──────────────────────────────────────────────────────────────
    if args.data_root is None:
        print("Downloading/locating dataset via kagglehub...")
        data_root = kagglehub.dataset_download("amanneo/diabetic-retinopathy-resized-arranged")
    else:
        data_root = args.data_root
    print(f"Dataset root: {data_root}")

    train_paths, train_labels, val_paths, val_labels, test_paths, test_labels = \
        split_dataset(data_root, seed=args.seed)

    # Class weights computed once from training labels
    class_weights = compute_class_weights(train_labels).to(device)
    print("Class weights:", [f"{w:.3f}" for w in class_weights.cpu().numpy()])

    # ── Experiments ──────────────────────────────────────────────────────────
    experiments = [
        (model_name, bool(use_clahe))
        for model_name in args.models
        for use_clahe in args.clahe
    ]

    all_results = {}
    os.makedirs(RESULTS_DIR, exist_ok=True)

    for model_name, use_clahe in experiments:
        exp_key = f"{model_name}_clahe={use_clahe}"
        exp_dir = os.path.join(RESULTS_DIR, exp_key)
        os.makedirs(exp_dir, exist_ok=True)

        print(f"\n{'='*70}")
        print(f"Experiment: {model_name}  |  CLAHE: {use_clahe}")
        print(f"{'='*70}")

        train_loader, val_loader, test_loader = build_dataloaders(
            train_paths, train_labels,
            val_paths, val_labels,
            test_paths, test_labels,
            use_clahe=use_clahe,
            batch_size=args.batch_size,
            num_workers=args.num_workers,
        )

        model = create_model(model_name, num_classes=5, pretrained=True).to(device)
        criterion = nn.CrossEntropyLoss(weight=class_weights)
        optimizer = torch.optim.AdamW(
            model.parameters(), lr=args.lr, weight_decay=args.weight_decay
        )
        scheduler = CosineAnnealingLR(optimizer, T_max=args.epochs)

        # Train
        train_results = train_model(
            model, train_loader, val_loader,
            criterion, optimizer, scheduler,
            device,
            num_epochs=args.epochs,
            patience=args.patience,
            model_name=exp_key,
        )

        # Save best checkpoint
        ckpt_path = os.path.join(exp_dir, 'best_model.pth')
        torch.save(train_results['best_model_state'], ckpt_path)
        print(f"  Checkpoint saved: {ckpt_path}")

        # Load best weights for evaluation
        model.load_state_dict(train_results['best_model_state'])

        # Evaluate on test set
        test_metrics = evaluate_model(model, test_loader, device)

        print(f"\n  Test results — Weighted F1: {test_metrics['weighted_f1']:.4f} | "
              f"Macro F1: {test_metrics['macro_f1']:.4f} | "
              f"Accuracy: {test_metrics['accuracy']:.4f} | "
              f"Kappa: {test_metrics['cohen_kappa']:.4f}")
        print(test_metrics['classification_report'])

        # Save per-class report
        report_path = os.path.join(exp_dir, 'classification_report.txt')
        with open(report_path, 'w') as f:
            f.write(f"Experiment: {exp_key}\n\n")
            f.write(test_metrics['classification_report'])

        # Plots
        plot_confusion_matrix(
            test_metrics['confusion_matrix'],
            save_path=os.path.join(exp_dir, 'confusion_matrix.png'),
        )
        plot_training_curves(
            train_results['train_losses'],
            train_results['val_losses'],
            train_results['val_f1s'],
            save_path=os.path.join(exp_dir, 'training_curves.png'),
        )

        all_results[exp_key] = {**train_results, **test_metrics}
        # Remove bulky state dict from summary dict (already saved to disk)
        all_results[exp_key].pop('best_model_state', None)

    # ── Summary ───────────────────────────────────────────────────────────────
    if len(all_results) > 1:
        generate_summary(all_results, save_dir=RESULTS_DIR)


if __name__ == '__main__':
    main()