import os
import sys
import signal
import time
from pathlib import Path
import argparse
import numpy as np
import OpenGL.GL as gl
try:
    from PIL import Image
    from .glfwBackend import glfwApp
    import imgui
    from imgui.integrations.glfw import GlfwRenderer
except ImportError as e:
    raise SystemExit("%s:\nAppliaction needs extra packages to run.\n"
            "Please install Pillow, imgui and glfw." % e)
signal.signal(signal.SIGINT, signal.SIG_DFL)

# Local imports
from fieldanimation import FieldAnimation, glInfo

CHOICES = (
    'epole',
    "Duffing's equation",
    'Structurally stable system',
    'Transcritical bifurcation',
    'Reaction diffusion',
    'Positive invariant set',
    'Spiral ccw',
    'Spiral cw',
    'wind',
    'gmod',
    )

#------------------------------------------------------------------------------
def ElectricField(q, r0, x, y):
    """ Return the electric field vector E=(Ex, Ey) due to charge q at r0.
    """
    den = np.hypot(x - r0[0], y - r0[1]) ** 1.5

    return q * (x - r0[0]) / den, q * (y - r0[1]) / den

#------------------------------------------------------------------------------
def createField(eq='Spiral ccw', m=64, n=64):
# def createField(eq='Spiral ccw', m=81, n=201):
    """ Equations are taken from
        http://reference.wolfram.com/language/ref/StreamPlot.html

        Args:
            eq (string): equation name
            m (integer): number of rows
            n (integer): number of columns
    """
    mj = m * 1j
    nj = n * 1j
    Y, X = np.mgrid[-3:3:mj, -3:3:nj]

    if eq == 'epole':
        # Return n x m x 2 2D array with Electric field values
        # generated by a (2 * charges) electric charges configuration
        charges = 1
        nq = 2 ** int(charges)
        charges = []
        for i in range(nq):
            q = i % 2 * 2 - 1
            charges.append((q,
                    (np.cos(2 * np.pi * i / nq), np.sin(2 * np.pi * i / nq))))
        U, V = np.zeros((m, n)), np.zeros((m, n))
        for charge in charges:
            ex, ey = ElectricField(*charge, x=X, y=Y)
            U += ex
            V += ey

    elif eq == "Duffing's equation":
        # Duffing's equation
        U = Y.copy()
        V = X - X**3
    elif eq == 'Structurally stable system':
        # Structurally stable system
        U = 2 * X + Y - X * (X**2 + Y**2)
        V = -Y - Y * (X**2 + Y**2)
    elif eq == 'Transcritical bifurcation':
        # Transcritical bifurcation
        U = X**2
        V = -Y
    elif eq == 'Reaction diffusion':
        U = 2 * (Y - X) + X * (1 - X ** 2)
        V = -2 * (Y - X) + Y * (1 - Y ** 2)
    elif eq == 'Positive invariant set':
        U = Y.copy()
        V = - X + Y * (1 - X ** 2 - 2 * Y ** 2)
    elif eq == 'Spiral ccw':
        origin = (0.8, 0.8)
        X, Y = X - origin[0], Y - origin[1]
        U = X - Y
        V = X + Y
    elif eq == 'Spiral cw':
        origin = (0.8, 0.8)
        X, Y = X - origin[0], Y - origin[1]
        U = Y - X
        V = -X - Y
    elif eq == 'wind':
        mydir = Path(__file__).parent
        field = np.load(mydir / "wind_2016-11-20T00-00Z.npy")
        r, c, bands = field.shape
        Y, X = np.mgrid[0:r, 0:c]
        U = field[:, :, 0][::-1]
        V = - field[:, :, 1][::-1]
    elif eq == 'gmod':
        mydir = Path(__file__).parent
        U = np.load(mydir / 'vx.npy')[::-1]
        V = np.load(mydir / 'vy.npy')[::-1]
        r,c = U.shape
        Y, X = np.mgrid[0:r, 0:c]
    else:
        raise SystemExit("Unknown field. Giving up...")

    return np.flipud(np.dstack((U, -V)))

#------------------------------------------------------------------------------
def userInterface(graphicItem, app):
    """ Control graphicItem parameters interactively
    """
    imgui.new_frame()
    expanded, opened = imgui.begin('Controls', closable=True,
            flags=imgui.WINDOW_ALWAYS_AUTO_RESIZE)
    if opened:
        # Field to show
        current = app.ifield
        changed, current = imgui.combo('Field', current, list(CHOICES))
        if changed:
            app.setField(current)

        # Speed Rate
        changed, speed = imgui.drag_float('Speed',
            graphicItem.speedFactor, 0.01, 0.0, 10.0)
        if changed:
            graphicItem.speedFactor = speed

        # Decay
        changed, decay = imgui.drag_float('Decay',
            graphicItem.decay, 0.001, 0.001, 1.0)
        if changed:
            graphicItem.decay = decay

        # Drop Rate Bump
        changed, decayBoost = imgui.drag_float('Decay boost',
            graphicItem.decayBoost, 0.01, 0.001, 1.0)
        if changed:
            graphicItem.decayBoost = decayBoost

        # Unknown const
        changed, opacity = imgui.drag_float('Opacity',
            graphicItem.fadeOpacity, 0.001, 0.900, 0.999, '%.4f')
        if changed:
            graphicItem.fadeOpacity = opacity

        # Palette
        changed, color = imgui.color_edit3('Color', *graphicItem.color)
        if changed:
            graphicItem.color = color
        imgui.same_line()
        changed, palette = imgui.checkbox("Palette", graphicItem.palette)
        if changed:
            graphicItem.palette = palette

        changed, bg_color = imgui.color_edit4('Background color',
                *app.bg_color)
        if changed:
            app.bg_color = bg_color

        # Point size
        changed, pointSize = imgui.input_int("Point size",
            graphicItem.pointSize, 1, 1, 1)
        if changed:
            if pointSize > 5:
                pointSize = 5
            elif pointSize < 1:
                pointSize = 1
            graphicItem.pointSize = pointSize

        # Number of Points
        changed, tracersCount = imgui.drag_int("Number of "
            "Tracers", graphicItem.tracersCount, 1000.0, 4000, 10000000)
        if changed:
            graphicItem.tracersCount = tracersCount

        # Periodic border
        changed, periodic = imgui.checkbox("Periodic", graphicItem.periodic)
        if changed:
            graphicItem.periodic = periodic

        # Draw field
        changed, drawfield = imgui.checkbox("Draw Field",
                graphicItem.drawField)
        if changed:
            graphicItem.drawField = drawfield
    imgui.end()
    imgui.render()
    return opened

#==============================================================================
class GLApp(glfwApp):
    def __init__(self, title, width, height, options):
        super(GLApp, self).__init__(title, width, height)

        if options.image is not None:
            options.image = np.flipud(
                    np.asarray(Image.open(options.image), np.uint8))

        self.ifield = CHOICES.index(options.choose)
        field = createField(CHOICES[self.ifield])

        glversion = glInfo()['glversion']
        if not options.use_fragment and glversion < 4.3:
            print("WARNING..... Compute shaders not available with OpenGL"
                    " ver. %.1f." % glversion)
            useCompute = False
        else:
            useCompute = True
        # Add Field Animation overlay
        self._fa = FieldAnimation(width, height, field, useCompute,
                options.image)
        if options.draw_field:
            self._fa.drawField = True
        self._t0 = time.time()
        self._fps = 0
        self.options = options
        self.makeGUI(options.gui)

    def makeGUI(self, gui):
        """ Create the GUI context and GlfwRenderer

            Args:
                gui (bool): if True create the GUI context and GlfwRenderer
        """
        if gui:
            imgui.create_context()
            self._renderer = GlfwRenderer(self.window(), True)
            self._opened = True
        else:
            self._renderer = None
            self._opened = False

    def renderScene(self):
        super(GLApp, self).renderScene()
        self._fa.draw()
        self._fps += 1
        if self._opened:
            self._renderer.process_inputs()
            self._opened = userInterface(self._fa, self)
            self._renderer.render(imgui.get_draw_data())
            if not self._opened:
                app.restoreKeyCallback()

        now = time.time()
        if now - self._t0 >= 1:
            if self.options.fps:
                self.setTitle("%s - %s FPS" % (self.title(), self._fps))
            self._fps = 0
            self._t0 = time.time()

    def onKeyboard(self, window, key, scancode, action, mode):
        if key == GLApp.KEY_G and action == GLApp.PRESS:
            # Draw the GUI
            self.makeGUI(not self._opened)

        super(GLApp, self).onKeyboard(window, key, scancode, action, mode)

    def onResize(self, window, width, height):
        gl.glViewport(0, 0, width, height)
        self._fa.setSize(width, height)

    def setField(self, ifield):
        field = createField(CHOICES[ifield])
        self._fa.setField(field)
        self.setTitle('%s' % CHOICES[ifield])
        self.ifield = ifield

#------------------------------------------------------------------------------
def makeParser():
    """ Return the ArgumentParser instance for this script.

        Returns:
            ArgumentParser: instance
    """
    parser = argparse.ArgumentParser(
            description='\nField Animation example. Press "G" to display the'
                    ' GUI while running.',
            add_help=True,
            formatter_class=argparse.ArgumentDefaultsHelpFormatter,
            )

    parser.add_argument('-f', '--draw_field', action='store_true',
            default=False,
            help=("Draw vector field as background image ")
            )
    parser.add_argument('-c', '--choose',
            choices=CHOICES,
            default="epole",
            help=("Choose field to animate ")
            )
    parser.add_argument('-p', '--fps', action='store_true', default=False,
            help=("Count Frames Per Second ")
            )
    parser.add_argument('-u', '--use-fragment', action='store_true',
            default=False, help=("Use fragment instead of compute shader ")
            )
    parser.add_argument('-g', '--gui', action='store_true', default=False,
            help=("Add gui control window ")
            )

    parser.add_argument('-i', '--image', action='store', default=None,
            help=("Load image as background texture")
            )
    return parser

#------------------------------------------------------------------------------
if __name__ == "__main__":
    options = makeParser().parse_args(sys.argv[1:])

    app = GLApp('Field Animation', 800, 800, options)
    app.run()
