import os
from types import SimpleNamespace

import torch
from torchvision.utils import save_image

from utils import load_config, psnr, resolve_device, ssim
from model import build_model
from dataloader import get_dataloaders, VARIANT_TO_IDX


@torch.no_grad()
def evaluate(
    config_path: str = "config.yaml",
    ckpt_path: str | None = None,
    num_samples: int | None = None,
    output_dir: str | None = None,
) -> None:
    cfg = load_config(config_path)
    data_cfg = cfg.data
    training_cfg = cfg.training
    eval_cfg = getattr(
        cfg,
        "evaluation",
        SimpleNamespace(save_dir="eval_samples", output_dir="eval_out", num_workers=2, save_images=False),
    )

    device = resolve_device(getattr(training_cfg, "device", None))
    ckpt = ckpt_path or os.path.join(training_cfg.out_dir, "model_ema.pt")
    if not os.path.exists(ckpt):
        fallback_ckpt = os.path.join(training_cfg.out_dir, "model_last.pt")
        if os.path.exists(fallback_ckpt):
            print(f"[Eval] Checkpoint {ckpt} not found. Falling back to {fallback_ckpt}.")
            ckpt = fallback_ckpt
        else:
            raise FileNotFoundError(f"No checkpoint found at {ckpt} or {fallback_ckpt}")
    samples = num_samples
    save_images = bool(getattr(eval_cfg, "save_images", False))

    out_dir = output_dir or getattr(eval_cfg, "output_dir", "eval_out")
    if not os.path.isabs(out_dir):
        out_dir = os.path.join(training_cfg.out_dir, out_dir)
    os.makedirs(out_dir, exist_ok=True)

    print(f"[Eval] device={device}, output_dir={out_dir}")

    _, test_loader = get_dataloaders(
        data_root=data_cfg.root,
        batch_size=training_cfg.batch_size,
        image_size=data_cfg.image_size,
        train_variant=data_cfg.train_variant,
        test_variant=data_cfg.test_variant,
        train_num_workers=data_cfg.train_num_workers,
        test_num_workers=getattr(eval_cfg, "num_workers", data_cfg.test_num_workers),
        pin_memory=data_cfg.pin_memory,
    )

    model = build_model(cfg, device=device)
    state = torch.load(ckpt, map_location=device)
    model.load_state_dict(state)
    model.eval()

    psnr_vals, ssim_vals = [], []

    for i, batch in enumerate(test_loader):
        if samples is not None and i >= samples:
            break

        y = batch["y"].to(device)
        x = batch["gt"].to(device)
        variant_idx = torch.tensor(
            [VARIANT_TO_IDX[v] for v in batch["variant"]],
            device=device,
            dtype=torch.long,
        )

        x_hat = model.sample(y, variant_idx)

        if i == 0:
            print(
                f"[Debug] x_hat stats: min={x_hat.min().item():.4f}, "
                f"max={x_hat.max().item():.4f}, mean={x_hat.mean().item():.4f}"
            )

        p = psnr(x_hat, x).item()
        s = ssim(x_hat, x).item()
        psnr_vals.append(p)
        ssim_vals.append(s)

        name = f"sample_{i:05d}"
        if save_images:
            cat = torch.cat([y, x_hat, x], dim=3)
            save_path = os.path.join(out_dir, f"{name}_cmp.png")
            save_image(cat, save_path)

        print(f"{name}: PSNR={p:.2f} dB, SSIM={s:.4f}")

    mean_psnr = sum(psnr_vals) / max(1, len(psnr_vals))
    mean_ssim = sum(ssim_vals) / max(1, len(ssim_vals))
    print(f"Average over {len(psnr_vals)} images -> PSNR={mean_psnr:.2f} dB, SSIM={mean_ssim:.4f}")


if __name__ == "__main__":
    evaluate()
