import os
import time
import torch
import numpy as np
import torch.nn.functional as F
import matplotlib.pyplot as plt

from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from transformers import ViTForImageClassification, ViTConfig
from sklearn.metrics import (
    accuracy_score,
    precision_score,
    recall_score,
    f1_score,
    roc_auc_score,
    confusion_matrix,
    roc_curve
)
from openpyxl import Workbook


def main():
    # ----------------------------------
    # PERFORMANCE TUNING
    # ----------------------------------
    torch.backends.cudnn.benchmark = True
    torch.set_num_threads(8)

    # ----------------------------------
    # Device (SAFE)
    # ----------------------------------
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    use_amp = device.type == "cuda"

    if use_amp:
        print(f"🚀 Using device: cuda ({torch.cuda.get_device_name(0)})")
    else:
        print("🚀 Using device: cpu")

    # ----------------------------------
    # Paths
    # ----------------------------------
    base_data_dir = "data"

    # 🔴 UPDATED MODEL PATH (NEW MODEL)
    model_path = "model/vit_face_final_best.pth"

    datasets_config = {
        "FF++": f"{base_data_dir}/ff++/test",
        "Celeb-DF": f"{base_data_dir}/celeb-df/test",
        "DFDC": f"{base_data_dir}/dfdc/test",
    }

    os.makedirs("outputs", exist_ok=True)

    # ----------------------------------
    # Transforms
    # ----------------------------------
    test_tfms = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        )
    ])

    # ----------------------------------
    # Model
    # ----------------------------------
    config = ViTConfig.from_pretrained(
        "google/vit-base-patch16-224",
        num_labels=2
    )

    model = ViTForImageClassification.from_pretrained(
        "google/vit-base-patch16-224",
        config=config,
        ignore_mismatched_sizes=True
    )

    model.load_state_dict(torch.load(model_path, map_location=device))
    model.to(device)
    model.eval()

    print("✅ Model loaded successfully")

    # ----------------------------------
    # Evaluation
    # ----------------------------------
    all_results = {}
    print("\n🧪 Running Cross-Dataset Evaluation...\n")

    for ds_name, test_dir in datasets_config.items():
        if not os.path.exists(test_dir):
            print(f"⚠️ {ds_name}: path not found — skipping")
            continue

        test_ds = datasets.ImageFolder(test_dir, transform=test_tfms)

        test_dl = DataLoader(
            test_ds,
            batch_size=32,
            shuffle=False,
            num_workers=2,      # SAFE (inside main)
            pin_memory=True
        )

        class_to_idx = test_ds.class_to_idx
        real_idx = class_to_idx["real"]
        fake_idx = class_to_idx["fake"]

        print(f"📂 {ds_name}: {len(test_ds)} samples | {class_to_idx}")

        y_true, y_pred, y_probs = [], [], []
        total_images = 0
        total_time = 0.0

        with torch.no_grad():
            for imgs, labels in test_dl:
                imgs = imgs.to(device, non_blocking=True)
                labels = labels.to(device, non_blocking=True)

                start = time.time()
                with torch.amp.autocast(device_type="cuda", enabled=use_amp):
                    logits = model(imgs).logits
                end = time.time()

                probs = F.softmax(logits, dim=1)[:, fake_idx]
                preds = torch.argmax(logits, dim=1)

                total_time += (end - start)
                total_images += imgs.size(0)

                y_true.extend(labels.cpu().numpy())
                y_pred.extend(preds.cpu().numpy())
                y_probs.extend(probs.cpu().numpy())

        # ---------------- Metrics ----------------
        acc = accuracy_score(y_true, y_pred)
        prec = precision_score(y_true, y_pred, zero_division=0)
        rec = recall_score(y_true, y_pred, zero_division=0)
        f1 = f1_score(y_true, y_pred, zero_division=0)

        try:
            auc = roc_auc_score(
                (np.array(y_true) == fake_idx).astype(int),
                y_probs
                )
        except ValueError:
            auc = float("nan")

        avg_time_ms = (total_time / total_images) * 1000
        fps = total_images / total_time

        all_results[ds_name] = {
            "acc": acc,
            "prec": prec,
            "rec": rec,
            "f1": f1,
            "auc": auc,
            "time": avg_time_ms,
            "fps": fps,
        }

        print(
            f"🎯 {ds_name} | Acc: {acc:.4f} | F1: {f1:.4f} | "
            f"AUC: {auc:.4f} | {avg_time_ms:.2f} ms/img | {fps:.2f} FPS"
        )

        # ---------------- Confusion Matrix ----------------
        cm = confusion_matrix(y_true, y_pred, labels=[real_idx, fake_idx])

        plt.figure(figsize=(5, 4))
        plt.imshow(cm)
        plt.title(f"{ds_name} - Confusion Matrix")
        plt.xlabel("Predicted")
        plt.ylabel("True")
        plt.xticks([0, 1], ["Real", "Fake"])
        plt.yticks([0, 1], ["Real", "Fake"])

        for i in range(2):
            for j in range(2):
                plt.text(j, i, cm[i, j], ha="center", va="center")

        plt.tight_layout()
        plt.savefig(f"outputs/cm_{ds_name}.png")
        plt.close()

        # ---------------- ROC Curve ----------------
        fpr, tpr, _ = roc_curve(y_true, y_probs, pos_label=fake_idx)

        plt.figure(figsize=(5, 4))
        plt.plot(fpr, tpr, label=f"AUC = {auc:.4f}")
        plt.plot([0, 1], [0, 1], linestyle="--")
        plt.xlabel("False Positive Rate")
        plt.ylabel("True Positive Rate")
        plt.title(f"{ds_name} - ROC Curve")
        plt.legend(loc="lower right")
        plt.tight_layout()
        plt.savefig(f"outputs/roc_{ds_name}.png")
        plt.close()

    # ----------------------------------
    # EXPORT (EXCEL + TEXT)
    # ----------------------------------
    avg_metrics = {
        "acc": np.mean([m["acc"] for m in all_results.values()]),
        "prec": np.mean([m["prec"] for m in all_results.values()]),
        "rec": np.mean([m["rec"] for m in all_results.values()]),
        "f1": np.mean([m["f1"] for m in all_results.values()]),
        "auc": np.nanmean([m["auc"] for m in all_results.values()]),
        "time": np.mean([m["time"] for m in all_results.values()]),
        "fps": np.mean([m["fps"] for m in all_results.values()]),
    }

    all_results["AVERAGE"] = avg_metrics

    wb = Workbook()
    ws = wb.active
    ws.title = "Evaluation Results"
    ws.append(["Dataset", "Accuracy", "Precision", "Recall", "F1", "AUC", "ms/img", "FPS"])

    for ds, m in all_results.items():
        ws.append([
            ds,
            round(m["acc"], 4),
            round(m["prec"], 4),
            round(m["rec"], 4),
            round(m["f1"], 4),
            round(m["auc"], 4),
            round(m["time"], 2),
            round(m["fps"], 2),
        ])

    wb.save("outputs/evaluation_results.xlsx")

    with open("outputs/summary.txt", "w") as f:
        for ds, m in all_results.items():
            f.write(f"Dataset: {ds}\n")
            for k, v in m.items():
                f.write(f"  {k.upper():<6}: {v}\n")
            f.write("-" * 45 + "\n")

    print("✅ Evaluation complete (new model evaluated successfully)")

    # ----------------------------------
    # COMBINED ROC CURVE (IEEE-FRIENDLY)
    # ----------------------------------
    plt.figure(figsize=(5, 4))

    for ds_name in datasets_config.keys():
        if ds_name not in all_results:
            continue

        # Reload dataset to recompute ROC cleanly
        test_dir = datasets_config[ds_name]
        test_ds = datasets.ImageFolder(test_dir, transform=test_tfms)

        class_to_idx = test_ds.class_to_idx
        fake_idx = class_to_idx["fake"]

        test_dl = DataLoader(
            test_ds,
            batch_size=32,
            shuffle=False,
            num_workers=2,
            pin_memory=True
        )

        y_true, y_probs = [], []

        with torch.no_grad():
            for imgs, labels in test_dl:
                imgs = imgs.to(device, non_blocking=True)
                labels = labels.to(device, non_blocking=True)

                with torch.amp.autocast(device_type="cuda", enabled=use_amp):
                    logits = model(imgs).logits

                probs = F.softmax(logits, dim=1)[:, fake_idx]

                y_true.extend((labels == fake_idx).cpu().numpy())
                y_probs.extend(probs.cpu().numpy())

        fpr, tpr, _ = roc_curve(y_true, y_probs)
        auc_val = roc_auc_score(y_true, y_probs)

        # Line styles for IEEE (print-safe)
        if ds_name == "FF++":
            style = "-"
        elif ds_name == "Celeb-DF":
            style = "--"
        else:  # DFDC
            style = "-."

        plt.plot(
            fpr,
            tpr,
            linestyle=style,
            linewidth=2,
            label=f"{ds_name} (AUC = {auc_val:.4f})"
        )

    # Random baseline
    plt.plot([0, 1], [0, 1], linestyle=":", linewidth=1)

    plt.xlabel("False Positive Rate")
    plt.ylabel("True Positive Rate")
    plt.legend(loc="lower right")
    plt.tight_layout()

    # Save IEEE-ready figure
    plt.savefig("outputs/roc_combined.png", dpi=300, bbox_inches="tight")
    plt.savefig("outputs/roc_combined.pdf", dpi=300, bbox_inches="tight")
    plt.close()


if __name__ == "__main__":
    main()