"""

 Deep_genomic.py allows predicting complex traits by using Deep Learning and Penalized Liner models

 Authors: Laura M Zingaretti (m.lau.zingaretti@gmail.com) and Miguel Perez-Enciso (miguel.perez@uab.es)

"""

'''

    Copyright (C) 2019  Laura Zingaretti

    This program is free software: you can redistribute it and/or modify

    it under the terms of the GNU General Public License as published by

    the Free Software Foundation, either version 3 of the License, or

    (at your option) any later version.

    This program is distributed in the hope that it will be useful,

    but WITHOUT ANY WARRANTY; without even the implied warranty of

    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License

    along with this program.  If not, see <https://www.gnu.org/licenses/>.

'''

import matplotlib as mpl

mpl.use('Agg')

from CNN.cnn import acc_pearson_r  as acc_pearson_r

import tensorflow as tf

from keras.backend.tensorflow_backend import set_session

import pandas as pd

import numpy as np

from run_cnn import run_cnn_main

from run_mlp import run_mlp_main

from run_ridge import run_ridge_main

from run_lasso import run_lasso_main

#from run_rnn import run_rnn_main

import seaborn as sns

from sklearn.model_selection import train_test_split

from scipy import stats

from sklearn.decomposition import PCA

from sklearn.preprocessing import StandardScaler

from sklearn.preprocessing import scale

from sklearn.preprocessing import LabelEncoder

from sklearn.preprocessing import OneHotEncoder

from keras.utils import to_categorical

import os

import sys

import argparse

import logging

if __name__ == '__main__':

    parser = argparse.ArgumentParser()

    parser.add_argument("--Xtr", required=True, help="train Predictor matrix")

    parser.add_argument("--ytr", required=True,help= "train response matrix")

    parser.add_argument("--cat", required=False, help="Should be y_train matrix normalized by categories? If yes, category should be provided as cat argument, e.g. --cat T ")

    # if you only want to train the model, you don't need a validation set

    parser.add_argument("--Xval", default=None,required=False, help="validation Predictor")

    parser.add_argument("--yval",default=None, required=False,help= "validation response matrix")

    parser.add_argument("--scale",help="Boolean indicating if y should be scaled", action="store_true")

    parser.add_argument("--dummy", help="Convert SNPs into OneHot encoding", action="store_true")

    parser.add_argument("--categorical", help="Are outputs categorical variables?", action="store_true")

    parser.add_argument("--gridProp", help="proportion of random search", default=1, type=float)

    parser.add_argument("--trait", help="name of trait to be used in the analysis", default=None,required=False)

    parser.add_argument("--output", help="path to output dir",required=True)

    Method = parser.add_mutually_exclusive_group()

    Method.add_argument('--mlp', action='store_true')

    Method.add_argument('--cnn', action='store_true')

    Method.add_argument('--lasso', action='store_true')

    Method.add_argument('--ridge', action='store_true')

    parser.add_argument('--lr',type=float, nargs='+', help='the learning rate parameter to be used in NN configuration. It can be a list to tune', required=False,default=0.025)

    parser.add_argument('--dr1',type=float, nargs='+', help='the dropout to be used in  the first layer of  NN configuration. It can be a list to tune', required=False, default=0)

    parser.add_argument('--dr2',type=float, nargs='+', help='the dropout to be used in  the hidden_layers of the NN configuration. It can be a list to tune', required=False, default=0)

    parser.add_argument('--reg1',type=float, nargs='+', help='the L2 regularization to be used in  the first layer of  NN configuration. It can be a list to tune', required=False,default=0)

    parser.add_argument('--nconv',type=int, nargs='+',help='number of convolutions layers to be considered in convolutional operation. It only works if Method is CNN and it can be a list to tune', default=1)

    parser.add_argument('--act1', action='store',type=str, nargs='+', default=['relu', 'elu', 'linear', 'softmax', 'tanh', 'softplus'],help="Examples: --act1 relu elu, --act1 linear")

    parser.add_argument('--act2', action='store',type=str, nargs='+', default=['relu', 'elu', 'linear', 'softmax', 'tanh', 'softplus'],help="Examples: --act2 relu elu, --act2 linear")

    parser.add_argument('--reg2',type=float, nargs='+', help='the L2 regularization to be used in  the hidden_layers of  NN configuration. It can be a list to tune', required=False,default=0)

    parser.add_argument('--reg3',type=float, nargs='+', help='the L2 regularization to be used in  the output layers of the  NN configuration. It can be a list to tune', required=False,default=0)

    parser.add_argument('--hn',type=int, nargs='+', help='Number of hidden neurons to be used on the NN configuration. It can be a list to tune', required=False, default=8)

    parser.add_argument('--ps',type=int, nargs='+', help='Pool size to be used on the CNN configuration. It can be a list to tune', required=False,default=3)

    parser.add_argument('--hl',type=int, nargs='+', help='Number of hidden layers to be used in the NN configuration. It can be a list to tune', required=False,default=1)

    parser.add_argument('--ks',type=int, nargs='+', help='kernel_size to be used in CNN configuration. It can be a list to tune', required=False,default=3)

    parser.add_argument('--ns',type=int, nargs='+', help='stride to be used in CNN configuration. It can be a list to tune', required=False,default=1)

    parser.add_argument('--epochs',type=int, nargs='+', help='nepochs to be used for training. It can be a list to tune', required=False,default=50)

    parser.add_argument('--nfilters',type=int, nargs='+', help='number of filter (#convolutions in each convolutional layer). It can be a list to tune', required=False,default=8)

    parser.add_argument('--optimizer', action='store',type=str, nargs='+', default=['Adam', 'Nadam', 'sgd'],help="Examples: -optimizer Adam")

    parser.add_argument('--bs',type=int, nargs='+', help='batch size to be used in NN configuration. It can be a list to tune', required=False,default=16)

    parser.add_argument('--N_neurons_FL',type=int, nargs='+', help='NNeurons to be used at the first layer of MLP configuration. It can be a list to tune', required=False,default=8)

    args = parser.parse_args()

    if not isinstance(args.hn, list):

           args.hn = [args.hn]

    if not isinstance(args.nconv, list):

           args.nconv = [args.nconv]

    if not isinstance(args.ps, list):

           args.ps = [args.ps]

    if not isinstance(args.hl, list):

           args.hl = [args.hl]

    if not isinstance(args.ks, list):

           args.ks = [args.ks]

    if not isinstance(args.ns, list):

           args.ns = [args.ns]

    if not isinstance(args.epochs, list):

           args.epochs = [args.epochs]

    if not isinstance(args.nfilters, list):

           args.nfilters = [args.nfilters]

    if not isinstance(args.bs, list):

           args.bs = [args.bs]

    if not isinstance(args.N_neurons_FL, list):

        args.N_neurons_FL = [args.N_neurons_FL]

    if not isinstance(args.reg1, list):

           args.reg1 = [args.reg1]

    if not isinstance(args.reg2, list):

           args.reg2 = [args.reg2]

    if not isinstance(args.reg3, list):

           args.reg3 = [args.reg3]

    if not isinstance(args.dr1, list):

           args.dr1 = [args.dr1]

    if not isinstance(args.dr2, list):

           args.dr2 = [args.dr2]

    if not isinstance(args.lr, list):

           args.lr = [args.lr]

    if args.dummy:

        if args.Xval is not None:

            Xtr=pd.read_csv(args.Xtr,sep='\s+')

            Xval=pd.read_csv(args.Xval,sep='\s+')

            All_X = pd.concat([Xtr,Xval])

            All_X = All_X.round(decimals=0)

            #All_X = All_X.apply(pd.to_numeric)

            le = OneHotEncoder(sparse=False)

            label_encoder = LabelEncoder()

            X_enco = All_X.apply(label_encoder.fit_transform)

            onehot_encoder = OneHotEncoder(sparse=False)

            All_X_oh = onehot_encoder.fit_transform(X_enco)

            X_tr = All_X_oh[0:Xtr.shape[0], :]

            X_vl = All_X_oh[Xtr.shape[0]:All_X_oh.shape[0], ]

        else:

            X_tr = pd.read_csv(args.Xtr,sep='\s+').round(decimals=0)

            #X_tr = X_tr.apply(pd.to_numeric)

            le = OneHotEncoder(sparse=False)

            label_encoder = LabelEncoder()

            X_enco = X_tr.apply(label_encoder.fit_transform)

            onehot_encoder = OneHotEncoder(sparse=False)

            X_tr = onehot_encoder.fit_transform(X_enco)

            y_tr = pd.read_csv(args.ytr,sep='\s+').apply(pd.to_numeric)

            X_tr, X_vl, y_tr, y_val = train_test_split(X_tr, y_tr, test_size=0.2)

    else:

        if args.Xval is not None:

            X_tr= pd.read_csv(args.Xtr,sep='\s+').apply(pd.to_numeric)

            X_vl= pd.read_csv(args.Xval,sep='\s+').apply(pd.to_numeric)

        else:

            X_tr=pd.read_csv(args.Xtr,sep='\s+').apply(pd.to_numeric)

            y_tr=pd.read_csv(args.ytr,sep='\s+').apply(pd.to_numeric)

            X_tr, X_vl, y_tr, y_val = train_test_split(X_tr, y_tr, test_size=0.2)

    if args.categorical is False:

        if args.scale is True:

            if args.cat is not None:

                if args.yval is not None:

                    y_tr = pd.read_csv(args.ytr,sep='\s+')

                    y_val = pd.read_csv(args.yval,sep='\s+').apply(pd.to_numeric)

                    cat = args.cat

                    loc = y_tr.columns.get_loc(cat)

                    Factors=np.unique(y_tr.iloc[:, loc])

                    scaled=pd.DataFrame([])

                    for i in Factors:

                        M = y_tr.iloc[np.where(y_tr.iloc[:, loc] == i)]

                        M = M.iloc[:, :-1].apply(scale)

                        scaled=scaled.append(M)

                    y_val = y_val.apply(scale)

                    scaled.columns=y_tr.columns[:-1]

                    y_val.columns=y_tr.columns[:-1]

                    y_tr = scaled

                else:

                    y_tr = pd.read_csv(args.ytr,sep='\s+').apply(pd.to_numeric)

                    y_tr = y_tr.apply(scale)

                    y_val = y_val.apply(scale)

            else:

                if args.yval is not None:

                    y_tr = pd.read_csv(args.ytr,sep='\s+').apply(pd.to_numeric)

                    y_val = pd.read_csv(args.yval,sep='\s+').apply(pd.to_numeric)

                    y_val = y_val.apply(scale)

                    y_tr = y_tr.apply(scale)

                else:

                    y_tr = pd.read_csv(args.ytr).apply(pd.to_numeric)

                    y_tr = y_tr.apply(scale)

                    y_val = y_val.apply(scale)

        else:

            if args.yval is not None:

                y_tr = pd.read_csv(args.ytr,sep='\s+').apply(pd.to_numeric)

                y_val = pd.read_csv(args.yval,sep='\s+').apply(pd.to_numeric)

            else:

                y_tr = pd.read_csv(args.ytr,sep='\s+').apply(pd.to_numeric)

        if args.trait is not None:

            trait= args.trait

            loc=y_tr.columns.get_loc(trait)

            y_tr=y_tr.iloc[:,loc]

            y_val=y_val.iloc[:,loc]

        else:

            trait=None

        y_tr = np.asarray(y_tr)

        y_val = np.asarray(y_val)

    else:

        if args.yval is not None:

            y_tr = pd.read_csv(args.ytr,sep='\s+')

            y_val = pd.read_csv(args.yval,sep='\s+')

        else:

            y_tr = pd.read_csv(args.ytr,sep='\s+')

        if args.trait is not None:

            trait = args.trait

            loc = y_tr.columns.get_loc(trait)

            y_tr = y_tr.iloc[:, loc]

            y_val = y_val.iloc[:, loc]

            y_tr = np.asarray(y_tr)

            y_val = np.asarray(y_val)

            encoder = LabelEncoder()

            encoder.fit(y_tr)

            y_tr = encoder.transform(y_tr)

            y_tr = to_categorical(y_tr)

            encoder = LabelEncoder()

            encoder.fit(y_val)

            y_val = encoder.transform(y_val)

            y_val=to_categorical(y_val)

        else:

            if len(y_tr.columns) > 1:

                sys.exit("To categorical Values,Only one trait is allowed ")

            else:

                y_tr = np.asarray(y_tr)

                y_val = np.asarray(y_val)

                encoder = LabelEncoder()

                encoder.fit(y_tr)

                y_tr = encoder.transform(y_tr)

                y_tr = to_categorical(y_tr)

                encoder = LabelEncoder()

                encoder.fit(y_val)

                y_val = encoder.transform(y_val)

                y_val = to_categorical(y_val)

    X_vl=np.asarray(X_vl)

    X_tr=np.asarray(X_tr)

    main=os.getcwd()

    diro=args.output

###to cat

    if args.cnn:

        if trait is not None:

            run_cnn_main(X_tr,y_tr,X_vl,y_val,output=diro,main=main,prop=args.gridProp,trait=trait,cat=args.categorical,lr=args.lr,dr_1=args.dr1,dr_2=args.dr2,reg_1=args.reg1,reg_2=args.reg2,reg_3=args.reg3,nconv=args.nconv,act_1=args.act1,act_2=args.act2,hn=args.hn,ps=args.ps,hl=args.hl,ks=args.ks,ns=args.ns,epochs=args.epochs,nf=args.nfilters,op=args.optimizer,bs=args.bs)

        else:

            run_cnn_main(X_tr, y_tr, X_vl, y_val, output=diro, main=main, prop=args.gridProp,trait=None,cat=args.categorical,lr=args.lr,dr_1=args.dr1,dr_2=args.dr2,reg_1=args.reg1,reg_2=args.reg2,reg_3=args.reg3,nconv=args.nconv,act_1=args.act1,act_2=args.act2,hn=args.hn,ps=args.ps,hl=args.hl,ks=args.ks,ns=args.ns,epochs=args.epochs,nf=args.nfilters,op=args.optimizer,bs=args.bs)

    if args.mlp:

        if trait is not None:

            run_mlp_main(X_tr,y_tr,X_vl,y_val,output=diro,main=main,prop=args.gridProp,trait=trait,cat=args.categorical,lr=args.lr,dr_1=args.dr1,dr_2=args.dr2,reg_1=args.reg1,reg_2=args.reg2,act_1=args.act1,hn=args.hn,hl=args.hl,epochs=args.epochs,op=args.optimizer,bs=args.bs,N_neurons_FL=args.N_neurons_FL)

        else:

            run_mlp_main(X_tr, y_tr, X_vl, y_val, output=diro, main=main, prop=args.gridProp,trait=None,cat=args.categorical,lr=args.lr,dr_1=args.dr1,dr_2=args.dr2,reg_1=args.reg1,reg_2=args.reg2,act_1=args.act1,hn=args.hn,hl=args.hl,epochs=args.epochs,op=args.optimizer,bs=args.bs,N_neurons_FL=args.N_Neurons_FL)

    if args.ridge:

        if args.categorical is True:

            sys.exit("ridge only allowing to continous outputs")

        else:

            run_ridge_main(X_tr,y_tr,X_vl,y_val,output=diro,main=main)

    if args.lasso:

        if args.categorical is True:

            sys.exit("ridge only allowing to continous outputs")

        else:

            run_lasso_main(X_tr,y_tr,X_vl,y_val,output=diro,main=main)