# mypy: disable-error-code="attr-defined"

__all__ = ["sanitize"]

from itertools import product
import rdkit
import rdkit.Chem.AllChem as Chem


def sanitize(mol: Chem.Mol, max_fix_iterations: int = 1000) -> None:
    """
    Fix small issues with RDKit SanitizeMol and sanitize molecule.

    This is an alternative to using :func:`SanitizeMol()` directly, in cases it fails
    due to issues with the molecule which are fixed by this function.

    Parameters
    ----------
    mol : Chem.rdchem.Mol
        The molecule to sanitize.
    max_fix_iterations : int, optional
        The maximum number of iterations to fix problems.

    Notes
    -----
    Deals with cases:

    - ``AtomValenceException``: add charge if total valence is exceeding default (for B, N, O atoms only)
    - ``KekulizeException``: for ring N atoms with unspecified protonation
    """
    mol.UpdatePropertyCache(strict=False)
    prev_problems: list[Exception] = []

    # RDKit detects problems iteratively,
    # i.e. new problems may show up as soon as previous problems are fixed
    encountered_same_problem = False
    for _ in range(max_fix_iterations):
        if encountered_same_problem:
            # Fixing this problem previously failed,
            # so there is no need to try again
            break
        with rdkit.rdBase.BlockLogs():
            # Temporarily disable RDKit verbosity while in scope
            problems = Chem.DetectChemistryProblems(mol)
        if not problems:
            # All problematic places have been fixed
            break
        for problem in problems:
            for prev_problem in prev_problems:
                if _is_same_problem(problem, prev_problem):
                    encountered_same_problem = True
        for problem in problems:
            _fix_problem(mol, problem)
        prev_problems = list(problems)

    with rdkit.rdBase.BlockLogs():
        # if any issue remains - it will be caught as an Error
        Chem.SanitizeMol(mol)


def _fix_problem(mol: Chem.Mol, problem: Exception) -> None:
    """
    Apply fixes for common input problems with RDKit SanitizeMol.

    Parameters
    ----------
    mol : Chem.rdchem.Mol
        The molecule to fix.
    problem : Exception
        The problem to fix.
    """
    pt = Chem.GetPeriodicTable()
    if problem.GetType() == "AtomValenceException":
        at = mol.GetAtomWithIdx(problem.GetAtomIdx())
        elem = at.GetSymbol()
        default_valence = pt.GetDefaultValence(at.GetAtomicNum())
        if elem == "C":
            # we do not like charged carbons, thus pass without modification
            return
        elif elem in ["N", "O"]:
            # only process positively charged N, O atoms
            # note: the gain of extra valence interpreted through sharing
            # their lone el. pairs, thus a positive formal charge
            opt_charge = at.GetTotalValence() - default_valence
            formal_charge = at.GetFormalCharge()
            if opt_charge > formal_charge:
                if abs(opt_charge) > 1:
                    raise AssertionError(
                        f"expected N/O atom net charge > 1: {opt_charge}"
                    )
                else:
                    # fix explicit valence issue - setd to +1
                    at.SetFormalCharge(opt_charge)
        elif elem == "B":
            # or negatively charged B atoms
            opt_charge = (
                default_valence - at.GetTotalValence()
            )  # note: reversed as B gets extra valence by accepting el. pair
            formal_charge = at.GetFormalCharge()
            if opt_charge < formal_charge:
                if abs(opt_charge) > 1:
                    raise AssertionError(
                        f"expected B atom net charge > 1: {opt_charge}"
                    )
                else:
                    # fix explicit valence issue - setd to -1
                    at.SetFormalCharge(opt_charge)
        else:
            # for higher elements like S - may need to use pt.GetValenceList()
            # currently not implemented
            return
    if problem.GetType() == "KekulizeException":
        # get all atoms in the same ring as the problem atoms
        problem_indices = set(problem.GetAtomIndices())
        # need to extract ring info
        Chem.SanitizeMol(
            mol=mol, sanitizeOps=Chem.rdmolops.SanitizeFlags.SANITIZE_SYMMRINGS
        )
        ring_info = mol.GetRingInfo()
        for ring_indices in ring_info.AtomRings():
            if problem_indices.intersection(ring_indices):
                problem_indices |= set(ring_indices)
        # if KekulizeException is related to ring N atoms - it could be due to
        # either incorrectly specified explicit protonation or charge
        original_nH_state = []
        original_charge_state = []
        possible_nHs = []
        for atidx in problem_indices:
            at = mol.GetAtomWithIdx(atidx)
            num_heavy_neighbors = at.GetTotalDegree() - at.GetTotalNumHs()
            if at.GetSymbol() == "N" and num_heavy_neighbors == 2:
                original_nH_state.append(at.GetNumExplicitHs())
                original_charge_state.append(at.GetFormalCharge())
                possible_nHs.append(at)

        if not len(possible_nHs):
            # no candidate nitrogens to fix the problem
            # note this early return is unnecessary as the loops below will
            # not be entered, but it makes the intention clearer
            return

        # first test if the issue is the formal charge due to explicit protonation
        for ni, nH, charge in zip(
            possible_nHs, original_nH_state, original_charge_state
        ):
            if nH > 0 and charge == 0:
                # if we found a protonated nitrogen, we can fix the charge
                ni.SetFormalCharge(nH)
                # and check if the problem is resolved
                check_problems = Chem.DetectChemistryProblems(
                    mol, sanitizeOps=Chem.rdmolops.SanitizeFlags.SANITIZE_KEKULIZE
                )
                if not [p for p in check_problems if _is_same_problem(p, problem)]:
                    # problem is resolved - no need to try other combinations
                    return
                else:
                    # if the problem is not resolved - reset the charge to try other candidates
                    ni.SetFormalCharge(charge)

        # alternatively, attempt to fix by setting one or more ring nitrogens as "[nH1]"
        # if there is one or more candidates - iterate through all the possible
        # combinations of explicit Hs for the candidate nitrogens
        for combo in product([0, 1], repeat=len(possible_nHs)):
            for ni, nH_explicit in zip(possible_nHs, combo):
                ni.SetNumExplicitHs(nH_explicit)
            # and check if the problem is resolved
            check_problems = Chem.DetectChemistryProblems(
                mol, sanitizeOps=Chem.rdmolops.SanitizeFlags.SANITIZE_KEKULIZE
            )
            if not [p for p in check_problems if _is_same_problem(p, problem)]:
                # problem is resolved - no need to try other combinations
                return
        # if all combos failed - reset explicit Hs back to the original values
        for ni, nH_explicit in zip(possible_nHs, original_nH_state):
            ni.SetNumExplicitHs(nH_explicit)


def _is_same_problem(problem1: Exception, problem2: Exception) -> bool:
    """
    Check if two exceptions related to chemistry problems are the same.

    Parameters
    ----------
    problem1, problem2 : Exception
        The problems to compare.

    Returns
    -------
    same : bool
        True if the two problems are the same.
    """
    problem_type = problem1.GetType()
    if problem_type != problem2.GetType():
        return False
    elif problem_type == "AtomValenceException":
        if problem1.GetAtomIdx() != problem2.GetAtomIdx():
            return False
    elif problem_type == "KekulizeException":
        if (
            len(
                set(problem1.GetAtomIndices()).intersection(
                    set(problem2.GetAtomIndices())
                )
            )
            == 0
        ):
            return False
    return True
