import math
import xml.etree.ElementTree as ET

import numpy as np
import pyqtgraph as pg
from pyqtgraph import debug
from pyqtgraph import functions as fn
from pyqtgraph.graphicsItems.ScatterPlotItem import ScatterPlotItem, SymbolAtlas
from qgis.PyQt import QtCore, QtGui
from qgis.PyQt.QtGui import QColor
from qgis.PyQt.QtSvg import QSvgRenderer

from openlog.__about__ import DIR_PLUGIN_ROOT
from openlog.datamodel.assay.generic_assay import AssaySeriesType, GenericAssay
from openlog.gui.assay_visualization.plot_item_utils import _convert_x_values
from openlog.gui.pyqtgraph.ItemSwitcher import DiscreteItemSwitcher

LINE_PATHS = {
    "s": str(DIR_PLUGIN_ROOT / "resources" / "structural_symbology" / "arrow.svg"),
    "h": str(
        DIR_PLUGIN_ROOT / "resources" / "structural_symbology" / "double_arrow.svg"
    ),
    "v": str(
        DIR_PLUGIN_ROOT / "resources" / "structural_symbology" / "vertical_line.svg"
    ),
}

PLANE_PATHS = {
    "s": str(
        DIR_PLUGIN_ROOT / "resources" / "structural_symbology" / "inclined_plane.svg"
    ),
    "h": str(
        DIR_PLUGIN_ROOT / "resources" / "structural_symbology" / "horizontal_plane.svg"
    ),
    "v": str(
        DIR_PLUGIN_ROOT / "resources" / "structural_symbology" / "vertical_plane.svg"
    ),
}


def rotate_svg(
    file: str,
    angle: float,
    label: str = None,
    label_y_pos: float = None,
    color: str = None,
) -> str:
    """
    Rotate an SVG symbol given azimuth angle and add a text (dip).
    Args:
        - file : path to the original SVG file
        - angle : rotation angle in degree (0-360) - clockwise
        - label : text to insert
        - label_y_pos : initial Y coordinates of label

    Return XML definition of transformed symbol.

    """

    namespaces = {
        "dc": "http://purl.org/dc/elements/1.1/",
        "cc": "http://creativecommons.org/ns#",
        "rdf": "http://www.w3.org/1999/02/22-rdf-syntax-ns#",
        "svg": "http://www.w3.org/2000/svg",
        "myns": "http://www.w3.org/2000/svg",
    }

    tree = ET.parse(file)
    root = tree.getroot()
    # viewbox dimensions
    dimensions_elt = root.find(
        'svg:g[@id="Square_x0020_frame"]/svg:rect', namespaces=namespaces
    )
    width, height = float(dimensions_elt.get("width")), float(
        dimensions_elt.get("height")
    )

    # add transformation
    symbol_elt = root.find('svg:g[@id="symbol"]', namespaces=namespaces)
    c_x, c_y = width / 2, height / 2
    rotation = f"rotate({angle} {c_x} {c_y})"
    symbol_elt.set("transform", rotation)

    if label is not None:
        # compute label coordinates
        radius = c_y - label_y_pos
        x_label = c_x + radius * math.cos(math.radians(90 - angle))
        y_label = (
            c_y - radius * math.sin(math.radians(90 - angle)) + int(width / 4) / 10
        )
        anchor = "middle"
        # add label
        text_element = ET.Element("text")
        text_element.set("x", str(x_label))
        text_element.set("y", str(y_label))
        text_element.set("font-family", "Arial")
        text_element.set("stroke-width", "2")
        text_element.set("font-size", str(int(width / 10)))
        text_element.set("stroke", "black")
        text_element.set("text-anchor", anchor)
        text_element.text = str(label)
        root.append(text_element)

    xml_str = ET.tostring(root, encoding="utf8")
    if color:
        xml_str = xml_str.replace(b"black", color.encode("utf8"))

    return xml_str


class SvgSymbolAtlas(SymbolAtlas):
    def __init__(self):
        super().__init__()
        # symbol dictionnary key : parameters
        self.parameters = {
            "o": {
                "file": LINE_PATHS["s"],
                "angle": 0,
                "label": None,
                "label_y_pos": None,
            }
        }

    def drawSvgSymbol(
        self, painter: QtGui.QPainter, symbol: str, size: int, pen, brush
    ):
        if self.parameters.get(symbol) is None:
            params = {
                "file": LINE_PATHS["s"],
                "angle": 45,
                "label": 258,
                "label_y_pos": 300,
            }
        else:
            params = self.parameters[symbol]
        painter.setPen(pen)
        painter.setBrush(brush)
        xml_str = rotate_svg(**params)
        renderer = QSvgRenderer(xml_str)
        painter.setRenderHint(painter.RenderHint.Antialiasing)
        renderer.render(painter)
        painter.end()

    def renderSvgSymbol(self, symbol, size, pen, brush, device=None):

        """
        Render a symbol specification to QImage.
        Symbol may be either a QPainterPath or one of the keys in the Symbols dict.
        If *device* is None, a new QPixmap will be returned. Otherwise,
        the symbol will be rendered into the device specified (See QPainter documentation
        for more information).
        """
        ## Render a spot with the given parameters to a pixmap

        if device is None:
            device = QtGui.QImage(
                int(size), int(size), QtGui.QImage.Format.Format_ARGB32_Premultiplied
            )
            device.fill(QtCore.Qt.GlobalColor.transparent)
        p = QtGui.QPainter(device)
        try:
            self.drawSvgSymbol(p, symbol, size, pen, brush)
        finally:
            p.end()
        return device

    def _extend(self, styles):

        profiler = debug.Profiler()

        images = []
        data = []
        for key, style in styles.items():
            img = self.renderSvgSymbol(*style)
            arr = fn.ndarray_from_qimage(img)
            images.append(img)  # keep these to delay garbage collection
            data.append((key, arr))

        profiler("render")
        self._extendFromData(data)
        profiler("insert")


class StructuralSymbolItem(ScatterPlotItem):
    def __init__(
        self,
        assay: GenericAssay,
        column: str,
        coordinates: np.ndarray = None,
        planned_coordinates: np.ndarray = None,
        planned_eoh: float = None,
        *args,
        **kargs,
    ):

        self.assay = assay
        self.column = column
        super().__init__(*args, **kargs)
        self.fragmentAtlas = SvgSymbolAtlas()

        self.symbols = self.create_symbol_parameters()

        self.switcher = DiscreteItemSwitcher(
            self,
            coordinates=coordinates,
            planned_coordinates=planned_coordinates,
            planned_eoh=planned_eoh,
        )

        self.setSymbol(self.symbols)

    def update_symbols_color(self, color: list[str]):

        self.symbols = self.create_symbol_parameters(color)
        self.setSymbol(self.symbols)

    def create_symbol_parameters(self, color: list[str] = None) -> list[str]:
        """
        Create each symbol with parameters (svg file, angle, label) and associated IDs.
        Store parameters in fragmentAtlas attribute.
        Returns IDs
        """
        symbols = []
        _, y = self.assay.get_all_values(self.column)

        if color is None:
            color_list = [None for value in y]
        else:
            color_list = color

        if (
            self.assay.assay_definition.columns[self.column].series_type
            == AssaySeriesType.POLAR
        ):

            for value in y:
                d = {
                    "file": LINE_PATHS["s"],
                    "angle": float(value),
                    "label": None,
                    "label_y_pos": None,
                    "color": None,
                }
                symbols.append(d)

        else:

            for value in y:

                type_, dip, azimuth = (
                    value.get("type"),
                    value.get("dip"),
                    value.get("azimuth"),
                )
                dict_path = LINE_PATHS if type_ == "LINE" else PLANE_PATHS
                label_y_pos = 1500 if type_ == "LINE" else 3500
                dict_key = "s"
                dip_label = str(int(dip))
                angle = float(azimuth)

                if type_ == "PLANE":
                    angle = angle + 90

                if dip == 0:
                    dict_key = "h"
                    dip_label = None
                    angle = 0

                if dip == 90:
                    dict_key = "v"
                    dip_label = None
                    if type_ == "LINE":
                        angle = 0

                d = {
                    "file": dict_path[dict_key],
                    "angle": angle,
                    "label": dip_label,
                    "label_y_pos": label_y_pos,
                }

                symbols.append(d)

        symbol_ids = [
            f"{d['file']}_{d['angle']}_{d['label']}_{d['label_y_pos']}" for d in symbols
        ]

        # set color
        for i, col in enumerate(color_list):
            symbol_ids[i] = symbol_ids[i] + f"_{col}"
            symbols[i]["color"] = col

        self.fragmentAtlas.parameters.update(
            {
                symbol_id: symbol_def
                for symbol_id, symbol_def in zip(symbol_ids, symbols)
            }
        )

        return symbol_ids

    def updateSpots(self, dataSet=None):
        if hasattr(self, "symbols"):
            dataSet = self.data
            for i in range(len(dataSet)):
                dataSet[i][3] = self.symbols[i]

        return super().updateSpots(dataSet)


def create_discrete_structural_plot_items(
    plot: pg.PlotWidget, assay: GenericAssay, column: str, config
) -> StructuralSymbolItem:

    (x, _) = assay.get_all_values(column)
    x_val = _convert_x_values(assay, x)

    item = StructuralSymbolItem(
        assay,
        column,
        x=np.full(x_val.shape, 0.0),
        y=x_val,
        coordinates=assay.get_coordinates(),
        planned_coordinates=assay.get_coordinates(planned=True),
        planned_eoh=assay.get_planned_eoh(),
    )

    plot.addItem(item)
    # set transparent color to axis labels to keep aligned plots
    color = QColor(255, 255, 255, 0)
    for side in ["top", "bottom"]:
        axe = plot.plotItem.getAxis(side)
        pen = axe.textPen()
        pen.setColor(color)
        axe.setTextPen(pen)

    config.set_plot_item(item)


def create_discrete_plot_items_from_extended_structural(
    plot: pg.PlotWidget,
    assay: GenericAssay,
    column: str,
    config,
) -> StructuralSymbolItem:

    (x, _) = assay.get_all_values(column)
    x_val = _convert_x_values(assay, x[:, 0])
    coords = assay.get_coordinates()
    if coords is not None:
        coords = coords[:, :3]
    planned_coords = assay.get_coordinates(planned=True)
    if planned_coords is not None:
        planned_coords = planned_coords[:, :3]

    item = StructuralSymbolItem(
        assay,
        column,
        x=np.full(x_val.shape, 0.0),
        y=x_val,
        coordinates=coords,
        planned_coordinates=planned_coords,
        planned_eoh=assay.get_planned_eoh(),
    )

    plot.addItem(item)

    # set transparent color to axis labels to keep aligned plots
    color = QColor(255, 255, 255, 0)
    for side in ["top", "bottom"]:
        axe = plot.plotItem.getAxis(side)
        pen = axe.textPen()
        pen.setColor(color)
        axe.setTextPen(pen)

    config.set_plot_item(item)
