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_wave(nsnaps=101, L=20.0):
    gamma  = 5.0 / 3.0
    rho_ic = 0.125 / L
    P_ic   = 1.0   / L
    csnd   = np.sqrt(gamma * P_ic / rho_ic)

    for i in range(nsnaps):
        snr  = f"{i:03d}"
        name = f"snap_{snr}"
        f    = g3.GadgetFile(name)
        try:
            h = f.read_new("HEAD", 0)
        except Exception:
            h = getattr(f, "header", {})
        x    = f.read_new("POS ", 0)   # (N,3)
        v    = f.read_new("VEL ", 0)   # (N,3)
        rho  = f.read_new("RHO ", 0)   # (N,)
        u    = f.read_new("U   ", 0)   # (N,)

        # pressure
        p = (gamma - 1.0) * u * rho

        # get time from header
        if isinstance(h, dict) and "time" in h:
            t = float(h["time"])
        else:
            t = float(getattr(h, "time", 0.0))

        # compute marker positions
        x0 = 1.0 + csnd * t
        while x0 > L:
            x0 -= L
        x1 = L - x0

        # set up 4 stacked panels
        fig, axes = plt.subplots(4, 1, figsize=(8, 10), sharex=True)
        plt.subplots_adjust(hspace=0.25)

        # common x-range
        xmin, xmax = 0.0, L

        # panel 1: pressure
        ax = axes[0]
        ax.plot(x[:, 0], p / P_ic, 'o', ms=2, color='k')
        ax.set_xlim(xmin, xmax)
        ax.set_ylim(1.0 - 0.03, 1.0 + 0.03)
        ax.set_ylabel('P', fontsize=12)
        ax.set_title('x-waves', fontsize=13)
        ax.axvline(x0, color='0.6', linewidth=1)
        ax.axvline(x1, color='0.6', linewidth=1)

        # panel 2: density
        ax = axes[1]
        ax.plot(x[:, 0], rho / rho_ic, 'o', ms=2, color='k')
        ax.set_xlim(xmin, xmax)
        ax.set_ylim(1.0 - 0.02, 1.0 + 0.02)
        ax.set_ylabel('rho', fontsize=12)
        ax.axvline(x0, color='0.6', linewidth=1)
        ax.axvline(x1, color='0.6', linewidth=1)

        # panel 3: vx
        ax = axes[2]
        ax.plot(x[:, 0], v[:, 0], 'o', ms=2, color='k')
        ax.set_xlim(xmin, xmax)
        ax.set_ylim(-0.06, 0.06)   # IDL: [-0.06e, 0.06] ~ [-0.06, 0.06]
        ax.set_ylabel('vx', fontsize=12)
        ax.axvline(x0, color='0.6', linewidth=1)
        ax.axvline(x1, color='0.6', linewidth=1)

        # panel 4: vy
        ax = axes[3]
        ax.plot(x[:, 0], v[:, 1], 'o', ms=2, color='k')
        ax.set_xlim(xmin, xmax)
        ax.set_ylim(-0.1, 0.1)
        ax.set_xlabel('x', fontsize=12)
        ax.set_ylabel('vy', fontsize=12)
        ax.axvline(x0, color='0.6', linewidth=1)
        ax.axvline(x1, color='0.6', linewidth=1)

        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, L=20 per your IDL
    show_wave(nsnaps=101, L=20.0)