import matplotlib.pyplot as plt
import numpy as np
from matplotlib.ticker import ScalarFormatter
from pySDC.implementations.problem_classes.FastWaveSlowWave_Scalar import swfw_scalar
from pylab import rcParams

from pySDC.core import CollocationClasses as collclass
from pySDC.core import Hooks as hookclass
from pySDC.core import Level as lvl
from pySDC.core import Step as stepclass
from pySDC.implementations.datatype_classes import mesh, rhs_imex_mesh
from pySDC.implementations.sweeper_classes.imex_1st_order import imex_1st_order as imex

if __name__ == "__main__":
    mvals = np.arange(2, 10)
    kvals = [3, 5, 7]
    lambda_fast = 10j
    lambda_slow = 3j
    stabval = np.zeros((np.size(kvals), np.size(mvals)))

    for i in range(0, np.size(mvals)):
        pparams = {}
        # the following are not used in the computation
        pparams['lambda_s'] = np.array([0.0])
        pparams['lambda_f'] = np.array([0.0])
        pparams['u0'] = 1.0

        swparams = {}
        #    swparams['collocation_class'] = collclass.CollGaussLobatto
        #    swparams['collocation_class'] = collclass.CollGaussLegendre
        swparams['collocation_class'] = collclass.CollGaussRadau_Right
        swparams['num_nodes'] = mvals[i]
        do_coll_update = True

        #
        # ...this is functionality copied from test_imexsweeper. Ideally, it should be available in one place.
        #
        step = stepclass.step(params={})
        L = lvl.level(
            problem_class=swfw_scalar,
            problem_params=pparams,
            dtype_u=mesh,
            dtype_f=rhs_imex_mesh,
            sweeper_class=imex,
            sweeper_params=swparams,
            level_params={},
            hook_class=hookclass.hooks,
            id="stability",
        )
        step.register_level(L)
        step.status.dt = 1.0  # Needs to be 1.0, change dt through lambdas
        step.status.time = 0.0
        u0 = step.levels[0].prob.u_exact(step.status.time)
        step.init_step(u0)
        nnodes = step.levels[0].sweep.coll.num_nodes
        level = step.levels[0]
        problem = level.prob

        QE = level.sweep.QE[1:, 1:]
        QI = level.sweep.QI[1:, 1:]
        Q = level.sweep.coll.Qmat[1:, 1:]
        LHS, RHS = level.sweep.get_scalar_problems_sweeper_mats(lambdas=[lambda_fast, lambda_slow])

        for k in range(0, np.size(kvals)):
            Kmax = kvals[k]
            Mat_sweep = level.sweep.get_scalar_problems_manysweep_mat(nsweeps=Kmax, lambdas=[lambda_fast, lambda_slow])
            if do_coll_update:
                stab_fh = 1.0 + (lambda_fast + lambda_slow) * level.sweep.coll.weights.dot(
                    Mat_sweep.dot(np.ones(nnodes))
                )
            else:
                q = np.zeros(nnodes)
                q[nnodes - 1] = 1.0
                stab_fh = q.dot(Mat_sweep.dot(np.ones(nnodes)))
            stabval[k, i] = np.absolute(stab_fh)

    rcParams['figure.figsize'] = 2.5, 2.5
    fig = plt.figure()
    fs = 8
    plt.plot(mvals, stabval[0, :], 'o-', color='b', label=(r"K=%1i" % kvals[0]))
    plt.plot(mvals, stabval[1, :], 's-', color='r', label=(r"K=%1i" % kvals[1]))
    plt.plot(mvals, stabval[2, :], 'd-', color='g', label=(r"K=%1i" % kvals[2]))
    plt.plot(mvals, 1.0 + 0.0 * mvals, '--', color='k')
    plt.xlabel('Number of nodes M', fontsize=fs)
    plt.ylabel(r'Modulus of stability function $\left| R \right|$', fontsize=fs)
    plt.ylim([0.0, 1.8])
    plt.legend(loc='lower right', fontsize=fs, prop={'size': fs})
    plt.gca().get_xaxis().get_major_formatter().labelOnlyBase = False
    plt.gca().get_xaxis().set_major_formatter(ScalarFormatter())
    plt.show()

#  filename = 'stablimit-M'+str(mvals[0])+'.pdf'
#  fig.savefig(filename, bbox_inches='tight')
#  call(["pdfcrop", filename, filename])
