from functools import reduce
from io import StringIO

import numpy as np
from ase import Atoms
from ase.data import chemical_symbols as AN2ELEM
from ase.io.xyz import read_xyz
from rdkit import Chem, RDLogger
from rdkit.Chem import AllChem, rdDetermineBonds, rdFreeSASA
from rdkit.Chem import Mol as RDMol
from scipy import sparse as sp


def _handle_with_bond_order(value) -> int:
    if abs(value - int(value)) < 1e-3:
        value = int(value)
    elif abs(value - 1.5) < 1e-3:  # aromatic
        value = 4
    else:
        raise ValueError(f"Invalid bond order {value} for bond.")
    return value


#############################################################################
#          The module for converting data into file (or string)s            #
#############################################################################


@staticmethod
def get_sdf_str(
    numbers: np.ndarray,
    source: np.ndarray,
    target: np.ndarray,
    geometry: np.ndarray | None = None,
    order: np.ndarray | None = None,
) -> str:
    """Convert SD file v3000 string from the given data."""
    source = np.asarray(source, dtype=int)
    target = np.asarray(target, dtype=int)
    numbers = np.asarray(numbers, dtype=int)
    assert numbers.ndim == source.ndim == target.ndim == 1
    assert source.shape[0] == target.shape[0]
    if order is None:
        order = np.zeros_like(source)
    else:
        order = np.vectorize(_handle_with_bond_order)(order)
    assert isinstance(order, np.ndarray), (type(order), order)
    n_atoms = numbers.shape[0]
    _bonds = np.sort(np.column_stack([source, target]), axis=1)
    _idx = _bonds[:, 0] * n_atoms + _bonds[:, 1]
    _, i = np.unique(_idx, return_index=True)
    source = source[i]
    target = target[i]
    order = order[i]
    n_bonds = source.shape[0]
    assert all(i != 0 for i in (n_atoms, n_bonds))
    na, nb = np.array(np.log10([n_atoms, n_bonds]) + 1, dtype=int)
    if geometry is None:
        geometry = np.random.rand(n_atoms, 3)

    # Header
    content = ""
    content += "\n"
    content += "   Generated by Python.\n"
    content += "\n"

    # Counts
    content += "  0  0  0     0  0            999 V3000\n"
    content += "M  V30 BEGIN CTAB\n"
    content += f"M  V30 COUNTS {n_atoms} {n_bonds} 0 0 0\n"

    # Atom block
    content += "M  V30 BEGIN ATOM\n"
    elem = np.vectorize(lambda z: AN2ELEM[z].ljust(2))(numbers)
    x = np.char.rjust(np.char.mod("%.4f", geometry[:, 0]), 10)
    y = np.char.rjust(np.char.mod("%.4f", geometry[:, 1]), 10)
    z = np.char.rjust(np.char.mod("%.4f", geometry[:, 2]), 10)
    idxs = np.char.mod(f"%{na}d ", np.arange(1, n_atoms + 1))
    _data = ["M  V30 ", idxs, elem, x, y, z, " 0\n"]
    content += "".join(reduce(np.char.add, _data))
    content += "M  V30 END ATOM\n"

    # Bond block
    content += "M  V30 BEGIN BOND\n"
    idxs = np.char.mod(f"%{nb}d ", np.arange(1, n_bonds + 1))
    i = np.char.mod(f"%{na}d ", source + 1)
    j = np.char.mod(f"%{na}d ", target + 1)
    o = np.char.mod("%d ", order)
    _data = ["M  V30 ", idxs, o, i, j, " \n"]
    content += "".join(reduce(np.char.add, _data))
    content += "M  V30 END BOND\n"

    # Footer
    content += "M  V30 END CTAB\n"
    content += "M  END\n"
    content += "$$$$\n"

    return content


@staticmethod
def get_xyz_str(
    numbers: np.ndarray,
    geometry: np.ndarray,
) -> str:
    """Convert XYZ file string from the given data."""
    content = ""
    content += f"{numbers.shape[0]}\n"
    content += "\n"  # title

    # Atom block
    elem = np.vectorize(lambda z: AN2ELEM[z].ljust(2))(numbers)
    x = np.char.rjust(np.char.mod("%.4f", geometry[:, 0]), 10)
    y = np.char.rjust(np.char.mod("%.4f", geometry[:, 1]), 10)
    z = np.char.rjust(np.char.mod("%.4f", geometry[:, 2]), 10)
    content += "".join(
        reduce(np.char.add, [" ", elem, "   ", x, y, z, "\n"])
    )  # + '\n'

    return content


#############################################################################
#             RDKit implementation for mol, graph and networkx              #
#############################################################################
RDLogger.DisableLog("rdApp.*")  # type: ignore


def _get_rdmol_without_bonds(
    numbers: np.ndarray,
    geometry: np.ndarray,
    infer_order: bool = False,
    infer_bond: bool = False,
    cov_factor: float = 1.3,
    charge: int = 0,
) -> RDMol:
    """Convert RDKit Mol object from the given data."""
    xyz_block: str = get_xyz_str(
        numbers=numbers,
        geometry=geometry,
    )
    rdmol = Chem.MolFromXYZBlock(xyz_block)
    if rdmol is None:
        raise RuntimeError("The `Mol` object's rdmol is None.")

    if infer_order:
        rdDetermineBonds.DetermineBonds(
            rdmol,
            covFactor=cov_factor,
            charge=charge,
        )
    elif infer_bond:
        rdDetermineBonds.DetermineConnectivity(
            rdmol,
            covFactor=cov_factor,
        )

    return rdmol


def _get_rdmol_with_bonds(
    source: np.ndarray,
    target: np.ndarray,
    numbers: np.ndarray,
    geometry: np.ndarray | None = None,
    order: np.ndarray | None = None,
    infer_order: bool = False,
    charge: int = 0,
) -> RDMol:
    """Convert RDKit Mol object from the given data."""
    sdf_block: str = get_sdf_str(
        numbers=numbers,
        geometry=geometry,
        source=source,
        target=target,
        order=order,
    )
    rdmol = Chem.MolFromMolBlock(sdf_block, removeHs=False)
    if rdmol is None:
        raise RuntimeError("The `Mol` object's rdmol is None.")

    if infer_order:
        assert order is None, (
            "Please set order is None when you want infer bond order by RDKit."
        )
        rdDetermineBonds.DetermineBondOrders(rdmol, charge=charge)

    return rdmol


def get_rdmol(
    numbers: np.ndarray,
    order: np.ndarray | None = None,
    source: np.ndarray | None = None,
    target: np.ndarray | None = None,
    geometry: np.ndarray | None = None,
    infer_order: bool = False,
    infer_bond: bool = False,
    cov_factor: float = 1.3,
    charge: int = 0,
) -> RDMol:
    """Convert RDKit Mol object from the given data."""
    if (source is not None) and (target is not None) and (not infer_bond):
        if infer_order:
            return _get_rdmol_with_bonds(
                numbers=numbers,
                geometry=geometry,
                source=source,
                target=target,
                infer_order=infer_order,
                charge=charge,
            )
        else:
            return _get_rdmol_with_bonds(
                numbers=numbers,
                geometry=geometry,
                source=source,
                target=target,
                order=order,
            )
    else:
        assert isinstance(geometry, np.ndarray), (type(geometry), geometry)
        return _get_rdmol_without_bonds(
            numbers=numbers,
            geometry=geometry,
            infer_bond=infer_bond,
            infer_order=infer_order,
            cov_factor=cov_factor,
            charge=charge,
        )


def rdmol2smiles(rdmol: RDMol, canonical: bool = True) -> str:
    """Convert RDKit Mol object as SMILES string."""
    result = Chem.MolToSmiles(rdmol)
    if canonical:
        result = Chem.CanonSmiles(result)
    return result


def smiles2rdmol(smi: str) -> RDMol:
    """Convert SMILES string as RDKit Mol object."""
    rdmol = Chem.AddHs(Chem.MolFromSmiles(str(smi)))
    AllChem.EmbedMolecule(rdmol, useRandomCoords=True)  # type: ignore
    AllChem.MMFFOptimizeMolecule(rdmol)  # type: ignore
    return rdmol


def get_smiles(
    numbers: np.ndarray,
    source: np.ndarray | None = None,
    target: np.ndarray | None = None,
    geometry: np.ndarray | None = None,
    order: np.ndarray | None = None,
    canonical: bool = True,
    charge: int = 0,
) -> str:
    """Convert SMILES string from the given data."""
    if (source is not None) and (target is not None):
        if order is None:
            rdmol = _get_rdmol_with_bonds(
                numbers=numbers,
                geometry=geometry,
                source=source,
                target=target,
                infer_order=True,
                charge=charge,
            )
        else:
            rdmol = _get_rdmol_with_bonds(
                numbers=numbers,
                geometry=geometry,
                source=source,
                target=target,
                order=order,
            )
    else:
        assert geometry is not None, (
            "Geomtry cannot be None when source or target is None."
        )
        rdmol = _get_rdmol_without_bonds(
            numbers=numbers,
            geometry=geometry,
            infer_bond=True,
            infer_order=True,
            charge=charge,
        )

    return rdmol2smiles(rdmol=rdmol, canonical=canonical)


def _get_structure(rdmol: RDMol) -> tuple[np.ndarray, np.ndarray]:
    xyz = Chem.MolToXYZBlock(mol=rdmol, precision=15)
    atoms = Atoms(list(read_xyz(StringIO(xyz), -1)))
    return atoms.numbers, atoms.positions


def _get_adjacency(rdmol: RDMol) -> sp.coo_array:
    m = Chem.GetAdjacencyMatrix(rdmol, useBO=True)
    return sp.coo_array(m)  # type: ignore


def get_adjacency_by_rdkit(
    numbers: np.ndarray,
    order: np.ndarray | None = None,
    source: np.ndarray | None = None,
    target: np.ndarray | None = None,
    geometry: np.ndarray | None = None,
    infer_order: bool = False,
    infer_bond: bool = False,
    cov_factor: float = 1.3,
    charge: int = 0,
) -> sp.coo_array:
    """Return the adjacency matrix of the given rdmol."""
    rdmol: RDMol = get_rdmol(
        numbers=numbers,
        order=order,
        source=source,
        target=target,
        geometry=geometry,
        infer_bond=infer_bond,
        infer_order=infer_order,
        cov_factor=cov_factor,
        charge=charge,
    )
    return _get_adjacency(rdmol=rdmol)


def get_atomic_sasa(rdmol: RDMol) -> list[float]:
    """Return the atomic SASA of the given rdmol."""
    rdmol = rdmol.__copy__()  # copy the rdmol object
    rdFreeSASA.CalcSASA(rdmol, rdFreeSASA.classifyAtoms(rdmol))
    return [float(a.GetProp("SASA")) for a in rdmol.GetAtoms()]
