Jyotisha/scripts/research/precision_gate_lib.py

"""Pure helpers for the precision-adaptive probe-gate research sweep.

Does not change production defaults in event_probes.py or the scoring engine.
"""

from __future__ import annotations

from contextlib import contextmanager
from copy import deepcopy
from dataclasses import dataclass, field
from datetime import date, timedelta
from random import Random
from typing import Any, Iterator, Sequence

from scripts.active_rectification_event_engine import (
    AUXILIARY_DOMAINS,
    AUXILIARY_SCORE_FACTOR,
    DOMAIN_CONFIG,
    OBSERVATION_ONLY_LAYERS,
    _active_narayana,
    _active_vimshottari,
    _ashtakavarga_auxiliary,
    _controlled_transit_rules,
    _event_datetime,
    _house_lords,
    _planet_house,
    _relative_house,
    _shadbala_verified_components_auxiliary,
    _varga_chart,
    _varga_house,
)
from scripts.active_rectification_events import precision_weight
from scripts.rectification.event_probes import (
    REFRESH_MIN_BOUNDARY_DAYS,
    _boundary_windows,
)
import functional_benefics
import varga

GATES: dict[str, dict[str, int]] = {
    "G0": {"year": 45, "month": 45, "day": 45},
    "G1": {"year": 45, "month": 30, "day": 10},
    "G2": {"year": 45, "month": 30, "day": 7},
    "G3": {"year": 45, "month": 21, "day": 5},
    "G4": {"year": 60, "month": 30, "day": 3},
}

TREATMENTS = ("A", "B", "C")
JITTER_SPANS = (3, 7, 14)
VARGA_CAP = 1.0
UPSTREAM_VARGA_MINUTES: dict[str, float] = {
    "D1": 120.0,
    "D9": 13.3,
    "D10": 12.0,
    "D12": 10.0,
    "D4": 7.5,
    "D24": 5.0,
    "D30": 4.0,
    "D60": 2.0,
}
PRODUCTION_VARGA_PREFIXES = (
    "D2", "D3", "D4", "D5", "D7", "D9", "D10", "D11", "D12", "D24", "D30",
)
JITTER_SEED = 20260914


@dataclass
class VargaPolicy:
    name: str
    window_minutes: float
    changing: frozenset[str] = field(default_factory=frozenset)
    use_d60: bool = False
    cap: float = VARGA_CAP
    hits: dict[str, int] = field(default_factory=dict)
    points: dict[str, float] = field(default_factory=dict)
    points_by_time: dict[str, dict[str, float]] = field(default_factory=dict)

    @property
    def is_baseline(self) -> bool:
        return self.name in {"V0", "baseline", ""}


def window_minutes_for_radius(radius: int) -> float:
    return float(2 * int(radius))


def varga_minutes(prefix: str) -> float:
    if prefix in UPSTREAM_VARGA_MINUTES:
        return UPSTREAM_VARGA_MINUTES[prefix]
    number = int(str(prefix)[1:])
    return 120.0 / max(number, 1)


def varga_factor(prefix: str, window_minutes: float, cap: float = VARGA_CAP) -> float:
    minutes = varga_minutes(prefix)
    if minutes <= 0 or window_minutes <= 0:
        return 0.0
    return min(float(window_minutes) / minutes, float(cap))


def finest_precision(events: Sequence[dict[str, Any]]) -> str:
    ranks = {"day": 3, "month": 2, "quarter": 2, "year": 1, "range": 1, "unknown": 0}
    best = "year"
    best_rank = 0
    for event in events:
        precision = str(event.get("precision") or "year")
        rank = ranks.get(precision, 0)
        if rank > best_rank:
            best = precision if precision in {"day", "month", "year"} else "year"
            best_rank = rank
    if best_rank >= 3:
        return "day"
    if best_rank >= 2:
        return "month"
    return "year"


def participating_precision(events: Sequence[dict[str, Any]], domain: str | None = None) -> str:
    pool = [
        event for event in events
        if domain is None or str(event.get("domain") or "") == domain
    ]
    return finest_precision(pool or events)


def threshold_for(gate: str, precision: str, *, refresh: bool) -> int:
    spec = GATES[gate]
    key = precision if precision in spec else "year"
    if refresh and gate == "G0":
        return REFRESH_MIN_BOUNDARY_DAYS
    return int(spec[key])


def count_precision(events: Sequence[dict[str, Any]]) -> dict[str, int]:
    tallies = {"day": 0, "month": 0, "year": 0, "other": 0}
    for event in events:
        precision = str(event.get("precision") or "year")
        if precision in tallies:
            tallies[precision] += 1
        else:
            tallies["other"] += 1
    return tallies


def treat_events(events: Sequence[dict[str, Any]], treatment: str) -> list[dict[str, Any]]:
    if treatment == "A":
        return [dict(event) for event in events]
    rows: list[dict[str, Any]] = []
    for event in events:
        item = dict(event)
        precision = str(item.get("precision") or "year")
        if treatment == "B" and precision == "day":
            item["precision"] = "month"
        elif treatment == "C":
            item["precision"] = "year"
        rows.append(item)
    return rows


def jitter_day_events(
    events: Sequence[dict[str, Any]],
    span: int,
    *,
    case_id: str,
    seed: int = JITTER_SEED,
) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    for event in events:
        item = dict(event)
        if str(item.get("precision") or "") != "day":
            rows.append(item)
            continue
        raw = str(item.get("date") or "")
        try:
            original = date.fromisoformat(raw[:10])
        except ValueError:
            rows.append(item)
            continue
        rng = Random(f"{seed}:{case_id}:{item.get('id')}:{span}")
        offset = 0
        while offset == 0:
            offset = rng.randint(-int(span), int(span))
        item["date"] = (original + timedelta(days=offset)).isoformat()
        item["jitter_days"] = offset
        rows.append(item)
    return rows


def varga_sign_index(context: dict[str, Any], prefix: str) -> int | None:
    charts = context.get("varga_charts") if isinstance(context.get("varga_charts"), dict) else {}
    chart = charts.get(prefix)
    if not isinstance(chart, dict):
        return None
    raw = (chart.get("Ascendant") or {}).get("sign_idx")
    return int(raw) if isinstance(raw, int) else None


def changing_vargas(
    contexts: Sequence[dict[str, Any]],
    prefixes: Sequence[str] = PRODUCTION_VARGA_PREFIXES,
) -> frozenset[str]:
    changed: set[str] = set()
    for prefix in prefixes:
        previous = None
        for context in contexts:
            current = varga_sign_index(context, prefix)
            if previous is not None and current is not None and current != previous:
                changed.add(prefix)
                break
            previous = current
    return frozenset(changed)


def ensure_d60(context: dict[str, Any]) -> dict[str, Any] | None:
    charts = context.setdefault("varga_charts", {})
    existing = charts.get("D60")
    if isinstance(existing, dict):
        return existing
    planets = context.get("planet_longitudes") or {}
    natal = context.get("chart") or {}
    ascendant = natal.get("ascendant") if isinstance(natal.get("ascendant"), dict) else {}
    lon = ascendant.get("lon")
    if not isinstance(lon, (int, float)) or not planets:
        return None
    computed = varga.calc_all_vargas(planets, float(lon), divisions=[60])
    chart = _varga_chart(computed, "D60")
    if isinstance(chart, dict):
        charts["D60"] = chart
    return chart


def attach_d60(contexts: Sequence[dict[str, Any]]) -> None:
    for context in contexts:
        ensure_d60(context)


@contextmanager
def patched_boundary_gate(*, initial: int, refresh: int) -> Iterator[None]:
    import scripts.rectification.event_probes as ep
    previous = (ep.MIN_BOUNDARY_DAYS, ep.REFRESH_MIN_BOUNDARY_DAYS)
    ep.MIN_BOUNDARY_DAYS = int(initial)
    ep.REFRESH_MIN_BOUNDARY_DAYS = int(refresh)
    try:
        yield
    finally:
        ep.MIN_BOUNDARY_DAYS, ep.REFRESH_MIN_BOUNDARY_DAYS = previous


def production_defaults_intact() -> bool:
    import scripts.rectification.event_probes as ep
    return ep.MIN_BOUNDARY_DAYS == 45 and ep.REFRESH_MIN_BOUNDARY_DAYS == 30


def close_same_year_windows(left: date, right: date, min_days: int) -> list[date]:
    return list(_boundary_windows([left], [right], min_days=min_days))


def _record_varga(policy: VargaPolicy, time: str, prefix: str, delta: float) -> None:
    policy.hits[prefix] = policy.hits.get(prefix, 0) + 1
    policy.points[prefix] = round(policy.points.get(prefix, 0.0) + delta, 6)
    by_time = policy.points_by_time.setdefault(time, {})
    by_time[prefix] = round(by_time.get(prefix, 0.0) + delta, 6)


def score_event_with_policy(
    *,
    candidate_time: str,
    event: dict[str, Any],
    natal_chart: dict[str, Any],
    varga_by_prefix: dict[str, dict[str, Any]],
    vimshottari: tuple[str, str, str],
    narayana: tuple[int | None, int | None],
    arudha_padas: dict[str, Any],
    policy: VargaPolicy,
) -> dict[str, Any]:
    _, target_houses = DOMAIN_CONFIG[event["domain"]]
    ascendant_index = int(natal_chart["ascendant"]["lon"] // 30)
    target_lords = _house_lords(ascendant_index, target_houses)
    functional = functional_benefics.derive_functional_benefic_malefic(
        natal_chart["ascendant"].get("sign"),
        # 第五轮裁定 3 后默认是 PVR 表 30；校正打分在冻结身份里，明确选中原产品表，分数不变（BUG-1746）。
        functional_profile=functional_benefics.FUNCTIONAL_PROFILE_BPHS,
    )
    functional_benefics_set = set(functional.get("functional_benefics") or [])
    functional_malefics_set = set(functional.get("functional_malefics") or [])
    major_lord, minor_lord, pratyantar_lord = vimshottari
    rules: list[str] = []
    points = 0.0
    active = list(varga_by_prefix.items())
    count = max(len(active), 1)

    for lord, weight, label in (
        (major_lord, 2.0, "vim_md"),
        (minor_lord, 1.5, "vim_ad"),
        (pratyantar_lord, 0.75, "vim_pd"),
    ):
        if _planet_house(natal_chart, lord) in target_houses:
            rules.append(f"{label}_domain_house")
            points += weight
        if lord in target_lords:
            rules.append(f"{label}_domain_lord")
            points += weight
        if not active:
            pass
        else:
            for prefix, varga_chart in active:
                if _varga_house(varga_chart, lord) not in target_houses:
                    continue
                factor = 1.0 if policy.is_baseline else varga_factor(
                    prefix, policy.window_minutes, policy.cap,
                )
                delta = weight * factor / (2 * count)
                rules.append(f"{label}_domain_varga")
                points += delta
                _record_varga(policy, candidate_time, prefix, delta)
        if lord in functional_benefics_set:
            rules.append(f"{label}_functional_benefic_auxiliary")
            points += 0.2
        elif lord in functional_malefics_set:
            rules.append(f"{label}_functional_malefic_auxiliary")
            points -= 0.1

    for sign_index, weight, label in (
        (narayana[0], 2.0, "narayana_md"),
        (narayana[1], 1.0, "narayana_ad"),
    ):
        if sign_index is not None and _relative_house(sign_index, ascendant_index) in target_houses:
            rules.append(f"{label}_domain_house")
            points += weight
    arudha_keys = (
        ("A7", "UL") if event["domain"] == "relationship"
        else ("A10",) if event["domain"] in {"career", "occupation"}
        else ()
    )
    arudha_signs = {
        value.get("sign_idx") for key in arudha_keys
        if isinstance((value := arudha_padas.get(key)), dict) and isinstance(value.get("sign_idx"), int)
    }
    if arudha_signs:
        for lord, label in ((major_lord, "vim_md"), (minor_lord, "vim_ad"), (pratyantar_lord, "vim_pd")):
            planet = natal_chart.get("planets", {}).get(lord) or {}
            if isinstance(planet.get("lon"), (int, float)) and int(planet["lon"] // 30) in arudha_signs:
                rules.append(f"{label}_arudha_auxiliary")
                points += 0.35

    event_kind = event.get("event_kind", event["domain"])
    if event["domain"] in AUXILIARY_DOMAINS:
        points *= AUXILIARY_SCORE_FACTOR
        rules.append(
            "occupation_auxiliary_not_primary"
            if event["domain"] == "occupation"
            else "appearance_auxiliary_not_primary"
        )
    if not rules:
        rules.append("no_domain_activation")
    rules.append(f"event_kind:{event_kind}")
    weighted_points = round(points * precision_weight(event["precision"]), 4)
    return {
        "event_id": event["id"],
        "domain": event["domain"],
        "candidate_time": candidate_time,
        "rule_ids": rules,
        "points": weighted_points,
    }


def domain_varga_prefixes(domain: str, policy: VargaPolicy) -> tuple[str, ...]:
    prefixes = list(DOMAIN_CONFIG[domain][0])
    if policy.use_d60 and "D60" not in prefixes:
        prefixes.append("D60")
    if policy.name in {"V2", "V3"}:
        prefixes = [item for item in prefixes if item in policy.changing]
    return tuple(prefixes)


def candidate_row_with_policy(
    request: dict[str, Any],
    context: dict[str, Any],
    policy: VargaPolicy,
) -> dict[str, Any]:
    candidate_at = context["candidate_at"]
    chart = context["chart"]
    planet_longitudes = context["planet_longitudes"]
    ascendant_index = context["ascendant_index"]
    arudha_padas = context["arudha_padas"]
    varga_charts = context["varga_charts"]
    moon_longitude = planet_longitudes["Moon"]
    evidence: list[dict[str, Any]] = []
    missing_layers: list[str] = []
    stamp = candidate_at.strftime("%H:%M")

    for event in request["events"]:
        event_at = _event_datetime(event)
        prefixes = domain_varga_prefixes(event["domain"], policy)
        selected = {prefix: varga_charts.get(prefix) for prefix in prefixes}
        if prefixes and any(chart is None for chart in selected.values()):
            missing_layers.extend(prefixes)
            continue
        try:
            vimshottari = _active_vimshottari(candidate_at, moon_longitude, event_at)
        except (KeyError, TypeError, ValueError):
            missing_layers.append("Vimshottari_MD_AD_PD")
            continue
        try:
            narayana = _active_narayana(
                ascendant_index, planet_longitudes, candidate_at, event_at,
            )
        except (KeyError, TypeError, ValueError):
            missing_layers.append("Narayana_MD_AD")
            continue
        if narayana[0] is None or narayana[1] is None:
            missing_layers.append("Narayana_MD_AD")
            continue
        usable = {prefix: chart for prefix, chart in selected.items() if isinstance(chart, dict)}
        evidence.append(score_event_with_policy(
            candidate_time=stamp,
            event=event,
            natal_chart=chart,
            varga_by_prefix=usable,
            vimshottari=vimshottari,
            narayana=narayana,
            arudha_padas=arudha_padas,
            policy=policy,
        ))
        transit_rules = _controlled_transit_rules(
            request, event, ascendant_index, DOMAIN_CONFIG[event["domain"]][1],
        )
        if transit_rules:
            evidence[-1]["rule_ids"].extend(transit_rules)
            evidence[-1]["points"] = round(
                evidence[-1]["points"] + 0.25 * len(transit_rules) * precision_weight(event["precision"]),
                4,
            )
        av_rules, av_points = _ashtakavarga_auxiliary(
            chart, ascendant_index, DOMAIN_CONFIG[event["domain"]][1],
        )
        if av_rules:
            evidence[-1]["rule_ids"].extend(av_rules)
            evidence[-1]["points"] = round(
                evidence[-1]["points"] + av_points * precision_weight(event["precision"]), 4,
            )
        shadbala_rules, shadbala_points = _shadbala_verified_components_auxiliary(
            chart, candidate_at.hour + candidate_at.minute / 60, vimshottari,
        )
        if shadbala_rules:
            evidence[-1]["rule_ids"].extend(shadbala_rules)
            evidence[-1]["points"] = round(
                evidence[-1]["points"] + shadbala_points * precision_weight(event["precision"]), 4,
            )

    return {
        "time": stamp,
        "score": round(sum(item["points"] for item in evidence), 4),
        "evidence": evidence,
        "missing_layers": sorted(set(
            missing_layers
            + [layer for layer in context["feature"]["blocked_layers"] if layer not in OBSERVATION_ONLY_LAYERS]
        )),
    }


def make_row_provider(static_contexts: Sequence[dict[str, Any]], policy: VargaPolicy):
    def provider(request: dict[str, Any]) -> list[dict[str, Any]]:
        return [candidate_row_with_policy(request, context, policy) for context in static_contexts]
    return provider


def ablate_top1(
    scores: dict[str, float],
    varga_points_by_time: dict[str, dict[str, float]],
    true_time: str,
) -> dict[str, Any]:
    if not scores:
        return {"baseline_top1": None, "flips": {}}
    baseline = max(scores, key=lambda time: (scores[time], time))
    flips: dict[str, int] = {}
    prefixes = sorted({prefix for row in varga_points_by_time.values() for prefix in row})
    for prefix in prefixes:
        adjusted = {
            time: scores[time] - float((varga_points_by_time.get(time) or {}).get(prefix) or 0.0)
            for time in scores
        }
        leader = max(adjusted, key=lambda time: (adjusted[time], time))
        flips[prefix] = int(leader != baseline)
    return {
        "baseline_top1": baseline,
        "true_is_top1": baseline == true_time[:5],
        "flips": flips,
    }


def coverage_ok(candidate: dict[str, Any], baseline: dict[str, Any]) -> bool:
    if candidate.get("coverage") is None or baseline.get("coverage") is None:
        return False
    return float(candidate["coverage"]) + 1e-9 >= float(baseline["coverage"])


def squeezed_ok(candidate: dict[str, Any], baseline: dict[str, Any]) -> bool:
    return int(candidate.get("squeezed") or 0) <= int(baseline.get("squeezed") or 0)


def gate_verdict(
    baseline: dict[str, Any],
    candidate: dict[str, Any],
    *,
    jitter7_squeezed: int | None = None,
) -> str:
    if not baseline.get("n") or not candidate.get("n"):
        return "uncertain"
    if jitter7_squeezed is not None and int(jitter7_squeezed) > 0:
        return "no_benefit"
    if not coverage_ok(candidate, baseline) or not squeezed_ok(candidate, baseline):
        return "no_benefit"
    hit_same_or_up = float(candidate["top1"]) + 1e-9 >= float(baseline["top1"])
    width_down = (
        candidate.get("width_median") is not None
        and baseline.get("width_median") is not None
        and float(candidate["width_median"]) < float(baseline["width_median"]) - 1e-9
    )
    extra = float(candidate.get("refresh_mean") or 0) - float(baseline.get("refresh_mean") or 0)
    width_gain = 0.0
    if candidate.get("width_median") is not None and baseline.get("width_median") is not None:
        width_gain = float(baseline["width_median"]) - float(candidate["width_median"])
    if not hit_same_or_up:
        return "no_benefit"
    if width_down:
        if extra >= 8 and width_gain <= 2:
            return "no_benefit"
        return "benefit"
    return "no_benefit"


def clone_events(events: Sequence[dict[str, Any]]) -> list[dict[str, Any]]:
    return [deepcopy(dict(event)) for event in events]