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

ACTION_TYPES = [
    "CLICK",
    "DOUBLE_CLICK",
    "RIGHT_CLICK",
    "TYPE",
    "SCROLL_UP",
    "SCROLL_DOWN",
    "HOVER",
    "WAIT"
]
NUM_ACTIONS = len(ACTION_TYPES)

def tokenize(text: str, max_len: int = 32) -> torch.Tensor:
    tokens = [ord(c) % 256 for c in text.lower()[:max_len]]
    tokens += [0] * (max_len - len(tokens))
    return torch.tensor([tokens], dtype=torch.long)


class InstructionEncoder(nn.Module):
    def __init__(self, vocab_size=2048, embed_dim=256):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.proj = nn.Sequential(
            nn.Linear(embed_dim, embed_dim),
            nn.GELU(),
            nn.Linear(embed_dim, embed_dim)
        )

    def forward(self, token_ids):
        embeds = self.embedding(token_ids)
        pooled = embeds.mean(dim=1)
        return self.proj(pooled)


class AevanisOperatorHead(nn.Module):
    def __init__(self, visual_dim=256, hidden_dim=256, num_actions=NUM_ACTIONS):
        super().__init__()
        self.visual_dim = visual_dim
        self.hidden_dim = hidden_dim

        # Cross-Attention
        self.q_proj = nn.Linear(hidden_dim, hidden_dim)
        self.k_proj = nn.Linear(visual_dim, hidden_dim)
        self.v_proj = nn.Linear(visual_dim, hidden_dim)
        self.scale = 1.0 / math.sqrt(hidden_dim)

        # 14x14 Heatmap
        self.heatmap_head = nn.Sequential(
            nn.Linear(hidden_dim, 64),
            nn.GELU(),
            nn.Linear(64, 1)
        )

        # Action Classifier
        self.action_head = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, num_actions)
        )

        # High-Precision (X, Y) Coordinates (sub-pixel sigmoid)
        self.coord_head = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, 64),
            nn.GELU(),
            nn.Linear(64, 2),
            nn.Sigmoid()
        )

    def forward(self, visual_tokens, instr_embed):
        B, N, _ = visual_tokens.shape
        q = self.q_proj(instr_embed).unsqueeze(1)
        k = self.k_proj(visual_tokens)
        v = self.v_proj(visual_tokens)

        attn_scores = torch.bmm(q, k.transpose(1, 2)) * self.scale
        attn_weights = F.softmax(attn_scores, dim=-1)
        context = torch.bmm(attn_weights, v).squeeze(1)

        heatmap_logits = self.heatmap_head(k).squeeze(-1)
        heatmap_2d = heatmap_logits.reshape(B, 14, 14)

        multimodal = torch.cat([context, instr_embed], dim=-1)
        action_logits = self.action_head(multimodal)
        pred_coords = self.coord_head(multimodal)

        action_idx = torch.argmax(action_logits, dim=-1)
        action_names = [ACTION_TYPES[idx.item()] for idx in action_idx]

        return {
            "coords": pred_coords,
            "action_logits": action_logits,
            "action_name": action_names,
            "heatmap": heatmap_2d,
            "confidence": torch.softmax(action_logits, dim=-1).max(dim=-1)[0]
        }


class AevanisOperatorAgent(nn.Module):
    """Standalone Agent with self-contained head and optional backbone loading."""
    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.instr_encoder = InstructionEncoder(vocab_size=2048, embed_dim=256)
        self.operator_head = AevanisOperatorHead(visual_dim=visual_dim, hidden_dim=256)
        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("operator_head.", ""): v for k, v in state.items()}
        self.operator_head.load_state_dict(cleaned, strict=False)
        self.eval()

    def forward(self, screen_input, token_ids):
        # screen_input can be:
        # 1. 14x14 Spatial Tokens [B, 196, D] directly
        # 2. Raw RGB Image [B, 3, 224, 224] if backbone is attached
        if screen_input.dim() == 4:
            if self.backbone is not None:
                with torch.no_grad():
                    features = self.backbone(screen_input)
                    visual_tokens = features["spatial_tokens"]
            else:
                # Standalone fallback: dummy projection for self-contained testing
                B = screen_input.shape[0]
                visual_tokens = torch.zeros(B, 196, self.operator_head.visual_dim)
        else:
            visual_tokens = screen_input

        instr_embed = self.instr_encoder(token_ids)
        return self.operator_head(visual_tokens, instr_embed)

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