# Plot the density-weighted temperature projection. [AI-Codex]
import argparse

import thorplot as tp
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm
from matplotlib import patheffects
import numpy as np
import zarr

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--dark", action="store_true")
args = parser.parse_args()
with tp.style(dark=args.dark):

    store = zarr.open("temperature.zr", mode="r")
    projection = store["projections/temperature"]
    # 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
    extent = [-width_kpc[0] / 2, width_kpc[0] / 2,
              -width_kpc[1] / 2, width_kpc[1] / 2]

    fig, ax = plt.subplots(figsize=(tp.DOC_FULL_WIDTH, 7.625), layout="constrained")
    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)
    stem = "temperature" + ("_dark" if args.dark else "")
    tp.save_doc_figure(fig, stem, fmt="png", output_dir=".", dpi=200)
    plt.close(fig)