# -*- coding: utf-8 -*-
"""
Unified Table2 evaluation script (Retarget Error), supporting VAE and RVQ models with sliding window strategy.
Based on the sliding window implementation from table1_unified_sliding.py.
"""
import os, argparse, tqdm, torch, gc, time
import numpy as np
from sata.mypath import *
from sata.mydataset import PairedDataset, get_mi_src_tgt_all_graph, PairedGraph_collate_fn
from sata.skel_pose_graph import SkelPoseGraph
from sata.conversions.graph_to_motion import hatD_recon_motion, gt_recon_motion


def load_model_by_type(model_type, model_epoch, device):
    """
    Dynamically load the corresponding model based on model type.
    
    Args:
        model_type: "vae" or "rvq"
        model_epoch: model checkpoint path
        device: device
    
    Returns:
        model, cfg, ms_dict
    """
    if model_type == "vae":
        from sata.test import prepare_model_test
        print(f"[Model] Loading VAE model from: {model_epoch}")
    elif model_type == "rvq":
        from sata.test_vq import prepare_model_test
        print(f"[Model] Loading RVQ model from: {model_epoch}")
    else:
        raise ValueError(f"Unknown model_type: {model_type}. Must be 'vae' or 'rvq'")
    
    model, cfg, ms_dict = prepare_model_test(model_epoch, device)
    return model, cfg, ms_dict


def extract_window_from_dataset(dataset, mi, src_ri, tgt_ri, start_frame, window_size, total_frames):
    """
    Extract a window from the dataset directly from its internal structure (faster than loading frame by frame).
    
    Args:
        dataset: PairedDataset instance
        mi: motion index
        src_ri: source rig index
        tgt_ri: target rig index
        start_frame: window start frame
        window_size: window size
        total_frames: total number of frames
    
    Returns:
        (src_window, tgt_window), actual_window_size
    """
    end_frame = min(start_frame + window_size, total_frames)
    actual_window_size = end_frame - start_frame
    
    # get file indices
    src_fi = dataset.mi_ri_2_fi[mi][src_ri]
    tgt_fi = dataset.mi_ri_2_fi[mi][tgt_ri]
    
    # get skeleton data (shared across all frames)
    src_skel = dataset.skel_list[src_fi]
    tgt_skel = dataset.skel_list[tgt_fi]
    
    # get pose data for window frames (from pose_list)
    src_start_idx = dataset.start_frames[src_fi] + start_frame
    tgt_start_idx = dataset.start_frames[tgt_fi] + start_frame
    
    src_pose_window = dataset.pose_list[src_start_idx:src_start_idx + actual_window_size]
    tgt_pose_window = dataset.pose_list[tgt_start_idx:tgt_start_idx + actual_window_size]
    
    # build graph list for window
    src_graphs = [SkelPoseGraph(src_skel, pose) for pose in src_pose_window]
    tgt_graphs = [SkelPoseGraph(tgt_skel, pose) for pose in tgt_pose_window]
    
    # collate into batch
    window_data = list(zip(src_graphs, tgt_graphs))
    src_window, tgt_window = PairedGraph_collate_fn(window_data)
    
    return (src_window, tgt_window), actual_window_size


def process_with_sliding_window(model, model_type, ds, mi, src_ri, tgt_ri, total_frames, 
                                window_size, overlap, device, debug=False):
    """
    Process a motion with sliding windows, supporting overlap.
    
    Strategy: in overlapping regions, keep only the prediction from the earlier window.
    - First window: keep all frames
    - Subsequent windows: discard the first `overlap` frames (the overlapping part), keep only the non-overlapping part
    
    Args:
        model: trained model
        model_type: "vae" or "rvq"
        ds: PairedDataset instance
        mi: motion index
        src_ri: source rig index
        tgt_ri: target rig index
        total_frames: total number of frames
        window_size: window size
        overlap: overlap size between consecutive windows
        device: device
        debug: print debug information
    
    Returns:
        z_full: concatenated z across all windows (overlaps discarded)
        hatD_full: concatenated hatD across all windows (overlaps discarded)
        tgt_batch_full: full target batch for GT computation
    """
    stride = window_size - overlap
    
    # compute number of windows
    if total_frames <= window_size:
        num_windows = 1
    else:
        num_windows = (total_frames - window_size + stride - 1) // stride + 1
    
    if debug:
        print(f"\n[DEBUG] Processing motion {mi} (src_ri={src_ri}, tgt_ri={tgt_ri}):")
        print(f"  Total frames: {total_frames}")
        print(f"  Window size: {window_size}, Overlap: {overlap}, Stride: {stride}")
        print(f"  Number of windows: {num_windows}")
        print(f"  Model type: {model_type}")
    
    # determine number of nodes per frame from dataset structure
    # important: use tgt_skel node count since hatD is output for the target skeleton
    tgt_fi = ds.mi_ri_2_fi[mi][tgt_ri]
    tgt_skel = ds.skel_list[tgt_fi]
    num_nodes_per_frame = tgt_skel.lo.shape[0]  # number of joints in target skeleton
    
    if debug:
        print(f"  Detected {num_nodes_per_frame} nodes per frame")
    
    # collect non-overlapping parts from each window
    z_parts = []
    hatD_parts = []
    
    # process each window
    for window_idx in range(num_windows):
        start_frame = window_idx * stride
        
        # skip if start_frame exceeds total_frames
        if start_frame >= total_frames:
            if debug:
                print(f"  Window {window_idx}: start_frame {start_frame} >= {total_frames}, skipping")
            break
            
        # extract window from dataset
        (src_window, tgt_window), window_frames = extract_window_from_dataset(
            ds, mi, src_ri, tgt_ri, start_frame, window_size, total_frames
        )
        
        # move to device
        src_window = src_window.to(device)
        tgt_window = tgt_window.to(device)
        
        # call model according to type (handles different return value counts)
        with torch.no_grad():
            if model_type == "vae":
                # VAE model returns 2 values
                z_win, hatD_win = model(src_window, tgt_window, window_frames)
            elif model_type == "rvq":
                # RVQ model returns 4 values
                z_win, hatD_win, _, _ = model(src_window, tgt_window, window_frames)
            else:
                raise ValueError(f"Unknown model_type: {model_type}")
        
        # hatD shape is [T*N, D] where T=frames, N=nodes_per_frame
        hatD_dim = hatD_win.shape[1]
        
        # reshape hatD from [T*N, D] to [T, N, D]
        hatD_win_reshaped = hatD_win.view(window_frames, num_nodes_per_frame, hatD_dim)
        
        # determine which frames to keep from this window
        if window_idx == 0:
            # first window: keep all frames
            keep_start_idx = 0
            keep_end_idx = window_frames
        else:
            # subsequent windows: discard overlapping part, keep only non-overlapping part
            keep_start_idx = overlap
            keep_end_idx = window_frames
        
        # extract non-overlapping part
        z_keep = z_win[keep_start_idx:keep_end_idx]
        hatD_keep = hatD_win_reshaped[keep_start_idx:keep_end_idx]
        
        if debug:
            print(f"  Window {window_idx}: frames [{start_frame}, {start_frame + window_frames})")
            print(f"    Keep: [{keep_start_idx}, {keep_end_idx}), z_keep: {z_keep.shape}, hatD_keep: {hatD_keep.shape}")
        
        # append to lists
        z_parts.append(z_keep)
        hatD_parts.append(hatD_keep)
        
        # cleanup
        del src_window, tgt_window, z_win, hatD_win, hatD_win_reshaped
    
    # concatenate all parts
    z_full = torch.cat(z_parts, dim=0)  # [T, D]
    hatD_full_3d = torch.cat(hatD_parts, dim=0)  # [T, N, D]
    
    # reshape hatD back to [T*N, D] format (expected by hatD_recon_motion)
    actual_frames = z_full.shape[0]
    hatD_full = hatD_full_3d.view(actual_frames * num_nodes_per_frame, -1)
    
    # verify frame count is correct
    if actual_frames != total_frames:
        if debug:
            print(f"  WARNING: Expected {total_frames} frames, got {actual_frames} frames")
    
    # check for NaN or Inf
    if torch.isnan(z_full).any() or torch.isinf(z_full).any():
        raise ValueError("NaN or Inf detected in z_full")
    if torch.isnan(hatD_full).any() or torch.isinf(hatD_full).any():
        raise ValueError("NaN or Inf detected in hatD_full")
    
    if debug:
        print(f"  Final shapes: z_full {z_full.shape}, hatD_full {hatD_full.shape}")
        print(f"  Total frames collected: {actual_frames} (expected: {total_frames})")
    
    # load the full target batch for GT computation
    (_, tgt_batch_full), _ = get_mi_src_tgt_all_graph(ds, mi, src_ri, tgt_ri, device)
    tgt_batch_full = tgt_batch_full.to(device)
    
    return z_full, hatD_full, tgt_batch_full


def compute_retarget_error_with_sliding_window(
    model, model_type, cfg, ms_dict, ds, mi, src_ri, tgt_ri, 
    total_frames, window_size, overlap, device, debug=False
):
    """
    Compute retarget positional error using a sliding window (consistent with table2 metrics).
    
    Args:
        model: trained model
        model_type: "vae" or "rvq"
        cfg: model config
        ms_dict: mean/std dictionary
        ds: PairedDataset instance
        mi: motion index
        src_ri: source rig index
        tgt_ri: target rig index
        total_frames: total number of frames
        window_size: window size
        overlap: overlap size
        device: device
        debug: whether to print debug information
    
    Returns:
        error: positional error (computed consistently with table2)
    """
    # 1. get full z and hatD using sliding window
    z_full, hatD_full, tgt_batch = process_with_sliding_window(
        model, model_type, ds, mi, src_ri, tgt_ri, 
        total_frames, window_size, overlap, device, debug=debug
    )
    
    # 2. reconstruct motion
    # note: GT tgt_motion is needed to compute the error
    out_motion_list, _ = hatD_recon_motion(
        hatD_full, tgt_batch, cfg["representation"]["out"], 
        ms_dict, total_frames
    )
    
    tgt_motion_list, _ = gt_recon_motion(tgt_batch, total_frames)
    
    # 3. compute positional error (consistent with table2.py)
    # Metric from <Skeleton-Aware Networks for Deep Motion Retargeting>
    # https://github.com/DeepMotionEditing/deep-motion-editing/blob/master/retargeting/get_error.py#L47-L55
    
    # character height: Max(joints' height at t-pose)
    height = tgt_batch.go[:, 1].max().item()
    
    pos_ref = tgt_motion_list[0].positions(local=False)  # GT positions
    pos = out_motion_list[0].positions(local=False)      # Predicted positions
    
    # compute error
    err = (pos - pos_ref) * (pos - pos_ref)
    err /= height**2
    
    # cleanup
    del z_full, hatD_full, tgt_batch, out_motion_list, tgt_motion_list
    
    return err.mean() * 1000


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_type", type=str, required=True, choices=["vae", "rvq"],
                        help="Model type: vae or rvq (REQUIRED)")
    parser.add_argument("--model_epoch", type=str, default="ckpt0")
    parser.add_argument("--device", type=str, default="cuda:0")
    parser.add_argument("--data_dir", type=str, default="evaluation/motion/processed")
    parser.add_argument("--window_size", type=int, default=64,
                        help="Sliding window size (number of frames)")
    parser.add_argument("--overlap", type=int, default=16,
                        help="Overlap size between consecutive windows")
    parser.add_argument("--num_samples", type=int, default=-1,
                        help="Number of samples to evaluate (-1 for all)")
    args = parser.parse_args()
    
    print("="*80)
    print("Unified Table2 Evaluation (Retarget Error, with sliding window)")
    print("="*80)
    print(f"Model type: {args.model_type.upper()}")
    print(f"Model: {args.model_epoch}")
    print(f"Device: {args.device}")
    print(f"Data dir: {args.data_dir}")
    print(f"Window size: {args.window_size}")
    print(f"Overlap: {args.overlap}")
    print(f"Stride: {args.window_size - args.overlap}")
    print("="*80)
    
    # load model by type
    model, cfg, ms_dict = load_model_by_type(args.model_type, args.model_epoch, args.device)
    
    # dataset
    ds = PairedDataset(min_motion_lens=8)
    data_dir = os.path.join(DATA_DIR, args.data_dir)
    ds.load_data_dir_pairs(data_dir)
    
    # determine test range based on num_samples
    total_samples = len(ds.mi_ri_2_fi)
    if args.num_samples > 0 and args.num_samples < total_samples:
        N = args.num_samples
    else:
        N = total_samples
    
    print(f"\nEvaluating {N} samples out of {total_samples} total samples\n")
    
    int_err, cross_err = 0.0, 0.0
    int_cnt, cross_cnt = 0, 0
    
    model.eval()  # set model to evaluation mode
    
    st = time.time()
    
    for mi in tqdm.tqdm(range(N)):
        # get total frame count and number of rigs for this motion
        fi = ds.mi_ri_2_fi[mi][0]
        total_frames = ds.frame_cnts[fi]
        R = len(ds.mi_ri_2_fi[mi])
        
        # enable debug only for first sample
        debug_mode = (mi == 0)
        
        # Comparison with <Skeleton-Aware Networks for Deep Motion Retargeting>
        # cross: BigVegas -> Goblin_m, Mousey_m, Mremireh_m, Vampire_m
        # internal: Goblin_m, Mousey_m, Mremireh_m, Vampire_m <->
        # https://github.com/DeepMotionEditing/deep-motion-editing/blob/master/retargeting/test.py
        
        # Cross retargeting: src_ri=0 -> tgt_ri=1,2,...,R-1
        for tgt_ri in range(1, R):
            err = compute_retarget_error_with_sliding_window(
                model, args.model_type, cfg, ms_dict, ds, mi, 
                src_ri=0, tgt_ri=tgt_ri, total_frames=total_frames,
                window_size=args.window_size, overlap=args.overlap, 
                device=args.device, debug=(debug_mode and tgt_ri == 1)
            )
            cross_err += err.item()
            cross_cnt += 1
            
            # # free memory promptly
            # gc.collect()
            # torch.cuda.empty_cache()
        
        # Internal retargeting: src_ri=1,2,...,R-1 -> tgt_ri=1,2,...,R-1
        for src_ri in range(1, R):
            for tgt_ri in range(1, R):
                err = compute_retarget_error_with_sliding_window(
                    model, args.model_type, cfg, ms_dict, ds, mi, 
                    src_ri=src_ri, tgt_ri=tgt_ri, total_frames=total_frames,
                    window_size=args.window_size, overlap=args.overlap, 
                    device=args.device, debug=False
                )
                int_err += err.item()
                int_cnt += 1
                
                # # free memory promptly
                # gc.collect()
                # torch.cuda.empty_cache()
    
    elapsed_time = time.time() - st
    
    # print results
    print("\n" + "="*80)
    print("Results:")
    print("="*80)
    print(f"Internal error: {int_err/int_cnt:.4f}")
    print(f"Cross error: {cross_err/cross_cnt:.4f}")
    print(f"Time elapsed: {elapsed_time:.2f}s")
    print("="*80)
    
    # save results to file
    metric_fp = os.path.join(RESULT_DIR, args.model_epoch.split("/")[0], f"table2_unified_{args.model_type}_sliding.txt")
    os.makedirs(os.path.dirname(metric_fp), exist_ok=True)
    
    if not os.path.exists(metric_fp):
        with open(metric_fp, "w") as f:
            f.write("model_epoch,\tmodel_type,\twindow_size,\toverlap,\tinternal,\tcross\n")
    
    with open(metric_fp, "a") as f:
        f.write(
            f"{args.model_epoch},\t{args.model_type},\t{args.window_size},\t{args.overlap},\t"
            f"{int_err/int_cnt:.5f},\t{cross_err/cross_cnt:.5f}\n"
        )
    
    print(f"\nResults saved to: {metric_fp}")


'''
Usage examples:

# RVQ model
python src/task/evaluation/table2_unified_sliding.py \
--model_type rvq \
--model_epoch same_rvq_lowerLR_512_wtextV3_nomask_newLoss_gps_t64_temporal_trans_trainFaceZMirror \
--data_dir evaluation/motion/processed \
--window_size 64 \
--overlap 16

# VAE model
python src/task/evaluation/table2_unified_sliding.py \
--model_type vae \
--model_epoch same_vae_lowerLR_512_wtextv3_nomask_newLoss_gps_t64_temporal_trans_trainFaceZMirror \
--data_dir evaluation/motion/processed \
--window_size 64 \
--overlap 16

# quick test (small number of samples)
python src/task/evaluation/table2_unified_sliding.py \
--model_type vae \
--model_epoch xxx \
--data_dir evaluation/motion/processed \
--num_samples 2 \
--window_size 64 \
--overlap 16

# different window size and overlap configurations
python src/task/evaluation/table2_unified_sliding.py \
--model_type rvq \
--model_epoch xxx \
--data_dir evaluation/motion/processed \
--window_size 128 \
--overlap 32
'''