"""Matplotlib backend for mesh visualization."""

from typing import Literal, TYPE_CHECKING
import torch
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection, PolyCollection
from matplotlib.colors import Normalize
from matplotlib.cm import ScalarMappable

if TYPE_CHECKING:
    from torchmesh import Mesh


def draw_mesh_matplotlib(
    mesh: "Mesh",
    point_scalar_values: torch.Tensor | None,
    cell_scalar_values: torch.Tensor | None,
    active_scalar_source: Literal["points", "cells", None],
    show: bool,
    cmap: str,
    vmin: float | None,
    vmax: float | None,
    alpha_points: float,
    alpha_cells: float,
    alpha_edges: float,
    show_edges: bool,
    ax=None,
):
    """Draw mesh using matplotlib backend.

    Supports 0D, 1D, 2D, and 3D spatial dimensions with appropriate matplotlib primitives.

    Args:
        mesh: Mesh object to visualize
        point_scalar_values: Processed point scalar values (1D tensor or None)
        cell_scalar_values: Processed cell scalar values (1D tensor or None)
        active_scalar_source: Which scalar source is active ("points", "cells", or None)
        show: Whether to call plt.show()
        cmap: Colormap name
        vmin: Minimum value for colormap normalization
        vmax: Maximum value for colormap normalization
        alpha_points: Opacity for points (0-1)
        alpha_cells: Opacity for cells (0-1)
        alpha_edges: Opacity for edges (0-1)
        show_edges: Whether to draw cell edges
        ax: Existing matplotlib axes (if None, creates new figure)

    Returns:
        matplotlib.axes.Axes object
    """
    ### Convert mesh data to numpy
    points_np = mesh.points.cpu().numpy()
    cells_np = mesh.cells.cpu().numpy()

    ### Determine neutral colors based on active_scalar_source
    point_neutral_color = "black"
    if active_scalar_source is None:
        cell_neutral_color = "lightblue"
    elif active_scalar_source == "points":
        cell_neutral_color = "lightgray"
    else:  # active_scalar_source == "cells"
        cell_neutral_color = None  # Will be colored by scalars

    ### Create figure and axes if not provided
    if ax is None:
        if mesh.n_spatial_dims == 3:
            fig = plt.figure(figsize=(8, 8))
            ax = fig.add_subplot(111, projection="3d")
        else:
            fig, ax = plt.subplots(figsize=(8, 8))
    else:
        fig = ax.get_figure()

    ### Determine scalar colormap normalization
    if active_scalar_source == "points" and point_scalar_values is not None:
        scalar_values_for_norm = point_scalar_values.cpu().numpy()
    elif active_scalar_source == "cells" and cell_scalar_values is not None:
        scalar_values_for_norm = cell_scalar_values.cpu().numpy()
    else:
        scalar_values_for_norm = None

    if scalar_values_for_norm is not None:
        norm = Normalize(
            vmin=vmin if vmin is not None else scalar_values_for_norm.min(),
            vmax=vmax if vmax is not None else scalar_values_for_norm.max(),
        )
        scalar_mapper = ScalarMappable(norm=norm, cmap=cmap)
    else:
        norm = None
        scalar_mapper = None

    ### Draw based on spatial dimensionality
    if mesh.n_spatial_dims == 0:
        _draw_0d(
            ax,
            points_np,
            point_scalar_values,
            active_scalar_source,
            scalar_mapper,
            point_neutral_color,
            alpha_points,
        )
    elif mesh.n_spatial_dims == 1:
        _draw_1d(
            ax,
            points_np,
            cells_np,
            point_scalar_values,
            cell_scalar_values,
            active_scalar_source,
            scalar_mapper,
            point_neutral_color,
            cell_neutral_color,
            alpha_points,
            alpha_cells,
        )
    elif mesh.n_spatial_dims == 2:
        _draw_2d(
            ax,
            points_np,
            cells_np,
            point_scalar_values,
            cell_scalar_values,
            active_scalar_source,
            scalar_mapper,
            point_neutral_color,
            cell_neutral_color,
            alpha_points,
            alpha_cells,
            alpha_edges,
            show_edges,
        )
    elif mesh.n_spatial_dims == 3:
        _draw_3d(
            ax,
            points_np,
            cells_np,
            point_scalar_values,
            cell_scalar_values,
            active_scalar_source,
            scalar_mapper,
            point_neutral_color,
            cell_neutral_color,
            alpha_points,
            alpha_cells,
            alpha_edges,
            show_edges,
        )
    else:
        raise ValueError(
            f"Cannot visualize mesh with {mesh.n_spatial_dims=}.\n"
            f"Supported spatial dimensions: 0, 1, 2, 3."
        )

    ### Add colorbar if we have active scalars
    if scalar_mapper is not None:
        plt.colorbar(scalar_mapper, ax=ax, label="Scalar Value")

    ### Set labels and make axes equal
    if mesh.n_spatial_dims == 1:
        ax.set_xlabel("x")
        ax.set_aspect("equal", adjustable="box")
    elif mesh.n_spatial_dims == 2:
        ax.set_xlabel("x")
        ax.set_ylabel("y")
        ax.set_aspect("equal", adjustable="box")
    elif mesh.n_spatial_dims == 3:
        ax.set_xlabel("x")
        ax.set_ylabel("y")
        ax.set_zlabel("z")

        ### Make 3D axes equal by adjusting limits to have same range
        ax.set_box_aspect((1, 1, 1))

        xlim = ax.get_xlim3d()
        ylim = ax.get_ylim3d()
        zlim = ax.get_zlim3d()

        x_range = abs(xlim[1] - xlim[0])
        x_middle = np.mean(xlim)
        y_range = abs(ylim[1] - ylim[0])
        y_middle = np.mean(ylim)
        z_range = abs(zlim[1] - zlim[0])
        z_middle = np.mean(zlim)

        # Use the maximum range to ensure all axes have equal scale
        plot_radius = 0.5 * max([x_range, y_range, z_range])

        ax.set_xlim3d([x_middle - plot_radius, x_middle + plot_radius])
        ax.set_ylim3d([y_middle - plot_radius, y_middle + plot_radius])
        ax.set_zlim3d([z_middle - plot_radius, z_middle + plot_radius])

    if show:
        plt.show()

    return ax


def _draw_0d(
    ax,
    points_np,
    point_scalar_values,
    active_scalar_source,
    scalar_mapper,
    point_neutral_color,
    alpha_points,
):
    """Draw 0D manifold (point cloud) in 0D space."""
    # For 0D spatial dimensions, all points are at the origin
    # We can represent them as points at x=0
    n_points = len(points_np)

    if active_scalar_source == "points" and point_scalar_values is not None:
        colors = scalar_mapper.to_rgba(point_scalar_values.cpu().numpy())
    else:
        colors = point_neutral_color

    # Draw points at the origin
    ax.scatter(
        np.zeros(n_points), np.zeros(n_points), c=colors, alpha=alpha_points, s=5
    )
    ax.set_xlim(-1, 1)
    ax.set_ylim(-0.5, 0.5)


def _draw_1d(
    ax,
    points_np,
    cells_np,
    point_scalar_values,
    cell_scalar_values,
    active_scalar_source,
    scalar_mapper,
    point_neutral_color,
    cell_neutral_color,
    alpha_points,
    alpha_cells,
):
    """Draw 1D manifold (edges) in 1D or 2D space."""
    # Points are 1D, so plot along x-axis (or in 2D if embedded in 2D)
    if points_np.shape[1] == 1:
        # Truly 1D: plot on x-axis
        x = points_np[:, 0]
    else:
        # Should not happen for n_spatial_dims=1, but handle gracefully
        raise ValueError(
            f"Expected 1D points for 1D spatial dimension, got shape {points_np.shape}"
        )

    ### Draw cells (line segments)
    if cells_np.shape[0] > 0 and alpha_cells > 0:
        segments = points_np[cells_np[:, :2]]  # Shape: (n_cells, 2, 1)
        segments = np.stack(
            [segments[:, :, 0], np.zeros((len(segments), 2))], axis=-1
        )  # Add y=0

        if active_scalar_source == "cells" and cell_scalar_values is not None:
            colors = scalar_mapper.to_rgba(cell_scalar_values.cpu().numpy())
        else:
            colors = cell_neutral_color

        lc = LineCollection(
            segments, colors=colors, alpha=alpha_cells, linewidths=2, zorder=1
        )
        ax.add_collection(lc)

    ### Draw points
    if alpha_points > 0:
        if active_scalar_source == "points" and point_scalar_values is not None:
            colors = scalar_mapper.to_rgba(point_scalar_values.cpu().numpy())
        else:
            colors = point_neutral_color

        ax.scatter(x, np.zeros_like(x), c=colors, alpha=alpha_points, s=5, zorder=2)


def _draw_2d(
    ax,
    points_np,
    cells_np,
    point_scalar_values,
    cell_scalar_values,
    active_scalar_source,
    scalar_mapper,
    point_neutral_color,
    cell_neutral_color,
    alpha_points,
    alpha_cells,
    alpha_edges,
    show_edges,
):
    """Draw 2D manifold (triangles) in 2D space."""
    ### Draw cells (filled polygons)
    if cells_np.shape[0] > 0 and alpha_cells > 0:
        # Create polygons from cells
        verts = points_np[cells_np]  # Shape: (n_cells, n_vertices_per_cell, 2)

        if active_scalar_source == "cells" and cell_scalar_values is not None:
            facecolors = scalar_mapper.to_rgba(cell_scalar_values.cpu().numpy())
        else:
            facecolors = cell_neutral_color

        if show_edges and alpha_edges > 0:
            edgecolors = "black"
            linewidths = 0.25
        else:
            edgecolors = "none"
            linewidths = 0

        pc = PolyCollection(
            verts,
            facecolors=facecolors,
            edgecolors=edgecolors,
            linewidths=linewidths,
            alpha=alpha_cells,
            zorder=1,
        )
        # Set edge alpha separately if needed
        if show_edges and alpha_edges > 0 and alpha_edges != alpha_cells:
            pc.set_edgecolor([(0, 0, 0, alpha_edges)] * len(verts))

        ax.add_collection(pc)

    ### Draw points
    if alpha_points > 0:
        if active_scalar_source == "points" and point_scalar_values is not None:
            colors = scalar_mapper.to_rgba(point_scalar_values.cpu().numpy())
        else:
            colors = point_neutral_color

        ax.scatter(
            points_np[:, 0],
            points_np[:, 1],
            c=colors,
            alpha=alpha_points,
            s=5,
            zorder=2,
        )

    ### Set axis limits based on data
    if len(points_np) > 0:
        margin = 0.05 * (points_np.max() - points_np.min())
        ax.set_xlim(points_np[:, 0].min() - margin, points_np[:, 0].max() + margin)
        ax.set_ylim(points_np[:, 1].min() - margin, points_np[:, 1].max() + margin)


def _draw_3d(
    ax,
    points_np,
    cells_np,
    point_scalar_values,
    cell_scalar_values,
    active_scalar_source,
    scalar_mapper,
    point_neutral_color,
    cell_neutral_color,
    alpha_points,
    alpha_cells,
    alpha_edges,
    show_edges,
):
    """Draw mesh in 3D space using mpl_toolkits.mplot3d."""
    from mpl_toolkits.mplot3d.art3d import Line3DCollection, Poly3DCollection

    ### Draw cells based on manifold dimension
    if cells_np.shape[0] > 0 and alpha_cells > 0:
        n_manifold_dims = cells_np.shape[1] - 1

        if n_manifold_dims == 0:
            # 0D manifold in 3D: just points (handled below)
            pass

        elif n_manifold_dims == 1:
            # 1D manifold (edges) in 3D: use Line3DCollection
            segments = points_np[cells_np[:, :2]]  # Shape: (n_cells, 2, 3)

            if active_scalar_source == "cells" and cell_scalar_values is not None:
                colors = scalar_mapper.to_rgba(cell_scalar_values.cpu().numpy())
            else:
                colors = cell_neutral_color

            lc = Line3DCollection(
                segments, colors=colors, alpha=alpha_cells, linewidths=2, zorder=1
            )
            ax.add_collection3d(lc)

        elif n_manifold_dims == 2:
            # 2D manifold (triangles) in 3D: use Poly3DCollection
            verts = points_np[cells_np]  # Shape: (n_cells, 3, 3)

            if active_scalar_source == "cells" and cell_scalar_values is not None:
                facecolors = scalar_mapper.to_rgba(cell_scalar_values.cpu().numpy())
            else:
                facecolors = cell_neutral_color

            if show_edges and alpha_edges > 0:
                edgecolors = [(0, 0, 0, alpha_edges)] * len(verts)
                linewidths = 0.25
            else:
                edgecolors = "none"
                linewidths = 0

            pc = Poly3DCollection(
                verts,
                facecolors=facecolors,
                edgecolors=edgecolors,
                linewidths=linewidths,
                alpha=alpha_cells,
                zorder=1,
            )
            ax.add_collection3d(pc)

        elif n_manifold_dims == 3:
            # 3D manifold (tetrahedra) in 3D: extract surface triangles
            # For solid tetrahedra, we need to draw the 4 triangular faces
            # This is complex; for now, we'll just show the vertices
            # A proper implementation would extract the boundary surface
            pass  # Handle tetrahedra by extracting surface in future version

    ### Draw points
    if alpha_points > 0:
        if active_scalar_source == "points" and point_scalar_values is not None:
            colors = scalar_mapper.to_rgba(point_scalar_values.cpu().numpy())
        else:
            colors = point_neutral_color

        ax.scatter(
            points_np[:, 0],
            points_np[:, 1],
            points_np[:, 2],
            c=colors,
            alpha=alpha_points,
            s=5,
            zorder=2,
        )

    ### Set axis limits based on data
    if len(points_np) > 0:
        margin = 0.01 * (points_np.max() - points_np.min())
        ax.set_xlim(points_np[:, 0].min() - margin, points_np[:, 0].max() + margin)
        ax.set_ylim(points_np[:, 1].min() - margin, points_np[:, 1].max() + margin)
        ax.set_zlim(points_np[:, 2].min() - margin, points_np[:, 2].max() + margin)
