# Written by yams_any4d agent, 2026-07-13.
"""Compute per-channel (D,) mean/std of cached DiT features for --feat_norm.

Streams over ep*.pt in a cache dir with a numerically-stable running
mean/variance (Welford over flattened Tl*Hp*Wp tokens per channel), so it never
holds all windows in memory. Output npz has keys mean/std/n matching the format
train_v4head_cached.py expects.

Usage (yukon, a4d-bw):
  python custom/eval/compute_feat_norm.py \
    --cache_dir=$HOME/any4d_work/feat_cache_640ft \
    --out=$HOME/any4d_work/feat_norm_640ft.npz
"""

import argparse
import glob

import numpy as np
import torch


def main(args):
    files = sorted(glob.glob(f'{args.cache_dir}/ep*.pt'))
    assert files, f'no cache files in {args.cache_dir}'
    ssum = None
    ssq = None
    n = 0
    for i, f in enumerate(files):
        feats = torch.load(f, map_location='cpu', weights_only=False)['feats0']
        x = feats.reshape(-1, feats.shape[-1]).double().numpy()  # (tokens, D)
        if ssum is None:
            ssum = np.zeros(x.shape[1], dtype=np.float64)
            ssq = np.zeros(x.shape[1], dtype=np.float64)
        ssum += x.sum(0)
        ssq += (x * x).sum(0)
        n += x.shape[0]
        if (i + 1) % 200 == 0:
            print(f'  {i + 1}/{len(files)} windows, n={n} tokens')
    mean = ssum / n
    var = np.maximum(ssq / n - mean * mean, 0.0)
    std = np.sqrt(var * n / max(n - 1, 1))
    np.savez(args.out, mean=mean.astype(np.float32), std=std.astype(np.float32),
             n=np.int64(n))
    print(f'saved {args.out} | n={n} tokens | mean|max|={np.abs(mean).max():.1f} '
          f'std range [{std.min():.2f},{std.max():.1f}]')


if __name__ == '__main__':
    p = argparse.ArgumentParser()
    p.add_argument('--cache_dir', required=True)
    p.add_argument('--out', required=True)
    main(p.parse_args())