import numpy as np
import g3read as g3
from shutil import copyfile
import os

def setup_slab(
    fin="glass_10x10x10",
    fout="box.ic",
    Ndup=20,
    gamma=5.0/3.0,
    rho=0.125,
    P=1.0,
    dv=0.1,         # wave velocity amplitude
    l=2.0,          # wavelength parameter for end waves
    ic_header_path="ic_header",
):
    # --- read input glass and normalize by its box size ---
    f_in = g3.GadgetFile(fin)
    xin  = f_in.read_new("POS ", 0).astype(np.float32)   # (Nin, 3)
    try:
        h_in = f_in.read_new("HEAD", 0)
        L_in = float(h_in["BOXSIZE"]) if isinstance(h_in, dict) else float(getattr(h_in, "boxsize"))
    except Exception:
        L_in = float(getattr(f_in.header, "boxsize", 1.0))
    xin /= L_in                                         # now xin in [0,1)^3
    Nin  = int(np.sum(f_in.header.npart))

    # --- duplicate along x only (Ndup slabs) ---
    Np = Ndup * Nin
    x  = np.zeros((Np, 3), dtype=np.float32)
    i  = 0
    for ix in range(Ndup):
        # positions for this slab: x = xin + (ix, 0, 0)
        x[i:i+Nin, 0] = xin[:, 0] + ix
        x[i:i+Nin, 1] = xin[:, 1]
        x[i:i+Nin, 2] = xin[:, 2]
        i += Nin
    assert i == Np

    # --- fields ---
    v   = np.zeros((Np, 3), dtype=np.float32)
    u   = np.full(Np, P / (gamma - 1.0) / rho, dtype=np.float32)
    m   = np.full(Np, (rho * 1.0**3) / Np,     dtype=np.float32)  # h.BOXSIZE will be 1
    ids = np.arange(Np, dtype=np.int32)  # IDL: lindgen(Np)

    # --- add end “wave” in vx (mirrors the IDL logic) ---
    # Note: IDL had 'du' defined but unused; we replicate only 'dv' usage.
    x0 = x[:, 0]
    left_mask  = x0 < l
    right_mask = x0 > (Ndup - l)
    # left end
    v[left_mask, 0]  = dv * np.sin(np.pi * x0[left_mask] / l)
    # right end (with a minus sign)
    v[right_mask, 0] = -dv * np.sin(np.pi * (Ndup - x0[right_mask]) / l)

    # --- write output using a reference header file ---
    if not os.path.exists(ic_header_path):
        raise FileNotFoundError(f"Reference header '{ic_header_path}' not found.")
    copyfile(ic_header_path, fout)

    f_out = g3.GadgetFile(fout, is_snap=False)
    # Match IDL: h.BOXSIZE = 1 and zero massarr; npart set only in slot 0
    f_out.header.npart       = [Np, 0, 0, 0, 0, 0]
    f_out.header.npartTotal  = [Np, 0, 0, 0, 0, 0]
    f_out.header.time        = 0.0
    f_out.header.boxsize     = 1.0
    try:
        f_out.header.massarr = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
    except Exception:
        pass
    f_out.write_header(f_out.header, filename=fout)

    # add blocks sized for Np
    f_out.add_file_block('POS ',  Np * 4 * 3, partlen=4 * 3)
    f_out.add_file_block('VEL ',  Np * 4 * 3, partlen=4 * 3)
    f_out.add_file_block('ID  ',  Np * 4,     partlen=4,     dtype=np.int32)
    f_out.add_file_block('MASS',  Np * 4,     partlen=4)
    f_out.add_file_block('U   ',  Np * 4,     partlen=4)

    # write data
    f_out.write_block('POS ', 0, x,    filename=fout)
    f_out.write_block('VEL ', 0, v,    filename=fout)
    f_out.write_block('ID  ', 0, ids,  filename=fout)
    f_out.write_block('MASS', 0, m,    filename=fout)
    f_out.write_block('U   ', 0, u,    filename=fout)

    print(f"Wrote '{fout}': Nin={Nin}, Ndup={Ndup}, Np={Np}, BoxSize=1")

if __name__ == "__main__":
    setup_slab(
        fin="glass_10x10x10",
        fout="box.ic",
        Ndup=20,
        gamma=5.0/3.0,
        rho=0.125,
        P=1.0,
        dv=0.1,
        l=2.0,
        ic_header_path="ic_header",
    )