from __future__ import annotations

from typing import Dict, Tuple
import pandas as pd
from rdkit import Chem
from rdkit import RDLogger

RDLogger.DisableLog("rdApp.*") # disable RDKit warnings

def canonicalize_smiles(smiles: str):
    """
    convert smiles to canonical smiles
    :param smiles: smiles string
    :return: returns canonical smiles string, or None if invalid
    """

    if not isinstance(smiles, str):
        print('not a valid smiles string')
        return None

    try:
        mol = Chem.MolFromSmiles(smiles, sanitize=True) # create rdkit mol object from smiles

        if mol is None:
            return None

        return Chem.MolToSmiles(mol, canonical=True)

    except Exception:
        print('exception occurred during canonicalization')
        return None

def canonicalize_smiles_df(df, smiles_column_name, canonical_column_name: str = "canonical_smiles", drop_bad_smiles: bool = True):
    """
    canonicalize a dataframe column of smiles strings
    :param df: pandas dataframe
    :param smiles_column_name: column name to look for smiles strings
    :param canonical_column_name: name of new column to store canonical smiles
    :param drop_bad_smiles: remove the rows with invalid smiles
    :return: original dataframe with new column of canonical smiles and invalid smiles removed (if specified)
    """

    df2 = df.copy()

    df2[canonical_column_name] = df2[smiles_column_name].apply(canonicalize_smiles)

    # find out how many smiles were sanitized by comparing the original and canonical columns
    df2['changed'] = df2[canonical_column_name] != df2[smiles_column_name]
    n_changed = df2['changed'].sum()
    print(len(df2) - n_changed, 'smiles were already canonical,', n_changed, 'smiles were changed during canonicalization')

    if drop_bad_smiles:

        n_before = len(df2)
        df2 = df2.dropna(subset=[canonical_column_name]).reset_index(drop=True)
        n_after = len(df2)

        print(f"dropped {n_before - n_after} invalid smiles")

    return df2

def count_stereocenters_from_smiles(smiles: str) -> Dict[str, int]:

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

    # Ensure RDKit has identified potential stereo bonds
    Chem.FindPotentialStereoBonds(mol)

    stereocenters = Chem.FindMolChiralCenters(mol, includeUnassigned=True)
    stereobonds = [
        bond
        for bond in mol.GetBonds()
        if bond.GetStereo() is not Chem.rdchem.BondStereo.STEREONONE
    ]

    atom_assigned = sum(1 for _, tag in stereocenters if tag != "?")
    atom_unassigned = sum(1 for _, tag in stereocenters if tag == "?")

    bond_assigned = sum(
        1
        for bond in stereobonds
        if bond.GetStereo() is not Chem.rdchem.BondStereo.STEREOANY
    )
    bond_unassigned = sum(
        1 for bond in stereobonds if bond.GetStereo() is Chem.rdchem.BondStereo.STEREOANY
    )

    return {
        "atom_assigned": atom_assigned,
        "atom_unassigned": atom_unassigned,
        "bond_assigned": bond_assigned,
        "bond_unassigned": bond_unassigned,
    }


def filter_unassigned_stereo_rows(df: pd.DataFrame, smiles_col: str = "original_smiles", *, drop_invalid_smiles: bool = True, keep_stereo_counts: bool = False) -> pd.DataFrame:

    counts = []
    bad_smiles_idx = []

    for idx, smi in df[smiles_col].items():
        if not isinstance(smi, str) or not smi.strip():
            bad_smiles_idx.append(idx)
            counts.append({"atom_assigned": 0, "atom_unassigned": 0, "bond_assigned": 0, "bond_unassigned": 0})
            continue

        try:
            c = count_stereocenters_from_smiles(smi)
            counts.append(c)
        except Exception:
            if drop_invalid_smiles:
                bad_smiles_idx.append(idx)
                counts.append({"atom_assigned": 0, "atom_unassigned": 0, "bond_assigned": 0, "bond_unassigned": 0})
            else:
                raise

    counts_df = pd.DataFrame(counts, index=df.index)

    if keep_stereo_counts:
        df2 = df.join(counts_df)
    else:
        df2 = df.copy()

    # mask rows to keep: no unassigned stereo AND not bad smiles (if dropping invalid)
    ok = (counts_df["atom_unassigned"] == 0) & (counts_df["bond_unassigned"] == 0)

    if drop_invalid_smiles and bad_smiles_idx:
        ok = ok & (~df2.index.isin(bad_smiles_idx))

    return df2.loc[ok].reset_index(drop=True)


if __name__ == '__main__':

    df = pd.read_csv('/Users/mimis_stuff/PycharmProjects/PythonProject/philicity_prediction/data/all_philicities.csv')

    df = canonicalize_smiles_df(df, 'radical_smiles', 'canonical_radical_smiles', True)

    df.to_csv('/Users/mimis_stuff/PycharmProjects/PythonProject/philicity_prediction/data/all_philicities_canonicalized.csv', index=False)