##
# File: CyanaNOAParserListener.py
# Date: 21-Nov-2024
#
# Updates:
""" ParserLister class for CYANA NOA 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 copy
import itertools
import sys
from typing import IO, List, Optional

from antlr4 import ParseTreeListener

try:
    from wwpdb.utils.nmr.NmrDpConstant import (REPRESENTATIVE_MODEL_ID,
                                               REPRESENTATIVE_ALT_ID,
                                               DIST_AMBIG_LOW,
                                               DIST_AMBIG_UP)
    from wwpdb.utils.nmr.AlignUtil import deepcopy
    from wwpdb.utils.nmr.nef.NefTranslator import NefTranslator
    from wwpdb.utils.nmr.io.CifReader import CifReader
    from wwpdb.utils.nmr.mr.CyanaNOAParser import CyanaNOAParser
    from wwpdb.utils.nmr.mr.BaseLinearMRParserListener import BaseLinearMRParserListener
    from wwpdb.utils.nmr.mr.ParserListenerUtil import (isIdenticalRestraint,
                                                       hasIntraChainRestraint,
                                                       isAmbigAtomSelection,
                                                       getAltProtonIdInBondConstraint,
                                                       getRow,
                                                       getStarAtom,
                                                       resetMemberId,
                                                       getDistConstraintType,
                                                       getPotentialType)
except ImportError:
    from nmr.NmrDpConstant import (REPRESENTATIVE_MODEL_ID,
                                   REPRESENTATIVE_ALT_ID,
                                   DIST_AMBIG_LOW,
                                   DIST_AMBIG_UP)
    from nmr.AlignUtil import deepcopy
    from nmr.nef.NefTranslator import NefTranslator
    from nmr.io.CifReader import CifReader
    from nmr.mr.CyanaNOAParser import CyanaNOAParser
    from nmr.mr.BaseLinearMRParserListener import BaseLinearMRParserListener
    from nmr.mr.ParserListenerUtil import (isIdenticalRestraint,
                                           hasIntraChainRestraint,
                                           isAmbigAtomSelection,
                                           getAltProtonIdInBondConstraint,
                                           getRow,
                                           getStarAtom,
                                           resetMemberId,
                                           getDistConstraintType,
                                           getPotentialType)


class CyanaNOAParserListener(ParseTreeListener, BaseLinearMRParserListener):
    """ This class defines a complete listener for a parse tree produced by CyanaNOAParser.
    """
    __slots__ = ('noeAssignments',
                 'asisList',
                 'weights',
                 'dstFunc')

    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,
                 reasons: Optional[dict] = None) -> None:
        self.__class_name__ = self.__class__.__name__
        self.__version__ = __version__

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

        self.file_type = 'nm-res-noa'
        self.software_name = 'CYANA'

        self.noeAssignments = []
        self.asisList = []
        self.weights = []

        # potential
        self.dstFunc = None

    def enterCyana_noa(self, ctx: CyanaNOAParser.Cyana_noaContext):  # pylint: disable=unused-argument
        """ Enter a parse tree produced by CyanaNOAParser#cyana_noa.
        """

    def exitCyana_noa(self, ctx: CyanaNOAParser.Cyana_noaContext):  # pylint: disable=unused-argument
        """ Exit a parse tree produced by CyanaNOAParser#cyana_noa.
        """

        self.exit()

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

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

    def enterNoe_peaks(self, ctx: CyanaNOAParser.Noe_peaksContext):  # pylint: disable=unused-argument
        """ Enter a parse tree produced by CyanaNOAParser#noe_peaks.
        """

        self.cur_subtype = 'dist'

    def exitNoe_peaks(self, ctx: CyanaNOAParser.Noe_peaksContext):  # pylint: disable=unused-argument
        """ Exit a parse tree produced by CyanaNOAParser#noe_peaks.
        """

    def enterPeak_header(self, ctx: CyanaNOAParser.Peak_headerContext):
        """ Enter a parse tree produced by CyanaNOAParser#peak_header.
        """

        if ctx.Angstrome(0):
            self.dstFunc = self.validateDistanceRange(1.0, None, None, float(str(ctx.Angstrome(0))[:-2]),
                                                      None, self.omitDistLimitOutlier)
        else:
            self.dstFunc = None

    def exitPeak_header(self, ctx: CyanaNOAParser.Peak_headerContext):  # pylint: disable=unused-argument
        """ Exit a parse tree produced by CyanaNOAParser#peak_header.
        """

    def enterPeak_quality(self, ctx: CyanaNOAParser.Peak_qualityContext):
        """ Enter a parse tree produced by CyanaNOAParser#peak_quality.
        """

        if ctx.Float() and self.dstFunc is not None:
            quality = float(str(ctx.Float()))

            if 0.0 < quality <= 1.0:

                if quality < 1.0:
                    self.dstFunc['weight'] = quality

                self.distRestraints += 1

                self.noeAssignments.clear()
                self.asisList.clear()
                self.weights.clear()

    def exitPeak_quality(self, ctx: CyanaNOAParser.Peak_qualityContext):  # pylint: disable=unused-argument
        """ Exit a parse tree produced by CyanaNOAParser#peak_quality.
        """

    def enterNoe_assignments(self, ctx: CyanaNOAParser.Noe_assignmentsContext):  # pylint: disable=unused-argument
        """ Enter a parse tree produced by CyanaNOAParser#noe_assignments.
        """

    def exitNoe_assignments(self, ctx: CyanaNOAParser.Noe_assignmentsContext):  # pylint: disable=unused-argument
        """ Exit a parse tree produced by CyanaNOAParser#noe_assignments.
        """

        if len(self.noeAssignments) > 0 and self.dstFunc is not None:

            try:

                if not self.hasPolySeq and not self.hasNonPolySeq:
                    return

                if self.createSfDict:
                    sf = self.getSf(constraintType=getDistConstraintType(self.noeAssignments[0], self.dstFunc,
                                                                         self.csStat, self.originalFileName),
                                    potentialType=getPotentialType(self.file_type, self.cur_subtype, self.dstFunc))
                    sf['id'] += 1

                for idx, atomSelectionSet in enumerate(self.noeAssignments):
                    self.dstFunc['weight'] = self.weights[idx]
                    memberLogicCode = 'OR' if len(atomSelectionSet[0]) * len(atomSelectionSet[1]) > 1 else '.'
                    has_intra_chain, rep_chain_id_set = hasIntraChainRestraint(atomSelectionSet)

                    if self.createSfDict:
                        if memberLogicCode == 'OR' and has_intra_chain and len(rep_chain_id_set) == 1:
                            memberLogicCode = '.'
                        memberId = '.'
                        if len(self.noeAssignments) > 1:
                            memberId = idx + 1
                            _atom1 = _atom2 = None
                        elif memberLogicCode == 'OR':
                            if len(atomSelectionSet[0]) * len(atomSelectionSet[1]) > 1\
                               and (isAmbigAtomSelection(atomSelectionSet[0], self.csStat)
                                    or isAmbigAtomSelection(atomSelectionSet[1], self.csStat)):
                                memberId = 0
                                _atom1 = _atom2 = None

                    for atom1, atom2 in itertools.product(atomSelectionSet[0],
                                                          atomSelectionSet[1]):
                        atoms = [atom1, atom2]
                        if isIdenticalRestraint(atoms, self.nefT):
                            continue
                        if self.createSfDict and isinstance(memberId, int):  # pylint: disable=used-before-assignment
                            star_atom1 = getStarAtom(self.authToStarSeq, self.authToOrigSeq, self.offsetHolder, copy.copy(atom1))
                            star_atom2 = getStarAtom(self.authToStarSeq, self.authToOrigSeq, self.offsetHolder, copy.copy(atom2))
                            if None in (star_atom1, star_atom2) or isIdenticalRestraint([star_atom1, star_atom2], self.nefT):
                                continue
                        if has_intra_chain and (atom1['chain_id'] != atom2['chain_id']
                                                or atom1['chain_id'] not in rep_chain_id_set):
                            continue
                        if self.createSfDict and memberLogicCode == '.':
                            altAtomId1, altAtomId2 = getAltProtonIdInBondConstraint(atoms, self.csStat)
                            if altAtomId1 is not None or altAtomId2 is not None:
                                atom1, atom2 =\
                                    self.selectRealisticBondConstraint(atom1, atom2,
                                                                       altAtomId1, altAtomId2,
                                                                       self.dstFunc)
                        if self.debug:
                            print(f"subtype={self.cur_subtype} id={self.distRestraints} "
                                  f"atom1={atom1} atom2={atom2} {self.dstFunc}")
                        if self.createSfDict and sf is not None:
                            if isinstance(memberId, int):
                                if _atom1 is None or isAmbigAtomSelection([_atom1, atom1], self.csStat)\
                                   or isAmbigAtomSelection([_atom2, atom2], self.csStat):
                                    if memberId != idx + 1:
                                        memberId += 1
                                    _atom1, _atom2 = atom1, atom2
                            sf['index_id'] += 1
                            row = getRow(self.cur_subtype, sf['id'], sf['index_id'],
                                         '.', memberId, memberLogicCode,
                                         sf['list_id'], self.entryId, self.dstFunc,
                                         self.authToStarSeq, self.authToOrigSeq, self.authToInsCode, self.offsetHolder,
                                         atom1, atom2, asis1=self.asisList[idx][0], asis2=self.asisList[idx][1])
                            sf['loop'].add_data(row)

                            if sf['constraint_subsubtype'] == 'ambi':
                                continue

                            if self.cur_constraint_type is not None and self.cur_constraint_type.startswith('ambiguous'):
                                sf['constraint_subsubtype'] = 'ambi'

                            if memberLogicCode == 'OR'\
                               and (isAmbigAtomSelection(atomSelectionSet[0], self.csStat)
                                    or isAmbigAtomSelection(atomSelectionSet[1], self.csStat)):
                                sf['constraint_subsubtype'] = 'ambi'

                            if 'upper_limit' in self.dstFunc and self.dstFunc['upper_limit'] is not None:
                                upperLimit = float(self.dstFunc['upper_limit'])
                                if upperLimit <= DIST_AMBIG_LOW or upperLimit >= DIST_AMBIG_UP:
                                    sf['constraint_subsubtype'] = 'ambi'

                if self.createSfDict and sf is not None and isinstance(memberId, int) and memberId == 1:
                    sf['loop'].data[-1] = resetMemberId(self.cur_subtype, sf['loop'].data[-1])

            finally:
                self.noeAssignments.clear()
                self.asisList.clear()
                self.weights.clear()

    def enterNoe_assignment(self, ctx: CyanaNOAParser.Noe_assignmentContext):  # pylint: disable=unused-argument
        """ Enter a parse tree produced by CyanaNOAParser#noe_assignment.
        """

    def exitNoe_assignment(self, ctx: CyanaNOAParser.Noe_assignmentContext):
        """ Exit a parse tree produced by CyanaNOAParser#noe_assignment.
        """

        if ctx.Ok():

            try:

                seqId1 = int(str(ctx.Integer(0)))
                compId1 = str(ctx.Simple_name(1)).upper()
                atomId1 = str(ctx.Simple_name(0)).upper()
                seqId2 = int(str(ctx.Integer(1)))
                compId2 = str(ctx.Simple_name(3)).upper()
                atomId2 = str(ctx.Simple_name(2)).upper()

                if not self.hasPolySeq and not self.hasNonPolySeq:
                    return

                self.retrieveLocalSeqScheme()

                chainAssign1, asis1 = self.assignCoordPolymerSequence(seqId1, compId1, atomId1.split('|', 1)[0])
                chainAssign2, asis2 = self.assignCoordPolymerSequence(seqId2, compId2, atomId2.split('|', 1)[0])

                if 0 in (len(chainAssign1), len(chainAssign2)):
                    return

                self.selectCoordAtoms(chainAssign1, seqId1, compId1, atomId1)
                self.selectCoordAtoms(chainAssign2, seqId2, compId2, atomId2)

                if len(self.atomSelectionSet) < 2:
                    return

                self.noeAssignments.append(deepcopy(self.atomSelectionSet))
                self.asisList.append([asis1, asis2])
                self.weights.append(float(str(ctx.Integer(2))) / 100.0)

            finally:
                self.atomSelectionSet.clear()

    def enterNumerical_report(self, ctx: CyanaNOAParser.Numerical_reportContext):  # pylint: disable=unused-argument
        """ Enter a parse tree produced by CyanaNOAParser#numerical_report.
        """

    def exitNumerical_report(self, ctx: CyanaNOAParser.Numerical_reportContext):  # pylint: disable=unused-argument
        """ Exit a parse tree produced by CyanaNOAParser#numerical_report.
        """

    def enterExtended_report(self, ctx: CyanaNOAParser.Extended_reportContext):  # pylint: disable=unused-argument
        """ Enter a parse tree produced by CyanaNOAParser#extended_report.
        """

    def exitExtended_report(self, ctx: CyanaNOAParser.Extended_reportContext):  # pylint: disable=unused-argument
        """ Exit a parse tree produced by CyanaNOAParser#extended_report.
        """

    def enterNoe_stat(self, ctx: CyanaNOAParser.Noe_statContext):  # pylint: disable=unused-argument
        """ Enter a parse tree produced by CyanaNOAParser#noe_stat.
        """

    def exitNoe_stat(self, ctx: CyanaNOAParser.Noe_statContext):  # pylint: disable=unused-argument
        """ Exit a parse tree produced by CyanaNOAParser#noe_stat.
        """

    def enterList_of_proton(self, ctx: CyanaNOAParser.List_of_protonContext):  # pylint: disable=unused-argument
        """ Enter a parse tree produced by CyanaNOAParser#list_of_proton.
        """

    def exitList_of_proton(self, ctx: CyanaNOAParser.List_of_protonContext):  # pylint: disable=unused-argument
        """ Exit a parse tree produced by CyanaNOAParser#list_of_proton.
        """

    def enterPeak_stat(self, ctx: CyanaNOAParser.Peak_statContext):  # pylint: disable=unused-argument
        """ Enter a parse tree produced by CyanaNOAParser#peak_stat.
        """

    def exitPeak_stat(self, ctx: CyanaNOAParser.Peak_statContext):  # pylint: disable=unused-argument
        """ Exit a parse tree produced by CyanaNOAParser#peak_stat.
        """
