"""
Core Python interface for GCONCORD algorithms.

This module provides a clean, Pythonic interface to the underlying C++ implementations
while maintaining backward compatibility with the original gconcord package.
"""

import numpy as np
import scipy.sparse
from typing import Union, Optional, Tuple, Dict, Any

try:
    from . import _gconcord
    _cpp_available = True
except ImportError:
    _cpp_available = False

# Algorithm mapping functions for lazy evaluation
def _get_algorithm_enum(algorithm: str):
    """Get algorithm enum safely with lazy evaluation."""
    if not _cpp_available:
        raise RuntimeError("C++ bindings not available")
    
    mapping = {
        'coordinate': _gconcord.Algorithm.COORDINATE_DESCENT,
        'coordinate_high_dim': _gconcord.Algorithm.COORDINATE_HIGH_DIM,
        'coordinate_blas': _gconcord.Algorithm.COORDINATE_BLAS,
        'ista': _gconcord.Algorithm.ISTA,
        'auto': _gconcord.Algorithm.AUTO
    }
    
    if algorithm not in mapping:
        raise ValueError(f"Unknown algorithm: {algorithm}. Must be one of {list(mapping.keys())}")
    
    return mapping[algorithm]

def _get_symlasso_enum(algorithm: str):
    """Get symmetric lasso algorithm enum safely with lazy evaluation."""
    if not _cpp_available:
        raise RuntimeError("C++ bindings not available")
    
    mapping = {
        'symlasso': _gconcord.Algorithm.SYMLASSO,
        'symlasso_iterates': _gconcord.Algorithm.SYMLASSO_ITERATES,
        'auto': _gconcord.Algorithm.SYMLASSO  # Default to standard symlasso
    }
    
    if algorithm not in mapping:
        raise ValueError(f"Unknown symmetric lasso algorithm: {algorithm}. Must be one of {list(mapping.keys())}")
    
    return mapping[algorithm]


def _validate_inputs(S: np.ndarray, lambda_param: Union[float, np.ndarray]) -> None:
    """Validate input parameters."""
    if not isinstance(S, np.ndarray):
        raise TypeError("S must be a numpy array")
    
    if S.ndim != 2 or S.shape[0] != S.shape[1]:
        raise ValueError("S must be a square matrix")
    
    if not np.allclose(S, S.T):
        raise ValueError("S must be symmetric")
    
    # Check for finite values
    if not np.all(np.isfinite(S)):
        raise ValueError("S contains non-finite values (NaN or Inf)")
    
    # Check positive definiteness by attempting Cholesky decomposition
    try:
        np.linalg.cholesky(S)
    except np.linalg.LinAlgError:
        # If not positive definite, at least check for positive semi-definiteness
        eigenvalues = np.linalg.eigvals(S)
        if np.any(eigenvalues < -1e-10):  # Allow small numerical errors
            raise ValueError("S must be positive semi-definite")
    
    # Validate lambda parameter
    if isinstance(lambda_param, np.ndarray):
        if lambda_param.shape != S.shape:
            raise ValueError("Lambda matrix must have same shape as S")
        if not np.allclose(lambda_param, lambda_param.T):
            raise ValueError("Lambda matrix must be symmetric")
        if np.any(lambda_param < 0):
            raise ValueError("Lambda values must be non-negative")
    elif isinstance(lambda_param, (int, float)):
        if lambda_param < 0:
            raise ValueError("Lambda must be non-negative")
    else:
        raise TypeError("lambda_param must be a scalar or array")


def concord(
    S: Optional[np.ndarray] = None,
    lambda_param: Union[float, np.ndarray] = None,
    data: Optional[np.ndarray] = None,
    algorithm: str = 'auto',
    tolerance: float = 1e-5,
    max_iterations: int = 100,
    save_iterates: bool = False,
    penalize_diagonal: bool = False,
    verbose: bool = False,
    lambda_scale: float = 2.0,
    **kwargs  # Accept additional parameters for compatibility
) -> Dict[str, Any]:
    """
    Solve sparse precision matrix estimation using CONCORD algorithms.
    
    This function provides access to coordinate descent and ISTA algorithms
    for estimating sparse precision matrices in graphical models.
    
    Parameters
    ----------
    S : np.ndarray, optional
        Sample covariance matrix (symmetric, positive definite).
        Either S or data must be provided, but not both.
    lambda_param : float or np.ndarray
        Penalty parameter(s). If scalar, same penalty for all off-diagonal elements.
        If matrix, element-wise penalties.
    data : np.ndarray, optional
        Raw data matrix (n x p). If provided, sample covariance is computed internally.
        Either S or data must be provided, but not both.
    algorithm : str, default='auto'
        Algorithm to use: 'coordinate', 'coordinate_high_dim', 'coordinate_blas', 'ista', 'auto'
    tolerance : float, default=1e-5
        Convergence tolerance
    max_iterations : int, default=100
        Maximum number of iterations
    save_iterates : bool, default=False
        Whether to save iteration history
    penalize_diagonal : bool, default=False
        Whether to penalize diagonal elements
    verbose : bool, default=False
        Whether to print convergence information
    lambda_scale : float, default=2.0
        Lambda scaling factor (for compatibility with original R interface)
        
    Returns
    -------
    dict
        Dictionary containing:
        - 'omega': Estimated precision matrix (scipy.sparse or np.ndarray)
        - 'converged': Whether algorithm converged (bool)
        - 'iterations': Number of iterations (int)
        - 'algorithm_used': Algorithm that was used (str)
        - 'objective_values': Objective function values (list, if save_iterates=True)
        - 'convergence_metrics': Convergence metrics (list, if save_iterates=True)
        
    Raises
    ------
    ValueError
        If inputs are invalid
    RuntimeError  
        If C++ bindings are not available
        
    Examples
    --------
    >>> import numpy as np
    >>> import gconcord as gc
    >>> S = np.eye(3) + 0.5 * np.random.randn(3, 3)
    >>> S = (S + S.T) / 2  # Make symmetric
    >>> result = gc.concord(S, lambda_param=0.1)
    >>> print(f"Converged: {result['converged']}")
    """
    if not _cpp_available:
        raise RuntimeError("C++ bindings not available. Please install with: pip install gconcord")
    
    # Check that exactly one of S or data is provided
    if (S is None) == (data is None):
        raise ValueError("Exactly one of 'S' or 'data' must be provided")
    
    # If data is provided, compute S for validation and use
    if data is not None:
        data = np.asarray(data, dtype=np.float64)
        if data.ndim != 2:
            raise ValueError("Data must be a 2D array (n x p)")
        S_for_validation = np.cov(data, rowvar=False)
    else:
        S_for_validation = S
        S = np.asarray(S, dtype=np.float64)
    
    # Validate inputs
    _validate_inputs(S_for_validation, lambda_param)
    
    # Create options
    options = _gconcord.ConcordOptions()
    options.algorithm = _get_algorithm_enum(algorithm)
    options.tolerance = tolerance
    options.max_iterations = max_iterations
    options.save_iterates = save_iterates
    options.penalize_diagonal = penalize_diagonal
    options.verbose = verbose
    options.lambda_scale = lambda_scale
    
    # Create solver and solve
    solver = _gconcord.ConcordSolver(options)
    
    # Use data matrix if provided (for high-dimensional algorithms)
    if data is not None:
        if isinstance(lambda_param, (int, float)):
            result = solver.solve_from_data(data, float(lambda_param))
        else:
            # solve_from_data only supports scalar lambda
            raise ValueError("Matrix lambda is not supported when using raw data")
    else:
        if isinstance(lambda_param, (int, float)):
            result = solver.solve(S, float(lambda_param))
        else:
            result = solver.solve(S, lambda_param)
    
    # Convert result to Python dict
    return {
        'omega': scipy.sparse.csc_matrix(result.omega),
        'converged': result.converged,
        'iterations': result.total_iterations,
        'algorithm_used': str(result.algorithm_used).split('.')[-1].lower(),
        'final_tolerance': result.final_tolerance,
        'convergence_reason': result.convergence_reason,
        'objective_values': result.objective_values if save_iterates else None,
        'convergence_metrics': result.convergence_metrics if save_iterates else None,
        'inner_iterations': result.inner_iterations if save_iterates else None,
    }


def symlasso(
    S: np.ndarray,
    lambda_param: float,
    algorithm: str = 'auto', 
    tolerance: float = 1e-5,
    max_iterations: int = 100,
    save_iterates: bool = False,
    verbose: bool = False
) -> Dict[str, Any]:
    """
    Solve sparse precision matrix estimation using Symmetric Lasso algorithms.
    
    This function provides access to symmetric lasso algorithms which use
    auxiliary variables for efficient diagonal updates.
    
    Parameters
    ----------
    S : np.ndarray
        Sample covariance matrix (symmetric, positive definite)
    lambda_param : float
        Penalty parameter (scalar only for symmetric lasso)
    algorithm : str, default='auto'
        Algorithm variant: 'symlasso', 'symlasso_iterates', 'auto'
    tolerance : float, default=1e-5
        Convergence tolerance
    max_iterations : int, default=100
        Maximum number of iterations
    save_iterates : bool, default=False
        Whether to save iteration history
    verbose : bool, default=False
        Whether to print convergence information
        
    Returns
    -------
    dict
        Dictionary with same structure as concord()
        
    Examples
    --------
    >>> import numpy as np
    >>> import gconcord as gc
    >>> S = np.eye(3) + 0.5 * np.random.randn(3, 3) 
    >>> S = (S + S.T) / 2
    >>> result = gc.symlasso(S, lambda_param=0.1)
    >>> print(f"Converged: {result['converged']}")
    """
    if not _cpp_available:
        raise RuntimeError("C++ bindings not available. Please install with: pip install gconcord")
    
    # Validate inputs
    _validate_inputs(S, lambda_param)
    
    if not isinstance(lambda_param, (int, float)):
        raise TypeError("Symmetric Lasso requires scalar lambda parameter")
    
    # Create options
    options = _gconcord.ConcordOptions()
    options.algorithm = _get_symlasso_enum(algorithm)
    options.tolerance = tolerance
    options.max_iterations = max_iterations
    options.save_iterates = save_iterates
    options.verbose = verbose
    
    # Create solver and solve
    solver = _gconcord.ConcordSolver(options)
    result = solver.solve(S, float(lambda_param))
    
    # Convert result to Python dict
    return {
        'omega': scipy.sparse.csc_matrix(result.omega),
        'converged': result.converged,
        'iterations': result.total_iterations,
        'algorithm_used': str(result.algorithm_used).split('.')[-1].lower(),
        'final_tolerance': result.final_tolerance,
        'convergence_reason': result.convergence_reason,
        'objective_values': result.objective_values if save_iterates else None,
        'convergence_metrics': result.convergence_metrics if save_iterates else None,
        'inner_iterations': result.inner_iterations if save_iterates else None,
    }


def ccista(
    S: np.ndarray,
    lambda1: Union[float, np.ndarray],
    epstol: float = 1e-5,
    maxitr: int = 100,
    penalize_diagonal: bool = False
) -> Tuple[scipy.sparse.csc_matrix, np.ndarray]:
    """
    ISTA algorithm for sparse precision matrix estimation (backward compatibility).
    
    This function maintains backward compatibility with the original gconcord package.
    It provides the same interface as the original ccista function.
    
    Parameters
    ----------
    S : np.ndarray
        Sample covariance matrix
    lambda1 : float or np.ndarray
        Penalty parameter(s)
    epstol : float, default=1e-5
        Convergence tolerance
    maxitr : int, default=100
        Maximum iterations
    penalize_diagonal : bool, default=False
        Whether to penalize diagonal elements
        
    Returns
    -------
    tuple
        (omega, hist) where:
        - omega: Sparse precision matrix estimate
        - hist: Convergence history array with columns [inner_iter_count, delta_subg, objective]
        
    Examples
    --------
    >>> import numpy as np
    >>> import gconcord as gc
    >>> S = np.eye(3)
    >>> omega, hist = gc.ccista(S, lambda1=0.1)
    >>> print(f"Converged in {len(hist)} iterations")
    """
    if not _cpp_available:
        raise RuntimeError("C++ bindings not available. Please install with: pip install gconcord")
    
    # Use the unified interface but format output for backward compatibility
    result = concord(
        S=S,
        lambda_param=lambda1,
        algorithm='ista',
        tolerance=epstol,
        max_iterations=maxitr,
        save_iterates=True,
        penalize_diagonal=penalize_diagonal,
        verbose=False
    )
    
    # Format history for backward compatibility
    hist = []
    if result['objective_values'] is not None:
        inner_iters = result['inner_iterations'] or [0] * len(result['objective_values'])
        conv_metrics = result['convergence_metrics'] or [0.0] * len(result['objective_values'])
        objectives = result['objective_values']
        
        for i in range(len(objectives)):
            hist.append([
                inner_iters[i] if i < len(inner_iters) else 0,
                conv_metrics[i] if i < len(conv_metrics) else 0.0,
                objectives[i]
            ])
    
    return result['omega'], np.array(hist) if hist else np.array([[0, 0.0, 0.0]])