trig.py

"""Trigonometric helpers shared across metrics and helpers.

This module focuses on mathematical utilities (means, compensated sums, etc.).
Caching of cosine/sine values lives in :mod:`tnfr.metrics.trig_cache`.
"""

from __future__ import annotations

import math
from collections.abc import Iterable, Iterator, Sequence
from itertools import tee
from typing import TYPE_CHECKING, Any, cast, overload

from ..errors import TNFRValueError
from ..mathematics.unified_numerical import np
from ..types import NodeId, Phase, TNFRGraph
from ..utils import cached_import, kahan_sum_nd

if TYPE_CHECKING:  # pragma: no cover - typing only
    from ..node import NodeProtocol

__all__ = (
    "accumulate_cos_sin",
    "_phase_mean_from_iter",
    "_neighbor_phase_mean_core",
    "_neighbor_phase_mean_generic",
    "neighbor_phase_mean_bulk",
    "neighbor_phase_mean_list",
    "neighbor_phase_mean",
)


def accumulate_cos_sin(
    it: Iterable[tuple[float, float] | None],
) -> tuple[float, float, bool]:
    """Accumulate cosine and sine pairs with compensated summation.

    ``it`` yields optional ``(cos, sin)`` tuples. Entries with ``None``
    components are ignored. The returned values are the compensated sums of
    cosines and sines along with a flag indicating whether any pair was
    processed.
    """

    processed = False

    def iter_real_pairs() -> Iterator[tuple[float, float]]:
        nonlocal processed
        for cs in it:
            if cs is None:
                continue
            c, s = cs
            if c is None or s is None:
                continue
            try:
                c_val = float(c)
                s_val = float(s)
            except (TypeError, ValueError):
                continue
            if not (math.isfinite(c_val) and math.isfinite(s_val)):
                continue
            processed = True
            yield (c_val, s_val)

    sum_cos, sum_sin = kahan_sum_nd(iter_real_pairs(), dims=2)

    if not processed:
        return 0.0, 0.0, False

    return sum_cos, sum_sin, True


def _phase_mean_from_iter(
    it: Iterable[tuple[float, float] | None], fallback: float
) -> float:
    """Return circular mean from an iterator of cosine/sine pairs.

    ``it`` yields optional ``(cos, sin)`` tuples. ``fallback`` is returned if
    no valid pairs are processed.
    """

    sum_cos, sum_sin, processed = accumulate_cos_sin(it)
    if not processed:
        return fallback
    return math.atan2(sum_sin, sum_cos)


def _neighbor_phase_mean_core(
    neigh: Sequence[Any],
    cos_map: dict[Any, float],
    sin_map: dict[Any, float],
    fallback: float,
) -> float:
    """Return circular mean of neighbour phases given trig mappings."""

    def _iter_pairs() -> Iterator[tuple[float, float]]:
        for v in neigh:
            c = cos_map.get(v)
            s = sin_map.get(v)
            if c is not None and s is not None:
                yield c, s

    pairs = _iter_pairs()

    if np is not None:
        cos_iter, sin_iter = tee(pairs, 2)
        cos_arr = np.fromiter((c for c, _ in cos_iter), dtype=float)
        sin_arr = np.fromiter((s for _, s in sin_iter), dtype=float)
        if cos_arr.size:
            mean_cos = float(np.mean(cos_arr))
            mean_sin = float(np.mean(sin_arr))
            return float(np.arctan2(mean_sin, mean_cos))
        return fallback

    sum_cos, sum_sin, processed = accumulate_cos_sin(pairs)
    if not processed:
        return fallback
    return math.atan2(sum_sin, sum_cos)


def _neighbor_phase_mean_generic(
    obj: "NodeProtocol" | Sequence[Any],
    cos_map: dict[Any, float] | None = None,
    sin_map: dict[Any, float] | None = None,
    fallback: float = 0.0,
) -> float:
    """Compute the neighbour phase mean via :func:`_neighbor_phase_mean_core`.

    ``obj`` may be either a node bound to a graph or a sequence of neighbours.
    When ``cos_map`` and ``sin_map`` are ``None`` the function assumes ``obj`` is
    a node and obtains the required trigonometric mappings from the cached
    structures. Otherwise ``obj`` is treated as an explicit neighbour
    sequence and ``cos_map``/``sin_map`` must be provided.
    """

    if cos_map is None or sin_map is None:
        node = cast("NodeProtocol", obj)
        if getattr(node, "G", None) is None:
            raise TypeError("neighbor_phase_mean requires nodes bound to a graph")
        from .trig_cache import get_trig_cache

        trig = get_trig_cache(node.G)
        fallback = trig.theta.get(node.n, fallback)
        cos_map = trig.cos
        sin_map = trig.sin
        neigh = node.G[node.n]
    else:
        neigh = cast(Sequence[Any], obj)

    return _neighbor_phase_mean_core(neigh, cos_map, sin_map, fallback)


def neighbor_phase_mean_list(
    neigh: Sequence[Any],
    cos_th: dict[Any, float],
    sin_th: dict[Any, float],
    fallback: float = 0.0,
) -> float:
    """Return circular mean of neighbour phases from cosine/sine mappings.

    This is a thin wrapper over :func:`_neighbor_phase_mean_generic` that
    operates on explicit neighbour lists.
    """

    return _neighbor_phase_mean_generic(
        neigh, cos_map=cos_th, sin_map=sin_th, fallback=fallback
    )


def neighbor_phase_mean_bulk(
    edge_src: Any,
    edge_dst: Any,
    *,
    cos_values: Any,
    sin_values: Any,
    theta_values: Any,
    node_count: int,
    neighbor_cos_sum: Any | None = None,
    neighbor_sin_sum: Any | None = None,
    neighbor_counts: Any | None = None,
    mean_cos: Any | None = None,
    mean_sin: Any | None = None,
) -> tuple[Any, Any]:
    """Vectorised neighbour phase means for all nodes in a graph.

    Parameters
    ----------
    edge_src, edge_dst:
        Arrays describing the source (neighbour) and destination (node) indices
        for each edge contribution. They must have matching shapes.
    cos_values, sin_values:
        Arrays containing the cosine and sine values of each node's phase. The
        arrays must be indexed using the same positional indices referenced by
        ``edge_src``.
    theta_values:
        Array with the baseline phase for each node. Positions that do not have
        neighbours reuse this baseline as their mean phase.
    node_count:
        Total number of nodes represented in ``theta_values``.

    Optional buffers
    -----------------
    neighbor_cos_sum, neighbor_sin_sum, neighbor_counts, mean_cos, mean_sin:
        Preallocated arrays sized ``node_count`` reused to accumulate the
        neighbour cosine/sine sums, neighbour sample counts, and the averaged
        cosine/sine vectors. When omitted, the helper materialises fresh
        buffers that match the previous semantics.

    Returns
    -------
    tuple[Any, Any]
        tuple ``(mean_theta, has_neighbors)`` where ``mean_theta`` contains the
        circular mean of neighbour phases for every node and ``has_neighbors``
        is a boolean mask identifying which nodes contributed at least one
        neighbour sample.
    """

    if node_count <= 0:
        empty_mean = np.zeros(0, dtype=float)
        return empty_mean, empty_mean.astype(bool)

    edge_src_arr = np.asarray(edge_src, dtype=np.intp)
    edge_dst_arr = np.asarray(edge_dst, dtype=np.intp)

    if edge_src_arr.shape != edge_dst_arr.shape:
        raise TNFRValueError(
            "edge_src and edge_dst must share the same shape",
            context={"src_shape": edge_src_arr.shape, "dst_shape": edge_dst_arr.shape},
        )

    theta_arr = np.asarray(theta_values, dtype=float)
    if theta_arr.ndim != 1 or theta_arr.size != node_count:
        raise TNFRValueError(
            "theta_values must be a 1-D array matching node_count",
            context={"shape": theta_arr.shape, "expected_size": node_count},
        )

    cos_arr = np.asarray(cos_values, dtype=float)
    sin_arr = np.asarray(sin_values, dtype=float)
    if cos_arr.ndim != 1 or cos_arr.size != node_count:
        raise TNFRValueError(
            "cos_values must be a 1-D array matching node_count",
            context={"shape": cos_arr.shape, "expected_size": node_count},
        )
    if sin_arr.ndim != 1 or sin_arr.size != node_count:
        raise TNFRValueError(
            "sin_values must be a 1-D array matching node_count",
            context={"shape": sin_arr.shape, "expected_size": node_count},
        )

    edge_count = edge_dst_arr.size

    def _coerce_buffer(buffer: Any | None, *, name: str) -> tuple[Any, bool]:
        if buffer is None:
            return None, False
        arr = np.array(buffer, dtype=float, copy=False)
        if arr.ndim != 1 or arr.size != node_count:
            raise TNFRValueError(
                f"{name} must be a 1-D array sized node_count",
                context={"name": name, "shape": arr.shape, "expected_size": node_count},
            )
        arr.fill(0.0)
        return arr, True

    neighbor_cos_sum, has_cos_buffer = _coerce_buffer(
        neighbor_cos_sum, name="neighbor_cos_sum"
    )
    neighbor_sin_sum, has_sin_buffer = _coerce_buffer(
        neighbor_sin_sum, name="neighbor_sin_sum"
    )
    neighbor_counts, has_count_buffer = _coerce_buffer(
        neighbor_counts, name="neighbor_counts"
    )

    if edge_count:
        cos_bincount = np.bincount(
            edge_dst_arr,
            weights=cos_arr[edge_src_arr],
            minlength=node_count,
        )
        sin_bincount = np.bincount(
            edge_dst_arr,
            weights=sin_arr[edge_src_arr],
            minlength=node_count,
        )
        count_bincount = np.bincount(
            edge_dst_arr,
            minlength=node_count,
        ).astype(float, copy=False)

        if not has_cos_buffer:
            neighbor_cos_sum = cos_bincount
        else:
            np.copyto(neighbor_cos_sum, cos_bincount)

        if not has_sin_buffer:
            neighbor_sin_sum = sin_bincount
        else:
            np.copyto(neighbor_sin_sum, sin_bincount)

        if not has_count_buffer:
            neighbor_counts = count_bincount
        else:
            np.copyto(neighbor_counts, count_bincount)
    else:
        if neighbor_cos_sum is None:
            neighbor_cos_sum = np.zeros(node_count, dtype=float)
        if neighbor_sin_sum is None:
            neighbor_sin_sum = np.zeros(node_count, dtype=float)
        if neighbor_counts is None:
            neighbor_counts = np.zeros(node_count, dtype=float)

    has_neighbors = neighbor_counts > 0.0

    mean_cos, _ = _coerce_buffer(mean_cos, name="mean_cos")
    mean_sin, _ = _coerce_buffer(mean_sin, name="mean_sin")

    if mean_cos is None:
        mean_cos = np.zeros(node_count, dtype=float)
    if mean_sin is None:
        mean_sin = np.zeros(node_count, dtype=float)

    if edge_count:
        with np.errstate(divide="ignore", invalid="ignore"):
            np.divide(
                neighbor_cos_sum,
                neighbor_counts,
                out=mean_cos,
                where=has_neighbors,
            )
            np.divide(
                neighbor_sin_sum,
                neighbor_counts,
                out=mean_sin,
                where=has_neighbors,
            )

    mean_theta = np.where(has_neighbors, np.arctan2(mean_sin, mean_cos), theta_arr)
    return mean_theta, has_neighbors


@overload
def neighbor_phase_mean(obj: "NodeProtocol", n: None = ...) -> Phase: ...


@overload
def neighbor_phase_mean(obj: TNFRGraph, n: NodeId) -> Phase: ...


def neighbor_phase_mean(
    obj: "NodeProtocol" | TNFRGraph, n: NodeId | None = None
) -> Phase:
    """Circular mean of neighbour phases for ``obj``.

    Parameters
    ----------
    obj:
        Either a :class:`~tnfr.node.NodeProtocol` instance bound to a graph or a
        :class:`~tnfr.types.TNFRGraph` from which the node ``n`` will be wrapped.
    n:
        Optional node identifier. Required when ``obj`` is a graph. Providing a
        node identifier for a node object raises :class:`TypeError`.
    """

    NodeNX = cached_import("tnfr.node", "NodeNX")
    if NodeNX is None:
        raise ImportError("NodeNX is unavailable")
    if n is None:
        if hasattr(obj, "nodes"):
            raise TypeError(
                "neighbor_phase_mean requires a node identifier when passing a graph"
            )
        node = obj
    else:
        if hasattr(obj, "nodes"):
            node = NodeNX(obj, n)
        else:
            raise TypeError(
                "neighbor_phase_mean received a node and an explicit identifier"
            )
    return _neighbor_phase_mean_generic(node)