"""Goal-conditioned behavior cloning (compact BC baseline) on OGBench state data.

Trains a small MLP policy pi(obs, goal) -> action with MSE loss and hindsight
future goals sampled within the same trajectory. Trajectory-level nested demo
subsets (10 subset of 100 subset of 1000 ...) via a fixed permutation seed.
"""
import argparse, os, json, time
import numpy as np, torch, torch.nn as nn
import ogbench
from common import get_traj_bounds, nested_traj_order, seed_everything


class MLPPolicy(nn.Module):
    def __init__(self, in_dim, act_dim, hidden=512, layers=3):
        super().__init__()
        mods, d = [], in_dim
        for _ in range(layers):
            mods += [nn.Linear(d, hidden), nn.ReLU()]
            d = hidden
        mods += [nn.Linear(d, act_dim)]
        self.net = nn.Sequential(*mods)

    def forward(self, obs, goal):
        return torch.tanh(self.net(torch.cat([obs, goal], dim=-1)))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dataset", default="cube-single-play-v0")
    ap.add_argument("--demos", type=int, default=1000)
    ap.add_argument("--steps", type=int, default=100000)
    ap.add_argument("--batch", type=int, default=1024)
    ap.add_argument("--hidden", type=int, default=512)
    ap.add_argument("--layers", type=int, default=3)
    ap.add_argument("--lr", type=float, default=3e-4)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--split_seed", type=int, default=0)
    ap.add_argument("--outdir", default="runs")
    args = ap.parse_args()

    seed_everything(args.seed)
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    t0 = time.time()

    env, train_ds, _ = ogbench.make_env_and_datasets(args.dataset, compact_dataset=False)
    obs = train_ds["observations"].astype(np.float32)
    act = train_ds["actions"].astype(np.float32)
    term = train_ds["terminals"].astype(np.float32)
    bounds = get_traj_bounds(term)
    order = nested_traj_order(len(bounds), seed=args.split_seed)
    sel = np.sort(order[:args.demos])           # selected trajectory ids (nested by prefix)

    # per selected transition: global index gi and its trajectory end te (for future-goal sampling)
    gi, te = [], []
    for ti in sel:
        s, e = bounds[int(ti)]
        gi.append(np.arange(s, e)); te.append(np.full(e - s, e))
    gi = np.concatenate(gi); te = np.concatenate(te)

    # normalization stats from the SELECTED observations only
    obs_mean = obs[gi].mean(0); obs_std = obs[gi].std(0) + 1e-6

    # move to GPU
    obs_t = torch.tensor((obs - obs_mean) / obs_std, device=dev)   # full, normalized
    act_t = torch.tensor(act, device=dev)
    gi_t = torch.tensor(gi, device=dev, dtype=torch.long)
    te_t = torch.tensor(te, device=dev, dtype=torch.long)
    N = gi_t.numel()

    model = MLPPolicy(obs.shape[1] * 2, act.shape[1], args.hidden, args.layers).to(dev)
    opt = torch.optim.Adam(model.parameters(), lr=args.lr)

    print(f"[train] dataset={args.dataset} demos={args.demos} trajs_selected={len(sel)} "
          f"transitions={N} obs_dim={obs.shape[1]} act_dim={act.shape[1]} device={dev}")

    model.train()
    for step in range(1, args.steps + 1):
        p = torch.randint(0, N, (args.batch,), device=dev)
        cur = gi_t[p]                                   # current global idx
        end = te_t[p]                                   # trajectory end (exclusive)
        u = torch.rand(args.batch, device=dev)
        gj = (cur + (u * (end - cur).float()).long()).clamp(max=obs_t.shape[0] - 1)  # future goal idx
        o = obs_t[cur]; g = obs_t[gj]; a = act_t[cur]
        pred = model(o, g)
        loss = ((pred - a) ** 2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
        if step % 5000 == 0 or step == 1:
            print(f"  step {step:6d}/{args.steps}  mse={loss.item():.4f}  "
                  f"elapsed={time.time()-t0:.0f}s", flush=True)

    run = f"bc_{args.dataset}_demos{args.demos}_seed{args.seed}"
    outdir = os.path.join(args.outdir, run); os.makedirs(outdir, exist_ok=True)
    torch.save({
        "model": model.state_dict(),
        "obs_mean": obs_mean, "obs_std": obs_std,
        "config": vars(args),
        "obs_dim": int(obs.shape[1]), "act_dim": int(act.shape[1]),
        "hidden": args.hidden, "layers": args.layers,
    }, os.path.join(outdir, "ckpt.pt"))
    # immutable split manifest (trajectory ids used) -> reproducibility deliverable
    with open(os.path.join(outdir, "manifest.json"), "w") as f:
        json.dump({"dataset": args.dataset, "demos": args.demos,
                   "split_seed": args.split_seed, "trajectory_ids": sel.tolist(),
                   "n_transitions": int(N)}, f)
    print(f"[train] saved -> {outdir}  ({time.time()-t0:.0f}s total)")


if __name__ == "__main__":
    main()