from scipy.stats import spearmanr, pearsonr, rankdata

import deampy.format_functions as F
import deampy.in_out_functions as IO
from deampy.statistics import partial_corr


class SensitivityAnalysis:
    def __init__(self, dic_parameter_values, output_values):
        """
        :param dic_parameter_values: (dictionary) of parameter values with parameter names as the key
        :param output_values: (list) of output values (e.g. cost or QALY observations)
        """

        self.dicParameterValues = dic_parameter_values
        self.outputValues = output_values

        for paramName, paramValues in dic_parameter_values.items():

            if len(paramValues) != len(output_values):
                raise ValueError('Number of parameter values should be equal to the number of output values. '
                                 'Error in parameter "{}".'.format(paramName))

    def get_pearson_corr(self):
        """ :returns (list) of [parameter name, Pearson's correlation coefficients, p-value] """

        return self._get_corr(f=pearsonr)

    def get_spearman_corr(self):
        """ :returns (list) of [parameter name, Spearman's rank correlation coefficients, p-value] """

        return self._get_corr(f=spearmanr)

    def get_partial_corr(self):
        """ :returns (list) of [parameter name, partial correlation coefficients, p-value] """

        return self._get_partial_corr(dic_param_values=self.dicParameterValues, output_values=self.outputValues)

    def get_partial_rank_corr(self):
        """ :returns (list) of [parameter name, partial rank correlation coefficients, p-value] """

        # find ranked output values
        ranked_outputs = rankdata(self.outputValues)
        # find ranked parameter values
        dic_ranked_param_values = dict()
        for paramName, paramValues in self.dicParameterValues.items():
            dic_ranked_param_values[paramName] = rankdata(paramValues)

        return self._get_partial_corr(dic_param_values=dic_ranked_param_values, output_values=ranked_outputs)

    def print_corr(self, corr='r'):
        """
        :param corr: 'r' for Pearson's,
                     'rho' for Spearman's,
                     'p' for partial correlation, and
                     'pr' for partial rank correlation
        :return:
        """

        results, text = self._get_results_text(corr=corr)
        print(text)
        for par, values in results.items():
            print('\t', par, '{} , {}'.format(values[0], values[1]))

    def export_to_csv(self, corrs='r', file_name='Correlations.csv', decimal=3, delimiter=','):
        """
        formats the correlation coefficients and p-value to the specified decimal point and exports to a csv file
        :param corrs: (string) or (list of strings) from
                     'r' for Pearson's,
                     'rho' for Spearman's rank correlation,
                     'p' for partial correlation, and
                     'pr' for partial rank correlation
        :param file_name: file name
        :param decimal: decimal points to round the estimates to
        :param delimiter: to separate by comma, use ',' and by tab, use '\t'
        """

        if not isinstance(corrs, list):
            corrs = [corrs]

        # make the header
        tile_row = ['Parameter']
        for corr in corrs:
            tile_row.append(self._full_name(corr=corr))
            tile_row.append('P-value')

        # add the header
        formatted_results = [tile_row]

        # add the names of parameters (first column)
        for name in self.dicParameterValues:
            formatted_results.append([name])

        # calculate all forms of correlation requested
        for corr in corrs:
            results, text = self._get_results_text(corr=corr)

            i = 1
            for par, values in results.items():
                coef = F.format_number(number=values[0], deci=decimal)
                p = F.format_number(number=values[1], deci=decimal)
                formatted_results[i].extend([coef, p])
                i += 1

        IO.write_csv(file_name=file_name, rows=formatted_results, delimiter=delimiter)

    def _get_results_text(self, corr):

        if corr == 'r':
            text = "Pearson's correlation coefficients and p-values"
            results = self.get_pearson_corr()
        elif corr == 'rho':
            text = "Spearman's rank correlation coefficients and p-values"
            results = self.get_spearman_corr()
        elif corr == 'p':
            text = "Partial correlation coefficients and p-values"
            results = self.get_partial_corr()
        elif corr == 'pr':
            text = "Partial rank correlation coefficients and p-values"
            results = self.get_partial_rank_corr()
        else:
            raise ValueError('Invalid correlation type is provided.')

        return results, text

    @staticmethod
    def _full_name(corr):

        if corr == 'r':
            return "Pearson's correlation"
        elif corr == 'rho':
            return "Spearman's correlation"
        elif corr == 'p':
            return 'Partial correlation'
        elif corr == 'pr':
            return 'Partial rank correlation'
        else:
            raise ValueError('Invalid correlation type is provided.')

    def _get_corr(self, f):
        """
        f: correlation function
        :returns (list) of [parameter name, correlation coefficients, p-value] """

        result = dict()  # each row [parameter name, correlation, p-value]
        for paramName, paramValues in self.dicParameterValues.items():
            # calculate Spearman rank-order correlation coefficient
            coef, p = f(paramValues, self.outputValues)
            # store [parameter name, Spearman coefficient, and p-value
            result[paramName] = [coef, p]

        return result

    @staticmethod
    def _get_partial_corr(dic_param_values, output_values):

        result = dict()  # each row [parameter name, correlation, p-value]
        for paramName, paramValues in dic_param_values.items():

            z = []
            for otherParamName, otherParamValues in dic_param_values.items():
                if paramName != otherParamName:
                    z.append(otherParamValues)

            # calculate partial correlation coefficient
            coef, p = partial_corr(x=paramValues, y=output_values, z=z)
            # store [parameter name, Spearman coefficient, and p-value
            result[paramName] = [coef, p]

        return result
