from __future__ import annotations

import hashlib
import json
import re
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


class BasicVLAPolicyNet(nn.Module):
    def __init__(self, text_dim: int = 128, proprio_dim: int = 32, action_dim: int = 7):
        super().__init__()
        self.image_encoder = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=5, stride=2, padding=2),
            nn.SiLU(),
            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),
            nn.SiLU(),
            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
            nn.SiLU(),
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(),
            nn.Linear(128, 256),
            nn.SiLU(),
        )
        self.proprio_encoder = nn.Sequential(nn.Linear(proprio_dim, 96), nn.SiLU())
        self.text_encoder = nn.Sequential(nn.Linear(text_dim, 96), nn.SiLU())
        self.action_head = nn.Sequential(
            nn.Linear(448, 256),
            nn.SiLU(),
            nn.LayerNorm(256),
            nn.Linear(256, 128),
            nn.SiLU(),
            nn.Linear(128, action_dim),
            nn.Tanh(),
        )

    def forward(self, images, proprio, text_features):
        if images.ndim != 4:
            raise ValueError(f"expected image batch shaped [B,H,W,3] or [B,3,H,W], got {tuple(images.shape)}")
        if images.shape[-1] == 3:
            images = images.permute(0, 3, 1, 2)
        images = images.float()
        if images.max() > 2.0:
            images = images / 255.0
        if images.shape[-2:] != (96, 96):
            images = F.interpolate(images, size=(96, 96), mode="bilinear", align_corners=False)
        image_emb = self.image_encoder(images)
        proprio_emb = self.proprio_encoder(proprio.float())
        text_emb = self.text_encoder(text_features.float())
        return self.action_head(torch.cat([image_emb, proprio_emb, text_emb], dim=-1))


def _text_vector(text: str, dim: int) -> np.ndarray:
    vec = np.zeros(dim, dtype=np.float32)
    tokens = re.findall(r"[a-z0-9_]+", text.lower())
    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 _proprio_vector(obs: dict, dim: int) -> np.ndarray:
    raw = np.asarray(obs.get("proprio", np.zeros(25, dtype=np.float32)), dtype=np.float32).reshape(-1)
    step = float(obs.get("step", 0))
    horizon = float(obs.get("horizon", 320) or 320)
    step_feature = np.asarray([step / max(horizon, 1.0)], dtype=np.float32)
    combined = np.concatenate([raw, step_feature], axis=0)
    if combined.size < dim:
        combined = np.pad(combined, (0, dim - combined.size))
    return combined[:dim].astype(np.float32)


class BasicVLAPolicy:
    def __init__(self, model_dir: str, device: str, dtype: str):
        self.model_dir = Path(model_dir)
        self.config = json.loads((self.model_dir / "vla_config.json").read_text())
        if device == "cuda" and torch.cuda.is_available():
            self.device = torch.device("cuda")
        elif device == "mps" and hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            self.device = torch.device("mps")
        else:
            self.device = torch.device("cpu")
        self.model = BasicVLAPolicyNet(
            text_dim=int(self.config.get("text_dim", 128)),
            proprio_dim=int(self.config.get("proprio_dim", 32)),
            action_dim=int(self.config.get("action_dim", 7)),
        ).to(self.device)
        checkpoint = torch.load(self.model_dir / "model.pt", map_location=self.device, weights_only=True)
        state_dict = checkpoint.get("state_dict", checkpoint)
        self.model.load_state_dict(state_dict)
        self.model.eval()

    def act(self, obs: dict) -> np.ndarray:
        image = np.asarray(obs.get("image", np.zeros((96, 96, 3), dtype=np.uint8)), dtype=np.uint8)
        if image.ndim == 2:
            image = np.repeat(image[..., None], 3, axis=-1)
        if image.shape[-1] > 3:
            image = image[..., :3]
        task = str(obs.get("task", ""))
        difficulty = str(obs.get("difficulty", ""))
        instruction = str(obs.get("instruction", ""))
        text = f"task {task} difficulty {difficulty} instruction {instruction}"
        text_features = _text_vector(text, int(self.config.get("text_dim", 128)))
        proprio = _proprio_vector(obs, int(self.config.get("proprio_dim", 32)))
        with torch.no_grad():
            action = self.model(
                torch.from_numpy(image).unsqueeze(0).to(self.device),
                torch.from_numpy(proprio).unsqueeze(0).to(self.device),
                torch.from_numpy(text_features).unsqueeze(0).to(self.device),
            )
        return np.clip(action.squeeze(0).detach().cpu().numpy(), -1.0, 1.0).astype(np.float32)


def load_policy(model_dir: str, device: str, dtype: str):
    return BasicVLAPolicy(model_dir=model_dir, device=device, dtype=dtype)