import sys
import time
import threading
import warnings
import numpy as np

from PyQt5.QtCore import QSizeF, QLoggingCategory
from PyQt5.QtGui import QColor, QPixmap, QPainter, QPen, QPdfWriter, QGuiApplication
from PyQt5.QtWidgets import QApplication

# Assuming you separate this into its own file
from .qt_viz_window import QtVizWindow


class WorldVisuals:
    """
    A rendering class for the sfctools.core.World where agents can be plotted visually
    Thisis based on a QApplication 
    """

    def __init__(self, geometry=None):
        # set up world visuals

        self.engine = None  # visualization engine ('Qt' or 'Pdf')

        if geometry is None:
            geometry = (800, 800)
        self.geometry = geometry  # dimensions of the drawing canvas
        self.handle_hover = True

        self.window = None  # reference to a window instance (UI, if any)

        self.running = False  # are the visuals currently running?
        self._is_open = False  # are the visuals 'opened'?

        # background
        self.bg_color = "white"  # background color as str
        self.bg_pixmap = None  # background image pixel map (placeholder)

        # objects
        self.object_buffer = []  # object buffer containing all objects to plot
        self.selection = None

        # coordinates
        self.X, self.Y = geometry[0] * 0.3, geometry[1] * 0.3
        self.Z = 0

        self.dX = 0  # shift in x-y direction
        self.dY = 0

        # for Qt
        self.qt_stylesheet = ""

        # for PDF
        self.pdf_filepath = None

        # application instances
        self.app_thread = None
        self.instance = None

    def enable_hover(self):
        self.handle_hover = True

    def disable_hover(self):
        self.handle_hover = False

    def is_paused(self):
        return self.window and self.window.paused

    def scale(self, val):
        self.window.scale = val

    def close(self):
        """Safely close the visualization."""
        instance = QApplication.instance()
        if instance:
            instance.quit()

        if self.window:
            self.window.close()

        if self.app_thread and self.app_thread.is_alive():
            self.app_thread.join()

        self._is_open = False
        self.running = False

    def set_center(self, x, y, z=0):
        self.X, self.Y, self.Z = x, y, z

    def get_scaled_obj_coord(self, obj):
        # returns the coordinates of a reference object
        return obj.vx(), obj.vy(), obj.z()

    def get_scaled_coord(self, x, y, z=None):
        # returns the coordinates of a point (x,y) or (x,y,z)
        s = self.window.scale
        cx, cy, cz = self.X, self.Y, self.Z
        scaled_x = (x - cx) * s + cx + self.dX
        scaled_y = (y - cy) * s + cy + self.dY
        if z is None:
            return scaled_x, scaled_y
        scaled_z = (z - cz) * s + cz
        return scaled_x, scaled_y, scaled_z

    def set_background_pixmap(self, pixmap):
        self.bg_pixmap = pixmap

    def set_pixmap(self, pixmap):
        self.window.pixmap = pixmap
        self.window.frame.setPixmap(pixmap)
        self.window.update()

    def clear_buffer(self):
        self.object_buffer.clear()

    def set_engine(self, engine="Qt"):
        if engine not in ("Qt", "Pdf"):
            raise ValueError(f"Unsupported engine: {engine}")
        self.engine = engine
        self.qt_stylesheet = ""

    def set_geometry(self, width, height):
        self.geometry = (width, height)

    def isOpen(self):
        return self._is_open

    def render(self, verbose=False, pixmap=None, painter=None, end_painter=True, pos_only=False, callback_bg=None, callback=None):
        """update the visualization.

        Args:
            verbose (bool, optional): print output for each visualized object. Defaults to False.
            pos_only (bool): update positions only, without painting
            callback_bg (callable): function to call in background. must take argument 'painter'
            callback (callable): function to call in foreground. must take argument 'painter'
        Returns:
            painter, pixmap 
        """

        if not self._is_open:
            return
        if self.engine is None:
            raise RuntimeError(
                "Cannot visualize World without engine. Please use World().visuals.set_engine(...) first.")
        if not self.running:
            if verbose:
                warnings.warn("WorldVisuals not running, returning.")
            return None, None

        # retrieve QApplication instance
        instance = self.instance
        while not instance:
            instance = QApplication.instance()
            if verbose:
                print("wait for instance", instance)
            try:
                time.sleep(.001)
            except:
                pass
        if instance is None:
            if verbose:
                print("WARNING: no instance for plotting")
            return None, None
        else:
            if verbose:
                print("instance found.", instance)

        # -- PREPARE engine COMPATIBILITY --
        # run Qt engine
        if self.engine == "Qt":
            if verbose:
                print("engine Qt")
            if pixmap is None:
                pixmap = QPixmap(*self.geometry)
                pixmap.fill(QColor(255, 255, 255, 255))
                if verbose:
                    print("create new pixmap")
            if painter is None:
                painter = QPainter(pixmap)
                painter.setRenderHint(QPainter.Antialiasing)
            if verbose:
                print("___ (START)")
            if self.bg_pixmap is not None:
                painter.drawPixmap(self.dX, self.dY, self.bg_pixmap)

        # run Pdf engine
        elif self.engine == "Pdf":
            if verbose:
                print("engine: Pdf")
            file_path = self.pdf_filepath
            pdf_writer = QPdfWriter(file_path)
            pdf_writer.setResolution(self.window.dpi)
            pdf_writer.setPageSizeMM(QSizeF(210, int(0.5*297)))
            if painter is None:
                painter = QPainter(pdf_writer)
                painter.setRenderHint(QPainter.Antialiasing)

        # -- TRANSFORM COORDINATES  --
        v, v2 = 0, 0
        if self.window and self.window.count_rot_xy is not None:
            v = self.window.count_rot_xy.value
            v2 = self.window.count_rot_z.value
        if verbose:
            print("Rotation Values:", v, v2)

        # -- RUN CALLBACK PROCEDURE (IF ANY) --
        if callback_bg is not None:
            callback_bg(painter)

        # -- VISUALIZE OBEJCTS IN BUFFER --
        # print("number of objects in buffer", len(self.object_buffer))
        for obj in self.object_buffer:
            if verbose:
                print("Rendering object:", obj)
            if not obj.block_rotation:
                R_xy = np.array([[1, 0, 0],
                                [0, np.cos(v), -np.sin(v)],
                                [0, np.sin(v), np.cos(v)]])
                R_z = np.array([[np.cos(v2), -np.sin(v2), 0],
                                [np.sin(v2), np.cos(v2), 0],
                                [0, 0, 1]])
                R = np.matmul(R_xy, R_z)
            else:
                R = None
            self.visualize_agent(obj, painter, verbose=verbose, R=R,
                                 center=[self.X, self.Y, self.Z], pos_only=pos_only)
            obj.update()

        # -- RUN CALLBACK PROCEDURE (IF ANY) --
        if callback is not None:
            callback(painter)

        # -- END PAINTER --
        if self.engine == "Qt" and end_painter:
            painter.end()
            self.set_pixmap(pixmap)

        if self.window:
            self.window.update()

        if self.engine == "Pdf" and end_painter:
            painter.end()
            self.engine = "Qt"
        if verbose:
            print("___ (END)")

        return painter, pixmap

    def rotate(self, rxy, rz):
        self.window.count_rot_xy.value = rxy
        self.window.count_rot_z.value = rz

    def visualize_agent(self, obj, painter, verbose=False, R=None, center=None, pos_only=False):
        """
        Visualizes an sfctools Agent
        Args:
            obj (Agent): an agent to visualize
            paitner (QPainter): the Qt painter to use
            R (np.ndarray): 3x3 transformation matrix
            center (np.ndarray or tuple): center of rotation 
            pos_only: (boolean, default False) only compute the object positions without rendering in Qt (experimental)
        """

        if not (self.engine == "Qt" or self.engine == "Pdf"):
            return

        # setup painter
        painter.setRenderHint(QPainter.Antialiasing)
        if isinstance(obj.fill_color, tuple):
            painter.setBrush(QColor(*obj.fill_color))
        else:
            try:
                painter.setBrush(QColor(obj.fill_color))
            except Exception as e:
                print("Could not paint %s: %s" % (obj, str(e)))

        # setup pen
        pen = None
        if obj.edge_width > 0:
            pen = QPen(QColor(*obj.edge_color))
            pen.setWidth(obj.edge_width)
            painter.setPen(pen)

        # transform coordinates
        w, h = self.geometry
        w_j = int(obj.size[0]*0.5)
        h_j = int(obj.size[1]*0.5)
        u, v, z = int(obj.x()-w_j), int(obj.y()-h_j), int(obj.z())
        u = int(min(max(0, u), w))
        v = int(min(max(0, v), h))

        # rotation transform
        if center is None:
            center = [self.X, self.Y, self.Z]
        if R is not None:  # transform by rotation?
            u0, v0, z0 = center[0], center[1], center[2]
            q = np.dot(R, np.array([u-u0, v-v0, z-z0]))
            # extract transformed coordinates
            u, v, z = q[0]+u0, q[1]+v0, q[2]+z0
            # NOTE z coordinate is neglected, is just projected onto a plane

        scale = self.window.scale if not obj.block_scaling else 1.0
        u = (u-center[0])*scale + center[0] + self.dX
        v = (v-center[1])*scale + center[1] + self.dY
        if np.isnan(u) or np.isnan(v):
            return  # no valid position? skip...

        u, v = int(u), int(v)  # convert coordinates to integer
        obj.visual_pos = (u, v)

        if not pos_only:
            if obj.shape == "rect":  # rectangle
                if obj.fill_color is None:
                    painter.drawRect(u, v, w_j, h_j)
                else:
                    if obj.edge_width > 0:
                        painter.drawRect(u-1, v-1, w_j+2, h_j+2)
                    painter.fillRect(u, v, w_j, h_j, QColor(
                        *obj.fill_color))
            elif obj.shape == "ell":  # elliptic
                if obj.fill_color is None:
                    painter.drawEllipse(u, v, w_j, h_j)
                else:
                    painter.drawEllipse(u-1, v-1, w_j+2, h_j+2)

            # draw center of coordinate system
            redpen = QPen(QColor(*(255, 0, 0)))
            redpen.setWidth(4)
            painter.setPen(redpen)
            painter.setBrush(QColor(0, 0, 0, 0))
            u0, v0, z0 = center[0]*scale, center[1]*scale, center[2]*scale

            # draw selection radius
            if obj.hover:
                W = 15  # width of selection circle in pixels
                painter.drawEllipse(u - W, v - W, w_j+2*W, h_j+2*W)

            # reset pen
            if pen is not None:
                painter.setPen(pen)
            else:
                pen = QPen(QColor(0, 0, 0, 255))
                pen.setWidth(1)
                painter.setPen(pen)

    def open(self, app=None, add_navigation=False):
        """
        Open the visuals

        Args:
            app (optional, QApplication): give an app instance if required
            add_navigation (optional, boolean): add a navigation bar inside the visualization window
        """
        # print("OPENING")

        if self._is_open:
            return

        self._is_open = True
        app_thread = threading.Thread(
            target=lambda: self.run_app(app, add_nav=add_navigation))
        self.app_thread = app_thread
        app_thread.start()

    def run_app(self, app=None, add_nav=False):

        self.app = app
        self._is_open = True

        if self.engine == "Qt":

            # NOTE activate to suppress warning
            QLoggingCategory.setFilterRules(
                '*.debug=false\n*.warning=false\n*.critical=true\n')

            # print("app", app)
            self.running = True

            window = QtVizWindow(
                self, *self.geometry, add_navigation=add_nav)
            self.window = window
            self.window.setStyleSheet(self.qt_stylesheet)

            screen = app.primaryScreen()  # Get the primary screen
            dpi = screen.logicalDotsPerInch()  # Get DPI of the screen
            self.window.dpi = dpi

            QLoggingCategory.setFilterRules('')

            window.show()
            # app.exec_()
            print("executing app...")

        else:
            raise RuntimeError("Unknown engine.")
