from __future__ import annotations

import math
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Literal

import torch

from mjlab.entity import Entity
from mjlab.managers.command_manager import CommandTerm, CommandTermCfg
from mjlab.utils.lab_api.math import (
  quat_from_euler_xyz,
  sample_uniform,
)

if TYPE_CHECKING:
  from mjlab.envs.manager_based_rl_env import ManagerBasedRlEnv
  from mjlab.viewer.debug_visualizer import DebugVisualizer


class LiftingCommand(CommandTerm):
  cfg: LiftingCommandCfg

  def __init__(self, cfg: LiftingCommandCfg, env: ManagerBasedRlEnv):
    super().__init__(cfg, env)

    self.object: Entity = env.scene[cfg.entity_name]
    self.target_pos = torch.zeros(self.num_envs, 3, device=self.device)
    self.episode_success = torch.zeros(self.num_envs, device=self.device)
    self._pending_forward = False

    self.metrics["object_height"] = torch.zeros(self.num_envs, device=self.device)
    self.metrics["position_error"] = torch.zeros(self.num_envs, device=self.device)
    self.metrics["at_goal"] = torch.zeros(self.num_envs, device=self.device)
    self.metrics["episode_success"] = torch.zeros(self.num_envs, device=self.device)

  @property
  def command(self) -> torch.Tensor:
    return self.target_pos

  def _update_metrics(self) -> None:
    object_pos_w = self.object.data.root_link_pos_w
    object_height = object_pos_w[:, 2]
    position_error = torch.norm(self.target_pos - object_pos_w, dim=-1)
    at_goal = (position_error < self.cfg.success_threshold).float()

    # Latch episode_success to 1 once goal is reached.
    self.episode_success = torch.maximum(self.episode_success, at_goal)

    self.metrics["object_height"] = object_height
    self.metrics["position_error"] = position_error
    self.metrics["at_goal"] = at_goal
    self.metrics["episode_success"] = self.episode_success

  def compute_success(self) -> torch.Tensor:
    position_error = self.metrics["position_error"]
    return position_error < self.cfg.success_threshold

  def _resample_command(self, env_ids: torch.Tensor) -> None:
    n = len(env_ids)

    # Reset episode success for resampled envs.
    self.episode_success[env_ids] = 0.0

    # Set target position based on difficulty mode.
    if self.cfg.difficulty == "fixed":
      target_pos = torch.tensor(
        [0.4, 0.0, 0.3], device=self.device, dtype=torch.float32
      ).expand(n, 3)
      self.target_pos[env_ids] = target_pos + self._env.scene.env_origins[env_ids]
    else:
      assert self.cfg.difficulty == "dynamic"
      r = self.cfg.target_position_range
      lower = torch.tensor([r.x[0], r.y[0], r.z[0]], device=self.device)
      upper = torch.tensor([r.x[1], r.y[1], r.z[1]], device=self.device)
      target_pos = sample_uniform(lower, upper, (n, 3), device=self.device)
      self.target_pos[env_ids] = target_pos + self._env.scene.env_origins[env_ids]

    # Reset object to new position.
    if self.cfg.object_pose_range is not None:
      r = self.cfg.object_pose_range
      lower = torch.tensor([r.x[0], r.y[0], r.z[0]], device=self.device)
      upper = torch.tensor([r.x[1], r.y[1], r.z[1]], device=self.device)
      pos = sample_uniform(lower, upper, (n, 3), device=self.device)
      pos = pos + self._env.scene.env_origins[env_ids]

      # Sample orientation (yaw only, keep upright).
      yaw = sample_uniform(r.yaw[0], r.yaw[1], (n,), device=self.device)
      quat = quat_from_euler_xyz(
        torch.zeros(n, device=self.device),  # roll
        torch.zeros(n, device=self.device),  # pitch
        yaw,
      )
      pose = torch.cat([pos, quat], dim=-1)

      velocity = torch.zeros(n, 6, device=self.device)

      self.object.write_root_link_pose_to_sim(pose, env_ids=env_ids)
      self.object.write_root_link_velocity_to_sim(velocity, env_ids=env_ids)
      self._pending_forward = True

  def reset(self, env_ids: torch.Tensor | slice | None) -> dict[str, float]:
    extras = super().reset(env_ids)
    # Reset-path resamples are followed by the env's own forward(); only
    # compute-path resamples (timer expiry) need our refresh.
    self._pending_forward = False
    return extras

  def _update_command(self, env_ids: torch.Tensor | None = None) -> None:
    # Reset-vs-step context is handled via _pending_forward, not env_ids.
    del env_ids
    # A timer-expiry resample teleports the object after the step's single
    # forward(); refresh kinematics so observations and sensors see the new
    # pose instead of the pre-teleport one.
    if self._pending_forward:
      self._pending_forward = False
      self._env.sim.forward()

  def _debug_vis_impl(self, visualizer: DebugVisualizer) -> None:
    env_indices = visualizer.get_env_indices(self.num_envs)
    if not env_indices:
      return

    for batch in env_indices:
      target_pos = self.target_pos[batch].cpu().numpy()
      visualizer.add_sphere(
        center=target_pos,
        radius=0.03,
        color=self.cfg.viz.target_color,
        label=f"target_position_{batch}",
      )


class MultiCubeLiftingCommand(CommandTerm):
  """Selects one of N cubes as the target at each reset."""

  cfg: MultiCubeLiftingCommandCfg

  def __init__(
    self,
    cfg: MultiCubeLiftingCommandCfg,
    env: ManagerBasedRlEnv,
  ):
    super().__init__(cfg, env)

    self.cubes = [env.scene[name] for name in cfg.entity_names]
    self._num_cubes = len(self.cubes)

    geom_ids = [c.indexing.geom_ids for c in self.cubes]
    max_geoms = max(g.shape[0] for g in geom_ids)
    self._padded_geom_ids = torch.full(
      (self._num_cubes, max_geoms),
      -999,
      device=self.device,
      dtype=geom_ids[0].dtype,
    )
    for i, g in enumerate(geom_ids):
      self._padded_geom_ids[i, : g.shape[0]] = g

    self.target_pos = torch.zeros(self.num_envs, 3, device=self.device)
    self.episode_success = torch.zeros(self.num_envs, device=self.device)
    self.target_selection = torch.zeros(
      self.num_envs, dtype=torch.long, device=self.device
    )

    self._env_arange = torch.arange(self.num_envs, device=self.device)
    self._cached_target_obj_pos = torch.zeros(self.num_envs, 3, device=self.device)
    self._pending_forward = False

    self.metrics["position_error"] = torch.zeros(self.num_envs, device=self.device)
    self.metrics["at_goal"] = torch.zeros(self.num_envs, device=self.device)
    self.metrics["episode_success"] = torch.zeros(self.num_envs, device=self.device)

  @property
  def command(self) -> torch.Tensor:
    return self.target_pos

  @property
  def target_geom_ids(self) -> torch.Tensor:
    """Geom IDs of the target cube per env. Shape: (B, K)."""
    return self._padded_geom_ids[self.target_selection]

  def target_object_pos(self) -> torch.Tensor:
    """Position of the target cube per env.

    Cached per step — updated in _update_metrics which runs before rewards.
    """
    return self._cached_target_obj_pos

  def _target_object_pos_from_sim(self) -> torch.Tensor:
    all_pos = torch.stack([c.data.root_link_pos_w for c in self.cubes])
    return all_pos[self.target_selection, self._env_arange]

  def _update_metrics(self) -> None:
    self._cached_target_obj_pos = self._target_object_pos_from_sim()
    obj_pos = self._cached_target_obj_pos
    error = torch.norm(self.target_pos - obj_pos, dim=-1)
    at_goal = (error < self.cfg.success_threshold).float()
    self.episode_success = torch.maximum(self.episode_success, at_goal)
    self.metrics["position_error"] = error
    self.metrics["at_goal"] = at_goal
    self.metrics["episode_success"] = self.episode_success

  def compute_success(self) -> torch.Tensor:
    return self.metrics["position_error"] < self.cfg.success_threshold

  def _resample_command(self, env_ids: torch.Tensor) -> None:
    n = len(env_ids)
    self.episode_success[env_ids] = 0.0

    self.target_selection[env_ids] = torch.randint(
      0, self._num_cubes, (n,), device=self.device
    )

    if self.cfg.difficulty == "fixed":
      target = torch.tensor([0.4, 0.0, 0.3], device=self.device).expand(n, 3)
      self.target_pos[env_ids] = target + self._env.scene.env_origins[env_ids]
    else:
      r = self.cfg.target_position_range
      lo = torch.tensor([r.x[0], r.y[0], r.z[0]], device=self.device)
      hi = torch.tensor([r.x[1], r.y[1], r.z[1]], device=self.device)
      target = sample_uniform(lo, hi, (n, 3), device=self.device)
      self.target_pos[env_ids] = target + self._env.scene.env_origins[env_ids]

    r = self.cfg.object_pose_range
    lo = torch.tensor([r.x[0], r.y[0], r.z[0]], device=self.device)
    hi = torch.tensor([r.x[1], r.y[1], r.z[1]], device=self.device)
    for cube in self.cubes:
      pos = sample_uniform(lo, hi, (n, 3), device=self.device)
      pos = pos + self._env.scene.env_origins[env_ids]
      yaw = sample_uniform(r.yaw[0], r.yaw[1], (n,), device=self.device)
      quat = quat_from_euler_xyz(
        torch.zeros(n, device=self.device),
        torch.zeros(n, device=self.device),
        yaw,
      )
      pose = torch.cat([pos, quat], dim=-1)
      velocity = torch.zeros(n, 6, device=self.device)
      cube.write_root_link_pose_to_sim(pose, env_ids=env_ids)
      cube.write_root_link_velocity_to_sim(velocity, env_ids=env_ids)
    self._pending_forward = True

  def reset(self, env_ids: torch.Tensor | slice | None) -> dict[str, float]:
    extras = super().reset(env_ids)
    # Reset-path resamples are followed by the env's own forward(); only
    # compute-path resamples (timer expiry) need our refresh.
    self._pending_forward = False
    return extras

  def _update_command(self, env_ids: torch.Tensor | None = None) -> None:
    # Reset-vs-step context is handled via _pending_forward, not env_ids.
    del env_ids
    # A timer-expiry resample teleports the cubes after the step's single
    # forward(); refresh kinematics and the reward-path cache (updated in
    # _update_metrics before the resample) so both see post-teleport state.
    if self._pending_forward:
      self._pending_forward = False
      self._env.sim.forward()
      self._cached_target_obj_pos = self._target_object_pos_from_sim()

  def _debug_vis_impl(self, visualizer: DebugVisualizer) -> None:
    env_indices = visualizer.get_env_indices(self.num_envs)
    if not env_indices:
      return
    for batch in env_indices:
      target_pos = self.target_pos[batch].cpu().numpy()
      visualizer.add_sphere(
        center=target_pos,
        radius=0.03,
        color=(1.0, 0.5, 0.0, 0.3),
        label=f"target_position_{batch}",
      )
      cube_pos = self.target_object_pos()[batch].cpu().numpy()
      marker = cube_pos.copy()
      marker[2] += 0.04
      visualizer.add_sphere(
        center=marker,
        radius=0.01,
        color=(1.0, 0.0, 0.0, 1.0),
        label=f"target_cube_marker_{batch}",
      )


@dataclass(kw_only=True)
class MultiCubeLiftingCommandCfg(CommandTermCfg):
  entity_names: tuple[str, ...] = ()
  success_threshold: float = 0.05
  difficulty: Literal["fixed", "dynamic"] = "fixed"

  @dataclass
  class TargetPositionRangeCfg:
    x: tuple[float, float] = (0.3, 0.5)
    y: tuple[float, float] = (-0.2, 0.2)
    z: tuple[float, float] = (0.2, 0.4)

  target_position_range: TargetPositionRangeCfg = field(
    default_factory=TargetPositionRangeCfg
  )

  @dataclass
  class ObjectPoseRangeCfg:
    x: tuple[float, float] = (0.25, 0.40)
    y: tuple[float, float] = (-0.15, 0.15)
    z: tuple[float, float] = (0.02, 0.05)
    yaw: tuple[float, float] = (-math.pi, math.pi)

  object_pose_range: ObjectPoseRangeCfg = field(default_factory=ObjectPoseRangeCfg)

  def build(self, env: ManagerBasedRlEnv) -> MultiCubeLiftingCommand:
    return MultiCubeLiftingCommand(self, env)


@dataclass(kw_only=True)
class LiftingCommandCfg(CommandTermCfg):
  entity_name: str
  success_threshold: float = 0.05
  difficulty: Literal["fixed", "dynamic"] = "fixed"

  @dataclass
  class TargetPositionRangeCfg:
    """Configuration for target position sampling in dynamic mode."""

    x: tuple[float, float] = (0.3, 0.5)
    y: tuple[float, float] = (-0.2, 0.2)
    z: tuple[float, float] = (0.2, 0.4)

  # Only used in dynamic mode.
  target_position_range: TargetPositionRangeCfg = field(
    default_factory=TargetPositionRangeCfg
  )

  @dataclass
  class ObjectPoseRangeCfg:
    """Configuration for object pose sampling when resampling commands."""

    x: tuple[float, float] = (0.3, 0.35)
    y: tuple[float, float] = (-0.1, 0.1)
    z: tuple[float, float] = (0.02, 0.05)
    yaw: tuple[float, float] = (-math.pi, math.pi)

  object_pose_range: ObjectPoseRangeCfg | None = field(
    default_factory=ObjectPoseRangeCfg
  )

  @dataclass
  class VizCfg:
    target_color: tuple[float, float, float, float] = (1.0, 0.5, 0.0, 0.3)

  viz: VizCfg = field(default_factory=VizCfg)

  def build(self, env: ManagerBasedRlEnv) -> LiftingCommand:
    return LiftingCommand(self, env)
