from __future__ import annotations

import time

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_14() -> None:
    #%%%  ICA using s^4 on linear waves: K and the optimal solution
    #%%%  random data
    start_time = time.perf_counter()
    # generate data
    tmax = 10
    NN = 75
    theta0 = np.zeros(NN)
    theta00 = np.zeros(NN)
    # IR=number of cases
    IR = 10
    # IRR=the paricular case plotted
    IRR = np.random.randint(1, IR + 1)
    # m=number of data points used (this is N in text)
    m = 3
    num = np.zeros(NN)
    for ic in range(NN):
        # if m<40:
        #     m=m+1
        # else:
        #     m=m+40
        m = m + 2
        num[ic] = m
        val = 0
        for ir in range(1, IR + 1):
            t = tmax * np.random.rand(m)
            w1 = np.pi
            w2 = 1.7 * np.pi
            x = np.zeros((m, 2))
            del x

            # mixing coefficients
            M11 = 2; M12 = 1
            M21 = 3; M22 = -2
            S = np.zeros((m, 2))
            for it in range(m):
                S[it, 0] = np.sin(w1 * t[it])
                S[it, 1] = np.sin(w2 * t[it] - 5)

            #%%% 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]
            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 = X @ VX @ D5

            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)

            iz = 0
            ia_found = ntheta - 2
            for ia in range(1, ntheta - 1):
                if K[ia - 1] < K[ia] and K[ia + 1] < K[ia]:
                    iz = iz + 1
                    # thetaMM[iz]=theta[ia]
                    ia_found = ia
                    break
            del iz

            if ir == IRR:
                theta00[ic] = theta[ia_found]
            val = val + theta[ia_found]
        theta0[ic] = val / IR

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

    ax.plot(num, theta0, linewidth=1.6)
    ax.plot(num, theta00, "--r", linewidth=2)
    ax.axis([0, 150, 0, np.pi / 2])
    ax.set_xlabel("N-axis")
    ax.set_ylabel("Theta")
    ax.set_yticks([0, np.pi / 4, np.pi / 2])
    ax.set_yticklabels(["0", r"$\pi/4$", r"$\pi/2$"])
    ax.grid(True)
    ax.set_frame_on(True)
    set_axes_font(ax, 14, "bold")

    print(f"Elapsed time is {time.perf_counter() - start_time:.6f} seconds.")

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


if __name__ == "__main__":
    fig10_14()
