"""Multi-GPU (torchrun) from-scratch training of FACT on a whole T-Rex dataset.

Mirrors the operational contract of the T-Rex Track-Force trainer:
* W&B scalar logging (`dynamics` -> `state` here) and a periodic monitoring
  video on a fixed episode (``eval/predicted_video``) that overlays the
  sampled action / future-state chunks on the ground truth.  FACT predicts
  future *states* instead of future frames, so the video shows the RGB
  observation next to the predicted-vs-GT trajectories.
* checkpoints every ``--save-steps`` under ``OUTPUT_DIR/checkpoint-N``
  (model, optimizer, trainer state) with ``--save-total-limit`` rotation and
  ``--resume auto``.
"""

from __future__ import annotations

import argparse
import datetime
import json
import math
import os
import random
import shutil
import time
from pathlib import Path

import numpy as np
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler

from fact import FactConfig, FactPolicy, count_parameters
from fact.data import ACTION_RATE_HZ, FactMultiEpisodeDataset, collate


# ----------------------------------------------------------------------------
def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser()
    p.add_argument("--dataset-root", required=True)
    p.add_argument("--cache-dir", required=True)
    p.add_argument("--output-dir", required=True)
    p.add_argument("--config", required=True, help="FactConfig JSON file")
    p.add_argument("--config-overrides", default="", help="JSON dict of FactConfig overrides")
    p.add_argument("--episodes", default="", help="comma list of episode indices (default all)")
    p.add_argument("--blocks", type=int, default=4)
    p.add_argument("--max-steps", type=int, default=100000)
    p.add_argument("--per-device-batch", type=int, default=8)
    p.add_argument("--lr", type=float, default=3e-4)
    p.add_argument("--min-lr-ratio", type=float, default=0.1)
    p.add_argument("--warmup-steps", type=int, default=2000)
    p.add_argument("--weight-decay", type=float, default=0.01)
    p.add_argument("--clip", type=float, default=1.0)
    p.add_argument("--force-lr-mult", type=float, default=0.3, help="lr multiplier for the force stage (it diverged late in training at the full lr)")
    p.add_argument("--skip-grad-norm", type=float, default=20.0, help="skip the optimizer step when the pre-clip grad norm exceeds this (or is non-finite)")
    p.add_argument("--force-fp32", type=int, default=1, help="run the force stage outside bf16 autocast")
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--workers", type=int, default=4)
    p.add_argument("--bf16", type=int, default=1)
    p.add_argument("--log-steps", type=int, default=20)
    p.add_argument("--save-steps", type=int, default=10000)
    p.add_argument("--save-total-limit", type=int, default=5)
    p.add_argument("--resume", default="auto", help="'auto', 'none' or a checkpoint dir")
    # W&B + monitoring video
    p.add_argument("--report-to", default="wandb", choices=["wandb", "none"])
    p.add_argument("--wandb-project", default="fact")
    p.add_argument("--wandb-run-name", default="")
    p.add_argument("--video-steps", type=int, default=1000, help="monitoring video every N steps")
    p.add_argument("--video-episode", type=int, default=0, help="dataset episode index for the video")
    p.add_argument("--video-num-chunks", type=int, default=4)
    p.add_argument("--video-save-local", type=int, default=1)
    return p.parse_args()


def load_config(args: argparse.Namespace) -> FactConfig:
    cfg = FactConfig.from_json(Path(args.config).read_text())
    if args.config_overrides:
        for key, value in json.loads(args.config_overrides).items():
            if not hasattr(cfg, key):
                raise KeyError(key)
            setattr(cfg, key, value)
    return cfg


def lr_at(step: int, args: argparse.Namespace) -> float:
    if step < args.warmup_steps:
        return args.lr * (step + 1) / args.warmup_steps
    progress = (step - args.warmup_steps) / max(1, args.max_steps - args.warmup_steps)
    cosine = 0.5 * (1 + math.cos(math.pi * min(1.0, progress)))
    return args.lr * (args.min_lr_ratio + (1 - args.min_lr_ratio) * cosine)


def is_main() -> bool:
    return (not dist.is_initialized()) or dist.get_rank() == 0


def log(msg: str) -> None:
    if is_main():
        print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)


# ----------------------------------------------------------------------------
def save_checkpoint(model: FactPolicy, optimizer, step: int, args, out: Path, extra: dict) -> None:
    path = out / f"checkpoint-{step}"
    tmp = out / f".checkpoint-{step}.tmp"
    if tmp.exists():
        shutil.rmtree(tmp)
    model.save(tmp, extra=extra)
    torch.save(optimizer.state_dict(), tmp / "optimizer.pt")
    (tmp / "trainer_state.json").write_text(json.dumps({"step": step, "args": vars(args)}, indent=1))
    if path.exists():
        shutil.rmtree(path)
    tmp.rename(path)
    (out / "latest").write_text(path.name)
    ckpts = sorted(
        (p for p in out.glob("checkpoint-*") if p.is_dir()),
        key=lambda p: int(p.name.split("-")[1]),
    )
    while len(ckpts) > args.save_total_limit:
        shutil.rmtree(ckpts.pop(0))


def find_resume(args, out: Path) -> Path | None:
    if args.resume == "none":
        return None
    if args.resume != "auto":
        return Path(args.resume)
    latest = out / "latest"
    if latest.exists():
        candidate = out / latest.read_text().strip()
        if (candidate / "fact_model.pt").exists():
            return candidate
    ckpts = sorted((p for p in out.glob("checkpoint-*") if (p / "fact_model.pt").exists()),
                   key=lambda p: int(p.name.split("-")[1]))
    return ckpts[-1] if ckpts else None


# ----------------------------------------------------------------------------
class VideoMonitor:
    """Sample consecutive chunks of one episode and render pred-vs-GT video."""

    def __init__(self, dataset: FactMultiEpisodeDataset, args, out: Path) -> None:
        self.episode = dataset.episodes[min(args.video_episode, len(dataset.episodes) - 1)]
        anchors = self.episode.valid_anchor_times(1)
        stride = 16 / ACTION_RATE_HZ
        self.anchors = [anchors[0] + i * stride for i in range(args.video_num_chunks)]
        self.anchors = [a for a in self.anchors if a <= anchors[-1] + 1e-9]
        self.out = out / "eval_videos"
        self.save_local = bool(args.video_save_local)
        self._frames: np.ndarray | None = None

    def _head_frames(self) -> np.ndarray:
        if self._frames is None:
            import decord

            ep = self.episode
            path = (ep.root / "videos" / f"chunk-{ep.episode_index // 1000:03d}"
                    / "observation.images.head_left" / f"episode_{ep.episode_index:06d}.mp4")
            reader = decord.VideoReader(str(path), num_threads=2)
            self._frames = reader.get_batch(list(range(len(reader)))).asnumpy()
        return self._frames

    @torch.no_grad()
    def run(self, model: FactPolicy, step: int) -> dict[str, float | str]:
        """Rank-0-only sampling: the tactile VQ's perplexity all-reduce must be
        disabled here or the other ranks would never join the collective."""

        from fact import tactile_vq as _tq

        original = _tq._is_distributed
        _tq._is_distributed = lambda: False
        try:
            return self._run(model, step)
        finally:
            _tq._is_distributed = original

    def _run(self, model: FactPolicy, step: int) -> dict[str, float | str]:
        import cv2
        import matplotlib

        matplotlib.use("Agg")
        import matplotlib.pyplot as plt

        model.eval()
        device = model.device
        frames_rgb = self._head_frames()
        video_frames: list[np.ndarray] = []
        metrics: dict[str, list[float]] = {}
        groups = {
            "L arm xyz": [0, 1, 2],
            "R arm xyz": [31, 32, 33],
            "L hand (5 joints)": [9, 10, 11, 12, 13],
            "R hand (5 joints)": [40, 41, 42, 43, 44],
        }
        for chunk_idx, anchor in enumerate(self.anchors):
            sample = self.episode.build(anchor, blocks=1)
            batch = {
                k: torch.as_tensor(v)[None].to(device)
                for k, v in sample.items()
                if k not in ("action_delta_raw", "reference_state", "obs_frame_index")
            }
            gt_action = sample["action"][0, :, :62]
            gt_state = sample["future_state"][0, :, :62]
            result = model.sample(batch, seed=500 + chunk_idx, mode="cascade")
            preds = {
                "cascade": result["action_pred"][0, :, :62].float().cpu().numpy(),
                "coarse_only": result["coarse_action_at_split"][0, :, :62].float().cpu().numpy(),
            }
            for mode, pred in preds.items():
                metrics.setdefault(f"eval/openloop_{mode}_mse", []).append(float(((pred - gt_action) ** 2).mean()))
            pred_state = None
            if "future_state_pred" in result:
                pred_state = result["future_state_pred"][0, :, :62].float().cpu().numpy()
                metrics.setdefault("eval/future_state_mse", []).append(float(((pred_state - gt_state) ** 2).mean()))
            # one figure per chunk: action (cascade vs coarse vs GT) and future state
            fig, axes = plt.subplots(2, 4, figsize=(12, 4.2), dpi=80)
            t = np.arange(16)
            for col, (name, dims) in enumerate(groups.items()):
                ax = axes[0, col]
                for d in dims:
                    ax.plot(t, gt_action[:, d], color="k", lw=1.2, alpha=0.8)
                    ax.plot(t, preds["cascade"][:, d], color="tab:red", lw=1.0)
                    ax.plot(t, preds["coarse_only"][:, d], color="tab:blue", lw=0.8, ls="--", alpha=0.6)
                ax.set_title(f"action {name}", fontsize=8)
                ax.set_ylim(-1.1, 1.1)
                ax.tick_params(labelsize=6)
                ax = axes[1, col]
                for d in dims:
                    ax.plot(t, gt_state[:, d], color="k", lw=1.2, alpha=0.8)
                    if pred_state is not None:
                        ax.plot(t, pred_state[:, d], color="tab:green", lw=1.0)
                ax.set_title(f"future state {name}", fontsize=8)
                ax.set_ylim(-1.1, 1.1)
                ax.tick_params(labelsize=6)
            fig.suptitle(
                f"step {step} chunk {chunk_idx} t={anchor:.2f}s  black=GT red=cascade blue=coarse(τ=.4) green=future-state",
                fontsize=8,
            )
            fig.tight_layout()
            fig.canvas.draw()
            plot = np.asarray(fig.canvas.buffer_rgba())[..., :3].copy()
            plt.close(fig)
            # 16 video frames per chunk: RGB frame at each action timestamp + plot with a step marker
            step_frames = [
                int(i) for i in np.clip(np.round((anchor - self.episode.timestamps[0]) * 30 + np.arange(16) * 1.5), 0, len(frames_rgb) - 1)
            ]
            for i, fi in enumerate(step_frames):
                rgb = cv2.resize(frames_rgb[fi], (plot.shape[0] * 16 // 9, plot.shape[0]))
                marker = plot.copy()
                x = int(marker.shape[1] * (0.07 + 0.86 * i / 15))
                marker[:, max(0, x - 1) : x + 1] = (255, 140, 0)
                video_frames.append(np.concatenate([rgb, marker], axis=1))
        model.train()
        result: dict[str, float | str] = {k: float(np.mean(v)) for k, v in metrics.items()}
        if video_frames:
            self.out.mkdir(parents=True, exist_ok=True)
            path = self.out / f"train_step_{step:06d}.mp4"
            h, w = video_frames[0].shape[:2]
            writer = cv2.VideoWriter(str(path), cv2.VideoWriter_fourcc(*"mp4v"), 20, (w, h))
            for frame in video_frames:
                writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
            writer.release()
            result["video_path"] = str(path)
        return result


# ----------------------------------------------------------------------------
def main() -> None:
    args = parse_args()
    # Long collective timeout: rank 0 renders the monitoring video while the
    # other ranks wait, and first-time PPU kernel compilation can take minutes.
    dist.init_process_group(backend="nccl", timeout=datetime.timedelta(hours=3))
    rank, world = dist.get_rank(), dist.get_world_size()
    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)
    torch.manual_seed(args.seed + rank)
    np.random.seed(args.seed + rank)
    random.seed(args.seed + rank)
    out = Path(args.output_dir)
    if is_main():
        out.mkdir(parents=True, exist_ok=True)
    dist.barrier()

    cfg = load_config(args)
    episodes = [int(e) for e in args.episodes.split(",")] if args.episodes else None
    dataset = FactMultiEpisodeDataset(
        args.dataset_root, args.cache_dir, episodes=episodes, blocks=args.blocks, use_text=cfg.use_text
    )
    if cfg.use_text:
        cache_tokens = dataset.episodes[0].text_cache.max_tokens
        if cache_tokens != cfg.max_text_tokens:
            raise ValueError(f"text cache has {cache_tokens} tokens, config expects {cfg.max_text_tokens}")
    log(f"dataset: {len(dataset.episodes)} episodes, {len(dataset)} anchors (K={args.blocks}), "
        f"global batch {args.per_device_batch * world} x {args.blocks} blocks")

    model = FactPolicy(cfg).to(device)
    n_total = count_parameters(model)
    log(f"FACT params: total {n_total/1e6:.1f}M coarse {count_parameters(model.coarse)/1e6:.1f}M "
        f"force {(count_parameters(model.force) if model.force else 0)/1e6:.1f}M (all trainable, from scratch)")
    groups: dict[tuple[bool, bool], list] = {}
    for name, param in model.named_parameters():
        is_force = name.startswith("force.")
        no_decay = param.ndim < 2 or "embedding" in name or "modulation" in name
        groups.setdefault((is_force, no_decay), []).append(param)
    optimizer = torch.optim.AdamW(
        [{"params": params, "weight_decay": 0.0 if nd else args.weight_decay,
          "lr_mult": args.force_lr_mult if is_force else 1.0}
         for (is_force, nd), params in groups.items()],
        lr=args.lr, betas=(0.9, 0.95),
    )
    if args.force_fp32 and model.force is not None:
        force_forward = model.force.forward

        def force_forward_fp32(*a, **k):
            with torch.autocast("cuda", enabled=False):
                a = [x.float() if torch.is_tensor(x) and x.is_floating_point() else x for x in a]
                k = {n: (x.float() if torch.is_tensor(x) and x.is_floating_point() else x) for n, x in k.items()}
                return force_forward(*a, **k)

        model.force.forward = force_forward_fp32
    step = 0
    resume = find_resume(args, out)
    if resume is not None:
        state = torch.load(resume / "fact_model.pt", map_location="cpu")
        model.load_state_dict(state, strict=True)
        opt_path = resume / "optimizer.pt"
        if opt_path.exists():
            optimizer.load_state_dict(torch.load(opt_path, map_location="cpu"))
        step = int(json.loads((resume / "trainer_state.json").read_text())["step"])
        log(f"resumed from {resume} at step {step}")
    ddp = DDP(model, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=False)

    sampler = DistributedSampler(dataset, num_replicas=world, rank=rank, shuffle=True, seed=args.seed, drop_last=True)
    loader = DataLoader(
        dataset, batch_size=args.per_device_batch, sampler=sampler, num_workers=args.workers,
        collate_fn=collate, drop_last=True, persistent_workers=args.workers > 0, pin_memory=False,
    )
    steps_per_epoch = len(loader)

    wandb_run = None
    if is_main():
        (out / "fact_config.json").write_text(cfg.to_json())
        (out / "train_args.json").write_text(json.dumps(vars(args), indent=1))
        if args.report_to == "wandb":
            import wandb

            wandb_run = wandb.init(
                project=args.wandb_project,
                name=args.wandb_run_name or out.name,
                dir=str(out),
                config={**vars(args), "fact_config": json.loads(cfg.to_json()), "params_total": n_total,
                        "world_size": world, "steps_per_epoch": steps_per_epoch},
                resume="allow",
                id=(json.loads((out / "wandb_config.json").read_text())["run_id"]
                    if (out / "wandb_config.json").exists() else None),
            )
            (out / "wandb_config.json").write_text(json.dumps({"project": args.wandb_project, "run_id": wandb_run.id}))
        monitor = VideoMonitor(dataset, args, out)
        log_file = (out / "loss_log.jsonl").open("a")

    # Warm up the monitor once before training (compiles the sampling kernels
    # while every rank waits, instead of inside the DDP step loop).
    if is_main() and args.video_steps > 0:
        try:
            t0 = time.time()
            warm = monitor.run(ddp.module, step)
            if "video_path" in warm and not args.video_save_local:
                os.remove(warm["video_path"])
            log(f"monitor warm-up done in {time.time() - t0:.0f}s: "
                + " ".join(f"{k.split('/')[-1]} {v:.4f}" for k, v in warm.items() if k != "video_path"))
        except Exception as exc:
            log(f"monitor warm-up failed: {exc!r}")
    dist.barrier()

    ddp.train()
    epoch = step // max(1, steps_per_epoch)
    running: dict[str, float] = {}
    start = time.time()
    done_steps_at_start = step
    while step < args.max_steps:
        sampler.set_epoch(epoch)
        for batch in loader:
            if step >= args.max_steps:
                break
            lr = lr_at(step, args)
            for group in optimizer.param_groups:
                group["lr"] = lr * group.get("lr_mult", 1.0)
            with torch.autocast("cuda", dtype=torch.bfloat16, enabled=bool(args.bf16)):
                losses = ddp(batch)  # FactPolicy.forward -> forward_core; DDP syncs grads
            loss = losses["loss"]
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            grad_norm = torch.nn.utils.clip_grad_norm_(ddp.parameters(), args.clip)
            skipped = (not torch.isfinite(grad_norm)) or float(grad_norm) > args.skip_grad_norm
            if skipped:
                optimizer.zero_grad(set_to_none=True)  # identical decision on every rank (synced grads)
            else:
                optimizer.step()
            step += 1

            if is_main():
                record = {k: float(v) for k, v in losses.items()}
                record.update({"step": step, "lr": lr, "grad_norm": float(grad_norm), "epoch": epoch, "skipped": int(skipped)})
                log_file.write(json.dumps(record) + "\n")
                for key, value in record.items():
                    running[key] = running.get(key, 0.0) + value
                if step % args.log_steps == 0:
                    avg = {k: v / args.log_steps for k, v in running.items()}
                    running = {}
                    elapsed = time.time() - start
                    rate = (step - done_steps_at_start) / max(elapsed, 1e-6)
                    log(f"step {step}/{args.max_steps} ep {epoch} loss {avg['loss']:.4f} act {avg['action_loss']:.4f} "
                        f"state {avg['state_loss']:.4f} force {avg['force_loss']:.4f} vq {avg['vq_loss']:.4f} "
                        f"gn {avg['grad_norm']:.2f} skip {avg['skipped']:.2f} lr {lr:.2e} {1/rate:.2f}s/it eta {(args.max_steps-step)/rate/3600:.1f}h")
                    log_file.flush()
                    if wandb_run is not None:
                        wandb_run.log({f"train/{k}": v for k, v in avg.items() if k not in ("step", "epoch")}
                                      | {"train/epoch": epoch, "train/sec_per_it": 1 / rate}, step=step)
                if args.video_steps > 0 and step % args.video_steps == 0:
                    try:
                        result = monitor.run(ddp.module, step)
                        scalars = {k: v for k, v in result.items() if k != "video_path"}
                        log("monitor " + " ".join(f"{k.split('/')[-1]} {v:.4f}" for k, v in scalars.items()))
                        if wandb_run is not None:
                            import wandb

                            payload = dict(scalars)
                            if "video_path" in result:
                                payload["eval/predicted_video"] = wandb.Video(result["video_path"], fps=20, format="mp4")
                            wandb_run.log(payload, step=step)
                        if not args.video_save_local and "video_path" in result:
                            os.remove(result["video_path"])
                    except Exception as exc:  # monitoring must never kill training
                        log(f"monitor failed at step {step}: {exc!r}")
            if step % args.save_steps == 0 or step == args.max_steps:
                if is_main():
                    save_checkpoint(ddp.module, optimizer, step, args, out,
                                    extra={"step": step, "params_total": n_total, "epoch": epoch})
                    log(f"saved checkpoint-{step}")
                dist.barrier()
        epoch += 1

    if is_main():
        ddp.module.save(out / "final", extra={"step": step, "params_total": n_total})
        log_file.close()
        if wandb_run is not None:
            wandb_run.finish()
        log(f"done: {step} steps in {(time.time()-start)/3600:.2f} h -> {out/'final'}")
    dist.barrier()
    dist.destroy_process_group()


if __name__ == "__main__":
    main()