"""V4-upward helpers: sign-aware upright reward, reward clamps and sustained fall.

This module is used only by the ``s10-v4-upward`` task. V4 itself keeps its
baseline exp-upright reward and immediate 70-degree fall termination.
"""

from __future__ import annotations

import math

import torch
from mjlab.envs import mdp
from mjlab.managers import ManagerTermBase
from mjlab.managers.scene_entity_config import SceneEntityCfg

from .world_model import GRAVITY

# Raw-term caps (before weight and dt) for the recovery task.
# Chosen from v4 random-orientation samples: p99/max were approximately
# leg_posture .78/.86, standstill 5.37/7.90, joint_torques 4772/5016,
# unnecessary_wheel_lift 750/826, rolling_leg_motion 2.38/4.47.
UPWARD_REWARD_CLAMPS = {
  "leg_posture": 2.0,
  "standstill": 4.0,
  "action_rate_l2": 4.0,
  "joint_torques_l2": 5000.0,
  "dof_pos_limits": 0.5,
  "unnecessary_wheel_lift": 2.0,
  "rolling_leg_motion": 5.0,
}

# IsaacLab upward: upward max 4 (raw 4 x weight 1), bad-orientation -80,
# i.e. the bad penalty is 20x the upward maximum. Here upward max is 2
# (raw 1 x weight 2), so the same ratio is -40.
BAD_ORIENTATION_WEIGHT = -40.0
BAD_ORIENTATION_XY = 0.7
# The old max(gz, |gx|-0.7, |gy|-0.7) severity saturated for gz >= 0.3 and had
# a local maximum near 105 deg, where the |gx|-0.7 term crossed gz. The reward
# severity is now a function of the tilt angle only: 0 at 45 deg, 1 at 180 deg,
# monotonic in between. The correction term steepens the near-threshold part
# while keeping a positive slope at 180 deg and a finite derivative at 45 deg.
BAD_ORIENTATION_START_DEG = 45.0
BAD_ORIENTATION_CORRECTION = 2.5

# Per-environment recovery curriculum, IsaacLab terrain-curriculum style.
# Each env keeps the outcomes of its last HISTORY_LENGTH episodes and advances
# when at least PROMOTE_SUCCESSES of them recovered (16/20 = 80% > 75%).
# There is no global batch gate: terrain and tilt levels are per environment.
TILT_STAGES_DEG = (60.0, 135.0, 180.0)
HISTORY_LENGTH = 20
PROMOTE_SUCCESSES = 16
DEMOTE_SUCCESSES = 5


def upright_upward_value(gravity: torch.Tensor) -> torch.Tensor:
  """Normalized upward reward in [0, 1].

  ``projected_gravity_b[:, 2]`` is -1 upright, 0 sideways and +1 upside down.
  The old v4 exp-upright term is sign-blind: upside-down also scores 1. This
  linear form is monotonic through the whole range and has a constant gradient
  with respect to gravity_z, which is what the recovery task needs.
  """
  return (0.5 * (1.0 - gravity[:, 2])).clamp(0.0, 1.0)


def upright_upward(env, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
  """Environment reward term replacing v4 ``upright`` for this task."""
  gravity = env.scene[asset_cfg.name].data.projected_gravity_b
  return upright_upward_value(gravity)


def bad_orientation_value(gravity: torch.Tensor) -> torch.Tensor:
  """Monotonic tilt-based bad-orientation severity in [0, 1].

  ``projected_gravity_b[:, 2]`` is -1 upright and +1 upside down, so
  ``tilt = acos(-gz)`` goes from 0 to pi. Severity is 0 below
  ``BAD_ORIENTATION_START_DEG``, then increases monotonically to 1 at 180 deg:

    x        = clamp((tilt - 45 deg) / 135 deg, 0, 1)
    severity = x + 2.5 * x * (1 - x)^2

  Unlike the previous max(gz, |gx|-0.7, |gy|-0.7) formula, this has no
  saturation plateau for deep inversion and no local maximum near 105 deg.
  The polynomial correction is used instead of a fractional power because
  ``x ** (1/3)`` has an infinite derivative at x=0 (upright), which would
  produce non-finite gradients during imagination.
  """
  gz = gravity[:, 2].clamp(-1.0 + 1e-4, 1.0 - 1e-4)
  tilt = torch.acos(-gz)
  start = math.radians(BAD_ORIENTATION_START_DEG)
  x = ((tilt - start) / (math.pi - start)).clamp(0.0, 1.0)
  return x + BAD_ORIENTATION_CORRECTION * x * (1.0 - x).square()


def bad_orientation_mask(gravity: torch.Tensor) -> torch.Tensor:
  """Binary bad set used by the 4-second termination.

  This intentionally keeps the IsaacLab-style boundary: upside-down
  hemisphere or |gx|/|gy| beyond 0.7. The reward severity above is a smooth
  tilt-based shaping term and is not used to decide termination.
  """
  return ((gravity[:, 2] > 0.0)
          | (gravity[:, :2].abs() > BAD_ORIENTATION_XY).any(-1))


def bad_orientation_penalty(env, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
  """Continuous per-step bad-orientation penalty used by the environment."""
  gravity = env.scene[asset_cfg.name].data.projected_gravity_b
  return bad_orientation_value(gravity)


def is_upward_task(env) -> bool:
  """Identify the task from its upright reward function, not from a new flag."""
  rewards = getattr(env.cfg, "rewards", {})
  term = rewards.get("upright") if hasattr(rewards, "get") else None
  return term is not None and getattr(term, "func", None) is upright_upward


def _sample_offsets(range_dict, count: int, device):
  """Sample x/y/z/roll/pitch/yaw offsets; missing keys are zero."""
  offsets = []
  for key in ("x", "y", "z", "roll", "pitch", "yaw"):
    lo, hi = (range_dict or {}).get(key, (0.0, 0.0))
    lo, hi = float(lo), float(hi)
    if hi <= lo:
      offsets.append(torch.full((count,), lo, device=device))
    else:
      offsets.append(torch.empty(count, device=device).uniform_(lo, hi))
  return torch.stack(offsets, -1)


def reset_root_state_upward(env, env_ids, pose_range, velocity_range=None,
                            asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")):
  """Reset with a per-environment roll/pitch limit.

  ``env._upward_tilt_limits`` is a per-env tensor updated by the curriculum.
  Play/evaluation has no such tensor and falls back to the full range in
  ``pose_range`` (the task sets +/-pi there).
  """
  env_ids = mdp.resolve_env_ids(env, env_ids)
  asset = env.scene[asset_cfg.name]
  count = len(env_ids)
  pose = _sample_offsets(pose_range, count, env.device)
  limits = getattr(env, "_upward_tilt_limits", None)
  if limits is None:
    roll = pose_range.get("roll", (-math.pi, math.pi))
    limit = max(abs(float(roll[0])), abs(float(roll[1])))
    limits = torch.full((count,), limit, device=env.device)
  else:
    limits = limits[env_ids]
  pose[:, 3] = (2.0 * torch.rand(count, device=env.device) - 1.0) * limits
  pose[:, 4] = (2.0 * torch.rand(count, device=env.device) - 1.0) * limits
  default = asset.data.default_root_state[env_ids].clone()
  positions = default[:, 0:3] + pose[:, 0:3] + env.scene.env_origins[env_ids]
  delta = mdp.quat_from_euler_xyz(pose[:, 3], pose[:, 4], pose[:, 5])
  orientations = mdp.quat_mul(default[:, 3:7], delta)
  velocities = default[:, 7:13] + _sample_offsets(velocity_range, count, env.device)
  asset.write_root_link_pose_to_sim(
    torch.cat((positions, orientations), dim=-1), env_ids=env_ids)
  asset.write_root_link_velocity_to_sim(velocities, env_ids=env_ids)


def init_upward_state(env):
  """Allocate the per-env recovery ring and tilt levels once."""
  if getattr(env, "_upward_history", None) is not None:
    return
  n, device = env.num_envs, env.device
  env._upward_history = torch.zeros(
    n, HISTORY_LENGTH, dtype=torch.bool, device=device)
  env._upward_traverse_history = torch.zeros_like(env._upward_history)
  env._upward_history_cursor = torch.zeros(n, dtype=torch.long, device=device)
  env._upward_history_filled = torch.zeros(n, dtype=torch.long, device=device)
  env._upward_tilt_levels = torch.zeros(n, dtype=torch.long, device=device)
  env._upward_tilt_limit_table = torch.tensor(
    [math.radians(d) for d in TILT_STAGES_DEG], device=device)
  env._upward_tilt_limits = env._upward_tilt_limit_table[0].expand(n).clone()


def upward_curriculum(env, env_ids):
  """Per-env rolling recovery curriculum (IsaacLab terrain-curriculum style).

  Each environment stores the outcomes of its last ``HISTORY_LENGTH``
  episodes. ``recovered`` means the episode did not end in a bad-orientation
  termination and the final projected gravity is not in the bad set. An env
  advances when at least ``PROMOTE_SUCCESSES`` of the window recovered and
  drops when at most ``DEMOTE_SUCCESSES`` did. Terrain and tilt levels are
  updated independently per environment; no global batch rate is involved.
  """
  init_upward_state(env)
  env_ids = mdp.resolve_env_ids(env, env_ids)
  terrain = env.scene.terrain
  data = env.scene["robot"].data
  active = env.episode_length_buf[env_ids] > 0
  fallen = env.termination_manager.terminated[env_ids]
  gravity = data.projected_gravity_b[env_ids]
  offset = data.root_link_pos_w[env_ids] - env.scene.env_origins[env_ids]
  recovered = active & ~fallen & ~bad_orientation_mask(gravity)
  traversed = (active & ~fallen & (offset[:, 0].abs() > 5.4)
               & (offset[:, 1].abs() < 6.5))
  ids = env_ids[active]
  if len(ids):
    cursor = env._upward_history_cursor[ids]
    env._upward_history[ids, cursor] = recovered[active]
    env._upward_traverse_history[ids, cursor] = traversed[active]
    env._upward_history_cursor[ids] = (cursor + 1) % HISTORY_LENGTH
    env._upward_history_filled[ids] = (
      env._upward_history_filled[ids] + 1).clamp_max(HISTORY_LENGTH)
  ready = env._upward_history_filled[env_ids] >= HISTORY_LENGTH
  successes = env._upward_history[env_ids].sum(-1)
  promote = ready & (successes >= PROMOTE_SUCCESSES)
  demote = ready & (successes <= DEMOTE_SUCCESSES)
  delta = promote.long() - demote.long()
  old_levels = terrain.terrain_levels[env_ids].clone()
  levels = (old_levels + delta).clamp(0, terrain.terrain_origins.shape[0] - 1)
  old_tilt_levels = env._upward_tilt_levels[env_ids].clone()
  tilt_levels = (old_tilt_levels + delta).clamp(0, len(TILT_STAGES_DEG) - 1)
  # A changed environment must prove itself for a fresh 20-episode window at
  # the new level; otherwise a long success streak could chain promotions in
  # consecutive episodes. A clamped promote/demote at the min/max level is not
  # a change and must keep its window so the logged rate stays meaningful.
  changed = (levels != old_levels) | (tilt_levels != old_tilt_levels)
  if changed.any():
    changed_ids = env_ids[changed]
    env._upward_history[changed_ids] = False
    env._upward_traverse_history[changed_ids] = False
    env._upward_history_filled[changed_ids] = 0
    env._upward_history_cursor[changed_ids] = 0
  terrain.terrain_levels[env_ids] = levels
  terrain.env_origins[env_ids] = terrain.terrain_origins[
    levels, terrain.terrain_types[env_ids]]
  env._upward_tilt_levels[env_ids] = tilt_levels
  env._upward_tilt_limits[env_ids] = env._upward_tilt_limit_table[tilt_levels]
  return terrain.terrain_levels.float().mean()


def upward_curriculum_stats(env):
  """Summary for logging and for choosing the training-video tilt level."""
  init_upward_state(env)
  levels = env._upward_tilt_levels
  terrain_levels = env.scene.terrain.terrain_levels
  ready = env._upward_history_filled >= HISTORY_LENGTH
  recovery = env._upward_history.sum(-1)
  traverse = env._upward_traverse_history.sum(-1)
  ready_count = int(ready.sum())
  recovery_rate = (float(recovery[ready].float().mean().item()) / HISTORY_LENGTH
                   if ready_count else 0.0)
  traverse_rate = (float(traverse[ready].float().mean().item()) / HISTORY_LENGTH
                   if ready_count else 0.0)
  def fractions(values, count):
    return [float((values == i).float().mean().item()) for i in range(count)]
  q75 = int(torch.quantile(levels.float(), 0.75).round().item())
  return {
    "tilt_level_mean": float(levels.float().mean().item()),
    "tilt_level_max": int(levels.max().item()),
    "tilt_level_q75": q75,
    "tilt_level_frac": fractions(levels, len(TILT_STAGES_DEG)),
    "tilt_limit_mean_deg": float(
      env._upward_tilt_limits.mean().item() * 180.0 / math.pi),
    "terrain_level_frac": fractions(
      terrain_levels, int(env.scene.terrain.terrain_origins.shape[0])),
    "recovery_window_rate": recovery_rate,
    "traverse_window_rate": traverse_rate,
    "recovery_window_ready": float(ready.float().mean().item()),
  }


def upward_video_tilt_limit(env):
  """Tilt limit reached by at least 75% of environments, for the video scene."""
  q75 = upward_curriculum_stats(env)["tilt_level_q75"]
  return math.radians(TILT_STAGES_DEG[q75])


def _cap(value: torch.Tensor, max_value: float) -> torch.Tensor:
  return value.clamp(max=float(max_value))


def clamped_leg_posture(env, asset_cfg: SceneEntityCfg, max_value: float) -> torch.Tensor:
  from .terrain_posture import env_posture_cost
  return _cap(env_posture_cost(env, asset_cfg), max_value)


def clamped_standstill(env, wheel_cfg: SceneEntityCfg, max_value: float) -> torch.Tensor:
  from s10_supermimic.env.flat import standstill_cost
  return _cap(standstill_cost(env, wheel_cfg), max_value)


def clamped_action_rate(env, max_value: float) -> torch.Tensor:
  return _cap(mdp.action_rate_l2(env), max_value)


def clamped_joint_torques(env, max_value: float) -> torch.Tensor:
  return _cap(mdp.joint_torques_l2(env), max_value)


def clamped_dof_pos_limits(env, asset_cfg: SceneEntityCfg, max_value: float) -> torch.Tensor:
  return _cap(mdp.joint_pos_limits(env, asset_cfg), max_value)


def clamped_rolling_cost(env, term: str, max_value: float) -> torch.Tensor:
  from .rolling_rewards import env_rolling_cost
  return _cap(env_rolling_cost(env, term), max_value)


def apply_reward_clamps(terms: dict[str, torch.Tensor], constants: dict) -> None:
  """Apply the same raw caps to imagined rewards used by the env task."""
  for name, max_value in constants.get("reward_clamps", {}).items():
    if name in terms:
      terms[name] = terms[name].clamp(max=float(max_value))


class BadOrientationSustained(ManagerTermBase):
  """Terminate only after bad orientation persists for ``duration_s``.

  This is the recovery-task state machine: a random roll/pitch spawn is allowed
  to be bad for a few seconds, but the episode ends if the robot never rights
  itself. Functionally it replaces v4's immediate 70-degree ``fell_over``.
  """

  def __init__(self, cfg, env):
    super().__init__(env)
    self.steps = torch.zeros(env.num_envs, dtype=torch.long, device=env.device)

  def __call__(self, env, duration_s: float, asset_cfg: SceneEntityCfg) -> torch.Tensor:
    gravity = env.scene[asset_cfg.name].data.projected_gravity_b
    # Termination keeps the binary IsaacLab-style bad set; the reward severity
    # above is only a smooth shaping term.
    bad = bad_orientation_mask(gravity)
    self.steps = torch.where(bad, self.steps + 1, torch.zeros_like(self.steps))
    duration_steps = math.ceil(duration_s / env.step_dt)
    return self.steps >= duration_steps

  def reset(self, env_ids=None) -> None:
    if env_ids is None:
      self.steps.zero_()
    else:
      self.steps[env_ids] = 0