reader.py

# coding=utf-8
"""
                            Reader objects
"""
from __future__ import absolute_import, division, print_function

import os
import re
import logging
import os.path
import collections
import scipy.spatial.distance as distance
from .base import sort_2dict
from .protmap import (ResMap, ResAtmMap)

LOG = logging.getLogger(__name__)
Atom = collections.namedtuple("Atom", ["name", "coords"])


class RegexFile(object):
    """
    File which can be parsed with a regex
    """

    def __init__(self, filepath, filetype='', regex='', sort=''):
        self.regex = regex
        self.sort = sort
        self.filepath = filepath
        self.filetype = filetype
        self.filename = os.path.basename(os.path.splitext(filepath)[0])
        self.lines = {}

    def load(self):
        """
        Fill lines with dictionary. Each key is a line number in the given file
        :return: None
        """
        lines_dict = {}

        if not self.regex:
            LOG.error("Can't parse file %s", self.filepath)

        with open(self.filepath) as f:
            for index, line in enumerate(f):
                match = self.regex.match(line)
                if match:
                    lines_dict[index] = match.groupdict()

        if self.sort:
            lines_dict = sort_2dict(lines_dict, self.sort)

        self.lines = lines_dict


class MapFile(RegexFile):
    # List of 3tuples ("regex_file", "filetype", "sort_field")
    # sort_field allow sorting lines with values into this field
    # TODO: wrong regex for native_full ?
    # TODO: smarter dict ...
    """
    Map file class
    """
    types = {
        "plmdca": {
            "regex": re.compile(r"^(?P<res1_nb>\d+)\s+(?P<res1_name>\w)\s+"
                                r"(?P<res2_nb>\d+)\s+(?P<res2_name>\w)\s+"
                                r"(?P<mi_score>\d)\s+"
                                r"(?P<plm_score>\-?\d+\.\d+)\s*$"),
            "score_field": "plm_score"
        },
        "evfold": {
            "regex": re.compile(
                r'^(?P<res1_nb>\d+),(?P<res2_nb>\d+),'
                r'(?P<ec_score>\-?\d+\.\d+e?\-?\d*),'
                r'(?P<placeholder>\d),(?P<res1_cons>\d+),'
                r'(?P<res2_cons>\d+),(?P<ss_filter>\d|\d{3}),'
                r'(?P<high_cons_filter>\d|\d{3}),'
                r'(?P<cc_filter>\d|\d{3}),(?P<res1_1l_code>\w),'
                r'(?P<res2_1l_code>\w)$'),
            "score_field": "ec_score"
        },
        "pconsc": {
            "regex": re.compile(
                r'^(?P<res1_nb>\d+) (?P<res2_nb>\d+) '
                r'(?P<ec_score>\-?\d+\.\d+e?\-?\d*)$'),
            "score_field": "ec_score"
        },
        "pconsc1": {
            "regex": re.compile(
                r'^(?P<res1_nb>\d+) (?P<res2_nb>\d+) '
                r'(?P<ec_score>\-?\d+\.\d+e?\-?\d*)$'),
            "score_field": "ec_score"
        },
        "pconsc2": {
            "regex": re.compile(
                r'^(?P<res1_nb>\d+) (?P<res2_nb>\d+) '
                r'(?P<ec_score>\-?\d+\.\d+e?\-?\d*)$'),
            "score_field": "ec_score"
        },
        "metapsicov_stg1": {
            "regex": re.compile(
                r'^(?P<res1_nb>\d+) (?P<res2_nb>\d+) \d (?P<res_dist>-?\d+.?\d*) '
                r'(?P<ec_score>\-?\d+.\d+e?\-?\d*)$'),
            "score_field": "ec_score"
        },
        "metapsicov_stg2": {
            "regex": re.compile(
                r'^(?P<res1_nb>\d+) (?P<res2_nb>\d+) \d (?P<res_dist>-?\d+.?\d*) '
                r'(?P<ec_score>\-?\d+.\d+e?\-?\d*)$'),
            "score_field": "ec_score"
        },
        "psicov": {
            "regex": re.compile(
                r'^(?P<res1_nb>\d+) (?P<res2_nb>\d+) \d (?P<res_dist>-?\d+.?\d*) '
                r'(?P<ec_score>\-?\d+.\d+e?\-?\d*)$'),
            "score_field": "ec_score"
        },
        "gremlin": {
            "regex": re.compile(
                r'^(?P<res1_nb>\d+)\t(?P<res2_nb>\d+)\t'
                r'(?P<res1_id>\d+_[AC-IK-NP-TVWYZ])\t'
                r'(?P<res2_id>\d+_[AC-IK-NP-TVWYZ])\t'
                r'(?P<raw_score>\-?\d+\.\d+e?\-?\d*)\t'
                r'(?P<scale_score>\-?\d+\.\d+e?\-?\d*)\t'
                r'(?P<prob>\-?\d+\.\d+e?\-?\d*)'),
            "score_field": "scale_score"
        },
        "native": {
            "regex": re.compile(
                r'^\s*(?P<res1_nb>\d+)\s+(?P<res2_nb>\d+)\s+('
                r'?P<ca_ca>\d+\.\d+)\s+(?P<cb_cb>\d+\.\d+)\s+'
                r'(?P<sc_sc>\d+\.\d+)\s+(?P<valid>\w+)'),
            "score_field": None
        },
        "native_full": {
            "regex": re.compile(
                r'^\s*(?P<res1_nb>\d+)\s+(?P<res2_nb>\d+)\s+('
                r'?P<ca_ca>\d+\.\d+)\s+(?P<cb_cb>\d+\.\d+)\s+'
                r'(?P<sc_sc>\d+\.\d+)\s+(?P<valid>\w+)'),
            "score_field": None
        },
        "bbcontacts": {
            "regex": re.compile(
                r'\s*(?P<identifier>\w+)\s+(?P<diversity>-?\d+\.?\d*)'
                r'\s+(?P<direction>Parallel|Antiparallel)'
                r'\s+(?P<viterbiscore>-?\d+\.?\d*)'
                r'\s+(?P<indexpred>\d+)'
                r'\s+(?P<state>first|internal|last)'
                r'\s+(?P<res1_nb>\d+)\s+(?P<res2_nb>\d+)'),
            "score_field": None
        },
        "contactlist": {
            "regex": re.compile(
                r'^\s*(?P<res1_nb>\d+)[\s,;]+(?P<res2_nb>\d+)[\s,'
                r';]*(?P<con_flag>\w*)'),
            "score_field": None
        },
        "metapsicovhb": {
            "regex": re.compile(
                r'^\s*(?P<res_donor>\d+)[\s,;]+'
                r'(?P<res_acceptor>\d+)[\s,;]+\d[\s,;]+'
                r'(?P<res_dist>-?\d+.?\d*)[\s,;]+'
                r'(?P<hbscore>-?\d+\.?\d*)'),
            "score_field": "hbscore"
        },
        "default_1": {
            "regex": re.compile(
                r'^\s*(?P<res1_nb>\d+)\s+(?P<res2_nb>\d+)\s+'
                r'(?P<resn1>\w+)\s+'
                r'(?P<resn2>\w+)\s+'
                r'(?P<score>[\w\d\.\+\-]+)'),
            "score_field": "score"
        },
        "default_2": {
            "regex": re.compile(
                r'^\s*(?P<res1_nb>\d+)\s+(?P<res2_nb>\d+)\s+'
                r'(?P<score>[\w\d\.\+\-]+)'),
            "score_field": "score"
        },
        "default_3": {
            "regex": re.compile(
                r'^\s*(?P<res1_nb>\d+)\s+(?P<res2_nb>\d+)\s+'
                r'(?P<resn1>\w+)\s+'
                r'(?P<resn2>\w+)\s+'),
            "score_field": None
        },
        "default_4": {
            "regex": re.compile(
                r'^\s*(?P<res1_nb>\d+)\s+(?P<res2_nb>\d+)\s+'),
            "score_field": None
        },
        "empty": {
            "regex": re.compile(r'^\s*$'),
            "score_field": None
        }
    }
    check_type = True

    def __init__(self, *args, **kwargs):
        """

        :param args:
        :param kwargs:
        :return:
        """
        super(MapFile, self).__init__(*args, **kwargs)
        if self.check_type:
            self.regex, self.filetype, self.sort = self.check_maptype()
        self.mapdict = {"alldistmap": None,
                        "allcontactmap": None,
                        "distmap": None,
                        "contactmap": None,
                        "scoremap": None}
        self.clashlist = None
        self.contactlist = None
        self.flaglist = None
        # self.contactmap = None
        # self.distmap = None

    def create_map(self, protein, contactdef, flaglist=None, offset=0, sym=True,
                   path="", **kwargs):
        """

        :param path:
        :param protein:
        :param contactdef:
        :param flaglist:
        :param offset:
        :param sym:
        :param kwargs:
        :return:
        """
        raise NotImplementedError("Class %s doesn't implement create_map" %
                                  self.__class__.__name__)

    def check_contacts(self, aa_seq):
        """
        Check if plm_dict is consistent with input sequence
        :param aa_seq:
        """
        LOG.info("Checking consistency of contacts with input sequence")
        for line in self.lines:
            if self.lines[line]['res1_name'] != aa_seq[int(self.lines[line]['res1_nb']) - 1] \
                    or self.lines[line]['res2_name'] != aa_seq[int(self.lines[line]['res2_nb']) - 1]:
                LOG.error("Difference between given sequence and residu "
                          "names in contact file at line %d !", line)

    def update_map(self, resmap):
        """

        :param resmap:
        :return:
        """
        raise NotImplementedError("Class %s doesn't implement update_map" %
                                  self.__class__.__name__)

    def check_maptype(self):
        """

        :return:
        """
        LOG.info("Checking if file %s correspond to %s format", self.filepath,
                 self.filetype)
        # Check if given type is supported
        # TODO: report this check into commands section
        defaults = ("default_1", "default_2", "default_3", "default_4", "empty")
        if os.stat(self.filepath).st_size == 0:
            LOG.warning("File %s is empty !", self.filepath)
            return [
                self.types["empty"].get("regex"),
                self.filetype,
                self.types["empty"].get("score_field")
            ]
        with open(self.filepath) as infile:
            # Check first and second line of file
            for index, line in enumerate(infile):
                if self.filetype in self.types:
                    LOG.info("Given format (%s) should be supported",
                             self.filetype)
                    match = self.types[self.filetype].get("regex").match(line)
                else:
                    match = None
                    LOG.error("Format %s not supported !", self.filetype)
                if match:
                    LOG.debug("Format type correct")
                    return [
                        self.types[self.filetype].get("regex"),
                        self.filetype,
                        self.types[self.filetype].get("score_field")
                    ]
                else:
                    LOG.warning("Given type do not correspond, checking default"
                                " format for contactlist or empty file...")
                    for subformat in defaults:
                        if self.types.get(subformat)["regex"].match(line):
                            LOG.debug("Format type correct %s", subformat)
                            return [
                                self.types[subformat].get("regex"),
                                self.filetype,
                                self.types[subformat].get("score_field")
                            ]
                if index > 2:
                    # Stop checking after second line
                    LOG.error("Error reading %s file.", self.filetype)
                    break
        LOG.error("Wrong format type given ...")
        return [None] * 3

    def read(self, protein=None, contactdef=5.0, groupby_method="min",
             scsc=None):
        """

        :param protein:
        :param contactdef:
        :param groupby_method:
        :param scsc:
        :return:
        """
        LOG.info("Reading %s file", self.filepath)
        if self.filetype:
            # Read file with regex related to filetype
            self.load()
            LOG.debug(self.lines)
        if protein:
            LOG.info("Loading contact list")
            if self.filetype == "contactlist":
                self.flaglist = {
                    tuple(sorted(
                        [int(self.lines[contact].get("res1_nb")),
                         int(self.lines[contact].get("res2_nb"))]
                    )): self.lines[contact].get("con_flag")
                    for contact in self.lines if
                    self.lines[contact].get("res1_nb") and
                    self.lines[contact].get("res2_nb")}
            if self.filetype == "metapsicovhb":
                # HB contacts aren't sorted since the first res correspond to
                #  donor and second to acceptor
                self.contactlist = [
                    tuple([int(self.lines[contact].get("res_donor")),
                           int(self.lines[contact].get("res_acceptor"))])
                    for contact in self.lines
                    if self.lines[contact].get("res_donor") and
                    self.lines[contact].get("res_acceptor")]
            else:
                self.contactlist = [
                    tuple(sorted(
                        [int(self.lines[contact].get("res1_nb")),
                         int(self.lines[contact].get("res2_nb"))]))
                    for contact in self.lines
                    if self.lines[contact].get("res1_nb") and
                    self.lines[contact].get("res2_nb")]
            LOG.debug(self.contactlist)
            sym = False if self.filetype == "metapsicovhb" else True
            self.create_map(protein, contactdef,
                            groupby_method=groupby_method, scsc=scsc,
                            flaglist=self.flaglist, sym=sym, path=self.filepath)
            if self.filetype == "plmdca":
                # If contact filetype contain residues name, check if it is
                # consistent with given sequence
                self.check_contacts(protein.aa_sequence.sequence)
            if self.filetype == "evfold":
                LOG.info("Loading evfold clash list")
                self.clashlist = [
                    next((el for el in (self.lines[contact].get("ss_filter"),
                                        self.lines[contact].get("high_cons_filter"),
                                        self.lines[contact].get("cc_filter")) if el != "0"), "0")
                    for contact in self.lines if
                    self.lines[contact].get("res1_nb") and
                    self.lines[contact].get("res2_nb")]
                if len(self.contactlist) != len(self.clashlist):
                    LOG.error("When reading input file, clash list is not "
                              "the same length than contactlist")
                LOG.debug(self.clashlist)


class ContactMapFile(MapFile):
    # "plmdca", "evfold", "bbcontacts", "pconsc", "gremlin", "metapsicov",
    """
    Contact map file
    """

    def __init__(self, filepath, filetype):
        """

        :param filepath:
        :param filetype:
        :return:
        """
        super(self.__class__, self).__init__(filepath, filetype)

    def update_map(self, resmap):
        """

        :param resmap:
        :return:
        """
        # TODO: swap dataframe factory here
        raise NotImplementedError

    def create_map(self, protein, *args, **kwargs):
        """
        Get res - res map
        :param protein:
        :param args:
        :param kwargs:
        :return:
        """
        offset = min(protein.index)  # Should be 1 or upper (human_idx)
        idxnames = ["residuex"] if self.filetype != "metapsicovhb" else [
            "donor"]
        colnames = ["residuey"] if self.filetype != "metapsicovhb" else [
            "acceptor"]
        contactmap = ResMap(protein.aa_sequence.sequence, mtype='contact',
                            flaglist=kwargs['flaglist'],
                            seqidx=protein.index, index=idxnames,
                            columns=colnames, sym=kwargs['sym'],
                            desc=self.filetype, path=kwargs.get("path"))
        # DataFrame containing ec scores
        scoremap = ResMap(protein.aa_sequence.sequence, mtype='score',
                          seqidx=protein.index, index=idxnames,
                          columns=colnames, sym=kwargs['sym'],
                          path=kwargs.get("path"),
                          desc=self.filetype) if self.sort else None
        distmap = ResMap(protein.aa_sequence.sequence, mtype='distance',
                         seqidx=protein.index, index=idxnames, path=kwargs.get("path"),
                         columns=colnames, sym=kwargs['sym'],
                         desc=self.filetype) if self.filetype == "metapsicovhb" else None

        for contact_id in self.lines:
            if self.filetype == "metapsicovhb":
                resid1 = int(self.lines[contact_id].get('res_donor'))
                resid2 = int(self.lines[contact_id].get('res_acceptor'))
                dist = float(self.lines[contact_id].get('res_dist'))
            else:
                resid1 = int(self.lines[contact_id].get('res1_nb'))
                resid2 = int(self.lines[contact_id].get('res2_nb'))
                dist = None
            # Res id start from 0 in res-res map
            residx1 = contactmap.index[resid1 - offset]
            residx2 = contactmap.index[resid2 - offset]

            if (int(residx1.split("-")[0]) != resid1) or \
                    (resid2 != int(residx2.split("-")[0])):
                LOG.error("Wrong resid humanidx (%d, %d) in contact (%d) is "
                          "not the same in resmap (%d, %d)",
                          resid1, resid2, contact_id,
                          int(residx1.split("-")[0]),
                          int(residx2.split("-")[0]))

            contactmap.set_value(residx1, residx2, True)
            if self.sort:
                scoremap.sort_list.append((resid1 - offset, resid2 - offset))
                scoremap.set_value(residx1, residx2,
                                   float(self.lines[contact_id].get(self.sort)))
            if distmap is not None:
                distmap.set_value(residx1, residx2, dist)

        LOG.debug("%s contact map:\n%s", self.filetype, contactmap)
        self.mapdict["contactmap"] = contactmap
        LOG.debug("%s score map:\n%s", self.filetype, scoremap)
        self.mapdict["scoremap"] = scoremap
        if distmap is not None:
            self.mapdict["distmap"] = distmap


class PDBFile(MapFile):
    """
    PDB file
    """
    pdbreg = re.compile(r'^(?P<record>ATOM  |HETATM)(?P<serial>[\s\w]{5})'
                        r'\s(?P<name>[\s\w]{4})'
                        r'(?P<altLoc>[\s\w])'
                        r'(?P<resName>\w{3})\s(?P<chainID>\w)'
                        r'(?P<resSeq>[\s\w]{4})(?P<iCode>[\s\w])'
                        r'\s{3}(?P<x>[\s\d-]{4}\.\d{3})(?P<y>[\s\d-]{4}\.\d{3})'
                        r'(?P<z>[\s\d-]{4}\.\d{3})'
                        r'(?P<occupancy>[\s\d-]{3}\.\d{2})'
                        r'(?P<tempFactor>[\s\d-]{3}\.\d{2})'
                        r'\s{10}(?P<element>[\s\w]{2})'
                        r'(?P<charge>[\s\w]{2})')

    def __init__(self, *args, **kwargs):
        # TODO: use PDB object in aria
        # TODO: write dataframe in a separated file
        self.check_type = False
        super(PDBFile, self).__init__(*args, regex=self.pdbreg, filetype="pdb",
                                      **kwargs)

    def create_map(self, protein, contactdef, groupby_method="min", scsc=None,
                   flaglist=None, sym=True, path=""):
        """

        :param protein:
        :param contactdef:
        :param groupby_method:
        :param scsc:
        :param flaglist:
        :param sym:
        :param path:
        """
        resmap = ResAtmMap(protein.aa_sequence.sequence, mtype='distance',
                           flaglist=flaglist, path=path,
                           seqidx=protein.index, desc=self.filetype)
        # noinspection PyTypeChecker
        resmap[:] = self.update_map(resmap, sym=sym)
        LOG.debug("pdb distance map:\n%s", resmap)
        self.mapdict["alldistmap"] = resmap
        self.mapdict["distmap"] = resmap.reduce(groupby=groupby_method)
        self.mapdict["allcontactmap"] = resmap.contact_map(
            contactdef=contactdef, scsc_min=scsc)
        self.mapdict["contactmap"] = self.mapdict["allcontactmap"].reduce()

    def update_map(self, resmap, sym=True):
        """

        :param resmap:
        :param sym:
        :return:
        """
        # Map only on heavy atoms
        # TODO: check if same sequence in pdb file
        LOG.info("Updating distance map with pdb file")
        newmap = resmap.copy()
        heavylist = []
        error_list = set()
        for atom in self.lines:
            if resmap.heavy_reg.match(self.lines[atom]['name'].strip()):
                heavylist.append(atom)

        # For each heavy atom
        for x, atomx in enumerate(heavylist):
            for atomy in heavylist[x:]:
                # TODO: Check first residue number in pdb file
                indx = "%03d-%s" % (int(self.lines[atomx]['resSeq']),
                                    self.lines[atomx]['resName'].strip()), \
                       self.lines[atomx]['name'].strip()
                indy = "%03d-%s" % (int(self.lines[atomy]['resSeq']),
                                    self.lines[atomy]['resName'].strip()), \
                       self.lines[atomy]['name'].strip()
                coordx = (float(self.lines[atomx]['x']),
                          float(self.lines[atomx]['y']),
                          float(self.lines[atomx]['z']))
                coordy = (float(self.lines[atomy]['x']),
                          float(self.lines[atomy]['y']),
                          float(self.lines[atomy]['z']))

                dist = distance.euclidean(coordx, coordy)
                if indx[0] in list(resmap.index.get_level_values("residuex"))\
                        and indy[0] in list(resmap.index.get_level_values("residuex")):
                    LOG.debug("Update distance value (%s, %s)", indx, indy)
                    newmap.at[indx, indy] = dist
                    if sym:
                        # If symmetric matrix
                        newmap.at[indy, indx] = dist
                elif indx[0] not in list(resmap.index.get_level_values("residuex")):
                    error_list.add(indx[0])
                elif indy[0] not in list(resmap.index.get_level_values("residuex")):
                    error_list.add(indy[0])
        if error_list:
            # Listing related humanidx in the initial df
            idxlist = list(resmap.index.get_level_values("residuex"))
            erridx = [idx.split("-")[0] for idx in list(error_list)]
            missidx = list(set([idx for idx in idxlist
                                if idx.split("-")[0] in erridx]))
            for idx in missidx:
                newmap.loc[idx] = None
                if sym:
                    newmap.loc[:][idx] = None
            LOG.error("Can't update pdb distance map for pos in pdb file "
                      "%s with %s", list(error_list), missidx)
        return newmap


class DistanceMapFile(MapFile):
    """
    Distance matrix file
    """
    def __init__(self, filepath, filetype):
        super(MapFile).__init__(filepath, filetype)
        raise NotImplementedError

    def create_map(self, aa_seq, contactdef, **kwargs):
        """

        :param aa_seq:
        :param contactdef:
        :param kwargs:
        :return:
        """
        pass

    # Native dist
    def update_map(self, resmap):
        """

        :param resmap:
        :return:
        """
        pass
        # Construit map avec la liste de residus +  infos de distance du fichier
        # return DistanceMap


class ProtFileListReader(object):
    """
    List of file object
    """
    def __init__(self, cont_def=5.0):
        self.filelist = []
        self.contactdef = cont_def

    def clear(self):
        """
        Initiatise from scratch object
        :return:
        """
        # TODO: Init supprime bien les fichiers du cache ?
        self.__init__(self.contactdef)

    def add_file(self, filepathlist, filetypelist=None):
        """

        :param filepathlist:
        :param filetypelist:
        :return:
        """
        filepathlist = [filepathlist] if type(
            filepathlist) != list else filepathlist
        filetypelist = [filetypelist] if type(
            filetypelist) != list else filetypelist
        if not filetypelist or len(filepathlist) != len(filetypelist):
            filetypelist = [os.path.splitext(_)[1][1:] for _ in filepathlist]
        LOG.info("Reader focused on file(s) %s %s", filepathlist,
                 filetypelist)
        for i, filepath in enumerate(filepathlist):
            if os.path.exists(filepath):
                # TODO: check_type functionstr
                LOG.debug("Adding %s file to watchlist", filetypelist[i])
                if filetypelist[i].lower() == "pdb" and \
                        os.path.splitext(filepath)[1][1:] == "pdb":
                        self.filelist.append(PDBFile(filepath))
                else:
                    self.filelist.append(ContactMapFile(filepath,
                                                        filetypelist[i]))
                # else:
                #     self.filelist.append(DistanceMapFile(filepath,
                #                                          filetypelist[i]))
                # TODO: DistanceMapFile condition
                if not self.filelist[-1].regex:
                    LOG.warning("Can't read %s", filepath)
                    self.filelist.pop()

    def read(self, filepathlist, filetypelist=None, protein=None, scsc=None,
             **kwargs):
        """

        :param filepathlist:
        :param filetypelist:
        :param protein:
        :param scsc:
        :param kwargs:
        :return:
        """
        self.clear()
        self.add_file(filepathlist, filetypelist=filetypelist)
        for fo in self.filelist:
            fo.read(protein=protein, contactdef=self.contactdef,
                    scsc=scsc, **kwargs)