#!/usr/bin/env python
"""Cost of d(loss)/d(action) -- the inner loop of gradient-based planning.

Three regimes:
  dino-wm full    encode context frames, predict, backprop to the action chunk
  dino-wm cached  context features precomputed once (what a planner actually does across
                  CEM/GD iterations, since the context is fixed), predict + backprop only
  mlp 3x256       a 3-layer 256-wide state-space dynamics model, (state, action) -> next state

    policy_learning/lerobot/.venv/bin/python dynamics_model/bench_action_grad.py
"""

import time

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModel

from dino_dynamics import DINO_ID, DinoDynamics
from sd3_dynamics import HORIZON

V, ACT, ST, HIST, SZ = 2, 7, 9, 2, 224
BATCHES = (1, 32, 256, 1024)
REPS = 10


def timeit(fn, reps=REPS, warmup=3):
    for _ in range(warmup):
        fn()
    torch.cuda.synchronize()
    t0 = time.time()
    for _ in range(reps):
        fn()
    torch.cuda.synchronize()
    return (time.time() - t0) / reps * 1000  # ms


def main():
    dev = "cuda"
    torch.manual_seed(0)
    dino = AutoModel.from_pretrained(DINO_ID, dtype=torch.float32)
    model = DinoDynamics(dino, None, V, ACT, ST, history=HIST).to(dev).eval()
    for p in model.parameters():
        p.requires_grad_(False)  # planning optimises actions, not weights

    mlp = nn.Sequential(
        nn.Linear(ST + HORIZON * ACT, 256), nn.ReLU(),
        nn.Linear(256, 256), nn.ReLU(),
        nn.Linear(256, 256), nn.ReLU(),
        nn.Linear(256, ST),
    ).to(dev).eval()
    for p in mlp.parameters():
        p.requires_grad_(False)
    print(f"dino-wm predictor params: {sum(p.numel() for p in model.parameters()) / 1e6:.1f}M "
          f"(incl. frozen DINOv2)   mlp params: {sum(p.numel() for p in mlp.parameters()) / 1e3:.0f}k")

    print(f"\n{'batch':>6} {'dino full':>12} {'dino cached':>13} {'mlp 3x256':>11} "
          f"{'cached/mlp':>11} {'full/cached':>12}")
    for b in BATCHES:
        ctx = torch.rand(b, HIST, V, 3, SZ, SZ, device=dev)
        state = torch.randn(b, ST, device=dev)
        goal = torch.randn(b, V, model.n_patches, model.dim, device=dev)
        goal_s = torch.randn(b, ST, device=dev)

        def dino_full():
            act = torch.randn(b, HORIZON, ACT, device=dev, requires_grad=True)
            loss = F.mse_loss(model({"context": ctx, "action": act, "state": state}), goal)
            return torch.autograd.grad(loss, act)[0]

        # cached: encode once, reuse the features across every planner iteration
        with torch.no_grad():
            z = model.encode_views(ctx)

        def dino_cached():
            act = torch.randn(b, HORIZON, ACT, device=dev, requires_grad=True)
            loss = F.mse_loss(model.predict_from_features(z, act, state), goal)
            return torch.autograd.grad(loss, act)[0]

        def mlp_grad():
            act = torch.randn(b, HORIZON, ACT, device=dev, requires_grad=True)
            loss = F.mse_loss(mlp(torch.cat([state, act.flatten(1)], 1)), goal_s)
            return torch.autograd.grad(loss, act)[0]

        def safe(fn):
            try:
                return timeit(fn)
            except torch.OutOfMemoryError:
                torch.cuda.empty_cache()
                return float("nan")

        t_full, t_cached, t_mlp = safe(dino_full), safe(dino_cached), safe(mlp_grad)
        fmt = lambda x, w, p: ("OOM" if x != x else f"{x:.{p}f}ms").rjust(w)
        ratio = lambda x, y: "-" if (x != x or y != y) else f"{x / y:.0f}x"
        print(f"{b:>6} {fmt(t_full, 12, 2)} {fmt(t_cached, 13, 2)} {fmt(t_mlp, 11, 3)} "
              f"{ratio(t_cached, t_mlp):>11} {ratio(t_full, t_cached):>12}")
        del ctx, z
        torch.cuda.empty_cache()


if __name__ == "__main__":
    main()