#!/usr/bin/env python3
"""
D2Q9 Regression Test — Poiseuille Channel + Cylinder Flow
==========================================================
Uses the ORIGINAL kernel.cu with same grid / BCs.
Produces matplotlib figures for visual validation.

Usage:
    python tests/test_d2q9_visual.py --device 2
    python tests/test_d2q9_visual.py --device 2 --cylinder
"""

import sys, os, argparse, time

# Ensure CelerisLab package is importable
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 matplotlib.colors import Normalize

from CelerisLab.cuda import compiler
from CelerisLab.common import preprocess as preproc

# ━━━━━━━━━━━━━━━━━━━━━━━  Configuration  ━━━━━━━━━━━━━━━━━━━━━━━
NX, NY = 1280, 512
NQ = 9
NT = 128
DIM = 2
VIS = 0.002
U0 = 0.01
RHO = 1.0
TOTAL = NX * NY

# Original direction vectors (const.h ordering)
E = np.array([[0,0],[1,0],[0,1],[-1,0],[0,-1],[1,1],[-1,1],[-1,-1],[1,-1]], dtype=np.int32)
OPP = np.array([0, 3, 4, 1, 2, 7, 8, 5, 6], dtype=np.int32)

FLUID_FLAG = 0b00000001
SOLID_FLAG = 0b00000010
INTERFACE_FLAG = 0b00001000

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

def configure_macros(n_objs=0):
    """Write macros.h to match original kernel settings."""
    lines = compiler.read_lines(compiler.kernel_path("macros.h"))
    defs = {
        'MULT_GPU': 'False', 'NT': NT,
        'X_1U': 128, 'Y_1U': 32, 'Z_1U': 1,
        'LBtype': 'float',
        'UX': 10, 'UY': 16, 'UZ': 1,
        'NX': NX, 'NY': NY, 'NZ': 1,
        'DIM': DIM, 'NQ': NQ,
        'VIS': VIS, 'RHO': f'{RHO}', 'U0': U0,
        'N_OBJS': n_objs,
    }
    for name, val in defs.items():
        lines = compiler.modify_macro(lines, name, val)
    compiler.write_lines(compiler.kernel_path("macros.h"), lines)


def extract_fields(ddf_host):
    """Compute rho, u, v from host DDF (original const.h ordering).

    The original kernel uses DDF-shifting: stores f_shifted = f - w_i*RHO.
    So sum(f_shifted) = rho - RHO  (~0 for incompressible flow),
    and momentum = sum(e_x * f_shifted) works because sum(w_i * e_x) = 0.
    """
    f = ddf_host.reshape(NQ, NY, NX)
    rho = np.sum(f, axis=0) + RHO          # un-shift: rho = sum(f_shifted) + RHO
    u = (f[1] + f[5] + f[8] - f[3] - f[6] - f[7]) / RHO
    v = (f[2] + f[5] + f[6] - f[4] - f[7] - f[8]) / RHO
    return rho, u, v


def analytical_poiseuille(y_arr):
    """Analytical parabolic profile matching InitTubeFlow."""
    yy = (y_arr - 0.5 * (NY - 1)) / (NY - 2.0)
    return U0 * 1.5 * (1 - 4 * yy**2)


def build_cylinder_data(cx, cy, radius):
    """Replicate driver.py add_cylinder logic for flag / delta / indx."""
    flag = np.ones(TOTAL, dtype=np.uint8)   # init all FLUID
    indx = np.zeros(TOTAL, dtype=np.int32)
    delta_list = []
    index_offset = 0

    # Build Poiseuille flag first (walls + solid borders)
    for y in range(NY):
        for x in range(NX):
            k = x + y * NX
            if y == 0 or y == NY - 1 or x == 0 or x == NX - 1:
                flag[k] = SOLID_FLAG

    # Add cylinder
    for x in range(int(cx - radius) - 1, int(cx + radius) + 1):
        for y in range(int(cy - radius) - 1, int(cy + radius) + 1):
            if (x - cx)**2 + (y - cy)**2 < radius**2:
                k = x + y * NX
                flag[k] = SOLID_FLAG
                dt = np.zeros(11, dtype=np.float32)
                dt[0] = np.int32(0).view(np.float32)  # id_object = 0
                has_interface = False
                for i in range(NQ):
                    xn = x + E[i][0]
                    yn = y + E[i][1]
                    if (xn - cx)**2 + (yn - cy)**2 >= radius**2:
                        has_interface = True
                        xi, yi = preproc.find_circle_intersection(
                            x, y, xn, yn, cx, cy, radius)
                        d_neb = np.sqrt((xi - xn)**2 + (yi - yn)**2)
                        e_len = np.sqrt(E[i][0]**2 + E[i][1]**2)
                        if e_len > 0:
                            dt[i] = d_neb / e_len
                if has_interface:
                    flag[k] |= INTERFACE_FLAG
                    dt[9] = (cy - y) / radius
                    dt[10] = (x - cx) / radius
                    indx[k] = index_offset
                    delta_list.append(dt)
                    index_offset += 11

    delta = np.concatenate(delta_list) if delta_list else np.zeros(1, dtype=np.float32)
    return flag, indx, delta


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

def run_simulation(device_id, n_steps, n_objs, flag_host, indx_host, delta_host):
    """Compile kernel, run LBM, return DDF on host."""
    cuda.init()
    dev = cuda.Device(device_id)
    ctx = dev.make_context()
    print(f"[GPU {device_id}] {dev.name()}")

    try:
        configure_macros(n_objs)
        compiler.compile_kernel()
        ptx_path = compiler.kernel_path("kernel.ptx")
        mod = cuda.module_from_file(ptx_path)
        step_fn = mod.get_function("OneStep")
        init_fn = mod.get_function("InitTubeFlow")

        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(indx_host.nbytes)
        delta_gpu = cuda.mem_alloc(max(delta_host.nbytes, 4))

        action_host = np.zeros(max(n_objs, 1), dtype=np.float32)
        obs_host    = np.zeros(max(n_objs * DIM, 1), dtype=np.float32)
        action_gpu = cuda.mem_alloc(action_host.nbytes)
        obs_gpu    = cuda.mem_alloc(obs_host.nbytes)

        cuda.memcpy_htod(action_gpu, action_host)
        cuda.memcpy_htod(obs_gpu, obs_host)

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

        # Init Poiseuille
        init_fn(flag_gpu, ddf_gpu, block=block, grid=grid)
        ctx.synchronize()

        # Overwrite flag / indx / delta for cylinder case
        cuda.memcpy_htod(flag_gpu, flag_host)
        cuda.memcpy_htod(indx_gpu, indx_host)
        cuda.memcpy_htod(delta_gpu, delta_host)

        # Step loop
        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
        ctx.synchronize()
        dt = time.time() - t0
        mlups = TOTAL * n_steps / dt / 1e6
        print(f"  {n_steps} steps in {dt:.2f}s  ({mlups:.1f} MLUPS)")

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

        return ddf, flag_out
    finally:
        ctx.pop()


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

def plot_poiseuille(ddf, flag, out_path):
    """3-panel figure: velocity mag, u(y) profile, pressure along centerline."""
    rho, u, v = extract_fields(ddf)
    vel_mag = np.sqrt(u**2 + v**2)

    # Mask solid cells for display
    mask = (flag.reshape(NY, NX) & SOLID_FLAG).astype(bool)
    vel_mag_masked = np.ma.array(vel_mag, mask=mask)

    fig, axes = plt.subplots(1, 3, figsize=(18, 5))

    # (a) Velocity magnitude heatmap
    ax = axes[0]
    im = ax.imshow(vel_mag_masked, origin='lower', aspect='auto',
                   cmap='jet', extent=[0, NX, 0, NY])
    plt.colorbar(im, ax=ax, label='|u|')
    ax.set_title('Velocity magnitude')
    ax.set_xlabel('x'); ax.set_ylabel('y')

    # (b) u(y) profile at x = NX/2  vs. analytical
    ax = axes[1]
    x_mid = NX // 2
    y_arr = np.arange(NY, dtype=float)
    ax.plot(u[:, x_mid], y_arr, 'b-', lw=2, label='LBM')
    ax.plot(analytical_poiseuille(y_arr), y_arr, 'r--', lw=1.5, label='Analytical')
    ax.set_xlabel('u_x'); ax.set_ylabel('y')
    ax.set_title(f'u(y) at x={x_mid}')
    ax.legend()
    ax.grid(True, alpha=0.3)

    # (c) Pressure along centerline y = NY/2
    ax = axes[2]
    y_mid = NY // 2
    p = rho / 3.0
    ax.plot(np.arange(NX), p[y_mid, :], 'g-', lw=1.5)
    ax.set_xlabel('x'); ax.set_ylabel('p = ρ/3')
    ax.set_title(f'Pressure along centerline (y={y_mid})')
    ax.grid(True, alpha=0.3)

    fig.suptitle(f'D2Q9 Poiseuille – NX={NX}, NY={NY}, VIS={VIS}, U0={U0}', fontsize=13)
    fig.tight_layout()
    fig.savefig(out_path, dpi=150)
    print(f"  Saved: {out_path}")
    plt.close(fig)


def plot_cylinder(ddf, flag, cx, cy, radius, out_path):
    """3-panel figure: velocity magnitude (zoom), vorticity, streamlines."""
    rho, u, v = extract_fields(ddf)
    vel_mag = np.sqrt(u**2 + v**2)
    mask = (flag.reshape(NY, NX) & SOLID_FLAG).astype(bool)

    # Zoom window around cylinder
    pad = int(radius * 8)
    x0 = max(int(cx - pad), 0)
    x1 = min(int(cx + pad * 2), NX)
    y0 = max(int(cy - pad), 0)
    y1 = min(int(cy + pad), NY)

    fig, axes = plt.subplots(1, 3, figsize=(20, 6))

    # (a) Velocity magnitude (zoomed)
    ax = axes[0]
    vm_z = np.ma.array(vel_mag[y0:y1, x0:x1], mask=mask[y0:y1, x0:x1])
    im = ax.imshow(vm_z, origin='lower', aspect='equal',
                   cmap='jet', extent=[x0, x1, y0, y1])
    circ = plt.Circle((cx, cy), radius, fill=True, color='gray', alpha=0.7)
    ax.add_patch(circ)
    plt.colorbar(im, ax=ax, label='|u|')
    ax.set_title('Velocity magnitude')

    # (b) Vorticity
    ax = axes[1]
    dvdx = np.gradient(v, axis=1)
    dudy = np.gradient(u, axis=0)
    omega = dvdx - dudy
    om_z = np.ma.array(omega[y0:y1, x0:x1], mask=mask[y0:y1, x0:x1])
    vmax = np.percentile(np.abs(omega[~mask]), 99)
    im = ax.imshow(om_z, origin='lower', aspect='equal',
                   cmap='RdBu_r', extent=[x0, x1, y0, y1],
                   vmin=-vmax, vmax=vmax)
    circ2 = plt.Circle((cx, cy), radius, fill=True, color='gray', alpha=0.7)
    ax.add_patch(circ2)
    plt.colorbar(im, ax=ax, label='ω')
    ax.set_title('Vorticity')

    # (c) Streamlines
    ax = axes[2]
    X, Y = np.meshgrid(np.arange(x0, x1), np.arange(y0, y1))
    u_z = u[y0:y1, x0:x1].copy()
    v_z = v[y0:y1, x0:x1].copy()
    u_z[mask[y0:y1, x0:x1]] = 0
    v_z[mask[y0:y1, x0:x1]] = 0
    speed = np.sqrt(u_z**2 + v_z**2)
    ax.streamplot(X, Y, u_z, v_z, color=speed, cmap='jet',
                  density=2.0, linewidth=0.8)
    circ3 = plt.Circle((cx, cy), radius, fill=True, color='gray', alpha=0.7)
    ax.add_patch(circ3)
    ax.set_xlim(x0, x1); ax.set_ylim(y0, y1)
    ax.set_aspect('equal')
    ax.set_title('Streamlines')

    fig.suptitle(f'D2Q9 Cylinder Flow – Re_D={U0*1.5*2*radius/VIS:.0f}, D={2*radius}', 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='D2Q9 Regression Test')
    parser.add_argument('--device', type=int, default=2,
                        help='CUDA device ID (default: 2)')
    parser.add_argument('--cylinder', action='store_true',
                        help='Also run cylinder flow test')
    parser.add_argument('--steps-pois', type=int, default=5000,
                        help='Steps for Poiseuille (default: 5000)')
    parser.add_argument('--steps-cyl', type=int, default=30000,
                        help='Steps for cylinder (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)

    # ---- Test 1: Poiseuille ----
    print("\n===== Test 1: Poiseuille Channel Flow =====")
    flag_pois = np.ones(TOTAL, dtype=np.uint8)
    indx_pois = np.zeros(TOTAL, dtype=np.int32)
    delta_pois = np.zeros(1, dtype=np.float32)

    ddf, flag = run_simulation(args.device, args.steps_pois, 0,
                               flag_pois, indx_pois, delta_pois)
    plot_poiseuille(ddf, flag, os.path.join(out_dir, 'poiseuille_d2q9.png'))

    # Error metric
    rho, u, v = extract_fields(ddf)
    y_arr = np.arange(NY, dtype=float)
    u_ana = analytical_poiseuille(y_arr)
    x_mid = NX // 2
    u_num = u[:, x_mid]
    # Interior cells only (skip walls)
    err = np.max(np.abs(u_num[2:-2] - u_ana[2:-2])) / np.max(np.abs(u_ana[2:-2]))
    print(f"  L∞ relative error at x={x_mid}: {err:.2e}")

    # ---- Test 2: Cylinder ----
    if args.cylinder:
        print("\n===== Test 2: Flow Around Cylinder =====")
        cyl_cx, cyl_cy, cyl_r = 256.0, 256.0, 32.0
        flag_cyl, indx_cyl, delta_cyl = build_cylinder_data(cyl_cx, cyl_cy, cyl_r)

        ddf2, flag2 = run_simulation(args.device, args.steps_cyl, 1,
                                     flag_cyl, indx_cyl, delta_cyl)
        plot_cylinder(ddf2, flag2, cyl_cx, cyl_cy, cyl_r,
                      os.path.join(out_dir, 'cylinder_d2q9.png'))

    print("\nDone.")


if __name__ == '__main__':
    main()
