"""S10 residual dynamics with position legs and velocity wheels.

Canonical order: fl, fr, hl, hr; hipx/hipy/knee per leg, then four wheels.
Actions are normalized policy outputs: leg targets = default + scale*a[:12],
wheel targets = clamp(35*a[12:], -35, 35) rad/s. Wheel phase is excluded.
This module is self-contained and depends only on torch.
"""

import torch
from torch import nn

STATE_DIM = 38
ACTION_DIM = 16
CONTEXT_DIM = 7  # four wheel sliding frictions, base COM xyz offset (metres)
POSE_DIM = 3
Z, GRAVITY, LIN_VEL, ANG_VEL = slice(0, 1), slice(1, 4), slice(4, 7), slice(7, 10)
LEG_POS, LEG_VEL, WHEEL_VEL = slice(10, 22), slice(22, 34), slice(34, 38)
FORMAT = "s10-flat-dynamics-v1"


class RunningNorm(nn.Module):
  def __init__(self, dim):
    super().__init__()
    self.register_buffer("mean", torch.zeros(dim))
    self.register_buffer("var", torch.ones(dim))
    self.register_buffer("count", torch.tensor(1e-4))

  @torch.no_grad()
  def update(self, x):
    x = x.reshape(-1, self.mean.numel())
    if not len(x):
      return
    delta = x.mean(0) - self.mean
    total = self.count + len(x)
    self.var.copy_((self.var * self.count + x.var(0, unbiased=False) * len(x)
                    + delta.square() * self.count * len(x) / total) / total)
    self.mean.add_(delta * len(x) / total)
    self.count.copy_(total)

  @property
  def scale(self):
    return self.var.clamp_min(1e-6).sqrt()

  def forward(self, x):
    return ((x - self.mean) / self.scale).clamp(-10, 10)


class S10WorldModel(nn.Module):
  """(state38, action16, previous16, context7) -> (next state38, pose delta3).

  Flat ground only. Linear and angular velocity use the base frame. Planar
  displacement uses the previous yaw frame. No target/chase or terrain scan.
  Policy history is separate from physical state and is not fed to this model.
  """
  def __init__(self, hidden=(512, 512, 512)):
    super().__init__()
    self.hidden = tuple(hidden)
    dims = (STATE_DIM + 2 * ACTION_DIM + CONTEXT_DIM, *hidden, STATE_DIM + POSE_DIM)
    layers = []
    for i, (a, b) in enumerate(zip(dims[:-1], dims[1:])):
      layers.append(nn.Linear(a, b))
      if i < len(dims) - 2:
        layers.append(nn.SiLU())
    self.net = nn.Sequential(*layers)
    self.input_norm = RunningNorm(dims[0])
    self.output_norm = RunningNorm(dims[-1])

  @staticmethod
  def inputs(state, action, previous_action, context):
    for name, x, dim in (("state", state, STATE_DIM), ("action", action, ACTION_DIM),
                         ("previous_action", previous_action, ACTION_DIM),
                         ("context", context, CONTEXT_DIM)):
      if x.ndim != 2 or x.shape != (state.shape[0], dim):
        raise ValueError(f"{name}: expected (batch, {dim}), got {tuple(x.shape)}")
    return torch.cat((state, action, previous_action, context), -1)

  @torch.no_grad()
  def update_stats(self, state, action, previous_action, context, next_state, pose_delta):
    self.input_norm.update(self.inputs(state, action, previous_action, context))
    self.output_norm.update(torch.cat((next_state - state, pose_delta), -1))

  def forward(self, state, action, previous_action, context):
    x = self.inputs(state, action, previous_action, context)
    delta = self.net(self.input_norm(x)) * self.output_norm.scale + self.output_norm.mean
    return state + delta[:, :STATE_DIM], delta[:, STATE_DIM:]

  def loss(self, state, actions, previous_actions, context, targets, poses, valid):
    """Autoregressive window loss; validity must be false at/reset after done.

    Tensor shapes: state[B,38], actions/previous[B,K,16], context[B,7],
    targets[B,K,38], poses[B,K,3], valid[B,K]. A reset invalidates the rest
    of that window, even if later transitions are individually valid.
    """
    mask = valid.bool().cumprod(1).bool()
    total = next(self.parameters()).sum() * 0.0
    for t in range(actions.shape[1]):
      # Do not execute invalid rows: NaN terminal/reset data must not leak
      # into either the loss or gradients through multiplication by zero.
      keep = mask[:, t]
      if not keep.any():
        break
      pred, pose = self(state[keep], actions[keep, t], previous_actions[keep, t], context[keep])
      error = torch.cat((pred - targets[keep, t], pose - poses[keep, t]), -1)
      total = total + (error / self.output_norm.scale).square().mean(-1).sum()
      state = state.clone()
      state[keep] = pred
    return total / mask.sum().clamp_min(1)

  def checkpoint(self):
    return {"format": FORMAT, "hidden": self.hidden, "model": self.state_dict()}

  @classmethod
  def from_checkpoint(cls, checkpoint, device="cpu"):
    if checkpoint.get("format") != FORMAT:
      raise ValueError("Expected a S10 dynamics checkpoint; other-robot checkpoints are incompatible")
    model = cls(checkpoint["hidden"]).to(device)
    model.load_state_dict(checkpoint["model"])
    return model