import numpy as np
from math import exp, sqrt, sin, cos, pi
from diffeq_heun  import diffeq_heun
from diffeq_euler import diffeq_euler

"""
  Solve first order diffrential equation by three-stage third-order Runge-kutta method
"""


# dx/dt = force(x,t)
# define function to be integrated
def force(t, x):
    return -x*x

# solution: x = 1 / (C + t), C = 1 for x(0) = 1.0
def fsolution(t):
    return 1.0 / (1.0 + t)

#===================
# parameters
#===================
x0 = 1.0
dt = 0.1
nt = 501
iprint_interval = 20

#===================
# main routine
#===================
def main(x0, dt, nt):
    print("Solve first order diffrential equation by three-stage third-order Runge-kutta method")

    print("{:^5}  {:^12}  {:^12}".format('t', 'x(cal)', 'x(exact)'))
    t0 = 0.0
    f0 = force(t0, x0)
# The next x (x1) must be predicted by Euler or Heum method
#    x1 = diffeq_euler(force, t0, x0, dt)
    x1 = diffeq_heun(force, t0, x0, dt)
    xexact = fsolution(t0)
    print("t={:5.2f}  {:12.6f}  {:12.6f}".format(t0, x0, xexact))
    for i in range(1, nt):
        t1 = i * dt
        k0 = dt * force(t1-dt, x0   )
        k1 = dt * force(t1,    x0+k0)
        k2 = dt * force(t1+dt, x0+4.0*k1-2.0*k0)
        x2 = x0 + 1.0 / 3.0 * (k0 + 4.0 * k1 + k2)
        xexact = fsolution(t1)
        if i % iprint_interval == 0:
            print("t={:5.2f}  {:12.6f}  {:12.6f}".format(t1, x1, xexact))
        x0 = x1
        x1 = x2

if __name__ == '__main__':
    main(x0, dt, nt)