"""Make a perspective flythrough temperature movie. [AI-Codex]"""
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, FFMpegWriter
from matplotlib.colors import LogNorm
from matplotlib import patheffects
import numpy as np
import zarr

import thorplot as tp

store = zarr.open("flythrough/output.zr", mode="r")
projections = store["projections"]
frames = sorted(
    (name for name in projections.keys() if name.startswith("temperature_")),
    key=lambda name: int(name.rsplit("_", 1)[1]),
)
if not frames:
    raise ValueError("No temperature frames found; run flythrough.yaml first.")
projection = projections[frames[0]]
# THOR stores camera axes as (horizontal, vertical); imshow uses (row, column).
temperature = projection[:].T
parameters = projection.attrs["Parameters"]
# Perspective rays have an angular field of view, not a fixed physical scale.
fov = np.deg2rad(parameters["fov_up"])
aspect = temperature.shape[1] / temperature.shape[0]
extent = [-aspect, aspect, -1, 1]

# Keep the poster and movie canvas identical; all labels sit inside the map.
with tp.style(), plt.rc_context({"savefig.bbox": None, "savefig.pad_inches": 0}):
    # Match the canvas to the map so videos have no outer margins.
    fig = plt.figure(figsize=(tp.DOC_FULL_WIDTH, tp.DOC_FULL_WIDTH / aspect))
    ax = fig.add_axes([0, 0, 1, 1])
    im = ax.imshow(temperature, origin="lower", extent=extent,
                   norm=LogNorm(vmin=1e4, vmax=1e6), cmap="inferno",
                   interpolation="nearest")
    ax.set_xticks([])
    ax.set_yticks([])
    for spine in ax.spines.values():
        spine.set_visible(False)

    label = ax.text(0.06, 0.08, f"Vertical field of view: {np.rad2deg(fov):g}°",
                    transform=ax.transAxes, color="white", ha="left", va="bottom",
                    fontsize=11)
    label.set_path_effects([patheffects.withStroke(linewidth=2, foreground="black")])

    # Inset colorbar adapted from thor_rendering.postprocess.projections.add_colorbar.
    cax = ax.inset_axes([0.59, 0.12, 0.35, 0.025])
    cbar = fig.colorbar(im, cax=cax, orientation="horizontal", ticks=[1e4, 1e5, 1e6])
    outline = [patheffects.withStroke(linewidth=1.5, foreground="black")]
    cbar.set_label("Temperature [K]", color="white", fontsize=11, path_effects=outline)
    cbar.ax.xaxis.set_label_position("top")
    cbar.minorticks_off()
    cbar.ax.tick_params(colors="white", labelsize=10, length=3)
    cbar.outline.set_edgecolor("white")
    cbar.outline.set_path_effects(outline)
    for tick_label in cbar.ax.get_xticklabels():
        tick_label.set_path_effects(outline)


    def update(frame):
        """Update pixels with a fixed field of view and color scale. [AI-Codex]"""
        im.set_data(projections[frame][:].T)
        return (im,)


    tp.save_doc_figure(fig, "flythrough-poster", fmt="png", output_dir=".", dpi=150)

    movie = FuncAnimation(fig, update, frames=frames, interval=1000 / 24,
                          cache_frame_data=False)
    movie.save("flythrough.mp4", writer=FFMpegWriter(
        fps=24, codec="libx264", extra_args=["-pix_fmt", "yuv420p", "-movflags", "+faststart"]
    ), dpi=150)
    plt.close(fig)
    print("Saved flythrough.mp4 and flythrough-poster.png")