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

Implements SVAF-inspired fusion where the agent evaluates each peer's
hidden state and autonomously decides whether to couple, how much,
and which neurons to couple.

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

import math
import time
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 _PeerState:
    h1: list[float]
    h2: list[float]
    confidence: float
    timestamp: float


class CfCMeshCoupler:
    """Autonomous coupling engine for CfC hidden state exchange.

    The agent evaluates each peer's hidden state against its own cognitive
    state and autonomously decides whether to couple, how much to couple,
    and which neurons to couple. This is NOT mechanical blending.
    """

    def __init__(self, config: Optional[CouplingConfig] = None):
        self.config = config or CouplingConfig()
        self._peers: dict[str, _PeerState] = {}
        self.last_decisions: dict[str, CouplingDecision] = {}

    # -- Peer Management --

    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

    # -- Autonomous Coupling --

    def couple(self, local_h1: list[float], local_h2: list[float]) -> tuple[list[float], list[float]]:
        """Autonomously couple local hidden state with mesh peers.

        Returns coupled (h1, h2). Returns local state unchanged if
        no peers or all peers rejected.
        """
        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]] = []  # (peer, weight, alpha)

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

            drift_h1 = 1.0 - self._cosine_similarity(local_h1, peer.h1)
            drift_h2 = 1.0 - self._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)

        # Compute weighted mesh aggregate
        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

        # Effective alpha
        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 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

            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

    # -- Kuramoto Metrics --

    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 kuramoto_order_parameter(self, local_h1: list[float], local_h2: list[float]) -> Optional[float]:
        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():
            dim = len(local_h1)
            if len(peer.h1) != dim or len(peer.h2) != dim:
                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

    # -- Private --

    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)

    @staticmethod
    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
