import matplotlib
import numpy as np
import scipy.linalg as LA

matplotlib.use('Agg')
import matplotlib.pylab as plt
from matplotlib import rc

from pySDC.core.Collocation import CollBase


def compute_and_plot_specrad(Nnodes, lam):
    """
    Compute and plot the spectral radius of the smoother for different step-sizes

    Args:
        Nnodes: number of collocation nodes
        lam: test parameter representing the spatial problem
    """

    coll = CollBase(Nnodes, 0, 1, quad_type='RADAU-RIGHT')
    Qmat = coll.Qmat[1:, 1:]

    # do LU decomposition of QT (St. Martin's trick)
    QT = coll.Qmat[1:, 1:].T
    [_, _, U] = LA.lu(QT, overwrite_a=True)
    QDmat = U.T

    Nmat = np.zeros((Nnodes, Nnodes))
    Nmat[:, -1] = 1

    Nsteps_list = [64, 256]
    color_list = ['red', 'blue']
    marker_list = ['s', 'o']

    setup_list = zip(Nsteps_list, color_list, marker_list)

    xlist = [0.1**i for i in range(11)]

    rc('font', **{"sans-serif": ["Arial"], "size": 24})
    plt.subplots(figsize=(15, 10))

    for Nsteps, color, marker in setup_list:
        Emat = np.zeros((Nsteps, Nsteps))
        np.fill_diagonal(Emat[1:, :], 1)

        Prho_list = []
        predict_list = []
        for x in xlist:
            mat = np.linalg.inv(np.eye(Nnodes * Nsteps) - x * lam * np.kron(np.eye(Nsteps), QDmat)).dot(
                x * lam * np.kron(np.eye(Nsteps), (Qmat - QDmat)) + np.kron(Emat, Nmat)
            )

            Prho_list.append(max(abs(np.linalg.eigvals(mat))))
            # predict_list.append((1 + x) ** (1.0 - 1.0 / (Nnodes * Nsteps)) * x ** (1.0 / (Nnodes * Nsteps)))
            predict_list.append(x ** (1.0 / (Nsteps)))

            if len(predict_list) > 1:
                print(
                    x,
                    predict_list[-1],
                    Prho_list[-1],
                    Prho_list[-2] / Prho_list[-1],
                    predict_list[-2] / predict_list[-1],
                )

        plt.loglog(
            xlist,
            Prho_list,
            linestyle='-',
            linewidth=3,
            color=color,
            marker=marker,
            markersize=10,
            label='spectral radius, L=' + str(Nsteps),
        )
        plt.loglog(
            xlist,
            predict_list,
            linestyle='--',
            linewidth=2,
            color=color,
            marker=marker,
            markersize=10,
            label='estimate, L=' + str(Nsteps),
        )

    ax = plt.gca()
    ax.invert_xaxis()

    plt.xlabel('time-step size')
    plt.ylabel('spectral radius')
    plt.legend(loc=3, numpoints=1)
    plt.grid()
    plt.ylim([1e-02, 1e01])

    if type(lam) is complex:
        fname = 'data/smoother_specrad_to0_L64+256_M' + str(Nnodes) + 'LU_imag.png'
    else:
        fname = 'data/smoother_specrad_to0_L64+256_M' + str(Nnodes) + 'LU_real.png'
    plt.savefig(fname, transparent=True, bbox_inches='tight')


if __name__ == "__main__":
    compute_and_plot_specrad(Nnodes=3, lam=-1)
    compute_and_plot_specrad(Nnodes=3, lam=1j)
