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

"""
  Solve first order diffrential equation by Heun 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)

def diffeq_heun(force, t0, x0, dt):
   k0 = dt * force(t0, x0)
   k1 = dt * force(t0+dt, x0+k0)
   x1 = x0 + (k0 + k1) / 2.0
   return x1


#===================
# 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 Heun method")

    print("{:^5}  {:^12}  {:^12}".format('t', 'x(cal)', 'x(exact)'))
    for i in range(nt):
        t = i * dt
        x1 = diffeq_heun(force, t, x0, dt)
        xexact = fsolution(t)
        if i == 1 or i % iprint_interval == 0:
            print("t={:5.2f}  {:12.6f}  {:12.6f}".format(t, x0, xexact))
        x0 = x1

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