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
from subaxis import subaxis


def fig10_10() -> None:
    #%%%  data plots in s and y planes

    #%% this requires the subaxis and parseArgs files

    # generate data
    m = 400
    t = np.linspace(0, 10, m)
    w1 = np.pi
    w2 = 1.7 * np.pi
    x = np.zeros((m, 2))
    del x
    # mixing coefficients A=[M11 M12; M21 M22]
    M11 = 1; M12 = 2
    M21 = 1; M22 = -2
    S = np.zeros((m, 2))
    for it in range(m):
        S[it, 0] = np.sin(w1 * t[it])
        S[it, 1] = (1 / 5) * (2.5 * np.sin(w2 * t[it] - 1) + 0.3 * t[it] * (10 - t[it]))

    #%%% center S
    sM = np.sum(S, axis=0) / m
    S = S - sM

    #%% whiten S
    UR, singular_values, VRh = np.linalg.svd(S, full_matrices=False)
    SR = np.diag(singular_values)
    VR = VRh.T
    D = np.array([[np.sqrt(m) / SR[0, 0], 0], [0, np.sqrt(m) / SR[1, 1]]])
    Fac = D @ VR.T
    W1 = np.zeros(m)
    W2 = np.zeros(m)
    for it in range(m):
        W1[it] = Fac[0, 0] * S[it, 0] + Fac[0, 1] * S[it, 1]
        W2[it] = Fac[1, 0] * S[it, 0] + Fac[1, 1] * S[it, 1]

    X = np.zeros((m, 2))
    for it in range(m):
        X[it, 0] = M11 * W1[it] + M12 * W2[it]
        X[it, 1] = M21 * W1[it] + M22 * W2[it]

    xM = np.sum(X, axis=0) / m
    X = X - xM

    XX = X.T @ X / m
    UX, singular_values_x, VXh = np.linalg.svd(XX)
    SX = np.diag(singular_values_x)
    VX = VXh.T
    D5 = np.array([[1 / np.sqrt(SX[0, 0]), 0], [0, 1 / np.sqrt(SX[1, 1])]])

    y = np.zeros((m, 2))
    for it in range(m):
        y[it, :] = D5 @ VX.T @ X[it, :]

    set_color_order()
    fig = matlab_figure([17, 1132, 385, 205], num=1)

    ax = subaxis(1, 2, 1, "MT", 0.04, "MB", 0.25, "MR", -0.03, "ML", 0.06, "P", 0.03)
    assert ax is not None
    ax.plot(W1, W2, ".b", markersize=9)
    # ax.plot(S[:,0],S[:,1],'.b',markersize=10)
    # ax.axis([-2.02,2.02,-2.02,2.02])
    ax.set_aspect("equal", adjustable="box")
    ax.set_xlabel(r"$s_1^*$ -axis")
    ax.set_ylabel(r"$s_2^*$ -axis")
    ax.grid(True)
    ax.set_frame_on(True)
    set_axes_font(ax, 10, "bold")

    ax = subaxis(1, 2, 2)
    assert ax is not None
    ax.plot(y[:, 0], y[:, 1], ".b", markersize=9)
    # ax.plot(y[:,1],y[:,0],'.b',markersize=10)
    ax.set_aspect("equal", adjustable="box")
    # ax.axis([-2,2,-2,2])
    ax.set_xlabel(r"$y_1^*$ -axis")
    ax.set_ylabel(r"$y_2^*$ -axis")
    ax.grid(True)
    ax.set_frame_on(True)
    set_axes_font(ax, 10, "bold")

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


if __name__ == "__main__":
    fig10_10()
