Jyotisha/scripts/rectification/candidate_contrast.py

"""Feature-signature clustering and discriminator contracts for rectification.

Public candidate sets are signature clusters over the full birth window,
not the three highest-scoring adjacent minutes. Distinguish probes must
carry a real split, positive information gain, and two expected outcomes.
"""

from __future__ import annotations

import hashlib
import json
from collections.abc import Sequence
from typing import Any

from scripts.rectification.contracts import AUXILIARY_EVENT_KINDS, BACKGROUND_EVENT_KINDS

PROBE_PHASE_EVIDENCE_COLLECTION = "evidence_collection"
PROBE_PHASE_EVENT_CLARIFICATION = "event_clarification"
PROBE_PHASE_CANDIDATE_DISCRIMINATOR = "candidate_discriminator"
PROBE_PHASE_HOLDOUT_VALIDATION = "holdout_validation"

MIN_DISCRIMINATOR_EVENTS = 3
MIN_DISCRIMINATOR_DOMAINS = 2
MAX_PUBLIC_CLUSTERS = 12
SIGNATURE_LAYERS = ("d1", "d9", "d10", "d24", "d4", "d12", "md")
NAKSHATRA_SPAN = 360.0 / 27.0
LAYER_VARGA = {
    "d9": "D9",
    "d10": "D10",
    "d24": "D24",
    "d4": "D4",
    "d12": "D12",
}


def _clock(value: str) -> int:
    return int(value[:2]) * 60 + int(value[3:5])


def _hhmm(value: object) -> str | None:
    text = str(value or "")[:5]
    return text if len(text) == 5 and text[2] == ":" else None


def _birth_year(value: object) -> int | None:
    text = str(value or "").strip()
    if len(text) < 4 or not text[:4].isdigit():
        return None
    year = int(text[:4])
    return year if 1900 <= year <= 2100 else None


def event_year(event: dict[str, Any]) -> int | None:
    for key in ("date", "date_start", "occurred_from"):
        year = _birth_year(event.get(key))
        if year is not None:
            return year
    return None


def is_primary_scoreable_dict(event: dict[str, Any]) -> bool:
    kind = str(event.get("event_kind") or event.get("kind") or "")
    domain = str(event.get("domain") or "")
    if not domain:
        return False
    if kind in BACKGROUND_EVENT_KINDS or kind in AUXILIARY_EVENT_KINDS:
        return False
    return event_year(event) is not None


def scoreable_event_stats(events: Sequence[dict[str, Any]] | None) -> tuple[int, int, frozenset[str]]:
    scoreable = [
        event for event in (events or [])
        if isinstance(event, dict) and is_primary_scoreable_dict(event)
    ]
    domains = frozenset(str(event["domain"]) for event in scoreable)
    return len(scoreable), len(domains), domains


def training_scoreable_stats(events: Sequence[dict[str, Any]] | None) -> tuple[int, int, frozenset[str]]:
    from scripts.rectification.case_holdout import holdout_event_ids

    holdout = holdout_event_ids(events)
    training = [
        event for event in (events or [])
        if isinstance(event, dict)
        and is_primary_scoreable_dict(event)
        and str(event.get("id") or "") not in holdout
    ]
    domains = frozenset(str(event["domain"]) for event in training)
    return len(training), len(domains), domains


def discriminator_gate_open(events: Sequence[dict[str, Any]] | None) -> bool:
    count, domain_count, _ = training_scoreable_stats(events)
    return count >= MIN_DISCRIMINATOR_EVENTS and domain_count >= MIN_DISCRIMINATOR_DOMAINS


def missing_collection_domains(
    events: Sequence[dict[str, Any]] | None,
    catalog_domains: Sequence[str],
    volunteer_only: frozenset[str],
) -> list[str]:
    _, _, present = scoreable_event_stats(events)
    missing: list[str] = []
    for domain in catalog_domains:
        if domain in volunteer_only:
            continue
        if domain not in present:
            missing.append(domain)
    return missing


def layer_value(context: dict[str, Any], layer: str) -> int | None:
    feature = context.get("feature") if isinstance(context.get("feature"), dict) else {}
    if layer == "md":
        raw = feature.get("moon_degree")
        if not isinstance(raw, (int, float)):
            planets = context.get("planet_longitudes") if isinstance(context.get("planet_longitudes"), dict) else {}
            raw = planets.get("Moon")
        if not isinstance(raw, (int, float)):
            return None
        return int((float(raw) % 360.0) / NAKSHATRA_SPAN) % 27
    if layer == "d1":
        raw = feature.get("ascendant_sign_index")
        if isinstance(raw, int):
            return raw
        index = context.get("ascendant_index")
        return index if isinstance(index, int) else None
    name = LAYER_VARGA.get(layer)
    if not name:
        return None
    vargas = feature.get("varga_ascendants") if isinstance(feature.get("varga_ascendants"), dict) else {}
    raw = vargas.get(name)
    if isinstance(raw, int):
        return raw
    charts = context.get("varga_charts") if isinstance(context.get("varga_charts"), dict) else {}
    chart = charts.get(name) if isinstance(charts.get(name), dict) else {}
    ascendant = chart.get("Ascendant") if isinstance(chart.get("Ascendant"), dict) else {}
    index = ascendant.get("sign_idx")
    return index if isinstance(index, int) else None


def feature_signature(context: dict[str, Any]) -> tuple[int | None, ...]:
    return tuple(layer_value(context, layer) for layer in SIGNATURE_LAYERS)


def signature_key(signature: tuple[int | None, ...]) -> str:
    return ",".join("x" if value is None else str(value) for value in signature)


def context_time(context: dict[str, Any]) -> str | None:
    feature = context.get("feature") if isinstance(context.get("feature"), dict) else {}
    raw = feature.get("time")
    parsed = _hhmm(raw)
    if parsed:
        return parsed
    at = context.get("candidate_at")
    if hasattr(at, "strftime"):
        return at.strftime("%H:%M")
    return _hhmm(context.get("time"))


def cluster_contexts_by_signature(
    contexts: Sequence[dict[str, Any]],
) -> list[dict[str, Any]]:
    buckets: dict[tuple[int | None, ...], list[dict[str, Any]]] = {}
    for context in contexts:
        if not isinstance(context, dict):
            continue
        time = context_time(context)
        if not time:
            continue
        buckets.setdefault(feature_signature(context), []).append(context)
    clusters: list[dict[str, Any]] = []
    for signature, members in buckets.items():
        ordered = sorted(members, key=lambda item: _clock(str(context_time(item))))
        times = [str(context_time(item)) for item in ordered]
        clusters.append({
            "signature": signature,
            "signature_key": signature_key(signature),
            "contexts": ordered,
            "times": times,
            "representative_time": times[len(times) // 2],
            "representative": ordered[len(ordered) // 2],
        })
    clusters.sort(key=lambda item: _clock(item["representative_time"]))
    return clusters


def candidate_set_version(groups: Sequence[Sequence[str]]) -> str:
    canonical = [sorted({str(time)[:5] for time in group if _hhmm(time)}) for group in groups]
    canonical = [group for group in canonical if group]
    canonical.sort(key=lambda group: (group[0], len(group)))
    raw = json.dumps(canonical, separators=(",", ":"))
    return hashlib.sha256(raw.encode("utf-8")).hexdigest()[:16]


def candidate_split_hash(
    *,
    candidate_set_version_value: str,
    domain: str,
    year: int,
    groups: Sequence[Sequence[str]],
) -> str:
    version = candidate_set_version_value or candidate_set_version(groups)
    grouped = "|".join(
        ",".join(sorted({str(time)[:5] for time in group if _hhmm(time)}))
        for group in groups
        if group
    )
    payload = f"{version}:{domain}:{year}:{grouped}"
    return hashlib.sha256(payload.encode("utf-8")).hexdigest()[:24]


def expand_times_through_clusters(
    times: Sequence[str],
    clusters: Sequence[dict[str, Any]],
) -> list[str]:
    wanted = {str(time)[:5] for time in times if _hhmm(time)}
    expanded: list[str] = []
    seen: set[str] = set()
    for cluster in clusters:
        if not wanted.intersection(cluster["times"]):
            continue
        for time in cluster["times"]:
            if time not in seen:
                seen.add(time)
                expanded.append(time)
    for time in sorted(wanted, key=_clock):
        if time not in seen:
            expanded.append(time)
    return expanded


def select_signature_representatives(
    rows: Sequence[dict[str, Any]],
    static_contexts: Sequence[dict[str, Any]] | None = None,
) -> list[dict[str, Any]]:
    """One public candidate per feature-signature cluster, never three adjacent minutes from one peak."""
    by_time = {_hhmm(row.get("time")): row for row in rows if _hhmm(row.get("time"))}
    if not by_time:
        return []
    contexts = [
        context for context in (static_contexts or [])
        if isinstance(context, dict) and context_time(context) in by_time
    ]
    if len(contexts) >= 2:
        clusters = cluster_contexts_by_signature(contexts)
    else:
        clusters = _adjacent_score_clusters(list(by_time.values()))
    representatives: list[dict[str, Any]] = []
    for cluster in clusters:
        members = [by_time[time] for time in cluster["times"] if time in by_time]
        if not members:
            continue
        representatives.append(max(members, key=lambda row: (float(row.get("score") or 0), str(row.get("time")))))
        if len(representatives) >= MAX_PUBLIC_CLUSTERS:
            break
    representatives.sort(key=lambda row: (-float(row.get("score") or 0), str(row.get("time"))))
    return representatives or list(rows)[:1]


def _adjacent_score_clusters(rows: Sequence[dict[str, Any]]) -> list[dict[str, Any]]:
    ordered = sorted(rows, key=lambda row: _clock(str(_hhmm(row.get("time")))))
    groups: list[list[dict[str, Any]]] = []
    for row in ordered:
        current = groups[-1] if groups else None
        previous = current[-1] if current else None
        time = str(_hhmm(row.get("time")))
        prev_time = str(_hhmm(previous.get("time"))) if previous else ""
        adjacent = bool(previous) and (_clock(time) - _clock(prev_time)) % (24 * 60) <= 2
        if not current or not adjacent:
            groups.append([row])
        else:
            current.append(row)
    clusters = []
    for members in groups:
        times = [str(_hhmm(item.get("time"))) for item in members]
        clusters.append({
            "signature": tuple(),
            "signature_key": "adjacent",
            "contexts": members,
            "times": times,
            "representative_time": times[len(times) // 2],
            "representative": members[len(members) // 2],
        })
    return clusters


def candidate_ids_from_outcomes(outcomes: Sequence[dict[str, Any]]) -> list[str]:
    ids: list[str] = []
    seen: set[str] = set()
    for row in outcomes:
        if not isinstance(row, dict):
            continue
        for key in ("supports", "conflicts", "supportsCandidateIds", "conflictsCandidateIds"):
            values = row.get(key) or []
            if not isinstance(values, list):
                continue
            for item in values:
                time = _hhmm(item)
                if time and time not in seen:
                    seen.add(time)
                    ids.append(time)
    return ids


def distinguish_contract_errors(probe: dict[str, Any]) -> list[str]:
    if not isinstance(probe, dict):
        return ["distinguish_not_an_object"]
    if probe.get("role") != "distinguish" and probe.get("phase") != PROBE_PHASE_CANDIDATE_DISCRIMINATOR:
        return []
    errors: list[str] = []
    outcomes = probe.get("expected_outcomes") or probe.get("expectedOutcomes") or []
    if not isinstance(outcomes, list) or len(outcomes) < 2:
        errors.append("distinguish_empty_expected_outcomes")
        outcome_rows: list[dict[str, Any]] = []
    else:
        outcome_rows = [row for row in outcomes if isinstance(row, dict)]
        if len(outcome_rows) < 2:
            errors.append("distinguish_empty_expected_outcomes")
    ids = probe.get("candidate_ids") or probe.get("candidateIds") or candidate_ids_from_outcomes(outcome_rows)
    if not isinstance(ids, list) or len([item for item in ids if _hhmm(item)]) < 2:
        errors.append("distinguish_empty_candidate_ids")
    try:
        gain = float(probe.get("information_gain") if "information_gain" in probe else probe.get("informationGain") or 0)
    except (TypeError, ValueError):
        gain = 0.0
    if gain <= 0:
        errors.append("distinguish_non_positive_information_gain")
    return errors


def assert_distinguish_contract(probes: Sequence[dict[str, Any]]) -> None:
    for probe in probes:
        errors = distinguish_contract_errors(probe)
        if errors:
            raise AssertionError(f"{errors}: {probe.get('semantic_key') or probe.get('source')}")


def opportunity_from_probe(probe: dict[str, Any]) -> dict[str, Any]:
    outcomes = [row for row in (probe.get("expected_outcomes") or []) if isinstance(row, dict)]
    groups = [
        list(row.get("supports") or [])
        for row in outcomes
        if row.get("answer_class") in {"yes", "no", "weak_yes"} and row.get("supports")
    ]
    year = probe.get("year")
    return {
        "domain": probe.get("domain"),
        "time_window": {
            "year": year,
            "year_label": probe.get("year_label") or (f"{year} 年前后" if year else None),
        },
        "candidate_groups": groups,
        "expected_outcomes": outcomes,
        "information_gain": float(probe.get("information_gain") or 0),
        "source_features": list(probe.get("source_features") or [{
            "technique": probe.get("source") or "event_probe",
            "layers": list(SIGNATURE_LAYERS),
        }]),
        "semantic_key": probe.get("semantic_key"),
        "candidate_split_hash": probe.get("candidate_split_hash"),
        "candidate_set_version": probe.get("candidate_set_version"),
        "event_family": probe.get("event_family"),
        "phase": PROBE_PHASE_CANDIDATE_DISCRIMINATOR,
    }