#!/usr/bin/env python3
"""
Benchmark script to measure actual inference times for CVPR CRF reporting.

This script runs your pipeline on sample images and measures precise timing
for each component to provide accurate compute metrics for the CRF form.

Usage:
    python benchmark_for_cvpr_crf.py --images 10 --dataset cityscapes
    python benchmark_for_cvpr_crf.py --images 5 --dataset custom --custom-path ./test_data/

    !!!
    cd /path/to/repository && python benchmark_for_cvpr_crf.py --images 1 --dataset cityscapes
"""

import argparse
import time
import subprocess
import json
import os
from pathlib import Path
from typing import List, Dict, Tuple
import sys

def measure_pipeline_time(image_path: Path, output_dir: Path) -> Tuple[float, bool]:
    """
    Measure time to process a single image through the full pipeline.
    
    Returns:
        (elapsed_time, success)
    """
    output_path = output_dir / f"output_{image_path.stem}.jpeg"
    
    start = time.time()
    
    # Note: Removed --cleanup-tmp to preserve all intermediate files for benchmark analysis
    command = [
        sys.executable, '-m', 'mcp.clients.agents.main',
        str(image_path),
        str(output_path),
    ]

    try:
        env = os.environ.copy()
        env.setdefault('DETECTOR_BACKEND', 'grounded-sam-2')
        result = subprocess.run(command, capture_output=True, text=True, timeout=300, env=env)
        
        elapsed = time.time() - start
        success = result.returncode == 0
        
        if not success:
            log_path = output_dir / f"{image_path.stem}_pipeline.log"
            with open(log_path, 'w') as log_file:
                log_file.write('COMMAND:\n')
                log_file.write(' '.join(command) + '\n\n')
                log_file.write('STDOUT:\n')
                log_file.write(result.stdout or '<empty>')
                log_file.write('\n\nSTDERR:\n')
                log_file.write(result.stderr or '<empty>')
                log_file.write('\n')

            stderr_preview = (result.stderr or '').strip().splitlines()
            stdout_preview = (result.stdout or '').strip().splitlines()
            preview_lines = stderr_preview if stderr_preview else stdout_preview
            preview_lines = preview_lines[:8]

            if preview_lines:
                print('  └─ Pipeline error (first lines):')
                for line in preview_lines:
                    print(f'      {line}')
            else:
                print('  └─ Pipeline exited with non-zero status but produced no output.')

            print(f'  └─ Full stdout/stderr saved to: {log_path}')

        return elapsed, success
        
    except subprocess.TimeoutExpired:
        print(f"  ✗ TIMEOUT processing {image_path.name}")
        return 300.0, False
    except Exception as e:
        print(f"  ✗ ERROR processing {image_path.name}: {e}")
        return -1, False


def get_dataset_images(dataset: str, max_images: int = 10) -> List[Path]:
    """Get sample images from specified dataset"""
    
    dataset_paths = {
        'cityscapes': Path('data/CityScapes/leftImg8bit/val/frankfurt'),
        'mapillary': Path('data/Mapillary-Vistas-v2.0/validation/images'),
        'market1501': Path('data/Market-1501-v15.09.15/bounding_box_test'),
        'test': Path('test_data'),
    }
    
    if dataset not in dataset_paths:
        raise ValueError(f"Unknown dataset: {dataset}. Choose from: {list(dataset_paths.keys())}")
    
    dataset_path = dataset_paths[dataset]
    
    if not dataset_path.exists():
        raise FileNotFoundError(f"Dataset path not found: {dataset_path}")
    
    # Find image files
    image_extensions = {'.png', '.jpg', '.jpeg', '.JPG', '.PNG', '.JPEG'}
    images = [f for f in dataset_path.iterdir() if f.suffix in image_extensions]
    
    if not images:
        raise ValueError(f"No images found in {dataset_path}")
    
    # Sort for reproducibility
    images.sort()
    
    # Return first N images
    return images[:max_images]


def benchmark_component_times(image_path: Path) -> Dict[str, float]:
    """
    Measure individual component times (approximate, based on logs).
    
    Note: This is a simplified version. For precise per-component timing,
    you would need to instrument your code with timing decorators.
    """
    
    # TODO: Implement detailed component timing if needed
    # For now, return estimates based on configuration
    
    return {
        'yolo_person_detection': 0.05,  # Estimate
        'yolo_license_plates': 0.08,  # Estimate
        'vision_description': 1.5,  # Estimate per person
        'sdxl_inpainting': 5.0,  # Estimate per person
        'llm_orchestration': 2.0,  # Estimate for complex PII
        'grounded_sam2': 0.5,  # Estimate
    }


def run_benchmark(dataset: str, num_images: int, custom_path: str = None):
    """Run benchmark and generate CRF-ready report"""
    
    print("\n" + "="*80)
    print("CVPR CRF COMPUTE BENCHMARK")
    print("="*80)
    print(f"\nDataset: {dataset}")
    print(f"Number of images: {num_images}")
    
    # Get images
    if custom_path:
        custom_dir = Path(custom_path)
        if not custom_dir.exists():
            print(f"Error: Custom path not found: {custom_path}")
            return
        image_extensions = {'.png', '.jpg', '.jpeg', '.JPG', '.PNG', '.JPEG'}
        images = [f for f in custom_dir.iterdir() if f.suffix in image_extensions][:num_images]
    else:
        try:
            images = get_dataset_images(dataset, num_images)
        except Exception as e:
            print(f"Error: {e}")
            return
    
    print(f"Found {len(images)} images")
    
    # Create output directory
    output_dir = Path('artifacts/cvpr_benchmark')
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Run benchmark
    print("\nRunning pipeline benchmark...")
    print("-" * 80)
    
    times = []
    successes = 0
    
    for i, image in enumerate(images, 1):
        print(f"\n[{i}/{len(images)}] Processing: {image.name}")
        
        elapsed, success = measure_pipeline_time(image, output_dir)
        
        if success:
            times.append(elapsed)
            successes += 1
            print(f"  ✓ Completed in {elapsed:.2f}s")
        else:
            print(f"  ✗ Failed (time: {elapsed:.2f}s)")
    
    # Calculate statistics
    if not times:
        print("\n✗ No successful runs. Cannot generate report.")
        return
    
    avg_time = sum(times) / len(times)
    min_time = min(times)
    max_time = max(times)
    
    # Generate report
    print("\n" + "="*80)
    print("BENCHMARK RESULTS")
    print("="*80)
    
    print(f"\nSuccess Rate: {successes}/{len(images)} ({100*successes/len(images):.1f}%)")
    print(f"\nTiming Statistics:")
    print(f"  Average: {avg_time:.2f}s per image")
    print(f"  Min:     {min_time:.2f}s")
    print(f"  Max:     {max_time:.2f}s")
    
    # Dataset-scale estimates
    print(f"\nDataset-Scale Estimates (using average time):")
    
    datasets = {
        'Cityscapes': 5000,
        'Mapillary Vistas': 25000,
        'Market-1501': 32668,
    }
    
    for name, size in datasets.items():
        total_seconds = avg_time * size
        gpu_hours = total_seconds / 3600
        print(f"  {name} ({size:,} images): {gpu_hours:.1f} GPU-hours")
    
    # CRF-ready format
    print("\n" + "="*80)
    print("FOR CVPR CRF FORM (Copy these values)")
    print("="*80)
    
    print(f"""
Per-Image Processing Time: {avg_time:.2f} seconds

GPU-Hours for Standard Datasets:
- Cityscapes (5,000 images): {(avg_time * 5000 / 3600):.1f} GPU-hours
- Mapillary Vistas (25,000 images): {(avg_time * 25000 / 3600):.1f} GPU-hours

Hardware: 2x NVIDIA L40S (44GB VRAM)

Success Rate: {100*successes/len(images):.1f}%

Notes:
- All models are pre-trained (no training compute required for reproduction)
- Times measured on NVIDIA L40S GPUs with dual-GPU setup
- Pipeline includes: Person detection, license plate detection, vision descriptions,
  SDXL inpainting, LLM orchestration for complex PII
- Optimizations: Flash Attention, model CPU offload, VAE slicing
""")
    
    # Save JSON report
    report = {
        'benchmark_config': {
            'dataset': dataset,
            'num_images': num_images,
            'custom_path': custom_path,
        },
        'results': {
            'success_rate': successes / len(images),
            'times_seconds': times,
            'avg_time_seconds': avg_time,
            'min_time_seconds': min_time,
            'max_time_seconds': max_time,
        },
        'dataset_estimates': {
            'cityscapes_5k_gpu_hours': avg_time * 5000 / 3600,
            'mapillary_25k_gpu_hours': avg_time * 25000 / 3600,
            'market1501_33k_gpu_hours': avg_time * 32668 / 3600,
        }
    }
    
    report_path = output_dir / 'cvpr_benchmark_report.json'
    with open(report_path, 'w') as f:
        json.dump(report, f, indent=2)
    
    print(f"\n✓ Detailed report saved to: {report_path}")
    print(f"✓ Output images saved to: {output_dir}")
    
    print("\n" + "="*80)
    print("NEXT STEPS")
    print("="*80)
    print("""
1. Review the timing results above
2. Copy the 'Per-Image Processing Time' and 'GPU-Hours' to your CRF form
3. Update cvpr_compute_report.json with these measured values
4. Re-run: python cvpr_compute_reporting_helper.py

If times seem off:
- Check GPU load (nvidia-smi) - other processes may be interfering
- Try with more images for better statistical confidence
- Verify Ollama models are loaded (first run is slower)
- Check if cleanup-tmp is affecting times (disk I/O)
""")


def main():
    parser = argparse.ArgumentParser(
        description='Benchmark inference times for CVPR CRF reporting',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Benchmark on 10 Cityscapes images
  python benchmark_for_cvpr_crf.py --images 10 --dataset cityscapes
  
  # Benchmark on 5 test images
  python benchmark_for_cvpr_crf.py --images 5 --dataset test
  
  # Benchmark on custom directory
  python benchmark_for_cvpr_crf.py --images 5 --dataset custom --custom-path ./my_images/

Datasets:
  cityscapes  - data/CityScapes/leftImg8bit/val/frankfurt/
  mapillary   - data/Mapillary-Vistas-v2.0/validation/images/
  market1501  - data/Market-1501-v15.09.15/bounding_box_test/
  test        - test_data/
  custom      - specify with --custom-path
"""
    )
    
    parser.add_argument(
        '--images',
        type=int,
        default=10,
        help='Number of images to benchmark (default: 10)'
    )
    
    parser.add_argument(
        '--dataset',
        type=str,
        default='test',
        choices=['cityscapes', 'mapillary', 'market1501', 'test', 'custom'],
        help='Dataset to benchmark on (default: test)'
    )
    
    parser.add_argument(
        '--custom-path',
        type=str,
        help='Path to custom image directory (required if --dataset=custom)'
    )
    
    args = parser.parse_args()
    
    # Validation
    if args.dataset == 'custom' and not args.custom_path:
        parser.error("--custom-path is required when --dataset=custom")
    
    # Run benchmark
    try:
        run_benchmark(args.dataset, args.images, args.custom_path)
    except KeyboardInterrupt:
        print("\n\n✗ Benchmark interrupted by user")
        sys.exit(1)
    except Exception as e:
        print(f"\n✗ Error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == '__main__':
    main()
