import os
import struct
import numpy as np

def write_fortran_record(f, data_bytes):
    length = len(data_bytes)
    f.write(struct.pack('i', length))
    f.write(data_bytes)
    f.write(struct.pack('i', length))

def generate_synthetic_grid(output_path, nx=16, ny=16, nz=16):
    """Generates a 16x16x16 synthetic flow field binary file matching Fortran unformatted record format.
    
    Variables:
      u(x,y,z) = sin(x) * cos(y) * sin(z)
      v(x,y,z) = cos(x) * sin(y) * cos(z)
      w(x,y,z) = sin(x) * sin(y) * cos(z)
      c(x,y,z) = cos(x) * cos(y) * cos(z) (Stored in new_scalar(:,:,:,2))
    """
    lx = 2.0 * np.pi
    ly = 2.0 * np.pi
    lz = 2.0 * np.pi
    
    hx = lx / nx
    hy = ly / ny
    hz = lz / nz
    
    # 1-based indexing coordinates for Fortran mesh mapping
    x = np.arange(1, nx + 1) * hx
    y = np.arange(1, ny + 1) * hy
    z = np.arange(1, nz + 1) * hz
    
    # Meshgrid with Fortran-contiguous indexing (x, y, z)
    X, Y, Z = np.meshgrid(x, y, z, indexing='ij')
    
    # Analytical functions
    u_field = np.sin(X) * np.cos(Y) * np.sin(Z)
    v_field = np.cos(X) * np.sin(Y) * np.cos(Z)
    w_field = np.sin(X) * np.sin(Y) * np.cos(Z)
    c_field = np.cos(X) * np.cos(Y) * np.cos(Z)
    
    # Convert fields to double precision (Fortran real*8)
    u_data = u_field.astype(np.float64)
    v_data = v_field.astype(np.float64)
    w_data = w_field.astype(np.float64)
    
    # new_scalar is shaped (nx, ny, nz, 2) in Fortran.
    # new_scalar(:,:,:,1) = 0.0
    # new_scalar(:,:,:,2) = c_field
    new_scalar = np.zeros((nx, ny, nz, 2), dtype=np.float64)
    new_scalar[:, :, :, 1] = c_field + 2.0
    
    # Serialize arrays in Fortran order (Column-major)
    u_bytes = u_data.tobytes(order='F')
    v_bytes = v_data.tobytes(order='F')
    w_bytes = w_data.tobytes(order='F')
    new_scalar_bytes = new_scalar.tobytes(order='F')
    
    # Combined data block record
    data_block_bytes = u_bytes + v_bytes + w_bytes + new_scalar_bytes
    
    # Write to Fortran unformatted file
    with open(output_path, 'wb') as f:
        # Record 1: tnow (double), nx (int64), ny (int64), nz (int64), tmp1 (double), tmp2 (double)
        rec1 = struct.pack('dqqqdd', 0.0, nx, ny, nz, 0.0, 0.0)
        write_fortran_record(f, rec1)
        
        # Record 2: ncyc (int64), dt (double), dummyu (double)
        rec2 = struct.pack('qdd', 1, 0.01, 0.0)
        write_fortran_record(f, rec2)
        
        # Record 3, 4, 5: tmpr(1:2) (double * 2)
        rec3 = struct.pack('dd', 0.0, 0.0)
        write_fortran_record(f, rec3)
        write_fortran_record(f, rec3)
        write_fortran_record(f, rec3)
        
        # Record 6: data arrays (u, v, w, new_scalar)
        write_fortran_record(f, data_block_bytes)
        
    print(f"Successfully generated synthetic grid data file at {output_path}")

if __name__ == '__main__':
    generate_synthetic_grid("fort.1000")