# Written by yams_any4d agent, 2026-07-12. Diagnostic: PCA of DiT features across
# (denoising sigma x video frame) for ONE window, teacher-forced. Shared PCA basis.
import argparse
import numpy as np
import cv2
import torch

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

SIGMAS = [80.0, 40.0, 20.0, 10.0, 5.0, 2.0, 1.0, 0.5, 0.1, 0.02]


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: {}

    dataloader = base.load_dataset(args, exp_cfg)
    batch = next(iter(dataloader))

    orig_draw = model.draw_training_noise
    feats_by_sigma = {}
    import copy as _copy
    for sig in SIGMAS:
        def fixed_draw(y0_shapes, harmonize_streams=True, harmonize_frames=True, _s=sig):
            (sigmas, epsilons) = orig_draw(y0_shapes, harmonize_streams=harmonize_streams,
                                           harmonize_frames=harmonize_frames)
            return ({k: torch.full_like(v, _s) for (k, v) in sigmas.items()}, epsilons)
        model.draw_training_noise = fixed_draw
        with torch.no_grad():
            model.training_step(_copy.deepcopy(batch), iteration=0)
        f = model.pipe.dit._v4head_feats[0][0].detach().float().cpu().numpy()  # (Tl,Hp,Wp,D)
        feats_by_sigma[sig] = f
        print(f'sigma {sig}: feats {f.shape} std {f.std():.3f}')

    # shared PCA basis across ALL (sigma, frame) tokens
    allf = np.concatenate([f.reshape(-1, f.shape[-1]) for f in feats_by_sigma.values()])
    allf = allf - allf.mean(0, keepdims=True)
    sub = allf[:: max(1, len(allf) // 40000)]
    _, _, Vt = np.linalg.svd(sub, full_matrices=False)
    W3 = Vt[:3].T
    lo = None; hi = None
    projs = {}
    for (sig, f) in feats_by_sigma.items():
        p = (f.reshape(-1, f.shape[-1]) - allf.mean(0, keepdims=True) * 0) @ W3
        projs[sig] = p
        lo = p.min(0) if lo is None else np.minimum(lo, p.min(0))
        hi = p.max(0) if hi is None else np.maximum(hi, p.max(0))

    (Tl, Hp, Wp, D) = feats_by_sigma[SIGMAS[0]].shape
    cw, ch = 120, int(120 * Hp / Wp)
    label_w = 90
    rows = []
    for sig in SIGMAS:
        p = (projs[sig] - lo) / (hi - lo + 1e-8)
        imgs = (p.reshape(Tl, Hp, Wp, 3) * 255).astype(np.uint8)
        cells = [np.full((ch, label_w, 3), 30, np.uint8)]
        cv2.putText(cells[0], f's={sig:g}', (4, ch // 2 + 4),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255, 255, 255), 1, cv2.LINE_AA)
        for t in range(Tl):
            cells.append(cv2.resize(imgs[t], (cw, ch), interpolation=cv2.INTER_NEAREST))
            cells.append(np.full((ch, 2, 3), 255, np.uint8))
        rows.append(np.concatenate(cells[:-1], axis=1))
        rows.append(np.full((2, rows[-1].shape[1], 3), 255, np.uint8))
    grid = np.concatenate(rows[:-1], axis=0)
    cv2.imwrite(args.out, cv2.cvtColor(grid, cv2.COLOR_RGB2BGR))
    print('saved', args.out, grid.shape)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--exp_cfg', required=True)
    parser.add_argument('--ckpt_path', required=True)
    parser.add_argument('--dset_cfg', required=True)
    parser.add_argument('--out', default='/tmp/pca_sigma_grid.png')
    parser.add_argument('--output_dir', default='/tmp/pca_grid_tmp')
    parser.add_argument('--batch_size', type=int, default=1)
    parser.add_argument('--device', default='cuda')
    parser.add_argument('--gpu_id', type=int, default=1)
    parser.add_argument('--debug', action='store_true')
    parser.add_argument('--data_overrides', default='{"subsample": null, "single_sample": "first"}')
    parser.add_argument('--data_defaults', default=None)
    parser.add_argument('--any4d_overrides', default=None)
    _a = parser.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)