"""Visualisation utilities used by the optimal-control examples."""

from __future__ import annotations

import matplotlib.animation as animation
import matplotlib.pyplot as plt
from matplotlib.patches import Circle, FancyArrowPatch, FancyBboxPatch
import numpy as np


COLORS = {
    "background": "white",
    "ink": "black",
    "muted": "0.40",
    "rail": "0.50",
    "cart": "C0",
    "cart_edge": "C0",
    "force": "C3",
    "link_1": "C1",
    "link_2": "C2",
    "mass_1": "black",
    "mass_2": "black",
    "trace": "0.5",
}


def double_cartpole_trajectory(
    x,
    tend,
    dt,
    segparms,
    pause=0.0,
    *,
    control=None,
    colors=None,
    width=None,
    height=None,
    show=True,
    repeat=True,
    hud=True,
    static=False,
):
    """Visualize a double cart-pole trajectory.

    Parameters
    ----------
    x : array-like, shape (6, n_frames)
        Rows contain ``x_c, theta_1, theta_2`` and their velocities.
    tend : float
        Duration of the optimized trajectory in seconds.
    dt : float
        Sampling interval between the supplied states.
    segparms : mapping
        Must contain ``l1`` and ``l2``; masses are used to scale the markers.
    pause : float, optional
        Number of seconds for which the final pose remains visible.
    control : array-like, optional
        Horizontal cart force. Supply one value per state or per interval.
    colors : mapping, optional
        Colour palette. Missing entries are taken from ``COLORS``.
    width, height : float, optional
        Figure size in inches. Set either width or height; the other dimension
        is inferred from the trajectory bounds. If neither is set, width is
        10 inches. The supplied dimension is respected exactly.
    show : bool, optional
        Call ``plt.show()`` before returning.
    repeat : bool, optional
        Repeat the animation after the last frame. Enabled by default.
    hud : bool, optional
        Show the time and force badge.
    static : bool, optional
        Render the first frame as a normal Matplotlib figure without creating
        a ``FuncAnimation``. Returns ``(figure, axes)`` in this mode.

    Returns
    -------
    matplotlib.animation.FuncAnimation or tuple
        Animation by default; ``(figure, axes)`` when ``static=True``.
    """

    states = np.asarray(x, dtype=float)
    duration = float(tend)
    parameters = segparms
    palette = COLORS | ({} if colors is None else dict(colors))
    if states.ndim != 2 or states.shape[0] < 3:
        raise ValueError("states must have shape (at least 3, n_frames)")
    if states.shape[1] < 2:
        raise ValueError("at least two trajectory frames are required")
    if dt <= 0:
        raise ValueError("dt must be positive")

    l1 = float(parameters["l1"])
    l2 = float(parameters["l2"])
    m1 = float(parameters.get("m1", 1.0))
    m2 = float(parameters.get("m2", 1.0))

    original_frame_count = states.shape[1]

    if control is None:
        force = np.zeros(original_frame_count)
        show_force = False
    else:
        force = np.asarray(control, dtype=float).reshape(-1)
        if force.size == original_frame_count - 1:
            force = np.append(force, force[-1])
        elif force.size != original_frame_count:
            raise ValueError(
                "control must contain one value per state or per interval"
            )
        show_force = True

    # Repeat the final state and force to create a deliberate end pause.
    extra_frames = max(0, int(round(float(pause) / dt)))
    if extra_frames:
        final_state = np.repeat(states[:, -1:], extra_frames, axis=1)
        states = np.hstack((states, final_state))
        force = np.append(force, np.repeat(force[-1], extra_frames))

    x_c = states[0]
    theta1 = states[1]
    theta2 = states[2]

    x1 = x_c + l1 * np.cos(theta1)
    y1 = l1 * np.sin(theta1)
    x2 = x1 + l2 * np.cos(theta2)
    y2 = y1 + l2 * np.sin(theta2)

    cart_width = 0.48
    cart_height = 0.22
    wheel_radius = 0.075
    rail_y = -cart_height / 2 - 2.1 * wheel_radius

    # Leave room for the force arrow when the cart reaches a trajectory bound.
    horizontal_margin = 1.25
    x_min = min(np.min(x_c), np.min(x1), np.min(x2)) - horizontal_margin
    x_max = max(np.max(x_c), np.max(x1), np.max(x2)) + horizontal_margin
    y_min = min(rail_y - 0.35, np.min(y1), np.min(y2)) - 0.25
    y_max = max(cart_height, np.max(y1), np.max(y2)) + 0.35

    data_width = max(x_max - x_min, 1e-9)
    data_height = max(y_max - y_min, 1e-9)
    data_aspect = data_width / data_height
    # Preserve enough canvas for the badges and axis labels even when
    # the physical trajectory itself is extremely narrow or wide.
    figure_aspect = float(np.clip(data_aspect, 1.35, 2.0))

    # Match the data limits to the canvas aspect ratio. Without this padding,
    # ``aspect="equal"`` shrinks the axes box when the horizontal trajectory
    # range is small, leaving unused space at either side of the figure.
    if data_aspect < figure_aspect:
        extra_width = data_height * figure_aspect - data_width
        x_min -= extra_width / 2
        x_max += extra_width / 2
    elif data_aspect > figure_aspect:
        extra_height = data_width / figure_aspect - data_height
        y_min -= extra_height / 2
        y_max += extra_height / 2

    if width is None and height is None:
        width = 10.0
    if width is not None and width <= 0:
        raise ValueError("width must be positive")
    if height is not None and height <= 0:
        raise ValueError("height must be positive")
    if width is None:
        width = float(height) * figure_aspect
    elif height is None:
        height = float(width) / figure_aspect

    fig, ax = plt.subplots(
        figsize=(float(width), float(height)),
        constrained_layout=True,
    )
    # Transparent canvas: it appears white in a normal window and blends into
    # the background when embedded in a webpage or exported with transparency.
    fig.patch.set_facecolor(palette["background"])
    fig.patch.set_alpha(0)
    ax.set_facecolor("none")
    ax.set_xlim(x_min, x_max)
    ax.set_ylim(y_min, y_max)
    ax.set_aspect("equal", adjustable="box")
    ax.set_xlabel(r"cart position $x_c$ [m]")
    ax.set_ylabel("height [m]")
    ax.tick_params(colors=palette["muted"], length=0)
    ax.grid(axis="y", color=palette["ink"], alpha=0.07, linewidth=0.8)
    for spine in ax.spines.values():
        spine.set_visible(False)

    # A single neutral rail provides a clean reference for cart motion.
    ax.axhline(rail_y, color=palette["rail"], linewidth=3.0, zorder=1)

    trace, = ax.plot(
        [], [],
        color=palette["trace"],
        linewidth=1.8,
        alpha=0.6,
        linestyle=(0, (2, 3)),
        zorder=2,
    )
    link1, = ax.plot(
        [], [],
        color=palette["link_1"],
        linewidth=5,
        solid_capstyle="round",
        zorder=5,
    )
    link2, = ax.plot(
        [], [],
        color=palette["link_2"],
        linewidth=5,
        solid_capstyle="round",
        zorder=5,
    )

    mass_scale = 190
    mass1 = ax.scatter(
        [], [], s=mass_scale * np.sqrt(max(m1, 0.05)),
        color=palette["mass_1"], edgecolor="white", linewidth=1.8, zorder=7,
    )
    mass2 = ax.scatter(
        [], [], s=mass_scale * np.sqrt(max(m2, 0.05)),
        color=palette["mass_2"], edgecolor="white", linewidth=1.8, zorder=7,
    )
    pivot = ax.scatter(
        [], [], s=38, color=palette["ink"], edgecolor="white",
        linewidth=1.0, zorder=8,
    )

    cart = FancyBboxPatch(
        (0, 0),
        cart_width,
        cart_height,
        boxstyle="round,pad=0.025,rounding_size=0.055",
        linewidth=2,
        edgecolor=palette["cart_edge"],
        facecolor=palette["cart"],
        zorder=6,
    )
    ax.add_patch(cart)

    wheel_left = Circle((0, 0), wheel_radius, color=palette["ink"], zorder=7)
    wheel_right = Circle((0, 0), wheel_radius, color=palette["ink"], zorder=7)
    ax.add_patch(wheel_left)
    ax.add_patch(wheel_right)

    force_arrow = FancyArrowPatch(
        (0, 0),
        (0, 0),
        arrowstyle="-|>",
        mutation_scale=16,
        linewidth=2.5,
        color=palette["force"],
        # Above the cart body, but below the black center-of-mass marker.
        zorder=7.5,
    )
    force_arrow.set_visible(show_force)
    ax.add_patch(force_arrow)

    time_text = ax.text(
        0.98, 0.90, "",
        transform=ax.transAxes,
        ha="right",
        va="top",
        fontsize=11,
        color=palette["ink"],
        bbox={
            "boxstyle": "round,pad=0.3",
            "facecolor": "white",
            "edgecolor": "none",
            "alpha": 0.72,
        },
    )
    force_text = ax.text(
        0.98, 0.82, "",
        transform=ax.transAxes,
        ha="right",
        va="top",
        fontsize=11,
        fontweight="bold",
        color=palette["force"],
    )
    force_text.set_visible(show_force)

    if not hud:
        for artist in (
            time_text,
            force_text,
            trace,
        ):
            artist.set_visible(False)

    max_force = max(float(np.max(np.abs(force))), np.finfo(float).eps)
    max_arrow_length = (x_max - x_min) / 3

    def update(frame):
        cart_x = x_c[frame]
        cart.set_bounds(
            cart_x - cart_width / 2,
            -cart_height / 2,
            cart_width,
            cart_height,
        )
        wheel_y = -cart_height / 2 - wheel_radius
        wheel_left.center = (cart_x - 0.145, wheel_y)
        wheel_right.center = (cart_x + 0.145, wheel_y)

        link1.set_data([cart_x, x1[frame]], [0, y1[frame]])
        link2.set_data([x1[frame], x2[frame]], [y1[frame], y2[frame]])
        mass1.set_offsets([[x1[frame], y1[frame]]])
        mass2.set_offsets([[x2[frame], y2[frame]]])
        pivot.set_offsets([[cart_x, 0]])
        trace.set_data(x2[: frame + 1], y2[: frame + 1])

        current_force = force[frame]
        arrow_length = max_arrow_length * current_force / max_force
        if arrow_length > 0:
            arrow_length = min(arrow_length, x_max - cart_x - 0.05)
        elif arrow_length < 0:
            arrow_length = -min(-arrow_length, cart_x - x_min - 0.05)
        # Apply the force at the marked cart center of mass (x_c, 0).
        arrow_y = 0.0
        force_arrow.set_visible(show_force and abs(current_force) > 1e-8)
        force_arrow.set_positions(
            (cart_x, arrow_y),
            (cart_x + arrow_length, arrow_y),
        )
        force_text.set_text(rf"$F_x = {current_force:+5.1f}\,\mathrm{{N}}$")

        elapsed = min(frame * dt, float(duration))
        time_text.set_text(f"t = {elapsed:4.2f} s")

        return (
            cart,
            wheel_left,
            wheel_right,
            link1,
            link2,
            mass1,
            mass2,
            pivot,
            trace,
            time_text,
            force_arrow,
            force_text,
        )

    if static:
        update(0)
        if show:
            plt.show()
        return fig, ax

    ani = animation.FuncAnimation(
        fig,
        update,
        frames=states.shape[1],
        interval=1000 * dt,
        blit=True,
        repeat=repeat,
        repeat_delay=700,
    )

    # Draw the first pose immediately, which also makes static notebook output
    # and screenshots useful before the animation starts.
    update(0)
    if show:
        plt.show()

    return ani


def save_animation_gif(ani, filename, *, fps, dpi=300):
    """Save a Matplotlib animation as a looping GIF on a white canvas."""

    if fps <= 0:
        raise ValueError("fps must be positive")
    fig = ani._fig
    
    fig.patch.set_alpha(1)
    fig.patch.set_facecolor("white")
    writer = animation.PillowWriter(
        fps=fps,
        metadata={"title": "Double cart-pole swing-up"},
    )
    ani.save(filename, writer=writer, dpi=dpi)