#! /usr/bin/env python3

"""
Contains the Residu class and Alignment class.
"""

import numpy as np


class Residu(object):
    """
    A class containing the coordinates of one amino acid, its name.
    """
    def __init__(self, line):
        """
        Initialize the Residu object

        Args:
            line: Coordinates are saved in a numpy array

        Attributes:
            coord:  x, y and z coordinates of the Residu
            resname:    Residu name
            resid:  Residu id
            chain:  chain which the Residu belongs to
        """
        self.coord = np.array([float(line[30:38]), float(line[38:46]),
                               float(line[46:54])])
        self.resname = line[17:20].strip()
        self.resid = int(line[22:26].strip())
        self.chain = line[21:22]

    def dist_btw_two_res(self, residu2):
        """
        Calculate distance between two Residu objects.

        Args:
            residu2: Residu object

        Returns:
            The distance between these two residues
        """
        return np.linalg.norm(self.coord - residu2.coord)


class Alignment(object):
    """
    A class containing the sequence of the alignment, its start and end index.
    It contains the different calculated scores to know if the structure of the
    template can be the potential structure of the query.
    """

    def __init__(self, name):
        """
        Initialize the Alignment object

        Args:
            name:   name of the alignment (studied template)

        Attributes:
            query:  list of residue objects which represents the sequence
            template:   list of residue objects which represents the sequence
            template_res:
            start_query:    start of the query for the given Alignment
            end_query:  end of the query for the given Alignment
            start_template: start of the template for the given Alignment
            end_template:   end of the template for the given Alignment
            name:   name of the alignment (studied template)
            pdb_id: pdb id
            score:  final score
        """
        self.query = None
        self.template = None
        self.template_res = None
        self.start_query = None
        self.end_query = None
        self.start_template = None
        self.end_template = None
        self.name = name
        self.pdb_id = None
        self.score = None

    def load_dist_matrix(self):
        """
        Calculate distance between all residus pair by pair.

        Returns:
            A matrix of distance between all residues
        """
        size_query = len(self.query)
        # Create the empty matrix distance, which will contain distances
        # between all residus
        matrix_dist = np.empty((size_query, size_query), dtype="float")
        # Fill the matrix distance with NaN
        matrix_dist[:] = np.nan
        for i in range(0, size_query):
            # Only the half matrix is completed because it's symetric
            for j in range(i + 1, size_query):
                # A window of 5 residus is taken because it's not interesting
                # to calculate their distance, because they are too close.
                if j not in range(i, i + 5):
                    # If the two residus aren't a gap
                    if (self.query[i] != "-" and
                            self.query[j] != "-" and
                            self.template_res[i].resname != "-"):
                        # calculate the distance between the two residus
                        dist = self.template_res[i].dist_btw_two_res(self.template_res[j])
                        # fill the matrix with the calculated distance
                        matrix_dist[i, j] = dist
                    else:
                        # -1 is added when one residu is a gap
                        matrix_dist[i, j] = -1
        return matrix_dist

    def clean_algn(self):
        """
        Remove gaps in the template alignment and reciprocally in the
        query alignment.

        Returns:
            The number of gaps in the template to penalize them because
            it means a portion of the query does not have a structure
        """
        tmpl = list(self.template)
        query = list(self.query)
        # Keep index to remove
        idx_remove = []
        for i, letter in enumerate(tmpl):
            if letter == '-':
                idx_remove.append(i)
        # Remove index with gap in these two alignment
        tmpl = [res for i, res in enumerate(tmpl) if i not in idx_remove]
        query = [res for i, res in enumerate(query) if i not in idx_remove]
        self.template = "".join(tmpl)
        self.query = "".join(query)
        return len(idx_remove)

    def init_structure(self, files_to_take, common_path):
        """
        Add the pdb id and the list of Residu object in the alignment

        Args:
            files_to_take: Dict containing the reference pdb for the family
            common_path: Path to the folder containing the pdb
        """
        self.pdb_id = files_to_take[self.name]
        self.set_template_res((common_path + self.pdb_id + ".atm"))

    def set_template_res(self, name_pdb):
        """
        Read a pdb or atm file and create a list of Residu objects which will
        be affected to the Alignment object.

        Args:
            name_pdb: The name of the pdb or atm file
        """

        list_res = []
        # Id for residu because in some pdb files the num residu
        # dosen't start to 1
        res_id = 1
        with open(name_pdb, "r") as filin:
            for line in filin:
                if line[0:6].strip() == "ATOM" and line[12:16].strip() == "CA":
                    if res_id >= self.start_template:
                        list_res.append(Residu(line))
                    res_id += 1
        self.template_res = list_res
