# Written by yams_any4d agent, 2026-07-12. Bug probe: are the stashed DiT features
# NUMERICALLY identical across (a) video frames and (b) denoising sigma, or do they
# only LOOK alike under top-3 PCA? Reports cosine similarities on the raw features.
import argparse
import numpy as np
import torch

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

SIGMAS = [10.0, 2.0, 0.02]


def cos(a, b):
    a = a.reshape(-1); b = b.reshape(-1)
    return float((a @ b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-8))


def feats_for_batch(model, batch, sig):
    orig = model.draw_training_noise

    def fixed_draw(y0_shapes, harmonize_streams=True, harmonize_frames=True, _s=sig):
        (s, e) = orig(y0_shapes, harmonize_streams=harmonize_streams,
                      harmonize_frames=harmonize_frames)
        return ({k: torch.full_like(v, _s) for (k, v) in s.items()}, e)
    model.draw_training_noise = fixed_draw
    import copy as _c
    with torch.no_grad():
        model.training_step(_c.deepcopy(batch), iteration=0)
    model.draw_training_noise = orig
    return model.pipe.dit._v4head_feats[0][0].detach().float().cpu().numpy()  # (T,H,W,D)


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, _ = merge_lora_state_dict(sd)
    model.pipe.dit.load_state_dict(sd, strict=False)
    if getattr(model.pipe.dit, 'v4head', None) is None:
        model.pipe.dit.v4head = torch.nn.Identity()
        model.pipe.dit._v4head_feats = None
    import custom.any4d.a4d_model as _a4dm
    _a4dm.create_save_visuals_a4d = lambda *a, **k: {}
    _a4dm.calculate_metrics_a4d = lambda *a, **k: {}

    dl = base.load_dataset(args, exp_cfg)
    it = iter(dl)
    batchA = next(it)
    batchB = next(it)

    # Fix RNG so epsilon is identical across sigma rows -> any feature diff is purely sigma.
    torch.manual_seed(0)
    fA = {s: feats_for_batch(model, batchA, s) for s in SIGMAS}
    fB = {s: feats_for_batch(model, batchB, s) for s in [2.0]}

    T = fA[SIGMAS[0]].shape[0]
    print('=' * 70)
    print(f'feature tensor per (sigma): shape {fA[SIGMAS[0]].shape}  (T,H,W,D)')
    print('per-sigma global std:', {s: round(float(fA[s].std()), 4) for s in SIGMAS})
    print('-' * 70)
    print('[A] CROSS-SIGMA cosine (whole map, batch A) — 1.0 == identical == BUG:')
    for i in range(len(SIGMAS)):
        for j in range(i + 1, len(SIGMAS)):
            print(f'    s={SIGMAS[i]:<5g} vs s={SIGMAS[j]:<5g} : {cos(fA[SIGMAS[i]], fA[SIGMAS[j]]):.4f}')
    print('-' * 70)
    print('[B] CROSS-FRAME cosine within sigma=2 (frame t vs t+1) — 1.0 == frames identical:')
    f2 = fA[2.0]
    for t in range(T - 1):
        print(f'    frame {t} vs {t+1} : {cos(f2[t], f2[t+1]):.4f}')
    print(f'    frame 0 vs {T-1} (ends) : {cos(f2[0], f2[T-1]):.4f}')
    print('-' * 70)
    print('[C] CONTENT TEST — same frame, DIFFERENT video (batch A vs B), sigma=2:')
    print('    if these are ~as high as cross-frame within A, features are NOT content-discriminative')
    for t in [0, T // 2, T - 1]:
        print(f'    frame {t}: cos(A,B) = {cos(fA[2.0][t], fB[2.0][t]):.4f}')
    print('-' * 70)
    # Per-position variance across frames: how much does each token move over the video?
    var_over_t = fA[2.0].reshape(T, -1, fA[2.0].shape[-1]).std(0).mean()
    var_over_space = fA[2.0].reshape(T, -1, fA[2.0].shape[-1]).std(1).mean()
    print(f'[D] mean std ACROSS FRAMES (temporal signal): {float(var_over_t):.4f}')
    print(f'    mean std ACROSS SPACE  (spatial signal) : {float(var_over_space):.4f}')
    print('    (if temporal << spatial, frames barely differ -> static-scene explanation)')
    print('=' * 70)
    print('PROBE EXIT OK')


if __name__ == '__main__':
    p = argparse.ArgumentParser()
    p.add_argument('--exp_cfg', required=True)
    p.add_argument('--ckpt_path', required=True)
    p.add_argument('--dset_cfg', required=True)
    p.add_argument('--output_dir', default='/tmp/feat_probe_tmp')
    p.add_argument('--batch_size', type=int, default=1)
    p.add_argument('--device', default='cuda')
    p.add_argument('--gpu_id', type=int, default=1)
    p.add_argument('--debug', action='store_true')
    p.add_argument('--data_overrides', default='{"subsample": null, "single_sample": null}')
    p.add_argument('--data_defaults', default=None)
    p.add_argument('--any4d_overrides', default=None)
    a = p.parse_args()
    d = 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=0, viz_extra_modes=None, donor_rgb_path=None,
             donor_rgb_cam=0, extrapolation_strategy='backtrack', num_steps=25,
             stop_after=-1, seed=1234, cond_frames_raw=None)
    for (k, v) in d.items():
        if not hasattr(a, k):
            setattr(a, k, v)
    main(a)