# -*- coding: utf-8 -*-
"""This module defines miscellaneous functions dealing with protein data."""

import numpy as np

from prody.atomic import Atomic, Atom, AtomGroup, Selection, HierView
from prody.utilities import openFile, showFigure, createStringIO, wrapModes
from prody import SETTINGS, PY3K

__all__ = ['view3D', 'showProtein']

def view3D(*alist, **kwargs):
    """Return a py3Dmol view instance for interactive visualization in
    Jupyter notebooks. Available arguments are: width, height (of
    the viewer), backgroundColor, zoomTo (a py3Dmol selection to center
    around), and style, which is a py3Dmol style object that will be
    applied to all atoms in the scene. More complex styling can be achieved
    by manipulating the view object directly.
    
    The default style is to show the protein in a rainbow cartoon and
    hetero atoms in sticks/spheres.
    
    GNM/ANM Coloring
    
    An array of fluctuation values can be provided with the *data* kwarg
    for visualization of GNM/ANM calculations.  The array is assumed to 
    correpond to a calpha selection of the provided protein.
    The default color will be set to a RWB color scheme on a per-residue
    basis.  If the fluctuation vector contains negative values, the
    midpoint (white) will be at zero.  Otherwise the midpoint is the mean.
    
    An array of displacement vectors can be provided with the *mode* kwarg.
    The animation of these motions can be controlled with frames (number
    of frames to animate over), amplitude (scaling factor), and animate
    (3Dmol.js animate options).  If animation isn't enabled, by default
    arrows are drawn.  Drawing of arrows is controlled by the boolean arrows
    option and the arrowcolor option.
    
    If multiple structures are provided with the data or mode arguments, these
    arguments must be provided as lists of arrays of the appropriate dimension.

    If a 3Dmol.js viewer as specified as the view argument, that viewer will be
    modified and returned.  After modification, update instead of show should
    be called on the viewer object if it is desired to update in-place
    instead of instantiating a new viewer.
    """

    try:
        import py3Dmol
    except:
        raise ImportError('py3Dmol needs to be installed to use view3D')
    
    from .pdbfile import writePDBStream
    
    width = kwargs.get('width', 400)
    height = kwargs.get('height', 400)
    data_list = kwargs.pop('data', None)
    modes = kwargs.pop('mode', None)
    style = kwargs.pop('style', [])
    zoomto = kwargs.pop('zoomto', {})
    bgcolor = kwargs.pop('backgroundcolor', 'white')
    bgcolor = kwargs.pop('backgroundColor', bgcolor)
    frames = kwargs.pop('frames', 30)
    interval = kwargs.pop('interval', 1)
    anim = kwargs.pop('anim', False)
    scale = kwargs.pop('scale', 100)
    arrows = kwargs.pop('arrows',True)
    arrowcolor = kwargs.pop('arrowcolor', 'darkgrey')
    arrowcolor = kwargs.pop('arrowColor', arrowcolor)

    if modes is None:
        n_modes = 0
    else:
        modes = wrapModes(modes)
        n_modes = len(modes)

    if data_list is None:
        n_data = 0
    else:
        data_list = wrapModes(data_list)
        n_data = len(data_list)

    view = kwargs.get('view',py3Dmol.view(width=width, height=height, js=kwargs.get('js','http://3dmol.csb.pitt.edu/build/3Dmol-min.js')))

    def _mapData(atoms, data):
        # construct map from residue to data property
        propmap = []
        for j, a in enumerate(atoms.calpha):
            propmap.append({'model': -1, 'chain': a.getChid(), 'resi': int(a.getResnum()), 
                            'props': {'data': float(data[j]) } })
        # set the atom property 
        # TODO: implement something more efficient on the 3Dmol.js side (this is O(n*m)!)
        view.mapAtomProperties(propmap)
        
        # color by property using gradient
        extreme = np.max(np.abs(data))
        lo = -extreme if np.min(data) < 0 else 0
        mid = np.mean(data) if np.min(data) >= 0 else 0
        view.setColorByProperty({'model': -1}, 'data', 'rwb', [extreme,lo,mid])
        view.setStyle({'model': -1},{'cartoon':{'style':'oval'}})    


    for i, atoms in enumerate(alist):
        pdb = createStringIO()
        writePDBStream(pdb, atoms)
        view.addAsOneMolecule(pdb.getvalue(), 'pdb')
        view.setStyle({'model': -1}, {'cartoon': {'color':'spectrum'}})
        view.setStyle({'model': -1, 'hetflag': True}, {'stick':{}})
        view.setStyle({'model': -1, 'bonds': 0}, {'sphere':{'radius': 0.5}})    

        if n_data:
            data = data_list[i]
            try:
                data = data.getArray()
            except AttributeError:
                pass

            if atoms.calpha.numAtoms() != len(data):
                raise RuntimeError("Atom count mismatch: {} vs {}. data styling assumes a calpha selection."
                                   .format(atoms.calpha.numAtoms(), len(data)))
            _mapData(atoms, data)
            
        if n_modes:
            mode = modes[i]
            try:
                arr = mode.getArray()
                is3d = mode.is3d()
            except AttributeError:
                arr = mode
                is3d = len(arr) == atoms.calpha.numAtoms()*3

            if is3d:
                if atoms.calpha.numAtoms()*3 != len(arr):
                    raise RuntimeError("Atom count mismatch: {} vs {}. mode animation assumes a calpha selection."
                                       .format(atoms.calpha.numAtoms(), len(arr)//3))

                if anim:
                    # construct map from residue to anm property and dx,dy,dz vectors
                    propmap = []
                    for j, a in enumerate(atoms.calpha):
                        propmap.append({'model': -1, 'chain': a.getChid(), 'resi': int(a.getResnum()),
                                        'props': {'dx': arr[3*j], 'dy': arr[3*j+1], 'dz': arr[3*j+2] } })
                    # set the atom property 
                    # TODO: implement something more efficient on the 3Dmol.js side (this is O(n*m)!)
                    view.mapAtomProperties(propmap)
                elif arrows:
                    for j, a in enumerate(atoms.calpha):
                        start = a._getCoords()
                        dcoords = arr[3*j:3*j+3]
                        end = start + dcoords * scale
                        view.addArrow({'start': {'x':start[0], 'y':start[1], 'z':start[2]},
                                       'end': {'x':end[0], 'y':end[1], 'z':end[2]},
                                       'radius': 0.3, 'color': arrowcolor})
            else:
                if atoms.calpha.numAtoms() != len(arr):
                    raise RuntimeError("Atom count mismatch: {} vs {}. mode styling assumes a calpha selection."
                                       .format(atoms.calpha.numAtoms(), len(arr)))

                _mapData(atoms, arr)

                
    # setting styles ...
    view.setBackgroundColor(bgcolor)

    if n_modes:
        #create vibrations
        view.vibrate(frames, scale)
        
        animate = kwargs.get('animate', {'loop':'rock', 'interval':interval})
        view.animate(animate)     
        
    if isinstance(style, dict):
        style = ({}, style)
    if isinstance(style, tuple):
        styles = [style]
    else:
        styles = style
    for sel, style in styles:
        view.setStyle(sel, style)
        
    view.zoomTo(zoomto)
            
    return view

def showProtein(*atoms, **kwargs):
    """Show protein representation using :meth:`~mpl_toolkits.mplot3d.Axes3D`.
    This function is designed for generating a quick view of the contents of a
    :class:`~.AtomGroup` or :class:`~.Selection`.    

    Protein atoms matching ``"calpha"`` selection are displayed using solid
    lines by picking a random and unique color per chain.  Line with can
    be adjusted using *lw* argument, e.g. ``lw=12``. Default width is 4.
    Chain colors can be overwritten using chain identifier as in ``A='green'``.

    Water molecule oxygen atoms are represented by red colored circles.  Color
    can be changed using *water* keyword argument, e.g. ``water='aqua'``.
    Water marker and size can be changed using *wmarker* and *wsize* keywords,
    defaults values are ``wmarker='.', wsize=6``.

    Hetero atoms matching ``"hetero and noh"`` selection are represented by
    circles and unique colors are picked at random on a per residue basis.
    Colors can be customized using residue name as in ``NAH='purple'``.  Note
    that this will color all distinct residues with the same name in the same
    color.  Hetero atom marker and size can be changed using *hmarker* and
    *hsize* keywords, default values are ``hmarker='o', hsize=6``.

    ProDy will set the size of axis so the representation is not distorted when
    the shape of figure window is close to a square.  Colors are picked at
    random, except for water oxygens which will always be colored red.
    
    *** Interactive 3D Rendering in Jupyter Notebook ***
    
    If py3Dmol has been imported then it will be used instead to display 
    an interactive viewer.  See :func:`view3D`
    
    
    """

    from prody.dynamics.mode import Mode

    method = kwargs.pop('draw', None)
    modes = kwargs.get('mode', None)
    scale = kwargs.get('scale', 100)

    # modes need to be specifically a list or a tuple (cannot be an array)
    if modes is None:
        n_modes = 0
    else:
        modes = wrapModes(modes)
        n_modes = len(modes)

    if method is None:
        import sys        
        if 'py3Dmol' in sys.modules: 
            method = 'py3Dmol'
        else:
            method = 'matplotlib'
    method = method.lower()
        
    alist = atoms
    for atoms in alist:
        if not isinstance(atoms, Atomic):
            raise TypeError('atoms must be an Atomic instance')
            
    if n_modes and n_modes != len(alist):
        raise RuntimeError('the number of proteins ({0}) does not match that of the modes ({1}).'
                            .format(len(alist), n_modes))

    if '3dmol' in method:
        mol = view3D(*alist, **kwargs)
        return mol
    else:
        kwargs.pop('mode', None)
        kwargs.pop('scale', 100)

        import matplotlib.pyplot as plt
        from mpl_toolkits.mplot3d import Axes3D
        cf = plt.gcf()
        show = None
        for child in cf.get_children():
            if isinstance(child, Axes3D):
                show = child
                break
        if show is None:
            show = Axes3D(cf)
        from matplotlib import colors
        cnames = dict(colors.cnames)
        wcolor = kwargs.get('water', 'red').lower()
        avoid = np.array(colors.hex2color(cnames.pop(wcolor, cnames.pop('red'))))
        for cn, val in cnames.copy().items():  # PY3K: OK
            clr = np.array(colors.hex2color(val))
            if clr.sum() > 2.4:
                cnames.pop(cn)
            elif np.abs(avoid - clr).sum() <= 0.6:
                cnames.pop(cn)
        cnames = list(cnames)
        import random
        random.shuffle(cnames)
        cnames_copy = list(cnames)
        min_ = list()
        max_ = list()
        for i, atoms in enumerate(alist):
            if isinstance(atoms, AtomGroup):
                title = atoms.getTitle()
            else:
                title = atoms.getAtomGroup().getTitle()
            calpha = atoms.select('calpha')
            if calpha:
                partition = False
                mode = modes[i] if n_modes else None
                if mode is not None:
                    is3d = False
                    try:
                        arr = mode.getArray()
                        is3d = mode.is3d()
                        n_nodes = mode.numAtoms()
                    except AttributeError:
                        arr = mode
                        is3d = len(arr) == len(calpha)*3
                        n_nodes = len(arr)//3 if is3d else len(arr)
                    if n_nodes != len(calpha):
                        raise RuntimeError('size mismatch between the protein ({0} residues) and the mode ({1} nodes).'
                                            .format(len(calpha), n_nodes))
                    partition = not is3d

                if partition:
                    xyz = calpha._getCoords()
                    chids = calpha.getChids()
                    rbody = []
                    last_sign = np.sign(arr[0])
                    last_chid = chids[0]
                    rcolor = ['red', 'red', 'blue']
                    n = 1
                    for i,a in enumerate(arr):
                        s = np.sign(a)
                        ch = chids[i]
                        if s == 0: s = last_sign
                        if last_sign != s or i == len(arr)-1 or last_chid != ch:
                            if last_chid == ch:
                                rbody.append(i)
                            show.plot(xyz[rbody, 0], xyz[rbody, 1], xyz[rbody, 2],
                                      label=title + '_regid%d'%n,
                                      color=rcolor[int(last_sign+1)],
                                      lw=kwargs.get('lw', 4))
                            rbody = []
                            n += 1
                            last_sign = s
                            last_chid = ch
                        rbody.append(i)
                else:
                    for ch in HierView(calpha, chain=True):
                        xyz = ch._getCoords()
                        chid = ch.getChid()
                        if len(cnames) == 0:
                            cnames = list(cnames_copy)
                        show.plot(xyz[:, 0], xyz[:, 1], xyz[:, 2],
                                label=title + '_' + chid,
                                color=kwargs.get(chid, cnames.pop()).lower(),
                                lw=kwargs.get('lw', 4))
                    
                    if mode is not None:
                        from prody.utilities.drawtools import drawArrow3D
                        XYZ = calpha._getCoords()
                        arr = arr.reshape((n_nodes, 3))
                        XYZ2 = XYZ + arr * scale
                        for i, xyz in enumerate(XYZ):
                            xyz2 = XYZ2[i]
                            mutation_scale = kwargs.pop('mutation_scale', 10)
                            drawArrow3D(xyz, xyz2, mutation_scale=mutation_scale, **kwargs)

            water = atoms.select('water and noh')
            if water:
                xyz = atoms.select('water')._getCoords()
                show.plot(xyz[:, 0], xyz[:, 1], xyz[:, 2], label=title + '_water',
                          color=wcolor,
                          ls='None', marker=kwargs.get('wmarker', '.'),
                          ms=kwargs.get('wsize', 6))
            hetero = atoms.select('not protein and not nucleic and not water and not dummy')
            if hetero:
                for res in HierView(hetero).iterResidues():
                    xyz = res._getCoords()
                    resname = res.getResname()
                    resnum = str(res.getResnum())
                    chid = res.getChid()
                    if len(cnames) == 0:
                        cnames = list(cnames_copy)
                    show.plot(xyz[:, 0], xyz[:, 1], xyz[:, 2], ls='None',
                              color=kwargs.get(resname, cnames.pop()).lower(),
                              label=title + '_' + chid + '_' + resname + resnum,
                              marker=kwargs.get('hmarker', 'o'),
                              ms=kwargs.get('hsize', 6))
            xyz = atoms._getCoords()
            min_.append(xyz.min(0))
            max_.append(xyz.max(0))

        show.set_xlabel('x')
        show.set_ylabel('y')
        show.set_zlabel('z')
        min_ = np.array(min_).min(0)
        max_ = np.array(max_).max(0)
        center = (max_ + min_) / 2
        half = (max_ - min_).max() / 2
        show.set_xlim3d(center[0]-half, center[0]+half)
        show.set_ylim3d(center[1]-half, center[1]+half)
        show.set_zlim3d(center[2]-half, center[2]+half)
        if kwargs.get('legend', False):
            show.legend(prop={'size': 10})
        if SETTINGS['auto_show']:
            showFigure()
        return show
