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_13() -> None:
    #%%%  ICA solutions for various max theta values

    #%% this requires the subaxis and parseArgs files

    # 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] = np.sin(w2*t[it]-5)
        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

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

    set_color_order()
    fig = matlab_figure([2, 830, 626, 515])

    # compute and then plot for particular theta values
    theta = np.array([0.4344, 0.9349, 1.4354, 1.9359]) * np.pi
    for ia in range(4):
        V = np.array([[np.cos(theta[ia]), -np.sin(theta[ia])], [np.sin(theta[ia]), np.cos(theta[ia])]])
        s = np.zeros((m, 2))
        for it in range(m):
            s[it, 0] = V[0, 0] * y[it, 0] + V[0, 1] * y[it, 1]
            s[it, 1] = V[1, 0] * y[it, 0] + V[1, 1] * y[it, 1]

        if ia == 0:
            ax = subaxis(4, 2, 1, 1, "MT", 0.04, "MB", 0.07, "MR", -0.005, "ML", 0.02, "P", 0.02)
            assert ax is not None
            # ax.plot(t,S[:,0],'--b',linewidth=1.6)
            ax.axis([0, 10, -2, 2])
            ax.plot(t, s[:, 0], "r", linewidth=1.6)
            say2 = r"$s_1^*$"
            ax.text(4.5, 3.0, say2, fontsize=16, fontweight="bold")
            ax.grid(True)
            ax.set_frame_on(True)
            set_axes_font(ax, 12, "bold")

            ax = subaxis(4, 2, 2, 1)
            assert ax is not None
            # ax.plot(t,S[:,1],'--b',linewidth=1.6)
            ax.axis([0, 10, -2, 2])
            ax.plot(t, s[:, 1], "r", linewidth=1.6)
            say2 = r"$s_2^*$"
            ax.text(4.5, 3.0, say2, fontsize=16, fontweight="bold")
            ax.grid(True)
            ax.set_frame_on(True)
            set_axes_font(ax, 12, "bold")

        elif ia == 1:
            ax = subaxis(4, 2, 1, 2)
            assert ax is not None
            # ax.plot(t,S[:,0],'--b',linewidth=1.6)
            ax.plot(t, s[:, 0], "r", linewidth=1.6)
            ax.axis([0, 10, -2, 2])
            ax.grid(True)
            ax.set_frame_on(True)
            set_axes_font(ax, 12, "bold")

            ax = subaxis(4, 2, 2, 2)
            assert ax is not None
            # ax.plot(t,S[:,1],'--b',linewidth=1.6)
            ax.axis([0, 10, -2, 2])
            ax.plot(t, s[:, 1], "r", linewidth=1.6)
            ax.grid(True)
            ax.set_frame_on(True)
            set_axes_font(ax, 12, "bold")

        elif ia == 2:
            ax = subaxis(4, 2, 1, 3)
            assert ax is not None
            # ax.plot(t,S[:,0],'--b',linewidth=1.6)
            ax.axis([0, 10, -2, 2])
            ax.plot(t, s[:, 0], "r", linewidth=1.6)
            ax.grid(True)
            ax.set_frame_on(True)
            set_axes_font(ax, 12, "bold")

            ax = subaxis(4, 2, 2, 3)
            assert ax is not None
            # ax.plot(t,S[:,1],'--b',linewidth=1.6)
            ax.axis([0, 10, -2, 2])
            ax.plot(t, s[:, 1], "r", linewidth=1.6)
            ax.grid(True)
            ax.set_frame_on(True)
            set_axes_font(ax, 12, "bold")

        elif ia == 3:
            ax = subaxis(4, 2, 1, 4)
            assert ax is not None
            # ax.plot(t,S[:,0],'--b',linewidth=1.6)
            ax.axis([0, 10, -2, 2])
            ax.plot(t, s[:, 0], "r", linewidth=1.6)
            ax.set_xlabel("t-axis")
            ax.grid(True)
            ax.set_frame_on(True)
            set_axes_font(ax, 12, "bold")

            ax = subaxis(4, 2, 2, 4)
            assert ax is not None
            # ax.plot(t,S[:,1],'--b',linewidth=1.6)
            ax.axis([0, 10, -2, 2])
            ax.plot(t, s[:, 1], "r", linewidth=1.6)
            ax.set_xlabel("t-axis")
            ax.grid(True)
            ax.set_frame_on(True)
            set_axes_font(ax, 12, "bold")

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


if __name__ == "__main__":
    fig10_13()
