"""
evaluate.py — quantitative evaluation on held-out TEST maps (unseen cities).

Reports U-Net RMSE (gray + dB), a log-distance path-loss baseline, and the
relative improvement. dB conversion: RadioMapSeer scales an ~80 dB pathloss
range into [0,1], so RMSE_dB = RMSE_gray * 80 (Levie et al.; Yapar et al.).
"""

import argparse
import numpy as np
import torch
from torch.utils.data import DataLoader

from dataset import RadioMapSeerDataset
from model import UNet

DB_DYNAMIC_RANGE = 80.0  # dB spanned by the [0,1] gray scale


def tx_location(antenna_img):
    return np.unravel_index(int(np.argmax(antenna_img)), antenna_img.shape)


def fit_log_distance(dataset, n_fit=200, seed=0):
    """Fit gray ~ a + b*log10(distance) by least squares (distance-only baseline)."""
    rng = np.random.default_rng(seed)
    idxs = rng.choice(len(dataset), size=min(n_fit, len(dataset)), replace=False)
    H = W = dataset.img_size
    yy, xx = np.mgrid[0:H, 0:W]
    logd_all, gray_all = [], []
    for i in idxs:
        x, y = dataset[int(i)]
        r, c = tx_location(x[1].numpy())
        dist = np.sqrt((yy - r) ** 2 + (xx - c) ** 2) + 1.0
        logd_all.append(np.log10(dist).ravel())
        gray_all.append(y.squeeze().numpy().ravel())
    logd, gray = np.concatenate(logd_all), np.concatenate(gray_all)
    A = np.vstack([np.ones_like(logd), logd]).T
    (a, b), *_ = np.linalg.lstsq(A, gray, rcond=None)
    return float(a), float(b)


def baseline_predict(antenna_img, a, b):
    H, W = antenna_img.shape
    yy, xx = np.mgrid[0:H, 0:W]
    r, c = tx_location(antenna_img)
    dist = np.sqrt((yy - r) ** 2 + (xx - c) ** 2) + 1.0
    return np.clip(a + b * np.log10(dist), 0.0, 1.0)


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--ckpt", default="checkpoints/best.pt")
    p.add_argument("--img-size", type=int, default=256)
    p.add_argument("--batch-size", type=int, default=32)
    p.add_argument("--num-workers", type=int, default=4)
    p.add_argument("--max-test", type=int, default=None, help="limit test samples for speed")
    args = p.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"

    ckpt = torch.load(args.ckpt, map_location=device, weights_only=False)
    base = ckpt.get("args", {}).get("base", 64)
    model = UNet(2, 1, base=base).to(device)
    model.load_state_dict(ckpt["model_state"])
    model.eval()
    print(f"Loaded {args.ckpt} (epoch {ckpt.get('epoch', '?')})")

    test_ds = RadioMapSeerDataset(split="test", img_size=args.img_size,
                                  max_samples=args.max_test)
    print(f"Test samples: {len(test_ds)}")

    # --- U-Net error over the test set ---
    loader = DataLoader(test_ds, batch_size=args.batch_size, shuffle=False,
                        num_workers=args.num_workers, pin_memory=True)
    sq_err, n_px = 0.0, 0
    with torch.no_grad():
        for x, y in loader:
            x, y = x.to(device), y.to(device)
            with torch.autocast(device_type=device, dtype=torch.float16,
                                enabled=(device == "cuda")):
                pred = model(x)
            sq_err += torch.sum((pred.float() - y.float()) ** 2).item()
            n_px += y.numel()
    unet_rmse = (sq_err / n_px) ** 0.5

    # --- log-distance baseline: fit on train, evaluate on test ---
    train_ds = RadioMapSeerDataset(split="train", img_size=args.img_size)
    a, b = fit_log_distance(train_ds, n_fit=200)
    print(f"Baseline fit: gray = {a:.3f} + ({b:.3f}) * log10(dist)")

    bsq, bn = 0.0, 0
    for i in range(len(test_ds)):
        x, y = test_ds[i]
        pred = baseline_predict(x[1].numpy(), a, b)
        truth = y.squeeze().numpy()
        bsq += float(np.sum((pred - truth) ** 2))
        bn += truth.size
    base_rmse = (bsq / bn) ** 0.5

    print("\n" + "=" * 54)
    print("  EVALUATION ON UNSEEN TEST CITIES")
    print("=" * 54)
    print(f"  U-Net (deep learning)  RMSE: {unet_rmse:.4f} gray | {unet_rmse*DB_DYNAMIC_RANGE:5.2f} dB")
    print(f"  Log-distance baseline  RMSE: {base_rmse:.4f} gray | {base_rmse*DB_DYNAMIC_RANGE:5.2f} dB")
    print(f"  Improvement over baseline:   {100*(base_rmse-unet_rmse)/base_rmse:.1f}%")
    print("=" * 54)


if __name__ == "__main__":
    main()