#!/usr/bin/env python3
"""
Basic usage examples for GCONCORD package.

This script demonstrates the main functionality of the GCONCORD package
for sparse precision matrix estimation.
"""

import numpy as np
import matplotlib.pyplot as plt
import gconcord as gc


def generate_test_data(p=50, n=100, sparsity=0.1):
    """Generate test data with known sparse structure."""
    np.random.seed(42)
    
    # Create sparse precision matrix
    Omega_true = np.eye(p)
    
    # Add random sparse off-diagonal elements
    for i in range(p):
        for j in range(i+1, p):
            if np.random.rand() < sparsity:
                val = np.random.randn() * 0.3
                Omega_true[i, j] = val
                Omega_true[j, i] = val
    
    # Ensure positive definiteness
    eigenvals = np.linalg.eigvals(Omega_true)
    if np.min(eigenvals) <= 0:
        Omega_true += (0.1 - np.min(eigenvals)) * np.eye(p)
    
    # Generate covariance and data
    Sigma_true = np.linalg.inv(Omega_true)
    data = np.random.multivariate_normal(np.zeros(p), Sigma_true, n)
    S = np.cov(data.T)
    
    return data, S, Omega_true


def example_basic_usage():
    """Demonstrate basic usage of GCONCORD."""
    print("=== Basic Usage Example ===")
    
    # Generate test data
    data, S, Omega_true = generate_test_data(p=20, n=50)
    print(f"Data shape: {data.shape}")
    print(f"Sample covariance shape: {S.shape}")
    
    # Basic estimation with automatic algorithm selection
    result = gc.concord(S, lambda_param=0.1)
    
    print(f"Algorithm used: {result['algorithm_used']}")
    print(f"Converged: {result['converged']}")
    print(f"Iterations: {result['iterations']}")
    print(f"Final tolerance: {result['final_tolerance']:.2e}")
    
    # Check sparsity
    omega_est = result['omega']
    sparsity = 1 - omega_est.nnz / (omega_est.shape[0] ** 2)
    print(f"Estimated sparsity: {sparsity:.2%}")
    print()


def example_algorithm_comparison():
    """Compare different algorithms."""
    print("=== Algorithm Comparison ===")
    
    data, S, _ = generate_test_data(p=15, n=40)
    algorithms = ['coordinate', 'ista', 'auto']
    
    results = {}
    for alg in algorithms:
        result = gc.concord(S, lambda_param=0.1, algorithm=alg, verbose=False)
        results[alg] = result
        
        print(f"{alg.upper():12} | "
              f"Converged: {str(result['converged']):5} | "
              f"Iterations: {result['iterations']:3d} | "
              f"Sparsity: {1 - result['omega'].nnz / (S.shape[0]**2):5.1%}")
    print()


def example_symmetric_lasso():
    """Demonstrate Symmetric Lasso algorithms."""
    print("=== Symmetric Lasso Example ===")
    
    data, S, _ = generate_test_data(p=12, n=30)
    
    # Regular CONCORD
    result_concord = gc.concord(S, lambda_param=0.1, algorithm='ista')
    
    # Symmetric Lasso
    result_symlasso = gc.symlasso(S, lambda_param=0.1)
    
    print(f"CONCORD     | Iterations: {result_concord['iterations']:3d} | "
          f"Sparsity: {1 - result_concord['omega'].nnz / (S.shape[0]**2):5.1%}")
    print(f"Sym Lasso   | Iterations: {result_symlasso['iterations']:3d} | "
          f"Sparsity: {1 - result_symlasso['omega'].nnz / (S.shape[0]**2):5.1%}")
    print()


def example_backward_compatibility():
    """Demonstrate backward compatibility with original ccista."""
    print("=== Backward Compatibility (ccista) ===")
    
    data, S, _ = generate_test_data(p=10, n=25)
    
    # Original ccista interface
    omega, hist = gc.ccista(S, lambda1=0.1, maxitr=50)
    
    print(f"Omega shape: {omega.shape}")
    print(f"History shape: {hist.shape}")
    print(f"Converged in {len(hist)} iterations")
    print(f"Final objective: {hist[-1, 2]:.4f}")
    print(f"Final subgrad norm: {hist[-1, 1]:.2e}")
    print()


def example_convergence_analysis():
    """Demonstrate convergence analysis with saved iterates."""
    print("=== Convergence Analysis ===")
    
    data, S, _ = generate_test_data(p=15, n=40)
    
    # Run with saved iterates
    result = gc.concord(S, lambda_param=0.1, algorithm='ista', 
                       save_iterates=True, max_iterations=100)
    
    if result['objective_values']:
        print(f"Tracking {len(result['objective_values'])} iterations")
        
        # Plot convergence if matplotlib is available
        try:
            plt.figure(figsize=(12, 4))
            
            # Plot objective values
            plt.subplot(1, 2, 1)
            plt.plot(result['objective_values'])
            plt.title('Objective Function')
            plt.xlabel('Iteration')
            plt.ylabel('Objective Value')
            plt.grid(True)
            
            # Plot convergence metrics
            plt.subplot(1, 2, 2)
            plt.semilogy(result['convergence_metrics'])
            plt.title('Convergence Metric')
            plt.xlabel('Iteration')
            plt.ylabel('Convergence Metric (log scale)')
            plt.grid(True)
            
            plt.tight_layout()
            plt.savefig('convergence_analysis.png', dpi=150, bbox_inches='tight')
            print("Convergence plot saved as 'convergence_analysis.png'")
            
        except ImportError:
            print("matplotlib not available for plotting")
    print()


def example_lambda_tuning():
    """Demonstrate lambda parameter tuning.""" 
    print("=== Lambda Parameter Tuning ===")
    
    data, S, _ = generate_test_data(p=12, n=30)
    lambda_values = [0.01, 0.05, 0.1, 0.2, 0.5]
    
    print("Lambda   | Sparsity | Iterations | Converged")
    print("-" * 45)
    
    for lam in lambda_values:
        result = gc.concord(S, lambda_param=lam, algorithm='ista')
        sparsity = 1 - result['omega'].nnz / (S.shape[0] ** 2)
        
        print(f"{lam:6.2f}   | {sparsity:7.1%}  | {result['iterations']:9d}  | "
              f"{str(result['converged']):9}")
    print()


def main():
    """Run all examples."""
    print("GCONCORD Package Examples")
    print("=" * 50)
    print()
    
    # Check if C++ bindings are available
    if not gc._cpp_available:
        print("Warning: C++ bindings not available!")
        print("Please install with: pip install gconcord")
        return
    
    # Run examples
    example_basic_usage()
    example_algorithm_comparison()
    example_symmetric_lasso()
    example_backward_compatibility()
    example_convergence_analysis()
    example_lambda_tuning()
    
    print("All examples completed successfully!")


if __name__ == "__main__":
    main()