"""Plotting functions for optical setups and mode matching solutions and analyses."""

from dataclasses import dataclass
from io import BytesIO
from itertools import product
from typing import TYPE_CHECKING, Any, cast

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.axes import Axes
from matplotlib.collections import (
    FillBetweenPolyCollection,
    LineCollection,
    PathCollection,
)
from matplotlib.colorbar import Colorbar
from matplotlib.figure import Figure
from matplotlib.legend import Legend
from matplotlib.legend_handler import HandlerPatch
from matplotlib.lines import Line2D
from matplotlib.patches import Ellipse, Rectangle
from matplotlib.text import Annotation
from matplotlib.ticker import FuncFormatter
from scipy import stats

from .config import Config

# prevent circular import for type annotation in function signature
if TYPE_CHECKING:
    from .core import OpticalSetup
    from .solver import ModeMatchingSolution

RELATIVE_MARGIN = 0.1  #: Relative margin size for plotting optical setups

milli_formatter = FuncFormatter(lambda x, _: f"{x*1e3:.0f}")  #: Formatter for millimeter axes
micro_formatter = FuncFormatter(lambda x, _: f"{x*1e6:.0f}")  #: Formatter for micrometer axes


def fig_to_png(fig: Figure) -> bytes:
    """Convert a Matplotlib figure to a PNG as bytes.

    Args:
        fig: The figure to convert.

    Returns:
        The PNG representation of the figure as bytes.
    """
    buf = BytesIO()
    fig.savefig(buf, format="png", bbox_inches="tight")
    plt.close(fig)
    return buf.getvalue()


# helper classes for fancy corset legend, yes, kind of overkill for an easter egg
class _CorsetPatch(Rectangle):
    def __init__(self, *args, lace_color, **kwargs):
        super().__init__(*args, **kwargs)
        self.lace_color = lace_color


class _HandlerCorset(HandlerPatch):
    def create_artists(self, legend, orig_handle: _CorsetPatch, xdescent, ydescent, width, height, fontsize, trans):
        corset = _CorsetPatch(
            (xdescent, ydescent),
            width,
            height,
            color=orig_handle.get_facecolor(),
            lace_color=orig_handle.lace_color,
            transform=trans,
        )
        cross_edges = np.linspace(xdescent, xdescent + width, 4)
        cross_x = (cross_edges[1:] + cross_edges[:-1]) / 2
        cross_y = np.ones_like(cross_x) * (ydescent + height / 2)
        crosses = Line2D(
            cross_x, cross_y, color=orig_handle.lace_color, transform=trans, linestyle="none", marker="x", markersize=4
        )
        return [corset, crosses]


Legend.update_default_handler_map({_CorsetPatch: _HandlerCorset()})


@dataclass(frozen=True)
class OpticalSetupPlot:
    """Plot of an optical setup with references to the plot elements."""

    ax: Axes  #: Axes containing the plot
    zs: np.ndarray  #: Z-coordinates of the beam profile
    rs: np.ndarray  #: Radii of the beam profile
    beam: FillBetweenPolyCollection  #: Fill between collection for the beam
    beam_ci: list[Line2D]  #: List of lines representing beam confidence intervals
    lenses: list[tuple[LineCollection, Annotation, Rectangle | None]]  #: List of lens plot elements
    handles: list[tuple[Any, str]]  #: List of plot handles by their legend labels
    r_max: float  #: Maximum beam radius in the plot


@dataclass(frozen=True)
class ModeMatchingPlot:
    """Plot of a mode matching solution with references to the plot elements."""

    ax: Axes  #: Axes containing the plot
    setup: OpticalSetupPlot  #: Plot of the actual optical setup
    desired_beam: OpticalSetupPlot  #: Plot of the desired beam profile
    ranges: list[Rectangle]  #: List of ranges in the plot
    apertures: list[Line2D]  #: List of aperture lines in the plot
    passages: list[Rectangle]  #: List of passage rectangles in the plot
    handles: list[tuple[Any, str]]  #: List of plot handles by their legend labels


def get_handles(ax: Axes) -> list[tuple[Any, str]]:
    """Get a list of plot handles by their legend labels.

    Args:
        ax: The axes to get the handles from.

    Returns:
        A list of tuples containing legend handles and their corresponding labels, empty if no legend exists.
    """

    legend = ax.get_legend()
    if legend is None:
        return []
    return [(handle, text.get_text()) for handle, text in zip(legend.legend_handles, legend.texts, strict=True)]


# TODO refactor this function to reduce complexity
def plot_setup(  # noqa: C901
    self: "OpticalSetup",
    *,
    points: int | np.ndarray | None = None,
    limits: tuple[float, float] | None = None,
    beam_kwargs: dict | None = None,
    confidence_interval: float | bool | None = None,
    rayleigh_range_cap: float | None = None,
    free_lenses: list[int] = [],  # noqa: B006
    ax: Axes | None = None,
    show_legend: bool | None = None,
    legend_loc: str | None = None,
    # TODO add back other configuration options?
) -> OpticalSetupPlot:
    """Plot the the beam profile and optical elements of the setup.

    Args:
        self: The optical setup instance.
        points: Number of points or specific z-coordinates to evaluate the beam profile.
            If `None`, this defaults to :attr:`Config.PlotSetup.beam_points <corset.config.Config.PlotSetup.beam_points>`.
        limits: Z-coordinate limits for the plot.
            If `None`, limits are determined from the beam and elements.
        beam_kwargs: Additional keyword arguments passed to the beam plot.
            If `None` this defaults to :attr:`Config.PlotSetup.beam_kwargs <corset.config.Config.PlotSetup.beam_kwargs>`.
        confidence_interval: Confidence interval probability for beam uncertainty visualization.
            If `False`, no uncertainty is plotted.
            If `None` this defaults to :attr:`Config.PlotSetup.confidence_interval <corset.config.Config.PlotSetup.confidence_interval>`.
        rayleigh_range_cap: Maximum Rayleigh range to consider when determining plot limits.
            If `None` this defaults to :attr:`Config.PlotSetup.rayleigh_range_cap <corset.config.Config.PlotSetup.rayleigh_range_cap>`.
        free_lenses: Indices of lenses to treat as free elements in the plot.
        ax: The axes to plot on. If `None`, the current axes are used.
        show_legend: Whether to show a legend for the plot.
            If `None`, this defaults to :attr:`Config.PlotSetup.show_legend <corset.config.Config.PlotSetup.show_legend>`.
        legend_loc: Location of the legend in the plot.
            If `None`, this defaults to :attr:`Config.PlotSetup.legend_loc <corset.config.Config.PlotSetup.legend_loc>`.

    Returns:
        An :class:`OpticalSetupPlot` containing references to the plot elements.
    """

    ax = ax or plt.gca()
    show_legend = Config.get(show_legend, Config.PlotSetup.show_legend)
    legend_loc = Config.get(legend_loc, Config.PlotSetup.legend_loc)
    beam_kwargs = Config.get(beam_kwargs, Config.PlotSetup.beam_kwargs)
    if confidence_interval is True:
        raise ValueError("confidence_interval cannot be True, must be a float between 0 and 1 or False")
    confidence_interval = Config.get(confidence_interval, Config.PlotSetup.confidence_interval)
    rayleigh_range_cap = Config.get(rayleigh_range_cap, Config.PlotSetup.rayleigh_range_cap)

    # plot the beam as a corset it it is called "Corset", there are no lenses,
    # the limits are automatically determined and there is no beam covariance
    # all of this should be the case for the desired beam in mode matching solutions
    plot_corset = all(
        [
            limits is None,
            points is None,
            beam_kwargs.get("label") == "Corset",
            not self.elements,
            self.initial_beam.gauss_cov is None,
        ]
    )
    if plot_corset:
        beam_kwargs |= {"color": "#e6e6e6" if Config.mpl_is_dark() else "#1a1a1a", "lw": 6, "alpha": 1, "zorder": 150}

    lens_positions = [pos for pos, _ in self.elements]

    if isinstance(points, np.ndarray):
        zs = points
    else:
        if not limits:
            cap = rayleigh_range_cap
            all_bounds = [
                self.beams[0].focus - min(cap, self.beams[0].rayleigh_range),
                self.beams[-1].focus + min(cap, self.beams[-1].rayleigh_range),
                *(),
                *lens_positions,
            ]
            if self.elements:
                all_bounds.append(self.elements[0][0] - min(cap, self.beams[0].rayleigh_range))
                all_bounds.append(self.elements[-1][0] + min(cap, self.beams[-1].rayleigh_range))

            limits = (min(all_bounds), max(all_bounds))

        num_points = points if isinstance(points, int) else Config.PlotSetup.beam_points
        zs = np.linspace(limits[0], limits[1], num_points)

    handles = get_handles(ax)

    rs = cast(np.ndarray, self.radius(zs))
    fill_between = ax.fill_between(zs, -rs, rs, **{"zorder": 100, **beam_kwargs})
    beam_label = beam_kwargs.get("label", "Beam")
    handles.append((fill_between, beam_label))
    r_max = np.max(rs)

    if plot_corset:
        # plot the laces
        lace_color = "#333333" if Config.mpl_is_dark() else "#cccccc"
        n_laces = 3
        bounds = np.linspace(*limits, n_laces + 1)  # pyright: ignore[reportCallIssue]
        width = np.diff(bounds)[0] * 0.6
        height = self.initial_beam.waist
        positions = (bounds[:-1] + bounds[1:]) / 2
        for pos in positions:
            kwargs = {"color": lace_color, "lw": 3, "solid_capstyle": "round", "zorder": 150}
            ax.plot([pos - width / 2, pos + width / 2], [-height / 2, height / 2], **kwargs)
            ax.plot([pos - width / 2, pos + width / 2], [height / 2, -height / 2], **kwargs)

        # make the legend artist
        handles[-1] = (_CorsetPatch((0, 0), 1, 1, color=beam_kwargs["color"], lace_color=lace_color), beam_label)

    beam_deviation = []
    if confidence_interval is not False and self.beams[0].gauss_cov is not None:
        rs_ci = self.radius_dev(zs) * stats.norm.interval(confidence_interval)[1]
        ci_color = beam_kwargs.get("color")
        if ci_color is None or ci_color == "none":
            ci_color = beam_kwargs.get("edgecolor")
        for r0, ci in product([-rs, rs], [-rs_ci, rs_ci]):
            line = ax.plot(zs, r0 + ci, ls="--", color=ci_color, alpha=beam_kwargs.get("alpha"), zorder=105)[0]
            beam_deviation.append(line)
        r_max = np.max(rs + rs_ci)
        handles.append((beam_deviation[0], f"{round(confidence_interval*100)}% CI ({beam_label})"))
    # TODO make beam plot function?

    # TODO factor out into plot lens function?
    lenses = []
    for i, (pos, lens) in enumerate(self.elements):
        color = "C2"
        zorder = 200

        line = ax.vlines(
            pos,
            -r_max * (1 + RELATIVE_MARGIN),
            r_max * (1 + RELATIVE_MARGIN),
            color=color,
            zorder=zorder,
        )

        label_text = lens.name if lens.name is not None else f"f={round(lens.focal_length*1e3)}mm"
        if i in free_lenses:
            label_text = f"$L_{i}$: {label_text} @{round(pos*1e3)}mm"
        label = ax.text(
            pos,
            -r_max * (1 + RELATIVE_MARGIN),
            label_text,
            va="center",
            ha="left",
            rotation="vertical",
            rotation_mode="anchor",
            bbox={"fc": plt.rcParams["axes.facecolor"], "ec": "none", "alpha": 0.5},
            zorder=zorder,
        )

        rect = None
        if lens.left_margin or lens.right_margin:
            rect = Rectangle(
                (pos - lens.left_margin, -r_max * (1 + RELATIVE_MARGIN)),
                lens.left_margin + lens.right_margin,
                2 * r_max * (1 + RELATIVE_MARGIN),
                fc="none",
                ec=color,
                ls="--",
                zorder=zorder,
            )
            ax.add_patch(rect)

        lenses.append((line, label, rect))

    if lenses and not any(label == "Lens" for label, _ in handles):
        # fake handle with vertical line marker
        handles.append((Line2D([], [], color=color, label="Lens", marker="|", linestyle="None", markersize=10), "Lens"))
    rectangles = [r for _, _, r in lenses if r is not None]
    if rectangles and not any(label == "Lens Margin" for label, _ in handles):
        handles.append((rectangles[0], "Lens Margin"))

    ax.set_ylim(
        min(-r_max * (1 + 3 * RELATIVE_MARGIN), ax.get_ylim()[0]),
        max(r_max * (1 + 3 * RELATIVE_MARGIN), ax.get_ylim()[1]),
    )

    ax.set_xlabel("z in mm")
    ax.xaxis.set_major_formatter(milli_formatter)

    ax.set_ylabel(r"w(z) in $\mathrm{\mu m}$")
    ax.yaxis.set_major_formatter(micro_formatter)

    if show_legend:
        ax.legend(*zip(*handles, strict=True), loc=legend_loc).set_zorder(1000)

    return OpticalSetupPlot(
        ax=ax,
        zs=zs,
        rs=rs,
        beam=fill_between,
        beam_ci=beam_deviation,
        lenses=lenses,
        handles=handles,
        r_max=r_max,
    )


def plot_mode_match_solution_setup(
    self: "ModeMatchingSolution",
    *,
    setup_kwargs: dict | None = None,
    desired_kwargs: dict | None = None,
    ax: Axes | None = None,
    show_legend: bool | None = None,
    legend_loc: str | None = None,
) -> ModeMatchingPlot:
    """Plot the mode matching solution setup including the desired beam and constraints.

    Args:
        self: The mode matching solution instance.
        setup_kwargs: Additional keyword arguments passed to the setup plot function.
            If `None`, this defaults to :attr:`Config.PlotSolution.setup_kwargs <corset.config.Config.PlotSolution.setup_kwargs>`.
        desired_kwargs: Additional keyword arguments passed to the desired beam plot function.
            If `None`, this defaults to :attr:`Config.PlotSolution.desired_kwargs <corset.config.Config.PlotSolution.desired_kwargs>`.
        ax: The axes to plot on. If `None`, the current axes are used.
        show_legend: Whether to show a legend for the plot.
            If `None`, this defaults to :attr:`Config.PlotSolution.show_legend <corset.config.Config.PlotSolution.show_legend>`.
        legend_loc: Location of the legend in the plot.
            If `None`, this defaults to :attr:`Config.PlotSolution.legend_loc <corset.config.Config.PlotSolution.legend_loc>`.
    Returns:
        An :class:`ModeMatchingPlot` containing references to the plot elements.
    """
    from .core import OpticalSetup
    from .solver import Aperture, Passage

    ax = ax or plt.gca()
    setup_kwargs = Config.get(setup_kwargs, Config.PlotSolution.setup_kwargs)
    desired_kwargs = Config.get(desired_kwargs, Config.PlotSolution.desired_kwargs)
    show_legend = Config.get(show_legend, Config.PlotSolution.show_legend)
    legend_loc = Config.get(legend_loc, Config.PlotSolution.legend_loc)

    problem = self.candidate.problem

    setup_plot = plot_setup(
        self.setup,
        ax=ax,
        free_lenses=self.candidate.parametrized_setup.free_elements,
        show_legend=show_legend,
        **setup_kwargs,
    )
    desired_plot = plot_setup(
        OpticalSetup(problem.desired_beam, []),
        ax=ax,
        beam_kwargs=desired_kwargs,
        show_legend=show_legend,
    )

    handles = get_handles(ax)
    if show_legend:
        desired = handles.pop(-1)
        assert desired[1] == desired_kwargs["label"]  # noqa: S101
        handles.insert(1, desired)  # hack to show desired beam 2nd in legend

    ranges = []
    for range_ in problem.ranges:
        # TODO color configuratbility?
        rectangle = Rectangle(
            (range_.left, -setup_plot.r_max * (1 + 2 * RELATIVE_MARGIN)),
            range_.right - range_.left,
            2 * setup_plot.r_max * (1 + 2 * RELATIVE_MARGIN),
            fc="none",
            ec="lightgrey",
            ls="--",
            zorder=50,  # TODO
        )
        ranges.append(ax.add_patch(rectangle))
    if ranges:
        handles.append((ranges[0], "Shifting Range"))

    apertures = []
    passages = []
    for con in problem.constraints:
        if isinstance(con, Aperture):
            apertures.extend(
                ax.plot([con.position, con.position], [-con.radius, con.radius], marker="o", color="C4", zorder=80)
            )
        elif isinstance(con, Passage):
            rect = Rectangle(
                (con.left, -con.radius),
                con.right - con.left,
                2 * con.radius,
                fc="C4",
                ec="none",
            )
            passages.append(ax.add_patch(rect))

    if apertures:
        # TODO add dots to vertical line?
        handles.append(  # fake handle with line marker
            (Line2D([], [], color="C4", label="Aperture", marker="|", linestyle="None", markersize=10), "Aperture")
        )
    if passages:
        handles.append((passages[0], "Passage"))

    if show_legend:
        ax.legend(*zip(*handles, strict=True), loc=legend_loc).set_zorder(1000)

    ax.set_title(f"Optical Setup ({self.overlap*100:.2f}% mode overlap)")

    return ModeMatchingPlot(
        ax=ax,
        setup=setup_plot,
        desired_beam=desired_plot,
        ranges=ranges,
        apertures=apertures,
        passages=passages,
        handles=handles,
    )


@dataclass(frozen=True)
class ReachabilityPlot:
    """Plot of a reachability analysis with references to the plot elements."""

    ax: Axes  #: Axes containing the plot
    center: PathCollection  #: Point representing the optimal focus and waist
    center_ci: Ellipse | None  #: Uncertainty of the arriving beam around the optimal point
    lines: list[list[Line2D]]  #: Lines representing parameter variations
    contour: Any  #: Contour plot of the mode overlap
    colorbar: Colorbar | None  #: Colorbar for the contour plot


@dataclass(frozen=True)
class SensitivityPlot:
    """Plot of a sensitivity analysis with references to the plot elements."""

    ax: Axes  #: Axes containing the plot
    displacements_ci: Ellipse | None  #: Confidence ellipse in terms of required displacements
    contours: list[Any]  #: Contour plots of the sensitivity
    colorbar: Colorbar | None  #: Colorbar for the contour plots
    handles: list  #: Handles for the plot elements


def plot_ellipse(
    ax: Axes, mean: tuple[float, float], var: np.ndarray, confidence_interval: float, **ellipse_kwargs
) -> Ellipse:
    eigenvalues, eigenvectors = np.linalg.eigh(var)
    angle = np.degrees(np.arctan2(*eigenvectors[:, 0][::-1]))
    diameters = 2 * np.sqrt(eigenvalues) * stats.chi2.ppf(confidence_interval, df=2) ** 0.5
    base_kwargs = {"fill": False, "ls": "--", "zorder": 100, "color": plt.rcParams["text.color"]}
    patch = Ellipse(mean, *diameters, angle=angle, **(base_kwargs | ellipse_kwargs))
    ax.add_patch(patch)
    return patch


def plot_reachability(
    self: "ModeMatchingSolution",
    *,
    displacement: float | list[float] | None = None,
    num_samples: int | list[int] | None = None,
    line_step: int | list[int] | None = None,
    dimensions: list[int] | None = None,
    focus_range: tuple[float, float] | None = None,
    waist_range: tuple[float, float] | None = None,
    confidence_interval: float | None = None,
    grid_resolution: int | None = None,
    ax: Axes | None = None,
) -> ReachabilityPlot:
    """Plot a reachability analysis of the mode matching solution.

    Args:
        self: The mode matching solution instance.
        displacement: Maximum displacement from the optimal position(s) in meters.
            If `None`, this defaults to :attr:`Config.PlotReachability.displacement <corset.config.Config.PlotReachability.displacement>`.
        num_samples: Number of samples to take along each dimension.
            If `None`, this defaults to :attr:`Config.PlotReachability.num_samples <corset.config.Config.PlotReachability.num_samples>`.
        line_step: Step size to skip certain lines for increased smoothness while retaining clarity.
            If `None`, this defaults to :attr:`Config.PlotReachability.line_step <corset.config.Config.PlotReachability.line_step>`.
        dimensions: Indices of the dimensions to analyze. If `None`, all dimensions are used.
        focus_range: Range of focus values to display on the x-axis. If `None`, determined automatically.
        waist_range: Range of waist values to display on the y-axis. If `None`, determined automatically.
        confidence_interval: Confidence ellipse probability for focus position and waist radius.
            If `None`, this defaults to :attr:`Config.PlotReachability.confidence_interval <corset.config.Config.PlotReachability.confidence_interval>`.
        grid_resolution: Resolution of the grid for the background mode overlap contour plot.
            If `None`, this defaults to :attr:`Config.Overlap.grid_resolution <corset.config.Config.Overlap.grid_resolution>`.
        ax: The axes to plot on. If `None`, the current axes are used.

    Returns:
        A :class:`ReachabilityPlot` containing references to the plot elements.
    """

    from .analysis import vector_partial
    from .solver import mode_overlap

    ax = ax or plt.gca()
    displacement = Config.get(displacement, Config.PlotReachability.displacement)
    num_samples = Config.get(num_samples, Config.PlotReachability.num_samples)
    line_step = Config.get(line_step, Config.PlotReachability.line_step)
    confidence_interval = Config.get(confidence_interval, Config.PlotReachability.confidence_interval)
    grid_resolution = Config.get(grid_resolution, Config.Overlap.grid_resolution)

    if dimensions is None:
        dimensions = list(range(len(self.positions)))

    num_dof = len(dimensions)
    displacement = cast(list, [displacement] * num_dof if np.isscalar(displacement) else displacement)
    num_samples = cast(list, [num_samples] * num_dof if np.isscalar(num_samples) else num_samples)
    line_step = cast(list, [line_step] * num_dof if np.isscalar(line_step) else line_step)

    linspaces = [
        np.linspace(-d, d, num=n) + offset
        for d, n, offset in zip(
            displacement, num_samples, self.positions, strict=True  # pyright: ignore[reportArgumentType]
        )
    ]

    if focus_range is not None:
        ax.set_xlim(focus_range)

    if waist_range is not None:
        ax.set_ylim(waist_range)

    focus_and_waist = np.vectorize(
        vector_partial(self.candidate.parametrized_focus_and_waist, self.positions, dimensions), signature="(n)->(2)"
    )

    grids = np.moveaxis(np.meshgrid(*linspaces, indexing="ij"), 0, -1)  # pyright: ignore[reportArgumentType]
    focuses, waists = np.moveaxis(focus_and_waist(grids), -1, 0)
    delta_focuses = focuses - self.candidate.problem.desired_beam.focus

    contrast_color = plt.rcParams["text.color"]
    center = ax.scatter(0, self.candidate.problem.desired_beam.waist, color=contrast_color, marker="o", zorder=100)
    center_ci = None
    if confidence_interval is not False and self.setup.beams[-1].gauss_cov is not None:
        center_ci = plot_ellipse(
            ax,
            (0, self.candidate.problem.desired_beam.waist),
            self.setup.beams[-1].gauss_cov,
            confidence_interval,
        )

    def select_lines(arr: np.ndarray, dim: int, steps: list[int]) -> np.ndarray:
        steps = [step if i != dim else 1 for i, step in enumerate(steps)]
        indices = tuple(slice(None, None, step) for step in steps)
        return np.moveaxis(arr[indices], dim, -1).reshape(-1, arr.shape[dim])

    lines: list[list[Line2D]] = []
    for i, (dim, disp) in enumerate(zip(dimensions, displacement, strict=True)):
        focuses_flat = select_lines(delta_focuses, i, line_step)  # pyright: ignore[reportArgumentType]
        waists_flat = select_lines(waists, i, line_step)  # pyright: ignore[reportArgumentType]
        lines.append([])
        for focus, waist in zip(focuses_flat, waists_flat, strict=True):
            line = ax.plot(focus, waist, color=f"C{dim}", label=rf"$\Delta x_{i}$ ($\pm{disp*1e3:.1f}$ mm)")[0]
            lines[-1].append(line)
        ax.annotate(
            "",
            xy=(focuses_flat[0][1], waists_flat[0][1]),
            xytext=(focuses_flat[0][0], waists_flat[0][0]),
            arrowprops={"color": f"C{i}", "headwidth": 10, "width": 2},
        )

    ax.legend(handles=[line[0] for line in lines], loc="lower left").set_zorder(1000)

    focuses_grid, waists_grid = np.meshgrid(
        np.linspace(*ax.get_xlim(), grid_resolution), np.linspace(*ax.get_ylim(), grid_resolution)
    )
    overlap = mode_overlap(
        focuses_grid,  # pyright: ignore[reportArgumentType]
        self.candidate.problem.desired_beam.waist,
        waists_grid,  # pyright: ignore[reportArgumentType]
        self.candidate.problem.setup.initial_beam.wavelength,
    )

    levels = Config.Overlap.levels
    colors = Config.Overlap.colors()
    res = ax.contourf(focuses_grid, waists_grid, overlap * 100, levels=levels, colors=colors, alpha=0.5)
    cb = ax.figure.colorbar(res, label="Mode overlap (%)")

    jacobian = self.analysis.focus_and_waist_jacobian

    ax.set_xlabel(rf"$\Delta z_0$ in mm ($\nabla z_0$=[{' '.join(f'{x:.3f}' for x in jacobian[0])}])")
    ax.xaxis.set_major_formatter(milli_formatter)

    ax.set_ylabel(rf"$w_0$ in $\mathrm{{\mu m}}$ ($\nabla w_0$=[{' '.join(f'{x:.3f}' for x in jacobian[1]*1e3)}]e-3)")
    ax.yaxis.set_major_formatter(micro_formatter)

    ax.set_title("Reachability Analysis")

    return ReachabilityPlot(ax=ax, center=center, center_ci=center_ci, lines=lines, contour=res, colorbar=cb)


def plot_sensitivity(
    self: "ModeMatchingSolution",
    *,
    dimensions: tuple[int, int] | tuple[int, int, int] | None = None,
    worst_overlap: float | None = None,
    x_displacement: float | None = None,
    y_displacement: float | None = None,
    z_displacement: float | None = None,
    num_samples_z: int | None = None,
    confidence_interval: float | None = None,
    force_contour_lines: bool | None = None,
    grid_resolution: int | None = None,
    ax: Axes | None = None,
) -> SensitivityPlot:
    """Plot a sensitivity analysis of the mode matching solution.

    Args:
        self: The mode matching solution instance.
        dimensions: Indices of the dimensions to analyze. If `None`, the two least
            coupled dimensions are used as the x and y dimensions, and the remaining most sensitive
            dimension is used as the auxiliary z dimension if applicable.
        worst_overlap: Worst-case mode overlap contour line that should still be fully visible in the
            plot used for automatic range determination.
            If `None`, this defaults to :attr:`Config.PlotSensitivity.worst_overlap <corset.config.Config.PlotSensitivity.worst_overlap>`.
        x_displacement: Displacement magnitude for x-axis values. If `None`, determined automatically from worst_overlap.
        y_displacement: Displacement magnitude for y-axis values. If `None`, determined automatically from worst_overlap.
        z_displacement: Displacement magnitude for z-axis values. If `None`, determined automatically from worst_overlap.
        num_samples_z: Number of z-slices to plot if 3 dimensions are used.
            If `None`, this defaults to :attr:`Config.PlotSensitivity.num_samples_z <corset.config.Config.PlotSensitivity.num_samples_z>`.
        confidence_interval: Confidence ellipse probability in terms of the required x and y displacements.
            If `None`, this defaults to :attr:`Config.PlotSensitivity.confidence_interval <corset.config.Config.PlotSensitivity.confidence_interval>`.
        force_contour_lines: If `True`, always use contour lines instead of filled contours for plots with two
            degrees of freedom. If `None`, this defaults to :attr:`Config.PlotSensitivity.force_contour_lines <corset.config.Config.PlotSensitivity.force_contour_lines>`.
        grid_resolution: Resolution of the grid for the contour plots.
            If `None`, this defaults to :attr:`Config.Overlap.grid_resolution <corset.config.Config.Overlap.grid_resolution>`.
        ax: The axes to plot on. If `None`, the current axes are used.

    Returns:
        A :class:`SensitivityPlot` containing references to the plot elements.
    """

    from .analysis import vector_partial

    worst_overlap = Config.get(worst_overlap, Config.PlotSensitivity.worst_overlap)
    num_samples_z = Config.get(num_samples_z, Config.PlotSensitivity.num_samples_z)
    confidence_interval = Config.get(confidence_interval, Config.PlotSensitivity.confidence_interval)
    force_contour_lines = Config.get(force_contour_lines, Config.PlotSensitivity.force_contour_lines)
    grid_resolution = Config.get(grid_resolution, Config.Overlap.grid_resolution)
    ax = ax or plt.gca()

    if dimensions is None:
        dimensions = self.analysis.min_coupling_pair
        if len(self.positions) > 2:
            all_dims = set(range(len(self.positions)))
            remaining_dims = list(all_dims - set(dimensions))
            most_sensitive = max(remaining_dims, key=lambda dim: self.analysis.sensitivities[dim, dim])
            dimensions = (*dimensions, most_sensitive)

    mode_overlap = np.vectorize(
        vector_partial(
            self.candidate.parametrized_overlap, self.positions, dimensions  # pyright: ignore[reportArgumentType]
        ),
        signature="(n)->()",
    )

    # the hessian will always be badly conditioned for ndim > 2, so only use the subspace we care about
    # which should be well conditioned
    A_xy_inv = np.linalg.inv(-self.analysis.hessian[np.ix_(dimensions[:2], dimensions[:2])] / 2)  # noqa: N806
    # use Config.get so the type checker knows the displacements are not None after this point
    x_displacement = Config.get(x_displacement, np.sqrt(A_xy_inv[0, 0] * (1 - worst_overlap)) * 1.1)
    y_displacement = Config.get(y_displacement, np.sqrt(A_xy_inv[1, 1] * (1 - worst_overlap)) * 1.1)

    if z_displacement is None and len(dimensions) > 2:
        A_xyz = -self.analysis.hessian[np.ix_(dimensions, dimensions)]  # noqa: N806
        eigs = np.linalg.eigh(A_xyz)
        v_min = eigs.eigenvectors[:, np.argmin(eigs.eigenvalues)]
        z_displacement = min(abs(v_min[2] / v_min[0] * x_displacement), abs(v_min[2] / v_min[1] * y_displacement)) / 1.1

    base = np.array([self.positions[dim] for dim in dimensions])
    xs = np.linspace(-x_displacement, x_displacement, grid_resolution)
    ys = np.linspace(-y_displacement, y_displacement, grid_resolution)
    xsg, ysg = np.meshgrid(xs, ys)
    cmap = plt.get_cmap("coolwarm") if Config.mpl_is_dark() else plt.get_cmap("berlin")

    contours = []
    colorbar = None
    handles: list = []

    if len(dimensions) == 2:
        overlaps = mode_overlap(np.stack([xsg, ysg], axis=-1) + base)
        if force_contour_lines:
            cont = ax.contour(xsg, ysg, overlaps * 100, levels=Config.Overlap.levels, colors=cmap(0.5))
            ax.clabel(cont, fmt="%1.1f%%")
            contours.append(cont)
        else:
            cont = ax.contourf(xsg, ysg, overlaps * 100, levels=Config.Overlap.levels, colors=Config.Overlap.colors())
            colorbar = ax.figure.colorbar(cont, label="Mode overlap (%)")
            contours.append(cont)
    else:
        z_displacement = cast(float, z_displacement)  # make type checker happy
        z_colors = cmap(np.linspace(0, 1, num_samples_z))
        zs = np.linspace(-z_displacement, z_displacement, num_samples_z)
        for i, (color, z) in enumerate(zip(z_colors, zs, strict=True)):
            positions = np.stack([xsg, ysg, z * np.ones_like(xsg)], axis=-1) + base
            overlaps = mode_overlap(positions)
            zorder = 100 if i == num_samples_z // 2 else 10
            cont = ax.contour(
                xsg, ysg, overlaps * 100, levels=Config.Overlap.levels, colors=[color], alpha=0.7, zorder=zorder
            )
            contours.append(cont)
            if i == num_samples_z // 2:
                ax.clabel(cont, fmt="%1.1f%%", colors=[color])
        handles = [
            Line2D([], [], color=color, label=rf"$\Delta x_{{{dimensions[2]}}} = {value*1e3:.1f}$ mm")
            for color, value in zip(z_colors, zs, strict=True)
        ]
        ax.legend(handles=handles, loc="lower left").set_zorder(1000)

    displacements_ci = None
    if confidence_interval is not False and self.setup.beams[-1].gauss_cov is not None:
        jac_inv = np.linalg.inv(self.analysis.focus_and_waist_jacobian[:, dimensions[:2]])
        cov_xy = jac_inv @ self.setup.beams[-1].gauss_cov @ jac_inv.T
        displacements_ci = plot_ellipse(ax, (0, 0), cov_xy, confidence_interval)

    unit = Config.sensitivity_unit
    dims = dimensions
    sens_x = self.analysis.sensitivities[dims[0], dims[0]] * unit.value.factor
    sens_y = self.analysis.sensitivities[dims[1], dims[1]] * unit.value.factor
    ax.set_xlabel(rf"$\Delta x_{{{dims[0]}}}$ in mm ($s_{{{str(dims[0])*2}}}={sens_x:.2f}{unit.value.tex}$)")
    ax.set_ylabel(rf"$\Delta x_{{{dims[1]}}}$ in mm ($s_{{{str(dims[1])*2}}}={sens_y:.2f}{unit.value.tex}$)")
    ax.xaxis.set_major_formatter(milli_formatter)
    ax.yaxis.set_major_formatter(milli_formatter)

    ax.set_title(rf"Sensitivity Analysis ($r_{{{dims[0]}{dims[1]}}}={self.analysis.min_coupling*100:.2f}\%$)")

    return SensitivityPlot(
        ax=ax, displacements_ci=displacements_ci, contours=contours, colorbar=colorbar, handles=handles
    )


def plot_mode_match_solution_all(
    self: "ModeMatchingSolution",
    *,
    figsize: tuple[float, float] | None = None,
    width_ratios: tuple[float, float, float] | None = None,
    tight_layout: bool | None = None,
    setup_kwargs: dict | None = None,
    reachability_kwargs: dict | None = None,
    sensitivity_kwargs: dict | None = None,
) -> tuple[Figure, tuple[ModeMatchingPlot, ReachabilityPlot, SensitivityPlot]]:
    """Plot the mode matching solution setup, reachability analysis, and sensitivity analysis in a single figure.

    Args:
        self: The mode matching solution instance.
        figsize: Figure size for the combined plot.
            If `None`, this defaults to :attr:`Config.PlotAll.figsize <corset.config.Config.PlotAll.figsize>`.
        width_ratios: Width ratios for the three subplots.
            If `None`, this defaults to :attr:`Config.PlotAll.width_ratios <corset.config.Config.PlotAll.width_ratios>`.
        tight_layout: Whether to use tight layout for the figure.
            If `None`, this defaults to :attr:`Config.PlotAll.tight_layout <corset.config.Config.PlotAll.tight_layout>`.
        setup_kwargs: Additional keyword arguments for the setup plot.
            If `None`, this defaults to :attr:`Config.PlotAll.setup_kwargs <corset.config.Config.PlotAll.setup_kwargs>`.
        reachability_kwargs: Additional keyword arguments for the reachability plot.
            If `None`, this defaults to :attr:`Config.PlotAll.reachability_kwargs <corset.config.Config.PlotAll.reachability_kwargs>`.
        sensitivity_kwargs: Additional keyword arguments for the sensitivity plot.
            If `None`, this defaults to :attr:`Config.PlotAll.sensitivity_kwargs <corset.config.Config.PlotAll.sensitivity_kwargs>`.

    Returns:
        A tuple containing the figure and an inner tuple with the three plot objects.
    """

    figsize = Config.get(figsize, Config.PlotAll.figsize)
    width_ratios = Config.get(width_ratios, Config.PlotAll.width_ratios)
    tight_layout = Config.get(tight_layout, Config.PlotAll.tight_layout)
    setup_kwargs = Config.get(setup_kwargs, Config.PlotAll.setup_kwargs)
    reachability_kwargs = Config.get(reachability_kwargs, Config.PlotAll.reachability_kwargs)
    sensitivity_kwargs = Config.get(sensitivity_kwargs, Config.PlotAll.sensitivity_kwargs)

    fig, (axl, axr, axc) = plt.subplots(
        1, 3, figsize=figsize, gridspec_kw={"width_ratios": width_ratios}, tight_layout=tight_layout
    )
    sol_plot = self.plot_setup(ax=axl, **setup_kwargs)
    reach_plot = self.plot_reachability(ax=axr, **reachability_kwargs)
    sens_plot = self.plot_sensitivity(ax=axc, **sensitivity_kwargs)
    return fig, (sol_plot, reach_plot, sens_plot)
