"""
A generic Runge-Kutta type integrator for integrating CellCenterData2d.
We support a generic Butcher tableau for explicit the Runge-Kutta update::

   0   |
   c_2 | a_21
   c_3 | a_31 a_32
   :   |  :        .
   :   |  :          .
   c_s | a_s1 a_s2 ... a_s,s-1
   ----+---------------------------
       | b_1  b_2  ... b_{s-1}  b_s


the update is::

   y_{n+1} = y_n + dt sum_{i=1}^s {b_i k_i}

and the s increment is::

   k_s = f(t + c_s dt, y_n + dt (a_s1 k1 + a_s2 k2 + ... + a_s,s-1 k_{s-1})

"""

import numpy as np

from pyro.mesh import patch

a = {}
b = {}
c = {}

nstages = {}

# second-order standard
a["RK2"] = np.array([[0.0, 0.0],
                     [0.5, 0.0]])

b["RK2"] = np.array([0.0, 1.0])

c["RK2"] = np.array([0.0, 0.5])


# second-order TVD (Gottlieb & Shu)
a["TVD2"] = np.array([[0.0, 0.0],
                      [1.0, 0.0]])

b["TVD2"] = np.array([0.5, 0.5])

c["TVD2"] = np.array([0.0, 1.0])


# third-order TVD (Gottlieb & Shu)
a["TVD3"] = np.array([[0.0,  0.0,  0.0],
                      [1.0,  0.0,  0.0],
                      [0.25, 0.25, 0.0]])

b["TVD3"] = np.array([1./6., 1./6., 2./3.])

c["TVD3"] = np.array([0.0, 1.0, 0.5])


# fourth-order
a["RK4"] = np.array([[0.0, 0.0, 0.0, 0.0],
                     [0.5, 0.0, 0.0, 0.0],
                     [0.0, 0.5, 0.0, 0.0],
                     [0.0, 0.0, 1.0, 0.0]])

b["RK4"] = np.array([1./6., 1./3., 1./3., 1./6.])

c["RK4"] = np.array([0.0, 0.5, 0.5, 1.0])



[docs]
class RKIntegrator:
    """the integration class for CellCenterData2d, supporting RK
    integration"""

    def __init__(self, t, dt, method="RK4"):
        """t is the starting time, dt is the total timestep to advance, method
        = {2,4} is the temporal method"""
        self.method = method

        self.t = t
        self.dt = dt

        # storage for the intermediate stages
        self.k = [None]*len(b[self.method])

        self.start = None


[docs]
    def nstages(self):
        """return the number of stages"""
        return len(b[self.method])



[docs]
    def set_start(self, start):
        """store the starting conditions (should be a CellCenterData2d
        object)"""
        self.start = start



[docs]
    def store_increment(self, istage, k_stage):
        """store the increment for stage istage -- this should not have a dt
        weighting"""
        self.k[istage] = k_stage



[docs]
    def get_stage_start(self, istage):
        """get the starting conditions (a CellCenterData2d object) for stage
        istage"""
        if istage == 0:
            ytmp = self.start
        else:
            ytmp = patch.cell_center_data_clone(self.start)
            for n in range(ytmp.nvar):
                var = ytmp.get_var_by_index(n)
                for s in range(istage):
                    var.v()[:, :] += self.dt*a[self.method][istage, s]*self.k[s].v(n=n)[:, :]

            ytmp.t = self.t + c[self.method][istage]*self.dt

        return ytmp



[docs]
    def compute_final_update(self):
        """this constructs the final t + dt update, overwriting the initial data"""
        ytmp = self.start
        for n in range(ytmp.nvar):
            var = ytmp.get_var_by_index(n)
            for s in range(self.nstages()):
                var.v()[:, :] += self.dt*b[self.method][s]*self.k[s].v(n=n)[:, :]

        return ytmp


    def __str__(self):
        return f"integration method: {self.method}; number of stages: {self.nstages()}"