# Wrapping NumPy arrays and QuantityArrays.
#
# Author: Malte J. Ziebarth (malte.ziebarth@tum.de)
#
# Copyright (C) 2026 Technische Universität München
#
# Licensed under the EUPL, Version 1.2 or – as soon they will be approved by
# the European Commission - subsequent versions of the EUPL (the "Licence");
# You may not use this work except in compliance with the Licence.
# You may obtain a copy of the Licence at:
#
# https://joinup.ec.europa.eu/collection/eupl/eupl-text-eupl-12
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the Licence is distributed on an "AS IS" basis,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the Licence for the specific language governing permissions and
# limitations under the Licence.
#
# SPDX-FileCopyrightText: 2026 Technical University of Munich
# SPDX-License-Identifier: EUPL-1.2
# SPDX-FileContributor: Malte J. Ziebarth <malte.ziebarth@tum.de>





from typing import Callable
from warnings import warn
from ..language.abstract import Pointer, Output, Input, Temporary, Const
from .parameter import CTypeParameterByValue, PythonInputParameter, \
    PythonOutputParameter, TemporaryPythonVariable, CppParameter
from ..language.arrays import NDArray, QuantityArray
from ..language.conversion import NUMPY_DTYPE_MAP
from ..language.strings import cpp_type_string
from ..language.python import CTypeType
from .variablegraph import VariableGraph, VariableGraphNode
from .string import OutputStringBuilder
from ..parse.sizecontrol import SizeControl
from ..tools.views import append


#
# 1. NumPy ndarray
# ================
#
# Wrap a `funktionsklempner::ndarray::NDArray` into a NumPy ndarray.
# This consists of two main functions:
#
#   - `wrap_output_ndarray`, to create an output `ndarray`
#   - `wrap_input_ndarray`, to create an `ndarray` input parameter
#
# The following helpers are used:
#
#   - `_InitParams`: an argument struct for assembling the init code
#     of an `ndarray` parameter.
#   - `_wrap_ndarray`: the workhorse function for assembling all
#     Python, ctypes, and intermediary parameters, and inserting them
#     in right order into the variable graph.

class _InitParams:
    name: str
    np_dtype: str
    shape: str

    def __init__(self, name: str, np_dtype: str, shape: str):
        self.name = name
        self.np_dtype = np_dtype
        self.shape = shape

class NewArrayInitializer:
    """
    Helper class for obtaining the shape of an array.
    """
    sc_name: tuple[str, ...]
    vargraph: VariableGraph

    def __init__(self,
            controlled: str,
            size_control: SizeControl,
            vargraph: VariableGraph
        ):
        self.sc_name = size_control.get_control(controlled)
        self.vargraph = vargraph


    def shape(self) -> tuple[str, tuple[VariableGraphNode, ...]] | None:
        sc_vgn: list[VariableGraphNode] = []
        get_shape: list[str] = []
        for sc_name in self.sc_name:
            if sc_name not in self.vargraph:
                # Size control shall be a new variable.
                size_control = PythonInputParameter(
                    sc_name, Input("int"),
                    [
                        f"    {sc_name} = int({sc_name})"
                    ]
                )
                sc_vgn.append(self.vargraph.insert(size_control))
                get_shape.append(sc_name)
                warn(
                    f"Created a new size parameter '{sc_name}', which was given "
                    "as an argument to the FUNKTIONSKLEMPNER_SIZE_CONTROL macro, "
                    "because no parameter of this name was found in the "
                    "argument list."
                )
            else:
                # The size control is an existing parameter.
                # It has to be an array parameter!
                sc_vgn.append(self.vargraph.get_node(sc_name))
                size_control = sc_vgn[-1].parameter
                if not isinstance(size_control, PythonInputParameter):
                    raise RuntimeError(
                        f"Size control parameter '{sc_name}' has to be a "
                        "Python input parameter."
                    )
                size_control_type = size_control.dtype.val
                if not isinstance(size_control_type, (NDArray, QuantityArray)):
                    raise RuntimeError(
                        f"Size control parameter '{sc_name}' has to be a "
                        "NDArray or QuantityArray."
                    )
                get_shape.append(f"*{sc_name}.shape")

        return f"({"".join(append(get_shape, ","))})", tuple(sc_vgn)



class ForwardInitializer:
    def shape(self) -> None:
        return None


def empty_init_code(ip: _InitParams) -> list[str]:
    return [
        f"    {ip.name} = np.empty({ip.shape}, {ip.np_dtype}, 'C')"
    ]

def forward_init_code(ip: _InitParams) -> list[str]:
    return [
        f"    {ip.name} = np.asarray({ip.name}, {ip.np_dtype}, 'C')"
    ]


class NDArrayWrapper:
    """
    Helper class for wrapping an NDArray.
    """
    name: str
    buffer_name: str

    python_parameter: tuple[
        PythonInputParameter |PythonOutputParameter | TemporaryPythonVariable,
        VariableGraphNode
    ]
    buffer_parameter: tuple[CTypeParameterByValue, VariableGraphNode]
    size_parameter: tuple[CTypeParameterByValue, VariableGraphNode]
    cpp_parameter: tuple[CppParameter, VariableGraphNode] | None

    def __init__(
            self,
            name: str,
            parameter_type: Output[NDArray] | Input[NDArray] | Temporary[NDArray],
            vargraph: VariableGraph,
            initializer: NewArrayInitializer | ForwardInitializer,
            get_init_code: Callable[[_InitParams], list[str]],
            fk_namespace: str,
            cpp_argument_id: int | None,
            size_name: str | None = None
        ):
        self.name = name
        # 0. Get the buffer type:
        array_type = parameter_type.val
        self.buffer_name = vargraph.get_temporary_name(name, 'buf_')
        buffer_ctype = array_type.ctypes_type()
        np_dtype = NUMPY_DTYPE_MAP[buffer_ctype]

        # 1. Get the shape of the array.
        shape_node = initializer.shape()
        if shape_node is None:
            # No shape (output parameter.)
            shape = ""
            sc_vgn = tuple()
        else:
            shape, sc_vgn = shape_node

        # 2. The Python output parameter
        init_code: list[str] = get_init_code(
            _InitParams(name=name, np_dtype=np_dtype, shape=shape)
        )
        buffer_ptr: Pointer[CTypeType]
        if isinstance(parameter_type, Output):
            py = PythonOutputParameter(
                name, parameter_type, init_code, None
            )

            # Output parameters need to be initialized:
            assert shape != ""
            assert len(sc_vgn) > 0

            # Buffer type:
            buffer_ptr = Pointer(buffer_ctype)

        elif isinstance(parameter_type, Input):
            py = PythonInputParameter(
                name, parameter_type, init_code
            )

            # Input parameters shall not be initialized:
            assert shape == ""
            assert len(sc_vgn) == 0

            # Buffer type:
            buffer_ptr = Pointer(Const(buffer_ctype))

        else:
            py = TemporaryPythonVariable(
                name, parameter_type.val, init_code
            )

            # Buffer type:
            buffer_ptr = Pointer(buffer_ctype)

        vgn_py = vargraph.insert(py, *sc_vgn)
        self.python_parameter = (py, vgn_py)

        # 3. The buffer ctypes parameter
        # Now that we know the name, create the initialization.
        # First assert that we have a numpy array of correct data type
        # given as input:
        init_code = [
            f"    {self.buffer_name} = {name}.ctypes.data_as(POINTER({buffer_ctype}))",
        ]
        c_buffer = CTypeParameterByValue(
            self.buffer_name,
            buffer_ptr,
            init_code
        )
        vgn_c_buffer = vargraph.insert(c_buffer, vgn_py)
        self.buffer_parameter = (c_buffer, vgn_c_buffer)


        # 4. The buffer size ctypes parameter
        if size_name is None:
            size_name = vargraph.get_temporary_name(name, prefix='size_')
        init_code = [
            f"    {size_name} = c_size_t({name}.size)"
        ]
        c_size = CTypeParameterByValue(
            size_name,
            "c_size_t",
            init_code
        )
        vgn_last = vargraph.insert(c_size, vgn_c_buffer)
        self.size_parameter = (c_size, vgn_last)


        # 5. The C++ parameter that collects the two ctypes arguments:
        if not isinstance(parameter_type, Temporary):
            assert cpp_argument_id is not None
            cpp_name = vargraph.get_temporary_name(name, prefix='c++_')
            cpp_type = cpp_type_string(array_type.dtype)
            if isinstance(parameter_type, Input):
                ndtype = "InputNDArray"
            else:
                ndtype = "OutputNDArray"
            init_code = [
                "            "
                f"auto {name} = {fk_namespace}::{ndtype}<{cpp_type}>(\n"
                "            "
                f"    {self.buffer_name},\n"
                "            "
                f"    {size_name},\n"
                "            "
                f"    \"{cpp_type}\"\n"
                "            "
                ");\n"
            ]
            cpp = CppParameter(
                cpp_name,
                parameter_type,
                init_code,
                None,
                name,
                cpp_argument_id
            )
            vgn_last = vargraph.insert(cpp, vgn_last)
            self.cpp_parameter = (cpp, vgn_last)

        else:
            self.cpp_parameter = None


    def last_variable_graph_node(self) -> VariableGraphNode:
        if self.cpp_parameter is None:
            return self.size_parameter[1]
        return self.cpp_parameter[1]






def wrap_output_ndarray(
        name: str,
        parameter_type: Output[NDArray],
        vargraph: VariableGraph,
        size_control_dict: SizeControl,
        fk_namespace: str,
        cpp_argument_id: int,
        size_name: str | None = None
    ) -> VariableGraphNode:

    # 2. The Python output parameter
    array = NDArrayWrapper(
        name,
        parameter_type,
        vargraph,
        NewArrayInitializer(name, size_control_dict, vargraph),
        empty_init_code,
        fk_namespace,
        cpp_argument_id,
        size_name=size_name
    )
    assert array.cpp_parameter is not None
    return array.cpp_parameter[1]




def wrap_input_ndarray(
        name: str,
        parameter_type: Input[NDArray],
        vargraph: VariableGraph,
        fk_namespace: str,
        cpp_argument_id: int
    ):
    """
    Wraps a Python input NDArray.
    """
    # Heavy lifting in the backend:
    NDArrayWrapper(
        name,
        parameter_type,
        vargraph,
        ForwardInitializer(),
        forward_init_code,
        fk_namespace,
        cpp_argument_id
    )




#
# 2. QuantityArray
# ================
#
# Here we do the same for the `QuantityArray` class.
# We use the facility for wrapping the ndarray and augment
# it by passing the unit string.

def _generate_cpp_qa_parameter(
        name: str,
        buffer_name: str,
        size_name: str,
        unit_name: str | None,
        unit_id_name: str | None,
        parameter_type: Input[QuantityArray] | Output[QuantityArray],
        vargraph: VariableGraph,
        root: VariableGraphNode,
        fk_namespace: str,
        cpp_argument_id: int,
    ) -> VariableGraphNode:
    """
    Generates the C++ QuantityArray parameter.
    """
    # 0. Get the buffer type:
    array_type = parameter_type.val
    bu = array_type.boost_unit

    cpp_name = vargraph.get_temporary_name(name, prefix='c++_')
    cpp_type = array_type.dtype
    if isinstance(parameter_type, Input):
        assert unit_name is not None
        assert unit_id_name is None
        # Get rid of the 'Const' we infused in parse.variable.parse_variable:
        assert isinstance(cpp_type, Const)
        cpp_type = cpp_type.val
        qatype = "InputQuantityArray"
        unit_initializer = f",\n                {unit_name}\n"
        post_call_code = None
        init_code = []
    else:
        assert unit_name is None
        assert unit_id_name is not None
        qatype = "OutputQuantityArray"
        unit_initializer = "\n"
        post_call_code = [
            f"*{unit_id_name} = "
            f"{fk_namespace}::create_string("
            f"{fk_namespace}::default_unit<{bu}>"
             ")"
        ]
        init_code = [
             "            "
            f"if ({unit_id_name} == nullptr)\n",
             "            "
            f"    throw std::runtime_error(\n"
             "            "
            f"        \"Internal error: output array '{name}' \"\n"
             "            "
            f"        \"has been passed a nullptr as output size parameter.\"\n"
             "            "
             "    );\n"
        ]
    cpp_type_str = cpp_type_string(cpp_type)
    init_code.append(
         "            "
        f"auto {name} = {fk_namespace}::{qatype}<{bu}, {cpp_type_str}>(\n"
         "            "
        f"    {buffer_name},\n"
         "            "
        f"    {size_name},\n"
         "            "
        f"    \"{cpp_type_str}\""
        f"{unit_initializer}"
         "            "
         ");\n"
    )



    cpp = CppParameter(
        cpp_name,
        parameter_type,
        init_code,
        post_call_code,
        name,
        cpp_argument_id
    )
    vgn_last = vargraph.insert(cpp, root)
    return vgn_last




def wrap_input_quantityarray(
        name: str,
        parameter_type: Input[QuantityArray],
        vargraph: VariableGraph,
        fk_namespace: str,
        cpp_argument_id: int
    ):
    """
    Wrap an input QuantityArray.
    """
    py = PythonInputParameter(
        name,
        parameter_type,
        [
            f"    assert isinstance({name}, QuantityArray)"
        ]
    )
    vgn_qa = vargraph.insert(py)

    # Get a temporary variable to unwrap the QuantityArray's ndarray:
    nd_name = vargraph.get_temporary_name(name, "asarray_")
    nd_type = NDArray(parameter_type.val.dtype)
    size_name = vargraph.get_temporary_name(name, prefix='size_')

    def init_code(ip: _InitParams) -> list[str]:
        return [
            f"    {ip.name} = np.asarray({name}.array, {ip.np_dtype}, 'C')"
        ]

    # Heavy lifting in the backend:
    array = NDArrayWrapper(
        nd_name,
        Temporary(nd_type),
        vargraph,
        ForwardInitializer(),
        init_code,
        fk_namespace,
        None,
        size_name=size_name
    )

    # Now monkey-patch the input parameter as a parent:
    vargraph.add_dependency(
        array.python_parameter[1],
        vgn_qa
    )

    # Then add the unit string after the size parameter:
    unit_name = vargraph.get_temporary_name(name)
    unit = CTypeParameterByValue(
        unit_name,
        Const("c_char_p"),
        [
            f"    assert isinstance({name}.unit, str)",
            f"    {unit_name} = c_char_p({name}.unit.encode('utf-8'))"
        ]
    )
    vgn_last = vargraph.insert(unit, array.last_variable_graph_node())


    # The C++ parameter that collects the three ctypes arguments:
    _generate_cpp_qa_parameter(
        name,
        array.buffer_name,
        size_name,
        unit_name,
        None,
        parameter_type,
        vargraph,
        vgn_last,
        fk_namespace,
        cpp_argument_id
    )








def wrap_output_quantityarray(
        name: str,
        parameter_type: Output[QuantityArray],
        vargraph: VariableGraph,
        size_control_dict: SizeControl,
        fk_namespace: str,
        cpp_argument_id: int
    ):
    """
    Wrap an output QuantityArray.
    """

    # Get temporary variables
    #  - to generate the QuantityArray's ndarray,
    #  - to obtain the QuantityArray's unit.
    nd_name = vargraph.get_temporary_name(name, "nd_")
    nd_type = NDArray(parameter_type.val.dtype)
    unit_name = vargraph.get_temporary_name(name, "unit_")
    size_name = vargraph.get_temporary_name(name, prefix='size_')

    # Set the size control for the NDArray to the same controller
    # as the quantity array:
    size_control_dict.set_control(
        nd_name,
        size_control_dict.get_control(name)
    )



    # Create a temporary NDArray:
    array = NDArrayWrapper(
        nd_name,
        Temporary(nd_type),
        vargraph,
        NewArrayInitializer(nd_name, size_control_dict, vargraph),
        empty_init_code,
        fk_namespace,
        cpp_argument_id,
        size_name=size_name
    )
    vgn_last = array.last_variable_graph_node()

    # Output string:
    osb = OutputStringBuilder(unit_name, vargraph, vgn_last, 0)
    _, vgn_str = osb.create_ctypes()
    post_call = osb.ctypes_post_call_code()
    osb.id_name


    # After we have read the string, we need to assemble the final
    # QuantityArray:
    post_call.append(
        f"    {name} = QuantityArray({nd_name}, {unit_name})"
    )

    # Finally generate the QuantityArray.
    # This happens purely post-call.
    py_qa = PythonOutputParameter(
        name,
        parameter_type,
        [],
        post_call
    )
    vgn_last = vargraph.insert(py_qa, vgn_str)


    # The C++ parameter that collects the three ctypes arguments:
    _generate_cpp_qa_parameter(
        name,
        array.buffer_name,
        size_name,
        None,
        osb.id_name,
        parameter_type,
        vargraph,
        vgn_last,
        fk_namespace,
        cpp_argument_id
    )
