from __future__ import print_function

import re
from itertools import chain
from typing import Set

import rdkit.Chem as Chem


def canonicalize_outcome_smiles(smiles: str, ensure: bool = True) -> str:
    """
    Uniquify via SMILES string.

    Need a full SMILES->MOL->SMILES cycle to get a true canonical string.
    Split by '.' and sort when outcome contains multiple molecules.

    Args:
        smiles (str): SMILES string
        ensure (bool): Whether to do SMILES->MOL-SMILES cycle.
            Otherwise only sort molecules within string.

    Returns:
        str: canonicalized, sorted SMILES string

    """
    if ensure:
        outcome = Chem.MolFromSmiles(smiles)
        if outcome is None:
            raise ValueError(f"Invalid SMILES: {smiles}")

        smiles = Chem.MolToSmiles(outcome, True)

    return ".".join(sorted(smiles.split(".")))


def combine_enantiomers_into_racemic(final_outcomes: Set[str]) -> Set[str]:
    """
    Collapse simple enantiomeric/alkene-stereoisomer pairs into an achiral SMILES.

    For each outcome SMILES, this function looks for a corresponding variant that differs
    only by a single tetrahedral chirality inversion (``@`` vs ``@@``) or by flipping the
    alkene stereochemistry slashes (``/`` vs ``\\``) around one double bond. When such a
    pair is found, both stereospecific SMILES are removed and replaced with a re-parsed,
    canonicalized SMILES with that specific stereochemical annotation stripped.

    Note that this is intentionally heuristic and limited: it does not attempt to invert
    multiple stereocenters at once, and changes in stereochemistry can, in principle,
    affect canonicalization (e.g., bond ordering) in ways this logic does not model.

    Args:
        final_outcomes (Set[str]): Set of product SMILES to deduplicate. This set is
            mutated in place and also returned.

    Returns:
        Set[str]: The updated set of SMILES outcomes with certain enantiomeric/cis-trans
            duplicates collapsed.

    Raises:
        ValueError: If the intermediate SMILES generated by stripping stereochemistry
            cannot be parsed by RDKit.
    """

    for smiles in list(final_outcomes)[:]:
        # Look for @@ tetrahedral center
        for match in re.finditer(r"@@", smiles):
            smiles_inv = "%s@%s" % (smiles[: match.start()], smiles[match.end() :])
            if smiles_inv in final_outcomes:
                if smiles in final_outcomes:
                    final_outcomes.remove(smiles)
                final_outcomes.remove(smiles_inv)
                # Re-parse smiles so that hydrogens can become implicit
                smiles = smiles[: match.start()] + smiles[match.end() :]
                outcome = Chem.MolFromSmiles(smiles)
                if outcome is None:
                    raise ValueError("Horrible mistake when fixing duplicate!")
                smiles = ".".join(sorted(Chem.MolToSmiles(outcome, True).split(".")))
                final_outcomes.add(smiles)

        # Look for // or \\ trans bond
        # where [^=\.] is any non-double bond or period or slash
        for match in chain(
            re.finditer(r"(\/)([^=\.\\\/]+=[^=\.\\\/]+)(\/)", smiles),
            re.finditer(r"(\\)([^=\.\\\/]+=[^=\.\\\/]+)(\\)", smiles),
        ):
            # See if cis version is present in list of outcomes
            opposite = {"\\": "/", "/": "\\"}
            smiles_cis1 = "%s%s%s%s%s" % (
                smiles[: match.start()],
                match.group(1),
                match.group(2),
                opposite[match.group(3)],
                smiles[match.end() :],
            )
            smiles_cis2 = "%s%s%s%s%s" % (
                smiles[: match.start()],
                opposite[match.group(1)],
                match.group(2),
                match.group(3),
                smiles[match.end() :],
            )
            # Also look for equivalent trans
            smiles_trans2 = "%s%s%s%s%s" % (
                smiles[: match.start()],
                opposite[match.group(1)],
                match.group(2),
                opposite[match.group(3)],
                smiles[match.end() :],
            )
            # Kind of weird remove conditionals...
            remove = False
            if smiles_cis1 in final_outcomes:
                final_outcomes.remove(smiles_cis1)
                remove = True
            if smiles_cis2 in final_outcomes:
                final_outcomes.remove(smiles_cis2)
                remove = True
            if smiles_trans2 in final_outcomes and smiles in final_outcomes:
                final_outcomes.remove(smiles_trans2)
            if remove:
                final_outcomes.remove(smiles)
                smiles = (
                    smiles[: match.start()] + match.group(2) + smiles[match.end() :]
                )
                outcome = Chem.MolFromSmiles(smiles)
                if outcome is None:
                    raise ValueError("Horrible mistake when fixing duplicate!")
                smiles = ".".join(sorted(Chem.MolToSmiles(outcome, True).split(".")))
                final_outcomes.add(smiles)
    return final_outcomes
