from pathlib import Path
import sys

notebook_directory_parent = Path.cwd().resolve().parent.parent
if str(notebook_directory_parent) not in sys.path:
    sys.path.append(str(notebook_directory_parent))

%matplotlib inline

import numpy as np
import scipy
import sympy
from numpy import linspace
from scipy.integrate import solve_ivp

import matplotlib.pyplot as plt

from T1000.numerical.RungeKuttaMethod import RungeKuttaMethod
from T1000.numerical.RKMethods.bCoefficients import bCoefficients
from T1000.numerical.RKMethods.DOPRI5Coefficients import DOPRI5Coefficients

"""
Explanation: Setup Jupyter Notebook
End of explanation
"""


def exponential_decay(t, y):
    return -0.5 * y

sol = solve_ivp(exponential_decay, [0, 10], [2, 4, 8])
print("t\n", sol.t)
print("y\n", sol.y)

"""
Explanation: cf. Examples for solve_ivp
End of explanation
"""


cs_for_rk4 = [0.5, 0.5, 1]
as_for_rk4 = [0.5, 0, 0.5, 0, 0, 1]
bs_for_rk4 = [1/6., 1/3., 1/3., 1/6.]
rk4 = RungeKuttaMethod(4, cs_for_rk4, as_for_rk4, bs_for_rk4)

"""
Explanation: An example for unit tests
cf. Runge-Kutta method from OK state
Depending on one's notation, for "stage" $s=4$, or $m=4$, then for
$\alpha_2 = \frac{1}{2}, \, \alpha_3 = \frac{1}{2}, \, \alpha_4 = 1 \equiv \
c_2 = \frac{1}{2}, \, c_3 = \frac{1}{2}, \, c_4 = 1$
$\beta_{21} = \frac{1}{2}, \, \beta_{31} = 0, \, \beta_{32} = \frac{1}{2}, \, \beta_{41} = \beta_{42} = 0, \, \beta_{43} = 1 \equiv \
a_{21} = \frac{1}{2}, \, a_{31} =0, \, a_{32} = \frac{1}{2}, \, a_{41} = a_{42} = 0, \, a_{43} = 1$
$c_1 = c_4 = \frac{1}{6}, \, c_2 = c_3 = \frac{1}{3} \equiv b_1 = b_4 = \frac{1}{6}, \, b_2 = b_3 = \frac{1}{3}$ 
(authors having different notation sucks).
End of explanation
"""


def derivative(x, y):
    return y - np.power(x, 2) + 1.0

def exact_y(x):
    return np.power(x, 2) + 2.0 * x + 1.0 - 0.5 * np.exp(x)

"""
Explanation: Consider example
$\begin{cases}
    & y'  =y -x^2 + 1\
    & y(0) = 0.5
\end{cases}$
Its exact solution is 
$ y= y(x) = x^2 + 2x + 1 - \frac{1}{2} \exp{x}$
End of explanation
"""


ks0 = rk4._calculate_k_coefficients(0.5, 0.0, 0.5, derivative)
print(ks0)
y1 = rk4.calculate_next_step(0.5, 0.0, 0.5, derivative)
print(y1)
print(exact_y(0))
print(exact_y(0.5))

"""
Explanation: We first solve this problem using RK4 with $h = 0.5$ and from $t=0$ to $t=2$.

With step size $h = 0.5$, it takes 4 steps: $t_0 = 0, \, t_1 = 0.5, \, t_2 = 1, \, t_3 = 1.5, \, t_4 = 2$
Step 0 \quad \, $x_0 = 0, \, y_0 = 0.5$ 
Step 1 \quad \, $x_1 = 0.5, \, h = 0.5$
End of explanation
"""


ks1 = rk4._calculate_k_coefficients(y1, 0.5, 0.5, derivative)
print(ks1)
y2 = rk4.calculate_next_step(y1, 0.5, 0.5, derivative)
print(y2)
print(exact_y(1.))

"""
Explanation: Step 2 $\quad \, x_2 = 1$
End of explanation
"""


alpha_for_rk4 = [0.5, 0.5, 1.]
beta_for_rk4 = [0.5, 0., 0.5, 0., 0., 1.]
c_for_rk4 = [1./6., 1./3., 1./3., 1./6.]
rk4 = RungeKuttaMethod(4, alpha_for_rk4, beta_for_rk4, c_for_rk4)

m = 4
print(rk4._beta_coefficients)
for i in range(2, m + 1):
    for j in range(1, i):
        print(i, " ", j, " ", rk4.get_beta_ij(i, j))

print(rk4._alpha_coefficients)
for i in range(2, m + 1):
    print(i, " ", rk4.get_alpha_i(i))

print(rk4._c_coefficients)
for i in range(1, m + 1):
    print(i, " ", rk4.get_c_i(i))

x_n = [np.array([2, 4, 8]), np.array([1.88836035, 3.7767207, 7.5534414])]
t_n = [0., 0.11487653, 1.26364188]

kresult = rk4._calculate_k_coefficients(x_n[0], t_n[0], t_n[1] - t_n[0], exponential_decay)
print(kresult)

result1 = rk4.calculate_next_step(x_n[0], t_n[0], t_n[1] - t_n[0], exponential_decay)
print(result1)

result2 = rk4.calculate_next_step(x_n[1], t_n[1], t_n[2] - t_n[1], exponential_decay)
print(result2)

[1, 2, 3] + [ 4, 5, 6]

np.array([2])

"""
Explanation: Calculated values for unit tests
End of explanation
"""


# print(DOPRI5Coefficients.b_coefficients.get_ith_element(1))
deltas = []
for i in range(1, 8):
    delta = \
        DOPRI5Coefficients.b_coefficients.get_ith_element(i) - \
            DOPRI5Coefficients.bstar_coefficients.get_ith_element(i)
    deltas.append(delta)
    if (type(delta) != int):
        print(delta.simplify())
    else:
        print(delta)

"""
Explanation: Embedded Formulas of Order 5
cf. pp. 178, Table 5.2. Dormand-Prince 5(4) (DOPRI5) Ordinary Differential Equations, Vol. 1. Nonstiff-Problems
End of explanation
"""


bs = DOPRI5Coefficients.b_coefficients
ds = DOPRI5Coefficients.dense_output_coefficients
cstars = DOPRI5Coefficients.cstar_coefficients

-2 * (1 + 4 * bs.get_ith_element(1) - 4 * cstars.get_ith_element(1))

"""
Explanation: Coefficients for dense output
cf. https://math.stackexchange.com/questions/2947231/how-can-i-derive-the-dense-output-of-ode45
End of explanation
"""


from sympy import Matrix, Symbol, symbols, pprint

theta, y_n, y_np1, f_n, f_np1, h = symbols("theta y_n y_np1 f_n f_np1 h")

"""
Explanation: Hermite Interpolation
End of explanation
"""


hermite_interpolation = (1 - theta) * y_n + theta * y_np1 + \
  theta * (theta - 1) * ((1 - 2 * theta) * (y_np1 - y_n) + (theta - 1) * h *f_n + theta * h * f_np1)
pprint(hermite_interpolation)

pprint(hermite_interpolation.expand())
pprint(sympy.collect(hermite_interpolation.expand(), y_n, evaluate=False)[y_n])
pprint(sympy.collect(hermite_interpolation.expand(), y_np1, evaluate=False)[y_np1])
pprint(sympy.collect(hermite_interpolation.expand(), f_n, evaluate=False)[f_n])
pprint(sympy.collect(hermite_interpolation.expand(), f_np1, evaluate=False)[f_np1])

"""
Explanation: cf. pp. 916 17.2.2 Dense Output, Ch. 17. Integration of Ordinary Differential Equations of Numerical Recipes, 3rd Ed., and
End of explanation
"""


