"""
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 time
import gc
from pathlib import Path
from typing import Dict, Tuple, Optional, List, Any

# Import newly created datasets module for fetching data
from ..datasets import fetch_osf_spatial_voting_2022_a100
from ..models.examples import bjm_spatial_triangle

# 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.
    Delegates to shared datasets module.
    """
    return fetch_osf_spatial_voting_2022_a100()


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, **kwargs) -> 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
        import jax
        from ..datasets import OSF_CACHE_DIR
    except ImportError:
        return {
            'error': 'gridvoting_jax or jax not available',
            'results': [],
            # Need to get cache dir string from datasets if possible, else "unknown"
            'cache_dir': "unknown"
        }
    
    # Ensure OSF data is cached
    cache_dir = ensure_osf_data_cached()
    
    # Determine which configs to test
    if configs is None:
        configs = ALL_CONFIGS

    # Determine solvers
    # PB 2025.12.21 One solver per line for future editing convenience
    # AI 2025.12.26 Added lazy solver variants with consistent naming
    # AI 2025.12.29 Added outline-based solvers
    solvers = kwargs.get('solvers', [
        "power_method",
        "bifurcated_power_method",
        "power_method (lazy)",
        "bifurcated_power_method (lazy)",
        "grid_upscaling",
        "grid_upscaling_lazy_gmres",
        "grid_upscaling_lazy_power",
        "outline_and_fill",
        "outline_and_power",
        "outline_and_gmres",
        "gmres_matrix_inversion",
        "full_matrix_inversion"
    ])


    
    # Check precision
    try:
        # Try new JAX config API then fallback
        is_float64 = jax.config.q("jax_enable_x64") if hasattr(jax.config, "q") else jax.config.jax_enable_x64
        # Or even simpler:
        import jax.numpy as jnp
        is_float64 = (jnp.zeros(1).dtype == jnp.float64)
    except:
        is_float64 = False
        
    precision = "Float64" if is_float64 else "Float32"

    print("="*80)
    print(f"OSF Stationary Distribution Comparison Report")
    print(f"Precision: {precision}")
    print(f"Solvers: {', '.join(solvers)}")
    print("="*80)
    print(f"\nCache directory: {cache_dir}")
    print(f"Testing {len(configs)} configurations x {len(solvers)} solvers\n")
    
    results = []
    
    for g, zi in configs:
        mode = 'ZI' if zi else 'MI'
        
        # Load OSF reference once per config
        print(f"\n{'='*80}")
        print(f"Configuration: g={g}, mode={mode}")
        print(f"{'='*80}")
        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}")
            continue
            
        print(f"  ✓ Loaded {len(osf_df)} grid points")

        # Setup Model (Once per config)
        try:
            # Create SpatialVotingModel (v0.8.0 API)
            # This works for all solvers including grid_upscaling
            # vm = gv.SpatialVotingModel(
            #     voter_ideal_points=voter_ideal_points,
            #     grid=grid,
            #     number_of_voters=3,
            #     majority=2,
            #     zi=zi,
            #     distance_measure='sqeuclidean'
            # )
            vm = bjm_spatial_triangle(g=g, zi=zi)
        except Exception as e:
            print(f"  ✗ Error setting up model: {e}")
            continue

        for solver in solvers:


            print(f"\n  Running {solver}...")
            try:
                start_time = time.time()
                
                # Dispatch to appropriate method based on solver type
                if solver == "power_method (lazy)":
                    vm.analyze_lazy(
                        solver="power_method",
                        force_lazy=True,
                        tolerance=1e-6,
                        max_iterations=5000
                    )
                elif solver == "bifurcated_power_method (lazy)":
                    vm.analyze_lazy(
                        solver="bifurcated_power_method",
                        force_lazy=True,
                        tolerance=1e-6,
                        max_iterations=5000
                    )
                elif solver in ["grid_upscaling", "grid_upscaling_lazy_gmres", "grid_upscaling_lazy_power"]:
                    # Grid upscaling variants handled by analyze()
                    vm.analyze(
                        solver=solver,
                        tolerance=1e-6,
                        max_iterations=5000
                    )
                else:
                    # Standard solvers (power_method, gmres_matrix_inversion, full_matrix_inversion)
                    vm.analyze(
                        solver=solver,
                        tolerance=1e-6,
                        max_iterations=5000
                    )
                runtime = time.time() - start_time
                
                print(f"    ✓ Completed in {runtime:.2f} seconds")
                
                # Compare
                metrics = compare_distributions(osf_df, vm.stationary_distribution)
                
                # Print concise status here
                status = "✓ Excellent" if metrics['l1_norm'] < 1e-6 else \
                         "✓ Good" if metrics['l1_norm'] < 1e-4 else \
                         "⚠ Acceptable" if metrics['l1_norm'] < 1e-2 else "✗ Poor"
                print(f"    L1 Norm: {metrics['l1_norm']:.2e} ({status})")
                
                # Store results
                results.append({
                    'g': g,
                    'zi': zi,
                    'mode': mode,
                    'solver': solver,
                    'precision': precision,
                    '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'],
                })
                
                # Reset analyzed state? VotingModel caches.
                # Actually, vm.analyze() re-runs if called again? 
                # Currently: self.analyzed = True. It checks?
                # Let's check VotingModel.analyze source.
                # It does "if self.core_exists ... self.stationary_distribution = ..."
                # It does NOT check `if self.analyzed: return`. It overwrites.
                # BUT `self.MarkovChain` is created inside analyze.
                # So calling analyze again re-runs. Good.
            except Exception as e:
                print(f"    ✗ Error: {e}")
                results.append({
                    'g': g,
                    'zi': zi,
                    'mode': mode,
                    'solver': solver,
                    'precision': precision,
                    'error': str(e)
                })
            
    # Clean up model
    del vm
    gc.collect()

    # Print summary
    print(f"\n{'='*80}")
    print("SUMMARY")
    print(f"{'='*80}")
    
    successful = [r for r in results if 'l1_norm' in r]
    
    if successful:
        # Sort by g, then mode, then solver
        successful.sort(key=lambda x: (x['g'], x['mode'], x['solver']))
        
        print(f"\n{'Config':<15} {'Solver':<25} {'Time':<10} {'L1 Norm':<12} {'Status':<15}")
        print(f"{'-'*80}")
        
        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['solver']:<25} {r['runtime']:<10.2f} {r['l1_norm']:<12.2e} {status:<15}")
        
    failed = [r for r in results if 'error' in r]
    if failed:
        print(f"\nFailed configs: {len(failed)}")
        for r in failed:
            print(f"  g={r['g']}, {r['mode']}, {r['solver']}: {r['error']}")
    
    return {
        'results': results,
        'cache_dir': str(cache_dir),
        'summary': {
            'total': len(results),
            'successful': len(successful),
            'failed': len(failed)
        }
    }


# Example usage
if __name__ == '__main__':
    import argparse
    import sys
    
    parser = argparse.ArgumentParser(description='Run OSF validation benchmarks.')
    parser.add_argument('--max_g', type=int, default=None,
                      help='Maximum grid size to test (e.g. 40 for quick tests)')
    
    args = parser.parse_args()
    
    # Filter configs based on max_g
    configs_to_run = None
    if args.max_g:
        configs_to_run = [(g, zi) for g, zi in ALL_CONFIGS if g <= args.max_g]
        print(f"Filtering benchmarks: g <= {args.max_g}")
    
    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("\nRunning comparison report...")
    run_comparison_report(configs=configs_to_run)
