from typing import Optional

from chemrecon.chem.mol import feats, Mol
from chemrecon.chem.create_mol import mol_from_smiles
from chemrecon.schema import Column, FeatureEnum, ProceduralGeneratorError, ProceduralRelation, MolStructure


class MolStructureStandardization(ProceduralRelation[MolStructure, MolStructure]):
    """ Standardization of molecular structures according to a particular 'feature' (Fragment, Isotope, Charge,
        Tautomerism, Stereochemical).
    """

    # Attributes
    feat: FeatureEnum  #: The feature w.r.t. which the structure is standardized.

    # Database
    entrytype_name = 'MolStructure standardisation'
    _table_name = 'MolStructureStandardisation'
    ignore_generation_limit = True
    symmetric = False
    source_entrytype = MolStructure
    target_entrytype = MolStructure
    _attribute_columns = [
        Column('feat', FeatureEnum)
    ]
    _index = [0]

    def __init__(
        self,
        feat: FeatureEnum,
        recon_id_1: Optional[int] = None,
        recon_id_2: Optional[int] = None
    ):
        super().__init__(recon_id_1, recon_id_2)
        self.feat = feat

    @classmethod
    def generate(
        cls, take_entry: MolStructure
    ) -> list[tuple[ProceduralRelation[MolStructure, MolStructure], MolStructure]]:

        # Create mol object
        m: Mol
        try:
            m = mol_from_smiles(take_entry.smiles)
        except ValueError as e:
            # Mol creation fails
            raise ProceduralGeneratorError('Could not create mol for std.') from e

        feat_entries: list[tuple[FeatureEnum, MolStructure]] = list()
        for f in feats:
            if f not in m.features:
                try:
                    m_: Mol = Mol(f.standardise(m))
                except Exception as e:
                    # Failed standardisation
                    continue

                feat_entries.append((f.feature_enum, m_.to_database_struct()))

        return [
            (
                MolStructureStandardization(feat = f),
                e
            )
            for f, e in feat_entries
        ]

    # Visualisation
    def _vis_str(self) -> str:
        return f'std: {self.feat.name}'
