#!/usr/bin/env python3
"""Independent checks for the conditional card-removal experiment.

The checker does not import the simulation module.  It reads the published 169
class table, expands it to all 1,326 physical combinations, recomputes the exact
one-fold oracle, and checks the multi-fold acceptance/rejection estimates with
self-normalized importance sampling (SNIS) from unconditional complete deals.
"""

from __future__ import annotations

import argparse
import csv
import json
import math
from pathlib import Path

import numpy as np


RANKS = ("2", "3", "4", "5", "6", "7", "8", "9", "T", "J", "Q", "K", "A")
VALUE_TO_RANK = {value: rank for value, rank in zip(range(2, 15), RANKS)}
MODELS = (
    "random_fold_control",
    "rank_gradient",
    "pair_suited_gradient",
    "bucketed_top_band",
    "inverted_rank_diagnostic",
)
METRICS = (
    "at_least_one_ace",
    "any_pair",
    "premium_pair_tt_plus",
    "ace_king",
    "both_cards_ten_plus",
)
TARGET_MEAN = 0.60
Z_95 = 1.959963984540054


def canonical(first: int, second: int) -> str:
    first_rank, second_rank = first // 4 + 2, second // 4 + 2
    high, low = max(first_rank, second_rank), min(first_rank, second_rank)
    if high == low:
        return VALUE_TO_RANK[high] * 2
    suffix = "s" if first % 4 == second % 4 else "o"
    return f"{VALUE_TO_RANK[high]}{VALUE_TO_RANK[low]}{suffix}"


def exact_baselines() -> dict[str, float]:
    total = math.comb(52, 2)
    return {
        "at_least_one_ace": 1.0 - math.comb(48, 2) / total,
        "any_pair": 13 * math.comb(4, 2) / total,
        "premium_pair_tt_plus": 5 * math.comb(4, 2) / total,
        "ace_king": 16 / total,
        "both_cards_ten_plus": math.comb(20, 2) / total,
    }


def read_likelihoods(path: Path) -> tuple[list[str], dict[str, dict[str, float]]]:
    with path.open(newline="", encoding="utf-8") as handle:
        rows = list(csv.DictReader(handle))
    names = [row["hand_class"] for row in rows]
    if len(rows) != 169 or len(set(names)) != 169:
        raise AssertionError("Fold-likelihood table must contain 169 unique hand classes")
    table = {
        model: {row["hand_class"]: float(row[model]) for row in rows} for model in MODELS
    }
    return names, table


def matrices(
    names: list[str], table: dict[str, dict[str, float]]
) -> tuple[dict[str, np.ndarray], np.ndarray]:
    class_lookup = {name: index for index, name in enumerate(names)}
    fold_matrices = {model: np.zeros((52, 52), dtype=np.float64) for model in MODELS}
    class_matrix = np.full((52, 52), -1, dtype=np.int16)
    for first in range(52):
        for second in range(first + 1, 52):
            hand_class = canonical(first, second)
            class_index = class_lookup[hand_class]
            class_matrix[first, second] = class_matrix[second, first] = class_index
            for model in MODELS:
                value = table[model][hand_class]
                if not 0.0 < value <= 1.0:
                    raise AssertionError(f"Invalid {model} likelihood {value} for {hand_class}")
                fold_matrices[model][first, second] = value
                fold_matrices[model][second, first] = value
    return fold_matrices, class_matrix


def recompute_one_fold(
    names: list[str], fold_matrices: dict[str, np.ndarray], class_matrix: np.ndarray
) -> dict[str, object]:
    pairs = [(first, second) for first in range(52) for second in range(first + 1, 52)]
    target_total = math.comb(50, 2)
    metric_given_fold = np.zeros((len(pairs), len(METRICS)), dtype=np.float64)
    rank_given_fold = np.zeros((len(pairs), 13), dtype=np.float64)
    class_given_fold = np.zeros((len(pairs), 169), dtype=np.float64)

    for pair_index, (first, second) in enumerate(pairs):
        rank_counts = np.full(13, 4, dtype=np.int16)
        rank_counts[first // 4] -= 1
        rank_counts[second // 4] -= 1
        aces = int(rank_counts[12])
        broadways = int(rank_counts[8:].sum())
        metric_given_fold[pair_index] = (
            1.0 - math.comb(50 - aces, 2) / target_total,
            sum(math.comb(int(count), 2) for count in rank_counts) / target_total,
            sum(math.comb(int(count), 2) for count in rank_counts[8:]) / target_total,
            int(rank_counts[12]) * int(rank_counts[11]) / target_total,
            math.comb(broadways, 2) / target_total,
        )
        rank_given_fold[pair_index] = rank_counts / 50.0

        available = [card for card in range(52) if card not in (first, second)]
        counts = np.zeros(169, dtype=np.int32)
        for left in range(49):
            indices = class_matrix[available[left], available[left + 1 :]]
            counts += np.bincount(indices, minlength=169).astype(np.int32)
        class_given_fold[pair_index] = counts / target_total

    output: dict[str, object] = {}
    for model in MODELS:
        matrix = fold_matrices[model]
        weights = np.array([matrix[first, second] for first, second in pairs])
        posterior = weights / weights.sum()
        output[model] = {
            "metrics": dict(zip(METRICS, (posterior @ metric_given_fold).tolist())),
            "rank_shares": dict(zip(RANKS, (posterior @ rank_given_fold).tolist())),
            "classes": dict(zip(names, (posterior @ class_given_fold).tolist())),
        }
    return output


def random_prefixes(samples: int, length: int, rng: np.random.Generator) -> np.ndarray:
    """Independent deal generator: random keys, partial selection, then key ordering."""

    keys = rng.random((samples, 52))
    selected = np.argpartition(keys, kth=length - 1, axis=1)[:, :length]
    selected_keys = np.take_along_axis(keys, selected, axis=1)
    order = np.argsort(selected_keys, axis=1)
    return np.take_along_axis(selected, order, axis=1).astype(np.int16, copy=False)


def ratio_and_se(
    sum_w: float,
    sum_w2: float,
    sum_wy: float,
    sum_w2y: float,
    sum_w2y2: float,
) -> tuple[float, float]:
    estimate = sum_wy / sum_w
    residual = sum_w2y2 - 2.0 * estimate * sum_w2y + estimate * estimate * sum_w2
    return estimate, math.sqrt(max(0.0, residual)) / sum_w


def check_scenario(
    primary: dict[str, object],
    fold_matrix: np.ndarray,
    class_matrix: np.ndarray,
    raw_samples: int,
    batch_size: int,
    seed: int,
) -> dict[str, object]:
    folds = int(primary["observed_folds_before_button"])
    rng = np.random.Generator(np.random.Philox(seed))
    sum_w = 0.0
    sum_w2 = 0.0
    metric_sums = {metric: np.zeros(3, dtype=np.float64) for metric in METRICS}
    class_sums_wy = np.zeros(169, dtype=np.float64)
    class_sums_w2y = np.zeros(169, dtype=np.float64)
    rank_sums_wy = np.zeros(13, dtype=np.float64)
    rank_sums_w2y = np.zeros(13, dtype=np.float64)
    rank_sums_w2y2 = np.zeros(13, dtype=np.float64)
    completed = 0

    while completed < raw_samples:
        count = min(batch_size, raw_samples - completed)
        deals = random_prefixes(count, 2 * folds + 2, rng)
        folding = deals[:, : 2 * folds].reshape(count, folds, 2)
        probabilities = fold_matrix[folding[:, :, 0], folding[:, :, 1]]
        weights = probabilities.prod(axis=1, dtype=np.float64)
        weights2 = np.square(weights)
        sum_w += float(weights.sum())
        sum_w2 += float(weights2.sum())

        first, second = deals[:, 2 * folds], deals[:, 2 * folds + 1]
        first_rank, second_rank = first // 4 + 2, second // 4 + 2
        high, low = np.maximum(first_rank, second_rank), np.minimum(first_rank, second_rank)
        pair = high == low
        events = {
            "at_least_one_ace": high == 14,
            "any_pair": pair,
            "premium_pair_tt_plus": pair & (high >= 10),
            "ace_king": (high == 14) & (low == 13),
            "both_cards_ten_plus": low >= 10,
        }
        for metric, values in events.items():
            numeric = values.astype(np.float64)
            metric_sums[metric] += (
                float(np.dot(weights, numeric)),
                float(np.dot(weights2, numeric)),
                float(np.dot(weights2, numeric)),
            )

        class_indices = class_matrix[first, second]
        class_sums_wy += np.bincount(class_indices, weights=weights, minlength=169)
        class_sums_w2y += np.bincount(class_indices, weights=weights2, minlength=169)

        rank_counts = np.full((count, 13), 4, dtype=np.int16)
        folded_ranks = folding.reshape(count, 2 * folds) // 4
        for rank_index in range(13):
            rank_counts[:, rank_index] -= (folded_ranks == rank_index).sum(axis=1)
        rank_shares = rank_counts / (52 - 2 * folds)
        rank_sums_wy += np.einsum("i,ij->j", weights, rank_shares)
        rank_sums_w2y += np.einsum("i,ij->j", weights2, rank_shares)
        rank_sums_w2y2 += np.einsum("i,ij->j", weights2, np.square(rank_shares))
        completed += count

    metric_checks = {}
    all_agree = True
    maximum_z = 0.0
    for metric in METRICS:
        sums = metric_sums[metric]
        estimate, checker_se = ratio_and_se(sum_w, sum_w2, sums[0], sums[1], sums[2])
        primary_metric = primary["metrics"][metric]
        primary_estimate = float(primary_metric["estimate"])
        accepted = int(primary["accepted_deals"])
        primary_se = math.sqrt(primary_estimate * (1.0 - primary_estimate) / accepted)
        combined_se = math.sqrt(primary_se * primary_se + checker_se * checker_se)
        z_score = abs(primary_estimate - estimate) / combined_se if combined_se else 0.0
        agrees = bool(z_score <= 5.0)
        all_agree &= agrees
        maximum_z = max(maximum_z, z_score)
        metric_checks[metric] = {
            "primary_estimate": primary_estimate,
            "checker_estimate": estimate,
            "checker_delta_method_standard_error": checker_se,
            "combined_standardized_difference": z_score,
            "within_five_combined_standard_errors": agrees,
        }

    class_checks = {}
    for class_index, (hand_class, primary_class) in enumerate(
        primary["all_169_button_hand_classes"].items()
    ):
        estimate, checker_se = ratio_and_se(
            sum_w,
            sum_w2,
            class_sums_wy[class_index],
            class_sums_w2y[class_index],
            class_sums_w2y[class_index],
        )
        primary_estimate = float(primary_class["estimate"])
        accepted = int(primary["accepted_deals"])
        primary_se = math.sqrt(primary_estimate * (1.0 - primary_estimate) / accepted)
        combined_se = math.sqrt(primary_se * primary_se + checker_se * checker_se)
        z_score = abs(primary_estimate - estimate) / combined_se if combined_se else 0.0
        agrees = bool(z_score <= 5.0)
        all_agree &= agrees
        maximum_z = max(maximum_z, z_score)
        class_checks[hand_class] = {
            "primary_estimate": primary_estimate,
            "checker_estimate": estimate,
            "checker_delta_method_standard_error": checker_se,
            "combined_standardized_difference": z_score,
            "within_five_combined_standard_errors": agrees,
        }

    rank_checks = {}
    for rank_index, rank in enumerate(RANKS):
        estimate, checker_se = ratio_and_se(
            sum_w,
            sum_w2,
            rank_sums_wy[rank_index],
            rank_sums_w2y[rank_index],
            rank_sums_w2y2[rank_index],
        )
        primary_rank = primary["remaining_rank_composition"][rank]
        primary_estimate = float(primary_rank["expected_share_of_remaining_cards"])
        primary_se = float(primary_rank["monte_carlo_standard_error"])
        combined_se = math.sqrt(primary_se * primary_se + checker_se * checker_se)
        z_score = abs(primary_estimate - estimate) / combined_se if combined_se else 0.0
        agrees = bool(z_score <= 5.0)
        all_agree &= agrees
        maximum_z = max(maximum_z, z_score)
        rank_checks[rank] = {
            "primary_estimate": primary_estimate,
            "checker_estimate": estimate,
            "checker_delta_method_standard_error": checker_se,
            "combined_standardized_difference": z_score,
            "within_five_combined_standard_errors": agrees,
        }

    return {
        "model": primary["model"],
        "table_size": primary["table_size"],
        "observed_folds": folds,
        "raw_unconditional_deals": raw_samples,
        "rng": "NumPy Philox",
        "seed": seed,
        "sum_weights": sum_w,
        "effective_sample_size": sum_w * sum_w / sum_w2,
        "effective_sample_fraction": sum_w * sum_w / (sum_w2 * raw_samples),
        "metrics": metric_checks,
        "all_169_button_hand_classes": class_checks,
        "remaining_rank_composition": rank_checks,
        "maximum_combined_standardized_difference": maximum_z,
        "all_results_within_five_combined_standard_errors": all_agree,
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input-dir", type=Path, default=Path(__file__).resolve().parent)
    parser.add_argument("--raw-samples", type=int, default=500_000)
    parser.add_argument("--batch-size", type=int, default=50_000)
    parser.add_argument("--seed", type=int, default=902026090406)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    with (args.input_dir / "results.json").open(encoding="utf-8") as handle:
        primary = json.load(handle)
    with (args.input_dir / "one_fold_exact_oracle.json").open(encoding="utf-8") as handle:
        published_oracle = json.load(handle)
    names, table = read_likelihoods(args.input_dir / "fold_likelihoods.csv")
    fold_matrices, class_matrix = matrices(names, table)

    mean_checks = {}
    for model in MODELS:
        values = fold_matrices[model][np.triu_indices(52, k=1)]
        mean = float(values.mean())
        mean_checks[model] = {
            "physical_combo_mean": mean,
            "target": TARGET_MEAN,
            "absolute_error": abs(mean - TARGET_MEAN),
            "passes_1e-11_tolerance": bool(abs(mean - TARGET_MEAN) <= 1e-11),
        }

    recomputed_oracle = recompute_one_fold(names, fold_matrices, class_matrix)
    oracle_max_error = 0.0
    for model in MODELS:
        published = published_oracle["models"][model]
        for metric in METRICS:
            oracle_max_error = max(
                oracle_max_error,
                abs(recomputed_oracle[model]["metrics"][metric] - published["metrics"][metric]["estimate"]),
            )
        for rank in RANKS:
            oracle_max_error = max(
                oracle_max_error,
                abs(
                    recomputed_oracle[model]["rank_shares"][rank]
                    - published["remaining_rank_composition"][rank]["expected_share_of_remaining_cards"]
                ),
            )
        for hand_class in names:
            oracle_max_error = max(
                oracle_max_error,
                abs(
                    recomputed_oracle[model]["classes"][hand_class]
                    - published["all_169_button_hand_classes"][hand_class]["estimate"]
                ),
            )

    baseline = exact_baselines()
    random_oracle_error = max(
        abs(recomputed_oracle["random_fold_control"]["metrics"][metric] - baseline[metric])
        for metric in METRICS
    )

    scenario_checks = []
    for index, scenario in enumerate(primary["scenarios"]):
        model = scenario["model"]
        print(f"Checking {model}, {scenario['table_size']}-max...", flush=True)
        scenario_checks.append(
            check_scenario(
                scenario,
                fold_matrices[model],
                class_matrix,
                args.raw_samples,
                args.batch_size,
                args.seed + index * 10_003,
            )
        )

    all_pass = (
        all(item["passes_1e-11_tolerance"] for item in mean_checks.values())
        and oracle_max_error <= 1e-12
        and random_oracle_error <= 1e-12
        and all(item["all_results_within_five_combined_standard_errors"] for item in scenario_checks)
    )
    output = {
        "schema_version": 1,
        "checker_method": (
            "Independent CSV-driven expansion to 1,326 combos; exact one-fold enumeration; "
            "multi-fold SNIS with delta-method standard errors using NumPy Philox and a "
            "random-key deal generator"
        ),
        "fold_likelihood_combo_mean_checks": mean_checks,
        "one_fold_oracle_max_absolute_error": oracle_max_error,
        "random_control_exact_baseline_max_absolute_error": random_oracle_error,
        "agreement_rule": "all estimates within five combined Monte Carlo standard errors",
        "scenario_checks": scenario_checks,
        "all_checks_pass": all_pass,
    }
    with (args.input_dir / "checker_results.json").open("w", encoding="utf-8", newline="\n") as handle:
        json.dump(output, handle, indent=2)
        handle.write("\n")

    maximum_z = max(item["maximum_combined_standardized_difference"] for item in scenario_checks)
    lines = [
        "# Independent checker",
        "",
        f"Overall: **{'PASS' if all_pass else 'FAIL'}**",
        "",
        f"- All five 1,326-combo-weighted likelihood means match 0.60: {all(item['passes_1e-11_tolerance'] for item in mean_checks.values())}.",
        f"- Independently recomputed one-fold oracle maximum absolute discrepancy: `{oracle_max_error:.3g}`.",
        f"- Exact random-control baseline maximum absolute discrepancy: `{random_oracle_error:.3g}`.",
        f"- Maximum primary-versus-SNIS standardized difference across metrics, all 169 classes, and 13 rank shares: `{maximum_z:.3f}` (limit 5.0).",
        "",
        "The primary and checker use different RNG families and different uniform-deal generators. The primary accepts complete deals and uses Wilson intervals; this checker keeps all deals, weights them, and uses delta-method ratio-estimator uncertainty. Agreement addresses implementation error, not model realism.",
        "",
    ]
    (args.input_dir / "CHECKS.md").write_text("\n".join(lines), encoding="utf-8", newline="\n")
    if not all_pass:
        raise SystemExit("Independent checker failed; inspect checker_results.json")
    print(f"All checks passed; maximum standardized difference {maximum_z:.3f}")


if __name__ == "__main__":
    main()