import math, random
import pandas as pd
import numpy as np
from jax import jit
import jax.numpy as jnp

class iterativeFAMD(object):

    def __init__(self):
        self.explained_var = 0.95
        self.max_iter = 1000
        self.tol = 1e-4
        
    def calculateSVD(self, data):
        U, s, VT = jnp.linalg.svd(data, full_matrices=False)
        return U, s, VT
        
    def encodeFeatures(self, dataframe, features_to_encode):
        # Maintain dataframe for columns that do not need to be encoded
        dropped = dataframe.drop(columns=features_to_encode)
        dropped_cols = dropped.columns.values
        # Create encoded feature dataframe. Take slices of this by column and reindex, joining to the master df
        # This will impute null values for encoded columns where necessary while maintaining encodings for other categorical columns
        feature_df = dataframe[features_to_encode]
        encoded_columns = []
        df_list = [dropped]
        
        for col in features_to_encode:
            sliced = feature_df[col].dropna()
            one_hot = pd.get_dummies(sliced, prefix=col+"_", sparse=True)
            if len(np.unique(sliced.values)) == 1:
                pass
            else:
                one_hot = one_hot.reindex(feature_df[col].index)
            encoded_columns.extend(one_hot.columns.values)
            df_list.append(one_hot)
            one_hot = None
        dropped = pd.concat(df_list, axis=1)
        return dropped, encoded_columns
    
    def initializeValues(self, dataframe):
        proportions = np.zeros(len(self.categorical))
        stdevs = np.zeros(len(self.continuous))

        # Sustitute the column means as the value for continuous nans, then standardize column
        # Substitute the missing values of categorical nans as the proportion of their category
        for index, c in enumerate(self.continuous):
            prop = dataframe[c].mean()
            if np.isnan(dataframe[c].values).any():
                dataframe[c] = np.where(dataframe[c].isnull(), prop, dataframe[c])
            stdev = np.std(dataframe[c])
            if stdev == 0:
                stdev = 0.01
            else:
                dataframe[c] = dataframe[c] / stdev
            stdevs[index] = stdev
        for index, c in enumerate(self.categorical):
            prop = dataframe[c].mean()
            if np.isnan(dataframe[c].values).any():
                dataframe[c] = np.where(dataframe[c].isnull(), prop, dataframe[c])
            proportions[index] = prop
            dataframe[c] = dataframe[c] / np.sqrt(prop)        

        # Standardize by dividing each column by the square root of the column-wise proportion of each category
        variances = [s ** 2 for s in stdevs]
        variances.extend(proportions)
        diag = np.diag(variances)
        return dataframe, diag

    def famdImpute(self, dataframe):
        j = 0
        diag = None
        X = dataframe.copy()
        prev_diff = np.infty
        for i in range(self.max_iter):
            prev_df = X.copy()
            new_df, diag = self.initializeValues(X)
            new_diag = diag.copy()
            for j in range(diag.shape[0]):
                if diag[j][j] != 0:
                    new_diag[j][j] = diag[j][j] ** (-1/2)
            XD = new_df.dot(new_diag)
            #XD = pd.DataFrame(np.dot(new_df, new_diag), index=X.index, columns=X.columns)
            M = np.mean(XD, axis=0)
            means = np.broadcast_to(M, XD.shape)
            resultant = jnp.asarray((XD - means))
            U, s, VT = jit(self.calculateSVD)(resultant)
            n_elements = 0
            explained_var = np.cumsum((s**2)/np.sum(s**2))
            for index, v in enumerate(explained_var):
                if v > self.explained_var:
                    n_elements = index
                    break
            U = U[:, :n_elements]
            sigma = np.diag(s[:n_elements])
            VT = VT[: n_elements]
            lr = np.dot(np.dot(U, sigma), VT)
            if self.noise == "gaussian":
                mu, sig = 0, 0.1 
                noise = np.random.normal(mu, sig, [lr.shape[0], lr.shape[1]])
                lr = lr + noise
            mul = (lr + means).dot(diag)
            
            for index, c in enumerate(X.columns):
                if np.isnan(dataframe[c].values).any():
                    X[c] = np.where(dataframe[c].isnull(), mul[:, index], dataframe[c])
                else:
                    X[c] = dataframe[c]
            diff = np.sqrt(((X.values - prev_df.values)**2).sum())
            old = np.sqrt(((prev_df.values)**2).sum())
            if (diff/old) < self.tol:
                break
        return X
    
    
    def fillNa(self, to_fill, imputed, categorical):
        for col in to_fill.columns.values:
            if col in categorical:
                cols_to_max = [c for c in imputed.columns.values if c.startswith(col)]
                to_fill[col] = imputed[cols_to_max].idxmax(axis=1).str.split("__").apply(lambda x : x[1]).astype(to_fill[col].dtype)
            else:
                to_fill[col] = imputed[col]
        return to_fill
    
    def filterColumns(self, features_to_encode, exclude):
        new_features = []
        drop_cols = [e for e in exclude]
        date_col_substrings = ["dt", "date", "_date"]
        for col in features_to_encode:
            if any(substring in col.lower() for substring in date_col_substrings) or col in exclude:
                drop_cols.append(col)
                continue
            else:
                new_features.append(col)
        if len(drop_cols) == 0:
            drop_cols = ["index"]
        else:
            drop_cols.append("index")
        return new_features, drop_cols
            
    def impute(self, impute_df, dataframe, features_to_encode):
        impute_df = self.famdImpute(impute_df)
        dataframe = self.fillNa(dataframe, impute_df, features_to_encode)
        return dataframe
