[tab] utils/lam_session.py (python)
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Callable, Dict, Optional, Tuple, cast

import einops
import jax
import jax.numpy as jnp
import flax.nnx as nnx
import optax
import orbax.checkpoint as ocp
import numpy as np
import grain

from models.lam import LatentActionModel
from utils.dataloader import get_dataloader
from utils.lr_utils import get_lr_schedule


@dataclass
class LamConfig:
    # Experiment / data
    seed: int = 0
    seq_len: int = 16
    image_channels: int = 3
    image_height: int = 90
    image_width: int = 160
    batch_size: int = 36
    data_dir: str = ""
    ckpt_dir: str = ""

    # Optimization
    vq_beta: float = 0.25
    init_lr: float = 0.0
    max_lr: float = 3e-5
    decay_end: float = 0.0
    wsd_decay_steps: int = 10000
    warmup_steps: int = 5000
    lr_schedule: str = "wsd"  # supported options: wsd, cos
    vq_reset_thresh: int = 50

    # LAM model
    model_dim: int = 512
    ffn_dim: int = 2048
    latent_dim: int = 32
    num_latents: int = 6
    patch_size: int = 16
    num_blocks: int = 4
    num_heads: int = 8
    dropout: float = 0.0
    codebook_dropout: float = 0.0
    param_dtype: jnp.dtype = jnp.float32
    dtype: jnp.dtype = jnp.bfloat16
    use_flash_attention: bool = True


class LamSession:
    """Utility for interactive work with LAM in notebooks.

    Provides:
    - Model/optimizer/dataloader construction
    - One-step train with codebook reset logic
    - Forward pass and optional sowing controls
    - Checkpoint restore/save matching the training script
    """

    def __init__(
        self,
        config: LamConfig,
        build_dataloader: bool = True,
        build_checkpoint_manager: bool = True,
    ) -> None:
        self.config = config
        self.rng = jax.random.key(config.seed)

        # Model
        self.rng, _rng = jax.random.split(self.rng)
        rngs = nnx.Rngs(_rng)
        self.model = LatentActionModel(
            in_dim=config.image_channels,
            model_dim=config.model_dim,
            ffn_dim=config.ffn_dim,
            latent_dim=config.latent_dim,
            num_latents=config.num_latents,
            patch_size=config.patch_size,
            num_blocks=config.num_blocks,
            num_heads=config.num_heads,
            dropout=config.dropout,
            codebook_dropout=config.codebook_dropout,
            param_dtype=config.param_dtype,
            dtype=config.dtype,
            use_flash_attention=config.use_flash_attention,
            rngs=rngs,
        )

        # Optimizer & schedule
        self.lr_schedule: Callable[[int], float] = get_lr_schedule(
            config.lr_schedule,
            config.init_lr,
            config.max_lr,
            config.decay_end,
            1_000_000,  # Large horizon; not critical for notebook steps
            config.warmup_steps,
            config.wsd_decay_steps,
        )
        tx = optax.adamw(
            learning_rate=self.lr_schedule,
            b1=0.9,
            b2=0.9,
            weight_decay=1e-4,
            mu_dtype=config.param_dtype,
        )
        self.optimizer = nnx.Optimizer(self.model, tx)

        # Dataloader
        self.loader_iterator: Optional[grain.DataLoaderIterator] = None
        if build_dataloader and config.data_dir:
            self.loader_iterator = self._build_dataloader()

        # Checkpoint manager
        self.ckpt_mgr: Optional[ocp.CheckpointManager] = None
        if build_checkpoint_manager and config.ckpt_dir:
            self.ckpt_mgr = self._build_checkpoint_manager(config.ckpt_dir)

        # Training state helpers
        self.step: int = 0
        self.action_last_active = jnp.zeros(config.num_latents, dtype=jnp.int32)

    def _build_dataloader(self) -> grain.DataLoaderIterator:
        cfg = self.config
        image_shape = (cfg.image_height, cfg.image_width, cfg.image_channels)
        array_record_files = [
            os.path.join(cfg.data_dir, x)
            for x in os.listdir(cfg.data_dir)
            if x.endswith(".array_record")
        ]
        dataloader = get_dataloader(
            array_record_files,
            cfg.seq_len,
            cfg.batch_size,
            *image_shape,
            num_workers=8,
            prefetch_buffer_size=1,
            seed=cfg.seed,
        )
        initial_state = dataloader._create_initial_state()
        return grain.DataLoaderIterator(dataloader, initial_state)

    def _build_checkpoint_manager(self, ckpt_dir: str) -> ocp.CheckpointManager:
        handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
        handler_registry.add(
            "model_state", ocp.args.PyTreeSave, ocp.handlers.PyTreeCheckpointHandler
        )
        handler_registry.add(
            "model_state", ocp.args.PyTreeRestore, ocp.handlers.PyTreeCheckpointHandler
        )
        handler_registry.add(
            "dataloader_state",
            grain.checkpoint.CheckpointSave,
            cast(ocp.handlers.CheckpointHandler, grain.checkpoint.CheckpointHandler),
        )
        handler_registry.add(
            "dataloader_state",
            grain.checkpoint.CheckpointRestore,
            cast(ocp.handlers.CheckpointHandler, grain.checkpoint.CheckpointHandler),
        )

        options = ocp.CheckpointManagerOptions(step_format_fixed_length=6)
        return ocp.CheckpointManager(
            directory=ckpt_dir,
            options=options,
            handler_registry=handler_registry,
        )

    def restore_latest(self) -> int:
        if self.ckpt_mgr is None:
            raise ValueError("Checkpoint manager not initialized; set ckpt_dir in config.")
        abstract_optimizer = nnx.eval_shape(lambda: self.optimizer)
        abstract_optimizer_state = nnx.state(abstract_optimizer)
        if self.loader_iterator is None:
            raise ValueError("Dataloader not initialized; set data_dir in config.")
        restored = self.ckpt_mgr.restore(
            self.ckpt_mgr.latest_step(),
            args=ocp.args.Composite(
                model_state=ocp.args.PyTreeRestore(abstract_optimizer_state),  # type: ignore
                dataloader_state=grain.checkpoint.CheckpointRestore(self.loader_iterator),  # type: ignore
            ),
        )
        nnx.update(self.optimizer, restored["model_state"])  # type: ignore
        self.loader_iterator = restored["dataloader_state"]
        self.step = self.ckpt_mgr.latest_step() or 0
        return self.step

    def save(self) -> None:
        if self.ckpt_mgr is None:
            raise ValueError("Checkpoint manager not initialized; set ckpt_dir in config.")
        optimizer_state = nnx.state(self.optimizer)
        if self.loader_iterator is None:
            raise ValueError("Dataloader not initialized; set data_dir in config.")
        self.ckpt_mgr.save(
            self.step,
            args=ocp.args.Composite(
                model_state=ocp.args.PyTreeSave(optimizer_state),  # type: ignore
                dataloader_state=grain.checkpoint.CheckpointSave(self.loader_iterator),  # type: ignore
            ),
        )

    def next_batch(self) -> Dict[str, jax.Array]:
        if self.loader_iterator is None:
            raise ValueError("Dataloader not initialized; set data_dir in config.")
        videos = next(self.loader_iterator)
        gt = jnp.asarray(videos, dtype=jnp.float32) / 255.0
        videos = gt.astype(self.config.dtype)
        self.rng, _rng = jax.random.split(self.rng)
        return {"videos": videos, "rng": _rng}

    def enable_sowing(self) -> None:
        for model in [self.model.encoder, self.model.decoder]:
            model.sow_logits = True
            for blk in getattr(model, "blocks", []):
                blk.sow_weights = True
                blk.sow_activations = True

    def forward(self, batch: Dict[str, jax.Array], training: bool = True) -> Dict[str, jax.Array]:
        return self.model(batch, training=training)

    def _lam_loss_fn(self, model: LatentActionModel, inputs: Dict[str, jax.Array]) -> Tuple[jax.Array, Tuple[jax.Array, jax.Array, Dict[str, jax.Array]]]:
        cfg = self.config
        gt = jnp.asarray(inputs["videos"], dtype=jnp.float32) / 255.0
        inputs["videos"] = gt.astype(cfg.dtype)
        model.train()
        outputs = model(inputs, training=True)
        outputs["recon"] = outputs["recon"].astype(jnp.float32)
        gt_future_frames = gt[:, 1:]
        mse = jnp.square(gt_future_frames - outputs["recon"]).mean()
        q_loss = jnp.square(jax.lax.stop_gradient(outputs["emb"]) - outputs["z"]).mean()
        commitment_loss = jnp.square(outputs["emb"] - jax.lax.stop_gradient(outputs["z"])) .mean()
        loss = mse + q_loss + cfg.vq_beta * commitment_loss

        # Metrics
        gt_img = gt_future_frames.clip(0, 1).reshape(-1, *gt_future_frames.shape[2:])
        recon_img = outputs["recon"].clip(0, 1).reshape(-1, *outputs["recon"].shape[2:])
        # Avoid importing dm_pix here for simplicity; mse-based metrics suffice for notebooks
        psnr = -10.0 * jnp.log10(jnp.maximum(jnp.mean((gt_img - recon_img) ** 2), 1e-8))
        ssim = jnp.nan  # Placeholder; add pix-based SSIM if needed interactively
        count_fn = jax.vmap(lambda i: (outputs["indices"] == i).sum())
        index_counts = count_fn(jnp.arange(cfg.num_latents))
        metrics = dict(
            loss=loss,
            mse=mse,
            q_loss=q_loss,
            commitment_loss=commitment_loss,
            psnr=psnr,
            ssim=ssim,
            codebook_usage=(index_counts != 0).mean(),
        )
        return loss, (outputs["recon"], index_counts, metrics)

    def train_step(self, batch: Dict[str, jax.Array]) -> Tuple[jax.Array, jax.Array, Dict[str, jax.Array]]:
        cfg = self.config

        def loss_fn(model: LatentActionModel) -> Tuple[jax.Array, Tuple[jax.Array, jax.Array, Dict[str, jax.Array]]]:
            return self._lam_loss_fn(model, dict(batch))

        (loss, (recon, idx_counts, metrics)), grads = nnx.value_and_grad(loss_fn, has_aux=True)(self.model)
        self.optimizer.update(grads)

        # Reset inactive latent actions (same logic as training script, with small safety)
        codebook = self.model.vq.codebook
        num_codes = len(codebook)
        active_codes = idx_counts != 0.0
        self.action_last_active = jnp.where(active_codes, 0, self.action_last_active + 1)
        denom = jnp.maximum(active_codes.sum(), 1)
        p_code = active_codes / denom
        self.rng, choice_rng = jax.random.split(self.rng)
        reset_idxs = jax.random.choice(choice_rng, num_codes, shape=(num_codes,), p=p_code)
        do_reset = self.action_last_active >= cfg.vq_reset_thresh
        new_codebook = jnp.where(jnp.expand_dims(do_reset, -1), codebook[reset_idxs], codebook.value)
        self.model.vq.codebook.value = new_codebook
        self.action_last_active = jnp.where(do_reset, 0, self.action_last_active)

        metrics = dict(metrics)
        metrics["lr"] = self.lr_schedule(self.step)
        self.step += 1
        return loss, recon, metrics


# Local imports placed at bottom to avoid circular deps during type checking
import os  # noqa: E402