"""FACT: Force-Aware Causal Transformer policy.

A from-scratch block-causal policy that

* drops the 250-point track pathway entirely (no track tokens / encoder /
  decoder / loss),
* replaces future-video (Wan world-model) prediction with *future-state*
  prediction: the 16 proprioceptive EEF/hand states that follow the 16-step
  action chunk are denoised jointly with the actions,
* uses a two-stage flow schedule: a coarse joint flow
  (tau 1 -> 0.4) on [action, future state], followed by the force-only
  tactile refinement (tau 0.4 -> 0), which is re-used unchanged from
  ``fact.force.ForceOnlyTransformer``.

The coarse denoiser is a small DiT-style block-causal transformer.  Frozen
perception components run offline: the Wan2.2 VAE as a fixed image tokenizer
(48-channel latent of the current 3-view frame -> 50 tokens) and the Wan
umT5-XXL encoder whose token embeddings enter as clean prefix tokens
(``use_text``).
"""

from __future__ import annotations

import json
import math
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any

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

from .force import (
    ACTION_HORIZON,
    FORCE_OFFSETS,
    PADDED_ACTION_DIM,
    PHYSICAL_ACTION_DIM,
    ForceOnlyOutput,
    ForceOnlyTransformer,
    TwoStageFlowSchedule,
    euler_flow_step,
    pad_action_62_to_64,
    sample_trex_timestep,
    sinusoidal_tau_embedding,
)


# Token type ids inside the packed sequence.
TYPE_OBS, TYPE_STATE, TYPE_ACTION, TYPE_FSTATE, TYPE_TASK, TYPE_TEXT = range(6)


@dataclass
class FactConfig:
    # coarse transformer
    dim: int = 1024
    num_layers: int = 12
    num_heads: int = 16
    ffn_dim: int = 4096
    dropout: float = 0.0
    gradient_checkpointing: bool = False
    # observation tokenizer (Wan2.2 VAE latent of the 160x320 3-view grid)
    obs_channels: int = 48
    obs_patch: int = 2
    obs_grid: tuple[int, int] = (10, 20)
    # action / state
    action_dim: int = PADDED_ACTION_DIM
    physical_action_dim: int = PHYSICAL_ACTION_DIM
    action_horizon: int = ACTION_HORIZON
    state_dim: int = PADDED_ACTION_DIM
    physical_state_dim: int = PHYSICAL_ACTION_DIM
    # future-state prediction (replaces future-video prediction)
    predict_future_state: bool = True
    future_state_horizon: int = ACTION_HORIZON
    # autoregressive block memory
    max_blocks: int = 32
    window: int = 4
    use_task_token: bool = True
    time_dim: int = 256
    # language conditioning: frozen umT5-XXL token embeddings as clean prefix
    use_text: bool = False
    text_dim: int = 4096
    max_text_tokens: int = 64
    # tactile force stage (force-only refiner)
    use_force: bool = True
    force_dim: int = 768
    force_layers: int = 6
    force_heads: int = 12
    force_ffn: int = 3072
    vq_codebook_size: int = 64
    vq_code_dim: int = 256
    tactile_dropout: float = 0.1
    # two-stage schedule
    coarse_steps: int = 6
    force_steps: int = 4
    tau_split: float = 0.4
    # loss weights
    action_loss_weight: float = 1.0
    state_loss_weight: float = 1.0
    force_loss_weight: float = 1.0
    vq_loss_weight: float = 1.0
    vq_commitment_weight: float = 0.25

    @property
    def obs_tokens(self) -> int:
        h, w = self.obs_grid
        return (h // self.obs_patch) * (w // self.obs_patch)

    @property
    def obs_in_dim(self) -> int:
        return self.obs_channels * self.obs_patch * self.obs_patch

    @property
    def fstate_tokens(self) -> int:
        return self.future_state_horizon if self.predict_future_state else 0

    @property
    def tokens_per_block(self) -> int:
        return self.obs_tokens + 1 + self.action_horizon + self.fstate_tokens

    @property
    def prefix_tokens(self) -> int:
        return (1 if self.use_task_token else 0) + (self.max_text_tokens if self.use_text else 0)

    def to_json(self) -> str:
        return json.dumps(asdict(self), indent=1)

    @classmethod
    def from_json(cls, text: str) -> "FactConfig":
        data = json.loads(text)
        if "obs_grid" in data:
            data["obs_grid"] = tuple(data["obs_grid"])
        return cls(**data)


class FactLayout:
    """Index bookkeeping for one packed [task | text | K x block] sequence."""

    def __init__(self, cfg: FactConfig, num_blocks: int) -> None:
        self.cfg = cfg
        self.num_blocks = num_blocks
        self.task = 1 if cfg.use_task_token else 0
        self.text = cfg.max_text_tokens if cfg.use_text else 0
        self.prefix = self.task + self.text
        self.obs = cfg.obs_tokens
        self.action = cfg.action_horizon
        self.fstate = cfg.fstate_tokens
        self.per_block = cfg.tokens_per_block
        self.length = self.prefix + num_blocks * self.per_block

    @property
    def text_slice(self) -> slice:
        return slice(self.task, self.task + self.text)

    def block_start(self, k: int) -> int:
        return self.prefix + k * self.per_block

    def obs_slice(self, k: int) -> slice:
        s = self.block_start(k)
        return slice(s, s + self.obs)

    def state_slice(self, k: int) -> slice:
        s = self.block_start(k) + self.obs
        return slice(s, s + 1)

    def action_slice(self, k: int) -> slice:
        s = self.block_start(k) + self.obs + 1
        return slice(s, s + self.action)

    def fstate_slice(self, k: int) -> slice:
        s = self.block_start(k) + self.obs + 1 + self.action
        return slice(s, s + self.fstate)

    def token_metadata(self, device: torch.device) -> dict[str, torch.Tensor]:
        slot = torch.zeros(self.length, dtype=torch.long)
        block = torch.zeros(self.length, dtype=torch.long)
        ttype = torch.zeros(self.length, dtype=torch.long)
        if self.task:
            ttype[0] = TYPE_TASK
        if self.text:
            ttype[self.text_slice] = TYPE_TEXT
        per_slot_type = (
            [TYPE_OBS] * self.obs
            + [TYPE_STATE]
            + [TYPE_ACTION] * self.action
            + [TYPE_FSTATE] * self.fstate
        )
        per_slot_type_t = torch.tensor(per_slot_type, dtype=torch.long)
        for k in range(self.num_blocks):
            s = self.block_start(k)
            slot[s : s + self.per_block] = torch.arange(self.per_block)
            block[s : s + self.per_block] = k
            ttype[s : s + self.per_block] = per_slot_type_t
        return {
            "slot": slot.to(device),
            "block": block.to(device),
            "type": ttype.to(device),
        }

    def allow_matrix(
        self, device: torch.device, text_mask: torch.Tensor | None = None
    ) -> torch.Tensor:
        """Block-causal mask: block k reads blocks [k-window, k] and the prefix.

        Returns [L, L] (no text padding) or [B, 1, L, L] when ``text_mask``
        [B, T] marks valid text tokens.
        """

        meta = self.token_metadata(device)
        block = meta["block"]
        is_prefix = torch.zeros(self.length, dtype=torch.bool, device=device)
        is_prefix[: self.prefix] = True
        qb = block[:, None]
        kb = block[None, :]
        allow = (kb <= qb) & (kb >= qb - self.cfg.window)
        allow = allow | is_prefix[None, :]
        # Prefix tokens are clean conditions: they read only the prefix.
        allow[is_prefix] = False
        allow[is_prefix[:, None] & is_prefix[None, :]] = True
        if text_mask is None or not self.text:
            return allow
        valid = torch.ones(text_mask.shape[0], self.length, dtype=torch.bool, device=device)
        valid[:, self.text_slice] = text_mask.to(device, torch.bool)
        allow_b = allow[None] & valid[:, None, :]
        eye = torch.eye(self.length, dtype=torch.bool, device=device)
        allow_b = allow_b | eye[None]  # padded tokens still attend to themselves
        return allow_b[:, None]

    def memory_indices(self, k: int) -> torch.Tensor:
        first = max(0, k - self.cfg.window)
        pieces = []
        if self.prefix:
            pieces.append(torch.arange(0, self.prefix))
        pieces.append(
            torch.arange(self.block_start(first), self.block_start(k) + self.per_block)
        )
        return torch.cat(pieces)


class FactBlock(nn.Module):
    """Pre-norm transformer block with per-token adaLN-zero modulation."""

    def __init__(self, dim: int, num_heads: int, ffn_dim: int, dropout: float) -> None:
        super().__init__()
        if dim % num_heads:
            raise ValueError("dim must be divisible by num_heads")
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.qkv = nn.Linear(dim, 3 * dim)
        self.q_norm = nn.RMSNorm(self.head_dim, eps=1e-6)
        self.k_norm = nn.RMSNorm(self.head_dim, eps=1e-6)
        self.proj = nn.Linear(dim, dim)
        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.ffn = nn.Sequential(
            nn.Linear(dim, ffn_dim),
            nn.GELU(approximate="tanh"),
            nn.Linear(ffn_dim, dim),
        )
        self.dropout = nn.Dropout(dropout)
        self.modulation = nn.Parameter(torch.zeros(1, 1, 6, dim))

    def forward(
        self, x: torch.Tensor, e: torch.Tensor, allow: torch.Tensor
    ) -> torch.Tensor:
        batch, length, _ = x.shape
        m = (e + self.modulation).unbind(2)
        shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = m
        h = self.norm1(x) * (1 + scale_a) + shift_a
        qkv = self.qkv(h).view(batch, length, 3, self.num_heads, self.head_dim)
        q, k, v = qkv.unbind(2)
        q = self.q_norm(q).transpose(1, 2)
        k = self.k_norm(k).transpose(1, 2)
        v = v.transpose(1, 2)
        a = F.scaled_dot_product_attention(q, k, v, attn_mask=allow)
        a = a.transpose(1, 2).reshape(batch, length, self.dim)
        x = x + gate_a * self.dropout(self.proj(a))
        h = self.norm2(x) * (1 + scale_f) + shift_f
        x = x + gate_f * self.dropout(self.ffn(h))
        return x


@dataclass
class FactCoarseOutput:
    action_flow: torch.Tensor
    fstate_flow: torch.Tensor | None
    hidden: torch.Tensor
    memory_by_block: torch.Tensor  # [B, K, M, d]
    memory_mask_by_block: torch.Tensor  # [B, K, M]
    layout: FactLayout


class FactCoarseTransformer(nn.Module):
    """Block-causal joint denoiser for [action, future state] tokens."""

    def __init__(self, cfg: FactConfig) -> None:
        super().__init__()
        self.cfg = cfg
        dim = cfg.dim
        self.obs_encoder = nn.Linear(cfg.obs_in_dim, dim)
        self.state_encoder = nn.Sequential(
            nn.Linear(cfg.state_dim, dim), nn.SiLU(), nn.Linear(dim, dim)
        )
        self.action_encoder = nn.Linear(cfg.action_dim, dim)
        if cfg.predict_future_state:
            self.fstate_encoder = nn.Linear(cfg.state_dim, dim)
            self.fstate_decoder = nn.Linear(dim, cfg.state_dim)
            nn.init.zeros_(self.fstate_decoder.weight)
            nn.init.zeros_(self.fstate_decoder.bias)
        self.slot_embedding = nn.Embedding(cfg.tokens_per_block, dim)
        self.block_embedding = nn.Embedding(cfg.max_blocks, dim)
        self.type_embedding = nn.Embedding(6, dim)
        if cfg.use_task_token:
            self.task_token = nn.Parameter(torch.randn(1, 1, dim) * 0.02)
        if cfg.use_text:
            self.text_encoder = nn.Sequential(
                nn.LayerNorm(cfg.text_dim), nn.Linear(cfg.text_dim, dim)
            )
            self.text_position = nn.Embedding(cfg.max_text_tokens, dim)
            nn.init.normal_(self.text_position.weight, std=0.02)
        self.time_embedding = nn.Sequential(
            nn.Linear(cfg.time_dim, dim), nn.SiLU(), nn.Linear(dim, dim)
        )
        self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim))
        nn.init.zeros_(self.time_projection[1].weight)
        nn.init.zeros_(self.time_projection[1].bias)
        self.blocks = nn.ModuleList(
            [
                FactBlock(dim, cfg.num_heads, cfg.ffn_dim, cfg.dropout)
                for _ in range(cfg.num_layers)
            ]
        )
        self.final_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.final_modulation = nn.Sequential(nn.SiLU(), nn.Linear(dim, 2 * dim))
        nn.init.zeros_(self.final_modulation[1].weight)
        nn.init.zeros_(self.final_modulation[1].bias)
        self.action_decoder = nn.Linear(dim, cfg.action_dim)
        nn.init.zeros_(self.action_decoder.weight)
        nn.init.zeros_(self.action_decoder.bias)
        nn.init.normal_(self.slot_embedding.weight, std=0.02)
        nn.init.normal_(self.block_embedding.weight, std=0.02)
        nn.init.normal_(self.type_embedding.weight, std=0.02)

    # ------------------------------------------------------------------
    def patchify_obs(self, latents: torch.Tensor) -> torch.Tensor:
        """[B,K,C,H,W] latents -> [B,K,obs_tokens,C*p*p]."""

        b, k, c, h, w = latents.shape
        p = self.cfg.obs_patch
        if (h, w) != tuple(self.cfg.obs_grid):
            raise ValueError(f"obs latent grid {(h, w)} != {self.cfg.obs_grid}")
        x = latents.view(b, k, c, h // p, p, w // p, p)
        x = x.permute(0, 1, 3, 5, 2, 4, 6).reshape(b, k, (h // p) * (w // p), c * p * p)
        return x

    def forward(
        self,
        *,
        obs_latents: torch.Tensor,
        state: torch.Tensor,
        action: torch.Tensor,
        tau: torch.Tensor,
        future_state: torch.Tensor | None = None,
        fstate_tau: torch.Tensor | None = None,
        text_tokens: torch.Tensor | None = None,
        text_mask: torch.Tensor | None = None,
    ) -> FactCoarseOutput:
        cfg = self.cfg
        batch, num_blocks = action.shape[:2]
        device = action.device
        dtype = self.action_encoder.weight.dtype
        layout = FactLayout(cfg, num_blocks)
        if tau.shape != (batch, num_blocks):
            raise ValueError("tau must be [B,K]")
        if fstate_tau is None:
            fstate_tau = tau

        obs = self.obs_encoder(self.patchify_obs(obs_latents.to(device, dtype)))
        st = self.state_encoder(state.to(device, dtype)).unsqueeze(2)
        act = self.action_encoder(action.to(device, dtype))
        pieces = [obs, st, act]
        if cfg.predict_future_state:
            if future_state is None:
                raise ValueError("future_state is required when predict_future_state")
            pieces.append(self.fstate_encoder(future_state.to(device, dtype)))
        x = torch.cat(pieces, dim=2)  # [B,K,per_block,d]
        meta = layout.token_metadata(device)
        x = x + self.slot_embedding.weight[None, None]
        x = x + self.block_embedding(torch.arange(num_blocks, device=device))[None, :, None]
        x = x.reshape(batch, num_blocks * layout.per_block, cfg.dim)
        prefix: list[torch.Tensor] = []
        if layout.task:
            prefix.append(self.task_token.to(dtype).expand(batch, -1, -1))
        if layout.text:
            if text_tokens is None:
                raise ValueError("use_text=True requires text_tokens [B,T,text_dim]")
            tt = text_tokens.to(device, dtype)
            if tt.shape[1] != layout.text:
                raise ValueError(f"text_tokens must have {layout.text} tokens, got {tt.shape[1]}")
            if text_mask is None:
                text_mask = torch.ones(batch, layout.text, dtype=torch.bool, device=device)
            tt = self.text_encoder(tt) + self.text_position.weight[None]
            tt = tt * text_mask.to(device, dtype)[..., None]
            prefix.append(tt)
        if prefix:
            x = torch.cat(prefix + [x], dim=1)
        x = x + self.type_embedding(meta["type"])[None]

        # per-token flow time: clean conditions sit at tau=0
        token_tau = torch.zeros(batch, layout.length, device=device, dtype=torch.float32)
        for k in range(num_blocks):
            token_tau[:, layout.action_slice(k)] = tau[:, k : k + 1].float()
            if cfg.predict_future_state:
                token_tau[:, layout.fstate_slice(k)] = fstate_tau[:, k : k + 1].float()
        t = self.time_embedding(
            sinusoidal_tau_embedding(token_tau.flatten(), cfg.time_dim).to(dtype)
        ).view(batch, layout.length, cfg.dim)
        e = self.time_projection(t).view(batch, layout.length, 6, cfg.dim)
        allow = layout.allow_matrix(device, text_mask if layout.text else None)

        for block in self.blocks:
            if cfg.gradient_checkpointing and self.training and torch.is_grad_enabled():
                x = torch.utils.checkpoint.checkpoint(block, x, e, allow, use_reentrant=False)
            else:
                x = block(x, e, allow)
        hidden = x
        shift, scale = self.final_modulation(t).chunk(2, dim=-1)
        y = self.final_norm(x) * (1 + scale) + shift
        action_flow = torch.stack(
            [self.action_decoder(y[:, layout.action_slice(k)]) for k in range(num_blocks)],
            dim=1,
        )
        fstate_flow = None
        if cfg.predict_future_state:
            fstate_flow = torch.stack(
                [self.fstate_decoder(y[:, layout.fstate_slice(k)]) for k in range(num_blocks)],
                dim=1,
            )
        # causal coarse memory for the force stage: everything block k may read
        token_valid = torch.ones(batch, layout.length, dtype=torch.bool, device=device)
        if layout.text and text_mask is not None:
            token_valid[:, layout.text_slice] = text_mask.to(device, torch.bool)
        indices = [layout.memory_indices(k).to(device) for k in range(num_blocks)]
        max_len = max(int(i.numel()) for i in indices)
        memories, masks = [], []
        for idx in indices:
            mem = hidden.index_select(1, idx)
            mask = token_valid.index_select(1, idx)
            valid = idx.numel()
            if valid < max_len:
                mem = F.pad(mem, (0, 0, 0, max_len - valid))
                mask = F.pad(mask, (0, max_len - valid), value=False)
            memories.append(mem)
            masks.append(mask)
        return FactCoarseOutput(
            action_flow=action_flow,
            fstate_flow=fstate_flow,
            hidden=hidden,
            memory_by_block=torch.stack(memories, dim=1),
            memory_mask_by_block=torch.stack(masks, dim=1),
            layout=layout,
        )


class FactPolicy(nn.Module):
    """Coarse joint [action, future-state] flow + tactile suffix refinement."""

    def __init__(self, cfg: FactConfig) -> None:
        super().__init__()
        self.cfg = cfg
        self.coarse = FactCoarseTransformer(cfg)
        self.force: ForceOnlyTransformer | None = None
        if cfg.use_force:
            self.force = ForceOnlyTransformer(
                dim=cfg.force_dim,
                num_heads=cfg.force_heads,
                num_layers=cfg.force_layers,
                ffn_dim=cfg.force_ffn,
                action_horizon=cfg.action_horizon,
                action_dim=cfg.action_dim,
                vq_codebook_size=cfg.vq_codebook_size,
                vq_code_dim=cfg.vq_code_dim,
                coarse_memory_dim=cfg.dim,
            )
        self.schedule = TwoStageFlowSchedule(
            coarse_steps=cfg.coarse_steps,
            force_steps=cfg.force_steps,
            tau_split=cfg.tau_split,
            action_horizon=cfg.action_horizon,
        )

    # ------------------------------------------------------------------
    @property
    def device(self) -> torch.device:
        return self.coarse.action_encoder.weight.device

    @property
    def dtype(self) -> torch.dtype:
        return self.coarse.action_encoder.weight.dtype

    def _action_mask(self, like: torch.Tensor) -> torch.Tensor:
        mask = torch.zeros_like(like, dtype=torch.bool)
        mask[..., : self.cfg.physical_action_dim] = True
        return mask

    def _state_mask(self, like: torch.Tensor) -> torch.Tensor:
        mask = torch.zeros_like(like, dtype=torch.bool)
        mask[..., : self.cfg.physical_state_dim] = True
        return mask

    @staticmethod
    def _noise(
        clean: torch.Tensor, tau: torch.Tensor, noise: torch.Tensor | None = None
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        if noise is None:
            noise = torch.randn_like(clean)
        t = tau.to(clean.dtype)
        while t.ndim < clean.ndim:
            t = t.unsqueeze(-1)
        noisy = (1 - t) * clean + t * noise
        target = noise - clean
        return noisy, target, noise

    def _text_kwargs(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        if not self.cfg.use_text:
            return {}
        if "text_tokens" not in batch:
            raise KeyError("use_text=True but the batch has no text_tokens")
        out = {"text_tokens": batch["text_tokens"].to(self.device, self.dtype)}
        if "text_mask" in batch:
            out["text_mask"] = batch["text_mask"].to(self.device, torch.bool)
        return out

    # ------------------------------------------------------------------
    def forward(
        self, batch: dict[str, torch.Tensor], *, tau: torch.Tensor | None = None
    ) -> dict[str, torch.Tensor]:
        """nn.Module entry point (so DDP wrappers sync gradients): training losses."""

        return self.forward_core(batch, tau=tau)

    def forward_core(
        self, batch: dict[str, torch.Tensor], *, tau: torch.Tensor | None = None
    ) -> dict[str, torch.Tensor]:
        cfg = self.cfg
        dev, dt = self.device, self.dtype
        action_clean = batch["action"].to(dev, dt)
        state = batch["state"].to(dev, dt)
        obs = batch["obs_latents"].to(dev, dt)
        text_kwargs = self._text_kwargs(batch)
        B, K = action_clean.shape[:2]
        if tau is None:
            tau = sample_trex_timestep((B, K), device=dev, dtype=dt)
        else:
            tau = tau.to(dev, dt)
            if tau.ndim == 1:
                tau = tau[:, None].expand(-1, K)
        amask = self._action_mask(action_clean).to(dt)
        action_noise = torch.randn_like(action_clean) * amask
        action_noisy, action_target, _ = self._noise(action_clean, tau, action_noise)

        fstate_clean = fstate_noisy = fstate_target = None
        smask = None
        if cfg.predict_future_state:
            fstate_clean = batch["future_state"].to(dev, dt)
            smask = self._state_mask(fstate_clean).to(dt)
            fstate_noisy, fstate_target, _ = self._noise(
                fstate_clean, tau, torch.randn_like(fstate_clean) * smask
            )

        out = self.coarse(
            obs_latents=obs,
            state=state,
            action=action_noisy,
            tau=tau,
            future_state=fstate_noisy,
            **text_kwargs,
        )
        action_loss = ((out.action_flow - action_target).square() * amask).sum() / amask.sum().clamp_min(1)
        zero = action_loss.new_zeros(())
        state_loss = zero
        if cfg.predict_future_state:
            state_loss = ((out.fstate_flow - fstate_target).square() * smask).sum() / smask.sum().clamp_min(1)

        force_loss = vq_loss = commitment_loss = zero
        perplexity = active_codes = zero
        if self.force is not None and "force_history" in batch:
            current_force = batch["current_force"].to(dev, dt)  # [B,K,S,10,6]
            history = batch["force_history"].to(dev, dt)  # [B,K,S,16,10,6]
            valid = batch.get("force_history_valid")
            if valid is not None:
                valid = valid.to(dev, torch.bool)
            S = current_force.shape[2]
            force_tau = sample_trex_timestep((B, K, S), device=dev, dtype=dt, upper=cfg.tau_split)
            force_clean = action_clean[:, :, None].expand(-1, -1, S, -1, -1)
            noise_base = torch.randn_like(action_clean) * amask
            force_noisy, force_target, _ = self._noise(
                force_clean, force_tau, noise_base[:, :, None].expand_as(force_clean)
            )
            split_tau = torch.full_like(tau, cfg.tau_split)
            action_at_split, _, _ = self._noise(action_clean, split_tau, noise_base)
            fstate_at_split = None
            if cfg.predict_future_state:
                fstate_at_split, _, _ = self._noise(
                    fstate_clean, split_tau, torch.randn_like(fstate_clean) * smask
                )
            with torch.no_grad():
                split_out = self.coarse(
                    obs_latents=obs,
                    state=state,
                    action=action_at_split,
                    tau=split_tau,
                    future_state=fstate_at_split,
                    **text_kwargs,
                )
            memory = (
                split_out.memory_by_block.detach()[:, :, None]
                .expand(-1, -1, S, -1, -1)
                .flatten(0, 2)
            )
            memory_mask = (
                split_out.memory_mask_by_block[:, :, None]
                .expand(-1, -1, S, -1)
                .flatten(0, 2)
            )
            offsets = torch.as_tensor(self.schedule.offsets[:S], device=dev, dtype=torch.long)
            offsets = offsets.view(1, 1, S).expand(B, K, -1).flatten()
            keep = torch.rand(B * K * S, device=dev) >= cfg.tactile_dropout
            fo = self.force(
                force_noisy.flatten(0, 2),
                force_tau.flatten(),
                current_force.flatten(0, 2),
                history.flatten(0, 2),
                coarse_memory=memory,
                coarse_memory_mask=memory_mask,
                update_offset=offsets,
                tactile_keep_mask=keep,
                tactile_history_valid_mask=None if valid is None else valid.flatten(0, 2),
                return_details=True,
            )
            assert isinstance(fo, ForceOnlyOutput)
            pred = fo.flow.view_as(force_target)
            positions = torch.arange(cfg.action_horizon, device=dev).view(1, 1, 1, -1, 1)
            suffix = positions >= offsets.view(B, K, S, 1, 1)
            fmask = amask[:, :, None] * suffix
            force_loss = ((pred - force_target).square() * fmask).sum() / fmask.sum().clamp_min(1)
            vq_loss = fo.vq_loss
            commitment_loss = fo.commitment_loss
            if fo.tactile_perplexity is not None:
                perplexity = fo.tactile_perplexity
            if fo.tactile_active_codes is not None:
                active_codes = fo.tactile_active_codes.to(dt)

        loss = (
            cfg.action_loss_weight * action_loss
            + cfg.state_loss_weight * state_loss
            + cfg.force_loss_weight * force_loss
            + cfg.vq_loss_weight * (vq_loss + cfg.vq_commitment_weight * commitment_loss)
        )
        return {
            "loss": loss,
            "action_loss": action_loss,
            "state_loss": state_loss,
            "force_loss": force_loss,
            "vq_loss": vq_loss,
            "commitment_loss": commitment_loss,
            "tactile_perplexity": perplexity,
            "tactile_active_codes": active_codes,
        }

    # ------------------------------------------------------------------
    @torch.no_grad()
    def refine_action_suffix(
        self,
        coarse_action_at_split: torch.Tensor,
        coarse_memory: torch.Tensor,
        current_force: torch.Tensor,
        history: torch.Tensor,
        *,
        update_offset: int = 0,
        history_valid: torch.Tensor | None = None,
        coarse_memory_mask: torch.Tensor | None = None,
        tactile_keep_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        assert self.force is not None
        action = pad_action_62_to_64(coarse_action_at_split).to(self.device, self.dtype).clone()
        action[..., self.cfg.physical_action_dim :] = 0
        for step in self.schedule.iter_steps("force"):
            flow = self.force(
                action,
                step.tau,
                current_force.to(self.device, self.dtype),
                history.to(self.device),
                coarse_memory=coarse_memory,
                coarse_memory_mask=coarse_memory_mask,
                update_offset=update_offset,
                tactile_history_valid_mask=history_valid,
                tactile_keep_mask=tactile_keep_mask,
            )
            proposed = euler_flow_step(action, flow, step.tau, step.tau_next)
            action[:, update_offset:] = proposed[:, update_offset:]
            action[..., self.cfg.physical_action_dim :] = 0
        return action

    @torch.no_grad()
    def sample(
        self,
        batch: dict[str, torch.Tensor],
        *,
        seed: int | None = None,
        mode: str = "cascade",
    ) -> dict[str, torch.Tensor]:
        """One-block open-loop sampling.

        mode:
          cascade        6 coarse Euler steps to tau=0.4, then 4 tactile steps.
          coarse_only    the coarse action read out at tau=0.4 (the baseline
                         eval's ``coarse_only`` semantics: no refinement).
          coarse_to_zero coarse model carries the flow to tau=0 (10 steps).
        """

        cfg = self.cfg
        dev, dt = self.device, self.dtype
        state = batch["state"].to(dev, dt)
        obs = batch["obs_latents"].to(dev, dt)
        text_kwargs = self._text_kwargs(batch)
        if state.ndim == 2:
            state = state[:, None]
        if obs.ndim == 4:
            obs = obs[:, None]
        B = state.shape[0]
        if state.shape[1] != 1:
            raise ValueError("online sampling consumes exactly one block")
        gen = None
        if seed is not None:
            gen = torch.Generator(device=dev).manual_seed(seed)
        action = torch.randn(B, 1, cfg.action_horizon, cfg.action_dim, device=dev, dtype=dt, generator=gen)
        action[..., cfg.physical_action_dim :] = 0
        fstate = None
        if cfg.predict_future_state:
            fstate = torch.randn(B, 1, cfg.future_state_horizon, cfg.state_dim, device=dev, dtype=dt, generator=gen)
            fstate[..., cfg.physical_state_dim :] = 0

        def run(tau_value: float) -> FactCoarseOutput:
            tau = torch.full((B, 1), tau_value, device=dev, dtype=dt)
            return self.coarse(
                obs_latents=obs, state=state, action=action, tau=tau, future_state=fstate, **text_kwargs
            )

        for step in self.schedule.iter_steps("coarse"):
            out = run(step.tau)
            action = euler_flow_step(action, out.action_flow, step.tau, step.tau_next)
            action[..., cfg.physical_action_dim :] = 0
            if fstate is not None:
                fstate = euler_flow_step(fstate, out.fstate_flow, step.tau, step.tau_next)
                fstate[..., cfg.physical_state_dim :] = 0
        out = run(cfg.tau_split)
        coarse_action = action[:, 0]
        memory = out.memory_by_block[:, 0]
        memory_mask = out.memory_mask_by_block[:, 0]
        fstate_pred = None
        if fstate is not None:
            fstate_pred = euler_flow_step(fstate, out.fstate_flow, cfg.tau_split, 0.0)[:, 0]
        result: dict[str, torch.Tensor] = {
            "coarse_action_at_split": coarse_action,
            "coarse_memory": memory,
            "coarse_memory_mask": memory_mask,
        }
        if fstate_pred is not None:
            result["future_state_pred"] = fstate_pred
        if mode == "cascade":
            if self.force is None:
                raise RuntimeError("cascade sampling requires the force stage")
            current_force = batch["current_force"].to(dev, dt)
            history = batch["force_history"].to(dev)
            valid = batch.get("force_history_valid")
            # [B,K,S,...] -> slot 0 of block 0
            while current_force.ndim > 3:
                current_force = current_force[:, 0]
            while history.ndim > 4:
                history = history[:, 0]
            if valid is not None:
                valid = valid.to(dev, torch.bool)
                while valid.ndim > 2:
                    valid = valid[:, 0]
            result["action_pred"] = self.refine_action_suffix(
                coarse_action, memory, current_force, history,
                update_offset=0, history_valid=valid, coarse_memory_mask=memory_mask,
            )
        elif mode == "coarse_only":
            result["action_pred"] = coarse_action.clone()
        elif mode == "coarse_to_zero":
            for step in self.schedule.iter_steps("force"):
                out = run(step.tau)
                action = euler_flow_step(action, out.action_flow, step.tau, step.tau_next)
                action[..., cfg.physical_action_dim :] = 0
                if fstate is not None:
                    fstate = euler_flow_step(fstate, out.fstate_flow, step.tau, step.tau_next)
                    fstate[..., cfg.physical_state_dim :] = 0
            result["action_pred"] = action[:, 0]
        else:
            raise ValueError(f"unknown sampling mode {mode!r}")
        return result

    # ------------------------------------------------------------------
    def save(self, directory: str | Path, extra: dict[str, Any] | None = None) -> None:
        directory = Path(directory)
        directory.mkdir(parents=True, exist_ok=True)
        (directory / "fact_config.json").write_text(self.cfg.to_json())
        torch.save(self.state_dict(), directory / "fact_model.pt")
        if extra:
            (directory / "extra.json").write_text(json.dumps(extra, indent=1))

    @classmethod
    def load(cls, directory: str | Path, map_location: str = "cpu") -> "FactPolicy":
        directory = Path(directory)
        cfg = FactConfig.from_json((directory / "fact_config.json").read_text())
        model = cls(cfg)
        state = torch.load(directory / "fact_model.pt", map_location=map_location)
        # Backward compatibility: checkpoints saved before the text prefix was
        # added have 5 token-type rows; the TEXT row is appended (unused).
        key = "coarse.type_embedding.weight"
        current = model.state_dict()[key]
        if key in state and state[key].shape[0] < current.shape[0]:
            padded = current.clone()
            padded[: state[key].shape[0]] = state[key]
            state[key] = padded
        model.load_state_dict(state, strict=True)
        return model


def count_parameters(module: nn.Module, trainable_only: bool = False) -> int:
    return sum(
        p.numel() for p in module.parameters() if (p.requires_grad or not trainable_only)
    )


__all__ = [
    "FactBlock",
    "FactCoarseOutput",
    "FactCoarseTransformer",
    "FactConfig",
    "FactLayout",
    "FactPolicy",
    "count_parameters",
]