"""S10 actor history and actor/critic networks, independent of the quadruped learner.

Two Actor observation schemas coexist:

* v4 / v4-upward / new v7: 53 values/frame, including measured wheel speed;
* legacy v7: 49 values/frame, no measured wheel speed (the last action still
  contains the four wheel commands).

The learner, replay and deployment exporter select the schema explicitly. Old
checkpoints are recognized from their first-layer width.
"""

import torch
from torch import nn
from .world_model import RunningNorm, ANG_VEL, GRAVITY, LEG_POS, LEG_VEL, WHEEL_VEL

LEGACY_TERMS = (3, 3, 12, 12, 4, 16, 3)
TERMS = (3, 3, 12, 12, 16, 3)
LEGACY_FRAME_DIM = sum(LEGACY_TERMS)
FRAME_DIM = sum(TERMS)
CRITIC_DIM = 38 + 3 + 7 + 16
POLICY_FORMAT = "s10-flat-policy-v1"


def frame(state, last_env_action, command, bias, noise=False, wheel_vel=True):
  """Exactly the term order/scaling in env.flat (noise before scaling)."""
  values = [state[:, ANG_VEL], state[:, GRAVITY], state[:, LEG_POS] + bias,
            state[:, LEG_VEL]]
  widths = [0.2, 0.05, 0.01, 1.5]
  if wheel_vel:
    values.append(state[:, WHEEL_VEL])
    widths.append(1.5)
  values.extend((last_env_action, command))
  widths.extend((0., 0.))
  if noise:
    values = [v + (torch.rand_like(v) * 2 - 1) * w if w else v
              for v, w in zip(values, widths)]
  # Leg velocity is always index 3; wheel velocity, when present, is index 4.
  velocity_indices = (3, 4) if wheel_vel else (3,)
  return torch.cat([v * (0.05 if i in velocity_indices else 1.)
                    for i, v in enumerate(values)], -1)


def push_history(history, latest, length, terms):
  """Term-wise histories, oldest to newest, matching mjlab flattening."""
  result, hidx, fidx = [], 0, 0
  for dim in terms:
    result.append(torch.cat((history[:, hidx + dim:hidx + dim * length],
                             latest[:, fidx:fidx + dim]), -1))
    hidx += dim * length
    fidx += dim
  return torch.cat(result, -1)


def critic_input(state, command, context, previous):
  return torch.cat((state, command, context, previous), -1)


class Network(nn.Module):
  def __init__(self, inputs, outputs, hidden):
    super().__init__()
    self.norm = RunningNorm(inputs)
    layers = []
    for n in hidden:
      layers.extend((nn.Linear(inputs, n), nn.ELU()))
      inputs = n
    layers.append(nn.Linear(inputs, outputs))
    self.net = nn.Sequential(*layers)

  def forward(self, x):
    return self.net(self.norm(x))


class Actor(Network):
  def __init__(self, history=5, hidden=(256, 128, 64), actor_wheel_vel=True):
    self.actor_wheel_vel = bool(actor_wheel_vel)
    self.terms = LEGACY_TERMS if self.actor_wheel_vel else TERMS
    self.frame_dim = sum(self.terms)
    super().__init__(self.frame_dim * history, 16, hidden)
    self.history, self.hidden = history, tuple(hidden)
    nn.init.uniform_(self.net[-1].weight, -0.01, 0.01)
    nn.init.zeros_(self.net[-1].bias)

  def forward(self, obs):
    # Same bounded action distribution during collection and imagination.
    return super().forward(obs).tanh()


def checkpoint_wheel_vel(checkpoint):
  """Infer legacy observation schema from weights; validate explicit metadata."""
  flag = checkpoint.get("actor_wheel_vel")
  if "actor" in checkpoint:
    dim = checkpoint["actor"]["net.0.weight"].shape[1]
    history = checkpoint["history"]
    if dim not in (FRAME_DIM * history, LEGACY_FRAME_DIM * history):
      raise ValueError("Unsupported actor observation dimension")
    inferred = dim == LEGACY_FRAME_DIM * history
    if flag is not None and bool(flag) != inferred:
      raise ValueError("Actor observation metadata disagrees with weights")
    return inferred
  return bool(flag) if flag is not None else False


def load_actor(checkpoint, device="cpu"):
  if checkpoint.get("format") not in (POLICY_FORMAT, "s10-terrain-policy-v1"):
    raise ValueError("Expected a S10 policy checkpoint, not other-robot or dynamics-only weights")
  actor = Actor(checkpoint["history"], checkpoint["actor_hidden"],
                checkpoint_wheel_vel(checkpoint)).to(device)
  actor.load_state_dict(checkpoint["actor"])
  return actor.eval()


def evaluation_env_cfg(checkpoint, level=0, kind=None):
  """Select the matching task; terrain remains privileged at deployment."""
  from s10_supermimic.tasks import get_task
  actor_wheel_vel = checkpoint_wheel_vel(checkpoint)
  if checkpoint.get("format") == "s10-terrain-policy-v1":
    checkpoint_task = checkpoint.get("config", {}).get("task")
    task = get_task(checkpoint_task) if checkpoint_task else None
    if task is not None:
      if task.terrain_revision != checkpoint["config"].get("terrain_revision", 1):
        raise ValueError("Checkpoint task revision differs from its config")
      if task.actor.wheel_vel != actor_wheel_vel:
        raise ValueError("Checkpoint task Actor schema differs from its weights")
      cfg = task.build_env_cfg(play=True, history=checkpoint["history"],
                               seed=checkpoint["config"]["seed"])
    else:
      from s10_supermimic.env.terrain import s10_velocity_terrain_env_cfg
      cfg = s10_velocity_terrain_env_cfg(
        play=True, history=checkpoint["history"],
        seed=checkpoint["config"]["seed"],
        revision=checkpoint["config"].get("terrain_revision", 1),
        actor_wheel_vel=actor_wheel_vel)
    # A single staircase can overlap the robot's pre-reset origin at compile
    # time. Reserve contacts for that transient pose as well as normal rolling.
    cfg.sim.nconmax = max(cfg.sim.nconmax or 0, 512)
    cfg.sim.njmax = max(cfg.sim.njmax or 0, 2048)
    if level not in range(6):
      raise ValueError("Terrain level must be 0..5")
    gen = cfg.scene.terrain.terrain_generator
    # Evaluation fixes difficulty; one row, same seeded generator.
    gen.num_rows = 1
    gen.difficulty_range = (level/5,level/5)
    if kind is not None:
      gen.sub_terrains = {kind: gen.sub_terrains[kind]}
    return cfg
  from s10_supermimic.env.flat import s10_velocity_flat_env_cfg
  return s10_velocity_flat_env_cfg(play=True, history=checkpoint["history"],
                                   actor_wheel_vel=actor_wheel_vel)