"""
inspect_features.py — visualize U-Net feature maps at each encoder level for one
demo case, to check whether the transmitter ("tower") registers through the depth.

Run:  python src\\inspect_features.py
(or paste these blocks into Jupyter Lab cells to poke at individual channels)
"""
import numpy as np
import torch
import matplotlib.pyplot as plt
from pathlib import Path

from dataset import RadioMapSeerDataset
from model import UNet

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

ckpt = torch.load("checkpoints/best.pt", 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()

# capture each level's output with forward hooks
acts = {}
def hook(name):
    def fn(m, i, o): acts[name] = o.detach().float().cpu()
    return fn

model.inc.register_forward_hook(hook("inc 256"))
model.down1.register_forward_hook(hook("down1 128"))
model.down2.register_forward_hook(hook("down2 64"))
model.down3.register_forward_hook(hook("down3 32"))
model.down4.register_forward_hook(hook("down4 16 (bottleneck)"))

# one demo case from the test set
ds = RadioMapSeerDataset(split="test", img_size=256)
x, y = ds[0]
tx = np.unravel_index(int(np.argmax(x[1].numpy())), (256, 256))  # (row, col)
with torch.no_grad():
    pred = model(x.unsqueeze(0).to(device))

print("Transmitter at (row, col):", tx)
for k, v in acts.items():
    print(f"{k}: shape {tuple(v.shape)}")

levels = list(acts.keys())
fig, axes = plt.subplots(2, 5, figsize=(20, 8))

axes[0, 0].imshow(x[0].numpy(), cmap="gray")
axes[0, 0].scatter([tx[1]], [tx[0]], c="red", marker="x", s=50)
axes[0, 0].set_title("input: buildings + Tx"); axes[0, 0].axis("off")
axes[0, 1].imshow(y.squeeze().numpy(), cmap="viridis")
axes[0, 1].set_title("ground truth"); axes[0, 1].axis("off")
axes[0, 2].imshow(pred.squeeze().cpu().numpy(), cmap="viridis")
axes[0, 2].set_title("prediction"); axes[0, 2].axis("off")
axes[0, 3].axis("off"); axes[0, 4].axis("off")

# bottom row: max activation over channels at each level (best for spotting the Tx)
for ax, name in zip(axes[1], levels):
    fmap = acts[name][0].amax(dim=0).numpy()      # [h, w]
    h = fmap.shape[0]
    ax.imshow(fmap, cmap="magma")
    ax.scatter([tx[1] * h / 256], [tx[0] * h / 256],
               facecolors="none", edgecolors="cyan", s=70, linewidths=1.5)
    ax.set_title(name, fontsize=10); ax.axis("off")

fig.suptitle(f"U-Net feature maps (max over channels) — Tx at {tx}", fontsize=13)
fig.tight_layout()
Path("outputs").mkdir(exist_ok=True)
fig.savefig("outputs/feature_maps.png", dpi=110, bbox_inches="tight")
print("saved outputs/feature_maps.png")