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_sk(name="snap_010", outname="sk.png"):
    # --- Read snapshot ---
    f = g3.GadgetFile(name)

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

    # Gas only (part/type 0)
    x = f.read_new("POS ", 0)      # (N,3)
    m = f.read_new("MASS", 0)      # (N,)
    sfr = f.read_new("SFR ", 0)    # (N,)

    # Cylindrical radius in the disk plane
    rr = np.sqrt(x[:, 0]**2 + x[:, 1]**2)

    # --- Radial bins (IDL translation) ---
    NN = 20
    # IDL: rtab = 20.0*10.0d0^(findgen(NN)/(NN-1)*3-3)
    idx = np.arange(NN, dtype=np.float64)
    rtab = 20.0 * 10.0 ** (idx / (NN - 1.0) * 3.0 - 3.0)

    rho_tab = np.zeros(NN - 1, dtype=np.float64)
    sfr_tab = np.zeros(NN - 1, dtype=np.float64)

    # --- Bin loop (keep it explicit, like IDL) ---
    for i in range(NN - 1):
        mask = (rr > rtab[i]) & (rr <= rtab[i + 1])
        if np.any(mask):
            area = np.pi * (rtab[i + 1]**2 - rtab[i]**2)
            # IDL: rho_tab = TOTAL(m*1e10)/area/1e6
            rho_tab[i] = np.sum(m[mask] * 1e10) / area / 1e6
            # IDL: sfr_tab = TOTAL(sfr)/area
            sfr_tab[i] = np.sum(sfr[mask]) / area
        else:
            rho_tab[i] = np.nan
            sfr_tab[i] = np.nan

    # Only plot bins with valid positive sfr and rho (log10 and log-x need >0)
    good = np.isfinite(rho_tab) & np.isfinite(sfr_tab) & (rho_tab > 0.0) & (sfr_tab > 0.0)

    # --- Plot ---
    fig, ax = plt.subplots(figsize=(7, 5))

    ax.plot(rho_tab[good], np.log10(sfr_tab[good]), "o", ms=8, color="k")

    ax.set_xscale("log")  # IDL: /xlog
    ax.set_xlim(3, 3000)
    ax.set_ylim(-4, 1.7)

    ax.set_xlabel("rho_gas [Msol/pc^2]", fontsize=12)
    ax.set_ylabel("rho_sfr [Msol/year/kpc^2]", fontsize=12)

    ax.tick_params(labelsize=10)
    fig.tight_layout()
    fig.savefig(outname, dpi=150, bbox_inches="tight")
    plt.close(fig)

    print(f"Saved {outname}")


if __name__ == "__main__":
    show_sk()