#!/bin/bash
#SBATCH --job-name=bigbang-256
#SBATCH --partition=fast-nodes
#SBATCH --nodes=4
#SBATCH --ntasks-per-node=22
#SBATCH --time=03:15:00
#SBATCH --output=universe-%j.out
# Big Bang universe, full run: GADGET-4 dark-matter-only LCDM (Planck 2018),
# 256^3 = 16.8 million particles in a 200 Mpc/h box, z=63 -> z=0, 88 MPI ranks
# on 4 fast-nodes nodes. Same particle mass as the pilot job (64^3 in 50 Mpc/h).
# GADGET stops itself at 85% of TimeLimitCPU and writes restart files, so a run
# that is slower than estimated can be continued by a second job.
# After the run (python-sci):
#   1) large-scale P(k) growth vs linear theory, and z=0 P(k) vs the linear spectrum
#   2) FOF halo mass function vs the Watson et al. (2013) FOF fit
#   3) density frames for a movie (one per snapshot)
set -o pipefail
T0=$(date +%s)
stamp() { echo "[$(date +%T)] +$(( $(date +%s) - T0 ))s $*"; }
WORK=$PWD
echo "nodes=$SLURM_JOB_NODELIST ntasks=$SLURM_NTASKS job=$SLURM_JOB_ID dir=$WORK"

# ---------- 1. modules ----------
module purge
module load gcc/13.5.0 openmpi/5.0.10 fftw/3.3.11 hdf5/1.14.6 gsl/2.8 || { echo "FAIL: module load"; exit 2; }
prefix_of() {
  local p
  for p in ${CMAKE_PREFIX_PATH//:/ } ${CPATH//:/ }; do
    p=${p%/include}; p=${p%/.}
    [ -f "$p/include/$1" ] && { echo "$p"; return 0; }
  done
  return 1
}
FFTW=$(prefix_of fftw3.h) || { echo "FAIL: fftw3.h not found"; exit 3; }
GSL=$(prefix_of gsl/gsl_math.h) || { echo "FAIL: gsl not found"; exit 3; }
HDF5=$(prefix_of hdf5.h) || { echo "FAIL: hdf5.h not found"; exit 3; }
libdir() { [ -d "$1/lib64" ] && ls "$1/lib64" | grep -q . && echo "$1/lib64" || echo "$1/lib"; }
FL=$(libdir "$FFTW"); GL=$(libdir "$GSL"); HL=$(libdir "$HDF5")
stamp "modules ok"

# ---------- 2. get + build GADGET-4 (same commit as the pilot) ----------
git clone -q https://gitlab.mpcdf.mpg.de/vrs/gadget4.git gadget4 \
  || git clone -q https://github.com/weinberger/gadget4.git gadget4 \
  || { echo "FAIL: could not clone gadget4"; exit 4; }
cd gadget4
git checkout -q 6fb393b 2>/dev/null || echo "note: commit 6fb393b not found, using $(git log -1 --format=%h)"
echo "gadget4 commit: $(git log -1 --format='%h %cd')"

cat > Config.sh <<'EOF'
LEAN
PERIODIC
SELFGRAVITY
RANDOMIZE_DOMAINCENTER
PMGRID=512
TREEPM_NOTIMESPLIT
ASMTH=2.0
NSOFTCLASSES=1
NTYPES=2
POSITIONS_IN_32BIT
DOUBLEPRECISION=2
FOF
POWERSPEC_ON_OUTPUT
NGENIC=512
NGENIC_2LPT
CREATE_GRID
IDS_32BIT
EOF
echo 'SYSTYPE="Site"' > Makefile.systype
cat > buildsystem/Makefile.comp.site <<'EOF'
CPP      =  mpicxx -std=c++11
OPTIMIZE =  -O3 -march=x86-64-v3 -Wno-format-security
CFLAGS_VECTOR =
CPV      =  $(CPP)
EOF
cat > buildsystem/Makefile.path.site <<EOF
GSL_INCL   = -I$GSL/include
GSL_LIBS   = -L$GL -Wl,-rpath,$GL
FFTW_INCL  = -I$FFTW/include
FFTW_LIBS  = -L$FL -Wl,-rpath,$FL
HDF5_INCL  = -I$HDF5/include
HDF5_LIBS  = -L$HL -Wl,-rpath,$HL
HWLOC_INCL =
HWLOC_LIBS =
EOF
sed -i 's|^ifeq (\$(SYSTYPE),"Generic-intel")|ifeq ($(SYSTYPE),"Site")\ninclude buildsystem/Makefile.comp.site\ninclude buildsystem/Makefile.path.site\nendif\n&|' Makefile
PY3=$(command -v python3 || command -v python) || { echo "FAIL: no python3 for the build helpers"; exit 5; }
make -j 22 CONFIG=Config.sh EXEC=Gadget4 PYTHON="$PY3" > build.log 2>&1
if [ ! -x Gadget4 ]; then echo "FAIL: build"; tail -40 build.log; exit 5; fi
stamp "built Gadget4"

# ---------- 3. run ----------
RUN=$WORK/run256; mkdir -p "$RUN/output"; cd "$RUN"
# snapshots at z = 63 20 10 6 4 3 2 1.5 1 0.5 0.25 0
for z in 63 20 10 6 4 3 2 1.5 1 0.5 0.25 0; do
  awk -v z=$z 'BEGIN{printf "%.8f\n", 1/(1+z)}'
done > outputs.txt
cat > param.txt <<'EOF'
InitCondFile         ./dummy.dat
OutputDir            ./output
SnapshotFileBase     snapshot
OutputListFilename   outputs.txt
ICFormat             1
SnapFormat           3
TimeLimitCPU         12600
CpuTimeBetRestartFile 3600
MaxMemSize           3000
TimeBegin            0.015625
TimeMax              1.0
ComovingIntegrationOn 1
Omega0               0.3111
OmegaLambda          0.6889
OmegaBaryon          0.0490
HubbleParam          0.6766
Hubble               100.0
BoxSize              200.0
OutputListOn         1
TimeBetSnapshot      0.0
TimeOfFirstSnapshot  0.0
TimeBetStatistics    0.01
NumFilesPerSnapshot  1
MaxFilesWithConcurrentIO 1
ErrTolIntAccuracy    0.01
CourantFac           0.3
MaxSizeTimestep      0.005
MinSizeTimestep      0.0
TypeOfOpeningCriterion 1
ErrTolTheta          0.75
ErrTolThetaMax       1.0
ErrTolForceAcc       0.002
TopNodeFactor        3.0
ActivePartFracForNewDomainDecomp 0.01
ActivePartFracForPMinsteadOfEwald 0.05
DesNumNgb            64
MaxNumNgbDeviation   1
UnitLength_in_cm     3.085678e24
UnitMass_in_g        1.989e43
UnitVelocity_in_cm_per_s 1e5
GravityConstantInternal 0
SofteningComovingClass0  0.02
SofteningMaxPhysClass0   0.02
SofteningClassOfPartType0 0
SofteningClassOfPartType1 0
ArtBulkViscConst     1.0
MinEgySpec           0
InitGasTemp          0
NSample              256
GridSize             256
Seed                 181170
SphereMode           1
PowerSpectrumType    1
ReNormalizeInputSpectrum 1
PrimordialIndex      0.9665
ShapeGamma           0.21
Sigma8               0.8102
PowerSpectrumFile    powerspec
InputSpectrum_UnitLength_in_cm 3.085678e24
EOF
stamp "starting run on $SLURM_NTASKS ranks"
srun "$WORK/gadget4/Gadget4" param.txt > gadget.log 2>&1
RC=$?
stamp "run finished rc=$RC"
tail -8 gadget.log
grep -E "^Sync-Point" gadget.log | tail -1
echo "steps: $(grep -c '^Sync-Point' gadget.log)"
ls output/snapshot_*.hdf5 2>/dev/null | wc -l
ls output/restartfiles 2>/dev/null | head -3
[ -f output/snapshot_000.hdf5 ] || { echo "FAIL: no snapshots"; exit 6; }

# ---------- 4. check + movie frames ----------
module purge
module load python-sci
mkdir -p frames
cat > check.py <<'PY'
import glob, json, h5py, numpy as np
TRAPZ = getattr(np, 'trapezoid', None) or np.trapz
import matplotlib; matplotlib.use("Agg"); import matplotlib.pyplot as plt
from scipy.integrate import quad
L, Om, OL, Ob, h, s8, ns = 200.0, 0.3111, 0.6889, 0.0490, 0.6766, 0.8102, 0.9665
RHO_M = 2.775e11 * Om   # Msun/h per (Mpc/h)^3

def growth(a):
    H = lambda x: np.sqrt(Om / x**3 + OL)
    return 2.5 * Om * H(a) * quad(lambda x: 1.0 / (x * H(x))**3, 0, a)[0]

def tk_eh(k):  # same Eisenstein-Hu no-wiggle form GADGET-4 uses (k in h/Mpc)
    ombh2, ommh2, theta = Ob * h * h, Om * h * h, 2.728 / 2.7
    s = 44.5 * np.log(9.83 / ommh2) / np.sqrt(1 + 10 * ombh2**0.75) * h
    a = 1 - 0.328 * np.log(431 * ommh2) * ombh2 / ommh2 + 0.380 * np.log(22.3 * ommh2) * (ombh2 / ommh2)**2
    gam = (a + (1 - a) / (1 + (0.43 * k * s)**4)) * Om * h
    q = k * theta**2 / gam
    L0 = np.log(2 * np.e + 1.8 * q); C0 = 14.2 + 731 / (1 + 62.5 * q)
    return L0 / (L0 + C0 * q * q)

kk_int = np.logspace(-4, 2, 4000)
def sigma_R(R, norm=1.0):
    x = kk_int * R; W = 3 * (np.sin(x) - x * np.cos(x)) / x**3
    P = norm * kk_int**ns * tk_eh(kk_int)**2
    return np.sqrt(TRAPZ(P * W**2 * kk_int**3, np.log(kk_int)) / (2 * np.pi**2))
NORM = (s8 / sigma_R(8.0))**2
P_lin0 = lambda k: NORM * k**ns * tk_eh(k)**2

def load(path):
    with h5py.File(path) as f:
        return f["PartType1/Coordinates"][:].astype(np.float64), float(f["Header"].attrs["Time"])

def pk(x, n=256):
    g = (x / L * n) % n; i = np.floor(g).astype(np.int64); d = g - i
    rho = np.zeros(n**3)
    for dx in (0, 1):
        for dy in (0, 1):
            for dz in (0, 1):
                w = (d[:, 0] if dx else 1 - d[:, 0]) * (d[:, 1] if dy else 1 - d[:, 1]) * (d[:, 2] if dz else 1 - d[:, 2])
                idx = (((i[:, 0] + dx) % n) * n + (i[:, 1] + dy) % n) * n + (i[:, 2] + dz) % n
                rho += np.bincount(idx, weights=w, minlength=n**3)
    delta = rho.reshape(n, n, n) / rho.mean() - 1
    P = np.abs(np.fft.rfftn(delta))**2 * L**3 / n**6
    k1 = np.fft.fftfreq(n, d=L / n) * 2 * np.pi; kz = np.fft.rfftfreq(n, d=L / n) * 2 * np.pi
    kk = np.sqrt(k1[:, None, None]**2 + k1[None, :, None]**2 + kz[None, None, :]**2)
    b = np.floor(kk / (2 * np.pi / L)).astype(int)
    return {j: (float(kk[b == j].mean()), float(P[b == j].mean()), int((b == j).sum())) for j in range(1, 9)}

def frame(x, a, out):
    sl = x[:, 2] < 20.0   # 20 Mpc/h thick slab
    H, _, _ = np.histogram2d(x[sl, 0], x[sl, 1], bins=1024, range=[[0, L], [0, L]])
    fig = plt.figure(figsize=(8, 8), dpi=128); ax = fig.add_axes([0, 0, 1, 1]); ax.axis("off")
    ax.imshow(np.log10(1 + H.T), origin="lower", cmap="inferno", vmin=0, vmax=2.6)
    ax.text(0.02, 0.97, f"z = {1/a - 1:.2f}   ({200:.0f} Mpc/h box, 20 Mpc/h slab)", color="w",
            transform=ax.transAxes, fontsize=13, va="top")
    fig.savefig(out); plt.close(fig)

res = {}
snaps = sorted(glob.glob("output/snapshot_*.hdf5"))
x0, a0 = load(snaps[0]); p0 = pk(x0); frame(x0, a0, "frames/frame_000.png"); del x0
for s in snaps[1:-1]:
    x, a = load(s); frame(x, a, f"frames/frame_{s[-8:-5]}.png"); del x
x1, a1 = load(snaps[-1]); p1 = pk(x1); frame(x1, a1, f"frames/frame_{snaps[-1][-8:-5]}.png"); del x1
expect = (growth(a1) / growth(a0))**2
print(f"last snapshot a={a1:.4f} (z={1/a1-1:.3f}); {len(snaps)} snapshots")
print(f"P(k) growth a={a0:.4f}->{a1:.4f}: linear theory D^2 ratio = {expect:.1f}")
growth_rows = []
for j in range(1, 9):
    k, P1, nm = p1[j]; r = P1 / p0[j][1]
    lin = P_lin0(k) * (growth(a1) / growth(1.0))**2
    growth_rows.append((k, r / expect, P1 / lin, nm))
    print(f"  k={k:.3f} h/Mpc ({nm:5d} modes): growth {r/expect:.2f} x theory, P(z) / P_lin {P1/lin:.2f}")
res["pk"] = growth_rows
lo = float(np.mean([g[1] for g in growth_rows[1:4]]))

# halo mass function at the last snapshot vs Watson+2013 FOF fit
fof = sorted(glob.glob("output/fof_tab_*.hdf5"))[-1]
with h5py.File(fof) as f:
    n = int(f["Header"].attrs["Ngroups_Total"])
    M = f["Group/GroupMass"][:] * 1e10; N = f["Group/GroupLen"][:]
mp = RHO_M * L**3 / 256**3
print(f"FOF at {fof}: {n} groups; particle mass {mp:.3e} Msun/h; largest {M.max():.2e} Msun/h ({N.max()} particles)")
Dz = growth(a1) / growth(1.0)
def dndlnM(m):
    R = (3 * m / (4 * np.pi * RHO_M))**(1 / 3)
    s = sigma_R(R, NORM) * Dz; s2 = sigma_R(R * 1.01, NORM) * Dz
    dlns_dlnM = (np.log(s2) - np.log(s)) / (3 * np.log(1.01))
    f = 0.282 * ((1.406 / s)**2.163 + 1) * np.exp(-1.210 / s**2)
    return f * RHO_M / m * abs(dlns_dlnM)
edges = np.logspace(np.log10(100 * mp), np.log10(M.max()) + 0.01, 9)
rows, ok = [], []
for lo_m, hi_m in zip(edges[:-1], edges[1:]):
    cnt = int(((M >= lo_m) & (M < hi_m)).sum())
    pred = quad(lambda lnm: dndlnM(np.exp(lnm)), np.log(lo_m), np.log(hi_m))[0] * L**3
    rows.append((lo_m, hi_m, cnt, pred))
    flag = "" if cnt < 50 else ("ok" if abs(cnt / pred - 1) < 0.25 else "OFF")
    if cnt >= 50: ok.append(abs(cnt / pred - 1) < 0.25)
    print(f"  {lo_m:.2e}-{hi_m:.2e}: {cnt:6d} halos, Watson fit {pred:9.1f}  ratio {cnt/max(pred,1e-9):.2f} {flag}")
res["hmf"] = rows
verdict = "PASS" if (0.8 < lo < 1.25 and ok and all(ok) and abs(a1 - 1) < 1e-3) else ("PARTIAL (stopped early, restart files written)" if a1 < 0.999 else "CHECK")
print("RUN VERDICT:", verdict)
res.update(verdict=verdict, a_last=a1, n_groups=n, largest=float(M.max()), particle_mass=mp)
json.dump(res, open("check.json", "w"), indent=1)

# plots
fig, ax = plt.subplots(1, 2, figsize=(12, 5))
cen = np.sqrt(edges[:-1] * edges[1:])
ax[0].loglog(cen, [r[2] for r in rows], "o", label="simulation (FOF)")
ax[0].loglog(cen, [r[3] for r in rows], "-", label="Watson+2013 fit")
ax[0].set_xlabel("halo mass [Msun/h]"); ax[0].set_ylabel("halos per bin"); ax[0].legend(); ax[0].set_title(f"Halo mass function, z={1/a1-1:.2f}")
ks = np.logspace(-1.6, 0.6, 200)
ax[1].loglog(ks, P_lin0(ks) * Dz**2, "-", label="linear theory")
ax[1].loglog([p1[j][0] for j in p1], [p1[j][1] for j in p1], "o", label="simulation")
ax[1].set_xlabel("k [h/Mpc]"); ax[1].set_ylabel("P(k) [(Mpc/h)^3]"); ax[1].legend(); ax[1].set_title("Matter power spectrum")
fig.tight_layout(); fig.savefig("check.png", dpi=110)
PY
python check.py
stamp "check done"
ls frames | wc -l
stamp "done"