from __future__ import annotations
import hashlib
import json
import math
import re
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

VARIANTS = {
    "small_cnn": {"width": 32, "depth": 3, "chunk": 1, "dropout": 0.05},
    "base_cnn": {"width": 64, "depth": 4, "chunk": 1, "dropout": 0.10},
    "chunk_cnn": {"width": 48, "depth": 4, "chunk": 8, "dropout": 0.10},
    "ensemble_cnn": {"width": 64, "depth": 5, "chunk": 8, "dropout": 0.15},
}

class VisionBackbone(nn.Module):
    def __init__(self, width, depth):
        super().__init__()
        layers = [nn.Conv2d(3, width, 5, 2, 2), nn.GroupNorm(max(1, width // 8), width), nn.SiLU()]
        channels = width
        for index in range(depth - 1):
            next_channels = min(width * (2 ** min(index + 1, 2)), 256)
            layers += [nn.Conv2d(channels, next_channels, 3, 2, 1),
                       nn.GroupNorm(max(1, next_channels // 8), next_channels), nn.SiLU()]
            channels = next_channels
        layers += [nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(), nn.Linear(channels, 384), nn.SiLU()]
        self.net = nn.Sequential(*layers)
    def forward(self, images):
        return self.net(images)

class PolicyNet(nn.Module):
    def __init__(self, config):
        super().__init__()
        variant = config["variant"]
        cfg = VARIANTS[variant]
        self.chunk = int(cfg["chunk"])
        self.vision = VisionBackbone(int(cfg["width"]), int(cfg["depth"]))
        self.proprio = nn.Sequential(nn.Linear(int(config["proprio_dim"]), 128), nn.LayerNorm(128), nn.SiLU())
        self.text = nn.Sequential(nn.Linear(int(config["text_dim"]), 128), nn.LayerNorm(128), nn.SiLU())
        self.task = nn.Embedding(int(config["task_count"]), 32)
        self.difficulty = nn.Embedding(int(config["difficulty_count"]), 16)
        self.head = nn.Sequential(
            nn.Linear(384 + 128 + 128 + 32 + 16, 512), nn.LayerNorm(512), nn.SiLU(),
            nn.Dropout(float(cfg["dropout"])), nn.Linear(512, 256), nn.SiLU(),
            nn.Linear(256, self.chunk * 7))
    def forward(self, images, proprio, text, task_id, difficulty_id):
        if images.ndim != 4:
            raise ValueError("image batch must have four dimensions")
        if images.shape[-1] == 3:
            images = images.permute(0, 3, 1, 2)
        images = images.float()
        if images.max() > 2:
            images = images / 255.0
        images = F.interpolate(images, size=(96, 96), mode="bilinear", align_corners=False)
        fused = torch.cat([
            self.vision(images), self.proprio(proprio.float()), self.text(text.float()),
            self.task(task_id.clamp(0, self.task.num_embeddings - 1)),
            self.difficulty(difficulty_id.clamp(0, self.difficulty.num_embeddings - 1)),
        ], dim=-1)
        return torch.tanh(self.head(fused)).reshape(-1, self.chunk, 7)

def _text_vector(text, dim):
    result = np.zeros(dim, dtype=np.float32)
    for token in re.findall(r"[a-z0-9_]+", str(text).lower()):
        value = int.from_bytes(hashlib.blake2b(token.encode(), digest_size=8).digest(), "little")
        result[value % dim] += 1.0 if value & 1 else -1.0
    norm = float(np.linalg.norm(result))
    return result / norm if norm else result

def _proprio(obs, config):
    value = np.asarray(obs.get("proprio", np.zeros(25, dtype=np.float32)), dtype=np.float32).reshape(-1)
    if value.size == 21:
        value = np.pad(value, (0, 4))
    if value.size != 25:
        raise ValueError(f"proprio must have 25 values, got {value.size}")
    step = float(obs.get("step", 0))
    horizon = float(obs.get("horizon", 320) or 320)
    value = np.concatenate([value, np.asarray([step / max(horizon, 1.0)], dtype=np.float32)])
    if value.size < int(config["proprio_dim"]):
        value = np.pad(value, (0, int(config["proprio_dim"]) - value.size))
    value = value[: int(config["proprio_dim"])]
    mean = np.asarray(config["proprio_mean"], dtype=np.float32)
    std = np.asarray(config["proprio_std"], dtype=np.float32)
    return ((value - mean) / np.maximum(std, 1e-4)).astype(np.float32)

class Policy:
    def __init__(self, model, config, device):
        self.model, self.config, self.device = model, config, device
        self.chunk = int(config["chunk"])
        self.last_step = None
        self.chunks = {}
        self.last_action = np.zeros(7, dtype=np.float32)
        self.task_to_id = config["task_to_id"]
        self.difficulty_to_id = config["difficulty_to_id"]
    @torch.inference_mode()
    def act(self, obs):
        image = np.asarray(obs.get("image", np.zeros((96, 96, 3), dtype=np.uint8)))
        if image.ndim == 2:
            image = np.repeat(image[..., None], 3, axis=-1)
        if image.shape[-1] == 4:
            image = image[..., :3]
        if image.shape[-1] != 3:
            raise ValueError("image must have 3 or 4 channels")
        step = int(obs.get("step", 0))
        if self.last_step is None or step == 0 or step != self.last_step + 1:
            self.chunks = {}
        if step == self.last_step:
            return self.last_action.copy()
        task = str(obs.get("task", ""))
        difficulty = str(obs.get("difficulty", "") or "")
        text = f"task {task} difficulty {difficulty} instruction {obs.get('instruction', '')}"
        text_feature = _text_vector(text, int(self.config["text_dim"]))
        proprio = _proprio(obs, self.config)
        task_id = self.task_to_id.get(task, len(self.task_to_id))
        difficulty_id = self.difficulty_to_id.get(difficulty, len(self.difficulty_to_id))
        raw = self.model(
            torch.from_numpy(image).unsqueeze(0).to(self.device),
            torch.from_numpy(proprio).unsqueeze(0).to(self.device),
            torch.from_numpy(text_feature).unsqueeze(0).to(self.device),
            torch.tensor([task_id], device=self.device),
            torch.tensor([difficulty_id], device=self.device),
        )[0].float().cpu().numpy()
        self.chunks[step] = raw
        candidates, weights = [], []
        for start, chunk in list(self.chunks.items()):
            offset = step - start
            if 0 <= offset < self.chunk:
                candidates.append(chunk[offset])
                weights.append(math.exp(-0.55 * offset))
            else:
                del self.chunks[start]
        action = np.average(np.asarray(candidates), axis=0, weights=np.asarray(weights))
        action = np.clip(action, -1.0, 1.0).astype(np.float32)
        if action[6] > 0.12:
            action[6] = 1.0
        elif action[6] < -0.12:
            action[6] = -1.0
        self.last_step, self.last_action = step, action.copy()
        return action

def load_policy(model_dir: str, device: str, dtype: str):
    root = Path(model_dir)
    config = json.loads((root / "vla_config.json").read_text())
    selected = torch.device(device if str(device).startswith("cuda") and torch.cuda.is_available() else "cpu")
    model = PolicyNet(config).to(selected)
    checkpoint = torch.load(root / "model.pt", map_location=selected, weights_only=True)
    state = checkpoint["state_dict"] if "state_dict" in checkpoint else checkpoint
    model.load_state_dict(state, strict=True)
    model.eval()
    return Policy(model, config, selected)