#!/usr/bin/env python3
"""
YOLOv5-3D Model Evaluation Script

This script evaluates 2D and 3D detection performance of YOLOv5-3D model.
It computes metrics including Precision, Recall, AP, mAP for 2D detection,
and lateral/longitudinal errors and heading errors for 3D detection.

Usage:
    # Basic usage with paths
    python eval_tools/eval.py --det-path /path/to/detections --gt-path /path/to/labels --output-dir results
    
    # With configuration file
    python eval_tools/eval.py --config eval_tools/configs/eval_config.yaml
    
    # Evaluate 2D only
    python eval_tools/eval.py --det-path /path/to/detections --gt-path /path/to/labels --eval-2d-only
    
    # Evaluate 3D only
    python eval_tools/eval.py --det-path /path/to/detections --gt-path /path/to/labels --eval-3d-only
"""

import argparse
import sys
import yaml
from pathlib import Path
from datetime import datetime

# Add parent directory to path
sys.path.insert(0, str(Path(__file__).parent.parent))

from eval_tools.evaluator import Evaluator


def load_config(config_path):
    """Load configuration from YAML file."""
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)
    
    # Replace timestamp placeholders in paths
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    if 'output' in config and 'save_path' in config['output']:
        config['output']['save_path'] = config['output']['save_path'].replace('{timestamp}', timestamp)
    if 'dataset' in config:
        if 'det_path' in config['dataset']:
            config['dataset']['det_path'] = config['dataset']['det_path'].replace('{timestamp}', timestamp)
        if 'gt_path' in config['dataset']:
            config['dataset']['gt_path'] = config['dataset']['gt_path'].replace('{timestamp}', timestamp)
    
    return config


def parse_arguments():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(
        description='Evaluate YOLOv5-3D model detection results',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Basic evaluation
  python eval_tools/eval.py --det-path /data/detections --gt-path /data/labels --output-dir results
  
  # With custom image size
  python eval_tools/eval.py --det-path /data/detections --gt-path /data/labels --img-width 3840 --img-height 2160
  
  # Using config file
  python eval_tools/eval.py --config eval_tools/configs/eval_config.yaml
  
  # 2D evaluation only
  python eval_tools/eval.py --det-path /data/detections --gt-path /data/labels --eval-2d-only
        """
    )
    
    # Path arguments
    parser.add_argument('--det-path', type=str,
                        help='Path to detection results root directory (contains case folders)')
    parser.add_argument('--gt-path', type=str,
                        help='Path to ground truth labels root directory (contains case folders)')
    parser.add_argument('--path-depth', type=int, default=1, choices=[1, 2],
                        help='Directory depth: 1 = det_path/case/txt_results, 2 = det_path/level1/case/txt_results (default: 1)')
    parser.add_argument('--det-format', type=str, default='auto', choices=['auto', 'json', 'txt'],
                        help='Detection file format: auto (default, probe json_results/ then txt_results/), json, or txt')
    parser.add_argument('--gt-format', type=str, default='auto', choices=['auto', 'json', 'txt'],
                        help='Ground truth file format: auto (default, probe labels_json/ then labels/), json, or txt')
    parser.add_argument('--output-dir', type=str, default='eval_results',
                        help='Output directory for evaluation results (default: eval_results)')
    
    # Config file
    parser.add_argument('--config', type=str,
                        help='Path to YAML configuration file')
    
    # Image parameters
    parser.add_argument('--img-width', type=int, default=1920,
                        help='Image width in pixels (default: 1920)')
    parser.add_argument('--img-height', type=int, default=1080,
                        help='Image height in pixels (default: 1080)')
    
    # Evaluation options
    parser.add_argument('--eval-2d-only', action='store_true',
                        help='Evaluate 2D detection only')
    parser.add_argument('--eval-3d-only', action='store_true',
                        help='Evaluate 3D detection only')
    parser.add_argument('--iou-threshold', type=float, default=0.5,
                        help='IoU threshold for 2D matching (default: 0.5)')
    parser.add_argument('--conf-threshold', type=float, default=0.5,
                        help='Confidence threshold for precision/recall (default: 0.5)')
    parser.add_argument('--ap-method', type=str, default='voc2010',
                        choices=['voc2010', 'coco'],
                        help='AP calculation method (default: voc2010)')
    parser.add_argument('--heading-tolerance', type=str, default='strict',
                        choices=['strict', 'relaxed', 'both'],
                        help='Heading error calculation mode: strict (default), relaxed (180° symmetry), or both')
    parser.add_argument('--coord-system', type=str, default='camera',
                        choices=['camera', 'ego'],
                        help='3D evaluation coordinate system: camera (default) or ego')
    parser.add_argument('--num-workers', type=int, default=None,
                        help='Number of parallel workers for evaluation (default: auto-detect)')
    parser.add_argument('--roi-bottom-offset', type=int, default=None,
                        help='Pixels to trim from bottom edge of ROI (shifts y2 upward, default: from config or 0)')
    parser.add_argument('--save-detailed-matches', action='store_true',
                        help='Save detailed 3D match information for common-match comparison')
    
    return parser.parse_args()


def main():
    """Main evaluation function."""
    args = parse_arguments()
    
    # Load configuration
    config = {}
    if args.config:
        print(f"Loading configuration from: {args.config}")
        config = load_config(args.config)
        
        # Override with command line arguments if provided
        if args.det_path:
            config.setdefault('dataset', {})['det_path'] = args.det_path
        if args.gt_path:
            config.setdefault('dataset', {})['gt_path'] = args.gt_path
        if args.path_depth != 1:
            config.setdefault('dataset', {})['path_depth'] = args.path_depth
        if args.det_format != 'auto':
            config.setdefault('dataset', {})['det_format'] = args.det_format
        if args.gt_format != 'auto':
            config.setdefault('dataset', {})['gt_format'] = args.gt_format
        # Only override output_dir if explicitly provided (not default value)
        if args.output_dir != 'eval_results':
            config.setdefault('output', {})['save_path'] = args.output_dir
        
        # Override image size if explicitly provided
        if args.img_width != 1920:
            config.setdefault('image', {})['width'] = args.img_width
        if args.img_height != 1080:
            config.setdefault('image', {})['height'] = args.img_height
        
        # Override matching threshold if explicitly provided
        if args.iou_threshold != 0.5:
            config.setdefault('matching', {})['iou_threshold'] = args.iou_threshold
        
        # Override 2D metrics parameters if explicitly provided
        if args.conf_threshold != 0.5:
            config.setdefault('metrics_2d', {})['conf_threshold'] = args.conf_threshold
        if args.ap_method != 'voc2010':
            config.setdefault('metrics_2d', {})['ap_method'] = args.ap_method
        
        # Override 3D metrics parameters if explicitly provided
        if args.heading_tolerance != 'strict':
            config.setdefault('metrics_3d', {})['heading_tolerance'] = args.heading_tolerance
        if args.coord_system != 'camera':
            config.setdefault('metrics_3d', {})['coordinate_system'] = args.coord_system
        
        # Override roi_bottom_offset if explicitly provided
        if args.roi_bottom_offset is not None:
            config.setdefault('roi_gt', {})['roi_bottom_offset'] = args.roi_bottom_offset
    else:
        # Build config from command line arguments
        if not args.det_path or not args.gt_path:
            print("Error: --det-path and --gt-path are required when not using --config")
            sys.exit(1)
        
        config = {
            'dataset': {
                'det_path': args.det_path,
                'gt_path': args.gt_path,
                'path_depth': args.path_depth,
                'det_format': args.det_format,
                'gt_format': args.gt_format,
            },
            'image': {
                'width': args.img_width,
                'height': args.img_height
            },
            'matching': {
                'iou_threshold': args.iou_threshold
            },
            'metrics_2d': {
                'enabled': not args.eval_3d_only,
                'conf_threshold': args.conf_threshold,
                'ap_method': args.ap_method
            },
            'metrics_3d': {
                'enabled': not args.eval_2d_only,
                'heading_tolerance': args.heading_tolerance,
                'coordinate_system': args.coord_system,
            },
            'output': {
                'save_path': args.output_dir
            }
        }
        
        if args.roi_bottom_offset is not None:
            config.setdefault('roi_gt', {})['roi_bottom_offset'] = args.roi_bottom_offset
    
    # Set evaluation flags
    config['eval_2d'] = config.get('metrics_2d', {}).get('enabled', True)
    config['eval_3d'] = config.get('metrics_3d', {}).get('enabled', True)
    
    if args.eval_2d_only:
        config['eval_2d'] = True
        config['eval_3d'] = False
    elif args.eval_3d_only:
        config['eval_2d'] = False
        config['eval_3d'] = True
    
    # Extract paths
    det_path = config['dataset']['det_path']
    gt_path = config['dataset']['gt_path']
    path_depth = config['dataset'].get('path_depth', 1)  # Default to 1-level structure
    det_format = config['dataset'].get('det_format', 'auto')
    gt_format  = config['dataset'].get('gt_format',  'auto')
    det_subdir = config['dataset'].get('det_subdir')
    output_dir = config['output']['save_path']
    
    img_width = config.get('image', {}).get('width', 1920)
    img_height = config.get('image', {}).get('height', 1080)
    iou_threshold = config.get('matching', {}).get('iou_threshold', 0.5)
    
    # Get num_workers from config if not specified in command line
    num_workers = args.num_workers
    if num_workers is None and 'performance' in config:
        num_workers = config['performance'].get('num_workers', None)
    
    print("="*80)
    print("YOLOv5-3D Model Evaluation")
    print("="*80)
    print(f"Detection path: {det_path}")
    print(f"Ground truth path: {gt_path}")
    print(f"Path depth: {path_depth}")
    print(f"Detection format: {det_format}")
    print(f"Ground truth format: {gt_format}")
    if det_subdir:
        print(f"Detection subdirectory: {det_subdir}")
    print(f"Output directory: {output_dir}")
    print(f"Image size: {img_width}x{img_height}")
    print(f"IoU threshold: {iou_threshold}")
    print(f"Confidence threshold: {config.get('metrics_2d', {}).get('conf_threshold', 0.5)}")
    print(f"AP method: {config.get('metrics_2d', {}).get('ap_method', 'voc2010')}")
    heading_tolerance = config.get('metrics_3d', {}).get('heading_tolerance', 'strict')
    coord_system = config.get('metrics_3d', {}).get('coordinate_system', 'camera')
    print(f"Heading tolerance: {heading_tolerance}")
    print(f"3D coordinate system: {coord_system}")
    roi_bottom_offset_val = config.get('roi_gt', {}).get('roi_bottom_offset', 0)
    if config.get('roi_gt', {}).get('enabled', False):
        print(f"ROI bottom offset: {roi_bottom_offset_val}")
    if num_workers:
        print(f"Number of workers: {num_workers}")
    print(f"Evaluate 2D: {config['eval_2d']}")
    print(f"Evaluate 3D: {config['eval_3d']}")
    if args.save_detailed_matches:
        print(f"Save detailed matches: Yes")
    print("="*80)
    
    # Initialize evaluator
    evaluator = Evaluator(
        config=config,
        iou_threshold=iou_threshold,
        num_workers=num_workers,
        save_detailed_matches=args.save_detailed_matches
    )
    
    # Load data
    print("\nLoading data...")
    evaluator.load_data_from_paths(det_path, gt_path, img_width, img_height, path_depth,
                                    det_format=det_format, gt_format=gt_format)
    
    if len(evaluator.image_pairs) == 0:
        print("Error: No image pairs found for evaluation!")
        sys.exit(1)
    
    # Run evaluation
    results = evaluator.evaluate()
    
    # Generate and save report
    evaluator.generate_report(results, output_dir)
    
    print("\n✓ Evaluation completed successfully!")


if __name__ == '__main__':
    main()
