'''
Graphical display classes. Experimental tasks involving graphical displays
inherit from these classes.
'''
import os

import numpy as np
from OpenGL.GL import *

from ..experiment import LogExperiment
from ..experiment import traits

from .render import stereo, render, ssao, shadow_map
from .models import Group
from .xfm import Quaternion, Transform
from .primitives import Sphere, Cube, Chain
from .environment import Box, Grid
from .primitives import Cylinder, Sphere, Cone
import socket

try:
    os.environ['PYGAME_HIDE_SUPPORT_PROMPT'] = "hide"
    import pygame
except ImportError:
    import warnings
    warnings.warn('riglib/stereo_opengl/window.py: not importing name pygame')

from config.rig_defaults import window as defaults

# for WindowDispl2D only
from .primitives import Shape2D

class Window(LogExperiment):
    '''
    Generic stereo window
    '''
    status = dict(draw=dict(stop=None))
    state = "draw"
    stop = False

    # XPS computer
    # window_size = (1920*2, 1080)
    window_size = traits.Tuple(defaults['window_size'], desc='Window size, in pixels')
    background = traits.Tuple((0.,0.,0.,1.), desc="Background color (R,G,B,A)")
    fullscreen = traits.Bool(True, desc="Fullscreen window")

    #Screen parameters, all in centimeters -- adjust for monkey
    screen_dist = traits.Float(defaults['screen_dist'], desc="Screen to eye distance (cm)")
    screen_half_height = traits.Float(defaults['screen_half_height'], desc="Screen half height (cm)")
    iod = traits.Float(2.5, desc="Intraocular distance (cm)")     # intraocular distance
    stereo_mode = traits.OptionsList(['hmd', 'mirror', 'projection', 'anaglyph'], desc="Stereo mode", 
                                     bmi3d_input_options=['hmd', 'mirror', 'projection', 'anaglyph'])

    show_environment = traits.Bool(False, desc="Show wireframe box around environment")

    hidden_traits = ['stereo_mode', 'screen_dist', 'screen_half_height', 'iod', 'show_environment', 'background']

    def __init__(self, *args, **kwargs):
        self.display_start_pos = kwargs.pop('display_start_pos', "0,0")
        super(Window, self).__init__(*args, **kwargs)

        self.models = []
        self.world = None
        self.event = None

        # os.popen('sudo vbetool dpms on')
        self.fov = np.degrees(np.arctan(self.screen_half_height/self.screen_dist))*2
        self.screen_cm = [2 * self.screen_half_height * self.window_size[0]/self.window_size[1], 2 * self.screen_half_height]

    def set_os_params(self):
        os.environ['SDL_VIDEO_WINDOW_POS'] = self.display_start_pos
        os.environ['SDL_VIDEO_X11_WMCLASS'] = "monkey_experiment"

    def screen_init(self):
        self.set_os_params()
        pygame.init()

        pygame.display.gl_set_attribute(pygame.GL_DEPTH_SIZE, 24)
        pygame.display.gl_set_attribute(pygame.GL_FRAMEBUFFER_SRGB_CAPABLE, True)
        pygame.display.gl_set_attribute(pygame.GL_CONTEXT_MAJOR_VERSION, 3)
        pygame.display.gl_set_attribute(pygame.GL_CONTEXT_MINOR_VERSION, 2)
        pygame.display.gl_set_attribute(pygame.GL_CONTEXT_PROFILE_MASK, pygame.GL_CONTEXT_PROFILE_CORE)
        flags = pygame.DOUBLEBUF | pygame.HWSURFACE | pygame.OPENGL | pygame.NOFRAME
        if self.fullscreen:
            flags = flags | pygame.FULLSCREEN
        try:
            pygame.display.gl_set_attribute(pygame.GL_MULTISAMPLEBUFFERS,1)
            self.screen = pygame.display.set_mode(self.window_size, flags)
        except:
            pygame.display.gl_set_attribute(pygame.GL_MULTISAMPLEBUFFERS,0)
            self.screen = pygame.display.set_mode(self.window_size, flags)

        glDisable(GL_FRAMEBUFFER_SRGB) # disable gamma correction
        glEnable(GL_BLEND)
        glDepthFunc(GL_LESS)
        glEnable(GL_DEPTH_TEST)
        glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)
        glClearColor(*self.background)
        glClearDepth(1.0)
        glDepthMask(GL_TRUE)
        glEnable(GL_CULL_FACE) # temporary solution to alpha blending issue with spheres. just draw the front half of the sphere
        glCullFace(GL_BACK)


        #this effectively determines the modelview matrix
        self.world = Group(self.models)
        self.world.init()

        #up vector is always (0,0,1), why would I ever need to roll the camera?!
        self.set_eye((0, 0, 0), (0,0))
        self.modelview = np.eye(4)
        self.modelview[:3,-1] = [0,0,-self.screen_dist]
        self.renderer = self._get_renderer()

        pygame.mouse.set_visible(False)

        if self.show_environment:
            self.add_model(Box())
    
    def _get_renderer(self):
        near = 1
        far = 1024
        if self.stereo_mode == 'mirror':
            self.modelview = np.dot(self.modelview, np.diag([-1,1,1,1]))
            glFrontFace(GL_CW);  # Switch to clockwise winding for mirrored objects
            return stereo.MirrorDisplay(self.window_size, self.fov, near, far, self.screen_dist, self.iod)
        if self.stereo_mode == 'hmd':
            return stereo.MirrorDisplay(self.window_size, self.fov, near, far, self.screen_dist, self.iod)
        elif self.stereo_mode == 'projection':
            return shadow_map.ShadowMapper(self.window_size, self.fov, near, far)
        elif self.stereo_mode == 'anaglyph':
            return stereo.Anaglyph(self.window_size, self.fov, near, far, self.screen_dist, self.iod)
        else:
            raise ValueError("Unknown stereo mode: %s" % self.stereo_mode)

    def set_eye(self, pos, vec, reset=True):
        '''Set the eye's position and direction. Camera starts at (0,0,0), pointing towards positive y'''
        self.world.translate(pos[0], pos[2], pos[1], reset=True).rotate_x(-90)
        self.world.rotate_y(vec[0]).rotate_x(vec[1])

    def add_model(self, model):
        if self.world is None:
            #world doesn't exist yet, add the model to cache
            self.models.append(model)
        else:
            #We're already running, initialize the model and add it to the world
            model.init()
            self.world.add(model)

    def remove_model(self, model):
        '''remove the second model from the last'''
        self.world.remove(model)

    def show_object(self, obj, show=False):
        '''
        Show or hide an object. This function is an abstraction so that tasks don't need to know about attach/detach
        '''
        if show:
            obj.attach()
        else:
            obj.detach()

    def draw_world(self):
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
        self.renderer.draw(self.world, modelview=self.modelview)
        pygame.display.flip()
        self.renderer.draw_done()

    def _get_event(self):
        for e in pygame.event.get(pygame.KEYDOWN):
            return (e.key, e.type)

    def _start_None(self):
        pygame.display.quit()

    def _test_stop(self, ts):
        '''
        Stop the task if the escape key is pressed, or if the super _test_stop instructs a stop
        '''
        super_stop = super(Window, self)._test_stop(ts)
        from pygame import K_ESCAPE
        key_stop = self.event is not None and self.event[0] == K_ESCAPE
        if key_stop:
            print("Window closed. Stopping task")
        return super_stop or key_stop
    
    def update_report_stats(self):
        self.reportstats['FPS'] = round(self.clock.get_fps(),2)
        return super().update_report_stats()
    
    def requeue(self):
        self.renderer._queue_render(self.world)

    def _cycle(self):
        super(Window, self)._cycle() # is this order intentional? why not cycle first then draw the screen?
        self.event = self._get_event()
        self.requeue()
        self.draw_world()

    
class WindowSSAO(Window):
    '''Screen Space Ambient Occlusion (SSAO) window. Replaces the renderer with one using SSAO'''

    def _get_renderer(self):
        return ssao.SSAO(self.window_size, self.fov, 0.05, 1024)


class WindowWithExperimenterDisplay(Window):
    hostname = socket.gethostname()
    # This class has a hard-coded window size
    if hostname == 'lynx':
        _stereo_window_flip = True
        _stereo_main_flip_z = True
        window_size = (1280+480, 1024)
        window_size1 = (1280, 1024)
        window_size2 = (480, 270)
    else:
        _stereo_window_flip = False
        _stereo_main_flip_z = False
        window_size = (1920 + 480, 1080)
        window_size1 = (1920, 1080)
        window_size2 = (480, 270)

    def __init__(self, *args, **kwargs):
        super(WindowWithExperimenterDisplay, self).__init__(*args, **kwargs)
        # This class has a hard-coded window size
        # self.window_size = (1920 + 480, 1080)

    def set_os_params(self):
        # NOTE: in Ubuntu Unity, setting the SDL_VIDEO_WINDOW_POS seems to be largely ignored.
        # You can set which screen the window appears on if you have a dual display, but you cannot set the exact position
        # Instead, you have to hard-code a render start location in the compiz-config settings manager
        # http://askubuntu.com/questions/452995/how-to-adjust-window-placement-in-unity-ubuntu-14-04-based-on-overlapping-top-b
        os.environ['SDL_VIDEO_WINDOW_POS'] = self.display_start_pos
        os.environ['SDL_VIDEO_X11_WMCLASS'] = "monkey_experiment_with_mini"

    def _get_renderer(self):
        near = 1
        far = 1024
        return stereo.DualMultisizeDisplay(self.window_size1, self.window_size2, self.fov, near, far, self.screen_dist, self.iod, flip=self._stereo_window_flip,
            flip_main_z = self._stereo_main_flip_z)


TRANSPARENT = (106,0,42)

class WindowDispl2D():
    '''For testing only -- draws world on a 2D screen. May cause mild confusion -- transforms 
    incoming 3D coordinates (x,y,z) into 2D coordinates (x,y) by mapping z onto y'''

    def __init__(self, *args, **kwargs):
        self.is_pygame_display = True
        super().__init__(*args, **kwargs)

    def screen_init(self):
        os.environ['SDL_VIDEO_WINDOW_POS'] = self.display_start_pos
        os.environ['SDL_VIDEO_X11_WMCLASS'] = "monkey_experiment"
        pygame.init()
        self.clock = pygame.time.Clock()

        flags = pygame.NOFRAME | pygame.DOUBLEBUF | pygame.HWSURFACE

        if self.fullscreen:
            flags = flags | pygame.FULLSCREEN

        self.display_border = 10

        self.size = np.array(self.window_size, dtype=np.float64)
        self.screen = pygame.display.set_mode(self.window_size, flags)
        self.screen_background = pygame.Surface(self.screen.get_size()).convert()
        self.screen_background.fill(self.background)

        x_max, y_max = (self.screen_cm[0]/2, self.screen_cm[1]/2)
        x_min, y_min = (-self.screen_cm[0]/2, -self.screen_cm[1]/2)
        self.normalize = np.array(np.diag([1./(x_max-x_min), 1./(y_max-y_min), 1]))
        self.center_xform = np.array([[1., 0, -x_min], 
                                      [0, 1., -y_min],
                                      [0, 0, 1]])
        self.norm_to_screen = np.array(np.diag(np.hstack([self.size, 1])))

        # the y-coordinate in pixel space has to be swapped graphics convention reasons
        self.flip_y_coord = np.array([[1, 0, 0],
                                      [0, -1, self.size[1]],
                                      [0, 0, 1]])
        self.pos_space_to_pixel_space = np.dot(self.flip_y_coord, np.dot(self.norm_to_screen, np.dot(self.normalize, self.center_xform)))

        self.world = Group(self.models)
        # Dont 'init' self.world in this Window. Just allocates a bunch of OpenGL stuff which is not necessary (and may not work in some cases)
        # self.world.init()

        #initialize surfaces for translucent markers
        self.surf={}
        self.surf['0'] = pygame.Surface(self.screen.get_size())
        self.surf['0'].fill(TRANSPARENT)
        self.surf['0'].set_colorkey(TRANSPARENT)

        self.surf['1'] = pygame.Surface(self.screen.get_size())
        self.surf['1'].fill(TRANSPARENT)
        self.surf['1'].set_colorkey(TRANSPARENT)

         #values of alpha: higher = less translucent
        self.surf['0'].set_alpha(170) #Cursor
        self.surf['1'].set_alpha(130) #Targets

        self.surf_background = pygame.Surface(self.surf['0'].get_size()).convert()
        self.surf_background.fill(TRANSPARENT)

        pygame.mouse.set_visible(False)

        self.i = 0

    def pos2pix(self, kfpos):
        # re-specify the point in homogenous coordinates
        pt = np.hstack([kfpos, 1]).reshape(-1, 1)

        # perform the homogenous transformation
        pix_coords = np.dot(self.pos_space_to_pixel_space, pt)

        pix_pos = np.array(pix_coords[:2,0], dtype=int)
        return pix_pos

    def get_surf(self):
        return self.surf[str(np.min([self.i,1]))]

    def draw_model(self, model):
        '''
        Draw a single Model on the current surface, or recurse if the model is a composite model (i.e., a Group)
        '''
        color = tuple([int(255*x) for x in model.color[0:3]])
        if isinstance(model, Sphere):
            pos = model._xfm.move[[0,2]]
            pix_pos = self.pos2pix(pos)

            rad = model.radius
            pix_radius = self.pos2pix(np.array([model.radius, 0]))[0] - self.pos2pix([0,0])[0]

            #Draws cursor and targets on transparent surfaces
            pygame.draw.circle(self.get_surf(), color, pix_pos, pix_radius)

        elif isinstance(model, Shape2D):
            # model.draw() returns True if the object was drawn
            #   (which happens if the object's .visible attr is True)
            if model.draw(self.get_surf(), self.pos2pix):
                pass
        elif isinstance(model, Cube):
            pos = model.xfm.move[[0,2]]
            side_len = model.side_len

            left = pos[0] - side_len/2
            right = pos[0] + side_len/2
            top = pos[1] + side_len/2
            bottom = pos[1] - side_len/2

            top_left = np.array([left, top])
            bottom_right = np.array([right, bottom])
            top_left_pix_pos = self.pos2pix(top_left)
            bottom_right_pix_pos = self.pos2pix(bottom_right)

            rect = pygame.Rect(top_left_pix_pos, bottom_right_pix_pos - top_left_pix_pos)
            color = tuple([int(255*x) for x in model.color[0:3]])

            pygame.draw.rect(self.get_surf(), color, rect)

        elif isinstance(model, (Cylinder, Cone)):
            vec_st = np.array([0., 0, 0, 1])
            vec_end = np.array([0., 0, model.height, 1])

            cyl_xform = model._xfm.to_mat()
            cyl_start = np.dot(cyl_xform, vec_st)
            cyl_end = np.dot(cyl_xform, vec_end)
            pix_radius = self.pos2pix(np.array([model.radius, 0]))[0] - self.pos2pix([0,0])[0]

            pygame.draw.line(self.get_surf(), color, self.pos2pix(cyl_start[[0,2]]), self.pos2pix(cyl_end[[0,2]]), pix_radius)

            # print cyl_start, cyl_end

        elif isinstance(model, Group):
            for mdl in model:
                self.draw_model(mdl)

    def _draw_other(self):
        pass

    def draw_world(self):
        #Refreshes the screen with original background
        self.screen.blit(self.screen_background, (0, 0))
        self.surf['0'].blit(self.surf_background,(0,0))
        self.surf['1'].blit(self.surf_background,(0,0))

        # surface index
        self.i = 0

        for model in self.world.models:
            self.draw_model(model)
            self.i += 1

        #Renders the new surfaces
        self.screen.blit(self.surf['0'], (0,0))
        self.screen.blit(self.surf['1'], (0,0))

        self._draw_other()
        pygame.display.update()

    def requeue(self):
        '''
        Simulation 'requeue' does nothing because the simulation is lazy and
        inefficient and chooses to redraw the entire screen every loop
        '''
        pass


class Window2D():
    '''
    Render the world in 2D without lighting
    '''
    def _get_renderer(self):
        glCullFace(GL_BACK)
        return render.Renderer2D(self.screen_cm)
    
class FakeWindow(Window):
    '''
    A dummy class to secretly avoid rendering graphics without
    the graphics-based tasks knowing about it. Used e.g. for simulation
    purposes where the graphics only slow down the simulation.
    '''
    background = (1,1,1,1)
    def __init__(self, *args, **kwargs):
        self.models = []
        self.world = None
        self.event = None
        super(FakeWindow, self).__init__(*args, **kwargs)

    def screen_init(self):
        self.world = Group(self.models)
        # self.world.init()

    def draw_world(self):
        pass

    def requeue(self):
        pass

    def _start_reward(self, *args, **kwargs):
        n_rewards = self.calc_state_occurrences('reward')
        if n_rewards % 10 == 0:
            print(n_rewards)
        super(FakeWindow, self)._start_reward(*args, **kwargs)

    def show_object(self, *args, **kwargs):
        pass

    def _get_event(self):
        pass


class FPScontrol(Window):
    '''A mixin that adds a WASD + Mouse controller to the window.
    Use WASD to move in XY plane, q to go down, e to go up'''

    def screen_init(self):
        super(FPScontrol, self).screen_init()
        pygame.mouse.set_visible(False)
        self.eyepos = [0,-self.screen_dist, 0]
        self.eyevec = [0,0]
        self.wasd = [False, False, False, False, False, False]

    def _get_event(self):
        retme = None
        for e in pygame.event.get([pygame.MOUSEMOTION, pygame.KEYDOWN, pygame.KEYUP, pygame.QUIT]):
            moved = any(self.wasd)
            if e.type == pygame.MOUSEMOTION:
                self.world.xfm.rotate *= Quaternion.from_axisangle((1,0,0), np.radians(e.rel[1]*.1))
                self.world.xfm.rotate *= Quaternion.from_axisangle((0,0,1), np.radians(e.rel[0]*.1))
                self.world._recache_xfm()
            elif e.type == pygame.KEYDOWN:
                kn = pygame.key.name(e.key)
                if kn in ["escape", "q"]:
                    self.stop = True
                retme = (e.key, e.type)
            elif e.type == pygame.QUIT:
                self.stop = True

            if moved:
                self.set_eye(self.eyepos, self.eyevec, reset=True)
        return retme
