#!/usr/bin/env python3
"""Reproduce every compact arithmetic check in the statistical-test catalog."""
from __future__ import annotations

import argparse
import json
import math
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]


def ranks(values: list[float]) -> list[float]:
    ordered = sorted(enumerate(values), key=lambda pair: pair[1])
    result = [0.0] * len(values)
    start = 0
    while start < len(ordered):
        end = start + 1
        while end < len(ordered) and ordered[end][1] == ordered[start][1]:
            end += 1
        average_rank = ((start + 1) + end) / 2
        for index, _ in ordered[start:end]:
            result[index] = average_rank
        start = end
    return result


def calculate(kind: str, inputs: dict) -> float:
    if kind == "ratio":
        return inputs["numerator"] / inputs["denominator"]
    if kind == "difference":
        return inputs["left"] - inputs["right"]
    if kind == "sum":
        return sum(inputs["values"])
    if kind == "chi_square":
        return sum(
            (observed - expected) ** 2 / expected
            for observed, expected in zip(inputs["observed"], inputs["expectedCounts"], strict=True)
        )
    if kind == "kruskal":
        groups = inputs["groups"]
        flat = [value for group in groups for value in group]
        ranked = ranks(flat)
        offset = 0
        rank_sums = []
        for group in groups:
            rank_sums.append(sum(ranked[offset : offset + len(group)]))
            offset += len(group)
        n = len(flat)
        raw = 12 / (n * (n + 1)) * sum(
            rank_sum**2 / len(group) for rank_sum, group in zip(rank_sums, groups, strict=True)
        ) - 3 * (n + 1)
        tie_counts: dict[float, int] = {}
        for value in flat:
            tie_counts[value] = tie_counts.get(value, 0) + 1
        correction = 1 - sum(count**3 - count for count in tie_counts.values()) / (n**3 - n)
        return raw / correction
    if kind == "ks_one_uniform":
        sample = sorted(inputs["sample"])
        n = len(sample)
        d_plus = max((index + 1) / n - value for index, value in enumerate(sample))
        d_minus = max(value - index / n for index, value in enumerate(sample))
        return max(d_plus, d_minus)
    if kind == "ks_two":
        left, right = sorted(inputs["sampleA"]), sorted(inputs["sampleB"])
        support = sorted(set(left + right))
        return max(
            abs(sum(value <= point for value in left) / len(left) - sum(value <= point for value in right) / len(right))
            for point in support
        )
    raise ValueError(f"Unknown check kind: {kind}")


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args()
    catalog = json.loads((ROOT / "content" / "statistical_tests.json").read_text(encoding="utf-8"))
    results = {}
    for item in catalog:
        check = item["check"]
        actual = calculate(check["kind"], check["inputs"])
        expected = check["expected"]
        results[item["id"]] = {
            "actual": actual,
            "expected": expected,
            "passed": math.isclose(actual, expected, rel_tol=1e-9, abs_tol=1e-9),
        }
    if args.json:
        print(json.dumps(results, ensure_ascii=False, indent=2))
    else:
        for item_id, result in results.items():
            print(f"{item_id}: {'PASS' if result['passed'] else 'FAIL'} ({result['actual']:.10g})")
    return 0 if all(result["passed"] for result in results.values()) else 1


if __name__ == "__main__":
    raise SystemExit(main())