import torch
import torch.nn as nn
import torch.nn.functional as F

class AevanisTraverseHead(nn.Module):
    def __init__(self, visual_dim=256, hidden_dim=128):
        super().__init__()
        self.visual_dim = visual_dim

        # 14x14 tokens -> 28x28 dense foothold map
        self.dense_conv = nn.Sequential(
            nn.ConvTranspose2d(visual_dim, hidden_dim, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(hidden_dim),
            nn.GELU(),
            nn.Conv2d(hidden_dim, 32, kernel_size=3, padding=1),
            nn.GELU(),
            nn.Conv2d(32, 1, kernel_size=1),
            nn.Sigmoid()
        )

        self.global_pool = nn.AdaptiveAvgPool1d(1)

        # Slope & Incline Angle Head (degrees [-45, +45])
        self.slope_head = nn.Sequential(
            nn.Linear(visual_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, 2)
        )

        # Surface Slip Hazard Head [0, 1]
        self.slip_head = nn.Sequential(
            nn.Linear(visual_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, 1),
            nn.Sigmoid()
        )

    def forward(self, visual_tokens):
        B, N, C = visual_tokens.shape
        h = w = int(N ** 0.5)
        feat_2d = visual_tokens.permute(0, 2, 1).reshape(B, C, h, w)

        foothold_map = self.dense_conv(feat_2d).squeeze(1)
        global_feat = self.global_pool(visual_tokens.permute(0, 2, 1)).squeeze(-1)

        slope_deg = self.slope_head(global_feat) * 45.0
        slip_risk = self.slip_head(global_feat)

        return {
            "foothold_map": foothold_map,
            "slope_angles": slope_deg,
            "slip_risk": slip_risk
        }


class AevanisTraverseAgent(nn.Module):
    def __init__(self, model_scale="nano", head_weights=None, backbone=None):
        super().__init__()
        self.model_scale = model_scale
        visual_dim = 256 if model_scale == "nano" else 384
        self.traverse_head = AevanisTraverseHead(visual_dim=visual_dim)
        self.backbone = backbone

        if head_weights is not None:
            self.load_head(head_weights)

    def load_head(self, path: str):
        state = torch.load(path, map_location="cpu")
        cleaned = {k.replace("traverse_head.", ""): v for k, v in state.items()}
        self.traverse_head.load_state_dict(cleaned, strict=False)
        self.eval()

    def forward(self, visual_input):
        if visual_input.dim() == 4:
            if self.backbone is not None:
                with torch.no_grad():
                    features = self.backbone(visual_input)
                    tokens = features["spatial_tokens"]
            else:
                B = visual_input.shape[0]
                tokens = torch.zeros(B, 196, self.traverse_head.visual_dim)
        else:
            tokens = visual_input

        return self.traverse_head(tokens)

    @classmethod
    def from_pretrained(cls, model_scale="nano", head_weights=None):
        return cls(model_scale=model_scale, head_weights=head_weights)