# Copyright 2017, Rohan Dandage <rraadd_8@hotmail.com,rohan@igib.in>
# This program is distributed under General Public License v. 3.  

"""
================================
``io_mut_files``
================================
"""
from __future__ import division
import sys
from os.path import splitext,exists,basename,abspath,dirname
from os import makedirs,stat
import pandas as pd
import numpy as np
from glob import glob
import logging
import subprocess

from dms2dfe.lib.io_dfs import set_index

def getcolsmets(d,cs):
    """
    Get the total non-na items in column. 

    :param d: pandas dataframe
    :param cs: list of columns
    """
    cols=[c for c in d.columns if cs in c]
    counts=[len(d[c].dropna()) for c in cols]
    return dict(zip(cols,counts))

def get_2SD_cutoffs(d,reps,N=False):
    """
    Get 2SD cutoffs for values in given column

    :param d: pandas dataframe
    :param reps: names of replicate conditions
    """
    t=d.copy()
    mus=[]
    sigmas=[]
    for r1i,r1 in enumerate(reps):
        for r2i,r2 in enumerate(reps):
            if r1i<r2i:
                t.loc[:,'reps']=t.loc[:,r1]-t.loc[:,r2]
                if N and ('mut' in t):
                    t.loc[(t.loc[:,'mut']==t.loc[:,'ref']),'reps']=np.nan
                mus.append(t.loc[:,'reps'].mean())
                sigmas.append(t.loc[:,'reps'].std())                
    mu=np.mean(mus)
    sigma=np.sqrt(np.mean([s**2 for s in sigmas]))
    return mu+sigma*2,mu,sigma

def get_repli_FiA(d,csel='.NiA_tran.sel',cref='.NiA_tran.ref'):
    """
    Get replicates from data_fit

    :param d: pandas dataframe data_fit
    """
    sels=np.sort([c for c in d.columns if csel in c])
    refs=np.sort([c for c in d.columns if cref in c])
    FCAi=1
    cols_FCA=[]
    for refi,ref in enumerate(refs):
        for seli,sel in enumerate(sels):
            if refi==seli:
                coln='FCA%s_reps' % FCAi
                d.loc[:,coln]=d.loc[:,sel]-d.loc[:,ref]
                d.loc[(d.loc[:,'mut']==d.loc[:,'ref']),coln]=np.nan
                cols_FCA.append(coln)
                FCAi+=1
    return d.loc[:,cols_FCA]

def data_fit2cutoffs(d,sA,sB,N=True):
    """
    Get cutoffs of data fit to assign classes

    :param d: pandas dataframe with fold change values
    :param sA,sB: columns with two conditions to be compared
    """
    refs=[c for c in d.columns if sA in c]
    sels=[c for c in d.columns if sB in c]
    _,mu1,sigma1=get_2SD_cutoffs(d,refs)
    _,mu2,sigma2=get_2SD_cutoffs(d,sels)
    return np.mean([mu1,mu2])+2*np.sqrt(np.mean([sigma1,sigma2]))

def class_fit(d,col_fit='FiA',FC=True,zscore=False): #column of the data_fit
    """
    This classifies the fitness of mutants into beneficial, neutral or, deleterious.
    
    :param d: dataframe of `data_fit`.
    :returns d: classes of fitness written in 'class-fit' column based on values in column 'FiA'. 
    """
    cols_reps=[c for c in d if '.NiA_tran.ref' in c]
    if FC and (len(cols_reps)==2):
        up,_,_=get_2SD_cutoffs(d,cols_reps,N=True)
        dw=-1*up
    else:
        if zscore:
            up,dw=-2,2
        else:
            up,dw=0,0
    d.loc[d.loc[:,col_fit]>+up,    'class_fit']="enriched"
    d.loc[((d.loc[:,col_fit]>=dw) & (d.loc[:,col_fit]<=up)),'class_fit']="neutral"
    d.loc[d.loc[:,col_fit]<dw,    'class_fit']="depleted"
    return d

def class_comparison(dA,dB):
    """
    This classifies differences in fitness i.e. relative fitness into positive, negative or robust categories. 
    
    :param dc: dataframe with `dc`. 
    :returns dc: dataframe with `class__comparison` added according to fitness levels in input and selected samples in `dc`
    """
    dA=set_index(dA,'mutids')
    dB=set_index(dB,'mutids')
    dc=get_repli_FiA(dA).join(get_repli_FiA(dB),lsuffix='_test',rsuffix='_ctrl')
    up=data_fit2cutoffs(dc,sA='_reps_test',sB='_reps_ctrl',N=False)
    dw=-1*up

    diff=dA.loc[:,'FiA']-dB.loc[:,'FiA']
    diff.index.name='mutids'
    diff=diff.reset_index()
    mutids_up=diff.loc[(diff.loc[:,'FiA']>up),'mutids'].tolist()
    mutids_dw=diff.loc[(diff.loc[:,'FiA']<dw),'mutids'].tolist()

    dc.loc[mutids_up,'class_comparison']="positive"
    dc.loc[mutids_dw,'class_comparison']="negative"
    return dc.loc[:,'class_comparison']


def get_data_metrics(prj_dh):
    """
    Obtain metrics for fold change values from a experiment

    :param prj_dh: path to project directory
    """
    data_fit_fns_all=[basename(s) for s in glob('%s/data_fit/aas/*' % prj_dh)]
    data_fit_fns_all2labels=dict(zip(data_fit_fns_all,data_fit_fns_all))
    data_fit_metrics=pd.DataFrame(index=data_fit_fns_all)

    fit=pd.read_csv('%s/cfg/fit' % prj_dh).set_index('unsel')
    comparison=pd.read_csv('%s/cfg/comparison' % prj_dh).set_index('ctrl')

    for ctrl in comparison.index:
        fh_ctrls=glob('%s/data_fit/aas/%s_WRT_*' % (prj_dh,ctrl))
        for fh_ctrl in fh_ctrls:
            dctrl=pd.read_csv(fh_ctrl).set_index('mutids')

            up,_,_=get_2SD_cutoffs(dctrl,[c for c in dctrl if '.NiA_tran.ref' in c],N=True)
            dw=-1*up
            # print up,dw
            ctrl=basename(fh_ctrl)
            ctrls={'versus_%s' % ctrl:ctrl}

            data_fit_metrics.index.name='fn'
            for fni,fn in enumerate(data_fit_fns_all):
                fh= '%s/data_fit/aas/%s' % (prj_dh,fn)
                d=pd.read_csv(fh).set_index('mutids')
                data_fit_metrics.loc[fn,'$\mu+2\sigma$']=up
                d=d.reset_index()
                mutids_up=d.loc[(d.loc[:,'FiA']>up),'mutids'].tolist()
                mutids_dw=d.loc[(d.loc[:,'FiA']<dw),'mutids'].tolist()
                mutids_updw=mutids_up+mutids_dw
                d=d.set_index('mutids')
                data_fit_metrics.loc[fn,'$n$']=len(d.loc[:,'FiA'].dropna())
                ref_counts=getcolsmets(d,cs='.NiA_tran.ref')
                sel_counts=getcolsmets(d,cs='.NiA_tran.sel')
                for i,_ in enumerate(ref_counts.keys()):
                    data_fit_metrics.loc[fn,'$n$ ref%02d' % i]=ref_counts[ref_counts.keys()[i]]
                for i,_ in enumerate(sel_counts.keys()):
                    data_fit_metrics.loc[fn,'$n$ sel%02d' % i]=sel_counts[sel_counts.keys()[i]]
                data_fit_metrics.loc[fn,'$n_{neutral}$']=data_fit_metrics.loc[fn,'$n$']-len(mutids_updw)
                data_fit_metrics.loc[fn,'$n_{enriched}$']=len(mutids_up)
                data_fit_metrics.loc[fn,'$n_{depleted}$']=len(mutids_dw)
                data_fit_metrics.loc[fn,'$F$ mean']=d.loc[:,'FiA'].mean()
    
    if len(data_fit_metrics.columns)>0:
        data_fit_metrics['$n_{enriched}$%%']=data_fit_metrics['$n_{enriched}$']\
        /(data_fit_metrics['$n_{enriched}$']+data_fit_metrics['$n_{depleted}$'])*100
        data_fit_metrics['$n_{depleted}$%%']=data_fit_metrics['$n_{depleted}$']\
        /(data_fit_metrics['$n_{enriched}$']+data_fit_metrics['$n_{depleted}$'])*100
        data_fit_metrics['$n_{enriched}$%']=data_fit_metrics['$n_{enriched}$']/data_fit_metrics['$n$']*100
        data_fit_metrics['$n_{depleted}$%']=data_fit_metrics['$n_{depleted}$']/data_fit_metrics['$n$']*100
        data_fit_metrics['$n_{neutral}$%']=data_fit_metrics['$n_{neutral}$']/data_fit_metrics['$n$']*100

        data_fit_metrics['$E$']=data_fit_metrics['$n_{enriched}$']/\
        (data_fit_metrics['$n_{enriched}$']+data_fit_metrics['$n_{depleted}$'])
        # print data_fit_metrics.index
        # print data_fit_metrics.columns
        
        data_fit_metrics.loc[:,'labels']=[data_fit_fns_all2labels[i] for i in data_fit_metrics.index]

        for c in ctrls:
            fh= '%s/data_fit/aas/%s' % (prj_dh,ctrls[c])
            dA=pd.read_csv(fh).set_index('mutids')
            for fni,fn in enumerate(data_fit_fns_all):
                fh= '%s/data_fit/aas/%s' % (prj_dh,fn)
                dB=pd.read_csv(fh).set_index('mutids')            
                dc=get_repli_FiA(dA).join(get_repli_FiA(dB),lsuffix='_ctrl',rsuffix='_test')
                up=data_fit2cutoffs(dc,sA='_reps_test',sB='_reps_ctrl',N=False)
                dw=-1*up
                data_fit_metrics.loc[fn,'$\mu+2\sigma$ comparison %s' % c]=up

                diff=dB.loc[:,'FiA']-dA.loc[:,'FiA']
                diff=diff.reset_index()
                mutids_up=diff.loc[(diff.loc[:,'FiA']>up),'mutids'].tolist()
                mutids_dw=diff.loc[(diff.loc[:,'FiA']<dw),'mutids'].tolist()
                mutids_updw=mutids_up+mutids_dw
                diff=diff.set_index('mutids')
                data_fit_metrics.loc[fn,'$n$ %s' % c]=len(diff.dropna())
                data_fit_metrics.loc[fn,'$n_{positive}$ %s' % c]=len(mutids_up)
                data_fit_metrics.loc[fn,'$n_{negative}$ %s' % c]=len(mutids_dw)

            data_fit_metrics['$\Delta n$ %s' % c]=(data_fit_metrics['$n$']-data_fit_metrics.loc[ctrls[c],'$n$'])
            data_fit_metrics['$\Delta n_{enriched}$ %s' % c]=(data_fit_metrics['$n_{enriched}$']-data_fit_metrics.loc[ctrls[c],'$n_{enriched}$'])
            data_fit_metrics['$\Delta n_{enriched}$%% %s' % c]=(data_fit_metrics['$n_{enriched}$%']-data_fit_metrics.loc[ctrls[c],'$n_{enriched}$%'])
            data_fit_metrics['$\Delta F$ %s' % c]=(data_fit_metrics['$F$ mean']-data_fit_metrics.loc[ctrls[c],'$F$ mean'])
        data_fit_metrics_fh='%s/data_comparison/data_fit_metrics' % prj_dh
        data_fit_metrics.to_csv(data_fit_metrics_fh)
        return data_fit_metrics
    else: 
        logging.info('data_metrics has 0 cols')