from __future__ import absolute_import
from uuid import uuid4
from collections import defaultdict

import numpy as np

from IPython.display import display, Javascript, Image
from IPython.html.widgets import DOMWidget
from IPython.utils.traitlets import (Unicode, Bool, Bytes, CInt, Any,
                                     Dict, Enum, CFloat, List, Tuple, CUnicode,
                                     CBool)

from .utils import encode_numpy

__all__ = ['RepresentationViewer', "TrajectoryControls"]


class RepresentationViewer(DOMWidget):

    # Name of the javascript class which this widget syncs against on the
    # browser side. To work correctly, this javascript class has to be
    # registered and loaded in the browser before this widget is constructed
    # (that's what enable_notebook() does)
    _view_name = Unicode('MolecularView', sync=True)

    width = CInt(sync=True)
    height = CInt(sync=True)

    # Update Camera Hack
    camera_str = CUnicode(sync=True)
    static_moving = CBool(sync=True)

    def __init__(self, width=500, height=500):
        '''RepresentationViewer is an IPython notebook widget useful to display 3d scenes through webgl.

        Example:

        .. code::

            from IPython.display import display

            rv = RepresentationViewer()
            rv.add_representation('point', {'coordinates': coordinates, 'colors': colors, 'sizes': sizes})
            display(rv)

        .. py:attribute: width

            Width in pixels of the IPython widget

        .. py:attribute: height

            Height in pixels of the IPython widget

        .. py:attribute: camera_str

            A string-representation of camera position and orientation

        .. py:attribute: static_moving

            Set to True to make the camera lose the "bouncy" rotation.


        '''
        super(RepresentationViewer, self).__init__()
        self.displayed = False
        self.width = width
        self.height = height

        # Things to be called when the js harnessing is intialized
        self._displayed_callbacks = []
        def callback(widget):
            for cb in widget._displayed_callbacks:
                cb(widget)
        self.on_displayed(callback)

        # Store the events sent from the javascript side
        self._event_handlers = defaultdict(list)

        # What to do when we export
        def callback(content):
            display(Image(url=content.get('dataUrl')))
        self._connect_event('displayImg', callback)

    def add_representation(self, rep_type, options):
        '''Add a 3D representation to the viewer.  See User Guide for
        a complete description of the representations available.
        
        :return: An unique hexadecimal identifier for the representation.
        :rtype: str

        '''
        # Add our unique id to be able to refer to the representation
        rep_id = uuid4().hex
        self._remote_call('addRepresentation', type=rep_type, repId=rep_id, options=options)
        return rep_id

    def remove_representation(self, rep_id):
        '''Remove a representation from the viewer 

        :param str rep_id: the unique identifier generated by RepresentationViewer.add_representation

        '''
        self._remote_call('removeRepresentation', repId=rep_id)

    def update_representation(self, rep_id, options):
        '''Update a representation with new data.

        :param str rep_id: the unique identifier returned by RepresentationViewer.add_representation
        :param dict options: dictionary containing the updated data.

        '''
        self._remote_call('updateRepresentation', repId=rep_id, options=options)

    def _connect_event(self, event_name, callback):
        '''Respond to an event sent by the Javascript side.

        Events available:
        
            - displayImg
            - serialize


        '''
        self._event_handlers[event_name].append(callback)

    def _remote_call(self, method_name, **kwargs):
        '''Call a method remotely on the javascript side'''
        msg = {}
        msg['type'] = 'callMethod'
        msg['methodName'] = method_name
        msg['args'] = self._recursive_serialize(kwargs)

        if self.displayed is True:
            self.send(msg) # This will be received with View.on_msg
        else:
            # We should prepare a callback to be 
            # called when widget is displayed
            def callback(widget, msg=msg):
                widget.send(msg)
                widget.on_displayed(callback, remove=True) # Auto-unbind

            self._displayed_callbacks.append(callback)

    def _recursive_serialize(self, dictionary):
        '''Serialize a dictionary inplace'''
        for k, v in dictionary.iteritems():
            if isinstance(v, dict):
                self._recursive_serialize(v)
            else:
                # This is when custom serialization happens
                if isinstance(v, np.ndarray):
                    if v.dtype == 'float64':
                        # We don't support float64 on js side
                        v = v.astype('float32')
                    
                    dictionary[k] = encode_numpy(v)
        return dictionary

    def _handle_custom_msg(self, content):
        # Handle custom messages sent by the javascript counterpart
        event = content.get('event', '')
        for cb in self._event_handlers[event]:
            cb(content)


    def _ipython_display_(self, **kwargs):
        super(RepresentationViewer, self)._ipython_display_(**kwargs)
        self.displayed = True

class TrajectoryControls(DOMWidget):
    _view_name = Unicode('TrajectoryControls', sync=True)

    frame = CInt(sync=True)
    n_frames = CInt(sync=True)
    fps = CInt(sync=True)
    
    def __init__(self, n_frames, fps=30):
        '''Play/Pause controls useful for playing trajectories.

        Example:

        You can connect a callback to be executed every time the frame changes.

        .. code::

            from IPython.display import display

            controls = TrajectoryControls(10) # 10 frames
            
            def callback(frame):
                print("Current frame %d" % frame)

            controls.on_frame_change(callback)
            display(controls)

        .. py:attribute:: frame

            Current frame

        .. py:attribute:: n_frames

            Total number of frames

        .. py:attribute:: fps

            Frames per second (defaults to 30)

        '''
        super(TrajectoryControls, self).__init__()
        self.n_frames = n_frames - 1
        self.fps = fps

    def on_frame_change(self, callback):
        '''Connect a callback to be executed every time the frame attribute changes.'''
        self.on_trait_change(lambda name, old, new: callback(new), "frame")