"""Teach the robot to balance by showing it, instead of letting it guess.

The classical controller in the bench already balances. Reinforcement learning
from scratch has to stumble onto that behaviour by accident before it can be
rewarded for it, and almost every random policy falls in half a second, so
nearly every episode looks equally bad and there is very little to climb.

So we drive the robot with the controller, record what the robot could SEE
(its own noisy estimator, never ground truth) alongside what the controller
did, and fit a policy to copy it. That is behaviour cloning. The result is not
better than the controller - it cannot be - but it starts PPO somewhere
sensible, and unlike the controller it can then be improved.

Coverage is the whole game. A teacher that only ever demonstrates a robot
standing perfectly still teaches a policy that has never seen a lean and does
not know what to do about one, so every episode starts from a different tilt
and gets shoved at random intervals.

  python teach.py --demo            record demonstrations
  python teach.py --clone           record, then fit a policy to them
"""

import argparse
import math
from pathlib import Path

import numpy as np

import demo
from env import RobotEnv

ROOT = Path(__file__).parent.parent


def controller(env, odo, kp=30.0, kd=2.0, station=1.95, r_wheel=0.0625):
    """The bench's auto-balance, as a function.

    Sign is measured, not assumed: with both wheels driven the same way, a
    forward lean needs kp NEGATIVE to catch it. Positive gains fall over.
    Station keeping is a cascade - speed and odometry pick a small lean target
    and the balance loop chases that - because feeding speed straight into
    torque destabilises a wheeled pendulum, which has to lean the wrong way
    first in order to stop.
    """
    up = env.est_up
    pitch = math.atan2(-up[0], max(1e-6, up[2]))
    rate = float(env.sensors.noisy("gyro", env.data.sensordata, env._raw)[1])
    wv = 0.5 * sum(float(env.data.qvel[env.model.jnt_dofadr[
        __import__("mujoco").mj_name2id(env.model,
                   __import__("mujoco").mjtObj.mjOBJ_JOINT, n)]])
        for n in ("l_wheel_j", "r_wheel_j"))
    speed = wv * r_wheel
    odo += speed * env.control_dt
    lean_ref = max(-0.08, min(0.08, station * (0.5 * speed + 0.1 * odo)))
    u = max(-1.0, min(1.0, -kp * (pitch - lean_ref) + kd * rate))
    return u, odo


def record(spec, task, episodes, seconds, seed=0, shove=True):
    """Drive with the controller and keep every (what it saw, what it did)."""
    obs_all, act_all = [], []
    rng = np.random.default_rng(seed)
    survived = []
    for ep in range(episodes):
        env = RobotEnv(spec, seed=seed + ep, task=task, randomise=False)
        o = env.reset()
        odo = 0.0
        n = int(seconds / env.control_dt)
        # a shove every second or so, so the policy sees recoveries too
        shoves = (rng.integers(20, 60, size=8) if shove else [])
        nxt = list(np.cumsum(shoves)) if shove else []
        for k in range(n):
            u, odo = controller(env, odo)
            a = np.full(env.act_dim, 0.0)
            a[-2:] = u                     # both wheels the same way
            obs_all.append(list(map(float, o)))
            act_all.append(list(map(float, a)))
            o, _, fell, trunc, _ = env.step(a)
            if nxt and k == nxt[0]:
                nxt.pop(0)
                env.data.qvel[0] += float(rng.uniform(-0.35, 0.35))
            if fell or trunc:
                break
        survived.append(env.t)
    return np.array(obs_all), np.array(act_all), np.array(survived)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--spec", default=str(ROOT / "robots/wheeled_biped.json"))
    ap.add_argument("--task", default="balance")
    ap.add_argument("--name", default="teacher_balance")
    ap.add_argument("--episodes", type=int, default=40)
    ap.add_argument("--seconds", type=float, default=8.0)
    ap.add_argument("--epochs", type=int, default=300)
    ap.add_argument("--out", default=str(ROOT / "runs/johnny6__balance"))
    ap.add_argument("--clone", action="store_true")
    args = ap.parse_args()

    obs, act, up = record(args.spec, args.task, args.episodes, args.seconds)
    full = args.seconds
    print(f"teacher: {len(obs)} steps from {args.episodes} episodes")
    print(f"  upright {up.mean():.2f}s of {full:.0f}s, "
          f"{int((up >= full - 0.05).sum())}/{len(up)} full episodes")
    if up.mean() < full * 0.5:
        print("  the teacher itself is falling - fix that before cloning")
    p, n = demo.save(args.name, args.task, obs, act,
                     meta={"teacher": "classical balance"})
    print(f"  wrote {p} ({n} steps total)")

    if args.clone:
        pol = demo.clone(args.name, args.task, args.spec, args.out,
                         epochs=args.epochs)
        got = demo.evaluate(pol, args.spec, args.task, episodes=10)
        env = RobotEnv(args.spec, seed=0, task=args.task)
        ep = env.max_steps * env.control_dt
        print(f"\ncloned policy: upright {got.mean():.2f}s of {ep:.0f}s, "
              f"{int((got >= ep - 1e-6).sum())}/{len(got)} full episodes")


if __name__ == "__main__":
    main()