"""
This script has been used to compute the Dirichlet results of the paper

    Efficient spectral-Galerkin methods for second-order equations using different Chebyshev bases

The results have been computed using Python 3.9 and Shenfun 3.1.1.

The generalized Chebyshev-Tau results are computed with dedalus,
and are as such not part of this script.

"""
import sympy as sp
import numpy as np
import scipy.sparse.linalg as lin
import array_to_latex as a2l
from time import time

x = sp.Symbol('x', real=True)

fe = {}
rnd = {}
func = {}

def matvec(u_hat, f_hat, A, B, alpha, method):
    """Compute matrix vector product

    Parameters
    ----------
    u_hat : Function
        The solution array
    f_hat : Function
        The right hand side array
    A : SparseMatrix
        The stiffness matrix
    B : SparseMatrix
        The mass matrix
    alpha : number
        The weight of the mass matrix
    method : int
        The chosen method

    """
    from shenfun import chebyshev, la
    if method == 1:
        if alpha == 0:
            A.scale *= -1
            f_hat = A.matvec(u_hat, f_hat)
            A.scale *= -1
        else:
            sol = chebyshev.la.Helmholtz(A, B, -1, alpha)
            f_hat = sol.matvec(u_hat, f_hat)
    else:
        if alpha == 0:
            A.scale *= -1
            f_hat = A.matvec(u_hat, f_hat)
            A.scale *= -1
        else:
            M = alpha*B - A
            f_hat = M.matvec(u_hat, f_hat)
    return f_hat

def get_solver(A, B, alpha, method):
    """Return optimal solver for given method

    Parameters
    ----------
    A : SparseMatrix
        The stiffness matrix
    B : SparseMatrix
        The mass matrix
    alpha : number
        The weight of the mass matrix
    method : int
        The chosen method
    """
    from shenfun import chebyshev, la
    if method == 2:
        if alpha == 0:
            sol = la.TDMA(A*(-1))
        else:
            sol = la.PDMA(alpha*B - A)
    elif method == 1:
        if alpha == 0:
            A.scale = -1
            sol = chebyshev.la.ADD_Solve(A)
        else:
            sol = chebyshev.la.Helmholtz(A, B, -1, alpha)
    elif method in (0, 3, 4):
        if alpha == 0:
            sol = chebyshev.la.TwoDMA(A*(-1))
        else:
            sol = chebyshev.la.FDMA(alpha*B-A)
    elif method == 5:
        if alpha == 0:
            AA = A*(-1)
            sol = AA.solve
        else:
            sol = la.TDMA(alpha*B-A)

    else:
        raise NotImplementedError
    return sol

def solve(f_hat, u_hat, A, B, alpha, method):
    """Solve (alpha*B-A)u_hat = f_hat

    Parameters
    ----------
    f_hat : Function
        The right hand side array
    u_hat : Function
        The solution array
    A : SparseMatrix
        The stiffness matrix
    B : SparseMatrix
        The mass matrix
    alpha : number
        The weight of the mass matrix
    method : int
        The chosen method

    """
    from shenfun import extract_bc_matrices, Function
    if isinstance(B, list):
        u_hat.set_boundary_dofs()
        bc_mat = extract_bc_matrices([B])
        B = B[0]
        w0 = Function(u_hat.function_space())
        f_hat -= alpha*bc_mat[0].matvec(u_hat, w0)

    sol = get_solver(A, B, alpha, method)
    if method == 1 and alpha != 0:
        u_hat = sol(u_hat, f_hat)
    else:
        u_hat = sol(f_hat, u_hat)
    return u_hat

def main(N, method=0, alpha=0, returntype=0):
    from shenfun import FunctionSpace, TrialFunction, TestFunction, \
        inner, div, grad, chebyshev, SparseMatrix, Function, Array
    global fe
    basis = {0: ('ShenDirichlet', 'Heinrichs'),
             1: ('ShenDirichlet', 'ShenDirichlet'),
             2: ('Heinrichs', 'Heinrichs'),
             3: ('DirichletU', 'ShenDirichlet'),
             4: ('Orthogonal', 'ShenDirichlet'),    # Quasi-Galerkin
             5: ('ShenDirichlet', 'ShenDirichlet'), # Legendre
             }

    test, trial = basis[method]
    if returntype == 2:
        ue = sp.sin(100*sp.pi*x)

    family = 'C' if method < 5 else 'L'
    kw = {}
    scaled = True if method in (0, 5) else False
    if scaled:
        kw['scaled'] = True
    ST = FunctionSpace(N, family, basis=test, **kw)
    TS = FunctionSpace(N, family, basis=trial, **kw)
    wt = {0: 1, 1: 1, 2: 1, 3: 1-x**2, 4: 1, 5: 1}[method]
    u = TrialFunction(TS)
    v = TestFunction(ST)
    A = inner(v*wt, div(grad(u)))
    B = inner(v*wt, u)

    if method == 4:
        # Quasi
        Q2 = chebyshev.quasi.QIGmat(N)
        A = Q2*A
        B = Q2*B
    if method == 3:
        k = np.arange(N-2)
        K = SparseMatrix({0: 1/((k+1)*(k+2)*2)}, (N-2, N-2))
        A[0] *= K[0]
        A[2] *= K[0][:-2]
        B[-2] *= K[0][2:]
        B[0] *= K[0]
        B[2] *= K[0][:-2]
        B[4] *= K[0][:-4]

    if returntype == 0:
        M = alpha*B.diags()-A.diags()
        con = np.linalg.cond(M.toarray())

    elif returntype == 1:
        # Use rnd to get the same random numbers for all methods
        buf = rnd.get(N, np.random.random(N))
        if not N in rnd:
            rnd[N] = buf
        v = Function(TS, buffer=buf)
        v[-2:] = 0
        u_hat = Function(TS)
        f_hat = Function(TS)
        f_hat = matvec(v, f_hat, A, B, alpha, method)
        u_hat = solve(f_hat, u_hat, A, B, alpha, method)
        con = np.abs(u_hat-v).max()

    elif returntype == 2:
        fe = alpha*ue - ue.diff(x, 2)
        f_hat = Function(ST)
        fj = Array(ST, buffer=fe)
        if wt != 1:
            fj *= np.sin((np.arange(N)+0.5)*np.pi/N)**2
        f_hat = ST.scalar_product(fj, f_hat, fast_transform=True)

        if method == 4:
            f_hat[:-2] = Q2.diags('csc')*f_hat

        if method == 3:
            f_hat[:-2] *= K[0]

        sol = get_solver(A, B, alpha, method)
        u_hat = Function(TS)
        u_hat = solve(f_hat, u_hat, A, B, alpha, method)
        uj = Array(TS)
        uj = TS.backward(u_hat, uj, fast_transform=True)
        ua = Array(TS, buffer=ue)
        con = np.sqrt(inner(1, (uj-ua)**2))

    return con

if __name__ == '__main__':
    import matplotlib.pyplot as plt
    import argparse
    import os
    import sys

    parser = argparse.ArgumentParser(description='Solve the Helmholtz problem with Dirichlet boundary conditions')
    parser.add_argument('--return_type', action='store', type=int, required=True)
    parser.add_argument('--include_legendre', action='store_true')
    parser.add_argument('--verbose', '-v', action='count', default=0)
    parser.add_argument('--plot', action='store_true')
    parser.add_argument('--numba', action='store_true')
    args = parser.parse_args()

    if args.numba:
        try:
            import numba
            os.environ['SHENFUN_OPTIMIZATION'] = 'NUMBA'
        except ModuleNotFoundError:
            os.warning('Numba not found - using Cython')

    cond = []
    if args.return_type == 2:
        N = (2**4,2**6, 2**8, 2**12, 2**16, 2**20)
    elif args.return_type == 1:
        N = (2**4, 2**12, 2**20)
    else:
        N = (32, 64, 128, 256, 512, 1024, 2048)
    M = 6 if args.include_legendre else 5
    alphas = (0, 1000)

    if args.return_type in (0, 2):
        for alpha in alphas:
            cond.append([])
            if args.verbose > 0:
                print('alpha =', alpha)
            for basis in range(M): # To include Legendre use --include_legendre (takes hours for N=2**20)
                if args.verbose > 1:
                    print('Method =', basis)
                cond[-1].append([])
                for n in N:
                    if args.verbose > 2:
                        print('N =', n)
                    cond[-1][-1].append(main(n, basis, alpha, args.return_type))
        linestyle = {0: 'solid', 1: 'dashed', 2: 'dotted'}
        for i in range(len(cond)):
            plt.loglog(N, cond[i][0], 'b',
                       N, cond[i][1], 'r',
                       N, cond[i][2], 'k',
                       N, cond[i][3], 'm',
                       N, cond[i][4], 'y',
                       linestyle=linestyle[i])
            if args.include_legendre:
                plt.loglog(N, cond[i][5], 'y', linestyle=linestyle[i])
            a2l.to_ltx(np.array(cond)[i], frmt='{:6.2e}', print_out=True, mathform=False)
    else:
        for basis in range(M):
            cond.append([])
            if args.verbose > 1:
                print('Method =', basis)

            for alpha in alphas:
                if args.verbose > 0:
                    print('alpha =', alpha)
                for n in N:
                    if args.verbose > 2:
                        print('N =', n)
                    cond[-1].append(main(n, basis, alpha, args.return_type))

        a2l.to_ltx(np.array(cond), frmt='{:6.2e}', print_out=True, mathform=False)

    if args.plot:
        plt.show()
