"""Original noise-free tensor-design comparison; write the PNG to cwd.

This is an algebra illustration, not a reconstruction of Jones's sampling
schemes, Monte Carlo simulations or acquisition data.
"""
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np


def design(directions):
    x, y, z = np.asarray(directions).T
    return np.column_stack((x*x, y*y, z*z, 2*x*y, 2*x*z, 2*y*z))


def main():
    axes = np.eye(3)
    paired = np.vstack((axes, -axes))
    mixed = np.vstack((axes, [[1, 1, 0], [1, 0, 1], [0, 1, 1]]))
    mixed /= np.linalg.norm(mixed, axis=1)[:, None]
    first, second = design(paired), design(mixed)
    assert np.linalg.matrix_rank(first) == 3
    assert np.linalg.matrix_rank(second) == 6
    assert np.allclose(first[:3], first[3:])
    coefficients = np.array([1.8, 0.9, 0.7, 0.2, 0.1, 0.05])
    recovered = np.linalg.solve(second, second @ coefficients)
    assert np.allclose(recovered, coefficients)
    fig, panels = plt.subplots(1, 2, figsize=(10, 5.2), dpi=100,
                               facecolor="white")
    fig.subplots_adjust(left=0.075, right=0.98, top=0.77, bottom=0.23,
                        wspace=0.56)
    fig.suptitle("Six directions can still give only three independent equations",
                 y=0.97, fontsize=14, weight="bold", color="#172330")
    fig.text(0.5, 0.90,
             "Original noise-free design matrices: antipodal gradients have identical rows",
             ha="center", fontsize=11, color="#45535d")
    columns = [r"$D_{xx}$", r"$D_{yy}$", r"$D_{zz}$",
               r"$D_{xy}$", r"$D_{xz}$", r"$D_{yz}$"]
    labels = [[r"$x$", r"$y$", r"$z$", r"$-x$", r"$-y$", r"$-z$"],
              [r"$x$", r"$y$", r"$z$", r"$(x+y)/\sqrt{2}$",
               r"$(x+z)/\sqrt{2}$", r"$(y+z)/\sqrt{2}$"]]
    titles = ["Three antipodal pairs: rank 3", "Axes plus mixed directions: rank 6"]
    for panel, matrix, rows, title in zip(panels, [first, second], labels, titles):
        panel.imshow(matrix, cmap="Blues", vmin=0, vmax=1, interpolation="nearest")
        panel.set_xticks(np.arange(6), columns, fontsize=11)
        panel.set_yticks(np.arange(6), rows, fontsize=11)
        panel.set_xticks(np.arange(-0.5, 6), minor=True)
        panel.set_yticks(np.arange(-0.5, 6), minor=True)
        panel.grid(which="minor", color="#a8b5bf", linewidth=0.6)
        panel.tick_params(which="both", length=0)
        panel.set_title(title, fontsize=11, pad=12, color="#172330")
        for i in range(6):
            for j in range(6):
                v = matrix[i, j]
                text = "0" if np.isclose(v, 0) else "1" if np.isclose(v, 1) else r"$\frac{1}{2}$"
                panel.text(j, i, text, ha="center", va="center", fontsize=12,
                           color="white" if v > 0.75 else "#172330")
        for spine in panel.spines.values():
            spine.set_color("#a8b5bf")
    fig.text(0.5, 0.13, "Rows: measurements.  Columns: the six independent symmetric-tensor components.",
             ha="center", fontsize=11, color="#45535d")
    fig.text(0.5, 0.06,
             "Identifiability only: full rank does not establish low uncertainty or an optimal sampling scheme.",
             ha="center", fontsize=10, color="#45535d")
    fig.savefig(Path.cwd() / (Path(__file__).stem + ".png"), dpi=100,
                facecolor="white", transparent=False)
    plt.close(fig)


if __name__ == "__main__":
    main()