from rdkit import Chem
from ase import atoms
from typing import Optional
import numpy as np
from rdkit.Geometry import Point3D


def MolToAtoms(mol: Chem.Mol, confID=-1) -> atoms.Atoms:
    conf = mol.GetConformer(confID)
    symbols = [atom.GetSymbol() for atom in mol.GetAtoms()]
    positions = [
        (conf.GetAtomPosition(i).x,
         conf.GetAtomPosition(i).y,
         conf.GetAtomPosition(i).z)
        for i in range(mol.GetNumAtoms())
    ]
    aseAtoms = atoms.Atoms(symbols=symbols, positions=positions)
    charges = [atom.GetFormalCharge() for atom in mol.GetAtoms()]
    aseAtoms.set_initial_charges(charges)

    return aseAtoms
    

def AtomsToMol(
    atoms: atoms.Atoms,
    mol: Optional[Chem.Mol] = None,
    conf_id: int = 0,
    charge: Optional[int] = None,
    allow_reorder: bool = False,
    inplace = False
) -> Chem.Mol:
    """
    Convert ASE Atoms -> RDKit Mol.

    If mol is provided:
        - verify natoms and element symbols match (unless allow_reorder=True)
        - update (or add) conformer coordinates so mol matches atoms

    If mol is None:
        - create a new Mol with atoms only, set 3D coords
        - Determine bonds from geometry using rdDetermineBonds.DetermineBonds()

    Parameters
    ----------
    atoms : ase.Atoms
        Must have positions (Å).
    mol : rdkit.Chem.Mol | None
        Optional template mol.
    conf_id : int
        Conformer id to update/use. If mol has no conformer, one will be added.
    charge : int | None
        Total molecular charge used by DetermineBonds (recommended if known).
        If None, will try to infer from formal charges when mol is provided;
        otherwise defaults to 0 for new mol.
    allow_reorder : bool
        If True, will not enforce symbol-by-index equality (only checks counts).
        Most of the time you want False to guarantee consistency.
    inplace : bool
        If True, will modify input mol instead of make a copy.

    Returns
    -------
    rdkit.Chem.Mol
    """
    positions = np.asarray(atoms.get_positions(), dtype=float)
    symbols = list(atoms.get_chemical_symbols())
    n = len(symbols)

    # -------- case 1: update existing mol --------
    if mol is not None:
        if mol.GetNumAtoms() != n:
            raise ValueError(f"mol has {mol.GetNumAtoms()} atoms but ASE atoms has {n}.")

        mol_symbols = [a.GetSymbol() for a in mol.GetAtoms()]
        if not allow_reorder:
            if mol_symbols != symbols:
                raise ValueError(
                    "Element symbols mismatch by index between mol and atoms.\n"
                    f"mol:   {mol_symbols}\n"
                    f"atoms: {symbols}\n"
                    "If you REALLY know what you're doing, set allow_reorder=True "
                    "and handle mapping yourself."
                )
        else:
            # only check multiset counts
            if sorted(mol_symbols) != sorted(symbols):
                raise ValueError("Element symbol counts differ between mol and atoms.")

        if not inplace:
            mol = Chem.Mol(mol)  # copy
            if mol.GetNumConformers() == 0:
                conf = Chem.Conformer(n)
                conf.Set3D(True)
                mol.AddConformer(conf, assignId=True)

        # pick conformer
        try:
            conf = mol.GetConformer(conf_id)
        except ValueError:
            # create new conformer if requested id doesn't exist
            conf = Chem.Conformer(n)
            conf.Set3D(True)
            mol.AddConformer(conf, assignId=True)
            conf = mol.GetConformer(mol.GetNumConformers() - 1)
            if conf_id!=0:
                print("Warning: Failed to pick conformer.")

        for i in range(n):
            x, y, z = map(float, positions[i])
            conf.SetAtomPosition(i, Point3D(x, y, z))

        # charge inference if not given
        if charge is None:
            charge = int(sum(a.GetFormalCharge() for a in mol.GetAtoms()))

        return mol

    # -------- case 2: build mol + determine bonds --------
    rw = Chem.RWMol()
    for sym in symbols:
        rw.AddAtom(Chem.Atom(sym))

    new_mol = rw.GetMol()
    conf = Chem.Conformer(n)
    conf.Set3D(True)
    for i in range(n):
        x, y, z = map(float, positions[i])
        conf.SetAtomPosition(i, Point3D(x, y, z))
    new_mol.AddConformer(conf, assignId=True)

    # Determine bonds from geometry
    if charge is None:
        charge = 0

    try:
        from rdkit.Chem import rdDetermineBonds
        rdDetermineBonds.DetermineBonds(new_mol, charge=charge)
    except Exception as e:
        raise RuntimeError(
            "DetermineBonds failed. This can happen for metals/ions/odd geometries.\n"
            "Consider providing a template mol, or implement custom distance-based bonding.\n"
            f"Original error: {e}"
        )

    return new_mol