"""Courses: learn one thing at a time, in order.

A course is a sequence of stages. Each stage trains a single incentive, then
sits an exam. Pass and the next stage inherits the weights; fail and the course
stops there rather than piling a harder skill on a broken one. You cannot learn
to run before you can walk, and this is the thing that enforces it.

Two properties make the hand-over work:

  * Every incentive publishes the same fixed-width goal channels (see
    tasks.GOAL_CHANNELS), so the network shape never changes between stages.
  * A stage resumes from the previous stage's checkpoint, so stage 2 starts
    from a robot that can already stand rather than from noise.

  python course.py --course standard --run myrobot
  python course.py --list
"""

import argparse
import json
import subprocess
import sys
import time
from pathlib import Path

import numpy as np

from env import RobotEnv
from ppo import Policy
from rollout import record

ROOT = Path(__file__).parent.parent
PY = sys.executable


# A gate is deliberately about behaviour, not reward. Reward scales differ
# between incentives, so "mean return above X" would mean nothing across
# stages; "stays up for 80% of a full episode, 8 times in 10" always does.
COURSES = {
    "standard": {
        "label": "Learn to balance, then to move",
        "description": (
            "The natural order for a wheeled biped. Stand before you drive, "
            "drive before you turn on the spot, and only then go somewhere."),
        "stages": [
            {"task": "balance", "steps": 1_000_000,
             "gate": {"survive_frac": 0.80, "episodes": 20, "pass_rate": 0.75},
             "why": "Legs locked straight. Just roll the wheels to stay under "
                    "the body: two numbers instead of six."},
            {"task": "bob", "steps": 2_000_000,
             "gate": {"survive_frac": 0.75, "episodes": 20, "pass_rate": 0.70},
             "why": "Now change height while staying up, so the legs and the "
                    "wheels have to cooperate."},
            {"task": "recover", "steps": 2_000_000,
             "gate": {"survive_frac": 0.75, "episodes": 20, "pass_rate": 0.70},
             "why": "Survive a shove: recovery beyond what a gentle correction "
                    "can manage."},
            {"task": "drive",   "steps": 4_000_000,
             "gate": {"survive_frac": 0.80, "episodes": 20, "pass_rate": 0.75},
             "why": "It can hold itself up, so now ask it to go somewhere."},
            {"task": "spin",    "steps": 2_000_000,
             "gate": {"survive_frac": 0.80, "episodes": 20, "pass_rate": 0.70},
             "why": "Turning on the spot is a harder balance than driving straight."},
            {"task": "goto",    "steps": 4_000_000,
             "gate": {"survive_frac": 0.70, "episodes": 20, "pass_rate": 0.60,
                      "success_rate": 0.40},
             "why": "Combine driving and turning into reaching a place."},
            {"task": "climb",   "steps": 4_000_000,
             "gate": {"survive_frac": 0.60, "episodes": 20, "pass_rate": 0.50,
                      "success_rate": 0.30},
             "why": "Last, because it needs every earlier skill at once."},
        ],
    },
    "balance_only": {
        "label": "Just learn to stand up",
        "description": "One stage. The quickest way to see the whole loop work.",
        "stages": [
            {"task": "stand", "steps": 3_000_000,
             "gate": {"survive_frac": 0.80, "episodes": 20, "pass_rate": 0.80},
             "why": "Balance is the whole job here."},
        ],
    },
}


def examine(policy_path, task, episodes=20, survive_frac=0.8, seed=7000,
            spec_path=None):
    """Sit the exam: run N randomised episodes and measure behaviour."""
    policy = Policy.load(policy_path)
    spec_path = spec_path or policy.meta.get("spec")
    layout = policy.meta.get("obs_layout", "union")

    survived, wins, full = [], [], None
    for ep in range(episodes):
        env = RobotEnv(spec_path, seed=seed + ep, task=task, layout=layout)
        full = env.max_steps * env.control_dt
        rng = np.random.default_rng(seed + ep)
        obs = env.reset()
        for _ in range(env.max_steps):
            a, _, _ = policy.act(obs[None], rng, deterministic=True)
            obs, r, fell, trunc, info = env.step(np.clip(a[0], -1, 1))
            if fell:
                break
        survived.append(env.t)
        wins.append(bool(env.task_obj.success(env)))

    surv = np.array(survived)
    need = full * survive_frac
    return {
        "episodes": episodes,
        "episode_s": round(full, 2),
        "mean_upright_s": round(float(surv.mean()), 2),
        "median_upright_s": round(float(np.median(surv)), 2),
        "worst_upright_s": round(float(surv.min()), 2),
        "pass_rate": round(float(np.mean(surv >= need)), 3),
        "success_rate": round(float(np.mean(wins)), 3),
        "threshold_s": round(need, 2),
    }


def run_course(course_name, run_name, spec, envs, workers, steps_scale=1.0,
               snapshot_every=500_000, resume_from=None, only=None):
    """Run every stage, or just one.

    `only` (1-based) runs a single stage and seeds it from the previous stage's
    best policy if that stage has been passed. That is what lets the bench work
    one step at a time under the user's control instead of running the whole
    course unattended.
    """
    course = COURSES[course_name]
    run_dir = ROOT / "runs" / run_name
    run_dir.mkdir(parents=True, exist_ok=True)
    state_path = run_dir / "course.json"

    state = None
    if state_path.exists():
        try:
            state = json.load(open(state_path))       # keep earlier stages' results
        except (OSError, ValueError):
            state = None
    if state is None:
        state = {
            "course": course_name, "label": course["label"], "run": run_name,
            "spec": spec, "started": time.time(), "stages": [],
            "status": "running", "current": 0,
        }
        for st in course["stages"]:
            state["stages"].append({
                "task": st["task"], "steps": int(st["steps"] * steps_scale),
                "why": st["why"], "gate": st["gate"], "status": "pending",
                "exam": None, "policy": None, "run": f"{run_name}__{st['task']}",
            })
    state["status"] = "running"

    def save():
        with open(state_path, "w") as f:
            json.dump(state, f, indent=1)

    save()
    carry = resume_from

    # Running one stage: pick up the previous stage's policy so the robot keeps
    # what it already knows.
    todo = range(len(state["stages"]))
    if only is not None:
        idx = only - 1
        todo = [idx]
        state["stages"][idx]["steps"] = int(
            course["stages"][idx]["steps"] * steps_scale)
        own_latest = ROOT / "runs" / state["stages"][idx]["run"] / "policy_latest.json"
        if resume_from is None and own_latest.exists() and \
                state["stages"][idx].get("status") in ("stopped", "failed"):
            # carry on from where this stage was interrupted rather than
            # throwing away the practice already done
            carry = own_latest
        elif idx > 0:
            prev = state["stages"][idx - 1]
            prev_best = ROOT / "runs" / prev["run"] / "policy_best.json"
            if prev["status"] == "passed" and prev_best.exists():
                carry = prev_best

    for i in todo:
        st = state["stages"][i]
        state["current"] = i
        st["status"] = "training"
        st["started"] = time.time()
        save()

        stage_run = st["run"]
        stage_dir = ROOT / "runs" / stage_run
        argv = [PY, "train.py", "--task", st["task"], "--steps", str(st["steps"]),
                "--run", stage_run, "--spec", spec, "--envs", str(envs),
                "--workers", str(workers), "--layout", "union",
                "--snapshot-every", str(snapshot_every)]
        if carry:
            argv += ["--resume", str(carry)]

        print(f"\n=== stage {i+1}/{len(state['stages'])}: {st['task']} "
              f"({st['steps']/1e6:.1f}M steps){'  resuming' if carry else ''}",
              flush=True)
        print(f"    {st['why']}", flush=True)

        with open(stage_dir.parent / f"{stage_run}.out", "w") as log:
            stage_dir.mkdir(parents=True, exist_ok=True)
            rc = subprocess.call(argv, cwd=str(Path(__file__).parent),
                                 stdout=log, stderr=subprocess.STDOUT)
        if rc != 0:
            st["status"] = "error"
            state["status"] = "error"
            save()
            print(f"    stage failed (exit {rc}); see runs/{stage_run}.out")
            return state

        best = stage_dir / "policy_best.json"
        st["policy"] = f"../runs/{stage_run}/policy_best.json"
        st["status"] = "examining"
        save()

        g = st["gate"]
        exam = examine(str(best), st["task"], episodes=g.get("episodes", 20),
                       survive_frac=g.get("survive_frac", 0.8), spec_path=spec)
        st["exam"] = exam

        passed = exam["pass_rate"] >= g.get("pass_rate", 0.8)
        if "success_rate" in g:
            passed = passed and exam["success_rate"] >= g["success_rate"]
        st["status"] = "passed" if passed else "failed"
        st["finished"] = time.time()
        save()

        print(f"    exam: upright {exam['mean_upright_s']}s of {exam['episode_s']}s, "
              f"pass rate {exam['pass_rate']*100:.0f}% "
              f"(needs {g.get('pass_rate',0.8)*100:.0f}%)  -> {st['status'].upper()}",
              flush=True)

        if not passed:
            state["status"] = "stopped" if only is None else "idle"
            save()
            print(f"\nStage {i+1} ({st['task']}) did not pass. Give it more "
                  f"steps and run it again.")
            return state

        carry = best      # the next stage starts from what this one learned

    if only is not None:
        state["status"] = "idle"
        state["current"] = min(only, len(state["stages"]) - 1)
    else:
        state["status"] = "complete"
        state["current"] = len(state["stages"])
    save()
    print("\nDone.")
    return state


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--course", default="standard")
    ap.add_argument("--run", default=None)
    ap.add_argument("--spec", default=str(ROOT / "robots/wheeled_biped.json"))
    ap.add_argument("--envs", type=int, default=24)
    ap.add_argument("--workers", type=int, default=3)
    ap.add_argument("--steps-scale", type=float, default=1.0,
                    help="scale every stage's budget, e.g. 0.1 for a quick trial")
    ap.add_argument("--snapshot-every", type=int, default=500_000)
    ap.add_argument("--resume-from", default=None,
                    help="checkpoint to seed stage 1 from")
    ap.add_argument("--only", type=int, default=None,
                    help="run just this stage (1-based), seeded from the previous one")
    ap.add_argument("--list", action="store_true")
    args = ap.parse_args()

    if args.list:
        for name, c in COURSES.items():
            print(f"\n{name}: {c['label']}")
            print(f"  {c['description']}")
            for i, s in enumerate(c["stages"], 1):
                print(f"  {i}. {s['task']:8s} {s['steps']/1e6:4.1f}M  {s['why']}")
        return

    run = args.run or f"course_{args.course}_{time.strftime('%Y%m%d_%H%M%S')}"
    run_course(args.course, run, str(Path(args.spec).resolve()), args.envs,
               args.workers, args.steps_scale, args.snapshot_every,
               args.resume_from, args.only)


if __name__ == "__main__":
    main()