#!/usr/bin/env python3
"""Independent limiting checks on truncated Gibbs conditionals and normalization."""
from pathlib import Path
import argparse
import math
import random
import statistics
from timing_data import read_observations
from chopping_model import CoefficientGrid, predict
from timing_posterior import conditional_update, physical, log_probability

parser = argparse.ArgumentParser()
parser.add_argument("--input", type=Path, required=True)
arguments = parser.parse_args()
_, measured = read_observations(arguments.input)
grid = CoefficientGrid()
# Vanishing information: the conditional must reproduce its actual truncated
# prior, not an untruncated Gaussian projected/clipped onto the support.
uninformative = [{**row, "sigma_days": 1e6} for row in measured]
state = [
    math.log(6),
    math.log(150),
    0,
    0,
    1635.87,
    1698.99,
    66.028,
    125.865,
    math.log(0.01),
]
generator = random.Random(191919)
samples = [
    physical(conditional_update(state, uninformative, grid, True, 0.3, generator))
    for _ in range(10000)
]
normal = statistics.NormalDist()
boundary = 0.15 / 0.3
normalization = normal.cdf(boundary) - normal.cdf(-boundary)
expected_variance = 0.3**2 * (1 - 2 * boundary * normal.pdf(boundary) / normalization)
for column in [2, 3]:
    values = [sample[column] for sample in samples]
    assert all(-0.15 < value < 0.15 for value in values)
    assert abs(statistics.mean(values)) < 0.003
    assert abs(statistics.variance(values) / expected_variance - 1) < 0.05
for column, lower, upper in [(0, 0.05, 60), (1, 5, 600)]:
    values = [sample[column] for sample in samples]
    assert all(lower < value < upper for value in values)
    assert abs(statistics.mean(values) - (lower + upper) / 2) < (upper - lower) * 0.015
    assert abs(statistics.variance(values) / ((upper - lower) ** 2 / 12) - 1) < 0.05
small_jitter = state[:]
large_jitter = state[:]
small_jitter[8] = math.log(0.01)
large_jitter[8] = math.log(0.19)
perfect = [
    {**row, "observed_reduced_bjd": value}
    for row, value in zip(measured, predict(physical(state), measured, grid, True))
]
assert log_probability(small_jitter, perfect, grid, True, 0.03) > log_probability(
    large_jitter, perfect, grid, True, 0.03
)
print(
    "PASS: 10000 conditional draws recover truncated eccentricity-prior moments and uniform finite positive mass priors; support holds; large jitter is penalized in the zero-residual limit"
)