"""Make a rotating temperature movie with the projection tutorial styling. [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("rotation/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 rotation.yaml first.")
projection = projections[frames[0]]
# THOR stores camera axes as (horizontal, vertical); imshow uses (row, column).
temperature = projection[:].T
parameters = projection.attrs["Parameters"]
width_kpc = np.asarray(parameters["widths"]) * parameters["length_unit_cm"] / 3.0856775814913673e21
aspect = width_kpc[0] / width_kpc[1]
extent = [-width_kpc[0] / 2, width_kpc[0] / 2,
          -width_kpc[1] / 2, width_kpc[1] / 2]

# 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)

    # The extent is in physical kpc, so the bar length uses the same units.
    bar_kpc = 20
    bar_x = extent[0] + 0.06 * width_kpc[0]
    bar_y = extent[2] + 0.07 * width_kpc[1]
    bar, = ax.plot([bar_x, bar_x + bar_kpc], [bar_y, bar_y], color="white",
                   linewidth=3.5, solid_capstyle="projecting")
    bar.set_path_effects([
        patheffects.Stroke(linewidth=5, foreground="black", capstyle="projecting"),
        patheffects.Normal(),
    ])
    label = ax.text(bar_x + bar_kpc / 2, bar_y + 0.025 * width_kpc[1],
                    f"{bar_kpc} kpc", color="white", ha="center", va="bottom",
                    fontweight="bold")
    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 physical extent and color scale. [AI-Codex]"""
        im.set_data(projections[frame][:].T)
        return (im,)


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

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