"""
Maestro 3D arrows.
The arrow module allows creation and drawing of arrows.  Clients draw
using graphics3d.common.Group instances, aliased here as ArrowGroup,
not through Arrow instances.
Control over the endpoints, radius, color, resolution and transparency
of a arrow are provided.  See the Arrow class for more info.
To draw any number of arrows create Arrow instances, add them to a Group
instance.
Copyright Schrodinger LLC, All rights reserved.
"""
import copy
import math
import numpy
import schrodinger
from . import common
from .common import OPACITY_MAX
from .common import OPACITY_MIN
from .common import TRANSPARENCY_MAX
from .common import TRANSPARENCY_MIN
from .common import Group
# NOTE: ArrowGroup is deprecated; use Group class.
ArrowGroup = Group
maestro = schrodinger.get_maestro()
RESOLUTION_MIN = 4
RESOLUTION_MAX = 50
RESOLUTION_DEFAULT = 16
TRANSPARENCY_DEFAULT = 0.0  # for backwards-compatibility
OPACITY_DEFAULT = 1.0
# Constants used to calculate bounding box:
BOUNDING_BOX_INIT_VALUE = 100000000.0
EPSILON = 0.0001
# The body length as a percentage of the total length
BODY_PERCENTAGE = 0.75
#############################################################################
#                              CLASSES                                      #
#############################################################################
[docs]class MaestroArrow(common._MaestroPrimitiveMixin, common.Primitive):
    """
    Class for creating a 3D arrow in Maestro
    Arrows should be added to a Group and drawing done via the Group.
    See the Group documentation.
    API Example::
        import schrodinger.maestro.maestro as maestro
        import schrodinger.graphics3d.common as common
        import schrodinger.graphics3d.arrow as arrow
        arrow_grp = common.Group()
        st = maestro.workspace_get()
        for bond in st.bond:
            ar = arrow.MaestroArrow(
                xhead=bond.atom2.x, yhead=bond.atom2.y, zhead=bond.atom2.z,
                xtail=bond.atom1.x, ytail=bond.atom1.y, ztail=bond.atom1.z,
                color='red',
                radius=0.15,
                opacity=0.8,
                resolution=50
            )
            # Add the primative to the container.
            arrow_grp.add(ar)
        # Unlike Arrow MaestroArrow simply needs to be shown to be drawn.
        # No special callback like there is for Arrows is needed.
        arrow_grp.show()
        # Hide the markers.
        arrow_grp.hide()
        # Remove the markers and the callback.
        arrow_grp.clear()
    """
[docs]    def __init__(
            self,
            xhead=None,
            yhead=None,
            zhead=None,
            xtail=None,
            ytail=None,
            ztail=None,
            color=None,
            r=0.0,
            g=1.0,
            b=0.0,  # for backward-compatbility
            radius=0.15,
            transparency=None,  # for backwards-compatability
            opacity=OPACITY_DEFAULT,
            resolution=RESOLUTION_DEFAULT,
            body_percentage=BODY_PERCENTAGE,
            remove_endcaps=False):
        """
        Constructor requires:
        :param xhead: x coordinate of head position
        :param yhead: y coordinate of head position
        :param zhead: z coordinate of head position
        :param xtail: x coordinate of tail position
        :param ytail: y coordinate of tail position
        :param ztail: z coordinate of tail position
        :param color: One of Color object (Color class) or Color name (string)
            or Tuple of (R, G, B) (each a float in range 0.0-1.0)
        :param radius: radius of the arrow in Angstroms Default: .15
        :param opacity: 0.0 (invisible) through 1.0 (opaque) Default: 0.0
        :param resolution: Ranges from ?? to ??.  Default: ??
        :param remove_endcaps: Whether to remove cylinder endcaps on arrow body.
        """
        # The head is the cone portion of the arrow
        self.head = None
        # The body is the cylinder portion of the arrow
        self.body = None
        self._xhead = 0
        self._yhead = 0
        self._zhead = 0
        self._xtail = 0
        self._ytail = 0
        self._ztail = 0
        self._radius = 0
        self.remove_endcaps = remove_endcaps
        if xhead is None or yhead is None or zhead is None:
            raise ValueError(
                "Must specify xhead, yhead and zhead values to define the arrow head"
            )
        elif xtail is None or ytail is None or ztail is None:
            raise ValueError(
                "Must specify xtail, ytail and ztail values to define the arrow tail"
            )
        else:
            self.xhead = xhead
            self.yhead = yhead
            self.zhead = zhead
            self.xtail = xtail
            self.ytail = ytail
            self.ztail = ztail
        self.radius = radius
        # Clamp to range of 0.0 and 1.0, inclusive
        if transparency is not None:  # for backwards-compatability
            if transparency < TRANSPARENCY_MIN:
                self.opacity = OPACITY_MAX
            elif transparency > TRANSPARENCY_MAX:
                self.opacity = TRANSPARENCY_MIN
            else:
                self.opacity = 1.0 - transparency
        else:  # Use opacity
            if opacity < OPACITY_MIN:
                self.opacity = OPACITY_MIN
            elif opacity > OPACITY_MAX:
                self.opacity = OPACITY_MAX
            else:
                self.opacity = opacity
        # Clamp the resolution
        if resolution < RESOLUTION_MIN:
            self.resolution = RESOLUTION_MIN
        elif resolution > RESOLUTION_MAX:
            self.resolution = RESOLUTION_MAX
        else:
            self.resolution = resolution
        if color is not None:
            self.r, self.g, self.b = common.color_arg_to_rgb(color)
        else:  # Use r/g/b for backwards-compatability
            self.r = float(r)
            self.g = float(g)
            self.b = float(b)
        (x0, y0, z0) = self.calculateCoords(body_percentage)
        self.head = maestro.create_cone(self.xhead, self.yhead, self.zhead, x0,
                                        y0, z0, self.r, self.g, self.b,
                                        2 * self.radius, self.opacity,
                                        self.resolution)
        self.body = maestro.create_cylinder(self.xtail, self.ytail, self.ztail,
                                            x0, y0, z0, self.r, self.g, self.b,
                                            self.radius, self.opacity,
                                            self.resolution,
                                            self.remove_endcaps)
        maestro_objects = [self.head, self.body]
        common.Primitive.__init__(self, maestro_objects) 
    def _calculateBoundingBox(self, mat):
        xyzmin = []
        xyzmax = []
        for k in range(6):
            xyzmin.append(BOUNDING_BOX_INIT_VALUE)
            xyzmax.append(-BOUNDING_BOX_INIT_VALUE)
        atom1 = numpy.array([self.xhead, self.yhead, self.zhead])
        atom2 = numpy.array([self.xtail, self.ytail, self.ztail])
        # Perform the vector operation
        # We have coordinates for two positions (atom1, atom2) which
        # define the membrane planes. We need to figure out a vector
        # perpendiclar to the vector between those two points and a
        # series of points in the planes in order to draw a polygon
        # Assign the size of the planes
        l = 2.0 * self.radius
        # AB is the vector between the two points
        AB = atom2 - atom1
        # M is the midpoint between the two points
        M = atom1 + 0.5 * AB
        # N is an arbitrary point we will project onto M to get
        # a perpendicular vector:
        N = copy.copy(M)
        if math.fabs(AB[0]) < EPSILON and math.fabs(AB[1]) < EPSILON:
            N[0] += 10.0
        elif math.fabs(AB[1]) < EPSILON and math.fabs(AB[2]) < EPSILON:
            N[1] += 10.0
        elif math.fabs(AB[2]) < EPSILON and math.fabs(AB[0]) < EPSILON:
            N[2] += 10.0
        else:
            N[0] += 10.0
        #  MN is a vector crosses AB at M
        MN = N - M
        # O is the projection of N onto AB. The vector O-N perpendicular
        # to A-B
        dot1 = numpy.dot(AB, AB)
        dot2 = numpy.dot(AB, MN)
        K = dot2 / dot1
        O = M + numpy.dot(K, AB)
        ON = common.get_normalized(N - O)
        # P is perpendicular to both AB and ON
        P = common.get_normalized(common.get_cross(AB, ON))
        # Now we have everything we need to generate the planes
        plane1 = []
        plane2 = []
        # Op is used to generate a vertex at each corner of what will
        # be drawn as a square plane. Basically it's atom1+S, atom1+U,
        # atom1-S, atom1-U
        op = [[l, 0.0], [0.0, l], [-l, 0.0], [0.0, -l]]
        for i in range(0, 4):
            plane1.append(atom1 + op[i][0] * ON + op[i][1] * P)
            plane2.append(atom2 + op[i][0] * ON + op[i][1] * P)
        tmp = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        for i in range(4):
            for k in range(3):
                tmp[k + 3] = plane1[i][k]
            tmp[0] = mat[0][0] * tmp[3] + mat[0][1] * tmp[4] + mat[0][2] * tmp[
                5] + mat[0][3]
            tmp[1] = mat[1][0] * tmp[3] + mat[1][1] * tmp[4] + mat[1][2] * tmp[
                5] + mat[1][3]
            tmp[2] = mat[2][0] * tmp[3] + mat[2][1] * tmp[4] + mat[2][2] * tmp[
                5] + mat[2][3]
            for k in range(6):
                if xyzmin[k] > tmp[k]:
                    xyzmin[k] = tmp[k]
                if xyzmax[k] < tmp[k]:
                    xyzmax[k] = tmp[k]
        for i in range(4):
            for k in range(3):
                tmp[k + 3] = plane2[i][k]
            tmp[0] = mat[0][0] * tmp[3] + mat[0][1] * tmp[4] + mat[0][2] * tmp[
                5] + mat[0][3]
            tmp[1] = mat[1][0] * tmp[3] + mat[1][1] * tmp[4] + mat[1][2] * tmp[
                5] + mat[1][3]
            tmp[2] = mat[2][0] * tmp[3] + mat[2][1] * tmp[4] + mat[2][2] * tmp[
                5] + mat[2][3]
            for k in range(6):
                if xyzmin[k] > tmp[k]:
                    xyzmin[k] = tmp[k]
                if xyzmax[k] < tmp[k]:
                    xyzmax[k] = tmp[k]
        return (xyzmin, xyzmax)
[docs]    def calculateCoords(self, body_percentage=BODY_PERCENTAGE):
        """
        Calculate the intermediate coordinate between head and body
        """
        xdiff = self.xhead - self.xtail
        ydiff = self.yhead - self.ytail
        zdiff = self.zhead - self.ztail
        x0 = self.xtail + xdiff * body_percentage
        y0 = self.ytail + ydiff * body_percentage
        z0 = self.ztail + zdiff * body_percentage
        return (x0, y0, z0) 
    # Helper functions
    def _setCoords(self):
        (x0, y0, z0) = self.calculateCoords()
        if self.head:
            maestro.set_coords1(self.head, x0, y0, z0)
            maestro.set_coords2(self.head, self._xhead, self._yhead,
                                self._zhead)
        if self.body:
            maestro.set_coords1(self.body, self._xtail, self._ytail,
                                self._ztail)
            maestro.set_coords2(self.body, x0, y0, z0)
    # Accessors
    def _getXHead(self):
        return self._xhead
    def _setXHead(self, value):
        self._xhead = value
        self._setCoords()
    def _getYHead(self):
        return self._yhead
    def _setYHead(self, value):
        self._yhead = value
        self._setCoords()
    def _getZHead(self):
        return self._zhead
    def _setZHead(self, value):
        self._zhead = value
        self._setCoords()
    def _getXTail(self):
        return self._xtail
    def _setXTail(self, value):
        self._xtail = value
        self._setCoords()
    def _getYTail(self):
        return self._ytail
    def _setYTail(self, value):
        self._ytail = value
        self._setCoords()
    def _getZTail(self):
        return self._ztail
    def _setZTail(self, value):
        self._ztail = value
        self._setCoords()
    def _getRadius(self):
        return self._radius
    def _setRadius(self, value):
        self._radius = value
        if self.head:
            maestro.set_radius(self.head, 2 * self._radius)
        if self.body:
            maestro.set_radius(self.body, self._radius)
    # **********
    # Properties
    # **********
    xhead = property(_getXHead, _setXHead)
    yhead = property(_getYHead, _setYHead)
    zhead = property(_getZHead, _setZHead)
    xtail = property(_getXTail, _setXTail)
    ytail = property(_getYTail, _setYTail)
    ztail = property(_getZTail, _setZTail)
    radius = property(_getRadius, _setRadius)