from __future__ import print_function, absolute_import, division
from future.builtins import *
from future import standard_library
standard_library.install_aliases()

# Copyright 2017 Autodesk Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from past.builtins import basestring
import copy
import collections
import itertools

import numpy as np

import moldesign as mdt
from .. import units as u
from .. import utils, external, mathutils, widgets
from . import toplevel


class AtomGroup(object):
    """ Mixin functions for objects that have a ``self.atoms`` attribute with a list of atoms

    Attributes:
        atoms (List[Atom]): a list of atoms
    """
    draw2d = widgets.WidgetMethod('atomgroups.draw2d')
    draw3d = widgets.WidgetMethod('atomgroups.draw3d')
    draw = widgets.WidgetMethod('atomgroups.draw')

    def __init__(self, *args, **kwargs):
        """ This should never be called directly - it will be called by the `super` methods
        of its subclasses """
        super().__init__(*args, **kwargs)
        self._atom_attrs = None
        self.viz2d = None
        self.viz3d = None

    @property
    def num_atoms(self):
        """ int: number of atoms in this object """
        return len(self.atoms)
    natoms = num_atoms

    @property
    def heavy_atoms(self):
        """ AtomList: a list of all heavy atoms (i.e., non-hydrogen) in this object """
        return AtomList([a for a in self.atoms if a.atnum != 1])

    @property
    def mass(self):
        """ u.Scalar[mass]: total mass of this object
        """
        return u.unitsum(a.mass for a in self.atoms)

    @property
    def momentum(self):
        """ u.Vector[momentum]: total momentum of this object
        """
        return self.momenta.sum(axis=0)

    @property
    def velocity(self):
        """ u.Vector[velocity]: center of mass velocity of this object
        """
        return self.momentum/self.mass

    @property
    def kinetic_energy(self):
        r""" u.Scalar[energy]: Classical kinetic energy :math:`\sum_{\text{atoms}} \frac{p^2}{2m}`
        """
        from ..helpers import kinetic_energy
        return kinetic_energy(self.momenta, self.masses)

    def get_atoms(self, *keywords, **queries):
        """Allows keyword-based atom queries. Returns atoms that match ALL queries.

        Args:
            *keywords (list): pre-set keywords (currently, just selects by residue type)
            **queries (dict): attributes (or residue attributes) to match

        Examples:
            >>> mol.get_atoms('protein')  # returns all atoms in proteins
            >>> mol.get_atoms(name='CA')  # returns all alpha carbons
            >>> mol.get_atoms('dna', symbol='P')  # returns all phosphorus in DNA

        Returns:
            AtomList: the atoms matching this query
        """
        if not (queries or keywords):
            return mdt.AtomList(self.atoms)

        atoms = self.atoms

        KEYS = 'protein dna rna water unknown ion'.split()

        for key in keywords:
            if key in KEYS:
                atoms = mdt.AtomList(atom for atom in atoms
                                     if atom.residue.type == key)
            else:
                raise ValueError("Invalid keyword '%s': valid values are %s" % (key, KEYS))

        result = mdt.AtomList()
        for atom in atoms:
            for field, val in queries.items():
                if getattr(atom, field, None) != val and getattr(atom.residue, field, None) != val:
                    break
            else:
                result.append(atom)

        return result

    def calc_distance_array(self, other=None):
        """ Calculate an array of pairwise distance between all atoms in self and other

        Args:
            other (AtomContainer): object to calculate distances to (default: self)

        Returns:
            u.Array[length]: 2D array of pairwise distances between the two objects

        Example:
            >>> dists = self.calc_distance_array(other)
            >>> dists[i, j] == self.atoms[i].distance(other.atoms[j])
        """
        from scipy.spatial.distance import cdist

        other = utils.if_not_none(other, self)
        try:
            other_positions = other.positions.defunits_value()
        except AttributeError:
            other_positions = np.array([other.position.defunits_value()])

        distances = cdist(self.positions.defunits_value(), other_positions)
        return distances * u.default.length

    def calc_displacements(self):
        """ Calculate an array of displacements between all atoms in this object

        Returns:
            u.Array[length]: array of pairwise displacements between atoms

        Example:
            >>> displacements = self.calc_displacements(other)
            >>> displacements[i, j] == (self.atoms[i].position - self.atoms[j].position)
        """
        # TODO: allow other, similar to calc_distance array
        return utils.pairwise_displacements(self.positions)

    def distance(self, other):
        """Returns closest distance between this and the other entity

        Args:
            other (AtomContainer): object to calculate distance to

        Returns:
            u.Scalar[length]: closest distance between self and other

        Example:
            >>> distance = self.distance(other)
            >>> distance == self.calc_distance_array(other).min()
        """
        distance_array = self.calc_distance_array(other)
        return distance_array.min()

    @property
    def center_of_mass(self):
        """ units.Vector[length]: The (x,y,z) coordinates of this object's center of mass """
        if self.num_atoms == 0:  # nicer exception than divide-by-zero
            raise ValueError('"%s" has no atoms' % str(self))

        total_mass = 0.0 * u.default.mass
        com = np.zeros(3) * u.default.length * u.default.mass
        for atom in self.atoms:
            total_mass += atom.mass
            com += atom.position * atom.mass
        com = com / total_mass
        return com

    @center_of_mass.setter
    def center_of_mass(self, value):
        vec = value - self.com
        self.translate(vec)

    com = center_of_mass  # synonym

    def _getatom(self, a):
        """ Given an atom's name, index, or object, return the atom object
        """
        if a is None:
            return None
        elif isinstance(a, basestring) or isinstance(a, int):
            return self[a]
        else:
            return a

    def angle(self, a1, a2, a3):
        """ Calculate the angle between three atoms.

        Atoms can be passed as the atoms themselves or as the atom names

        Args:
            a1, a2, a3 (str OR int OR moldesign.Atom): atoms defining the angle

        Returns:
            units.Scalar[angle]
        """
        # TODO: use single dispatch to also accept two bonds
        return mdt.geom.angle(*list(map(self._getatom, (a1, a2, a3))))

    def dihedral(self, a1, a2, a3=None, a4=None):
        """ Calculate the dihedral angle between atoms a1, a2, a3, a4.

        Atoms can be passed as the atoms themselves or as the atom names

        Args:
            a1, a2, a3, a4 (str OR int OR moldesign.Atom): atoms defining the dihedral

        Returns:
            units.Scalar[angle]
        """
        return mdt.geom.dihedral(*list(map(self._getatom, (a1, a2, a3, a4))))

    def copy_atoms(self):
        """ Copy a group of atoms along and relevant topological information.

        This specifically copies:
          - the atoms themselves (along with their positions, momenta, and bond graphs)
          - any residues they are are part of
          - any chains they are part of
          - any bonds between them
        It does NOT copy:
          - other atoms that are part of these atoms' residues
          - other residues that are part of these atoms chains
          - the molecule these atoms are part of

        Returns:
            AtomList: list of copied atoms
        """
        graph = {}
        memo = {'bondgraph':graph}
        for atom in self.atoms:
            atom._subcopy(memo)
        tempatoms = [memo[atom] for atom in self.atoms]
        newatoms, newbonds = copy.deepcopy((tempatoms, graph))
        for atom in newatoms:
            atom.bond_graph = newbonds[atom]
        return AtomList(newatoms)

    ###########################################
    # Routines to modify the geometry
    def rotate(self, angle, axis, center=None):
        """Rotate this object in 3D space

        Args:
            angle (u.Scalar[angle]): angle to rotate by
            axis (u.Vector[length]): axis to rotate about (len=3)
            center (u.Vector[length]): center of rotation (len=3) (default: origin)
        """
        center = utils.if_not_none(center, self.com)

        if hasattr(angle, 'units'): angle = angle.value_in(u.radians)
        rotmat = external.transformations.rotation_matrix(angle, axis, point=center)
        self.transform(rotmat)

    def translate(self, vector):
        """Translate this object in 3D space

        Args:
            vector (u.Vector[length]): translation vector, len=3
        """
        for atom in self.atoms:
            atom.position += vector

    def transform(self, matrix):
        """ Transform this object's coordinates using the provided 4x4 matrix

        Args:
            matrix (numpy.ndarray): transformation matrix, shape=(4,4)
        """
        # TODO: deal with units ... hard because the matrix has diff units for diff columns
        assert matrix.shape == (4, 4)
        self.positions = mathutils.apply_4x4_transform(matrix, self.positions)

    def atoms_within(self, radius, other=None, include_self=False):
        """ Return all atoms in an object within a given radius of this object

        Args:
            radius (u.Scalar[length]): radius to search for atoms
            other (AtomContainer): object containing the atoms to search (default:self.molecule)
            include_self (bool): if True, include the atoms from this object (since, by definition,
               their distance from this object is 0)

        Returns:
            AtomList: list of the atoms within ``radius`` of this object
        """
        if other is None:
            other = self.atoms[0].molecule
        if not include_self:
            filter_atoms = set(self.atoms)
        else:
            filter_atoms = set()

        distances = self.calc_distance_array(other=other)
        mindists = distances.min(axis=0)

        close_atoms = AtomList(atom for dist,atom in zip(mindists, other.atoms)
                               if dist <= radius and atom not in filter_atoms)
        return close_atoms

    def residues_within(self, radius, other=None, include_self=False):
        """ Return all atoms in an object within a given radius of this object

        Args:
            radius (u.Scalar[length]): radius to search for atoms
            other (AtomContainer): object containing the atoms to search (default:self.molecule)
            include_self (bool): if True, include the atoms from this object (since, by definition,
               their distance from this object is 0)

        Returns:
            AtomList: list of the atoms within ``radius`` of this object
        """
        atoms = self.atoms_within(radius, other=other, include_self=include_self)
        residues = collections.OrderedDict((atom.residue, None) for atom in atoms)
        return list(residues.keys())


class _AtomArray(object):
    def __init__(self, attrname):
        self.attrname = attrname

    def __get__(self, instance, owner):
        return u.array([getattr(atom, self.attrname) for atom in instance.atoms])

    def __set__(self, instance, value):
        assert len(value) == instance.num_atoms
        for atom, atomval in zip(instance.atoms, value):
            setattr(atom, self.attrname, atomval)


class AtomContainer(AtomGroup):
    """
    Mixin functions for NON-MOLECULE objects that have a list of atoms at``self.atoms``
    """
    positions = _AtomArray('position')
    masses = _AtomArray('mass')
    momenta = _AtomArray('momentum')
    velocities = _AtomArray('velocity')

    def __add__(self, other):
        l = mdt.AtomList(self.atoms)
        l.extend(other.atoms)
        return l

    @property
    def bond_graph(self):
        """ Dict[moldesign.Atom: List[moldesign.Atom]]: bond graph for all atoms in this object
        """
        return {atom: atom.bond_graph for atom in self.atoms}

    @property
    def bonds(self):
        """ Iterable[moldesign.Bond]: iterator over bonds from this object's atoms
        """
        bg = self.bond_graph
        for atom, nbrs in bg.items():
            for nbr, order in nbrs.items():
                if atom.index < nbr.index or nbr not in bg:
                    yield mdt.Bond(atom, nbr)

    def get_bond(self, a1, a2):
        return mdt.Bond(a1, a2)

    @property
    def internal_bonds(self):
        """ Iterable[moldesign.Bond]: iterator over bonds that connect two atoms in this object
        """
        bg = self.bond_graph
        for atom, nbrs in bg.items():
            for nbr, order in nbrs.items():
                if atom.index < nbr.index and nbr in bg: