"""This module does jobs about pssm and tensor building"""
import os
import numpy
import torch
import torch.nn as nn
from torch.utils.data import TensorDataset
from torch.autograd import Variable


def convert2q3(one_hot_labels):
    """ PSS one hot encoding of Q8 To Q3 Convertion Function

    GHITEBSC(L) to HEC
    H=HGI
    E=EB
    C(L)=STC(L)

    Input is the one-hot encoding of Q8
    Sequence is ['L', 'B', 'E', 'G', 'I', 'H', 'S', 'T','NoSeq']
    """
    q8 = ['L', 'B', 'E', 'G', 'I', 'H', 'S', 'T','NoSeq']
    q3_h = ['H', 'G', 'I']  ## 0
    q3_e = ['E', 'B']       ## 1
    q3_l = ['S', 'T', "L"]  ## 2

    structures = []

    for label in one_hot_labels:
        q8_structure = ''
        q3_structure = 0

        try:
            location_of_one = label.tolist().index(1)
            q8_structure = q8[location_of_one]
            if q8_structure in q3_h:
                #q3_structure = [1, 0, 0]
                q3_structure = 1
            elif q8_structure in q3_e:
                #q3_structure = [0, 1, 0]
                q3_structure = 2
            elif q8_structure in q3_l:
                #q3_structure = [0, 0, 1]
                q3_structure = 3
            else:
                #q3_structure = [0, 0, 0]
                q3_structure = 0
        except ValueError:
            location_of_one = -1
            #q3_structure = [0, 0, 0]  # noseq
            q3_structure = 0

        structures.append(q3_structure)

    return structures


def interpret_pssm(filename):
    """ Interpretion PSSM file
    Get the matrix of pssm values from the pssm plain text file generated by psi-blast

    Parameters:
    filename -- path of pssm file

    Return:
    2d matrix of pssm values 
    """
    with open(filename) as opened_file:

        # ignore the first three lines and last 6 lines
        content = opened_file.readlines()[3:-6]

        # 2d list for storing all values in this file
        chunk_lines = []

        for line in content:
            chunks = []
            # content contains whitespaces
            untrimmed_chunks = line.split(' ')

            # remove space and newline symbol
            for chunk in untrimmed_chunks:
                if chunk != '' and chunk != '\n':
                    chunks.append(chunk)

            # only need column 2 to 23
            cut_chunks = []
            for chunk in chunks[2:22]:
                cut_chunks.append(int(chunk))

            chunk_lines.append(cut_chunks)

        # convert 2d list to numpy matrix
        matrix = numpy.array(chunk_lines)
        return matrix


def standardize_matrix(matrix):
    """ Standardization Matrix
    Including normalization via mean-centering and scaling by standard deviation

    Parameters:
    matrix -- 2d matrix of pssm values

    Return:
    standardized matrix
    """

    std_pssm = []

    # number of row and column
    row_num = matrix.shape[0]
    col_num = matrix.shape[1]

    for i in range(col_num):
        column = matrix[:, i]
        # standard deviation
        std = numpy.std(column)
        # mean
        column_sum = numpy.sum(column)
        mean = column_sum / row_num

        # standardization
        std_column = []
        for j in column:
            standardized_j = standardize(j, std, mean)
            std_column.append(standardized_j)

        std_pssm.append(std_column)

    std_pssm_matrix = numpy.array(std_pssm)
    # need transpose, column becomes row after standardization
    transposed_std_pssm_matrix = numpy.transpose(std_pssm_matrix)
    return transposed_std_pssm_matrix


def standardize(num, std, mean):
    """ Standardization Function

    Parameters:
    num -- number to be standardized
    std -- standard deviation
    mean -- mean
    """
    return (num - mean) / std


def build_datasets():
    """ Training Data and Label Tensor Building Function """

    numpy_file = os.path.abspath("../../Desktop/data/cullpdb+profile_6133.npy.gz")
    cullpdb = numpy.load(numpy_file)

    # row number and column number
    row_num = cullpdb.shape[0]
    subrow_num = 700  # 39990/57

    data = []

    for i in range(row_num):
        str_i = str(i)
        row = cullpdb[i]

        # interpret pssm first
        pssm_content = interpret_pssm('./pssm/' + str_i + ".fasta.txt")
        pssm_row_num = pssm_content.shape[0]
        # padding pssm matrix
        padding_pssm = numpy.zeros((700 - pssm_row_num, 20))  # 700*20
        pssm = numpy.concatenate((pssm_content, padding_pssm), axis=0)

        # reshape from (39900,) to (700, 57)
        row_matrix = numpy.reshape(row, (700, 57))
        acids = row_matrix[:, 0:22]  # 700*22 for acids
        labels = row_matrix[:, 22:31]   # 700*9 for labels

        # concatenate the amino acids, labels and pssm together
        sequence = numpy.concatenate((acids, pssm), axis=1) # 700*42
        sequence_with_labels = numpy.concatenate((sequence, labels), axis=1) # 700*51

        # append to python list
        data.append(sequence_with_labels)

    # convert the python list to numpy array and save to file
    data_array = numpy.array(data)

    return data_array


if __name__ == '__main__':
    datasets = build_datasets()
    numpy.save("datasets.npy", datasets)
