import json
import csv
import numpy as np
import os
import pandas as pd
import random
import tensorflow as tf
import matplotlib.pyplot as plt

import keras
from keras.models import Sequential,Model
from keras.utils.generic_utils import get_custom_objects
from keras import backend as K
from keras.layers import Input, Conv1D, Reshape,\
    GlobalAveragePooling2D, Dense, BatchNormalization, Activation,AveragePooling1D, \
    GlobalMaxPooling2D, Flatten, MaxPool1D, Conv2D,MaxPool2D,SeparableConv2D,Conv3D,Add,Dropout
from keras.utils.vis_utils import plot_model
#from keras.applications.mobilenet import DepthwiseConv2D
from keras.callbacks import ModelCheckpoint, ReduceLROnPlateau, EarlyStopping

from keras.layers.advanced_activations import LeakyReLU, PReLU
from functions import quan_detector, most_repeared_promoter,dataset
from sklearn.metrics import confusion_matrix
import argparse
from keras.callbacks import LearningRateScheduler, TensorBoard, ModelCheckpoint
from sklearn import metrics
from sklearn.metrics import roc_auc_score

np.random.seed(41)
tf.set_random_seed(41)
random.seed(41)


labels_file = './labes.csv'
labels_df = pd.read_csv(labels_file, index_col=0)
ids_csv = labels_df.FID.tolist()

files_num_chr7 = [1907,2,1908,2780,2428,3,1173,1291]
files_num_chr17 = [2264,865,66,69,1931,71,1932,70]
files_num_chr9 = [502, 1420, 1503, 1504, 1505, 1506, 1507, 1508]
files_num_chr22 = [65,66,0,80,15,16,59,5]

def fn(d_dir):
    file_list = os.listdir(d_dir)
    n = [i.split('.')[0] for i in file_list]
    num = [int(i.split('_')[1]) for i in n  if len(i)>3]
    return num



num_7 = len(files_num_chr7)
num_9 = len(files_num_chr9)
num_17 = len(files_num_chr17)
num_22 = len(files_num_chr22)

files_num = files_num_chr7 + files_num_chr9 + files_num_chr17 + files_num_chr22

all_HL_prom = np.zeros((11908,64*len(files_num)))
for idx in range(len(files_num)):
    promoter_num  = files_num[idx]
    if idx <num_7:
        #promoter_file = './chr7_hl_prom/chr22_'+str(promoter_num)+'.json'
        promoter_file = './chr7_hl_prom/chr22_'+str(promoter_num)+'.json'
    elif idx <num_7+num_9 and idx >= num_7:
        promoter_file = './chr9_hl_prom/chr9_'+str(promoter_num)+'.json'
    elif idx <num_7+num_9+ num_17 and idx >= num_7+num_9:
        promoter_file = './chr17_hl_prom/chr22_'+str(promoter_num)+'.json'
    else:
        promoter_file = './chr22_hl_prom/chr22_'+str(promoter_num)+'.json'
        
    # # read files
    with open(promoter_file) as json_data:
        ind_var = json.load(json_data)
    var_num = []
    for i in ids_csv:
        id_name = str(i)
        temp = ind_var[id_name]
        var_seq = [int(t) for t in temp]
        var_num.append(var_seq)
    all_HL_prom[:,idx*64:(idx+1)*64] =np.array(var_num)

    
print(all_HL_prom.shape)
n = len(files_num)

labels_df['vars'] = all_HL_prom.tolist()
lab_num = {1: [1, 0],  # control
           2: [0, 1]}  # ALS
lab_num_batch = {'c1': [1,0,0,0],  # control
           'c3': [0,1,0,0],
           'c5': [0,0,1,0],
           'c44':[0,0,0,1]}  # ALS

pheno_new = []
pheno_batch = []
for i in labels_df.Pheno.tolist():
    pheno_new.append(lab_num[i])

# for i in labels_df.Sex.tolist():
#     pheno_new.append(lab_num[i])
for i in labels_df.FID.tolist():
    l = i.split('-')[0]
    pheno_batch.append(lab_num_batch[l])

d = {"Pheno": pheno_new, "Vars": labels_df.vars}
dataset_ = pd.DataFrame(d)
dataset_X = np.array(dataset_.Vars.tolist())
dataset_Y = np.array(dataset_.Pheno.tolist())


N,M = dataset_X.shape
print(N,M)
# network accuracy

t_idx = [int(line.strip()) for line in open("train_id.txt", 'r')]
te_idx = [int(line.strip()) for line in open("test_id.txt", 'r')]
x_tv = dataset_X[t_idx]
y_tv = dataset_Y[t_idx]
x_test = dataset_X[te_idx]
y_test = dataset_Y[te_idx]



x_train, y_train, x_val, y_val,tr_idx,val_idx = dataset(x_tv, y_tv, test_ratio=0.05)


# print x_test.shape
num_classes = y_test.shape[-1]
x_train = x_train.astype('float32')  
x_test = x_test.astype('float32')
x_val = x_val.astype('float32')

x_train = x_train.astype('float32').reshape((len(x_train), M,1))
x_test = x_test.astype('float32').reshape((len(x_test),M,1))
x_val = x_val.astype('float32').reshape((len(x_val),M,1))

print(np.sum(y_test,axis=0))
print(np.sum(y_val,axis=0))
print(np.sum(y_train,axis=0))
conv_kwargs = {'padding':'same',
               'data_format':'channels_first'}

mp_kwargs = {'padding':'same',
               'data_format':'channels_first'}


def architecture_X(input_shape, num_classes):
    act = 'relu'
    
    x = Input(shape=(M, 1))
    input_Conv = Conv1D(filters = 256, kernel_size =64, strides=64)(x)
    input_Conv = Conv1D(filters = 256, kernel_size =1,strides=1)(input_Conv)
    input_BN = BatchNormalization(epsilon=1e-05)(input_Conv)
    input_Conv = Conv1D(filters = 256, kernel_size =1,strides=1)(input_Conv)
    input_act = Activation(act)(input_BN)
    
    
    reshape_h = Reshape((int(M/64), 16, 16))(input_act)
    # DepthwiseConv2D(32, strides=(2, 2), **conv_kwargs)(x)
    
    # add traditional conv
    conv1 = Conv2D(64, (3, 3),activation=act,**conv_kwargs)(reshape_h)
    conv1 = Conv2D(64, (2, 2),activation=act,**conv_kwargs)(conv1)
    conv1 = Conv2D(64, (2, 2),activation=act,**conv_kwargs)(conv1)
    mp1 = MaxPool2D(pool_size=(2, 2),**mp_kwargs)(conv1)
    
    Sconv1 = SeparableConv2D(128, (2, 2),activation=act,**conv_kwargs)(mp1)
    # Sconv1 = SeparableConv2D(32, (2, 2),activation=act,**conv_kwargs)(Sconv1)
    # Sconv1 = SeparableConv2D(32, (2, 2),activation=act,**conv_kwargs)(Sconv1)
    Smp1 = MaxPool2D(pool_size=(2, 2),**mp_kwargs)(Sconv1)
    
    Smp1 = Conv2D(256, (1, 1),activation=act,**conv_kwargs)(Smp1)
    
    Sconv2 = Conv2D(64, (2, 2),activation=act,**conv_kwargs)(Smp1)
    Sconv2 = Conv2D(64, (2, 2),activation=act,**conv_kwargs)(Sconv2)
    Sconv2 = SeparableConv2D(64, (2, 2),activation=act,**conv_kwargs)(Sconv2)
    Smp2 = Sconv2#MaxPool2D(pool_size=(2, 2),**mp_kwargs)(Sconv2)
    
    
    conv2 = Conv2D(64, (2, 2),activation=act,**conv_kwargs)(Smp1)
    conv2 = Conv2D(64, (2, 2),activation=act,**conv_kwargs)(conv2)
    conv2 = Conv2D(64, (2, 2),activation=act,**conv_kwargs)(conv2)
    mp2 = conv2#MaxPool2D(pool_size=(2, 2),**mp_kwargs)(conv2)
    
    sum_add =keras.layers.Concatenate(axis=-1)([mp2,Smp2])
    
    conv3 = Conv2D(128, (2, 2),activation=act,**conv_kwargs)(sum_add)
    conv3 = Conv2D(128, (2, 2),activation=act,**conv_kwargs)(conv3)
    conv3 = Conv2D(128, (2, 2),activation=act,**conv_kwargs)(conv3)
    conv3 = MaxPool2D(pool_size=(1, 2),**mp_kwargs)(conv3)
    
    conv3 = Conv2D(64, (2, 2),activation=act,**conv_kwargs)(conv3)
    conv3 = Conv2D(64, (2, 2),activation=act,**conv_kwargs)(conv3)
    conv3 = Conv2D(64, (2, 2),activation=act,**conv_kwargs)(conv3)
    
#     new = keras.layers.dot([conv2,conv3], axes = -2)
    conv4 = keras.layers.Add()([Sconv2,conv2,conv3])

    conv4 = Conv2D(64, (2, 2),activation=act,**conv_kwargs)(conv4)
    conv4 = Conv2D(64, (2, 2),activation=act,**conv_kwargs)(conv4)
    conv4 = Conv2D(64, (2, 2),activation=act,**conv_kwargs)(conv4)
#     conv4_ = Conv2D(64, (1, 1),activation=act,**conv_kwargs)(conv4)
    
    conv5 = keras.layers.Add()([conv4,conv3])
    conv5 = Conv2D(128, (2, 2),activation=act,**conv_kwargs)(conv5)
    conv5 = Conv2D(128, (2, 2),activation=act,**conv_kwargs)(conv5)
    conv5 = Conv2D(128, (2, 2),activation=act,**conv_kwargs)(conv5)
#     conv5 = keras.layers.Add()([conv4,conv5])
    
    mp3 = GlobalAveragePooling2D(data_format='channels_first')(conv5)
    
    flatten =mp3#Flatten()(mp3)

    d1 = Dense(64*4, activation='linear')(flatten)
    d1_act = Activation(act)(d1)

    d2 = Dense(16, activation='linear')(d1_act)
    flatten = Activation(act)(d2)

    # added 
    # flatten = BatchNormalization(epsilon=1e-05)(flatten)
    pred = Dense(num_classes, activation='softmax')(flatten)
    # Compile model
    model = Model(inputs=[x], outputs=[pred])
    model.compile(loss='categorical_crossentropy', 
                  optimizer=keras.optimizers.adagrad(lr=0.002,decay=0.0001),#0.0013
                  metrics=['accuracy'])
    return model
def scheduler(epoch):
    if epoch < 50:
        return 0.1
    if epoch < 122:
        return 0.01
    return 0.003# 0.001

cnn = architecture_X(M, num_classes)
reduce_lr = ReduceLROnPlateau(monitor = 'val_acc',
                                  factor = 0.7,
                                  patience = 50,
                                  min_lr = 0.00001,verbose=1)
earlystop = EarlyStopping(monitor='val_acc',patience=25,verbose=0,mode='auto')

# print cnn.summary()
history = cnn.fit(x_train, y_train,
                  batch_size=16*2,#640,
                  epochs=30,#300,
                  verbose=1,callbacks = [reduce_lr],
                  validation_data=(x_val, y_val))

print("=" * 5)
print(np.sum(y_test,axis=0))
print(cnn.evaluate(x_test,y_test))
y_pred = cnn.predict(x_test)
y_pred_ = np.argmax(y_pred, axis=1)
y_test_num = np.argmax(y_test, axis=1)
tn, fp, fn, tp = confusion_matrix(y_test_num, y_pred_).ravel()

acc = (tp + tn) * 1. / (tp + fp + tn + fn)

# from scikitplot as skplt
# skplt.metrics.plot_roc_curve(y_test_num,y_pred)

print('='*10)
print('='*5,'  ','Our network')
print('='*10)

print('test: ', acc)

ps = tp*1./(tp+fp)
rc = tp*1./(tp+fn)
print('Accuracy:',(tp+tn)*1./(tp+tn+fp+fn))
print("Pression: ", ps)
print("Recall:", rc)
print("F1: ",2*(ps*rc)/(ps+rc))

fp_rate, tp_rate, thresholds = metrics.roc_curve(y_test_num, np.max(y_pred,axis=-1))
print("auc: ",metrics.auc(fp_rate, tp_rate))

print('='*10)
print('='*5,'  ','LR')
print('='*10)
from sklearn.linear_model import LogisticRegression

from sklearn import svm
X = x_train.reshape((len(x_train), M))
y = y_train.argmax(axis=1).astype('float32')
logisticRegr = LogisticRegression(C=10,max_iter=10)
logisticRegr.fit(X, y)
y_pred = logisticRegr.predict(x_test.reshape((len(x_test), M)))

y_test_num = np.argmax(y_test, axis=1)
tn, fp, fn, tp = confusion_matrix(y_test_num, y_pred).ravel()

acc = (tp + tn) * 1. / (tp + fp + tn + fn)
ps = tp*1./(tp+fp)
rc = tp*1./(tp+fn)
print('Accuracy:',(tp+tn)*1./(tp+tn+fp+fn))
print("Pression: ", ps)
print("Recall:", rc)
print("F1: ",2*(ps*rc)/(ps+rc))


print('='*10)
print('='*5,'  ','SVM')
print('='*10)
from sklearn import svm
X = x_train.reshape((len(x_train), M))
y = y_train.argmax(axis=1).astype('float32')
clf = svm.SVC(gamma=0.001)
clf.fit(X, y) 
y_pred = clf.predict(x_test.reshape((len(x_test), M)))

y_test_num = np.argmax(y_test, axis=1)
tn, fp, fn, tp = confusion_matrix(y_test_num, y_pred).ravel()

acc = (tp + tn) * 1. / (tp + fp + tn + fn)
ps = tp*1./(tp+fp)
rc = tp*1./(tp+fn)
print('Accuracy:',(tp+tn)*1./(tp+tn+fp+fn))
print("Pression: ", ps)
print("Recall:", rc)
print("F1: ",2*(ps*rc)/(ps+rc))
#print("auc: ",roc_auc_score(y_test_num,clf.predict_proba(x_test.reshape((len(x_test), M)))))


print('='*10)
print('='*5,'  ','RF')
print('='*10)
from sklearn.ensemble import RandomForestClassifier, AdaBoostClassifier
X = x_train.reshape((len(x_train), M))
y = y_train.argmax(axis=1).astype('float32')
# Instantiate model with 1000 decision trees
rf = RandomForestClassifier(max_depth=5, n_estimators=1000, max_features=100)
# Train the model on training data
rf.fit(X, y)

y_pred = rf.predict(x_test.reshape((len(x_test), M)))

y_test_num = np.argmax(y_test, axis=1)
tn, fp, fn, tp = confusion_matrix(y_test_num, y_pred).ravel()

acc = (tp + tn) * 1. / (tp + fp + tn + fn)

ps = tp*1./(tp+fp)
rc = tp*1./(tp+fn)
print('Accuracy:',(tp+tn)*1./(tp+tn+fp+fn))
print("Pression: ", ps)
print("Recall:", rc)
print("F1: ",2*(ps*rc)/(ps+rc))

print('='*10)
print('='*5,'  ','AdaBoost')
print('='*10)
from sklearn.tree import DecisionTreeClassifier
X = x_train.reshape((len(x_train), M))
y = y_train.argmax(axis=1).astype('float32')
# Instantiate model with 1000 decision trees
bdt = AdaBoostClassifier(DecisionTreeClassifier(max_depth=3),
                         algorithm="SAMME.R",
                         n_estimators=1000)
# Train the model on training data
bdt.fit(X, y)

y_pred = bdt.predict(x_test.reshape((len(x_test), M)))

y_test_num = np.argmax(y_test, axis=1)
tn, fp, fn, tp = confusion_matrix(y_test_num, y_pred).ravel()

acc = (tp + tn) * 1. / (tp + fp + tn + fn)

ps = tp*1./(tp+fp)
rc = tp*1./(tp+fn)
print('Accuracy:',(tp+tn)*1./(tp+tn+fp+fn))
print("Pression: ", ps)
print("Recall:", rc)
print("F1: ",2*(ps*rc)/(ps+rc))