# CelerisLab/cuda/compiler_v2.py
"""
Kernel configuration, compilation, and module loading.

Workflow:
  1. generate_config(lbm_cfg, n_objects) — write config/*.h from LBMConfig
  2. compile_kernel(arch)               — nvcc -ptx kernel_v2.cu
  3. load_module(ptx_path)              — PyCUDA module from PTX
"""

import os
import subprocess
from typing import Optional

import pycuda.driver as cuda

from ..config import LBMConfig
from .obs_layout import obs_layout  # re-exported for body.manager import path

# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------
_KERNEL_DIR = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
    "lbm", "kernels",
)


def kernel_path(filename: str) -> str:
    return os.path.join(_KERNEL_DIR, filename)


# ---------------------------------------------------------------------------
# Header generation (f-string templates, no external deps)
# ---------------------------------------------------------------------------
_HEADER = "// AUTO-GENERATED by CelerisLab compiler – DO NOT EDIT MANUALLY\n"


def _write(path: str, content: str):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, "w") as f:
        f.write(content)


def generate_config(cfg: LBMConfig, n_objects: int = 0):
    """Render all config/*.h files from *cfg*."""
    m = cfg.to_macros()
    cdir = os.path.join(_KERNEL_DIR, "config")

    _write(os.path.join(cdir, "config_grid.h"), f"""{_HEADER}\
// Layer: Global / Grid
#ifndef CELERIS_CONFIG_GRID_H
#define CELERIS_CONFIG_GRID_H

#define NT   {m['NT']}
#define MULT_GPU {m['MULT_GPU']}

#define NX  {m['NX']}
#define NY  {m['NY']}
#define NZ  {m['NZ']}

// ---- Lattice model (single source of truth) ----
// Model IDs: 0=D2Q9, 1=D3Q19  (extend here for future models like D2Q19)
#define LATTICE_D2Q9   0
#define LATTICE_D3Q19  1

#define LATTICE_MODEL  {m['LATTICE_MODEL']}

// DIM and NQ are derived from LATTICE_MODEL — do NOT set independently.
#if LATTICE_MODEL == LATTICE_D2Q9
  #define DIM 2
  #define NQ  9
#elif LATTICE_MODEL == LATTICE_D3Q19
  #define DIM 3
  #define NQ  19
#else
  #error "Unknown LATTICE_MODEL. Supported: LATTICE_D2Q9 (0), LATTICE_D3Q19 (1)."
#endif

#endif
""")

    _write(os.path.join(cdir, "config_physics.h"), f"""{_HEADER}\
// Layer: Global / Physics
#ifndef CELERIS_CONFIG_PHYSICS_H
#define CELERIS_CONFIG_PHYSICS_H

#define LBtype {m['LBtype']}
#define VIS {m['VIS']}
#define RHO {m['RHO']}
#define U0  {m['U0']}

#define PI 3.141592653589793238

// Flag constants are defined in core/flags.cuh (uint16_t).
// Do NOT redefine them here.

#define V_TAYLOR  1

#endif
""")

    _write(os.path.join(cdir, "config_method.h"), f"""{_HEADER}\
// Layer: Method
#ifndef CELERIS_CONFIG_METHOD_H
#define CELERIS_CONFIG_METHOD_H

#define COLLISION_MODEL  {m['COLLISION_MODEL']}
#define STREAMING_MODEL  {m['STREAMING_MODEL']}
#define STORE_PRECISION  {m['STORE_PRECISION']}
#define USE_DDF_SHIFTING {m['USE_DDF_SHIFTING']}

#define USE_LES {m['USE_LES']}
#define LES_CS  {m['LES_CS']}
#define LES_CLOSED_FORM {m['LES_CLOSED_FORM']}

#define INLET_PROFILE        {m['INLET_PROFILE']}
#define INLET_SCHEME         {m['INLET_SCHEME']}
#define OUTLET_MODE          {m['OUTLET_MODE']}
#define OUTLET_BLEND_ALPHA   {m['OUTLET_BLEND_ALPHA']}
#define OUTLET_BACKFLOW_CLAMP {m['OUTLET_BACKFLOW_CLAMP']}
#define Y_WALL_BC            {m['Y_WALL_BC']}

#define OMEGA_COLLISION_MIN {m['OMEGA_COLLISION_MIN']}
#define OMEGA_COLLISION_MAX {m['OMEGA_COLLISION_MAX']}
#define TRT_MAGIC_PARAM     {m['TRT_MAGIC_PARAM']}

// NEQ damping coefficients for inlet/outlet BC reconstruction.
// TRT inlet: donor damping used by the channel_stabilized inlet.
// Regularized inlet: damping on incoming-direction donor NEQ.
// SRT outlet: damped outlet reconstruction to suppress checkerboard noise.
#define INLET_TRT_NEQ_DAMP         {m['INLET_TRT_NEQ_DAMP']}
#define INLET_REGULARIZED_NEQ_DAMP {m['INLET_REGULARIZED_NEQ_DAMP']}
#define OUTLET_SRT_NEQ_DAMP        {m['OUTLET_SRT_NEQ_DAMP']}

#endif
""")

    _write(os.path.join(cdir, "config_objects.h"), f"""{_HEADER}\
// Layer: Case / Objects
#ifndef CELERIS_CONFIG_OBJECTS_H
#define CELERIS_CONFIG_OBJECTS_H

#define N_OBJS {n_objects}

#endif
""")

    _write(os.path.join(cdir, "config_obs.h"), f"""{_HEADER}\
// Layer: Case / Objects — packed obs (force + torque + sensor), float indices not bytes
// Derived from N_OBJS (config_objects.h) and DIM (config_grid.h); include via config.h.
#ifndef CELERIS_CONFIG_OBS_H
#define CELERIS_CONFIG_OBS_H

#if (N_OBJS) >= 1
#define OBS_N_SLOTS (N_OBJS)
#else
#define OBS_N_SLOTS 1
#endif

#define OBS_SLOT_STRIDE_FLOATS ((OBS_N_SLOTS) * (DIM))
#define OBS_FORCE_FLOATS       (OBS_SLOT_STRIDE_FLOATS)

#if DIM == 2
#define OBS_TORQUE_COMPONENTS  1
#else
#define OBS_TORQUE_COMPONENTS  (DIM)
#endif

#define OBS_TORQUE_FLOATS      ((OBS_N_SLOTS) * (OBS_TORQUE_COMPONENTS))
#define OBS_BODY_SLOT_FLOATS   (3 * (DIM))

#define OBS_FORCE0_FLOATS      0
#define OBS_TORQUE0_FLOATS     (OBS_FORCE0_FLOATS + OBS_FORCE_FLOATS)
#define OBS_SENSOR0_FLOATS     (OBS_TORQUE0_FLOATS + OBS_TORQUE_FLOATS)
#define OBS_TOTAL_FLOATS       (OBS_SENSOR0_FLOATS + OBS_SLOT_STRIDE_FLOATS)

#endif
""")


# ---------------------------------------------------------------------------
# Compilation
# ---------------------------------------------------------------------------
def compile_kernel(arch: str = "sm_70",
                   source: str = "kernel_v2.cu",
                   output: Optional[str] = None) -> str:
    """Compile *source* to PTX.  Returns path to the .ptx file.

    The translation unit includes ``config/*.h``. Those headers must match the
    intended ``LBMConfig`` — call :func:`generate_config` before compiling.
    Invoking ``nvcc`` on ``kernel_v2.cu`` alone can pick up stale checked-in
    headers from ``lbm/kernels/config/``.
    """
    src = kernel_path(source)
    if output is None:
        output = source.replace(".cu", ".ptx")
    dst = kernel_path(output)
    import shutil
    nvcc = shutil.which("nvcc") or "/usr/local/cuda/bin/nvcc"
    subprocess.run(
        [nvcc, "-ptx", f"-arch={arch}", "--use_fast_math", src, "-o", dst],
        check=True,
    )
    return dst


# ---------------------------------------------------------------------------
# Module loading
# ---------------------------------------------------------------------------
def load_module(ptx_path: Optional[str] = None) -> cuda.Module:
    """Load a compiled PTX as a PyCUDA module."""
    if ptx_path is None:
        ptx_path = kernel_path("kernel_v2.ptx")
    return cuda.module_from_file(ptx_path)


# ---------------------------------------------------------------------------
# Back-compat helpers (used by test scripts)
# ---------------------------------------------------------------------------

def compile_kernel_v2(arch: str = "sm_70") -> str:
    """Alias for compile_kernel() kept for test-script compatibility."""
    return compile_kernel(arch=arch)
