from __future__ import annotations

from typing import Any

import torch
import torch.nn as nn
from transformers import Dinov2Config, Dinov2Model

from competitive_vla_model import (
    DEFAULT_PHASE_DIM,
    CompetitivePolicyHead,
    normalize_images,
    pool_dinov2_tokens,
    rgb_grid_tokens,
)


def _make_head(config: dict[str, Any]) -> CompetitivePolicyHead:
    vision_config = config["vision_config"]
    return 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)),
    )


class HybridCompetitiveVLAModel(nn.Module):
    """One frozen DINOv2 encoder shared by a stable and specialist policy head."""

    def __init__(self, config: dict[str, Any]) -> None:
        super().__init__()
        self.policy_config = config
        self.route_configs = {
            str(name): dict(value)
            for name, value in config.get("route_configs", {}).items()
        }
        if self.route_configs:
            first_config = next(iter(self.route_configs.values()))
            self.base_config = None
            self.specialist_config = first_config
        else:
            self.base_config = dict(config["base_config"])
            self.specialist_config = dict(config["specialist_config"])
            first_config = self.specialist_config
        self.spatial_grid = int(first_config["spatial_grid"])
        self.vision = Dinov2Model(
            Dinov2Config(**first_config["vision_config"])
        )
        if self.route_configs:
            self.heads = nn.ModuleDict(
                {name: _make_head(value) for name, value in self.route_configs.items()}
            )
        else:
            assert self.base_config is not None
            self.base_head = _make_head(self.base_config)
            self.specialist_head = _make_head(self.specialist_config)

    def encode_images(self, images: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        config = self.specialist_config
        image_size = int(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(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,
        route: str,
        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:
        if self.route_configs:
            if route not in self.heads:
                raise ValueError(f"unknown policy route: {route}")
            head = self.heads[route]
        elif route == "base":
            head = self.base_head
        elif route == "specialist":
            head = self.specialist_head
        else:
            raise ValueError(f"unknown policy route: {route}")
        return head.forward_cached(
            visual_tokens,
            rgb_tokens,
            proprio,
            phase,
            text_features,
            task_ids,
            difficulty_ids,
        )