from openride import BoundingBox, BoundingBoxCollection, Point, Polyline, Polygon, PointCloud
from openride.core.categories import Category, SubCategory
from openride.viewer import vtk_utils
from openride.viewer.camera import Camera
from openride.viewer.interactor import Interactor
from openride.viewer.hud import Hud
from openride.viewer.model_loader import ModelLoader
from openride.viewer.models import DEFAULT_MODELS

from typing import Iterable, List, Tuple, Union

import vtk



class Viewer:

    def __init__(self, resolution=(1600,900), background:Tuple[float,float,float]=(0.05,0.05,0.1), mouse_camera_interactions:bool=True):

        self.renderer = vtk.vtkRenderer()
        self.renderWindow = vtk.vtkRenderWindow()
        self.renderWindow.AddRenderer(self.renderer)

        self.set_resolution(resolution)
        self.set_background(background)

        self.camera = Camera(self)
        self.interactor = Interactor(self, mouse_camera_interactions)
        self.hud = Hud(self, position=(50,50))

        self._model_loader = ModelLoader()

        self._actors_to_clear = []
        self._text_actor_index = 0


    def set_resolution(self, resolution:Tuple[int, int]):
        self.renderWindow.SetSize(resolution[0], resolution[1])


    def set_background(self, color=Tuple[float,float,float]):
        self.renderer.SetBackground(*color)
        

    def update(self):
        self.interactor.update()
        self.hud.update()
        self.renderWindow.Render()
        self.__clear_actors()


    def render_actor(self, actor:vtk.vtkActor, persist:bool=False):
        self.renderer.AddActor(actor)
        if not persist: self._actors_to_clear.append(actor)


    def draw_point(self, point:Point, color:tuple=(1,1,1), size:int=0.15, alpha:float=1, resolution:int=10, persist:bool=False):
        actor = vtk_utils.get_point_actor(point.x, point.y, point.z, color, size, alpha, resolution)
        self.render_actor(actor, persist)


    def draw_point_cloud(self, pcloud:PointCloud, color:tuple=None, colors:Iterable[float]=None, colormap:str='jet', color_range:float=2**16, size:int=10, alpha:float=1, persist:bool=False):
        if color is None and colors is None: color = (1,1,1)
        actor = vtk_utils.get_point_cloud_actor(pcloud.get_point_cloud(), colors, color, colormap, size, alpha, color_range)
        self.render_actor(actor, persist)


    def draw_bounding_box(self, bounding_box:Union[BoundingBox, BoundingBoxCollection], color:tuple=(1,1,1), alpha:float=1, edge_color:tuple=None, persist:bool=False):

        if isinstance(bounding_box, BoundingBoxCollection):
            [self.draw_bounding_box(box, color, alpha, edge_color, persist) for box in bounding_box]
            return

        actor = vtk_utils.get_box_actor(
            bounding_box.position.x, bounding_box.position.y, bounding_box.position.z, 
            bounding_box.size.x, bounding_box.size.y, bounding_box.size.z,
            bounding_box.rotation.roll, bounding_box.rotation.pitch, bounding_box.rotation.yaw,
            color, alpha
        )

        if edge_color is not None:
            actor.GetProperty().SetEdgeVisibility(1)
            actor.GetProperty().SetLineWidth(10)
            actor.GetProperty().SetEdgeColor(*edge_color)

        self.render_actor(actor, persist)


    def draw_model(self, model:Union[str, Category, SubCategory], bounding_box:BoundingBox, number_polygons:int=100000, color:tuple=(1,1,1), alpha:float=1, persist:bool=False):
        """TODO"""
        p = bounding_box.position
        s = bounding_box.size
        r = bounding_box.rotation
        if isinstance(model, (Category, SubCategory)):
            file = DEFAULT_MODELS.get(model)
        else:
            file = model
        if not file: return
        actor = vtk_utils.get_model_actor(self._model_loader, file, number_polygons, p.x, p.y, p.z, s.x, s.y, s.z, r.roll, r.pitch, r.yaw, color, alpha)
        self.render_actor(actor, persist)


    def draw_polyline(self, line:Polyline, color:tuple=(1,1,1), alpha:float=1, linewidth:int=2, persist:bool=False):
        actor = vtk_utils.get_line_actor(line.vertices, color, alpha, linewidth)
        self.render_actor(actor, persist)


    def draw_polylines(self, lines:List[Polyline], color:tuple=(1,1,1), alpha:float=1, linewidth:int=2, persist:bool=False):
        actor = vtk_utils.get_lines_actor([line.vertices for line in lines], color=color, alpha=alpha, linewidth=linewidth)
        self.render_actor(actor, persist)


    def draw_polygon(self, polygon:Polygon, color:tuple=(1,1,1), alpha:float=1, persist:bool=False):
        actor = vtk_utils.get_polygon_actor(polygon.vertices, color=color, alpha=alpha)
        self.render_actor(actor, persist)


    def labelize(self, label:str, pos:Point, fontsize:int=16, persist:bool=False):
        if not self.camera.is_in_front(pos): return
        actor = vtk_utils.get_text_actor(label, pos.x, pos.y, pos.z, fontsize, self._text_actor_index)
        self.render_actor(actor, persist)
        self._text_actor_index += 1


    def draw_arrow(self, start:Point, end:Point, color:tuple=(1,1,1), alpha:float=1, linewidth:int=10, headsize:float=1, persist:bool=False):
        self.draw_line(Polyline([start, end]), color=color, alpha=alpha, linewidth=linewidth)
        actor = vtk_utils.get_cone_actor(
            start.x, start.y, start.z,
            end.x, end.y, end.z, 
            color, alpha, headsize,
        )
        self.render_actor(actor, persist)


    def draw_grid(self, origin:Point=Point(), delta:float=5, extent:float=200, alpha:float=1.0, persist:bool=True):
        lines = []

        for s in [0,1]:
            for i in range(int(2*extent/delta)+1):
                points = []
                if s == 0:
                    points.append(Point(-extent, -extent + i * delta) + origin)
                    points.append(Point( extent, -extent + i * delta) + origin)
                else:
                    points.append(Point(-extent + i * delta, -extent) + origin)
                    points.append(Point(-extent + i * delta,  extent) + origin)
                lines.append(Polyline(points))

        self.draw_polylines(lines, color=(0.3,0.3,0.3), alpha=alpha, linewidth=1, persist=persist)


    def __clear_actors(self):
        for actor in self._actors_to_clear:
            self.renderer.RemoveActor(actor)
        self._actors_to_clear = []
        self._text_actor_index = 0


    def close(self):
        del self.renderWindow
        self.interactor._interactor.InvokeEvent("ExitEvent")
