import datetime,pandas,numpy,scipy,os,pickle
import sklearn,sklearn.decomposition
from pkg_resources import Requirement, resource_filename
import miner2.coexpression

def axisTfs(axesDf,tfList,expressionData,correlationThreshold=0.3):
    
    axesArray = numpy.array(axesDf.T)
    if correlationThreshold > 0:
        tfArray = numpy.array(expressionData.loc[tfList,:])
    axes = numpy.array(axesDf.columns)
    tfDict = {}

    if type(tfList) is list:
        tfs = numpy.array(tfList)
    elif type(tfList) is not list:
        tfs = tfList
        
    if correlationThreshold == 0:
        for axis in range(axesArray.shape[0]):
            tfDict[axes[axis]] = tfs

        return tfDict       
    
    for axis in range(axesArray.shape[0]):
        tfCorrelation = miner2.coexpression.pearson_array(tfArray,axesArray[axis,:])
        tfDict[axes[axis]] = tfs[numpy.where(numpy.abs(tfCorrelation)>=correlationThreshold)[0]]
    
    return tfDict

def enrichment(axes,revisedClusters,expressionData,correlationThreshold=0.3,numCores=1,p=0.05,database="tfbsdb_tf_to_genes.pkl"):
    
    print(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S \t mechanistic inference"))
    
    tfToGenesPath = resource_filename(Requirement.parse("miner2"), 'miner2/data/{}'.format(database))
    with open(tfToGenesPath, 'rb') as f:
        tfToGenes = pickle.load(f)
     
    if correlationThreshold <= 0:
        allGenes = [int(len(expressionData.index))]
    elif correlationThreshold > 0:
        allGenes = list(expressionData.index)
        
    tfs = tfToGenes.keys()     
    tfMap = axisTfs(axes,tfs,expressionData,correlationThreshold=correlationThreshold)

    print('taskSplit',taskSplit)
    sys.exit()

    taskSplit = splitForMultiprocessing(revisedClusters.keys(),numCores)
    tasks = [[taskSplit[i],(allGenes,revisedClusters,tfMap,tfToGenes,p)] for i in range(len(taskSplit))]
    tfbsdbOutput = multiprocess(tfbsdbEnrichment,tasks)
    mechanisticOutput = condenseOutput(tfbsdbOutput)
        
    return mechanisticOutput

def principalDf(dict_,expressionData,regulons=None,subkey='genes',minNumberGenes=8,random_state=12):

    print(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S \t preparing mechanistic inference"))

    pcDfs = []
    setIndex = set(expressionData.index)
    
    if regulons is not None:
        dict_, df = regulonDictionary(regulons)
    for i in dict_.keys():
        if subkey is not None:
            genes = list(set(dict_[i][subkey])&setIndex)
            if len(genes) < minNumberGenes:
                continue
        elif subkey is None:
            genes = list(set(dict_[i])&setIndex)
            if len(genes) < minNumberGenes:
                continue
            
        pca = sklearn.decomposition.PCA(1,random_state=random_state)
        principalComponents = pca.fit_transform(expressionData.loc[genes,:].T)
        principalDf = pandas.DataFrame(principalComponents)
        principalDf.index = expressionData.columns
        principalDf.columns = [str(i)]
        
        normPC = numpy.linalg.norm(numpy.array(principalDf.iloc[:,0]))
        pearson = scipy.stats.pearsonr(principalDf.iloc[:,0],numpy.median(expressionData.loc[genes,:],axis=0))
        signCorrection = pearson[0]/numpy.abs(pearson[0])
        
        principalDf = signCorrection*principalDf/normPC
        
        pcDfs.append(principalDf)
    
    principalMatrix = pandas.concat(pcDfs,axis=1)
        
    return principalMatrix
