# CelerisLab/driver.py

import pycuda.driver as cuda
import numpy as np
from typing import List, Tuple, Union

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

FLUID = 0b00000001
SOLID = 0b00000010
GAS = 0b00000100
INTERFACE = 0b00001000
SENSOR = 0b00010000


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.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.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):
        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))
        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, "force_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 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.")

        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()