from __future__ import annotations

import hashlib
import math
import re
from typing import Any

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import Dinov2Config, Dinov2Model


ACTION_DIM = 7
DEFAULT_TEXT_DIM = 128
DEFAULT_PHASE_DIM = 4
DEFAULT_DIFFICULTIES = ["low", "medium", "hard", "very_high"]
TEXT_FEATURE_VERSION = "signed_hash_subtokens_v2"

# Keep train/eval text handling dependency-free, but collapse common instruction
# paraphrases so the tiny text branch can transfer across public-eval wording.
TOKEN_ALIASES = {
    "build": "stack",
    "clean": "wipe",
    "container": "bin",
    "grab": "pick",
    "grasp": "pick",
    "move": "place",
    "put": "place",
    "sweep": "wipe",
}


def text_vector(text: str, dim: int = DEFAULT_TEXT_DIM) -> np.ndarray:
    """Return the same stable, dependency-free text feature in train and eval."""
    vec = np.zeros(dim, dtype=np.float32)
    raw_tokens = re.findall(r"[a-z0-9_]+", str(text).lower())
    tokens: list[str] = []
    for raw_token in raw_tokens:
        token = TOKEN_ALIASES.get(raw_token, raw_token)
        tokens.append(token)
        # Task names such as pick_place_milk used to hash to one opaque token.
        # Retaining the full token and its parts gives unseen task names a useful
        # compositional overlap with the six training tasks.
        if "_" in token:
            tokens.extend(TOKEN_ALIASES.get(part, part) for part in token.split("_"))
    for token in tokens:
        digest = hashlib.blake2b(token.encode("utf-8"), digest_size=8).digest()
        value = int.from_bytes(digest, byteorder="little", signed=False)
        vec[value % dim] += 1.0 if value & 1 else -1.0
    norm = float(np.linalg.norm(vec))
    if norm > 0:
        vec /= norm
    return vec


def phase_vector(step: int, horizon: int) -> np.ndarray:
    horizon_f = max(float(horizon), 1.0)
    progress = float(np.clip(float(step) / horizon_f, 0.0, 1.0))
    return np.asarray(
        [
            progress,
            1.0 - progress,
            math.sin(math.pi * progress),
            math.cos(math.pi * progress),
        ],
        dtype=np.float32,
    )


def normalize_images(images: torch.Tensor, image_size: int = 224) -> torch.Tensor:
    """Convert NHWC/NCHW uint8-like images to DINOv2's normalized input."""
    if images.ndim != 4:
        raise ValueError(f"expected a 4D image tensor, got {tuple(images.shape)}")
    if images.shape[-1] in (1, 3, 4):
        images = images[..., :3].permute(0, 3, 1, 2)
    elif images.shape[1] != 3:
        raise ValueError(f"cannot determine image channels for {tuple(images.shape)}")
    images = images.float()
    if images.numel() and float(images.detach().max().cpu()) > 2.0:
        images = images / 255.0
    if tuple(images.shape[-2:]) != (image_size, image_size):
        images = F.interpolate(
            images, size=(image_size, image_size), mode="bilinear", align_corners=False
        )
    mean = images.new_tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
    std = images.new_tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
    return (images - mean) / std


def rgb_grid_tokens(
    images: torch.Tensor, spatial_grid: int, image_size: int = 224
) -> torch.Tensor:
    """Keep cheap absolute color cues that are useful in synthetic robot scenes."""
    if images.ndim != 4:
        raise ValueError(f"expected a 4D image tensor, got {tuple(images.shape)}")
    if images.shape[-1] in (1, 3, 4):
        images = images[..., :3].permute(0, 3, 1, 2)
    elif images.shape[1] != 3:
        raise ValueError(f"cannot determine image channels for {tuple(images.shape)}")
    images = images.float()
    if images.numel() and float(images.detach().max().cpu()) > 2.0:
        images = images / 255.0
    if tuple(images.shape[-2:]) != (image_size, image_size):
        images = F.interpolate(
            images, size=(image_size, image_size), mode="bilinear", align_corners=False
        )
    global_rgb = images.mean(dim=(-2, -1)).unsqueeze(1)
    grid_rgb = F.adaptive_avg_pool2d(images, (spatial_grid, spatial_grid))
    grid_rgb = grid_rgb.flatten(2).transpose(1, 2)
    return torch.cat([global_rgb, grid_rgb], dim=1)


def pool_dinov2_tokens(hidden_states: torch.Tensor, spatial_grid: int) -> torch.Tensor:
    """Keep CLS plus a small spatial grid instead of discarding local geometry."""
    if hidden_states.ndim != 3 or hidden_states.shape[1] < 2:
        raise ValueError(f"unexpected DINOv2 output: {tuple(hidden_states.shape)}")
    cls_token = hidden_states[:, :1]
    patches = hidden_states[:, 1:]
    side = int(round(math.sqrt(int(patches.shape[1]))))
    if side * side != int(patches.shape[1]):
        raise ValueError(f"patch count {patches.shape[1]} is not a square")
    patches = patches.transpose(1, 2).reshape(
        patches.shape[0], patches.shape[2], side, side
    )
    patches = F.adaptive_avg_pool2d(patches, (spatial_grid, spatial_grid))
    patches = patches.flatten(2).transpose(1, 2)
    return torch.cat([cls_token, patches], dim=1)


class CompetitivePolicyHead(nn.Module):
    """Task-conditioned spatial pooling followed by a short temporal policy."""

    def __init__(
        self,
        vision_dim: int,
        proprio_dim: int,
        text_dim: int,
        phase_dim: int,
        num_tasks: int,
        num_difficulties: int,
        hidden_dim: int = 256,
        history: int = 4,
        action_chunk: int = 8,
        ensemble_heads: int = 3,
        dropout: float = 0.10,
    ) -> None:
        super().__init__()
        self.vision_dim = int(vision_dim)
        self.proprio_dim = int(proprio_dim)
        self.text_dim = int(text_dim)
        self.phase_dim = int(phase_dim)
        self.hidden_dim = int(hidden_dim)
        self.history = int(history)
        self.action_chunk = int(action_chunk)
        self.ensemble_heads = int(ensemble_heads)

        self.vision_proj = nn.Sequential(
            nn.LayerNorm(vision_dim), nn.Linear(vision_dim, hidden_dim)
        )
        self.rgb_proj = nn.Sequential(nn.Linear(3, hidden_dim), nn.SiLU())
        self.proprio_proj = nn.Sequential(
            nn.LayerNorm(proprio_dim + phase_dim),
            nn.Linear(proprio_dim + phase_dim, hidden_dim),
            nn.SiLU(),
            nn.Dropout(dropout),
        )
        self.text_proj = nn.Sequential(
            nn.LayerNorm(text_dim), nn.Linear(text_dim, hidden_dim), nn.SiLU()
        )
        self.task_embedding = nn.Embedding(num_tasks + 1, hidden_dim)
        self.difficulty_embedding = nn.Embedding(num_difficulties + 1, hidden_dim)
        # Scalar gates stop categorical IDs from overwhelming visual/text cues.
        # The final embedding rows are trained fallback IDs via conditioning
        # dropout in the trainer.
        self.task_scale = nn.Parameter(torch.tensor(0.5))
        self.difficulty_scale = nn.Parameter(torch.tensor(0.25))
        self.condition_norm = nn.LayerNorm(hidden_dim)
        self.spatial_attention = nn.MultiheadAttention(
            embed_dim=hidden_dim,
            num_heads=8,
            dropout=dropout,
            batch_first=True,
        )
        self.frame_fusion = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.SiLU(),
            nn.LayerNorm(hidden_dim),
            nn.Dropout(dropout),
        )
        self.temporal_gru = nn.GRU(
            input_size=hidden_dim,
            hidden_size=hidden_dim,
            num_layers=2,
            dropout=dropout,
            batch_first=True,
        )
        self.output_heads = nn.ModuleList(
            [
                nn.Sequential(
                    nn.LayerNorm(hidden_dim),
                    nn.Linear(hidden_dim, hidden_dim),
                    nn.SiLU(),
                    nn.Dropout(dropout),
                    nn.Linear(hidden_dim, action_chunk * ACTION_DIM),
                )
                for _ in range(ensemble_heads)
            ]
        )

    def forward_cached(
        self,
        visual_tokens: torch.Tensor,
        rgb_tokens: torch.Tensor,
        proprio: torch.Tensor,
        phase: torch.Tensor,
        text_features: torch.Tensor,
        task_ids: torch.Tensor,
        difficulty_ids: torch.Tensor,
    ) -> torch.Tensor:
        """Return raw actions shaped [ensemble, batch, chunk, 7]."""
        if visual_tokens.ndim != 4:
            raise ValueError(
                f"expected visual tokens [B,T,V,D], got {tuple(visual_tokens.shape)}"
            )
        batch, timesteps, visual_count, _ = visual_tokens.shape
        if timesteps != self.history:
            raise ValueError(f"expected history={self.history}, got {timesteps}")
        if rgb_tokens.shape[:3] != visual_tokens.shape[:3] or rgb_tokens.shape[-1] != 3:
            raise ValueError(
                f"expected RGB tokens [B,T,V,3] aligned with vision, got {tuple(rgb_tokens.shape)}"
            )

        visual = self.vision_proj(visual_tokens) + self.rgb_proj(rgb_tokens)
        state = self.proprio_proj(torch.cat([proprio, phase], dim=-1))
        text = self.text_proj(text_features)
        task = self.task_embedding(task_ids)
        difficulty = self.difficulty_embedding(difficulty_ids)
        condition = self.condition_norm(
            state
            + text[:, None]
            + torch.tanh(self.task_scale) * task[:, None]
            + torch.tanh(self.difficulty_scale) * difficulty[:, None]
        )

        flat_visual = visual.reshape(batch * timesteps, visual_count, self.hidden_dim)
        flat_query = condition.reshape(batch * timesteps, 1, self.hidden_dim)
        attended, _ = self.spatial_attention(
            flat_query, flat_visual, flat_visual, need_weights=False
        )
        attended = attended.reshape(batch, timesteps, self.hidden_dim)
        frames = self.frame_fusion(torch.cat([attended, condition], dim=-1))
        temporal, _ = self.temporal_gru(frames)
        final = temporal[:, -1]
        outputs = [
            head(final).reshape(batch, self.action_chunk, ACTION_DIM)
            for head in self.output_heads
        ]
        return torch.stack(outputs, dim=0)


class CompetitiveVLAModel(nn.Module):
    """Submission model containing every counted parameter, including DINOv2."""

    def __init__(self, config: dict[str, Any]) -> None:
        super().__init__()
        self.policy_config = config
        self.spatial_grid = int(config["spatial_grid"])
        vision_config = Dinov2Config(**config["vision_config"])
        self.vision = Dinov2Model(vision_config)
        self.head = CompetitivePolicyHead(
            vision_dim=int(vision_config.hidden_size),
            proprio_dim=int(config["proprio_dim"]),
            text_dim=int(config["text_dim"]),
            phase_dim=int(config.get("phase_dim", DEFAULT_PHASE_DIM)),
            num_tasks=len(config["task_to_id"]),
            num_difficulties=len(config["difficulty_to_id"]),
            hidden_dim=int(config["hidden_dim"]),
            history=int(config["history"]),
            action_chunk=int(config["action_chunk"]),
            ensemble_heads=int(config["ensemble_heads"]),
            dropout=float(config.get("dropout", 0.10)),
        )

    def encode_images(self, images: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        image_size = int(self.policy_config.get("image_size", 224))
        rgb_tokens = rgb_grid_tokens(
            images, spatial_grid=self.spatial_grid, image_size=image_size
        )
        pixels = normalize_images(images, image_size=image_size)
        target_dtype = next(self.vision.parameters()).dtype
        pixels = pixels.to(dtype=target_dtype)
        rgb_tokens = rgb_tokens.to(dtype=target_dtype)
        layer_average = max(
            1, int(self.policy_config.get("vision_layer_average", 1))
        )
        output = self.vision(
            pixel_values=pixels,
            output_hidden_states=layer_average > 1,
        )
        if layer_average > 1:
            if output.hidden_states is None:
                raise RuntimeError("DINOv2 did not return requested hidden states")
            hidden = torch.stack(output.hidden_states[-layer_average:], dim=0).mean(dim=0)
        else:
            hidden = output.last_hidden_state
        return pool_dinov2_tokens(hidden, self.spatial_grid), rgb_tokens

    def forward_cached(
        self,
        visual_tokens: torch.Tensor,
        rgb_tokens: torch.Tensor,
        proprio: torch.Tensor,
        phase: torch.Tensor,
        text_features: torch.Tensor,
        task_ids: torch.Tensor,
        difficulty_ids: torch.Tensor,
    ) -> torch.Tensor:
        return self.head.forward_cached(
            visual_tokens,
            rgb_tokens,
            proprio,
            phase,
            text_features,
            task_ids,
            difficulty_ids,
        )