import numpy as np
from SALib.sample import fast_sampler, ff, finite_diff, latin, sobol

from .sensitivity_analysis import _get_SA_problem

def run_global_sampler(params_dict: dict, sampler: str, N: int = None, M: int = 4, seed: int = 42) -> dict:
    """Generate parameter samples using global sampling methods.
    
    This function creates parameter samples using various global sampling strategies
    implemented in SALib, suitable for global sensitivity analysis methods.
    
    Args:
        params_dict (dict): Dictionary containing parameter definitions with
            'lower_bound' and 'upper_bound' for each parameter.
        sampler (str): Sampling method to use. Supported methods include:
            - 'Fast': FAST sampling for Fourier Amplitude Sensitivity Test
            - 'Fractional Factorial': Fractional factorial design
            - 'Finite Difference': Finite difference sampling
            - 'Latin hypercube sampling (LHS)': Latin hypercube sampling
            - 'Sobol': Sobol sequence sampling
        N (int, optional): Number of samples to generate. If None, method-specific
            defaults are used. Defaults to None.
        M (int, optional): Number of harmonics for FAST sampler. Only used with
            'Fast' method. Defaults to 4.
        seed (int, optional): Random seed for reproducible sampling. Defaults to 42.
    
    Returns:
        dict: Dictionary with sample IDs as keys and parameter dictionaries as values.
            Each sample dictionary contains parameter names as keys and sampled
            values as values.
    
    Raises:
        ValueError: If an unsupported sampling method is specified or if
            required parameters are missing.
    
    """

    # Define the problem for SALib
    problem = _get_SA_problem(params_dict)
    if sampler == "Fast":
        # Generate model inputs for extended Fourier Amplitude Sensitivity Test.
        # Returns a NumPy matrix containing the model inputs required by the extended Fourier Amplitude sensitivity test. 
        # The resulting matrix contains N * D rows and D columns, where D is the number of parameters.
        # N: The number of samples to generate
        # M: The interference parameter, i.e., the number of harmonics to sum in the Fourier series decomposition (default 4)
        try:
            global_samples = fast_sampler.sample(problem, N, M=M, seed=seed)
        except Exception as e:
            raise ValueError(f"Error in Fast sampler: {e}")
    elif sampler == "Fractional Factorial":
        # Generates model inputs using a fractional factorial sample.
        # Returns a NumPy matrix containing the model inputs required for a fractional factorial analysis. 
        # The resulting matrix has D columns, where D is smallest power of 2 that is greater than the number of parameters. 
        # These model inputs are intended to be used with SALib.analyze.ff.analyze(). 
        try:
            global_samples = ff.sample(problem, seed=seed)
        except Exception as e:
            raise ValueError(f"Error in Fractional Factorial sampler: {e}")
    elif sampler == "Finite Difference":
        # Generate matrix of samples for Derivative-based Global Sensitivity Measure (DGSM).
        # Start from a QMC (Sobol’) sequence and finite difference with delta % steps
        try:
            global_samples = finite_diff.sample(problem, seed=seed)
        except Exception as e:
            raise ValueError(f"Error in Finite Difference sampler: {e}")
    elif sampler == "Latin hypercube sampling (LHS)":
        # Generate model inputs using Latin hypercube sampling (LHS).
        # Returns a NumPy matrix containing the model inputs generated by Latin hypercube sampling. 
        # The resulting matrix contains N rows and D columns, where D is the number of parameters.
        try:
            global_samples = latin.sample(problem, N, seed=seed)
        except Exception as e:
            raise ValueError(f"Error in Latin hypercube sampling sampler: {e}")
    elif sampler == "Sobol":
        # Generates model inputs using Saltelli’s extension of the Sobol’ sequence.
        # The Sobol’ sequence is a popular quasi-random low-discrepancy sequence used to generate uniform samples of parameter space.
        # If second-order indices: N*(2D+2) sample if only first-order indices: N*(D+2) samples, where D is the number of parameters
        try:
            global_samples = sobol.sample(problem, N, calc_second_order=True, seed=seed)
        except Exception as e:
            raise ValueError(f"Error in Sobol sampler: {e}")
    else:
        raise ValueError("Error: Invalid sampler selected.")
    
    # Convert the samples to a dictionary of dictionaries - sorted by sample ID
    global_samples_dict = {i: dict(zip(problem['names'], global_samples[i])) for i in range(len(global_samples))}
    
    return global_samples_dict

def run_local_sampler(params_dict: dict, sampler: str = 'OAT') -> dict:
    """Generate parameter samples using local sampling methods for sensitivity analysis.
    
    This function creates parameter samples using local sampling strategies,
    particularly the One-At-a-Time (OAT) method, where parameters are perturbed
    individually around reference values.
    
    Args:
        params_dict (dict): Dictionary containing parameter definitions. Each
            parameter should have:
            - 'ref_value': Reference value for the parameter
            - 'perturbation': Single value or list of perturbation percentages
        sampler (str, optional): Local sampling method to use. Currently only
            'OAT' (One-At-a-Time) is supported. Defaults to 'OAT'.
    
    Returns:
        dict: Dictionary with sample IDs as keys and parameter dictionaries as values.
            Sample 0 contains the reference values, and subsequent samples contain
            perturbations of individual parameters.
    
    Note:
        For OAT sampling, the first sample (ID 0) contains all reference values.
        Subsequent samples perturb one parameter at a time while keeping others
        at their reference values. If multiple perturbations are specified for
        a parameter, multiple samples are generated for that parameter.
    """
    params_ref = {key: params_dict[key]["ref_value"] for key in params_dict.keys()}
    # First sample is the reference value
    local_samples_dict = {0: params_ref}
    for id, par in enumerate(params_dict.keys()):
        perturbations = np.array(params_dict[par]["perturbation"])
        perturbations = np.concatenate((-perturbations, perturbations))  # Combine negative and positive perturbations
        for idx, var in enumerate(perturbations):
            sample_id = id * len(perturbations) + idx + 1  # Unique sample ID for each perturbation
            if sample_id not in local_samples_dict: # Avoid overwriting existing samples
                local_samples_dict[sample_id] = params_ref.copy()  # Start with reference values
            local_samples_dict[sample_id][par] = params_ref[par] * (1 + var / 100.0)

    return local_samples_dict