"""Enumerate two finite error laws with identical covariance; PNG in cwd."""
from itertools import product
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt


def main():
    signs = np.array(list(product([-1., 1.], repeat=2))).T
    factor = np.array([[1., 0.], [0.8, 0.6]])
    rotation = np.array([[1., -1.], [1., 1.]]) / np.sqrt(2)
    first, second = factor @ signs, factor @ rotation @ signs
    expected = np.array([[1., 0.8], [0.8, 1.]])
    for points in (first, second):
        np.testing.assert_allclose(points.mean(axis=1), 0, atol=1e-15)
        np.testing.assert_allclose(points @ points.T / 4, expected)
    np.testing.assert_allclose(rotation @ rotation.T, np.eye(2), atol=1e-15)
    assert np.mean(np.abs(first[0]) > 1.2) == 0
    assert np.mean(np.abs(second[0]) > 1.2) == 0.5
    error_h = .02 * signs[0]
    error_theta = .01 * (.9 * signs[0] + np.sqrt(.19) * signs[1])
    assert np.all(.5 + error_theta > 0)
    assert np.all(1 + error_h > .5 + error_theta)
    np.testing.assert_allclose(np.min(.5 + error_h - error_theta), .5 - .011 - .01 * np.sqrt(.19))
    fig, axes = plt.subplots(1, 2, figsize=(10, 4.3), dpi=100, facecolor='white')
    order = [0, 1, 3, 2, 0]
    axes[0].plot(first[0, order], first[1, order], color='#2475ad', alpha=.55)
    axes[0].scatter(*first, s=65, color='#2475ad', zorder=3)
    axes[0].set_title('Model A: each point has probability ¼')
    axes[1].scatter(*first, s=65, color='#2475ad', label='A: E = LS', zorder=3)
    axes[1].scatter(*second, s=70, marker='D', color='#df8525', label='B: E = LOS', zorder=3)
    for edge in (-1.2, 1.2):
        axes[1].axvline(edge, ls='--', color='#555555', lw=1)
    axes[1].axvspan(-1.65, -1.2, color='#dddddd', alpha=.5)
    axes[1].axvspan(1.2, 1.65, color='#dddddd', alpha=.5)
    axes[1].legend(loc='upper left', fontsize=9)
    axes[1].set_title('P(|X| > 1.2): A = 0; B = ½')
    for ax in axes:
        ax.axhline(0, color='#aaaaaa', lw=.7)
        ax.axvline(0, color='#aaaaaa', lw=.7)
        ax.set(xlabel='Error X (normalized units)', ylabel='Error Y (normalized units)', xlim=(-1.65, 1.65), ylim=(-1.65, 1.65))
        ax.grid(alpha=.2)
    fig.suptitle('Bounded laws with the same mean and covariance · constructed states', fontsize=12)
    fig.tight_layout(rect=(0, 0, 1, .94))
    fig.savefig(Path.cwd() / (Path(__file__).stem + '.png'), facecolor='white')
    plt.close(fig)


if __name__ == '__main__':
    main()