"""Practise: improve a watched rule by falling over, the way a toddler does.

Watching the teacher gives the general shape of the rule in one closed-form
fit, and that alone balances the robot on flat ground through a 16 degree lean.
What it cannot give is any idea what to do beyond where the teacher ever went,
because those states are simply not in the demonstrations.

So we practise. The rule is a small linear map - tens of numbers, not thousands
- which means we can improve it by directly trying variations on it rather than
by backpropagating a gradient. That is Augmented Random Search: jitter every
number a little in both directions, keep the directions that helped, step that
way. It is almost embarrassingly simple, it parallelises trivially, and on a
search space this small it is far faster than PPO.

Crucially the practice conditions are the ones that BREAK the watched rule -
big tilts and hard shoves - because there is no point drilling something it can
already do.

  python practise.py --iters 12
"""

import argparse
import json
import math
import time
from pathlib import Path

import numpy as np

from env import RobotEnv

ROOT = Path(__file__).parent.parent
SPEC = str(ROOT / "robots/wheeled_biped.json")


def watched_rule(demo_name="teacher_balance"):
    """One least-squares fit of what the teacher did to what it could see.

    Deliberately linear. A 6594-parameter network fitted these same
    demonstrations 17x more closely and then fell over in 2 seconds: it had
    enough capacity to memorise the teacher's comfortable states instead of
    being forced to extract the rule behind them.
    """
    d = json.load(open(ROOT / "demos" / f"{demo_name}.json"))
    O, A = np.array(d["obs"]), np.array(d["act"])
    X = np.hstack([O, np.ones((len(O), 1))])
    W, *_ = np.linalg.lstsq(X, A, rcond=None)
    return W


def episode(W, seed, secs, tilt=0.0, shove=0.0, task="balance"):
    env = RobotEnv(SPEC, seed=seed, task=task, randomise=False)
    o = env.reset()
    if tilt:
        env.data.qpos[3:7] = [math.cos(tilt / 2), 0, math.sin(tilt / 2), 0]
        env.data.qpos[2] = 0.4625
    n = int(secs / env.control_dt)
    rng = np.random.default_rng(seed)
    hits = set(rng.integers(10, max(11, n - 10), size=3)) if shove else set()
    for k in range(n):
        a = np.clip(np.append(o, 1.0) @ W, -1, 1)
        o, _, fell, trunc, _ = env.step(a)
        if k in hits:
            env.data.qvel[0] += float(rng.choice([-1.0, 1.0])) * shove
        if fell or trunc:
            break
    # upright time is the honest measure; drift is a small tie-breaker so that
    # two rules which both survive are separated by which one held its ground
    return env.t - 0.05 * min(2.0, abs(float(env.base_pos[0])))


# What actually beats the watched rule. No point practising what it can do.
HARD = [
    {"tilt": 0.22}, {"tilt": -0.22}, {"tilt": 0.30}, {"tilt": -0.30},
    {"shove": 1.5}, {"shove": 2.5},
]


def score(W, conds, secs, seed0=0):
    return float(np.mean([episode(W, seed0 + i, secs, **c)
                          for i, c in enumerate(conds)]))


def practise(W, iters=12, dirs=6, nu=0.25, alpha=0.30, secs=5.0, seed=0):
    rng = np.random.default_rng(seed)
    # Perturb each weight in proportion to ITS OWN size, not the average size.
    # Scaling every weight by the mean magnitude swamps the small ones with
    # noise far bigger than themselves and shreds the rule - every round came
    # back worse than the start. Relative jitter keeps the shape of the rule
    # and lets the dominant weights grow, which is what is needed here: the
    # teacher balanced with a pitch gain of 30 and a 30 degree recovery wants
    # nearer 200, so the search has to be able to scale a weight severalfold.
    scale = np.abs(W) + 0.02 * float(np.abs(W).mean())
    base = score(W, HARD, secs)
    print(f"  start: {base:.2f}s mean over {len(HARD)} hard conditions")
    best, bestW = base, W.copy()
    for it in range(iters):
        deltas, rp, rm = [], [], []
        for _ in range(dirs):
            d = rng.normal(0, 1, W.shape) * scale
            deltas.append(d)
            rp.append(score(W + nu * d, HARD, secs))
            rm.append(score(W - nu * d, HARD, secs))
        rp, rm = np.array(rp), np.array(rm)
        # keep only the directions that told us the most: averaging in the
        # useless ones just dilutes the step
        keep = np.argsort(-np.maximum(rp, rm))[:max(1, dirs // 2)]
        # step towards whichever half of each pair did better, weighted by how
        # much better, and normalised by the spread so the step size does not
        # depend on the units of the reward
        sd = np.concatenate([rp[keep], rm[keep]]).std() + 1e-6
        step = sum((rp[i] - rm[i]) * deltas[i] for i in keep)
        W = W + (alpha / (len(keep) * sd)) * step
        cur = score(W, HARD, secs)
        if cur > best:
            best, bestW = cur, W.copy()
        print(f"  round {it+1:2d}: {cur:5.2f}s   best {best:5.2f}s")
    return bestW, base, best


def report(W, label, secs=10.0):
    print(f"\n{label}")
    print(f"  {'condition':>20} {'upright':>9}")
    rows = [("flat", {}), ("tilt 16 deg", {"tilt": 0.28}),
            ("tilt 23 deg", {"tilt": 0.40}), ("tilt 31 deg", {"tilt": 0.55}),
            ("shove 1.5 m/s", {"shove": 1.5}), ("shove 2.5 m/s", {"shove": 2.5}),
            ("shove 4.0 m/s", {"shove": 4.0})]
    out = {}
    for name, c in rows:
        # held-out seeds: never used during practice
        v = np.mean([episode(W, 9000 + i, secs, **c) for i in range(4)])
        out[name] = v
        print(f"  {name:>20} {v:8.2f}s")
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--iters", type=int, default=12)
    ap.add_argument("--dirs", type=int, default=4)
    ap.add_argument("--secs", type=float, default=5.0)
    ap.add_argument("--out", default=str(ROOT / "runs/linear_balance.json"))
    args = ap.parse_args()

    W0 = watched_rule()
    print(f"watched rule: {W0.size} numbers\n")
    before = report(W0, "BEFORE practice (watched only)")

    print(f"\npractising on the {len(HARD)} conditions that break it:")
    t0 = time.time()
    W, base, best = practise(W0, iters=args.iters, dirs=args.dirs, secs=args.secs)
    print(f"  {time.time()-t0:.0f}s of practice")

    after = report(W, "AFTER practice")
    print(f"\n  {'condition':>20} {'before':>9} {'after':>9}  {'change':>8}")
    for k in before:
        d = after[k] - before[k]
        print(f"  {k:>20} {before[k]:8.2f}s {after[k]:8.2f}s  {d:+7.2f}s")
    Path(args.out).parent.mkdir(parents=True, exist_ok=True)
    json.dump({"W": W.tolist(), "kind": "linear", "task": "balance"},
              open(args.out, "w"))
    print(f"\nwrote {args.out}")


if __name__ == "__main__":
    main()


# ---------------------------------------------------------------------------
# Structured practice.
#
# The 72-number watched rule turned out to be unimprovable: least squares built
# it from huge cancelling coefficients on inputs that barely vary (condition
# number 1e68), so it works exactly where it was fitted and falls apart at the
# smallest nudge. There is nothing to climb.
#
# The behaviour underneath is four numbers. Searching those directly gives a
# landscape with actual slopes on it - and the result is four gains you could
# type into the robot's firmware, rather than 72 opaque weights.
GAIN_NAMES = ["pitch", "gyro", "speed", "odometry"]
# expanded from the teacher: u = -kp*(pitch - station*(0.5*v + 0.1*odo)) + kd*w
SEED = np.array([-30.0, 2.0, 29.25, 5.85])


def gain_episode(g, seed, secs, tilt=0.0, shove=0.0):
    import mujoco
    env = RobotEnv(SPEC, seed=seed, task="balance", randomise=False)
    env.reset()
    if tilt:
        env.data.qpos[3:7] = [math.cos(tilt / 2), 0, math.sin(tilt / 2), 0]
        env.data.qpos[2] = 0.4625
    jd = [env.model.jnt_dofadr[mujoco.mj_name2id(
        env.model, mujoco.mjtObj.mjOBJ_JOINT, n)] for n in ("l_wheel_j", "r_wheel_j")]
    n = int(secs / env.control_dt)
    rng = np.random.default_rng(seed)
    hits = set(rng.integers(10, max(11, n - 10), size=3)) if shove else set()
    odo = 0.0
    for k in range(n):
        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])
        v = 0.5 * sum(float(env.data.qvel[i]) for i in jd) * 0.0625
        odo += v * env.control_dt
        u = float(np.clip(g[0] * pitch + g[1] * rate + g[2] * v + g[3] * odo, -1, 1))
        _, _, fell, trunc, _ = env.step(np.array([u, u]))
        if k in hits:
            env.data.qvel[0] += float(rng.choice([-1.0, 1.0])) * shove
        if fell or trunc:
            break
    return env.t - 0.05 * min(2.0, abs(float(env.base_pos[0])))


def gain_score(g, conds, secs, seed0=0):
    return float(np.mean([gain_episode(g, seed0 + i, secs, **c)
                          for i, c in enumerate(conds)]))


def practise_gains(g, iters=25, dirs=6, nu=0.3, alpha=0.5, secs=5.0, seed=0):
    rng = np.random.default_rng(seed)
    best = gain_score(g, HARD, secs)
    bestg = g.copy()
    print(f"  start {best:5.2f}s   gains " +
          " ".join(f"{n}={v:+.1f}" for n, v in zip(GAIN_NAMES, g)))
    for it in range(iters):
        deltas = [rng.normal(0, 1, g.shape) * np.abs(g) for _ in range(dirs)]
        rp = np.array([gain_score(bestg + nu * d, HARD, secs) for d in deltas])
        rm = np.array([gain_score(bestg - nu * d, HARD, secs) for d in deltas])
        keep = np.argsort(-np.maximum(rp, rm))[:max(1, dirs // 2)]
        sd = np.concatenate([rp[keep], rm[keep]]).std() + 1e-6
        cand = bestg + (alpha / (len(keep) * sd)) * sum(
            (rp[i] - rm[i]) * deltas[i] for i in keep)
        cur = gain_score(cand, HARD, secs)
        # always step from the best rule so far, never from a worse one: the
        # 72-number search random-walked downhill because it did not do this
        if cur > best:
            best, bestg = cur, cand
        if (it + 1) % 5 == 0 or it == 0:
            print(f"  round {it+1:2d}: {cur:5.2f}s   best {best:5.2f}s   " +
                  " ".join(f"{v:+7.1f}" for v in bestg))
    return bestg, best