"""Train sumo policies with OpenAI-style evolution strategies (ES) and self-play.

Stages (each continues from the previous one's weights):
  L2  ES against the hand-coded charger only.
  L3  Self-play: ES against a growing pool of its own past checkpoints plus the
      hand-coded bots (a simple fictitious-self-play / "league" scheme).
  L4  L3 continued with domain randomization over friction, masses, motor lag,
      sensor noise and observation delay.

ES in one paragraph: perturb the parameter vector theta with Gaussian noise
eps_k (in antithetic pairs +eps/-eps), play matches with every perturbed
policy, replace the raw fitness by centred ranks (robust to outliers and reward
scale), and move theta along  sum_k rank_k * eps_k / (n * sigma)  -- a
finite-difference estimate of the gradient of the Gaussian-smoothed fitness --
using Adam. No backprop, and every population member is evaluated in the same
vectorized simulator call.

Usage:
  python train.py                 # full pipeline L2 -> L3 -> L4 into weights/
  python train.py --quick         # small budget into weights_quick/ (~2-3 min)
  python train.py --stage L3      # just one stage (needs the previous stage's weights)
"""

from __future__ import annotations

import argparse
import json
import os
import time

import numpy as np

from policies import (
    N_PARAMS,
    ChargerPolicy,
    DodgerPolicy,
    MLPPolicy,
    MuxPolicy,
    RandomPolicy,
    init_params,
    run_matches,
)
from sumo_env import MAX_STEPS, RING_RADIUS, PhysParams, SumoEnv, random_starts, tile_starts

HERE = os.path.dirname(os.path.abspath(__file__))


# ---------------------------------------------------------------- ES pieces
class Adam:
    def __init__(self, dim, lr=0.02, beta1=0.9, beta2=0.999, eps=1e-8):
        self.lr, self.b1, self.b2, self.eps = lr, beta1, beta2, eps
        self.m = np.zeros(dim)
        self.v = np.zeros(dim)
        self.t = 0

    def step(self, grad):
        """Returns the update to ADD to the parameters (gradient ascent)."""
        self.t += 1
        self.m = self.b1 * self.m + (1 - self.b1) * grad
        self.v = self.b2 * self.v + (1 - self.b2) * grad**2
        mhat = self.m / (1 - self.b1**self.t)
        vhat = self.v / (1 - self.b2**self.t)
        return self.lr * mhat / (np.sqrt(vhat) + self.eps)


def centered_ranks(x: np.ndarray) -> np.ndarray:
    """Map fitnesses to ranks scaled into [-0.5, 0.5] (ties broken arbitrarily)."""
    ranks = np.empty(len(x))
    ranks[np.argsort(x, kind="stable")] = np.arange(len(x))
    return ranks / (len(x) - 1) - 0.5


def es_gradient(fitness: np.ndarray, eps: np.ndarray, sigma: float) -> np.ndarray:
    """Antithetic ES gradient. fitness: (2P,) ordered [+eps_0..+eps_P-1, -eps_0..]."""
    r = centered_ranks(fitness)
    P = eps.shape[0]
    return (r[:P] - r[P:]) @ eps / (2 * P * sigma)


def shaped_fitness(res: dict, shape: float, tie_penalty: float = 0.0) -> np.ndarray:
    """Win/loss/tie (+1/-1/0) plus light shaping (only if shape > 0).

    * a win counts more the faster it happens, a loss hurts more the earlier
      it happens (up to +-shape);
    * on a tie (time cap) you get up to +-shape for ending closer to the ring
      center than your opponent, minus ``tie_penalty`` (discourages the
      "hide and run out the clock" strategy that otherwise appears under
      domain randomization).
    Final evaluation (evaluate.py) never uses shaping.
    """
    out = res["outcome"].astype(float)
    speed = 1.0 - res["end_step"] / MAX_STEPS
    d = res["final_dist"] / RING_RADIUS
    tie = out == 0
    f = out + shape * out * speed
    f[tie] = shape * np.clip(d[tie, 1] - d[tie, 0], -1, 1) - tie_penalty
    return f


# ---------------------------------------------------------------- opponents
class OpponentSampler:
    """Chooses an opponent for each of the M scenarios in a generation."""

    def __init__(self, mode: str, pool: list[np.ndarray], rng: np.random.Generator,
                 hand_frac: float = 0.3, recent_frac: float = 0.4):
        self.mode = mode
        self.pool = pool  # list of flat parameter vectors (shared, appended to during training)
        self.rng = rng
        self.hand_frac = hand_frac
        self.recent_frac = recent_frac
        # Prioritized fictitious self-play (PFSP, as in AlphaStar): keep a running
        # estimate of how often we beat each pool member and sample the ones we
        # struggle against more often.
        self.win_est: list[float] = []

    def _pfsp_probs(self):
        while len(self.win_est) < len(self.pool):
            self.win_est.append(0.5)  # new snapshots start "even"
        w = (1.0 - np.array(self.win_est)) ** 2 + 0.02
        return w / w.sum()

    def update(self, outcome: np.ndarray, M: int):
        """Update win estimates from this generation's outcomes (E,) of the population."""
        if self.mode == "charger":
            return
        score = (outcome.reshape(-1, M) == 1).mean(0) + 0.5 * (outcome.reshape(-1, M) == 0).mean(0)
        for i in np.unique(self.last_pool_idx[self.last_pool_idx >= 0]):
            s = score[self.last_pool_idx == i].mean()
            self.win_est[i] = 0.8 * self.win_est[i] + 0.2 * s

    def build(self, M: int, reps: int) -> MuxPolicy:
        """Opponent policy for E = reps*M envs; env e plays scenario e % M."""
        if self.mode == "charger":
            return ChargerPolicy()
        rng = self.rng
        kind = np.empty(M, dtype=object)
        u = rng.random(M)
        # hand-coded share: charger 50%, dodger 30%, random 20%
        h = self.hand_frac
        kind[u < h * 0.5] = "charger"
        kind[(u >= h * 0.5) & (u < h * 0.8)] = "dodger"
        kind[(u >= h * 0.8) & (u < h)] = "random"
        pool_idx = np.full(M, -1)
        is_pool = u >= h
        n_pool = int(is_pool.sum())
        # pool share: some from the most recent snapshots, the rest by PFSP over all history
        recent = rng.random(n_pool) < self.recent_frac
        last = len(self.pool)
        k_recent = min(3, last)
        pool_idx[is_pool] = np.where(
            recent, rng.integers(last - k_recent, last, n_pool), rng.choice(last, n_pool, p=self._pfsp_probs())
        )
        kind[is_pool] = "pool"
        self.last_kind, self.last_pool_idx = kind, pool_idx

        scen = np.arange(reps * M) % M
        groups = []
        for name, pol in (("charger", ChargerPolicy()), ("dodger", DodgerPolicy()), ("random", RandomPolicy())):
            groups.append((pol, np.where(kind[scen] == name)[0]))
        envs = np.where(kind[scen] == "pool")[0]
        if len(envs):
            params = np.stack(self.pool)[pool_idx[scen[envs]]]
            groups.append((MLPPolicy(params), envs))
        return MuxPolicy(groups)


# ---------------------------------------------------------------- evaluation helper
def quick_eval(theta, opponents: dict, n: int = 256, seed: int = 123) -> dict:
    """Win/loss/tie rates of theta vs each named opponent, nominal physics."""
    out = {}
    for name, opp in opponents.items():
        env = SumoEnv(n, seed=seed)
        o = run_matches(env, MLPPolicy(theta), opp, np.random.default_rng(seed))["outcome"]
        out[name] = [float((o == 1).mean()), float((o == -1).mean()), float((o == 0).mean())]
    return out


# ---------------------------------------------------------------- one ES stage
def run_stage(name, theta, cfg, out_dir, pool=None, dr=False, opp_mode="pool", seed=0, verbose=True):
    rng = np.random.default_rng(seed)
    P, M = cfg["pairs"], cfg["episodes"]
    sigma, gens = cfg["sigma"], cfg["gens"]
    adam = Adam(N_PARAMS, lr=cfg["lr"])
    pool = pool if pool is not None else []
    sampler = OpponentSampler(opp_mode, pool, rng, hand_frac=cfg.get("hand_frac", 0.3))
    eval_opps = {"charger": ChargerPolicy(), "dodger": DodgerPolicy()}
    if pool:
        eval_opps["pool0"] = MLPPolicy(pool[0])
    log = {"gen": [], "fitness": [], "win_rate": [], "eval": [], "eval_gen": [], "time": []}
    t0 = time.time()
    for g in range(1, gens + 1):
        shape = cfg["shape"] * max(0.0, 1.0 - (g - 1) / max(1, cfg.get("shape_decay_gens", 10**9)))
        eps = rng.standard_normal((P, N_PARAMS))
        members = np.concatenate([theta + sigma * eps, theta - sigma * eps])  # (2P, D)
        reps = 2 * P
        starts = tile_starts(random_starts(M, rng), reps)  # common random numbers
        phys = PhysParams.nominal(M)
        if dr:
            # domain randomization: most scenarios randomized, some kept nominal
            rand = PhysParams.randomized(M, rng)
            k = rng.random(M) >= cfg.get("nominal_frac", 0.0)
            for f in ("mu", "mass", "tau", "noise", "delay"):
                getattr(phys, f)[k] = getattr(rand, f)[k]
        phys = phys.tile(reps)
        env = SumoEnv(reps * M, phys, seed=int(rng.integers(1 << 31)))
        me = MLPPolicy(np.repeat(members, M, axis=0))  # env e -> member e // M
        opp = sampler.build(M, reps)
        res = run_matches(env, me, opp, rng, starts=starts)
        fit = shaped_fitness(res, shape, cfg.get("tie_penalty", 0.0)).reshape(reps, M).mean(1)
        sampler.update(res["outcome"], M)
        theta = theta + adam.step(es_gradient(fit, eps, sigma) - cfg["l2"] * theta)

        log["gen"].append(g)
        log["fitness"].append(float(fit.mean()))
        log["win_rate"].append(float((res["outcome"] == 1).mean()))
        log["time"].append(time.time() - t0)
        if opp_mode == "pool" and g % cfg["snap_every"] == 0:
            pool.append(theta.copy())
            save(os.path.join(out_dir, f"{name}_gen{g:04d}.npz"), theta, f"{name}_gen{g:04d}", cfg)
        if g % cfg["eval_every"] == 0 or g == gens:
            ev = quick_eval(theta, eval_opps, n=cfg.get("eval_n", 256))
            log["eval"].append(ev)
            log["eval_gen"].append(g)
            if verbose:
                s = "  ".join(f"{k}: W{v[0]:.2f}/L{v[1]:.2f}/T{v[2]:.2f}" for k, v in ev.items())
                print(f"[{name}] gen {g:4d}  fit {fit.mean():+.3f}  pop-win {log['win_rate'][-1]:.2f}  "
                      f"| {s}  ({time.time() - t0:.0f}s)", flush=True)
    save(os.path.join(out_dir, f"{name}.npz"), theta, name, cfg)
    with open(os.path.join(out_dir, f"log_{name}.json"), "w") as f:
        json.dump(log, f)
    return theta, pool, log


def save(path, theta, name, cfg):
    np.savez(path, params=theta, name=name, config=json.dumps(cfg))


# ---------------------------------------------------------------- configs
FULL = {
    "L2": dict(gens=120, pairs=24, episodes=32, sigma=0.05, lr=0.03, l2=0.005, shape=0.3,
               eval_every=20, snap_every=10**9),
    "L3": dict(gens=200, pairs=24, episodes=32, sigma=0.05, lr=0.02, l2=0.005, shape=0.15,
               shape_decay_gens=150, eval_every=20, snap_every=20, hand_frac=0.3),
    "L4": dict(gens=150, pairs=24, episodes=32, sigma=0.05, lr=0.015, l2=0.005, shape=0.1,
               tie_penalty=0.25, nominal_frac=0.25, eval_every=20, snap_every=20, hand_frac=0.3),
}
QUICK = {
    "L2": dict(FULL["L2"], gens=30, pairs=12, episodes=24, eval_every=10),
    "L3": dict(FULL["L3"], gens=30, pairs=12, episodes=24, eval_every=10, snap_every=10, shape_decay_gens=30),
    "L4": dict(FULL["L4"], gens=20, pairs=12, episodes=24, eval_every=10, snap_every=10),
}


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--quick", action="store_true", help="small budget (~2-3 min) for a smoke run")
    ap.add_argument("--stage", choices=["all", "L2", "L3", "L4"], default="all")
    ap.add_argument("--out", default=None, help="output dir (default weights/ or weights_quick/ with --quick)")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--gens", type=int, default=None, help="override generations for the stage(s)")
    args = ap.parse_args()

    cfgs = {k: dict(v) for k, v in (QUICK if args.quick else FULL).items()}
    if args.gens:
        for c in cfgs.values():
            c["gens"] = args.gens
    out = args.out or os.path.join(HERE, "weights_quick" if args.quick else "weights")
    os.makedirs(out, exist_ok=True)
    stages = ["L2", "L3", "L4"] if args.stage == "all" else [args.stage]
    t0 = time.time()

    def load(n):
        return np.load(os.path.join(out, f"{n}.npz"))["params"]

    def load_pool(prefix):
        files = sorted(f for f in os.listdir(out) if f.startswith(prefix + "_gen") and f.endswith(".npz"))
        return [np.load(os.path.join(out, f))["params"] for f in files]

    for st in stages:
        print(f"=== stage {st}: {cfgs[st]}", flush=True)
        if st == "L2":
            theta0 = init_params(np.random.default_rng(args.seed))
            run_stage("L2", theta0, cfgs["L2"], out, opp_mode="charger", seed=args.seed + 2)
        elif st == "L3":
            pool = [load("L2")]
            run_stage("L3", load("L2"), cfgs["L3"], out, pool=pool, seed=args.seed + 3)
        else:
            # L4 continues L3 with the whole L3 self-play pool, now under domain randomization
            pool = [load("L2")] + load_pool("L3")
            run_stage("L4", load("L3"), cfgs["L4"], out, pool=pool, dr=True, seed=args.seed + 4)
        print(f"=== {st} done, total {time.time() - t0:.0f}s", flush=True)


if __name__ == "__main__":
    main()