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

# -----------------------------------
# Parameters
# -----------------------------------
gamma = 5.0 / 3.0
snr   = "050"
snap  = f"snap_{snr}"

# -----------------------------------
# Read snapshot
# -----------------------------------
f   = g3.GadgetFile(snap)
x   = f.read_new("POS ", 0)      # positions (N,3)
v   = f.read_new("VEL ", 0)      # velocities (N,3)
rho = f.read_new("RHO ", 0)      # density (N,)
u   = f.read_new("U   ", 0)      # internal energy (N,)

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

# -----------------------------------
# Create figure with 3 vertically stacked subplots
# -----------------------------------
fig, axes = plt.subplots(3, 1, figsize=(6, 8), sharex=True)
plt.subplots_adjust(hspace=0.3)

# Plot 1: Pressure vs x
axes[0].plot(x[:, 0], p, 'o', ms=4)
axes[0].set_ylabel("P", fontsize=14)
axes[0].set_xlim(10, 30)
axes[0].set_ylim(0.4, 1.1)

# Plot 2: Density vs x
axes[1].plot(x[:, 0], rho, 'o', ms=4)
axes[1].set_ylabel("rho", fontsize=14)
axes[1].set_xlim(10, 30)
axes[1].set_ylim(0.2, 1.2)

# Plot 3: Velocity vx vs x
axes[2].plot(x[:, 0], v[:, 0], 'o', ms=4)
axes[2].set_xlabel("x", fontsize=14)
axes[2].set_ylabel("vx", fontsize=14)
axes[2].set_xlim(10, 30)
axes[2].set_ylim(0.0, 0.5)

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