import glob
import math
import os

import torch
from PIL import Image
from torchvision import transforms

from train_network import UnrolledNetwork
from train_tnrd_baseline import TNRDBaselineNetwork


SIGMA = 25.0 / 255.0
TEST_DIR = "./datasets/Test_Datasets/FFDNet-master/testsets/Set12"
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def calculate_psnr(img1, img2):
    mse = torch.mean((img1 - img2) ** 2)
    if mse == 0:
        return float("inf")
    return 20 * math.log10(1.0 / math.sqrt(mse))


def _make_model(spec):
    if spec["kind"] == "telegraph":
        return UnrolledNetwork(num_stages=spec["stages"], use_wave=spec["wave"]).to(DEVICE)
    if spec["kind"] == "tnrd":
        return TNRDBaselineNetwork(num_stages=spec["stages"]).to(DEVICE)
    raise ValueError(f"Unknown model kind: {spec['kind']}")


def _autocast_context():
    return torch.amp.autocast("cuda") if DEVICE.type == "cuda" else torch.autocast("cpu", enabled=False)


def evaluate_model(spec):
    model = _make_model(spec)
    model.load_state_dict(torch.load(spec["file"], map_location=DEVICE))
    model.eval()

    test_transform = transforms.Compose([transforms.Grayscale(), transforms.ToTensor()])
    test_paths = sorted(glob.glob(os.path.join(TEST_DIR, "*.png")))
    if not test_paths:
        return "Error: No images found in Set12 directory."

    torch.manual_seed(42)
    total_psnr = 0.0
    with torch.no_grad():
        for path in test_paths:
            clean = test_transform(Image.open(path)).unsqueeze(0).to(DEVICE)
            noisy = torch.clamp(clean + torch.randn_like(clean) * SIGMA, 0.0, 1.0)

            with _autocast_context():
                output = model(noisy)

            total_psnr += calculate_psnr(clean, output)

    return f"{total_psnr / len(test_paths):.2f} dB"


def main():
    print("[*] Evaluating Checkpoints on Set12 (Sigma = 25)...")
    print("-" * 60)

    models_to_test = [
        {
            "name": "3-Stage TDE (Proposed)",
            "kind": "telegraph",
            "stages": 3,
            "wave": True,
            "file": "model_3stages_waveTrue.pth",
        },
        {
            "name": "5-Stage TDE (Proposed)",
            "kind": "telegraph",
            "stages": 5,
            "wave": True,
            "file": "model_5stages_waveTrue.pth",
        },
        {
            "name": "5-Stage Telegraph w/o wave",
            "kind": "telegraph",
            "stages": 5,
            "wave": False,
            "file": "model_5stages_waveFalse.pth",
        },
        {
            "name": "5-Stage TNRD baseline",
            "kind": "tnrd",
            "stages": 5,
            "file": "tnrd_baseline_5stages.pth",
        },
    ]

    for spec in models_to_test:
        if os.path.exists(spec["file"]):
            score = evaluate_model(spec)
            print(f"{spec['name']:<30} | PSNR: {score}")
        else:
            print(f"{spec['name']:<30} | PSNR: [File not found]")

    print("-" * 60)
    print("PDE baseline: run classical_baseline.py for the fixed telegraph/PDE comparison.")


if __name__ == "__main__":
    main()