import math
from enum import Enum
from typing import Union, Dict, Literal, Tuple, Any

from adjustText import adjust_text # type: ignore[import-untyped]
from matplotlib import pyplot as plt
from matplotlib.patches import Polygon, Circle

from pycivil.EXAStructural.modeler import SectionModeler
from pycivil.EXAStructural.sections import ConcreteSection
from pycivil.EXAUtils.logging import log


class SectionPlotViewEnum(str, Enum):
    SP_VIEW_SECTION = "SP_VIEW_SECTION"
    SP_VIEW_SECTION_NOID = "SP_VIEW_SECTION_NOID"
    SP_VIEW_NODES = "SP_VIEW_NODES"
    SP_VIEW_SOLID_NODES = "SP_VIEW_SOLID_NODES"
    SP_HIDE_SOLID_NODES = "SP_HIDE_SOLID_NODES"
    SP_HIDE_NODES = "SP_HIDE_NODES"
    SP_HIDE_REBAR_ID = "SP_HIDE_REBAR_ID"
    SP_VIEW_GRID = "SP_VIEW_GRID"
    SP_HIDE_GRID = "SP_HIDE_GRID"
    SP_VIEW_TRIANGLES = "SP_VIEW_TRIANGLES"
    SP_VIEW_TRIANGLE_NOID = "SP_VIEW_TRIANGLE_NOID"
    SP_VIEW_MESH = "SP_VIEW_MESH"
    SP_VIEW_STRESSES = "SP_VIEW_STRESSES"


class SectionPlot:
    def __init__(self):
        self.__fig = plt.figure()
        self.__ax = self.__fig.add_subplot(1, 1, 1)
        self.__ax.set_aspect("equal", "datalim")

        # Data to plot
        self.__rcRect: Union[None, ConcreteSection] = None
        self.__modeler: Union[None, SectionModeler] = None
        self.__vertexStresses: Union[Dict[int, float], None] = None
        self.__rebarStresses: Union[Dict[int, float], None] = None

        # Aspect controls
        self.__colorEdge: str = "black"
        self.__textSize: float = 8
        self.__colorTriaLines: str = "blue"
        self.__colorTriaFill: str = "green"
        self.__colorNodes: str = "red"
        self.__plotShapeEdges: bool = True
        self.__plotTriaLines: bool = False
        self.__plotTriaFill: bool = True
        self.__plotTriaShrinkFill: bool = False
        self.__plotTriaShrink: bool = False
        self.__plotNodes: bool = False
        self.__plotNodesId: bool = False
        self.__plotSolidNodes: bool = True
        self.__plotSolidNodesId: bool = True
        self.__plotStresses: bool = False
        self.__plotRebar: bool = True
        self.__plotMeshTria: bool = False
        self.__plotRebarId: bool = True
        self.__plotGrid: bool = False

        self.__ll: Literal[0, 1, 2] = 0

    @property
    def logLevel(self):
        return self.__ll

    @logLevel.setter
    def logLevel(self, value):
        self.__ll = value

    @property
    def plotRebar(self):
        return self.__plotRebar

    @plotRebar.setter
    def plotRebar(self, value):
        self.__plotRebar = value

    @property
    def colorNodes(self):
        return self.__colorNodes

    @colorNodes.setter
    def colorNodes(self, value):
        self.__colorNodes = value

    @property
    def ax(self):
        return self.__ax

    @ax.setter
    def ax(self, value):
        self.__ax = value

    @property
    def colorTriaFill(self):
        return self.__colorTriaFill

    @colorTriaFill.setter
    def colorTriaFill(self, value):
        self.__colorTriaFill = value

    @property
    def plotTriaShrinkFill(self):
        return self.__plotTriaShrinkFill

    @plotTriaShrinkFill.setter
    def plotTriaShrinkFill(self, value):
        self.__plotTriaShrinkFill = value

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

    @textSize.setter
    def textSize(self, value):
        self.__textSize = value

    @property
    def vertexStresses(self) -> Union[Dict[int, float], None]:
        return self.__vertexStresses

    @vertexStresses.setter
    def vertexStresses(self, value: Union[Dict[int, float], None] = None) -> None:
        self.__vertexStresses = value

    @property
    def rebarStresses(self) -> Union[Dict[int, float], None]:
        return self.__rebarStresses

    @rebarStresses.setter
    def rebarStresses(self, value: Union[Dict[int, float], None] = None) -> None:
        self.__rebarStresses = value

    @property
    def plotGrid(self):
        return self.__plotGrid

    @plotGrid.setter
    def plotGrid(self, value):
        self.__plotGrid = value

    @property
    def plotSolidNodesId(self):
        return self.__plotSolidNodesId

    @plotSolidNodesId.setter
    def plotSolidNodesId(self, value):
        self.__plotSolidNodesId = value

    @property
    def modeler(self):
        return self.__modeler

    @modeler.setter
    def modeler(self, value: SectionModeler) -> None:
        self.__modeler = value

    def __plotScaled(self, xDataScaled, yDataScaled, refLenght):
        # Scale from reference lenght in data space
        dx = xDataScaled
        dy = yDataScaled
        headScale = 0.02
        i1 = 0
        i2 = 1
        self.__ax.arrow(
            dx[i1],
            dy[i1],
            dx[i2] - dx[i1],
            dy[i2] - dy[i1],
            head_width=refLenght * headScale,
            head_length=refLenght * headScale,
            length_includes_head=True,
        )
        i1 = 1
        i2 = 2
        self.__ax.arrow(
            dx[i1],
            dy[i1],
            dx[i2] - dx[i1],
            dy[i2] - dy[i1],
            head_width=refLenght * headScale,
            head_length=refLenght * headScale,
            length_includes_head=True,
        )
        i1 = 2
        i2 = 0
        self.__ax.arrow(
            dx[i1],
            dy[i1],
            dx[i2] - dx[i1],
            dy[i2] - dy[i1],
            head_width=refLenght * headScale,
            head_length=refLenght * headScale,
            length_includes_head=True,
        )

    @staticmethod
    def __genTriaData(coords: Tuple[float, float, float, float, float, float]) -> Any:
        xData = [coords[0], coords[2], coords[4], coords[0]]
        yData = [coords[1], coords[3], coords[5], coords[1]]

        xg = (xData[0] + xData[1] + xData[2]) / 3
        yg = (yData[0] + yData[1] + yData[2]) / 3

        xDataScaled = [
            (xData[0] - xg) * 0.8 + xg,
            (xData[1] - xg) * 0.8 + xg,
            (xData[2] - xg) * 0.8 + xg,
        ]

        yDataScaled = [
            (yData[0] - yg) * 0.8 + yg,
            (yData[1] - yg) * 0.8 + yg,
            (yData[2] - yg) * 0.8 + yg,
        ]

        polyData = [
            [xData[0], yData[0]],
            [xData[1], yData[1]],
            [xData[2], yData[2]],
        ]

        polyDataScaled = [
            [xDataScaled[0], yDataScaled[0]],
            [xDataScaled[1], yDataScaled[1]],
            [xDataScaled[2], yDataScaled[2]],
        ]

        return xData, yData, xDataScaled, yDataScaled, polyData, polyDataScaled

    def __plotModeler(self, modeler: SectionModeler) -> None:
        trianglesCoords = modeler.getTrianglesCoords()

        # We need this for operate in data space
        refLenght = math.sqrt(modeler.calcSolidAreaProperties().area)

        ########################################################################
        #                               TRIANGLES                              #
        ########################################################################
        for k in trianglesCoords:
            coords = trianglesCoords[k]
            (
                xData,
                yData,
                xDataScaled,
                yDataScaled,
                polyData,
                polyDataScaled,
            ) = self.__genTriaData(coords)

            ####################################################################
            #                               FILL                               #
            ####################################################################
            if self.__plotTriaFill:
                self.__ax.add_patch(
                    Polygon(polyData, facecolor=self.__colorTriaFill, alpha=0.3)
                )

            ####################################################################
            #                               LINES                              #
            ####################################################################
            if self.__plotTriaLines:
                self.__ax.plot(xData, yData, color=self.__colorTriaLines)

            ####################################################################
            #                       SHRINK LINES AND FILL                      #
            ####################################################################
            if self.__plotTriaShrinkFill:
                self.__ax.add_patch(
                    Polygon(polyDataScaled, facecolor=self.__colorTriaFill, alpha=0.3)
                )

            if self.__plotTriaShrink:
                self.__plotScaled(xDataScaled, yDataScaled, refLenght)

        ########################################################################
        #                               SHAPE                                  #
        ########################################################################
        trianglesIds = modeler.getTrianglesIds()
        if self.__plotShapeEdges:
            segsId = []
            for k in trianglesIds:
                t = trianglesIds[k]
                # Order by id segments
                seg1 = [t[1], t[0]]
                if t[1] > t[0]:
                    seg1 = [t[0], t[1]]
                if seg1 not in segsId:
                    segsId.append(seg1)
                else:
                    segsId.remove(seg1)

                seg2 = [t[2], t[1]]
                if t[2] > t[1]:
                    seg2 = [t[1], t[2]]
                if seg2 not in segsId:
                    segsId.append(seg2)
                else:
                    segsId.remove(seg2)

                seg3 = [t[0], t[2]]
                if t[0] > t[2]:
                    seg3 = [t[2], t[0]]
                if seg3 not in segsId:
                    segsId.append(seg3)
                else:
                    segsId.remove(seg3)

            for ids in segsId:
                self.__ax.plot(
                    [modeler.getNodeX(ids[0]), modeler.getNodeX(ids[1])],
                    [modeler.getNodeY(ids[0]), modeler.getNodeY(ids[1])],
                    color=self.__colorEdge,
                )

        ########################################################################
        #                                NODES                                 #
        ########################################################################
        if self.__plotNodes or self.__plotSolidNodes:
            if self.__plotNodes:
                nodes = modeler.getNodes()
            else:
                nodes = modeler.getSolidNodes()

            for k in nodes:
                node = nodes[k]
                sc = self.__ax.scatter(
                    node.xn,
                    node.yn,
                    marker="o",
                    color=self.__colorNodes,
                    # linewidths = 3.5
                )
                sizeOfAnnotation = math.sqrt(sc.get_sizes()) / 2 + 2

                text_to_put = ""
                if self.__vertexStresses is not None:
                    if self.__plotStresses and self.__plotSolidNodesId:
                        stress = self.__vertexStresses[node.id]
                        text_to_put = f"{stress:.1f}"
                    else:
                        text_to_put = f"{node.id}"

                if self.__plotNodesId or self.__plotSolidNodesId:
                    self.__ax.annotate(
                        text=text_to_put,
                        xy=(node.xn, node.yn),
                        fontsize=self.__textSize,
                        xytext=(sizeOfAnnotation, sizeOfAnnotation),
                        textcoords="offset points",
                        color="red",
                        ha="left",
                        va="center",
                    )

        ########################################################################
        #                                REBAR                                 #
        ########################################################################
        texts = []
        patches = []
        if self.__plotRebar:
            circles = modeler.getCircles()
            for k in circles:
                circle = circles[k]
                x = circle.center.xn
                y = circle.center.yn
                r = circle.radius
                patch = self.__ax.add_patch(Circle((x, y), r))
                self.__ax.add_patch(Circle((x, y), r, fill=False, edgecolor="black"))
                patches.append(patch)

                text_to_put = ""
                if self.__rebarStresses is not None:
                    if self.__plotStresses and self.__plotRebarId:
                        stress = self.__rebarStresses[circle.id]
                        text_to_put = f"{stress:.0f}"
                    else:
                        text_to_put = f"{circle.id}"

                if self.__plotRebarId:
                    # text = self.__ax.annotate(
                    #     text=f"{text_to_put}",
                    #     xy=(x + r, y + r),
                    #     fontsize=self.__textSize,
                    #     # xytext = (sizeOfAnnotation, sizeOfAnnotation),
                    #     # textcoords="offset points",
                    #     ha="left",
                    #     va="center",
                    # )
                    text = self.__ax.text(
                        s=f"{text_to_put}",
                        x=x, # + r,
                        y=y, # + r,
                        fontsize=self.__textSize,
                        # xytext = (sizeOfAnnotation, sizeOfAnnotation),
                        # textcoords="offset points",
                        ha="left",
                        va="center",
                    )
                    texts.append(text)
        #self.__fig.canvas.draw()
        adjust_text(
            texts=texts,
            objects=patches,
            ax=self.__ax,
            max_move=100,
            expand=(1.2, 1.2),
            arrowprops=dict(arrowstyle='->', color='red', linewidth=0.3),
            force_static=(2.0,2.0))

        ########################################################################
        #                               MESHES                                 #
        ########################################################################
        if self.__plotMeshTria:
            for i in range(modeler.meshSize()):
                trianglesCoords = modeler.getTrianglesCoordsAtMesh(i)
                for k in trianglesCoords:
                    coords = trianglesCoords[k]
                    (
                        xData_,
                        yData_,
                        xDataScaled_,
                        yDataScaled_,
                        polyData_,
                        polyDataScaled_,
                    ) = self.__genTriaData(coords)
                    self.__plotScaled(xDataScaled_, yDataScaled_, refLenght)
                    self.__ax.plot(xData_, yData_, color=self.__colorTriaLines)

    def setView(self, opt: SectionPlotViewEnum) -> None:
        if opt == SectionPlotViewEnum.SP_VIEW_MESH:
            self.__plotShapeEdges = False
            self.__plotTriaLines = False
            self.__plotTriaFill = False
            self.__plotTriaShrinkFill = False
            self.__plotTriaShrink = False
            self.__plotNodes = False
            self.__plotNodesId = False
            self.__plotRebar = False
            self.__plotMeshTria = True
            self.__plotRebarId = False
            self.__plotGrid = False

        if opt == SectionPlotViewEnum.SP_VIEW_SECTION:
            self.__plotShapeEdges = True
            self.__plotTriaLines = False
            self.__plotTriaFill = True
            self.__plotTriaShrinkFill = False
            self.__plotTriaShrink = False
            self.__plotNodes = False
            self.__plotNodesId = False
            self.__plotRebar = True
            self.__plotRebarId = True
            self.__plotStresses = False

        if opt == SectionPlotViewEnum.SP_VIEW_SECTION_NOID:
            self.__plotShapeEdges = True
            self.__plotTriaLines = False
            self.__plotTriaFill = True
            self.__plotTriaShrinkFill = False
            self.__plotTriaShrink = False
            self.__plotNodes = False
            self.__plotNodesId = False
            self.__plotRebar = True
            self.__plotRebarId = False

        if opt == SectionPlotViewEnum.SP_HIDE_REBAR_ID:
            self.__plotRebarId = False

        if opt == SectionPlotViewEnum.SP_VIEW_NODES:
            self.__plotNodes = True
            self.__plotNodesId = True

        if opt == SectionPlotViewEnum.SP_HIDE_NODES:
            self.__plotNodes = False
            self.__plotNodesId = False

        if opt == SectionPlotViewEnum.SP_VIEW_SOLID_NODES:
            self.__plotNodes = False
            self.__plotNodesId = False
            self.__plotSolidNodes = True
            self.__plotSolidNodesId = True
            self.__plotStresses = False

        if opt == SectionPlotViewEnum.SP_VIEW_STRESSES:
            self.__plotNodes = False
            self.__plotNodesId = False
            self.__plotRebarId = True
            self.__plotSolidNodes = True
            self.__plotSolidNodesId = True
            self.__plotStresses = True

        if opt == SectionPlotViewEnum.SP_HIDE_NODES:
            self.__plotNodes = False
            self.__plotNodesId = False
            self.__plotSolidNodes = False
            self.__plotSolidNodesId = False

        if opt == SectionPlotViewEnum.SP_VIEW_GRID:
            self.__plotGrid = True

        if opt == SectionPlotViewEnum.SP_HIDE_GRID:
            self.__plotGrid = False

        if opt == SectionPlotViewEnum.SP_VIEW_TRIANGLES:
            self.__plotShapeEdges = False
            self.__plotTriaLines = True
            self.__plotTriaFill = True
            self.__plotTriaShrinkFill = True
            self.__plotTriaShrink = True
            self.__plotNodes = True
            self.__plotNodesId = True
            self.__plotRebar = False
            self.__plotRebarId = False

        if opt == SectionPlotViewEnum.SP_VIEW_TRIANGLE_NOID:
            self.__plotShapeEdges = False
            self.__plotTriaLines = True
            self.__plotTriaFill = True
            self.__plotTriaShrinkFill = True
            self.__plotTriaShrink = True
            self.__plotNodes = True
            self.__plotNodesId = False
            self.__plotRebar = False
            self.__plotRebarId = False

    def plot(self):
        self.__ax.cla()
        #self.__ax.set_facecolor("white")
        if self.__plotGrid:
            self.__ax.grid(True, "both")
            self.__ax.minorticks_on()
            self.__ax.set_facecolor("#EBEBEB")

        if self.__modeler is not None:
            self.__plotModeler(self.__modeler)
        else:
            log("ERR", "Modeler does not assigned !!!", self.__ll)

    @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=None,
        edgecolor=None,
        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,
        )
