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_12() -> None:
    #%%%  ICA using s^4 on linear waves: K and the optimal solution

    # generate data
    m = 6000
    t = np.linspace(0, 10, m)
    w1 = np.pi
    w2 = 1.7 * np.pi
    x = np.zeros((m, 2))
    del x
    # mixing coefficients
    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

    #%% SVD of 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]
    del W1, W2

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

    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, :]

    # compute (and then plot) K as function of theta
    ntheta = 1000
    theta = np.linspace(0, 2 * np.pi, ntheta)
    K = np.zeros(ntheta)
    for ia in range(ntheta):
        V = np.array([[np.cos(theta[ia]), -np.sin(theta[ia])], [np.sin(theta[ia]), np.cos(theta[ia])]])
        ssk = y @ V.T
        kks = ssk**4 / m
        suk = np.sum(kks, axis=0)
        K[ia] = abs(suk[0] - 3) + abs(suk[1] - 3)

    # find max's
    thetaMM = []
    for ia in range(1, ntheta - 1):
        if K[ia - 1] < K[ia] and K[ia + 1] < K[ia]:
            thetaMM.append(theta[ia])
    thetaMM = np.asarray(thetaMM)
    print(thetaMM / np.pi)
    theta0 = thetaMM[2] / np.pi
    print(theta0)

    set_color_order()
    fig = matlab_figure([25, 1115, 658, 230])
    ax = subaxis(1, 1, 1, 1, "MT", 0.01, "MB", 0.18, "MR", -0.02, "ML", 0.06, "P", 0.04)
    assert ax is not None

    ax.plot(theta, K, linewidth=1.6)
    ax.axis([0, 2 * np.pi, 1.5, 2.3])
    ax.set_xlabel(r"$\theta$-axis")
    ax.set_ylabel("K -axis")

    ax.set_xticks([0, np.pi / 2, np.pi, 3 * np.pi / 2, 2 * np.pi])
    # ax.set_yticks([0,1,2,3])
    ax.set_xticklabels(["0", r"$\pi/2$", r"$\pi$", r"$3\pi/2$", r"$2\pi$"])
    # ax.set_yticklabels(['0',' ',' ',' '])

    ax.grid(True)
    ax.set_frame_on(True)
    set_axes_font(ax, 14, "bold")

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


if __name__ == "__main__":
    fig10_12()
