# -*- coding: utf-8 -*-
#
# Set Cython compiler directives. This section must appear before any code!
#
# For available directives, see:
#
# http://docs.cython.org/en/latest/src/reference/compilation.html
#
# cython: wraparound  = False
# cython: boundscheck = False
# cython: cdivision   = True
#
"""Integrators based on the Galerkin method, to update w by one timestep.

Also provided is a DataManager class that loads in the data generated by pydgq.utils.precalc,
required to use classes derived from GalerkinIntegrator.

Do not instantiate DataManager directly; instead, call the init() function exported by this module.
"""
# TODO: might be cleaner if the user could instantiate DataManager directly (would trivially allow multiple instances with different settings)

from __future__ import division, print_function, absolute_import

import sys

try:
    import cPickle as pickle  # Python 2.7
except ImportError:
    import pickle  # Python 3.x

import pkg_resources  # for accessing installed data file (pydgq_data.bin)

import numpy as np

import  pylu.dgesv as dgesv    # Cython-based LU decomposition and linear equation system solver, callable from inside nogil blocks -- Python interface (using np.arrays)
cimport pylu.dgesv as dgesv_c  # -"- -- Cython interface (using raw pointers, explicit sizes)

from pydgq.solver.types cimport DTYPE_t, RTYPE_t
from pydgq.solver.types import DTYPE, RTYPE
from pydgq.solver.kernel_interface cimport KernelBase
from pydgq.solver.integrator_interface cimport ImplicitIntegrator
from pydgq.solver.compsum cimport accumulate


cdef class GalerkinIntegrator(ImplicitIntegrator):

    def __init__(self, str name, KernelBase rhs, int maxit):
        """def __init__(self, str name, KernelBase rhs, int maxit):

Base class for integrators based on the Galerkin method.

Handles data array allocation and provides helper methods common to all Galerkin integrators.

Parameters:
    as in ancestor (pydgq.solver.integrator_interface.ImplicitIntegrator).
"""
        global datamanager

        if datamanager is None:
            raise RuntimeError("%s: init() must be called first." % name)
        if not datamanager.available:
            raise RuntimeError("%s: Cannot use Galerkin integrators because the auxiliary class did not initialize." % name)
        if datamanager.method != name:
            raise RuntimeError("%s: Trying to integrate with %s, but the auxiliary class has been initialized for %s." % (name, name, datamanager.method))

        # super
        ImplicitIntegrator.__init__(self, name, rhs, maxit)

        self.n_time_dofs  = datamanager.n_time_dofs
        self.n_quad       = datamanager.rule   # number of quadrature points (Gauss-Legendre integration points)
        self.n_vis        = datamanager.nt_vis

        # allocate instance-specific arrays
        #
        my_storage = datamanager.allocate_storage(solver_id=id(self), n_space_dofs=rhs.n)

        # retrieve instance arrays
        #
        cdef DTYPE_t[:,:,::1] g          = my_storage["g"]      # effective load vector, for each space DOF, for each time DOF, at each integration point
        cdef DTYPE_t[:,::1]   b          = my_storage["b"]      # right-hand sides (integral, over the timestep, of g*psi)
        cdef DTYPE_t[:,::1]   u          = my_storage["u"]      # Galerkin coefficients (unknowns)
        cdef DTYPE_t[:,::1]   uprev      = my_storage["uprev"]  # Galerkin coefficients from previous iteration
        cdef DTYPE_t[:,::1]   uass       = my_storage["uass"]   # u, assembled for integration
        cdef DTYPE_t[::1]     ucorr      = my_storage["ucorr"]  # correction for compensated summation in assemble() (for integration)
        cdef DTYPE_t[:,::1]   uvis       = my_storage["uvis"]   # u, assembled for visualization
        cdef DTYPE_t[::1]     ucvis      = my_storage["ucvis"]  # correction for compensated summation in assemble() (for visualization)
        cdef DTYPE_t[::1]     wrk_arr    = my_storage["wrk"]    # work space for dG(), cG()

        # get raw pointers to array data (needed by our C routines)
        #
        # FIXME/TODO: This is brittle, since the array shape information is carried separately from the data; and only implicitly
        # FIXME/TODO: by convention between this class and DataManager. Could also store the arrays here?
        #
        self.g       = &g[0,0,0]
        self.b       = &b[0,0]
        self.u       = &u[0,0]
        self.uprev   = &uprev[0,0]
        self.uass    = &uass[0,0]
        self.ucorr   = &ucorr[0]
        self.uvis    = &uvis[0,0]
        self.ucvis   = &ucvis[0]
        self.wrk     = &wrk_arr[0]

        # retrieve global arrays
        #
        # NOTE: our "t" = "x" in precalc.py; it is a general support module for Lobatto basis functions, not specific to its ODE solver use case here.
        #
        cdef RTYPE_t[:,::1]   LU         = datamanager.LU       # LU decomposed mass matrix (packed format), for one space DOF, shape (n_time_dofs, n_time_dofs)
        cdef int[::1]         p          = datamanager.p        # row permutation information, length n_time_dofs
        cdef int[::1]         mincols    = datamanager.mincols  # band information for L, length n_time_dofs
        cdef int[::1]         maxcols    = datamanager.maxcols  # band information for U, length n_time_dofs
        cdef RTYPE_t[::1]     qw         = datamanager.integ_w  # quadrature weights (Gauss-Legendre)
        cdef RTYPE_t[:,::1]   psi        = datamanager.integ_y  # basis function values at the quadrature points, qy[j,i] is N[j]( tquad[i] )
        cdef RTYPE_t[:,::1]   psivis     = datamanager.vis_y    # basis function values at the visualization points, qy[j,i] is N[j]( tvis[i] )
        cdef RTYPE_t[::1]     tvis       = datamanager.vis_x    # time values at the visualization points (on the reference element [-1,1])
        cdef RTYPE_t[::1]     tquad      = datamanager.integ_x  # Gauss-Legendre points of the chosen rule, in (-1,1)

        # get raw pointers to array data
        #
        self.LU      = &LU[0,0]
        self.p       = &p[0]
        self.mincols = &mincols[0]
        self.maxcols = &maxcols[0]
        self.qw      = &qw[0]
        self.psi     = &psi[0,0]
        self.psivis  = &psivis[0,0]
        self.tvis    = &tvis[0]
        self.tquad   = &tquad[0]

        # map tvis and tquad: [-1,1] --> [0,1] (reference element --> offset within timestep)
        cdef unsigned int j
        for j in range(self.n_vis):
            tvis[j] += 1.0
            tvis[j] *= 0.5
        for j in range(self.n_quad):
            tquad[j] += 1.0
            tquad[j] *= 0.5

    def __del__(self):
        global datamanager
        datamanager.free_storage(solver_id=id(self))


    # Assemble Galerkin series at given points inside the timestep.
    #
    # The points are specified implicitly, by evaluating the basis functions at the desired points (see the parameter psi, below).
    # This is parametrized to allow for different use cases (integration, and visualization/reporting).
    #
    # Space DOFs are independent, each has its own Galerkin series. This assembles them all in one go.
    #
    # psi:    rank-2 array, size [n_time_dofs, n_points]:     values of basis functions, evaluated at each of the points t[i] (at which the series is to be assembled); psi[j,i] is psi[j]( t[i] )
    # uass:   rank-2 array, size [n_points, n_space_dofs]:    output, values of u; uass[j,i] is u_i( t[j] )
    # ucorr:  rank-1 array, size [n_points]:                  work area for correction terms in compensated summation
    #
    # Additionally, this automatically uses self.u:
    #
    # self.u: rank-2 array, size [n_space_dofs, n_time_dofs]: Galerkin coefficients; u[j,i] is the coefficient of psi[i] for solution component u_j(t)
    #
    #                         basis         output         summ. corr. wrk
    cdef void assemble( self, RTYPE_t* psi, DTYPE_t* uass, DTYPE_t* ucorr, int n_points ) nogil:
        cdef unsigned int i,j,k
        cdef int n_space_dofs = self.rhs.n
        cdef int n_time_dofs  = self.n_time_dofs

#        # naive summation
#        for k in range(n_space_dofs):
#    
#            # zero out the assembled values for this space DOF
#            for i in range(n_points):
#                uass[ i*n_space_dofs + k ] = 0.0
#    
#            # sum contributions from each basis function
#            for j in range(n_time_dofs):
#                for i in range(n_points):
#                    uass[ i*n_space_dofs + k ] += self.u[ k*n_time_dofs + j ] * psi[ j*n_points + i ]

        # compensated summation
        cdef DTYPE_t* ps
        cdef DTYPE_t* pc
        for k in range(n_space_dofs):

            # Zero out the assembled values for this space DOF.
            #
            # Also zero out the correction array. Note that since we process one space DOF at a time, we only need one correction array.
            # The same position uass[i,k] gets written multiple times during the inner loops, but each time k changes, a new summation begins
            # (at that point, the sums for any previous k are not updated any more).
            #
            for i in range(n_points):
                uass[ i*n_space_dofs + k ] = 0.0
                ucorr[ i ] = 0.0

            # Sum contributions from each basis function.
            for i in range(n_points):
                # Compute the pointers only once for each i.
                ps = &uass[ i*n_space_dofs + k ]  # sum (accumulator)
                pc = &ucorr[ i ]                  # correction
                for j in range(n_time_dofs):
                    accumulate( ps, pc, self.u[ k*n_time_dofs + j ] * psi[ j*n_points + i ] )


    # Assemble Galerkin series at the end of the timestep.
    #
    # uass:   rank-1 array, size [n_space_dofs]:    output, values of u; uass[i] is u_i( t )
    #
    # Compared to the general-purpose assemble():
    #   - Only N[1] (the linear basis function associated with the endpoint) is nonzero at the endpoint.
    #   - Its value there is exactly 1, so we can just take the corresponding Galerkin coefficient from u[] as the function value
    #     (no need to actually sum anything).
    #   - There are always at least 2 basis functions (the linear ones), so we don't need to check whether this basis function exists.
    #   - Here n_points is always 1.
    #
    cdef void final_value( self, DTYPE_t* uass ) nogil:
        cdef unsigned int k
        cdef int n_space_dofs = self.rhs.n
        cdef int n_time_dofs  = self.n_time_dofs
        for k in range(n_space_dofs):
            # here only one visualization point, so uass[ n_space_dofs*0 + k ] --> uass[ k ]
            uass[ k ] = self.u[ k*n_time_dofs + 1 ]


    # Integrate a function over the timestep by applying a quadrature.
    #
    # Compensated summation is used to obtain maximal accuracy.
    #
    # funcvals:    array of values of the function to be integrated, evaluated at the quadrature points (rank-1 or C order, length self.n_quad)
    # dt:          size of timestep (will be used for scaling the result)
    #
    # Additionally, this uses:
    #
    # self.qw:     array of quadrature weights (rank-1 or C order); length of funcvals must match length of self.qw
    # self.n_quad: number of items in funcvals and qw
    #
    cdef DTYPE_t do_quadrature( self, DTYPE_t* funcvals, RTYPE_t dt ) nogil:
        cdef unsigned int i

#        # naive summation (not good for problems sensitive to initial conditions, such as vibration problems with low damping)
#        cdef DTYPE_t s
#        s = 0.0
#        for i in range(self.n_quad):
#            s += self.qw[i]*funcvals[i]

        # compensated summation
        cdef DTYPE_t s, c
        s = self.qw[0]*funcvals[0]  # first term in sum; no correction by Kahan algorithm (if we wanted to correct, we should take the roundoff from this multiplication, instead)
        c = 0.0
        for i in range(1,self.n_quad):
            accumulate( &s, &c, self.qw[i]*funcvals[i] )

        # Scale the result to account for the length of the timestep.
        #
        # This amounts to a change of variable in the integral, as is commonly used in finite element methods.
        # ( For a refresher, see https://en.wikipedia.org/wiki/Integration_by_substitution#Substitution_for_single_variable )
        #
        # We use a simple affine coordinate mapping so that the Jacobian is constant across the timestep.
        #
        # The reference interval is [-1,1], so its length is 2; the target length is dt; thus we must scale by dt/2.
        #
        # Alternatively, formally speaking, the affine mapping of tau: [-1,1] --> t: [a,b] is given by
        #
        #    t = a + (b - a) ( tau - (-1) ) / ( 1 - (-1) )
        #      = a + (b - a) ( tau + 1 ) / 2
        #      = a + (1/2) (b - a) ( tau + 1 )
        #
        # Hence by differentiating both sides with respect to tau,
        #
        #    dt/dtau = (b - a) / 2
        #
        # Noting that in our application, (b - a) = dt, we obtain the result.
        #
        s *= dt/2.0

        return s


cdef class DG(GalerkinIntegrator):
    def __init__(self, KernelBase rhs, int maxit):
        """def __init__(self, KernelBase rhs, int maxit):

Time-discontinuous Galerkin.
"""
        # super
        GalerkinIntegrator.__init__(self, name="dG", rhs=rhs, maxit=maxit)

    # Integrate w over one timestep.
    #
    # When call() returns, the Galerkin coefficients in self.u correspond to the
    # behavior of w over the last processed timestep. This is used in visualization
    # to support "interp".
    #
    #    # array shapes and types:
    #
    #    # instance arrays (see galerkin.DataManager.allocate_storage())
    #    cdef DTYPE_t[:,:,::1] g     = np.empty( [n_space_dofs,n_time_dofs,n_quad], dtype=DTYPE, order="C" )  # effective load vector, for each space DOF, for each time DOF, at each integration point
    #    cdef DTYPE_t[:,::1]   b     = np.empty( [n_space_dofs,n_time_dofs],        dtype=DTYPE, order="C" )  # right-hand sides (integral, over the timestep, of g*psi)
    #    cdef DTYPE_t[:,::1]   u     = np.empty( [n_space_dofs,n_time_dofs],        dtype=DTYPE, order="C" )  # Galerkin coefficients (unknowns)
    #    cdef DTYPE_t[:,::1]   uprev = np.empty( [n_space_dofs,n_time_dofs],        dtype=DTYPE, order="C" )  # Galerkin coefficients from previous iteration
    #    cdef DTYPE_t[:,::1]   uass  = np.empty( [n_quad,n_space_dofs],             dtype=DTYPE, order="C" )  # u, assembled for integration (this ordering needed for speed!)
    #    cdef DTYPE_t[::1]     ucorr = np.empty( [n_quad],                          dtype=DTYPE, order="C" )  # correction for compensated summation in assemble() (for integration)
    #    cdef DTYPE_t[:,::1]   uvis  = np.empty( [nt_vis,n_space_dofs],             dtype=DTYPE, order="C" )  # u, assembled for visualization
    #    cdef DTYPE_t[::1]     ucvis = np.empty( [nt_vis],                          dtype=DTYPE, order="C" )  # correction for compensated summation in assemble() (for visualization)
    #
    #    # global arrays, same for each solver instance (see galerkin.DataManager.load_data(), galerkin.DataManager.prep_solver())
    #    cdef RTYPE_t[:,::1] LU      = galerkin.datamanager.LU       # LU decomposed mass matrix (packed format), for one space DOF, shape (n_time_dofs, n_time_dofs)
    #    cdef int[::1]       p       = galerkin.datamanager.p        # row permutation information, length n_time_dofs
    #    cdef int[::1]       mincols = galerkin.datamanager.mincols  # band information for L, length n_time_dofs
    #    cdef int[::1]       maxcols = galerkin.datamanager.maxcols  # band information for U, length n_time_dofs
    #    cdef RTYPE_t[::1]   qw      = galerkin.datamanager.integ_w  # quadrature weights (Gauss-Legendre)
    #    cdef RTYPE_t[:,::1] psi     = galerkin.datamanager.integ_y  # basis function values at the quadrature points, psi[j,i] is N[j]( tquad[i] )
    #    cdef RTYPE_t[:,::1] psivis  = galerkin.datamanager.vis_y    # basis function values at the visualization points, psivis[j,i] is N[j]( tvis[i] )
    #
    cdef int call(self, DTYPE_t* w, RTYPE_t t, RTYPE_t dt) nogil:
        cdef DTYPE_t* up = self.wrk  # temporary storage for u'
        cdef unsigned int nequals  # for convergence check

        cdef int n_space_dofs = self.rhs.n
        cdef int n_time_dofs  = self.n_time_dofs
        cdef int n_quad       = self.n_quad

        cdef RTYPE_t tcurr  # t at current quadrature point

        # Loop counters.
        #
        # j = space DOF
        # k = time DOF
        # l = quadrature point
        # m = implicit iteration (Banach fixed point iteration)
        #
        cdef int j, k, l, m=-1

        # Initialize the Galerkin coefficients u: initially, we set u = u0, which is available in the array w.
        #
        # u :     current value of Galerkin coefficients
        # uprev : value of Galerkin coefficients at previous iteration, for convergence check
        #
        for j in range(n_space_dofs):
            # linears
            self.u[j*n_time_dofs + 0]     = w[j]  # value at start of timestep
            self.u[j*n_time_dofs + 1]     = w[j]  # value at end of timestep
            self.uprev[j*n_time_dofs + 0] = w[j]
            self.uprev[j*n_time_dofs + 1] = w[j]

            # bubbles
            for k in range(2,n_time_dofs):
                self.u[j*n_time_dofs + k]     = 0.0
                self.uprev[j*n_time_dofs + k] = 0.0

        # Implicit iteration (Banach/Picard fixed point iteration)
        #
        for m in range(self.maxit):
            self.rhs.begin_iteration(m)  # inform RHS kernel that a new iteration starts

            # Assemble u at the quadrature points for this iteration.
            #
            # This forms the core of the Banach fixed point iteration: here we use the
            # previous Galerkin coefficients u (latest known iterate) to model the behavior of u
            # inside the timestep.
            #
            # The old coefficients are then no longer needed; they get overwritten by the solver below.
            #
            self.assemble( self.psi, self.uass, self.ucorr, n_quad )

            # Form the right-hand side at each quadrature point.
            #
            for l in range(n_quad):
                tcurr = t + dt*self.tquad[l]

                # Get the instantaneous value of u' at this quadrature point, store result in "up".
                #
                self.rhs.call(&self.uass[n_space_dofs*l + 0], up, tcurr)

                # At the quadrature point l, compute the value of the RHS integrand
                # in the weak form (original RHS multiplied by test function).
                #
                # We do this for each space DOF independently.
                #
                # The time DOFs here correspond to the set of test functions.
                #
                for j in range(n_space_dofs):
                    for k in range(n_time_dofs):
                        # g: [n_space_dofs,n_time_dofs,n_quad], values of the RHS integrand at the quadrature points
                        # psi: [n_time_dofs, n_quad], values of the test functions at the quadrature points
                        self.g[j*(n_time_dofs*n_quad) + k*n_quad + l] = up[j] * self.psi[k*n_quad + l]

            # Solve the dG linear equation system (separately for each space DOF).
            #
            for j in range(n_space_dofs):

                # For this space DOF, integrate the right-hand side across the timestep (for each test function).
                #
                # We use the Gauss-Legendre rule that was set up by init(). (This information is implicit
                # in the array qw, and in the array psi, which was above used to evaluate the integrand.)
                #
                for k in range(n_time_dofs):
                    # b: [n_space_dofs,n_time_dofs], load vector for dG linear equation system
                    # g: [n_space_dofs,n_time_dofs,n_quad], values of the RHS integrand at the quadrature points
                    self.b[j*n_time_dofs + k] = self.do_quadrature( &self.g[j*(n_time_dofs*n_quad) + k*n_quad + 0], dt )

                # Add the RHS contribution from the jump term (see user manual). This accounts for the initial condition for this timestep.
                #
                self.b[j*n_time_dofs + 0] += w[j]

                # Solve the dG linear equation system for this space DOF.
                #
                # This updates the Galerkin coefficients self.u.
                #
                # TODO: need some variant of zgesv if DTYPE is ZTYPE
                #
                dgesv_c.solve_decomposed_banded_c( self.LU, self.p, self.mincols, self.maxcols, &self.b[j*n_time_dofs + 0], &self.u[j*n_time_dofs + 0], n_time_dofs )

            # Check convergence; break early if converged to within machine precision.
            #
            nequals = 0
            for j in range(n_space_dofs):
                for k in range(n_time_dofs):
                    if self.u[j*n_time_dofs + k] == self.uprev[j*n_time_dofs + k]:
                        nequals += 1
            if nequals == n_space_dofs*n_time_dofs:
                break

            # Store the current values for performing the convergence check at the next iteration.
            #
            for j in range(n_space_dofs):
                for k in range(n_time_dofs):
                    self.uprev[j*n_time_dofs + k] = self.u[j*n_time_dofs + k]

        # Store into w the final value at the end of this timestep.
        #
        self.final_value( w )

        return (m+1)


cdef class CG(GalerkinIntegrator):
    def __init__(self, KernelBase rhs, int maxit):
        """def __init__(self, KernelBase rhs, int maxit):

Time-continuous Galerkin.

This is almost the same code as dG, the only difference being in the handling of the load vector.

(Be aware that usually dG gives better results.)
"""
        # super
        GalerkinIntegrator.__init__(self, name="cG", rhs=rhs, maxit=maxit)

    # Integrate w over one timestep.
    #
    # When call() returns, the Galerkin coefficients in self.u correspond to the
    # behavior of w over the last processed timestep. This is used in visualization
    # to support "interp".
    #
    #    # array shapes and types:
    #
    #    # instance arrays (see galerkin.DataManager.allocate_storage())
    #    cdef DTYPE_t[:,:,::1] g     = np.empty( [n_space_dofs,n_time_dofs,n_quad], dtype=DTYPE, order="C" )  # effective load vector, for each space DOF, for each time DOF, at each integration point
    #    cdef DTYPE_t[:,::1]   b     = np.empty( [n_space_dofs,n_time_dofs],        dtype=DTYPE, order="C" )  # right-hand sides (integral, over the timestep, of g*psi)
    #    cdef DTYPE_t[:,::1]   u     = np.empty( [n_space_dofs,n_time_dofs],        dtype=DTYPE, order="C" )  # Galerkin coefficients (unknowns)
    #    cdef DTYPE_t[:,::1]   uprev = np.empty( [n_space_dofs,n_time_dofs],        dtype=DTYPE, order="C" )  # Galerkin coefficients from previous iteration
    #    cdef DTYPE_t[:,::1]   uass  = np.empty( [n_quad,n_space_dofs],             dtype=DTYPE, order="C" )  # u, assembled for integration (this ordering needed for speed!)
    #    cdef DTYPE_t[::1]     ucorr = np.empty( [n_quad],                          dtype=DTYPE, order="C" )  # correction for compensated summation in assemble() (for integration)
    #    cdef DTYPE_t[:,::1]   uvis  = np.empty( [nt_vis,n_space_dofs],             dtype=DTYPE, order="C" )  # u, assembled for visualization
    #    cdef DTYPE_t[::1]     ucvis = np.empty( [nt_vis],                          dtype=DTYPE, order="C" )  # correction for compensated summation in assemble() (for visualization)
    #
    #    # global arrays, same for each solver instance (see galerkin.DataManager.load_data(), galerkin.DataManager.prep_solver())
    #    cdef RTYPE_t[:,::1] LU      = galerkin.datamanager.LU       # LU decomposed mass matrix (packed format), for one space DOF, shape (n_time_dofs, n_time_dofs)
    #    cdef int[::1]       p       = galerkin.datamanager.p        # row permutation information, length n_time_dofs
    #    cdef int[::1]       mincols = galerkin.datamanager.mincols  # band information for L, length n_time_dofs
    #    cdef int[::1]       maxcols = galerkin.datamanager.maxcols  # band information for U, length n_time_dofs
    #    cdef RTYPE_t[::1]   qw      = galerkin.datamanager.integ_w  # quadrature weights (Gauss-Legendre)
    #    cdef RTYPE_t[:,::1] psi     = galerkin.datamanager.integ_y  # basis function values at the quadrature points, psi[j,i] is N[j]( tquad[i] )
    #    cdef RTYPE_t[:,::1] psivis  = galerkin.datamanager.vis_y    # basis function values at the visualization points, psivis[j,i] is N[j]( tvis[i] )
    #
    cdef int call(self, DTYPE_t* w, RTYPE_t t, RTYPE_t dt) nogil:
        cdef DTYPE_t* up = self.wrk  # temporary storage for u'
        cdef unsigned int nequals  # for convergence check

        cdef int n_space_dofs = self.rhs.n
        cdef int n_time_dofs  = self.n_time_dofs
        cdef int n_quad       = self.n_quad

        cdef RTYPE_t tcurr  # t at current quadrature point

        # Loop counters.
        #
        # j = space DOF
        # k = time DOF
        # l = quadrature point
        # m = implicit iteration (Banach fixed point iteration)
        #
        cdef int j, k, l, m=-1

        # Initialize the Galerkin coefficients u: initially, we set u = u0, which has been saved as the array w.
        #
        for j in range(n_space_dofs):
            self.u[j*n_time_dofs + 0]     = w[j]  # value at start of timestep
            self.u[j*n_time_dofs + 1]     = w[j]  # value at end of timestep

            # value of Galerkin coefficients at previous iteration, for convergence check
            self.uprev[j*n_time_dofs + 0] = w[j]
            self.uprev[j*n_time_dofs + 1] = w[j]

            # bubbles
            for k in range(2,n_time_dofs):
                self.u[j*n_time_dofs + k]     = 0.0
                self.uprev[j*n_time_dofs + k] = 0.0

        # Implicit iteration (Banach/Picard fixed point iteration)
        #
        for m in range(self.maxit):
            self.rhs.begin_iteration(m)  # inform RHS kernel that a new iteration starts

            # Assemble latest known u at the quadrature points for this iteration.
            #
            self.assemble( self.psi, self.uass, self.ucorr, n_quad )

            # Form the right-hand side at each quadrature point.
            #
            for l in range(n_quad):
                tcurr = t + dt*self.tquad[l]

                # Get the instantaneous value of u' at this quadrature point, store result in "up".
                #
                self.rhs.call(&self.uass[n_space_dofs*l + 0], up, tcurr)

                # Compute the value of the RHS integrand in the weak form.
                for j in range(n_space_dofs):
                    for k in range(1,n_time_dofs):  # note that in cG, the load vector for time DOF 0 will be filled in separately (initial condition)
                        # g: [n_space_dofs,n_time_dofs,n_quad], values of the RHS integrand at the quadrature points
                        # psi: [n_time_dofs, n_quad], values of the test functions at the quadrature points
                        self.g[j*(n_time_dofs*n_quad) + k*n_quad + l] = up[j] * self.psi[k*n_quad + l]

            # Solve the cG linear equation system (separately for each space DOF).
            #
            for j in range(n_space_dofs):

                # For this space DOF, integrate the right-hand side across the timestep (for each test function).
                #
                for k in range(1,n_time_dofs):  # note that in cG, the load vector for time DOF 0 will be filled in separately (initial condition)
                    # b: [n_space_dofs,n_time_dofs], load vector for dG linear equation system
                    # g: [n_space_dofs,n_time_dofs,n_quad], values of the RHS integrand at the quadrature points
                    self.b[j*n_time_dofs + k] = self.do_quadrature( &self.g[j*(n_time_dofs*n_quad) + k*n_quad + 0], dt )

                # Account for the initial condition in the load vector (see user manual).
                #
                # The first time basis function (index 0) is the only one that is nonzero at the start point of the timestep;
                # its value there is exactly 1. Hence the value of the first time DOF (index 0) for space DOF j can be set to w[j].
                #
                # The mass matrix on the LHS has already been adjusted in init() to have the multiplicative identity on the first row,
                # so that when the linear equation system is solved, u[j,0] will just pick up the value set here.
                #
                self.b[j*n_time_dofs + 0] = w[j]

                # Solve the cG linear equation system for this space DOF.
                #
                # This updates the Galerkin coefficients self.u.
                #
                # TODO: need some variant of zgesv if DTYPE is ZTYPE
                #
                dgesv_c.solve_decomposed_banded_c( self.LU, self.p, self.mincols, self.maxcols, &self.b[j*n_time_dofs + 0], &self.u[j*n_time_dofs + 0], n_time_dofs )

            # Check convergence; break early if converged to within machine precision.
            #
            nequals = 0
            for j in range(n_space_dofs):
                for k in range(n_time_dofs):
                    if self.u[j*n_time_dofs + k] == self.uprev[j*n_time_dofs + k]:
                        nequals += 1
            if nequals == n_space_dofs*n_time_dofs:
                break

            # Store the current values for performing the convergence check at the next iteration.
            #
            for j in range(n_space_dofs):
                for k in range(n_time_dofs):
                    self.uprev[j*n_time_dofs + k] = self.u[j*n_time_dofs + k]

        # Store into w the final value at the end of this timestep.
        #
        self.final_value( w )

        return (m+1)


#######################
# Python-level helpers
#######################

# When instantiated, the following object performs the common initialization that is needed only once per process,
# including the construction of the mass matrix for the first-order initial value problem
# for a single space DOF:
#
#   u'    = f( u(t), t ),   t in (-1,1)
#   u(-1) = u0
#
# This is done in generic form over the reference element [-1,1]. The LU decomposition
# of the matrix is also formed. The LU decomposition can then be used for fast solving
# of the linear equation systems in the main solver loop.
#
#
# Thread-safety:
#   Multiple dG (or cG) solvers may be running at the same time, but they *must* have the same settings (method and q).
#   Simultaneous processing with different settings is currently not supported and will cause errors and/or crashes. (TODO)
#
#
# Multiple "space" degrees of freedom in the semilinear ODE system
#
#   M u' = g( u(t), t )                         (*)
#
# are to be handled by the user-provided kernel by formally inverting M:
#
#   u' = M^(-1) g( u(t), t ) =: f( u(t), t )    (**)
#
# The problem (**) (first and last form) is of the standard form handled by the solver.
#
# From the last two forms of (**), we have (multiplying from the left by M)
#
#   M f = g
#
# Thus, the kernel can solve a linear equation system (different from the dG(q) one!) to
# obtain the effective load f whenever it is needed. See kernels.pyx for examples.
#
# This needs M, g and the latest u (all "space" degrees of freedom), and outputs the
# corresponding value of f (for all "space" degrees of freedom).
#
# This in turn allows us to solve the time evolution of each "space" degree of freedom separately,
# for each of them using the time integrator that handles a single "space" degree of freedom.
#
# (In addition to "space" degrees of freedom, i.e. degrees of freedom in the ODE system,
#  there are also "time" degrees of freedom, of which there are q+1 for *each* "space" degree of freedom.
#  These "time" degrees of freedom are the ones whose behavior is encoded in the Galerkin mass matrix C, below.
#
#  We use the label "space" degrees of freedom, because ODE systems like (*) typically arise
#  in space discretization of PDE problems.)
#
#
class DataManager:
    """Helper class to manage the Galerkin solver."""

    def __init__(self, q, method, nt_vis, rule=None):
        """def __init__(self, q, method, nt_vis, rule=None):

Perform global initialization of the Galerkin solver.

This object must be instantiated before using the integrators "dG" or "cG".
See the module-level global function init(), which does this.

During one session (run of the process), one initialization is enough for all simulations
using the same method (combination of "q" and "method"). To change the settings, simply initialize again.

The Galerkin solver uses the hierarchical basis functions (linears, plus bubbles
based on integrated Legendre polynomials).

Parameters:
    q : int, >= 1
        Polynomial degree of basis. In practice, 1 or 2 are usually good values.

        The upper limit available depends on data stored in the precalculated data file (pydgq_data.bin).
        To determine what your data file has, instantiate a DataManager and call its load_data() method.

    method : string
        Which Galerkin method to use:

         "dG" : time-discontinuous Galerkin (recommended)
            The solution is finitely discontinuous (C^{-1}) across timestep boundaries,
            with left-continuity at each boundary.

            Because no degrees of freedom are wasted on enforcing continuity,
            this allows better approximation at the same q.

         "cG" : continuous Galerkin
            Conforming finite elements in time.

            The solution is C0-continuous across timestep boundaries. This may be
            more intuitive for visualization, but note that usually the approximation
            is much worse than with dG.

    nt_vis : int
        Number of visualization points per computed timestep.

        This corresponds to the parameter "interp" of odesolve.result_len(),
        odesolve.make_tt() and odesolve.ivp().

        !!! Currently, the value of interp in all subsequent calls to odesolve.ivp() must be the same
            as the value of nt_vis passed here. This is a implementation-technical limitation
            that may be removed in the future. !!!

    rule : int or None
        Order of the Gauss-Legendre quadrature to use for integrating the load vector.
        If None (default), the highest available order is used (see pydqg.utils.precalc).

        To determine what your data file has, instantiate a DataManager and call its load_data() method.
"""

        if q < 1:
            raise ValueError( "q must be >= 1; got %d" % (q) )
        if nt_vis < 1:
            raise ValueError( "nt_vis must be >= 1; got %d" % (nt_vis) )
        if method not in ["dG", "cG"]:
            raise ValueError( "Unknown method '%s'; valid: 'dG', 'cG'" % (method) )

        self.method      = method
        self.q           = q       # degree of basis
        self.n_time_dofs = q+1     # the hierarchical basis of degree q has q+1 degrees of freedom.
        self.nt_vis      = nt_vis  # number of visualization points per timestep

        try:
            self.__load_data(rule)
            self.__build_C()
            self.__prep_solver()
            self.__storage = {}  # problem instance specific data will be stored here (key = solver id)
            self.available = True
        except Exception as e:
            # NOTE: our "t" = "x" in precalc.py; it is a general support module for Lobatto basis functions, not specific to its ODE solver use case here.

            # precomputed data for Gauss-Legendre integration (used for processing the load vector in the solver)
            self.integ_x     = None   # GL points
            self.integ_y     = None   # values of basis functions at GL points
            self.integ_w     = None   # GL weights
            self.rule        = None   # order of the chosen GL rule (integer)

            # precomputed data for timestep visualization
            self.vis_x       = None   # visualization points (time values on reference element [-1,1])
            self.vis_y       = None   # values of basis functions at visualization points
            self.nt_vis      = None   # number of visualization points within each timestep

            self.q           = None   # degree of basis
            self.n_time_dofs = None   # number of "time degrees of freedom" for each "space degree of freedom"

            self.C           = None   # Galerkin mass matrix
            self.LU          = None   # the LU decomposition of C
            self.p           = None   # row permutation for load vector ( L U = P C,  so if  C u = b,  then  P C u = P b  and thus  L U u = P b )
            self.mincols     = None   # band information for L
            self.maxcols     = None   # band information for U

            self.__storage     = None   # storage for solver instance arrays

            self.available   = False  # Galerkin integrators enabled/disabled

            print( "    ERROR: Initialization failed. Galerkin integrators not available.", file=sys.stderr )
            raise

    # Load precalculated data for the hierarchical basis functions.
    #
    def __load_data(self, rule):
        """Load precalculated data (pydgq_data.bin).

The hierarchical basis functions are numerically sensitive to cancellation especially at high degrees q,
hence must be computed using higher internal precision (double is not enough).

Because arbitrary precision floating point in pure software is slow, pydqg.utils.precalc makes
a precomputed array for this routine to load.

Parameters:
    rule : int or None
        Order of the Gauss-Legendre quadrature to use for integrating the load vector.
        None = highest available.
"""
        q           = self.q
        n_time_dofs = self.n_time_dofs
        nt_vis      = self.nt_vis

        # first try ./pydgq.data.bin (local override)
        data = None
        filename = 'pydgq_data.bin'
        try:
            with open(filename, 'rb') as infile:
                data = pickle.load( infile )
        except IOError as e:
            pass

        # next try ~/.config/pydgq/pydgq_data.bin (user override)
        if data is None:
            import os.path
            home = os.path.expanduser("~")
            url = os.path.join(home, ".config", "pydgq", filename)

            try:
                with open(url, 'rb') as infile:
                    data = pickle.load( infile )
            except IOError as e:
                pass

        # finally, use pkg_resources to find installed data file
        if data is None:
            if not pkg_resources.resource_exists("pydgq", "pydgq_data.bin"):
                raise IOError("Cannot find data file '%s'. The file can be generated by running the module pydqg.utils.precalc as the main program." % (filename))

            try:
                data = pickle.load( pkg_resources.resource_stream("pydgq", "pydgq_data.bin") )
            except:
                print("Cannot find data file '%s'. The file can be generated by running the module pydqg.utils.precalc as the main program." % (filename), file=sys.stderr)
                raise

        maxq      = data["maxq"]
        maxnt_vis = data["vis"]["maxnx"]
        maxrule   = data["integ"]["maxrule"]

        # save maximum values available in data file, for information
        #
        self.maxq      = maxq
        self.maxnt_vis = maxnt_vis
        self.maxrule   = maxrule

        # Sanity check that we have enough data available
        #
        if q > maxq:
            raise ValueError("Data file '%s' contains data up to degree maxq=%d, but the requested degree was higher, q=%d. Use a lower q or re-run pydqg.utils.precalc for higher maxq." % (filename, maxq, q))
        need_nt_vis = nt_vis+1 if self.method == "cG" else nt_vis  # "visualization" points are used for evaluating the result inside the timestep (see parameter "interp" of odesolve.result_len() and odesolve.ivp()).
        if need_nt_vis > maxnt_vis:
            raise ValueError("Data file '%s' contains data up to visualization maxnt_vis=%d, but the needed number of points is higher, %d. Use a lower nt_vis or re-run pydqg.utils.precalc for higher maxnx (same thing as our maxnt_vis)." % (filename, maxnt_vis, need_nt_vis))
        if rule is not None and rule > maxrule:
            raise ValueError("Data file '%s' contains data for integration rules up to order %d, but the requested order was higher, rule=%d. Use a lower rule." % (filename, maxrule, rule))

        # Choose the order of the Gauss-Legendre quadrature and the number of visualization points.
        #
        # The load is not a polynomial, so we default to the highest available integration rule.
        #
        # The Gauss-Legendre rule of order q (where q us the degree of the basis) would be enough for our mass matrix C;
        # the fact that the matrix has the integrand N'*N and not N*N saves us one order, compared to the usual q+1.
        #
        # But we handle the matrix analytically, so the required order needed by the matrix does not matter;
        # Gauss-Legendre is used only for processing the load vector.
        #
        # If we wanted to be exact, we should least-squares fit the load to the given basis before performing the quadrature; this would give us computable error.
        # However, doing that is expensive, since the quadrature runs inside the loop for implicit iterations (iterative linearization using the Banach fixed point theorem).
        #
        self.rule = rule if rule is not None else maxrule
        rule = self.rule

        # Choose the appropriate data arrays
        #
        integ = data["integ"][rule]

        if self.method == "dG":
            # In dG, the value of the solution at the end of the timestep is different
            # from its value at the beginning of the next timestep, even though these points
            # occur at the same value of t.
            #
            # (Technically speaking, the solution is left-continuous, i.e. the end-of-timestep
            #  point has the actual value of the solution at that value of t.
            #
            #  The start-of-next-timestep point represents the limit from the right.)
            #
            # Thus, we just take the subdivision that has nt_vis points (including both endpoints).
            #
            vis = data["vis"][nt_vis]

        else: # self.method == "cG":
            # In cG, the endpoint of timestep n is the start point of timestep n+1.
            # Avoid visualizing it twice.
            #
            # To get the desired number of points, take one more, and then discard the start point,
            # obtaining nt_vis unique points.
            #
            vis = data["vis"][nt_vis+1]
            vis["x"] = vis["x"][1:]
            vis["y"] = vis["y"][:,1:]

        # Chop off higher degree (>q) basis functions from the y arrays, if the arrays contain more data than what we actually need.
        #
        if data["maxq"] > q:
            integ["y"] = integ["y"][ :(q+1), : ]
            vis["y"]   = vis["y"][   :(q+1), : ]

        # NOTE: our "t" = "x" in precalc.py; it is a general support module for Lobatto basis functions, not specific to its ODE solver use case here.
        #
        cdef RTYPE_t[::1]   integ_x = integ["x"].copy(order="C")  # ...and hope for the correct dtype (our RTYPE) in the data file
        cdef RTYPE_t[::1]   integ_w = integ["w"].copy(order="C")
        cdef RTYPE_t[:,::1] integ_y = integ["y"].copy(order="C")
        cdef RTYPE_t[::1]   vis_x   = vis["x"].copy(order="C")
        cdef RTYPE_t[:,::1] vis_y   = vis["y"].copy(order="C")

        self.integ_x = integ_x
        self.integ_y = integ_y
        self.integ_w = integ_w
        self.vis_x   = vis_x
        self.vis_y   = vis_y

        # We extracted what we need, unload the rest.
        #
        del data

        # Now the loaded data can be accessed as follows:
        #
        # Gauss-Legendre integration (will be used for processing the load vector):
        #  - self.integ_x: Gauss-Legendre points of the chosen rule, in (-1,1)
        #  - self.integ_w: corresponding Gauss-Legendre weights
        #  - self.integ_y: values of basis functions at the Gauss-Legendre points; integ_y[j,i] is N[j]( tquad[i] )
        #  - self.rule:    order of the chosen Gauss-Legendre rule (for information only)
        #
        # Visualization:
        #  - self.vis_x:   x values for visualization (time value on reference element [-1,1])
        #  - self.vis_y:   values of basis functions at the visualization points; vis_y[j,i] is N[j]( tvis[i] )

    def __build_C(self):
        """Build the dG(q) mass matrix C.

C contains integrals of Ni'*Nj (i column, j row) over the reference element [-1,1],
where the Ns are the hierarchical (Lobatto) basis functions.

This routine uses an analytical result (see user manual), avoiding the need
to numerically evaluate the integrands to get this. This approach is both
much more accurate and much faster.
"""
        # Note that there is no need to scale C by the length of each timestep, because the one differentiation and the integration exactly cancel out the Jacobians.

        q = self.q
        n = q+1

        C = np.zeros( (n,n), dtype=RTYPE )
        C[0,0] = -1./2.
        C[0,1] =  1./2.
        C[1,0] = -1./2.
        C[1,1] =  1./2.

        if q >= 2:
            t = 1./np.sqrt(6.)
            C[0,2] = -t
            C[1,2] =  t
            C[2,0] =  t
            C[2,1] = -t

        # General formula for C_ji for j,i >= 2.
        for j in range(2,n):
            i = j + 1  # i-1 = j  <=>  i = j+1
            if i >= 2 and i < n:
                C[j,i] =  2. * np.sqrt( 1. / ( ( 2.*j - 1. ) * ( 2.*j + 1. ) ) )
            i = j - 1  # i+1 = j  <=>  i = j-1
            if i >= 2 and i < n:
                C[j,i] = -2. * np.sqrt( 1. / ( ( 2.*j - 1. ) * ( 2.*j - 3. ) ) )

        if self.method == "dG":
            # Jump term for discontinous Galerkin (dG). N[0] is the only basis function that is nonzero at x=-1.
            #
            # See the user manual for derivation.
            #
            C[0,0] += 1.0
        else: # self.method == "cG":
            # In continuous Galerkin (cG), we don't have the jump term; there we must instead:
            #
            #   - Eliminate the first degree of freedom by substituting in its known value (eliminating a row and column, transferring the known terms to the load), or
            #   - Replace the first row with the multiplicative identity and replace the first component of the load with the known value of the first degree of freedom.
            #
            # The first is more elegant (disturbing the matrix structure less), while the second is easier to code. Both approaches achieve the same result.
            # Note that both require implementation also in the timestepping code that fills in the load vector (not just in this helper class).
            #
            # We use the easier second approach, because then we won't need to remap DOF numbers. Also, the  problem is antisymmetric (skew-symmetric),
            # i.e. a symmetric solver could not be used in any case. Note also that the equation system will be rather small, and the hierarchical
            # basis functions are good for numerical linear algebra (far from linear dependence in the sense of the L2 inner product), so there
            # should not be much roundoff or cancellation.
            #
            C[0,:] = 0.0
            C[0,0] = 1.0

        self.C = C

    def __prep_solver(self):
        """Prepare the dG(q) linear equation system for solving."""
        self.LU,self.p = dgesv.lup_packed(self.C)
        self.mincols,self.maxcols = dgesv.find_bands(self.LU, tol=1e-15)  # this allows us to use banded solver (the bw of the LU decomposition of the dG(q) mass matrix is rather small)
#        r   = np.arange(self.q+1, dtype=int)
#        bwL = np.max(r - self.mincols)  # convention: diagonal matrix = zero bandwidth ( https://en.wikipedia.org/wiki/Band_matrix )
#        bwU = np.max(self.maxcols - r)
#        print( "        Banded solver detected bandwidths L=%d (lower bw), U=%d (upper bw)" % (bwL, bwU) )

    def allocate_storage(self, solver_id, n_space_dofs):
        """def allocate_storage(self, solver_id, n_space_dofs):

Allocate and return per-problem storage.

This must be called separately for each problem that will be running simultaneously (even if same settings),
as each problem instance needs its own work space.

This is handled automatically by GalerkinIntegrator.__init__.

Parameters:
    solver_id : hashable
        unique identifier (id(self) at the calling end is a good choice, or if not object-oriented, id(result_array_ww))
    n_space_dofs : int
        number of space DOFs (i.e. size of the 1st-order ODE system) for the problem to be solved.
"""

        assert self.__storage is not None
        assert self.n_time_dofs is not None
        assert self.rule is not None
        assert self.nt_vis is not None

        n_time_dofs = self.n_time_dofs
        n_quad      = self.rule  # Gauss-Legendre has as many quadrature points as the order of the rule being used.
        nt_vis      = self.nt_vis

        cdef DTYPE_t[:,:,::1] g     = np.empty( [n_space_dofs,n_time_dofs,n_quad], dtype=DTYPE, order="C" )  # effective load vector, for each space DOF, for each time DOF, at each integration point
        cdef DTYPE_t[:,::1]   b     = np.empty( [n_space_dofs,n_time_dofs],        dtype=DTYPE, order="C" )  # right-hand sides (integral, over the timestep, of g*psi)
        cdef DTYPE_t[:,::1]   u     = np.empty( [n_space_dofs,n_time_dofs],        dtype=DTYPE, order="C" )  # Galerkin coefficients (unknowns)
        cdef DTYPE_t[:,::1]   uprev = np.empty( [n_space_dofs,n_time_dofs],        dtype=DTYPE, order="C" )  # Galerkin coefficients from previous iteration
        cdef DTYPE_t[:,::1]   uass  = np.empty( [n_quad,n_space_dofs],             dtype=DTYPE, order="C" )  # u, assembled for integration (this ordering needed for speed!)
        cdef DTYPE_t[::1]     ucorr = np.empty( [n_quad],                          dtype=DTYPE, order="C" )  # correction for compensated summation in assemble() (for integration)
        cdef DTYPE_t[:,::1]   uvis  = np.empty( [nt_vis,n_space_dofs],             dtype=DTYPE, order="C" )  # u, assembled for visualization
        cdef DTYPE_t[::1]     ucvis = np.empty( [nt_vis],                          dtype=DTYPE, order="C" )  # correction for compensated summation in assemble() (for visualization)
        cdef DTYPE_t[::1]     wrk   = np.empty( [n_space_dofs],                    dtype=DTYPE, order="C" )  # work space for dG(), cG()

        items = {}

        items["n_space_dofs"] = n_space_dofs

        items["g"]     = g
        items["b"]     = b
        items["u"]     = u
        items["uprev"] = uprev
        items["uass"]  = uass
        items["ucorr"] = ucorr
        items["uvis"]  = uvis
        items["ucvis"] = ucvis
        items["wrk"]   = wrk

        self.__storage[solver_id] = items
        return items

    def get_storage(self, solver_id):
        """Return per-problem storage for given solver_id."""
        return self.__storage[solver_id]

    def free_storage(self, solver_id):
        """Deallocate per-problem storage."""
        self.__storage.pop(solver_id, None)   # http://stackoverflow.com/questions/11277432/how-to-remove-a-key-from-a-dictionary


datamanager = None
def init(q=3, method="dG", nt_vis=1, rule=None):
    """def init(q=3, method="dG", nt_vis=1, rule=None):

Initialize Galerkin solver.

This function loads common precalculated data.

This must be called before using the integrators "dG" or "cG".

During one session (run of the process), one initialization is enough for all simulations
using the same method (combination of "q" and "method"). To change the settings, simply initialize again.

The Galerkin solver uses the hierarchical (Lobatto) basis functions (linears, plus bubbles
based on integrated Legendre polynomials).

Parameters (see DataManager.__init__ for full description):
    q : int, >= 1
        Polynomial degree of basis.

    method : str
        Which Galerkin method to use.

    nt_vis : int, >= 1
        Number of visualization points per computed timestep.

         !!! Currently, the value of interp in all subsequent calls to odesolve.ivp() must be the same
             as the value of nt_vis passed here. This is a implementation-technical limitation
             that may be removed in the future. !!!

    rule : int or None
        Order of the Gauss-Legendre quadrature to use for integrating the load vector.
        None = highest available.
"""
    # TODO: relax the technical limitation on nt_vis <-> interp
    global datamanager
    datamanager = DataManager(q, method, nt_vis, rule)

