import numpy as np
import argparse
import tensorflow as tf
import matplotlib.pyplot as plt
from sklearn.utils import shuffle
from sklearn.model_selection import StratifiedKFold
import math
from datetime import datetime

NUM_LABELS = 47
rnd = np.random.RandomState(123)
tf.set_random_seed(123)

#hyper parameters
learning_rate = 0.001
momentum = 0.9
BATCH_SIZE = 256

learning_rate_tune = [ 0.01, 0.001,0.0001]
batch_size_tune = [32, 256, 512]
momentum_factor_tune = [0.1, 0.5, 0.9]

#settings
NUM_ITERATIONS= 50
STEP_VALIDATE_EPOCH=1 # check performance on the validation set every EPOCH_VALIDATE_STEP epochs
NUM_FOLD = 5

#MODEL_NAME = ""
CNN_MODEL_PATH = "../checkpoints/cnn/cnn_model"

# Following functions are helper functions that you can feel free to change
def convert_image_data_to_float(image_raw):
    #[None, 28, 28] -> [None, 28, 28, 1] and normalize 255 to 1
    img_float = tf.expand_dims(tf.cast(image_raw, tf.float32) / 255, axis=-1)
    return img_float

def visualize_ae(i, x, features, reconstructed_image):
    '''
    This might be helpful for visualizing your autoencoder outputs
    :param i: index
    :param x: original data
    :param features: feature maps
    :param reconstructed_image: autoencoder output
    :return:
    '''
    plt.figure("input img " + str(i))
    plt.imshow(x[i, :, :], cmap="gray")
    plt.title("input image", fontsize=16)
    plt.figure("reconstructed image"+str(i))
    plt.imshow(reconstructed_image[i, :, :, 0], cmap="gray")
    plt.title("reconsturcted image", fontsize=16)
    plt.figure("encoded features "+str(i))
    plt.imshow(np.reshape(features[i, :, :, :], (7, -1), order="F"), cmap="gray",)
    plt.title("encoded features", fontsize=16)

def build_cae_model(placeholder_x):
    with tf.variable_scope("cae") as scope:
        shapes=[]
        #input layer
        img_float = convert_image_data_to_float(placeholder_x)

        shapes.append(img_float.get_shape().as_list()) #shape[0], input shape of conv1

        #conv1
        w_conv1 = tf.get_variable("conv1_weight", shape=(5, 5, 1, 32),
                                  initializer=tf.contrib.layers.xavier_initializer())
        b_conv1 = tf.get_variable("conv1_bias", shape=(32), initializer=tf.contrib.layers.xavier_initializer())
        y_conv1 = tf.nn.conv2d(img_float, w_conv1, strides=[1, 2, 2, 1], padding='SAME') + b_conv1
        conv1= tf.nn.relu(y_conv1)
        shapes.append(conv1.get_shape().as_list()) #shape[1]

        #conv2
        w_conv2 = tf.get_variable("conv2_weight", shape=(5,5,32,64), initializer=tf.contrib.layers.xavier_initializer())
        b_conv2 = tf.get_variable("conv2_bias", shape=(64), initializer=tf.contrib.layers.xavier_initializer())
        y_conv2 = tf.nn.conv2d(conv1, w_conv2, strides=[1,2,2,1], padding='SAME') + b_conv2
        conv2 = tf.nn.relu(y_conv2)
        shapes.append(conv2.get_shape().as_list())  #shape[2]

        #conv3
        w_conv3 = tf.get_variable("conv3_weight", shape=(3,3,64,2), initializer=tf.contrib.layers.xavier_initializer())
        b_conv3 = tf.get_variable("conv3_bias", shape=(2), initializer=tf.contrib.layers.xavier_initializer())
        y_conv3 = tf.nn.conv2d(conv2, w_conv3, strides=[1,1,1,1], padding='SAME') + b_conv3
        conv3 = tf.nn.relu(y_conv3)

        print("encoder layer shape : %s" % conv3.get_shape())
        encode_result = conv3
        #shapes.reverse()

        #deconv1
        w_deconv1 = tf.get_variable("deconv1_weight", shape=(3, 3, 64, 2), initializer=tf.contrib.layers.xavier_initializer())
        b_deconv1 = tf.get_variable("deconv1_bias", shape=(64), initializer=tf.contrib.layers.xavier_initializer())
        shape=shapes[2]
        output_shape = tf.stack([tf.shape(img_float)[0], shape[1], shape[2], shape[3]]) #[batch_size, size_of_output_feature_map_of_conv2 x y, conv2_chs]
        deconv1 = tf.nn.relu( tf.nn.conv2d_transpose(conv3, w_deconv1, output_shape, strides=[1,1,1,1], padding='SAME') + b_deconv1 ) #TODO: strides

        #deconv2
        w_deconv2 = tf.get_variable("deconv2_weight", shape=(5, 5, 32, 64), initializer=tf.contrib.layers.xavier_initializer())
        b_deconv2 = tf.get_variable("deconv2_bias", shape=(32), initializer=tf.contrib.layers.xavier_initializer())
        shape=shapes[1]
        output_shape = tf.stack([tf.shape(img_float)[0], shape[1], shape[2], shape[3]]) #[batch_size, size_of_output_feature_map1_conv1 x y, conv1_chs]
        deconv2 = tf.nn.relu( tf.nn.conv2d_transpose(deconv1, w_deconv2, output_shape, strides=[1,2,2,1], padding='SAME') + b_deconv2 )

        #deconv3
        w_deconv3 = tf.get_variable("deconv3_weight", shape=(5, 5, 1, 32), initializer=tf.contrib.layers.xavier_initializer())
        b_deconv3 = tf.get_variable("deconv3_bias", shape=(1), initializer=tf.contrib.layers.xavier_initializer())
        shape=shapes[0]
        output_shape = tf.stack([tf.shape(img_float)[0], shape[1], shape[2], shape[3]]) #[batch_size, 28, 28, 1]
        deconv3 = tf.nn.relu( tf.nn.conv2d_transpose(deconv2, w_deconv3, output_shape, strides=[1,2,2,1], padding='SAME') + b_deconv3 )

        img_reconstructed = deconv3

        print("reconstruct layer shape : %s" % img_reconstructed.get_shape())

        #loss
        loss = tf.reduce_mean(tf.square(img_reconstructed - img_float))
        optimizer = tf.train.AdamOptimizer(learning_rate)
        #optimizer = tf.train.MomentumOptimizer(learning_rate, momentum) # doesn't work for cae model
        train_op = optimizer.minimize(loss,global_step = tf.train.get_global_step())

        return train_op, loss, encode_result, img_reconstructed

def build_cnn_model(placeholder_x, placeholder_y):
    with tf.variable_scope("cnn") as scope:

        img_float = convert_image_data_to_float(placeholder_x)
        #input_layer = tf.reshape(img_float, [-1, 28, 28, 1])
        #conv1
        w_conv1 = tf.get_variable("conv1_weight", shape=(3, 3, 1, 32),
                                  initializer=tf.contrib.layers.xavier_initializer())
        b_conv1 = tf.get_variable("conv1_bias", shape=(32), initializer=tf.contrib.layers.xavier_initializer())
        y_conv1 = tf.nn.conv2d(img_float, w_conv1, strides=[1, 1, 1, 1], padding='SAME') + b_conv1
        conv1= tf.nn.relu(y_conv1)
        #conv2
        w_conv2 = tf.get_variable("conv2_weight", shape=(5,5,32,32), initializer=tf.contrib.layers.xavier_initializer())
        b_conv2 = tf.get_variable("conv2_bias", shape=(32), initializer=tf.contrib.layers.xavier_initializer())
        y_conv2 = tf.nn.conv2d(conv1, w_conv2, strides=[1,2,2,1], padding='SAME') + b_conv2
        conv2 = tf.nn.relu(y_conv2)
        #conv3
        w_conv3 = tf.get_variable("conv3_weight", shape=(3,3,32,64), initializer=tf.contrib.layers.xavier_initializer())
        b_conv3 = tf.get_variable("conv3_bias", shape=(64), initializer=tf.contrib.layers.xavier_initializer())
        y_conv3 = tf.nn.conv2d(conv2, w_conv3, strides=[1,1,1,1], padding='SAME') + b_conv3
        conv3 = tf.nn.relu(y_conv3)
        #conv 4
        w_conv4 = tf.get_variable("conv4_weight", shape=(5,5,64,64), initializer=tf.contrib.layers.xavier_initializer())
        b_conv4 = tf.get_variable("conv4_bias", shape=(64), initializer=tf.contrib.layers.xavier_initializer())
        y_conv4 = tf.nn.conv2d(conv3, w_conv4, strides=[1,2,2,1], padding='SAME') + b_conv4
        conv4 = tf.nn.relu(y_conv4)

        # Flatten
        features_flattened = tf.reshape(conv4, [-1, np.prod(conv4.shape[1:])])

        # FC Layer
        w_fc = tf.get_variable("fc_weight", shape=(features_flattened.shape[1], NUM_LABELS),
                                 initializer=tf.contrib.layers.xavier_initializer())
        logits = tf.matmul(features_flattened, w_fc)
        # loss
        loss = tf.losses.sparse_softmax_cross_entropy(labels=placeholder_y, logits=logits)

        params = [w_conv1, b_conv1, w_conv2, b_conv2, w_conv3, b_conv3, w_conv4, b_conv4, w_fc]

        #optimizeation, SGD with momentum, using learning rate decay after x epoch
        #test 1: simple high-level simple SGD with momentum, learning rate fixed in each step.
        '''
        optimizer = tf.train.MomentumOptimizer(learning_rate,momentum)
        train_op = optimizer.minimize(loss, global_step = tf.train.get_global_step())
        '''
        #test 2: self-implement
        train_op = []
        v = []
        for i,param in enumerate(params):
            grad = tf.gradients(loss, param)[0]
            v.append(tf.Variable(tf.zeros(param.shape), dtype=tf.float32))
            v_temp = momentum * v[i] - learning_rate * grad
            train_op.append( tf.assign_add(param, v_temp) )
            train_op.append( tf.assign(v[i], v_temp))

        #calc accuracy
        predictions = tf.argmax(logits, 1)
        one_hot_y = tf.one_hot(placeholder_y, NUM_LABELS)
        correct_prediction = tf.equal(predictions, tf.argmax(one_hot_y, 1))
        correct_cnt =  tf.reduce_sum(tf.cast(correct_prediction, tf.int32)) #tf.reduce_mean(tf.cast(correct_prediction, tf.float31))
        #acc, acc_op = tf.metrics.accuracy(placeholder_y, predictions) #TODO: this API seems to give a wrong answer in debugging.

        return params, train_op, loss, correct_cnt,predictions

def build_linear_model(placeholder_x,placeholder_y):
    with tf.variable_scope("linear") as scope:
        img_float = convert_image_data_to_float(placeholder_x)

        # This is a simple fully connected network
        img_flattened = tf.reshape(img_float,[-1,np.prod(placeholder_x.shape[1:])])
        weight = tf.get_variable("fc_weight",shape=(img_flattened.shape[1],NUM_LABELS),
                                 initializer=tf.random_normal_initializer(stddev=0.01))
        logits = tf.matmul(img_flattened, weight)
        loss = tf.losses.sparse_softmax_cross_entropy(
            labels=placeholder_y, logits=logits)

        # gradient decent algorithm
        params = [weight]
        learning_rate = 0.001
        grad = tf.gradients(loss, weight)[0]
        train_op = tf.assign_add(weight, -learning_rate * grad)

        #calc accuracy
        predictions = tf.argmax(logits, 1)
        one_hot_y = tf.one_hot(placeholder_y, NUM_LABELS)
        correct_prediction = tf.equal(predictions, tf.argmax(one_hot_y, 1))
        correct_cnt =  tf.reduce_sum(tf.cast(correct_prediction, tf.int32)) #tf.reduce_mean(tf.cast(correct_prediction, tf.float31))
        #acc, acc_op = tf.metrics.accuracy(placeholder_y, predictions)

    return params, train_op, loss, correct_cnt, predictions

test_full_training_set = False

CAE_MODEL_PATH = "../checkpoints/cae/cae_model"
#lr_cae = 0.001
#bs_cae = 256

def train_ae(x, placeholder_x, x_evualate = None):
    # TODO: implement autoencoder training
    train_op, loss, encode_result, img_reconstructed = build_cae_model(placeholder_x)
    cae_saver = tf.train.Saver(max_to_keep=50)
    #x = x[0:512]
    NUM_ITERATIONS = 150    #it takes more time to converge
    NUM_BATCHES = int(math.ceil(x.shape[0]/BATCH_SIZE)) #ceil, make sure to ultilize all the data. Size of each fold may not perfectly align
    loss_changes = np.zeros(NUM_ITERATIONS)

    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        sess.run(tf.local_variables_initializer())
        start_time = datetime.now()
        print("---------------------{} Training CAE -------------------".format(start_time))
        print("Num samples to train: {}, num_batches: {}".format(x.shape[0], NUM_BATCHES))
        for epoch in range(NUM_ITERATIONS):
            # 1) sample-2: shuffle with fixed random seed in each epoch and split into batches. TODO: it's not efficiency to copy data in each epoch. indexing is more efficient.
            shuffledX = shuffle(x,  random_state=epoch)
            loss_val= 0.0; acc_val= 0.0; num_trained_sample = 0; end_time=[]

            for bi  in range(NUM_BATCHES):
                batch_x = shuffledX[bi * BATCH_SIZE : (bi+1)*BATCH_SIZE]
                feed_dict = {placeholder_x: batch_x}
                #2) training
                _, loss_batch, feature_maps, img_rec = sess.run([train_op, loss, encode_result, img_reconstructed],  feed_dict=feed_dict)
                loss_val += loss_batch
            loss_val = loss_val/NUM_BATCHES
            end_time.append(datetime.now())
            print("{}: Epoch {} finished, Training loss: {:.6f}".format(end_time[-1], epoch, loss_val))

            # validation
            if(( (epoch+1) % STEP_VALIDATE_EPOCH== 0) or (epoch+1 == NUM_ITERATIONS)):
                loss_eva, feature_maps, img_rec = sess.run([loss, encode_result, img_reconstructed], feed_dict={placeholder_x:x_evualate})
                loss_changes[epoch] = loss_eva
                print("**** Loss in evaluation set: {:.6f}".format(loss_eva))

            cae_saver.save(sess=sess, save_path=CAE_MODEL_PATH, global_step=epoch)

        print("\r\nFinish Training under setting: lr={} batch_size={}".format(learning_rate, BATCH_SIZE))
        print("==> Training time consumed: {}".format(end_time[-1] - start_time))
        print("==> Training loss: {:.6f}".format(loss_val))
        print("==> Best evaluation loss: {:.6f} in epoch {}" .format(np.min(loss_changes),np.argmin(loss_changes)))

def evaluate_ae(x,placeholder_x):
    # TODO: evaluate your autoencoder

    train_op, loss, encode_result, img_reconstructed = build_cae_model(placeholder_x)
    saver = tf.train.Saver()
    with tf.Session() as sess:
        #saver.restore(sess,tf.train.latest_checkpoint("../checkpoints/well_trained_cae"))
        saver.restore(sess, "../checkpoints/well_trained_cae/cae_model-140") # select the model generated in epoch 3

        loss_eva, feature_maps, img_rec = sess.run([loss, encode_result, img_reconstructed], feed_dict={placeholder_x:x})
        print("==>  Loss in evaluation set: {:.6f}".format(loss_eva))
        #idxs=[1,10,100] # random
        idxs = np.floor(np.random.random(2) * x.shape[0]).astype(int)
        for i in idxs:
            visualize_ae(i, x, feature_maps, img_rec)
        plt.show()
    # show feature map and reconstructed images

# Major interfaces
def train_cnn(x, y, placeholder_x, placeholder_y, cross_validate=False):
    #TODO: reduce size for debug
    #x = x[0:512]
    #y = y[0:512]

    params, train_op, loss, correct_cnt, predictions = build_cnn_model(placeholder_x, placeholder_y)
    cnn_saver = tf.train.Saver(max_to_keep=40) #(var_list=params)

    # 1) sample-1: for cross-validation, split into 5-fold.
    skf = StratifiedKFold(n_splits=NUM_FOLD, random_state=10, shuffle=True)
    fold_idx = 0
    if(cross_validate==True):
        fold_settings = list(skf.split(x,y)) # the index settings of the 5 folds
    else:
        fold_settings = list(skf.split(x,y))[0:1] #pick the first fold
    #set up variables to record performance change
    loss_train_changes = np.zeros((len(fold_settings), NUM_ITERATIONS), dtype = np.float32) #np.zeros((len(fold_settings), int(NUM_ITERATIONS/STEP_VALIDATE_EPOCH)), dtype = np.float32)
    acc_train_changes = np.zeros((len(fold_settings), NUM_ITERATIONS),  dtype = np.float32)
    loss_validate_changes = np.zeros((len(fold_settings), NUM_ITERATIONS),  dtype = np.float32)
    acc_validate_changes = np.zeros((len(fold_settings), NUM_ITERATIONS),  dtype = np.float32)

    print("\r\nStart training\r\nParams setting: learning_rate={}, batch size={}, momentum={}".format(learning_rate, BATCH_SIZE, momentum))
    start_time_cross = datetime.now()
    for fold_idx, [train_index, validation_index] in enumerate(fold_settings):
        #print("TRAIN:", train_index, len(train_index),  "TEST:", validation_index, len(validation_index))
        x_train, y_train = x[train_index],y[train_index]
        x_validate, y_validate = x[validation_index],y[validation_index]

        if(test_full_training_set==True):
            x_train = x
            y_train = y

        NUM_BATCHES = int(math.ceil(x_train.shape[0]/BATCH_SIZE)) #ceil, make sure to ultilize all the data. Size of each fold may not perfectly align

        with tf.Session() as sess:
            sess.run(tf.global_variables_initializer())
            sess.run(tf.local_variables_initializer())
            start_time = datetime.now()
            print("---------------------{} Training Fold {} -------------------".format(start_time, fold_idx))
            print("Num samples to train: {}, num_batches: {}".format(x_train.shape[0], NUM_BATCHES))
            for epoch in range(NUM_ITERATIONS):
                # 1) sample-2: shuffle with fixed random seed in each epoch and split into batches. TODO: it's not efficiency to copy data in each epoch. indexing is more efficient.
                shuffledX, shuffledY = shuffle(x_train, y_train, random_state=epoch)
                loss_val= 0.0; acc_val= 0.0; num_trained_sample = 0; end_time=[]

                for bi  in range(NUM_BATCHES):
                    batch_x = shuffledX[bi * BATCH_SIZE : (bi+1)*BATCH_SIZE]
                    batch_y = shuffledY[bi * BATCH_SIZE : (bi+1)*BATCH_SIZE]
                    feed_dict = {placeholder_x: batch_x, placeholder_y: batch_y}
                    #2) training
                    _,loss_batch, correct_batch , pred_val = sess.run([train_op, loss, correct_cnt, predictions],  feed_dict=feed_dict)

                    loss_val += loss_batch # or loss.eval(feed_dict = feed_dict)
                    acc_val += correct_batch # add up the number of correct predictions in each batch
                    num_trained_sample += batch_x.shape[0] #secure, since it may not equal to NUM_BATCHES * BATCH_SIZE
                    #print("debug: ",  batch_y, pred_val, acc_batch)

                #report performance in training set
                end_time.append(datetime.now())
                acc_val = acc_val/num_trained_sample
                loss_val = loss_val/NUM_BATCHES
                print("{}: Epoch {} finished, Training loss: {:.4f}, acc: {:.4f}".format(end_time[-1], epoch, loss_val, acc_val))

                #store the performance and the model
                loss_train_changes[fold_idx][epoch] = loss_val
                acc_train_changes[fold_idx][epoch] = acc_val
                if(epoch>3):
                    cnn_saver.save(sess=sess, save_path=CNN_MODEL_PATH, global_step=epoch)

                # 3) validate and save every 5 epoches
                if(( (epoch+1) % STEP_VALIDATE_EPOCH== 0) or (epoch+1 == NUM_ITERATIONS)):
                    #NUM_BATCHES_VALIDATE = int(math.ceil(x_train.shape[0]/BATCH_SIZE))
                    #for bi in range(NUM_BATCHES_VALIDATE):
                    if(test_full_training_set==True): #validate directly on the test set
                        x_validate, y_validate = get_test_set()
                    feed_dict_validate={placeholder_x: x_validate, placeholder_y: y_validate}
                    loss_validate, correct_cnt_validate  = sess.run([loss, correct_cnt],  feed_dict=feed_dict_validate)
                    acc_validate = correct_cnt_validate/x_validate.shape[0]
                    print("**** Validation loss: {:.4f}, acc: {:.4f}".format(loss_validate, acc_validate))

                    loss_validate_changes[fold_idx][epoch] = loss_validate
                    acc_validate_changes[fold_idx][epoch] = acc_validate

            print("==> Training time consumed: {}".format(end_time[-1] - start_time))
            print("==> Best acc: {} in epoch {}" .format(np.max(acc_validate_changes[fold_idx]),np.argmax(acc_validate_changes[fold_idx])))

    end_time_cross = datetime.now()
    # comput the performance of the training set
    print("\r\nFinish training in {} \r\nParams setting: learning_rate={}, batch size={}, momentum={}".format(end_time_cross - start_time_cross, learning_rate, BATCH_SIZE, momentum))
    #best_epoch = np.argmax(acc_train_changs[])
    print("==> Loss on training set of fold 0: {:.4f}".format(loss_train_changes[0][-1]))
    print("==> Accuarcy on training set of fold 0: {:.4f}".format(acc_train_changes[0][-1]))
    print("==> Loss on validation set of fold 0: {:.4f}".format(loss_validate_changes[0][-1]))
    print("==> Accuarcy on valiationc set of fold 0: {:.4f}".format(acc_validate_changes[0][-1]))

    # average performance on 5-folds
    if(cross_validate==True):
        #loss_folds = loss_validate_changes[:, -1] #TODO: acc in last epoch may not be the best and stable enough to represent the performance in current fold. Use average filter.
        loss_folds = np.min(loss_validate_changes, axis=1)
        # acc of each validatation set in  k-folds
        #acc_folds = acc_validate_changes[:,-1]
        acc_folds = np.max(acc_validate_changes, axis=1)
        print("\r\n==> {}-fold cross_validation result (best acc among each fold): {} \r\n==>Average acc over 5-fold: {:.4f} ".format(NUM_FOLD, acc_folds, np.mean(acc_folds)))
        print("==>Average loss over 5-fold: {:.4f}".format(np.mean(loss_folds)))
    # store the performance change over time into npz
    hyper_param_str = "lr{}bs{}mf{}".format(learning_rate, BATCH_SIZE, momentum)
    np.savez_compressed('./performance_change_'+hyper_param_str+'.npz',
                            loss_train = loss_train_changes, acc_train = acc_train_changes,
                            loss_validate = loss_validate_changes, acc_validate = acc_validate_changes)
    verbose=True
    if(verbose==True):
        # plot performance change over time for different param settings
        print("\r\nPerformace changes over time under: learning_rate={}, batch size={}, momentum={}".format(learning_rate, BATCH_SIZE, momentum))
        print("loss changes in training set:\r\n",loss_train_changes[0])
        print("acc changes in training set:\r\n",acc_train_changes[0])
        print("loss changes (every {} epochs) in validation set:\r\n{}".format(STEP_VALIDATE_EPOCH, loss_validate_changes[0]))
        print("acc changes (every {} epochs) in validation set:\r\n{}".format(STEP_VALIDATE_EPOCH, acc_validate_changes[0]))

def get_test_set():
    file_test = np.load("../datasets/data_classifier_test.npz")
    x_test = file_test["x_test"]
    y_test = file_test["y_test"]

    return x_test, y_test

def test_cnn(x, y, placeholder_x, placeholder_y):
    # TODO: implement CNN testing

    params, train_op, loss, correct_cnt, predictions = build_cnn_model(placeholder_x, placeholder_y)
    with tf.Session() as sess:
        #saver = tf.train.import_meta_graph(CNN_MODEL_PATH + "-4.meta")
        saver = tf.train.Saver()
        #saver.restore(sess,tf.train.latest_checkpoint("../checkpoints/well_trained"))
        saver.restore(sess, "../checkpoints/well_trained/cnn_model-12") # select the model generated in epoch 3
        #print(tf.train.latest_checkpoint("../checkpoints/"))

        NUM_BATCHES = int(math.ceil(x.shape[0]/BATCH_SIZE))
        loss_val = 0; acc_val=0; num_trained_sample=0
        feed_dict = {placeholder_x: x, placeholder_y: y}
        result_loss, correct_val = sess.run([loss, correct_cnt],feed_dict=feed_dict)
        result_accuracy = correct_val / x.shape[0]
        print("\r\nPerformance on test set.")
        print("==> Loss: {:.4f}, Acc: {:.4f}".format(result_loss, result_accuracy))

        #TODO: ensemble the 5 model of the 5-fold
        '''
        for bi in range(NUM_BATCHES):
            batch_x = x[bi * BATCH_SIZE : (bi+1)*BATCH_SIZE]
            batch_y = y[bi * BATCH_SIZE : (bi+1)*BATCH_SIZE]
            feed_dict = {placeholder_x: batch_x, placeholder_y: batch_y}

            loss_cur, acc_cur = sess.run([loss, acc_op],feed_dict=feed_dict)

            loss_val += loss_cur
            acc_val += acc_cur
            num_trained_sample += batch_x.shape[0]

        result_accuracy = acc_val/num_trained_sample
        '''
    return result_accuracy

def main():
    global BATCH_SIZE, learning_rate, momentum

    parser = argparse.ArgumentParser(description='COMP5212 Programming Project 2')
    parser.add_argument('--task', default="train_cnn", type=str,
                        help='Select the task, train_cnn, test_cnn, '
                             'train_ae, evaluate_ae, ')
    parser.add_argument('--datapath',default="../datasets",type=str, required=False, help='Select the path to the data directory')
    parser.add_argument('--cross_validate',default=1,type=int, required=False, help='Set 1 to use cross validation for tuning')
    parser.add_argument('--lr',default=0.001,type=float, required=False, help='learning rate')
    parser.add_argument('--momentum',default=0.9,type=float, required=False, help='momentum facotr of SGD optimizer')
    parser.add_argument('--batch_size',default=256,type=int, required=False, help='batch size')
    args = parser.parse_args()

    datapath = args.datapath
    flag_cross_validate = args.cross_validate
    BATCH_SIZE = args.batch_size
    learning_rate = args.lr
    momentum = args.momentum

    # set up the input image size, create the correspond place holder.
    with tf.variable_scope("placeholders"):
        img_var = tf.placeholder(tf.uint8, shape=(None, 28, 28), name="img")
        label_var = tf.placeholder(tf.int32, shape=(None,), name="true_label")

    if args.task == "train_cnn":
        file_train = np.load(datapath+"/data_classifier_train.npz")
        x_train = file_train["x_train"]
        y_train = file_train["y_train"]
        train_cnn(x_train, y_train, img_var, label_var, cross_validate=flag_cross_validate)

    elif args.task == "test_cnn":
        file_test = np.load(datapath+"/data_classifier_test.npz")
        x_test = file_test["x_test"]
        y_test = file_test["y_test"]
        accuracy = test_cnn(x_test, y_test,img_var,label_var)
        print("accuracy = {}\n".format(accuracy))
    elif args.task == "train_ae":
        file_unsupervised = np.load(datapath + "/data_autoencoder_train.npz")
        x_ae_train = file_unsupervised["x_ae_train"]
        file_unsupervised = np.load(datapath + "/data_autoencoder_eval.npz")
        x_ae_eval = file_unsupervised["x_ae_eval"]
        train_ae(x_ae_train, img_var, x_evualate = x_ae_eval)
    elif args.task == "evaluate_ae":
        file_unsupervised = np.load(datapath + "/data_autoencoder_eval.npz")
        x_ae_eval = file_unsupervised["x_ae_eval"]
        evaluate_ae(x_ae_eval, img_var)

if __name__ == "__main__":
    main()


import torch
import torchvision
import torch.nn as nn
import numpy as np

from problog.logic import Var
from utils.df_dpl_utils import neural_predicate
from DeepProbLog.train import train_model2
from DeepProbLog.data_loader import load
from DeepProbLog.model import Model
from DeepProbLog.optimizer import Optimizer
from DeepProbLog.network import Network
from metrics.metric_recording import save_data, clear_file, save_cm

#### NET

# Wardrobe specific parameters:
num_ftrs_wardrobe = 7
nr_output = 3
num_outputs_cm = 8

net = torchvision.models.resnet50(pretrained=True)

# Freeze feature extraction weights to speed up training (these parameters will not be changed during back propagation)
for param in net.parameters():
    param.requires_grad = False

# Parameters of newly constructed modules have requires_grad=True by default
num_ftrs_resnet = net.fc.in_features
# net.fc = nn.Linear(num_ftrs_resnet,num_ftrs_wardrobe)
net.fc = nn.Sequential(
            nn.Linear(num_ftrs_resnet, num_ftrs_wardrobe),
            # nn.ReLU(),
            # nn.Linear(120, 84),
            # nn.ReLU(),
            # nn.Linear(84, num_ftrs_wardrobe),
            nn.Softmax(1)
            )

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)
net.to(device)

clear_file("advanced_outfit_dist.txt")
clear_file("advanced_outfit_acc.txt")
clear_file("advanced_outfit_F1.txt")

def labelVector(bits) :
    return 1*int(bits[2]) + 2*int(bits[1]) + 4*int(bits[0])

labelMap={0:[0,0,0], 1:[0,0,1],2:[0,1,0],3:[0,1,1],4:[1,0,0],5:[1,0,1],6:[1,1,0],7:[1,1,1]}

def hamming_dist(l, c) :
    bl = labelMap[l]
    bc = labelMap[c]
    distance = 0
    for index in range(len(bl)):
        distance += abs(bc[index] - bl[index])
    return distance

def test(model,iteration):
    net.eval()
    for param in net.fc.parameters() :
        param.requires_grad = False
    n=0
    correct = 0
    N = len(test_queries)
    confusion = np.zeros((num_outputs_cm, num_outputs_cm), dtype=np.uint32)  # First index actual, second index predicted
    total_distance=0
    for d in test_queries:
        args = list(d.args)
        label = args[-nr_output:]
        args[-nr_output:] = [Var('X_{}'.format(i)) for i in range(nr_output)]
        q = d(*args)
        out = model.solve(q, None, True)
        out = max(out, key=lambda x: out[x][0])
        if out == d:
            correct += 1
        else:
            n += 1
        l = labelVector(label)
        c = labelVector(list(out.args)[-nr_output:])
        confusion[l, c] += 1
        total_distance += hamming_dist(l, c)

    save_cm(confusion,"advanced_outfit_cm.txt")
    print(confusion)
    F1 = 0
    for nr in range(num_outputs_cm):
            TP = confusion[nr, nr]
            FP = sum(confusion[:, nr]) - TP
            FN = sum(confusion[nr, :]) - TP
            F1 += 2 * TP / (2 * TP + FP + FN) * (FN + TP) / N

    acc = correct / N
    print("Acc : " + str(acc))
    avg_distance = total_distance / n
    save_data(iteration,avg_distance,"advanced_outfit_dist.txt")
    save_data(iteration,acc,"advanced_outfit_acc.txt")
    save_data(iteration,F1,"advanced_outfit_F1.txt")

    for param in net.fc.parameters() :
        param.requires_grad = True
    net.train()
    return

#### DPL

pl_file_path = 'pl_files/advanced_outfit.pl'

train_queries = load('data/generated_data/train_advanced_outfit_data.txt')
test_queries = load('data/generated_data/test_advanced_outfit_data.txt')

with open(pl_file_path) as f:
    problog_string = f.read()

net.optimizer = torch.optim.Adam(net.fc.parameters(), lr=0.001)

netwrk = Network(net,'fashion_df_net', neural_predicate)

model = Model(problog_string,[netwrk],caching=False)
optimizer = Optimizer(model,2)

train_model2(model,train_queries,3,optimizer,test_iter=5000,test=test,snapshot_iter=10000)


# Import Meteostat library and dependencies

from datetime import datetime
from dateutil.deltatime import deltatime
from meteostat import Point, Daily
import localpackage.modified_ozone as ooo
import pandas as pd
import localpackage.fetch_city as fc
from citipy import citipy

#API key AQICN
o3 = ooo.Ozone('d8a099d2492a444219a00713bddafb693699c53f')

#obtener las coordenadas elegidas en tableau en el point location
def obtain_params():
    sheet_id = '1rjx1D0tLXRb3c01ezuaQwfi2LAq5ybR1R1el9V1anMM'
    sheet_name = 'areaverde'
    url = f'https://docs.google.com/spreadsheets/d/{sheet_id}/gviz/tq?tqx=out:csv&sheet={sheet_name}'
    a=pd.read_csv(url)
    lat=a['lat'][0]
    lon=a['lon'][0]
    especiedict={'area': a['area'][1],a['especie1'][1]:a['especie1'][2],a['especie2'][1]:a['especie2'][2],a['especie3'][1]:a['especie3'][2]}
    return lat, lon, especiedict
def point(latitud, longitud):
    location = Point(latitud, longitud)
    city=o3.get_city_id(latitud,longitud)
    return location, city
#obtener la poblacion
def population_data(latitud, longitud):
    city = citipy.nearest_city(latitud, longitud).city_name
    if city== 'new york':
        city='New York City'
    data=fc.obtain_data(city)[0]
    return data
#obtener la fecha actual
def date():
    end = datetime.today()
    end = end.replace(hour=0, minute=0, second=0, microsecond=0)
    start = end- relativedelta(years=1)
    return start, end
#obtener datos diarios del último año
def last_year_weather(latitud,longitud):
    start=date()[0]
    end=date()[1]
    location = point(latitud,longitud)[0]
    data = Daily(location, start, end)
    data = data.fetch().drop(['tmin','tmax','snow','wdir','wpgt','pres','tsun'], axis=1)
    data=data[['wspd','prcp','tavg']]
    return data

def historical_Ct(latitud,longitud):
    city=point(latitud, longitud)[1]
    data = o3.get_historical_data(city_id=city)
    return data
def get_data(latitud,longitud):
    df=last_year_weather(latitud,longitud)
    start=date()[0]
    end=date()[1]
    rawdata=historical_Ct(latitud,longitud).fillna(method='ffill').fillna(method='bfill')
    Ct=rawdata['pm25']
    data=pd.concat([df,Ct],axis=1).fillna(method='ffill').fillna(method='bfill')[start:end]
    data=data[['wspd','prcp','pm25','tavg']]
    return data
def get_city_emission(latitud,longitud):
    city= list(population_data(latitud, longitud))
    contrydict=pd.read_csv('localpackage/Emission_inventory.csv', sep=';').set_index('Pais').T.to_dict('list')
    dict={}
    name=city[0]
    Chab=float(city[1])
    pais=city[4]
    Et=contrydict[pais][0]
    Et=float(Et.replace(',','.'))
    hab=float(city[5])
    Ehab=Et/hab
    Ecity=Ehab*Chab
    dict[name]=Ecity
    return dict, name


"""
给定一个字符串，找出不含有重复字符的最长子串的长度。
定 "abcabcbb" ，没有重复字符的最长子串是 "abc" ，那么长度就是3。
给定 "bbbbb" ，最长的子串就是 "b" ，长度是1。
给定 "pwwkew" ，最长子串是 "wke" ，长度是3。请注意答案必须是一个子串，"pwke" 是 子序列  而不是子串。
"""

def longest_substing(s):
    a, b = '', 0
    for i in range(len(s)):
        if i == 0:
            a += s[i]
            continue
        if s[i] not in a:
            a += s[i]
        else:
            if len(a) > b:
                b = len(a)
            a = s[i]
    return b


from numpy import packbits
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.modules import conv
from torchvision import models

class PartialConv(nn.Module):
    def weight_init(self, m):
        classname = m.__class__.__name__    ##获取当前结构名称
        if classname.find('Conv') != -1:
            nn.init.normal_(m.weight, 0.0, 0.002)  ##正态分布初始化
        elif classname.find('BatchNorm') != -1:
            nn.init.normal_(m.weight, 1.0, 0.02)   ##正态分布初始化
            nn.init.zeros_(m.bias)

    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
                padding=0,dilation = 1, group=1, bias = False ):
        super(PartialConv, self).__init__()
        self.input_conv = nn.Conv2d(in_channels, out_channels, kernel_size,
                                    stride, padding, dilation, group, bias)
        self.mask_conv = nn.Conv2d(in_channels, out_channels, kernel_size,
                                        stride, padding, dilation, group, False)
        torch.nn.init.constant_(self.mask_conv.weight, 1.0)
        self.input_conv.apply(self.weight_init)

        for param in self.mask_conv.parameters():
            param.requires_grad = False

    def forward(self, input, mask):
        output = self.input_conv(input*mask)
        if self.input_conv.bias is not None:
            output_bias = self.input_conv.bias.view(1,-1,1,1).expand_as(output)
        else:
            output_bias = torch.zeros_like(output)
        with torch.no_grad():
            output_mask = self.mask_conv(mask)
        no_update_holes = output_mask == 0
        mask_sum = output_mask.masked_fill_(no_update_holes, 1.0)

        output_pre = (output - output_bias) / mask_sum + output_bias
        output = output_pre.masked_fill_(no_update_holes, 0.0)

        new_mask = torch.ones_like(output)
        new_mask = new_mask.masked_fill_(no_update_holes, 0.0)

        return output, new_mask

class PCBActive(nn.Module):
    def __init__(self, in_channels, out_channels, stride = 1, activ='leaky', conv_bias = False):
        super(PCBActive, self).__init__()

        self.conv = PartialConv(in_channels, out_channels, 3, stride=stride, padding=1, bias= conv_bias)
        self.bn = nn.BatchNorm2d(out_channels)
        if activ == 'relu':
            self.activation = nn.ReLU()
        elif activ == 'leaky':
            self.activation = nn.LeakyReLU(negative_slope=0.2, inplace=True)

    def forward(self, input, mask):
        out, out_mask = self.conv(input, mask)
        out = self.bn(out)
        out = self.activation(out)

        return out, out_mask

class DoublePConv(nn.Module):
    def __init__(self, inchannels, outchannels, stride = 1, active='leaky', conv_bias = False):
        super(DoublePConv, self).__init__()
        self.pconv1 = PCBActive(inchannels, outchannels, stride=stride, activ= active, conv_bias=conv_bias)
        self.pconv2 = PCBActive(outchannels, outchannels, stride=stride, activ= active, conv_bias=conv_bias)

    def forward(self, input, mask):
        out1, mask1 = self.pconv1(input, mask)
        out, out_mask = self.pconv2(out1, mask1)

        return out, out_mask


import sys
input = sys.stdin.readline
from collections import deque

def dfs(graph, node):
    stack = [node]
    visited[node] = 1
    while stack:
        cur = stack.pop()
        for adj in graph[cur]:
            if visited[adj] == 0:
                stack.append(adj)
                visited[adj] = 1

n, m = map(int, input().split())

graph = {i:[] for i in range(1,n+1)}
for i in range(m):
    x, y = map(int, input().split())
    graph[x].append(y)
    graph[y].append(x)

count = 0
visited = [0]*(n+1)
for key in graph:
    if visited[key] == 0:
        count += 1
        dfs(graph, key)

print(count)

# visited = [] ---> if key not in visited 를 인덱스로 !~~!! 수정 시간이 10/1 로 줌 ㄷㄷ

# 유니온 파인드로도 풀 수 있다.
'''
n, m = map(int, input().split())

parent = [0]*(n+1)
for i in range(1, n+1):
    parent[i] = i

# 최상위 부모를 찾습니다
def getParent(parent, x):
    if parent[x] == x:
        return x
    return getParent(parent, parent[x])

# 최상위 부모끼리 비교해서, 더 작은 값으로 바꿉니다.
def unionParent(parent, x, y):
    x = getParent(parent, x)
    y = getParent(parent, y)

    # 둘중, 더 작은 부모를 갖습니다.
    if x < y:
        parent[y] = x
    else:
        parent[x] = y

for _ in range(m):
    a, b = map(int, input().split())
    unionParent(parent, a, b) # 노드끼리 연결합니다.

result = [] # 최상위 부모노드들만 담을 배열
for a in parent:  # parent의 각 원소마다 최상위 부모노드를 구합니다.
    result.append(getParent(parent, a))
#print(result)

result = result[1:] # 0번째 노드는 필요없으니 잘라줍니다.

print(len(set(result)))
'''


"""View for checkout"""
from django.shortcuts import render, get_object_or_404, redirect, reverse
from django.contrib.auth.decorators import login_required
from django.contrib import messages
from django.conf import settings
from django.utils import timezone
import stripe
from tickets.models import Ticket
from .models import NewFeatureUpvote
from .forms import MakePaymentForm

stripe.api_key = settings.STRIPE_SECRET

@login_required()
def checkout(request):
    """Checkout view"""
    if request.method == 'POST':
        payment_form = MakePaymentForm(request.POST)
        # If form is valid
        if payment_form.is_valid():
            cart = request.session.get('cart', {})
            total = 0
            new_feature_to_upvote = []
            # Add username and date and save order
            for id, amount in cart.items():
                new_feature = get_object_or_404(Ticket, pk=id)
                new_feature_to_upvote.append((new_feature, amount))
                total += amount
                new_feature_upvote = NewFeatureUpvote(username=request.user.username, amount=total, new_feature=new_feature, date=timezone.now())
                new_feature_upvote.save()
            # Try payment, fetch errors
            try:
                customer = stripe.Charge.create(
                    amount=int(total*100),
                    currency='EUR',
                    description=request.user.email,
                    card=payment_form.cleaned_data['stripe_id']
                    )
            except stripe.error.CardError:
                messages.error(request, "Your card was declined!")
            # If payment was successful
            if customer.paid:
                messages.error(request, "Your payment was successfull!")
                # Upvote payed features, add amount to payments
                for feature in new_feature_to_upvote:
                    ticket = Ticket.objects.get(pk=feature[0].pk)
                    ticket.votes.up(user_id=request.user.pk)
                    amount = feature[1]
                    ticket.payments += amount
                    ticket.save()
                request.session['cart'] = {}
                return redirect(reverse('tickets'))
            # If payment was not successful
            else:
                messages.error(request, "Payment was not successful!")
        # If form is not valid
        else:
            print(payment_form.errors)
            messages.error(request, "We were unable to make a payment with that card!")
    else:
        payment_form = MakePaymentForm

    return render(request, 'checkout.html', {'payment_form': payment_form, 'publishable': settings.STRIPE_PUBLISHABLE})


import matplotlib.pyplot as plt
import math
import random
from decimal import Decimal

# Recursive factorial
def factorial(n):
	if n < 0:
		return -1 # Error, defined only for positive numbers
	if n == 1 or n == 0:
		return 1
	return n*factorial(n-1)

# Heron's method
def sqrt(s):
	x = s/2 # initial guess
	# Iterate until error is small
	while x**2 - s > 0.01  or x**2 - s < -0.01:
		x = ((s/x)+x)/2
	return x

# Tylor sin x in degrees
def sin(x):
	x = (x*3.1416)/180 # assuming degrees as input
	s = 0
	for n in range(0,10):
		alpha = 2*n+1
		s = s + (((-1)**n)*(x**alpha))/factorial(alpha)
	return s

# Tylor cos x in degrees
def cos(x):
	x = (x*3.1416)/180
	s = 0
	for n in range(0,10):
		alpha = 2*n
		s = s + (((-1)**n)*(x**alpha))/factorial(alpha)
	return s

# linear congruential method
# x_i+1 = (a*X_n + c) mod m
def myrand(x):
	x = hash(str(x)) # using a hash to have a large starging number
	c = 1013904223
	a = 1664525
	m = 2**32
	x = (a*x+c) % m
	return x/m # to get number between 0 and 1

# Euclidean distance
def distance(v1,v2):
	# If one cluster is empty
	if len(v1)==0 or len(v2)==0:
		return 1000 # distance is infinite
	sum = 0
	for i in range(len(v1)):
		sum += ((v2[i]-v1[i])**2)
	return sqrt(sum)

fig = plt.figure()
t = []
for i in range(0,360):
	t.append(sin(i))

s = []
for i in range(0,360):
	s.append(math.sin(math.radians(i)))

plt.plot(t,"-",c="blue", label="Without import")
plt.plot(s,":", linewidth=3,c="red", label="With import")
plt.title("Sine")
plt.xlabel("Degree")
plt.ylabel("Sine")
plt.legend(loc="upper left")
plt.show()

fig = plt.figure()
t = []
for i in range(0,360):
	t.append(cos(i))

s = []
for i in range(0,360):
	s.append(math.cos(math.radians(i)))

plt.plot(t,"-",c="blue", label="Without import")
plt.plot(s,":", linewidth=3,c="red", label="With import")
plt.title("Cosine")
plt.xlabel("Degree")
plt.ylabel("Cosine")
plt.legend(loc="upper left")
plt.show()

fig = plt.figure()
t = []
for i in range(0,100):
	t.append(sqrt(i))

s = []
for i in range(0,100):
	s.append(math.sqrt(i))

plt.plot(t,"-",c="blue", label="Without import")
plt.plot(s,":", linewidth=3,c="red", label="With import")
plt.title("Square root")
plt.xlabel("Value")
plt.ylabel("Square root")
plt.legend(loc="upper left")
plt.show()

fig = plt.figure()
x = []
for j in range(0,50):
	for i in range(0,50):
		x.append(myrand(i+myrand(j)))

y = []
for j in range(0,50):
	for i in range(0,50):
		y.append(myrand(j+myrand(i)))

x1 = []
for j in range(0,50):
	for i in range(0,50):
		x1.append(random.random())

y1 = []
for j in range(0,50):
	for i in range(0,50):
		y1.append(random.random())

plt.scatter(x,y,c="blue",marker="o", label="Without import")
plt.scatter(x1,y1,c="red",marker="x", label="With import")
plt.title("Random")
plt.xlabel("x coordinate")
plt.ylabel("y coordinate")
plt.legend(loc="upper left")
plt.show()


from simulators.base_simulator import Prediction

def parse_traces_from_file(fname: str) -> [(str, "Prediction")]:
    """
    parse_traces_from_file is a wrapper around parse_traces_from_string. it returns
    a list of branch results.
    """
    with open(fname, "r") as f:
        return parse_traces_from_string("\n".join(f.readlines()))

def parse_traces_from_string(body: str) -> [(str, "Prediction")]:
    """
    given a string encoding of branch results, parse_traces_from_string will break
    up these values and return a sequence of branch results
    """
    lines = body.split("\n")
    branches = []
    for line in lines:
        if line != "":
            (addr, taken) = line.split(" ")
            if taken == "t":
                branches.append((addr, Prediction.TAKEN))
            elif taken == "n":
                branches.append((addr, Prediction.NOTTAKEN))

    return branches


import glob
import json
import sys
import numpy as np
import os

# this scrpit is used to find a proper color map factor used in 'building.py' by mapping the average density to average color

if __name__ == '__main__':
	graph = sys.argv[1]

	wavemaps = glob.glob(f'{graph}/{graph}_waves/wavemap_*.json')

	buildingList = []
	for wavemap in wavemaps:
		_, name = os.path.split(wavemap)
		if name.count('_') == 3:
			buildingList.append(wavemap)

	densityArray = []
	for waveFile in buildingList:
		with open(waveFile, 'r') as f:
			data = json.load(f)
		for wave, waveData in data.items():
			if int(wave) > 0:
				v = waveData['ss']
				vo = waveData['t'] + waveData['s']
				e = 2 * waveData['ie'] + waveData['ee']
				densityArray.append((e / (v * vo)) if e > 0 else 0)
			else:
				continue

	aveDensity = np.mean(densityArray)
	factor = np.square(aveDensity / (aveDensity - 1))
	print(factor)


import subprocess
import sys
import os
import fnmatch
import concurrent.futures
from commonFunction import deal_fastqc, find_files_by_suffix, process_fastq_files

BASE_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
software_path = os.path.join(BASE_DIR, "pipeline/software")
os.environ["PATH"] += os.pathsep + f"{software_path}/hlahd/hlahd.1.7.0/bin"
# print(os.environ["PATH"])

def trim_fastq(sample_name, sample_file, project_dir, software_path):
    print("I'm trim_fastq function!")
    seqkit = os.path.join(software_path, "seqkit/seqkit")
    trim_threads = os.cpu_count()
    trim_portion = 0.1
    seqkit_seed = 120
    trim_string = "1M"
    trim_number = int(float(trim_string.rstrip("M"))*1000000)  # 1000000

    fq1_path = sample_file['fq1']
    fq2_path = sample_file['fq2']

    trim_data_dir = os.path.join(project_dir, sample_name, f"downsample_{trim_string}")

    os.makedirs(trim_data_dir, exist_ok=True)
    subprocess.run(f"{seqkit} sample -j {trim_threads} -p {trim_portion} -s {seqkit_seed} {fq1_path} | {seqkit} head -n {trim_number} -o {trim_data_dir}/{sample_name}_1.fq.gz --quiet", shell=True)
    subprocess.run(f"{seqkit} sample -j {trim_threads} -p {trim_portion} -s {seqkit_seed} {fq2_path} | {seqkit} head -n {trim_number} -o {trim_data_dir}/{sample_name}_2.fq.gz --quiet", shell=True)

def analyze_sample(sample_name, sample_files, project_dir, software_path, database_path, script_path, ref_fa, software_list):
    print("I'm analyze_sample function!")
    fq1 = sample_files['fq1']
    fq2 = sample_files['fq2']
    threads = os.cpu_count()

    sample_dir = os.path.join(project_dir, sample_name, 'mapping')
    os.makedirs(sample_dir, exist_ok=True)

    # bwa
    bwa_command = (
        f"bwa mem -t {threads} -Y -H '@HD\\tVN:1.5\\tGO:none\\tSO:coordinate' "
        f"-R '@RG\\tID:TEST\\tSM:TEST\\tLB:Target\\tPL:BGI\\tPU:HVW2MCCXX:6:none' "
        f"{ref_fa} {fq1} {fq2} | "
        f"samtools view -buhS -t {ref_fa}.fai - | "
        f"samtools sort -o {sample_dir}/{sample_name}.hla.sortedByCoord.bam -"
    )
    subprocess.run(bwa_command, shell=True)
    subprocess.run(["samtools", "index", f"{sample_dir}/{sample_name}.hla.sortedByCoord.bam"])
    subprocess.run(f'samtools mpileup -f {ref_fa} {sample_dir}/{sample_name}.hla.sortedByCoord.bam -aa -A -Q 13 -o {sample_dir}/{sample_name}.pileup.Q13.out', shell=True)
    subprocess.run(f'python {script_path}/bwa_mileup_stats.py {sample_dir}/{sample_name}.pileup.Q13.out > {sample_dir}/{sample_name}.pileup.Q13.out.stats', shell=True)

    # extract fastq
    subprocess.run(f"samtools view  -@ {threads} -bF 4 {sample_dir}/{sample_name}.hla.sortedByCoord.bam > {sample_dir}/{sample_name}.mapped.hla.bam", shell=True)
    subprocess.run(f"samtools view  -@ {threads} -bf 4 {sample_dir}/{sample_name}.hla.sortedByCoord.bam > {sample_dir}/{sample_name}.Unmapped.hla.bam", shell=True)
    subprocess.run(f"samtools fastq -@ {threads} -1    {sample_dir}/{sample_name}.mapped.hla.1.fastq -2 {sample_dir}/{sample_name}.mapped.hla.2.fastq -n {sample_dir}/{sample_name}.mapped.hla.bam", shell=True)

    return_list = []
    if 'hlahd' in software_list:
        # hlahd
        hlahd_command = (
            f"hlahd.sh -t {threads} -m 100 -c 0.95 "
            f"-f {software_path}/hlahd/hlahd.1.7.0/freq_data "
            f"{sample_dir}/{sample_name}.mapped.hla.1.fastq "
            f"{sample_dir}/{sample_name}.mapped.hla.2.fastq "
            f"{software_path}/hlahd/hlahd.1.7.0/HLA_gene.split.txt "
            f"{software_path}/hlahd/hlahd.1.7.0/dictionary/ "
            "HLA-HD_Result "
            f"{project_dir}/{sample_name}/"
        )
        subprocess.run(hlahd_command, shell=True)
        subprocess.run(f"python {script_path}/hla.02.freq.py {database_path}/hla.cwd.xls {project_dir}/{sample_name}/HLA-HD_Result/result/HLA-HD_Result_final.result.txt", shell=True)
        return_list.append(f"{project_dir}/{sample_name}/HLA-HD_Result/result/HLA-HD_Result_final.result.txt")
    if 'hla_scan' in software_list:
        # hla_scan
        print("Start hla_scan!")
        hla_scan_dir = os.path.join(project_dir, sample_name, "HLAscan_Result")
        os.makedirs(hla_scan_dir, exist_ok=True)
        hla_genes = ["HLA-A","HLA-B","HLA-C","HLA-DMA","HLA-DMB","HLA-DOA","HLA-DOB","HLA-DPA1","HLA-DPB1","HLA-DQA1","HLA-DQB1","HLA-DRA","HLA-DRB1","HLA-DRB5","HLA-E","HLA-F","HLA-G","MICA","MICB","TAP1","TAP2"]
        with open(f"{hla_scan_dir}/{sample_name}.HLAscan.shell.sh", "w") as f:
            for gene in hla_genes:
                f.write(f"{software_path}/hla_scan/hla_scan -t {threads} -l {sample_dir}/{sample_name}.mapped.hla.1.fastq -r {sample_dir}/{sample_name}.mapped.hla.2.fastq -d {software_path}/hla_scan/db/HLA-ALL.IMGT -g {gene} > {hla_scan_dir}/{sample_name}.{gene}.out.txt\n")
                subprocess.run(f"{software_path}/hla_scan/hla_scan -t {threads} -l {sample_dir}/{sample_name}.mapped.hla.1.fastq -r {sample_dir}/{sample_name}.mapped.hla.2.fastq -d {software_path}/hla_scan/db/HLA-ALL.IMGT -g {gene} > {hla_scan_dir}/{sample_name}.{gene}.out.txt", shell=True)
        subprocess.run(f"python3 {script_path}/hla.04.merge_HLAscan_result.py {hla_scan_dir}/{sample_name}*.out.txt > {hla_scan_dir}/{sample_name}.HLAscan.results.txt\n", shell=True)
        return_list.append(f"{hla_scan_dir}/{sample_name}.HLAscan.results.txt")

    if 'OptiType' in software_list:
        # optitype
        print("Start OptiType!")
        opt_dir = os.path.join(project_dir, sample_name ,"OptiType_Result")
        os.makedirs(opt_dir)
        OptiType = os.path.join(software_path, "OptiType-1.3.5", "OptiTypePipeline.py")
        razers3 = os.path.join(software_path, "razers", "razers3")
        OptiType_ref = os.path.join(ref_path, "hla_reference_dna.fasta")
        subprocess.run(f"{razers3} -tc {threads} -i 95 -m 1 -dr 0 -o {opt_dir}/sample.fished_1.bam {OptiType_ref} {fq1}\n", shell=True)
        subprocess.run(f"{razers3} -tc {threads} -i 95 -m 1 -dr 0 -o {opt_dir}/sample.fished_2.bam {OptiType_ref} {fq2}\n", shell=True)
        subprocess.run(f"samtools bam2fq {opt_dir}/sample.fished_1.bam > {opt_dir}/sample.fished_1.fastq\n", shell=True)
        subprocess.run(f"samtools bam2fq {opt_dir}/sample.fished_2.bam > {opt_dir}/sample.fished_2.fastq\n", shell=True)
        subprocess.run(f"python2.7 {OptiType} -i {opt_dir}/sample.fished_1.fastq {opt_dir}/sample.fished_2.fastq --dna -v -o {opt_dir} -p {sample_name}\n", shell=True)
        return_list.append(f"{opt_dir}/{sample_name}.result.tsv")

    return return_list

def main(data_dir, project_dir, software_path, database_path, script_path, ref_fa, software_list):
    fastq_list  = find_files_by_suffix(".fq.gz", data_dir)
    sample_dict = process_fastq_files(fastq_list)

    # Step 1 对原始数据进行 fastqc， 不需要等待执行结果。
    with concurrent.futures.ProcessPoolExecutor() as executor:
        for sample_name, sample_files in sample_dict.items():
            executor.submit(deal_fastqc, sample_name, sample_files, project_dir, software_path, redirct=True)

    '''
    if 'hlahd' in software_list:
        ### Step 2 对原始数据进行downsample. 要等待执行结果。
        ##################### 处理downsample文件夹下的fastq文件 #####################################
        with concurrent.futures.ProcessPoolExecutor() as executor:
            futures = [executor.submit(trim_fastq, sample_name, sample_files, project_dir, software_path)
                    for sample_name, sample_files in sample_dict.items()]
            for future in concurrent.futures.as_completed(futures):
                future.result()

        ### Step 3 对downsample的数据进行fastqc， 不需要等待执行结果。
        fastq_list  = find_files_by_suffix(".fq.gz", project_dir)
        sample_dict = process_fastq_files(fastq_list)
        with concurrent.futures.ProcessPoolExecutor() as executor:
            for sample_name, sample_files in sample_dict.items():
                executor.submit(deal_fastqc, sample_name, sample_files, project_dir, software_path)
    else:
        print("No need to downsample!")

    '''
    ### Step 4 multiqc  # multiqc was installed through pip install.
    subprocess.run(f"multiqc {project_dir}  --outdir {project_dir}", shell=True)

    ### Step 5 对downsample数据进行 bwa + hlahd， 要等待执行结果。
    result_list = []
    with concurrent.futures.ProcessPoolExecutor() as executor:
        futures = [executor.submit(analyze_sample, sample_name, sample_files, project_dir, software_path, database_path, script_path, ref_fa, software_list)
                   for sample_name, sample_files in sample_dict.items()]

        for future in concurrent.futures.as_completed(futures):
            result_list.extend(future.result())

    print(f"result_list: {result_list}")
    ###########################################################################################
    # write file list that need to be packaged into a file.
    with open(os.path.join(project_dir, "need_to_be_packaged.txt"), "w") as f:
        if 'hlahd' in software_list:
            # 额外统计HLA-HD的结果，并写入文件
            subprocess.run(f"python {script_path}/hla.03.summary.hla-hd.py {project_dir}",shell=True)
            f.write(f"{project_dir}/HLA-HD_Result_final.summary.xls\n")

        # 添加其他要打包的文件， 需要确认！
        for result in result_list:
            print(f"result: {result}")
            f.write(f"{result}\n")
    print("Package result finished!")

if __name__ == "__main__":
    data_dir     = sys.argv[1]
    project_name = sys.argv[2]  # test1
    project_dir  = sys.argv[3]  # /data/storeData/ztron/rawdata/HLA/analysis/test3
    software_list = sys.argv[4].split(",")  # ['hla-hd', 'hla_scan', 'optitype']
    print(f"software_list in hla.main.py 1st : {software_list}")
    os.makedirs(project_dir, exist_ok=True)

    BASE_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
    software_path = os.path.join(BASE_DIR, "pipeline/software")
    database_path = os.path.join(BASE_DIR, "pipeline/database")
    ref_path = os.path.join(BASE_DIR, "pipeline/ref/hla")
    ref_fa = os.path.join(ref_path, "hla_gen.fasta")
    script_path = os.path.join(BASE_DIR, "tools/scripts")

    main(data_dir, project_dir, software_path, database_path, script_path, ref_fa, software_list)


import json
from pytube import YouTube
import ssl
import requests
from lxml import etree

ssl._create_default_https_context = ssl._create_unverified_context

def extra_video_data(video_url):
    if 'tiktok' in video_url:
        print(video_url)
        page=requests.get(url='https://www.tiktok.com/oembed?url='+video_url)
        data = json.loads(page.text)
        video_data = {"embed_url": video_url, "video_title": data['title'],"video_id":data['provider_url'],'articleImageLink': data['thumbnail_url'], 'articleUser' : data['author_name']}
        return video_data
        # html = etree.HTML(page)
    else:

        result = YouTube(video_url)
        embed_url = result.embed_url
        video_id=result.video_id
        video_title = result.title
        video_data = {"embed_url": embed_url, "video_title": video_title,"video_id":video_id}
        return video_data


# -*- coding: utf-8 -*-
"""
Created on Sat Apr 14 11:18:12 2018

@author: ziwenruo
"""

# -*- coding: utf-8 -*-
"""
Created on Sun Mar 18 09:35:34 2018

@author: ziwenruo
"""

import tensorflow as tf
import numpy as np
import pandas as pd
import os
import argparse
from collections import namedtuple

REC_DB_PATH = 'G:\\workarea\\scripts\\stock_tfexample\\train'
SRECORD = namedtuple('REC',"date high close tick")

record_info = pd.DataFrame({
        'name':['code','date','label_type','tick'],
        'type':[tf.string,tf.string,tf.int64,tf.float32],
        'shape':[(),(),(),(299,299,3)],
        'isbyte':[True,True,False,False],
        'length_type':['fixed','fixed','fixed','var'],
        }, columns=['name','type','shape','isbyte','length_type'])

def get_filenames(cate, path=REC_DB_PATH):
    # get all file names
    files= os.listdir(path)
    filepaths = [os.path.join(path,file) for file in files if not os.path.isdir(file) and '.tfrecord' in file]
    if cate in ("t0", "t1", "t2"):
        filepaths = [file for file in filepaths if "_%s_"%cate in file]
    filepaths.sort(key=os.path.getsize, reverse=True)
    return np.array(filepaths)

def read_stock_tfrecord(fname):
    rec_opt = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.GZIP)
    record_iterator = tf.python_io.tf_record_iterator(fname, options=rec_opt)
    example = tf.train.Example()
    TFREC = namedtuple('TFREC',"date code label tick")
    for i in iter(record_iterator):
        example.ParseFromString(i)
        mycode = example.features.feature['code'].bytes_list.value
        mydate = example.features.feature['date'].bytes_list.value
        mylabel = example.features.feature['label_type'].int64_list.value[0]
        mytick = example.features.feature['tick'].float_list.value
        yield TFREC(mydate, mycode, mylabel, mytick)

#
#tarfile = ["G:\\workarea\\scripts\\stock_tfexample\\train\\tick20180319_t1_603712.tfrecord",
#           "G:\\workarea\\scripts\\stock_tfexample\\train\\tick20180319_t0_002294.tfrecord",
#           "G:\\workarea\\scripts\\stock_tfexample\\train\\tick20180319_t0_002032.tfrecord",
#           ]

def provide_rec(flist):
    for f in flist:
        yield from read_stock_tfrecord(f)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--rtype', help='example: t0|t2|t1')
    parser.add_argument('--path', default=REC_DB_PATH, help='example: t0|t2|t1')

    args = parser.parse_args()

#    data_path = 'G:\\workarea\\scripts\\stock_tfexample\\train'
    db_file_list = get_filenames(args.rtype)

    count = 0
    for y in iter(provide_rec(db_file_list)):
        count += 1
        print("<{}> {} {} {}".format(count, y.date, y.code, y.label))

    print("Got {} records finally!!".format(count))

if __name__ == '__main__':
    main()


#!/usr/bin/python3
"""
Module containing ways of pronouncing a word,
using howjsay website (online) and pyttsx3 (local)
"""
from playsound import playsound
import pyttsx3

engine = pyttsx3.init()
engine.setProperty("rate", 125)

try:
    with open("blacklist", "r") as f:
        blacklist = f.read().split("\n")
    if blacklist[-1] == "":
        blacklist = blacklist[:-1]
except FileNotFoundError:
    print("Can't find 'blacklist' file")
    blacklist = []

def pronounce(raw_word):
    """pronounces the given word"""
    word = ""
    for i in raw_word:
        if i in "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz":
            word += i.lower()
        else:
            word += " "
    if " " in word:
        split_word = word.split(" ")
        for i in split_word:
            pronounce(i)
    else:
        if word != "" and word not in blacklist:
            url = "https://howjsay.com/mp3/" + word + ".mp3"
            playsound(url)

def tts(word):
    """says the selected word using pyttsx3 module"""
    engine.say(word)
    engine.runAndWait()


import copy

def merge(lst1, lst2):
    """
    (list,list) -> list
    function return one sorted list of two lists
    >>>merge([1,3],[2,4])
    [1,2,3,4]
    >>>merge([1,3,5],[1,2,6])
    [1,1,3,5,6]
    """
    i1 = 0
    i2 = 0
    newlst = []
    while (i1 != len(lst1)) and (i2 != len(lst2)):
        if lst1[i1] < lst2[i2]:
            newlst.append(lst1[i1])
            i1 += 1
        else:
            newlst.append(lst2[i2])
            i2 += 1
    newlst.extend(lst1[i1:])
    newlst.extend(lst2[i2:])
    return newlst

def mergesort(lst):
    """
    (list) -> list
    function return sorted list
    >>>mergesort([1,2,3,4,5])
    [1,2,3,4,5]
    >>>mergesort([])
    []
    >>>mergesort([5,7,1,3,-1])
    [-1,1,4,5,7]
    """
    workspace = []
    for i in range(len(lst)):
        workspace.append([lst[i]])
        if type(lst[i]) == str:
            raise ValueError
    i = 0
    while i < len(workspace) - 1:
        L1 = workspace[i]
        L2 = workspace[i + 1]
        new_lst = merge(L1, L2)
        workspace.append(new_lst)
        i = i + 2
    if len(workspace) != 0:
        lst[:] = workspace[-1][:]
    return lst

def is_sorted(lst):
    """
    (lst) -> bool
    check if list is sorted
    :return:
    """
    start_lst = copy.deepcopy(lst)
    sorted_lst = mergesort(lst)
    return sorted_lst == start_lst


import re

def answer(n):
    row = open('input.txt').read().replace('^', '1').replace('.', '0')
    w = len(row)
    row = int('0b'+row, 2)
    safe = 0
    for i in range(n):
        safe += w-bin(row).count('1')
        row = (row >> 1) ^ (row << 1) & ((2<<(w-1))-1)
    return safe

print('Answer #1:', answer(40))
print('Answer #2:', answer(400000))


#! /usr/bin/env python
import heapq
import itertools

from storygen.events.event import PlagueEvent, BirthEvent
from storygen.events.timestamp import Time, TimeDelta
from storygen.names.name import Name

def heappopwhile(predicate, heap):
    while True:
        try:
            e = heapq.heappop(heap)
        except IndexError:
            break
        if predicate(e):
            yield e
        else:
            heapq.heappush(heap, e)
            break

class World:

    def __init__(self):
        self.clock = Time(5000, 1, 1)
        self.events = []
        self.people = set()
        self.men = set()
        self.women = set()

    def schedule(self, e):
        if e.timestamp < self.clock:
            raise RuntimeError("Cannot schedule an event in the past")
        heapq.heappush(self.events, e)

    def step(self, dt=None):
        dt = dt or TimeDelta(days=1)
        self.clock += dt
        # self.log("A new day")
        population = len(self.people)
        for e in heappopwhile((lambda e: e.timestamp <= self.clock), self.events):
            e(self)
        # if population != len(self.people):
        #     print("WORLD POPULATION: {}".format(len(self.people)))
        return bool(self.events)

    def steps(self, n, dt=None):
        for i in range(n):
            yield self.step(dt=dt)

    def step_infinite(self, dt=None):
        while True:
            yield self.step(dt=dt)

    def step_all(self, dt=None):
        for e in itertools.takewhile((lambda has_events: has_events), self.step_infinite(dt=dt)):
            yield e

    def log(self, message, *args, timestamp=None, **kwargs):
        timestamp = timestamp or self.clock
        msg = message.format(*args, **kwargs)
        print("[{clock:%Y-%m-%d}] {msg}".format(clock=timestamp, msg=msg))

def main():
    initial_population = 100

    from storygen.names.phoneme import load_phonemes
    from storygen.languages.language import LanguageGenerator
    from storygen.culture.culture import Culture
    from storygen.people.person import Person

    language_gen = LanguageGenerator(load_phonemes())
    language = language_gen()
    culture = Culture(language)
    print("Language phonemes: {}".format(", ".join(sorted(str(p) for p in language.phonemes))))
    names = [Name(culture) for n in range(initial_population)]

    world = World()

    for i, name in enumerate(names):
        birthdate = world.clock
        person = Person(name, Person.MALE if (i % 2) else Person.FEMALE, birthdate, culture)
        birth = BirthEvent(person)
        world.schedule(birth)

    # world.schedule(PlagueEvent(world.clock + TimeDelta(days=100000)))

    for _ in world.step_all():
        pass

    world.log("History has ended after {} years", world.clock.year-5000)

if __name__ == "__main__":
    main()


###################################
# CS B551 Fall 2018, Assignment #3
#
#Yashaswini Dhatrika(yhatrik)

#
#
####
# The following are the process and models using for part of tagging the sentence:

#Step1:Training process:
#The required dictionaries are used which would be useful while implementing different models
#words_dic():stores probabilities of words in the trainingdataset where word is key and prob is value
#t2w() :nested dictionary which stores the count of word in a given tag and total count
#tag_dic():stores probabilities of tags, where each tag is key and prob is value
#transition():stores the probability of each tag starting the sentence
#start_tag_dic:stores the probability of each tag starting the sentence
#tri_transition:stores the 3 consecutive tag states as key and respective count

#Step2:
#Model1:Simplified Method:
#Calculation of pred tags to a sentence:
#1)P(Si/wi) for each tag for a given word is calculated using bayes law, which is equalivalent to P(wi/si)*p(si)/P(w1)
#2) And a tag for a given word is assigned based on maximum probability.

#Calculation of log(posterior):
#1:calculating the posterior as P(Si,Si+1,Si+2/wi,wi+1,w+2)=P(wi/si)*p(si)*P(wi+1/si+1).../P(wi,wi+1,wi+2..)
#2:In the calculation we are ignoring the denominator as it is constant

##Model2: HMM model using viterbi Algorithm

#Calculation of pred tags to a sentence:
#Step1: Emission probabilities(P(wi/si)), Transition probabilities(P(si/si-1)
# and most probable path in the prior state is calculated (v(i))
#Step2:for the 1st word in the sentence or the 1st word (emission probability* intital prob) is calculated
        #Intital prob=the prob of tag starting the sentence.
#Step3: for all the remaining words, by emission probab*max(transition probability*V(i))
#Step4: In order to track which prior state tag is resulting a max value of present state, a key called prev_tag is created to store the values.
#Step5: Once, the step 3 and Step4 is carried for all the words in the sentence, then tag for each word can be found using the backtracking mechanism

#Calculation of log(posterior):
#1.Calculating the posterior as P(s1,s2,s3../w1,w2,w3..)=P(w1/s1)*P(s1)*P(w2/s2)*P(s2/s1)*P(w3/s3)*P(s3/s2)..../P(w1,w2,w3,w4...)
#2.In the calculation we are ignoring the denominator as it is constant

##Model3: Complex model using gibbs samplings

#Calculation of pred tags to a sentence:
#Step1:  The idea in Gibbs sampling is to generate posterior samples by  sweeping through each variable (or block of variables) to sample from its conditional distribution with the remaining variables fixed to their current values.
#Step2: Calculation of  conditional probability (CP) can break down to 4 cases
    # Let i be the position of word in the sentence which we are sampling
    #Case1: if i==0 (i.e. 1st word in the sentence)
        #CP=P(s1)P(sw1/s1)*P(s2/s1)*P(s3/s2,S1)/marginalization over s1(P(s1)P(sw1/s1)*P(s2/s1)*P(s3/s2,S1))
    #Case2: if i==1(i.e. 2nd word in the sentence)
        #CP=P(s2/s1)*P(s3/s2,S1)*P(s4/s3,s2)*P(w2/s2)/marginalization over s2(P(s2/s1)*P(s3/s2,S1)*P(s4/s3,s2)*P(w2/s2))

    #Case3: if i=len(sentence)-1 ((i.e. 1st word in the sentence))
        #CP=P(wi/si)*P(si/si-1,si-2)/arginalization over si(P(wi/si)*P(si/si-1,si-2))
    #Case4: Other words in the sentence:
        #CP=P(si/si-1,si-2)*P(wi/si)*P(si+1/si,si-1)*P(si+2/Si,Si+1)/marginalization over Si(P(si/si-1,si-2)*P(wi/si)*P(si+1/si,si-1)*P(si+2/Si,Si+1))

#Step3: Now the samples are generated by using the conditional distribution.The process is just like a flipping a coin but which # of outputs and also baised.
        #A random output of a tag is assigned using the conditional distribution
#Step4: few thousands of samples are generated and for calculating the predictive tag we eliminate
#the initial sample generate(warm up period)

#Step5:Based on the max. occurrence in the reliable sample, the tag is assigned to each word.

#Calculation of log(posterior):
 # Calculating the posterior as P(s1,s2,s3../w1,w2,w3..)=P(w1/s1)*P(s1)*P(w2/s2)*P(s2/s1)*P(w3/s3)*P(s3/s1,s2)..../P(w1,w2,w3)
 ##In the calculation we are ignoring the denominator as it is constant

####Results (accuracies) for each technique on test data
##So far scored 2000 sentences with 29442 words.
                  ## Words correct:     Sentences correct:
 ## 0. Ground truth:      100.00%              100.00%
 ##       1. Simple:       91.73%               37.75%
 ##       2. HMM:          94.88%               53.55%
   ##     3. Complex:       91.74%               40.20%
####

#Observation of the results:
#We can see that HMM Model is outperforms the other models in accuracy
#But generally it is expected that complex model is better than others model as it consider a sequence of tags while computing
# But it might be possible that if the training data isn't huge, then we would be able to get the all the combination of transitions,so thereby missing out some details

import random
import math
import numpy as np

# We've set up a suggested code structure, but feel free to change it. Just
# make sure your code still works with the label.py and pos_scorer.py code
# that we've supplied.
#
class Solver:

    def __init__(self):
        # word_dic stores probabilities of words in the trainingdataset where word is key and prob is value
        self.words_dic = {}
        # tag_dic: stores probabilities of tags, where each tag is key and prob is value
        self.tag_dic = {}
        # t2w is nested dictionary which stores the count of word in a given tag and total count
        # for example {N:{word:4},N:{tag_count:100}}
        self.t2w = {}
        # transition is nested dictionary which stores the transition state of 2 consecutive pairs in the traindata
        self.transition = {}
        # start_tag_dic stores the probability of each tag starting the sentence
        self.start_tag_dic = {}
        # tri_transition is dictionary which stores the 3 consecutive tag states as key and respective count
        self.tri_transition = {}

    # Calculating the log of the posterior probability of a given sentence for given model
    def posterior(self, model, sentence, label):
        # if the tag and word pair isn't found in the training set, a very low probability is assigned
        if model == "Simple":
            log_posterior = 0
            for word, tag in zip(sentence, label):
                # calculating the posterior as P(Si,Si+1,Si+2/wi,wi+1,w+2)=P(wi/si)*p(si)*P(wi+1/si+1).../P(w1,w2,w3)
                # In the calculation we are ignoring the denominator as it is constant
                if word in self.t2w[tag].keys():
                    pos_prob = (self.t2w[tag][word] / self.t2w[tag]['tag_count']) \
                               * (self.tag_dic[tag])
                    log_posterior += math.log(pos_prob)

                else:
                    pos_prob = 0.0000001
                    log_posterior += math.log(pos_prob)
            return log_posterior
        elif model == "Complex":
            # Calculating the posterior as P(s1,s2,s3../w1,w2,w3..)=P(w1/s1)*P(s1)*P(w2/s2)*P(s2/s1)*P(w3/s3)*P(s3/s1,s2)..../P(w1,w2,w3)
            ##In the calculation we are ignoring the denominator as it is constant
            # if the tag and word pair isn't found in the training set, a very low probability is assigned
            # if an transition between states aren't found in the training data a very low probability is assigned.
            log_post = 0
            log_prob = 0
            for i in range(len(sentence)):
                # calculating p(wi/si) using the dictionary t2w

                prob_w_s = self.t2w.get(label[i]).get(sentence[i], 0.0000001)
                if prob_w_s != 0.0000001:
                    prob_w_s = prob_w_s / self.t2w.get(label[i]).get('tag_count')
                # for 1st word in the sentence, we need to calculate P(s1)
                if i == 0:
                    prob_trans = self.tag_dic[label[i]]
                    log_prob += math.log(prob_trans)
                # for 2nd word in the sentence, we need to calculate P(s2/s1)
                elif i == 1:
                    prob_trans = self.transition.get(label[i - 1]).get(label[i], 0.0000001)
                    if prob_trans != 0.0000001:
                        prob_trans = prob_trans / self.transition.get(label[i - 1]).get('trans_count')
                    log_prob += math.log(prob_trans)
                # for all other words in the sentence, we need to calculate P(si/si-1,si-2)
                else:
                    if label[i - 2] + "/" + label[i - 1] + "/" + label[i] in self.tri_transition.keys():
                        tri_trans_key = label[i - 2] + "/" + label[i - 1] + "/" + label[i]
                        prob_trans = self.tri_transition[tri_trans_key] / self.tri_transition['tot_count']
                    else:
                        prob_trans = 0.0000001
                    log_prob += math.log(prob_trans)
                log_prob += math.log(prob_w_s)
                log_post += log_prob
            return log_post

        elif model == "HMM":
            # Calculating the posterior as P(s1,s2,s3../w1,w2,w3..)=P(w1/s1)*P(s1)*P(w2/s2)*P(s2/s1)*P(w3/s3)*P(s3/s2)..../P(w1,w2,w3)
            ##In the calculation we are ignoring the denominator as it is constant
            # if the tag and word pair isn't found in the training set, a very low probability is assigned
            # if an transition between states aren't found in the training data a very low probability is assigned.
            log_post = 0
            for i in range(len(sentence)):
                # for the 1st word calculating the P(s1) and P(w1/s1)
                if i == 0:
                    s0 = self.start_tag_dic.get(label[i])
                    prob_w0_s0 = self.t2w.get(label[i]).get(sentence[i], 0.0000001)
                    if prob_w0_s0 != 0.0000001:
                        prob_w0_s0 = prob_w0_s0 / self.t2w[label[i]]['tag_count']
                    prob0 = math.log(s0 * prob_w0_s0)
                    log_post += prob0

                # for all other words, calculating the P(si/si-1) and P(wi/si)
                if i >= 1:
                    if self.transition[label[i - 1]].get(label[i]) != None:
                        prob_trans = self.transition[label[i - 1]][label[i]] / self.transition[label[i - 1]][
                            'trans_count']
                    else:
                        prob_trans = 0.0000001
                    prob_w_s = self.t2w.get(label[i]).get(sentence[i], 0.0000001)
                    if prob_w_s != 0.0000001:
                        prob_w_s = prob_w_s / self.t2w[label[i]]['tag_count']
                    log_prob = math.log(prob_trans * prob_w_s)
                    log_post += log_prob
            return log_post
        else:
            print("Unknown algo!")

    # Do the training!

    def train(self, data):
        # Calculation of P[Si]
        S = []
        for i in range(0, len(data)):
            x = data[i][1]
            S.append(x)
        S = [tag for sentence in S for tag in sentence]
        dict = {}
        for i in S:
            if (i in dict.keys()):
                dict[i] += 1
            else:
                dict[i] = 1

        total = sum(dict.values())

        self.tag_dic = {k: dict[k] / total for k in dict}

        # Calculation of p(wi)
        W = []
        for i in range(0, len(data)):
            x = data[i][0]
            W.append(x)
        W_new = ('&&'.join([data for ele in W for data in ele])).split('&&')
        dict1 = {}
        for i in W_new:
            if (i in dict1.keys()):
                dict1[i] += 1
            else:
                dict1[i] = 1

        total = sum(dict1.values())

        self.words_dic = {k: dict1[k] / total for k in dict1.keys()}

        # calculation of p(w/s)
        tags_per_sentence = [x[1] for x in data]
        words_per_sentence = [x[0] for x in data]

        W = [w for sentence in words_per_sentence for w in sentence]

        word_data = [(i, j) for i, j in zip(W, S)]

        S = set(S)

        for tag in S:
            temp = []
            for word in word_data:
                if word[1] == tag:
                    temp.append(word[0])
            if tag not in self.t2w.keys():
                self.t2w[tag] = {'tag_count': 0}
            for word in temp:
                if word not in self.t2w[tag].keys():
                    self.t2w[tag][word] = 1
                else:
                    self.t2w[tag][word] += 1
                self.t2w[tag]['tag_count'] += 1

        # calculation of p(Si+1/Si):
        for tags in tags_per_sentence:
            for i in range(0, len(tags) - 1):
                if tags[i] not in self.transition.keys():
                    self.transition[tags[i]] = {'trans_count': 0}
                if tags[i + 1] not in self.transition[tags[i]].keys():
                    self.transition[tags[i]][tags[i + 1]] = 1
                else:
                    self.transition[tags[i]][tags[i + 1]] += 1
                self.transition[tags[i]]['trans_count'] += 1

        # Calculating the prob of each tag starting a sentence
        start_tag = [tags_per_sentence[i][0] for i in range(len(tags_per_sentence))]
        start_dict = {}
        for i in start_tag:
            if (i in start_dict.keys()):
                start_dict[i] += 1
            else:
                start_dict[i] = 1
        total = sum(start_dict.values())
        self.start_tag_dic = {k: start_dict[k] / total for k in start_dict.keys()}

        # Calculation of p(Si+1/Si,Si-1)
        for tags in tags_per_sentence:
            self.tri_transition['tot_count'] = 0
            for i in range(0, len(tags) - 2):
                temp = tags[i] + "/" + tags[i + 1] + "/" + tags[i + 2]
                if temp not in self.tri_transition.keys():
                    self.tri_transition[temp] = 1
                else:
                    self.tri_transition[temp] += 1
                self.tri_transition['tot_count'] += 1

# Functions for each algorithm which predicts the the tags for each word in the sentence
# Following are the consideration taken:
# 1.if the tag and word pair isn't found in the training set, a very low probability is assigned
# 2.if an transition between states aren't found in the training data a very low probability is assigned.

    def simplified(self, sentence):
        list_pred = []
        tags = self.tag_dic.keys()
        for word in sentence:
            temp_tag = {}
            # Simplified model: We just need to calculate the P(si/wi) for all the
            # tag and assign the tag to the word which has max. prob
            for tag in tags:
                if word in self.t2w[tag].keys():
                    # P(si/wi)=P(si/w1)*p(si)/P(wi)
                    temp_tag[tag] = (self.t2w[tag][word] / self.t2w[tag]['tag_count']) \
                                    * (self.tag_dic[tag] / self.words_dic[word])
                else:
                    temp_tag[tag] = 0.0000001
            tag_pred = max(temp_tag, key=temp_tag.get)
            list_pred.append(tag_pred)

        return list_pred

#     The below function calculates the condition probabilities for each word in the sentence, this probabilies are useful gibbs sampling
#    Calculation of P(Si/all the other states except Si are fixed)
#    Calculation of  conditional probability (CP) can break down to 4 cases:
#     Let i be the position of word in the sentence which we are sampling
#    Case1: if i==0 (i.e. 1st word in the sentence)
#        CP=P(s1)P(sw1/s1)*P(s2/s1)*P(s3/s2,S1)/marginalization over s1(P(s1)P(sw1/s1)*P(s2/s1)*P(s3/s2,S1))
#    Case2: if i==1(i.e. 2nd word in the sentence)
#        CP=P(s2/s1)*P(s3/s2,S1)*P(s4/s3,s2)*P(w2/s2)/marginalization over s2(P(s2/s1)*P(s3/s2,S1)*P(s4/s3,s2)*P(w2/s2))
#
#    Case3: if i=len(sentence)-1 ((i.e. 1st word in the sentence))
#        CP=P(wi/si)*P(si/si-1,si-2)/arginalization over si(P(wi/si)*P(si/si-1,si-2))
#    Case4: Other words in the sentence:
#        CP=P(si/si-1,si-2)*P(wi/si)*P(si+1/si,si-1)*P(si+2/Si,Si+1)/marginalization over Si(P(si/si-1,si-2)*P(wi/si)*P(si+1/si,si-1)*P(si+2/Si,Si+1))

    def joint_dist(self, tag_sample, sentence, i):
        #Attributes:
        # 1.tag_sample: this is the list of all tags assigned to words in the sentence
        #2. Sentence: Sentence of dataset
        #3. i: i is the position on which sampling is done using gibbs sampling
        self.sentence = sentence
        self.sample = tag_sample
        tags = self.tag_dic.keys()
        list1 = []
        #Some of edge cases are considered, when a sentence as single word:
        #CP=P(s1)P(sw1/s1)/marginalization over s1(P(s1)P(sw1/s1))
        if len(sentence) == 1:
            for tag in tags:
                if sentence[i] in self.t2w[tag].keys():
                    joint_prob = self.tag_dic[tag] * (self.t2w[tag][sentence[i]] / self.t2w[tag]['tag_count'])
                    list1.append(joint_prob)
                else:
                    joint_prob = 0.0000001
                    list1.append(joint_prob)
            value = sum(list1)
            list1 = [x / value for x in list1]
            return list1
        # CP=P(s1)P(sw1/s1)*P(s2/s1)/Marginalization over s2(P(s1)P(sw1/s1)*P(s2/s1))
        if len(sentence) == 2 and i == 0:
            for tag in tags:
                if sentence[i] in self.t2w[tag].keys() and tag_sample[i + 1] in self.transition[tag].keys():
                    joint_prob = self.tag_dic[tag] * (self.t2w[tag][sentence[i]] / self.t2w[tag]['tag_count']) \
                                 * (self.transition[tag][tag_sample[i + 1]] / self.transition[tag]['trans_count'])
                    list1.append(joint_prob)
                else:
                    joint_prob = 0.0000001
                    list1.append(joint_prob)
            value = sum(list1)
            list1 = [x / value for x in list1]
            return list1
        # CP=P(s2/s1)*P(w2/s2)/marginalization over (P(s2/s1)*P(w2/s2))
        if len(sentence) == 2 and i == 1:
            for tag in tags:
                if sentence[i] in self.t2w[tag].keys() and tag in self.transition[tag_sample[i - 1]].keys():
                    joint_prob = (self.transition[tag_sample[i - 1]][tag] / self.transition[tag_sample[i - 1]][
                        'trans_count']) \
                                 * (self.t2w[tag][sentence[i]] / self.t2w[tag]['tag_count'])
                    list1.append(joint_prob)
                else:
                    joint_prob = 0.0000001
                    list1.append(joint_prob)
            value = sum(list1)
            list1 = [x / value for x in list1]
            return list1
        #CP=P(s1)P(sw1/s1)*P(s2/s1)*P(s3/s2,S1)/marginalization over s1(P(s1)P(sw1/s1)*P(s2/s1)*P(s3/s2,S1))
        if len(sentence) == 3 and i == 0:
            for tag in tags:
                if sentence[i] in self.t2w[tag].keys() and tag_sample[i + 1] in self.transition[tag].keys() \
                        and tag + "/" + tag_sample[i + 1] + "/" + tag_sample[i + 2] in self.tri_transition.keys():
                    tri_trans_key = tag + "/" + tag_sample[i + 1] + "/" + tag_sample[i + 2]
                    joint_prob = self.tag_dic[tag] * (self.t2w[tag][sentence[i]] / self.t2w[tag]['tag_count']) \
                                 * (self.transition[tag][tag_sample[i + 1]] / self.transition[tag]['trans_count']) \
                                 * (self.tri_transition[tri_trans_key] / self.tri_transition['tot_count'])
                    list1.append(joint_prob)
                else:
                    joint_prob = 0.0000001
                    list1.append(joint_prob)
            value = sum(list1)
            list1 = [x / value for x in list1]
            return list1
        # CP=P(s2/s1)*P(s3/s2,S1)*P(w2/s2)/marginalization over s2(P(s2/s1)*P(s3/s2,S1)*P(w2/s2))
        if len(sentence) == 3 and i == 1:
            for tag in tags:
                if sentence[i] in self.t2w[tag].keys() and tag in self.transition[tag_sample[i - 1]].keys() and \
                        tag_sample[i - 1] + "/" + tag + "/" + tag_sample[i + 1] in self.tri_transition.keys():
                    tri_trans_key1 = tag_sample[i - 1] + "/" + tag + "/" + tag_sample[i + 1]
                    joint_prob = (self.transition[tag_sample[i - 1]][tag] / self.transition[tag_sample[i - 1]][
                        'trans_count']) \
                                 * (self.t2w[tag][sentence[i]] / self.t2w[tag]['tag_count']) \
                                 * (self.tri_transition[tri_trans_key1] / self.tri_transition['tot_count'])
                    list1.append(joint_prob)
                else:
                    joint_prob = 0.0000001
                    list1.append(joint_prob)
            value = sum(list1)
            list1 = [x / value for x in list1]
            return list1
        #case3:CP=P(wi/si)*P(si/si-1,si-2)/arginalization over si(P(wi/si)*P(si/si-1,si-2))
        if (len(sentence) == 3 and i == 2) or (i == len(sentence) - 1):
            for tag in tags:
                tri_trans_key = tag_sample[i - 2] + "/" + tag_sample[i - 1] + "/" + tag_sample[i]
                if sentence[i] in self.t2w[tag].keys() and tri_trans_key in self.tri_transition.keys():
                    joint_prob = (self.t2w[tag][sentence[i]] / self.t2w[tag]['tag_count']) \
                                 * (self.tri_transition[tri_trans_key] / self.tri_transition['tot_count'])
                    list1.append(joint_prob)
                else:
                    joint_prob = 0.0000001
                    list1.append(joint_prob)
            value = sum(list1)
            list1 = [x / value for x in list1]
            return list1
        #Case1:CP=P(s1)P(sw1/s1)*P(s2/s1)*P(s3/s2,S1)/marginalization over s1(P(s1)P(sw1/s1)*P(s2/s1)*P(s3/s2,S1))
        if i == 0:
            for tag in tags:
                if sentence[i] in self.t2w[tag].keys() and tag_sample[i + 1] in self.transition[tag].keys() \
                        and tag + "/" + tag_sample[i + 1] + "/" + tag_sample[i + 2] in self.tri_transition.keys():
                    tri_trans_key = tag + "/" + tag_sample[i + 1] + "/" + tag_sample[i + 2]
                    joint_prob = self.tag_dic[tag] * (self.t2w[tag][sentence[i]] / self.t2w[tag]['tag_count']) \
                                 * (self.transition[tag][tag_sample[i + 1]] / self.transition[tag]['trans_count']) \
                                 * (self.tri_transition[tri_trans_key] / self.tri_transition['tot_count'])
                    list1.append(joint_prob)

                else:
                    joint_prob = 0.0000001
                    list1.append(joint_prob)
            value = sum(list1)
            list1 = [x / value for x in list1]
            return list1
        #Case2:#CP=P(s2/s1)*P(s3/s2,S1)*P(s4/s3,s2)*P(w2/s2)/marginalization over s2(P(s2/s1)*P(s3/s2,S1)*P(s4/s3,s2)*P(w2/s2))
        if i == 1:
            for tag in tags:
                if sentence[i] in self.t2w[tag].keys() and tag in self.transition[tag_sample[i - 1]].keys() and \
                        tag_sample[i - 1] + "/" + tag + "/" + tag_sample[i + 1] in self.tri_transition.keys() and \
                        tag + "/" + tag_sample[i + 1] + "/" + tag_sample[i + 2] in self.tri_transition.keys():
                    tri_trans_key1 = tag_sample[i - 1] + "/" + tag + "/" + tag_sample[i + 1]
                    tri_trans_key2 = tag + "/" + tag_sample[i + 1] + "/" + tag_sample[i + 2]
                    joint_prob = (self.transition[tag_sample[i - 1]][tag] / self.transition[tag_sample[i - 1]][
                        'trans_count']) \
                                 * (self.t2w[tag][sentence[i]] / self.t2w[tag]['tag_count']) \
                                 * (self.tri_transition[tri_trans_key1] / self.tri_transition['tot_count']) \
                                 * (self.tri_transition[tri_trans_key2] / self.tri_transition['tot_count'])
                    list1.append(joint_prob)
                else:
                    joint_prob = 0.0000001
                    list1.append(joint_prob)
            value = sum(list1)
            list1 = [x / value for x in list1]
            return list1
        #Case4:
        if i > 1:
            for tag in tags:
                if sentence[i] in self.t2w[tag].keys() and len(sentence) - 1 >= i + 2 and \
                        tag_sample[i - 2] + "/" + tag_sample[i - 1] + "/" + tag in self.tri_transition.keys() and \
                        tag_sample[i - 1] + "/" + tag + "/" + tag_sample[i + 1] in self.tri_transition.keys() and \
                        tag + "/" + tag_sample[i + 1] + "/" + tag_sample[i + 2] in self.tri_transition.keys():
                    tri_trans_key1 = tag_sample[i - 2] + "/" + tag_sample[i - 1] + "/" + tag
                    tri_trans_key2 = tag_sample[i - 1] + "/" + tag + "/" + tag_sample[i + 1]
                    tri_trans_key3 = tag + "/" + tag_sample[i + 1] + "/" + tag_sample[i + 2]
                    # print(tri_trans_key1,tri_trans_key2,tri_trans_key3,sentence[i],tag)
                    join_prob = (self.t2w[tag][sentence[i]] / self.t2w[tag]['tag_count']) \
                                * (self.tri_transition[tri_trans_key1] / self.tri_transition['tot_count']) \
                                * (self.tri_transition[tri_trans_key2] / self.tri_transition['tot_count']) \
                                * (self.tri_transition[tri_trans_key3] / self.tri_transition['tot_count'])
                    list1.append(join_prob)

                elif sentence[i] in self.t2w[tag].keys() and len(sentence) - 1 >= i + 1 and \
                        tag_sample[i - 2] + "/" + tag_sample[i - 1] + "/" + tag in self.tri_transition.keys() and \
                        tag_sample[i - 1] + "/" + tag + "/" + tag_sample[i + 1] in self.tri_transition.keys():
                    tri_trans_key1 = tag_sample[i - 2] + "/" + tag_sample[i - 1] + "/" + tag
                    tri_trans_key2 = tag_sample[i - 1] + "/" + tag + "/" + tag_sample[i + 1]
                    join_prob = (self.t2w[tag][sentence[i]] / self.t2w[tag]['tag_count']) \
                                * (self.tri_transition[tri_trans_key1] / self.tri_transition['tot_count']) \
                                * (self.tri_transition[tri_trans_key2] / self.tri_transition['tot_count'])
                    list1.append(join_prob)
                else:
                    joint_prob = 0.0000001
                    list1.append(joint_prob)
            value = sum(list1)
            list1 = [x / value for x in list1]

        return list1

    # This function performs the MCMC -Gibbs samplings
    #Flow of the code inside complex_mcmc is as follows:
    #Step1.A random tag is assigned to each words in the sentence
    #Step2.Gibbs sampling is performed across all the words using the the joint_dist function
    #Step3: Each tag for a word is sampled based on the distribution , and a tag is choosen from the 12 values just as flipping a coin but which as many outputs and baised
    #Step4: Few thousands of samples are generated by fixing all the tags and just sampling one.
    #Step5: Then the posterior for each word is calculated based on the max. occurrence on the samples excluding some of the inital samples generated(warm up period)
    def complex_mcmc(self, sentence):
        tags = list(self.tag_dic.keys())
        #Random tag of len of sentence is generated
        rand_tag = list(np.random.choice(list(tags), len(sentence)))
        all_sample = []
        all_sample.append(rand_tag)
        t = 0
        while t <= 500:
            for i in range(0, len(sentence)):
                #samples are generating by sampling one tag which would traverse from 1st word of sentence to last
                tags_prob = self.joint_dist(rand_tag, sentence, i)
                # a tag is assigned based on conditional distribution using random.choice()
                rand_tag[i] = np.random.choice(list(tags), 1, p=tags_prob)[0]
                all_sample.append(rand_tag)
            t += 1

        reliable_sample = all_sample[100:]
        dict1 = {}
        for i in range(0, len(sentence)):
            if i not in dict1.keys():
                dict1[i] = {}
        for sample in reliable_sample:
            for j in range(0, len(sample)):
                if sample[j] not in dict1[j].keys():
                    dict1[j][sample[j]] = 1
                else:
                    dict1[j][sample[j]] += 1
        pred_tag = []
        for i in range(len(sentence)):
            pred_tag.append(max(dict1[i], key=lambda x: dict1[i][x]))
        return pred_tag

#The below function perform the viterbi algorithm
#Flow of code inside the function is as follows:
#Step1: for the 1st word emission probability* intital prob is calculated
    #where emission probability=p(wi/si)
    #Intital prob=the prob of tag starting the sentence.
#Step2: for all the remaining words, by emission probab*max(transition probability*prob of prior state)
#Step3: In order to track which prior state tag is resulting a max value of present state, a key called prev_tag is created to store the values.
#Step4: Once, the step 2 and Step3 is carried for all the words in the sentence, then tag for each word can be found using the backtracking mechanism
    def hmm_viterbi(self, sentence):
        vit = []
        for i in range(0, len(sentence)):
            temp_dic = {tag: {'prob': 0, 'prev_tag': ''} for tag in self.tag_dic.keys()}
            # temp_lst=[]

            for cur_tag in self.tag_dic.keys():
                #
                max_prob = -9999999
                likely_tag = ''
                if sentence[i] not in self.t2w[cur_tag].keys():
                    e_prob = 0.000001
                else:
                    e_prob = self.t2w[cur_tag][sentence[i]] / self.t2w[cur_tag]['tag_count']

                if i == 0:

                    if cur_tag not in self.start_tag_dic.keys():
                        init_prob = 0.0000001
                    else:
                        init_prob = self.start_tag_dic[cur_tag]
                    temp_dic[cur_tag]['prob'] = e_prob * init_prob

                    # print(temp_dic)

                    # temp_dic[cur_tag]['prob']=0.000001
                else:
                    for prev_tag in vit[-1].keys():

                        if cur_tag not in self.transition[prev_tag].keys():
                            trans_prob = 0.0000001
                        else:
                            trans_prob = self.transition[prev_tag][cur_tag] / self.transition[prev_tag]['trans_count']

                        if max_prob < vit[-1][prev_tag]['prob'] * trans_prob:
                            max_prob = vit[-1][prev_tag]['prob'] * trans_prob
                            likely_tag = prev_tag

                    temp_dic[cur_tag]['prob'] = e_prob * max_prob
                    temp_dic[cur_tag]['prev_tag'] = likely_tag
            #                    print(temp_dic[cur_tag])
            vit.append(temp_dic)
        #            print(vit[-1])
        #Backtracking in order to assign the predictive tag to each word based on the max value obtained
        pred = []
        start = max(vit[-1], key=lambda x: vit[-1][x]['prob'])
        pred.append(start)
        vit = vit[::-1]
        for i in range(0, len(vit) - 1):
            pred.append(vit[i][start]["prev_tag"])
            start = vit[i][start]["prev_tag"]

        return pred[::-1]

    # This solve() method is called by label.py, so you should keep the interface the
    #  same, but you can change the code itself.
    # It should return a list of part-of-speech labelings of the sentence, one
    #  part of speech per word.
    #
    def solve(self, model, sentence):
        if model == "Simple":
            return self.simplified(sentence)
        elif model == "Complex":
            return self.complex_mcmc(sentence)
        elif model == "HMM":
            return self.hmm_viterbi(sentence)
        else:
            print("Unknown algo!")


from django.contrib import admin

from .models import Order, OrderLineItem, Coupon

class OrderLineItemAdminInline(admin.TabularInline):
    '''Customize the admin interface to display order line items'''
    model = OrderLineItem
    readonly_fields = ('line_item_total',)

class OrderAdmin(admin.ModelAdmin):
    '''Customize the admin interface to display order details'''

    inlines = (OrderLineItemAdminInline,)

    readonly_fields = ('order_id', 'created', 'updated', 'tax',
                       'delivery_cost', 'order_total', 'grand_total',)

    fields = ('order_id', 'user_profile', 'first_name', 'last_name',
              'email_address', 'phone_number', 'street_address1',
              'building_number1',
              'street_address2', 'building_number2',
              'city', 'postal_code',
              'country', 'created', 'updated', 'tax', 'paid',
              'delivery_cost', 'order_total', 'grand_total',)

    list_display = (
        'order_id',
        'first_name',
        'last_name',
        'email_address',
        'phone_number',
        'street_address1',
        'building_number1',
        'street_address2',
        'building_number2',
        'city',
        'postal_code',
        'country',
        'order_total',
        'tax',
        'delivery_cost',
        'grand_total',
        'created',
        'updated',
        'paid',
    )

    ordering = ('-created',)

class CouponAdmin(admin.ModelAdmin):
    '''Customize the admin interface to display gift coupon details'''
    list_display = (
        'coupon_id',
        'coupon_code',
        'coupon_value',
        'coupon_expiry',
        'created',
        'active',
        'description',
    )

admin.site.register(Order, OrderAdmin)
admin.site.register(Coupon, CouponAdmin)