"""
Utilities for comparing benchmark results with OSF stationary distributions.

This module provides functions to load and compare stationary distributions
from the OSF repository with those generated by gridvoting-jax benchmarks.

The OSF data is automatically downloaded and cached in /tmp on first use.
"""

import numpy as np
import zipfile
import io
import shutil
import time
from pathlib import Path
from typing import Dict, Tuple, Optional, List, Any


import os

# Cache directory
OSF_CACHE_DIR = Path(os.environ.get('GV_OSF_CACHE_DIR', '/tmp/gridvoting_osf_cache'))
OSF_FILE_ID = 'kms9z'  # Zip file with A100 GPU replication data

# All available configurations
ALL_CONFIGS = [
    (20, False), (20, True),
    (40, False), (40, True),
    (60, False), (60, True),
    (80, False), (80, True),
]


def ensure_osf_data_cached() -> Path:
    """
    Ensure OSF data is downloaded and cached in /tmp.
    
    Returns:
        Path to cache directory
    """
    # Check if cache exists and has data
    if OSF_CACHE_DIR.exists():
        csv_files = list(OSF_CACHE_DIR.glob('*_*_stationary_distribution.csv'))
        if len(csv_files) >= 8:  # All 8 configurations present
            return OSF_CACHE_DIR
    
    # Lazy import requests to avoid hard dependency
    try:
        import requests
    except ImportError:
        print("Warning: 'requests' module not installed. Cannot download OSF data.")
        return OSF_CACHE_DIR

    # Download and cache
    print(f"Downloading OSF data to cache: {OSF_CACHE_DIR}")
    OSF_CACHE_DIR.mkdir(parents=True, exist_ok=True)
    
    download_url = f"https://osf.io/download/{OSF_FILE_ID}/"
    
    try:
        response = requests.get(download_url, timeout=60)
        response.raise_for_status()
        
        # Extract zip file
        zip_file = zipfile.ZipFile(io.BytesIO(response.content))
        
        # Find Run2 directory in zip
        run2_files = [name for name in zip_file.namelist() 
                     if 'Run2' in name and name.endswith('_stationary_distribution.csv')]
        
        # Extract stationary distribution files
        for filename in run2_files:
            # Extract to cache with simplified name
            basename = Path(filename).name
            with zip_file.open(filename) as source:
                dest_path = OSF_CACHE_DIR / basename
                with open(dest_path, 'wb') as dest:
                    dest.write(source.read())
        
        print(f"  Downloaded {len(run2_files)} stationary distributions")
        
    except Exception as e:
        print(f"  Warning: Failed to download OSF data: {e}")
        print(f"  Comparison with OSF data will not be available")
    
    return OSF_CACHE_DIR


def load_osf_distribution(g: int, zi: bool) -> Optional[Any]:
    """
    Load a stationary distribution from the OSF cache.
    
    Automatically downloads and caches data on first use.
    
    Args:
        g: Grid size (20, 40, 60, or 80)
        zi: True for Zero Intelligence (ZI), False for Myopic Intelligence (MI)
    
    Returns:
        DataFrame with columns: r, x, y, theta, log10prob, log10diagonal, log10r
        Returns None if the file doesn't exist or pandas is not installed.
    
    Example:
        >>> df = load_osf_distribution(g=20, zi=False)  # Load 20_MI
        >>> probabilities = 10 ** df['log10prob']
    """
    try:
        import pandas as pd
    except ImportError:
        print("Warning: pandas not installed. Cannot load OSF distribution.")
        return None

    # Ensure data is cached
    cache_dir = ensure_osf_data_cached()
    
    mode = 'ZI' if zi else 'MI'
    filename = f'{g}_{mode}_stationary_distribution.csv'
    filepath = cache_dir / filename
    
    if not filepath.exists():
        return None
    
    df = pd.read_csv(filepath)
    # Drop the unnamed index column if it exists
    if 'Unnamed: 0' in df.columns:
        df = df.drop(columns=['Unnamed: 0'])
    
    return df


def get_available_distributions() -> Dict[Tuple[int, bool], str]:
    """
    Get a dictionary of available OSF distributions.
    
    Automatically downloads and caches data on first use.
    
    Returns:
        Dictionary mapping (g, zi) tuples to filenames
    
    Example:
        >>> available = get_available_distributions()
        >>> print(available)
        {(20, False): '20_MI_stationary_distribution.csv', ...}
    """
    cache_dir = ensure_osf_data_cached()
    available = {}
    
    for filepath in cache_dir.glob('*_stationary_distribution.csv'):
        filename = filepath.name
        parts = filename.split('_')
        if len(parts) >= 2:
            try:
                g = int(parts[0])
                mode = parts[1]
                zi = (mode == 'ZI')
                available[(g, zi)] = filename
            except ValueError:
                continue
    
    return available


def compare_distributions(
    osf_df: Any,
    benchmark_dist: np.ndarray,
    grid_coords: Optional[np.ndarray] = None
) -> Dict[str, float]:
    """
    Compare an OSF stationary distribution with a benchmark result.
    
    Args:
        osf_df: DataFrame from load_osf_distribution()
        benchmark_dist: Stationary distribution from benchmark (1D array)
        grid_coords: Optional grid coordinates for spatial comparison
    
    Returns:
        Dictionary with comparison metrics:
        - 'l1_norm': L1 norm of difference
        - 'l2_norm': L2 norm of difference
        - 'max_abs_diff': Maximum absolute difference
        - 'correlation': Correlation coefficient
    
    Example:
        >>> osf_df = load_osf_distribution(g=20, zi=False)
        >>> # Run benchmark and get stationary distribution
        >>> vm.analyze()
        >>> benchmark_dist = vm.stationary_distribution
        >>> metrics = compare_distributions(osf_df, benchmark_dist)
        >>> print(f"L1 norm: {metrics['l1_norm']:.6f}")
    """
    # Convert OSF log probabilities to actual probabilities
    osf_probs = 10 ** osf_df['log10prob'].values
    
    # Ensure same length
    if len(osf_probs) != len(benchmark_dist):
        raise ValueError(
            f"Distribution lengths don't match: OSF has {len(osf_probs)}, "
            f"benchmark has {len(benchmark_dist)}"
        )
    
    # Calculate metrics
    diff = osf_probs - benchmark_dist
    
    metrics = {
        'l1_norm': np.sum(np.abs(diff)),
        'l2_norm': np.sqrt(np.sum(diff ** 2)),
        'max_abs_diff': np.max(np.abs(diff)),
        'correlation': np.corrcoef(osf_probs, benchmark_dist)[0, 1],
        'osf_sum': np.sum(osf_probs),
        'benchmark_sum': np.sum(benchmark_dist),
    }
    
    return metrics


def format_comparison_report(
    g: int,
    zi: bool,
    metrics: Dict[str, float],
    runtime: Optional[float] = None
) -> str:
    """
    Format a comparison report as a string.
    
    Args:
        g: Grid size
        zi: Zero Intelligence mode
        metrics: Dictionary from compare_distributions()
        runtime: Optional runtime in seconds
    
    Returns:
        Formatted report string
    """
    mode = 'ZI' if zi else 'MI'
    
    report = [
        f"Comparison Report: g={g}, mode={mode}",
        "=" * 60,
        f"L1 Norm (sum of absolute differences): {metrics['l1_norm']:.8f}",
        f"L2 Norm (Euclidean distance):          {metrics['l2_norm']:.8f}",
        f"Maximum absolute difference:           {metrics['max_abs_diff']:.8f}",
        f"Correlation coefficient:               {metrics['correlation']:.8f}",
        f"OSF probability sum:                   {metrics['osf_sum']:.10f}",
        f"Benchmark probability sum:             {metrics['benchmark_sum']:.10f}",
    ]
    
    if runtime is not None:
        report.append(f"Benchmark runtime:                     {runtime:.2f} seconds")
    
    # Add interpretation
    report.append("")
    if metrics['l1_norm'] < 1e-6:
        report.append("✓ Excellent match (L1 < 1e-6)")
    elif metrics['l1_norm'] < 1e-4:
        report.append("✓ Good match (L1 < 1e-4)")
    elif metrics['l1_norm'] < 1e-2:
        report.append("⚠ Acceptable match (L1 < 1e-2)")
    else:
        report.append("✗ Poor match (L1 >= 1e-2)")
    
    return "\n".join(report)


def run_comparison_report(configs: Optional[List[Tuple[int, bool]]] = None) -> Dict:
    """
    Run complete OSF comparison report for specified configurations.
    
    This is the single entry point for running OSF comparisons. It will:
    1. Automatically download and cache OSF data if needed
    2. Run benchmarks for specified configurations
    3. Compare with OSF data using L1 norm
    4. Return comprehensive report
    
    Args:
        configs: List of (g, zi) tuples to test. If None, tests all available configs.
                 Example: [(20, False), (40, True)]
    
    Returns:
        Dictionary with:
        - 'results': List of comparison results for each config
        - 'summary': Summary statistics
        - 'cache_dir': Path to OSF cache directory
    
    Example:
        >>> from gridvoting_jax.benchmarks.osf_comparison import run_comparison_report
        >>> report = run_comparison_report([(20, False), (20, True)])
        >>> for result in report['results']:
        ...     print(f"g={result['g']}, {result['mode']}: L1={result['l1_norm']:.2e}")
    """
    # Import here to avoid circular dependency
    try:
        import gridvoting_jax as gv
    except ImportError:
        return {
            'error': 'gridvoting_jax not available',
            'results': [],
            'cache_dir': str(OSF_CACHE_DIR)
        }
    
    # Ensure OSF data is cached
    cache_dir = ensure_osf_data_cached()
    
    # Determine which configs to test
    if configs is None:
        configs = ALL_CONFIGS
    
    print("="*70)
    print("OSF Stationary Distribution Comparison Report")
    print("Using L1 norm as primary comparison metric")
    print("="*70)
    print(f"\nCache directory: {cache_dir}")
    print(f"Testing {len(configs)} configurations\n")
    
    results = []
    
    for g, zi in configs:
        mode = 'ZI' if zi else 'MI'
        print(f"\n{'='*70}")
        print(f"Configuration: g={g}, mode={mode}")
        print(f"{'='*70}")
        
        # Load OSF reference
        print(f"Loading OSF reference data...")
        osf_df = load_osf_distribution(g=g, zi=zi)
        
        if osf_df is None:
            print(f"  ✗ OSF data not available for g={g}, {mode}")
            results.append({
                'g': g,
                'zi': zi,
                'mode': mode,
                'error': 'OSF data not available'
            })
            continue
        
        print(f"  ✓ Loaded {len(osf_df)} grid points")
        
        # Run benchmark
        print(f"Running benchmark...")
        try:
            start_time = time.time()
            
            # Setup
            grid = gv.Grid(x0=-g, x1=g, y0=-g, y1=g)
            number_of_alternatives = (2*g+1)**2
            voter_ideal_points = [[-15, -9], [0, 17], [15, -9]]
            
            u = grid.spatial_utilities(
                voter_ideal_points=voter_ideal_points,
                metric='sqeuclidean'
            )
            
            vm = gv.VotingModel(
                utility_functions=u,
                majority=2,
                zi=zi,
                number_of_voters=3,
                number_of_feasible_alternatives=number_of_alternatives
            )
            
            vm.analyze()
            runtime = time.time() - start_time
            
            print(f"  ✓ Completed in {runtime:.2f} seconds")
            
            # Compare
            print(f"Comparing with OSF data (L1 norm)...")
            metrics = compare_distributions(osf_df, vm.stationary_distribution)
            
            # Print report
            print(f"\n{format_comparison_report(g, zi, metrics, runtime)}")
            
            # Store results
            results.append({
                'g': g,
                'zi': zi,
                'mode': mode,
                'runtime': runtime,
                'l1_norm': metrics['l1_norm'],
                'l2_norm': metrics['l2_norm'],
                'max_abs_diff': metrics['max_abs_diff'],
                'correlation': metrics['correlation'],
                'osf_sum': metrics['osf_sum'],
                'benchmark_sum': metrics['benchmark_sum'],
            })
            
        except Exception as e:
            print(f"  ✗ Error running benchmark: {e}")
            results.append({
                'g': g,
                'zi': zi,
                'mode': mode,
                'error': str(e)
            })
    
    # Print summary
    print(f"\n{'='*70}")
    print("SUMMARY")
    print(f"{'='*70}")
    
    successful = [r for r in results if 'l1_norm' in r]
    
    if successful:
        print(f"\n{'Config':<15} {'Runtime':<12} {'L1 Norm':<15} {'Status':<15}")
        print(f"{'-'*70}")
        
        for r in successful:
            config = f"g={r['g']}, {r['mode']}"
            status = "✓ Excellent" if r['l1_norm'] < 1e-6 else \
                     "✓ Good" if r['l1_norm'] < 1e-4 else \
                     "⚠ Acceptable" if r['l1_norm'] < 1e-2 else "✗ Poor"
            print(f"{config:<15} {r['runtime']:<12.2f} {r['l1_norm']:<15.2e} {status:<15}")
        
        # Overall statistics
        l1_norms = [r['l1_norm'] for r in successful]
        print(f"\nOverall L1 norm statistics:")
        print(f"  Mean:   {np.mean(l1_norms):.2e}")
        print(f"  Median: {np.median(l1_norms):.2e}")
        print(f"  Max:    {np.max(l1_norms):.2e}")
        print(f"  Min:    {np.min(l1_norms):.2e}")
    
    failed = [r for r in results if 'error' in r]
    if failed:
        print(f"\nFailed configurations: {len(failed)}")
        for r in failed:
            print(f"  g={r['g']}, {r['mode']}: {r['error']}")
    
    return {
        'results': results,
        'cache_dir': str(cache_dir),
        'summary': {
            'total': len(results),
            'successful': len(successful),
            'failed': len(failed),
            'mean_l1': np.mean([r['l1_norm'] for r in successful]) if successful else None,
        }
    }


# Example usage
if __name__ == '__main__':
    print("Available OSF distributions:")
    for (g, zi), filename in get_available_distributions().items():
        mode = 'ZI' if zi else 'MI'
        print(f"  g={g}, {mode}: {filename}")
    
    print("\nLoading g=20, MI distribution:")
    df = load_osf_distribution(g=20, zi=False)
    if df is not None:
        print(f"  Loaded {len(df)} rows")
        print(f"  Columns: {list(df.columns)}")
        probs = 10 ** df['log10prob']
        print(f"  Probability sum: {probs.sum():.10f}")
