"""The incentive library.

A task is what the robot is being paid to do. It owns four things: any scenery
it needs, the goal information the robot is allowed to observe, the reward, and
what counts as failing.

Nothing here tells the robot how to move. Every task is expressed as a payoff.

Note on goal observations: a task like `goto` has to tell the robot where the
goal is, and that information has to come from somewhere real. Treat those
channels as a localisation source on the robot (UWB beacon, vision, or wheel
odometry with a known origin). They are listed in the UI as goal sensing so the
distinction from "free knowledge of world state" stays visible. Base tilt,
velocity and absolute pose remain reward-only.
"""

import math
import numpy as np

DEG = math.pi / 180.0


def _quat_yaw(q):
    w, x, y, z = q
    return math.atan2(2 * (w * z + x * y), 1 - 2 * (y * y + z * z))


# Every goal channel any incentive can publish, in a fixed order. A policy
# trained under one incentive can carry into the next because the observation
# keeps this exact shape throughout: an incentive that does not use a channel
# leaves it at zero. Without this, "learn to stand, then learn to drive" would
# mean throwing the weights away between stages.
GOAL_CHANNELS = [
    "cmd_vx", "cmd_wz",           # drive
    "spin_wz", "spin_drift",      # spin
    "goal_bx", "goal_by", "goal_range",   # goto
    "step_dist",                  # climb
    "cmd_height",                 # bob
]
GOAL_INDEX = {c: i for i, c in enumerate(GOAL_CHANNELS)}


class Task:
    name = "base"
    label = "Base"
    description = ""
    actuators = None            # None = the policy drives every actuator
    pose = None                 # override the spec's standing pose for this task
    init_noise = None           # override how randomly it is dropped
    goal_sensing = []          # human-readable list of extra observation channels
    channels = []              # which GOAL_CHANNELS this incentive publishes
    episode_s = 12.0
    props = []                 # static scenery: dicts consumed by MJCF + the bench

    def __init__(self, spec, cfg=None, layout="union"):
        self.spec = spec
        self.cfg = cfg or {}
        self.curriculum = 1.0
        self.layout = layout

    @property
    def obs_extra(self):
        return len(GOAL_CHANNELS) if self.layout == "union" else len(self.channels)

    def channel_values(self, env):
        """Goal information this incentive is allowed to publish, by name."""
        return {}

    def obs(self, env):
        v = self.channel_values(env)
        names = GOAL_CHANNELS if self.layout == "union" else self.channels
        return np.array([float(v.get(n, 0.0)) for n in names])

    # -- lifecycle ---------------------------------------------------------
    def reset(self, env, rng):
        pass

    def on_step(self, env, rng):
        pass

    def goal_marker(self, env):
        """World position of the current goal, or None. Drawn by the bench."""
        return None

    def reward(self, env, action):
        raise NotImplementedError

    def terminated(self, env):
        return False

    def success(self, env):
        return False

    # -- shared reward pieces ---------------------------------------------
    @staticmethod
    def _posture(env, action, w):
        q_leg = env.data.qpos[env.leg_joint_ids]
        # Charge only the actuators this policy COMMANDS. Summing all six
        # meant the 2-action balance policy was billed for leg holding
        # torques it cannot influence - noise in the objective, not signal.
        tau = env.data.actuator_force[env.act_idx]
        return {
            "torque": -w.get("torque", 0.0) * float(np.sum(tau ** 2)),
            "action_rate": -w.get("action_rate", 0.0) * float(np.sum((action - env.prev_action) ** 2)),
            "joint_nominal": -w.get("joint_nominal", 0.0) * float(np.sum((q_leg - env.q_nominal) ** 2)),
        }

    @staticmethod
    def _stay_up(env, w):
        return {
            "alive": w.get("alive", 1.0),
            "upright": w.get("upright", 2.0) * env.up_z,
            "height": -w.get("height", 8.0) * (env.height - env.h_target) ** 2,
        }


class Stand(Task):
    name = "stand"
    label = "Stand still"
    description = ("Hold station on the wheels without falling. Staying put is "
                   "rewarded by position, not by stillness, because balancing "
                   "IS constant small wheel motion.")
    episode_s = 10.0
    # No velocity penalty. A wheeled inverted pendulum holds itself up by
    # driving the wheels back and forth under its own centre of mass, so
    # rewarding low speed rewards exactly the thing that makes it fall over.
    # "anchor" pays for ending up in the same place instead, which permits any
    # amount of corrective shuffling.
    # action_rate is 0.005 everywhere, down from 0.02: at full-swing wheel
    # commands the old weight cost ~27x the torque term, so the policy was
    # charged far more for moving the command fast than for torque itself -
    # backwards for a body whose reflexes ARE fast command changes. Smoothness
    # still pays, it just no longer outbids strength.
    W = {"alive": 1.0, "upright": 2.5, "height": 8.0, "anchor": 1.5,
         "torque": 0.0015, "action_rate": 0.005, "joint_nominal": 0.3}

    def reset(self, env, rng):
        self.origin = env.base_pos[:2].copy()

    def reward(self, env, action):
        drift = float(np.linalg.norm(env.base_pos[:2] - self.origin))
        terms = self._stay_up(env, self.W)
        terms["anchor"] = self.W["anchor"] * math.exp(-(drift ** 2) / 0.25)
        terms.update(self._posture(env, action, self.W))
        return sum(terms.values()), terms


class Balance(Stand):
    """Stand, but with the legs bolted at the standing pose.

    The whole problem becomes "drive the wheels forward and back to stay under
    the body" - two numbers instead of six. There is no reason to make the
    robot search leg poses before it can even stay upright, and a 2-D problem
    trains in minutes rather than half an hour.
    """
    name = "balance"
    label = "Balance on the wheels"
    description = ("Legs locked straight, so the only thing to learn is rolling "
                   "the wheels back and forth to stay under the body.")
    episode_s = 10.0
    actuators = ("l_wheel_a", "r_wheel_a")     # everything else is held
    # Straight, not the crouched standing pose. Bending the legs swings the
    # centre of mass 14.5mm BEHIND the wheel contact, so the robot topples
    # backwards continuously and the wheels burn their authority fighting a
    # constant bias instead of learning to balance. Straight legs put the mass
    # exactly over the contact patch, and cost no knee torque either.
    pose = {"l_hip": 0.0, "l_knee": 0.0, "r_hip": 0.0, "r_knee": 0.0}
    # Dropped square and level, with only a small pitch to catch. The wheels
    # cannot correct roll at all, so randomising roll injects outcomes the
    # policy can neither see coming nor do anything about - noise in the reward
    # that swamps the signal it is trying to learn from. Same for dropping it
    # from a height: bouncing is not what we are teaching here.
    init_noise = {"joint": 0.0, "base_z": 0.0, "base_roll": 0.0,
                  "base_pitch": 0.03, "base_yaw": 0.0, "base_vel": 0.02}


class Bob(Task):
    """Balance while deliberately changing height.

    Squatting and rising moves the centre of mass vertically, which changes the
    pendulum under the controller while it is running. Learning this after
    plain balancing, and before shoves, is the natural order: it is the first
    time the legs and the wheels have to cooperate.
    """
    name = "bob"
    label = "Bob up and down"
    description = ("Follow a height command while staying upright. The legs "
                   "change the pendulum the wheels are balancing, so the two "
                   "have to work together.")
    goal_sensing = ["commanded body height"]
    channels = ["cmd_height"]
    episode_s = 12.0
    pose = {"l_hip": 0.0, "l_knee": 0.0, "r_hip": 0.0, "r_knee": 0.0}
    W = {"alive": 1.0, "upright": 2.5, "height": 0.0, "track_h": 6.0,
         "anchor": 1.0, "torque": 0.0015, "action_rate": 0.005, "joint_nominal": 0.0}

    def __init__(self, spec, cfg=None, layout="union"):
        super().__init__(spec, cfg, layout)
        self.low = self.cfg.get("low", 0.36)
        self.high = self.cfg.get("high", 0.47)
        self.resample_s = self.cfg.get("resample_s", 3.0)
        # Sectioned, not a continuous wander: a fixed ladder of heights the
        # robot must settle ON and hold. A drifting target can be tracked by
        # lagging behind it, which looks like following and teaches nothing.
        # Discrete gates force it to arrive, stop, and stay there - and each
        # one is a checkpoint you can see it pass or miss.
        self.n_gates = int(self.cfg.get("gates", 5))
        self.gates = np.linspace(self.low, self.high, self.n_gates)

    def reset(self, env, rng):
        self.rng = rng
        self.origin = env.base_pos[:2].copy()
        self.mid = env.h_target
        self.target = float(self.gates[int(np.argmin(abs(self.gates - self.mid)))])
        self._pick(env)
        self._next = self.resample_s

    def _pick(self, env):
        # curriculum: open the ladder outwards from the height it already
        # stands at, so early on it only steps to a neighbouring gate
        # full span at full curriculum, measured from where it stands. Half the
        # range was not enough: the robot's natural stand height sits near the
        # top of the ladder, so the two lowest gates could never be asked for.
        span = (self.high - self.low) * self.curriculum
        near = self.mid + np.array([-span, span])
        allowed = [g for g in self.gates if near[0] - 1e-9 <= g <= near[1] + 1e-9]
        if not allowed:
            allowed = [float(self.gates[int(np.argmin(abs(self.gates - self.mid)))])]
        # never re-pick the gate it is already sitting on: the point is to move
        options = [g for g in allowed if abs(g - getattr(self, "target", -1)) > 1e-6]
        self.target = float(self.rng.choice(options or allowed))

    def gate_index(self):
        return int(np.argmin(abs(self.gates - self.target)))

    def on_step(self, env, rng):
        if env.t >= self._next:
            self._pick(env)
            self._next = env.t + self.resample_s

    def channel_values(self, env):
        return {"cmd_height": (self.target - self.mid) * 10.0}

    def reward(self, env, action):
        err = env.height - self.target
        drift = float(np.linalg.norm(env.base_pos[:2] - self.origin))
        terms = {"alive": self.W["alive"], "upright": self.W["upright"] * env.up_z}
        terms["track_h"] = self.W["track_h"] * math.exp(-(err ** 2) / 0.0025)
        terms["anchor"] = self.W["anchor"] * math.exp(-(drift ** 2) / 0.25)
        terms.update(self._posture(env, action, self.W))
        return sum(terms.values()), terms


class Drive(Task):
    name = "drive"
    label = "Drive to command"
    description = ("Follow a forward-speed and turn-rate command that changes "
                   "every few seconds. This is the general-purpose locomotion "
                   "incentive.")
    goal_sensing = ["commanded forward speed", "commanded turn rate"]
    channels = ["cmd_vx", "cmd_wz"]
    episode_s = 12.0
    W = {"alive": 1.0, "upright": 2.0, "height": 8.0, "track_vx": 1.5,
         "track_wz": 0.6, "torque": 0.0015, "action_rate": 0.005,
         "joint_nominal": 0.3, "yaw_drift": 0.2}

    def __init__(self, spec, cfg=None, layout="union"):
        super().__init__(spec, cfg, layout)
        c = self.cfg
        self.vx_range = c.get("vx_range", [-0.6, 0.6])
        self.wz_range = c.get("wz_range", [-1.0, 1.0])
        self.resample_s = c.get("resample_s", 4.0)

    def reset(self, env, rng):
        self.rng = rng
        self._resample()
        self._next = self.resample_s

    def _resample(self):
        s = self.curriculum
        self.cmd = np.array([
            self.rng.uniform(*self.vx_range) * s,
            self.rng.uniform(*self.wz_range) * s,
        ])

    def on_step(self, env, rng):
        if env.t >= self._next:
            self._resample()
            self._next = env.t + self.resample_s

    def channel_values(self, env):
        return {"cmd_vx": self.cmd[0], "cmd_wz": self.cmd[1]}

    def reward(self, env, action):
        v, g = env.vel_local, env.gyro_true
        terms = self._stay_up(env, self.W)
        terms["track_vx"] = self.W["track_vx"] * math.exp(-((v[0] - self.cmd[0]) ** 2) / 0.25)
        terms["track_wz"] = self.W["track_wz"] * math.exp(-((g[2] - self.cmd[1]) ** 2) / 0.25)
        terms["yaw_drift"] = -self.W["yaw_drift"] * float(v[1] ** 2)
        terms.update(self._posture(env, action, self.W))
        return sum(terms.values()), terms


class Spin(Task):
    name = "spin"
    label = "Spin on the spot"
    description = ("Turn at a commanded rate while holding position. Rewards "
                   "yaw rate and punishes translation, so it has to counter-"
                   "rotate the wheels rather than drive in a circle.")
    goal_sensing = ["commanded turn rate", "distance from start"]
    channels = ["spin_wz", "spin_drift"]
    episode_s = 10.0
    W = {"alive": 1.0, "upright": 2.0, "height": 8.0, "track_wz": 2.5,
         "anchor": 1.2, "torque": 0.0015, "action_rate": 0.005, "joint_nominal": 0.3}

    def reset(self, env, rng):
        self.rng = rng
        rng_lo, rng_hi = self.cfg.get("wz_range", [-3.0, 3.0])
        mag = rng.uniform(1.0, abs(rng_hi)) * self.curriculum
        self.target_wz = mag * (1 if rng.random() < 0.5 else -1)
        self.origin = env.data.sensordata[env._gt["gt_pos"]][:2].copy()

    def channel_values(self, env):
        d = env.base_pos[:2] - self.origin
        return {"spin_wz": self.target_wz, "spin_drift": float(np.linalg.norm(d))}

    def reward(self, env, action):
        g = env.gyro_true
        drift = float(np.linalg.norm(env.base_pos[:2] - self.origin))
        terms = self._stay_up(env, self.W)
        terms["track_wz"] = self.W["track_wz"] * math.exp(-((g[2] - self.target_wz) ** 2) / 1.0)
        terms["anchor"] = self.W["anchor"] * math.exp(-(drift ** 2) / 0.09)
        terms.update(self._posture(env, action, self.W))
        return sum(terms.values()), terms


class GoTo(Task):
    name = "goto"
    label = "Go to a point"
    description = ("Drive to marker A, then marker B, then A again. Rewards "
                   "closing the distance, with a bonus for arriving. Nothing "
                   "says how to get there.")
    goal_sensing = ["bearing to goal (sin, cos)", "range to goal"]
    channels = ["goal_bx", "goal_by", "goal_range"]
    episode_s = 20.0
    W = {"alive": 1.0, "upright": 1.5, "height": 8.0, "progress": 12.0,
         "arrive": 40.0, "heading": 0.4, "torque": 0.0015,
         "action_rate": 0.005, "joint_nominal": 0.3}

    def __init__(self, spec, cfg=None, layout="union"):
        super().__init__(spec, cfg, layout)
        self.radius = self.cfg.get("radius", 1.6)
        self.tolerance = self.cfg.get("tolerance", 0.20)

    def reset(self, env, rng):
        self.rng = rng
        r = self.radius * (0.35 + 0.65 * self.curriculum)
        ang = rng.uniform(0, 2 * math.pi)
        self.waypoints = [
            np.array([r * math.cos(ang), r * math.sin(ang)]),
            np.array([r * math.cos(ang + math.pi), r * math.sin(ang + math.pi)]),
        ]
        self.wp_i = 0
        self.reached = 0
        self.prev_dist = self._dist(env)

    def _goal(self):
        return self.waypoints[self.wp_i % len(self.waypoints)]

    def _dist(self, env):
        return float(np.linalg.norm(self._goal() - env.base_pos[:2]))

    def goal_marker(self, env):
        g = self._goal()
        return [float(g[0]), float(g[1]), 0.0]

    def channel_values(self, env):
        d = self._goal() - env.base_pos[:2]
        yaw = _quat_yaw(env.base_quat)
        # Bearing in the robot's own frame: this is what a beacon or a camera
        # would give it, not a world-frame coordinate.
        bx = math.cos(-yaw) * d[0] - math.sin(-yaw) * d[1]
        by = math.sin(-yaw) * d[0] + math.cos(-yaw) * d[1]
        rng_ = math.hypot(bx, by)
        return {"goal_bx": bx / (rng_ + 1e-6), "goal_by": by / (rng_ + 1e-6),
                "goal_range": min(rng_, 5.0) * 0.3}

    def reward(self, env, action):
        d = self._dist(env)
        terms = self._stay_up(env, self.W)
        terms["progress"] = self.W["progress"] * (self.prev_dist - d)
        self.prev_dist = d

        terms["heading"] = self.W["heading"] * self.channel_values(env)["goal_bx"]

        terms["arrive"] = 0.0
        if d < self.tolerance:
            terms["arrive"] = self.W["arrive"]
            self.wp_i += 1
            self.reached += 1
            self.prev_dist = self._dist(env)

        terms.update(self._posture(env, action, self.W))
        return sum(terms.values()), terms

    def success(self, env):
        return self.reached >= 2


class Climb(Task):
    name = "climb"
    label = "Climb a step"
    description = ("Get over a raised threshold and keep going. The step "
                   "height rises with the curriculum, so it learns the easy "
                   "lip before the hard one.")
    goal_sensing = ["distance to step edge"]
    channels = ["step_dist"]
    episode_s = 15.0
    W = {"alive": 1.0, "upright": 1.5, "height": 4.0, "progress": 10.0,
         "cross": 50.0, "torque": 0.0015, "action_rate": 0.005, "joint_nominal": 0.2}

    def __init__(self, spec, cfg=None, layout="union"):
        super().__init__(spec, cfg, layout)
        self.step_x = self.cfg.get("step_x", 1.1)
        self.max_h = self.cfg.get("max_height", 0.045)
        self.depth = self.cfg.get("depth", 1.2)

    @property
    def height(self):
        return self.max_h * (0.25 + 0.75 * self.curriculum)

    def build_props(self):
        h = self.height
        return [{
            "name": "step", "kind": "obstacle", "type": "box",
            "size": [self.depth / 2, 1.6, h / 2],
            "pos": [self.step_x + self.depth / 2, 0.0, h / 2],
            "rgba": [0.42, 0.45, 0.52, 1.0],
        }]

    def reset(self, env, rng):
        self.crossed = False
        self.prev_x = float(env.base_pos[0])

    def channel_values(self, env):
        return {"step_dist": min(3.0, self.step_x - float(env.base_pos[0])) * 0.5}

    def reward(self, env, action):
        x = float(env.base_pos[0])
        terms = self._stay_up(env, self.W)
        terms["progress"] = self.W["progress"] * (x - self.prev_x)
        self.prev_x = x
        terms["cross"] = 0.0
        if not self.crossed and x > self.step_x + self.depth * 0.5:
            self.crossed = True
            terms["cross"] = self.W["cross"]
        terms.update(self._posture(env, action, self.W))
        return sum(terms.values()), terms

    def success(self, env):
        return self.crossed


class Recover(Task):
    name = "recover"
    label = "Survive shoves"
    description = ("Stand while being pushed hard at random intervals. Trains "
                   "the recovery reflex specifically, rather than hoping it "
                   "falls out of the other incentives.")
    episode_s = 12.0
    W = {"alive": 1.5, "upright": 2.5, "height": 8.0, "settle": 1.5,
         "torque": 0.0015, "action_rate": 0.005, "joint_nominal": 0.3}

    def __init__(self, spec, cfg=None, layout="union"):
        super().__init__(spec, cfg, layout)
        self.every_s = self.cfg.get("every_s", 2.0)
        self.impulse = self.cfg.get("impulse", [0.8, 2.6])

    def reset(self, env, rng):
        self.rng = rng
        self._next = self.every_s

    def on_step(self, env, rng):
        if env.t >= self._next:
            lo, hi = self.impulse
            mag = rng.uniform(lo, lo + (hi - lo) * self.curriculum)
            ang = rng.uniform(0, 2 * math.pi)
            env.data.qvel[0] += mag * math.cos(ang)
            env.data.qvel[1] += mag * math.sin(ang)
            self._next = env.t + self.every_s

    def reward(self, env, action):
        v = env.vel_local
        terms = self._stay_up(env, self.W)
        terms["settle"] = self.W["settle"] * math.exp(-(v[0] ** 2 + v[1] ** 2) / 0.2)
        terms.update(self._posture(env, action, self.W))
        return sum(terms.values()), terms


REGISTRY = {t.name: t for t in (Balance, Bob, Stand, Drive, Spin, GoTo, Climb, Recover)}

# The old spec wrote its task type this way; keep it loading.
ALIASES = {"balance_and_drive": "drive"}


def make_task(spec, name=None, cfg=None, layout="union"):
    name = ALIASES.get(name or "", name) or "drive"
    if name not in REGISTRY:
        raise KeyError(f"unknown incentive {name!r}; have {sorted(REGISTRY)}")
    task_cfg = dict((spec.get("tasks", {}) or {}).get(name, {}))
    task_cfg.update(cfg or {})
    return REGISTRY[name](spec, task_cfg, layout=layout)


def catalogue():
    """Machine-readable list for the bench's incentive picker."""
    return [{
        "name": t.name,
        "label": t.label,
        "description": t.description.strip(),
        "episode_s": t.episode_s,
        "goal_sensing": list(t.goal_sensing),
        "channels": list(t.channels),
        "reward_terms": sorted(t.W.keys()),
        "has_scenery": hasattr(t, "build_props"),
        "actuators": list(t.actuators) if t.actuators else None,
    } for t in REGISTRY.values()]