import numpy as np
import g3read as g3
import matplotlib.pyplot as plt
import matplotlib as mpl
mpl.use("Agg")  # batch / no GUI


def show_galaxy(nsnaps=101, dx=30.0):
    # IDL loadct,5 -> just pick a reasonable fixed palette here
    # (IDL color indices don't map 1:1 to matplotlib, so we choose distinct colors.)
    col_disk  = "C1"
    col_bulge = "C3"
    col_gas   = "C2"
    col_star  = "C0"

    for i in range(nsnaps):
        snr = f"{i:03d}"
        name = f"snap_{snr}"

        f = g3.GadgetFile(name)

        # Required header-reading pattern
        try:
            h = f.read_new("HEAD", 0)
        except Exception:
            h = getattr(f, "header", {})

        # Determine if stars (type 4) exist
        if isinstance(h, dict):
            npart = h.get("npart", None)
            nstars = int(npart[4]) if npart is not None else 0
        else:
            npart = getattr(h, "npart", getattr(getattr(f, "header", {}), "npart", [0, 0, 0, 0, 0, 0]))
            nstars = int(npart[4])

        # Read positions by particle type
        # (g3read: type is the 2nd argument; older scripts used part=...)
        x_gas   = f.read_new("POS ", 0)  # gas
        x_disk  = f.read_new("POS ", 2)  # disk
        x_bulge = f.read_new("POS ", 3)  # bulge
        x_stars = None
        if nstars > 0:
            x_stars = f.read_new("POS ", 4)  # stars

        # --- Figure: 2 panels (x-y and x-z) ---
        fig, axes = plt.subplots(1, 2, figsize=(10, 5))
        plt.subplots_adjust(wspace=0.25)

        # Common axis formatting
        for ax in axes:
            ax.set_xlim(-dx, dx)
            ax.set_ylim(-dx, dx)
            ax.set_aspect("equal", adjustable="box")
            ax.tick_params(labelsize=10)
            ax.set_title(snr, fontsize=12)

        # Panel 1: x-y
        ax = axes[0]
        ax.set_xlabel("x [kpc]", fontsize=12)
        ax.set_ylabel("y [kpc]", fontsize=12)

        ax.plot(x_disk[:, 0],  x_disk[:, 1],  "o", ms=1.0, color=col_disk)
        ax.plot(x_bulge[:, 0], x_bulge[:, 1], "o", ms=1.0, color=col_bulge)
        ax.plot(x_gas[:, 0],   x_gas[:, 1],   "o", ms=1.0, color=col_gas)
        if x_stars is not None:
            ax.plot(x_stars[:, 0], x_stars[:, 1], "o", ms=1.0, color=col_star)

        # Panel 2: x-z
        ax = axes[1]
        ax.set_xlabel("x [kpc]", fontsize=12)
        ax.set_ylabel("z [kpc]", fontsize=12)

        ax.plot(x_disk[:, 0],  x_disk[:, 2],  "o", ms=1.0, color=col_disk)
        ax.plot(x_bulge[:, 0], x_bulge[:, 2], "o", ms=1.0, color=col_bulge)
        ax.plot(x_gas[:, 0],   x_gas[:, 2],   "o", ms=1.0, color=col_gas)
        if x_stars is not None:
            ax.plot(x_stars[:, 0], x_stars[:, 2], "o", ms=1.0, color=col_star)

        fig.tight_layout()
        fig.savefig(f"frame_{snr}.png", dpi=150, bbox_inches="tight")
        plt.close(fig)
        print(f"Saved frame_{snr}.png")


if __name__ == "__main__":
    # 0..100 inclusive -> 101 frames
    show_galaxy(nsnaps=101, dx=30.0)