import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

from factor_analyzer import FactorAnalyzer, Rotator

class ExploratoryFactorAnalyzer:
    """An Exploratory Factor Analyzer class which has the following methods:

    Parameters:
    -----------

    data: DataFrame
        The data that will be analyzed.

    Attributes:
    -----------

    factorAnalyzer: factor_analyzer.FactorAnalyzer object

    rotatedLoadings: numpy ndarray
        The rotated loadings matrix.

    communalities: numpy ndarray
        The communalities of each item.

    loadingsTable: pandas DataFrame
        The factor loadings with item and communalities. Any factor loading
        less than 0.3 is suppressed.

    factorItems: dict
        The list of items for each factor.

    factorLoadings: dict
        The list of loadings for each fator.

    factorScores: pandas DataFrame
        Generated factor scores by dot product of factor items with loadings.

    averageWithinFactorCorr: float
        Average of average Pearson R correlation coefficient for each factor.

    averageAcrossFactorCorr: float
        Average Pearson R correlation coefficient among factors

    cronbachsAlpha: pandas DataFrame
        The Cronbach's Alpha score for each factor.

    factorResults: dict
        The final version of each factor where the factor items are named.
    """

    def __init__(self, data):
        self.data = data

        # attributes are set to None

        self.factorAnalyzer = None
        self.rotatedLoadings = None
        self.communalities = None
        self.loadingsTable = None
        self.factorItems = None
        self.factorLoadings = None
        self.factorScores = None
        self.averageWithinFactorCorr = None
        self.averageAcrossFactorCorr = None
        self.cronbachsAlpha = None
        self.factorResults = None
        self.explainedVariance = None

    def _arg_checker(self):
        """Checks if the input is a pandas DataFrame"""
        if not isinstance(self.data, pd.DataFrame):
            raise ValueError('The data has to be a pandas DataFrame')

    def eliminatePairwiseCorrelation(self, threshold, getItems=False):
        """Eliminates a variable of a pair which are highly correlated (above
        the given threshold) based on Pearson R correlation coefficient.

        Parameters:
        -----------

        threshold: float
            A cut-off value between 0 and 1. One item from each pair of items
            that have at least threshold Pearson R correlation coefficient
            will be dropped.

        getItems: bool
            If True, it only returns the list of items where each item is
            picked arbitrarily from one of the pair of items that have at
            least threshold Pearson R correlation coefficient.
        """

        # a list to store items that have correlation coefficient above the
        # given threshold
        itemsDropped = []

        while True:

            # generates lower triangle of a correlation matrix
            # (excluding the diagonal - 1s)
            corrMatrix = self.data.drop(itemsDropped, axis=1).corr().abs()
            mask = np.tril(np.ones(corrMatrix.shape), k=-1).astype(np.bool)
            corrMatrix = corrMatrix.where(mask)

            # a list to store how many pair of variables a single variable has
            # that have correlation coefficient above the threshold
            t = []

            for col in corrMatrix.columns:
                columnSum = corrMatrix[col] > threshold
                columnSum = columnSum.astype(int).sum()
                t.append((columnSum, col))

            # sorts the list using the number of pairs as the key - descending
            t.sort(key=lambda x: x[0], reverse=True)

            # keep running th loop until there is no pair of variables that
            # have correlation coefficient above the given threshold
            if t[0][0] > 0:
                itemsDropped.append(t[0][1])
            else:
                break

        if getItems:
            return itemsDropped
        else:
            self.data = self.data.drop(itemsDropped, axis=1)
            return self

    def doFactorAnalysis(self,
                         n_factors,
                         method='principal',
                         rotation='varimax',
                         power=4,
                         kaiser_on=True):
        """Performs EFA based on the given method, rotation and power values.

        Parameters:
        -----------

        n_factors: int
            The number of factors to be retained.

        method : {'minres', 'ml', 'principal'}, optional
            The method for factor extraction. Defaults to 'principal'.

        rotation: {'varimax', 'promax'}
            The rotation method. Defaults to 'varimax'.

        power: int
            The power value for 'promax' rotation. It is recommend to be
            between 1 and 4.

        kaiser_on: bool
            If True, applies Kaiser normalization.
        """

        # initiate a FactorAnalyzer object and do the factor analysis without
        # any rotation. It uses principal factor extraction method.
        fa = FactorAnalyzer(n_factors=n_factors, method=method, rotation=None)
        fa.fit(self.data)

        loadings = fa.loadings_
        communalities = fa.get_communalities()

        # generate a Rotator object where it applies promax rotation
        # (power = 3) with kaiser normalization.

        rotator = Rotator(method=rotation, normalize=kaiser_on, power=power)
        rotatedLoadings = rotator.fit_transform(loadings)

        self.rotatedLoadings = rotatedLoadings
        self.communalities = communalities
        self.factorAnalyzer = fa
        self.explainedVariance = max(fa.get_factor_variance()[2])

        return self

    def generateLoadingsTable(self, maskLoadings=True):
        """Puts the loadings into a data frame alongside with communalities.

        Parameters:
        -----------

        maskLoadings: bool
            If True, suppress all loadings < 0.3.
        """

        itemEncodings = self.data.columns

        # generate a data frame with the rotated loadings.
        dataFrameLoadings = pd.DataFrame(self.rotatedLoadings)

        # infer the number of factors
        columnNames = ['Elements']

        n_factors = self.rotatedLoadings.shape[1]
        factorColumns = ['Factor {}'.format(i + 1) for i in range(n_factors)]
        columnNames.extend(factorColumns)

        # generate a column for element names

        dataFrameLoadings.reset_index(inplace=True)
        dataFrameLoadings.columns = columnNames

        # replace the index column with element encodings
        dataFrameLoadings.Elements = itemEncodings

        dataFrameLoadings['Communalities'] = self.communalities

        # generate max column for sorting the factor loadings from factor 1 to
        # factor n in descending

        dataFrameLoadings['max'] = dataFrameLoadings.filter(
            regex='(Factor)').apply(lambda x: max(abs(x)), axis=1)

        idx = []

        for col in dataFrameLoadings.drop(['Elements', 'Communalities'],
                                          axis=1).columns:
            if col != 'max':
                cond = dataFrameLoadings[col].apply(
                    lambda x: abs(x)) == dataFrameLoadings['max']
                t = list(dataFrameLoadings.loc[cond].sort_values('max',
                                                                 ascending=False).index)
                idx.extend(t)

        dataFrameLoadings = dataFrameLoadings.iloc[idx]

        dataFrameLoadings.drop('max', axis=1, inplace=True)

        # remove all loadings which has absolute value less than 0.3
        if maskLoadings:
            dataFrameLoadings.iloc[:, 1:-1] = dataFrameLoadings.iloc[:, 1:-1].apply(
                lambda x: np.where(
                    abs(x) < 0.3, '', round(x, 3)))

        dataFrameLoadings.reset_index(drop=True, inplace=True)

        # returns the loadings table
        self.loadingsTable = dataFrameLoadings

        return self

    def dropItem(self, items):
        """Drops given items from the data. Items that have crossloadings
        should be dropped from the data one at a time based on the magnitude
        of the loadings. Choose the lower magnitude loadings to drop.

        All items with crossloadings less than 75% ratio can be retained.

        Parameters:
        -----------

        items: str or list
            Item(s) to be dropped
        """

        self.data = self.data.drop(items, axis=1)

        return self

    def calculateFactorScores(self):
        """Calculates factor scores for each factor by dot product the items
        with the factor loadings. It is the weighted average of factor items
        where the factor loadings are the weights.
        """

        # generate a data frame which consist of only the factors -
        # the numeric data
        temp = self.loadingsTable.filter(regex='(Factor)')

        # for each factor replace the empty string value with 0
        for col in temp.columns:
            temp[col] = np.where(temp[col] != '', temp[col], 0).astype(float)

        # get the row wise max of the absoulute value of loadings
        temp['max'] = temp.apply(lambda x: max(abs(x)), axis=1)

        # add items to the data frame
        temp['Elements'] = self.loadingsTable.Elements

        # a dictionary to hold the items and the loadings for each factor
        loadings = {}

        # iterates over the each factor column and select the factor loadings
        # which are equal to the values on the max column.
        # Save those later in loadings dictionary.
        for col in temp.drop(['Elements', 'max'], axis=1).columns:
            temp.sort_values(col, ascending=False)
            cond = temp[col].apply(lambda x: abs(x)) == temp[max]
            d = {}
            for row in temp.loc[cond, ['Elements', col]].iterrows():
                d[row[1][0]] = row[1][1]
            loadings[col] = d

        # generates an empty data frame for factor scores
        factorScores = pd.DataFrame()

        # populates the empty data frame with the dot product of items and
        # the factor loadings
        for k, v in loadings.items():
            factorScores[k] = self.data.filter(items=list(v.keys())).dot(
                list(v.values())).div(sum(list(v.values())))

        self.factorScores = factorScores
        self.factorItems = {k: list(v.keys()) for k, v in loadings.items()}
        self.factorLoadings = {
            k: list(v.values())
            for k, v in loadings.items()
        }

        return self

    def averageWithinFactorCorrelations(self):
        """Calculates the average of the average Pearson R correlation
        coefficient among the items of each factor."""

        # dictionary to store the factors and their average correlations
        d_correlations = {}

        # for each factor generates a lower triangle correlation matrix
        # without a diagonal and calculates the mean correlation coefficient
        for factor, items in self.factorItems.items():
            corrMatrix = self.data.loc[:, items].corr()
            mask = np.tril(np.ones(corrMatrix.shape), k=-1).astype(np.bool)
            corrMatrix = corrMatrix.where(mask)
            corrMean = np.nanmean(corrMatrix.values.flatten())
            d_correlations[factor] = corrMean

        averageCorrMean = np.mean(list(d_correlations.values()))

        self.averageWithinFactorCorr = averageCorrMean

        return self

    def averageAcrossFactorsCorrelations(self):
        """Calculates the average Pearson R correlatin coefficient among the
        factors."""

        # generate correlation matrix of the factors
        corrMatrix = self.factorScores.corr()
        mask = np.tril(np.ones(corrMatrix.shape), k=-1).astype(np.bool)
        corrMatrix = corrMatrix.where(mask)
        # calculate the average without counting the nan values
        corrMean = np.nanmean(corrMatrix.values.flatten())

        self.averageAcrossFactorCorr = corrMean

        return self

    def calculateCronbachsAlpha(self):
        """Calculates Cronbach's Alpha score for each factor."""

        # dict to store factor and their cronbach's alpha scores
        d_alphaScores = {}

        # iterate over each factor and calculate the cronbach's alpha score
        # and store the value in the d_alphaScores dictionary.
        for factor, items in self.factorItems.items():

            dataframe = self.data.filter(items=items, axis=1)

            df_corr = dataframe.corr()

            # Calculate N
            # The number of variables equals the number of columns in the df
            N = dataframe.shape[1]

            # Calculate R
            # For this, we'll loop through the columns and append every
            # relevant correlation to an array calles "r_s". Then, we'll
            # calculate the mean of "r_s"
            rs = np.array([])
            for i, col in enumerate(df_corr.columns):
                sum_ = df_corr[col][i + 1:].values
                rs = np.append(sum_, rs)
            mean_r = np.mean(rs)

            # Use the formula to calculate Cronbach's Alpha
            cronbach_alpha = (N * mean_r) / (1 + (N - 1) * mean_r)

            d_alphaScores[factor] = cronbach_alpha

        temp = pd.DataFrame(d_alphaScores, index=np.arange(1))
        temp = temp.T
        temp.columns = ["Cronbach's Alpha"]

        self.cronbachsAlpha = temp

        return self

    def getElementNames(self, elementNames=None):
        """Map the element encodings with element names from a given
        encoding-name dictionary.

        Parameters:
        -----------

        elementNames: dict
            A container of key (element encoding) - value (element name)
            pairs.
        """

        d_names = {}

        for factor, items in self.factorItems.items():
            elements = [elementNames[item] for item in items]
            d_names[factor] = elements

        maxLenghtItems = max([len(v) for k, v in d_names.items()]) - 1

        for factor, items in d_names.items():
            itemLength = len(items)
            while itemLength <= maxLenghtItems:
                items.append('')
                itemLength += 1

        temp = pd.DataFrame(d_names)

        self.factorResults = temp

        return self

    def screePlot(self, title='Scree Plot'):
        """Scree plot of the data."""

        faScree = FactorAnalyzer(n_factors=self.data.shape[1], rotation=None)
        faScree.fit(self.data)
        eigenvalues, _ = faScree.get_eigenvalues()

        idx = np.arange(1, self.data.shape[1] + 1)

        fig, ax = plt.subplots(figsize=(10, 8))

        ax.scatter(idx, eigenvalues)
        ax.plot(idx, eigenvalues)
        ax.axhline(1, color='tab:red', linestyle='--')
        ax.set_ylabel('Eigenvalues')
        ax.set_xlabel('Factors')
        ax.set_title(title.capitalize())
        ax.grid()

        return fig