#!/usr/bin/env python
"""Measure the TinyVLA training step on real hardware, synthetic data.

Builds the model out of the raw modules (SemanticPath / FlowMatchingExpert /
projections) so lerobot isn't needed, and reproduces _conditioning + the flow
loss exactly as modeling_tinyvla.py does. Sweeps: precision, dedup of the cam0
vision-tower pass, torch.compile, batch size.

Prints samples/s -> hours per 1M frames and per the 8.64M-sample C-scaled budget.
"""

from __future__ import annotations

import argparse
import time

import torch
import torch.nn as nn
import torch.nn.functional as F

from tinyvla.modules.expert import FlowMatchingExpert
from tinyvla.modules.semantic import SemanticPath


class Bench(nn.Module):
    """_conditioning + forward of TinyVLAPolicy, minus the lerobot wrapper."""

    def __init__(self, d=512, max_state=64, max_action=64, chunk=50,
                 lm_layers=12, readout=8, freeze_lm=False, freeze_vision=True,
                 model_name="Qwen/Qwen3.5-0.8B"):
        super().__init__()
        self.semantic = SemanticPath(
            model_name=model_name, num_layers=lm_layers, num_readout=readout,
            out_dim=d, image_size=256, freeze_lm=freeze_lm, freeze_vision=freeze_vision,
        )
        self.spatial_proj = nn.Linear(self.semantic.visual_hidden_size, d)
        self.camera_emb = nn.Embedding(3, d)
        self.state_proj = nn.Linear(max_state, d)
        self.embodiment_emb = nn.Embedding(16, d)
        self.expert = FlowMatchingExpert(action_dim=max_action, chunk_size=chunk, d_model=d)
        self.chunk, self.max_action = chunk, max_action

    def forward(self, batch, dedup=True):
        cam0, cam1 = batch["cam0"], batch["cam1"]
        f0 = self.semantic.encode_image(cam0)
        f1 = self.semantic.encode_image(cam1)
        if dedup:
            latent = self.semantic(None, batch["tok"], batch["mask"], image_embeds=f0)
        else:
            latent = self.semantic(cam0, batch["tok"], batch["mask"])
        spatial = torch.cat([
            self.spatial_proj(f0) + self.camera_emb.weight[0][None, None],
            self.spatial_proj(f1) + self.camera_emb.weight[1][None, None],
        ], dim=1)
        state_tok = self.state_proj(batch["state"])[:, None]
        emb = self.embodiment_emb(batch["emb_id"])[:, None]
        cond = torch.cat([latent, spatial, state_tok, emb], dim=1)

        actions = batch["actions"]
        b = actions.shape[0]
        t = torch.rand(b, device=actions.device) * 0.999 + 0.001
        noise = torch.randn_like(actions)
        x_t = t[:, None, None] * noise + (1 - t[:, None, None]) * actions
        pred = self.expert(x_t, t, cond)
        return F.mse_loss(pred, noise - actions)


def make_batch(b, dev, max_state=64, max_action=64, chunk=50, lm_len=48):
    return {
        "cam0": torch.rand(b, 3, 256, 256, device=dev),
        "cam1": torch.rand(b, 3, 256, 256, device=dev),
        "tok": torch.randint(1000, 5000, (b, lm_len), device=dev),
        "mask": torch.ones(b, lm_len, dtype=torch.bool, device=dev),
        "state": torch.randn(b, max_state, device=dev),
        "actions": torch.randn(b, chunk, max_action, device=dev),
        "emb_id": torch.randint(0, 16, (b,), device=dev),
    }


def run(model, opt, batch, dedup, steps, warmup, autocast):
    for i in range(warmup + steps):
        if i == warmup:
            torch.cuda.synchronize()
            t0 = time.time()
        opt.zero_grad(set_to_none=True)
        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=autocast):
            loss = model(batch, dedup=dedup)
        loss.backward()
        opt.step()
    torch.cuda.synchronize()
    return (time.time() - t0) / steps


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model-name", default="Qwen/Qwen3.5-0.8B")
    ap.add_argument("--steps", type=int, default=30)
    ap.add_argument("--warmup", type=int, default=10)
    ap.add_argument("--batches", type=int, nargs="+", default=[48, 96, 144, 192])
    ap.add_argument("--compile", action="store_true")
    ap.add_argument("--freeze-lm", action="store_true")
    args = ap.parse_args()

    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
    dev = torch.device("cuda")

    model = Bench(freeze_lm=args.freeze_lm, model_name=args.model_name).to(dev)
    n_train = sum(p.numel() for p in model.parameters() if p.requires_grad)
    n_all = sum(p.numel() for p in model.parameters())
    print(f"params: {n_all/1e6:.1f}M total, {n_train/1e6:.1f}M trainable "
          f"(freeze_lm={args.freeze_lm}, vision frozen)")

    opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad],
                            lr=1e-4, betas=(0.9, 0.95), weight_decay=1e-10, fused=True)
    if args.compile:
        model.expert = torch.compile(model.expert, dynamic=False)

    BUDGET = 8_640_000  # C-scaled: 60k steps x eff batch 144
    print(f"\n{'batch':>6} {'dedup':>6} {'ms/step':>9} {'samples/s':>10} "
          f"{'ч на 1M кадров':>15} {'ч на 8.64M':>11} {'VRAM GB':>8}")
    for b in args.batches:
        batch = make_batch(b, dev)
        for dedup in (True, False):
            torch.cuda.reset_peak_memory_stats()
            try:
                dt = run(model, opt, batch, dedup, args.steps, args.warmup, autocast=True)
            except torch.cuda.OutOfMemoryError:
                print(f"{b:>6} {str(dedup):>6}      OOM")
                torch.cuda.empty_cache()
                continue
            sps = b / dt
            mem = torch.cuda.max_memory_allocated() / 1e9
            print(f"{b:>6} {str(dedup):>6} {dt*1000:>9.1f} {sps:>10.0f} "
                  f"{1e6/sps/3600:>15.2f} {BUDGET/sps/3600:>11.2f} {mem:>8.1f}", flush=True)


if __name__ == "__main__":
    main()