"""
Learned 40-DoF (arms-only) action model for the RX1  (v2).

Goal-conditioned action-chunking policy trained by behaviour cloning
(scripts/train_act.py) on expert demonstrations (scripts/collect_demos.py).

    input  obs   = [ proprio(80) , pick_pos0(3) , place_pos(3) ]   = 86
    output chunk = (chunk_len, 15)  →  [ 7 right-arm , 7 left-arm , 1 grasp ]

Because the policy is conditioned on (pick_pos0, place_pos), a single trained
model performs pick, place-in-box, stack and sort — the high-level planner just
supplies different place targets. Torso/head are held at rest and the 18 finger
joints follow the grasp bit, so the model only learns the 14 arm joints + grasp.

A plain MLP is enough (and runs in microseconds on CPU, so it is genuinely
real-time); normalisation stats are stored in the checkpoint so inference
matches training. Swap in a Transformer/diffusion head later without changing
the I/O contract.
"""
from __future__ import annotations

import logging

import numpy as np
import torch
import torch.nn as nn

log = logging.getLogger(__name__)

OBS_DIM = 86
ACT_DIM = 15            # 14 arm joints + 1 grasp
DEFAULT_CHUNK = 8


class ActionChunkPolicy(nn.Module):
    def __init__(self, obs_dim: int = OBS_DIM, act_dim: int = ACT_DIM,
                 chunk_len: int = DEFAULT_CHUNK, hidden: int = 512):
        super().__init__()
        self.obs_dim = obs_dim
        self.act_dim = act_dim
        self.chunk_len = chunk_len
        self.net = nn.Sequential(
            nn.Linear(obs_dim, hidden), nn.SiLU(),
            nn.Linear(hidden, hidden), nn.SiLU(),
            nn.Linear(hidden, hidden), nn.SiLU(),
            nn.Linear(hidden, chunk_len * act_dim),
        )

    def forward(self, obs: torch.Tensor) -> torch.Tensor:
        out = self.net(obs)
        return out.view(-1, self.chunk_len, self.act_dim)


class Normalizer:
    """Stores obs/act mean+std; normalises for the net, denormalises actions."""

    def __init__(self, obs_mean, obs_std, act_mean, act_std):
        self.obs_mean = np.asarray(obs_mean, np.float32)
        self.obs_std  = np.asarray(obs_std,  np.float32)
        self.act_mean = np.asarray(act_mean, np.float32)
        self.act_std  = np.asarray(act_std,  np.float32)

    @staticmethod
    def fit(obs: np.ndarray, act: np.ndarray) -> "Normalizer":
        eps = 1e-5
        return Normalizer(obs.mean(0), obs.std(0) + eps,
                          act.mean(0), act.std(0) + eps)

    def norm_obs(self, o):  return (o - self.obs_mean) / self.obs_std
    def norm_act(self, a):  return (a - self.act_mean) / self.act_std
    def denorm_act(self, a): return a * self.act_std + self.act_mean

    def state(self) -> dict:
        return dict(obs_mean=self.obs_mean, obs_std=self.obs_std,
                    act_mean=self.act_mean, act_std=self.act_std)


def save_checkpoint(path, model: ActionChunkPolicy, norm: Normalizer):
    torch.save({
        "state_dict": model.state_dict(),
        "cfg": dict(obs_dim=model.obs_dim, act_dim=model.act_dim,
                    chunk_len=model.chunk_len),
        "norm": norm.state(),
    }, path)


def load_policy(path, device="cpu"):
    """Load a trained policy + normaliser for inference. Returns (model, norm)."""
    # Degrade gracefully to CPU when a CUDA device is requested but unavailable
    # (e.g. configs/rx1_models_v2.yaml hardcodes device: cuda, but the host is
    # CPU-only). Without this torch.load(map_location="cuda") raises and the
    # whole brain server crashes at startup instead of running on CPU.
    if "cuda" in str(device) and not torch.cuda.is_available():
        log.warning("[load_policy] device=%r requested but CUDA unavailable — using CPU", device)
        device = "cpu"
    ckpt = torch.load(path, map_location=device, weights_only=False)
    cfg = ckpt["cfg"]
    model = ActionChunkPolicy(cfg["obs_dim"], cfg["act_dim"], cfg["chunk_len"])
    model.load_state_dict(ckpt["state_dict"])
    model.to(device).eval()
    n = ckpt["norm"]
    norm = Normalizer(n["obs_mean"], n["obs_std"], n["act_mean"], n["act_std"])
    return model, norm