# Copyright (c) 2026 Luigi Paone <ppc.luigi.paone@gmail.com>
# SPDX-License-Identifier: BSD-3-Clause

import copy
import math
from typing import List, Tuple, Union, TypedDict, Unpack

import numpy as np

import matplotlib
matplotlib.use('Agg')

import matplotlib.pyplot as plt
from matplotlib.axes import Axes
from matplotlib.patches import PathPatch
from matplotlib.path import Path
from matplotlib.ticker import FormatStrFormatter

from pycivil.EXAGeometry.clouds import PointCloud2d
from pycivil.EXAGeometry.geometry import Point2d
from adjustText import adjust_text # type: ignore[import-untyped]


class Geometry2DPlot:
    def __init__(
        self,
        nrows: int = 1,
        ncols: int = 1,
        figsize: Tuple[float, float] = (7, 7),
        sx: float = 1,
        sy: float = 1,
    ):
        self.__fig = plt.figure(figsize=figsize)

        # Adjust layout
        # self.__fig.subplots_adjust(hspace=0.9, wspace=0.9)

        self.__ax: List[Axes] = [
            self.__fig.add_subplot(nrows, ncols, i + 1) for i in range(nrows * ncols)
        ]
        self.__titles: List[str] = [""] * len(self.__ax)
        self.__xlabels: List[str] = [""] * len(self.__ax)
        self.__ylabels: List[str] = [""] * len(self.__ax)

        # Data to plot
        self.__points: List[Point2d] = []
        self.__pointArray: List[List[Point2d]] = []
        self.__pointClouds: List[PointCloud2d] = []

        # Scale Factor
        self.__sx = sx
        self.__sy = sy

        # Aspect controls
        self.__pointSize: float = 30
        self.__textSize: float = 12
        self.__showPointIndex: bool = True
        self.__showPointIndexArray: bool = True
        self.__showLineArray: bool = False
        self.__showArrow: bool = True
        self.__showLine: bool = True
        self.__colorLines: str = "black"
        self.__colorPoints: str = "red"
        self.__colorPointsArray: List[List[str]] = []

    @property
    def pointSize(self):
        return self.__pointSize

    @pointSize.setter
    def pointSize(self, value: float) -> None:
        self.__pointSize = value

    @property
    def textSize(self):
        return self.__textSize

    @textSize.setter
    def textSize(self, value: float) -> None:
        self.__textSize = value

    def setXLabel(self, xlabel: List[str]) -> None:
        self.__xlabels = xlabel

    def setYLabel(self, ylabel: List[str]) -> None:
        self.__ylabels = ylabel

    def setTitles(self, titles: List[str]) -> None:
        self.__titles = titles

    def showItems(
        self,
        pointIndex: bool = True,
        arrow: bool = True,
        line: bool = True,
        lineArray: bool = False,
        pointIndexArray: bool = True,
    ) -> None:
        self.__showPointIndex = pointIndex
        self.__showArrow = arrow
        self.__showLine = line
        self.__showLineArray = lineArray
        self.__showPointIndexArray = pointIndexArray

    def addPoint(self, point: Point2d) -> None:
        self.__points.append(point)

    def addPointArray(self, array: List[Point2d], colorArray: List[str] | None = None) -> None:
        if colorArray is None:
            colorArray = []
        self.__pointArray.append(array)

        if len(colorArray) == 0:
            self.__colorPointsArray.append(["black"] * len(array))
            return

        if len(colorArray) == 1:
            self.__colorPointsArray.append([colorArray[0] * len(array)])
            return

        if len(colorArray) not in [0, 1, len(array)]:
            raise ValueError(
                f"Len of array {len(array)} must be equal to \
                            len of colorArray {len(colorArray)}"
            )
        else:
            self.__colorPointsArray.append(colorArray)
            return

    def addPointCloud(self, cloud: PointCloud2d) -> None:
        self.__pointClouds.append(cloud)

    def __plotFromPointArray(self, points: List[List[Point2d]]) -> None:
        if len(self.__ax) != len(points):
            print("ERR: Can't plot many set of points with one axes")

        else:
            for idx, pts in enumerate(points):
                # save options
                showLineOld = self.__showLine
                self.__showLine = self.__showLineArray

                showPointIndexOld = self.__showPointIndex
                self.__showPointIndex = self.__showPointIndexArray

                self.__plotFromPoint2dLst(
                    pts, self.__ax[idx], self.__colorPointsArray[idx]
                )

                # restore options
                self.__showLine = showLineOld
                self.__showPointIndex = showPointIndexOld

    def __plotFromPoint2dLst(
        self, points: List[Point2d], axes: Axes, color: Union[str, List[str]]
    ) -> None:
        for i, p in enumerate(points):
            if isinstance(color, str):
                sc = axes.scatter(
                    [p.x * self.__sx],
                    [p.y * self.__sy],
                    color=color,
                    s=self.__pointSize,
                )
            else:
                sc = axes.scatter(
                    [p.x * self.__sx],
                    [p.y * self.__sy],
                    color=color[i],
                    s=self.__pointSize,
                )

            if self.__showPointIndex:
                sizeOfAnnotation = math.sqrt(sc.get_sizes()[0]) / 2 + 2
                axes.annotate(
                    text=f"{i}",
                    xy=(p.x * self.__sx, p.y * self.__sy),
                    xytext=(sizeOfAnnotation, 0),
                    fontsize=self.__textSize,
                    textcoords="offset points",
                    ha="left",
                    va="center",
                )

            if i < len(points) - 1:
                p0 = points[i]
                p1 = points[i + 1]

                if self.__showLine:
                    axes.plot(
                        [p0.x * self.__sx, p1.x * self.__sx],
                        [p0.y * self.__sy, p1.y * self.__sy],
                        color=self.__colorLines,
                    )

                if self.__showArrow:
                    centerOfArrow = [
                        (p0.x * self.__sx + p1.x * self.__sx) / 2,
                        (p0.y * self.__sy + p1.y * self.__sy) / 2,
                    ]
                    direction = [
                        (p1.x * self.__sx - p0.x * self.__sx),
                        (p1.y * self.__sx - p0.y),
                    ]
                    lenght = math.sqrt(
                        math.pow(direction[0], 2) + math.pow(direction[1], 2)
                    )
                    if not lenght == 0:
                        versor = [
                            direction[0] / lenght / 100,
                            direction[1] / lenght / 100,
                        ]
                        axes.annotate(
                            "",
                            xy=(centerOfArrow[0], centerOfArrow[1]),
                            xytext=(
                                centerOfArrow[0] + versor[0],
                                centerOfArrow[1] + versor[1],
                            ),
                            fontsize=self.__textSize * 2,
                            arrowprops={
                                "arrowstyle": "<-",
                            },
                        )
                    else:
                        print(
                            f"Point at coords {p0} duplicate with index {i} and {i + 1}"
                        )

    def __plotFromPointCloudLst(self, points: List[PointCloud2d]) -> None:

        if len(self.__titles) == len(self.__ax):
            for i, t in enumerate(self.__titles):
                self.__ax[i].set_title(t)
        else:
            if len(self.__titles) == 1:
                for a in self.__ax:
                    a.set_title(self.__titles[0])
            else:
                print("ERR: Can't set titles because lenght !!!")

        if len(self.__xlabels) == len(self.__ax):
            for i, t in enumerate(self.__xlabels):
                self.__ax[i].set_xlabel(t)
        else:
            if len(self.__xlabels) == 1:
                for a in self.__ax:
                    a.set_xlabel(self.__xlabels[0])
            else:
                print("ERR: Can't set xlabels because lenght !!!")

        if len(self.__ylabels) == len(self.__ax):
            for i, t in enumerate(self.__ylabels):
                self.__ax[i].set_ylabel(t)
        else:
            if len(self.__ylabels) == 1:
                for a in self.__ax:
                    a.set_ylabel(self.__ylabels[0])
            else:
                print("ERR: Can't set ylabels because lenght !!!")

        if len(self.__ax) == 1:
            for p in points:
                self.__plotFromPoint2dLst(
                    p.getPoints(), self.__ax[0], self.__colorPoints
                )

        else:
            if len(self.__ax) != len(points):
                print("ERR: Can't plot many set of points with one axes")

            else:
                for idx, p in enumerate(points):
                    self.__plotFromPoint2dLst(
                        p.getPoints(), self.__ax[idx], self.__colorPoints
                    )

    def plot(self):

        # Points
        self.__plotFromPoint2dLst(self.__points, self.__ax[0], self.__colorPoints)

        # Point cloud
        self.__plotFromPointCloudLst(self.__pointClouds)

        # Point array
        if len(self.__pointArray) > 0:
            self.__plotFromPointArray(self.__pointArray)

        self.__fig.tight_layout(pad=2)

    @staticmethod
    def show():
        plt.show()

    @staticmethod
    def save(
        fname,
        *,
        transparent=None,
        dpi="figure",
        format=None,
        metadata=None,
        bbox_inches=None,
        pad_inches=0.1,
        facecolor="auto",
        edgecolor="auto",
        backend=None,
        **kwargs,
    ):

        plt.savefig(
            fname,
            dpi=dpi,
            format=format,
            metadata=metadata,
            bbox_inches=bbox_inches,
            pad_inches=pad_inches,
            facecolor=facecolor,
            edgecolor=edgecolor,
            backend=backend,
            transparent=transparent,
            **kwargs,
        )

class KvargsPlot(TypedDict):
    export: str | None
    dpi: int | None
    xLabel: str
    yLabel: str
    titleAddStr: str
    tensionPoints: List[List[float]]
    scale_Nx: float
    scale_Mz: float
    lines: List[List[Point2d]]
    markers: bool
    savingSingleDomains: bool
    hotPoints: bool | List[bool]
    printDomains: List[int] | None

def interactionDomainBasePlot2d(
        vertices: List[List[float]],
        fields: List[float],
        **kwargs: Unpack[KvargsPlot]
) -> None:
    if "export" in kwargs:
        fileName = kwargs["export"]
    else:
        fileName = None

    if "dpi" in kwargs:
        dpi = kwargs["dpi"]
    else:
        dpi = None

    if "xLabel" in kwargs:
        xLabel = kwargs["xLabel"]
    else:
        xLabel = "Nz [N]"

    if "yLabel" in kwargs:
        yLabel = kwargs["yLabel"]
    else:
        yLabel = "Mz [Nmm]"

    if "titleAddStr" in kwargs:
        titleAddStr = kwargs["titleAddStr"]
    else:
        titleAddStr = None

    if "tensionPoints" in kwargs:
        tensionPoints = kwargs["tensionPoints"]
    else:
        tensionPoints = None

    if "scale_Nx" in kwargs:
        scale_Nx = kwargs["scale_Nx"]
    else:
        scale_Nx = 0.001

    if "scale_Mz" in kwargs:
        scale_Mz = kwargs["scale_Mz"]
    else:
        scale_Mz = 0.000001

    if "lines" in kwargs:
        lines = kwargs["lines"]
    else:
        lines = []

    if "markers" in kwargs:
        markers = kwargs["markers"]
    else:
        markers = True

    if "savingSingleDomains" in kwargs:
        savingSingleDomains = kwargs["savingSingleDomains"]
    else:
        savingSingleDomains = False

    if "hotPoints" in kwargs:
        hotPoints = kwargs["hotPoints"]
    else:
        hotPoints = False

    if "printDomains" in kwargs:
        printDomains = kwargs["printDomains"]
        if printDomains is None:
            printDomains = list(range(0, len(vertices)))
    else:
        printDomains = list(range(0, len(vertices)))

    """
    if "pointsScale" in kwargs:
        pointsScale = kwargs["pointsScale"]
    else:
        pointsScale = 40
    """
    if isinstance(vertices, list) and len(vertices) != 0:
        # Only one diagrammpython
        if isinstance(vertices[0][0], float) and isinstance(vertices[0][1], float):
            verticesList = [vertices]
            fieldsList = [fields]
        else:
            verticesList = vertices
            fieldsList = fields
    else:
        raise ValueError(
            "Vertices must be a list with lenght > 0 !!! Maybe you need build domain."
        )

    fig, ax = plt.subplots()

    if tensionPoints is not None:
        if len(tensionPoints) > 0:
            tp = np.array(tensionPoints, float)
            ax.scatter(
                tp[:, 0] * scale_Nx,
                tp[:, 1] * scale_Mz,
                marker="+",
                s=2000,
                color=[0, 0, 0],
            )

    # plt.autoscale(enable = None)
    plt.xlabel(xLabel)
    plt.ylabel(yLabel)
    plt.xticks(rotation="vertical")
    plt.margins(0.0)
    plt.subplots_adjust(bottom=0.25)
    plt.subplots_adjust(left=0.20)

    if len(lines) != 0:
        for ll in lines:
            x1 = ll[0].x * scale_Nx
            y1 = ll[0].y * scale_Mz
            x2 = ll[1].x * scale_Nx
            y2 = ll[1].y * scale_Mz
            plt.plot([x1, x2], [y1, y2], "k-", lw=2)

    strTitle = "Interaction Domain Nx - Mz"
    if titleAddStr is not None:
        plt.subplots_adjust(top=0.87)
        strTitle = strTitle + "\n" + titleAddStr

    ax.set_title(strTitle)
    ax.grid(True)
    ax.yaxis.set_major_formatter(FormatStrFormatter("%.2e"))
    ax.xaxis.set_major_formatter(FormatStrFormatter("%.2e"))

    # Cloning list only for plot only with deep
    verticesListCopy = copy.deepcopy(verticesList)

    # Removing columns from 2 to n that contains others values
    for _idx, vertices_i in enumerate(verticesListCopy):
        for idy, v in enumerate(vertices_i):
            vertices_i[idy] = [v[0], v[1]]

    for idx, vertices_i in enumerate(verticesListCopy):

        if idx not in printDomains:
            continue

        codes = [Path.MOVETO] + [Path.LINETO] * (len(vertices_i) - 1)

        vertices_i_for_plot = np.array(vertices_i, float)

        # setting scale
        vertices_i_for_plot[:, 0] = vertices_i_for_plot[:, 0] * scale_Nx
        vertices_i_for_plot[:, 1] = vertices_i_for_plot[:, 1] * scale_Mz

        # setting color for hot and cold section
        _edgecolor = "black"
        if isinstance(hotPoints, bool):
            if not hotPoints:
                _edgecolor = "black"
        if isinstance(hotPoints, list):
            if not hotPoints[idx]:
                _edgecolor = "black"
            else:
                _edgecolor = "red"

        path = Path(vertices_i_for_plot, codes)
        pathpatch = PathPatch(path, facecolor="None", edgecolor=_edgecolor)

        ax.add_patch(pathpatch)

        if markers:
            vertex = path.vertices
            assert isinstance(vertex,np.ndarray)
            x = vertex[:, 0]
            y = vertex[:, 1]
            colorsMap = np.array(
                [
                    [1, 0, 0, 0.5],
                    [0, 1, 0, 0.5],
                    [0, 0, 1, 0.5],
                    [1, 0, 1, 0.5],
                    [1, 0, 0, 1.0],
                    [0, 1, 0, 1.0],
                    [0, 0, 1, 1.0],
                    [1, 0, 1, 1.0],
                ]
            )

            # Displayng Domain Different colors
            scatt_colors = []
            for f in fieldsList[idx]:
                if f == 1.0:
                    scatt_colors.append(0)
                elif f == 2.0:
                    scatt_colors.append(1)
                elif f == 3.0:
                    scatt_colors.append(2)
                elif f == 4.0:
                    scatt_colors.append(3)
                elif f == 11.0:
                    scatt_colors.append(4)
                elif f == 12.0:
                    scatt_colors.append(5)
                elif f == 13.0:
                    scatt_colors.append(6)
                elif f == 14.0:
                    scatt_colors.append(7)
                else:
                    raise ValueError("The value of f is <%1.4f> not mapped!!!" % f)

            ax.scatter(x, y, marker="o", s=40, c=colorsMap[scatt_colors])

        if savingSingleDomains:
            ax.autoscale_view()
            if titleAddStr is not None:
                plt.savefig(titleAddStr + "_domain_" + str("%1.0i" % idx) + ".png")

    ax.autoscale_view()

    if fileName is not None:
        if dpi is not None:
            print("save figure to file name -->", fileName)
            plt.style.use("classic")
            plt.savefig(fileName, bbox_inches="tight", dpi=600)
        else:
            raise ValueError("If fileName is not None must provide dpi argument")
    else:
        plt.show()
