"""
Mesh Coupler — Autonomous drift-bounded coupling engine.

Provides the MeshCoupler protocol and SemanticCoupler implementation
for cosine-drift-based state blending (shared memory, LLM agents).

For CfC neural coupling with per-neuron time constants, see NeuralCoupler.

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

import math
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import Enum
from typing import Optional


class CouplingDecisionType(Enum):
    ALIGNED = "aligned"
    GUARDED = "guarded"
    REJECTED = "rejected"


@dataclass
class CouplingDecision:
    agent_id: str
    drift: float
    decision: CouplingDecisionType
    alpha: float
    weight: float
    timestamp: float = field(default_factory=time.time)


@dataclass
class CouplingConfig:
    """Configuration for autonomous drift-bounded coupling."""
    drift_threshold_aligned: float = 0.25
    drift_threshold_guarded: float = 0.5
    alpha_aligned: float = 0.4
    alpha_guarded: float = 0.15
    temporal_decay: float = 0.05
    peer_retention_seconds: float = 300.0


@dataclass
class NeuralCouplingConfig(CouplingConfig):
    """Configuration for neural coupling with per-neuron time constants.

    Extends CouplingConfig with tau (from trained CfC model) and
    global Kuramoto coupling strength K.
    """
    tau: Optional[list[float]] = None
    coupling_rate: float = 1.0


@dataclass
class _PeerState:
    h1: list[float]
    h2: list[float]
    confidence: float
    timestamp: float


class MeshCoupler(ABC):
    """Common interface for all coupling engines."""

    @abstractmethod
    def update_peer(self, agent_id: str, h1: list[float], h2: list[float], confidence: float) -> None: ...

    @abstractmethod
    def remove_peer(self, agent_id: str) -> None: ...

    @abstractmethod
    def remove_all_peers(self) -> None: ...

    @property
    @abstractmethod
    def active_peer_count(self) -> int: ...

    @property
    @abstractmethod
    def has_active_peers(self) -> bool: ...

    @abstractmethod
    def couple(self, local_h1: list[float], local_h2: list[float]) -> tuple[list[float], list[float]]: ...

    @abstractmethod
    def l2_distance(self, local_h1: list[float], local_h2: list[float]) -> Optional[tuple[float, float]]: ...

    @abstractmethod
    def coherence(self, local_h1: list[float], local_h2: list[float]) -> Optional[float]: ...


class SemanticCoupler(MeshCoupler):
    """Semantic Coupler — cosine-drift-based state blending.

    Designed for embedding-based agents, shared memory, and LLM applications.
    Uses uniform alpha across all neurons, modulated by per-neuron similarity.
    Coherence metric is average pairwise cosine similarity (not Kuramoto r(t)).
    """

    def __init__(self, config: Optional[CouplingConfig] = None):
        self.config = config or CouplingConfig()
        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)
        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

        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

            a1 = effective_alpha * max(sim_h1, 0.0)
            a2 = effective_alpha * max(sim_h2, 0.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]:
        """Cosine coherence: average pairwise cosine similarity."""
        self._prune_expired()
        if not self._peers:
            return None

        local_vec = local_h1 + local_h2
        local_norm = math.sqrt(sum(x * x for x in local_vec))
        if local_norm < 1e-8:
            return None

        total_r = 0.0
        count = 0

        for peer in self._peers.values():
            if len(peer.h1) != len(local_h1):
                continue
            peer_vec = peer.h1 + peer.h2
            peer_norm = math.sqrt(sum(x * x for x in peer_vec))
            if peer_norm < 1e-8:
                continue

            dot = sum(local_vec[i] * peer_vec[i] for i in range(len(local_vec)))
            total_r += dot / (local_norm * peer_norm)
            count += 1

        return total_r / count if count > 0 else None

    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)


# Deprecated alias
CfCMeshCoupler = SemanticCoupler


def cosine_similarity(a: list[float], b: list[float]) -> float:
    if len(a) != len(b) or not a:
        return 0.0
    dot = sum(x * y for x, y in zip(a, b))
    norm_a = math.sqrt(sum(x * x for x in a))
    norm_b = math.sqrt(sum(x * x for x in b))
    denom = norm_a * norm_b
    return dot / denom if denom > 1e-8 else 0.0
