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


def show_nozzle(nsnaps=121, 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)  # not used below, but kept for completeness

    # color scaling for u (like umin/umax and color indices in IDL)
    umin = 5.0
    umax = 13.0

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

        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)
            ids = f.read_new("ID  ", 0)   # (N,)
            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

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

        p = (gamma - 1.0) * u * rho

        # masks
        ids = np.asarray(ids)
        jjb  = np.where(ids >= 1)[0]                           # bulk particles
        jjw0 = np.where((ids == 0) & (v[:, 0] == 0.0))[0]      # fixed walls
        jjw1 = np.where((ids == 0) & (v[:, 0] != 0.0))[0]      # moving wall

        # central y-slice: |y - 0.5| < 0.1
        y = x[:, 1]
        jj  = np.where(np.abs(y - 0.5) < 0.1)[0]
        jjb0 = np.where((np.abs(y - 0.5) < 0.1) & (ids == 0) & (v[:, 0] == 0.0))[0]
        jjb1 = np.where((np.abs(y - 0.5) < 0.1) & (ids == 0) & (v[:, 0] != 0.0))[0]

        # color coding for u (bulk particles)
        u_norm = (u - umin) / (umax - umin)
        u_norm = np.clip(u_norm, 0.0, 1.0)  # clamp, analogous to icol clipping in IDL

        # set up figure with 3 vertical panels (like !P.MULTI=[0,1,3])
        fig, axes = plt.subplots(3, 1, figsize=(8, 10), sharex=True)
        plt.subplots_adjust(hspace=0.3)

        # ------------------------------------------------------------------
        # Panel 1: geometry in x–y with walls + color-coded bulk by u
        # ------------------------------------------------------------------
        ax = axes[0]
        ax.set_xlim(10.0, 40.0)
        ax.set_ylim(0.0, 1.0)
        ax.set_xlabel("x", fontsize=12)
        ax.set_ylabel("y", fontsize=12)
        ax.set_title(f"t = {t:5.2f}", fontsize=13)

        # walls (fixed and moving)
        if jjw0.size > 0:
            ax.scatter(x[jjw0, 0], x[jjw0, 1], s=30, color='C0', label='wall (fixed)')
        if jjw1.size > 0:
            ax.scatter(x[jjw1, 0], x[jjw1, 1], s=30, color='C3', label='wall (moving)')

        # bulk particles, colored by u
        if jjb.size > 0:
            sc = ax.scatter(x[jjb, 0], x[jjb, 1],
                            s=5, c=u_norm[jjb],
                            cmap='viridis', vmin=0.0, vmax=1.0)
            cbar = fig.colorbar(sc, ax=ax, pad=0.01)
            cbar.set_label("u (scaled)", fontsize=10)

        ax.legend(fontsize=9, loc='upper right')

        # ------------------------------------------------------------------
        # Panel 2: pressure vs x (central slice)
        # ------------------------------------------------------------------
        ax = axes[1]
        ax.set_xlim(10.0, 40.0)
        ax.set_ylim(0.0, 0.6)
        ax.set_ylabel("P", fontsize=12)

        # all slice particles
        ax.plot(x[jj, 0], p[jj], 'o', ms=3, color='k', label='bulk')

        # overplot walls in slice
        if jjb0.size > 0:
            ax.plot(x[jjb0, 0], p[jjb0], 'o', ms=5, color='C0', label='wall (fixed)')
        if jjb1.size > 0:
            ax.plot(x[jjb1, 0], p[jjb1], 'o', ms=5, color='C3', label='wall (moving)')

        ax.legend(fontsize=9, loc='upper right')

        # ------------------------------------------------------------------
        # Panel 3: vx vs x (central slice) + -c_s
        # ------------------------------------------------------------------
        ax = axes[2]
        ax.set_xlim(10.0, 40.0)
        ax.set_ylim(-5.0, 0.5)
        ax.set_xlabel("x", fontsize=12)
        ax.set_ylabel("vx", fontsize=12)

        # bulk slice
        ax.plot(x[jj, 0], v[jj, 0], 'o', ms=3, color='k', label='bulk vx')

        # walls
        if jjb0.size > 0:
            ax.plot(x[jjb0, 0], v[jjb0, 0], 'o', ms=5, color='C0', label='wall (fixed)')
        if jjb1.size > 0:
            ax.plot(x[jjb1, 0], v[jjb1, 0], 'o', ms=5, color='C3', label='wall (moving)')

        # local sound speed c(x) = sqrt(gamma * P / rho), plot -c
        cc = np.sqrt(gamma * p / rho)
        ax.plot(x[jj, 0], -cc[jj], 'o', ms=3, color='C2', label='-c_s')

        ax.legend(fontsize=9, loc='upper right')

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


if __name__ == "__main__":
    # 0..120 inclusive -> 121 frames
    show_nozzle(nsnaps=121, L=20.0)