fitness_coach/reference_processor.py

"""
Reference Video Processor
Processes reference videos once and saves noisy samples for scoring
"""

import os
import sys
import numpy as np
import json
from pathlib import Path

# Add parent directory and demo directory to path
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))
sys.path.insert(0, str(project_root / 'demo'))

from fitness_coach.noise_scoring import create_noisy_samples, calculate_statistical_bounds
from fitness_coach.body_parts import calculate_body_scale, get_joints_for_exercise


def process_reference_video(video_path, exercise_type='pushup', output_dir=None, n_samples=100):
    """
    Process a reference video and generate noisy samples for scoring
    
    Args:
        video_path: Path to reference video file
        exercise_type: Type of exercise (e.g., 'pushup', 'squat')
        output_dir: Directory to save processed data (default: references/{exercise_type}/)
        n_samples: Number of noisy samples to generate
    
    Returns:
        Dictionary with paths to saved files and metadata
    """
    # Change to project root for imports to work correctly
    original_cwd = os.getcwd()
    os.chdir(project_root)
    
    try:
        # Import after changing directory
        from demo.vis import get_pose2D, get_pose3D
    finally:
        os.chdir(original_cwd)
    
    video_path = Path(video_path)
    if not video_path.exists():
        raise FileNotFoundError(f"Video not found: {video_path}")
    
    # Set up output directory
    if output_dir is None:
        output_dir = Path('references') / exercise_type
    else:
        output_dir = Path(output_dir)
    
    output_dir.mkdir(parents=True, exist_ok=True)
    
    print(f"Processing reference video: {video_path.name}")
    print(f"Exercise type: {exercise_type}")
    print(f"Output directory: {output_dir}")
    
    # Create temporary output directory for processing
    temp_output = output_dir / 'temp_processing'
    temp_output.mkdir(exist_ok=True)
    
    # Format output directory string (get_pose3D expects trailing slash)
    # Use absolute path to avoid issues when changing directories
    temp_output_abs = temp_output.resolve()
    output_dir_str = str(temp_output_abs).replace('\\', '/')
    if not output_dir_str.endswith('/'):
        output_dir_str += '/'
    
    video_path_abs = video_path.resolve()
    
    # Change to project root for processing
    os.chdir(project_root)
    
    # Save original argv and temporarily clear it to avoid argparse conflicts
    original_argv = sys.argv.copy()
    sys.argv = [sys.argv[0]]  # Keep only script name
    
    try:
        # Step 1: Extract 2D poses
        print("\n[1/4] Extracting 2D poses...")
        try:
            # get_pose2D expects output_dir with trailing slash
            # It adds 'input_2D/' to it (line 95 in vis.py)
            get_pose2D(str(video_path_abs), output_dir_str)
        except Exception as e:
            print(f"Error in 2D pose extraction: {e}")
            raise
        
        # Step 2: Extract 3D poses
        print("\n[2/4] Extracting 3D poses...")
        try:
            # get_pose3D also expects output_dir with trailing slash
            # It looks for output_dir + 'input_2D/keypoints.npz' (line 190 in vis.py)
            get_pose3D(str(video_path_abs), output_dir_str)
        except Exception as e:
            print(f"Error in 3D pose extraction: {e}")
            raise
    finally:
        sys.argv = original_argv  # Restore original argv
        os.chdir(original_cwd)
    
    # Step 3: Load 3D poses
    # get_pose3D saves to output_dir + 'keypoints_3D.npz' (line 279 in vis.py)
    keypoints_3d_path = temp_output_abs / 'keypoints_3D.npz'
    
    if not keypoints_3d_path.exists():
        # Try alternative locations in case path handling differs
        alt_paths = [
            temp_output_abs / 'keypoints_3D.npz',
            temp_output_abs.parent / 'keypoints_3D.npz',
        ]
        for alt_path in alt_paths:
            if alt_path.exists():
                keypoints_3d_path = alt_path
                break
        else:
            # List what files actually exist to help debug
            print(f"\nDebug: Looking for keypoints_3D.npz")
            print(f"Expected location: {keypoints_3d_path}")
            print(f"Files in temp_processing:")
            if temp_output_abs.exists():
                for item in temp_output_abs.rglob('*'):
                    if item.is_file():
                        print(f"  {item}")
            raise FileNotFoundError(f"3D keypoints not found: {keypoints_3d_path}")
    
    keypoints_3d = np.load(str(keypoints_3d_path), allow_pickle=True)['reconstruction']
    print(f"Loaded {len(keypoints_3d)} frames of 3D poses")
    
    # Convert to numpy array if needed
    if isinstance(keypoints_3d, list):
        keypoints_3d = np.array(keypoints_3d)
    
    # Step 4: Generate noisy samples
    print(f"\n[3/4] Generating {n_samples} noisy samples...")
    noisy_samples = create_noisy_samples(keypoints_3d, n_samples=n_samples, per_joint_noise=True)
    print(f"Generated noisy samples shape: {noisy_samples.shape}")
    
    # Step 5: Calculate metadata
    print("\n[4/4] Calculating metadata...")
    body_scale = calculate_body_scale(keypoints_3d)
    relevant_body_parts = get_joints_for_exercise(exercise_type)
    
    # Calculate statistical bounds
    mean_poses, lower_bound, upper_bound, tolerance = calculate_statistical_bounds(
        keypoints_3d, noise_std=0.05
    )
    
    metadata = {
        'exercise_type': exercise_type,
        'video_path': str(video_path),
        'video_name': video_path.stem,
        'num_frames': len(keypoints_3d),
        'body_scale': float(body_scale),
        'relevant_body_parts': relevant_body_parts,
        'n_samples': n_samples,
        'timestamp': str(Path(video_path).stat().st_mtime) if video_path.exists() else None
    }
    
    # Step 6: Save everything
    print("\nSaving processed data...")
    
    # Save 3D poses
    poses_3d_path = output_dir / 'keypoints_3D.npz'
    np.savez_compressed(str(poses_3d_path), reconstruction=keypoints_3d)
    print(f"  Saved 3D poses: {poses_3d_path}")
    
    # Save noisy samples
    noisy_samples_path = output_dir / 'noisy_samples.npz'
    np.savez_compressed(str(noisy_samples_path), samples=noisy_samples)
    print(f"  Saved noisy samples: {noisy_samples_path}")
    
    # Save statistical bounds
    bounds_path = output_dir / 'statistical_bounds.npz'
    np.savez_compressed(
        str(bounds_path),
        mean=mean_poses,
        lower_bound=lower_bound,
        upper_bound=upper_bound,
        tolerance=tolerance
    )
    print(f"  Saved statistical bounds: {bounds_path}")
    
    # Save metadata
    metadata_path = output_dir / 'metadata.json'
    with open(metadata_path, 'w') as f:
        json.dump(metadata, f, indent=2)
    print(f"  Saved metadata: {metadata_path}")
    
    # Clean up temporary files (optional - keep 2D poses for debugging)
    # import shutil
    # shutil.rmtree(temp_output, ignore_errors=True)
    
    print(f"\n✓ Reference video processed successfully!")
    print(f"  Output directory: {output_dir}")
    
    return {
        'output_dir': str(output_dir),
        'poses_3d_path': str(poses_3d_path),
        'noisy_samples_path': str(noisy_samples_path),
        'bounds_path': str(bounds_path),
        'metadata_path': str(metadata_path),
        'metadata': metadata
    }


def load_reference(exercise_type, references_dir='references'):
    """
    Load a processed reference
    
    Args:
        exercise_type: Type of exercise (e.g., 'pushup')
        references_dir: Directory containing references
    
    Returns:
        Dictionary with loaded data
    """
    ref_dir = Path(references_dir) / exercise_type
    
    if not ref_dir.exists():
        raise FileNotFoundError(f"Reference not found: {ref_dir}")
    
    # Load metadata
    metadata_path = ref_dir / 'metadata.json'
    if not metadata_path.exists():
        raise FileNotFoundError(f"Metadata not found: {metadata_path}")
    
    with open(metadata_path, 'r') as f:
        metadata = json.load(f)
    
    # Load 3D poses
    poses_3d_path = ref_dir / 'keypoints_3D.npz'
    if not poses_3d_path.exists():
        raise FileNotFoundError(f"3D poses not found: {poses_3d_path}")
    
    poses_3d = np.load(str(poses_3d_path), allow_pickle=True)['reconstruction']
    if isinstance(poses_3d, list):
        poses_3d = np.array(poses_3d)
    
    # Load noisy samples
    noisy_samples_path = ref_dir / 'noisy_samples.npz'
    noisy_samples = None
    if noisy_samples_path.exists():
        noisy_samples = np.load(str(noisy_samples_path), allow_pickle=True)['samples']
    
    # Load statistical bounds
    bounds_path = ref_dir / 'statistical_bounds.npz'
    bounds = None
    if bounds_path.exists():
        bounds_data = np.load(str(bounds_path), allow_pickle=True)
        bounds = {
            'mean': bounds_data['mean'],
            'lower_bound': bounds_data['lower_bound'],
            'upper_bound': bounds_data['upper_bound'],
            'tolerance': bounds_data['tolerance']
        }
    
    return {
        'poses_3d': poses_3d,
        'noisy_samples': noisy_samples,
        'bounds': bounds,
        'metadata': metadata,
        'ref_dir': str(ref_dir)
    }


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='Process reference video for scoring')
    parser.add_argument('--video', type=str, required=True, help='Path to reference video')
    parser.add_argument('--exercise', type=str, default='pushup', help='Exercise type')
    parser.add_argument('--output', type=str, default=None, help='Output directory')
    parser.add_argument('--samples', type=int, default=100, help='Number of noisy samples')
    
    args = parser.parse_args()
    
    try:
        result = process_reference_video(
            args.video,
            exercise_type=args.exercise,
            output_dir=args.output,
            n_samples=args.samples
        )
        print("\n" + "="*50)
        print("SUCCESS!")
        print("="*50)
        print(f"Reference saved to: {result['output_dir']}")
    except Exception as e:
        print(f"\nERROR: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)