"""Episode-return policy search for the head-to-wheel recovery on the Arena's roller rig."""

import functools
import hashlib
import json
import multiprocessing as mp
import os
from pathlib import Path

import numpy as np
from arena import runner
from arena.course import compile_stage
from arena.manifest import load_event
from arena.policy import load_policy
from arena.timeline import Start, Timeline

OUT = Path(__file__).resolve().parent
RIG = "rollers"
EVENT = load_event("got-a-talent")
START = Start(x=0, y=-1.45, yaw_deg=90)
BASE_MODELS = {
    "skate": (
        "JackDH/microduck-roller-sprint-20261110-best",
        "43e9b57825fbfe08cf4b510368b643ff9ae47d93",
        "policy.onnx",
        "0b724f5cebe7f50815f6dd9631009e47c4b0a607f81b22d0f4048ffd6d5a77d6",
    ),
    "roll": (
        "pollen-robotics/microduck-policies",
        "1b56c396825c052a4e26e95cf2b8d8298af9e9b4",
        "roulade.onnx",
        "3d60da08fc13f29c1b57f41977aa898132c0d60042100149d8e775affcbca32b",
    ),
    "rise": (
        "pollen-robotics/microduck-policies",
        "1b56c396825c052a4e26e95cf2b8d8298af9e9b4",
        "velstand.onnx",
        "1c659be55da94bc5753b707de5c6a3e7c49931e05ca3b6991615cef1a8ba9a45",
    ),
}


@functools.lru_cache(maxsize=3)
def model_file(name):
    repo, revision, filename, digest = BASE_MODELS[name]
    path = None
    if project := os.environ.get("MICRODUCK_PROJECT"):
        rel = (
            "challenges/publish-roller-best-20261110/policy.onnx"
            if name == "skate"
            else f"policies/official/{filename}"
        )
        path = Path(project) / rel
    if path is None or not path.is_file():
        from huggingface_hub import hf_hub_download

        path = Path(hf_hub_download(repo, filename, revision=revision))
    if hashlib.sha256(path.read_bytes()).hexdigest() != digest:
        raise ValueError(f"Unexpected contents for {name} policy")
    return path


# The timeline gives the coarse cues. The network reads gravity and angular velocity to adjust its leg action.
TIMES = [0, 0.2, 1.37, 1.55, 1.87, 2.33, 2.85, 3.15]
MODES = [0, 0, 1, 2, 3, 4, 5, 0]
SPEEDS = [0, 0.32, 0, 0, 0, 0.54, 0.54, 0.65]
DURATION = 5.16


def timeline_dict():
    frames = [
        {"t": t, "twist": [v, 0, 0], "body": [0, 0, 0, 0, 0, m]}
        for t, v, m in zip(TIMES, SPEEDS, MODES)
    ]
    return {
        "timeline_version": 2,
        "duration_s": DURATION,
        "ducks": [{"policy": "self", "start": START.to_dict(), "keyframes": frames}],
    }


TIMELINE = Timeline(
    DURATION,
    tuple(TIMES),
    tuple((v,) + (0.0,) * 11 + (m,) for v, m in zip(SPEEDS, MODES)),
)


@functools.lru_cache(maxsize=4)
def actor(name):
    p = load_policy(model_file(name))
    return p.session, p._input_name, p._output_name


def actions(name, x):
    sess, i, o = actor(name)
    return sess.run([o], {i: x.reshape(1, 61)})[0][0]


# Symmetric sagittal legs, neck/head, and a lateral hip balance response.
MAP = np.zeros((6, 14), np.float32)
for n, (l, r) in enumerate([(2, 11), (3, 12), (4, 13)]):
    MAP[n, l] = 1
    MAP[n, r] = -1
MAP[3, 5] = 1
MAP[4, 6] = 1
MAP[5, [1, 10]] = [1, -1]
SHAPE = (4, 6)
SIZE = 32


def configure(speed=0.32, entry=1.37):
    global TIMES, SPEEDS, DURATION, TIMELINE
    offset = entry - 1.37
    TIMES = [0, 0.2] + [
        round(t + offset, 2) for t in (1.37, 1.55, 1.87, 2.33, 2.85, 3.15)
    ]
    SPEEDS = [0, speed, 0, 0, 0, 0.54, 0.54, 0.65]
    DURATION = round(5.16 + offset, 2)
    TIMELINE = Timeline(
        DURATION,
        tuple(TIMES),
        tuple((v,) + (0.0,) * 11 + (m,) for v, m in zip(SPEEDS, MODES)),
    )


class Landing:
    def __init__(self, theta):
        self.p = np.asarray(theta, np.float32)
        self.steps = self.p[:24].reshape(SHAPE)

    def act(self, obs):
        mode = int(obs[-1] + 0.5)
        x = obs.copy()
        x[49:] = 0
        if mode == 0:
            return actions("skate", x)
        if mode <= 3:
            x[34:48] /= np.float32(1.12)
            x[48] = 0
            base = actions("roll", x) * np.float32(1.12)
        else:
            base = actions("rise", x)
        if mode == 1:
            return base
        residual = self.steps[mode - 2].copy()
        # The last eight parameters adapt the joints to pitch, roll and their velocities.
        pitch_vel = np.clip(obs[1] * np.float32(0.14), -1, 1)
        residual[:3] += self.p[24:27] * obs[3] + self.p[27:30] * pitch_vel
        residual[5] += self.p[30] * obs[4] + self.p[31] * np.clip(
            obs[0] * np.float32(0.15), -1, 1
        )
        return base + np.tanh(residual) @ MAP


# Compile each stage once per process. Each new game resets friction/delay and carries the original physics.
@functools.lru_cache(maxsize=1)
def stage():
    return compile_stage(EVENT, [(RIG, START)])


def run(policy, seed):
    model = stage()
    sim, settle, _ = runner._start(
        EVENT,
        model,
        [policy],
        [START.yaw_deg],
        int(seed),
        [np.asarray(TIMELINE.command_at(0))],
    )
    trace = np.empty((round(DURATION * 50), 21), np.float32)
    for i in range(len(trace)):
        sim.step([np.asarray(TIMELINE.command_at(i * 0.02))])
        sim.record(trace[i : i + 1])
        if runner._reset_by_mujoco(sim.data):
            raise RuntimeError("Diverged")
    return trace, settle


def metrics(tr):
    w, x, y, z = tr[:, 3:7].T
    tilt = np.degrees(np.arccos(np.clip(1 - 2 * (x * x + y * y), -1, 1)))
    pitch = np.degrees(np.unwrap(np.arctan2(2 * (w * y - x * z), 2 * (x * y + w * z))))
    start = int(np.ceil(TIMES[2] * 50 - 1e-8)) - 1
    angle = pitch - np.median(pitch[start - 8 : start])
    theta = float(np.median(angle[-25:]))
    h = tr[:, 2]
    pos = tr[:, 1]
    upright = (tilt < 35) & (h > 0.085)
    inv = float(tilt[start:].max())
    glide = float(pos[-1] - pos[-51])
    inside = bool((abs(tr[:, 0]) < 1.28).all() and (abs(pos) < 1.87).all())
    success = bool(
        310 < theta < 405
        and inv > 150
        and np.all(upright[-45:])
        and glide > 0.09
        and inside
        and pos[start] - pos[start - 25] > 0.035
    )
    # Reward partial catches continuously and reserve the completion bonus for a genuine roll and rollout.
    turn = np.exp(-(((angle - 360) / 82) ** 2))
    orient = np.clip((1 + 1 - 2 * (x * x + y * y)) / 2, 0, 1)
    height = np.clip((h - 0.043) / 0.073, 0, 1)
    alignment = np.clip(2 * (x * y + w * z), 0, 1)
    quality = turn * orient * height * (0.45 + 0.55 * alignment)
    reward = (
        0.9 * np.clip((inv - 100) / 70, 0, 1)
        + 1.3 * quality[start + 50 : start + 100].mean()
        + 3.1 * quality[start + 100 :].mean()
        + 1.4 * np.clip(glide / 0.36, -0.5, 1) * quality[-50:].mean()
        + 2.5 * success
    )
    good = np.flatnonzero(~upright[start:])
    recovery = round((int(good[-1]) + 1) * 0.02, 2) if len(good) and success else None
    return {
        "return": round(float(reward), 6),
        "success": success,
        "turn_deg": round(theta, 1),
        "max_tilt": round(inv, 1),
        "entry_mps": round(float((pos[start] - pos[start - 25]) / 0.5), 3),
        "entry_distance_m": round(float(pos[start] - pos[0]), 3),
        "final_m": round(glide, 3),
        "recovery_s": recovery,
        "inside": inside,
    }


def test_job(args):
    theta, seed = args
    tr, _ = run(Landing(theta), seed)
    return metrics(tr)


def calc(pool, models, seeds):
    res = list(
        pool.imap(test_job, [(v, s) for v in models for s in seeds], chunksize=2)
    )
    return [res[i : i + len(seeds)] for i in range(0, len(res), len(seeds))]


def save(path, obj):
    tmp = path.with_suffix(".tmp")
    tmp.write_text(json.dumps(obj, indent=2) + "\n")
    tmp.replace(path)


def train(
    steps,
    pop,
    batch,
    warm=None,
    run_dir=OUT,
    random_seed=11062026,
    validation=None,
    search_scale=0.62,
    floor=0.075,
    workers=16,
):
    rng = np.random.default_rng(random_seed)
    mean = np.zeros(SIZE)
    sig = np.full(SIZE, 0.33)
    sig[24:] = 0.22
    best = []
    hist = []
    if warm is not None:
        mean = np.asarray(warm)
        sig *= search_scale
    with mp.Pool(
        workers, initializer=configure, initargs=(SPEEDS[1], TIMES[2])
    ) as pool:
        validation = list(range(31, 43)) if validation is None else list(validation)
        baseline = calc(pool, [mean], validation)[0]
        print(
            "Baseline validation:",
            sum(r["success"] for r in baseline),
            "/",
            len(baseline),
            flush=True,
        )
        for gen in range(steps):
            first = rng.normal(size=(pop // 2, SIZE))
            eps = np.concatenate((first, -first))
            candidates = np.clip(mean + sig * eps, -2.8, 2.8)
            candidates[0] = mean
            seeds = list(
                map(int, rng.choice(np.arange(80, 600), size=batch, replace=False))
            )
            results = calc(pool, candidates, seeds)
            returns = np.array([np.mean([a["return"] for a in rr]) for rr in results])
            elite = candidates[np.argsort(returns)[-max(5, pop // 6) :]]
            mean = 0.27 * mean + 0.73 * np.mean(elite, axis=0)
            sig = np.maximum(floor, 0.62 * sig + 0.38 * np.std(elite, axis=0))
            leaders = [mean, candidates[np.argmax(returns)]]
            if (gen + 1) % 2 == 0 or gen == 0:
                trials = calc(pool, leaders, [*validation, 1])
                for model, rr in zip(leaders, trials):
                    item = {
                        "gen": gen + 1,
                        "theta": model.tolist(),
                        "validation_success": sum(q["success"] for q in rr[:-1]),
                        "official_pass": rr[-1]["success"],
                        "validation_return": float(
                            np.mean([q["return"] for q in rr[:-1]])
                        ),
                    }
                    best.append(item)
                best.sort(
                    key=lambda o: (o["validation_success"], o["validation_return"]),
                    reverse=True,
                )
                best = best[:5] + [o for o in best[5:] if o["official_pass"]][:3]
                save(run_dir / "checkpoints.json", best)
            hist.append(
                {
                    "generation": gen + 1,
                    "training_seeds": seeds,
                    "mean_return": float(returns.mean()),
                    "best_training_return": float(returns.max()),
                }
            )
            save(run_dir / "learning.json", hist)
            print(
                "Generation",
                gen + 1,
                "best validation",
                best[0]["validation_success"],
                "/",
                len(validation),
                "training return",
                round(returns.max(), 2),
                flush=True,
            )