"""Run the fixed 128-frame HMP on overlapping windows of a full sequence."""
import copy
import json
from pathlib import Path
import numpy as np
import torch
from scipy.spatial.transform import Rotation


def window_starts(length, size=128, overlap=16):
    if length < 1 or not 0 <= overlap < size:
        raise ValueError('Invalid sequence length or overlap')
    starts = [0]
    while starts[-1] + size < length:
        starts.append(starts[-1] + size - overlap)
    return starts


def slice_window(data, start, length, size=128):
    result = {}
    for key, value in data.items():
        if isinstance(value, (torch.Tensor, np.ndarray)) and value.ndim >= 2 and value.shape[1] == length:
            clip = value[:, start:min(start + size, length)].clone() if isinstance(value, torch.Tensor) else value[:, start:min(start + size, length)].copy()
            valid = clip.shape[1]
            if valid < size:
                if isinstance(clip, torch.Tensor):
                    pad = clip[:, -1:].repeat_interleave(size-valid, dim=1)
                    clip = torch.cat([clip, pad], dim=1)
                else:
                    clip = np.concatenate([clip, np.repeat(clip[:, -1:], size-valid, axis=1)], axis=1)
                if key == 'vis_mask':
                    clip[:, valid:] = 0
                elif key == 'joints2d':
                    clip[:, valid:, :, -1] = 0
            result[key] = clip
        else:
            result[key] = copy.deepcopy(value)
    return result


def blend_windows(original, clips, size=128, overlap=16):
    result = {k: np.array(v, copy=True) for k, v in original.items()}
    B, T = result['trans'].shape[:2]
    denom = np.zeros(T, dtype=np.float64)
    trans = np.zeros((B, T, 3), dtype=np.float64)
    rotation_keys = ['root_orient', 'pose_body']
    if all('decode_root' in clip for _, clip in clips):
        result['decode_root'] = np.zeros((B, T, 3), dtype=np.float32)
        rotation_keys.append('decode_root')
    accum, refs = {}, {}
    for key in rotation_keys:
        shape = result[key].shape
        refs[key] = Rotation.from_rotvec(result[key].reshape(-1, 3)).as_quat().reshape(*shape[:-1], 4)
        accum[key] = np.zeros_like(refs[key])
    for start, clip in clips:
        n = min(size, T-start)
        w = np.ones(n, dtype=np.float64)
        if start > 0:
            m = min(overlap, n)
            w[:m] *= np.arange(1, m+1)/(m+1)
        if start + size < T:
            m = min(overlap, n)
            w[-m:] *= np.arange(m, 0, -1)/(m+1)
        denom[start:start+n] += w
        trans[:, start:start+n] += clip['trans'][:, :n] * w[None, :, None]
        for key in rotation_keys:
            a = clip[key][:, :n].reshape(result[key][:, start:start+n].shape)
            q = Rotation.from_rotvec(a.reshape(-1, 3)).as_quat().reshape(*a.shape[:-1], 4)
            reference = refs[key][:, start:start+n]
            q *= np.where(np.sum(q*reference, axis=-1, keepdims=True) < 0, -1., 1.)
            weights = w.reshape((1,n) + (1,)*(q.ndim-2))
            accum[key][:, start:start+n] += q * weights
    if np.any(denom <= 0):
        raise ValueError('Uncovered frames in HMP windows')
    result['trans'] = (trans/denom[None,:,None]).astype(np.float32)
    for key, q in accum.items():
        if np.any(np.linalg.norm(q, axis=-1) < 1e-8):
            raise ValueError('Degenerate quaternion blend')
        result[key] = Rotation.from_quat(q.reshape(-1,4)).as_rotvec().reshape(result[key].shape).astype(np.float32)
    # HMP currently fixes shape. Reject unhandled per-window shape changes.
    for _, clip in clips:
        if not np.allclose(clip['betas'], result['betas'], atol=1e-5):
            raise ValueError('Windowed HMP requires fixed shape (opt_betas=false)')
    result['latent_pose'] = result['pose_body'].copy()
    return result


def fit_windowed(obs_data, res_dicts, opt, data_args, out_dir, device, fit_fn):
    from body_model import MANO
    original = {k: v.detach().cpu().numpy() if isinstance(v, torch.Tensor) else np.asarray(v) for k,v in res_dicts[0].items()}
    T = original['trans'].shape[1]
    starts = window_starts(T)
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    mano_cfg = {k.lower():v for k,v in dict(opt.MANO).items()}
    hand_model = MANO(batch_size=128, pose2rot=True, **mano_cfg).to(device)
    clips = []
    manifest = {'frames':T, 'window_size':128, 'overlap':16, 'windows':[]}
    for index, start in enumerate(starts):
        clip_dir = out / 'windows' / f'{start:06d}'
        path = clip_dir / f'{Path(opt.HMP.vid_path).stem}_000000_world_results.npz'
        print(f'HMP window {index+1}/{len(starts)}: frames {start}:{min(start+128,T)}', flush=True)
        if path.is_file():
            with np.load(path) as cached:
                compatible = np.allclose(cached['betas'], original['betas'], atol=1e-5)
            if not compatible:
                path.rename(path.with_suffix('.variable_shape_backup.npz'))
        if not path.is_file():
            obs = slice_window(obs_data, start, T)
            res = slice_window(res_dicts[0], start, T)
            fit_fn(obs, (res,), hand_model, opt, data_args, str(clip_dir), device)
        with np.load(path) as saved:
            clip = {k:saved[k].copy() for k in saved.files}
        if clip['trans'].shape[1] != 128 or not all(np.isfinite(v).all() for v in clip.values() if v.dtype.kind in 'fc'):
            raise ValueError(f'Invalid HMP window output: {path}')
        clips.append((start,clip))
        manifest['windows'].append({'start':start,'valid_frames':min(128,T-start),'path':str(path)})
        (out/'window_manifest.json').write_text(json.dumps(manifest,indent=2)+'\n')
    result = blend_windows(original,clips)
    target = out/f'{Path(opt.HMP.vid_path).stem}_000000_world_results.npz'
    temporary = target.with_suffix('.partial.npz')
    np.savez(temporary,**result)
    temporary.replace(target)
    print(f'Saved full-length windowed HMP: {target}', flush=True)