# Graph of parameter dependencies.
#
# 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 .parameter import (
    PythonInputParameter, PythonOutputParameter,
    TemporaryPythonVariable, CTypeParameterByValue,
    CTypeParameterByRef, CppParameter, Forward
)

_VG_Parameter = PythonInputParameter | PythonOutputParameter \
    | TemporaryPythonVariable | CTypeParameterByValue | CTypeParameterByRef \
    | CppParameter



class VariableGraphNode:
    parameter: _VG_Parameter
    depends_on: "list[VariableGraphNode | str] | None"
    dependents: "list[VariableGraphNode]"
    vid: int

    def __init__(self,
            parameter: _VG_Parameter,
            depends_on: "list[VariableGraphNode | str] | None",
            vid: int
        ):
        self.parameter = parameter
        self.depends_on = depends_on
        self.dependents = []
        self.vid = vid


    def __hash__(self) -> int:
        return self.vid

    def __eq__(self, other: object) -> bool:
        if isinstance(other, VariableGraphNode):
            return self.vid == other.vid
        return False


class VariableGraph:
    """
    A graph of variable interdependencies.
    """
    root: list[VariableGraphNode]

    # Index the parameters by their names.
    name2p: dict[
        str,
        VariableGraphNode
    ]

    # Unbound variables, i.e. those whose parents
    # are known by name only so far.
    # Index them by their unknown parent's name.
    unbound: dict[str, list[VariableGraphNode]]

    def __init__(self):
        self.root = []
        self.name2p = {}
        self.unbound = {}


    def __len__(self) -> int:
        return len(self.name2p)


    def __getitem__(self, name: str) -> _VG_Parameter:
        res = self.name2p.get(name)
        if res is None:
            raise KeyError(
                f"Parameter '{name}' unknown."
            )
        return res.parameter


    def __contains__(self, name: str) -> bool:
        return name in self.name2p


    def get_node(self, name: str) -> VariableGraphNode:
        res = self.name2p.get(name)
        if res is None:
            raise KeyError(
                f"Parameter '{name}' unknown."
            )
        return res


    def insert(self,
            parameter: _VG_Parameter,
            *depends_on: VariableGraphNode | str
        ) -> VariableGraphNode:
        # The node's name:
        name = parameter.name
        if isinstance(name, Forward):
            # Forward a parameter:
            pname = name
            name = self.get_temporary_name(pname.name, "fwd_")
            pname.varname = name

        if name in self.name2p:
            raise RuntimeError(
                f"Parameter '{name}' collides with previously defined "
                "parameter of same name."
            )

        # Resolve the dependencies:
        #  - str,  if has unknown parent
        #  - VariableGraphNode, if has known parent
        def try_resolve_dependency(
                dependency: VariableGraphNode | str
            ) -> VariableGraphNode | str:
            if isinstance(dependency, VariableGraphNode):
                return dependency
            return self.name2p.get(dependency, dependency)

        dependencies = [try_resolve_dependency(dep) for dep in depends_on]

        # Create the node and insert it by name:
        node = VariableGraphNode(parameter, dependencies, len(self))
        self.name2p[name] = node

        # Act on the parent node (or make it a root node):
        if len(depends_on) == 0:
            self.root.append(node)
        else:
            # Handle the dependencies:
            for dep in depends_on:
                if isinstance(dep, str):
                    self._add_unbound_dependecy(dep, node)
                else:
                    dep.dependents.append(node)

        # Check if any nodes depend on this node.
        # If so, these node(s) now have one further known
        # dependency, and we can swap the name string with
        # this actual node.
        dependents = self.unbound.pop(name, None)
        if dependents is not None:
            for dep in dependents:
                assert dep.depends_on is not None
                i = dep.depends_on.index(name)
                dep.depends_on[i] = node


        # Return the newly inserted node.
        return node


    def _add_unbound_dependecy(self,
            dependency: str,
            dependent: VariableGraphNode
        ):
        siblings = self.unbound.get(dependency)
        if siblings:
            siblings.append(dependent)
        else:
            siblings = [dependent]
            self.unbound[dependency] = siblings


    def add_dependency(self,
            dependent: str | VariableGraphNode,
            dependency: str | VariableGraphNode
        ):
        # Resolve dependent and dependency.
        # If the latter is not possible, work with an
        # unbound dependency.
        if isinstance(dependent, str):
            dependent = self.get_node(dependent)
        if isinstance(dependency, str):
            resolved = self.name2p.get(dependency)
            if resolved is None:
                # If we cannot resolve, we have an unbound
                # dependency:
                self._add_unbound_dependecy(dependency, dependent)
            else:
                dependency = resolved

        # Add the dependency to the dependent:
        if dependent.depends_on is None:
            dependent.depends_on = [dependency]
        else:
            dependent.depends_on.append(dependency)




    def derive_initialization_order(self) -> list[_VG_Parameter]:
        """
        Derive an order in which the parameters can be initialized.
        """
        current_batch: list[VariableGraphNode] = [*self.root]
        next_batch: list[VariableGraphNode] = []
        result: list[_VG_Parameter] = []
        nodes_handled: set[VariableGraphNode] = set()
        nodes_seen: set[VariableGraphNode] = set()

        # To check whether the current batch changed:
        def batch_hash() -> int:
            return hash(tuple(hash(vgn) for vgn in current_batch))

        # Adding new nodes only to the next batch:
        def add_new(
                dependents: list[VariableGraphNode],
                next_batch: list[VariableGraphNode]
            ):
            newly_seen = [
                dep for dep in dependents if dep not in nodes_seen
            ]
            nodes_seen.update(newly_seen)
            next_batch.extend(newly_seen)


        # This is a fairly simple, O(N²) worst case algorithm
        # to create a valid initialization order. If many nodes
        # depend only on a single parent, the algorithm will run
        # in O(N).
        last_hash = batch_hash()
        while len(current_batch) > 0:
            for vgn in current_batch:
                # We might have already handled this parameter:
                if vgn.depends_on is None:
                    result.append(vgn.parameter)
                    nodes_handled.add(vgn)
                    add_new(vgn.dependents, next_batch)
                    continue
                if all(s in nodes_handled for s in vgn.depends_on):
                    result.append(vgn.parameter)
                    nodes_handled.add(vgn)
                    add_new(vgn.dependents, next_batch)
                    continue
                # Node still has an undiscovered dependency.
                next_batch.append(vgn)

            # Now clear the current batch and swap with the next batch:
            current_batch.clear()
            current_batch, next_batch = next_batch, current_batch

            # Check if we made any kind of progress:
            next_hash = batch_hash()
            if last_hash == next_hash:
                raise RuntimeError(
                    "Did not make progress deriving the initialization "
                    "order. This might indicate that there is a loop-dependency "
                    "in the variable initializatoin."
                )
            last_hash = next_hash

        return result


    def name_in_use(self, name: str) -> bool:
        """
        Check whether a variable name is already in use.
        """
        return name in self.name2p


    def get_temporary_name(self,
            basename: str,
            prefix: str = "tmp_"
        ) -> str:
        name = prefix + basename
        i = 0
        while name in self.name2p:
            name = f"{basename}_{i}"
            i += 1
        return name
