"""Enumerate a uniform two-of-four sample and its inclusion indicators."""
from itertools import combinations
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np


def main():
    samples = list(combinations(range(1, 5), 2))
    values = np.array([[int(1 in s), int(2 in s)] for s in samples])
    assert np.allclose(np.cov(values.T, bias=True), [[.25, -1/12], [-1/12, .25]])
    counts = np.zeros((2, 2), dtype=int)
    for i, j in values:
        counts[i, j] += 1
    assert counts.tolist() == [[1, 2], [2, 1]]
    fig, (ax, note) = plt.subplots(1, 2, figsize=(7.2, 3.6), dpi=100,
                                   layout="constrained", width_ratios=[1, 1])
    fig.set_facecolor("white")
    ax.imshow(counts/6, vmin=0, vmax=.5, cmap="Blues", origin="lower")
    for i in range(2):
        for j in range(2):
            ax.text(j, i, "1/6" if counts[i, j] == 1 else "1/3",
                    ha="center", va="center", fontsize=18)
    ax.set(xticks=[0, 1], yticks=[0, 1], xlabel="Indicator I₂", ylabel="Indicator I₁",
           title="Joint probabilities")
    note.axis("off")
    note.text(.05, .95, "Six equally likely samples\n\n" +
              "\n".join(f"{{{s[0]}, {s[1]}}}  →  ({v[0]}, {v[1]})"
                        for s, v in zip(samples, values)) +
              "\n\nCovariance = −1/12\nCorrelation = −1/3",
              transform=note.transAxes, va="top", fontsize=11, linespacing=1.35)
    fig.savefig(Path(__file__).with_suffix(".png").name, facecolor="white")
    plt.close(fig)


if __name__ == "__main__":
    main()