from .solvers import FixedGridODESolver
from .rk_common import rk4_alt_step_func
from .misc import Perturb


class Euler(FixedGridODESolver):
    order = 1

    def _step_func(self, func, t0, dt, t1, y0):
        f0 = func(t0, y0, perturb=Perturb.NEXT if self.perturb else Perturb.NONE)
        return dt * f0, f0


class Midpoint(FixedGridODESolver):
    order = 2

    def _step_func(self, func, t0, dt, t1, y0):
        half_dt = 0.5 * dt
        f0 = func(t0, y0, perturb=Perturb.NEXT if self.perturb else Perturb.NONE)
        y_mid = y0 + f0 * half_dt
        return dt * func(t0 + half_dt, y_mid), f0

class Bosh3(FixedGridODESolver):
    order = 3

    def _step_func(self, func, t0, dt, t1, y0):
        f0 = func(t0, y0, perturb=Perturb.NEXT if self.perturb else Perturb.NONE)
        k1 = dt * f0
        k2 = dt * func(t0 + dt/2.0, y0 + 1.0/2.0*k1)
        k3 = dt * func(t0 + dt*3.0/ 4.0, y0 + 3.0/4.0*k2)
        out = y0 + 2.0/9.0*k1 + 1.0/3.0*k2 + 4.0/9.0*k3
        return out, f0

class RK4(FixedGridODESolver):
    order = 4

    def _step_func(self, func, t0, dt, t1, y0):
        f0 = func(t0, y0, perturb=Perturb.NEXT if self.perturb else Perturb.NONE)
        return rk4_alt_step_func(func, t0, dt, t1, y0, f0=f0, perturb=self.perturb), f0

class Dopri5(FixedGridODESolver):
    order = 5
    alpha=[1./5., 3./10., 4./5., 8./9., 1., 1.]
    beta=[
         [1./5.],
         [3./40., 9./40.],
         [44./45., -56./15., 32./9.],
         [19372./6561., -25360./2187., 64448./6561., -212./729.],
         [9017./3168., -355./33., 46732./5247., 49./176., -5103./18656.],
         [35./384., 0, 500./1113., 125./192., -2187./6784., 11./84.],
         ]
    c_sol=[35./384., 0, 500./1113., 125./192., -2187./6784., 11./84., 0]
    c_error=[
            35./384. - 1951./21600.,
            0,
            500./1113. - 22642./50085.,
            125./192. - 451./720.,
            -2187./6784. - -12231./42400.,
            11./84. - 649./6300.,
            -1. / 60.,
            ]

    def _step_func(self, func, t0, dt, t1, y0):
        f0 = func(t0, y0, perturb=Perturb.NEXT if self.perturb else Perturb.NONE)
        k = [dt * f0]
        stage = 6
        for i in range(stage):
            ti = t0 + self.alpha[i] * dt
            yi = y0
            for j in range(i+1):
                yi = yi + self.beta[i][j]*k[j]
            if i < 5:
                k = k + [ dt * func(ti, yi) ]
        out = yi
        #for i in range(stage):
        #    out = out + self.c_sol[i]*k[i]
        return out, f0
