from __future__ import annotations

import matplotlib.pyplot as plt
import numpy as np

from plot_helpers import finish_figure, matlab_figure, set_axes_font, set_color_order


def fig10_3() -> None:
    #  the data is taken to be within the ellipsoid x^2+4y^2+8z^2=1

    n = 200

    # ellipsoidal data
    x = np.zeros(n)
    y = np.zeros(n)
    z = np.zeros(n)
    for i in range(n):
        x[i] = 2 * np.random.rand() - 1
        y[i] = (2 * np.random.rand() - 1) * 0.5 * np.sqrt(1 - x[i] ** 2)
        z[i] = (2 * np.random.rand() - 1) * 0.5 * np.sqrt(1 - x[i] ** 2 - 4 * y[i] ** 2) / 3

    #  rotate data
    theta = np.pi / 3; phi = -np.pi / 5
    Ry = np.array([[np.cos(theta), 0, np.sin(theta)], [0, 1, 0], [-np.sin(theta), 0, np.cos(theta)]])
    Rz = np.array([[np.cos(phi), -np.sin(phi), 0], [np.sin(phi), np.cos(phi), 0], [0, 0, 1]])
    print(Ry)
    print(Rz)
    X = np.zeros((n, 3))
    Q = np.zeros((n, 3))
    for i in range(n):
        #    [x(i) y(i) z(i)]'
        X[i, :] = Ry @ Rz @ np.array([x[i], y[i], z[i]])
        Q[i, :] = Ry @ Rz @ np.array([0, y[i], z[i]])
    xx = 1.2
    P1 = Ry @ Rz @ np.array([xx, 0, 0])
    P2 = Ry @ Rz @ np.array([-xx, 0, 0])
    yy = 1.2
    Q1 = Ry @ Rz @ np.array([0, yy, 0])
    Q2 = Ry @ Rz @ np.array([0, -yy, 0])

    set_color_order()
    fig1 = matlab_figure([8, 925, 560, 420], num=1)
    ax = fig1.add_subplot(111, projection="3d")
    ax.scatter(
        X[:, 0], X[:, 1], X[:, 2],
        edgecolors="k", facecolors=[3 / 255, 192 / 255, 60 / 255], linewidths=1.2,
    )

    ax.plot([P1[0], P2[0]], [P1[1], P2[1]], [P1[2], P2[2]], color=[255 / 255, 8 / 255, 0 / 255], linewidth=2.4)
    ax.set_box_aspect((1, 1, 1))
    ax.set_xlim(-1, 1); ax.set_ylim(-1, 1); ax.set_zlim(-1, 1)
    ax.view_init(elev=10, azim=-46)
    ax.set_xticks([-1, 0, 1])
    ax.set_yticks([-1, 0, 1])
    ax.set_zticks([-1, 0, 1])
    ax.set_xlabel(r"$p_1$")
    ax.set_ylabel(r"$p_2$")
    ax.set_zlabel(r"$p_3$")
    set_axes_font(ax, 16, "bold")
    # ax.GridAlpha=1
    # H=gca
    # H.LineWidth=1
    ax.tick_params(width=0.5)
    ax.plot([-1, -1], [-1, 1], [-1, -1], color="k", linewidth=1.5)
    ax.plot([-1, 1], [-1, -1], [-1, -1], color="k", linewidth=1.5)
    ax.plot([-1, -1], [1, 1], [-1, 1], color="k", linewidth=1.5)
    # box on

    # fig1.savefig('/Users/mark/Desktop/dat1.eps')

    fig2 = matlab_figure([11, 416, 560, 420], num=2)
    ax = fig2.add_subplot(111)
    ax.plot(
        Q[:, 1], Q[:, 2], "o", markersize=7,
        markeredgecolor="k", markerfacecolor=[3 / 255, 192 / 255, 60 / 255], linewidth=1.2,
    )
    ax.plot([Q1[1], Q2[1]], [Q1[2], Q2[2]], "--k", linewidth=2.4)
    ax.plot(0, 0, ".r", markersize=30)
    ax.set_aspect("equal", adjustable="box")
    ax.grid(True)
    ax.axis([-0.5, 0.5, -0.5, 0.5])
    ax.set_xticks([-0.5, 0, 0.5])
    ax.set_yticks([-0.5, 0, 0.5])
    set_axes_font(ax, 16, "bold")
    ax.tick_params(width=0.8)

    # fig2.savefig('/Users/mark/Desktop/dat2.eps')
    finish_figure(fig1)
    finish_figure(fig2)
    plt.show()


if __name__ == "__main__":
    fig10_3()
