model/tactile_vae.py

"""ViT-based tactile VAE (no MAE masking).

Architecture:
- Encoder: PatchEmbed + Transformer blocks + fixed sin-cos positional embedding.
- Latent: mu/logvar heads + reparameterization.
- Decoder: latent-conditioned transformer decoder that predicts image patches.
- Reconstruction: unpatchify patch predictions into image space.

Training objective: regular VAE loss (reconstruction + beta * KL), with optional SSIM.
"""
from __future__ import annotations

from pathlib import Path
from typing import Any, Optional
from PIL import Image

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from timm.models.vision_transformer import Block, PatchEmbed


DEFAULT_TACTILE_VAE_DIR = Path("/group2/ct/weihanx/tactile_world_model/tactile_wm/pretrained_models")
DEFAULT_CHECKPOINT_NAME = "ckpt_best.pt"


def _get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: np.ndarray) -> np.ndarray:
    assert embed_dim % 2 == 0
    omega = np.arange(embed_dim // 2, dtype=float)
    omega /= embed_dim / 2.0
    omega = 1.0 / (10000**omega)
    pos = pos.reshape(-1)
    out = np.einsum("m,d->md", pos, omega)
    emb_sin = np.sin(out)
    emb_cos = np.cos(out)
    return np.concatenate([emb_sin, emb_cos], axis=1)


def _get_2d_sincos_pos_embed_from_grid(embed_dim: int, grid: np.ndarray) -> np.ndarray:
    assert embed_dim % 2 == 0
    emb_h = _get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])
    emb_w = _get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])
    return np.concatenate([emb_h, emb_w], axis=1)


def get_2d_sincos_pos_embed(embed_dim: int, grid_size: int, cls_token: bool = False) -> np.ndarray:
    grid_h = np.arange(grid_size, dtype=np.float32)
    grid_w = np.arange(grid_size, dtype=np.float32)
    grid = np.meshgrid(grid_w, grid_h)
    grid = np.stack(grid, axis=0)
    grid = grid.reshape([2, 1, grid_size, grid_size])
    pos_embed = _get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
    if cls_token:
        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
    return pos_embed


def unpatchify(pred_patches: torch.Tensor, patch_size: int, in_chans: int) -> torch.Tensor:
    """Convert patch predictions (B, L, p*p*C) to images (B, C, H, W)."""
    h = w = int(pred_patches.shape[1] ** 0.5)
    assert h * w == pred_patches.shape[1], "number of patches must be a square"
    x = pred_patches.reshape(pred_patches.shape[0], h, w, patch_size, patch_size, in_chans)
    x = torch.einsum("nhwpqc->nchpwq", x)
    return x.reshape(pred_patches.shape[0], in_chans, h * patch_size, h * patch_size)


class ViTEncoder(nn.Module):
    """PatchEmbed + transformer blocks + fixed sin-cos positional embeddings."""

    def __init__(
        self,
        img_size: int,
        patch_size: int,
        in_chans: int,
        embed_dim: int,
        depth: int,
        num_heads: int,
        mlp_ratio: float,
        norm_layer: type[nn.Module] = nn.LayerNorm,
    ):
        super().__init__()
        self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim)
        num_patches = self.patch_embed.num_patches

        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim), requires_grad=False)
        self.blocks = nn.ModuleList(
            [Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer) for _ in range(depth)]
        )
        self.norm = norm_layer(embed_dim)

        self._initialize_weights()

    def _initialize_weights(self) -> None:
        pos_embed = get_2d_sincos_pos_embed(
            self.pos_embed.shape[-1], int(self.patch_embed.num_patches**0.5), cls_token=True
        )
        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))

        w = self.patch_embed.proj.weight.data
        torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
        torch.nn.init.normal_(self.cls_token, std=0.02)
        self.apply(self._init_weights)

    @staticmethod
    def _init_weights(m: nn.Module) -> None:
        if isinstance(m, nn.Linear):
            torch.nn.init.xavier_uniform_(m.weight)
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.patch_embed(x)
        cls_tokens = self.cls_token.expand(x.shape[0], -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)
        x = x + self.pos_embed
        for blk in self.blocks:
            x = blk(x)
        return self.norm(x)


class ViTDecoder(nn.Module):
    """Latent-conditioned transformer decoder that predicts image patches."""

    def __init__(
        self,
        img_size: int,
        patch_size: int,
        in_chans: int,
        latent_dim: int,
        embed_dim: int,
        depth: int,
        num_heads: int,
        mlp_ratio: float,
        norm_layer: type[nn.Module] = nn.LayerNorm,
    ):
        super().__init__()
        self.patch_size = patch_size
        self.in_chans = in_chans
        self.num_patches = (img_size // patch_size) ** 2

        self.z_token = nn.Linear(latent_dim, embed_dim)
        self.patch_tokens = nn.Parameter(torch.zeros(1, self.num_patches, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches, embed_dim), requires_grad=False)

        self.blocks = nn.ModuleList(
            [Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer) for _ in range(depth)]
        )
        self.norm = norm_layer(embed_dim)
        self.pred = nn.Linear(embed_dim, patch_size * patch_size * in_chans)

        self._initialize_weights()

    def _initialize_weights(self) -> None:
        pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.num_patches**0.5), cls_token=False)
        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
        torch.nn.init.normal_(self.patch_tokens, std=0.02)
        self.apply(ViTEncoder._init_weights)

    def forward(self, z: torch.Tensor) -> torch.Tensor:
        b = z.shape[0]
        ztok = self.z_token(z).unsqueeze(1)
        ptok = self.patch_tokens.expand(b, -1, -1)
        x = ztok + ptok
        x = x + self.pos_embed
        for blk in self.blocks:
            x = blk(x)
        x = self.norm(x)
        return self.pred(x)


class TactileVAE(nn.Module):
    """Regular ViT-based VAE for tactile image reconstruction."""

    def __init__(
        self,
        img_size: int = 128,
        patch_size: int = 16,
        in_chans: int = 3,
        embed_dim: int = 256,
        encoder_depth: int = 4,
        encoder_heads: int = 8,
        decoder_embed_dim: int = 192,
        decoder_depth: int = 4,
        decoder_heads: int = 8,
        mlp_ratio: float = 4.0,
        latent_dim: int = 128,
    ):
        super().__init__()
        self.patch_size = patch_size
        self.in_chans = in_chans

        self.encoder = ViTEncoder(
            img_size=img_size,
            patch_size=patch_size,
            in_chans=in_chans,
            embed_dim=embed_dim,
            depth=encoder_depth,
            num_heads=encoder_heads,
            mlp_ratio=mlp_ratio,
        )

        self.mu_head = nn.Linear(embed_dim, latent_dim)
        self.logvar_head = nn.Linear(embed_dim, latent_dim)

        self.decoder = ViTDecoder(
            img_size=img_size,
            patch_size=patch_size,
            in_chans=in_chans,
            latent_dim=latent_dim,
            embed_dim=decoder_embed_dim,
            depth=decoder_depth,
            num_heads=decoder_heads,
            mlp_ratio=mlp_ratio,
        )

    def encode(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]:
        enc_tokens = self.encoder(x)
        cls = enc_tokens[:, 0]
        mu = self.mu_head(cls)
        logvar = self.logvar_head(cls)
        return mu, logvar, {"enc_tokens": enc_tokens}

    @staticmethod
    def reparameterize(mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def decode(self, z: torch.Tensor, enc_ctx: Optional[dict[str, torch.Tensor]] = None) -> torch.Tensor:
        del enc_ctx  # decoder is latent-conditioned only for regular VAE.
        pred_patches = self.decoder(z)
        return unpatchify(pred_patches, patch_size=self.patch_size, in_chans=self.in_chans)

    def reconstruct(self, x: torch.Tensor, use_mean: bool = True) -> torch.Tensor:
        mu, logvar, enc_ctx = self.encode(x)
        z = mu if use_mean else self.reparameterize(mu, logvar)
        return self.decode(z, enc_ctx=enc_ctx)

    def forward(self, x: torch.Tensor, sample: bool = True) -> dict[str, torch.Tensor]:
        mu, logvar, enc_ctx = self.encode(x)
        z = self.reparameterize(mu, logvar) if sample else mu
        pred_patches = self.decoder(z)
        x_hat = unpatchify(pred_patches, patch_size=self.patch_size, in_chans=self.in_chans)
        return {
            "x_hat": x_hat,
            "mu": mu,
            "logvar": logvar,
            "z": z,
            "pred_patches": pred_patches,
            "enc_ctx": enc_ctx,
        }


class SSIMLoss(nn.Module):
    """Simple differentiable SSIM loss (1 - SSIM mean)."""

    def __init__(self, window_size: int = 11, channels: int = 3):
        super().__init__()
        self.window_size = window_size
        self.channels = channels
        self.padding = window_size // 2
        self.register_buffer(
            "kernel",
            torch.ones((channels, 1, window_size, window_size), dtype=torch.float32)
            / (window_size * window_size),
            persistent=False,
        )

    def _filter(self, x: torch.Tensor) -> torch.Tensor:
        if x.shape[1] == self.channels:
            kernel = self.kernel.to(device=x.device, dtype=x.dtype)
        else:
            kernel = torch.ones(
                (x.shape[1], 1, self.window_size, self.window_size),
                device=x.device,
                dtype=x.dtype,
            ) / (self.window_size * self.window_size)
        return F.conv2d(x, kernel, padding=self.padding, groups=x.shape[1])

    def forward(self, x_hat: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
        c1 = 0.01**2
        c2 = 0.03**2
        mu_x = self._filter(x)
        mu_y = self._filter(x_hat)
        sigma_x = self._filter(x * x) - mu_x * mu_x
        sigma_y = self._filter(x_hat * x_hat) - mu_y * mu_y
        sigma_xy = self._filter(x * x_hat) - mu_x * mu_y
        ssim_map = ((2 * mu_x * mu_y + c1) * (2 * sigma_xy + c2)) / (
            (mu_x * mu_x + mu_y * mu_y + c1) * (sigma_x + sigma_y + c2) + 1e-8
        )
        return 1.0 - ssim_map.mean()


class VAELoss(nn.Module):
    """Regular VAE loss: reconstruction + beta * KL."""

    def __init__(
        self,
        beta: float = 1.0,
        recon_type: str = "l1",
        ssim_weight: float = 0.0,
        ssim_window_size: int = 11,
    ):
        super().__init__()
        if recon_type not in {"l1", "mse"}:
            raise ValueError(f"recon_type must be 'l1' or 'mse', got: {recon_type}")
        self.beta = beta
        self.recon_type = recon_type
        self.ssim_weight = ssim_weight
        self.ssim_loss = SSIMLoss(window_size=ssim_window_size)

    def forward(
        self,
        x_hat: torch.Tensor,
        x: torch.Tensor,
        mu: torch.Tensor,
        logvar: torch.Tensor,
    ) -> dict[str, torch.Tensor]:
        recon = F.l1_loss(x_hat, x) if self.recon_type == "l1" else F.mse_loss(x_hat, x)
        ssim_term = self.ssim_loss(x_hat, x) if self.ssim_weight > 0 else x_hat.new_zeros(())
        recon_total = recon + self.ssim_weight * ssim_term
        kl = -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())
        total = recon_total + self.beta * kl
        return {
            "total": total,
            "recon": recon,
            "ssim": ssim_term,
            "recon_total": recon_total,
            "kl": kl,
        }


class BetaVAELoss(VAELoss):
    """Backward-compatible alias for VAELoss."""


class TactileVAEWrapper(nn.Module):
    """Inference wrapper around TactileVAE with a mode-based interface.

    Usage:
        wrapper = TactileVAEWrapper(ckpt_path, device)
        z     = wrapper(x, mode="encode")   # (B, latent_dim)
        x_hat = wrapper(z, mode="decode")   # (B, C, H, W)
    Or equivalently:
        z     = wrapper.encode(x)
        x_hat = wrapper.decode(z)
    """

    def __init__(
        self,
        ckpt_path: str | Path,
        device: str | torch.device = "cpu",
        model_kwargs: Optional[dict[str, Any]] = None,
    ):
        super().__init__()
        self.device = torch.device(device)
        self.vae = self._load(ckpt_path, model_kwargs or {})

    def _load(self, ckpt_path: str | Path, model_kwargs: dict) -> TactileVAE:
        ckpt_path = Path(ckpt_path)
        if not ckpt_path.exists():
            raise FileNotFoundError(f"TactileVAE checkpoint not found: {ckpt_path}")
        state = torch.load(str(ckpt_path), map_location=self.device, weights_only=False)
        state_dict = _unwrap_state(state)
        vae = TactileVAE(**model_kwargs)
        vae.load_state_dict(state_dict, strict=True)
        vae.eval().to(self.device)
        vae.requires_grad_(False)
        return vae

    def to(self, device):
        super().to(device)
        self.device = torch.device(device)
        self.vae = self.vae.to(device)
        return self

    @torch.no_grad()
    def encode(self, x: torch.Tensor) -> torch.Tensor:
        """x: (B, C, H, W) float in [-1, 1]. Returns z: (B, latent_dim) using mu."""
        mu, _logvar, _ctx = self.vae.encode(x.to(self.device))
        return mu

    @torch.no_grad()
    def decode(self, z: torch.Tensor) -> torch.Tensor:
        """z: (B, latent_dim). Returns x_hat: (B, C, H, W) float in [-1, 1]."""
        return self.vae.decode(z.to(self.device))

    @torch.no_grad()
    def forward(self, x: torch.Tensor, mode: str) -> torch.Tensor:
        if mode == "encode":
            return self.encode(x)
        if mode == "decode":
            return self.decode(x)
        raise ValueError(f"mode must be 'encode' or 'decode', got {mode!r}")


def _resolve_checkpoint(checkpoint: Optional[str | Path], vae_dir: str | Path) -> Path:
    if checkpoint is None:
        return Path(vae_dir) / DEFAULT_CHECKPOINT_NAME
    p = Path(checkpoint)
    return p if p.is_absolute() else Path(vae_dir) / p


def _unwrap_state(state: Any) -> dict[str, torch.Tensor]:
    if isinstance(state, dict):
        if "state_dict" in state and isinstance(state["state_dict"], dict):
            return state["state_dict"]
        if "model" in state and isinstance(state["model"], dict):
            return state["model"]
        return state
    raise TypeError(f"Unsupported checkpoint payload type: {type(state)!r}")


def load_pretrained(
    checkpoint: Optional[str | Path] = None,
    vae_dir: str | Path = DEFAULT_TACTILE_VAE_DIR,
    map_location: str | torch.device = "cpu",
    freeze: bool = True,
    strict: bool = True,
    model_kwargs: Optional[dict[str, Any]] = None,
) -> TactileVAE:
    ckpt_path = _resolve_checkpoint(checkpoint, vae_dir)
    if not ckpt_path.exists():
        raise FileNotFoundError(f"Tactile VAE checkpoint not found at: {ckpt_path}")

    state = torch.load(str(ckpt_path), map_location=map_location)
    state_dict = _unwrap_state(state)

    model = TactileVAE(**(model_kwargs or {}))
    model.load_state_dict(state_dict, strict=strict)

    if freeze:
        model.eval()
        for p in model.parameters():
            p.requires_grad_(False)
    return model


if __name__ == "__main__":
    # Architecture roundtrip test (no real checkpoint required).
    # To test with a real checkpoint, set CKPT_PATH below.
    _CKPT_PATH = Path("/group2/ct/weihanx/tactile_world_model/tactile_wm/pretrained_models/ckpt_best.pt")
    _OUT_DIR = Path("/group2/ct/weihanx/tactile_world_model/tactile_vae/test_output")
    _EPISODE_PATH = Path("/group2/ct/weihanx/tactile_world_model/mode1_v1/0323_episode_000.pt")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"device: {device}")

    # ── 1. Architecture test with random weights ─────────────────────────────
    print("\n[1] Architecture roundtrip with random weights")
    vae = TactileVAE().eval().to(device)
    x_rand = torch.randn(2, 3, 128, 128, device=device)
    with torch.no_grad():
        out = vae(x_rand)
    print(f"  input:  {tuple(x_rand.shape)}")
    print(f"  z:      {tuple(out['z'].shape)}")
    print(f"  x_hat:  {tuple(out['x_hat'].shape)}")

    # Separate encode → decode
    with torch.no_grad():
        mu, logvar, _ = vae.encode(x_rand)
        x_hat2 = vae.decode(mu)
    print(f"  encode → z (mu): {tuple(mu.shape)}")
    print(f"  decode → x_hat:  {tuple(x_hat2.shape)}")

    # ── 2. TactileVAEWrapper with real checkpoint (if available) ──────────────
    print(f"\n[2] TactileVAEWrapper from checkpoint: {_CKPT_PATH}")
    if not _CKPT_PATH.exists():
        print("  checkpoint not found — skipping pretrained test")
    else:
        ep = torch.load(str(_EPISODE_PATH), map_location="cpu", weights_only=False)
        views = ep["view"]  # (T, 3, H, W) uint8
        sample_indices = [0, 100, 500, 1000, 2000]
        frames_u8 = views[sample_indices]  # (N, 3, H, W)
        frames = frames_u8.float() / 127.5 - 1.0  # [-1, 1]

        wrapper = TactileVAEWrapper(str(_CKPT_PATH), device=device)
        z = wrapper.encode(frames)
        x_hat = wrapper.decode(z)
        print(f"  frames: {tuple(frames.shape)}  z: {tuple(z.shape)}  x_hat: {tuple(x_hat.shape)}")

        _OUT_DIR.mkdir(parents=True, exist_ok=True)
        for i, idx in enumerate(sample_indices):
            orig_np = ((frames[i].permute(1, 2, 0).clamp(-1, 1) + 1) * 127.5).byte().numpy()
            recon_np = ((x_hat[i].permute(1, 2, 0).clamp(-1, 1) + 1) * 127.5).byte().cpu().numpy()
            diff_np = (np.abs(orig_np.astype(int) - recon_np.astype(int))).astype(np.uint8)
            h, w = orig_np.shape[:2]
            panel = Image.new("RGB", (3 * w + 16, h), (20, 20, 20))
            panel.paste(Image.fromarray(orig_np), (0, 0))
            panel.paste(Image.fromarray(recon_np), (w + 8, 0))
            panel.paste(Image.fromarray(diff_np), (2 * w + 16, 0))
            panel.save(_OUT_DIR / f"vae_frame_{idx:05d}_panel.png")
            mse = float(((orig_np.astype(float) - recon_np.astype(float)) ** 2).mean())
            psnr = 10 * np.log10(255.0 ** 2 / mse) if mse > 0 else float("inf")
            print(f"  frame {idx:5d}  PSNR={psnr:.2f} dB")
        print(f"  saved panels to {_OUT_DIR}")

    print("\nAll tests passed.")