"""Exact enumeration of a three-unit unequal-probability sampling design."""
from fractions import Fraction as F
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np


def main():
    samples = [(0, 1), (0, 2), (1, 2)]
    probabilities = [F(1, 2), F(1, 3), F(1, 6)]
    inclusion = [F(5, 6), F(2, 3), F(1, 2)]
    values = [F(2), F(4), F(6)]
    weighted = [sum(values[i]/inclusion[i] for i in s) for s in samples]
    naive = [F(3, 2)*sum(values[i] for i in s) for s in samples]
    assert sum(p*t for p, t in zip(probabilities, weighted)) == 12
    assert sum(p*t for p, t in zip(probabilities, naive)) == 11
    assert sum(p*(t-12)**2 for p, t in zip(probabilities, weighted)) == F(72, 5)
    estimates = []
    for sample, pij in zip(samples, probabilities):
        i, j = sample
        v = sum((1-inclusion[k])*(values[k]/inclusion[k])**2 for k in sample)
        v += 2*(pij-inclusion[i]*inclusion[j])/pij*(values[i]/inclusion[i])*(values[j]/inclusion[j])
        estimates.append(v)
    assert estimates == [F(244, 25), F(1464, 25), F(-60)]
    assert sum(p*v for p, v in zip(probabilities, estimates)) == F(72, 5)
    fig, (chance, totals) = plt.subplots(1, 2, figsize=(9.6, 4.2), dpi=100,
                                       layout="constrained", width_ratios=[1, 1.5])
    fig.set_facecolor("white")
    x = np.arange(3)
    labels = ["{1, 2}", "{1, 3}", "{2, 3}"]
    bars = chance.bar(x, [float(p) for p in probabilities], color="#577d9d", width=.65)
    for bar, label in zip(bars, ["1/2", "1/3", "1/6"]):
        chance.text(bar.get_x()+bar.get_width()/2, bar.get_height()+.015,
                    label, ha="center")
    chance.set(xticks=x, xticklabels=labels, xlabel="Selected sample",
               ylabel="Selection probability", ylim=(0, .62),
               title="Unequal chances of selecting each pair")
    totals.bar(x-.18, [float(t) for t in weighted], width=.36,
               label="Horvitz–Thompson total", color="#1565a5")
    totals.bar(x+.18, [float(t) for t in naive], width=.36,
               label="Three × sample mean", color="#c56427")
    totals.axhline(12, color="#555555", ls="--", lw=1.2, label="True total = 12")
    totals.set(xticks=x, xticklabels=labels, xlabel="Selected sample", ylabel="Estimated total",
               ylim=(0, 22), title="Expected totals: 12 versus 11")
    totals.legend(fontsize=8, loc="upper left")
    for ax in [chance, totals]:
        ax.grid(axis="y", alpha=.15)
        ax.set_axisbelow(True)
    fig.savefig(Path(__file__).with_suffix(".png").name, facecolor="white")
    plt.close(fig)


if __name__ == "__main__":
    main()