##
# File: BarePDBParserListener.py
# Date: 19-Sep-2025
#
# Updates:
""" ParserLister class for Bare PDB files.
    @author: Masashi Yokochi
"""
__docformat__ = "restructuredtext en"
__author__ = "Masashi Yokochi"
__email__ = "yokochi@protein.osaka-u.ac.jp"
__license__ = "Apache License 2.0"
__version__ = "1.1.1"

import re
import sys
from typing import IO, List, Optional

from antlr4 import ParseTreeListener

try:
    from wwpdb.utils.nmr.NmrDpConstant import (STD_MON_DICT,
                                               PROTON_BEGIN_CODE,
                                               REPRESENTATIVE_MODEL_ID,
                                               REPRESENTATIVE_ALT_ID)
    from wwpdb.utils.nmr.AlignUtil import (letterToDigit,
                                           indexToLetter,
                                           retrieveAtomIdentFromMRMap)
    from wwpdb.utils.nmr.nef.NefTranslator import NefTranslator
    from wwpdb.utils.nmr.io.CifReader import CifReader
    from wwpdb.utils.nmr.mr.BarePDBParser import BarePDBParser
    from wwpdb.utils.nmr.mr.BaseTopologyParserListener import BaseTopologyParserListener
    from wwpdb.utils.nmr.mr.ParserListenerUtil import translateToStdResName
except ImportError:
    from nmr.NmrDpConstant import (STD_MON_DICT,
                                   PROTON_BEGIN_CODE,
                                   REPRESENTATIVE_MODEL_ID,
                                   REPRESENTATIVE_ALT_ID)
    from nmr.AlignUtil import (letterToDigit,
                               indexToLetter,
                               retrieveAtomIdentFromMRMap)
    from nmr.nef.NefTranslator import NefTranslator
    from nmr.io.CifReader import CifReader
    from nmr.mr.BarePDBParser import BarePDBParser
    from nmr.mr.BaseTopologyParserListener import BaseTopologyParserListener
    from nmr.mr.ParserListenerUtil import translateToStdResName


class BarePDBParserListener(ParseTreeListener, BaseTopologyParserListener):
    """ This class defines a complete listener for a parse tree produced by BarePDBParser.
    """
    __slots__ = ('coordinatesStatements', )

    # total appearances of TER clause
    __ter_count = 0
    # offset value for atom number due to TER clause
    __ter_offset = 0
    # END clause
    __end = False

    # collection of atom name selection
    atomNameSelection = []

    def __init__(self, verbose: bool = True, log: IO = sys.stdout,
                 representativeModelId: int = REPRESENTATIVE_MODEL_ID,
                 representativeAltId: str = REPRESENTATIVE_ALT_ID,
                 mrAtomNameMapping: Optional[List[dict]] = None,
                 cR: Optional[CifReader] = None, caC: Optional[dict] = None,
                 nefT: NefTranslator = None) -> None:
        self.__class_name__ = self.__class__.__name__
        self.__version__ = __version__

        super().__init__(verbose, log, representativeModelId, representativeAltId, mrAtomNameMapping,
                         cR, caC, nefT)

        self.file_type = 'nm-aux-pdb'

        self.coordinatesStatements = 0

    def enterBare_pdb(self, ctx: BarePDBParser.Bare_pdbContext):  # pylint: disable=unused-argument
        """ Enter a parse tree produced by BarePDBParser#bare_pdb.
        """

        self.__ter_count = 0
        self.__ter_offset = 0
        self.__end = False

    def exitBare_pdb(self, ctx: BarePDBParser.Bare_pdbContext):  # pylint: disable=unused-argument
        """ Exit a parse tree produced by BarePDBParser#bare_pdb.
        """

        if not self.hasPolySeqModel:
            return

        if len(self.atoms) == 0:
            return

        # set tentative chain_id from label_asym_id, which will be assigned to coordinate auth_asym_id
        chainIndex = letterToDigit(self.polySeqModel[0]['chain_id']) - 1
        chainId = indexToLetter(chainIndex)

        terminus = [atom['auth_atom_id'].endswith('T') for atom in self.atoms if isinstance(atom, dict)]

        atomTotal = len(terminus)

        if atomTotal == 0:
            return

        if terminus[0]:
            terminus[0] = False
        for i in range(0, atomTotal - 1):
            j = i + 1
            if terminus[i] and terminus[j]:
                terminus[i] = False
        if terminus[-1]:
            terminus[-1] = False

        seqIdList, compIdList, retrievedAtomNumList = [], [], []

        hasSegCompId = False
        ancAtomName = prevAtomName = ''
        prevAsymId = prevSeqId = prevCompId = None
        offset = 0
        for atom in self.atoms:

            if isinstance(atom, str):
                self.polySeqPrmTop.append({'chain_id': chainId,
                                           'seq_id': seqIdList,
                                           'auth_comp_id': compIdList})
                seqIdList, compIdList = [], []
                chainIndex += 1
                chainId = indexToLetter(chainIndex)
                offset = 0
                continue

            atomNum = atom['atom_number']
            atomName = atom['auth_atom_id']
            asymId = atom['auth_chain_id']
            _seqId = atom['auth_seq_id']
            compId = atom['auth_comp_id']
            if self.noWaterMol and (compId in ('HOH', 'H2O', 'WAT') or (len(compId) > 3 and compId[:3] in ('HOH', 'H2O', 'WAT'))):
                break
            if not hasSegCompId and (compId.endswith('5') or compId.endswith('3')):
                hasSegCompId = True
            if not hasSegCompId and compId not in STD_MON_DICT and self.mrAtomNameMapping is not None\
               and atomName[0] in PROTON_BEGIN_CODE:
                _, compId, _atomName = retrieveAtomIdentFromMRMap(self.ccU, self.mrAtomNameMapping, _seqId, compId, atomName)
                if _atomName != atomName:
                    atomName = _atomName
                    retrievedAtomNumList.append(atomNum)

            if (offset > 0 or self.__ter_count == 0)\
               and ((0 < atomNum < len(terminus) + 1
                     and ((terminus[atomNum - 1] and ancAtomName.endswith('T'))
                          or (terminus[atomNum - 2] and prevAtomName.endswith('T')
                              and self.csStat.getTypeOfCompId(prevCompId) != self.csStat.getTypeOfCompId(compId))))
                    or self.isSegmentWithAsymHint(prevAsymId, prevCompId, prevAtomName, asymId, compId, atomName)
                    or self.isLigand(prevCompId, compId)
                    or self.isMetalIon(compId, atomName)
                    or self.isMetalIon(prevCompId, prevAtomName)
                    or self.isMetalElem(prevAtomName, prevSeqId, _seqId)):

                self.polySeqPrmTop.append({'chain_id': chainId,
                                           'seq_id': seqIdList,
                                           'auth_comp_id': compIdList})
                seqIdList, compIdList = [], []
                chainIndex += 1
                chainId = indexToLetter(chainIndex)
                offset = 1 - _seqId

            seqId = _seqId + offset
            if seqId not in seqIdList:
                seqIdList.append(seqId)
                compIdList.append(compId)
            self.atomNumberDict[atomNum] = {'chain_id': chainId,
                                            'seq_id': seqId,
                                            'auth_comp_id': compId,
                                            'auth_atom_id': atomName}
            ancAtomName = prevAtomName
            prevAtomName = atomName
            prevAsymId = asymId
            prevSeqId = _seqId
            prevCompId = compId

        if len(self.polySeqPrmTop) == 0:
            if len(seqIdList) > 0:
                self.polySeqPrmTop.append({'chain_id': chainId,
                                           'seq_id': seqIdList,
                                           'auth_comp_id': compIdList})

        self.exit(retrievedAtomNumList)

    def enterComment(self, ctx: BarePDBParser.CommentContext):  # pylint: disable=unused-argument
        """ Enter a parse tree produced by BarePDBParser#comment.
        """

    def exitComment(self, ctx: BarePDBParser.CommentContext):  # pylint: disable=unused-argument
        """ Exit a parse tree produced by BarePDBParser#comment.
        """

    def enterCoordinates(self, ctx: BarePDBParser.CoordinatesContext):  # pylint: disable=unused-argument
        """ Enter a parse tree produced by BarePDBParser#coordinates.
        """

        self.coordinatesStatements += 1

    def exitCoordinates(self, ctx: BarePDBParser.CoordinatesContext):  # pylint: disable=unused-argument
        """ Exit a parse tree produced by BarePDBParser#coordinates.
        """

    def enterAtom_coordinate(self, ctx: BarePDBParser.Atom_coordinateContext):  # pylint: disable=unused-argument
        """ Enter a parse tree produced by BarePDBParser#atom_coordinate.
        """

    def exitAtom_coordinate(self, ctx: BarePDBParser.Atom_coordinateContext):
        """ Exit a parse tree produced by BarePDBParser#atom_coordinate.
        """

        try:

            if self.__end:
                return

            nr = self.cur_nr

            if nr < 0 or nr <= self.prev_nr:
                return

            self.prev_nr = nr

            if ctx.Integer(1):
                chainId = str(ctx.Integer(0))
                seqId = int(str(ctx.Integer(1)))
            elif ctx.Integer(0):
                seqId = int(str(ctx.Integer(0)))
                chainId = str(ctx.Simple_name()) if ctx.Simple_name() else None
            elif ctx.Integer_concat_alt():
                seqId = int(str(ctx.Integer_concat_alt())[:-1])
                chainId = str(ctx.Simple_name()) if ctx.Simple_name() else None
            else:
                concat_string = str(ctx.Simple_name())
                integers = re.findall(r'\d+', concat_string)
                seqId = int(integers[0])
                chainId = concat_string[:concat_string.index(integers[0])]

            atomId = self.atomNameSelection[0]

            if atomId is None:
                return

            compId = self.atomNameSelection[1]

            if atomId.endswith('*'):
                _, nucleotide, _ = self.csStat.getTypeOfCompId(translateToStdResName(compId, ccU=self.ccU))
                atomId = atomId[:-1] + ("'" if nucleotide and not atomId[0].isdigit() else "")

            atom = {'atom_number': nr + self.__ter_offset,
                    'auth_chain_id': chainId,
                    'auth_seq_id': seqId,
                    'auth_comp_id': compId,
                    'auth_atom_id': atomId}

            if atom not in self.atoms:
                self.atoms.append(atom)

        except (ValueError, IndexError):
            pass

        finally:
            self.atomNameSelection.clear()

    def enterAtom_num(self, ctx: BarePDBParser.Atom_numContext):
        """ Enter a parse tree produced by BarePDBParser#atom_num.
        """

        if ctx.Atom() and ctx.Integer():
            self.cur_nr = int(str(ctx.Integer()))
        elif ctx.Hetatm() and ctx.Integer():
            self.cur_nr = int(str(ctx.Integer()))
        elif ctx.Hetatm_decimal():
            concat_string = str(ctx.Hetatm_decimal())
            self.cur_nr = int(concat_string[6:])
        else:
            self.cur_nr = -1

    def exitAtom_num(self, ctx: BarePDBParser.Atom_numContext):  # pylint: disable=unused-argument
        """ Exit a parse tree produced by BarePDBParser#atom_num.
        """

    def enterAtom_name(self, ctx: BarePDBParser.Atom_nameContext):
        """ Enter a parse tree produced by BarePDBParser#atom_name.
        """

        if ctx.Simple_name():
            self.atomNameSelection.append(str(ctx.Simple_name()))
        else:
            self.atomNameSelection.append(str(ctx.Integer_concat_alt()))

    def exitAtom_name(self, ctx: BarePDBParser.Atom_nameContext):  # pylint: disable=unused-argument
        """ Exit a parse tree produced by BarePDBParser#atom_name.
        """

    def enterXyz(self, ctx: BarePDBParser.XyzContext):  # pylint: disable=unused-argument
        """ Enter a parse tree produced by BarePDBParser#xyz.
        """

    def exitXyz(self, ctx: BarePDBParser.XyzContext):  # pylint: disable=unused-argument
        """ Exit a parse tree produced by BarePDBParser#xyz.
        """

    def enterX_yz(self, ctx: BarePDBParser.X_yzContext):  # pylint: disable=unused-argument
        """ Enter a parse tree produced by BarePDBParser#x_yz.
        """

    def exitX_yz(self, ctx: BarePDBParser.X_yzContext):  # pylint: disable=unused-argument
        """ Exit a parse tree produced by BarePDBParser#x_yz.
        """

    def enterXy_z(self, ctx: BarePDBParser.Xy_zContext):  # pylint: disable=unused-argument
        """ Enter a parse tree produced by BarePDBParser#xy_z.
        """

    def exitXy_z(self, ctx: BarePDBParser.Xy_zContext):  # pylint: disable=unused-argument
        """ Exit a parse tree produced by BarePDBParser#xy_z.
        """

    def enterX_y_z(self, ctx: BarePDBParser.X_y_zContext):  # pylint: disable=unused-argument
        """ Enter a parse tree produced by BarePDBParser#x_y_z.
        """

    def exitX_y_z(self, ctx: BarePDBParser.X_y_zContext):  # pylint: disable=unused-argument
        """ Exit a parse tree produced by BarePDBParser#x_y_z.
        """

    def enterUndefined(self, ctx: BarePDBParser.UndefinedContext):  # pylint: disable=unused-argument
        """ Enter a parse tree produced by BarePDBParser#undefined.
        """

    def exitUndefined(self, ctx: BarePDBParser.UndefinedContext):  # pylint: disable=unused-argument
        """ Exit a parse tree produced by BarePDBParser#undefined.
        """

    def enterNumber(self, ctx: BarePDBParser.NumberContext):  # pylint: disable=unused-argument
        """ Enter a parse tree produced by BarePDBParser#number.
        """

    def exitNumber(self, ctx: BarePDBParser.NumberContext):  # pylint: disable=unused-argument
        """ Exit a parse tree produced by BarePDBParser#number.
        """

    def enterTerminal(self, ctx: BarePDBParser.TerminalContext):  # pylint: disable=unused-argument
        """ Enter a parse tree produced by BarePDBParser#terminal.
        """

    def exitTerminal(self, ctx: BarePDBParser.TerminalContext):
        """ Exit a parse tree produced by BarePDBParser#terminal.
        """

        self.atoms.append('TER')
        self.__ter_count += 1

        if ctx.Any_name(0):
            try:
                if int(str(ctx.Any_name(0))) == self.cur_nr + 1:
                    self.__ter_offset -= 1
            except ValueError:
                pass

    def enterEnd(self, ctx: BarePDBParser.EndContext):  # pylint: disable=unused-argument
        """ Enter a parse tree produced by BarePDBParser#end.
        """

        self.__end = True

    def exitEnd(self, ctx: BarePDBParser.EndContext):  # pylint: disable=unused-argument
        """ Exit a parse tree produced by BarePDBParser#end.
        """

    def getContentSubtype(self) -> dict:
        """ Return content subtype of Bare PDB file.
        """

        contentSubtype = {'coordinates': self.coordinatesStatements}

        return {k: 1 for k, v in contentSubtype.items() if v > 0}
