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

def setup_grid_2d(
    Ndump=20,
    Ny_left=16,
    Ny_right=10,
    P_left=1.0,
    P_right=0.5,
    Rho_left=1.0,
    gamma=5.0/3.0,
    boxlen_y=1.0,
    close_pack=False,  # toggles the commented "close pack" shift in IDL
):
    """
    Build a 2D test problem like the IDL 'setup_grid':
      - Left block:  Ndump slabs, each with Ny_left x Ny_left particles
      - Right block: Ndump slabs, each with Ny_right x Ny_right particles
      - y in [0, boxlen_y], z = 0; x runs from [0, 2*Ndump)
    """
    # Choose Rho_right so that particle masses match across sides:
    # mass_left  = Rho_left / Ny_left^2
    # mass_right = Rho_right / Ny_right^2  -> set equal to mass_left
    Rho_right = Rho_left * (Ny_right / Ny_left) ** 2

    # Total particles
    N_left  = Ndump * Ny_left  * Ny_left
    N_right = Ndump * Ny_right * Ny_right
    Np = N_left + N_right

    # Allocate arrays
    pos  = np.zeros((Np, 3), dtype=np.float32)
    vel  = np.zeros((Np, 3), dtype=np.float32)
    mass = np.zeros(Np,       dtype=np.float32)
    U    = np.zeros(Np,       dtype=np.float32)
    ids  = np.arange(1, Np + 1, dtype=np.int32)

    # Common: velocities zero (as in IDL)
    # (already zero-initialized)

    ip = 0

    # --- Left block ---
    # k = 0..Ndump-1, i,j on Ny_left grid
    for k in range(Ndump):
        for i in range(Ny_left):
            for j in range(Ny_left):
                # x position spans [k, k+1) with Ny_left resolution
                x = k + (i + 0.5) / Ny_left
                # (optional) close-packing like the commented IDL line
                if close_pack:
                    if (j % 2) == 0:
                        x += 0.25 / Ny_left
                    else:
                        x -= 0.25 / Ny_left

                y = (j + 0.5) * boxlen_y / Ny_left
                pos[ip, 0] = x
                pos[ip, 1] = y
                pos[ip, 2] = 0.0

                # thermodynamics
                mass[ip] = Rho_left / (Ny_left * Ny_left)
                U[ip]    = P_left / (gamma - 1.0) / Rho_left
                ip += 1

    # --- Right block ---
    # k = 0..Ndump-1, i,j on Ny_right grid; shifted by Ndump along x
    for k in range(Ndump):
        for i in range(Ny_right):
            for j in range(Ny_right):
                x = Ndump + k + (i + 0.5) / Ny_right
                if close_pack:
                    if (j % 2) == 0:
                        x += 0.25 / Ny_right
                    else:
                        x -= 0.25 / Ny_right

                y = (j + 0.5) * boxlen_y / Ny_right
                pos[ip, 0] = x
                pos[ip, 1] = y
                pos[ip, 2] = 0.0

                mass[ip] = Rho_right / (Ny_right * Ny_right)
                U[ip]    = P_right / (gamma - 1.0) / Rho_right
                ip += 1

    assert ip == Np

    # For reference, you can print ranges like in IDL:
    # print("Left mass,u:", mass[:N_left].min(), mass[:N_left].max(), U[:N_left].min(), U[:N_left].max())
    # print("Right mass,u:", mass[N_left:].min(), mass[N_left:].max(), U[N_left:].min(), U[N_left:].max())

    # Domain: x spans [0, 2*Ndump), y in [0, 1], z=0
    Lx = float(2 * Ndump)
    Ly = float(boxlen_y)
    Lz = 1.0  # arbitrary; positions have z=0

    return pos, vel, ids, mass, U, (Lx, Ly, Lz), Np


def write_output(output_file, pos, vel, ids, mass, U, boxsize_x, ic_header_path="ic_header"):
    """
    Write Gadget IC file using a reference header file, like your original script.
    """
    if not os.path.exists(ic_header_path):
        raise FileNotFoundError(f"Reference header file '{ic_header_path}' not found.")

    # Copy reference header to output
    copyfile(ic_header_path, output_file)

    Np = pos.shape[0]

    # Open and patch header
    f = g3.GadgetFile(output_file, is_snap=False)  # is_snap=False so we can add POS block etc.
    f.header.npart = [Np, 0, 0, 0, 0, 0]
    f.header.npartTotal = [Np, 0, 0, 0, 0, 0]
    f.header.time = 0.0
    f.header.boxsize = 1.0

    f.write_header(f.header, filename=output_file)

    # Add data blocks sized for Np, not N**3
    f.add_file_block('POS ', Np * 4 * 3, partlen=4 * 3)                 # 3 floats per particle
    f.add_file_block('VEL ', Np * 4 * 3, partlen=4 * 3)
    f.add_file_block('ID  ', Np * 4,     partlen=4, dtype=np.int32)     # 1 int32 per particle
    f.add_file_block('MASS', Np * 4,     partlen=4)                     # 1 float per particle
    f.add_file_block('U   ', Np * 4,     partlen=4)                     # 1 float per particle

    # Write data
    f.write_block('POS ',  0, pos,  filename=output_file)
    f.write_block('VEL ',  0, vel,  filename=output_file)
    f.write_block('ID  ',  0, ids,  filename=output_file)
    f.write_block('MASS',  0, mass, filename=output_file)
    f.write_block('U   ',  0, U,    filename=output_file)


def main():
    # Mirror the IDL defaults
    Ndump    = 20
    Ny_left  = 16
    Ny_right = 10
    P_left   = 1.0
    P_right  = 0.5
    Rho_left = 1.0
    gamma    = 5.0 / 3.0
    close_pack = False
    output_file = "box.ic"

    pos, vel, ids, mass, U, (Lx, Ly, Lz), Np = setup_grid_2d(
        Ndump=Ndump, Ny_left=Ny_left, Ny_right=Ny_right,
        P_left=P_left, P_right=P_right,
        Rho_left=Rho_left, gamma=gamma,
        boxlen_y=1.0, close_pack=close_pack
    )

    # Write IC file (use Lx for header.boxsize to reflect x-extent)
    write_output(output_file, pos, vel, ids, mass, U, boxsize_x=Lx, ic_header_path="ic_header")

    # Quick sanity prints (like the IDL log)
    N_left = Ndump * Ny_left * Ny_left
    print("Left:")
    print("  mass min/max =", float(mass[:N_left].min()), float(mass[:N_left].max()))
    print("  U    min/max =", float(U[:N_left].min()),    float(U[:N_left].max()))
    print("Right:")
    print("  mass min/max =", float(mass[N_left:].min()), float(mass[N_left:].max()))
    print("  U    min/max =", float(U[N_left:].min()),    float(U[N_left:].max()))
    print(f"Wrote {output_file} with Np = {Np}, boxsize (x) = {Lx}")

if __name__ == "__main__":
    main()