#!/usr/bin/env python
# -*- coding: utf-8 -*-

import warnings

import numpy as np

from .utils import pdbutils


class AtomTyper:

    @classmethod
    def type_everything(
        cls,
        molsetup,
        atom_params,
        charge_model,
        offatom_params=None,
        dihedral_params=None,
    ):

        cls._type_atoms(molsetup, atom_params)

        # offatoms must be typed after charges, because offsites pull charge
        if offatom_params is not None:
            cached_offatoms = cls._cache_offatoms(molsetup, offatom_params)
            coords = {atom.index: atom.coord for atom in molsetup.atoms if not atom.is_dummy}
            cls._set_offatoms(molsetup, cached_offatoms, coords)

        if dihedral_params not in (None, "espaloma"):
            cls._type_dihedrals(molsetup, dihedral_params)

        return

    @staticmethod
    def _type_atoms(molsetup, atom_params):
        # ensure every "atompar" is defined in a single "smartsgroup"
        ensure = {}
        # go over all "smartsgroup"s
        for smartsgroup in atom_params:
            if smartsgroup == "comment":
                continue
            for line in atom_params[
                smartsgroup
            ]:  # line is a dict, e.g. {"smarts": "[#1][#7,#8,#9,#15,#16]","atype": "HD"}
                smarts = str(line["smarts"])
                # get indices of the atoms in the smarts to which the parameters will be assigned
                idxs = [
                    0
                ]  # by default, the first atom in the smarts gets parameterized
                if "IDX" in line:
                    idxs = [i for i in line["IDX"]]
                # match SMARTS
                hits = molsetup.find_pattern(smarts)
                for atompar in line:
                    if atompar in ["smarts", "comment", "IDX"]:
                        continue
                    if atompar not in molsetup.atom_params:
                        molsetup.atom_params[atompar] = [None] * len(molsetup.atoms)
                    value = line[atompar]
                    # keep track of every "smartsgroup" that modified "atompar"
                    ensure.setdefault(atompar, [])
                    ensure[atompar].append(smartsgroup)
                    # Each "hit" is a tuple of atom indices that matched the smarts
                    # The length of each "hit" is the number of atoms in the smarts
                    for hit in hits:
                        # Multiple atoms may be targeted by a single smarts:
                        # For example: both oxygens in NO2 are parameterized by a single smarts pattern.
                        # "idxs" are 1-indices of atoms in the smarts to which parameters are to be assigned.
                        for idx in idxs:
                            if atompar == "atype":
                                molsetup.set_atom_type(
                                    hit[idx], value
                                )  # overrides previous calls
                            molsetup.atom_params[atompar][hit[idx]] = value

        # guarantee that each atompar is exclusive of a single group
        for atompar in ensure:
            if len(set(ensure[atompar])) > 1:
                msg = "%s is modified in multiple smartsgroups: %s" % (
                    atompar,
                    set(ensure[atompar]),
                )
                warnings.warn(msg)
        return

    @staticmethod
    def _cache_offatoms(molsetup, offatom_params):
        """precalculate off-site atoms"""
        cached_offatoms = {}
        n_offatoms = 0
        atoms_with_offchrg = set()
        # each parent atom can only be matched once in each smartsgroup
        for smartsgroup in offatom_params:
            if smartsgroup == "comment":
                continue
            tmp = {}
            for line in offatom_params[smartsgroup]:
                # SMARTS
                smarts = str(line["smarts"])
                hits = molsetup.find_pattern(smarts)
                # atom indexes in smarts string
                smarts_idxs = [0]
                if "IDX" in line:
                    smarts_idxs = [i for i in line["IDX"]]
                for smarts_idx in smarts_idxs:
                    for hit in hits:
                        parent_idx = hit[smarts_idx]
                        tmp.setdefault(
                            parent_idx, []
                        )  # TODO tmp[parent_idx] = [], yeah?
                        for offatom in line["OFFATOMS"]:
                            # set defaults
                            tmp[parent_idx].append(
                                {
                                    "offatom": {
                                        "distance": 1.0,
                                        "x90": False,
                                        "phi": 0.0,
                                        "theta": 0.0,
                                        "z": [],
                                        "x": [],
                                    },
                                    "atom_params": {},
                                }
                            )
                            for key in offatom:
                                if key in ["distance", "x90"]:
                                    tmp[parent_idx][-1]["offatom"][key] = offatom[key]
                                # replace SMARTS indexes by the atomic index
                                elif key in ["z", "x"]:
                                    for i in offatom[key]:
                                        idx = hit[i]
                                        tmp[parent_idx][-1]["offatom"][key].append(idx)
                                # convert degrees to radians
                                elif key in ["theta", "phi"]:
                                    tmp[parent_idx][-1]["offatom"][key] = np.radians(
                                        offatom[key]
                                    )
                                # ignore comments
                                elif key in ["comment"]:
                                    pass
                                elif key == "atype":
                                    tmp[parent_idx][-1]["atom_params"][key] = offatom[
                                        key
                                    ]
                                elif key == "pull_charge_fraction":
                                    if parent_idx in atoms_with_offchrg:
                                        raise RuntimeError(
                                            "atom %d has charge pulled more than once"
                                            % parent_idx
                                        )
                                    atoms_with_offchrg.add(parent_idx)
                                    tmp[parent_idx][-1]["atom_params"][key] = offatom[
                                        key
                                    ]
                                else:
                                    pass
            for parent_idx in tmp:
                for offatom_dict in tmp[parent_idx]:
                    atom_params = offatom_dict["atom_params"]
                    offatom = offatom_dict["offatom"]
                    atomgeom = AtomicGeometry(
                        parent_idx,
                        neigh=offatom["z"],
                        xneigh=offatom["x"],
                        x90=offatom["x90"],
                    )
                    if "pull_charge_fraction" in atom_params:
                        pull_charge_fraction = atom_params["pull_charge_fraction"]
                    else:
                        pull_charge_fraction = 0.0
                    args = (
                        atom_params["atype"],
                        offatom["distance"],
                        offatom["theta"],
                        offatom["phi"],
                        pull_charge_fraction,
                    )
                    # number of coordinates (before adding new offatom)
                    cached_offatoms[n_offatoms] = (atomgeom, args)
                    n_offatoms += 1
        return cached_offatoms

    @staticmethod
    def _set_offatoms(molsetup, cached_offatoms, coords):
        """add cached offatoms"""
        for k, (atomgeom, args) in cached_offatoms.items():
            (atom_type, dist, theta, phi, pull_charge_fraction) = args
            offatom_coords = atomgeom.calc_point(dist, theta, phi, coords)
            tmp = molsetup.get_pdbinfo(atomgeom.parent + 1)
            pdbinfo = pdbutils.PDBAtomInfo(
                "G", tmp.resName, tmp.resNum, tmp.icode, tmp.chain
            )
            q_parent = (1 - pull_charge_fraction) * molsetup.get_charge(atomgeom.parent)
            q_offsite = pull_charge_fraction * molsetup.get_charge(atomgeom.parent)
            pseudo_atom = {
                "coord": offatom_coords,
                "anchor_list": [atomgeom.parent],
                "charge": q_offsite,
                "pdbinfo": pdbinfo,
                "atom_type": atom_type,
                "rotatable": False,
            }
            molsetup.atoms[atomgeom.parent].charge = q_parent
            molsetup.add_pseudoatom(**pseudo_atom)
        return

    @staticmethod
    def _type_dihedrals(molsetup, dihedral_params):

        dihedrals = {}

        for line in dihedral_params:
            smarts = str(line["smarts"])
            hits = molsetup.find_pattern(smarts)
            if len(hits) == 0:
                continue  # non-rotatable bonds get dihedrals
            idxs = [i for i in line["IDX"]]
            tid = line["id"] if "id" in line else None
            fourier_series = []
            term_indices = {}
            for key in line:
                for keyword in ["phase", "k", "periodicity", "idivf"]:
                    if key.startswith(keyword):
                        t = int(key.replace(keyword, ""))  # e.g. "phase2" -> int(2)
                        if t not in term_indices:
                            term_indices[t] = len(fourier_series)
                            fourier_series.append({})
                        index = term_indices[t]
                        fourier_series[index][keyword] = line[key]
                        break

            for index in range(len(fourier_series)):
                if "idivf" in fourier_series[index]:
                    idivf = fourier_series[index].pop("idivf")
                    fourier_series[index]["k"] /= idivf

            dihedral_index = molsetup.add_dihedral_interaction(fourier_series)

            for hit in hits:
                atom_idxs = tuple([hit[j] for j in idxs])
                molsetup.dihedral_partaking_atoms[atom_idxs] = dihedral_index
                molsetup.dihedral_labels[atom_idxs] = tid


class AtomicGeometry:
    """generate reference frames and add extra sites"""

    def __init__(self, parent, neigh, xneigh=[], x90=False, planar_tol=0.1):
        """arguments are indices of atoms"""

        self.planar_tol = planar_tol  # angstroms, length of neighbor vecs for z-axis

        # real atom hosting extra sites
        if type(parent) != int:
            raise RuntimeError("parent must be int")
        self.parent = parent

        # list of bonded atoms (used to define z-axis)
        self.neigh = []
        for i in neigh:
            if type(i) != int:
                raise RuntimeError("neigh indices must be int")
            self.neigh.append(i)

        # list of atoms that
        self.xneigh = []
        for i in xneigh:
            if type(i) != int:
                raise RuntimeError("xneigh indices must be int")
            self.xneigh.append(i)

        self.calc_x = len(self.xneigh) > 0
        self.x90 = x90  # y axis becomes x axis (useful to rotate in-plane by 90 deg)

    def calc_point(self, distance, theta, phi, coords):
        """return coordinates of point specified in spherical coordinates"""

        z = self._calc_z(coords)

        # return pt aligned with z-axis
        if phi == 0.0:
            return z * distance + np.array(coords[self.parent])

        # need x-vec if phi != 0
        elif self.calc_x == False:
            raise RuntimeError("phi must be zero if X undefined")

        else:
            x = self._calc_x(coords)
            if self.x90:
                x = np.cross(self.z, x)
            y = np.cross(z, x)
            pt = z * distance
            pt = self.rot3D(pt, y, phi)
            pt = self.rot3D(pt, z, theta)
            pt += np.array(coords[self.parent])
            return pt

    def _calc_z(self, coords):
        """maximize distance from neigh"""
        z = np.zeros(3)
        cumsum = np.zeros(3)
        for i in self.neigh:
            v = np.array(coords[self.parent]) - np.array(coords[i])
            cumsum += v
            z += self.normalized(v)
        z = self.normalized(z)
        if np.sum(cumsum**2) < self.planar_tol**2:
            raise RuntimeError("Refusing to place Z axis on planar atom")
        return z

    def _calc_x(self, coords):
        x = np.zeros(3)
        for i in self.xneigh:
            v = np.array(coords[self.parent]) - np.array(coords[i])
            x += self.normalized(v)
        x = self.normalized(x)
        return x

    @staticmethod
    def rot3D(pt, ax, rad):
        """
        Rotate point:
        pt = (x,y,z) coordinates to be rotated
        ax = vector around wich rotation is performed
        rad = rotate by "rad" radians
        """
        # If axis has len=0, rotate by 0.0 rad on any axis
        # Make sure ax has unitary length
        len_ax = (ax[0] ** 2 + ax[1] ** 2 + ax[2] ** 2) ** 0.5
        if len_ax == 0.0:
            u, v, w = (1, 0, 0)
            rad = 0.0
        else:
            u, v, w = [i / len_ax for i in ax]
        x, y, z = pt
        ux, uy, uz = u * x, u * y, u * z
        vx, vy, vz = v * x, v * y, v * z
        wx, wy, wz = w * x, w * y, w * z
        sa = np.sin(rad)
        ca = np.cos(rad)
        p0 = (
            u * (ux + vy + wz)
            + (x * (v * v + w * w) - u * (vy + wz)) * ca
            + (-wy + vz) * sa
        )
        p1 = (
            v * (ux + vy + wz)
            + (y * (u * u + w * w) - v * (ux + wz)) * ca
            + (wx - uz) * sa
        )
        p2 = (
            w * (ux + vy + wz)
            + (z * (u * u + v * v) - w * (ux + vy)) * ca
            + (-vx + uy) * sa
        )
        return (p0, p1, p2)

    def normalized(self, vec):
        l = sum([x**2 for x in vec]) ** 0.5
        if type(vec) == list:
            return [x / l for x in vec]
        else:
            # should be np.array
            return vec / l
