#!/usr/bin/env python3

"""
    This module deals with input and output files to the program.
    It contains parsing and writing functions.
"""

import re
import itertools as it
import numpy as np
from Bio.PDB.Polypeptide import three_to_one, one_to_three
import src.residu_algn as resa
import src.score as sc
import src.checking_argument as cA


def parse_foldrec(filin):
    """
    Parses a .foldrec file.

    Args:
        filin: .foldrec file of the query, as an open file python object

    Returns:
        A dict of all 413 alignments of the query sequence onto the Homstrad
        template sequences and a dict of all secondary strutures for each
        alignment
    Keys:
        Names of the templates on which the query is aligned
    Values:
        Alignement objects containing info about the alignments
    """
    # Regexs for template protein name, alignment score,
    # query sequence & template sequence.
    regex_tplname = re.compile('^Alignment.+vs\s+([\w-]+)')
    regex_score = re.compile('^Score :\s+(-{0,1}\d*\.{0,1}\d+)')
    regex_query = re.compile('^Query\s+(\d+)\s+([ARDNCEQGHILKMFPSTWYV-]+)\s+'
                             '(\d+)')
    regex_template = re.compile('^Template\s+(\d+)\s+([ARDNCEQGHILKMFPSTWYV-]+'
                                ')\s+(\d+)')
    #secondary structure for each residue for query and template
    regex_ss_query = re.compile('^Query\s+(\d+)\s+([GHIEBSTC-]+)\s+(\d+)')
    regex_ss_tpl = re.compile('^Template\s+(\d+)\s+([GHIEBSTC-]+)\s+([1-9]\d+|\d+)')
    #confidence score (/9) for each residue of the query
    regex_conf_query = re.compile('^Query\s+(\d+)\s+([\d-]+)\s+(([1-9]\d+|\d+))')
    # Dictionnay of Alignment objects to return
    dico_aligns = {}
    # Dictionnary of SecStruct objects to return
    dico_ss = {}
    # Loop through the file
    line = filin.readline()
    while line != "":
        # Match template name
        if regex_tplname.search(line):
            tplname = regex_tplname.search(line).group(1)
            dico_aligns[tplname] = resa.Alignment(tplname)
            dico_ss[tplname] = sc.SecStruct()
        # Match alignment score
        elif regex_score.search(line):
            dico_aligns[tplname].score = float(regex_score.search(line).group(1))

        # Match query sequence
        elif regex_query.search(line):
            dico_aligns[tplname].start_query = int(regex_query.search(line).group(1))
            dico_aligns[tplname].query = regex_query.search(line).group(2)
            dico_aligns[tplname].end_query = int(regex_query.search(line).group(3))
            line = filin.readline()
            # Query secondary structure
            dico_ss[tplname].query_ss = regex_ss_query.search(line).group(2)
            line = filin.readline()
            # Confidence score and alignement length
            dico_ss[tplname].query_ss_conf = regex_conf_query.search(line).group(2)
            dico_ss[tplname].query_len = int(regex_conf_query.search(line).group(4))
        # Match template sequence
        elif regex_template.search(line):
            dico_aligns[tplname].start_template = int(regex_template.search(line).group(1))
            dico_aligns[tplname].template = regex_template.search(line).group(2)
            dico_aligns[tplname].end_template = int(regex_template.search(line).group(3))
            line = filin.readline()
            # Template secondary structure
            dico_ss[tplname].tpl_ss = regex_ss_tpl.search(line).group(2)
        line = filin.readline()
    return dico_aligns, dico_ss



def dope_to_arr_ca(input_dope_file):
    """
    Parse the dope.par file, in order to extract only the potentials between
    the CA.
    All the values are stored in a 2D Numpy array (20 amino acids
    in line and columns and 30 values of pseudo-energy).

    Args:
        input_dope_file:    the original dope.par file, with all interactions
        between all atoms of the 20 amino acids (when we need only
        the one between the CA)

    Returns:
        A 2D numpy array and a dict that makes correspond 1-letter amino acid
        code with indexes of the array
    """

    nb_aa = 20
    nb_coef = int(nb_aa*(nb_aa+1)/2)  # Low-memory storage of symetric array
    dope_arr = np.full((nb_coef, 30), np.NaN)  # 2D array
    k = 0
    index_aa = {}  # To make correspond amino acids with index inside the array
    iterat = it.product(range(nb_aa), range(nb_aa))

    for one_line in input_dope_file:
        splitted_line = one_line.split()

        if splitted_line[1] == "CA" and splitted_line[3] == "CA":
            # We get the aa name with its corresponding index:
            if k < 20:
                index_aa[three_to_one(splitted_line[2])] = k
                k += 1

            # Corresponding index in list with index in array (file):
            idx_l, idx_c = next(iterat)
            i, j = min(idx_l, idx_c), max(idx_l, idx_c)
            ind_list = i*nb_aa - int((i-1)*i/2) + j-i

            # Assignation only if it was not already assigned
            if np.isnan(dope_arr[ind_list][0]):
                nrgies = list(map(float, splitted_line[4:]))
                dope_arr[ind_list] = nrgies

    return (dope_arr, index_aa)


def parse_metafold_list(metafold_file):
    """
    Parse the metafold file to retrieve the pdb id for each metafold family.

    Args:
        metafold_file:  name of the metafold file

    Returns:
        A dict of the family name of the metafold with the associated pdb id
    """

    # {Family_name : pdb_id}
    files_to_take = {}

    for line in metafold_file:
        splitted_line = line.split()
        files_to_take[splitted_line[0]] = splitted_line[1].split('.')[0]

    return files_to_take


def parse_needed_atm(name_pdb, start):
    """
    Read a pdb or atm file and extract the coordinates of the CA, the N and the
    C atoms for each residues.
    We need these 3 types of atoms, to make pyMoL connect them to each other.

    Args:
        name_pdb:   name of the pdb or atm file
        start:  resID of the residue in the template
                (according to the aln)

    Returns:
        A list of dicts containing coordinates of residues
    """

    dict_dict = {}
    # Id for residu because in some pdb files the num residu
    # does not start to 1
    res_id = 1
    flag = False

    with open(name_pdb, "r") as filin:
        for line in filin:
            if line[0:6].strip() == "ATOM":
                atom_type = line[12:16].strip()
                if atom_type == "N":
                    if res_id < start:
                        res_id += 1
                    else:
                        flag = True
                if flag:
                    if (atom_type == "N" or
                            atom_type == "C" or
                            atom_type == "CA"):
                        res_name = line[17:20].strip()
                        coord_x = float(line[30:38])
                        coord_y = float(line[38:46])
                        coord_z = float(line[46:54])

                        if atom_type == "N":
                            # If it is a N, we have to create the list:
                            dict_dict[res_id] = [{"atom_type": atom_type,
                                                  "resName": res_name,
                                                  "res_id": res_id,
                                                  "x": coord_x,
                                                  "y": coord_y,
                                                  "z": coord_z}]
                            res_id += 1

                        else:
                            # Else if it is a CA or C, we can append it:
                            # But the resID to put are the one
                            # of the previous CA:
                            res_id_c_star = res_id - 1
                            dict_dict[res_id_c_star].append({"atom_type": atom_type,
                                                             "resName": res_name,
                                                             "res_id": res_id_c_star,
                                                             "x": coord_x,
                                                             "y": coord_y,
                                                             "z": coord_z})

    return dict_dict


def write_pdb_out(aln):
    """
    Takes the different information of the best model, in order to generate
    a pdb with the query threaded on it.
    Takes also the N and C atoms (in addition of the CA), in order to have a
    nice pyMOL display, with atoms linked to each other.

    Args:
        aln:    Alignment instance
    """

    pdb_path = cA.BASE_DIR + "Data/HOMSTRAD/" + aln.name + "/" + aln.pdb_id + ".atm"
    dict_dict_atm = parse_needed_atm(pdb_path, aln.start_template)

    # Write pdb model file with best hit and only CA
    with open("results/" + aln.name + "_threaded.pdb", "w") as filout:
        res_id = aln.start_query
        nb_atom = 1
        for i, res in enumerate(aln.query):
            if res != "-":
                for atom in dict_dict_atm[i+1]:
                    filout.write("ATOM  {:5d} {:^4s} {:3s} A{:4d}    {:8.3f}"
                                 "{:8.3f}{:8.3f}{:6.2f}{:6.2f}\n".format(nb_atom,
                                                                         atom["atom_type"],
                                                                         one_to_three(res),
                                                                         atom["res_id"],
                                                                         atom["x"],
                                                                         atom["y"],
                                                                         atom["z"],
                                                                         1.0, 10.0))
                    nb_atom += 1
            res_id += 1
