"""Fused ΔNFR computation kernels for optimized backend.

This module provides optimized implementations of ΔNFR computation that fuse
multiple operations to reduce memory traffic and improve cache locality.

The fused kernels maintain exact TNFR semantics while achieving better
performance through:

1. Combined gradient computation (phase + EPI + topology in single pass)
2. Reduced intermediate array allocations
3. Better cache locality through sequential memory access
4. Optional Numba JIT compilation for critical loops

All implementations preserve TNFR structural invariants:
- ΔNFR = w_phase·g_phase + w_epi·g_epi + w_vf·g_vf + w_topo·g_topo
- Isolated nodes receive ΔNFR = 0
- Deterministic results with fixed topology
"""

from __future__ import annotations

import math
from typing import Any, Mapping

from ..utils import get_numpy, get_logger

logger = get_logger(__name__)

# Try to import Numba for JIT acceleration
_NUMBA_AVAILABLE = False
_numba = None
try:
    import numba

    _numba = numba
    _NUMBA_AVAILABLE = True
    logger.debug("Numba JIT available for fused gradient acceleration")
except ImportError:
    logger.debug("Numba not available, using pure NumPy implementation")


def _compute_fused_gradients_jit_kernel(
    edge_src,
    edge_dst,
    phase,
    epi,
    vf,
    w_phase: float,
    w_epi: float,
    w_vf: float,
    w_topo: float,
    delta_nfr,
):
    """Numba JIT kernel for fused gradient computation.

    This function is designed to be JIT-compiled by Numba for maximum
    performance. It operates directly on arrays with explicit loops.

    Parameters
    ----------
    edge_src : ndarray
        Source node indices (int)
    edge_dst : ndarray
        Destination node indices (int)
    phase : ndarray
        Phase values (float)
    epi : ndarray
        EPI values (float)
    vf : ndarray
        νf values (float)
    w_phase : float
        Phase gradient weight
    w_epi : float
        EPI gradient weight
    w_vf : float
        νf gradient weight
    w_topo : float
        Topology gradient weight
    delta_nfr : ndarray
        Output array for ΔNFR (modified in-place)
    """
    n_edges = edge_src.shape[0]

    for i in range(n_edges):
        src = edge_src[i]
        dst = edge_dst[i]

        # Compute gradients for this edge
        phase_diff = math.sin(phase[dst] - phase[src])
        epi_diff = epi[dst] - epi[src]
        vf_diff = vf[dst] - vf[src]

        # Fused contribution
        contrib = (
            w_phase * phase_diff + w_epi * epi_diff + w_vf * vf_diff + w_topo * 1.0
        )

        # Accumulate to destination node
        delta_nfr[dst] += contrib


# Create JIT-compiled version if Numba is available
if _NUMBA_AVAILABLE:
    try:
        _compute_fused_gradients_jit = _numba.njit(
            _compute_fused_gradients_jit_kernel,
            parallel=False,
            fastmath=True,
            cache=True,
        )
        logger.debug("Numba JIT compilation successful for fused gradients")
    except Exception as e:
        logger.warning(f"Numba JIT compilation failed: {e}")
        _compute_fused_gradients_jit = _compute_fused_gradients_jit_kernel
else:
    _compute_fused_gradients_jit = _compute_fused_gradients_jit_kernel


def compute_fused_gradients(
    *,
    edge_src: Any,
    edge_dst: Any,
    phase: Any,
    epi: Any,
    vf: Any,
    weights: Mapping[str, float],
    np: Any,
    use_jit: bool = True,
) -> Any:
    """Compute all ΔNFR gradients in a fused kernel.

    This function combines phase, EPI, and νf gradient computations into
    a single pass over the edge list, reducing memory traffic by ~50%
    compared to separate accumulations.

    When Numba is available and `use_jit=True`, uses JIT-compiled inner
    loop for additional 2-3x speedup on large graphs.

    Parameters
    ----------
    edge_src : array-like
        Source node indices for each edge (shape: [E])
    edge_dst : array-like
        Destination node indices for each edge (shape: [E])
    phase : array-like
        Phase values for each node (shape: [N])
    epi : array-like
        EPI values for each node (shape: [N])
    vf : array-like
        Structural frequency νf for each node (shape: [N])
    weights : Mapping[str, float]
        ΔNFR component weights (w_phase, w_epi, w_vf, w_topo)
    np : module
        NumPy module
    use_jit : bool, default=True
        Whether to use JIT compilation if available

    Returns
    -------
    ndarray
        Fused gradient vector (shape: [N])

    Notes
    -----
    The fused kernel computes:

    ```
    for each edge (i → j):
        g_phase = sin(phase[j] - phase[i])
        g_epi = epi[j] - epi[i]
        g_vf = vf[j] - vf[i]

        delta_nfr[dst] += w_phase * g_phase + w_epi * g_epi + w_vf * g_vf + w_topo
    ```

    Where dst is the destination node index (j) for each edge.

    By processing all components simultaneously, we:
    - Access edge arrays once instead of 3+ times
    - Reduce cache misses on node attribute arrays
    - Minimize temporary array allocations
    - Enable JIT compilation for 2-3x additional speedup

    Examples
    --------
    >>> import numpy as np
    >>> edge_src = np.array([0, 0, 1])
    >>> edge_dst = np.array([1, 2, 2])
    >>> phase = np.array([0.0, 0.1, 0.2])
    >>> epi = np.array([0.5, 0.6, 0.7])
    >>> vf = np.array([1.0, 1.0, 1.0])
    >>> weights = {'w_phase': 0.3, 'w_epi': 0.3, 'w_vf': 0.2, 'w_topo': 0.2}
    >>> result = compute_fused_gradients(
    ...     edge_src=edge_src, edge_dst=edge_dst,
    ...     phase=phase, epi=epi, vf=vf,
    ...     weights=weights, np=np
    ... )
    >>> result.shape
    (3,)
    """
    n_nodes = phase.shape[0]
    n_edges = edge_src.shape[0]

    # Extract weights
    w_phase = float(weights.get("w_phase", 0.0))
    w_epi = float(weights.get("w_epi", 0.0))
    w_vf = float(weights.get("w_vf", 0.0))
    w_topo = float(weights.get("w_topo", 0.0))

    # Allocate result
    delta_nfr = np.zeros(n_nodes, dtype=float)

    if n_edges == 0:
        return delta_nfr

    # Choose implementation based on JIT availability and preference
    if use_jit and _NUMBA_AVAILABLE and n_edges > 100:
        # Use JIT-compiled kernel for large graphs
        _compute_fused_gradients_jit(
            edge_src, edge_dst, phase, epi, vf, w_phase, w_epi, w_vf, w_topo, delta_nfr
        )
    else:
        # Use vectorized NumPy implementation
        # Compute edge contributions (fused)
        phase_src = phase[edge_src]
        phase_dst = phase[edge_dst]
        phase_diff = np.sin(phase_dst - phase_src)

        epi_src = epi[edge_src]
        epi_dst = epi[edge_dst]
        epi_diff = epi_dst - epi_src

        vf_src = vf[edge_src]
        vf_dst = vf[edge_dst]
        vf_diff = vf_dst - vf_src

        # Fused contribution per edge
        edge_contrib = (
            w_phase * phase_diff + w_epi * epi_diff + w_vf * vf_diff + w_topo * 1.0
        )

        # Accumulate contributions to destination nodes
        np.add.at(delta_nfr, edge_dst, edge_contrib)

    return delta_nfr


def compute_fused_gradients_symmetric(
    *,
    edge_src: Any,
    edge_dst: Any,
    phase: Any,
    epi: Any,
    vf: Any,
    weights: Mapping[str, float],
    np: Any,
) -> Any:
    """Compute ΔNFR gradients using TNFR canonical formula with neighbor means.

    For undirected graphs, this function implements the canonical TNFR gradient
    computation where each node's gradient is computed from the circular mean
    of its neighbors' phases (divided by π), and arithmetic means for EPI/νf.

    The canonical TNFR formulas are:
    - Phase gradient: g_phase = sin(θ_node - θ̄_neighbors) / π
    - EPI/νf gradient: g = x̄_neighbors - x_node
    - Topology gradient: g_topo = degree · w_topo

    This implementation performs the computation in two passes:
    1. Accumulate neighbor cos/sin sums and value sums for means
    2. Compute gradients from means for each node

    Parameters
    ----------
    edge_src : array-like
        Source node indices (shape: [E])
    edge_dst : array-like
        Destination node indices (shape: [E])
    phase : array-like
        Phase values (shape: [N])
    epi : array-like
        EPI values (shape: [N])
    vf : array-like
        νf values (shape: [N])
    weights : Mapping[str, float]
        Component weights (w_phase, w_epi, w_vf, w_topo)
    np : module
        NumPy module

    Returns
    -------
    ndarray
        ΔNFR gradient vector (shape: [N])

    Notes
    -----
    The division by π in the phase gradient is part of the canonical TNFR
    formulation and ensures proper scaling of phase-based reorganization.

    For isolated nodes (no neighbors), all gradients are zero.

    Examples
    --------
    >>> import numpy as np
    >>> edge_src = np.array([0, 0, 1])
    >>> edge_dst = np.array([1, 2, 2])
    >>> phase = np.array([0.0, 0.1, 0.2])
    >>> epi = np.array([0.5, 0.6, 0.7])
    >>> vf = np.array([1.0, 1.0, 1.0])
    >>> weights = {'w_phase': 0.3, 'w_epi': 0.3, 'w_vf': 0.2, 'w_topo': 0.2}
    >>> result = compute_fused_gradients_symmetric(
    ...     edge_src=edge_src, edge_dst=edge_dst,
    ...     phase=phase, epi=epi, vf=vf,
    ...     weights=weights, np=np
    ... )
    >>> result.shape
    (3,)
    """
    n_nodes = phase.shape[0]
    n_edges = edge_src.shape[0]

    w_phase = float(weights.get("w_phase", 0.0))
    w_epi = float(weights.get("w_epi", 0.0))
    w_vf = float(weights.get("w_vf", 0.0))
    w_topo = float(weights.get("w_topo", 0.0))

    delta_nfr = np.zeros(n_nodes, dtype=float)

    if n_edges == 0:
        return delta_nfr

    # Pass 1: Accumulate neighbor statistics for computing means
    # For phase: accumulate cos/sin sums for circular mean
    # For EPI/vf: accumulate value sums for arithmetic mean
    neighbor_cos_sum = np.zeros(n_nodes, dtype=float)
    neighbor_sin_sum = np.zeros(n_nodes, dtype=float)
    neighbor_epi_sum = np.zeros(n_nodes, dtype=float)
    neighbor_vf_sum = np.zeros(n_nodes, dtype=float)
    neighbor_count = np.zeros(n_nodes, dtype=float)

    # For undirected graphs, each edge contributes to both endpoints
    # Extract neighbor values
    phase_src_vals = phase[edge_src]
    phase_dst_vals = phase[edge_dst]
    epi_src_vals = epi[edge_src]
    epi_dst_vals = epi[edge_dst]
    vf_src_vals = vf[edge_src]
    vf_dst_vals = vf[edge_dst]

    # Accumulate from src to dst (dst's neighbors include src)
    np.add.at(neighbor_cos_sum, edge_dst, np.cos(phase_src_vals))
    np.add.at(neighbor_sin_sum, edge_dst, np.sin(phase_src_vals))
    np.add.at(neighbor_epi_sum, edge_dst, epi_src_vals)
    np.add.at(neighbor_vf_sum, edge_dst, vf_src_vals)
    np.add.at(neighbor_count, edge_dst, 1.0)

    # Accumulate from dst to src (src's neighbors include dst)
    np.add.at(neighbor_cos_sum, edge_src, np.cos(phase_dst_vals))
    np.add.at(neighbor_sin_sum, edge_src, np.sin(phase_dst_vals))
    np.add.at(neighbor_epi_sum, edge_src, epi_dst_vals)
    np.add.at(neighbor_vf_sum, edge_src, vf_dst_vals)
    np.add.at(neighbor_count, edge_src, 1.0)

    # Pass 2: Compute gradients from means
    # Avoid division by zero for isolated nodes
    has_neighbors = neighbor_count > 0

    # Compute circular mean phase for nodes with neighbors
    phase_mean = np.zeros(n_nodes, dtype=float)
    phase_mean[has_neighbors] = np.arctan2(
        neighbor_sin_sum[has_neighbors], neighbor_cos_sum[has_neighbors]
    )

    # Compute arithmetic means for EPI and νf
    epi_mean = np.zeros(n_nodes, dtype=float)
    vf_mean = np.zeros(n_nodes, dtype=float)
    epi_mean[has_neighbors] = (
        neighbor_epi_sum[has_neighbors] / neighbor_count[has_neighbors]
    )
    vf_mean[has_neighbors] = (
        neighbor_vf_sum[has_neighbors] / neighbor_count[has_neighbors]
    )

    # Compute gradients using TNFR canonical formula
    # Phase: g_phase = -angle_diff(θ_node, θ_mean) / π
    # angle_diff(a, b) = (a - b + π) % 2π - π (minimal angular difference)
    phase_diff = (phase_mean - phase + np.pi) % (2 * np.pi) - np.pi
    g_phase = phase_diff / np.pi
    g_phase[~has_neighbors] = 0.0  # Isolated nodes have no gradient

    # EPI/νf: g = mean - node_value
    g_epi = epi_mean - epi
    g_epi[~has_neighbors] = 0.0

    g_vf = vf_mean - vf
    g_vf[~has_neighbors] = 0.0

    # Topology: each neighbor contributes w_topo
    g_topo = neighbor_count * w_topo

    # Combine gradients
    delta_nfr = (
        w_phase * g_phase
        + w_epi * g_epi
        + w_vf * g_vf
        + g_topo  # Already scaled by w_topo
    )

    return delta_nfr


def apply_vf_scaling(
    *,
    delta_nfr: Any,
    vf: Any,
    np: Any,
) -> None:
    """Apply structural frequency scaling to ΔNFR in-place.

    Applies the fundamental TNFR transformation:
    ΔNFR_final = νf · ΔNFR_gradient

    This completes the nodal equation: ∂EPI/∂t = νf · ΔNFR(t)

    Parameters
    ----------
    delta_nfr : array-like
        Gradient vector to scale in-place (shape: [N])
    vf : array-like
        Structural frequency for each node (shape: [N])
    np : module
        NumPy module

    Notes
    -----
    Modified in-place to avoid allocating result array.

    Examples
    --------
    >>> import numpy as np
    >>> delta_nfr = np.array([1.0, 2.0, 3.0])
    >>> vf = np.array([0.5, 1.0, 1.5])
    >>> apply_vf_scaling(delta_nfr=delta_nfr, vf=vf, np=np)
    >>> delta_nfr
    array([0.5, 2. , 4.5])
    """
    np.multiply(delta_nfr, vf, out=delta_nfr)


__all__ = [
    "compute_fused_gradients",
    "compute_fused_gradients_symmetric",
    "apply_vf_scaling",
]
