# CelerisLab/kernels/compiler.py

import subprocess
import re
import os

from .utils import FlowFieldConfig, CudaConfig


def kernel_path(file_name: str) -> str:
    current_dir = os.path.dirname(os.path.abspath(__file__))
    return os.path.join(current_dir, "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 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_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)