"""
Creation:
    Author: Martin Grunnill
    Date: 2024-03-08
Description: Generate Latin-Hypercube Sample (LHS), run simulations and calculate Partial
Correlation Coefficients (PRCCs). For a description of LHS and PRCCs use in model sensitivity analyses see:
    Marino, S., Hogue, I. B., Ray, C. J., & Kirschner, D. E. (2008). A methodology for performing global uncertainty and
    sensitivity analysis in systems biology. In Journal of Theoretical Biology (Vol. 254, Issue 1, pp. 178–196).
    https://doi.org/10.1016/j.jtbi.2008.04.011

Notes
-----
Serial processing of LHS can be slow. but parallel processing of LHS can take up a lot of computing resources.

"""
import pandas as pd
import pingouin as pg
from dask.distributed import Client
from scipy.stats import qmc
from tqdm.auto import tqdm


def _format_sample(parameters_df, LH_samples, other_samples_to_repeat=None):
    """
    Formats Latin Hypercube sample generated by scipy.stats.qmc.

    Scales LH_samples in with boundaries outlined in parameters_df.

    Parameters
    ----------
    parameters_df : pd.DataFrame
        DataFrame outlining the boundaries for each parameter. Must contain fields 'Lower Bound' and
        'Upper Bound'. Name of the parameters is assumed to be in the index.
    LH_samples : numpy.array
        Output from scipy.stats.qmc.LatinHypercube.
    other_samples_to_repeat : pandas.DataFrame
        Samples to resampled and merged with LH samples.

    Returns
    -------
    samples_df : pandas.Dataframe
        Fully formatted samples.
    parameters_sampled : list
        A list of samples being sampled.
    """
    # if any(~parameters_df['Distribution'].isin(['boolean','uniform'])):
    #     raise ValueError('Only Boolean and Uniform distributions currently supported.')
    samples_df = pd.DataFrame(qmc.scale(LH_samples,
                                        parameters_df['Lower Bound'],
                                        parameters_df['Upper Bound']),
                              columns=parameters_df.index)
    if other_samples_to_repeat is not None:
        multiple = len(samples_df) / len(other_samples_to_repeat)
        if not multiple.is_integer():
            raise ValueError(
                'LHS sample size devided by length of other_samples_to_repeat must be expressable as an interger value.')
        multiple = int(multiple)
        repeated_samples = pd.concat([other_samples_to_repeat] * multiple, ignore_index=True)
        samples_df = pd.concat([samples_df, repeated_samples], axis=1, ignore_index=True)
        parameters_sampled = parameters_df.index.to_list() + other_samples_to_repeat.columns.to_list()
        samples_df.columns = parameters_sampled
    else:
        parameters_sampled = samples_df.columns.to_list()

    # convert_to_bool = parameters_df.Parameter[parameters_df['Distribution'] == 'boolean']
    # for parameter in convert_to_bool:
    #     samples_df[parameter] = samples_df[parameter] >= 0.5
    return samples_df, parameters_sampled


def calculate_prcc(results_and_sample_df, parameter, output, covariables, method='spearman'):
    """
    Calculate Partial Correlation Coefficient PCC (default Rank 'Spearman' PRCC).

    A wrapper for pingouin.partial_corr. Partial Rank Correlation Coefficients (PRCCS) can be used to evaluate
    sensitivity of a model to a parameter[1].

    Parameters
    ----------
    results_and_sample_df : pandas.DataFrame
        DataFrame of results and parameters for calculating PCC.
    parameter : string
        Parameter for which PCC will be calculated.
    output : string
        Model output for which PCC will be calculated.
    covariables : list of strings
        Parameters whose effects will be discounted.
    method : string, default 'spearman' (Rank correlations)
        Form of PCC see documentation of pingouin.partial_corr.


    Returns
    -------
    pandas.DataFrame
        Partial Corelation Coefficient of parameter and output.

    References
    ----------
    [1] Marino, S., Hogue, I. B., Ray, C. J., & Kirschner, D. E. (2008). A methodology for performing global uncertainty
        and sensitivity analysis in systems biology. In Journal of Theoretical Biology (Vol. 254, Issue 1, pp. 178–196).
        https://doi.org/10.1016/j.jtbi.2008.04.011
    """
    param_pcor = pg.partial_corr(results_and_sample_df,
                                 x=parameter,
                                 y=output,
                                 covar=covariables,
                                 method=method)
    param_pcor.rename(index={method: parameter + ' on ' + output}, inplace=True)
    param_pcor['Number_of_Covariables'] = len(covariables)
    confidence_interval = pd.DataFrame(param_pcor['CI95%'].tolist())
    param_pcor['lower_CI_0.95'] = confidence_interval[0][0]
    param_pcor['upper_CI_0.95'] = confidence_interval[1][0]
    param_pcor.drop(columns=['CI95%'], inplace=True)
    return param_pcor


def lhs_prcc(parameters_df,
             sample_size,
             model_run_method,
             client=None,
             lhs_obj=None,
             other_samples_to_repeat=None,
             **kwargs):
    """
    Generate a Latin Hypercube sample, run model with sample and calculate PRCC for sampled parameters.

    Latin Hypercube Sampling with Partial Rank Correlation Coefficients (PRCCS) can be used to evaluate sensitivity of a
    model to a parameter[1]. Note currently only supports uniform distribution.

    Parameters
    ----------
    parameters_df : pandas.DataFrame
        DataFrame outlining the boundaries for each parameter. Must contain fields 'Lower Bound' and
        'Upper Bound'. The name of the parameters is assumed to be in the index of the DataFrame.
    sample_size : int
        Sample size of Latin Hypercube.
    model_run_method : function
        Method of running model's simulations. Must accept parameters as a single dictionary. Must output dictionary of 
        input parameters and model results.      
    client: dask.distributed.Client (default None)
        Dask client for running simulations in parallel. If not given simulations are run serially with a tqdm progress
         bar.
    lhs_obj : scipy.stats.qmc.LatinHypercube, optional
        Pre-initialised Latin Hypercube sample generator. If not provided one is generated by within this function.
    other_samples_to_repeat : pandas.DataFrame
        Samples to resampled and merged with LH samples.
    kwargs: 
        Key word arguments to be passed to model_run_method.

    Returns
    -------
    results_df : pandas.DataFrame
        Results of simulations preceded by parameters used in simulations.
    prccs : pandas.DataFrame
        PRCC summary of the model's sensitivity to its parameters.

    References
    ----------
    [1] Marino, S., Hogue, I. B., Ray, C. J., & Kirschner, D. E. (2008). A methodology for performing global uncertainty
        and sensitivity analysis in systems biology. In Journal of Theoretical Biology (Vol. 254, Issue 1, pp. 178–196).
        https://doi.org/10.1016/j.jtbi.2008.04.011
    """
    if not pd.api.types.is_object_dtype(parameters_df.index):
        raise TypeError('Index of parameters_df should be the models parameters. Therefore, they should be strings.')
    if lhs_obj is None:
        num_LH_parameters_sampled = len(parameters_df)
        lhs_obj = qmc.LatinHypercube(num_LH_parameters_sampled)
    LH_sample = lhs_obj.random(sample_size)
    sample_df, parameters_sampled = _format_sample(parameters_df, LH_sample, other_samples_to_repeat)
    if client is None:
        results = run_samples_serially(sample_df, model_run_method, **kwargs)
    else:
        results = run_samples_in_parallel(sample_df, model_run_method, client=client, **kwargs)

    results_df = pd.DataFrame.from_records(results)
    result_columns = [column
                      for column in list(results_df.columns)
                      if column not in parameters_sampled]
    results_df = results_df[parameters_sampled + result_columns]

    prccs = []
    ### FUTURE IMPROVEMENT ###
    # Is it worth running nested for loops below in parallel if given dask client? It already runs very quickly.
    # However, I have never had to many calls of calculate_prcc to run.
    for parameter in parameters_sampled:
        covariables = [item
                       for item in parameters_sampled
                       if item != parameter]
        for column in result_columns:
            param_rank_pcor = calculate_prcc(results_and_sample_df=results_df,
                                             parameter=parameter,
                                             output=column,
                                             covariables=covariables,
                                             method='spearman')
            prccs.append(param_rank_pcor)

    prccs = pd.concat(prccs)
    prccs.sort_index(inplace=True)
    return results_df, prccs


def run_samples_serially(parameters_df, model_run_method, **kwargs):
    results = []
    for parameters in tqdm(parameters_df.to_dict('records'),
                           desc='Simulating LH Sample',
                           position=1,
                           leave=False,
                           colour='green'):
        results.append(model_run_method(parameters, **kwargs))
    return results


def run_samples_in_parallel(parameters_df, model_run_method, client, **kwargs):
    """
    Runs model simulations using parameters in sample_df using parallel processing.

    Parameters
    ----------
    parameters_df : pandas.Dataframe
        Parameter samples being run. Column fields are parameters.
    model_run_method : function
        Method of running model's simulations. Must accept parameters as a single dictionary in the first argument. Must
         output dictionary of input parameters and model results.
    kwargs : 
        Key word arguments to pass to model_run_method.

    Returns
    -------
    If return_focused_results is True a pandas DataFrame of results is returned.
    """
    if not isinstance(client, Client):
        raise TypeError('client must be a dask.distributed.Client.')
    ### FUTURE IMPROVEMENT ###
    # Add progress bar  below so users do not have to use the one on the dask client dashboard.
    futures = []
    for parameters in parameters_df.to_dict('records'):
        future = client.submit(model_run_method, parameters, **kwargs)
        futures.append(future)

    results = client.gather(futures)
    return results

if __name__ == "__main__":
    pass

