"""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

import numpy as np
from pydantic import BaseModel, ConfigDict

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

logger = logging.getLogger(__name__)


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)

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

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