"""Exact intervals in an identifiable cubic normal-mean model; output to cwd."""

from pathlib import Path
from statistics import NormalDist

import matplotlib.pyplot as plt
import numpy as np


def main():
    z = NormalDist().inv_cdf(0.975)
    sample_sizes = np.array([8, 64, 512, 4096], dtype=float)
    mean_halfwidth = z / np.sqrt(sample_sizes)
    parameter_halfwidth = np.cbrt(mean_halfwidth)
    assert np.allclose(parameter_halfwidth[:-1] / parameter_halfwidth[1:], np.sqrt(2))
    assert np.isclose(parameter_halfwidth[1] / parameter_halfwidth[3], 2)
    assert np.isclose(np.cbrt(z / 8), parameter_halfwidth[1])
    theta = np.linspace(-1.2, 1.2, 801)
    half_mean = z / 8
    half_theta = np.cbrt(half_mean)
    assert np.isclose(half_theta**3, half_mean)

    fig, axes = plt.subplots(1, 2, figsize=(10, 4.2), dpi=100, facecolor="white")
    ax = axes[0]
    ax.plot(theta, theta**3, color="#245e8d", linewidth=2, label=r"$\mu=\theta^3$")
    ax.axhspan(-half_mean, half_mean, color="#e8a34a", alpha=0.22)
    ax.axvspan(-half_theta, half_theta, color="#357bb3", alpha=0.10)
    ax.scatter([-half_theta, half_theta], [-half_mean, half_mean], color="#a04e0c", zorder=3)
    ax.axhline(0, color="0.6", linewidth=0.7)
    ax.axvline(0, color="0.6", linewidth=0.7)
    ax.set(xlabel=r"Parameter $\theta$", ylabel=r"Mean $\mu$", title="Exact interval by monotone inversion")
    ax.text(0.04, 0.94, "Observed mean = 0; n = 64\n95% interval for mean: ±0.245\n95% interval for parameter: ±0.626", transform=ax.transAxes, va="top", fontsize=9)
    ax.grid(alpha=0.18)
    ax.legend(loc="lower right", frameon=False)

    ax = axes[1]
    ax.loglog(sample_sizes, mean_halfwidth, "o-", label=r"Mean half-width $z/\sqrt{n}$", color="#be6c24")
    ax.loglog(sample_sizes, parameter_halfwidth, "s-", label=r"Parameter half-width $(z/\sqrt{n})^{1/3}$", color="#245e8d")
    ax.set(xlabel="Sample size n", ylabel="95% interval half-width", title="A flat mean map slows interval shrinkage")
    ax.set_xticks(sample_sizes, ["8", "64", "512", "4096"])
    ax.grid(which="both", alpha=0.18)
    ax.legend(loc="lower left", fontsize=8.5, frameon=False)
    fig.suptitle(r"Independent $Y_i\sim N(\theta^3,1)$: identifiable even though $I(0)=0$", fontsize=13)
    fig.tight_layout(rect=(0, 0.055, 1, 0.94))
    fig.text(0.5, 0.016, "Original model; right panel uses an observed sample mean of zero, not an expected interval width.", ha="center", fontsize=8)
    fig.savefig(Path.cwd() / "flat-mean-map-confidence.png", facecolor="white", transparent=False)
    plt.close(fig)
    print(f"Exact 95% half-width at n=64: {half_theta:.9f}; checks passed")


if __name__ == "__main__":
    main()