"""World-model future preview for the play scripts.

The actor checkpoint already contains the world model (``world_model``) and the
training config. We use it directly; no separate ONNX export is required.

The preview has two independent outputs:

* a kinematic ghost trail: every predicted state is mapped back to MJCF qpos
  and transparently added to the offscreen renderer;
* a prediction-error ledger: predictions are stored with their target step and
  compared with the real state when that step arrives (xy/z/yaw/gravity/leg q).

The ghost is only a visualization. It never writes into ``env.sim.data`` or the
physics model.
"""
from __future__ import annotations

from dataclasses import dataclass
import math

import mujoco
import numpy as np
import torch

from .imagine import frame, push_history
from .policy import Actor
from .terrain import advance_pose
from .terrain_world_model import TerrainWorldModel
from .world_model import GRAVITY, LEG_POS, S10WorldModel

DEFAULT_HORIZON = 32
DEFAULT_STRIDE = 4
GHOST_NEAR = (1.0, 0.86, 0.08, 0.20)
GHOST_FAR = (1.0, 0.24, 0.02, 0.055)


@dataclass
class GhostPrediction:
  horizon: int
  state: np.ndarray       # 38
  pose: np.ndarray        # x, y, yaw
  qpos: np.ndarray        # MJCF qpos


def load_world_model(checkpoint, device="cpu"):
  payload = checkpoint.get("world_model")
  if not isinstance(payload, dict):
    raise ValueError("Checkpoint has no world_model payload")
  fmt = payload.get("format")
  if fmt == "s10-terrain-dynamics-v1":
    model = TerrainWorldModel.from_checkpoint(payload, device)
  elif fmt == "s10-flat-dynamics-v1":
    model = S10WorldModel.from_checkpoint(payload, device)
  else:
    raise ValueError(f"Unsupported world-model format: {fmt}")
  return model.to(device).eval()


def _quat_mul(a, b):
  w1, x1, y1, z1 = a
  w2, x2, y2, z2 = b
  return np.array([
    w1*w2 - x1*x2 - y1*y2 - z1*z2,
    w1*x2 + x1*w2 + y1*z2 - z1*y2,
    w1*y2 - x1*z2 + y1*w2 + z1*x2,
    w1*z2 + x1*y2 - y1*x2 + z1*w2,
  ], dtype=np.float64)


def _quat_from_up(up_b, yaw):
  """Rotation mapping body-frame up to world +z, followed by world yaw."""
  up = np.asarray(up_b, dtype=np.float64)
  up /= max(np.linalg.norm(up), 1e-9)
  target = np.array([0., 0., 1.])
  dot = float(np.clip(up @ target, -1., 1.))
  if dot < -0.999999:
    align = np.array([0., 1., 0., 0.])
  else:
    axis = np.cross(up, target)
    align = np.concatenate(([1. + dot], axis))
    align /= max(np.linalg.norm(align), 1e-9)
  yaw_q = np.array([math.cos(yaw / 2), 0., 0., math.sin(yaw / 2)])
  return _quat_mul(yaw_q, align)


class WorldModelPreview:
  def __init__(self, checkpoint, env, bridge, field, actor, horizon=DEFAULT_HORIZON,
               device="cpu", env_action_order=None):
    self.env = env
    self.bridge = bridge
    self.field = field
    self.actor = actor
    self.model = load_world_model(checkpoint, device)
    self.horizon = int(horizon)
    self.device = device
    self.env_action_order = list(env_action_order or checkpoint["env_action_order"])
    self.default_q = bridge.robot.data.default_joint_pos[0, bridge.legs].detach().cpu().numpy()
    model = env.sim.mj_model
    free = np.flatnonzero(model.jnt_type == mujoco.mjtJoint.mjJNT_FREE)
    if len(free) != 1:
      raise ValueError(f"Expected one free root joint, found {len(free)}")
    self.root_qadr = int(model.jnt_qposadr[int(free[0])])
    # ``bridge.legs``/``bridge.wheels`` are entity joint indices, not MuJoCo
    # joint ids (the free joint is not in the entity list). Use the entity's
    # own qpos-address table so the ghost writes the intended joints.
    qadr = bridge.robot.indexing.joint_q_adr.cpu().numpy()
    self.leg_qadr = [int(qadr[int(j)]) for j in bridge.legs]
    self.wheel_qadr = [int(qadr[int(j)]) for j in bridge.wheels]
    self.qpos_template = env.sim.data.qpos[0].detach().cpu().numpy().copy()
    # root_link_pos_w may be env-local while qpos is world. Store the constant
    # offset so predicted future positions land in the renderer's world frame.
    self.world_offset = (self.qpos_template[self.root_qadr:self.root_qadr + 2]
                         - bridge.pose()[0, :2].detach().cpu().numpy())
    self.tracker = PredictionTracker(self.horizon)

  def _ghost(self, state, pose, horizon):
    s = state[0].detach().cpu().numpy().astype(np.float64)
    p = pose[0].detach().cpu().numpy().astype(np.float64)
    qpos = self.qpos_template.copy()
    z = float(s[0])
    if self.field is not None:
      z += float(self.field.sample(pose[:, :2])[0].detach().cpu())
    quat = _quat_from_up(-s[1:4], float(p[2]))
    qpos[self.root_qadr:self.root_qadr + 3] = (
      p[0] + self.world_offset[0], p[1] + self.world_offset[1], z)
    qpos[self.root_qadr + 3:self.root_qadr + 7] = quat
    qpos[self.leg_qadr] = s[LEG_POS] + self.default_q
    qpos[self.wheel_qadr] = self.qpos_template[self.wheel_qadr]
    return GhostPrediction(horizon, s, p, qpos)

  @torch.no_grad()
  def predict(self, history, previous, command, context, bias, state, pose):
    """Roll the actor + world model forward from the current real state."""
    ghosts = []
    history = history.clone()
    previous = previous.clone()
    state = state.clone()
    pose = pose.clone()
    scan = None if self.field is None else self.field.scan(state, pose)
    for i in range(self.horizon):
      action = self.actor(history)
      model_context = context if self.field is None else torch.cat((context, scan), -1)
      if self.field is not None:
        nxt, delta, _, _ = self.model.forward_all(state, action, previous, model_context)
      else:
        nxt, delta = self.model(state, action, previous, model_context)
      pose = advance_pose(pose, delta)
      if self.field is not None:
        scan = self.field.scan(nxt, pose)
      latest = frame(nxt, action[:, self.env_action_order], command, bias,
                     noise=False, wheel_vel=self.actor.actor_wheel_vel)
      history = push_history(history, latest, self.actor.history, self.actor.terms)
      previous = action
      state = nxt
      ghosts.append(self._ghost(state, pose, i + 1))
    return ghosts

  def update_errors(self, step, state, pose):
    self.tracker.update(step, state[0].detach().cpu().numpy(),
                        pose[0].detach().cpu().numpy())

  def push_errors(self, step, ghosts):
    self.tracker.push(step, ghosts)

  def summary(self):
    return self.tracker.summary()

  def clear(self):
    self.tracker.clear()

  def render(self, ghosts, stride=DEFAULT_STRIDE, rgba=None):
    """Actual frame plus transparent ghost geoms, using the env's renderer."""
    return render_with_ghosts(self.env, ghosts, stride=stride, rgba=rgba)


class PredictionTracker:
  """Compare predictions with the state that actually arrives later."""
  def __init__(self, horizon, keep=200):
    self.horizon = int(horizon)
    self.keep = int(keep)
    self.pending = []
    self.errors = {h: [] for h in (1, 4, 8, 16, 32) if h <= self.horizon}

  def clear(self):
    self.pending.clear()
    for values in self.errors.values():
      values.clear()

  @staticmethod
  def _error(pred_state, pred_pose, state, pose):
    dxy = float(np.linalg.norm(pred_pose[:2] - pose[:2]))
    dz = float(abs(pred_state[0] - state[0]))
    dyaw = float(abs(math.atan2(math.sin(pred_pose[2] - pose[2]),
                                math.cos(pred_pose[2] - pose[2]))))
    pg = pred_state[GRAVITY]
    ag = state[GRAVITY]
    pg = pg / max(np.linalg.norm(pg), 1e-9)
    ag = ag / max(np.linalg.norm(ag), 1e-9)
    dg = float(math.degrees(math.acos(float(np.clip(pg @ ag, -1., 1.)))))
    dq = float(np.sqrt(np.mean((pred_state[LEG_POS] - state[LEG_POS]) ** 2)))
    return dxy, dz, dyaw, dg, dq

  def push(self, step, ghosts):
    for ghost in ghosts:
      if ghost.horizon in self.errors:
        self.pending.append((step + ghost.horizon, ghost.horizon,
                             ghost.state, ghost.pose))

  def update(self, step, state, pose):
    keep = []
    for target, horizon, pred_state, pred_pose in self.pending:
      if target == step:
        values = self.errors[horizon]
        values.append(self._error(pred_state, pred_pose, state, pose))
        if len(values) > self.keep:
          del values[:-self.keep]
      elif target > step:
        keep.append((target, horizon, pred_state, pred_pose))
    self.pending = keep

  def summary(self):
    parts = []
    for horizon in sorted(self.errors):
      values = self.errors[horizon]
      if not values:
        continue
      mean = np.mean(values, axis=0)
      parts.append(f"h{horizon} xy {mean[0]:.3f}m yaw {math.degrees(mean[2]):.1f}deg "
                   f"leg {mean[4]:.3f}")
    return " | ".join(parts) if parts else "collecting..."


def annotate_rgb(rgb, lines, color=(255, 185, 90)):
  """Small bottom-left text overlay for MP4 frames (Pillow optional)."""
  if not lines:
    return rgb
  try:
    from PIL import Image, ImageDraw, ImageFont
  except ImportError:
    return rgb
  image = Image.fromarray(rgb)
  draw = ImageDraw.Draw(image)
  font = ImageFont.load_default()
  y = image.height - 8 - 14 * len(lines)
  for line in lines:
    draw.rectangle((6, y - 2, 6 + 8 * len(line), y + 13), fill=(8, 10, 14))
    draw.text((10, y), line, fill=color, font=font)
    y += 14
  return np.asarray(image)


def render_with_ghosts(env, ghosts, stride=DEFAULT_STRIDE, rgba=None):
  """Add ghost robot geoms to the existing offscreen scene and render once."""
  rgb = env.render()
  renderer = env._offline_renderer
  if renderer is None or not ghosts:
    return rgb
  import mujoco

  model, data = renderer._model, renderer._data
  scene = renderer.renderer.scene
  old_qpos, old_qvel = data.qpos.copy(), data.qvel.copy()
  old_group = renderer._opt.geomgroup.copy()
  # S10 visual meshes are group 2; group 0 is terrain and group 1 collision
  # geometry. Adding only group 2 gives robot ghosts without duplicating terrain.
  visual = np.flatnonzero(model.geom_group == 2)
  old_geom_rgba = model.geom_rgba[visual].copy()
  mat_ids = np.unique(model.geom_matid[visual])
  mat_ids = mat_ids[mat_ids >= 0]
  old_mat_rgba = {int(i): model.mat_rgba[int(i)].copy() for i in mat_ids}
  near = np.asarray(rgba if rgba is not None else GHOST_NEAR, dtype=np.float32)
  far = np.asarray(GHOST_FAR, dtype=np.float32)
  try:
    renderer._opt.geomgroup[:] = 0
    renderer._opt.geomgroup[2] = 1
    for ghost in ghosts[max(0, int(stride) - 1)::max(1, int(stride))]:
      # Near future is yellow and stronger; the far tail fades to red.
      t = ghost.horizon / max(1, ghosts[-1].horizon)
      tint = near * (1. - t) + far * t
      for mat_id in mat_ids:
        model.mat_rgba[int(mat_id)] = tint
      model.geom_rgba[visual] = tint
      data.qpos[:] = ghost.qpos
      data.qvel[:] = 0.
      mujoco.mj_forward(model, data)
      mujoco.mjv_addGeoms(model, data, renderer._opt, renderer._pert,
                          renderer._catmask.value, scene)
  finally:
    renderer._opt.geomgroup[:] = old_group
    model.geom_rgba[visual] = old_geom_rgba
    for mat_id, value in old_mat_rgba.items():
      model.mat_rgba[mat_id] = value
    data.qpos[:] = old_qpos
    data.qvel[:] = old_qvel
    mujoco.mj_forward(model, data)
  return renderer.render()