# Written by yams_any4d agent, 2026-07-12. Probe v2: are the DiT features content-poor
# because a few massive-activation channels dominate the norm? Test whether per-channel
# normalization / outlier removal recovers content discriminability between two videos.
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


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(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)
    A = feats(model, next(it), 2.0)   # (T,H,W,D)
    B = feats(model, next(it), 2.0)
    T, H, W, D = A.shape
    Xa = A.reshape(-1, D); Xb = B.reshape(-1, D)   # (N, D)

    # --- channel norm distribution: are a few channels huge? ---
    ch_absmean = np.abs(np.concatenate([Xa, Xb])).mean(0)   # (D,)
    order = np.argsort(ch_absmean)[::-1]
    tot = ch_absmean.sum()
    print('=' * 70)
    print(f'shape per video: {A.shape}  N_tokens={Xa.shape[0]}')
    top = order[:10]
    print('top-10 channels by |mean| :', [int(c) for c in top])
    print('their |mean| values       :', [round(float(ch_absmean[c]), 1) for c in top])
    print(f'top-10 channels carry {100*ch_absmean[top].sum()/tot:.1f}% of total |activation|')
    print(f'top-50 channels carry {100*ch_absmean[order[:50]].sum()/tot:.1f}% of total |activation|')
    print('-' * 70)

    # --- cross-video cosine under different normalizations ---
    print('CROSS-VIDEO cosine (batch A vs B, whole map). Lower = more content-discriminative.')
    print(f'  raw                         : {cos(Xa, Xb):.4f}')
    for k in (5, 20, 50):
        drop = order[:k]
        Ma = np.delete(Xa, drop, axis=1); Mb = np.delete(Xb, drop, axis=1)
        print(f'  drop top-{k:<3d} channels        : {cos(Ma, Mb):.4f}')
    # per-channel standardization using combined stats
    both = np.concatenate([Xa, Xb])
    mu = both.mean(0, keepdims=True); sd_ = both.std(0, keepdims=True) + 1e-6
    Za = (Xa - mu) / sd_; Zb = (Xb - mu) / sd_
    print(f'  per-channel z-score (center+scale): {cos(Za, Zb):.4f}')
    print(f'  per-channel center only     : {cos(Xa-mu, Xb-mu):.4f}')
    print('-' * 70)

    # --- does content contrast appear after normalization? per-token cross-video cosine ---
    def tokwise_mean_cos(Pa, Pb):
        na = Pa / (np.linalg.norm(Pa, axis=1, keepdims=True) + 1e-8)
        nb = Pb / (np.linalg.norm(Pb, axis=1, keepdims=True) + 1e-8)
        return float((na * nb).sum(1).mean())
    print('TOKEN-WISE mean cross-video cosine (per spatial token, averaged):')
    print(f'  raw                         : {tokwise_mean_cos(Xa, Xb):.4f}')
    print(f'  z-score                     : {tokwise_mean_cos(Za, Zb):.4f}')
    print('  (raw~1 but z-score low => content lives under the outlier channels)')
    print('=' * 70)
    print('PROBE2 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_probe2_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)