# CelerisLab/utils.py

import pycuda.driver as cuda
import subprocess
import json

from typing import NamedTuple, Optional, List, Tuple, Union


class CudaDeviceInfo(NamedTuple):
    name: str
    compute_capability: str
    multiprocessors: int
    total_global_memory: int
    max_shared_memory_per_block: int
    max_threads_per_block: int
    max_blocks_per_multiprocessor: int
    device_interconnect: Optional[str] = None


class FlowFieldConfig(NamedTuple):
    data_type: str
    dimensionality: int
    lattice: int
    field_dim_in_U: Tuple[int, int, int]
    viscosity: float
    velocity: float
    boundary_conditions: Tuple[str, str, str, str, str, str]


class CudaConfig(NamedTuple):
    multi_gpu: bool
    gpu_connection: str
    required_cuda_capability: str
    threads_per_block: int
    unit_dimensions: Tuple[int, int, int]


def check_cuda_device_availability(device_id=0):
    if cuda.Device.count() == 0:
        raise RuntimeError("No CUDA device is available.")

    if device_id < 0 or device_id >= cuda.Device.count():
        raise ValueError(
            f"Invalid device_id {device_id}. Must be between 0 and {cuda.Device.count() - 1}."
        )

    try:
        subprocess.check_output(["nvidia-smi", "--version"])
    except subprocess.CalledProcessError:
        raise RuntimeError("nvidia-smi is not available or not installed correctly.")


def query_cuda_device_info(device_id=0) -> CudaDeviceInfo:
    check_cuda_device_availability(device_id)

    try:
        output = subprocess.check_output(
            ["nvidia-smi", "-q", "-d", "TOPOLOGY", "-i", str(device_id)], text=True
        )
        if "NVLink" in output:
            device_interconnect = "NVLink"
        elif "PCIe" in output:
            device_interconnect = "PCIe"
        else:
            device_interconnect = "Unknown"
    except Exception as e:
        device_interconnect = None

    device = cuda.Device(device_id)

    return CudaDeviceInfo(
        name=device.name(),
        compute_capability=f"{device.compute_capability()[0]}.{device.compute_capability()[1]}",
        multiprocessors=device.get_attribute(
            cuda.device_attribute.MULTIPROCESSOR_COUNT
        ),
        total_global_memory=device.total_memory(),
        max_shared_memory_per_block=device.get_attribute(
            cuda.device_attribute.MAX_SHARED_MEMORY_PER_BLOCK
        ),
        max_threads_per_block=device.get_attribute(
            cuda.device_attribute.MAX_THREADS_PER_BLOCK
        ),
        max_blocks_per_multiprocessor=device.get_attribute(
            cuda.device_attribute.MAX_BLOCKS_PER_MULTIPROCESSOR
        ),
        device_interconnect=device_interconnect,
    )


def load_flow_field_config(config_path: str) -> FlowFieldConfig:
    try:
        with open(config_path, "r") as file:
            config = json.load(file)

        required_keys = [
            "data_type",
            "dimensionality",
            "lattice",
            "field_dim_in_U",
            "viscosity",
            "boundary_conditions",
        ]
        if not all(key in config for key in required_keys):
            raise ValueError("Missing required configuration items.")

        if config["data_type"] not in ["FP32", "FP64"]:
            raise ValueError("Data type must be either FP32 or FP64.")

        if config["dimensionality"] not in [2, 3]:
            raise ValueError("Dimensionality must be either 2 or 3.")

        if config["dimensionality"] == 2 and config["field_dim_in_U"][2] != 1:
            raise ValueError(
                "Field dimensions must be 1 in the third dimension for 2D simulations."
            )

        if config["lattice"] not in [9]:
            raise ValueError("Lattice must be either 9 or 19.")

        boundary_conditions = tuple(
            condition
            for key in ["x", "y", "z"]
            for condition in config["boundary_conditions"].get(key, [])
        )
        if len(boundary_conditions) != 6:
            raise ValueError("Boundary conditions must contain exactly six elements.")

        return FlowFieldConfig(
            data_type=config["data_type"],
            dimensionality=config["dimensionality"],
            lattice=config["lattice"],
            field_dim_in_U=tuple(config["field_dim_in_U"]),
            viscosity=config["viscosity"],
            velocity=config["velocity"],
            boundary_conditions=boundary_conditions,
        )
    except Exception as e:
        raise RuntimeError(f"Failed to load or parse the flow field configuration: {e}")


def load_cuda_config(config_path: str) -> CudaConfig:
    try:
        with open(config_path, "r") as file:
            config = json.load(file)

        required_keys = [
            "multi_gpu",
            "gpu_connection",
            "required_cuda_capability",
            "threads_per_block",
            "X_1U",
            "Y_1U",
            "Z_1U",
        ]

        if not all(key in config for key in required_keys):
            raise ValueError("Missing required configuration items.")

        return CudaConfig(
            multi_gpu=config["multi_gpu"],
            gpu_connection=config["gpu_connection"],
            required_cuda_capability=config["required_cuda_capability"],
            threads_per_block=config["threads_per_block"],
            unit_dimensions=(config["X_1U"], config["Y_1U"], config["Z_1U"]),
        )
    except Exception as e:
        raise RuntimeError(f"Failed to load or parse the CUDA configuration: {e}")


def check_cuda_capability(
    field_config: FlowFieldConfig,
    cuda_config: CudaConfig,
    device_id: Union[int, List[int]] = None,
):
    SAFE_FACTOR = 0.8

    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.")
        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
        device_info = query_cuda_device_info(device_id)

        if device_info.compute_capability != cuda_config.required_cuda_capability:
            raise ValueError(
                f"Device {device_info.name} has compute capability {device_info.compute_capability}, but {cuda_config.required_cuda_capability} is required."
            )

        field_size = sum(
            size * unit
            for size, unit in zip(
                field_config.field_dim_in_U, cuda_config.unit_dimensions
            )
        )
        if (
            device_info.total_global_memory * SAFE_FACTOR
            < calc_field_memory_consumption(
                field_size,
                field_config.dimensionality,
                field_config.lattice,
                field_config.data_type,
            )
        ):
            raise ValueError(
                f"Device {device_info.name} does not have enough memory to store the flow field."
            )

        if (
            device_info.max_threads_per_block * SAFE_FACTOR
            < cuda_config.threads_per_block
        ):
            raise ValueError(
                f"Device {device_info.name} does not have enough threads per block to run the simulation."
            )

        block_size = cuda_config.threads_per_block
        if (
            device_info.max_shared_memory_per_block * SAFE_FACTOR
            < 2
            * calc_field_memory_consumption(
                block_size,
                field_config.dimensionality,
                field_config.lattice,
                field_config.data_type,
            )
        ):
            raise ValueError(
                f"Device {device_info.name} does not have enough shared memory per block to run the simulation."
            )


def calc_field_memory_consumption(
    field_size: int, dimensionality: int, directions: int, data_type: str
) -> int:
    if data_type == "FP32":
        data_size = 4
    elif data_type == "FP64":
        data_size = 8
    else:
        raise ValueError(f"Unsupported data type {data_type}.")

    return (
        field_size * directions * data_size * 2
        + field_size * dimensionality * data_size
        + field_size
    )
