"""Encode every PushT frame with the frozen LeWM encoder into a latent cache.

Controller training never needs pixels: the encoder is frozen and no image
augmentation is used, so latents can be computed once. The full dataset is
2.34M frames x 192 dims = ~0.9 GB in fp16, which fits in RAM.

Also stores the action z-score statistics. LeWM was trained on z-scored
actions (``column_normalizer`` defaults to ``method='zscore'``), so the
controller must emit actions in that same normalized space.
"""

import argparse
import json
import time
from pathlib import Path

import h5py
import hdf5plugin  # noqa: F401  -- registers the blosc filter used by the h5
import numpy as np
import torch

from lejepa_control.world_model import load_lewm

IMAGENET_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
IMAGENET_STD = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--h5', default='data/swm_home/datasets/pusht_expert_train.h5'
    )
    parser.add_argument('--out', default='data/latents')
    parser.add_argument('--batch-size', type=int, default=512)
    parser.add_argument('--limit-episodes', type=int, default=None)
    parser.add_argument('--wm-name', default='quentinll/lewm-pusht')
    args = parser.parse_args()

    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model = load_lewm(name=args.wm_name, device=device)
    encoder, projector = model.encoder, model.projector
    D = model.predictor.input_dim

    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    mean = IMAGENET_MEAN.to(device)
    std = IMAGENET_STD.to(device)

    with h5py.File(args.h5, 'r') as f:
        lengths = f['ep_len'][:].astype(np.int64)
        offsets = f['ep_offset'][:].astype(np.int64)

        if args.limit_episodes is not None:
            lengths = lengths[: args.limit_episodes]
            offsets = offsets[: args.limit_episodes]

        n_frames = int(lengths.sum())
        end = int(offsets[-1] + lengths[-1])
        print(f'{len(lengths)} episodes, {n_frames} frames -> {D}-dim latents')

        actions = f['action'][:end].astype(np.float32)
        # the frame that ends an episode has no outgoing action (there is no
        # next state to leave it for) and is NaN-padded; ignore those rows for
        # the stats, then fill them with the mean so a block that straddles an
        # episode boundary (LatentGoalDataset's real_action, at the last valid
        # start) never hands training a NaN.
        nan_mask = np.isnan(actions).any(axis=1)
        a_mean = np.nanmean(actions, axis=0)
        a_std = np.nanstd(actions, axis=0)
        if nan_mask.any():
            print(f'{int(nan_mask.sum())} NaN (terminal-step) action rows -> filled with mean')
            actions[nan_mask] = a_mean
        print(f'action mean={a_mean} std={a_std}')

        latents = np.lib.format.open_memmap(
            out_dir / 'latents.npy',
            mode='w+',
            dtype=np.float16,
            shape=(n_frames, D),
        )

        t0 = time.perf_counter()
        for start in range(0, end, args.batch_size):
            stop = min(start + args.batch_size, end)
            frames = f['pixels'][start:stop]  # (B, 224, 224, 3) uint8

            x = torch.from_numpy(frames).to(device, non_blocking=True)
            x = x.permute(0, 3, 1, 2).float().div_(255.0).sub_(mean).div_(std)

            with torch.no_grad(), torch.autocast('cuda', dtype=torch.bfloat16):
                out = encoder(x, interpolate_pos_encoding=True)
                emb = projector(out.last_hidden_state[:, 0].float())

            latents[start:stop] = emb.float().cpu().numpy().astype(np.float16)

            if start % (args.batch_size * 200) == 0:
                done = stop / end
                rate = stop / (time.perf_counter() - t0)
                eta = (end - stop) / rate / 60
                print(
                    f'  {done:6.1%}  {rate:6.0f} img/s  eta {eta:5.1f} min',
                    flush=True,
                )

        latents.flush()

    np.save(out_dir / 'actions.npy', actions)
    np.save(out_dir / 'ep_len.npy', lengths)
    np.save(out_dir / 'ep_offset.npy', offsets)

    stats = {
        'action_mean': a_mean.tolist(),
        'action_std': a_std.tolist(),
        'latent_dim': int(D),
        'n_frames': n_frames,
        'n_episodes': len(lengths),
    }
    (out_dir / 'stats.json').write_text(json.dumps(stats, indent=2))

    z = np.asarray(latents[:10000], dtype=np.float32)
    print(f'latent per-coord std (mean) = {z.std(0).mean():.4f}')
    print(f'done in {(time.perf_counter() - t0) / 60:.1f} min -> {out_dir}')


if __name__ == '__main__':
    main()