"""
Neural Coupler — CfC per-neuron tau-modulated coupling with real Kuramoto r(t).

Designed for applications running trained CfC neural networks on-device.
Each neuron's coupling rate is modulated by its learned time constant τ.

Copyright (c) 2026 SYM.BOT Ltd. Apache 2.0 License.
"""

import math
import time
from typing import Optional

from .coupler import (
    CouplingDecision,
    CouplingDecisionType,
    MeshCoupler,
    NeuralCouplingConfig,
    _PeerState,
    cosine_similarity,
)


class NeuralCoupler(MeshCoupler):
    """Neural Coupler — per-neuron tau-modulated coupling with real Kuramoto r(t).

    Fast neurons (small τ) couple readily.
    Slow neurons (large τ) couple conservatively.
    Coherence is real Kuramoto r(t) = |⟨e^{iθ}⟩|.
    """

    def __init__(self, config: NeuralCouplingConfig):
        if config.tau is None:
            raise ValueError("NeuralCouplingConfig requires tau (per-neuron time constants)")
        self.config = config
        self.tau = list(config.tau)
        self.K = config.coupling_rate
        self._peers: dict[str, _PeerState] = {}
        self.last_decisions: dict[str, CouplingDecision] = {}

    def update_peer(self, agent_id: str, h1: list[float], h2: list[float], confidence: float):
        self._peers[agent_id] = _PeerState(
            h1=list(h1), h2=list(h2),
            confidence=max(confidence, 0.01),
            timestamp=time.time(),
        )

    def remove_peer(self, agent_id: str):
        self._peers.pop(agent_id, None)
        self.last_decisions.pop(agent_id, None)

    def remove_all_peers(self):
        self._peers.clear()
        self.last_decisions.clear()

    @property
    def active_peer_count(self) -> int:
        self._prune_expired()
        return len(self._peers)

    @property
    def has_active_peers(self) -> bool:
        self._prune_expired()
        return len(self._peers) > 0

    def couple(self, local_h1: list[float], local_h2: list[float]) -> tuple[list[float], list[float]]:
        self._prune_expired()
        if not self._peers:
            return list(local_h1), list(local_h2)

        dim = len(local_h1)
        if len(self.tau) != dim:
            raise ValueError(f"Tau dimension mismatch: tau[{len(self.tau)}] vs hiddenDim[{dim}]")

        decisions: dict[str, CouplingDecision] = {}
        accepted: list[tuple[_PeerState, float, float]] = []

        for agent_id, peer in self._peers.items():
            if len(peer.h1) != dim or len(peer.h2) != dim:
                continue

            drift_h1 = 1.0 - cosine_similarity(local_h1, peer.h1)
            drift_h2 = 1.0 - cosine_similarity(local_h2, peer.h2)
            drift = (drift_h1 + drift_h2) / 2.0

            cfg = self.config
            if drift <= cfg.drift_threshold_aligned:
                decision = CouplingDecisionType.ALIGNED
                alpha = cfg.alpha_aligned
            elif drift <= cfg.drift_threshold_guarded:
                decision = CouplingDecisionType.GUARDED
                alpha = cfg.alpha_guarded
            else:
                decision = CouplingDecisionType.REJECTED
                alpha = 0.0

            age = time.time() - peer.timestamp
            recency = math.exp(-cfg.temporal_decay * age)
            weight = peer.confidence * (1.0 - drift) * recency

            decisions[agent_id] = CouplingDecision(
                agent_id=agent_id, drift=drift, decision=decision,
                alpha=alpha, weight=weight,
            )

            if decision != CouplingDecisionType.REJECTED:
                accepted.append((peer, weight, alpha))

        self.last_decisions = decisions

        if not accepted:
            return list(local_h1), list(local_h2)

        mesh_h1 = [0.0] * dim
        mesh_h2 = [0.0] * dim
        total_weight = 0.0

        for peer, w, _ in accepted:
            for i in range(dim):
                mesh_h1[i] += peer.h1[i] * w
                mesh_h2[i] += peer.h2[i] * w
            total_weight += w

        if total_weight <= 0:
            return list(local_h1), list(local_h2)

        for i in range(dim):
            mesh_h1[i] /= total_weight
            mesh_h2[i] /= total_weight

        alpha_sum = sum(a * w for _, w, a in accepted)
        weight_sum = sum(w for _, w, _ in accepted)
        effective_alpha = alpha_sum / weight_sum if weight_sum > 0 else 0.0

        # Per-neuron TAU-MODULATED coupling
        coupled_h1 = [0.0] * dim
        coupled_h2 = [0.0] * dim

        for i in range(dim):
            sim_h1 = (1.0 - abs(local_h1[i] - mesh_h1[i]) / max(abs(local_h1[i]), abs(mesh_h1[i]))) \
                if abs(local_h1[i]) > 1e-8 and abs(mesh_h1[i]) > 1e-8 else 0.0
            sim_h2 = (1.0 - abs(local_h2[i] - mesh_h2[i]) / max(abs(local_h2[i]), abs(mesh_h2[i]))) \
                if abs(local_h2[i]) > 1e-8 and abs(mesh_h2[i]) > 1e-8 else 0.0

            tau_i = max(self.tau[i], 0.01)
            a1 = min(effective_alpha * self.K * max(sim_h1, 0.0) / tau_i, 1.0)
            a2 = min(effective_alpha * self.K * max(sim_h2, 0.0) / tau_i, 1.0)

            coupled_h1[i] = (1.0 - a1) * local_h1[i] + a1 * mesh_h1[i]
            coupled_h2[i] = (1.0 - a2) * local_h2[i] + a2 * mesh_h2[i]

        return coupled_h1, coupled_h2

    def l2_distance(self, local_h1: list[float], local_h2: list[float]) -> Optional[tuple[float, float]]:
        self._prune_expired()
        if not self._peers:
            return None

        dim = len(local_h1)
        mesh_h1 = [0.0] * dim
        mesh_h2 = [0.0] * dim
        total_w = 0.0

        for peer in self._peers.values():
            if len(peer.h1) != dim:
                continue
            w = peer.confidence
            for i in range(dim):
                mesh_h1[i] += peer.h1[i] * w
                mesh_h2[i] += peer.h2[i] * w
            total_w += w

        if total_w <= 0:
            return None
        for i in range(dim):
            mesh_h1[i] /= total_w
            mesh_h2[i] /= total_w

        d1 = math.sqrt(sum((local_h1[i] - mesh_h1[i]) ** 2 for i in range(dim)))
        d2 = math.sqrt(sum((local_h2[i] - mesh_h2[i]) ** 2 for i in range(dim)))
        return d1, d2

    def coherence(self, local_h1: list[float], local_h2: list[float]) -> Optional[float]:
        """Real Kuramoto order parameter: r(t) = |⟨e^{iθ}⟩|"""
        self._prune_expired()
        if not self._peers:
            return None

        dim = len(local_h1)
        oscillators = [{"h1": local_h1, "h2": local_h2}]
        for peer in self._peers.values():
            if len(peer.h1) == dim and len(peer.h2) == dim:
                oscillators.append({"h1": peer.h1, "h2": peer.h2})

        if len(oscillators) < 2:
            return None

        N = len(oscillators)
        r_sum = 0.0

        for i in range(dim):
            real_sum = 0.0
            imag_sum = 0.0
            for osc in oscillators:
                theta = math.atan2(osc["h2"][i], osc["h1"][i])
                real_sum += math.cos(theta)
                imag_sum += math.sin(theta)
            r_i = math.sqrt((real_sum / N) ** 2 + (imag_sum / N) ** 2)
            r_sum += r_i

        return r_sum / dim

    def kuramoto_order_parameter(self, local_h1: list[float], local_h2: list[float]) -> Optional[float]:
        """Deprecated: use coherence() instead."""
        return self.coherence(local_h1, local_h2)

    def _prune_expired(self):
        now = time.time()
        expired = [k for k, v in self._peers.items()
                   if now - v.timestamp > self.config.peer_retention_seconds]
        for k in expired:
            self._peers.pop(k, None)
            self.last_decisions.pop(k, None)
