src/sata/conversions/same2bvh.py

"""
SAME to BVH Conversion Module

This module provides functionality to convert SAME format data (stored in .npz files)
to BVH format through the following pipeline:
    SAME (.npz) -> SkelPoseGraph -> Fairmotion Motion -> BVH (.bvh)

The conversion process:
1. Load SAME data from .npz file
2. Convert to SkelPoseGraph format (internal graph representation)
3. Convert SkelPoseGraph to Fairmotion Motion object
4. Save Motion object as BVH file using fairmotion's bvh.save()
"""

import os
import numpy as np
import torch
from fairmotion.data import bvh
from fairmotion.core import motion as motion_class

from sata.mydataset import SkelData, PoseData
from sata.skel_pose_graph import SkelPoseGraph


def _log(message, verbose=True):
    """Helper function to control verbose output"""
    if verbose:
        print(message)


def load_same_from_npz(npz_path, verbose=True, load_tf=True):
    """
    Load SAME format data from .npz file
    
    Args:
        npz_path: str, path to the .npz file containing SAME format data
        
    Returns:
        tuple: (SkelPoseGraph, num_frames)
            - SkelPoseGraph: The loaded SAME data as a graph object
            - num_frames: int, number of frames in the motion
        
    The .npz file should contain the following arrays:
        - lo: local offsets [nJ, 3]
        - go: global offsets [nJ, 3]
        - qb: quaternion boolean [nJ]
        - edges: edge indices [nE, 4] (parent, child, depth, reverse_depth)
        - q: joint rotations [nF, nJ, 6] (rotation 6D representation)
        - p: joint positions [nF, nJ, 3]
        - qv: joint rotation velocities [nF, nJ, 6]
        - pv: joint position velocities [nF, nJ, 3]
        - pprev: previous joint positions [nF, nJ, 3]
        - c: contact labels [nF, nJ, 1]
        - r: root motion [nF, 4] (theta, dx, dz, height)
    """
    if not os.path.exists(npz_path):
        raise FileNotFoundError(f"NPZ file not found: {npz_path}")
    
    _log(f"Loading SAME data from: {npz_path}", verbose)
    data = np.load(npz_path)
    
    # Extract skeleton data
    lo = torch.from_numpy(data['lo']).float()  # [nJ, 3]
    go = torch.from_numpy(data['go']).float()  # [nJ, 3]
    qb = torch.from_numpy(data['qb']).bool()   # [nJ]
    edges = data['edges']  # [nE, 4]
    
    # Build edge_index [2, nE]
    edge_index = torch.from_numpy(edges[:, :2].T).long()
    edge_feature = torch.from_numpy(edges[:, 2:]).long()  # [nE, 2] (depth, reverse_depth)
    
    # Extract pose data
    q = torch.from_numpy(data['q']).float()      # [nF, nJ, 6]
    p = torch.from_numpy(data['p']).float()      # [nF, nJ, 3]
    qv = torch.from_numpy(data['qv']).float()    # [nF, nJ, 6]
    pv = torch.from_numpy(data['pv']).float()    # [nF, nJ, 3]
    if 'pprev' in data:
        pprev = torch.from_numpy(data['pprev']).float()  # [nF, nJ, 3]
        pprev_flat = pprev.reshape(-1, pprev.shape[-1])
    else:
        pprev = pprev_flat = None
    c = torch.from_numpy(data['c']).float()      # [nF, nJ, 1]
    r = torch.from_numpy(data['r']).float()      # [nF, 4]
    
    nF, nJ = q.shape[0], q.shape[1]
    _log(f"  Loaded {nF} frames, {nJ} joints", verbose)
    
    # Flatten temporal dimension: [nF, nJ, D] -> [nF*nJ, D]
    q_flat = q.reshape(-1, q.shape[-1])
    p_flat = p.reshape(-1, p.shape[-1])
    qv_flat = qv.reshape(-1, qv.shape[-1])
    pv_flat = pv.reshape(-1, pv.shape[-1])
    c_flat = c.reshape(-1, c.shape[-1])
    
    # Create SkelData - check if text features exist
    if 'tf' in data:
        tf = torch.from_numpy(data['tf']).float()  # [nJ, 768]
    elif not load_tf:
        # _log(f"  Warning: 'tf' not found in {npz_path}, using zeros.", verbose)
        tf = torch.zeros((nJ, 768), dtype=torch.float32)
    else:
        # load tf from joint_text_features
        # Replace only the last 'processed' directory with 'joint_text_features'
        tf_path = npz_path.rsplit('/processed/', 1)[0] + '/joint_text_features/' + npz_path.rsplit('/', 1)[1]
        tf = torch.from_numpy(np.load(tf_path)['tf'])
    
    skel_data = SkelData(
        lo=lo, 
        go=go, 
        qb=qb, 
        edge_index=edge_index,
        edge_feature=edge_feature,
        tf = tf
    )
    
    pose_data = PoseData(
        q=q_flat, 
        p=p_flat, 
        qv=qv_flat, 
        pv=pv_flat, 
        pprev=pprev_flat, 
        c=c_flat, 
        r=r
    )
    
    # Create SkelPoseGraph
    same_graph = SkelPoseGraph(skel_data, pose_data)
    same_graph.tf = tf  # Add text features directly to graph
    
    return same_graph, nF


def same_graph_to_motion_direct(same_graph, num_frames, first_frame_zero=False, 
                                 contact_cleanup=False, cids=None, fps=30, verbose=True):
    """
    Directly convert single SAME graph to Fairmotion Motion without batch mechanism.
    This is a simplified version that avoids the complexity of batch processing.
    
    Args:
        same_graph: SkelPoseGraph, the SAME format graph data (single motion)
        num_frames: int, number of frames in the motion sequence
        first_frame_zero: bool, whether to zero out the first frame root position
        contact_cleanup: bool, whether to apply contact cleanup
        cids: list, contact joint IDs for cleanup (optional, will auto-detect feet)
        fps: int, frames per second for the motion (default: 30)
        
    Returns:
        motion: fairmotion.core.Motion object
        contact: contact information (if contact_cleanup is True)
    """
    from sata.utils.motion_utils import make_motion
    from sata.utils.tensor_utils import cdn, tensor_q2qR
    from sata.mymodel import accum_root
    from fairmotion.ops import conversions
    from sata.skel_pose_graph import find_feet
    
    _log(f"Converting SAME graph to Motion (direct method)...", verbose)
    
    nJ = same_graph.lo.shape[0]
    
    # Reshape flattened data back to [nF, nJ, D]
    q = same_graph.q.reshape(num_frames, nJ, -1)  # [nF, nJ, 6]
    c = same_graph.c.reshape(num_frames, nJ, -1)  # [nF, nJ, 1]
    r = same_graph.r_nopad                         # [nF, 4]
    
    # Convert 6D rotation to rotation matrix
    qR = cdn(tensor_q2qR(q))  # [nF, nJ, 3, 3]
    
    # Accumulate root motion
    r_expanded = r.unsqueeze(1)  # [nF, 1, 4]
    ra_T = cdn(accum_root(r_expanded, num_frames, apply_height=True))  # [nF, 1, 4, 4]
    ra_T = ra_T.squeeze(1)  # [nF, 4, 4]
    
    # Create skeleton from graph
    lo = same_graph.lo
    qb = same_graph.qb
    edge_index = same_graph.edge_index
    
    skel = motion_class.Skeleton()
    root_joint = motion_class.Joint(dof=6)
    skel.add_joint(root_joint, None)
    
    # Build skeleton hierarchy
    for j_idx in range(nJ):
        for pid, jid in edge_index.transpose(1, 0):
            if pid == jid:  # Skip self-loop (root dummy edge)
                continue
            if jid == j_idx:
                dof = 3 if qb[jid] else 0
                new_joint = motion_class.Joint(
                    dof=dof, 
                    xform_from_parent_joint=conversions.p2T(cdn(lo[jid]))
                )
                new_joint.set_parent_joint(skel.joints[pid])
                skel.add_joint(new_joint, skel.joints[pid])
                break
    
    # Auto-detect contact points if needed
    if contact_cleanup and cids is None:
        cids = list(find_feet(same_graph))
    
    # Create motion
    motion, contact = make_motion(
        skel=skel,
        qR=qR,
        ra_T=ra_T,
        c=cdn(c),
        first_frame_zero=first_frame_zero,
        contact_cleanup=contact_cleanup,
        cid=cids
    )
    
    # Set fps
    motion.set_fps(fps)
    
    _log(f"  ✓ Motion: {motion.num_frames()} frames, {motion.skel.num_joints()} joints, {motion.fps} fps", verbose)
    
    return motion, contact


def save_motion_as_bvh(motion, output_path, rot_order="XYZ", scale=1.0, 
                       ee_as_joint=True, verbose=True):
    """
    Save Fairmotion Motion object as BVH file
    
    Args:
        motion: fairmotion.core.Motion object
        output_path: str, path to save the BVH file
        rot_order: str, rotation order for BVH (default: "XYZ")
        scale: float, scale factor for the motion (default: 1.0)
        ee_as_joint: bool, whether to treat end effectors as joints (default: True)
        verbose: bool, whether to print progress (default: True)
    """
    # Create output directory if it doesn't exist
    output_dir = os.path.dirname(output_path)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)
        _log(f"Created output directory: {output_dir}", verbose)
    
    # Save using fairmotion's bvh.save()
    bvh.save(
        motion=motion,
        filename=output_path,
        scale=scale,
        rot_order=rot_order,
        verbose=verbose,
        ee_as_joint=ee_as_joint
    )
    
    _log(f"✓ BVH file saved to: {output_path}", verbose)


def same_npz_to_bvh(npz_path, output_path, rot_order="XYZ", 
                    scale=1.0, first_frame_zero=False, contact_cleanup=False,
                    cids=None, ee_as_joint=True, verbose=True, fps=30):
    """
    Complete pipeline: Convert SAME .npz file to BVH file
    
    This is the main function that combines all steps:
    1. Load SAME data from .npz
    2. Convert to Motion object using direct method
    3. Save as BVH file
    
    Args:
        npz_path: str, path to input .npz file
        output_path: str, path to output .bvh file
        rot_order: str, rotation order for BVH (default: "XYZ")
        scale: float, scale factor (default: 1.0)
        first_frame_zero: bool, zero out first frame root (default: False)
        contact_cleanup: bool, apply contact cleanup (default: False)
        cids: list, contact joint IDs (optional, auto-detected if not provided)
        ee_as_joint: bool, treat end effectors as joints (default: True)
        verbose: bool, print progress (default: True)
        fps: int, frames per second for the motion (default: 30)
        
    Returns:
        motion: fairmotion.core.Motion object that was saved
        
    Example:
        >>> from sata.conversions.same2bvh import same_npz_to_bvh
        >>> motion = same_npz_to_bvh(
        ...     npz_path="data/motion_001.npz",
        ...     output_path="output/motion_001.bvh"
        ... )
    """
    _log("="*60, verbose)
    _log("SAME to BVH Conversion Pipeline", verbose)
    _log("="*60, verbose)
    
    # Step 1: Load SAME data
    same_graph, num_frames = load_same_from_npz(npz_path, verbose=verbose)
    
    # Step 2: Convert to Motion (using direct method)
    motion, contact = same_graph_to_motion_direct(
        same_graph=same_graph,
        num_frames=num_frames,
        first_frame_zero=first_frame_zero,
        contact_cleanup=contact_cleanup,
        cids=cids,
        fps=fps,
        verbose=verbose
    )
    
    # Step 3: Save as BVH
    save_motion_as_bvh(
        motion=motion,
        output_path=output_path,
        rot_order=rot_order,
        scale=scale,
        ee_as_joint=ee_as_joint,
        verbose=verbose
    )
    
    _log("="*60, verbose)
    _log("Conversion completed successfully!", verbose)
    _log("="*60, verbose)
    
    return motion


def batch_convert_directory(input_dir, output_dir, rot_order="XYZ", 
                           scale=1.0, first_frame_zero=False, contact_cleanup=False,
                           cids=None, ee_as_joint=True, verbose=True, fps=30):
    """
    Batch convert all .npz files in a directory to BVH format
    
    This function recursively finds all .npz files in input_dir and converts
    them to BVH format, preserving the directory structure in output_dir.
    
    Args:
        input_dir: str, path to directory containing .npz files
        output_dir: str, path to directory where .bvh files will be saved
        rot_order: str, rotation order for BVH (default: "XYZ")
        scale: float, scale factor (default: 1.0)
        first_frame_zero: bool, zero out first frame root (default: False)
        contact_cleanup: bool, apply contact cleanup (default: False)
        cids: list, contact joint IDs (optional)
        ee_as_joint: bool, treat end effectors as joints (default: True)
        verbose: bool, print progress (default: True)
        fps: int, frames per second for the motion (default: 30)
        
    Returns:
        dict: Statistics dictionary with keys:
            - 'total': total number of files processed
            - 'success': number of successful conversions
            - 'failed': number of failed conversions
            - 'failed_files': list of failed file paths
            - 'elapsed_time': elapsed time in seconds
            
    Example:
        >>> from sata.conversions.same2bvh import batch_convert_directory
        >>> stats = batch_convert_directory(
        ...     input_dir="data/same_motions",
        ...     output_dir="output/bvh_motions",
        ...     fps=60
        ... )
        >>> print(f"Converted {stats['success']}/{stats['total']} files")
    """
    import glob
    import time
    from pathlib import Path
    
    start_time = time.time()
    
    # Find all .npz files
    input_path = Path(input_dir)
    if not input_path.exists():
        raise FileNotFoundError(f"Input directory not found: {input_dir}")
    
    npz_files = sorted(input_path.rglob("*.npz"))
    
    if not npz_files:
        _log(f"No .npz files found in {input_dir}", verbose)
        return {
            'total': 0,
            'success': 0,
            'failed': 0,
            'failed_files': [],
            'elapsed_time': 0
        }
    
    # Create output directory
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    
    _log("="*60, verbose)
    _log(f"Batch Converting {len(npz_files)} .npz files", verbose)
    _log("="*60, verbose)
    _log(f"Input:  {input_dir}", verbose)
    _log(f"Output: {output_dir}", verbose)
    _log("="*60, verbose)
    
    # Statistics
    success_count = 0
    failed_count = 0
    failed_files = []
    
    # Process each file
    for idx, npz_file in enumerate(npz_files, 1):
        try:
            # Compute relative path and output path
            relative_path = npz_file.relative_to(input_path)
            output_file = output_path / relative_path.with_suffix(".bvh")
            
            # Create output subdirectory
            output_file.parent.mkdir(parents=True, exist_ok=True)
            
            # Print progress
            _log(f"\n[{idx}/{len(npz_files)}] Converting: {relative_path}", verbose)
            
            # Convert
            same_npz_to_bvh(
                npz_path=str(npz_file),
                output_path=str(output_file),
                rot_order=rot_order,
                scale=scale,
                first_frame_zero=first_frame_zero,
                contact_cleanup=contact_cleanup,
                cids=cids,
                ee_as_joint=ee_as_joint,
                verbose=verbose,
                fps=fps
            )
            
            success_count += 1
            _log(f"✓ Saved to: {output_file}", verbose)
            
        except Exception as e:
            failed_count += 1
            failed_files.append(str(relative_path))
            _log(f"✗ Failed to convert {relative_path}", verbose)
            _log(f"  Error: {str(e)}", verbose)
    
    # Summary
    elapsed_time = time.time() - start_time
    minutes = int(elapsed_time) // 60
    seconds = int(elapsed_time) % 60
    
    _log("\n" + "="*60, verbose)
    _log("Batch Conversion Summary", verbose)
    _log("="*60, verbose)
    _log(f"Total:   {len(npz_files)} files", verbose)
    _log(f"Success: {success_count} files", verbose)
    _log(f"Failed:  {failed_count} files", verbose)
    _log(f"Time:    {minutes}m {seconds}s", verbose)
    
    if failed_files:
        _log(f"\nFailed files:", verbose)
        for f in failed_files:
            _log(f"  - {f}", verbose)
    
    _log("="*60, verbose)
    
    return {
        'total': len(npz_files),
        'success': success_count,
        'failed': failed_count,
        'failed_files': failed_files,
        'elapsed_time': elapsed_time
    }


# Example usage
if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(
        description="Convert SAME format (.npz) to BVH format (.bvh)",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Convert single file
  python same2bvh.py --input data/motion.npz --output result/motion.bvh
  
  # With custom rotation order and scale
  python same2bvh.py --input data/motion.npz --output result/motion.bvh \\
      --rot_order ZXY --scale 0.01
  
  # With contact cleanup
  python same2bvh.py --input data/motion.npz --output result/motion.bvh \\
      --contact_cleanup --first_frame_zero
  
  # Batch convert entire directory
  python same2bvh.py --input data/same_motions --output result/bvh_motions --batch
  
  # Batch convert with custom parameters
  python same2bvh.py --input data/ --output output/ --batch \\
      --fps 60 --rot_order ZXY --contact_cleanup
        """
    )
    
    # Input/output arguments
    parser.add_argument("--input", type=str, required=True, help="Input .npz file or directory path")
    parser.add_argument("--output", type=str, required=True, help="Output .bvh file or directory path")
    
    # Conversion options
    parser.add_argument("--rot_order", type=str, default="XYZ", 
                       choices=["XYZ", "ZXY", "ZYX"],
                       help="Rotation order for BVH (default: XYZ)")
    parser.add_argument("--scale", type=float, default=1.0,
                       help="Scale factor for the motion (default: 1.0)")
    parser.add_argument("--fps", type=int, default=20,
                       help="Frames per second for the motion (default: 20)")
    parser.add_argument("--first_frame_zero", action="store_true",
                       help="Zero out the first frame root position")
    parser.add_argument("--contact_cleanup", action="store_true",
                       help="Apply contact cleanup to the motion")
    parser.add_argument("--no_ee_as_joint", action="store_true",
                       help="Don't treat end effectors as joints")
    parser.add_argument("--quiet", action="store_true",
                       help="Suppress verbose output")
    parser.add_argument("--batch", action="store_true",
                       help="Batch convert all .npz files in input directory")
    
    args = parser.parse_args()
    
    # Batch conversion mode
    if args.batch:
        batch_convert_directory(
            input_dir=args.input,
            output_dir=args.output,
            rot_order=args.rot_order,
            scale=args.scale,
            fps=args.fps,
            first_frame_zero=args.first_frame_zero,
            contact_cleanup=args.contact_cleanup,
            ee_as_joint=not args.no_ee_as_joint,
            verbose=not args.quiet
        )
    else:
        # Single file conversion
        same_npz_to_bvh(
            npz_path=args.input,
            output_path=args.output,
            rot_order=args.rot_order,
            scale=args.scale,
            fps=args.fps,
            first_frame_zero=args.first_frame_zero,
            contact_cleanup=args.contact_cleanup,
            ee_as_joint=not args.no_ee_as_joint,
            verbose=not args.quiet
        )