#!/usr/bin/env python3
"""Pre-registered D12 parent-probe scoring. Research only.

The reading rules live in
``docs/research/d12_parent_probes_preregistration_2026_10_12.json``.
This module does not compute an ephemeris. Callers pass longitudes from
``domain_calculation_service.compute_chart`` and signs from ``varga.calc_varga``.
"""

from __future__ import annotations

import hashlib
import json
import sys
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from decimal import Decimal
from pathlib import Path
from random import Random
from typing import Any

_SCRIPTS = str(Path(__file__).resolve().parents[1])
if _SCRIPTS not in sys.path:
    sys.path.insert(0, _SCRIPTS)

import varga

ROOT = Path(__file__).resolve().parents[2]
PREREG_PATH = ROOT / "docs" / "research" / "d12_parent_probes_preregistration_2026_10_12.json"
HOLDOUT_PATH = ROOT / "references" / "real_case_calibration" / "minute_rectification_holdout_v5.json"
FACTS_PATH = ROOT / "references" / "real_case_calibration" / "parent_facts_v1.json"
RESULTS_PATH = ROOT / "docs" / "research" / "d12_parent_probes_results_2026_10_12.json"

PROBES: tuple[str, ...] = (
    "mother_more_daily_care",
    "father_financial_lead",
    "father_guided_study",
    "father_work_institutional",
    "father_work_commerce",
    "parents_long_apart",
)
KENDRA = frozenset({1, 4, 7, 10})
SEPARATED_DELTAS = frozenset({1, 5, 7, 11})
ANSWERS = frozenset({"yes", "no", "unclear"})
PREDICTIONS = frozenset({"yes", "no", "abstain"})
GRADES = frozenset({"biography", "interview", "encyclopedia"})
ALPHA = Decimal("0.0125")
ROBUST_MAX = Decimal("0.2")
CONTROL_P_MAX = Decimal("0.01")
MIN_DECISIVE = 4
MIN_INCLUDED = 60
MIN_DIFFERING = 20
PERMUTATIONS = 999
PERM_SEED = 20261012
CONTROL_SEED = 20261013


@dataclass(frozen=True)
class Placement:
    d12_lagna: int
    d1_lagna: int
    d12_signs: Mapping[str, int]
    d1_signs: Mapping[str, int]


def canonical_bytes(payload: Any) -> bytes:
    return (
        json.dumps(payload, ensure_ascii=False, sort_keys=True, indent=2) + "\n"
    ).encode("utf-8")


def case_id_digest(case_ids: Sequence[str]) -> str:
    text = "\n".join(sorted(case_ids))
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


def load_prereg(path: Path = PREREG_PATH) -> dict[str, Any]:
    return json.loads(path.read_text(encoding="utf-8"))


def assert_prereg_locked(prereg: Mapping[str, Any]) -> None:
    if prereg.get("id") != "d12_parent_probes_preregistration_2026_10_12":
        raise SystemExit("preregistration id mismatch")
    if list(prereg.get("probes") or []) != list(PROBES):
        raise SystemExit("preregistration probe list mismatch")
    if int(prereg["permutation"]["n"]) != PERMUTATIONS:
        raise SystemExit("preregistration permutation count mismatch")
    if int(prereg["permutation"]["seed"]) != PERM_SEED:
        raise SystemExit("preregistration permutation seed mismatch")
    if int(prereg["positive_control"]["seed"]) != CONTROL_SEED:
        raise SystemExit("preregistration control seed mismatch")
    if int(prereg["positive_control"]["minimum_differing_cases"]) != MIN_DIFFERING:
        raise SystemExit("preregistration control minimum mismatch")
    if int(prereg["birth_records"]["expected_case_count"]) != 77:
        raise SystemExit("preregistration case count mismatch")
    if "at least 60" not in str(prereg["gates"]["sample"]):
        raise SystemExit("preregistration sample gate mismatch")
    if format(prereg["family"]["bonferroni_alpha"], ".4f") != "0.0125":
        raise SystemExit("preregistration alpha mismatch")
    if format(prereg["robustness"]["gate_max_change_rate"], ".1f") != "0.2":
        raise SystemExit("preregistration robustness line mismatch")
    if format(prereg["positive_control"]["permutation_p_max"], ".2f") != "0.01":
        raise SystemExit("preregistration control p mismatch")


def house_of(body_sign: int, lagna: int) -> int:
    return (int(body_sign) - int(lagna)) % 12 + 1


def lord_of(lagna: int, house: int) -> str:
    sign_index = (int(lagna) + int(house) - 1) % 12
    return varga.SIGN_LORDS[varga.SIGNS[sign_index]]


def career_house(father_house: int) -> int:
    return ((int(father_house) - 1 + 9) % 12) + 1


def separated(sign_a: int, sign_b: int) -> bool:
    return (int(sign_b) - int(sign_a)) % 12 in SEPARATED_DELTAS


def _unique(planets: Sequence[str]) -> tuple[str, ...]:
    return tuple(dict.fromkeys(planets))


def _in_kendra(signs: Mapping[str, int], lagna: int, planet: str) -> bool:
    return house_of(signs[planet], lagna) in KENDRA


def read_probes(
    signs: Mapping[str, int],
    lagna: int,
    father_house: int,
) -> dict[str, str]:
    """Return yes / no / abstain for the six pre-registered probes."""
    mother = _unique(("Moon", lord_of(lagna, 4)))
    father = _unique(("Sun", lord_of(lagna, father_house)))
    mother_kendra = sum(_in_kendra(signs, lagna, planet) for planet in mother)
    father_kendra = sum(_in_kendra(signs, lagna, planet) for planet in father)
    if mother_kendra > father_kendra:
        care = "yes"
    elif father_kendra > mother_kendra:
        care = "no"
    elif _in_kendra(signs, lagna, "Moon") and not _in_kendra(signs, lagna, "Sun"):
        care = "yes"
    elif _in_kendra(signs, lagna, "Sun") and not _in_kendra(signs, lagna, "Moon"):
        care = "no"
    else:
        care = "abstain"

    lord2 = lord_of(lagna, 2)
    financial = (
        any(house_of(signs[planet], lagna) in {2, 11} for planet in father)
        or house_of(signs[lord2], lagna) == father_house
        or any(signs[lord2] == signs[planet] for planet in father)
    )
    lord5 = lord_of(lagna, 5)
    study = (
        any(house_of(signs[planet], lagna) == 5 for planet in father)
        or house_of(signs[lord5], lagna) == father_house
        or house_of(signs["Jupiter"], lagna) == father_house
        or any(signs["Jupiter"] == signs[planet] for planet in father)
    )
    career = career_house(father_house)
    career_lord = lord_of(lagna, career)
    institutional = (
        house_of(signs["Saturn"], lagna) == career
        or house_of(signs["Sun"], lagna) == career
        or career_lord in {"Saturn", "Sun"}
    )
    commerce = (
        house_of(signs["Mercury"], lagna) == career
        or house_of(signs["Venus"], lagna) == career
        or career_lord in {"Mercury", "Venus"}
    )
    lord4 = lord_of(lagna, 4)
    father_lord = lord_of(lagna, father_house)
    lords_apart = lord4 != father_lord and separated(signs[lord4], signs[father_lord])
    rahu_house = house_of(signs["Rahu"], lagna) if "Rahu" in signs else None
    apart = (
        separated(signs["Sun"], signs["Moon"])
        or lords_apart
        or rahu_house in {4, father_house}
    )
    return {
        "mother_more_daily_care": care,
        "father_financial_lead": "yes" if financial else "no",
        "father_guided_study": "yes" if study else "no",
        "father_work_institutional": "yes" if institutional else "no",
        "father_work_commerce": "yes" if commerce else "no",
        "parents_long_apart": "yes" if apart else "no",
    }


def placement_from_longitudes(
    ascendant_longitude: float,
    planet_longitudes: Mapping[str, float],
    planet_d1_signs: Mapping[str, int],
) -> Placement:
    d12_signs = {
        name: int(varga.calc_varga(float(longitude), 12)["sign_idx"])
        for name, longitude in planet_longitudes.items()
    }
    d1_signs = {name: int(sign) for name, sign in planet_d1_signs.items()}
    ascendant = float(ascendant_longitude)
    return Placement(
        d12_lagna=int(varga.calc_varga(ascendant, 12)["sign_idx"]),
        d1_lagna=int(ascendant // 30) % 12,
        d12_signs=d12_signs,
        d1_signs=d1_signs,
    )


def predict(placement: Placement, chart: str, father_house: int) -> dict[str, str]:
    if chart == "D12":
        return read_probes(placement.d12_signs, placement.d12_lagna, father_house)
    if chart == "D1":
        return read_probes(placement.d1_signs, placement.d1_lagna, father_house)
    raise ValueError(f"unknown chart {chart}")


def cuts(sign_by_offset: Mapping[int, int]) -> list[tuple[int, int]]:
    offsets = sorted(int(item) for item in sign_by_offset)
    found: list[tuple[int, int]] = []
    for left, right in zip(offsets, offsets[1:]):
        if right == left + 1 and sign_by_offset[left] != sign_by_offset[right]:
            found.append((left, right))
    return found


def nearest_cut(pairs: Sequence[tuple[int, int]]) -> tuple[int, int] | None:
    if not pairs:
        return None
    return min(pairs, key=lambda pair: (min(abs(pair[0]), abs(pair[1])), abs(pair[0]), pair[0]))


def true_side(
    sign_by_offset: Mapping[int, int],
    pair: tuple[int, int],
) -> tuple[int, int]:
    truth = sign_by_offset[0]
    left, right = pair
    left_true = sign_by_offset[left] == truth
    right_true = sign_by_offset[right] == truth
    if left_true == right_true:
        raise RuntimeError("boundary does not separate the recorded sign")
    if left_true:
        return left, right
    return right, left


def segments(sign_by_offset: Mapping[int, int]) -> list[list[int]]:
    offsets = sorted(int(item) for item in sign_by_offset)
    if not offsets:
        return []
    runs: list[list[int]] = [[offsets[0]]]
    for offset in offsets[1:]:
        if offset == runs[-1][-1] + 1 and sign_by_offset[offset] == sign_by_offset[runs[-1][-1]]:
            runs[-1].append(offset)
        else:
            runs.append([offset])
    return runs


def segment_representative(run: Sequence[int]) -> int:
    return min(run, key=lambda offset: (abs(offset), offset))


def side_score(prediction: Mapping[str, str], answers: Mapping[str, str]) -> int:
    score = 0
    for probe in PROBES:
        answer = answers[probe]
        label = prediction[probe]
        if answer not in {"yes", "no"} or label not in {"yes", "no"}:
            continue
        score += 1 if label == answer else -1
    return score


def fisher_yates(items: Sequence[str], rng: Random) -> list[str]:
    out = list(items)
    for index in range(len(out) - 1, 0, -1):
        swap = rng.randrange(index + 1)
        out[index], out[swap] = out[swap], out[index]
    return out


def permutation_maps(case_ids: Sequence[str], seed: int, count: int = PERMUTATIONS) -> list[list[str]]:
    rng = Random(seed)
    ordered = sorted(case_ids)
    return [fisher_yates(ordered, rng) for _ in range(count)]


def assign_answers(
    destinations: Sequence[str],
    sources: Sequence[str],
    answers: Mapping[str, Mapping[str, str]],
) -> dict[str, Mapping[str, str]]:
    if len(destinations) != len(sources):
        raise RuntimeError("permutation length mismatch")
    return {dest: answers[source] for dest, source in zip(destinations, sources)}


def selection_count(
    case_ids: Sequence[str],
    answers: Mapping[str, Mapping[str, str]],
    true_predictions: Mapping[str, Mapping[str, str]],
    other_predictions: Mapping[str, Mapping[str, str]],
) -> int:
    selected = 0
    for case_id in case_ids:
        true_score = side_score(true_predictions[case_id], answers[case_id])
        other_score = side_score(other_predictions[case_id], answers[case_id])
        if true_score > other_score:
            selected += 1
    return selected


def ratio_below(numer: int, denom: int, threshold: Decimal) -> bool:
    if denom <= 0:
        return False
    return Decimal(numer) / Decimal(denom) < threshold


def display_ratio(numer: int, denom: int) -> str:
    if denom <= 0:
        return "0.000000"
    return format(numer / denom, ".6f")


def permutation_summary(observed: int, denom: int, perm_counts: Sequence[int]) -> dict[str, Any]:
    extreme = sum(1 for count in perm_counts if count >= observed)
    below = sum(1 for count in perm_counts if count < observed)
    null_sum = sum(perm_counts)
    n_perm = len(perm_counts)
    greater = observed * n_perm > null_sum if denom > 0 else False
    return {
        "observed_numer": observed,
        "observed_denom": denom,
        "observed_display": display_ratio(observed, denom),
        "extreme": extreme,
        "p_numer": extreme + 1,
        "p_denom": n_perm + 1,
        "p_display": display_ratio(extreme + 1, n_perm + 1),
        "null_sum": null_sum,
        "null_mean_display": display_ratio(null_sum, n_perm * denom) if denom else "0.000000",
        "above_null_mean": greater,
        "percentile_numer": below,
        "percentile_denom": n_perm,
        "percentile_display": display_ratio(below, n_perm),
    }


def cell_is_significant(summary: Mapping[str, Any]) -> bool:
    return bool(summary["above_null_mean"]) and ratio_below(
        int(summary["p_numer"]),
        int(summary["p_denom"]),
        ALPHA,
    )


def vectors_differ(
    true_prediction: Mapping[str, str],
    other_prediction: Mapping[str, str],
) -> bool:
    return any(
        true_prediction[probe] in {"yes", "no"} and true_prediction[probe] != other_prediction[probe]
        for probe in PROBES
    )


def plant_answers(true_prediction: Mapping[str, str]) -> dict[str, str]:
    planted: dict[str, str] = {}
    for probe in PROBES:
        label = true_prediction[probe]
        planted[probe] = label if label in {"yes", "no"} else "unclear"
    return planted


def change_count(
    baseline: Mapping[str, Mapping[int, Mapping[str, str]]],
    alternate: Mapping[str, Mapping[int, Mapping[str, str]]],
    sides: Mapping[str, tuple[int, int]],
) -> int:
    changes = 0
    for case_id, pair in sides.items():
        for offset in pair:
            for probe in PROBES:
                if baseline[case_id][offset][probe] != alternate[case_id][offset][probe]:
                    changes += 1
    return changes


def decisive_count(answers: Mapping[str, str]) -> int:
    return sum(answer in {"yes", "no"} for answer in answers.values())


def validate_fact_case(case: Mapping[str, Any]) -> dict[str, str]:
    probes = case.get("probes")
    if not isinstance(probes, dict) or set(probes) != set(PROBES):
        raise SystemExit(f"fact probes incomplete for {case.get('case_id')}")
    answers: dict[str, str] = {}
    for probe in PROBES:
        row = probes[probe]
        if not isinstance(row, dict):
            raise SystemExit(f"fact row missing for {case.get('case_id')} {probe}")
        answer = str(row.get("answer") or "")
        grade = str(row.get("source_grade") or "")
        url = str(row.get("url") or "")
        excerpt = str(row.get("excerpt") or "").strip()
        if answer not in ANSWERS or grade not in GRADES or not url.startswith("https://") or not excerpt:
            raise SystemExit(f"fact source incomplete for {case.get('case_id')} {probe}")
        answers[probe] = answer
    return answers


def answer_tally(cases: Sequence[Mapping[str, str]]) -> dict[str, dict[str, int]]:
    tally = {probe: {"yes": 0, "no": 0, "unclear": 0} for probe in PROBES}
    for answers in cases:
        for probe in PROBES:
            tally[probe][answers[probe]] += 1
    return tally


def window_offsets(radius: int) -> list[int]:
    return list(range(-int(radius), int(radius) + 1))


def _slice_map(values: Mapping[int, Any], offsets: Sequence[int]) -> dict[int, Any]:
    return {int(offset): values[int(offset)] for offset in offsets}


def case_window(
    sign_by_offset: Mapping[int, int],
    prediction_by_offset: Mapping[int, Mapping[str, str]],
    answers: Mapping[str, str],
    radius: int,
) -> dict[str, Any]:
    signs = _slice_map(sign_by_offset, window_offsets(radius))
    pair = nearest_cut(cuts(signs))
    if pair is None:
        return {
            "boundary": False,
            "cut": None,
            "true_offset": None,
            "other_offset": None,
            "true_score": None,
            "other_score": None,
            "selected": False,
            "true_predictions": None,
            "other_predictions": None,
            "segments": 1 if signs else 0,
            "head_hit": None,
            "retained": None,
        }
    true_offset, other_offset = true_side(signs, pair)
    true_prediction = dict(prediction_by_offset[true_offset])
    other_prediction = dict(prediction_by_offset[other_offset])
    true_score = side_score(true_prediction, answers)
    other_score = side_score(other_prediction, answers)
    runs = segments(signs)
    scored: list[tuple[list[int], int]] = []
    for run in runs:
        representative = segment_representative(run)
        scored.append((run, side_score(prediction_by_offset[representative], answers)))
    true_segment = next(score for run, score in scored if 0 in run)
    other_scores = [score for run, score in scored if 0 not in run]
    return {
        "boundary": True,
        "cut": [pair[0], pair[1]],
        "true_offset": true_offset,
        "other_offset": other_offset,
        "true_score": true_score,
        "other_score": other_score,
        "selected": true_score > other_score,
        "true_predictions": true_prediction,
        "other_predictions": other_prediction,
        "segments": len(runs),
        "head_hit": true_segment > max(other_scores),
        "retained": true_segment >= max(other_scores),
        "differ": vectors_differ(true_prediction, other_prediction),
    }


def evaluate_cell(
    case_ids: Sequence[str],
    signs: Mapping[str, Mapping[int, int]],
    predictions: Mapping[str, Mapping[int, Mapping[str, str]]],
    answers: Mapping[str, Mapping[str, str]],
    radius: int,
    perm_maps: Sequence[Sequence[str]],
) -> dict[str, Any]:
    ordered = sorted(case_ids)
    rows = {
        case_id: case_window(signs[case_id], predictions[case_id], answers[case_id], radius)
        for case_id in ordered
    }
    boundary_ids = [case_id for case_id in ordered if rows[case_id]["boundary"]]
    true_predictions = {case_id: rows[case_id]["true_predictions"] for case_id in boundary_ids}
    other_predictions = {case_id: rows[case_id]["other_predictions"] for case_id in boundary_ids}
    observed = selection_count(boundary_ids, answers, true_predictions, other_predictions)
    perm_counts = []
    for sources in perm_maps:
        assigned = assign_answers(ordered, sources, answers)
        perm_counts.append(
            selection_count(boundary_ids, assigned, true_predictions, other_predictions)
        )
    summary = permutation_summary(observed, len(boundary_ids), perm_counts)
    head_hits = sum(1 for case_id in boundary_ids if rows[case_id]["head_hit"])
    retained = sum(1 for case_id in boundary_ids if rows[case_id]["retained"])
    return {
        "summary": summary,
        "significant": cell_is_significant(summary),
        "head_hit_numer": head_hits,
        "head_hit_denom": len(boundary_ids),
        "head_hit_display": display_ratio(head_hits, len(boundary_ids)),
        "retained_numer": retained,
        "retained_denom": len(boundary_ids),
        "retained_display": display_ratio(retained, len(boundary_ids)),
        "no_boundary": len(ordered) - len(boundary_ids),
        "cases": rows,
        "not_a_gate_segment_rates": True,
    }


def positive_control_cell(
    case_ids: Sequence[str],
    signs: Mapping[str, Mapping[int, int]],
    predictions: Mapping[str, Mapping[int, Mapping[str, str]]],
    radius: int,
) -> dict[str, Any]:
    """Plant the true side's own labels. Permute only cases that have a cut.

    Those are the cases with a planted vector. The stream starts at the
    pre-registered control seed for every cell.
    """
    blank = {probe: "unclear" for probe in PROBES}
    boundary_ids: list[str] = []
    true_predictions: dict[str, Mapping[str, str]] = {}
    other_predictions: dict[str, Mapping[str, str]] = {}
    planted: dict[str, dict[str, str]] = {}
    missed: list[str] = []
    differing = 0
    for case_id in sorted(case_ids):
        row = case_window(signs[case_id], predictions[case_id], blank, radius)
        if not row["boundary"]:
            continue
        boundary_ids.append(case_id)
        true_predictions[case_id] = row["true_predictions"]
        other_predictions[case_id] = row["other_predictions"]
        planted[case_id] = plant_answers(row["true_predictions"])
        if row["differ"]:
            differing += 1
            checked = case_window(signs[case_id], predictions[case_id], planted[case_id], radius)
            if not checked["selected"]:
                missed.append(case_id)
    maps = permutation_maps(boundary_ids, CONTROL_SEED, PERMUTATIONS)
    observed = selection_count(boundary_ids, planted, true_predictions, other_predictions)
    perm_counts = [
        selection_count(
            boundary_ids,
            assign_answers(boundary_ids, sources, planted),
            true_predictions,
            other_predictions,
        )
        for sources in maps
    ]
    summary = permutation_summary(observed, len(boundary_ids), perm_counts)
    return {
        "summary": summary,
        "differing": differing,
        "pipeline_exact": not missed,
        "missed": missed,
        "p_below": ratio_below(int(summary["p_numer"]), int(summary["p_denom"]), CONTROL_P_MAX),
    }


def decide_status(
    *,
    included: int,
    control_exact: bool,
    differing: int,
    control_p_below: bool,
    any_robust_significant: bool,
) -> str:
    if not control_exact:
        return "control_invalid"
    if differing < MIN_DIFFERING or not control_p_below:
        return "control_underpowered"
    if included < MIN_INCLUDED:
        return "underpowered_sample"
    if any_robust_significant:
        return "pass"
    return "fail_closed"