import sys
import models
import numpy as np
import os
import argparse
from Bio import SeqIO
from torch import optim
import re
import inspect
import torch
#from fastText import train_unsupervised,load_model

def get_roc(y_true, y_pred, cutoff):
    '''    
    :param y_true: 
    :param y_pred: 
    :param cutoff: 
    :return: 
    '''
    score = []
    label = []

    for i in range(y_pred.shape[0]):
        label.append(y_true[i])
        score.append(y_pred[i])

    index = np.argsort(score)
    index = index[::-1]
    t_score = []
    t_label = []
    for i in index:
        t_score.append(score[i])
        t_label.append(label[i])

    score = xroc(t_label, cutoff)
    return  score

def getSorted(dict, flag):
    proteins = []
    scores = []
    for name in dict:
        proteins.append(name)
        scores.append(float(dict[name]))
    index = np.argsort(scores)
    if flag == 1: # sorted descendant
        index = index[::-1]
    t_score = []
    t_protein = []
    for i in index:
        t_score.append(scores[i])
        t_protein.append(proteins[i])
    return (t_score, t_protein)

def getAccession(record):
    record = record.split(' ')
    info = record[1].split('.')
    sfam = info[0] + '.' + info[1] + '.' + info[2]
    return sfam

def getTrueLable(proteins, fam, testing_samples):
    fam = fam.split('.')
    sfam = fam[0] + '.' + fam[1] + '.' + fam[2]
    true_labels = []
    for x in proteins:
        x_fam = getAccession(testing_samples[x].description)
        if cmp(x_fam, sfam) == 0:
            true_labels.append(1)
        else:
            true_labels.append(0)
    return true_labels

def xroc(res, cutoff):
    area, height, fp, tp = 0.0, 0.0, 0.0, 0.0
    for x in res:
        label = x
        if cutoff > fp:
            if label == 1:
                height += 1
                tp += 1
            else:
                area += height
                fp += 1
        else:
            if label == 1:
                tp += 1
    lroc = 0
    if fp != 0 and tp !=0:
        lroc = area / (fp * tp)
    elif fp == 0 and tp != 0:
        lroc = 1
    elif fp != 0 and tp == 0:
        lroc = 0
    return lroc

def getROC(proteins, predicted_labels, label, fam):
    writer = open('results/' + fam + '/result.txt', 'w')
    roc = xroc(label, len(label))
    roc50 = xroc(label, 50)
    writer.write('protein\tpredicted_label\ttrue label\n')
    for i in range(len(proteins)):
        writer.write(proteins[i] + '\t' + str(predicted_labels[proteins[i]])+'\t'+str(label[i]) + '\n')
    writer.write(str(roc) + '\t' + str(roc50) + '\n')
    writer.close()
    return(roc, roc50)

def obtainTestingData(target_fam):    
    #read the testing samples
    pos_test = list(SeqIO.parse('data/pos-test.' + target_fam + '.fasta', 'fasta'))
    neg_test = list(SeqIO.parse('data/neg-test.' + target_fam + '.fasta', 'fasta'))
    pos_test.extend(neg_test)
    test = pos_test
    testing_samples = SeqIO.to_dict(test)
    return testing_samples

def getModelPredictions(q_name):
    #obtain the predictions generated by CNN-BLSTM-PSSM 
    path = 'results/'+target_fam+'/CNN-BLSTM-PSSM/result/' + q_name + '.txt'
    reader = open(path, 'r')
    line = reader.readline()
    line = reader.readline()
    line = line.strip('\n').split('\t')
    return line[-2],line[-1]

def parseArguments(parser):
    
    #parser the input argument 
    parser.add_argument('-e', type=float, default=0.0005, help='the threshold of HHblits')
    parser.add_argument('-family_index', type=str, help = 'family index')
    parser.add_argument('-train', type = bool, default = False, help='train a CNN-BLSTM-PSSM')
    parser.add_argument('-test', type = bool, default = False, help=' load the trained CNN-BLSTM-PSSM model')
    parser.add_argument('-model_dir', type=str, help='the directory of the trained model json file of CNN-BLSTM-PSSM. If test is false, this argument can be empty.')
    parser.add_argument('-weights_dir', type=str, help='the directory of the trained model weights file of CNN-BLSTM-PSSM. If test is false, this argument can be empty.')
    parser.add_argument('-pos_train_dir', type=str,help='the directory of positive training dataset')
    parser.add_argument('-neg_train_dir', type=str,help='the directory of negative training dataset')
    parser.add_argument('-pos_test_dir', type = str,help='the directory of positive testing dataset')
    parser.add_argument('-neg_test_dir', type=str,help='the directory of negative testing dataset')
    
    # training
    parser.add_argument("--n_epochs", type=int, default=100)
    parser.add_argument("--batch_size", type=int, default=64)
    parser.add_argument("--dpout_model", type=float, default=0., help="encoder dropout")
    parser.add_argument("--dpout_fc", type=float, default=0.5, help="classifier dropout")
    parser.add_argument("--nonlinear_fc", type=float, default=1, help="use nonlinearity in fc")
    parser.add_argument("--optimizer", type=str, default="sgd,lr=0.1", help="adam or sgd,lr=0.1")
    parser.add_argument("--lrshrink", type=float, default=5, help="shrink factor for sgd")
    parser.add_argument("--decay", type=float, default=0.99, help="lr decay")
    parser.add_argument("--minlr", type=float, default=1e-5, help="minimum lr")
    parser.add_argument("--max_norm", type=float, default=5., help="max norm (grad clipping)")

    # model
    parser.add_argument("--encoder_type", type=str, default='ConvNetEncoder', help="see list of encoders")
    parser.add_argument("--enc_lstm_dim", type=int, default=20, help="encoder nhid dimension")
    parser.add_argument("--n_enc_layers", type=int, default=1, help="encoder num layers")
    parser.add_argument("--fc_dim", type=int, default=100, help="nhid of fc layers")
    parser.add_argument("--n_classes", type=int, default=2, help="entailment/neutral/contradiction")
    parser.add_argument("--pool_type", type=str, default='max', help="max or mean")

    # gpu
    parser.add_argument("--gpu_id", type=int, default=1, help="GPU ID")
    parser.add_argument("--seed", type=int, default=1234, help="seed")

    #add kmer
    parser.add_argument('-kmer', type = bool, default = False, help='train a CNN-BLSTM-PSSM')

    args = parser.parse_args()
    return args

def get_optimizer(s):
    """
    Parse optimizer parameters.
    Input should be of the form:
        - "sgd,lr=0.01"
        - "adagrad,lr=0.1,lr_decay=0.05"
    """
    if "," in s:
        method = s[:s.find(',')]
        optim_params = {}
        for x in s[s.find(',') + 1:].split(','):
            split = x.split('=')
            assert len(split) == 2
            assert re.match("^[+-]?(\d+(\.\d*)?|\.\d+)$", split[1]) is not None
            optim_params[split[0]] = float(split[1])
    else:
        method = s
        optim_params = {}

    if method == 'adadelta':
        optim_fn = optim.Adadelta
    elif method == 'adagrad':
        optim_fn = optim.Adagrad
    elif method == 'adam':
        optim_fn = optim.Adam
    elif method == 'adamax':
        optim_fn = optim.Adamax
    elif method == 'asgd':
        optim_fn = optim.ASGD
    elif method == 'rmsprop':
        optim_fn = optim.RMSprop
    elif method == 'rprop':
        optim_fn = optim.Rprop
    elif method == 'sgd':
        optim_fn = optim.SGD
        assert 'lr' in optim_params
    else:
        raise Exception('Unknown optimization method: "%s"' % method)

    # check that we give good parameters to the optimizer
    expected_args = inspect.getargspec(optim_fn.__init__)[0]
    assert expected_args[:2] == ['self', 'params']
    if not all(k in expected_args[2:] for k in optim_params.keys()):
        raise Exception('Unexpected parameters: expected "%s", got "%s"' % (
            str(expected_args[2:]), str(optim_params.keys())))

    return optim_fn, optim_params

def get_batch_embedding(batch, index2char):
    # sent in batch in decreasing order of lengths (bsize, max_len, word_dim)
    lengths = np.array([len(x) for x in batch])
    max_len = np.max(lengths)
    embed = np.zeros((max_len, len(batch), 100)) #20 for PSSM, and 300 for word embedding

    f = load_model("bio_protein_embedding.bin")
    words, freq = f.get_words(include_freq=True) # the character embedding

    #we need the reverse index map
    #for w, frequency in zip(words, freq):
    #    try:
    #        print(w + "\t" + str(frequency)+'\t'+ str(f.get_word_vector(w)))
    #    except IOError as e:
    #        if e.errno == errno.EPIPE:
    #            pass

    #it should be the index I guess?
    #we need to reverse the index

    for i in range(len(batch)):
        for j in range(len(batch[i])):
            #embed[j, i, :] = word_vec[batch[i][j]]
            print(index2char[batch[i][j]])
            print(f.get_word_vector(index2char[batch[i][j]]).shape)
            embed[j, i, :] = f.get_word_vector(index2char[batch[i][j]])

    return torch.from_numpy(embed).float(), lengths

#since PSSM no need mapping
def get_batch_pssm(batch):
    # sent in batch in decreasing order of lengths (bsize, max_len, word_dim)
    #print (batch.shape)
    lengths = np.array([len(x) for x in batch])
    max_len = np.max(lengths)
    embed = np.zeros((max_len, len(batch), batch.shape[2])) #20 for PSSM

    for i in range(len(batch)):
        for j in range(len(batch[i])):
            embed[j, i, :] = batch[i][j]

    return torch.from_numpy(embed).float(), lengths

#one hot encoding
def get_batch_one_hot(batch):
    # sent in batch in decreasing order of lengths (bsize, max_len, word_dim)
    lengths = np.array([len(x) for x in batch])
    max_len = np.max(lengths)
    embed = np.zeros((max_len, len(batch), 20)) #20 for PSSM

    for i in range(len(batch)):
        for j in range(len(batch[i])):
            embed[j, i, :] = batch[i][j] #to do modify here

    return torch.from_numpy(embed).float(), lengths