#!/usr/bin/env python3
"""
D3Q19 SRT — Cylinder Wake Flow (Periodic Z)
=============================================
Tests the v2 modular kernel (kernel_v2.cu) in 3D.
Cylinder axis along z, parabolic inlet, pressure outlet, no-slip y-walls.
Produces cross-section visualizations at z=NZ/2.

Usage:
    python tests/test_d3q19_cylinder.py --device 3
    python tests/test_d3q19_cylinder.py --device 3 --re 200 --steps 50000
"""

import sys, os, argparse, time, struct

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'src'))

import numpy as np
import pycuda.driver as cuda

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

from CelerisLab.cuda import compiler

# ━━━━━━━━━━━━━━━━━━━━━━━  Configuration  ━━━━━━━━━━━━━━━━━━━━━━━
# Reasonable 3D grid — fits in < 500 MB GPU memory
NX, NY, NZ = 256, 128, 32
NQ  = 19
NT  = 128
DIM = 3
RHO = 1.0
CYL_CX, CYL_CY = 64.0, 64.0   # Cylinder center (x,y)
CYL_R = 12.0                   # Cylinder radius
TOTAL = NX * NY * NZ

# D3Q19 paired direction ordering (from descriptors.cuh)
#   0:rest  (1,2)±x  (3,4)±y  (5,6)±z
#   (7,8)±(x+y) (9,10)±(x+z) (11,12)±(y+z)
#   (13,14)±(x-y) (15,16)±(x-z) (17,18)±(y-z)
CX = np.array([0, 1,-1, 0, 0, 0, 0, 1,-1, 1,-1, 0, 0, 1,-1, 1,-1, 0, 0], dtype=np.int32)
CY = np.array([0, 0, 0, 1,-1, 0, 0, 1,-1, 0, 0, 1,-1,-1, 1, 0, 0, 1,-1], dtype=np.int32)
CZ = np.array([0, 0, 0, 0, 0, 1,-1, 0, 0, 1,-1, 1,-1, 0, 0,-1, 1,-1, 1], dtype=np.int32)
W  = np.array([1/3] + [1/18]*6 + [1/36]*12, dtype=np.float32)

FLUID_FLAG = 0x01
SOLID_FLAG = 0x02
OBSTACLE_FLAG = 0x04   # triggers half-way BB at adjacent fluid nodes

# ━━━━━━━━━━━━━━━━━━━━━━━  Helpers  ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def compute_vis_omega(Re, D, U0):
    """Compute viscosity and omega from target Reynolds number."""
    vis = U0 * D / Re
    omega = 1.0 / (3.0 * vis + 0.5)
    return vis, omega


def configure_macros_3d(vis, u0, n_objs=0):
    """Write macros.h for D3Q19 SRT."""
    lines = compiler.read_lines(compiler.kernel_path("macros.h"))
    defs = {
        'MULT_GPU': 'False', 'NT': NT,
        'X_1U': NX, 'Y_1U': NY, 'Z_1U': NZ,
        'LBtype': 'float',
        'UX': 1, 'UY': 1, 'UZ': 1,
        'NX': NX, 'NY': NY, 'NZ': NZ,
        'DIM': DIM, 'NQ': NQ,
        'VIS': f'{vis:.10f}', 'RHO': f'{RHO}', 'U0': u0,
        'N_OBJS': n_objs,
        'COLLISION_MODEL': 0,     # SRT
        'STREAMING_MODEL': 0,     # double-buffer
        'STORE_PRECISION': 0,     # FP32
        'USE_DDF_SHIFTING': 0,
    }
    for name, val in defs.items():
        lines = compiler.modify_macro(lines, name, val)
    compiler.write_lines(compiler.kernel_path("macros.h"), lines)


def build_cylinder_3d(cx, cy, radius):
    """Build flag array for 3D cylinder (axis along z, periodic z)."""
    flag = np.ones(TOTAL, dtype=np.uint8) * FLUID_FLAG
    for z in range(NZ):
        for y in range(NY):
            for x in range(NX):
                k = z * NY * NX + y * NX + x
                # Channel walls
                if y == 0 or y == NY - 1 or x == 0 or x == NX - 1:
                    flag[k] = SOLID_FLAG
                # Cylinder body (obstacle — triggers BB at fluid neighbors)
                elif (x - cx)**2 + (y - cy)**2 < radius**2:
                    flag[k] = OBSTACLE_FLAG
    return flag


def extract_fields_3d(ddf_host, z_slice):
    """Extract rho, u, v, w at a given z-slice from D3Q19 DDF (v2 ordering)."""
    # DDF layout: f[i * TOTAL + k]  where k = z*NY*NX + y*NX + x
    f = ddf_host.reshape(NQ, NZ, NY, NX)
    fz = f[:, z_slice, :, :]  # shape (NQ, NY, NX)
    rho = np.sum(fz, axis=0)
    ux = np.zeros_like(rho)
    uy = np.zeros_like(rho)
    uz_field = np.zeros_like(rho)
    for i in range(NQ):
        ux += CX[i] * fz[i]
        uy += CY[i] * fz[i]
        uz_field += CZ[i] * fz[i]
    ux /= rho
    uy /= rho
    uz_field /= rho
    return rho, ux, uy, uz_field


# ━━━━━━━━━━━━━━━━━━━━━━━  Simulation  ━━━━━━━━━━━━━━━━━━━━━━━━━━

def run_d3q19(device_id, n_steps, vis, u0, flag_host):
    """Compile v2 kernel, run D3Q19 SRT, return DDF."""
    omega = 1.0 / (3.0 * vis + 0.5)
    cuda.init()
    dev = cuda.Device(device_id)
    ctx = dev.make_context()
    print(f"[GPU {device_id}] {dev.name()}")

    try:
        configure_macros_3d(vis, u0)
        compiler.compile_kernel_v2()
        ptx_path = compiler.kernel_path("kernel_v2.ptx")
        mod = cuda.module_from_file(ptx_path)

        # Get kernels (extern "C" entries from kernel_v2.cu)
        init_fn = mod.get_function("InitTubeFlow_v2")
        step_fn = mod.get_function("OneStep")

        # Set d_params.omega via __constant__ memory
        params_ptr, params_size = mod.get_global("d_params")
        # LBMParams struct layout (see params.cuh):
        # Nx(4) Ny(4) Nz(4) N(8) omega(4) omega_bulk(4) fx(4) fy(4) fz(4)
        # rho_ref(4) u_inlet(4) n_objects(4)
        # Pack: unsigned int Nx, Ny, Nz; unsigned long N; float omega, omega_bulk, fx, fy, fz, rho_ref, u_inlet; unsigned int n_objects
        params_data = struct.pack('IIIQfffffffI',
                                  NX, NY, NZ,
                                  TOTAL,
                                  omega, 0.0,      # omega, omega_bulk
                                  0.0, 0.0, 0.0,   # fx, fy, fz
                                  RHO, u0,          # rho_ref, u_inlet
                                  0)                 # n_objects
        # Pad to match struct size
        if len(params_data) < params_size:
            params_data += b'\x00' * (params_size - len(params_data))
        cuda.memcpy_htod(params_ptr, params_data)

        # Allocate
        nbytes_ddf = TOTAL * NQ * 4
        ddf_gpu  = cuda.mem_alloc(nbytes_ddf)
        temp_gpu = cuda.mem_alloc(nbytes_ddf)
        flag_gpu = cuda.mem_alloc(flag_host.nbytes)
        indx_gpu = cuda.mem_alloc(TOTAL * 4)
        delta_gpu = cuda.mem_alloc(4)
        action_gpu = cuda.mem_alloc(4)
        obs_gpu = cuda.mem_alloc(4)

        # Dummy arrays
        cuda.memset_d32(indx_gpu, 0, TOTAL)
        cuda.memset_d32(delta_gpu, 0, 1)
        cuda.memset_d32(action_gpu, 0, 1)
        cuda.memset_d32(obs_gpu, 0, 1)

        block = (NT, 1, 1)
        grid  = (NX // NT, NY, NZ)

        # Initialize parabolic flow
        init_fn(flag_gpu, ddf_gpu, block=block, grid=grid)
        ctx.synchronize()

        # Overwrite flags with cylinder geometry
        cuda.memcpy_htod(flag_gpu, flag_host)

        # Step loop
        print(f"  Running {n_steps} steps (NX={NX}, NY={NY}, NZ={NZ}, omega={omega:.4f})...")
        t0 = time.time()
        for i in range(n_steps):
            step_fn(flag_gpu, ddf_gpu, temp_gpu, indx_gpu, delta_gpu,
                    action_gpu, obs_gpu,
                    block=block, grid=grid)
            ddf_gpu, temp_gpu = temp_gpu, ddf_gpu
            if (i + 1) % 5000 == 0:
                ctx.synchronize()
                elapsed = time.time() - t0
                mlups = TOTAL * (i + 1) / elapsed / 1e6
                print(f"    step {i+1}/{n_steps}  ({mlups:.1f} MLUPS)")
        ctx.synchronize()
        dt = time.time() - t0
        mlups = TOTAL * n_steps / dt / 1e6
        print(f"  Done: {dt:.1f}s, {mlups:.1f} MLUPS")

        # Copy back
        ddf = np.zeros(TOTAL * NQ, dtype=np.float32)
        cuda.memcpy_dtoh(ddf, ddf_gpu)
        return ddf

    finally:
        ctx.pop()


# ━━━━━━━━━━━━━━━━━━━━━━━  Visualization  ━━━━━━━━━━━━━━━━━━━━━━━

def plot_d3q19_cylinder(ddf, flag, Re, u0, out_path):
    """4-panel figure at z=NZ/2: vel-mag, vorticity, streamlines, u(y) profile."""
    z_mid = NZ // 2
    rho, ux, uy, uz = extract_fields_3d(ddf, z_mid)
    vel_mag = np.sqrt(ux**2 + uy**2 + uz**2)

    mask2d = ((flag.reshape(NZ, NY, NX)[z_mid] & (SOLID_FLAG | OBSTACLE_FLAG)) != 0)
    vel_masked = np.ma.array(vel_mag, mask=mask2d)

    fig, axes = plt.subplots(2, 2, figsize=(16, 10))

    # (a) Velocity magnitude
    ax = axes[0, 0]
    im = ax.imshow(vel_masked, origin='lower', aspect='auto',
                   cmap='jet', extent=[0, NX, 0, NY])
    circ = plt.Circle((CYL_CX, CYL_CY), CYL_R, fill=True, color='gray', alpha=0.7)
    ax.add_patch(circ)
    plt.colorbar(im, ax=ax, label='|u|')
    ax.set_title(f'Velocity magnitude (z={z_mid})')
    ax.set_xlabel('x'); ax.set_ylabel('y')

    # (b) Vorticity  ω_z = ∂v/∂x − ∂u/∂y
    ax = axes[0, 1]
    dvdx = np.gradient(uy, axis=1)
    dudy = np.gradient(ux, axis=0)
    omega_z = dvdx - dudy
    om_masked = np.ma.array(omega_z, mask=mask2d)
    vmax = np.percentile(np.abs(omega_z[~mask2d]), 99) if np.any(~mask2d) else 1e-3
    im = ax.imshow(om_masked, origin='lower', aspect='auto',
                   cmap='RdBu_r', extent=[0, NX, 0, NY],
                   vmin=-vmax, vmax=vmax)
    circ2 = plt.Circle((CYL_CX, CYL_CY), CYL_R, fill=True, color='gray', alpha=0.7)
    ax.add_patch(circ2)
    plt.colorbar(im, ax=ax, label='ω_z')
    ax.set_title('Vorticity ω_z')
    ax.set_xlabel('x'); ax.set_ylabel('y')

    # (c) Streamlines
    ax = axes[1, 0]
    X, Y = np.meshgrid(np.arange(NX), np.arange(NY))
    ux_c = ux.copy(); ux_c[mask2d] = 0
    uy_c = uy.copy(); uy_c[mask2d] = 0
    speed = np.sqrt(ux_c**2 + uy_c**2)
    ax.streamplot(X, Y, ux_c, uy_c, color=speed, cmap='jet',
                  density=2.5, linewidth=0.7)
    circ3 = plt.Circle((CYL_CX, CYL_CY), CYL_R, fill=True, color='gray', alpha=0.7)
    ax.add_patch(circ3)
    ax.set_xlim(0, NX); ax.set_ylim(0, NY)
    ax.set_aspect('auto')
    ax.set_title('Streamlines')
    ax.set_xlabel('x'); ax.set_ylabel('y')

    # (d) u_x(y) profiles at different x stations
    ax = axes[1, 1]
    y_arr = np.arange(NY)
    # Analytical parabolic inlet
    yy = (y_arr - 0.5 * (NY - 1)) / (NY - 2.0)
    u_ana = u0 * 1.5 * (1 - 4 * yy**2)

    x_stations = [NX // 8, NX // 4, NX // 2, 3 * NX // 4]
    for xs in x_stations:
        ax.plot(ux[:, xs], y_arr, label=f'x={xs}')
    ax.plot(u_ana, y_arr, 'k--', lw=1.5, label='Analytical inlet')
    ax.set_xlabel('u_x'); ax.set_ylabel('y')
    ax.set_title('u_x(y) profiles')
    ax.legend(fontsize=8)
    ax.grid(True, alpha=0.3)

    fig.suptitle(f'D3Q19 SRT Cylinder — Re={Re:.0f}, D={2*CYL_R:.0f}, '
                 f'Grid={NX}×{NY}×{NZ}', fontsize=13)
    fig.tight_layout()
    fig.savefig(out_path, dpi=150)
    print(f"  Saved: {out_path}")
    plt.close(fig)


# ━━━━━━━━━━━━━━━━━━━━━━━  Main  ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def main():
    parser = argparse.ArgumentParser(description='D3Q19 SRT Cylinder Flow')
    parser.add_argument('--device', type=int, default=0,
                        help='CUDA device ID (default: 0)')
    parser.add_argument('--re', type=float, default=100.0,
                        help='Reynolds number based on diameter (default: 100)')
    parser.add_argument('--u0', type=float, default=0.04,
                        help='Inlet characteristic velocity (default: 0.04)')
    parser.add_argument('--steps', type=int, default=30000,
                        help='Number of LBM steps (default: 30000)')
    args = parser.parse_args()

    out_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'output')
    os.makedirs(out_dir, exist_ok=True)

    D = 2 * CYL_R
    vis, omega = compute_vis_omega(args.re, D, args.u0)
    print(f"\n===== D3Q19 SRT Cylinder Flow =====")
    print(f"  Re = {args.re:.0f},  D = {D:.0f},  U0 = {args.u0}")
    print(f"  ν  = {vis:.6f},  ω = {omega:.4f}")
    if omega > 1.95:
        print(f"  WARNING: omega={omega:.4f} is close to 2.0, stability may be poor.")
        print(f"  Consider reducing U0 or Re.")

    flag = build_cylinder_3d(CYL_CX, CYL_CY, CYL_R)
    n_solid = np.sum(flag == SOLID_FLAG)
    n_fluid = np.sum(flag == FLUID_FLAG)
    print(f"  Grid: {NX}×{NY}×{NZ} = {TOTAL} cells  (fluid: {n_fluid}, solid: {n_solid})")
    print(f"  Memory: ~{2 * TOTAL * NQ * 4 / 1e6:.0f} MB for double-buffer DDF")

    ddf = run_d3q19(args.device, args.steps, vis, args.u0, flag)

    plot_d3q19_cylinder(ddf, flag, args.re, args.u0,
                        os.path.join(out_dir, f'cylinder_d3q19_re{int(args.re)}.png'))

    print("\nDone.")


if __name__ == '__main__':
    main()
