#!/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
  - stamp_uncertainty_fields() — initialize E-fields on new items
"""
import math


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