import pandas as pd
import numpy as np
import scipy.stats as stats
import matplotlib.pyplot as plt
from matplotlib.ticker import FormatStrFormatter

import bluebelt.analysis.ci as ci
import bluebelt.helpers.boxplot
import bluebelt.helpers.text

import bluebelt.styles


class Correlation():
        
    def __init__(self, frame, columns=None, confidence=0.95, **kwargs):
        
        # check arguments
        if not isinstance(frame, pd.DataFrame):
            raise ValueError('frame is not a Pandas DataFrame')

        if columns is not None and not isinstance(columns, list):
            raise ValueError('columns is not a list')
        
        self.frame = frame if columns is None else frame[columns]
        self.columns = columns if columns is not None else self.frame.columns.to_list()
        self.nrows = self.frame.shape[0]
        self.ncols = self.frame.shape[1]
        self.confidence = confidence
            
        self.calculate()

    def calculate(self):
        self.z = stats.norm.ppf(1-(1-self.confidence)/2)
        self.rci = self._r_ci()

    def __str__(self):
        precision = 3

        result =f'Pearson correlation coefficient with {self.confidence*100:1.1f}% confidence interval\n\n'
        result += self._r_ci().applymap(lambda x: np.round(x,precision) if isinstance(x, (np.floating, int, float, tuple)) else x).to_string()
        
        return result
    
    def __repr__(self):
        return (f'{self.__class__.__name__}(columns={self.columns}, nrows={self.nrows}, ncols={self.ncols}, confidence={self.confidence})')
    
    def result(self, precision=3):
        print(self._r_ci().applymap(lambda x: np.round(x,precision) if isinstance(x, (np.floating, int, float, tuple)) else x).to_string())
    
    def plot(self, **kwargs):
        
        style = kwargs.get('style', bluebelt.styles.paper)
        kwargs.pop('style', None)
        path = kwargs.pop('path', None)

        fig, ax = plt.subplots(nrows=self.ncols, ncols=self.ncols, **kwargs)
        for row, row_name in enumerate(self.columns):
            for col, col_name in enumerate(self.columns):
                if row==col:
                    # plot a histogram where row and column are equal
                    ax[row, col].hist(self.frame[row_name], **style.correlation.histogram)
                else:
                    # plot a scatter plot
                    ax[row, col].plot(self.frame[col_name], self.frame[row_name], **style.correlation.scatter)
                    
                    # fetch r and ci
                    r_value = self.rci.iloc[row, col*2]
                    rho_l = self.rci.iloc[row, 1+col*2][0]
                    rho_u = self.rci.iloc[row, 1+col*2][1]
                    ax[row, col].text(0.01, 0.99, f'r :{r_value:1.2f} CI: ({rho_l:1.2f} - {rho_u:1.2f})', transform=ax[row, col].transAxes, **style.correlation.text)
                    
                    # share X-axis
                    if col==0 and row>1:
                        ax[row, col].sharex(ax[1, col])
                    elif col>0 and row>0:
                        ax[row, col].sharex(ax[0, col])
                        
                    # share Y-axis
                    if row==0 and col>1:
                        ax[row, col].sharey(ax[row, 1])
                    elif row>0 and col>0:
                        ax[row, col].sharey(ax[row, 0])
                                        
                        
                if row==self.ncols-1:
                    ax[row, col].set_xlabel(col_name)
                
                
                if col==0:
                    ax[row, col].set_ylabel(row_name)
                
                
                ax[row, col].set_xticks([])
                ax[row, col].set_yticks([])
                
        plt.subplots_adjust(hspace = 0)
        plt.subplots_adjust(wspace = 0)
        ax[0, 0].set_title(f'correlation plot for {", ".join(self.columns[:-1])+" and "+self.columns[-1]}', **style.correlation.title)
        plt.close()
        return fig

    def _r_ci(self):
        
        cols = []
        for col in range(0, self.ncols):
            rows = []
            for row in range(0, self.ncols):

                # calculate r and ci
                if row != col:
                    r_value = (np.corrcoef(self.frame[self.columns[row]], self.frame[self.columns[col]])[0,1])
                        
                    z_l = 0.5 * (np.log(1+r_value) - np.log(1-r_value)) - self.z / (self.frame[[self.columns[row], self.columns[col]]].dropna(axis=0).shape[0]-3)
                    z_u = 0.5 * (np.log(1+r_value) - np.log(1-r_value)) + self.z / (self.frame[[self.columns[row], self.columns[col]]].dropna(axis=0).shape[0]-3)
                    rho_l = (np.exp(2 * z_l) - 1) / (np.exp(2 * z_l) + 1)
                    rho_u = (np.exp(2 * z_u) - 1) / (np.exp(2 * z_u) + 1)

                    ci = [r_value, (rho_l, rho_u)]
                else:
                    ci = ["", ""]#[1., 1., 1.]

                rows+=ci
            cols.append(rows)

        columns_array = np.array([item for sublist in [[col]*2 for col in self.frame.columns] for item in sublist] + ['r', 'CI'] * self.ncols).reshape(2, self.ncols*2)
        columns = pd.MultiIndex.from_tuples(list(zip(*columns_array)))#, names=["array", "value"])

        return pd.DataFrame(cols, index=self.columns, columns=columns)

    def _ci(self):
        
        cols = []
        for col in range(0, self.ncols):
            rows = []
            for row in range(0, self.ncols):

                # calculate r and ci
                if row != col:
                    r_value = (np.corrcoef(self.frame[self.columns[row]], self.frame[self.columns[col]])[0,1])
                        
                    z_l = 0.5 * (np.log(1+r_value) - np.log(1-r_value)) - self.z / (self.frame[[self.columns[row], self.columns[col]]].dropna(axis=0).shape[0]-3)
                    z_u = 0.5 * (np.log(1+r_value) - np.log(1-r_value)) + self.z / (self.frame[[self.columns[row], self.columns[col]]].dropna(axis=0).shape[0]-3)
                    rho_l = (np.exp(2 * z_l) - 1) / (np.exp(2 * z_l) + 1)
                    rho_u = (np.exp(2 * z_u) - 1) / (np.exp(2 * z_u) + 1)

                    ci = (rho_l, rho_u)
                else:
                    ci = (1.,1.)

                rows.append(ci)
            cols.append(rows)

        return pd.DataFrame(cols, index=self.columns, columns=self.columns)

    
    def _r(self):
        return pd.DataFrame([[(np.corrcoef(self.frame[col], self.frame[row])[0,1]) for col in self.columns] for row in self.columns], index=self.columns, columns=self.columns)
