#!/usr/bin/env python3
"""
Algorithm comparison example for GCONCORD package.

This script compares different algorithms on various problem sizes and structures.
"""

import numpy as np
import time
import pandas as pd
import gconcord as gc


def generate_structured_precision(p, structure='random', sparsity=0.1):
    """Generate precision matrix with specific structure."""
    np.random.seed(123)
    
    if structure == 'random':
        # Random sparse structure
        Omega = np.eye(p)
        for i in range(p):
            for j in range(i+1, p):
                if np.random.rand() < sparsity:
                    val = np.random.uniform(-0.4, 0.4)
                    Omega[i, j] = val
                    Omega[j, i] = val
                    
    elif structure == 'banded':
        # Banded structure
        Omega = np.eye(p)
        bandwidth = max(1, int(p * sparsity))
        for i in range(p):
            for j in range(i+1, min(i+bandwidth+1, p)):
                val = 0.3 * (bandwidth - (j - i)) / bandwidth
                Omega[i, j] = val
                Omega[j, i] = val
                
    elif structure == 'hub':
        # Hub structure (first few nodes connected to many)
        Omega = np.eye(p)
        n_hubs = max(1, int(p * sparsity / 5))
        for hub in range(n_hubs):
            for j in range(hub+1, p):
                if np.random.rand() < 0.3:
                    val = np.random.uniform(-0.3, 0.3)
                    Omega[hub, j] = val
                    Omega[j, hub] = val
    
    # Ensure positive definiteness
    eigenvals = np.linalg.eigvals(Omega)
    if np.min(eigenvals) <= 0:
        Omega += (0.1 - np.min(eigenvals)) * np.eye(p)
        
    return Omega


def benchmark_algorithms(p_values, n_samples, lambda_val=0.1, structure='random'):
    """Benchmark different algorithms across problem sizes."""
    
    results = []
    
    for p in p_values:
        print(f"\nTesting p={p}, n={n_samples}")
        
        # Generate data
        Omega_true = generate_structured_precision(p, structure)
        Sigma_true = np.linalg.inv(Omega_true)
        data = np.random.multivariate_normal(np.zeros(p), Sigma_true, n_samples)
        S = np.cov(data.T)
        
        # Algorithms to test
        algorithms = ['coordinate', 'ista']
        if p <= 50:  # Add coordinate_blas for medium problems
            algorithms.append('coordinate_blas')
            
        for alg in algorithms:
            print(f"  Testing {alg}...", end='', flush=True)
            
            try:
                start_time = time.time()
                result = gc.concord(S, lambda_param=lambda_val, algorithm=alg, 
                                  max_iterations=500, tolerance=1e-5)
                elapsed_time = time.time() - start_time
                
                # Calculate metrics
                omega_est = result['omega']
                sparsity = 1 - omega_est.nnz / (p ** 2)
                
                # Recovery error (if converged)
                if result['converged']:
                    omega_dense = omega_est.toarray()
                    recovery_error = np.linalg.norm(omega_dense - Omega_true, 'fro') / np.linalg.norm(Omega_true, 'fro')
                else:
                    recovery_error = np.nan
                
                results.append({
                    'p': p,
                    'algorithm': alg,
                    'time': elapsed_time,
                    'iterations': result['iterations'],
                    'converged': result['converged'],
                    'sparsity': sparsity,
                    'recovery_error': recovery_error,
                    'final_tolerance': result['final_tolerance']
                })
                
                print(f" Done ({elapsed_time:.2f}s)")
                
            except Exception as e:
                print(f" Failed: {e}")
                results.append({
                    'p': p,
                    'algorithm': alg,
                    'time': np.nan,
                    'iterations': 0,
                    'converged': False,
                    'sparsity': np.nan,
                    'recovery_error': np.nan,
                    'final_tolerance': np.nan
                })
    
    return pd.DataFrame(results)


def compare_on_structures():
    """Compare algorithms on different graph structures."""
    
    structures = ['random', 'banded', 'hub']
    p = 30
    n = 100
    lambda_val = 0.1
    
    print("Comparing algorithms on different structures")
    print("=" * 50)
    
    for struct in structures:
        print(f"\nStructure: {struct.upper()}")
        print("-" * 30)
        
        # Generate data
        Omega_true = generate_structured_precision(p, struct, sparsity=0.2)
        Sigma_true = np.linalg.inv(Omega_true)
        data = np.random.multivariate_normal(np.zeros(p), Sigma_true, n)
        S = np.cov(data.T)
        
        # Test algorithms
        algorithms = ['coordinate', 'ista', 'symlasso']
        
        for alg in algorithms:
            if alg == 'symlasso':
                result = gc.symlasso(S, lambda_param=lambda_val)
            else:
                result = gc.concord(S, lambda_param=lambda_val, algorithm=alg)
                
            omega_est = result['omega']
            sparsity = 1 - omega_est.nnz / (p ** 2)
            
            print(f"{alg:12} | Iter: {result['iterations']:3d} | "
                  f"Sparsity: {sparsity:5.1%} | "
                  f"Converged: {result['converged']}")


def test_high_dimensional():
    """Test performance on high-dimensional problems (n < p)."""
    
    print("\nHigh-dimensional problem (n < p)")
    print("=" * 50)
    
    n = 50
    p_values = [100, 200]
    
    for p in p_values:
        print(f"\nn={n}, p={p}")
        print("-" * 30)
        
        # Generate high-dimensional data
        data = np.random.randn(n, p)
        S = np.cov(data.T)
        
        # Add small diagonal to ensure positive definiteness
        S += 0.01 * np.eye(p)
        
        # Test different algorithms
        algorithms = ['coordinate', 'coordinate_high_dim']
        
        for alg in algorithms:
            try:
                start_time = time.time()
                result = gc.concord(S, lambda_param=0.5, algorithm=alg,
                                  max_iterations=100)
                elapsed_time = time.time() - start_time
                
                omega_est = result['omega']
                sparsity = 1 - omega_est.nnz / (p ** 2)
                
                print(f"{alg:20} | Time: {elapsed_time:6.2f}s | "
                      f"Iter: {result['iterations']:3d} | "
                      f"Sparsity: {sparsity:5.1%}")
                      
            except Exception as e:
                print(f"{alg:20} | Failed: {e}")


def main():
    """Run all comparison examples."""
    
    print("GCONCORD Algorithm Comparison")
    print("=" * 60)
    
    # Check if C++ bindings are available
    if not gc._cpp_available:
        print("Error: C++ bindings not available!")
        return
    
    # 1. Benchmark on different problem sizes
    print("\n1. PERFORMANCE BENCHMARK")
    print("=" * 60)
    
    p_values = [10, 20, 50]
    n_samples = 100
    
    df_results = benchmark_algorithms(p_values, n_samples)
    
    print("\nBenchmark Results:")
    print(df_results.to_string(index=False))
    
    # 2. Compare on different structures
    print("\n\n2. STRUCTURE COMPARISON")
    print("=" * 60)
    
    compare_on_structures()
    
    # 3. High-dimensional problems
    print("\n\n3. HIGH-DIMENSIONAL PROBLEMS")
    print("=" * 60)
    
    test_high_dimensional()
    
    # Summary
    print("\n\nSUMMARY")
    print("=" * 60)
    print("- Coordinate descent: Best for large, sparse problems")
    print("- ISTA: Best convergence for small problems")
    print("- Symmetric Lasso: Alternative formulation with different properties")
    print("- Auto selection: Chooses based on problem characteristics")


if __name__ == "__main__":
    main()