from numpy import asarray

import matplotlib.pyplot as plt

from matplotlib.patches import Circle
from matplotlib.patches import Polygon

from matplotlib.collections import LineCollection
from matplotlib.collections import PatchCollection
from matplotlib.collections import PolyCollection

from mpl_toolkits.mplot3d.art3d import Line3DCollection

from compas.geometry import centroid_points_xy
from compas.geometry import midpoint_line_xy
from compas.utilities import color_to_rgb


ZORDER_POLYGONS = 1000
ZORDER_LINES = 2000
ZORDER_POINTS = 3000
ZORDER_LABELS = 4000


# ==============================================================================
# axes
# ==============================================================================


def create_axes_xy(
    figsize=(8.0, 6.0),
    dpi=100,
    xlabel=None,
    ylabel=None,
    fontname="Times New Roman",
    fontsize=10,
    grid=False,
    xlim=None,
    ylim=None,
    ticklength=20,
    tickfontsize=10,
    xscale="linear",
    yscale="linear",
    bgcolor="#ffffff",
):
    """Initializes plot axes object for matplotlib plotting.

    Parameters
    ----------
    figsize : 2-tuple of float, optional
        Size of the figure.
        Default is ``(8.0, 6.0)``
    dpi : int, optional
        Resolution of the plot.
        Default is ``100``.
    xlabel : str, optional
        Label for the x-axis.
        Default is None.
    ylabel : str, optional
        Label for the y-axis.
        Default is None.
    fontname : str, optional
        Fontname of the main labels and text.
        Default is ``'Times New Roman'``.
    fontsize : int, optional
        Fontsize of the main labels and text.
        Default is ``10``.
    grid : bool, optional
        Display grid.
        Default is False.
    xlim : 2-tuple, optional
        Limits of the X-axis.
        Default is None.
    ylim : 2-tuple, optional
        Limits of the Y-axis.
        Default is None.
    ticklength : float, optional
        Length of the ticks.
        Default is ``20``.
    tickfontsize : int, optional
        Fontsize of the ticks.
        Default is ``10``.
    xscale : {'linear', 'log'}
        Scale of the X axis.
    yscale : {'linear', 'log'}
        Scale of the Y axis.
    bgcolor : str | list, optional
        Background color as hex string or rgb tuple.
        Default is white.

    Returns
    -------
    object
        Matplotlib axes.

    """
    # mpl.rcParams['figure.figsize'] = figsize
    # mpl.rcParams['figure.dpi'] = dpi
    # mpl.rcParams['savefig.dpi'] = dpi
    fig = plt.figure(facecolor=bgcolor, figsize=figsize, dpi=dpi)
    axes = fig.add_subplot(111, aspect="equal")
    axes.grid(b=grid)
    axes.set_frame_on(False)
    if xlabel:
        axes.set_xlabel(xlabel, fontname=fontname, fontsize=fontsize)
    if ylabel:
        axes.set_ylabel(ylabel, fontname=fontname, fontsize=fontsize)
    if xlim:
        axes.set_xlim(xlim[0], xlim[1])
    if ylim:
        axes.set_ylim(ylim[0], ylim[1])
    # plt.minorticks_on()
    # plt.tick_params(which='major', length=ticklength, labelsize=tickfontsize)
    # plt.tick_params(which='minor', length=ticklength * 0.33)
    axes.set_xscale(xscale)
    axes.set_yscale(yscale)
    axes.set_xticks([])
    axes.set_yticks([])
    # axes.set_xmargin(0.05)
    # axes.set_ymargin(0.05)
    # axes.autoscale()
    return axes


def create_axes_3d(
    size=(10, 7),
    xlabel="$x$",
    ylabel="$y$",
    zlabel="$z$",
    fontname="Times New Roman",
    fontsize=20,
    grid=True,
    limits=None,
    ticklength=20,
    tickfontsize=15,
    angle=(30, 45),
):
    """Initializes plot axes object for matplotlib plotting.

    Parameters
    ----------
    size : 2-tuple of float, optinoal
        Size of the figure.
        Default is ``(10.0, 7.0)``.
    xlabel : str, optional
        Label for the x-axis.
        Default is ``'$x$'``.
    ylabel : str, optional
        Label for the y-axis.
        Default is ``'$y$'``.
    zlabel : str, optional
        Label for the z-axis.
        Default is ``'$z$'``.
    fontname : str, optional
        Fontname of the main labels and text.
        Default is ``'Times New Roman'``.
    fontsize : int, optional
        Fontsize of the main labels and text.
        Default is ``10``.
    grid : bool, optional
        Display grid.
        Default is False.
    limits : dict, optional
        Axis limits and tick spacing.
        Default is None.
    ticklength : float, optional
        Length of the ticks.
        Default is ``20``.
    tickfontsize : int, optional
        Fontsize of the ticks.
        Default is ``15``.
    angle : 2-tuple of float
        Elevation and azimuth angles for 3D plots.
        Default is ``'30.0, 45.0'``.

    Returns
    -------
    object
        Matplotlib axes.

    """
    fig = plt.figure(facecolor="white", figsize=size)
    axes = fig.add_subplot(111, projection="3d", aspect="equal")
    # axes.w_xaxis.set_pane_color((1, 1, 1, 1))
    # axes.w_yaxis.set_pane_color((1, 1, 1, 1))
    # axes.w_zaxis.set_pane_color((1, 1, 1, 1))
    axes.grid(b=grid)
    axes.set_xlabel(xlabel, fontname=fontname, fontsize=fontsize)
    axes.set_ylabel(ylabel, fontname=fontname, fontsize=fontsize)
    axes.set_zlabel(zlabel, fontname=fontname, fontsize=fontsize)
    # axes.view_init(elev=angle[0], azim=angle[1])
    axes.set_xticks([])
    axes.set_yticks([])
    axes.set_zticks([])
    axes.set_xmargin(0.05)
    axes.set_ymargin(0.05)
    axes.autoscale()
    return axes


# ==============================================================================
# points
# ==============================================================================


def draw_points_xy(points, axes, facecolor="#ffffff", edgecolor="#000000", linewidth=0.5, radius=1.0):
    """Creates an XY point collection and adds it to the axis.

    Parameters
    ----------
    points : list
        XY(Z) coordinates of the points.
    axes : object
        Matplotlib axes.
    facecolor : str | list, optional
        Color of the point face.
        Default is white.
    edgecolor : str | list, optional
        Color of the point edge.
        Default is black.
    linewidth : float or list, optional
        Width of the point edge.
        Default is ``0.5``.
    radius : float or list, optional
        The radius of the points.
        Default is ``1.0``.

    Returns
    -------
    object
        A collection of points.

    """
    p = len(points)
    # preprocess patch parameters
    if isinstance(facecolor, str):
        facecolor = [facecolor] * p
    if isinstance(edgecolor, str):
        edgecolor = [edgecolor] * p
    if isinstance(linewidth, (int, float)):
        linewidth = float(linewidth)
        linewidth = [linewidth] * p
    if isinstance(radius, (int, float)):
        radius = float(radius)
        radius = [radius] * p
    # --------------------------------------------------------------------------
    circles = []
    for i in range(p):
        point = points[i]
        circle = Circle(point[0:2], radius=radius[i])
        circles.append(circle)
    coll = PatchCollection(
        circles,
        facecolors=facecolor,
        edgecolors=edgecolor,
        linewidhts=linewidth,
        alpha=1.0,
        zorder=ZORDER_POINTS,
    )
    axes.add_collection(coll)
    return coll


def draw_xpoints_xy(points, axes):
    """Creates an XY point collection and adds it to the axis.

    Parameters
    ----------
    points : list of dict
        List of dictionaries containing the point properties.
        Each point is represented by a circle with a given radius.
        The following properties of the circle can be specified in the point dict.

        * pos (list): XY(Z) coordinates
        * radius (float, optional): the radius of the circle. Default is 0.1.
        * text (str, optional): the text of the label. Default is None.
        * facecolor (rgb or hex color, optional): The color of the face of the circle. Default is white.
        * edgecolor (rgb or hex color, optional): The color of the edge of the cicrle. Default is black.
        * edgewidth (float, optional): The width of the edge of the circle. Default is 1.0.
        * textcolor (rgb or hex color, optional): Color of the text label. Default is black.
        * fontsize (int, optional): Font size of the text label. Default is 12.

    axes : object
        Matplotlib axes.

    Returns
    -------
    object
        The matplotlib point collection object.

    """
    circles = []
    facecolors = []
    edgecolors = []
    linewidths = []
    for point in points:
        pos = point["pos"]
        radius = point["radius"]
        text = point.get("text")
        fcolor = point.get("facecolor") or "#ffffff"
        ecolor = point.get("edgecolor") or "#000000"
        lwidth = point.get("edgewidth") or 1.0
        textcolor = point.get("textcolor") or "#000000"
        fontsize = point.get("fontsize") or 12
        circles.append(Circle(pos[0:2], radius=radius))
        facecolors.append(color_to_rgb(fcolor, normalize=True))
        edgecolors.append(color_to_rgb(ecolor, normalize=True))
        linewidths.append(lwidth)
        if text is not None:
            axes.text(
                pos[0] - 0.01,
                pos[1] - 0.01,
                text,
                fontsize=fontsize,
                zorder=ZORDER_LABELS,
                ha="center",
                va="center",
                color=textcolor,
            )
    coll = PatchCollection(
        circles,
        linewidths=linewidths,
        facecolors=facecolors,
        edgecolors=edgecolors,
        alpha=1.0,
        zorder=ZORDER_POINTS,
    )
    axes.add_collection(coll)
    return coll


def draw_points_3d(points, axes, facecolor="#ffffff", edgecolor="#000000"):
    """Creates a 3D point collection and adds it to the axis.

    Parameters
    ----------
    points : list
        XYZ coordinates of the points.
    axes : object
        Matplotlib axes.
    facecolor : str | list, optional
        Color of the face of the points.
        Default is white.
    edgecolor : str | list, optional
        Color of the edge of the points.
        Default is black.

    Returns
    -------
    object
        The matplotlib point collection object.

    """
    p = len(points)
    points = asarray(points)
    if isinstance(facecolor, str):
        facecolor = [facecolor] * p
    if isinstance(edgecolor, str):
        edgecolor = [edgecolor] * p
    x = points[:, 0]
    y = points[:, 1]
    z = points[:, 2]
    coll, _ = axes.plot(x, y, z, "o", color=(1.0, 1.0, 1.0))
    return coll


# ==============================================================================
# lines
# ==============================================================================


def draw_lines_xy(lines, axes, linewidth=1.0, linestyle="-", color="#000000", alpha=1.0):
    """Creates an XY line collection and adds it to the axis.

    Parameters
    ----------
    lines : list
        List of ((X1, Y1), (X2, X2)) lines.
    axes : object
        Matplotlib axes.
    linewidth : float or list of float, optional
        Width of the lines.
        Default is ``1.0``.
    linestyle : str | list of str, optional
        Matplotlib line style strings.
        Default is ``'-'``.
    color : str | list of str, optional
        Color of the lines.
        Default is black.
    alpha : float or list of float, optional
        Opacity of the lines.
        Default is ``1.0``.

    Returns
    -------
    object
        The matplotlib point collection object.

    """
    n = len(lines)
    if isinstance(linewidth, (int, float)):
        linewidth = float(linewidth)
        linewidth = [linewidth] * n
    if isinstance(color, str):
        color = [color] * n
    # --------------------------------------------------------------------------
    coll = LineCollection(
        [(start[0:2], end[0:2]) for start, end in lines],
        linewidths=linewidth,
        colors=color,
        linestyle=linestyle,
        alpha=alpha,
        zorder=ZORDER_LINES,
    )
    axes.add_collection(coll)
    return coll


def draw_xlines_xy(lines, axes, alpha=1.0, linestyle="solid"):
    """Creates an XY line collection and adds it to the axis.

    Parameters
    ----------
    lines : list
        List of dictionaries containing the line properties.
        The following properties of a line can be specified in the dict.

        * start (list): XY(Z) coordinates of the start point.
        * end (list): XY(Z) coordinatesof the end point.
        * width (float, optional): The width of the line. Default is ``1.0``.
        * color (rgb tuple or hex string, optional): The color of the line. Default is black.
        * text (str, optional): The text of the label. Default is None.
        * textcolor (rgb tuple or hex string, optional): Color of the label text. Default is black.
        * fontsize (int, optional): The size of the font of the label text. Default is ```12``.

    axes : object
        Matplotlib axes.
    alpha : float, optional
        Opacity of the lines.
        Default is ``1.0``.
    linestyle : str, optional
        Matplotlib line style strings.
        Default is ``'solid'``.

    Returns
    -------
    object
        The matplotlib line collection object.

    """
    fromto = []
    widths = []
    colors = []
    for line in lines:
        sp = line["start"]
        ep = line["end"]
        width = line.get("width", 1.0)
        color = line.get("color", "#000000")
        text = line.get("text", None)
        textcolor = line.get("textcolor") or "#000000"
        fontsize = line.get("fontsize") or 6
        fromto.append((sp[0:2], ep[0:2]))
        widths.append(width)
        colors.append(color_to_rgb(color, normalize=True))
        if text:
            x, y, z = midpoint_line_xy((sp, ep))
            t = axes.text(
                x,
                y,
                text,
                fontsize=fontsize,
                zorder=ZORDER_LABELS,
                ha="center",
                va="center",
                color=color_to_rgb(textcolor, normalize=True),
            )
            t.set_bbox({"color": "#ffffff", "alpha": 1.0})
    coll = LineCollection(
        fromto,
        linewidths=widths,
        colors=colors,
        linestyle=linestyle,
        alpha=alpha,
        zorder=ZORDER_LINES,
    )
    axes.add_collection(coll)
    return coll


def draw_lines_3d(lines, axes, linewidth=1.0, linestyle="solid", color="#000000"):
    """Creates an 3D line collection and adds it to the axis.

    Parameters
    ----------
    lines : list
        Pairs of XYZ coordinates defining start and end points of the lines.
    axes : object
        Matplotlib axes.
    linewidth : float or list of float, optional
        Width for the lines.
        Default is ``1.0``.
    linestyle : str, optional
        Matplotlib line style strings.
        Default is ``'solid'``.
    color : str | list of str, optional
        Color of the lines.
        Default is black.

    Returns
    -------
    object
        The matplotlib line collection object.

    """
    n = len(lines)
    if isinstance(linewidth, (int, float)):
        linewidth = float(linewidth)
        linewidth = [linewidth] * n
    if isinstance(color, str):
        color = [color] * n

    coll = Line3DCollection(
        lines,
        linewidths=linewidth,
        colors=color,
        linestyle=linestyle,
        zorder=ZORDER_LINES,
    )
    axes.add_collection(coll)
    return coll


# ==============================================================================
# polylines
# ==============================================================================


def draw_xpolylines_xy(polylines, axes):
    paths = []

    widths = []
    colors = []

    for polyline in polylines:
        points = polyline["points"]
        width = polyline.get("width", 1.0)
        color = polyline.get("color", "#000000")
        text = polyline.get("text", None)
        textcolor = polyline.get("textcolor") or "#000000"
        fontsize = polyline.get("fontsize") or 6

        path = [point[0:2] for point in points]
        paths.append(path)

        widths.append(width)
        colors.append(color_to_rgb(color, normalize=True))

        if text:
            p = len(points)

            if p % 2 == 0:
                a = points[p // 2]
                b = points[p // 2 + 1]
                x, y, z = midpoint_line_xy((a, b))
            else:
                x, y = points[p // 2 + 1]

            t = axes.text(
                x,
                y,
                text,
                fontsize=fontsize,
                zorder=ZORDER_LABELS,
                ha="center",
                va="center",
                color=color_to_rgb(textcolor, normalize=True),
            )

            t.set_bbox({"color": "#ffffff", "alpha": 1.0, "edgecolor": "#ffffff"})

    coll = PolyCollection(
        paths,
        closed=False,
        linewidths=widths,
        edgecolors=colors,
        facecolors="none",
        zorder=ZORDER_LINES,
    )

    axes.add_collection(coll)

    return coll


# ==============================================================================
# arrows
# ==============================================================================


def draw_xarrows_xy(lines, axes):
    """Creates an XY arrow collection and adds it to the axis.

    Parameters
    ----------
    lines : list of dict
        List of dictionaries containing the arrow line properties.
        The following properties of an arrow can be specified in the dict.

        * start (list): XY(Z) coordinates of the starting point.
        * end (list): XY(Z) coordinates of the end point.
        * text (str, optional): The text of the label. Default is None.
        * textcolor (rgb tuple or hex string, optional): Color of the label text. Default is black.
        * fontsize (int, optional): The size of the font of the label text. Default is ```6``.
        * color (rgb tuple or hex string, optional): Color of the arrow. Default is black.
        * width (float): Width of the arrow. Default is ``1.0``.

    axes : object
        Matplotlib axes.

    """
    arrowprops = {
        "arrowstyle": "-|>,head_length=0.6,head_width=0.2",
        "connectionstyle": "arc3,rad=0.0",
        "linewidth": 1.0,
        "color": "#000000",
        "shrinkB": 0.0,
        "shrinkA": 0.0,
    }
    xys = []
    for line in lines:
        sp = line["start"][:2]
        ep = line["end"][:2]
        text = line.get("text", None)
        textcolor = line.get("textcolor") or "#000000"
        fontsize = line.get("fontsize") or 6
        arrowprops["color"] = color_to_rgb(line.get("color", "#000000"), normalize=True)
        arrowprops["linewidth"] = line.get("width", 1.0)
        axes.annotate(
            "",
            xy=ep,
            xytext=sp,
            arrowprops=arrowprops,
            zorder=ZORDER_LINES,
        )
        xys.append(sp)
        xys.append(ep)
        if text:
            x, y, z = midpoint_line_xy((sp, ep))
            t = axes.text(
                x,
                y,
                text,
                fontsize=fontsize,
                zorder=ZORDER_LABELS,
                ha="center",
                va="center",
                color=color_to_rgb(textcolor, normalize=True),
            )
            t.set_bbox({"color": "#ffffff", "alpha": 1.0, "edgecolor": "#ffffff"})
    axes.update_datalim(xys)


# ==============================================================================
# labels
# ==============================================================================


def draw_xlabels_xy(labels, axes):
    """Creates a label collection and adds it to the axis.

    Parameters
    ----------
    labels : list of dict
        List of dictionaries containing the label properties.
    axes : object
        Matplotlib axes.

    """
    for label in labels:
        x, y = label["pos"]
        text = label["text"]
        fontsize = label["fontsize"]
        color = label.get("color") or "#ffffff"
        textcolor = label.get("textcolor") or "#000000"
        bbox = dict(
            color=color_to_rgb(color, normalize=True),
            edgecolor=color_to_rgb(color, normalize=True),
            alpha=1.0,
            pad=0.0,
        )
        t = axes.text(
            x,
            y,
            text,
            fontsize=fontsize,
            zorder=ZORDER_LABELS,
            ha="center",
            va="center",
            color=color_to_rgb(textcolor, normalize=True),
        )
        t.set_bbox(bbox)


# ==============================================================================
# faces
# ==============================================================================


def draw_xpolygons_xy(polygons, axes):
    """Creates a polygon collection and adds it to the axis.

    Parameters
    ----------
    polygons : list of dict
        List of dictionaries containing the polygon properties.
        The following properties can be specified in the dict.

        * points (list): XY(Z) coordinates of the polygon vertices.
        * text (str, optional): The text of the label. Default is None.
        * textcolor (rgb tuple or hex string, optional): Color of the label text. Default is black.
        * fontsize (int, optional): The size of the font of the label text. Default is ```12``.
        * facecolor (rgb tuple or hex string, optional): Color of the polygon face. Default is white.
        * edgecolor (rgb tuple or hex string, optional): Color of the edge of the polygon. Default is black.
        * edgewidth (float): Width of the polygon edge. Default is ``1.0``.

    axes : object
        Matplotlib axes.

    Returns
    -------
    object
        The matplotlib polygon collection object.

    """
    facecolors = []
    edgecolors = []
    linewidths = []
    patches = []

    for attr in polygons:
        points = attr["points"]
        text = attr.get("text")
        textcolor = color_to_rgb(attr.get("textcolor", "#000000"), normalize=True)

        facecolors.append(color_to_rgb(attr.get("facecolor", "#ffffff"), normalize=True))
        edgecolors.append(color_to_rgb(attr.get("edgecolor", "#000000"), normalize=True))
        linewidths.append(attr.get("edgewidth", 1.0))

        patches.append(Polygon([point[0:2] for point in points]))

        if text:
            c = centroid_points_xy(points)
            axes.text(
                c[0],
                c[1],
                text,
                fontsize=attr.get("fontsize", 10.0),
                zorder=ZORDER_LABELS,
                ha="center",
                va="center",
                color=textcolor,
            )

    coll = PatchCollection(
        patches,
        facecolors=facecolors,
        edgecolors=edgecolors,
        lw=linewidths,
        zorder=ZORDER_POLYGONS,
    )

    axes.add_collection(coll)

    return coll
