import numpy as np
import g3read as g3
from shutil import copyfile
import matplotlib.pyplot as plt
import matplotlib as mpl
mpl.use('Agg')  # batch / no GUI
import os

def setup_blast_3D(
    fin="glass_10x10x10",
    fout="box.ic",
    Ndup=10,
    gamma=5.0/3.0,
    rho=0.125,
    Pback=1e-20,
    Pexpl=1.0,
    r_blast=0.013,          # radius of hot sphere (unit-box coords)
    ic_header_path="ic_header",
    preview_png="blast_preview.png",
):
    # --- 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
    Nin  = int(np.sum(f_in.header.npart))

    # --- tile to Ndup^3 in the unit box ---
    Ntiles = Ndup**3
    Np = Nin * Ntiles

    # 3D offsets on a regular Ndup grid: (ix,iy,iz)/Ndup
    ix, iy, iz = np.meshgrid(
        np.arange(Ndup, dtype=np.float32),
        np.arange(Ndup, dtype=np.float32),
        np.arange(Ndup, dtype=np.float32),
        indexing="ij"
    )
    offsets = np.stack([ix, iy, iz], axis=-1).reshape(-1, 3)  # (Ntiles,3)

    # Broadcast: x = (xin + offset)/Ndup → shape (Ntiles,Nin,3) → (Np,3)
    x = ((xin[None, :, :] + offsets[:, None, :]) / float(Ndup)).reshape(-1, 3).astype(np.float32)

    # --- fields ---
    v   = np.zeros((Np, 3), dtype=np.float32)
    u   = np.full(Np, Pback / (gamma - 1.0) / rho, dtype=np.float32)
    m   = np.full(Np, (rho * 1.0**3) / Np, dtype=np.float32)  # BoxSize(out)=1
    ids = np.arange(Np, dtype=np.int32)  # IDL: lindgen(Np) (0-based). Use +1 if you want 1-based.

    # --- central hot sphere (around (0.5,0.5,0.5) in unit box) ---
    r2 = np.sum((x - 0.5)**2, axis=1)
    mask_blast = r2 < (r_blast**2)
    nfound = int(mask_blast.sum())
    print("Particles in hot sphere:", nfound)
    u[mask_blast] = Pexpl / (gamma - 1.0) / rho

    # --- 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)
    # Mirror IDL header edits
    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 data blocks and write them
    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)

    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")

    # --- quick preview plot (xy projection + highlight hot particles) ---
    try:
        fig, ax = plt.subplots(figsize=(6, 6))
        ax.plot(x[:, 0], x[:, 1], 'o', ms=0.02, color='0.3')
        hot = np.where(mask_blast)[0]
        if hot.size > 0:
            ax.plot(x[hot, 0], x[hot, 1], 'o', ms=2.0, color='C3')
        ax.axvline(0.5, linestyle='--', linewidth=1.2, color='0.4')
        ax.axhline(0.5, linestyle='--', linewidth=1.2, color='0.4')
        ax.set_xlim(0, 1)
        ax.set_ylim(0, 1)
        ax.set_aspect('equal')
        fig.savefig(preview_png, dpi=150, bbox_inches='tight')
        plt.close(fig)
        print(f"Saved preview to '{preview_png}'")
    except Exception as e:
        print(f"[preview skipped] {e}")

if __name__ == "__main__":
    setup_blast_3D(
        fin="glass_10x10x10",
        fout="box.ic",
        Ndup=10,
        gamma=5.0/3.0,
        rho=0.125,
        Pback=1e-20,
        Pexpl=1.0,
        r_blast=0.013,
        ic_header_path="ic_header",
        preview_png="blast_preview.png",
    )