# show_kernels.py
import numpy as np
import matplotlib as mpl
mpl.use("Agg")  # headless, like IDL's window-less save
import matplotlib.pyplot as plt
import g3read as g3

def _get_boxsize(gfile, head_obj):
    """
    Try a few common places/keys for BoxSize, mirroring IDL's h.BOXSIZE.
    """
    # If HEAD block came back as a dict-like
    if head_obj is not None:
        for k in ("BOXSIZE", "BoxSize", "boxsize"):
            if isinstance(head_obj, dict) and k in head_obj:
                return float(head_obj[k])
            if hasattr(head_obj, k):
                return float(getattr(head_obj, k))
    # Fallback to GadgetFile header
    hdr = getattr(gfile, "header", None)
    if hdr is not None:
        if isinstance(hdr, dict):
            for k in ("BoxSize", "boxsize", "BOXSIZE"):
                if k in hdr:
                    return float(hdr[k])
        else:
            for k in ("BoxSize", "boxsize", "BOXSIZE"):
                if hasattr(hdr, k):
                    return float(getattr(hdr, k))
    raise KeyError("BoxSize not found in header.")

def show_kernels():
    # --- IDL: win & initial plot with axis config ---
    fig = plt.figure(figsize=(5, 5), dpi=100)
    ax = fig.add_subplot(111)

    # xlog; xr=[9,500]; yr=[-0.005,0.07]
    ax.set_xscale("log")
    ax.set_xlim(9, 500)
    ax.set_ylim(-0.005, 0.07)

    # xtit='N', ytit='log(rho_sim/rho_ini)'
    ax.set_xlabel("N")
    ax.set_ylabel("log(rho_sim/rho_ini)")

    # Rough equivalents to IDL's char size/tick size params
    ax.tick_params(axis="both", which="both", labelsize=10)

    # --- IDL: oplot,[1,1000],[0,0],col=96 (horizontal reference line at y=0) ---
    ax.plot([1, 1000], [0.0, 0.0], linestyle="-", linewidth=1)

    # --- IDL: FOR ingb=32,96 DO BEGIN ... END ---
    for ingb in range(32, 97):
        snr = f"wc2_{ingb}_000"   # 'wc2_'+strtrim(string(ingb,form='(i)'),1)+'_000'

        # readnew,snr,h,'HEAD'
        f = g3.GadgetFile(snr)
        try:
            h = f.read_new("HEAD", 0)
        except Exception:
            h = None

        # readnew,snr,mass,'MASS'
        mass = f.read_new("MASS", 0)

        # rho_ini = TOTAL(mass) / h.BOXSIZE^3
        boxsize = _get_boxsize(f, h)
        rho_ini = np.sum(mass, dtype=np.float64) / float(boxsize**3)

        # readnew,snr,rho,'RHO'
        rho = f.read_new("RHO ", 0)

        # rho_sim = TOTAL(rho) / N_ELEMENTS(rho)
        rho_sim = float(np.mean(rho))

        # oplot,[ingb],[alog10(rho_sim/rho_ini)],psym=4
        yval = np.log10(rho_sim / rho_ini)
        ax.plot([ingb], [yval], marker="o", linestyle="none")

    # --- IDL: save_screen,'plot' ---
    fig.tight_layout()
    fig.savefig("plot.png")  # saved as PNG; change extension if you prefer PDF/SVG

if __name__ == "__main__":
    show_kernels()