import sys
import csv
import numpy as np
from math import exp, sqrt, sin, cos, pi
import matplotlib.pyplot as plt


"""
  Solve first order diffrential equation by Heun method
"""


#===================
# parameters
#===================
outfile = 'diffeq_euler_heun.csv'
x0 = 1.0
dt = 0.1
nt = 501
iprint_interval = 20

argv = sys.argv
n = len(argv)
if n >= 2:
    x0 = float(argv[1])
if n >= 3:
    dt = float(argv[2])
if n >= 4:
    nt = int(argv[3])
if n >= 5:
    iprint_interval = int(argv[4])


# dx/dt = dxdt(x,t)
# define function to be integrated
def dxdt(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_euler(diff1func, t0, x0, dt):
   k1 = dt * diff1func(t0, x0)
   x1 = x0 + k1
   return x1

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


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

# time
    xt      = [0.0]
# analytical x(t)
    yxex    = [x0]
# x(t) by Euler and Heun methods
    yxeuler = [x0]
    yxheun  = [x0]
# error
    yeeuler = [0.0]
    yeheun  = [0.0]

# prepare for graph
    fig = plt.figure(figsize = (8, 8))
    ax1 = fig.add_subplot(3, 1, 1)
    ax2 = fig.add_subplot(3, 1, 2)
    ax3 = fig.add_subplot(3, 1, 3)
# 凡例を表示させるため、とりあえずplot()を呼び出す
# 後でプロット毎にデータリストを再設定するので、lineオブジェクトを受け取っておく
    line11, = ax1.plot(xt, yxeuler, label = 'Euler')
    line12, = ax1.plot(xt, yxheun,  label = 'Heun')
    line13, = ax1.plot(xt, yxex,    label = 'exact')
    line21, = ax2.plot(xt, yeeuler, label = 'Euler')
    line31, = ax3.plot(xt, yeheun,  label = 'Heun')
#    ax1.set_xscale('log')
#    ax1.set_yscale('log')
    ax1.set_xlabel("t")
    ax1.set_ylabel("x(t)")
    ax1.legend()
    ax2.set_xlabel("t")
    ax2.set_ylabel("error")
    ax2.legend()
    ax3.set_xlabel("t")
    ax3.set_ylabel("error")
    ax3.legend()

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

    xeuler = x0
    xheun  = x0
    print("{:^5}  {:^12}  {:^12}  {:^12}".format('t', 'x(cal)', 'x(euler)', 'x(heun)'))
    for i in range(1, nt):
        t0 = i * dt
        xeuler = diffeq_euler(dxdt, t0, xeuler, dt)
        xheun  = diffeq_heun(dxdt, t0, xheun, dt)
        xexact = fsolution(t0)

        xt.append(t0)
        yxex.append(xexact)
        yxeuler.append(xeuler)
        yxheun.append(xheun)
        yeeuler.append(xeuler - xexact)
        yeheun.append(xheun - xexact)

# graphをupdateするには、プロットデータ line1/line2 に set_data() でデータリストを設定し、plt.pause()を呼び出す
# set_data() ではグラフの表示範囲は更新されないので、データの最小・最大値で設定する
        line11.set_data(xt, yxeuler)
        line12.set_data(xt, yxheun)
        line13.set_data(xt, yxex)
        line21.set_data(xt, yeeuler)
        line31.set_data(xt, yeheun)
        ax1.set_xlim((min(xt), max(xt)))
        ax1.set_ylim((min(yxex), max(yxex)))
        ax2.set_xlim((min(xt), max(xt)))
        ax2.set_ylim((min(yeeuler), max(yeeuler)))
        ax3.set_xlim((min(xt), max(xt)))
        ax3.set_ylim((min(yeheun), max(yeheun)))
        plt.pause(0.1)

        if i == 1 or i % iprint_interval == 0:
            print("t={:5.2f}  {:12.6f}  {:12.6f}  {:12.6f}".format(t0, xexact, xeuler, xheun))

        fout.writerow([t0, x0, xeuler, xheun])

    f.close()

    print("Press ENTER to exit>>", end = '')
    input()

    exit()

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