"""Pure ForkScore computation over (parent, child) BatteryResult bundles.

No validator state, no consensus protocol. The output is a single
scalar plus a per-term breakdown so downstream tooling (the v8 paper's
acceptance protocol) can reproduce the verdict deterministically.

Inf-aware aggregation (per the 2026-05-08 fbd99a8 review on
``adaptation_efficiency`` zero-cost guards):

- ``infinite_count > 0`` in child but not parent → bonus capped at
  :data:`INFINITE_BONUS_CAP` (don't let +inf pollute the ranking).
- ``neg_infinite_count > 0`` in child → automatic REJECT regardless
  of other terms (free regression is a hard fail).

Verdict thresholds:

- ``ACCEPT`` when ``score >= SIGNAL_FLOOR`` (5pp default per
  ``docs/bench-maxing.md``).
- ``REJECT`` when ``score < 0`` OR free-regression detected OR any
  sub-battery rate dropped by ≥ ``HARD_REGRESSION_FLOOR``.
- ``MARGINAL`` for the band ``[0, SIGNAL_FLOOR)``.
"""

from __future__ import annotations

import json
import math
import sqlite3
import time
from dataclasses import asdict, dataclass, field
from typing import Iterable, Optional

from arborist.substrate.weights import DEFAULT_WEIGHTS, WeightSet


#: Phase 1c — version pin for rows persisted to ``fork_score_branches``.
#: Bump when a code change would produce a different :class:`ScoredFork`
#: from the same inputs (algorithm change, weight semantics shift,
#: hard-regression policy change). Pinning lets a reader filter out
#: rows produced under a prior estimator without re-scoring.
ESTIMATOR_VERSION = "fork-score-v1"


SIGNAL_FLOOR = 0.05
"""5-pp signal floor (matches docs/bench-maxing.md). Below this,
score is MARGINAL — not strong enough to commit a fork."""


HARD_REGRESSION_FLOOR = 0.05
"""Drop of ≥5pp on any sub-battery rate triggers an automatic REJECT
even if other axes improved."""


INFINITE_BONUS_CAP = 1.0
"""Per-fork cap on contribution from ``infinite_count`` increases.
Prevents a single +inf efficiency bucket from swamping the score."""


# Sub-battery metric keys per battery — the keys the bench runner
# emits in BatteryResult.metrics. Phase 1a maps each to the (single)
# rate-style metric that drives the Δ-rate computation.
_BATTERY_RATE_KEYS = {
    "5s": {
        "syntax": "parse_pass_rate",
        "semantics": "equivalence_recovery_rate",
        "syllogism": "step_validity_rate",
        "synthesis": "derivation_pass_rate",
        "semiotics": "invariance_under_swap",
    },
    "5t": {
        "transfer": "transfer_pass_rate",
        "transfer-learning": "transfer_learning_success_rate",
        "triangulation": "triangulation_agreement_rate",
        "truthtables": "truth_table_coverage_rate",
        "transitivity": "full_chain_pass_rate",
        "time": "temporal_context_preservation_rate",
    },
    "5f": {
        "function": "function_pass_rate",
        "finetuning": "adaptation_improvement_rate",
        "falsification": "error_detection_rate",
        "formulate": "structural_match_rate",
        "feedback-loop": "integration_coverage_rate",
    },
}


_EFFICIENCY_KEYS = {
    "finetuning": (
        "adaptation_efficiency_mean_finite",
        "adaptation_efficiency_infinite_count",
        "adaptation_efficiency_neg_infinite_count",
    ),
    "feedback-loop": (
        "feedback_efficiency_mean_finite",
        "feedback_efficiency_infinite_count",
        # No neg_inf for feedback-loop in Phase 1a (cost = chain length, > 0).
    ),
}


@dataclass
class ScoredFork:
    """Output of :func:`fork_score`. All fields JSON-serializable."""

    score: float
    verdict: str  # "ACCEPT" | "REJECT" | "MARGINAL"
    breakdown: dict[str, float] = field(default_factory=dict)
    flags: list[str] = field(default_factory=list)
    weights: dict[str, float] = field(default_factory=dict)

    def to_dict(self) -> dict:
        return asdict(self)


def bench_result_to_metrics(payload: dict) -> dict[str, dict[str, dict]]:
    """Adapter: ``bench.batteries.runner`` JSON → nested metrics dict.

    Output shape::

        {
          "5s": {"syntax": {<metrics>}, "semantics": {<metrics>}, ...},
          "5t": {"transfer-learning": {<metrics>}, ...},
          "5f": {"function": {<metrics>}, ...},
        }

    The keys exactly mirror ``BatteryResult.metrics`` so the scorer
    can deal with either a directly-passed metrics dict or a
    ``runner --all`` payload.
    """
    out: dict[str, dict[str, dict]] = {"5s": {}, "5t": {}, "5f": {}}
    for r in payload.get("results", []):
        battery = r.get("battery")
        sub = r.get("sub_battery")
        if battery in out and sub:
            out[battery][sub] = dict(r.get("metrics", {}))
    return out


def _aggregate(deltas: list[float], how: str) -> float:
    """Collapse a per-sub-battery Δ-rate vector into one term (ticket
    #000047). ``how`` ∈ :data:`arborist.substrate.weights.DELTA_AGGREGATORS`:

    - ``"mean"`` — arithmetic mean (default). Broad, balanced,
      noise-robust; matches ``docs/bench-maxing.md``'s per-rate 5-pp
      floor (a single sub-battery's gain is diluted 1/n).
    - ``"max"`` — the largest single-sub gain, floored at 0
      (``max(0, max_i Δ_i)``). Rewards specialization; weighs a
      single-sub improvement at face value. Silent on the rest of the
      vector — but per-sub regressions are still caught by the
      independent ``HARD_REGRESSION_FLOOR`` flag, so the verdict
      doesn't soften.
    - ``"sum"`` — total improvement (``Σ Δ_i``). Rewards breadth *and*
      magnitude; can over-reward many sub-noise-floor gains.

    Empty vector → 0.0 under every aggregator.
    """
    if not deltas:
        return 0.0
    if how == "mean":
        return sum(deltas) / len(deltas)
    if how == "max":
        return max(0.0, max(deltas))
    if how == "sum":
        return sum(deltas)
    raise ValueError(f"unknown delta aggregator {how!r}")


def _delta_5s(
    parent: dict[str, dict], child: dict[str, dict], aggregator: str = "mean"
) -> tuple[float, list[str]]:
    """Aggregated Δ-rate across 5S sub-batteries. Returns (delta, regression_subs)."""
    deltas: list[float] = []
    regressed: list[str] = []
    for sub, key in _BATTERY_RATE_KEYS["5s"].items():
        p = float(parent.get(sub, {}).get(key, 0.0))
        c = float(child.get(sub, {}).get(key, 0.0))
        d = c - p
        deltas.append(d)
        if d <= -HARD_REGRESSION_FLOOR:
            regressed.append(f"5s/{sub}: -{abs(d):.3f}")
    return _aggregate(deltas, aggregator), regressed


def _delta_5t(
    parent: dict[str, dict], child: dict[str, dict], aggregator: str = "mean"
) -> tuple[float, list[str]]:
    """Aggregated Δ-rate across 5T sub-batteries (using canonical
    Dav1DPrometheus names — transfer-learning, not legacy transfer)."""
    deltas: list[float] = []
    regressed: list[str] = []
    canonical = {
        k: v for k, v in _BATTERY_RATE_KEYS["5t"].items() if k != "transfer"
    }
    for sub, key in canonical.items():
        p = float(parent.get(sub, {}).get(key, 0.0))
        c = float(child.get(sub, {}).get(key, 0.0))
        d = c - p
        deltas.append(d)
        if d <= -HARD_REGRESSION_FLOOR:
            regressed.append(f"5t/{sub}: -{abs(d):.3f}")
    return _aggregate(deltas, aggregator), regressed


def _delta_5f(
    parent: dict[str, dict], child: dict[str, dict], aggregator: str = "mean"
) -> tuple[float, list[str], list[str]]:
    """Aggregated Δ-rate across 5F sub-batteries + efficiency-aware bonus.

    Returns (delta, regression_subs, efficiency_flags). The
    efficiency_flags include ``"NEG_INF_REGRESSION"`` when the child
    has ``adaptation_efficiency_neg_infinite_count > 0`` — a hard
    reject signal regardless of other terms. The efficiency bonus is
    added *after* the aggregated base, so it's aggregator-independent.
    """
    deltas: list[float] = []
    regressed: list[str] = []
    flags: list[str] = []

    for sub, key in _BATTERY_RATE_KEYS["5f"].items():
        p = float(parent.get(sub, {}).get(key, 0.0))
        c = float(child.get(sub, {}).get(key, 0.0))
        d = c - p
        deltas.append(d)
        if d <= -HARD_REGRESSION_FLOOR:
            regressed.append(f"5f/{sub}: -{abs(d):.3f}")

    # Efficiency-aware bonus on finetuning + feedback-loop.
    bonus = 0.0
    for sub, keys in _EFFICIENCY_KEYS.items():
        p_metrics = parent.get(sub, {})
        c_metrics = child.get(sub, {})
        # Mean-finite delta — direct ratio improvement.
        p_finite = float(p_metrics.get(keys[0], 0.0))
        c_finite = float(c_metrics.get(keys[0], 0.0))
        if not (math.isnan(p_finite) or math.isnan(c_finite)):
            bonus += (c_finite - p_finite) * 0.1  # damped — efficiencies are unbounded
        # Infinite-bonus increase: count delta * 0.1, capped at INFINITE_BONUS_CAP.
        p_inf = float(p_metrics.get(keys[1], 0.0))
        c_inf = float(c_metrics.get(keys[1], 0.0))
        inf_delta = c_inf - p_inf
        if inf_delta > 0:
            bonus += min(inf_delta * 0.1, INFINITE_BONUS_CAP)
        # Negative-inf regression — hard reject signal.
        if len(keys) > 2:
            p_neg = float(p_metrics.get(keys[2], 0.0))
            c_neg = float(c_metrics.get(keys[2], 0.0))
            if c_neg > p_neg:
                flags.append(
                    f"NEG_INF_REGRESSION:{sub}: free regression count "
                    f"increased {p_neg}→{c_neg}"
                )

    base = _aggregate(deltas, aggregator)
    return base + bonus, regressed, flags


def fork_score(
    parent: dict[str, dict[str, dict]],
    child: dict[str, dict[str, dict]],
    *,
    weights: WeightSet = DEFAULT_WEIGHTS,
    capital_delta: float = 0.0,
    selfmodel_calibration_gain: float = 0.0,
    audit_completeness: float = 0.0,
    validator_diversity: float = 0.0,
    security_risk: float = 0.0,
    complexity_delta: float = 0.0,
    memory_invalidation_count: float = 0.0,
) -> ScoredFork:
    """Compute the v8 fork score.

    Inputs are nested dicts ``{battery: {sub_battery: metrics_dict}}``
    matching :func:`bench_result_to_metrics` output. All non-bench
    inputs default to zero so single-validator Phase 1a deployments
    can score a (parent, child) pair from bench output alone.

    Returns a :class:`ScoredFork` with score + verdict + per-term
    breakdown + flags. The per-sub-battery Δ-rate vectors are
    collapsed by ``weights.delta_aggregator`` (#000047 — ``"mean"`` by
    default; see :func:`_aggregate`); the chosen aggregator is recorded
    in ``ScoredFork.weights["delta_aggregator"]``.
    """
    agg = weights.delta_aggregator
    delta_5s, regressions_5s = _delta_5s(parent.get("5s", {}), child.get("5s", {}), agg)
    delta_5t, regressions_5t = _delta_5t(parent.get("5t", {}), child.get("5t", {}), agg)
    delta_5f, regressions_5f, flags_5f = _delta_5f(
        parent.get("5f", {}), child.get("5f", {}), agg
    )

    regression_penalty = sum(
        max(0.0, -d) for d in (delta_5s, delta_5t, delta_5f)
    )

    # Per-term contributions. Sign baked in here.
    breakdown = {
        "alpha_x_delta_5s": weights.alpha * delta_5s,
        "beta_x_delta_5t": weights.beta * delta_5t,
        "gamma_x_delta_5f": weights.gamma * delta_5f,
        "delta_x_selfmodel_calibration": weights.delta * selfmodel_calibration_gain,
        "epsilon_x_audit_completeness": weights.epsilon * audit_completeness,
        "zeta_x_validator_diversity": weights.zeta * validator_diversity,
        "minus_eta_x_regression_penalty": -weights.eta * regression_penalty,
        "minus_theta_x_capital_cost_penalty": -weights.theta * capital_delta,
        "minus_iota_x_security_risk": -weights.iota * security_risk,
        "minus_kappa_x_complexity_penalty": -weights.kappa * complexity_delta,
        "minus_lambda_x_memory_invalidation": (
            -weights.lambda_ * memory_invalidation_count
        ),
    }
    score = sum(breakdown.values())

    # Hard-reject flags surface even when the score is positive.
    flags = list(flags_5f)
    flags.extend(f"REGRESSION_5S:{r}" for r in regressions_5s)
    flags.extend(f"REGRESSION_5T:{r}" for r in regressions_5t)
    flags.extend(f"REGRESSION_5F:{r}" for r in regressions_5f)

    # Verdict.
    has_neg_inf = any(f.startswith("NEG_INF_REGRESSION:") for f in flags)
    has_hard_regression = any(
        f.startswith("REGRESSION_") for f in flags
    )
    if has_neg_inf or has_hard_regression:
        verdict = "REJECT"
    elif score < 0:
        verdict = "REJECT"
    elif score >= SIGNAL_FLOOR:
        verdict = "ACCEPT"
    else:
        verdict = "MARGINAL"

    return ScoredFork(
        score=score,
        verdict=verdict,
        breakdown=breakdown,
        flags=flags,
        weights=weights.as_dict(),
    )


# ---------------------------------------------------------------------
# Phase 1c — branch-set persistence (#000012 §7 Phase 1c)
# ---------------------------------------------------------------------


def persist_branch_score(
    conn: sqlite3.Connection,
    *,
    branch_set_id: str,
    branch_id: str,
    parent_root: str,
    child_root: Optional[str],
    scored: ScoredFork,
    weights_id: str = "default",
    estimator_version: str = ESTIMATOR_VERSION,
    ts: Optional[int] = None,
) -> None:
    """Write one :class:`ScoredFork` to ``fork_score_branches``.

    Upserts on ``(branch_set_id, branch_id)`` so re-scoring the same
    fork under the same checkpoint is a clean overwrite, not a
    duplicate row. Sibling table — never enters
    ``audit_events.event_hash`` preimage. Caller wraps in
    :func:`arborist.store.transaction` when batching.
    """
    if ts is None:
        ts = int(time.time())
    breakdown_blob = json.dumps(
        scored.breakdown, sort_keys=True, separators=(",", ":")
    )
    conn.execute(
        "INSERT INTO fork_score_branches ("
        " branch_set_id, branch_id, parent_root, child_root,"
        " score, verdict, breakdown_blob, weights_id,"
        " estimator_version, recorded_at"
        ") VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"
        " ON CONFLICT(branch_set_id, branch_id) DO UPDATE SET"
        "   parent_root = excluded.parent_root,"
        "   child_root = excluded.child_root,"
        "   score = excluded.score,"
        "   verdict = excluded.verdict,"
        "   breakdown_blob = excluded.breakdown_blob,"
        "   weights_id = excluded.weights_id,"
        "   estimator_version = excluded.estimator_version,"
        "   recorded_at = excluded.recorded_at",
        (
            branch_set_id,
            branch_id,
            parent_root,
            child_root,
            float(scored.score),
            scored.verdict,
            breakdown_blob,
            weights_id,
            estimator_version,
            ts,
        ),
    )


def branch_set_density(
    conn: sqlite3.Connection, branch_set_id: str
) -> int:
    """Count distinct branches recorded under a checkpoint id.

    Used by the #000037 §12 Trigger 1 probe to satisfy the
    "ForkScore regularly receives ≥ 4 candidate branches per
    checkpoint" gate.
    """
    row = conn.execute(
        "SELECT COUNT(*) FROM fork_score_branches"
        " WHERE branch_set_id = ?",
        (branch_set_id,),
    ).fetchone()
    return int(row[0]) if row else 0
