from __future__ import annotations

import matplotlib.pyplot as plt
import numpy as np

from arrow3 import arrow3
from plot_helpers import finish_figure, matlab_figure, set_axes_font, set_color_order


def fig10_4() -> None:
    #  this requires the file arrow3.m

    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])
    del Q, P1, P2, Q1, Q2

    # calculate svd
    U, singular_values, Vh = np.linalg.svd(X, full_matrices=False)
    V = Vh.T
    print(V)

    set_color_order()
    fig = matlab_figure([8, 925, 560, 420])
    ax = fig.add_subplot(111, projection="3d")

    # ax.scatter(X[:,0],X[:,1],X[:,2],edgecolors='k',facecolors=[0,0.75,0.75])
    acol = 0 / 255; bcol = 200 / 255; ccol = 0
    ax.scatter(X[:, 0], X[:, 1], X[:, 2], edgecolors="k", facecolors=[acol, bcol, ccol], linewidths=1)

    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=110)
    ax.view_init(elev=10, azim=35)
    # surf(A1,A2,F)
    # axis([xa xb ya yb 0 Emax])
    ax.set_xticks([-1, 0, 1])
    ax.set_yticks([-1, 0, 1])
    ax.set_zticks([-1, 0, 1])
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    ax.set_zlabel("z")

    p0 = [0, 0, 0]
    a1 = 1.4 * np.sign(V[2, 0])
    p1 = [a1 * V[0, 0], a1 * V[1, 0], a1 * V[2, 0]]
    a2 = 0.9 * np.sign(V[2, 1])
    p2 = [a2 * V[0, 1], a2 * V[1, 1], a2 * V[2, 1]]
    a3 = -0.8 * np.sign(V[2, 2])
    p3 = [a3 * V[0, 2], a3 * V[1, 2], a3 * V[2, 2]]
    arrow3(p0, p1, "3_v", [1.4, 1.4, 1.4], ax=ax)
    arrow3(p0, p2, "3_v", [1.4, 1.4, 1.4], ax=ax)
    arrow3(p0, p3, "3_v", [1.4, 1.4, 1.4], ax=ax)

    # H=gca
    # H.LineWidth=1
    ax.tick_params(width=1)
    # box on

    set_axes_font(ax, 16, "bold")

    # fig.savefig('/Users/mark/Desktop/dataA.eps')
    finish_figure(fig)
    plt.show()


if __name__ == "__main__":
    fig10_4()
