#!/usr/bin/env python3
# Companion script for the article at https://www.dr-lex.be/info-stuff/zocchihedron.html
# Run in an environment with scipy installed.
import sys
from scipy import stats
import numpy as np

# For the histogram plots
MAX_LINE_WIDTH = 80

# Mapping for the second version of the Zocchihedron number arrangement.
# The very first edition and the latest editions have different arrangements.
# The mapping specifies how the original faces need to be reordered to end up
# with a more or less sequential ordering when going from one side to the other.
MAPPING = [
    1, 2, 65, 18, 92, 56, 27, 3, 89, 71, 49, 33, 15, 95, 77, 59, 39, 21, 4, 94,
    76, 58, 34, 16, 88, 70, 46, 28, 10, 82, 64, 40, 22, 5, 81, 63, 51, 35, 17, 93,
    75, 57, 47, 29, 11, 87, 69, 53, 41, 23, 44, 54, 72, 90, 14, 32, 48, 60, 78, 96,
    20, 38, 50, 66, 84, 8, 26, 31, 55, 73, 91, 19, 37, 61, 79, 97, 7, 25, 43, 67,
    85, 13, 6, 24, 42, 62, 80, 98, 12, 30, 52, 68, 86, 9, 45, 74, 99, 36, 83, 100
]


# ── Start of script ──────────────────────────────────────────────────────────

file_path = "zocchi.txt"
if len(sys.argv) > 1:
    file_path = sys.argv[1]

with open(file_path, encoding="utf-8") as in_file:
    throws = [int(line) for line in in_file.readlines() if not line.startswith("#")]

# Count occurrences of each face
observed = np.array([throws.count(i) for i in range(1, 101)])

# Expected count per face
expected = np.full(100, len(throws) / 100)  # 51.5 each

# Run chi-square test
chi2_stat, p_value = stats.chisquare(f_obs=observed, f_exp=expected)

print(f"χ² statistic: {chi2_stat:.4f}")
print(f"p-value:      {p_value:.6f}")
print(f"Degrees of freedom: 99")

if p_value < 0.001:
    print("Result: Die is UNFAIR (99.9% confidence)")
elif p_value < 0.01:
    print("Result: Die is UNFAIR (99% confidence)")
elif p_value < 0.05:
    print("Result: Die is UNFAIR (95% confidence)")
elif p_value < 0.10:
    print("Result: Die is UNFAIR (90% confidence)")
else:
    print("Result: No evidence of unfairness")


# ── Histogram helper ──────────────────────────────────────────────────────────
def print_histogram(counts, title):
    """Print a horizontal ASCII-art bar chart for a list of 100 counts."""
    n          = len(counts)
    label_w    = len(str(n))          # width of the widest face label ("100" → 3)
    count_w    = len(str(max(counts)))# width of the widest count value
    # prefix: "NNN [CCCCC] " before the bar
    prefix_w   = label_w + 1 + count_w + 3   # "NNN " + "CCCCC" + " |"  (+space)
    bar_area   = MAX_LINE_WIDTH - prefix_w
    max_count  = max(counts)

    print(f"\n{'─' * MAX_LINE_WIDTH}")
    print(title)
    print(f"{'─' * MAX_LINE_WIDTH}")
    print(f"{'Expected per face:':>{prefix_w}}{len(throws) // 100} "
          f"({'=' * (len(throws) // 100 * bar_area // max_count)})")
    print()

    for i, count in enumerate(counts, start=1):
        bar_len = round(count * bar_area / max_count)
        bar     = '=' * bar_len
        label   = str(i).rjust(label_w)
        cnt     = str(count).rjust(count_w)
        print(f"{label} {cnt} |{bar}")

# ── Histogram 1: raw observed frequencies ────────────────────────────────────
print_histogram(
    observed.tolist(),
    "Observed frequency per face (die-face order)"
)

# ── Histogram 2: remapped frequencies (spiral / physical order) ───────────────
assert len(MAPPING) == 100, "MAPPING must contain exactly 100 entries"
assert sorted(MAPPING) == list(range(1, 101)), \
    "MAPPING must be a permutation of 1–100"

# MAPPING[i] is the die face whose count should appear at position i+1.
# observed is 0-indexed: observed[f-1] = count for face f.
remapped = [observed[face - 1] for face in MAPPING]

print_histogram(
    remapped,
    "Observed frequency remapped to physical (spiral) order"
)