#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
"""
Dimension E: Uncertainty Quantification for Brain Knowledge.

Replaces flat +0.05 confidence bumps with proper Bayesian updating.
Provides statistical confidence intervals for search result ranking.

Functions:
  - bayesian_update() — Beta-Binomial conjugate prior update
  - wilson_interval() — Wilson score confidence bounds
  - propagate_confidence_update() — graph BFS confidence propagation
  - stamp_uncertainty_fields() — initialize E-fields on new items
"""
import math
from collections import deque


def bayesian_update(prior_confidence: float, evidence_count: int,
                    observation_positive: bool) -> tuple[float, int]:
    """Update confidence using Beta-Binomial conjugate prior.

    Maps prior confidence to Beta(alpha, beta) where:
      alpha = prior * (evidence_count + 2)
      beta = (1 - prior) * (evidence_count + 2)

    Then incorporates the new observation:
      positive → alpha += 1, negative → beta += 1

    Returns (new_confidence, new_evidence_count).

    This replaces the old flat `+0.05` bump in reconsolidation.
    A lesson confirmed by 15 sessions will have meaningfully higher
    confidence than one seen once.
    """
    prior_confidence = max(0.01, min(0.99, prior_confidence))
    n = evidence_count + 2  # Pseudocount gives prior weight
    alpha = prior_confidence * n
    beta = (1 - prior_confidence) * n

    if observation_positive:
        alpha += 1
    else:
        beta += 1

    new_confidence = alpha / (alpha + beta)
    return round(new_confidence, 4), evidence_count + 1


def wilson_interval(confidence: float, evidence_count: int,
                    z: float = 1.96) -> tuple[float, float]:
    """Wilson score confidence interval.

    Works well even for small n (unlike normal approximation).
    Returns (lower_bound, upper_bound) — use lower_bound for
    conservative ranking in search results.

    z=1.96 gives a 95% confidence interval.
    z=1.645 gives a 90% interval (less conservative).
    """
    if evidence_count <= 0:
        return (0.0, 1.0)
    n = evidence_count
    p = max(0.0, min(1.0, confidence))
    z2 = z ** 2

    denominator = 1 + z2 / n
    center = (p + z2 / (2 * n)) / denominator
    spread = z * math.sqrt((p * (1 - p) / n + z2 / (4 * n * n))) / denominator

    lower = max(0.0, round(center - spread, 4))
    upper = min(1.0, round(center + spread, 4))
    return (lower, upper)


def propagate_confidence_update(graph: dict, node_id: str, delta: float,
                                attenuation: float = 0.3,
                                max_hops: int = 2) -> int:
    """Propagate confidence change through causal graph edges.

    When a node's confidence changes, neighbors connected by causal
    edges (CAUSES, LEADS_TO, REQUIRES) receive an attenuated update.
    Each hop multiplies delta by attenuation * edge_confidence.

    Args:
        graph: Knowledge graph dict with "nodes" and "edges".
        node_id: The node whose confidence changed.
        delta: The confidence change (+/- float).
        attenuation: Decay factor per hop (0.3 = 30% propagation).
        max_hops: Maximum BFS depth.

    Returns:
        Number of nodes updated.
    """
    nodes = graph.get("nodes", {})
    edges = graph.get("edges", [])
    causal_types = {"CAUSES", "LEADS_TO", "REQUIRES"}

    # Build adjacency for causal edges only
    adj: dict[str, list[tuple[str, float]]] = {}
    for edge in edges:
        if edge.get("type") in causal_types:
            src = edge.get("from", "")
            dst = edge.get("to", "")
            edge_conf = edge.get("confidence", 0.5)
            if src and dst:
                adj.setdefault(src, []).append((dst, edge_conf))

    # BFS with attenuation
    queue: deque[tuple[str, float, int]] = deque()
    queue.append((node_id, delta, 0))
    visited = {node_id}
    updated = 0

    while queue:
        nid, current_delta, depth = queue.popleft()
        if depth >= max_hops:
            continue
        for neighbor, edge_conf in adj.get(nid, []):
            if neighbor in visited:
                continue
            visited.add(neighbor)
            propagated = current_delta * attenuation * edge_conf

            if abs(propagated) < 0.001:
                continue  # Too small to matter

            node = nodes.get(neighbor, {})
            old_conf = node.get("confidence", 0.5)
            node["confidence"] = max(0.0, min(1.0,
                                    round(old_conf + propagated, 4)))
            updated += 1
            queue.append((neighbor, propagated, depth + 1))

    return updated


def stamp_uncertainty_fields(item: dict) -> dict:
    """Add uncertainty metadata to a new knowledge item.

    Called by save_writers at write time to initialize all E-fields.
    """
    item.setdefault("evidence_count", 1)
    conf = item.get("confidence", 0.5)
    ec = item.get("evidence_count", 1)
    lower, upper = wilson_interval(conf, ec)
    item["confidence_interval"] = [lower, upper]
    return item
