"""
Author:
    <NAME>, PhD Student
Institution:
    Ishida Laboratory, Department of Computer Science, School of Computing, Tokyo Institute of Technology
Updated on:
    May, 2021
"""

from typing import Union

from rdkit.Chem.AllChem import Mol, MolFromSmarts, MolFromSmiles, MolToSmarts, MolToSmiles, SanitizeMol


class CompoundRepresentationUtils:
    """
    Description:
        Group of methods for handling chemical compound representation conversion.
    """

    @staticmethod
    def string_to_mol(compound: str, str_format: str = "smiles", sanitize_result: bool = True,
                      verbose: bool = False) -> Union[Mol, None]:
        """
        Description:
            Convert a chemical compound string representation to a Mol object. The currently supported string
            representations are SMILES and SMARTS.
        Input:
            compound (str): The string representation of the chemical compound structure.
            str_format (str): The indicator of the string representation format.
            verbose (bool): The flag signalizing whether feedback messages should be displayed in the console.
        Output:
            (Mol or None): The chemical compound structure as a Mol object or None if the conversion fails.
        """

        mol_object = None

        try:
            if str_format.lower() == "smiles":
                mol_object = MolFromSmiles(compound)
            elif str_format.lower() == "smarts":
                mol_object = MolFromSmarts(compound)
            else:
                raise Exception("Currently supported 'str_format' parameter values are 'smiles' and 'smarts'. "
                                "Got: '{}'.".format(str_format))

            if sanitize_result:
                SanitizeMol(mol_object)

            return mol_object

        except Exception as exc_msg:
            if verbose:
                if mol_object is None:
                    print("Exception occurred during the conversion of the string representation to the Mol object. ",
                          end="")
                else:
                    print("Exception occurred during the sanitization of the constructed Mol object. ", end="")

                print("Detailed exception message: \n{}".format(exc_msg))

            return None

    @staticmethod
    def mol_to_string(compound: Mol, str_format: str = "smiles", additional_transform: bool = True,
                      verbose: bool = False) -> Union[str, None]:
        """
        Description:
            Convert a chemical compound Mol object to a string representation. The currently supported string
            representations are SMILES and SMARTS.
        Input:
            compound (Mol): The Mol object of the chemical compound structure.
            str_format (str): The indicator of the string representation format.
            additional_transform (bool): The flag signalizing whether the output string representation will be canonical
                                         or isomeric in the case of SMILES and SMARTS strings, respectively.
            verbose (bool): The flag signalizing whether feedback messages should be displayed in the console.
        Output:
            (str or None): The chemical compound structure as a SMILES or SMARTS string or None if the conversion fails.
        """

        try:
            if str_format.lower() == "smiles":
                return MolToSmiles(compound, canonical=additional_transform)
            elif str_format.lower() == "smarts":
                return MolToSmarts(compound, isomericSmiles=additional_transform)
            else:
                raise Exception("Currently supported 'str_format' parameter values are 'smiles' and 'smarts'. "
                                "Got: '{}'.".format(str_format))

        except Exception as exc_msg:
            if verbose:
                print("Exception occurred during the conversion of the Mol object to the string representation. "
                      "Detailed exception message: \n{}".format(exc_msg))

            return None

    @staticmethod
    def string_to_string(compound: str, in_str_format: str = "smiles", out_str_format: str = "smiles",
                         sanitize_result: bool = True, additional_transform: bool = True,
                         verbose: bool = False) -> Union[str, None]:
        """
        Description:
            Convert a given chemical compound string representation to another string representation via the conversion
            to a Mol object. The currently supported string representations are SMILES and SMARTS. This primarily
            enables the sanitization of the compound, and additionally the re-conversion to a canonical or isomeric
            string representation in the case of SMILES and SMARTS strings, respectively.
        Input:
            compound (str): The string representation of the chemical compound structure.
            in_str_format (str): The indicator of the input string representation format.
            out_str_format (str): The indicator of the output string representation format.
            additional_transform (bool): The flag signalizing whether the output string representation will be canonical
                                         or isomeric in the case of SMILES and SMARTS strings, respectively.
            verbose (bool): The flag signalizing whether feedback messages should be displayed in the console.
        Output:
            (str or None): The chemical compound structure as a SMILES or SMARTS string or None if the conversion fails.
        """

        return CompoundRepresentationUtils.mol_to_string(
            compound=CompoundRepresentationUtils.string_to_mol(
                compound=compound,
                str_format=in_str_format,
                sanitize_result=sanitize_result,
                verbose=verbose
            ),
            str_format=out_str_format,
            additional_transform=additional_transform,
            verbose=verbose
        )
