from collections import OrderedDict
import torch
import torch.nn as nn
import torch.nn.functional as F


class SceneCNN(nn.Module):
    def __init__(self, base=16, num_classes=6, conv_dropout=0.1, fc_dropout=0.4):
        super().__init__()
        channels = [base, base * 2, base * 4, base * 8]

        layers = OrderedDict()
        in_ch = 3
        for i, out_ch in enumerate(channels, 1):
            layers[f"conv{i}a"] = nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False)
            layers[f"bn{i}a"]   = nn.BatchNorm2d(out_ch)
            layers[f"relu{i}a"] = nn.ReLU()
            layers[f"conv{i}b"] = nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False)
            layers[f"bn{i}b"]   = nn.BatchNorm2d(out_ch)
            layers[f"relu{i}b"] = nn.ReLU()
            layers[f"pool{i}"]  = nn.MaxPool2d(2)
            layers[f"drop{i}"]  = nn.Dropout2d(conv_dropout)
            in_ch = out_ch
        self.features = nn.Sequential(layers)

        self.gap = nn.AdaptiveAvgPool2d(1)
        self.flatten = nn.Flatten()
        self.drop_fc = nn.Dropout(fc_dropout)
        self.fc1 = nn.Linear(channels[-1], 128)
        self.relu_fc = nn.ReLU()
        self.fc2 = nn.Linear(128, num_classes)

    def forward(self, x):
        x = self.features(x)
        x = self.flatten(self.gap(x))
        x = self.drop_fc(x)
        x = self.relu_fc(self.fc1(x))
        return self.fc2(x)


def grad_cam(model, x):
    """x: normalized tensor (3,H,W). Returns (cam HxW in [0,1], softmax probs)."""
    model.eval()
    store = {}

    def hook(_, __, out):
        store["act"] = out
        out.register_hook(lambda g: store.__setitem__("grad", g))

    h = model.features.relu4b.register_forward_hook(hook)
    with torch.enable_grad():
        logits = model(x.unsqueeze(0))
        h.remove()
        cls = logits.argmax(1).item()
        model.zero_grad()
        logits[0, cls].backward()

    act, grad = store["act"][0], store["grad"][0]
    weights = grad.mean(dim=(1, 2))
    cam = F.relu((weights[:, None, None] * act).sum(0))
    cam = F.interpolate(cam[None, None], size=x.shape[1:], mode="bilinear", align_corners=False)[0, 0]
    cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)
    probs = F.softmax(logits, 1)[0]
    return cam.detach().numpy(), probs.detach().numpy()