import argparse
import os
import torch
from models.dou import DoUNet, load_clip_to_cpu
from utils.options import TrainCfg
import numpy as np
import torch.nn.functional as F
from sklearn.metrics import accuracy_score, average_precision_score, f1_score
from torch.utils.data import DataLoader
from dataset.dataset import ImageFolderWithPaths
import pandas as pd

DEFAULT_TEST_DATASETS = [
    "biggan", "cyclegan", "gaugan", "progan", "stargan", "stylegan", "dalle", "guided",
    "deepfake", "seeingdark", "san", "crn", "imle",
    "glide_50_27", "glide_100_10", "glide_100_27", "ldm_100", "ldm_200", "ldm_200_cfg",
    "biggan_genimage",
    "ADM", "Glide", "Midjourney", "stable_diffusion_v_1_4", "stable_diffusion_v_1_5", "VQDM",
    "wukong",
    "Infinity", "LlamaGen", "PixArt-XL", "SD35-L", "Show-o", "VAR",
    "FLUX", "Janus", "Janus-Pro-1B", "Janus-Pro-7B",
]


def parse_args():
    parser = argparse.ArgumentParser(description="Evaluate DoU on one or more test datasets")
    parser.add_argument("--method", default="DoUNet", help="Model method name stored in TrainCfg")
    parser.add_argument(
        "--data-root",
        help="Root directory containing test dataset folders",
    )
    parser.add_argument(
        "--datasets",
        nargs="+",
        default=DEFAULT_TEST_DATASETS,
        help="Dataset folder names under --data-root, separated by spaces",
    )
    parser.add_argument("--ckpt", default="runs/best.pth", help="Path to the saved model checkpoint")
    parser.add_argument("--output", default="result.xlsx", help="Path to save the Excel result file")
    parser.add_argument(
        "--device",
        default="auto",
        help="Device, e.g. auto, cuda, cuda:0, or cpu",
    )
    parser.add_argument("--batch-size", type=int, default=256, help="Evaluation batch size")
    parser.add_argument("--num-workers", type=int, default=4, help="Number of dataloader workers")
    parser.add_argument("--model-name", default="ViT-L/14", help="CLIP backbone name")
    parser.add_argument("--disable-amp", action="store_true", help="Disable AMP during evaluation")
    parser.add_argument("--skip-missing", action="store_true", help="Skip dataset folders that do not exist")
    return parser.parse_args()


def build_model(args, device):
    cfg = TrainCfg(
        method=args.method,
        model_name=args.model_name,
        classnames=("real", "generated"),
        learn_logit_scale=False,
        amp=not args.disable_amp,
    )

    clip_model = load_clip_to_cpu(cfg)
    clip_model.eval()
    for param in clip_model.parameters():
        param.requires_grad_(False)

    model = DoUNet(cfg, clip_model).to(device)
    ckpt = torch.load(args.ckpt, map_location="cpu", weights_only=False)
    model.load_state_dict(ckpt["model"], strict=False)
    model.eval()
    return model, cfg


def evaluate_dataset(model, cfg, dataset_path, batch_size, num_workers, device):
    val_dataset = ImageFolderWithPaths(dataset_path)
    val_dataloader = DataLoader(
        val_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=device.type == "cuda",
    )

    preds, truths, scores = [], [], []
    for imgs, y in val_dataloader:
        imgs = imgs.to(device)
        y = y.to(device)

        with torch.no_grad(), torch.amp.autocast("cuda", enabled=cfg.amp and device.type == "cuda"):
            logits = model(imgs)
            probs = F.softmax(logits, dim=-1)
            pred = probs.argmax(dim=-1)

        preds.extend(pred.flatten().tolist())
        truths.extend(y.flatten().tolist())
        scores.extend(probs[:, 1].detach().cpu().numpy().tolist())

    y_true = np.array(truths)
    y_pred = np.array(preds)
    y_score = np.array(scores)

    acc = accuracy_score(y_true, y_pred)
    r_acc = accuracy_score(y_true[y_true == 0], y_pred[y_true == 0])
    f_acc = accuracy_score(y_true[y_true == 1], y_pred[y_true == 1])
    f1 = f1_score(y_true, y_pred)
    ap_score = average_precision_score(y_true, y_score)

    return {
        "ACC": acc * 100,
        "Real ACC": r_acc * 100,
        "Fake ACC": f_acc * 100,
        "F1": f1 * 100,
        "AP": ap_score * 100,
    }


def main():
    args = parse_args()

    if args.device == "auto":
        args.device = "cuda" if torch.cuda.is_available() else "cpu"
    device = torch.device(args.device)
    model, cfg = build_model(args, device)

    max_len = max(len(dataset) for dataset in args.datasets) + 2
    results = []
    for dataset in args.datasets:
        dataset_path = os.path.join(args.data_root, dataset)
        if not os.path.isdir(dataset_path):
            message = f"Dataset folder not found: {dataset_path}"
            if args.skip_missing:
                print(f"Skip {dataset}: {message}")
                continue
            raise FileNotFoundError(message)

        metrics = evaluate_dataset(
            model=model,
            cfg=cfg,
            dataset_path=dataset_path,
            batch_size=args.batch_size,
            num_workers=args.num_workers,
            device=device,
        )

        print(
            f"{dataset:<{max_len}} | "
            f"{metrics['ACC']:6.2f}% | "
            f"{metrics['Real ACC']:7.2f}% | "
            f"{metrics['Fake ACC']:7.2f}% | "
            f"{metrics['F1']:6.2f}% | "
            f"{metrics['AP']:6.2f}%"
        )
        results.append({"Dataset": dataset, **metrics})

    if not results:
        raise RuntimeError("No datasets were evaluated")

    df = pd.DataFrame(results)
    df = df.set_index("Dataset").transpose()
    df.to_excel(args.output)

    acc_mean = sum(item["ACC"] for item in results) / len(results)
    print(f"Average ACC: {acc_mean:.2f}%")
    print(f"Done! Results saved to {args.output}")


if __name__ == "__main__":
    main()