"""
Compare temporal stability of 3D predictions between two models.

Usage:
    python compare_temporal_stability.py \
        --input1 runs/model1/tracking.json \
        --input2 runs/model2/tracking.json \
        --track-mapping 1:5 2:10 3:15 \
        --class-id 0 \
        --output comparison_report.json
"""

import argparse
import json
from pathlib import Path
import numpy as np
from collections import defaultdict
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg')


def compute_angle_diff(angle1, angle2):
    """Compute the smallest difference between two angles in radians.
    
    Args:
        angle1, angle2: Angles in radians
        
    Returns:
        Angle difference in range [-pi, pi]
    """
    diff = angle1 - angle2
    while diff > np.pi:
        diff -= 2 * np.pi
    while diff < -np.pi:
        diff += 2 * np.pi
    return diff


def extract_trajectories(tracking_data):
    """Extract trajectories for each tracked object, organized by class.
    
    Args:
        tracking_data: List of frame data from tracking.json
        
    Returns:
        Dictionary mapping class_id to dict of {track_id: list of (frame_idx, detection) tuples}
    """
    trajectories_by_class = defaultdict(lambda: defaultdict(list))
    
    for frame_idx, frame_data in enumerate(tracking_data):
        for det in frame_data.get('detections', []):
            track_id = det.get('track_id')
            class_id = det.get('class_id')
            if track_id is not None and class_id is not None and 'object_3d' in det:
                trajectories_by_class[class_id][track_id].append((frame_idx, det))
    
    return trajectories_by_class


def extract_3d_data(trajectory):
    """Extract 3D data from trajectory.
    
    Args:
        trajectory: List of (frame_idx, detection) tuples
        
    Returns:
        Dictionary with extracted data arrays
    """
    # Use dictionary to handle potential duplicate frames (take last occurrence)
    frame_data = {}
    
    for frame_idx, det in trajectory:
        object_3d = det['object_3d']
        
        # Support both list and dict formats
        if isinstance(object_3d, dict):
            x, y, z = object_3d['location']
            l, h, w = object_3d['dimensions']
            rot_y = object_3d['rotation_y']
        else:
            x, y, z = object_3d[0], object_3d[1], object_3d[2]
            l, h, w = object_3d[3], object_3d[4], object_3d[5]
            rot_y = object_3d[6]
        
        frame_data[frame_idx] = {
            'position': [x, y, z],
            'dimension': [l, h, w],
            'rotation': rot_y
        }
    
    # Sort by frame index and extract arrays
    sorted_frames = sorted(frame_data.keys())
    frames = []
    positions = []
    dimensions = []
    rotations = []
    
    for frame_idx in sorted_frames:
        data = frame_data[frame_idx]
        frames.append(frame_idx)
        positions.append(data['position'])
        dimensions.append(data['dimension'])
        rotations.append(data['rotation'])
    
    return {
        'frames': np.array(frames),
        'positions': np.array(positions),
        'dimensions': np.array(dimensions),
        'rotations': np.array(rotations),
    }


def compute_trajectory_metrics(data1, data2):
    """Compute comparison metrics between two trajectories.
    
    Args:
        data1, data2: Extracted 3D data dictionaries
        
    Returns:
        Dictionary with comparison metrics
    """
    # Ensure arrays have the same length
    if len(data1['positions']) != len(data2['positions']):
        raise ValueError(f"Data arrays have different lengths: {len(data1['positions'])} vs {len(data2['positions'])}")
    
    # Position differences
    pos_diff_mean = np.mean(np.linalg.norm(data1['positions'] - data2['positions'], axis=1))
    pos_diff_std = np.std(np.linalg.norm(data1['positions'] - data2['positions'], axis=1))
    pos_diff_max = np.max(np.linalg.norm(data1['positions'] - data2['positions'], axis=1))
    
    # Dimension differences
    dim_diff = np.abs(data1['dimensions'] - data2['dimensions'])
    dim_diff_mean = np.mean(dim_diff, axis=0)
    dim_diff_std = np.std(dim_diff, axis=0)
    dim_diff_max = np.max(dim_diff, axis=0)
    
    # Rotation differences
    rot_diffs = []
    for r1, r2 in zip(data1['rotations'], data2['rotations']):
        diff = abs(compute_angle_diff(r1, r2))
        rot_diffs.append(diff)
    rot_diffs = np.array(rot_diffs)
    
    rot_diff_mean = np.mean(rot_diffs)
    rot_diff_std = np.std(rot_diffs)
    rot_diff_max = np.max(rot_diffs)
    
    # Position jitter (frame-to-frame variation)
    pos_jitter1 = np.linalg.norm(np.diff(data1['positions'], axis=0), axis=1)
    pos_jitter2 = np.linalg.norm(np.diff(data2['positions'], axis=0), axis=1)
    
    # Dimension jitter
    dim_jitter1 = np.abs(np.diff(data1['dimensions'], axis=0))
    dim_jitter2 = np.abs(np.diff(data2['dimensions'], axis=0))
    
    # Rotation jitter
    rot_jitter1 = np.abs(np.diff(data1['rotations']))
    rot_jitter2 = np.abs(np.diff(data2['rotations']))
    
    metrics = {
        'position_difference': {
            'mean': float(pos_diff_mean),
            'std': float(pos_diff_std),
            'max': float(pos_diff_max),
        },
        'dimension_difference': {
            'length_mean': float(dim_diff_mean[0]),
            'height_mean': float(dim_diff_mean[1]),
            'width_mean': float(dim_diff_mean[2]),
            'length_std': float(dim_diff_std[0]),
            'height_std': float(dim_diff_std[1]),
            'width_std': float(dim_diff_std[2]),
            'length_max': float(dim_diff_max[0]),
            'height_max': float(dim_diff_max[1]),
            'width_max': float(dim_diff_max[2]),
        },
        'rotation_difference': {
            'mean': float(rot_diff_mean),
            'std': float(rot_diff_std),
            'max': float(rot_diff_max),
        },
        'position_jitter': {
            'model1_mean': float(np.mean(pos_jitter1)),
            'model2_mean': float(np.mean(pos_jitter2)),
            'model1_std': float(np.std(pos_jitter1)),
            'model2_std': float(np.std(pos_jitter2)),
        },
        'dimension_jitter': {
            'model1_length_mean': float(np.mean(dim_jitter1[:, 0])),
            'model2_length_mean': float(np.mean(dim_jitter2[:, 0])),
            'model1_height_mean': float(np.mean(dim_jitter1[:, 1])),
            'model2_height_mean': float(np.mean(dim_jitter2[:, 1])),
            'model1_width_mean': float(np.mean(dim_jitter1[:, 2])),
            'model2_width_mean': float(np.mean(dim_jitter2[:, 2])),
        },
        'rotation_jitter': {
            'model1_mean': float(np.mean(rot_jitter1)),
            'model2_mean': float(np.mean(rot_jitter2)),
            'model1_std': float(np.std(rot_jitter1)),
            'model2_std': float(np.std(rot_jitter2)),
        },
    }
    
    return metrics


def plot_comparison(data1, data2, track_id1, track_id2, model1_name, model2_name, output_dir, class_id):
    """Generate comparison plots for two trajectories.
    
    Args:
        data1, data2: Extracted 3D data dictionaries
        track_id1, track_id2: Track IDs from each model
        model1_name, model2_name: Model names for labels
        output_dir: Output directory for plots
        class_id: Object class ID
    """
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Plot 1: Position (x, y, z) comparison
    fig, axes = plt.subplots(3, 1, figsize=(14, 12))
    
    axes[0].plot(data1['frames'], data1['positions'][:, 0], marker='o', markersize=4, 
                label=f'{model1_name} (Track {track_id1})', alpha=0.7, linewidth=2, color='blue')
    axes[0].plot(data2['frames'], data2['positions'][:, 0], marker='s', markersize=4, 
                label=f'{model2_name} (Track {track_id2})', alpha=0.7, linewidth=2, color='red')
    axes[0].set_ylabel('X Position (m)', fontsize=11)
    axes[0].set_title(f'Class {class_id}: Position Comparison - {model1_name} vs {model2_name}', fontsize=13, fontweight='bold')
    axes[0].legend(fontsize=10)
    axes[0].grid(True, alpha=0.3)
    
    axes[1].plot(data1['frames'], data1['positions'][:, 1], marker='o', markersize=4, 
                label=f'{model1_name} (Track {track_id1})', alpha=0.7, linewidth=2, color='blue')
    axes[1].plot(data2['frames'], data2['positions'][:, 1], marker='s', markersize=4, 
                label=f'{model2_name} (Track {track_id2})', alpha=0.7, linewidth=2, color='red')
    axes[1].set_ylabel('Y Position (m)', fontsize=11)
    axes[1].legend(fontsize=10)
    axes[1].grid(True, alpha=0.3)
    
    axes[2].plot(data1['frames'], data1['positions'][:, 2], marker='o', markersize=4, 
                label=f'{model1_name} (Track {track_id1})', alpha=0.7, linewidth=2, color='blue')
    axes[2].plot(data2['frames'], data2['positions'][:, 2], marker='s', markersize=4, 
                label=f'{model2_name} (Track {track_id2})', alpha=0.7, linewidth=2, color='red')
    axes[2].set_ylabel('Z Position (m)', fontsize=11)
    axes[2].set_xlabel('Frame Index', fontsize=11)
    axes[2].legend(fontsize=10)
    axes[2].grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig(output_dir / f'class_{class_id}_track_{track_id1}_{track_id2}_position_comparison.png', dpi=150, bbox_inches='tight')
    plt.close()
    
    # Plot 2: Position difference over time
    fig, ax = plt.subplots(figsize=(14, 5))
    pos_diff = np.linalg.norm(data1['positions'] - data2['positions'], axis=1)
    ax.plot(data1['frames'], pos_diff, marker='o', markersize=4, linewidth=2, color='purple', alpha=0.7)
    ax.fill_between(data1['frames'], 0, pos_diff, alpha=0.3, color='purple')
    ax.set_xlabel('Frame Index', fontsize=11)
    ax.set_ylabel('Position Difference (m)', fontsize=11)
    ax.set_title(f'Class {class_id}: Position Difference - {model1_name} vs {model2_name}', fontsize=13, fontweight='bold')
    ax.grid(True, alpha=0.3)
    ax.axhline(y=np.mean(pos_diff), color='red', linestyle='--', label=f'Mean: {np.mean(pos_diff):.3f}m')
    ax.legend(fontsize=10)
    plt.tight_layout()
    plt.savefig(output_dir / f'class_{class_id}_track_{track_id1}_{track_id2}_position_difference.png', dpi=150, bbox_inches='tight')
    plt.close()
    
    # Plot 3: Dimensions (l, h, w) comparison
    fig, axes = plt.subplots(3, 1, figsize=(14, 12))
    dim_labels = ['Length', 'Height', 'Width']
    
    for i, label in enumerate(dim_labels):
        axes[i].plot(data1['frames'], data1['dimensions'][:, i], marker='o', markersize=4, 
                    label=f'{model1_name} (Track {track_id1})', alpha=0.7, linewidth=2, color='blue')
        axes[i].plot(data2['frames'], data2['dimensions'][:, i], marker='s', markersize=4, 
                    label=f'{model2_name} (Track {track_id2})', alpha=0.7, linewidth=2, color='red')
        axes[i].set_ylabel(f'{label} (m)', fontsize=11)
        if i == 0:
            axes[i].set_title(f'Class {class_id}: Dimension Comparison - {model1_name} vs {model2_name}', fontsize=13, fontweight='bold')
        axes[i].legend(fontsize=10)
        axes[i].grid(True, alpha=0.3)
    
    axes[2].set_xlabel('Frame Index', fontsize=11)
    plt.tight_layout()
    plt.savefig(output_dir / f'class_{class_id}_track_{track_id1}_{track_id2}_dimension_comparison.png', dpi=150, bbox_inches='tight')
    plt.close()
    
    # Plot 4: Dimension difference over time
    fig, axes = plt.subplots(3, 1, figsize=(14, 10))
    dim_diff = np.abs(data1['dimensions'] - data2['dimensions'])
    
    for i, label in enumerate(dim_labels):
        axes[i].plot(data1['frames'], dim_diff[:, i], marker='o', markersize=4, linewidth=2, color='purple', alpha=0.7)
        axes[i].fill_between(data1['frames'], 0, dim_diff[:, i], alpha=0.3, color='purple')
        axes[i].set_ylabel(f'{label} Diff (m)', fontsize=11)
        if i == 0:
            axes[i].set_title(f'Class {class_id}: Dimension Difference - {model1_name} vs {model2_name}', fontsize=13, fontweight='bold')
        axes[i].grid(True, alpha=0.3)
        axes[i].axhline(y=np.mean(dim_diff[:, i]), color='red', linestyle='--', 
                       label=f'Mean: {np.mean(dim_diff[:, i]):.4f}m')
        axes[i].legend(fontsize=10)
    
    axes[2].set_xlabel('Frame Index', fontsize=11)
    plt.tight_layout()
    plt.savefig(output_dir / f'class_{class_id}_track_{track_id1}_{track_id2}_dimension_difference.png', dpi=150, bbox_inches='tight')
    plt.close()
    
    # Plot 5: Rotation comparison
    fig, ax = plt.subplots(figsize=(14, 5))
    ax.plot(data1['frames'], data1['rotations'], marker='o', markersize=4, 
           label=f'{model1_name} (Track {track_id1})', alpha=0.7, linewidth=2, color='blue')
    ax.plot(data2['frames'], data2['rotations'], marker='s', markersize=4, 
           label=f'{model2_name} (Track {track_id2})', alpha=0.7, linewidth=2, color='red')
    ax.set_ylabel('Rotation Y (rad)', fontsize=11)
    ax.set_xlabel('Frame Index', fontsize=11)
    ax.set_title(f'Class {class_id}: Rotation Comparison - {model1_name} vs {model2_name}', fontsize=13, fontweight='bold')
    ax.legend(fontsize=10)
    ax.grid(True, alpha=0.3)
    ax.axhline(y=0, color='k', linestyle='--', alpha=0.3)
    ax.axhline(y=np.pi, color='k', linestyle='--', alpha=0.3)
    ax.axhline(y=-np.pi, color='k', linestyle='--', alpha=0.3)
    plt.tight_layout()
    plt.savefig(output_dir / f'class_{class_id}_track_{track_id1}_{track_id2}_rotation_comparison.png', dpi=150, bbox_inches='tight')
    plt.close()
    
    # Plot 6: Rotation difference over time
    fig, ax = plt.subplots(figsize=(14, 5))
    rot_diffs = [abs(compute_angle_diff(r1, r2)) for r1, r2 in zip(data1['rotations'], data2['rotations'])]
    ax.plot(data1['frames'], rot_diffs, marker='o', markersize=4, linewidth=2, color='purple', alpha=0.7)
    ax.fill_between(data1['frames'], 0, rot_diffs, alpha=0.3, color='purple')
    ax.set_xlabel('Frame Index', fontsize=11)
    ax.set_ylabel('Rotation Difference (rad)', fontsize=11)
    ax.set_title(f'Class {class_id}: Rotation Difference - {model1_name} vs {model2_name}', fontsize=13, fontweight='bold')
    ax.grid(True, alpha=0.3)
    ax.axhline(y=np.mean(rot_diffs), color='red', linestyle='--', label=f'Mean: {np.mean(rot_diffs):.4f} rad')
    ax.legend(fontsize=10)
    plt.tight_layout()
    plt.savefig(output_dir / f'class_{class_id}_track_{track_id1}_{track_id2}_rotation_difference.png', dpi=150, bbox_inches='tight')
    plt.close()
    
    # Plot 7: 3D trajectory comparison (bird's eye view)
    fig, ax = plt.subplots(figsize=(12, 12))
    
    # Plot trajectories
    ax.plot(data1['positions'][:, 0], data1['positions'][:, 2], marker='o', markersize=5, 
           label=f'{model1_name} (Track {track_id1})', alpha=0.7, linewidth=2.5, color='blue')
    ax.plot(data2['positions'][:, 0], data2['positions'][:, 2], marker='s', markersize=5, 
           label=f'{model2_name} (Track {track_id2})', alpha=0.7, linewidth=2.5, color='red')
    
    # Mark start points
    ax.scatter(data1['positions'][0, 0], data1['positions'][0, 2], s=150, marker='o', 
              edgecolors='darkblue', facecolors='blue', linewidths=3, zorder=5, label=f'{model1_name} Start')
    ax.scatter(data2['positions'][0, 0], data2['positions'][0, 2], s=150, marker='s', 
              edgecolors='darkred', facecolors='red', linewidths=3, zorder=5, label=f'{model2_name} Start')
    
    # Mark end points
    ax.scatter(data1['positions'][-1, 0], data1['positions'][-1, 2], s=150, marker='^', 
              edgecolors='darkblue', facecolors='cyan', linewidths=3, zorder=5, label=f'{model1_name} End')
    ax.scatter(data2['positions'][-1, 0], data2['positions'][-1, 2], s=150, marker='^', 
              edgecolors='darkred', facecolors='orange', linewidths=3, zorder=5, label=f'{model2_name} End')
    
    # Draw lines connecting corresponding points
    for i in range(len(data1['frames'])):
        ax.plot([data1['positions'][i, 0], data2['positions'][i, 0]], 
               [data1['positions'][i, 2], data2['positions'][i, 2]], 
               color='gray', alpha=0.3, linewidth=1, linestyle='--', zorder=1)
    
    ax.set_xlabel('X Position (m)', fontsize=11)
    ax.set_ylabel('Z Position (m)', fontsize=11)
    ax.set_title(f'Class {class_id}: 3D Trajectory Comparison (Bird\'s Eye View)\n{model1_name} vs {model2_name}', 
                fontsize=13, fontweight='bold')
    ax.legend(fontsize=9, loc='best')
    ax.grid(True, alpha=0.3)
    ax.axis('equal')
    plt.tight_layout()
    plt.savefig(output_dir / f'class_{class_id}_track_{track_id1}_{track_id2}_trajectory_birds_eye.png', dpi=150, bbox_inches='tight')
    plt.close()
    
    # Plot 8: Jitter comparison (position stability)
    fig, ax = plt.subplots(figsize=(14, 5))
    pos_jitter1 = np.linalg.norm(np.diff(data1['positions'], axis=0), axis=1)
    pos_jitter2 = np.linalg.norm(np.diff(data2['positions'], axis=0), axis=1)
    frames_jitter = data1['frames'][1:]  # Skip first frame since we're using diff
    
    ax.plot(frames_jitter, pos_jitter1, marker='o', markersize=4, 
           label=f'{model1_name} Position Jitter', alpha=0.7, linewidth=2, color='blue')
    ax.plot(frames_jitter, pos_jitter2, marker='s', markersize=4, 
           label=f'{model2_name} Position Jitter', alpha=0.7, linewidth=2, color='red')
    ax.set_xlabel('Frame Index', fontsize=11)
    ax.set_ylabel('Position Jitter (m/frame)', fontsize=11)
    ax.set_title(f'Class {class_id}: Position Stability Comparison - {model1_name} vs {model2_name}', fontsize=13, fontweight='bold')
    ax.legend(fontsize=10)
    ax.grid(True, alpha=0.3)
    ax.axhline(y=np.mean(pos_jitter1), color='blue', linestyle='--', alpha=0.5, label=f'{model1_name} Mean: {np.mean(pos_jitter1):.4f}')
    ax.axhline(y=np.mean(pos_jitter2), color='red', linestyle='--', alpha=0.5, label=f'{model2_name} Mean: {np.mean(pos_jitter2):.4f}')
    plt.tight_layout()
    plt.savefig(output_dir / f'class_{class_id}_track_{track_id1}_{track_id2}_jitter_comparison.png', dpi=150, bbox_inches='tight')
    plt.close()
    
    print(f"Comparison plots saved to {output_dir}")


def parse_track_mapping(mapping_strings):
    """Parse track mapping from command line arguments.
    
    Args:
        mapping_strings: List of strings in format "id1:id2"
        
    Returns:
        List of tuples (id1, id2)
    """
    mappings = []
    for s in mapping_strings:
        try:
            id1, id2 = s.split(':')
            mappings.append((int(id1), int(id2)))
        except ValueError:
            print(f"Warning: Invalid mapping format '{s}', expected 'id1:id2'")
            continue
    return mappings


def main():
    parser = argparse.ArgumentParser(description='Compare temporal stability between two models')
    parser.add_argument('--input1', type=str, required=True, help='First tracking.json file path (model 1)')
    parser.add_argument('--input2', type=str, required=True, help='Second tracking.json file path (model 2)')
    parser.add_argument('--track-mapping', type=str, nargs='+', required=True,
                       help='Track ID mappings in format "id1:id2" (e.g., --track-mapping 1:5 2:10)')
    parser.add_argument('--class-id', type=int, default=None,
                       help='Filter by class ID (optional)')
    parser.add_argument('--model1-name', type=str, default='Model 1',
                       help='Name for first model (for plot labels)')
    parser.add_argument('--model2-name', type=str, default='Model 2',
                       help='Name for second model (for plot labels)')
    parser.add_argument('--output', type=str, default='comparison_report.json',
                       help='Output comparison report JSON file path')
    parser.add_argument('--output-dir', type=str, default=None,
                       help='Output directory for plots (default: same as output JSON)')
    
    args = parser.parse_args()
    
    # Read input tracking data
    input1_path = Path(args.input1)
    input2_path = Path(args.input2)
    
    if not input1_path.exists():
        print(f"Error: Input file 1 not found: {input1_path}")
        return
    if not input2_path.exists():
        print(f"Error: Input file 2 not found: {input2_path}")
        return
    
    print(f"Loading tracking data from:")
    print(f"  Model 1: {input1_path}")
    print(f"  Model 2: {input2_path}")
    
    with open(input1_path, 'r', encoding='utf-8') as f:
        tracking_data1 = json.load(f)
    with open(input2_path, 'r', encoding='utf-8') as f:
        tracking_data2 = json.load(f)
    
    print(f"Loaded {len(tracking_data1)} frames from model 1")
    print(f"Loaded {len(tracking_data2)} frames from model 2")

    print("Sorting frames by image name...")
    tracking_data1.sort(key=lambda x: x.get('image_name', ''))
    print("Sorting frames by image name...")
    tracking_data2.sort(key=lambda x: x.get('image_name', ''))
    
    # Extract trajectories
    print("\nExtracting trajectories...")
    trajectories_by_class1 = extract_trajectories(tracking_data1)
    trajectories_by_class2 = extract_trajectories(tracking_data2)
    
    total_tracks1 = sum(len(tracks) for tracks in trajectories_by_class1.values())
    total_tracks2 = sum(len(tracks) for tracks in trajectories_by_class2.values())
    print(f"Model 1: {total_tracks1} unique tracks across {len(trajectories_by_class1)} classes")
    for class_id, tracks in sorted(trajectories_by_class1.items()):
        print(f"  Class {class_id}: {len(tracks)} tracks")
    print(f"Model 2: {total_tracks2} unique tracks across {len(trajectories_by_class2)} classes")
    for class_id, tracks in sorted(trajectories_by_class2.items()):
        print(f"  Class {class_id}: {len(tracks)} tracks")
    
    # Parse track mappings
    track_mappings = parse_track_mapping(args.track_mapping)
    print(f"\nTrack mappings: {track_mappings}")
    
    # Filter by class ID if specified and flatten the trajectories dict
    if args.class_id is not None:
        trajectories1 = trajectories_by_class1.get(args.class_id, {})
        trajectories2 = trajectories_by_class2.get(args.class_id, {})
        print(f"\nFiltered to class_id={args.class_id}:")
        print(f"  Model 1: {len(trajectories1)} tracks")
        print(f"  Model 2: {len(trajectories2)} tracks")
    else:
        # Flatten all classes into single dict if no class filter specified
        trajectories1 = {}
        for class_tracks in trajectories_by_class1.values():
            trajectories1.update(class_tracks)
        trajectories2 = {}
        for class_tracks in trajectories_by_class2.values():
            trajectories2.update(class_tracks)
        print(f"\nUsing all classes:")
        print(f"  Model 1: {len(trajectories1)} total tracks")
        print(f"  Model 2: {len(trajectories2)} total tracks")
    
    # Determine output directory
    output_path = Path(args.output)
    if args.output_dir:
        output_dir = Path(args.output_dir)
    else:
        output_dir = output_path.parent / 'comparison_plots'
    
    # Process each track mapping
    comparison_results = []
    
    for track_id1, track_id2 in track_mappings:
        print(f"\n{'='*80}")
        print(f"Processing track mapping: Model 1 Track {track_id1} <-> Model 2 Track {track_id2}")
        print('='*80)
        
        # Check if tracks exist
        if track_id1 not in trajectories1:
            print(f"Warning: Track {track_id1} not found in model 1, skipping")
            continue
        if track_id2 not in trajectories2:
            print(f"Warning: Track {track_id2} not found in model 2, skipping")
            continue
        
        traj1 = trajectories1[track_id1]
        traj2 = trajectories2[track_id2]
        
        # Get class IDs
        class_id1 = traj1[0][1].get('class_id')
        class_id2 = traj2[0][1].get('class_id')
        
        # Filter by class if specified
        if args.class_id is not None:
            if class_id1 != args.class_id:
                print(f"Warning: Track {track_id1} has class_id={class_id1}, expected {args.class_id}, skipping")
                continue
            if class_id2 != args.class_id:
                print(f"Warning: Track {track_id2} has class_id={class_id2}, expected {args.class_id}, skipping")
                continue
        
        # Check if classes match
        if class_id1 != class_id2:
            print(f"Warning: Class mismatch - Track {track_id1} (class {class_id1}) vs Track {track_id2} (class {class_id2})")
            print("Continuing with comparison, but results may not be meaningful")
        
        print(f"Model 1 Track {track_id1}: {len(traj1)} frames, class {class_id1}")
        print(f"Model 2 Track {track_id2}: {len(traj2)} frames, class {class_id2}")
        
        # Find common frame range
        frames1 = set([f for f, _ in traj1])
        frames2 = set([f for f, _ in traj2])
        common_frames = sorted(frames1.intersection(frames2))
        
        if len(common_frames) < 2:
            print(f"Warning: Only {len(common_frames)} common frames, need at least 2 for comparison, skipping")
            continue
        
        print(f"Common frames: {len(common_frames)} ({min(common_frames)} to {max(common_frames)})")
        
        # Filter trajectories to common frames and remove duplicates
        # Use dict to keep only one detection per frame (last occurrence)
        traj1_dict = {f: d for f, d in traj1 if f in common_frames}
        traj2_dict = {f: d for f, d in traj2 if f in common_frames}
        
        # Ensure both have exactly the same frames
        final_common_frames = sorted(set(traj1_dict.keys()).intersection(set(traj2_dict.keys())))
        
        if len(final_common_frames) < 2:
            print(f"Warning: After deduplication, only {len(final_common_frames)} common frames, skipping")
            continue
        
        # Convert back to sorted list of tuples
        traj1_filtered = [(f, traj1_dict[f]) for f in final_common_frames]
        traj2_filtered = [(f, traj2_dict[f]) for f in final_common_frames]
        
        print(f"After alignment: {len(traj1_filtered)} frames for comparison")
        
        # Extract 3D data
        data1 = extract_3d_data(traj1_filtered)
        data2 = extract_3d_data(traj2_filtered)
        
        # Compute metrics
        metrics = compute_trajectory_metrics(data1, data2)
        
        # Generate plots
        plot_comparison(data1, data2, track_id1, track_id2, 
                       args.model1_name, args.model2_name, output_dir, class_id1)
        
        # Store results
        comparison_results.append({
            'track_id_model1': track_id1,
            'track_id_model2': track_id2,
            'class_id_model1': class_id1,
            'class_id_model2': class_id2,
            'num_common_frames': len(common_frames),
            'frame_range': [min(common_frames), max(common_frames)],
            'metrics': metrics,
        })
        
        # Print summary for this pair
        print(f"\nComparison Metrics:")
        print(f"  Position Difference: {metrics['position_difference']['mean']:.4f} ± {metrics['position_difference']['std']:.4f} m (max: {metrics['position_difference']['max']:.4f})")
        print(f"  Dimension Difference (L/H/W): {metrics['dimension_difference']['length_mean']:.4f} / {metrics['dimension_difference']['height_mean']:.4f} / {metrics['dimension_difference']['width_mean']:.4f} m")
        print(f"  Rotation Difference: {metrics['rotation_difference']['mean']:.4f} ± {metrics['rotation_difference']['std']:.4f} rad (max: {metrics['rotation_difference']['max']:.4f})")
        print(f"  Position Jitter (Model1/Model2): {metrics['position_jitter']['model1_mean']:.4f} / {metrics['position_jitter']['model2_mean']:.4f} m/frame")
        print(f"  Rotation Jitter (Model1/Model2): {metrics['rotation_jitter']['model1_mean']:.4f} / {metrics['rotation_jitter']['model2_mean']:.4f} rad/frame")
    
    # Save comparison report
    report = {
        'model1': {
            'input_file': str(input1_path),
            'name': args.model1_name,
            'total_frames': len(tracking_data1),
            'total_tracks': len(trajectories1),
        },
        'model2': {
            'input_file': str(input2_path),
            'name': args.model2_name,
            'total_frames': len(tracking_data2),
            'total_tracks': len(trajectories2),
        },
        'comparison_results': comparison_results,
        'summary': {
            'num_comparisons': len(comparison_results),
            'filter_class_id': args.class_id,
        }
    }
    
    print(f"\n{'='*80}")
    print("Saving comparison report...")
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(report, f, indent=2, ensure_ascii=False)
    
    print(f"Report saved to: {output_path}")
    print(f"Plots saved to: {output_dir}")
    
    # Print overall summary
    print(f"\n{'='*80}")
    print("COMPARISON SUMMARY")
    print('='*80)
    print(f"Model 1: {args.model1_name}")
    print(f"Model 2: {args.model2_name}")
    print(f"Successful comparisons: {len(comparison_results)}")
    
    if comparison_results:
        # Compute aggregate statistics
        avg_pos_diff = np.mean([r['metrics']['position_difference']['mean'] for r in comparison_results])
        avg_rot_diff = np.mean([r['metrics']['rotation_difference']['mean'] for r in comparison_results])
        
        print(f"\nAggregate Statistics:")
        print(f"  Average position difference: {avg_pos_diff:.4f} m")
        print(f"  Average rotation difference: {avg_rot_diff:.4f} rad ({np.rad2deg(avg_rot_diff):.2f}°)")
    
    print("\nComparison completed!")


if __name__ == '__main__':
    main()