# CelerisLab/lbm/driver.py

import pycuda.driver as cuda
import numpy as np
import struct
from scipy.special import jv, expi
from typing import List, Tuple, Union, Optional

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

FLUID = 0b00000001
SOLID = 0b00000010
GAS = 0b00000100
INTERFACE = 0b00001000
SENSOR = 0b00010000
V_TAYLOR = np.int32(1)

class FlowField:
    def __init__(
        self,
        field_config: utils.FlowFieldConfig,
        cuda_config: utils.CudaConfig,
        device_id: Union[int, List[int]] = None,
    ):
        self.field_config = field_config
        self.cuda_config = cuda_config
        cuda.init()

        # Sanity checks
        if cuda_config.multi_gpu:
            if device_id is None or isinstance(device_id, int):
                raise ValueError("Multi-GPU support requires a list of device IDs.")
            # self.devices = [cuda.Device(id) for id in device_id]
            raise NotImplementedError("Multi-GPU support is not implemented yet.")
        else:
            if isinstance(device_id, list):
                if len(device_id) > 1:
                    raise ValueError(
                        "Single-GPU mode does not support multiple device IDs."
                    )
                device_id = device_id[0]
            elif device_id is None:
                device_id = 0
            utils.check_cuda_device_availability(device_id)
            self.device = cuda.Device(device_id)
            self.context = self.device.make_context()

        utils.check_cuda_capability(field_config, cuda_config, device_id)

        # Config kernel
        compiler.config_kernal(cuda_config, field_config)
        compiler.config_object(int(0))
        # compiler.config_sensor(int(0))

        # Set constants
        if field_config.data_type == "FP32":
            self.DATA_TYPE = np.float32
        else:
            raise ValueError(f"Unsupported data type {field_config.data_type}.")

        self.FIELD_SHAPE = tuple(
            size * unit
            for size, unit in zip(
                field_config.field_dim_in_U, cuda_config.unit_dimensions
            )
        )
        self.FIELD_SIZE = np.prod(self.FIELD_SHAPE)
        self.LATTICE = field_config.lattice
        self.DIM = field_config.dimensionality
        if field_config.lattice == 9 and field_config.dimensionality == 2:
            self.E = np.array(
                [0, 0, 1, 0, 0, 1, -1, 0, 0, -1, 1, 1, -1, 1, -1, -1, 1, -1],
                dtype=np.int32,
            ).reshape(9, 2)
            self.OPP = np.array([0, 3, 4, 1, 2, 7, 8, 5, 6], dtype=np.int32)
            self.WW = np.array(
                [4 / 9, 1 / 9, 1 / 9, 1 / 9, 1 / 9, 1 / 36, 1 / 36, 1 / 36, 1 / 36],
                dtype=self.DATA_TYPE,
            )
        else:
            raise NotImplementedError(
                f"Unsupported lattice type {field_config.lattice} in {field_config.dimensionality} dimensions."
            )

        # Compile kernel
        compiler.compile_kernel()
        self.ptx = cuda.module_from_file(compiler.kernel_path("kernel.ptx"))
        self.step = self.ptx.get_function("OneStep")
        initflow = self.ptx.get_function("InitTubeFlow")

        # Initialize memory
        self.ddf = np.zeros(self.FIELD_SIZE * self.LATTICE, dtype=self.DATA_TYPE)
        self.ddf_save = np.zeros(self.FIELD_SIZE * self.LATTICE, dtype=self.DATA_TYPE)
        self.flag = np.ones(self.FIELD_SIZE, dtype=np.uint8)
        self.indx = np.zeros(self.FIELD_SIZE, dtype=np.int32)
        self.delta_curve = np.zeros(0, dtype=self.DATA_TYPE)
        self.vortex_config = np.zeros(7, dtype=float)

        self.ddf_gpu = cuda.mem_alloc(self.ddf.nbytes)
        self.temp_gpu = cuda.mem_alloc(self.ddf.nbytes)
        self.flag_gpu = cuda.mem_alloc(self.flag.nbytes)
        self.indx_gpu = cuda.mem_alloc(self.indx.nbytes)
        self.delta_gpu = cuda.mem_alloc(1)
        self.vortex_gpu = cuda.mem_alloc(self.vortex_config.nbytes)

        self.objects = {}
        self.action = np.zeros(0, dtype=self.DATA_TYPE)
        self.obs = np.zeros(0, dtype=self.DATA_TYPE)

        initflow(
            self.flag_gpu,
            self.ddf_gpu,
            block=(self.cuda_config.threads_per_block, 1, 1),
            grid=(
                int(self.FIELD_SHAPE[0] / self.cuda_config.threads_per_block),
                int(self.FIELD_SHAPE[1]),
                int(self.FIELD_SHAPE[2]),
            ),
        )
        cuda.memcpy_dtoh(self.flag, self.flag_gpu)
        cuda.memcpy_dtoh(self.ddf, self.ddf_gpu)

    def add_cylinder(self, center: Tuple[float, float, float], radius: float, id_obj: Optional[int] = None):
        x_c, y_c, z_c = center

        if (
            x_c - radius <= 0
            or x_c + radius >= self.FIELD_SHAPE[0] - 1
            or y_c - radius <= 0
            or y_c + radius >= self.FIELD_SHAPE[1] - 1
        ):
            raise ValueError("Cylinder is out of bounds.")

        index = self.delta_curve.size if self.delta_curve.size > 0 else 0

        if self.DATA_TYPE == np.float32:
            id_object = np.int32(len(self.objects))
            # max_id = max(self.objects.keys())
        else:
            raise ValueError(f"Unsupported data type {self.DATA_TYPE}.")
        
        for x in range(int(x_c - radius) - 1, int(x_c + radius) + 1):
            for y in range(int(y_c - radius) - 1, int(y_c + radius) + 1):
                if (x - x_c) ** 2 + (y - y_c) ** 2 < radius**2:
                    k = x + y * self.FIELD_SHAPE[0]
                    self.flag[k] = SOLID
                    delta_temp = np.zeros(11, dtype=self.DATA_TYPE)
                    delta_temp[0] = id_object.view(self.DATA_TYPE)
                    for i in range(self.LATTICE):
                        x_neb = x + self.E[i][0]
                        y_neb = y + self.E[i][1]
                        if (x_neb - x_c) ** 2 + (y_neb - y_c) ** 2 >= radius**2:
                            self.flag[k] |= INTERFACE
                            x_i, y_i = preproc.find_circle_intersection(
                                x, y, x_neb, y_neb, x_c, y_c, radius
                            )
                            d_neb = np.sqrt((x_i - x_neb) ** 2 + (y_i - y_neb) ** 2)
                            delta_temp[i] = d_neb / np.sqrt(
                                self.E[i][0] ** 2 + self.E[i][1] ** 2
                            )
                    if self.flag[k] & INTERFACE:
                        delta_temp[9] = (y_c - y) / radius
                        delta_temp[10] = (x - x_c) / radius
                        self.delta_curve = np.concatenate(
                            (self.delta_curve, delta_temp)
                        )
                        self.indx[k] = index
                        index += delta_temp.size

        self.objects[id_object] = {
            "type": "cylinder",
            "center": center,
            "radius": radius,
        }

        if hasattr(self, "delta_gpu"):
            self.delta_gpu.free()
        self.delta_gpu = cuda.mem_alloc(self.delta_curve.nbytes)

        self.action = np.zeros(len(self.objects), dtype=self.DATA_TYPE)
        if hasattr(self, "action_gpu"):
            self.action_gpu.free()
        self.action_gpu = cuda.mem_alloc(self.action.nbytes)

        self.obs = np.zeros(len(self.objects) * self.DIM, dtype=self.DATA_TYPE)
        if hasattr(self, "obs_gpu"):
            self.obs_gpu.free()
        self.obs_gpu = cuda.mem_alloc(self.obs.nbytes)

        cuda.memcpy_htod(self.delta_gpu, self.delta_curve)
        cuda.memcpy_htod(self.flag_gpu, self.flag)
        cuda.memcpy_htod(self.indx_gpu, self.indx)

        compiler.config_object(len(self.objects))
        compiler.compile_kernel()
        self.ptx = cuda.module_from_file(compiler.kernel_path("kernel.ptx"))
        self.step = self.ptx.get_function("OneStep")

    def add_sensor(self, center: Tuple[float, float, float], radius: float):
        x_c, y_c, z_c = center

        if (
            x_c - radius <= 0
            or x_c + radius >= self.FIELD_SHAPE[0] - 1
            or y_c - radius <= 0
            or y_c + radius >= self.FIELD_SHAPE[1] - 1
        ):
            raise ValueError("Sensor is out of bounds.")

        id_object = len(self.objects)
        for x in range(int(x_c - radius) - 1, int(x_c + radius) + 1):
            for y in range(int(y_c - radius) - 1, int(y_c + radius) + 1):
                if (x - x_c) ** 2 + (y - y_c) ** 2 < radius**2:
                    k = x + y * self.FIELD_SHAPE[0]
                    self.flag[k] |= SENSOR
                    self.indx[k] = id_object

        self.objects[id_object] = {
            "type": "sensor",
            "center": center,
        }

        self.action = np.zeros(len(self.objects), dtype=self.DATA_TYPE)
        if hasattr(self, "action_gpu"):
            self.action_gpu.free()
        self.action_gpu = cuda.mem_alloc(self.action.nbytes)

        self.obs = np.zeros(len(self.objects) * self.DIM, dtype=self.DATA_TYPE)
        if hasattr(self, "force_gpu"):
            self.obs_gpu.free()
        self.obs_gpu = cuda.mem_alloc(self.obs.nbytes)

        cuda.memcpy_htod(self.flag_gpu, self.flag)
        cuda.memcpy_htod(self.indx_gpu, self.indx)

        compiler.config_object(len(self.objects))
        compiler.compile_kernel()
        self.ptx = cuda.module_from_file(compiler.kernel_path("kernel.ptx"))
        self.step = self.ptx.get_function("OneStep")

    def add_vortex(self, center: Tuple[float, float, float], radius: float, strength: float, direction: float, type: str):
        x_c, y_c, z_c = center

        if (
            x_c - radius <= 0
            or x_c + radius >= self.FIELD_SHAPE[0] - 1
            or y_c - radius <= 0
            or y_c + radius >= self.FIELD_SHAPE[1] - 1
        ):
            raise ValueError("Vortex is out of bounds.")
        
        if type not in ["lamb", "oseen", "taylor"]:
            raise ValueError("Vortex type" + type + " is not supported.")
        
        x = np.linspace(-x_c, self.FIELD_SHAPE[0] - 1 - x_c, self.FIELD_SHAPE[0])
        y = np.linspace(-y_c, self.FIELD_SHAPE[1] - 1 - y_c, self.FIELD_SHAPE[1])
        X, Y = np.meshgrid(x, y)
        r = np.sqrt(X**2 + Y**2)
        nu = self.field_config.viscosity
        theta = np.arctan2(Y, X)
        psi = np.zeros_like(r)

        if type == "lamb":
            b = 3.831705970207512
            n = b / radius
            u0 = strength
            inside = r <= radius
            outside = r > radius

            psi[inside] = (2 * u0 / n / jv(0, b) * jv(1, n * r[inside]) - u0 * r[inside]) * np.sin(theta[inside])
            psi[outside] = -u0 * radius**2 / r[outside] * np.sin(theta[outside])

            u_vor = np.gradient(psi, axis=0)
            v_vor = -np.gradient(psi, axis=1)
            p_vor = -2 * (np.gradient(v_vor, axis=1) - np.gradient(u_vor, axis=0)) * psi - (u_vor**2 + v_vor**2) / 2
        elif type == "oseen":
            # 4 nu t = radius^2 / 4
            kappa = 2 * np.pi * radius **2 * strength
            u_vor = - kappa / (2 * np.pi * r) * (1 - np.exp(-4 * r**2 / radius**2)) * np.sin(theta)
            v_vor = kappa / (2 * np.pi * r) * (1 - np.exp(-4 * r**2 / radius**2)) * np.cos(theta)
            zeta = 4 * r**2 / radius**2
            p_vor = -kappa**2 / 8 / np.pi**2 / r**2 * (-2 * zeta * (expi(-zeta) - expi(-2 * zeta)) + (1 - np.exp(-zeta))**2)
        elif type == "taylor":
            # 4 nu t = radius^2
            M = strength * np.pi * radius**4 / 8 / nu
            u_vor = - M * r * 4 * nu / radius**4 * np.exp(-r**2 / radius**2) * np.sin(theta)
            v_vor = M * r * 4 * nu / radius**4 * np.exp(-r**2 / radius**2) * np.cos(theta)
            p_vor = -4 * M**2 * nu**2 * np.exp(-2 * r**2 / radius**2) / np.pi**2 / radius**6

        cuda.memcpy_dtoh(self.ddf, self.ddf_gpu)
        ddf_temp = self.ddf.copy().reshape((self.LATTICE, self.FIELD_SHAPE[1], self.FIELD_SHAPE[0])).transpose(2, 1, 0)
        u_ddf = ddf_temp[:, :, 1] + ddf_temp[:, :, 5] + ddf_temp[:, :, 8] - ddf_temp[:, :, 3] - ddf_temp[:, :, 6] - ddf_temp[:, :, 7]
        v_ddf = ddf_temp[:, :, 2] + ddf_temp[:, :, 5] + ddf_temp[:, :, 6] - ddf_temp[:, :, 4] - ddf_temp[:, :, 7] - ddf_temp[:, :, 8]
        p_ddf = np.sum(ddf_temp, axis=2) / 3

        for i in range(self.FIELD_SHAPE[0]):
            for j in range(self.FIELD_SHAPE[1]):
                k = i + j * self.FIELD_SHAPE[0]
                if (j == 0 or j == self.FIELD_SHAPE[1] - 1) or (i == 0 or i == self.FIELD_SHAPE[0] - 1):
                    continue
                else:
                    for e in range(self.LATTICE):
                        u = u_ddf[i, j] + u_vor[j, i]
                        v = v_ddf[i, j] + v_vor[j, i]
                        p = p_ddf[i, j] + p_vor[j, i]
                        eu = self.E[e][0] * u + self.E[e][1] * v
                        u2 = u ** 2 + v ** 2
                        self.ddf[k + e * self.FIELD_SIZE] = self.WW[e] * (3 * p + 3 * eu + 4.5 * eu ** 2 - 1.5 * u2)
            
        cuda.memcpy_htod(self.ddf_gpu, self.ddf)

    # def add_vortex_gpu(self, center: Tuple[float, float, float], radius: float, strength: float, direction: float, type: str):
    #     x_c, y_c, z_c = center

    #     if (
    #         x_c - radius <= 0
    #         or x_c + radius >= self.FIELD_SHAPE[0] - 1
    #         or y_c - radius <= 0
    #         or y_c + radius >= self.FIELD_SHAPE[1] - 1
    #     ):
    #         raise ValueError("Vortex is out of bounds.")
        
    #     if type not in ["lamb", "oseen", "taylor"]:
    #         raise ValueError("Vortex type" + type + " is not supported.")
        
    #     add_vortex = self.ptx.get_function("AddVortex")

    #     self.vortex_config[0:3] = np.array(center, dtype=float)
    #     self.vortex_config[3] = radius
    #     self.vortex_config[4] = strength
    #     self.vortex_config[5] = direction
    #     if type == "taylor":
    #         self.vortex_config[6] = 

    def run(self, num_steps: int, action_target: np.ndarray):
        if (
            action_target.size != len(self.objects)
            or action_target.dtype != self.DATA_TYPE
        ):
            raise ValueError("action data type or size does not match the objects.")
        elif len(self.objects) == 0:
            raise ValueError("No objects have been added to the flow field.")

        weight = 0.1
        stream = cuda.Stream()
        action_pinned = cuda.pagelocked_empty_like(self.action)
        action_pinned[:] = self.action
        obs_pinned = cuda.pagelocked_empty_like(self.obs)
        self.obs[:] = 0
        for i in range(num_steps):
            action_pinned = (1 - weight) * action_pinned + weight * action_target
            cuda.memcpy_htod_async(self.action_gpu, action_pinned, stream)
            self.step(
                self.flag_gpu,
                self.ddf_gpu,
                self.temp_gpu,
                self.indx_gpu,
                self.delta_gpu,
                self.action_gpu,
                self.obs_gpu,
                block=(self.cuda_config.threads_per_block, 1, 1),
                grid=(
                    int(self.FIELD_SHAPE[0] / self.cuda_config.threads_per_block),
                    int(self.FIELD_SHAPE[1]),
                    int(self.FIELD_SHAPE[2]),
                ),
                stream=stream,
            )
            self.ddf_gpu, self.temp_gpu = self.temp_gpu, self.ddf_gpu
            cuda.memcpy_dtoh_async(obs_pinned, self.obs_gpu, stream)
            cuda.memset_d32_async(self.obs_gpu, 0, self.obs.size, stream)
            self.obs += obs_pinned
        stream.synchronize()
        self.obs = (self.obs / num_steps).astype(self.DATA_TYPE)

    def apply_ddf(self):
        cuda.memcpy_htod(self.ddf_gpu, self.ddf)

    def get_ddf(self):
        cuda.memcpy_dtoh(self.ddf, self.ddf_gpu)

    def save_ddf(self):
        self.ddf_save = self.ddf.copy()

    def restore_ddf(self):
        self.ddf = self.ddf_save.copy()

    def __del__(self):
        self.context.pop()