# CelerisLab/cuda/compiler.py

import subprocess
import re
import os

from ..common.utils import FlowFieldConfig, CudaConfig


def kernel_path(file_name: str) -> str:
    # kernels are in lbm/kernels/
    current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    return os.path.join(current_dir, "lbm", "kernels", file_name)


def read_lines(file_path):
    with open(file_path, "r") as file:
        lines = file.readlines()
    return lines


def write_lines(file_path, lines):
    with open(file_path, "w") as file:
        file.writelines(lines)


def modify_macro(lines, macro_name, new_value):
    pattern = re.compile(rf"(#define\s+{macro_name}\s+)(\S+)")
    for i, line in enumerate(lines):
        if pattern.match(line):
            lines[i] = pattern.sub(rf"\g<1>{new_value}", line)
            break
    return lines

def modify_const(lines, const_name, new_type, new_value):
    pattern = re.compile(rf"(__constant__\s+)(\S+\s+{const_name}\s*=\s*)([^;]+)(;)")
    for i, line in enumerate(lines):
        if pattern.match(line):
            lines[i] = pattern.sub(rf"\g<1>{new_type} {const_name} = {new_value}\4", line)
            break
    return lines

def compile_kernel():
    subprocess.run(
        [
            "nvcc",
            "-ptx",
            kernel_path("kernel.cu"),
            "-o",
            kernel_path("kernel.ptx"),
        ]
    )


def compile_kernel_v2():
    """Compile the new modular kernel (kernel_v2.cu → kernel_v2.ptx)."""
    subprocess.run(
        [
            "nvcc",
            "-ptx",
            kernel_path("kernel_v2.cu"),
            "-o",
            kernel_path("kernel_v2.ptx"),
        ]
    )

def config_kernal(config_cuda: CudaConfig, config_field: FlowFieldConfig):
    lines = read_lines(kernel_path("macros.h"))
    lines = modify_macro(lines, "MULT_GPU", config_cuda.multi_gpu)
    lines = modify_macro(lines, "NT", config_cuda.threads_per_block)
    lines = modify_macro(lines, "X_1U", config_cuda.unit_dimensions[0])
    lines = modify_macro(lines, "Y_1U", config_cuda.unit_dimensions[1])
    lines = modify_macro(lines, "Z_1U", config_cuda.unit_dimensions[2])

    if config_field.data_type == "FP32":
        lb_type = "float"
    else:
        raise ValueError(f"Unsupported data type {config_field.data_type}.")
    lines = modify_macro(lines, "LBtype", lb_type)
    lines = modify_macro(lines, "UX", config_field.field_dim_in_U[0])
    lines = modify_macro(lines, "UY", config_field.field_dim_in_U[1])
    lines = modify_macro(lines, "UZ", config_field.field_dim_in_U[2])
    lines = modify_macro(lines, "NX", config_field.field_dim_in_U[0] * config_cuda.unit_dimensions[0])
    lines = modify_macro(lines, "NY", config_field.field_dim_in_U[1] * config_cuda.unit_dimensions[1])
    lines = modify_macro(lines, "NZ", config_field.field_dim_in_U[2] * config_cuda.unit_dimensions[2])
    lines = modify_macro(lines, "DIM", config_field.dimensionality)
    lines = modify_macro(lines, "NQ", config_field.lattice)
    lines = modify_macro(lines, "VIS", config_field.viscosity)
    lines = modify_macro(lines, "U0", config_field.velocity)

    write_lines(kernel_path("macros.h"), lines)


def config_kernal_v2(config_cuda: CudaConfig, config_field: FlowFieldConfig,
                     collision_model: int = 2,
                     streaming_model: int = 0,
                     store_precision: int = 0,
                     use_ddf_shifting: int = 0,
                     use_les: int = 0,
                     les_cs: float = 0.16):
    """Configure macros.h for the new modular kernel architecture.

    Args:
        collision_model:  0=SRT, 1=TRT, 2=MRT (default)
        streaming_model:  0=double-buffer (default), 1=Esoteric-Pull
        store_precision:  0=FP32 (default), 1=FP16S, 2=FP16C
        use_ddf_shifting: 0=off (default), 1=on
        use_les:          0=off (default), 1=Smagorinsky LES
        les_cs:           Smagorinsky constant C_s
    """
    # First apply legacy config
    config_kernal(config_cuda, config_field)

    # Then apply new architecture macros
    lines = read_lines(kernel_path("macros.h"))
    lines = modify_macro(lines, "COLLISION_MODEL", collision_model)
    lines = modify_macro(lines, "STREAMING_MODEL", streaming_model)
    lines = modify_macro(lines, "STORE_PRECISION", store_precision)
    lines = modify_macro(lines, "USE_DDF_SHIFTING", use_ddf_shifting)
    lines = modify_macro(lines, "USE_LES", use_les)
    lines = modify_macro(lines, "LES_CS", f"{les_cs:.6f}f")
    write_lines(kernel_path("macros.h"), lines)

def config_object(n_obj: int):
    lines = read_lines(kernel_path("macros.h"))
    lines = modify_macro(lines, "N_OBJS", n_obj)
    write_lines(kernel_path("macros.h"), lines)

def config_sensor(n_sen: int):
    lines = read_lines(kernel_path("macros.h"))
    lines = modify_macro(lines, "N_SENS", n_sen)
    write_lines(kernel_path("macros.h"), lines)