import matplotlib.pyplot as plt
import numpy as np
import scipy.stats as stats
import seaborn as sns
import statsmodels.api as sm
from statsmodels.stats.power import TTestPower, TTestIndPower

class Utilities:
    @staticmethod
    def t_to_d(t, n):
        return t / (n ** 0.5)

    @staticmethod
    def calc_required_sample_size(d, power, alpha=0.05, tail='two-sided'):
        analysis = TTestPower()
        return analysis.solve_power(effect_size=d, alpha=alpha, nobs=None, power=power, alternative=tail)

    @staticmethod
    def calc_post_hoc_power(n, d, alpha=0.05, tail='two-sided'):
        analysis = TTestPower()
        return analysis.solve_power(effect_size=d, alpha=alpha, nobs=n, power=None, alternative=tail)

    @staticmethod
    def dereferencify(x):
        if 'C(' in x:
            return f'{x.split("C(")[1].split(",")[0]}[{x.split("[")[1].split("]")[0]}]'
        else:
            return x

    @staticmethod
    def represents_int(s):
        try:
            int(s)
        except ValueError:
            return False
        else:
            return True

    @staticmethod
    def first_matching(iterable, condition=lambda x: True):
        return next(x for x in iterable if condition(x))

class StatsHandler:
    def __init__(self, data=None):
        self._test_objective = "independent"
        self._tail = "two-sided"
        self.reg_results = None
        self._alpha = 0.05
        self._permutation_count = 10000
        self._test_type = "rank"
        self._data = None
        self._query_column = None
        self._query_on_column = False
        self._identifier2 = None
        self._identifier1 = None
        self._measuring_column = None
        self._english_measuring = None
        self._is_regression_spec = False
        self._regression_lhs = None
        self._regression_rhs = []
        self._regression_rhs_names = []
        self._override_cohens_d = None
        self._override_n = None
        self._override_t = None
        self._override_n2 = None
        self._override_power = None
        self._sigfig = 3

        if data is not None:
            self.set_data(data)

    def set_data(self, data):
        self._data = data
        return self

    def set_p(self, value):
        self._alpha = value
        return self

    def is_paired(self):
        self._test_objective = "paired"
        self._data['Difference'] = self.data2 - self.data1
        return self

    def using_rank(self):
        self._test_type = "rank"
        return self

    def measuring(self, english_name):
        self._english_measuring = english_name
        return self

    def using_permutation(self, permutation_count=10000):
        self._test_type = "permutation"
        self._permutation_count = permutation_count
        return self

    def _get_data(self, identifier):
        if self._query_on_column:
            data_here = self._data.query(self._query_column + ' == "' + identifier + '"')

            if self._measuring_column is not None:
                return data_here[self._measuring_column]
            else:
                return data_here
        else:
            return self._data[identifier]

    @property
    def sse(self):
        if not self.reg_results:
            return None

        return sum([i ** 2 for i in self.reg_results.resid])

    @property
    def intercept(self):
        x_bar = self.data1.mean()
        y_bar = self.data2.mean()

        return y_bar - self.slope * x_bar

    @property
    def slope(self):
        s_x = self.data1.std()
        s_xy = self.data1.cov(self.data2)

        return s_xy / (s_x ** 2)

    @property
    def equation(self):
        return f'{self.round(self.intercept)} = ȳ - {self.round(self.slope)}x̄'

    @property
    def safe_value_map(self):
        return np.array([not np.isnan(self.data1[i]) and not np.isnan(self.data2[i]) for i in range(len(self.data1))])

    def regress_to(self, column):
        self._is_regression_spec = True
        self._regression_lhs = column
        self._regression_rhs = []
        self._regression_rhs_names = []

        return self

    def add_regression_component(self, column, reference=None):
        self._regression_rhs_names.append(column)

        if reference is not None:
            self._regression_rhs.append(f'C({column}, Treatment(reference="{reference}"))')
        else:
            self._regression_rhs.append(column)

        return self

    def add_interaction_term(self, column1, column2, add_main_effects = True):
        self._regression_rhs.append(f'{column1}{"*" if add_main_effects else ":"}{column2}')
        self._regression_rhs_names.append(None)
        return self

    def apply_regression_formula(self, formula):
        reg_formula = sm.regression.linear_model.OLS.from_formula(data=self._data, formula=formula)
        reg_results = reg_formula.fit()
        self._data['resid'] = reg_results.resid
        self._data['fitted'] = reg_results.fittedvalues
        self.reg_results = reg_results

        return reg_results.summary()

    @property
    def regression_formula_stub(self):
        return f'{self._regression_lhs} ~ B0 + {" ".join([f"B{i + 1}*{e}" for i, e in enumerate(self._regression_rhs_names)])}'

    def predict_using_regression(self, values, allow_missing_values = False):
        data_source = self.reg_results.summary().tables[1].data[1:]
        columns = [[Utilities.dereferencify(x[0].strip()), float(x[1].strip())] for x in data_source if x[0].strip() != '' and '*' not in x[0] and ':' not in x[0] and 'Intercept' not in x[0]]
        column_lookup = { k[0]: k[1] for k in columns }

        intercept = float(data_source[0][1].strip())

        for entry in values:
            if entry not in self._regression_rhs_names:
                raise ValueError(f'Invalid/unknown column specified: {entry}')

        if not allow_missing_values:
            for entry in self._regression_rhs_names:
                if entry not in values:
                    raise ValueError(f'Missing column: {entry}')

        print('Using equation of form:')
        print(self.regression_formula_stub)

        equation_components = []
        for entry in values:
            if Utilities.represents_int(values[entry]):
                equation_components.append([column_lookup[entry], values[entry], f'coef({entry})'])
            else:
                entry_name = f"{entry}[T.{values[entry]}]"
                if entry_name in column_lookup:
                    equation_components.append([column_lookup[entry_name], 1, f'coef({entry_name})'])
                else:
                    pulling_instead_on = None
                    for key in column_lookup:
                        if key.startswith(f'{entry}[T.'):
                            pulling_instead_on = key
                            break

                    if pulling_instead_on is None:
                        raise ValueError(f'Unresolved value for {entry}: {values[entry]}')

                    equation_components.append([column_lookup[pulling_instead_on], 0, f'coef({pulling_instead_on})'])

        print('Using:')
        print(f"{self._regression_lhs} = intercept + {' + '.join([f'{x[2]}*{x[1]}' for x in equation_components])}")
        print(f"   = {intercept} + {' + '.join([f'{x[0]}*{x[1]}' for x in equation_components])}")

        res = intercept
        for x in equation_components:
            res += x[0] * x[1]
        print(f'\n= {self.round(res)}')

        return res

    def do_regression(self):
        na_map = self.safe_value_map

        x = np.array(self.data1)[na_map].tolist()
        y = np.array(self.data2)[na_map].tolist()

        X = sm.add_constant(x)
        reg_formula = sm.OLS(y, X)
        reg_results = reg_formula.fit()

        fitted_col = []
        resid_col = []

        i = 0
        for j in na_map:
            if not j:
                fitted_col.append(reg_results.fittedvalues[i])
                resid_col.append(reg_results.resid[i])
                i += 1
            else:
                fitted_col.append(np.nan)
                resid_col.append(np.nan)

        self.reg_results = reg_results
        self._data['resid'] = resid_col
        self._data['fitted'] = fitted_col

        return reg_results.summary()

    def count(self):
        return len(self.data1.dropna())

    def set_min(self, min_value):
        return self.allow_datapoints_if(lambda d: d >= min_value)

    def set_max(self, max_value):
        return self.allow_datapoints_if(lambda d: d <= max_value)
        pass

    def ban_value(self, values):
        self._data.where(values, inplace=True)
        return self

    def allow_datapoints_if(self, filter_func):
        if self._query_on_column:
            self._data.where((self._data[self._query_column] != self._identifier1 or filter_func(
                self._data[self._measuring_column])), inplace=True)
            self._data.where((self._data[self._query_column] != self._identifier2 or filter_func(
                self._data[self._measuring_column])), inplace=True)
        else:
            self._data.where(filter_func(self._data[self._identifier1]), inplace=True)
            self._data.where(filter_func(self._data[self._identifier2]), inplace=True)

        return self

    def ban_if_column(self, column, filter_func):
        mask = self._data[column].apply(filter_func)
        self._data.loc[mask, :] = np.nan
        return self

    def ban_if_column_value(self, column, banned_values):
        self.ban_if_column(column, lambda b: b in banned_values)
        return self

    def scatterplot(self):
        return sns.scatterplot(x=self.data1, y=self.data2)

    def plot_variance(self):
        print(self.reg_results.fittedvalues.shape)
        print(self.reg_results.resid.shape)

        ax = sns.scatterplot(
            x=self.reg_results.fittedvalues,
            y=self.reg_results.resid,
        )
        ax.set(xlabel="y-hat", ylabel="Residual")

    def plot_residuals(self):
        return sns.histplot(self.reg_results.resid)

    def plot(self):
        if self.reg_results:
            self.scatterplot()
            sns.lineplot(x=self.data1[self.safe_value_map], y=list(self.reg_results.fittedvalues))
        elif self._test_objective == "paired":
            plt.figure(figsize=(8, 4))

            plt.subplot(1, 2, 1)
            sns.scatterplot(data=self._data, x=self._identifier1, y=self._identifier2)
            plt.plot([self.data1.min(), self.data1.max()], [self.data2.min(), self.data2.max()], 'r--')

            plt.subplot(1, 2, 2)
            sns.barplot(data=self._data[[self._identifier1, self._identifier2]], errorbar=None, color=[0.9, 0.9, 0.9])
            sns.lineplot(data=self._data[[self._identifier1, self._identifier2]].T, legend=False, marker='o')
        else:
            sns.kdeplot(data=self._data, x=self._measuring_column, hue=self._query_column, fill=True)
            sns.rugplot(data=self._data, x=self._measuring_column, hue=self._query_column, height=0.1)

            plt.xlabel(self.measuring_var, fontsize=12)
            plt.ylabel("Density", fontsize=12)

        plt.show()

    @property
    def n(self):
        if self._override_n:
            return self._override_n
        elif self.data1:
            return len(self.data1)
        else:
            analysis = TTestPower()
            return analysis.solve_power(effect_size=self.cohens_d, alpha=self._alpha, nobs=None, power=self.power,
                                        alternative=self._tail)

    @n.setter
    def n(self, value):
        self._override_n = value

    @property
    def n2(self):
        if self._override_n2:
            return self._override_n2
        else:
            return len(self.data2)

    @n2.setter
    def n2(self, value):
        self._override_n2 = value

    @property
    def data1(self):
        if self._identifier1:
            return self._get_data(self._identifier1)
        else:
            return None

    @property
    def data2(self):
        if self._identifier1:
            return self._get_data(self._identifier2)
        else:
            return None

    @property
    def hypothesis_summary_str(self):
        if self._tail == "two-sided":
            return "This is a two-sided (non-directional) alternative hypothesis."
        else:
            return "This is a one-sided (directional) alternative hypothesis."

    @property
    def measuring_var(self):
        return self._english_measuring if self._english_measuring is not None else self._measuring_column

    @property
    def average_type(self):
        return "median" if self._test_type == "rank" else "mean"

    @property
    def null_hypothesis(self):
        if self._test_objective == "correlation":
            return f"There is no {self.test_statistic} (the correlation, Pearson's r = 0)"
        else:
            return f"The {self.test_statistic} is zero."

    @property
    def alt_hypothesis(self):
        if self._test_objective == "correlation":
            if self._tail == "greater":
                return f"Those with higher {self._identifier1} have higher {self._identifier2} (the correlation, Pearson's r > 0)"
            elif self._tail == "less":
                return f"Those with higher {self._identifier1} have lower {self._identifier2} (the correlation, Pearson's r < 0)"
            else:
                return f"There is a correlation between {self._identifier1} and {self._identifier2} (the correlation, Pearson's r ≠ 0)"

        return f"The {self.test_statistic} is {self._tail} than zero."

    @property
    def test_statistic(self):
        if self._test_objective == "correlation":
            return f"correlation between {self._identifier1} and {self._identifier2}"
        if self._test_objective == "independent":
            return f"difference in {self.average_type} {self.measuring_var} between group {self._identifier1} and group {self._identifier2}"
        if self._test_objective == "paired":
            return f"{self.average_type} difference in {self.measuring_var} between group {self._identifier1} and group {self._identifier2}"

    def get_conclusion(self, pvalue):
        if pvalue < self._alpha:
            return (f"{self.round(pvalue)} < {self.round(self._alpha)}.\n"
                    f"There is sufficient evidence to reject the null hypothesis.")
        else:
            return (f"{self.round(pvalue)} > {self.round(self._alpha)}.\n"
                    f"There is insufficient evidence to reject the null hypothesis.")

    def get_hypothesis_str(self, results, extra_info="", statistic_is_arr=False, use_df=False):
        statistic = results.statistic[0] if statistic_is_arr else results.statistic
        pvalue = results.pvalue[0] if statistic_is_arr else results.pvalue

        self.p = pvalue

        df_str = ''
        if use_df:
            df = results.df[0] if statistic_is_arr else results.df
            df_str = f'Degrees of freedom: {df}\n'

        return f"""H0: {self.null_hypothesis}
H1: {self.alt_hypothesis}
{self.hypothesis_summary_str}
{df_str}
Test statistic: {self.test_statistic} ({self.round(statistic)})
Alpha: {self._alpha}

P-value: {self.round(pvalue) if pvalue >= 0.001 else '<0.001'} {f"({extra_info})" if extra_info else ""}

{self.get_conclusion(pvalue)}
"""

    def _do_test_is_normal(self, column_now):
        identifier = self._identifier1 if column_now == 1 else self._identifier2
        result = stats.shapiro(self.data1 if column_now == 1 else self.data2)

        group_name = f'{self.measuring_var} of group {identifier}'
        if column_now == -1:
            result = stats.shapiro(self.data_diff)
            group_name = f'difference between the {self.measuring_var} of group {self._identifier1} and {self._identifier2}'

        return f"""H0: The {group_name} is normally distributed.
H1: The {group_name} is not normally distributed.

{self.best_fitting_normal(column_now)}

Test statistic: Shapiro-Wilk W ({self.round(result.statistic)})
Alpha: {self._alpha}

P-value: {self.round(result.pvalue)} (using Shapiro-Wilk Test)

{self.get_conclusion(result.pvalue)}\n\n"""

    def test_is_normal(self, column=1):
        if self._test_objective == "paired":
            return self._do_test_is_normal(-1)

        columns = [column] if column != 0 else [1, 2]
        res = ''
        for column_now in columns:
            res += self._do_test_is_normal(column_now)

        return res.strip()

    def rounding_to(self, n):
        self._sigfig = n

    def round(self, value):
        return '{:g}'.format(float('{:.{p}g}'.format(value, p=self._sigfig)))

    def exec(self):
        if self._is_regression_spec:
            return self.apply_regression_formula(f'{self._regression_lhs} ~ {" + ".join(self._regression_rhs)}')
        elif self._test_type == "permutation":
            if self._test_objective == "correlation":
                def correlate(x, y):
                    tmp = np.corrcoef(x, y)
                    c = tmp[0][1]
                    return c

                function = correlate
                permutation_type = "pairings"
            elif self._test_objective == "independent":
                def dMeans(x, y):
                    return np.mean(x) - np.mean(y)

                function = dMeans
                permutation_type = "independent"
            elif self._test_objective == "paired":
                def mDiff(x, y):
                    return np.mean(x - y)

                function = mDiff
                permutation_type = "samples"
            else:
                raise ValueError('Unsupported test type for permutation test: ' + self._test_objective)

            results = stats.permutation_test((self.data1, self.data2), function, permutation_type=permutation_type,
                                             alternative=self._tail, n_resamples=self._permutation_count)
            return self.get_hypothesis_str(results, extra_info=f"after {self._permutation_count} permutations")
        elif self._test_type == "rank":
            if self._test_objective == "correlation":
                test_type = 'Spearman'
                results = stats.spearmanr(self.data1, self.data2, alternative=self._tail)
            elif self._test_objective == "independent":
                test_type = 'Mann Whiney U'
                results = stats.mannwhitneyu(self.data1, self.data2, alternative=self._tail)
            elif self._test_objective == "paired":
                test_type = 'Wilcoxon'
                results = stats.wilcoxon(self.data1, self.data2, alternative=self._tail)
            else:
                raise ValueError('Unsupported test type for permutation test: ' + self._test_objective)

            return self.get_hypothesis_str(results, extra_info=f"using {test_type} test")
        elif self._test_type == "t-test":
            if self._test_objective == "independent":
                results = stats.ttest_ind(self.data1, self.data2, alternative=self._tail)
            elif self._test_objective == "paired":
                results = stats.ttest_rel(self.data1, self.data2, alternative=self._tail)
            elif self._test_objective == "correlation":
                results = stats.pearsonr(self.data1, self.data2, alternative=self._tail)
            else:
                raise ValueError('Unsupported test type for t-test: ' + self._test_objective)

            self._override_t = results.statistic

            return self.get_hypothesis_str(results, extra_info=f"using t-test", use_df=True)

    @property
    def t(self):
        if self._override_t:
            return self._override_t
        elif self.data1 and self.data2:
            self.exec()
            return self._override_t

    @t.setter
    def t(self, t):
        self._override_t = t

    def using_correlation(self):
        self._test_objective = "correlation"
        return self

    def using_t_test(self):
        self._test_type = "t-test"
        return self

    def t_test_on_one(self, value):
        results = stats.ttest_1samp(self.data1, value, alternative=self._tail)
        return self.get_hypothesis_str(results, extra_info=f"using t-test", statistic_is_arr=True)

    def query_on_column(self, column):
        self._query_on_column = True
        self._query_column = column

        return self

    def datapoint(self, identifier1):
        self._identifier1 = identifier1
        return self

    def is_greater_than(self, identifier2=None):
        self._identifier2 = identifier2
        self._tail = "greater"
        return self

    def is_less_than(self, identifier2=None):
        self._identifier2 = identifier2
        self._tail = "less"
        return self

    def measuring_column(self, column):
        self._measuring_column = column
        return self

    @property
    def data_diff(self):
        if self._test_objective == "paired":
            return self._data["Difference"]

        return None

    def best_fitting_normal(self, column=1):
        if self._test_objective == "paired" and column == 0:
            return f'Best fitting normal for differences: μ={self.round(self.data_diff.mean())}, σ={self.round(self.data_diff.std())}\n'

        res = ''
        if column != 2:
            res += f'Best fitting normal ({self._identifier1}): μ={self.round(self.data1.mean())}, σ={self.round(self.data1.std())}\n'

        if self.data2 is not None and column != 1:
            res += f'Best fitting normal ({self._identifier2}): μ={self.round(self.data2.mean())}, σ={self.round(self.data2.std())}'

        return res.strip()

    def set_datapoints(self, identifier1, identifier2):
        self._identifier1 = identifier1
        self._identifier2 = identifier2
        self._tail = "two-sided"

        return self

    @property
    def cohens_d(self):
        if self._override_cohens_d:
            return self._override_cohens_d
        elif self.data1 and self.data2:
            if self._test_objective == 'correlation':
                r = stats.pearsonr(self.data1, self.data2, alternative=self._tail).statistic
                t = (r * (len(self.data1) - 2) ** 0.5) / ((1 - r ** 2) ** 0.5)
                d = t / (len(self.data1) ** 0.5)

                return d
            else:
                xP = self.data1.mean()
                xG = self.data2.mean()

                sP = self.data1.std()
                sG = self.data2.std()

                nP = self.data1.count()
                nG = self.data2.count()

                s = (((nP - 1) * (sP ** 2) + (nG - 1) * (sG ** 2)) / (nP + nG - 2)) ** 0.5
                return (xG - xP) / s
        elif self.t and self.n:
            return Utilities.t_to_d(self.t, self.n)

    @cohens_d.setter
    def cohens_d(self, x):
        self._override_cohens_d = x

    @property
    def power(self):
        if self._override_power:
            return self._override_power

        if self._test_objective == 'correlation':
            analysis = TTestPower()
            return analysis.solve_power(effect_size=self.cohens_d, alpha=self._alpha, nobs=self.n, power=None,
                                        alternative=self._tail)
        if self._test_objective == 'paired':
            analysis = TTestPower()
            return analysis.solve_power(effect_size=self.cohens_d, alpha=self._alpha, nobs=self.n, power=None,
                                        alternative=self._tail)
        if self._test_objective == 'independent':
            analysis = TTestIndPower()
            return analysis.solve_power(effect_size=self.cohens_d, alpha=self._alpha, nobs1=self.n,
                                        ratio=self.n2 / self.n, power=None, alternative=self._tail)

    @power.setter
    def power(self, power):
        self._override_power = power
