import numpy as np
from OpenGL import GL
import math

from .__prelude__ import (
    SelectedBand, Qt, Shader, Point, Matrix, MapPoint, MapProfile, Bound,
    inOpenGLContext, OpenGLPainter,
    dynamic, SELF, EACH,
    GLProgram, PureShader,
    logger
)
from .Scene import Scene, SceneDrawing
from .layer import LayerStack
from .TileCache import TileCache

MAP_SHADER = """
#version 330
float {geo_height}(vec3 geo_coords);
vec4 {geo_color}(vec3 geo_coords, inout float up_opacity);

float geo_height(vec3 geo_coords) {{
  return {geo_height}(geo_coords);
}}
vec4 geo_color(vec3 geo_coords) {{
  float up_opacity = 1.0;
  return {geo_color}(geo_coords, up_opacity);
}}
"""

def location_path():
    path = Qt.QPainterPath()
    path.moveTo(38., 80.)
    path.cubicTo(35., 62., 26., 60., 20., 50.)
    path.cubicTo(0., 0., 80., 0., 60., 50.)
    path.cubicTo(54., 60., 48., 62., 45., 80.)
    path.lineTo(38., 80.)
    return path

class MapScene(Scene):
    selected_band      = dynamic.readonly()
    band_colormap      = dynamic.readonly()
    tile_cache         = dynamic.readonly()
    points             = dynamic.readonly()
    temp_point         = dynamic.variable(None)
    profiles           = dynamic.readonly()
    temp_profile       = dynamic.variable(None)
    overlay_layer      = dynamic.readonly()

    flip_horizontally  = dynamic.variable(False)
    flip_vertically    = dynamic.variable(False)
    geo_focus          = dynamic.variable(None)

    center             = dynamic.variable(Point(0.0, 0.0))
    heightUnits        = dynamic.variable(0.0)
    distance           = dynamic.variable(0.0)
    yaw                = dynamic.variable(0.0)
    pitch              = dynamic.variable(0.0)

    def __init__(self, selected_band: SelectedBand, points, profiles, layers, band_colormap):
        super().__init__()
        self._selected_band = selected_band
        self._shaders = [
            Shader(Qt.QOpenGLShader.ShaderTypeBit.Vertex,"map/vertex.glsl"),
            Shader(Qt.QOpenGLShader.ShaderTypeBit.TessellationControl,"map/tess_control.glsl"),
            Shader(Qt.QOpenGLShader.ShaderTypeBit.TessellationEvaluation, "map/tess_eval.glsl"),
            Shader(Qt.QOpenGLShader.ShaderTypeBit.TessellationEvaluation, "map/utils.glsl"),
            Shader(Qt.QOpenGLShader.ShaderTypeBit.Fragment,"map/fragment.glsl"),
            Shader(Qt.QOpenGLShader.ShaderTypeBit.Fragment,"hsv_rgb.glsl"),
            Shader(Qt.QOpenGLShader.ShaderTypeBit.Fragment,"map/utils.glsl")
        ]
        self._tile_shaders = [
            Shader(Qt.QOpenGLShader.ShaderTypeBit.Vertex,"tile/vertex.glsl"),
            Shader(Qt.QOpenGLShader.ShaderTypeBit.Fragment,"tile/fragment.glsl"),
        ]
        self._band_colormap = band_colormap

        self._tile_cache = TileCache()
        self.tile_cache.new_tile_available.connect(self.renderChanged.emit)

        self.color = (0.5, 0.5, 0.5, 1.0)

        self.distance = 5.0
        self._points = points
        self._point_counter = 1
        self._profiles = profiles
        self._profile_counter = 1

        self._overlay_layer = LayerStack(
            SELF.selected_band.dataset[self],
            SELF.temp_point[self],
            SELF.temp_profile[self],
            layers
        )
        self._overlay_layer.renderChanged.connect(self.renderChanged.emit)
        self.temp_point = None
        self.temp_profile = None

        self.dynamic_attribute("_rerender").value_changed.connect(lambda: self._rerender())

    def init_from_dict(self, dct):
        self.yaw = dct['yaw']
        self.pitch = dct['pitch']
        self.distance = dct['distance']
        self.center = Point(*dct['center'])
        self.heightUnits = dct['heightUnits']
        self.flip_horizontally = dct.get('flip_horizontally', False)
        self.flip_vertically = dct.get('flip_vertically', False)

    def to_dict(self):
        return {
            'yaw': self.yaw,
            'pitch': self.pitch,
            'distance': self.distance,
            'center': [self.center.x, self.center.y],
            'heightUnits': self.heightUnits,
            'flip_horizontally': self.flip_horizontally,
            'flip_vertically': self.flip_vertically,
        }

    def next_point_name(self):
        counter = self._point_counter
        self._point_counter += 1
        return f"Point {counter}"

    @dynamic.memo(SELF.distance, SELF.pitch, SELF.yaw, SELF.center)
    def model_to_world(self):
        return Matrix.product(
            Matrix.translate((0.0,0.0,self.distance)),
            Matrix.rotate3(self.pitch, (-1.0, 0.0, 0.0, 0.0)),
            Matrix.rotate3(self.yaw, (0.0, 0.0, 1.0, 0.0)),
            Matrix.translate((self.center.x, self.center.y, 0.0))
        )
    @dynamic.memo(SELF.flip_horizontally, SELF.flip_vertically)
    def model_to_texture(self):
        h_flip = -1.0 if self.flip_horizontally else 1.0
        v_flip = -1.0 if self.flip_vertically else 1.0
        return Matrix.product(
            Matrix.scale((h_flip, v_flip, 1.0)),
            Matrix.scale((0.5,0.5,1.0)),
            Matrix.translate((h_flip,v_flip))
        )

    @dynamic.memo(SELF.band_colormap.xzero, SELF.band_colormap.xone)
    def value_to_elevation(self):
        lower, upper = self.band_colormap.xzero, self.band_colormap.xone
        scale = max(abs(lower), abs(upper))
        return Matrix.scale((1./scale, 1.0))
    @dynamic.memo(SELF.selected_band.dataset)
    def image_to_geo(self):
        w,h = self.selected_band.size
        return Matrix.product(
            self.selected_band.dataset.pixel_to_crs,
            Matrix.scale((w,h,1))
        )

    @dynamic.memo(SELF.selected_band.dataset)
    def show_background(self):
        return self.selected_band.dataset.is_georeferenced

    @dynamic.method(SELF.selected_band.image, SELF.selected_band.reference_image,
                    SELF.heightUnits, SELF.model_to_world, SELF.temp_point, SELF.temp_profile,
                    SELF.model_to_texture, SELF.geo_focus,
                    SELF.band_colormap.image, SELF.band_colormap.value_to_color,
                    SELF.points / EACH(SELF.show_marker),
                    SELF.profiles / EACH(SELF.show_focus_marker))
    def _rerender(self):
        self.renderChanged.emit()

    def create_point(self, dragStart: tuple[int, int], dragEnd: tuple[int, int]):
        xa,ya = dragStart
        xb,yb = dragEnd
        model_to_image = Matrix.product(
            self.selected_band.texture.texture_to_image,
            self.model_to_texture
        )
        def sq(x):
            return x*x
        dx_image, dy_image, _ = model_to_image.transform_vect((xb-xa, yb-ya, 0.0))
        w,h = self.selected_band.size
        # Clamp the radius to 100 pixels, to avoid users accidentally
        # computing the mean of the whole dataset
        radius = min(int(math.sqrt(sq(dx_image*w) + sq(dy_image*h))), 100)
        x_image, y_image = model_to_image.transform_point(dragStart)
        return MapPoint("Point", x_image*w, y_image*h, radius, Qt.QColor("blue"))
    def create_profile(self, dragStart):
        dragEnd = dragStart
        x0, y0 = self.__model_to_pixel(*dragStart)
        x1, y1 = self.__model_to_pixel(*dragEnd)
        counter = self._profile_counter
        self._profile_counter = counter + 1
        return MapProfile(f"Profile {counter}", 10, Qt.QColor("orange"), True, x0, y0, [(x1, y1)])
    def set_last_temp_profile_point(self, position):
        self.temp_profile.end_points[-1] = self.__model_to_pixel(*position)
        self.dynamic_attribute("temp_profile").value_changed.emit(self.temp_profile)
    def append_temp_profile_point(self, position):
        new_point = self.__model_to_pixel(*position)
        if len(self.temp_profile.end_points) < 2 or new_point != self.temp_profile.end_points[-2]:
            self.temp_profile.end_points.append(new_point)
        self.dynamic_attribute("temp_profile").value_changed.emit(self.temp_profile)

    def __model_to_pixel(self, x_model, y_model):
        model_to_image = Matrix.product(
            self.selected_band.texture.texture_to_image,
            self.model_to_texture
        )
        x_image, y_image = model_to_image.transform_point((x_model, y_model))
        w , h = self.selected_band.size
        return (int(x_image*w), int(y_image*h))

    def initializeDrawing(self, context):
        gl = context.functions()
        logger.debug("Loading map scene (max texture size : %d)", gl.glGetIntegerv(GL.GL_MAX_TEXTURE_SIZE))
        return InitializedMapScene(context, self)

class InitializedMapScene(SceneDrawing["MapScene"]):
    class IdleMode(SceneDrawing.IdleMode):
        def on_mouse_wheel(self, position, angle_x, angle_y):
            zoom = angle_y / 32.0
            def doit():
                self.drawing.distance *= math.exp(math.log(0.25)*math.radians(zoom))
            self.initialized_drawing.preserving_coords(position, doit)
            return self

        def on_mouse_button_pressed(self, position, button):
            if button == Qt.Qt.MouseButton.LeftButton:
                return self.initialized_drawing.DraggingMode(self.initialized_drawing, position)
            return self

        def on_key_pressed(self, button):
            if button == Qt.Qt.Key_Control:
                return self.initialized_drawing.ControlMode(self.initialized_drawing)
            elif button == Qt.Qt.Key_Shift:
                return self.initialized_drawing.ShiftMode(self.initialized_drawing)
            elif button == Qt.Qt.Key_Right or button == Qt.Qt.Key_Up:
                self.drawing.selected_band.band_number += 1
                return self
            elif button == Qt.Qt.Key_Left or button == Qt.Qt.Key_Down:
                self.drawing.selected_band.band_number -= 1
                return self
            return self

    class DraggingMode(SceneDrawing.IdleMode):
        def __init__(self, initialized_map_scene, start_pos):
            super().__init__(initialized_map_scene)
            self.__start_position = start_pos

        def on_mouse_move(self, new_pos):
            xa,ya = self.initialized_drawing.clip_to_model(*self.__start_position)
            xb,yb = self.initialized_drawing.clip_to_model(*new_pos)
            self.drawing.center += Point(xb-xa, yb-ya)
            return self.initialized_drawing.DraggingMode(self.initialized_drawing, new_pos)

        def on_mouse_button_released(self, pos, button):
            return self.initialized_drawing.IdleMode(self.initialized_drawing)

    class ControlMode(SceneDrawing.IdleMode):
        def on_mouse_wheel(self, position, angle_x, angle_y):
            angle = angle_y / 32.0
            def doit():
                self.drawing.yaw += angle
            self.initialized_drawing.preserving_coords(position, doit)
            return self

        def on_key_released(self, button):
            return self.initialized_drawing.IdleMode(self.initialized_drawing)
        def on_key_pressed(self, key):
            if key == Qt.Qt.Key_Shift:
                return self.initialized_drawing.ControlShiftMode(self.initialized_drawing)
            return self

        def on_mouse_button_pressed(self, position, button):
            if button == Qt.Qt.MouseButton.LeftButton:
                return self.initialized_drawing.CreatePointMode(self.initialized_drawing, position, self.__class__)
            return self

    class CreatePointMode(SceneDrawing.IdleMode):
        def __init__(self, initialized_scene, start_pos, return_mode):
            super().__init__(initialized_scene)
            self.__start_position = start_pos
            self.__return_mode = return_mode

        def on_mouse_move(self, position):
            model_center = self.initialized_drawing.clip_to_model(*self.__start_position)
            model_edge = self.initialized_drawing.clip_to_model(*position)
            self.drawing.temp_point = self.drawing.create_point(model_center, model_edge)
            return self

        def on_mouse_button_released(self, position, button):
            model_center = self.initialized_drawing.clip_to_model(*self.__start_position)
            model_edge = self.initialized_drawing.clip_to_model(*position)
            point = self.drawing.create_point(model_center, model_edge)
            point.name = self.drawing.next_point_name()
            self.drawing.points.append(point)
            self.drawing.temp_point = None
            return self.__return_mode(self.initialized_drawing)

    class ShiftMode(SceneDrawing.IdleMode):
        def on_mouse_wheel(self, position, angle_x, angle_y):
            angle = angle_y / 32.0
            def doit():
                self.drawing.pitch = max(0.0,min(85.0,self.drawing.pitch + angle))
            self.initialized_drawing.preserving_coords(position, doit)
            return self

        def on_key_released(self, button):
            return self.initialized_drawing.IdleMode(self.initialized_drawing)
        def on_key_pressed(self, key):
            if key == Qt.Qt.Key_Control:
                return self.initialized_drawing.ControlShiftMode(self.initialized_drawing)
            return self

        def on_mouse_button_pressed(self, position, button):
            if button == Qt.Qt.MouseButton.LeftButton:
                return self.initialized_drawing.CreateProfileMode(self.initialized_drawing, position)
            return self

    class CreateProfileMode(SceneDrawing.IdleMode):
        def __init__(self, initialized_scene, position):
            super().__init__(initialized_scene)
            self.drawing.temp_profile = self.drawing.create_profile(self.initialized_drawing.clip_to_model(*position))

        def __finalize_profile(self):
            profile = self.drawing.temp_profile
            self.drawing.temp_profile = None
            if len(profile.end_points) > 1:
                profile.end_points[-1:] = []
                self.drawing.profiles.append(profile)
            return self.initialized_drawing.IdleMode(self.initialized_drawing)

        def on_key_released(self, key):
            if key == Qt.Qt.Key_Shift:
                return self.__finalize_profile()
            return self

        def on_mouse_button_released(self, position, button):
            if button == Qt.Qt.MouseButton.LeftButton:
                self.drawing.append_temp_profile_point(self.initialized_drawing.clip_to_model(*position))
            return self
        def on_mouse_button_pressed(self, position, button):
            if button == Qt.Qt.MouseButton.LeftButton:
                self.drawing.append_temp_profile_point(self.initialized_drawing.clip_to_model(*position))
            if button == Qt.Qt.MouseButton.RightButton:
                return self.__finalize_profile()
            return self

        def on_mouse_wheel(self, position, angle_x, angle_y):
            self.drawing.temp_profile.r = max(0, self.drawing.temp_profile.r + int(angle_y / 64.0))
            self.drawing.dynamic_attribute("temp_profile").value_changed.emit(self.drawing.temp_profile)
            return self

        def on_mouse_move(self, position):
            self.drawing.set_last_temp_profile_point(self.initialized_drawing.clip_to_model(*position))
            return self

    class ControlShiftMode(SceneDrawing.IdleMode):
        def on_mouse_wheel(self, position, angle_x, angle_y):
            angle = angle_y / 32.0
            def doit():
                self.drawing.heightUnits = min(1.0, max(0, self.drawing.heightUnits + angle/360.0))
            self.initialized_drawing.preserving_coords(position, doit)
            return self

        def on_key_released(self, button):
            if button == Qt.Qt.Key_Control:
                return self.initialized_drawing.ShiftMode(self.initialized_drawing)
            elif button == Qt.Qt.Key_Shift:
                return self.initialized_drawing.ControlMode(self.initialized_drawing)
            return self

    def __init__(self, context, map_scene):
        super().__init__(context, map_scene)

        self.__tile_program: Qt.QOpenGLShaderProgram = GLProgram(map_scene._tile_shaders)
        self.__tile_mesh = self.__tile_program.create_square_mesh(GL.GL_TRIANGLES, split=0, texture_to_model = Matrix.identity(3))

        self.__initialized_overlay = self.drawing.overlay_layer.GL_initialize(context)
        GL.glPatchParameteri(GL.GL_PATCH_VERTICES, 3)

    def idle_mode(self):
        return self.IdleMode(self)

    def preserving_coords(self, coords, doit):
        x_b, y_b = self.clip_to_model(*coords)
        doit()
        x_a, y_a = self.clip_to_model(*coords)
        self.drawing.center += Point(x_a - x_b, y_a - y_b)

    @dynamic.memo(SELF.drawing.overlay_layer.shaders)
    def GL_program(self):
        return GLProgram(self.drawing._shaders + [
            PureShader(Qt.QOpenGLShader.ShaderTypeBit.Fragment, shader)
            for shader in self.drawing.overlay_layer.shaders
        ] + [
            PureShader(Qt.QOpenGLShader.ShaderTypeBit.TessellationEvaluation, shader)
            for shader in self.drawing.overlay_layer.shaders
        ] + [
            PureShader(Qt.QOpenGLShader.ShaderTypeBit.Fragment, MAP_SHADER.format(
                geo_color = self.drawing.overlay_layer.local_ident("geo_color"),
                geo_height = self.drawing.overlay_layer.local_ident("geo_height")
            )),
            PureShader(Qt.QOpenGLShader.ShaderTypeBit.TessellationEvaluation, MAP_SHADER.format(
                geo_color = self.drawing.overlay_layer.local_ident("geo_color"),
                geo_height = self.drawing.overlay_layer.local_ident("geo_height")
            ))

        ])
    @dynamic.memo(SELF.GL_program, SELF.drawing.model_to_texture)
    def GL_mesh(self):
        return self.GL_program.create_square_mesh(
            GL.GL_PATCHES, split=8,
            texture_to_model = self.drawing.model_to_texture.inverse())

    @dynamic.memo(SELF.viewport_size)
    def world_to_clip(self):
        w,h = self.viewport_size
        hfactor = max(1.0, w/h)/4.0
        vfactor = max(1.0, h/w)/4.0
        near = 4.0*0.005
        return Matrix.frustum(-hfactor*near, hfactor*near,
                              -vfactor*near, vfactor*near,
                              -near, -1000.0)
    @dynamic.memo(SELF.viewport_size)
    def tiles_per_screen(self):
        w,h = self.viewport_size
        return max(w/256.0, h/256.0) # Each tile is 256 pixels wide
    @dynamic.memo(SELF.viewport_size)
    def clip_to_pixel(self):
        w,h = self.viewport_size
        return Matrix.scale((w/2.0, -h/2.0, 1.0)) * Matrix.translate((1.0,-1.0))

    def paintGL(self):
        program = self.GL_program
        mesh = self.GL_mesh
        glfunc = self.context.functions()
        setup_overlay = self.__initialized_overlay.GL_setup

        glfunc.glEnable(GL.GL_DEPTH_TEST)

        glfunc.glClearColor(*self.drawing.color)
        glfunc.glClear(int(GL.GL_COLOR_BUFFER_BIT) | int(GL.GL_DEPTH_BUFFER_BIT))

        if self.drawing.show_background:
            if self.drawing.pitch < 60.0:
                self.drawing.tile_cache.GL_paint_background(
                    glfunc,
                    model_to_crs = Matrix.product(
                        self.drawing.selected_band.image_to_crs,
                        self.drawing.selected_band.texture.texture_to_image,
                        self.drawing.model_to_texture
                    ),
                    clip_to_model = self.clip_to_model,
                    model_to_world = self.drawing.model_to_world,
                    world_to_clip = self.world_to_clip,
                    tiles_per_screen = self.tiles_per_screen,
                )

        dem_texture = self.drawing.selected_band.texture
        ref_texture = self.drawing.selected_band.reference_texture
        tex_to_image = dem_texture.texture_to_image

        glfunc.glEnable(GL.GL_BLEND)
        glfunc.glBlendFuncSeparate(GL.GL_SRC_ALPHA, GL.GL_ONE_MINUS_SRC_ALPHA, GL.GL_ONE, GL.GL_ONE)
        with program as p:
            p.uniforms.image = dem_texture
            p.uniforms.ref_image = ref_texture

            p.uniforms.texture_to_image = tex_to_image
            p.uniforms.image_to_texture = tex_to_image.inverse()

            p.uniforms.image_to_geo = self.drawing.image_to_geo
            p.uniforms.geo_to_model = Matrix.product(
                self.drawing.image_to_geo,
                tex_to_image,
                self.drawing.model_to_texture,
            ).inverse()
            setup_overlay(p)

            p.uniforms.heightUnits = float(self.drawing.heightUnits)
            p.uniforms.value_to_elevation = self.drawing.value_to_elevation
            _,_,z_near = self.world_to_clip.inverse().transform_point((0.0,0.0,-1.0))
            p.uniforms.near = z_near

            p.uniforms.world_to_clip = self.world_to_clip
            p.uniforms.model_to_world = self.drawing.model_to_world

            mesh.draw()

        with OpenGLPainter(*self.viewport_size) as painter:
            if self.drawing.geo_focus is not None:
                center = (tex_to_image * self.drawing.model_to_texture).inverse().transform_point(self.drawing.geo_focus)
                clip_cx, clip_cy, _ = (self.world_to_clip * self.drawing.model_to_world).transform_point((*center, 0.0))
                pixel_cx, pixel_cy = self.clip_to_pixel.transform_point((clip_cx, -clip_cy))
                painter.setBrush(Qt.QColor('red'))
                painter.drawEllipse(Qt.QPointF(pixel_cx, pixel_cy), 10 , 10)

            marker_path = location_path()
            path_x = marker_path.boundingRect().center().x()
            path_y = marker_path.boundingRect().bottom()

            w,h = self.drawing.selected_band.size

            point_to_model = Matrix.product(
                self.drawing.model_to_texture.inverse(),
                tex_to_image.inverse(),
                Matrix.scale((1/w, 1/h, 1)),
            )
            model_to_clip = Matrix.product(
                self.world_to_clip,
                self.drawing.model_to_world,
            )
            def marker_clip(point_x, point_y):
                model_x, model_y = point_to_model.transform_point((point_x, point_y))
                return model_to_clip.transform_point((model_x, model_y, 0.0))
            def show_marker_at(color, clip_x, clip_y):
                pixel_x, pixel_y = self.clip_to_pixel.transform_point((clip_x, -clip_y))
                painter.setBrush(color)
                painter.save()
                painter.translate(pixel_x, pixel_y)
                painter.scale(1,-1)
                painter.translate(-path_x,-path_y)
                painter.drawPath(marker_path)
                painter.restore()

            markers = []
            for point in self.drawing.points:
                if point.show_marker:
                    markers.append((point.color, marker_clip(point.x, point.y)))

            for profile in self.drawing.profiles:
                if profile.show_focus_marker:
                    markers.append((profile.color, marker_clip(*profile.focus_center)))

            def marker_z(marker):
                match marker:
                    case _, coords:
                        return coords[2]
            for color, coords in sorted(markers, key=marker_z, reverse=True):
                clip_x, clip_y, _ = coords
                show_marker_at(color, clip_x, clip_y)

    def freeGL(self):
        self.__initialized_overlay.GL_free()
