"""Behaviour-cloning warm-start from the classical controller.

Collects (observation, classical action) pairs by running the classical
PD + potential-field policy on GENERAL layouts only (trap=False, where
the teacher succeeds 10/10 in sanity checks). Trap layouts are excluded
deliberately: imitating the teacher there would teach the stall.
Trains only the actor (mean head) of ActorCritic with MSE on the
tanh-bounded action. The critic is left for PPO fine-tuning.
Outputs: bc_policy.pt, bc.log via stdout.
"""
import argparse

import numpy as np
import torch
import torch.nn as nn

from sim import AMAX, BatchedSim, generate_layout
from policies import ActorCritic, classical_action


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--episodes", type=int, default=400)
    ap.add_argument("--n-agents", type=int, default=4)
    ap.add_argument("--epochs", type=int, default=30)
    ap.add_argument("--seed", type=int, default=123)
    ap.add_argument("--out", default="bc_policy.pt")
    args = ap.parse_args()
    torch.manual_seed(args.seed)
    torch.set_num_threads(2)
    rng = np.random.default_rng(args.seed)
    sim = BatchedSim(16, args.n_agents, seed=args.seed)
    obs_list, act_list = [], []
    layouts = [generate_layout(rng, args.n_agents, trap=False) for _ in range(16)]
    obs = sim.reset(layouts)
    done_eps = 0
    while done_eps < args.episodes:
        act = classical_action(obs)
        obs_list.append(obs.reshape(-1, obs.shape[-1]).copy())
        act_list.append(act.reshape(-1, 2).copy())
        obs, _, ep_done, _ = sim.step(act)
        fin = np.where(ep_done)[0]
        if len(fin):
            done_eps += len(fin)
            obs = sim.reset_envs(fin.tolist(),
                                 [generate_layout(rng, args.n_agents, trap=False) for _ in fin])
    O = torch.as_tensor(np.concatenate(obs_list, axis=0))
    A = torch.as_tensor(np.concatenate(act_list, axis=0))
    print(f"BC dataset: {O.shape[0]} samples, {done_eps} episodes", flush=True)
    net = ActorCritic()
    opt = torch.optim.Adam(list(net.actor.parameters()) + list(net.mean.parameters()), lr=1e-3)
    n = O.shape[0]
    for ep in range(args.epochs):
        perm = torch.randperm(n)
        tot = 0.0
        for s in range(0, n, 1024):
            b = perm[s:s + 1024]
            mean, _ = net.forward(O[b])
            pred = torch.tanh(mean) * AMAX
            loss = nn.functional.mse_loss(pred, A[b])
            opt.zero_grad(); loss.backward(); opt.step()
            tot += float(loss) * len(b)
        print(f"bc epoch {ep + 1}/{args.epochs} mse={tot / n:.4f}", flush=True)
    # action agreement on a held-out-style check: mean |pred - teacher|
    with torch.no_grad():
        mean, _ = net.forward(O[:4096])
        pred = torch.tanh(mean) * AMAX
        print(f"BC final: mean |action error| = {float((pred - A[:4096]).abs().mean()):.3f} "
              f"(accel units)", flush=True)
    torch.save(net.state_dict(), args.out)
    print(f"saved {args.out}")


if __name__ == "__main__":
    main()