from __future__ import print_function

import matplotlib.pyplot as plt
import numpy as np
import pySDC.core.deprecated.PFASST_stepwise as mp

from pySDC.implementations.problem_classes.FastWaveSlowWave_Scalar import swfw_scalar
from pySDC.implementations.sweeper_classes.imex_1st_order import imex_1st_order
from pySDC.core import CollocationClasses as collclass
from pySDC.core import Log


if __name__ == "__main__":

    # set global logger (remove this if you do not want the output at all)
    logger = Log.setup_custom_logger('root')

    num_procs = 1

    N_s = 100
    N_f = 125

    # This comes as read-in for the level class
    lparams = {}
    lparams['restol'] = 1e-12

    sparams = {}
    sparams['maxiter'] = 4

    # This comes as read-in for the problem class
    pparams = {}
    pparams['lambda_s'] = 1j * np.linspace(0.0, 4.0, N_s)
    pparams['lambda_f'] = 1j * np.linspace(0.0, 8.0, N_f)
    pparams['u0'] = 1

    # Fill description dictionary for easy hierarchy creation
    description = {}
    description['problem_class'] = swfw_scalar
    description['problem_params'] = pparams
    description['collocation_class'] = collclass.CollGaussLobatto
    description['num_nodes'] = [3]
    description['do_LU'] = False
    description['sweeper_class'] = imex_1st_order
    description['level_params'] = lparams

    # quickly generate block of steps
    MS = mp.generate_steps(num_procs, sparams, description)

    # setup parameters "in time"
    t0 = 0
    dt = 1.0
    Tend = 1.0

    P = MS[0].levels[0].prob
    uinit = P.u_exact(t0)

    print('Init:', uinit.values)

    # call main function to get things done...
    uend, stats = mp.run_pfasst(MS, u0=uinit, t0=t0, dt=dt, Tend=Tend)

    uex = P.u_exact(Tend)

    fig = plt.figure(figsize=(4, 4))
    # pcol = plt.pcolor(pparams['lambda_s'].imag, pparams['lambda_f'].imag, np.absolute(uend.values).T, vmin=0.99, vmax=1.01)
    pcol = plt.pcolor(
        pparams['lambda_s'].imag, pparams['lambda_f'].imag, np.absolute(uend.values).T, vmin=1.0, vmax=2.0
    )
    # plt.pcolor(np.imag(uend.values))
    pcol.set_edgecolor('face')
    plt.plot([0, 4], [0, 4], color='w', linewidth=2.5)
    # plt.colorbar()
    plt.gca().set_xticks([0.0, 1.0, 2.0, 3.0])
    plt.xlabel('$\Delta t \lambda_{slow}$', fontsize=18, labelpad=0.0)
    plt.ylabel('$\Delta t \lambda_{fast}$', fontsize=18)

    # plt.show()
    fig.savefig(
        'sdc-fwsw-stability-K' + str(sparams['maxiter']) + '-M' + str(description['num_nodes'][0]) + '.pdf',
        bbox_inches='tight',
    )
    print(
        'error at time %s: %s'
        % (Tend, np.linalg.norm(uex.values - uend.values, np.inf) / np.linalg.norm(uex.values, np.inf))
    )
