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_nozzle(
    N_x=1000,
    N_y=50,
    gamma=5.0/3.0,
    rho=0.125,
    P=0.5,
    fout="box.ic",
    ic_header_path="ic_header",
    preview_png="nozzle_geometry.png",
):
    # ------------------------------------------------------------------
    # Basic setup and geometry
    # ------------------------------------------------------------------
    L = N_x / N_y  # domain length in x (IDL: L = N_x / N_y)
    print(f"Region: {N_x} x {N_y}, Length = {L}")

    Np = N_x * N_y

    # positions
    x = np.zeros((Np, 3), dtype=np.float32)
    i = 0
    for ix in range(N_x):
        for iy in range(N_y):
            # base x coordinate
            xx = (ix + 0.5) / N_y
            # staggered (close-packed) shift in x depending on row parity
            if (iy % 2) == 0:
                xx += 0.25 / N_y
            else:
                xx -= 0.25 / N_y

            yy = (iy + 0.5) / N_y

            x[i, 0] = xx
            x[i, 1] = yy
            x[i, 2] = 0.0
            i += 1

    if i != Np:
        raise RuntimeError("Incorrect particle count in position setup")

    # velocities, internal energy, mass, IDs
    v   = np.zeros((Np, 3), dtype=np.float32)
    u   = np.zeros(Np,       dtype=np.float32)
    m   = np.zeros(Np,       dtype=np.float32)
    ids = np.arange(1, Np + 1, dtype=np.int32)  # LINDGEN(Np)+1

    # thermodynamics
    u[:] = P / (gamma - 1.0) / rho
    mpart = rho / (N_y * N_y)
    m[:] = mpart

    # ------------------------------------------------------------------
    # Mark boundaries and nozzle
    # ------------------------------------------------------------------

    # Wall at end of slab: mark as boundary (id=0) for x in (L - dL, L)
    dL = 0.5
    mask_end = x[:, 0] > (L - dL)
    ids[mask_end] = 0

    # Nozzle in left part with length dL = 8
    x0 = 0.5
    dL_noz = 8.0
    dmin = 0.15
    amp = (0.5 - dmin) / 2.0

    mask_noz_region = (x[:, 0] > x0) & (x[:, 0] < x0 + dL_noz)
    if np.any(mask_noz_region):
        dy = np.abs(x[:, 1] - 0.5)
        # dy_nozzle = 0.5 + amp*(sin((x - x0)/dL * 2π + π/2) - 1)
        theta = (x[:, 0] - x0) / dL_noz * 2.0 * np.pi + np.pi / 2.0
        dy_nozzle = 0.5 + amp * (np.sin(theta) - 1.0)
        mask_outside = mask_noz_region & (dy > dy_nozzle)
        ids[mask_outside] = 0

    # Wall at beginning of slab (x in [0, dL]): mark as boundary and assign velocity
    dL = 0.5
    dv = -1.0
    mask_start = x[:, 0] < dL
    ids[mask_start] = 0
    v[mask_start, 0] = dv

    # Move everything to the right by L
    x[:, 0] += L

    # ------------------------------------------------------------------
    # Write Gadget IC using 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)
    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  # as in IDL
    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
    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}' with Np={Np}, L={L}")

    # ------------------------------------------------------------------
    # Quick geometry plot (like the IDL check)
    # ------------------------------------------------------------------
    try:
        fig, ax = plt.subplots(figsize=(10, 3))
        ax.plot(x[:, 0], x[:, 1], 'o', ms=0.5, color='0.6')

        # fixed particles (id==0 and v_x == 0)
        mask_fixed = (ids == 0) & (v[:, 0] == 0.0)
        if np.any(mask_fixed):
            ax.plot(x[mask_fixed, 0], x[mask_fixed, 1],
                    'o', ms=4, color='C0')  # like mycolor(1)

        # moving wall particles (id==0 and v_x != 0)
        mask_moving = (ids == 0) & (v[:, 0] != 0.0)
        if np.any(mask_moving):
            ax.plot(x[mask_moving, 0], x[mask_moving, 1],
                    'o', ms=4, color='C3')  # like mycolor(4)

        ax.set_xlim(0, 40)
        ax.set_xlabel("x")
        ax.set_ylabel("y")
        ax.set_aspect('equal', adjustable='box')
        fig.tight_layout()
        fig.savefig(preview_png, dpi=150, bbox_inches='tight')
        plt.close(fig)
        print(f"Saved geometry preview to '{preview_png}'")
    except Exception as e:
        print(f"[preview skipped] {e}")


if __name__ == "__main__":
    setup_nozzle()