""" Functions for plotting and animation. """

import logging
import subprocess
import warnings
import itertools
from pathlib import Path
from typing import Tuple, List, Dict, Optional, Iterable
from tqdm import tqdm

import cv2
import numpy as np
from matplotlib import animation
from matplotlib.gridspec import GridSpec
import matplotlib.pyplot as plt
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas

from seqikpy.utils import load_file

# Ignore the warnings
warnings.filterwarnings("ignore")

# Change the logging level here
logging.basicConfig(format=" %(asctime)s - %(levelname)s- %(message)s")
# Get the logger of the module
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)



[docs]
def generate_color_map(cmap: str, n: int) -> List:
    """Generates a list of colors from a given colormap."""
    cmap = plt.get_cmap(cmap)
    colors = cmap(np.linspace(0, 1, n))
    return colors




[docs]
def get_video_writer(video_path, fps, output_shape):
    """Initialize and return a VideoWriter object."""
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    return cv2.VideoWriter(video_path, fourcc, fps, output_shape[::-1])




[docs]
def resize_frame(original_size, target_size):
    """
    Resize a frame to the target size
    while preserving the width length ratio.
    """
    if target_size[0] == -1 and target_size[1] == -1:
        new_size = original_size
    elif target_size[0] == -1:
        ratio = original_size[0] / original_size[1]
        new_size = (int(target_size[1] * ratio), target_size[1])
    elif target_size[1] == -1:
        ratio = original_size[1] / original_size[0]
        new_size = (target_size[0], int(target_size[0] * ratio))
    else:
        new_size = target_size

    if new_size[0] > original_size[0] and new_size[1] > original_size[1]:
        return original_size

    return new_size




[docs]
def resize_rgb_frame(frame, output_shape):
    """resize the frame and convert it to RGB"""
    resized = cv2.resize(frame, output_shape[::-1])
    return cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)




[docs]
def make_video(
    video_path: str,
    frame_generator: Iterable,
    fps: int,
    output_shape: Tuple = (-1, 2880),
    n_frames: int = -1,
):
    """Makes videos from a generator of images."""
    first_frame = next(frame_generator)
    frame_generator = itertools.chain([first_frame], frame_generator)

    adjusted_output_shape = resize_frame(first_frame.shape[:2], output_shape)
    video_writer = get_video_writer(video_path, fps, adjusted_output_shape)

    for frame_count, frame in tqdm(enumerate(frame_generator)):
        rgb_frame = resize_rgb_frame(frame, adjusted_output_shape)
        video_writer.write(rgb_frame)
        if 0 < n_frames == frame_count + 1:
            break

    video_writer.release()
    logging.info("Video is saved at %s", video_path)




[docs]
def video_frames_generator(
    video_path: Path,
    start_frame: int,
    end_frame: int,
    stim_lines: List[int],
    radius=30,
    center=(50, 50),
    color=(255, 0, 0),
):
    """Returns the frames as a generator in a given interval.
    Modifies the brightness and contrast of the images.

    Parameters
    ----------
    video_path : Path
        Video path.
    start_frame : int
        Starting frame.
    end_frame : int
        End frame.

    Yields
    ------
    Frame
        Generator containing all the frames in the specified interval.
    """
    cap = cv2.VideoCapture(str(video_path))
    cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)

    # Define the contrast and brightness values
    alpha = 1.15  # Contrast control
    beta = -5  # Brightness control

    for t in range(start_frame, end_frame):
        ret, frame = cap.read()
        adjusted_frame = cv2.convertScaleAbs(frame, alpha=alpha, beta=beta)
        # if stimulation, add a red dot
        if stim_lines[0] <= t <= stim_lines[-1]:
            adjusted_frame = cv2.circle(adjusted_frame, center, radius, color, -1)

        if not ret:
            break
        yield adjusted_frame

    cap.release()




[docs]
def get_plot_config(data_path: Path):
    """Get experimental conditions from the data path.
    Data path should look like:
    "/mnt/nas2/GO/7cam/220810_aJO-GAL4xUAS-CsChr/Fly001/001_RLF/behData/pose-3d"
    """
    assert data_path.parts[-1] == "pose-3d", "The data path should end with pose-3d"
    assert data_path.parts[-2] == "behData", "The data path should contain behData"

    plot_config = {}
    trial_type = data_path.parts[-3]
    if "_RLF_coxa" in trial_type:
        exp_type = "coxa"
        plot_config["plot_head"] = True
        plot_config["plot_right_leg"] = True
        plot_config["plot_left_leg"] = True
        plot_config["azim"] = 23
    elif "_RLF" in trial_type:
        exp_type = "RLF"
        plot_config["plot_head"] = True
        plot_config["plot_right_leg"] = False
        plot_config["plot_left_leg"] = False
    elif "_LF" in trial_type:
        exp_type = "LF"
        plot_config["plot_head"] = True
        plot_config["plot_right_leg"] = True
        plot_config["plot_left_leg"] = False
    elif "_RF" in trial_type:
        exp_type = "RF"
        plot_config["plot_head"] = True
        plot_config["plot_right_leg"] = False
        plot_config["plot_left_leg"] = True
    else:
        exp_type = "Beh"
        plot_config["plot_head"] = True
        plot_config["plot_right_leg"] = True
        plot_config["plot_left_leg"] = True

    return exp_type, plot_config




[docs]
def get_frames_from_video_ffmpeg(path):
    """Saves frames of a video using FFMPEG.

    Parameters
    ----------
    path : Path
        Video path.
        This path appended with a `_frames` folder will be used to save the frames.
    """
    write_path = path.parents[0] / str(path.name).replace(".mp4", "_frames")
    write_path.mkdir()
    cmd = [
        "ffmpeg",
        "-i",
        str(path),
        "-r",
        "1",
        str(write_path / "frame_%d.jpg"),
    ]
    subprocess.run(cmd, check=True)




[docs]
def load_grid_plot_data(data_path: Path) -> [Dict, Dict]:
    """Loads the set of data necessary for plotting the grid.

    Parameters
    ----------
    data_path : Path
        Data path where the pose3d and inverse kinematics are saved.

    Returns
    -------
    Tuple
        Returns joint angles (head and leg) and aligned pose as a tuple.
    """
    if (data_path / "body_joint_angles.pkl").is_file():
        joint_angles = load_file(data_path / "body_joint_angles.pkl")
    else:
        head_joint_angles = load_file(data_path / "head_joint_angles.pkl")
        leg_joint_angles = (
            load_file(data_path / "leg_joint_angles.pkl")
            if (data_path / "leg_joint_angles.pkl").is_file()
            else {}
        )
        joint_angles = {**head_joint_angles, **leg_joint_angles}
    aligned_pose = load_file(data_path / "pose3d_aligned.pkl")

    return joint_angles, aligned_pose




[docs]
def animate_3d_points(
    points3d: Dict[str, np.ndarray],
    # key_points: Dict[str, Tuple[np.ndarray, str]],
    export_path: Path,
    points3d_second: Dict[str, np.ndarray] = None,
    # key_points_second: Dict[str, Tuple[np.ndarray, str]] = None,
    fps: int = 100,
    frame_no: int = 1000,
    format_video: str = "mp4",
    elev: int = 10,
    azim: int = 90,
    title: str = "",
    marker_types: Dict[str, str] = None,
) -> None:
    """Makes an animation of 3D pose.
    This code is intended for animating the raw 3D pose and
    forward kinematics from seqikpy.


    Parameters
    ----------
    points3d : Dict[str, np.ndarray]
        Dictionary containing the 3D pose,
        usually this is the raw 3D pose.
    export_path : Path
        Path where the animation will be saved.
    points3d_second : Dict[str, np.ndarray], optional
        Dictionary containing the 3D pose
        usually this is the forward kinematics, by default None
    fps : int, optional
        Frames per second, by default 100
    frame_no : int, optional
        Maximum number of frames of the animation,
        by default 1000
    format_video : str, optional
        Video format, by default "mp4"
    elev : int, optional
        Elevation of the point of view, by default 10
    azim : int, optional
        Azimuth of the point of view, by default 90
    title : str, optional
        Title of the video, by default ""
    marker_types : Dict[str, str], optional
        Marker types for each key point, by default None
    """
    # Dark background
    plt.rcParams.update(
        {
            "axes.facecolor": "black",
            "axes.edgecolor": "black",
            "axes.labelcolor": "white",
            "xtick.color": "white",
            "ytick.color": "white",
            "grid.color": "lightgray",
            "figure.facecolor": "black",
            "figure.edgecolor": "black",
            "savefig.facecolor": "black",
            "savefig.edgecolor": "black",
        }
    )

    if marker_types is None:
        marker_types = {
            "R_head": "o",
            "L_head": "o",
            "Neck": "x",
        }

    fig = plt.figure()
    ax3d = fig.add_subplot(projection="3d")
    ax3d.view_init(azim=azim, elev=elev)
    # First remove fill
    ax3d.xaxis.pane.fill = False
    ax3d.yaxis.pane.fill = False
    ax3d.zaxis.pane.fill = True

    # Now set color to white (or whatever is "invisible")
    ax3d.xaxis.pane.set_edgecolor("black")
    ax3d.yaxis.pane.set_edgecolor("black")
    ax3d.zaxis.pane.set_edgecolor("black")

    color_map_right = generate_color_map(cmap="Reds", n=len(points3d))
    color_map_left = generate_color_map(cmap="Blues", n=len(points3d))
    color_map_scatter = generate_color_map(cmap="RdBu", n=len(points3d))
    i, j, k = 0, 0, 0
    line_data = []
    line_data_second = []

    for kp, points3d_array in points3d.items():
        order = points3d_array.shape[1]

        if order > 3:
            if "L" in kp:
                color = color_map_left[j]
                j += 1
            else:
                color = color_map_right[k]
                k += 1
            line_data.append(
                ax3d.plot(
                    points3d_array[0, :, 0],
                    points3d_array[0, :, 1],
                    points3d_array[0, :, 2],
                    label=kp,
                    linestyle="solid",
                    linewidth=4,
                    color=color,
                    alpha=0.85,
                )[0]
            )
        else:
            line_data.append(
                ax3d.plot(
                    points3d_array[0, :, 0],
                    points3d_array[0, :, 1],
                    points3d_array[0, :, 2],
                    lw=2.5,
                    label=kp,
                    marker=marker_types[kp],
                    markersize=9,
                    color=color_map_scatter[i],
                    alpha=0.7,
                )[0]
            )
            i += 1

    i, j, k = 0, 0, 0

    if points3d_second is not None:

        for kp, points3d_second_array in points3d_second.items():
            order = points3d_second_array.shape[1]

            if order > 4:
                if "L" in kp:
                    color = color_map_left[j]
                    j += 1
                else:
                    color = color_map_right[k]
                    k += 1
                line_data_second.append(
                    ax3d.plot(
                        points3d_second_array[0, :, 0],
                        points3d_second_array[0, :, 1],
                        points3d_second_array[0, :, 2],
                        label=kp,
                        linestyle="--",
                        linewidth=4,
                        color=color,
                        alpha=0.85,
                    )[0]
                )
            else:
                line_data_second.append(
                    ax3d.plot(
                        points3d_second_array[0, :, 0],
                        points3d_second_array[0, :, 1],
                        points3d_second_array[0, :, 2],
                        lw=2.5,
                        label=kp,
                        marker=marker_types[kp],
                        markersize=9,
                        color=color_map_scatter[j],
                        alpha=0.7,
                    )[0]
                )

    # Setting the axes properties
    # ax3d.set_xlim((-0.0, 1.5))
    # ax3d.set_ylim((-1., 1.))
    # ax3d.set_zlim((0.2, 1.8))

    ax3d.set_xticks([])
    ax3d.set_yticks([])
    ax3d.set_zticks([])

    ax3d.tick_params(axis="x", color="black")
    ax3d.tick_params(axis="y", color="black")
    ax3d.tick_params(axis="z", color="black")
    # ax3d.set_axis_off("z")

    ax3d.set_xticklabels([])
    ax3d.set_yticklabels([])
    ax3d.set_zticklabels([])

    # ax3d.set_xlabel("x")
    # ax3d.set_ylabel("y")
    # ax3d.set_zlabel("z")
    ax3d.set_title(title, loc="center")
    # ax3d.legend(bbox_to_anchor=(1.2, 0.9), frameon=False)

    def update(frame, lines, points3d, lines_second, points3d_second):
        i = 0
        for kp, points3d_array in points3d.items():
            lines[i].set_data(points3d_array[frame, :, 0], points3d_array[frame, :, 1])
            lines[i].set_3d_properties(points3d_array[frame, :, 2])
            i += 1
        if lines_second:
            j = 0
            for kp, points3d_second_array in points3d_second.items():
                lines_second[j].set_data(
                    points3d_second_array[frame, :, 0],
                    points3d_second_array[frame, :, 1],
                )
                lines_second[j].set_3d_properties(points3d_second_array[frame, :, 2])
                j += 1

    # Creating the Animation object
    line_ani = animation.FuncAnimation(
        fig,
        update,
        frame_no,
        fargs=(line_data, points3d, line_data_second, points3d_second),
        interval=10,
        blit=False,
    )
    logger.info("Making animation...")
    export_path = str(export_path)
    export_path += (
        f".{format_video}" if not export_path.endswith((".mp4", ".avi", ".mov")) else ""
    )
    line_ani.save(export_path, fps=fps, dpi=300)
    logger.info(f"Animation is saved at {export_path}")




[docs]
def plot_3d_points(
    ax3d, points3d, export_path=None, t=0, marker_types=None, line_style="solid"
):
    """Plots 3D points at time t."""

    if marker_types is None:
        marker_types = {
            "R_head": "o",
            "L_head": "o",
            "Neck": "x",
        }

    color_map_right = generate_color_map(
        cmap="Reds", n=len([kp for kp in points3d if "R" in kp]) + 1
    )
    color_map_left = generate_color_map(
        cmap="Blues", n=len([kp for kp in points3d if "L" in kp]) + 1
    )

    i, j = 1, 1

    for kp, points3d_array in points3d.items():
        order = points3d_array.shape[1]

        if "R" in kp:
            color = color_map_right[i]
            i += 1
        elif "L" in kp:
            color = color_map_left[j]
            j += 1
        else:
            color = "lightgrey"

        if order > 3:
            ax3d.plot(
                points3d_array[t, :, 0],
                points3d_array[t, :, 1],
                points3d_array[t, :, 2],
                label=kp,
                linestyle=line_style,
                linewidth=1.7,
                color=color,
            )
        else:
            ax3d.plot(
                points3d_array[t, :, 0],
                points3d_array[t, :, 1],
                points3d_array[t, :, 2],
                label=kp,
                marker=marker_types[kp],
                markersize=4.5,
                color=color,
            )

    if export_path is not None:
        plt.savefig(export_path, bbox_inches="tight")




[docs]
def plot_trailing_kp(
    ax3d,
    points3d,
    segments_to_plot,
    export_path=None,
    t=0,
    trail=5,
    marker_type="x",
):
    """Plots the traces of key points from t-trail to t."""

    color_map_right = generate_color_map(cmap="Reds", n=len(points3d) + 1)
    color_map_left = generate_color_map(cmap="Blues", n=len(points3d) + 1)

    i, j = 1, 1
    for kp, ind in segments_to_plot.items():

        if "R" in kp:
            color = color_map_right[i]
            i += 1
        elif "L" in kp:
            color = color_map_left[j]
            j += 1
        else:
            color = "grey"

        ax3d.scatter(
            points3d[kp][max(0, t - trail) : t, ind, 0],
            points3d[kp][max(0, t - trail) : t, ind, 1],
            points3d[kp][max(0, t - trail) : t, ind, 2],
            label=kp,
            marker=marker_type,
            #             markersize=9,
            color=color,
        )

    if export_path is not None:
        plt.savefig(export_path, bbox_inches="tight")




[docs]
def plot_joint_angle(
    ax: plt.Axes,
    kinematics_data: Dict[str, np.ndarray],
    angles_to_plot: List,
    degrees: bool = True,
    until_t: int = -1,
    stim_lines: List[int] = None,
    show_legend: bool = True,
    export_path: Path = None,
):
    """Plot joint angles from a given kinematics data.

    Parameters
    ----------
    ax : plt.Axes
        Axis where the plot will be displayed.
    kinematics_data : Dict[str, np.ndarray]
        Dictionary containing the kienmatics, pose and angle
    angles_to_plot : List
        Angles to plot. Exact column name should be given.
    degrees : bool, optional
        Convert to degrees, by default True
    until_t : int, optional
        Plots the angles until t^th frame, by default -1
    stim_lines : List[int], optional
        Plots vertical lines in indicated locations,
        by default None
    show_legend : bool, optional
        Shows legend, by default True
    export_path : Path, optional
        Path where the plot will be saved, by default None
    """
    colors = generate_color_map(cmap="Set2", n=len(angles_to_plot))

    for i, joint_name in enumerate(angles_to_plot):
        joint_angles = kinematics_data[joint_name]

        label = " ".join((joint_name.split("_")[-2], joint_name.split("_")[-1]))

        if label in [
            "pitch R",
            "pitch L",
            "yaw R",
            "yaw L",
            "roll R",
            "roll L",
        ]:
            label = "ant. " + label

        convert2deg = 180 / np.pi if degrees else 1

        ax.plot(
            np.array(joint_angles[:until_t]) * convert2deg,
            "o-",
            ms=4.5,
            markevery=[-1],
            label=label,
            color=colors[i],
        )

    if stim_lines is not None:
        ax.vlines(stim_lines, -200, 200, "red", lw=0.5)

    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    if show_legend:
        ax.legend(bbox_to_anchor=(1.2, 1), frameon=False, borderaxespad=0.0)

    if export_path is not None:
        plt.savefig(export_path, bbox_inches="tight")




[docs]
def plot_grid(
    img_front: np.ndarray,
    img_side: np.ndarray,
    aligned_pose: Dict[str, np.ndarray],
    joint_angles: Dict[str, np.ndarray],
    leg_angles_to_plot: List[str],
    head_angles_to_plot: List[str],
    t: int,
    t_start: int,
    t_end: int,
    t_interval: int = 20,
    fps: int = 100,
    trail: int = 30,
    key_points_to_trail: Optional[Dict[str, Iterable]] = None,
    marker_types_3d: Optional[Dict[str, str]] = None,
    marker_trail: Optional[str] = None,
    stim_lines: Optional[List[int]] = None,
    export_path: Optional[Path] = None,
    **kwargs,
):
    """
    Plots an instance of the animal recording, 3D pose,
    head and leg joint angles in a grid layout. This code
    is intended to use in a for loop to plot and save all
    the frames from `t_start` to `t_end` to make a video afterwards.

    NOTE: This function is intended to plot leg and head joint angles together,
    it will not work if any of these data is missing.

    Parameters
    ----------
    img_front : Path
        Image of the fly at frame t on the front camera.
    img_side : Path
        Image of the fly at frame t on the side camera.
    aligned_pose : Dict[str, np.ndarray]
        Aligned 3D pose.
    joint_angles : Dict[str, np.ndarray]
        Joint angles.
    leg_angles_to_plot : List[str]
        List containing leg joint angle names without the side.
        Examples

        >>> leg_joint_angles = [
                "ThC_yaw",
                "ThC_pitch",
                "ThC_roll",
                "CTr_pitch",
                "CTr_roll",
                "FTi_pitch",
                "TiTa_pitch",
                ]
    head_angles_to_plot : List[str]
        List containing exact names of head joint angle names.
    t : int
        Frame number t.
    t_start : int
        Start of the time series, i.e., joint angles.
    t_end : int
        End of the time series, i.e., joint angles.
    t_interval : int, optional
        Interval of frame numbers in between x ticks, by default 20
    fps : int, optional
        Frames per second, by default 100
    trail : int, optional
        Number of previous frames where the key point will be visible,
        by default 30
    marker_types_3d : Dict[str, Tuple[np.ndarray, str]]
        Dictionary mapping key points names to marker styles.
    marker_trail : Dict[str, Tuple[np.ndarray, str]]
        Dictionary mapping key points names to their indices
        and line styles for trailing key points.
    stim_lines : List[int], optional
        Stimulation indicators, by default None
    export_path : Path, optional
        Path where the plot will be saved, by default None

    Returns
    -------
    Fig
        Figure.
    """
    plot_right_leg = kwargs.pop("plot_right_leg", True)
    plot_left_leg = kwargs.pop("plot_left_leg", True)
    plot_head = kwargs.pop("plot_head", True)
    azim = kwargs.pop("azim", 7)

    assert (
        t_start <= t <= t_end
    ), "t_start should be smaller than t_end, t should be in between"
    # import pylustrator
    # pylustrator.start()

    plt.style.use("dark_background")

    fig = plt.figure(figsize=(14, 6), dpi=120)

    gs = GridSpec(3, 4, figure=fig)
    # 7cam recording
    ax_img_side = fig.add_subplot(gs[0, :2])
    ax_img_front = fig.add_subplot(gs[1, :2])
    # 3D pose
    ax1 = fig.add_subplot(gs[2, :2], projection="3d")
    # head, right leg, left leg joint angles
    ax2 = fig.add_subplot(gs[0, 2:])
    ax3 = fig.add_subplot(gs[1, 2:])
    ax4 = fig.add_subplot(gs[2, 2:])

    # load the image

    # img = cv2.imread(str(img_path / f"frame_{t}.jpg"), 0)
    ax_img_side.imshow(img_side, vmin=0, vmax=255, cmap="gray")
    ax_img_front.imshow(img_front, vmin=0, vmax=255, cmap="gray")

    plot_3d_points(ax1, aligned_pose, marker_types=marker_types_3d, t=t)
    if key_points_to_trail is not None:
        plot_trailing_kp(
            ax1,
            aligned_pose,
            key_points_to_trail,
            marker_type=marker_trail,
            trail=trail,
            t=t,
        )
    if plot_head:
        plot_joint_angle(
            ax2,
            joint_angles,
            angles_to_plot=head_angles_to_plot,
            until_t=t,
            stim_lines=stim_lines,
        )
    if plot_left_leg:
        plot_joint_angle(
            ax3,
            joint_angles,
            angles_to_plot=[f"Angle_LF_{ja}" for ja in leg_angles_to_plot],
            until_t=t,
            stim_lines=stim_lines,
        )
    if plot_right_leg:
        plot_joint_angle(
            ax4,
            joint_angles,
            angles_to_plot=[f"Angle_RF_{ja}" for ja in leg_angles_to_plot],
            until_t=t,
            show_legend=False,
            stim_lines=stim_lines,
        )

    ax_img_side.axis("off")
    ax_img_front.axis("off")
    # ax1 properties
    ax1.view_init(azim=azim, elev=10)
    ax1.set_xlim3d([-0.45, 0.45])
    ax1.set_ylim3d([-0.6, 0.6])
    ax1.set_zlim3d([0.42, 1.1])
    ax1.set_xticks([])
    ax1.set_yticks([])
    ax1.set_zticks([])
    ax1.axis("off")

    # ax2 properties
    ax2.set_xlim((t_start, t_end))
    ax2.spines["bottom"].set_visible(False)
    ax2.spines["top"].set_visible(False)
    ax2.spines["right"].set_visible(False)
    ax2.set_xlim((t_start, t_end))
    ax2.set_ylim((-90, 70))
    ax2.set_xticks([])
    ax2.set_yticks(ticks=[-90, 0, 70])
    ax2.set_yticklabels(labels=[-90, 0, 70])

    # ax3 properties
    ax3.set_xlim((t_start, t_end))
    ax3.spines["bottom"].set_visible(False)
    ax3.spines["top"].set_visible(False)
    ax3.spines["right"].set_visible(False)
    ax3.set_ylim((-160, 160))
    ax3.set_xticks([])
    ax3.set_yticks(ticks=[-160, 0, 160])
    ax3.set_yticklabels(labels=[-160, 0, 160])

    # ax4 properties
    ax4.set_xlim((t_start, t_end))
    ax4.set_ylim((-160, 160))
    ax4.spines["top"].set_visible(False)
    ax4.spines["right"].set_visible(False)
    ax4.set_yticks(ticks=[-160, 0, 160])
    ax4.set_yticklabels(labels=[-160, 0, 160])

    ax4.set_xticks(ticks=np.arange(t_start, t_end + t_interval, t_interval))
    ax4.set_xticklabels(labels=np.arange(t_start, t_end + t_interval, t_interval) / fps)
    ax4.set_xlabel("Time (s)")

    # #% start: automatic generated code from pylustrator
    fig.set_size_inches(22.710000 / 2.54, 11.430000 / 2.54, forward=True)
    fig.text(
        0.3865,
        0.9184,
        "Head and antennae joint angles (deg)",
        transform=fig.transFigure,
    )
    fig.text(
        0.3865,
        0.6346,
        "Left front leg joint angles (deg)",
        transform=fig.transFigure,
    )  # id=fig.texts[0].new
    fig.text(
        0.3865,
        0.3502,
        "Right front leg joint angles (deg)",
        transform=fig.transFigure,
    )  # id=fig.texts[1].new

    # #% start: automatic generated code from pylustrator
    fig.set_size_inches(21.890000 / 2.54, 11.420000 / 2.54, forward=True)
    fig.axes[0].set_position([0.049668, 0.665633, 0.277935, 0.266469])
    fig.axes[1].set_position([0.049668, 0.382601, 0.277935, 0.266469])
    fig.axes[2].set(position=[0.1194, 0.06658, 0.1318, 0.2619])
    fig.axes[2].set_position([0.123600, 0.069459, 0.136149, 0.261065])
    fig.axes[3].set_position([0.400791, 0.677390, 0.383414, 0.225807])
    fig.axes[4].set_position([0.400791, 0.396851, 0.383414, 0.225807])
    fig.axes[5].set_position([0.400791, 0.112723, 0.383414, 0.225807])
    fig.texts[0].set_position([0.399756, 0.918650])
    fig.texts[1].set_position([0.399756, 0.635718])
    fig.texts[2].set_position([0.399756, 0.352188])

    fig.axes[3].legend(loc=(1.068, -0.1324), frameon=False)
    fig.axes[4].legend(loc=(1.064, -0.569), frameon=False)
    # % end: automatic generated code from pylustrator

    if export_path is not None:
        fig.savefig(export_path, bbox_inches="tight")
        print(f"Figure saved at {str(export_path)}")

    return fig




[docs]
def plot_grid_generator(
    fly_frames_front: np.ndarray,
    fly_frames_side: np.ndarray,
    aligned_pose: Dict[str, np.ndarray],
    joint_angles: Dict[str, np.ndarray],
    leg_angles_to_plot: List[str],
    head_angles_to_plot: List[str],
    t_start: int,
    t_end: int,
    t_interval: int = 20,
    fps: int = 100,
    trail: int = 30,
    key_points_to_trail: Optional[List[str]] = None,
    marker_types_3d: Dict[str, str] = None,
    marker_trail: Optional[str] = "x",
    stim_lines: List[int] = None,
    export_path: Path = None,
    **kwargs,
):
    """Generator for plotting grid."""

    for t, (fly_img_front, fly_img_side) in enumerate(
        zip(fly_frames_front, fly_frames_side)
    ):
        fig = plot_grid(
            img_front=fly_img_front,
            img_side=fly_img_side,
            aligned_pose=aligned_pose,
            joint_angles=joint_angles,
            leg_angles_to_plot=leg_angles_to_plot,
            head_angles_to_plot=head_angles_to_plot,
            marker_types_3d=marker_types_3d,
            marker_trail=marker_trail,
            key_points_to_trail=key_points_to_trail,
            t=t + t_start,
            t_start=t_start,
            t_end=t_end,
            t_interval=t_interval,
            fps=fps,
            trail=trail,
            stim_lines=stim_lines,
            export_path=export_path,
            **kwargs,
        )
        plt.close(fig)
        yield fig_to_array(fig)




[docs]
def fig_to_array(fig):
    """Converts a matplotlib figure into an array."""
    canvas = FigureCanvas(fig)
    canvas.draw()
    data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
    data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
    return data