"""A classical balance controller, and the case for not learning from scratch.

A wheeled inverted pendulum is one of the best-understood systems in control.
Its balance law is four numbers:

    wheel torque = -(k1*pitch + k2*pitch_rate + k3*speed + k4*distance)

The first two catch the fall. The second two stop it running away doing so,
which is the part people forget: a controller that only looks at pitch balances
beautifully while accelerating across the room forever.

There is no reason to make reinforcement learning rediscover this. What RL is
actually good for here is everything the linear law cannot express: what to do
with the legs, how to recover from a shove that exceeds the linear region, how
to change height while staying up, how to handle a step. So the sensible
architecture is this controller as an inner loop with RL learning a correction
on top, rather than RL from noise.

  python baseline.py --tune          # search for gains
  python baseline.py --episodes 20   # score the tuned controller
"""

import argparse
import math

import numpy as np

from env import RobotEnv

# Tuned by the search below on the Johnny 6 spec.
GAINS = dict(k_pitch=9.0, k_rate=1.2, k_speed=2.2, k_pos=0.9, k_posture=0.0)


def pitch_from_up(up):
    """Lean angle about the pitch axis, from the IMU's up vector.

    Sign convention measured against the simulator, not assumed: leaning
    forward in +x drives up_x NEGATIVE, so the minus sign here is the
    difference between catching the fall and driving into it.
    """
    return math.atan2(-up[0], max(1e-6, up[2]))


class BalanceController:
    def __init__(self, gains=None, spec=None):
        self.g = dict(GAINS if gains is None else gains)

    def reset(self):
        pass

    def act(self, env):
        """Return an action in the env's normalised [-1, 1] space."""
        up = env.est_up if getattr(env, "est_on", False) else \
            np.asarray(env.data.sensordata[env._gt["gt_up"]])
        pitch = pitch_from_up(up)
        rate = float(env.sensors.noisy("gyro", env.data.sensordata, env._raw)[1])
        speed = float(env.vel_local[0])
        pos = float(env.base_pos[0])

        u = -(self.g["k_pitch"] * pitch + self.g["k_rate"] * rate
              + self.g["k_speed"] * speed + self.g["k_pos"] * pos)
        u = float(np.clip(u, -1.0, 1.0))

        a = np.zeros(env.act_dim)
        a[4] = a[5] = u          # both wheels, same torque
        return a

    @staticmethod
    def describe():
        return ("tau = -(k_pitch*pitch + k_rate*pitch_rate "
                "+ k_speed*speed + k_pos*distance)")


def run(env, ctrl, seconds=None, seed=0):
    obs = env.reset()
    ctrl.reset()
    n = env.max_steps if seconds is None else int(seconds / env.control_dt)
    total = 0.0
    for _ in range(n):
        a = ctrl.act(env)
        obs, r, fell, trunc, info = env.step(a)
        total += r
        if fell:
            break
    return env.t, total


def score(gains, spec, episodes=8, randomise=False):
    ups, drifts = [], []
    for ep in range(episodes):
        env = RobotEnv(spec, seed=5000 + ep, task="stand", randomise=randomise)
        t, _ = run(env, BalanceController(gains), seed=5000 + ep)
        ups.append(t)
        drifts.append(abs(float(env.base_pos[0])))
    # reward standing for the full episode, then penalise wandering off
    return float(np.mean(ups)) - 0.3 * float(np.mean(drifts)), np.mean(ups), np.mean(drifts)


def tune(spec, iters=250, seed=0):
    rng = np.random.default_rng(seed)
    best = dict(GAINS)
    best_s, up, dr = score(best, spec)
    print(f"start  {best_s:6.2f}  upright {up:5.2f}s  drift {dr:5.2f}m  {best}")
    scale = 0.6
    for i in range(iters):
        cand = {k: max(0.0, v * math.exp(rng.normal(0, scale)))
                for k, v in best.items() if k != "k_posture"}
        cand["k_posture"] = 0.0
        s, up, dr = score(cand, spec)
        if s > best_s:
            best, best_s = cand, s
            print(f"  {i:3d}  {s:6.2f}  upright {up:5.2f}s  drift {dr:5.2f}m  "
                  + " ".join(f"{k}={v:.2f}" for k, v in cand.items() if v))
        scale = max(0.12, scale * 0.99)
    return best


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--spec", default="../robots/wheeled_biped.json")
    ap.add_argument("--tune", action="store_true")
    ap.add_argument("--iters", type=int, default=250)
    ap.add_argument("--episodes", type=int, default=20)
    ap.add_argument("--dr", action="store_true", help="score with randomisation on")
    args = ap.parse_args()

    if args.tune:
        g = tune(args.spec, args.iters)
        print("\nGAINS =", {k: round(v, 3) for k, v in g.items()})
        return

    env = RobotEnv(args.spec, seed=1, task="stand", randomise=args.dr)
    full = env.max_steps * env.control_dt
    ups, drifts = [], []
    for ep in range(args.episodes):
        env = RobotEnv(args.spec, seed=9000 + ep, task="stand", randomise=args.dr)
        t, _ = run(env, BalanceController())
        ups.append(t)
        drifts.append(abs(float(env.base_pos[0])))
    ups = np.array(ups)
    print(f"classical controller, {args.episodes} episodes of {full:.0f}s, "
          f"randomisation {'on' if args.dr else 'off'}")
    print(f"  upright        {ups.mean():5.2f}s mean, {ups.min():.2f}s worst")
    print(f"  full episodes  {int((ups >= full - 1e-6).sum())}/{args.episodes}")
    print(f"  drift          {np.mean(drifts):.2f} m mean")


if __name__ == "__main__":
    main()