import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# データ取得
data_x, data_psi_r, data_psi_i, data_pot = np.loadtxt("./output.dat", comments='!', unpack=True)

# グリッド数、スナップショット数取得
num_grid = 1202
num_tot = len(data_x)
num_snapshots = int(num_tot/num_grid)

data_prob = (data_psi_r*data_psi_r + data_psi_i*data_psi_i)*10

#print(num_tot)
#print(num_snapshots)

# 図、枠（軸）の設定
fig = plt.figure()
ax = fig.add_subplot(111)


def plot(num):
    # clear previous frames
    ax.cla()
    # Axesのタイトルを設定
    ax.set_title("1D wave fn.", size = 20)

    # 軸ラベルを設定
    ax.set_xlabel("x", size = 14, color = "k")
    ax.set_ylabel("$\psi$", size = 14, color = "k")

    # 軸目盛を設定
#    ax.set_xticks([-2, -1, 0, 1, 2])
#    ax.set_yticks([-2, -1, 0, 1, 2])


    # 軸範囲の設定
    ax.set_xlim(-10, 310)
    ax.set_ylim(-1, 1)


    ax.plot(data_x[num*num_grid:(num*num_grid+num_grid)], data_psi_r[num*num_grid:(num*num_grid+num_grid)],  "-", color="r", label="Real part")
    ax.plot(data_x[num*num_grid:(num*num_grid+num_grid)], data_psi_i[num*num_grid:(num*num_grid+num_grid)],  "-", color="b", label="Imaginary part")
    ax.plot(data_x[num*num_grid:(num*num_grid+num_grid)], data_pot[num*num_grid:(num*num_grid+num_grid)],  "-", color="gray", label="potential")
    ax.plot(data_x[num*num_grid:(num*num_grid+num_grid)], data_prob[num*num_grid:(num*num_grid+num_grid)],  "-", color="green", label="probability")
    ax.legend(loc="upper left")


ani = FuncAnimation(fig, plot, frames=num_snapshots, interval=50)

#print(frames)
#ani.save("quantum1.gif", writer="pillow")
ani.save('quantum1.mp4', writer="ffmpeg")

plt.show()