import numpy as np
import g3read as g3
import matplotlib.pyplot as plt
import matplotlib as mpl
mpl.use('Agg')   # non-interactive / batch mode


def show_blast(nsnaps=6):
    gamma = 5.0 / 3.0
    rho0 = 0.125
    rho_max = (gamma + 1.0) / (gamma - 1.0) * rho0

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

        # --- read snapshot ---
        try:
            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,)
        except Exception as e:
            print(f"Skipping {name}: {e}")
            continue

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

        # --- derived fields ---
        p = (gamma - 1.0) * u * rho

        rr = np.sqrt((x[:, 0] - 0.5)**2 +
                     (x[:, 1] - 0.5)**2 +
                     (x[:, 2] - 0.5)**2)

        vv = np.sqrt(v[:, 0]**2 + v[:, 1]**2 + v[:, 2]**2)

        # analytic scaling (guard t=0 to avoid inf/nan)
        if t > 0.0:
            r1 = 0.21 * t**(2.0 / 5.0)
            u1 = 0.08 * t**(-3.0 / 5.0)
            P1 = 1e-3 * t**(-6.0 / 5.0)
        else:
            r1 = u1 = P1 = None

        # --- figure: 3 panels in 1 row (like !P.MULTI=[0,3,1]) ---
        fig, axes = plt.subplots(1, 3, figsize=(15, 4))
        plt.subplots_adjust(wspace=0.3)

        # 1) Pressure vs radius (log y)
        ax = axes[0]
        ax.plot(rr, p, 'o', ms=2, color='k')
        ax.set_xlabel('r', fontsize=12)
        ax.set_ylabel('p', fontsize=12)
        ax.set_yscale('log')
        ax.set_ylim(1e-20, 0.1)
        ax.set_xlim(0.0, 1.0)
        if r1 is not None:
            ax.plot([r1, r1], [1e-20, 1.0], linestyle='--', color='0.5')
        if P1 is not None:
            ax.plot([0.0, 1.0], [P1, P1], linestyle='--', color='0.5')

        # 2) Density vs radius
        ax = axes[1]
        ax.plot(rr, rho, 'o', ms=2, color='k')
        ax.set_xlabel('r', fontsize=12)
        ax.set_ylabel('rho', fontsize=12)
        ax.set_ylim(0.0, 0.55)
        ax.set_xlim(0.0, 1.0)
        ax.plot([0.0, 1.0], [rho_max, rho_max], linestyle='--', color='0.5')
        ax.plot([0.0, 1.0], [rho0, rho0],     linestyle='--', color='0.5')
        if r1 is not None:
            ax.plot([r1, r1], [0.0, 1.0], linestyle='--', color='0.5')

        # 3) Velocity vs radius
        ax = axes[2]
        ax.plot(rr, vv, 'o', ms=2, color='k')
        ax.set_xlabel('r', fontsize=12)
        ax.set_ylabel('v', fontsize=12)
        ax.set_ylim(0.0, 0.5)
        ax.set_xlim(0.0, 1.0)
        if r1 is not None:
            ax.plot([r1, r1], [0.0, 1.0], linestyle='--', color='0.5')
        if u1 is not None:
            ax.plot([0.0, 1.0], [u1, u1], linestyle='--', color='0.5')

        # optional: title with time
        fig.suptitle(f"snap_{snr}, t={t:.3e}", fontsize=14)

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


if __name__ == "__main__":
    show_blast(nsnaps=6)