"""A small perceptron node, kept as a worked example of numeric logic."""

from __future__ import annotations

import logging
import threading
from collections.abc import Iterable
from typing import Any

from pydantic import BaseModel, ConfigDict

from fluksio.flow.messages import MessageSpec
from fluksio.flow.nodes.base import Node

logger = logging.getLogger(__name__)


def _numpy() -> Any:
    """numpy, which is the `ml` extra.

    Imported when a node is built rather than at boot, so an instance without
    it still offers every other type, and a node placed without it says so.
    """
    try:
        import numpy
    except ImportError:
        raise RuntimeError(
            "the mlp node needs the ml extra: pip install 'fluksio[ml]', "
            "or build the image with EXTRAS=ml"
        ) from None
    return numpy


class MLPNode(Node):
    """
    Multi-Layer Perceptron node for neural network processing in pipelines.

    This node implements a simple single-layer neural network that applies
    weights and biases to input values. Weights and biases are randomly
    initialized using the provided random number generator.

    The computation follows the standard neural network formula:
    output = weights @ inputs + biases

    :param requires: Input messages consumed by this node.
    :type requires: MessageSpec | list[MessageSpec]
    :param provides: Output messages produced by this node.
    :type provides: MessageSpec | list[MessageSpec]
    :param params: Parameters dict containing:
        - ``rng`` (numpy.random.Generator): Random number generator for weight initialization
        - Additional node parameters
    :type params: dict
    :param name: Name for this node.
    :type name: str

    :example:
        >>> import numpy as np
        >>> rng = np.random.default_rng(seed=42)
        >>> mlp = MLPNode(
        ...     requires=[MessageSpec(name="input1", dtype=DType.FLOAT), MessageSpec(name="input2", dtype=DType.FLOAT)],
        ...     provides=[MessageSpec(name="output", dtype=DType.FLOAT)],
        ...     params={"rng": rng},
        ...     name="mlp_layer1",
        ... )
    """

    class Params(BaseModel):
        """Weights are drawn from ``seed``, so a node reloads identically."""

        model_config = ConfigDict(extra="allow")

        seed: int = 0

    def __init__(
        self,
        requires: MessageSpec | Iterable[MessageSpec] = (),
        provides: MessageSpec | Iterable[MessageSpec] = (),
        params: dict[str, Any] | None = None,
        name: str | None = None,
    ):
        super().__init__(
            self._forward,
            requires=requires,
            provides=provides,
            params=params,
            name=name or "mlp",
        )
        cfg = self.Params.model_validate(self.params)

        np = self._np = _numpy()
        rng = np.random.default_rng(seed=cfg.seed)
        num_inputs = max(1, len(self.input_ports))
        num_outputs = max(1, len(self.output_ports))
        self.weights = rng.normal(loc=1, size=(num_outputs, num_inputs))
        self.biases = rng.normal(loc=0, size=(num_outputs,))

    def _forward(
        self, params: dict[str, Any], **kwargs: float
    ) -> dict[str, Any] | None:
        """Apply ``weights @ inputs + biases`` to the incoming values."""
        if not self.output_ports:
            return None

        logger.info(
            "Executing MLP node in thread %s: %s",
            threading.current_thread().name,
            self.id,
        )

        np = self._np
        if kwargs:
            input_array = np.array([float(v) for v in kwargs.values()])
        else:
            input_array = np.array([1.0])  # Bias only, for source nodes.

        outputs = np.dot(self.weights, input_array) + self.biases
        return {p.port: float(outputs[i]) for i, p in enumerate(self.output_ports)}