#################################################################
#
# The R matrix on INPUT is expected to contain the covariates as well.
# The R matrix is usually wither the original RAW R data OR the df_new_R
# data object form running this Class a prior iteration. The best choice is TBD
# Then the final imputted R wil contain the same covariates
# The format of the NN matrices from the preceding Clustering class is
# CELL NN1 NN2 .... NNk+1
# The values can either be the indexing OR the distances depending on the file you are reading 
# Feb 15. Change OLS to a weighted LS.
# NOTE: WLS yield BLUE if weights are inverse variances

# CHECK if the reshape(-1,numGenes is correct)

# TODO read batchfilename as an option for covariate correction

# Might consider usingh KNeighborsRegressor directly !
import time as tm
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import math 
import scipy 
from scipy import stats
import sys
import sklearn.preprocessing
from scipy import stats
import scipy.stats as stats
from scipy import linalg
from sklearn.linear_model import LinearRegression
from sklearn.externals.joblib import Parallel, delayed

#########################################################################
# Keep the known covariates of input R

def getKnownGroupEffects(fulldata,batchfilename):
    """ Check the input file for KNOWN possible group effects
    Exclude those that may not exist such as TREATMENT. Make no
    assumptions about order or sorting. Return list in arbitrary 
    order. Independent of list provided by the user
    """
    foundColumns = list()
    listOneHotEncoded = list()
    if (batchfilename == False):
        print('Using prechosen batch effect names')
        currentPossibleColumnEffects = ['TREATMENT','MOUSE','SEX','DATE','NODEID']
        allColumns = fulldata.columns.values
        listOneHotEncoded = ['MOUSE'] # By default for use with our data testing
        # Grab only those fulldata columns that exist inb currentPossible
        foundColumns = list(set(allColumns) & set(currentPossibleColumnEffects))
    else:
        print('Generating effects list from batchfile')
        foundColumnsCombined = list()
        with open(batchfilename) as f:
            tempColumns = f.read().splitlines()
            for effect in tempColumns:
                foundColumnsCombined.append(effect.upper())
            for ipair in foundColumnsCombined:
                #print('effect is '+ipair)
                name,typeEnc = ipair.strip().split(' ')
                foundColumns.append(name)
                if (typeEnc=='ONEHOT'):
                    listOneHotEncoded.append(name)
    print('Number of Possible batch effects found/specified is '+str(len(foundColumns)))
    print('Number of groupeffect columns in data is '+str(foundColumns))
    print('They have been converted to UPPER CASE: '+str(foundColumns))
    return (foundColumns,listOneHotEncoded)

def buildCellIndexMap(NNindex):
    """Build a dataframe with integer IDs as index and cell names as values
    One would prefer to have grabbed the IDs from the indices file but NOT!
    """
    cellNames=NNindex.index.values
    #cellIDs=NNindex['CELLTARGET'].values
    # Because of NN behavior this previous step doesn't always work
    # SO just give them integer values assuming input order is retained
    dataMap = pd.DataFrame(cellNames)
    ##dataMap.index = cellIDs
    dataMap.columns=['CELL']
    return (dataMap)

########################################################
class regressDataClass(object):
    """Using Regression,predict EVERY cell read from the input (unscaled) reads
    using the top k nearest neighbors are identified in the input INDICES object.
    For the regression we INCLUDE the target cell meaning that we actually regress k+1 terms
    On return is an unscaled new R matrix (nxp) and includes all the covariates and S which is a 
    scaled version of R.
    R is returns as reads with per-cell counts normalized
    S is returned with per-cell counts normalized, features scaled to N(0,1). and then log(read+1)

    Attributes
    ----------
    self.kNN = Sp[ecify number of nearest neighbors (Excluding self term) ( so k not k+1)
            I.e., same value as passed to the generateCluster methods
    self.batchfilename (opt) file that contains n rows and 2 column withs relevant batch effect names
            col 1 effect name: col two either ONEHOT or CAT for type of inclusion. Same file as used for project
    self.rawfilename. The input read data from which regression will proceeed. Values must be >=0 as the total 
            cell counts is determined from the file ands used for X in the OLS
    self.xtildeNeighborIndicesFilename. Input neighbor matrix (n by k+1)
    self.xtildeNeighborIndicesFilename. Input neighbor distance matrix (n by k+1) (not used at this time)
    self.outputRfilename. Output imputed R matrix
    self.outputSfilename. Output scaled version of imputed R
    """

    def __init__(self, kNN, inrawfilename, indexfilename, distancefilename, outputRfilename, outputSfilename, batchfilename=False):
        self.kNN = kNN
        self.rawfilename = inrawfilename
        self.xtildeNeighborIndicesFilename = indexfilename
        self.xtildeNeighborDistancesFilename = distancefilename
        self.df_new_R = pd.DataFrame()
        self.df_new_S = pd.DataFrame()
        self.df_Rraw = pd.DataFrame()
        self.df_Regress_Pvalues = pd.DataFrame()
        self.outputRfilename = outputRfilename
        self.outputSfilename = outputSfilename
        self.batchfilename = batchfilename
        
    def getRmatrix(self):
        return (self.df_new_R)

    def getSmatrix(self):
        return (self.df_new_S)

    def reportParameters(self):
        print('Regress Data set')
        print('Input reads matrix '+self.rawfilename)
        print('Input Indices '+self.xtildeNeighborIndicesFilename)
        print('Input distances '+self.xtildeNeighborDistancesFilename)
        print('Output R filename '+self.outputRfilename)
        print('Output S filename '+self.outputSfilename)
     
    def regressCellValues(self, df_R, NNindex, cellIndexMap, datacelltotals):
        debug = False
        wlsStyleRegression=True
        #
        if (wlsStyleRegression):
            print('Perform regression with sample weighting and zero-zero point registration')
        timein = tm.time()
        RT = pd.DataFrame()  # TRansposed to allow efficient per-cell updates to be columns
        RTPvalues = pd.DataFrame() # Carry down P-values for double checking        
        icountCells=0
        numGenes = df_R.shape[1]
        #print('Truncated REGRESS list for testing')
        additionalMultiplier = 10.0 # Kirk wanted an additional weight applied to zero point
        numNNs = NNindex.shape[1] # NOTE this INCLUDES the target cell
        for icell in df_R.index.values:
            #print('icel is '+icell)
            NNlist = NNindex.loc[icell].tolist()[0:] # return integer IDs: Also keep the first entry which is redundant to cell name
            NNlistNames = cellIndexMap.loc[NNlist].values.flatten() # Cells for performing regressin
            if (len(NNlist) != len(NNlistNames)):
                print('mismatch: NNlist and NNListNames')
                print('cell maps')
                print(str(cellIndexMap))
                cellIndexMap.to_csv('cellMapIndex.tsv')
                print('detailed output')
                for ival in NNlist:
                    print('ival '+str(ival))
                    print('cell val '+str(cellIndexMap.loc[ival])+'XXX')
            NNlistReadData = df_R.loc[NNlistNames] # Get the actual current reads
            icellTotal = datacelltotals.loc[icell].values # For predicting new R for the target cell/gene pair
            Xdata = datacelltotals.loc[NNlistNames].values.flatten()
            Ydata = NNlistReadData.values
            newGeneValues = list()
            model = LinearRegression(n_jobs=8)
            icellTotal = datacelltotals.loc[icell].values
            if (wlsStyleRegression):
                Xdata = np.insert(Xdata, 0, 0.0)
                Yzeros = np.zeros(shape=(numGenes,1))  # Introduce a "fake" cell with total=0 and read=0 for each gene
                Ydata = np.insert(Ydata, 0, Yzeros.reshape(1,-1),axis=0) 
                sample_weights = [np.sqrt(x) for x in Xdata]  # Large valued Xs have better Ys
                sample_weights[0] = additionalMultiplier * np.max(sample_weights)
                model.fit(Xdata.reshape(-1,1), Ydata, sample_weights) 
            else:
                Xdata = np.insert(Xdata, 0, 0.0)
                Yzeros = np.zeros(shape=(numGenes,1)) 
                Ydata = np.insert(Ydata, 0, Yzeros.reshape(1,-1),axis=0)
                model.fit(Xdata.reshape(-1,1), Ydata)
            #
            # Could predict all K+1 cells at ones and update the input R matrix
            # What if any of the reads are negative ?
            YdataPredict = model.predict(icellTotal.reshape(-1,1)) # Expect only one column of data
            #print('predict '+str(YdataPredict))
            ydatamin = YdataPredict.min()
            if (ydatamin < 0.0):
                print('Negative fit new WLS with new weights found for '+str(YdataPredict.shape))
                zeroout_indices = np.argwhere(YdataPredict < 0.0)
                YdataPredict[ YdataPredict < 0.0] = 0.0
                #YdataPredict[zeroout_indices]=0.0
                #print(str(zeroout_indices))
                print('applied a builtin floor to cell fit '+str(zeroout_indices.shape)+' '+str(icell))
            # Some values can be < 0 probably because the cluster neighbors are poorly chosen
            # Should we set a floor of zero?
            RT[icell]=YdataPredict.flatten()
            RT.index = df_R.columns.values
        print('completed regressions '+str(tm.time()-timein))
        return (RT.T)

    def writeNewR(self):
        self.df_new_R.to_csv(self.outputRfilename,sep=' ')
        print('wrote out new R matrix '+self.outputRfilename)

    def writeNewS(self):
        self.df_new_S.to_csv(self.outputSfilename,sep=' ')
        print('wrote out new S matrix '+self.outputSfilename)
 
    def getOutputRfilename(self):
        return (self.outputRfilename)

    def getOutputSfilename(self):
        return (self.outputSfilename)

    def scaleResults(self,df_R):
        """ For the impute R (which includes group effects) scale 
        data for read counts, and then studentize the features to N(0,1)
        (value+1)*median-all-cell-read / call total read
        The post averaged cell totals is approx 2100 for a factor of 1
        """
        factor=1.0
        # First ave number of reads per cell (or do we want median values?)
        total_reads = float(df_R.values.sum())
        num_reads = float(df_R.values.size)
        ave_total_reads = total_reads / num_reads
        # Alternative method use the median
        ## total_sum_c = pd.DataFrame(df_R.sum(axis=1))
        ## median_total_reads = total_sum_c.median()
        #
        ## total_median_c = pd.DataFrame(df_R.median(axis=1)) # or do we want median values
        ## total_median_c.columns=['TOTALC']
        total_mean_c = pd.DataFrame(df_R.mean(axis=1)) # or do we want median values
        total_mean_c.columns=['TOTALC']
        # Ave number of reads per cells divided by cell reads
        total_correction_factor = factor * ave_total_reads / total_mean_c['TOTALC']
        ## total_correction_factor = factor * median_total_reads / total_median_c['TOTALC']
        # 
        df_logS = np.log((df_R+1.0).mul(total_correction_factor,axis=0))
        df_new_R = pd.DataFrame(sklearn.preprocessing.scale(df_logS, with_mean=True, with_std=True, copy=True,axis=0).tolist())
        df_new_R.index = df_R.index.values
        df_new_R.columns = df_R.columns.values
        print('Scale the R to S function')
        print(str(df_R.shape))
        #print(str(df_new_R)) 
        print('---------------')
        return(df_new_R)

    def runProcess(self):
        print(' read matrix '+self.rawfilename)
        print(' Index neighbors '+self.xtildeNeighborIndicesFilename)
        print(' Distance neighbors '+self.xtildeNeighborDistancesFilename)
        t0=tm.time()
        RinData = pd.read_csv(self.rawfilename,delim_whitespace=True)
        NNindex = pd.read_csv(self.xtildeNeighborIndicesFilename,delim_whitespace=True,index_col=0)
        NNdistance = pd.read_csv(self.xtildeNeighborDistancesFilename,delim_whitespace=True,index_col=0)
        print(tm.time()-t0)
        # Fetch covariates to carry around separately
        columnEffects,oneHotList = getKnownGroupEffects(RinData,self.batchfilename)  # oneHotList not used at this time
        df_Z = RinData[columnEffects]
        df_Z.index = RinData['CELL']
        df_R = RinData.drop(columnEffects,axis=1,inplace=False)
        df_R.set_index('CELL',inplace=True)
        # Total reads per CELL
        datacelltotals = pd.DataFrame(df_R.sum(axis=1))
        datacelltotals.reset_index(inplace=True)
        datacelltotals.columns=['CELL','CELLTOTAL']
        datacelltotals.set_index('CELL',inplace=True)
        # Integer cell ID to cell name
        cellIndexMap = buildCellIndexMap(NNindex)
        # Build regression
        print('Call regression')
        # Can this result in negative values for R?
        df_new_R = self.regressCellValues(df_R, NNindex,cellIndexMap,datacelltotals)  # All Z columns have been removed
        print('min value oif regressed R is '+str(df_new_R.values.min()))
        # Build the properly formatted output matrix including Z covariates
        numRcells = df_R.shape[0]
        numNNcells = NNindex.shape[0]
        if (numRcells != numNNcells):
            print('Cell rank of input R data and input K data are not equal. The resulting data object will take intersection')
            print('Need to investigate: should ABORT')
            print('R and NN cell ranks are '+str(numRcells)+' '+str(numNNcells))
            #sys.exit()
        print('R and NN cell ranks are '+str(numRcells)+' '+str(numNNcells))
        # Compute a feature scaled and log transformed version log the R matrix (called S)

        df_new_S = self.scaleResults(df_new_R)

        # Add effects back into the data object
        finalR = pd.merge(df_Z,df_new_R,left_index=True,right_index=True)
        self.df_new_R = finalR 
        finalS = pd.merge(df_Z,df_new_S,left_index=True,right_index=True)
        self.df_new_S = finalS
        print('Write to disk the new R matrix which is the imputed data and the S matrix which scales R as follows')
        print('S = log( (R+1)*cellCorrectionFactor)')
        self.writeNewR()
        self.writeNewS()
        return ()
