# Written by yams_any4d agent, 2026-07-11.
"""
Cache DiT features for standalone action-head training (scene-only derisk pipeline).

For each dataset window: run the FROZEN video model's training-style forward once,
teacher-forced at a FIXED LOW SIGMA (monkeypatched draw_training_noise), and save the
stashed pre-final-layer features + GT action + an rgb thumbnail.

Output: <cache_dir>/ep<EEE>_w<NNNNN>.pt with
  feats0 (Tl, Hp, Wp, D) bf16 [scene], action (41, 20) fp32, sample_id str,
  ep int, rgb0_thumb (H, W, 3) uint8.

Usage (yukon, a4d-bw):
  PYTHONPATH=$PWD:$PWD/externals/AnyData python custom/eval/extract_v4feats.py \
    --exp_cfg=custom/experiment/basile/rwm4_pp70_scene_lora.py \
    --ckpt_path=<merged scene ckpt> \
    --dset_cfg=custom/config/anydata/web/debug/yamyukon_pp70_scene.yaml \
    --output_dir=$HOME/any4d_work/feat_cache --max_windows 4000 --sigma 0.05
"""

import argparse
import glob
import os
import re

import numpy as np
import torch

import custom.eval.infer_anydata as base
from custom.eval.infer_v4head_rerun import merge_lora_state_dict, gt_action_from_batch


def main(args):
    (model, config, _, device, exp_cfg) = base.setup_runtime(args)

    sd = torch.load(args.ckpt_path, map_location='cpu', weights_only=False)
    sd, merged = merge_lora_state_dict(sd)
    missing, unexpected = model.pipe.dit.load_state_dict(sd, strict=False)
    print(f'[cyan]ckpt loaded (merged={merged}); {len(missing)} missing / {len(unexpected)} unexpected')

    # The a4d_network feature stash is gated on `self.v4head is not None` — attach a
    # stub directly (the any4d_overrides plumbing does not reach the built DiT).
    if getattr(model.pipe.dit, 'v4head', None) is None:
        model.pipe.dit.v4head = torch.nn.Identity()
        model.pipe.dit._v4head_feats = None
        print('[cyan]attached v4head stub to enable feature stash')

    # Fix the training-noise sigma so every window's features come from the same
    # low-noise (deploy-like) readout point.
    fixed = float(args.sigma)
    orig_draw = model.draw_training_noise

    def fixed_draw(y0_shapes, harmonize_streams=True, harmonize_frames=True):
        (sigmas, epsilons) = orig_draw(y0_shapes, harmonize_streams=harmonize_streams,
                                       harmonize_frames=harmonize_frames)
        sigmas = {k: torch.full_like(v, fixed) for (k, v) in sigmas.items()}
        return (sigmas, epsilons)

    model.draw_training_noise = fixed_draw

    # visuals off: gallery rendering makes validation_step ~15x slower, and
    # detail=0 makes create_save_visuals_a4d ERROR (its failure-rate guard then
    # kills every batch after 10). No-op the function at the call site instead.
    import custom.any4d.a4d_model as _a4dm
    _a4dm.create_save_visuals_a4d = lambda *a, **k: {}
    _a4dm.calculate_metrics_a4d = lambda *a, **k: {}
    dataloader = base.load_dataset(args, exp_cfg)
    os.makedirs(args.output_dir, exist_ok=True)

    n_saved, ep_counts, epoch = 0, {}, 0
    # resume: count existing files so a restart skips covered episodes cheaply
    for f0 in glob.glob(f'{args.output_dir}/ep*.pt'):
        ep0 = int(os.path.basename(f0)[2:6])
        ep_counts[ep0] = ep_counts.get(ep0, 0) + 1
        n_saved += 1
    if n_saved:
        print(f'[cyan]resume: {n_saved} existing windows across {len(ep_counts)} eps')
    batches = iter(dataloader)
    batch_idx = -1
    while n_saved < args.max_windows:
        try:
            data_batch = next(batches)
        except StopIteration:
            epoch += 1
            if epoch > args.max_epochs:
                print(f'[yellow]exhausted {epoch} epochs at {n_saved} windows; stopping')
                break
            print(f'[cyan]epoch {epoch} done, re-iterating ({n_saved} windows so far)')
            batches = iter(dataloader)
            continue
        batch_idx += 1
        try:
            sample_id = str(data_batch.get('sample_id', ['unk'])[0])
            m = re.search(r'(\d{3,4})', sample_id)
            ep = int(m.group(1)) if m else -1
            if ep_counts.get(ep, 0) >= args.per_ep_cap:
                continue                                # quota met: no GPU spent
            gt_action = gt_action_from_batch(data_batch)                 # (41, 20)

            # rgb thumbnail (frame 0, scene cam) for later panel rendering
            rgb = data_batch['anydata']['rgb']
            k0 = sorted(rgb.keys())[0]
            fr = rgb[k0]
            fr = fr[0] if fr.ndim == 4 else fr                            # (3,H,W)
            fr = fr.float()
            # FIXED RANGE mapping (assume [-1,1] video convention), NOT per-frame
            # min-max stretch — stretched thumbs created a 3rd distribution that
            # broke bridge A/B probes (2026-07-12).
            if fr.min() < -0.01:
                fr = (fr + 1.0) / 2.0
            thumb = (fr.clamp(0, 1).permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)

            if args.generate:
                # DEPLOY-MATCHED extraction: full sampling conditioned on frame 0
                # only; features = last denoise step over GENERATED latents.
                gdir = dict(getattr(model.config, 'val_directives', {}) or {})
                gdir['cond_frames_raw'] = 1
                with torch.no_grad():
                    model.validation_step(
                        data_batch=data_batch, iteration=-1, dataloader_key='genx',
                        local_path=args.output_dir, directives=gdir,
                        val_iter=0, seed=batch_idx)
            else:
                with torch.no_grad():
                    model.training_step(data_batch, iteration=0)
            feats = model.pipe.dit._v4head_feats
            assert feats is not None and 0 in feats, 'no stash — enable v4head via any4d_overrides'

            out = {
                'feats0': feats[0][0].detach().to(torch.bfloat16).cpu(),  # (Tl,Hp,Wp,D)
                'action': gt_action.float().cpu(),
                'sample_id': sample_id, 'ep': ep,
                'rgb0_thumb': thumb, 'sigma': fixed,
            }
            torch.save(out, f'{args.output_dir}/ep{ep:04d}_w{n_saved:05d}.pt')
            ep_counts[ep] = ep_counts.get(ep, 0) + 1
            n_saved += 1
            if n_saved % 100 == 0:
                print(f'[green]{n_saved} windows cached ({len(ep_counts)} episodes)')
        except Exception as e:
            print(f'[yellow]skip batch {batch_idx}: {type(e).__name__}: {e}')

    print(f'[green]DONE: {n_saved} windows, episodes: {sorted(ep_counts.keys())}')
    print(f'[green]per-ep counts: {ep_counts}')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--exp_cfg', type=str, required=True)
    parser.add_argument('--ckpt_path', type=str, required=True)
    parser.add_argument('--dset_cfg', type=str, required=True)
    parser.add_argument('--output_dir', type=str, required=True)
    parser.add_argument('--max_windows', type=int, default=4000)
    parser.add_argument('--max_epochs', type=int, default=200)
    parser.add_argument('--per_ep_cap', type=int, default=50)
    parser.add_argument('--sigma', type=float, default=0.05)
    parser.add_argument('--generate', action='store_true',
                        help='extract from full generation (deploy-matched) instead of teacher-forced')
    parser.add_argument('--num_steps', type=int, default=25)
    parser.add_argument('--batch_size', type=int, default=1)
    parser.add_argument('--device', type=str, default='cuda')
    parser.add_argument('--gpu_id', type=int, default=0)
    parser.add_argument('--debug', action='store_true')
    parser.add_argument('--data_overrides', type=str, default='{"subsample": null, "single_sample": "random"}')
    parser.add_argument('--data_defaults', type=str, default=None)
    parser.add_argument('--any4d_overrides', type=str, default='{"v4head_enabled": true}')
    _args = parser.parse_args()
    _defaults = dict(exp_overrides=None, cond_aug_sigma=None, guidance=None, infer_metrics=False,
                     metric_resolution=None, num_samples=1, num_segments=4, perturb_action=None,
                     perturb_per_sample=False, perturb_traj=None, run_autoregressive=False,
                     shard=None, visual_detail=1, viz_extra_modes=None, donor_rgb_path=None,
                     donor_rgb_cam=0, extrapolation_strategy='backtrack',
                     stop_after=-1, seed=1234, cond_frames_raw=None)
    for (_k, _v) in _defaults.items():
        if not hasattr(_args, _k):
            setattr(_args, _k, _v)
    main(_args)