pc.py

from __future__ import annotations
from dataclasses import dataclass
from typing import Dict, Iterable, List, Tuple
import math
import random

from .contexts import Context, ContextId


@dataclass
class PCParams:
    # theta[cid] = reliability parameter in (0,1)
    theta: Dict[ContextId, float]
    # training metadata
    alpha: float = 1.0  # Laplace smoothing


def initialize_params(contexts: Dict[ContextId, Context]) -> PCParams:
    theta = {cid: min(0.9, max(0.1, ctx.support_facts / max(1, ctx.support_facts + ctx.negative_hits))) for cid, ctx in contexts.items()}
    return PCParams(theta=theta)


def estimate_with_negatives(
    contexts: Dict[ContextId, Context],
    neg_samples_per_ctx: int = 2,
    seed: int = 0,
) -> PCParams:
    # Simple frequency-based estimation with Laplace smoothing
    random.seed(seed)
    theta: Dict[ContextId, float] = {}
    for cid, ctx in contexts.items():
        pos = ctx.support_facts
        neg = ctx.negative_hits
        # Estimate: (pos + alpha) / (pos + neg + 2*alpha)
        alpha = 1.0
        th = (pos + alpha) / (pos + neg + 2 * alpha) if (pos + neg) > 0 else 0.5
        theta[cid] = max(0.01, min(0.99, th))
    return PCParams(theta=theta)


def combine_noisy_or(supporting_contexts: Iterable[ContextId], params: PCParams) -> float:
    p = 0.0
    prod = 1.0
    for cid in supporting_contexts:
        prod *= (1.0 - params.theta.get(cid, 0.2))
    p = 1.0 - prod
    return max(0.0, min(1.0, p))