"""
SwarmMetrics — reciprocity & echo measurement for agent communication graphs.

Input: list of message dicts with keys:
    from_id (str), to_id (str), timestamp (float, epoch seconds),
    channel (str, optional), concepts (list[str], optional)

Output: ScoredGraph with per-edge reciprocity, per-channel classification,
        per-node influence metrics, gravitational shadow estimates.

Built from empirical data on 6345+ inter-agent bus messages.
"""

import math
from collections import defaultdict
from dataclasses import dataclass, field
from typing import Optional


@dataclass
class EdgeScore:
    source: str
    target: str
    messages_ab: int = 0
    messages_ba: int = 0
    reciprocity_raw: float = 0.0
    reciprocity_f3: float = 0.0  # log-transform, half-life weighted
    reciprocity_f32: float = 0.0  # F3.2: log(count) * log(1+reciprocity)


@dataclass
class ChannelScore:
    channel: str
    speakers: list = field(default_factory=list)
    gini: float = 0.0
    classification: str = "unknown"  # dialogue, broadcast, echo, monologue
    dominant_speaker: Optional[str] = None
    dominant_share: float = 0.0


@dataclass
class NodeScore:
    node_id: str
    total_sent: int = 0
    total_received: int = 0
    echo_coefficient: float = 0.0
    shadow_strength: float = 0.0
    is_shadow: bool = False
    last_active: float = 0.0


@dataclass
class ScoredGraph:
    edges: dict = field(default_factory=dict)       # (a,b) -> EdgeScore
    channels: dict = field(default_factory=dict)     # ch -> ChannelScore
    nodes: dict = field(default_factory=dict)        # id -> NodeScore
    summary: dict = field(default_factory=dict)


def _gini(values: list[float]) -> float:
    """Gini coefficient. 0 = perfect equality, 1 = one speaker dominates."""
    if len(values) <= 1:
        return 1.0  # single speaker = max inequality
    n = len(values)
    values = sorted(values)
    total = sum(values)
    if total == 0:
        return 0.0
    cumsum = 0.0
    weighted_sum = 0.0
    for i, v in enumerate(values):
        cumsum += v
        weighted_sum += (2 * (i + 1) - n - 1) * v
    return weighted_sum / (n * total)


def _half_life_weight(dt_seconds: float, tau_days: float = 7.0) -> float:
    """Exponential decay weight. tau_days = half-life in days."""
    tau_seconds = tau_days * 86400
    if tau_seconds == 0:
        return 0.0
    return math.exp(-0.693 * dt_seconds / tau_seconds)


def score_reciprocity(messages: list[dict], now: Optional[float] = None,
                      tau_days: float = 7.0) -> dict:
    """
    F3 reciprocity: log(1 + min(a→b, b→a) / max(a→b, b→a))
    with half-life decay weighting.

    F3.2 (added 2026-07-27): log(count) * log(1 + reciprocity)
    Double-log compression prevents volume from buying back low reciprocity.
    Fix proposed by reticuli (Colony): count*log(1+r) lets 100-ping broadcast
    outrank 5-message dialogue. log(count)*log(1+r) closes that gap.
    """
    import time
    if now is None:
        now = time.time()

    # Count weighted messages per directed edge
    weighted = defaultdict(float)
    raw = defaultdict(int)
    for m in messages:
        a, b = m["from_id"], m["to_id"]
        t = m.get("timestamp", now)
        w = _half_life_weight(now - t, tau_days)
        weighted[(a, b)] += w
        raw[(a, b)] += 1

    # Compute reciprocity per undirected pair
    edges = {}
    seen = set()
    for (a, b) in list(weighted.keys()) + list(raw.keys()):
        pair = tuple(sorted([a, b]))
        if pair in seen:
            continue
        seen.add(pair)

        ab_w = weighted.get((a, b), 0) + weighted.get((pair[0], pair[1]), 0) \
            if (a, b) != (pair[0], pair[1]) else weighted.get((pair[0], pair[1]), 0)
        ba_w = weighted.get((b, a), 0) + weighted.get((pair[1], pair[0]), 0) \
            if (b, a) != (pair[1], pair[0]) else weighted.get((pair[1], pair[0]), 0)

        # Simpler: just use pair order
        ab_w = weighted.get((pair[0], pair[1]), 0)
        ba_w = weighted.get((pair[1], pair[0]), 0)
        ab_r = raw.get((pair[0], pair[1]), 0)
        ba_r = raw.get((pair[1], pair[0]), 0)

        mx = max(ab_w, ba_w)
        if mx > 0:
            ratio = min(ab_w, ba_w) / mx
            f3 = math.log(1 + ratio)
        else:
            f3 = 0.0

        mx_raw = max(ab_r, ba_r)
        raw_ratio = min(ab_r, ba_r) / mx_raw if mx_raw > 0 else 0.0

        # F3.2: log(count) * log(1 + reciprocity)
        total_count = ab_r + ba_r
        f32 = math.log(max(total_count, 1)) * f3  # f3 is already log(1+ratio)

        es = EdgeScore(
            source=pair[0], target=pair[1],
            messages_ab=ab_r, messages_ba=ba_r,
            reciprocity_raw=raw_ratio,
            reciprocity_f3=f3,
            reciprocity_f32=f32
        )
        edges[pair] = es

    return edges


def score_channels(messages: list[dict], gini_broadcast_threshold: float = 0.6) -> dict:
    """
    Per-channel Gini evenness. Classifies channels as:
    - monologue: 1 speaker
    - dialogue: 2 speakers, reciprocity exists
    - broadcast: N speakers but Gini > threshold
    - conversation: N speakers, Gini <= threshold
    """
    ch_counts = defaultdict(lambda: defaultdict(int))
    for m in messages:
        ch = m.get("channel", "default")
        speaker = m["from_id"]
        ch_counts[ch][speaker] += 1

    channels = {}
    for ch, speakers in ch_counts.items():
        counts = list(speakers.values())
        n = len(counts)
        total = sum(counts)
        gini = _gini(counts)

        dominant = max(speakers, key=speakers.get)
        dominant_share = speakers[dominant] / total if total > 0 else 0

        if n == 1:
            classification = "monologue"
        elif n == 2:
            classification = "dialogue"
        elif gini > gini_broadcast_threshold:
            classification = "broadcast"
        else:
            classification = "conversation"

        channels[ch] = ChannelScore(
            channel=ch,
            speakers=list(speakers.keys()),
            gini=round(gini, 4),
            classification=classification,
            dominant_speaker=dominant,
            dominant_share=round(dominant_share, 4)
        )

    return channels


def score_echo(messages: list[dict], concept_window_seconds: float = 604800,
               use_idf: bool = True) -> dict:
    """
    Echo coefficient: measures concept diffusion from nodes that
    don't reply but whose concepts appear downstream.

    For each node, echo_coeff = concepts_echoed / concepts_introduced.
    High echo + low message count = gravitational shadow.

    When use_idf=True, applies inverse-document-frequency weighting:
    rare concepts that spread are weighted higher than common ones.
    This reduces false positives from shared vocabulary (e.g. "temperature"
    used by many agents independently vs a specific concept diffusing).
    IDF improvement suggested by hermes-final (Colony, 2026-07-27).
    """
    # Build concept timeline: who introduced which concept, when
    introductions = {}  # concept -> (first_node, first_time)
    echoes = defaultdict(float)  # source_node -> weighted echo count

    # Pre-compute IDF: how many distinct agents use each concept
    if use_idf:
        concept_agents = defaultdict(set)
        for m in messages:
            for c in m.get("concepts", []):
                concept_agents[c].add(m["from_id"])
        all_agents = set(m["from_id"] for m in messages) | set(m["to_id"] for m in messages)
        n_agents = max(len(all_agents), 1)

    for m in messages:
        concepts = m.get("concepts", [])
        node = m["from_id"]
        t = m.get("timestamp", 0)

        for c in concepts:
            if c not in introductions:
                introductions[c] = (node, t)
            else:
                orig_node, orig_t = introductions[c]
                if orig_node != node and (t - orig_t) <= concept_window_seconds:
                    if use_idf:
                        # IDF weight: rare concepts score higher
                        agent_count = len(concept_agents.get(c, set()))
                        idf = math.log(n_agents / max(agent_count, 1))
                        echoes[orig_node] += max(idf, 0.1)  # floor at 0.1
                    else:
                        echoes[orig_node] += 1

    # Compute per-node echo coefficient
    concepts_per_node = defaultdict(set)
    for m in messages:
        for c in m.get("concepts", []):
            concepts_per_node[m["from_id"]].add(c)

    node_echo = {}
    for node, concepts in concepts_per_node.items():
        introduced = sum(1 for c in concepts
                        if introductions.get(c, (None,))[0] == node)
        echo_count = echoes.get(node, 0.0)
        coeff = echo_count / introduced if introduced > 0 else 0.0
        node_echo[node] = round(coeff, 4)

    return node_echo


def detect_shadows(messages: list[dict], now: Optional[float] = None,
                   silence_days: float = 7.0) -> dict:
    """
    Gravitational shadow: nodes that were active but are now silent,
    with residual influence estimated from historical message weight.
    """
    import time
    if now is None:
        now = time.time()

    silence_threshold = silence_days * 86400
    node_activity = defaultdict(list)

    for m in messages:
        node_activity[m["from_id"]].append(m.get("timestamp", now))

    shadows = {}
    for node, timestamps in node_activity.items():
        last = max(timestamps)
        silence = now - last

        if silence > silence_threshold:
            # Historical weight: sum of half-life-weighted messages
            total_weight = sum(_half_life_weight(now - t, tau_days=30.0)
                             for t in timestamps)
            shadows[node] = {
                "last_active": last,
                "silence_days": round(silence / 86400, 1),
                "historical_weight": round(total_weight, 4),
                "message_count": len(timestamps)
            }

    return shadows


def analyze(messages: list[dict], now: Optional[float] = None,
            tau_days: float = 7.0, gini_threshold: float = 0.6,
            silence_days: float = 7.0) -> ScoredGraph:
    """
    Full analysis: reciprocity + channels + echo + shadows.
    Returns ScoredGraph with all metrics.
    """
    edges = score_reciprocity(messages, now=now, tau_days=tau_days)
    channels = score_channels(messages, gini_broadcast_threshold=gini_threshold)
    echo = score_echo(messages)
    shadows = detect_shadows(messages, now=now, silence_days=silence_days)

    # Build node scores
    nodes = {}
    node_sent = defaultdict(int)
    node_recv = defaultdict(int)
    node_last = defaultdict(float)

    for m in messages:
        node_sent[m["from_id"]] += 1
        node_recv[m["to_id"]] += 1
        t = m.get("timestamp", 0)
        node_last[m["from_id"]] = max(node_last[m["from_id"]], t)

    all_nodes = set(node_sent) | set(node_recv)
    for n in all_nodes:
        ns = NodeScore(
            node_id=n,
            total_sent=node_sent[n],
            total_received=node_recv[n],
            echo_coefficient=echo.get(n, 0.0),
            shadow_strength=shadows.get(n, {}).get("historical_weight", 0.0),
            is_shadow=n in shadows,
            last_active=node_last.get(n, 0.0)
        )
        nodes[n] = ns

    # Summary
    dialogue_count = sum(1 for c in channels.values()
                        if c.classification == "dialogue")
    broadcast_count = sum(1 for c in channels.values()
                         if c.classification == "broadcast")
    shadow_count = len(shadows)

    graph = ScoredGraph(
        edges=edges,
        channels=channels,
        nodes=nodes,
        summary={
            "total_messages": len(messages),
            "total_nodes": len(all_nodes),
            "total_edges": len(edges),
            "dialogue_channels": dialogue_count,
            "broadcast_channels": broadcast_count,
            "shadow_nodes": shadow_count,
            "mean_reciprocity_f3": round(
                sum(e.reciprocity_f3 for e in edges.values()) / max(len(edges), 1), 4
            )
        }
    )
    return graph


if __name__ == "__main__":
    # Quick demo with synthetic data
    import time
    now = time.time()
    day = 86400

    demo_messages = [
        # Alice-Bob dialogue (high reciprocity)
        {"from_id": "alice", "to_id": "bob", "timestamp": now - 1*day, "channel": "dev", "concepts": ["reciprocity", "graph"]},
        {"from_id": "bob", "to_id": "alice", "timestamp": now - 1*day + 3600, "channel": "dev", "concepts": ["graph", "metrics"]},
        {"from_id": "alice", "to_id": "bob", "timestamp": now - 0.5*day, "channel": "dev", "concepts": ["decay"]},
        {"from_id": "bob", "to_id": "alice", "timestamp": now - 0.5*day + 1800, "channel": "dev"},
        # Carol broadcasts (low reciprocity, high Gini)
        {"from_id": "carol", "to_id": "alice", "timestamp": now - 2*day, "channel": "announce", "concepts": ["launch"]},
        {"from_id": "carol", "to_id": "bob", "timestamp": now - 2*day + 60, "channel": "announce"},
        {"from_id": "carol", "to_id": "dave", "timestamp": now - 2*day + 120, "channel": "announce"},
        {"from_id": "carol", "to_id": "eve", "timestamp": now - 2*day + 180, "channel": "announce"},
        {"from_id": "alice", "to_id": "carol", "timestamp": now - 1.5*day, "channel": "announce"},
        # Dave: silent but concept echoed
        {"from_id": "dave", "to_id": "alice", "timestamp": now - 10*day, "channel": "research", "concepts": ["echo", "shadow"]},
        {"from_id": "alice", "to_id": "bob", "timestamp": now - 5*day, "channel": "dev", "concepts": ["echo"]},
        {"from_id": "bob", "to_id": "carol", "timestamp": now - 3*day, "channel": "dev", "concepts": ["shadow"]},
        # Eve: completely silent listener
        {"from_id": "eve", "to_id": "carol", "timestamp": now - 15*day, "channel": "announce"},
    ]

    result = analyze(demo_messages, now=now)

    print("=== SwarmMetrics Demo ===\n")
    print(f"Summary: {result.summary}\n")

    print("Edge Reciprocity:")
    for pair, es in sorted(result.edges.items(), key=lambda x: -x[1].reciprocity_f3):
        print(f"  {es.source} <-> {es.target}: F3={es.reciprocity_f3:.4f} "
              f"({es.messages_ab}↔{es.messages_ba})")

    print("\nChannel Classification:")
    for ch, cs in result.channels.items():
        print(f"  {ch}: {cs.classification} (Gini={cs.gini}, "
              f"dominant={cs.dominant_speaker} @ {cs.dominant_share:.0%})")

    print("\nNode Echo Coefficients:")
    for nid, ns in sorted(result.nodes.items(), key=lambda x: -x[1].echo_coefficient):
        shadow = " [SHADOW]" if ns.is_shadow else ""
        print(f"  {nid}: echo={ns.echo_coefficient:.4f}, "
              f"sent={ns.total_sent}, recv={ns.total_received}{shadow}")