"""Monomer class with parameters and methods."""

from typing import Tuple, Union

import pandas as pd
from rdkit import Chem
from rdkit.Chem.Descriptors import ExactMolWt
from rdkit import rdBase
from rdkit.Chem.rdchem import Mol

functionality_smarts = {
    "ols": "[C,c;!$(C=O)][OH]",
    "aliphatic_ols": "[C;!$(C=O);!$([a])][OH]",
    "acids": "[#6][#6](=[#8:4])([F,Cl,Br,I,#8H,O-])",
    "prime_amines": "[#6;!$(C=O)][NH2;!$([NH2+])]",
    "carbonates": "[O]=[C]([F,Cl,Br,I,O])([F,Cl,Br,I,O])",
    "acidanhydrides": "[#8]([#6](=[#8]))([#6](=[#8]))",
    "prime_thiols": "[#6;!$(C=O)][SH]",
}


def molecule_from_smiles(smiles: str) -> Union[Mol, None]:
    """Generate rdkit mol from smiles

    Parameters
    ----------
    smiles : str
        SMILES string

    Returns
    -------
    Union[Mol, None]
        RDKit Molecule, or None if can't generate
    """
    try:
        mol = Chem.MolFromSmiles(smiles)
    except:
        mol = None

    return mol


def get_functionality(reactants, distribution=[]):
    """gets the functional groups from a list of reactants

    inputs: list of smiles
    output: dataframe with count of functional groups
    """

    def id_functionality(r):
        mol = Chem.MolFromSmiles(r.name)
        r.ols = len(
            mol.GetSubstructMatches(Chem.MolFromSmarts(functionality_smarts["ols"]))
        )
        r.aliphatic_ols = len(
            mol.GetSubstructMatches(
                Chem.MolFromSmarts(functionality_smarts["aliphatic_ols"])
            )
        )
        r.acids = len(
            mol.GetSubstructMatches(Chem.MolFromSmarts(functionality_smarts["acids"]))
        )
        r.prime_amines = len(
            mol.GetSubstructMatches(
                Chem.MolFromSmarts(functionality_smarts["prime_amines"])
            )
        )
        r.carbonates = len(
            mol.GetSubstructMatches(
                Chem.MolFromSmarts(functionality_smarts["carbonates"])
            )
        )
        r.acidanhydrides = len(
            mol.GetSubstructMatches(
                Chem.MolFromSmarts(functionality_smarts["acidanhydrides"])
            )
        )
        return r

    df_func = pd.DataFrame(
        data=0,
        index=reactants,
        columns=[
            "ols",
            "acids",
            "prime_amines",
            "carbonates",
            "aliphatic_ols",
            "acidanhydrides",
        ],
    )
    df_func = df_func.apply(lambda r: id_functionality(r), axis=1)

    # appends distribution to dataframe

    if len(distribution) == 0:
        df_func["distribution"] = [1] * df_func.shape[0]
    else:
        df_func["distribution"] = list(distribution)
    return df_func


def enumerate_ester_enantiomers(smiles: str) -> Tuple[Tuple[str]]:
    """Generate enantiomer pairs for a monomer that would participate in an esterification reaction

    Parameters
    ----------
    smiles : str
        SMILES string of the monomer

    Returns
    -------
    Tuple[Tuple[str]]
        The enantiomer pairs, or a single string if no enantiomers are created
    """
    try:
        mol = Chem.MolFromSmiles(smiles)
        # Get the atom ids for the acid and ol functionalities
        acid = Chem.MolFromSmarts("[O]-[C]=[O]")
        ol = Chem.MolFromSmarts("[C;X4]-[O]")
        acid_atoms = mol.GetSubstructMatch(acid)
        ol_atoms = mol.GetSubstructMatch(ol)

        # Get the carbon atoms for both, doesn't matter for acid, but does for ol
        atoms = mol.GetAtoms()

        for atom_i in acid_atoms:
            if atoms[atom_i].GetAtomicNum() == 8:
                acid_O = atom_i

        for atom_i in ol_atoms:
            if atoms[atom_i].GetAtomicNum() == 8:
                ol_O = atom_i

        # Get shortest path (backbone) and make stereo based on that
        # Does shortest path make sense?
        atom_path = Chem.GetShortestPath(mol, acid_O, ol_O)

        # Get sites that can be stereo
        stereo_sites = [
            sinfo.centeredOn for sinfo in Chem.rdmolops.FindPotentialStereo(mol)
        ]

        # Find the enantiomers based on stereo sites within the shortest path
        enantiomers = []

        for atom_i in set(atom_path).intersection(stereo_sites):
            atoms[atom_i].SetChiralTag(Chem.rdchem.ChiralType.CHI_TETRAHEDRAL_CW)
            smilesR = Chem.MolToSmiles(mol)
            atoms[atom_i].SetChiralTag(Chem.rdchem.ChiralType.CHI_TETRAHEDRAL_CCW)
            smilesS = Chem.MolToSmiles(mol)

            enantiomers.append((smilesR, smilesS))
            atoms[atom_i].SetChiralTag(Chem.rdchem.ChiralType.CHI_UNSPECIFIED)

    except BaseException as e:
        enantiomers = [[smiles]]

    if not enantiomers:
        enantiomers = [[smiles]]

    return enantiomers


class Monomer:
    def __init__(self, smiles: str):
        self.smiles = smiles
        self.canonical_smiles = Chem.CanonSmiles(smiles)
        self.molecule = molecule_from_smiles(self.smiles)
        self.molecular_weight = ExactMolWt(self.molecule)

        df = pd.DataFrame([self.smiles], columns=["smiles"])
        self.functionality = get_functionality(df.smiles)

    def __repr__(self) -> str:
        if self.is_valid:
            return f"Valid Monomer with smiles {self.smiles}"
        else:
            return f"Invalid Monomer with smiles {self.smiles}"

    @property
    def esterification_enantiomers(self) -> Tuple[Tuple[str]]:
        """Get possible enantiomers that would participate in an esterification reaction

        Returns
        -------
        Tuple[Tuple[str, str]]
            Tuple containing diads of enantiomers
        """
        return enumerate_ester_enantiomers(self.smiles)
