# AUTOGENERATED! DO NOT EDIT! File to edit: ../notebooks/api/04_visualization.ipynb.

# %% ../notebooks/api/04_visualization.ipynb 3
from __future__ import annotations

import tempfile

import imageio.v3 as iio
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from tqdm import tqdm

# %% auto 0
__all__ = ['plot_drr', 'animate']

# %% ../notebooks/api/04_visualization.ipynb 4
import torch


def plot_drr(
    img: torch.Tensor,
    title: str | None = None,
    ticks: bool | None = True,
    axs: matplotlib.axes._axes.Axes | None = None,
):
    """Plot an image generated by a DRR module."""

    if axs is None:
        fig, axs = plt.subplots(ncols=len(img), figsize=(10, 5))
    if len(img) == 1:
        axs = [axs]
    for img, ax in zip(img, axs):
        ax.imshow(img.squeeze().cpu().detach(), cmap="gray")
        _, height, width = img.shape
        ax.xaxis.tick_top()
        ax.set(
            title=title,
            xticks=[0, width - 1],
            xticklabels=[1, width],
            yticks=[0, height - 1],
            yticklabels=[1, height],
        )
        if ticks is False:
            ax.set_xticks([])
            ax.set_yticks([])
    return axs

# %% ../notebooks/api/04_visualization.ipynb 5
import pathlib

import pandas

from .drr import DRR


def animate(
    out: str | pathlib.Path,  # Savepath
    df: pandas.DataFrame,
    drr: DRR,
    ground_truth: torch.Tensor | None = None,
    verbose: bool = True,
    device="cpu",
    **kwargs,  # To pass to imageio.v3.imwrite
):
    """Animate the optimization of a DRR."""
    # Make the axes
    if ground_truth is None:

        def make_fig():
            fig, ax_opt = plt.subplots(
                figsize=(3, 3),
                constrained_layout=True,
            )
            return fig, ax_opt

    else:

        def make_fig(ground_truth):
            fig, (ax_fix, ax_opt) = plt.subplots(
                ncols=2,
                figsize=(6, 3),
                constrained_layout=True,
            )
            plot_drr(ground_truth, axs=ax_fix)
            ax_fix.set(xlabel="Fixed DRR")
            return fig, ax_opt

    # Compute DRRs, plot, and save to temporary folder
    if verbose:
        itr = tqdm(df.iterrows(), desc="Precomputing DRRs", total=len(df))
    else:
        itr = df.iterrows()

    with tempfile.TemporaryDirectory() as tmpdir:
        idxs = []
        for idx, row in itr:
            fig, ax_opt = make_fig() if ground_truth is None else make_fig(ground_truth)
            params = row[["theta", "phi", "gamma", "bx", "by", "bz"]].values
            rotations = (
                torch.tensor(row[["theta", "phi", "gamma"]].values)
                .unsqueeze(0)
                .to(device)
            )
            translations = (
                torch.tensor(row[["bx", "by", "bz"]].values).unsqueeze(0).to(device)
            )
            itr = drr(rotations, translations)
            _ = plot_drr(itr, axs=ax_opt)
            ax_opt.set(xlabel="Moving DRR")
            fig.savefig(f"{tmpdir}/{idx}.png")
            plt.close(fig)
            idxs.append(idx)
        frames = np.stack(
            [iio.imread(f"{tmpdir}/{idx}.png") for idx in idxs],
            axis=0,
        )

    # Make the animation
    return iio.imwrite(out, frames, **kwargs)
