# (c) 2015-2022 Acellera Ltd http://www.acellera.com
# All Rights Reserved
# Distributed under HTMD Software License Agreement
# No redistribution in whole or part
#
from typing import TYPE_CHECKING

from moleculekit.util import sequenceID
import numpy as np
import logging

if TYPE_CHECKING:
    from moleculekit.molecule import Molecule

logger = logging.getLogger(__name__)


class BuildError(Exception):
    def __init__(self, text, errors=()):
        self.text = text
        self.errors = errors

    def __str__(self):
        if isinstance(self.text, str):
            return repr(self.text)
        elif isinstance(self.text, list):
            return "\n".join([v if isinstance(v, str) else repr(v) for v in self.text])


class _MissingErrorType(Exception):
    def __init__(self, text, values=()):
        self.text = text
        self.values = values
        if isinstance(self.values, np.ndarray):
            self.values = self.values.tolist()

    def __str__(self):
        return self.text


class MixedSegmentError(Exception):
    pass


class ResidueInsertionError(Exception):
    pass


class MissingResidueError(_MissingErrorType):
    pass


class MissingParameterError(_MissingErrorType):
    pass


class MissingTorsionError(_MissingErrorType):
    pass


class MissingBondError(_MissingErrorType):
    pass


class MissingAngleError(_MissingErrorType):
    pass


class MissingAtomTypeError(_MissingErrorType):
    pass


def embed(mol1: "Molecule", mol2: "Molecule", gap: float = 1.3) -> "Molecule":
    """Embed one molecule into another, removing overlapping residues.

    Removes residues of mol2 that have collisions with atoms of mol1, then
    appends mol1 into mol2.

    Parameters
    ----------
    mol1 : :class:`Molecule <moleculekit.molecule.Molecule>`
        The first Molecule object (embedded into mol2).
    mol2 : :class:`Molecule <moleculekit.molecule.Molecule>`
        The second Molecule object (residues clashing with mol1 are removed).
    gap : float, optional
        Minimum distance in Angstroms between atoms of the two molecules below
        which a residue is considered to clash.

    Returns
    -------
    newmol : :class:`Molecule <moleculekit.molecule.Molecule>`
        The resulting Molecule object with mol1 embedded into mol2.

    Examples
    --------
    >>> merged = embed(memb, prot)
    """
    from scipy.spatial.distance import cdist

    dists = cdist(mol2.coords[:, :, 0], mol1.coords[:, :, 0])
    s2close = np.where(dists < gap)[0]

    mol2res = sequenceID(mol2.resid)
    s2closeres = np.isin(mol2res, np.unique(mol2res[s2close]))

    mol2 = mol2.copy()
    mol2.remove(s2closeres, _logger=False)
    mol2.append(mol1)
    return mol2


def convertDisulfide(mol: "Molecule", disu: list) -> list:
    """Convert disulfide selection-string pairs to UniqueResidueID pairs.

    Parameters
    ----------
    mol : :class:`Molecule <moleculekit.molecule.Molecule>`
        The molecule the selections refer to.
    disu : list
        A list of pairs of atom selection strings, each pair identifying the two
        residues forming a disulfide bond.

    Returns
    -------
    newdisu : list
        The same pairs with each selection string resolved to a
        :class:`UniqueResidueID <moleculekit.molecule.UniqueResidueID>`.
    """
    from moleculekit.molecule import UniqueResidueID

    newdisu = []
    for d in disu:
        if not isinstance(d[0], str) or not isinstance(d[1], str):
            raise RuntimeError("All disulfide selections should be strings")
        newdisu.append(
            [
                UniqueResidueID.fromMolecule(mol, d[0]),
                UniqueResidueID.fromMolecule(mol, d[1]),
            ]
        )
    return newdisu


def detectDisulfideBonds(mol: "Molecule", thresh: float = 3) -> list:
    """Automatically detect disulfide bonds in a molecule.

    Finds all SG atoms in cysteine-like residues (resnames starting with "CY")
    and returns pairs whose inter-sulfur distance is below `thresh`.

    Parameters
    ----------
    mol : :class:`Molecule <moleculekit.molecule.Molecule>`
        The molecule for which to detect disulfide bonds.
    thresh : float, optional
        Distance threshold in Angstroms below which two sulfur atoms are
        considered bonded.

    Returns
    -------
    disubonds : list
        A list of pairs of
        :class:`UniqueResidueID <moleculekit.molecule.UniqueResidueID>` objects
        representing the detected disulfide bonds, sorted by residue ID.

    Raises
    ------
    RuntimeError
        If segment names are not defined, if multiple SG atoms are found in the
        same residue, or if a sulfur atom has more than one possible bond partner.
    """
    from scipy.spatial.distance import pdist, squareform
    from moleculekit.molecule import UniqueResidueID

    disubonds = []

    # Find all SG atoms belonging to resnames starting with CY
    # 'resname "CY.*" and name SG'
    idx = np.where(
        [(rn[0:2] == "CY") and (n == "SG") for rn, n in zip(mol.resname, mol.name)]
    )[0]
    if len(idx) == 0:
        return disubonds

    if np.any([len(s) == 0 for s in mol.segid[idx]]):
        raise RuntimeError(
            "Cannot detect disulfide bonds without segment names defined."
        )

    residues = [UniqueResidueID.fromMolecule(mol, idx=i) for i in idx]
    for r1 in range(len(residues)):
        for r2 in range(r1 + 1, len(residues)):
            if residues[r1] == residues[r2]:
                raise RuntimeError(
                    f"Multiple SG atoms detected in the same residue {residues[r1]}. "
                    "Can't guess disulfide bridges."
                )

    sd = squareform(pdist(mol.coords[idx, :, mol.frame]))
    sd[np.diag_indices(sd.shape[0])] = thresh + 1  # Set the diagonal over threshold
    close = sd < thresh
    rows, cols = np.where(close)

    numbonds = np.sum(close, axis=0)
    if np.any(numbonds > 1):
        multibonded_idx1 = np.where(numbonds > 1)[0]
        multibonded_indexes = np.where(close[multibonded_idx1])
        multibonded_idx1 = multibonded_idx1[multibonded_indexes[0]]
        multibonded_idx2 = multibonded_indexes[1]
        pairs = [
            (str(residues[r]), str(residues[c]))
            for r, c in zip(multibonded_idx1, multibonded_idx2)
        ]
        raise RuntimeError(
            f"Sulphur atoms between pairs {pairs} have multiple possible bonds. Cannot guess disulfide bonds. "
            "Please specify them manually."
        )

    uniquerowcols = list(set([tuple(sorted((r, c))) for r, c in zip(rows, cols)]))
    for rc in uniquerowcols:
        disubonds.append([residues[rc[0]], residues[rc[1]]])
        msg = (
            f"Disulfide Bond between: {residues[rc[0]]}\n"
            f"                   and: {residues[rc[1]]}\n"
        )
        print(msg)

    if len(disubonds) == 1:
        logger.info("One disulfide bond was added")
    else:
        logger.info(f"{len(disubonds)} disulfide bonds were added")
    return sorted(disubonds, key=lambda x: x[0].resid)


def detectCisPeptideBonds(mol: "Molecule", respect_bonds: bool = False) -> None:
    """Detect and warn about cis peptide bonds in a protein.

    Projects the protein omega dihedrals and logs a warning for every frame and
    residue whose omega angle indicates a cis peptide bond.

    Parameters
    ----------
    mol : :class:`Molecule <moleculekit.molecule.Molecule>`
        The molecule to check for cis peptide bonds.
    respect_bonds : bool, optional
        If True, only report cis bonds whose backbone atoms form a connected
        component in the molecule's bond graph.
    """
    from moleculekit.projections.metricdihedral import MetricDihedral, Dihedral
    import networkx as nx

    protsel = mol.atomselect("protein and backbone and name C CA N")
    if np.sum(protsel) < 4:  # Less atoms than dihedral
        return

    dih = Dihedral.proteinDihedrals(mol, sel=protsel, dih=("omega",))

    metr = MetricDihedral(dih=dih, sincos=False)
    data = metr.project(mol)
    mapping = metr.getMapping(mol)

    frames, idxs = np.where(np.abs(data) < 120)
    if respect_bonds:
        bond_graph = nx.Graph()
        bond_graph.add_edges_from(mol.bonds)

    for ii in np.unique(idxs):
        currframes = frames[idxs == ii]
        nframes = len(currframes)
        description = mapping.loc[ii].description
        atomIndexes = [int(x) for x in mapping.loc[ii].atomIndexes]

        if respect_bonds:
            sub_bond_graph = bond_graph.subgraph(atomIndexes)
            if not nx.is_connected(sub_bond_graph):
                continue

        currframes_str = "{}".format(currframes)
        if nframes > 5:
            currframes_str = f"[{currframes[0]} ... {currframes[-1]}]"

        logger.warning(
            f'Found cis peptide bond in {nframes} frames: {currframes_str} in the omega diheral "{description}" with indexes {atomIndexes}'
        )


def _checkMixedSegment(mol):
    prot = mol.atomselect("protein")
    acenme = (mol.resname == "ACE") | (mol.resname == "NME")
    sel1 = prot | acenme  # 'protein or resname ACE NME'
    sel2 = ~prot & ~acenme  # 'not protein and not resname ACE NME'
    segsProt = np.unique(mol.segid[sel1])
    segsNonProt = np.unique(mol.segid[sel2])
    intersection = np.intersect1d(segsProt, segsNonProt)
    if len(intersection) != 0:
        logger.warning(
            f"Segments {intersection} contain both protein and non-protein atoms. "
            "Please assign separate segments to them or the build procedure might fail."
        )


def _checkLongResnames(mol, aliasresidues):
    for resname in np.unique(mol.resname):
        if len(resname) > 4 and resname not in aliasresidues:
            raise RuntimeError(
                "Too long residue names in Molecule. Please give a 4-letter alias to these residues with the aliasresidues option."
            )


def removeLipidsInProtein(prot, memb, lipidsel="lipids"):
    """Calculates the convex hull of the protein. If a lipid lies inside the hull it gets removed.

    This does not work well for lipids crossing out of the hull. If even one atom of the lipid is outside it will
    change the hull and will not get removed. I assume it will get removed by the clashes with the protein though.
    """
    return removeAtomsInHull(prot, memb, "name CA", lipidsel)


def removeAtomsInHull(
    mol1: "Molecule",
    mol2: "Molecule",
    hullsel: str | np.ndarray,
    removesel: str | np.ndarray,
) -> tuple:
    """Calculate the convex hull of an atom selection in mol1 and remove atoms within that hull in mol2.

    Parameters
    ----------
    mol1 : :class:`Molecule <moleculekit.molecule.Molecule>`
        Molecule for which to calculate the convex hull.
    mol2 : :class:`Molecule <moleculekit.molecule.Molecule>`
        Molecule containing atoms to check for hull membership.
    hullsel : str or np.ndarray
        An atom selection string, a boolean mask, or an integer index array (see :meth:`Molecule.atomselect <moleculekit.molecule.Molecule.atomselect>`) for
        atoms in mol1 from which to calculate the convex hull.
    removesel : str or np.ndarray
        An atom selection string, a boolean mask, or an integer index array (see :meth:`Molecule.atomselect <moleculekit.molecule.Molecule.atomselect>`) for
        atoms in mol2 from which to remove those located within the hull.

    Returns
    -------
    newmol2 : :class:`Molecule <moleculekit.molecule.Molecule>`
        mol2 without atoms located within the convex hull.
    numrem : int
        Number of fragments removed.
    """
    # TODO: Look into Morphological Snakes
    from scipy.spatial import ConvexHull

    mol2 = mol2.copy()
    # Convex hull of the protein
    hullcoords = mol1.get("coords", hullsel)
    hull = ConvexHull(hullcoords)

    sequence = sequenceID((mol2.resid, mol2.segid))
    uqres = np.unique(sequence)

    toremove = np.zeros(len(sequence), dtype=bool)
    numlipsrem = 0
    for (
        res
    ) in uqres:  # For each fragment check if it's atoms lie within the convex hull
        atoms = np.where(sequence == res)[0]
        newhull = ConvexHull(np.vstack((hullcoords, mol2.get("coords", sel=atoms))))

        # If the hull didn't change by adding the fragment, it lies within convex hull. Remove it.
        if list(hull.vertices) == list(newhull.vertices):
            toremove[atoms] = True
            numlipsrem += 1

    rematoms = mol2.atomselect(removesel)
    mol2.remove(toremove & rematoms)
    return mol2, numlipsrem


def removeHET(prot: "Molecule") -> "Molecule":
    """Remove all HETATM residues from a structure.

    Each unique HETATM residue name is removed, assuming it is a bound ligand or
    other heteroatom group.

    Parameters
    ----------
    prot : :class:`Molecule <moleculekit.molecule.Molecule>`
        The molecule to clean.

    Returns
    -------
    prot : :class:`Molecule <moleculekit.molecule.Molecule>`
        A copy of the input molecule with all HETATM residues removed.
    """
    prot = prot.copy()
    hetatoms = np.unique(prot.resname[prot.record == "HETATM"])
    for het in hetatoms:
        logger.info(
            "Found resname "
            "{}"
            " in structure. Removed assuming it is a bound ligand.".format(het)
        )
        prot.remove("resname {}".format(het))
    return prot


def tileMembrane(
    memb: "Molecule",
    xmin: float,
    ymin: float,
    xmax: float,
    ymax: float,
    buffer: float = 1.5,
) -> "Molecule":
    """Tile a membrane in the X and Y dimensions to reach a specific size.

    Parameters
    ----------
    memb : :class:`Molecule <moleculekit.molecule.Molecule>`
        The membrane to be tiled.
    xmin : float
        Minimum x coordinate.
    ymin : float
        Minimum y coordinate.
    xmax : float
        Maximum x coordinate.
    ymax : float
        Maximum y coordinate.
    buffer : float, optional
        Buffer distance in Angstroms between tiles.

    Returns
    -------
    megamemb : :class:`Molecule <moleculekit.molecule.Molecule>`
        A tiled membrane Molecule covering the specified dimensions.
    """
    from tqdm import tqdm

    memb = memb.copy()
    memb.resid = sequenceID((memb.resid, memb.insertion, memb.chain, memb.segid))

    minmemb = np.min(memb.get("coords", "water"), axis=0).flatten()

    size = np.max(memb.get("coords", "water"), axis=0) - np.min(
        memb.get("coords", "water"), axis=0
    )
    size = size.flatten()
    xreps = int(np.ceil((xmax - xmin) / size[0]))
    yreps = int(np.ceil((ymax - ymin) / size[1]))

    logger.info("Replicating Membrane {}x{}".format(xreps, yreps))

    from moleculekit.molecule import Molecule

    megamemb = Molecule()
    bar = tqdm(total=xreps * yreps, desc="Replicating Membrane")
    k = 0
    for x in range(xreps):
        for y in range(yreps):
            tmpmemb = memb.copy()
            xpos = xmin + x * (size[0] + buffer)
            ypos = ymin + y * (size[1] + buffer)

            tmpmemb.moveBy([-float(minmemb[0]) + xpos, -float(minmemb[1]) + ypos, 0])
            tmpmemb.remove(
                "same resid as (x > {} or y > {})".format(xmax, ymax), _logger=False
            )
            if tmpmemb.numAtoms == 0:
                continue

            tmpmemb.set("segid", "M{}".format(k), sel="not water")
            tmpmemb.set("segid", "MW{}".format(k), sel="water")

            megamemb.append(tmpmemb)
            k += 1
            bar.update(1)
    bar.close()

    # Membranes don't tile perfectly. Need to remove waters that clash with lipids of other tiles
    # Some clashes will still occur between periodic images however
    megamemb.remove("same resid as water and within 1.5 of not water", _logger=False)
    return megamemb


def minimalRotation(prot):
    """Find the rotation around Z that minimizes the X and Y dimensions of the protein to best fit in a box.

    Essentially PCA in 2D
    """
    from numpy.linalg import eig
    from numpy import cov

    xycoords = prot.coords[:, 0:2]

    c = cov(np.transpose(np.squeeze(xycoords)))
    values, vectors = eig(c)
    idx = np.argsort(values)

    xa = vectors[0, idx[-1]]
    ya = vectors[1, idx[-1]]

    def cart2pol(x, y):
        # Cartesian to polar coordinates. Rho is the rotation angle
        rho = np.sqrt(x**2 + y**2)
        phi = np.arctan2(y, x)
        return rho, phi

    angle, _ = cart2pol(xa, ya)
    return angle + np.radians(45)
