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

"""
  Solve second order diffrential equation by Verlet method
"""

# d2x/dt2 = force(x,t)
# define function to be integrated
def force(t, x):
    return -x

# solution: x(0.0) = 0.0 v(0.0) = 1.0
def xsolution(t):
    return sin(t)

#===================
# parameters
#===================
outfile = 'diffeq2nd_verlet.csv'
x0 = 0.0
v0 = 1.0
dt = 0.01
nt = 501
iprint_interval = 20

#===================
# main routine
#===================
def main(x0, v0, dt, nt):
    print("Solve second order diffrential equation by Verlet method")

    print("Write to [{}]".format(outfile))

# open outfile to write a csv file
    f = open(outfile, 'w')
    fout = csv.writer(f, lineterminator='\n')
    fout.writerow([
        't', 'x(cal)', 'x(exact)', 'v(cal)'
        ])

    print("{:^5}  {:^12}  {:^12}  {:^12}".format('t', 'x(cal)', 'x(exact)', 'v(cal)'))
    kv0 = dt * force(0.0, x0)
    kv1 = dt * force(0.0+dt, x0+kv0)
    v1 = v0 + (kv0 + kv1) / 2.0
    kx0 = dt * v0
    kx1 = dt * v1
    x1  = x0 + (kx0 + kx1) / 2.0
    xexact = xsolution(0.0)
    print("t={:5.2f}  {:12.6f}  {:12.6f}  {:12.6f}".format(0.0, x0, xexact, v0))
    fout.writerow([0.0, x0, xexact, v0])
    for i in range(1, nt):
        t1 = i * dt
        f1 = force(t1, x1)
        x2 = 2.0 * x1 - x0 + dt * dt * f1
        v1 = (x2 - x0) / 2.0 / dt
        xexact = xsolution(t1)
        if i == 1 or i % iprint_interval == 0:
            print("t={:5.2f}  {:12.6f}  {:12.6f}  {:12.6f}".format(t1, x1, xexact, v1))
        fout.writerow([t1, x1, xexact, v1])
        x0 = x1
        v0 = v1
        x1 = x2

    f.close()

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