Dochia v0.1.0: model card, params and reference implementations
e215d3c 9 days ago
27.6 kB
"""
Dochia v1 — pure-Python reference implementation.
This module is the analytical reference used to fit parameters, evaluate the
model and generate golden vectors for the PHP implementation. It is not app
code; it is the "prepared model" artefact (see ../README.md).
No third-party dependencies (works with the stdlib only). All formulas are
implemented to match the PHP reference in ../src.
Model summary
-------------
Cycle length y ~ two-component lognormal mixture:
    y = (1 - pi) * LogNormal(m_main, s_main) + pi * LogNormal(m_tail, s_tail)
with moments matched to the personalised (mean, sd) and a heavy-tail component
for skipped / irregular cycles. Onset hazard is the discrete hazard of that
predictive law conditioned on "not yet". Ovulation day O = y - L with
L ~ N(14.13, 1.41); Var(O) = Var(y) - Var(L) (the corrected identity).
"""

from __future__ import annotations

import json
import math
from dataclasses import dataclass, field
from pathlib import Path
from typing import Iterable, Sequence

SQRT2 = math.sqrt(2.0)
SQRT2PI = math.sqrt(2.0 * math.pi)
MAE_FLOOR_FACTOR = math.sqrt(2.0 / math.pi)

# ---------------------------------------------------------------------------
# Normal / lognormal helpers
# ---------------------------------------------------------------------------


def norm_cdf(z: float) -> float:
    return 0.5 * (1.0 + math.erf(z / SQRT2))


def norm_pdf(z: float) -> float:
    return math.exp(-0.5 * z * z) / SQRT2PI


def norm_ppf(p: float) -> float:
    """Inverse normal CDF by bisection (stdlib only, ~1e-13 accuracy)."""
    if p <= 0.0:
        return -math.inf
    if p >= 1.0:
        return math.inf
    lo, hi = -12.0, 12.0
    for _ in range(200):
        mid = 0.5 * (lo + hi)
        if norm_cdf(mid) < p:
            lo = mid
        else:
            hi = mid
    return 0.5 * (lo + hi)


def _lognormal_params(mean: float, sd: float) -> tuple[float, float]:
    """Return (m, s) in log space so LogNormal(m, s) has the given mean/sd."""
    mean = max(mean, 1e-6)
    sd = max(sd, 1e-6)
    s2 = math.log1p((sd * sd) / (mean * mean))
    s = math.sqrt(s2)
    m = math.log(mean) - 0.5 * s2
    return m, s


# ---------------------------------------------------------------------------
# Parameters
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class ModelParams:
    version: str = "dochia-v1-literature"
    # population prior
    pop_mean: float = 29.0
    pop_between_sd: float = 3.5
    within_sd_prior: float = 2.0
    prior_weight: float = 4.0
    min_mean: float = 18.0
    max_mean: float = 90.0
    min_within_var: float = 1.44
    max_within_var: float = 36.0
    # luteal / fertile window
    luteal_mean: float = 14.13
    luteal_sd: float = 1.41
    fertile_window_days: int = 6
    ovulation_sd_floor: float = 0.35
    # predictive mixture / trend / regime
    tail_prob_base: float = 0.05
    tail_prob_max: float = 0.40
    tail_prob_sd_gain: float = 0.06
    tail_prob_range_gain: float = 0.02
    tail_prob_gap_bonus: float = 0.20
    tail_shift_days: float = 2.0
    tail_sd_multiplier: float = 1.8
    variance_blend: float = 0.5
    moment_match: bool = False
    trend_ewma: float = 0.4
    trend_threshold: float = 1.5
    trend_max_shift: float = 3.0
    # calibration: body scale + separate tail stretching around anchors
    body_scale: float = 1.0
    tail_scale_low: float = 1.0
    tail_scale_high: float = 1.0
    tail_anchor_low: float = 0.10
    tail_anchor_high: float = 0.90
    calibration_method: str = "none"
    # plausibility
    min_length: int = 15
    max_length: int = 120
    # hazard / ovulation grids
    hazard_max_day: int = 120
    ovulation_grid_step: float = 0.25
    # abstention thresholds
    abstain_gap_days: int = 90
    abstain_recent_sd: float = 6.0
    abstain_min_cycles: int = 0

    @classmethod
    def from_dict(cls, raw: dict) -> "ModelParams":
        flat: dict = {"version": raw.get("version", cls.version)}

        pop = raw.get("population", {})
        flat.update(
            pop_mean=pop.get("mean", 29.0),
            pop_between_sd=pop.get("between_sd", 3.5),
            within_sd_prior=pop.get("within_sd_prior", 2.0),
            prior_weight=pop.get("prior_weight", 4.0),
            min_mean=pop.get("min_mean", 18.0),
            max_mean=pop.get("max_mean", 90.0),
            min_within_var=pop.get("min_within_var", 1.44),
            max_within_var=pop.get("max_within_var", 36.0),
        )
        luteal = raw.get("luteal", {})
        flat.update(
            luteal_mean=luteal.get("mean", 14.13),
            luteal_sd=luteal.get("sd", 1.41),
            fertile_window_days=raw.get("fertile_window_days", 6),
        )
        pred = raw.get("predictive", {})
        flat.update(
            tail_prob_base=pred.get("tail_prob_base", 0.05),
            tail_prob_max=pred.get("tail_prob_max", 0.40),
            tail_prob_sd_gain=pred.get("tail_prob_sd_gain", 0.06),
            tail_prob_range_gain=pred.get("tail_prob_range_gain", 0.02),
            tail_prob_gap_bonus=pred.get("tail_prob_gap_bonus", 0.20),
            tail_shift_days=pred.get("tail_shift_days", 2.0),
            tail_sd_multiplier=pred.get("tail_sd_multiplier", 1.8),
            variance_blend=pred.get("variance_blend", 0.5),
            moment_match=pred.get("moment_match", False),
            trend_ewma=pred.get("trend_ewma", 0.4),
            trend_threshold=pred.get("trend_threshold", 1.5),
            trend_max_shift=pred.get("trend_max_shift", 3.0),
        )
        cal = raw.get("calibration", {})
        flat.update(
            body_scale=cal.get("body_scale", 1.0),
            tail_scale_low=cal.get("tail_scale_low", 1.0),
            tail_scale_high=cal.get("tail_scale_high", 1.0),
            tail_anchor_low=cal.get("tail_anchor_low", 0.10),
            tail_anchor_high=cal.get("tail_anchor_high", 0.90),
            calibration_method=cal.get("method", "none"),
        )
        plaus = raw.get("plausible", {})
        flat.update(
            min_length=plaus.get("min", 15),
            max_length=plaus.get("max", 120),
        )
        ov = raw.get("ovulation", {})
        flat.update(
            ovulation_sd_floor=ov.get("sd_floor", 0.35),
            ovulation_grid_step=ov.get("grid_step", 0.25),
        )
        ab = raw.get("abstain", {})
        flat.update(
            abstain_gap_days=ab.get("gap_days", 90),
            abstain_recent_sd=ab.get("recent_sd", 6.0),
            abstain_min_cycles=ab.get("min_cycles", 0),
        )
        return cls(**flat)

    @classmethod
    def from_json(cls, path: str | Path) -> "ModelParams":
        return cls.from_dict(json.loads(Path(path).read_text()))


# ---------------------------------------------------------------------------
# Lognormal mixture
# ---------------------------------------------------------------------------


@dataclass
class LognormalMixture:
    """Two-component lognormal mixture with exact moments."""

    mean: float
    sd: float
    tail_prob: float
    components: list[tuple[float, float, float]] = field(default_factory=list)

    @classmethod
    def from_moments(
        cls,
        mean: float,
        sd: float,
        pi: float,
        tail_shift: float,
        tail_sd_multiplier: float,
        moment_match: bool = False,
    ) -> "LognormalMixture":
        pi = min(max(pi, 0.0), 0.45)
        mean = max(mean, 1.0)
        sd = max(sd, 0.2)
        if pi <= 1e-6:
            m, s = _lognormal_params(mean, sd)
            return cls(mean=mean, sd=sd, tail_prob=0.0, components=[(1.0, m, s)])

        tail_mean = max(mean + tail_shift, mean * 1.02)
        tail_sd = max(sd * tail_sd_multiplier, 0.3)
        if not moment_match:
            # Keep the body distribution intact (protects the 80% interval) and
            # add a right-tail component; the overall mean/sd rise slightly.
            m_main, s_main = _lognormal_params(mean, sd)
            m_tail, s_tail = _lognormal_params(tail_mean, tail_sd)
            return cls(
                mean=mean,
                sd=sd,
                tail_prob=pi,
                components=[(1.0 - pi, m_main, s_main), (pi, m_tail, s_tail)],
            )
        # Solve the main component moments so the mixture matches (mean, var).
        main_mean = (mean - pi * tail_mean) / (1.0 - pi)
        var_target = sd * sd
        main_var = (
            (var_target + mean * mean - pi * (tail_sd * tail_sd + tail_mean * tail_mean)) / (1.0 - pi)
            - main_mean * main_mean
        )
        if main_mean <= 0.5 or main_var <= 0.04:
            # Degenerate tail: fall back to the single-component distribution.
            m, s = _lognormal_params(mean, sd)
            return cls(mean=mean, sd=sd, tail_prob=0.0, components=[(1.0, m, s)])
        m_main, s_main = _lognormal_params(main_mean, math.sqrt(main_var))
        m_tail, s_tail = _lognormal_params(tail_mean, tail_sd)
        return cls(
            mean=mean,
            sd=sd,
            tail_prob=pi,
            components=[(1.0 - pi, m_main, s_main), (pi, m_tail, s_tail)],
        )

    def cdf(self, x: float) -> float:
        if x <= 0.0:
            return 0.0
        total = 0.0
        for w, m, s in self.components:
            total += w * norm_cdf((math.log(x) - m) / s)
        return min(max(total, 0.0), 1.0)

    def pdf(self, x: float) -> float:
        if x <= 0.0:
            return 0.0
        total = 0.0
        for w, m, s in self.components:
            z = (math.log(x) - m) / s
            total += w * norm_pdf(z) / (x * s)
        return total

    def quantile(self, q: float) -> float:
        if q <= 0.0:
            return 0.0
        if q >= 1.0:
            return 1000.0
        lo, hi = 1.0, 400.0
        if self.cdf(lo) >= q:
            return lo
        if self.cdf(hi) <= q:
            return hi
        for _ in range(40):
            mid = 0.5 * (lo + hi)
            if self.cdf(mid) < q:
                lo = mid
            else:
                hi = mid
        return 0.5 * (lo + hi)

    def crps(self, observed: float, grid: Sequence[float] | None = None) -> float:
        """CRPS = integral (F(z) - 1{z >= observed})^2 dz, numerically."""
        if grid is None:
            grid = [1.0 + 0.5 * i for i in range(int((365.0 - 1.0) / 0.5) + 1)]
        total = 0.0
        for i in range(len(grid) - 1):
            z = 0.5 * (grid[i] + grid[i + 1])
            f = self.cdf(z)
            indicator = 1.0 if observed <= z else 0.0
            total += (f - indicator) ** 2 * (grid[i + 1] - grid[i])
        return total


# ---------------------------------------------------------------------------
# Calibrated predictive distribution
# ---------------------------------------------------------------------------


@dataclass
class CalibratedPredictive:
    """Predictive distribution with body scale and tail-only stretching.
    The body is a lognormal mixture with sd multiplied by `body_scale`. The
    lower tail below the `anchor_low` quantile and the upper tail above the
    `anchor_high` quantile are stretched by `tail_scale_low` / `tail_scale_high`
    respectively. This lets the 80% interval (body) and the 95% interval (tails)
    be calibrated separately without changing the likelihood in between.
    """

    base: LognormalMixture
    anchor_low: float
    anchor_high: float
    scale_low: float
    scale_high: float
    q_low: float
    q_high: float

    @classmethod
    def from_base(cls, base: LognormalMixture, params: "ModelParams") -> "CalibratedPredictive":
        return cls(
            base=base,
            anchor_low=params.tail_anchor_low,
            anchor_high=params.tail_anchor_high,
            scale_low=max(params.tail_scale_low, 1e-6),
            scale_high=max(params.tail_scale_high, 1e-6),
            q_low=base.quantile(params.tail_anchor_low),
            q_high=base.quantile(params.tail_anchor_high),
        )

    def quantile(self, p: float) -> float:
        if p <= 0.0:
            return 0.0
        if p >= 1.0:
            return 1000.0
        if p <= self.anchor_low:
            return self.q_low + (self.base.quantile(p) - self.q_low) * self.scale_low
        if p >= self.anchor_high:
            return self.q_high + (self.base.quantile(p) - self.q_high) * self.scale_high
        return self.base.quantile(p)

    def cdf(self, x: float) -> float:
        if x <= self.q_low:
            y = self.q_low + (x - self.q_low) / self.scale_low
            return min(max(self.base.cdf(y), 0.0), 1.0)
        if x >= self.q_high:
            y = self.q_high + (x - self.q_high) / self.scale_high
            return min(max(self.base.cdf(y), 0.0), 1.0)
        return self.base.cdf(x)

    def pdf(self, x: float) -> float:
        if x <= 0.0:
            return 0.0
        if x < self.q_low:
            return self.base.pdf(self.q_low + (x - self.q_low) / self.scale_low) / self.scale_low
        if x > self.q_high:
            return self.base.pdf(self.q_high + (x - self.q_high) / self.scale_high) / self.scale_high
        return self.base.pdf(x)

    @property
    def mean(self) -> float:
        return self.base.mean

    @property
    def sd(self) -> float:
        return self.base.sd

    def crps(self, observed: float, grid: Sequence[float] | None = None) -> float:
        if grid is None:
            grid = [1.0 + 0.5 * i for i in range(int((365.0 - 1.0) / 0.5) + 1)]
        total = 0.0
        for i in range(len(grid) - 1):
            z = 0.5 * (grid[i] + grid[i + 1])
            f = self.cdf(z)
            indicator = 1.0 if observed <= z else 0.0
            total += (f - indicator) ** 2 * (grid[i + 1] - grid[i])
        return total


# ---------------------------------------------------------------------------
# Personal distribution & predictive
# ---------------------------------------------------------------------------


@dataclass
class PersonalState:
    mean: float
    within_var: float
    cycle_sd: float
    n: int
    tail_prob: float
    trend: float
    recent_sd: float
    recent_range: int
    gap: bool
    abstain: bool
    confidence: str


def _median(values: Sequence[float]) -> float:
    s = sorted(values)
    n = len(s)
    if n == 0:
        return 0.0
    return s[n // 2] if n % 2 == 1 else 0.5 * (s[n // 2 - 1] + s[n // 2])


def personal_state(lengths: Sequence[int], params: ModelParams) -> PersonalState:
    lengths = [int(y) for y in lengths if params.min_length <= int(y) <= params.max_length]
    n = len(lengths)
    if n == 0:
        return PersonalState(
            mean=params.pop_mean,
            within_var=params.within_sd_prior**2,
            cycle_sd=math.sqrt(params.within_sd_prior**2 + (params.pop_between_sd**2 * params.within_sd_prior**2) / max(1e-9, params.within_sd_prior**2)),
            n=0,
            tail_prob=params.tail_prob_base,
            trend=0.0,
            recent_sd=0.0,
            recent_range=0,
            gap=False,
            abstain=False,
            confidence="new",
        )

    mean = sum(lengths) / n
    sample_var = 0.0
    if n > 1:
        sample_var = sum((y - mean) ** 2 for y in lengths) / (n - 1)
        # robust blend (MAD) so a single outlier does not explode the interval
        med = _median(lengths)
        mad = _median([abs(y - med) for y in lengths])
        robust_sd = 1.4826 * mad
        robust_var = robust_sd * robust_sd
        sample_var = params.variance_blend * sample_var + (1.0 - params.variance_blend) * robust_var

    prior_var = params.within_sd_prior**2
    prior_weight = params.prior_weight
    # If the observed variance clearly exceeds the prior, trust the user's data
    # more; this widens intervals for genuinely variable users instead of
    # shrinking them back to the population prior.
    if sample_var > prior_var:
        prior_weight = max(0.5, params.prior_weight * prior_var / sample_var)
    within_var = (n * sample_var + prior_weight * prior_var) / (n + prior_weight)
    within_var = min(max(within_var, params.min_within_var), params.max_within_var)

    between_var = params.pop_between_sd**2
    weight = (n * between_var) / max(1e-9, n * between_var + within_var)
    shrunk_mean = weight * mean + (1.0 - weight) * params.pop_mean
    shrunk_mean = min(max(shrunk_mean, params.min_mean), params.max_mean)

    posterior_var = (within_var * between_var) / max(1e-9, n * between_var + within_var)
    cycle_sd = math.sqrt(within_var + posterior_var)

    # trend: EWMA of recent deviations from the personal mean
    recent = lengths[-6:]
    trend = 0.0
    if n >= 4 and len(recent) >= 3:
        weights = []
        w = 1.0
        for _ in reversed(recent):
            weights.append(w)
            w *= 1.0 - params.trend_ewma
        weights.reverse()
        total_w = sum(weights)
        trend = sum((y - shrunk_mean) * wt for y, wt in zip(recent, weights)) / total_w
        if abs(trend) < params.trend_threshold:
            trend_shift = 0.0
        else:
            trend_shift = max(-params.trend_max_shift, min(params.trend_max_shift, trend))
        shrunk_mean = min(max(shrunk_mean + trend_shift, params.min_mean), params.max_mean)

    recent_vals = lengths[-6:]
    recent_mean = sum(recent_vals) / len(recent_vals)
    recent_sd = math.sqrt(sum((y - recent_mean) ** 2 for y in recent_vals) / max(1, len(recent_vals) - 1))
    recent_range = max(recent_vals) - min(recent_vals)
    gap = max(recent_vals) > params.abstain_gap_days

    pi = params.tail_prob_base
    pi += params.tail_prob_sd_gain * max(0.0, recent_sd - 2.0)
    pi += params.tail_prob_range_gain * max(0.0, recent_range - 7)
    if gap:
        pi += params.tail_prob_gap_bonus
    pi = min(max(pi, params.tail_prob_base), params.tail_prob_max)

    abstain = gap or recent_sd > params.abstain_recent_sd or n < params.abstain_min_cycles

    confidence = (
        "abstain"
        if abstain
        else "established"
        if n >= 6
        else "building"
        if n >= 3
        else "learning"
        if n >= 1
        else "new"
    )

    return PersonalState(
        mean=shrunk_mean,
        within_var=within_var,
        cycle_sd=cycle_sd,
        n=n,
        tail_prob=pi,
        trend=trend,
        recent_sd=recent_sd,
        recent_range=recent_range,
        gap=gap,
        abstain=abstain,
        confidence=confidence,
    )


def predictive(state: PersonalState, params: ModelParams) -> CalibratedPredictive:
    sd = state.cycle_sd * params.body_scale
    base = LognormalMixture.from_moments(
        mean=state.mean,
        sd=sd,
        pi=state.tail_prob,
        tail_shift=params.tail_shift_days,
        tail_sd_multiplier=params.tail_sd_multiplier,
        moment_match=params.moment_match,
    )
    return CalibratedPredictive.from_base(base, params)


# ---------------------------------------------------------------------------
# Hazard and ovulation
# ---------------------------------------------------------------------------


def hazard_curve(dist: CalibratedPredictive, max_day: int = 120) -> dict[int, float]:
    """Discrete hazard h(d) = P(Y=d | Y>=d) for d = 1..max_day."""
    curve: dict[int, float] = {}
    prev_survival = 1.0
    for day in range(1, max_day + 1):
        if prev_survival <= 1e-12:
            curve[day] = 1.0
            continue
        f = dist.cdf(day) - dist.cdf(day - 1.0)
        h = max(0.0, f) / prev_survival
        curve[day] = min(max(h, 0.0), 1.0)
        prev_survival -= max(0.0, f)
        if prev_survival < 0.0:
            prev_survival = 0.0
    return curve


def onset_probabilities(dist: CalibratedPredictive, current_day: int, params: ModelParams) -> dict:
    """Probabilities conditioned on 'no onset yet' at current_day."""
    current_day = max(1, current_day)
    survival_now = max(1e-12, 1.0 - dist.cdf(current_day - 1.0))

    def within(k: int) -> float:
        survival_then = max(0.0, 1.0 - dist.cdf(current_day - 1 + k))
        return min(max(1.0 - survival_then / survival_now, 0.0), 1.0)

    curve = hazard_curve(dist, params.hazard_max_day)
    p_today = curve.get(current_day, 0.0)
    return {
        "current_day": current_day,
        "p_today": p_today,
        "p_within_3": within(3),
        "p_within_7": within(7),
        "p_within_14": within(14),
        "late_probability": (1.0 - dist.cdf(current_day - 1.0)),
        "curve": {d: curve[d] for d in range(current_day, min(params.hazard_max_day, current_day + 13) + 1)},
    }


def ovulation_distribution(dist: CalibratedPredictive, params: ModelParams) -> dict:
    """Discrete P(O=d) and derived fertile-window probabilities.
    O = y - L with L ~ N(mu_L, sigma_L). Under independence:
        P(O <= d) = E_y[P(L >= y - d)] = E_y[1 - Phi((y - d - mu_L)/sigma_L)],
    integrated over the predictive distribution of y by bin midpoint rule.
    Var(O) = Var(y) - Var(L), floored (the corrected variance identity).
    """
    sd_o = math.sqrt(max(dist.sd * dist.sd - params.luteal_sd**2, params.ovulation_sd_floor**2))
    mean_o = dist.mean - params.luteal_mean

    step = params.ovulation_grid_step
    max_y = 365.0
    edges = []
    yv = 0.0
    while yv <= max_y:
        edges.append(yv)
        yv += step
    cdf_edges = [dist.cdf(b) for b in edges]

    def cdf_o(d: float) -> float:
        total = 0.0
        for i in range(len(edges) - 1):
            mass = cdf_edges[i + 1] - cdf_edges[i]
            if mass <= 0.0:
                continue
            y_mid = 0.5 * (edges[i] + edges[i + 1])
            z = (d - y_mid + params.luteal_mean) / params.luteal_sd
            total += mass * norm_cdf(z)
        return total

    probs: dict[int, float] = {}
    max_day = params.hazard_max_day
    prev = 0.0
    for day in range(1, max_day + 1):
        cur = cdf_o(day)
        probs[day] = min(max(cur - prev, 0.0), 1.0)
        prev = cur
    # residual tail beyond the grid is assigned to the last day (honest tail)
    probs[max_day] += max(0.0, 1.0 - prev)

    fertile: dict[int, float] = {}
    for day in range(1, max_day + 1):
        fertile[day] = min(
            max(
                sum(probs[o] for o in range(day, min(max_day, day + params.fertile_window_days - 1) + 1)),
                0.0,
            ),
            1.0,
        )

    max_prob = max(fertile.values()) if fertile else 0.0
    peak_days = [d for d, p in fertile.items() if p >= max(0.25, 0.5 * max_prob)]
    window_days = [d for d, p in fertile.items() if p >= 0.05]

    return {
        "mean_day": mean_o,
        "sd_day": sd_o,
        "probabilities": probs,
        "fertile_probabilities": fertile,
        "peak_days": peak_days,
        "window_start": min(window_days) if window_days else None,
        "window_end": max(window_days) if window_days else None,
    }


# ---------------------------------------------------------------------------
# Forecast facade
# ---------------------------------------------------------------------------


def forecast(lengths: Sequence[int], *, current_day: int = 1, params: ModelParams | None = None) -> dict:
    params = params or ModelParams()
    state = personal_state(lengths, params)
    dist = predictive(state, params)
    hazard = onset_probabilities(dist, current_day, params)
    ovulation = ovulation_distribution(dist, params)

    quantiles = {str(q): dist.quantile(q) for q in (0.025, 0.05, 0.10, 0.25, 0.50, 0.75, 0.90, 0.95, 0.975)}

    explanations: list[str] = []
    if state.abstain:
        explanations.append("Insufficient or highly variable history for a confident forecast.")
    if state.gap:
        explanations.append("A long gap was detected; the forecast is intentionally wide.")
    if state.recent_sd > 4:
        explanations.append("Recent cycle lengths are unusually variable; the interval is wide.")
    if state.n < 3:
        explanations.append("Only a few cycles recorded; the population prior dominates.")
    if state.trend != 0:
        explanations.append("A recent trend was applied to the expected length.")

    return {
        "params_version": params.version,
        "mean_length": state.mean,
        "sd_length": state.cycle_sd * params.body_scale,
        "cycles_used": state.n,
        "confidence": state.confidence,
        "abstain": state.abstain,
        "tail_prob": state.tail_prob,
        "trend": state.trend,
        "quantiles": quantiles,
        "hazard": hazard,
        "ovulation": {
            k: v for k, v in ovulation.items() if k != "probabilities"
        },
        "prompt": _prompt(state, hazard),
        "explanations": explanations,
    }


def _prompt(state: PersonalState, hazard: dict) -> str:
    if state.abstain:
        return "Not enough signal for a reliable forecast."
    p3 = hazard["p_within_3"]
    if p3 < 0.05:
        return "Your period is unlikely in the next three days."
    if p3 < 0.25:
        return "There is a small chance your period starts in the next three days."
    if p3 < 0.5:
        return "Your period is fairly likely within three days."
    if p3 < 0.8:
        return "Your period is likely within three days."
    return "Your period is very likely within three days."


# ---------------------------------------------------------------------------
# Data helpers (Creighton public dataset; research validation only)
# ---------------------------------------------------------------------------


def load_creighton(path: str | Path) -> dict[int, list[list[int]]]:
    """Load Creighton cycle lengths by woman as consecutive complete segments.
    The released file has one row per cycle with a blank cycle_length for
    conception cycles; sequences are split at censoring so no forecast spans it.
    Returns {woman_id: [segment, ...]} where each segment is a list of lengths.
    """
    import csv

    rows_by_user: dict[int, list[dict]] = {}
    with open(path, encoding="utf-8-sig", newline="") as fh:
        for row in csv.DictReader(fh):
            try:
                user = int(row["new_id"])
                cycle_no = int(row["cycle_number"])
            except (KeyError, TypeError, ValueError):
                continue
            rows_by_user.setdefault(user, []).append(
                {
                    "cycle_number": cycle_no,
                    "cycle_length": row.get("cycle_length", "").strip(),
                    "conception": row.get("conception_cycle", "").strip().lower() == "yes",
                }
            )

    sequences: dict[int, list[list[int]]] = {}
    for user, rows in rows_by_user.items():
        rows.sort(key=lambda r: r["cycle_number"])
        segments: list[list[int]] = []
        current: list[int] = []
        for row in rows:
            raw = row["cycle_length"]
            if raw == "" or row["conception"]:
                if current:
                    segments.append(current)
                current = []
                continue
            try:
                y = int(float(raw))
            except ValueError:
                if current:
                    segments.append(current)
                current = []
                continue
            if 15 <= y <= 120:
                current.append(y)
            else:
                if current:
                    segments.append(current)
                current = []
        if current:
            segments.append(current)
        if segments:
            sequences[user] = segments
    return sequences


def load_json(path: str | Path) -> dict:
    return json.loads(Path(path).read_text())