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

import seaborn as sns 
params = {
    "font.family": "Serif",
    "font.serif": "Roman",
    "text.usetex": True,
    "axes.titlesize": "large",
    "axes.labelsize": "large",
    "xtick.labelsize": "large",
    "ytick.labelsize": "large",
    "legend.fontsize": "medium",
}
plt.rcParams.update(params)

wave_arr = []
fig = plt.figure()
ax = plt.gca()
img = ax.imshow([[]])

def plot():
    with open("test.txt") as f:
        lines = f.readlines();
        size = int(lines[0])
        for line in lines[1:]:
            arr = line.strip().split("\t")
            arr = np.asarray(list(map(lambda x: ((a := complex(*ast.literal_eval(x)))*a.conjugate()).real, arr)))

            # print(sum(arr))
            arr = arr.reshape(size,size)
            wave_arr.append(arr.T)
            # print(arr)

            # plt.imshow(arr, cmap="hot", interpolation="nearest")
            # plt.show()


def animation(i):
    norm = matplotlib.cm.colors.Normalize(vmin=0, vmax=np.max(wave_arr[i]))
    img.set_norm(norm)
    img.set_data(wave_arr[i])


if __name__ == "__main__":
    plot()
    
    norm = matplotlib.cm.colors.Normalize(vmin=0, vmax=np.max(wave_arr[0]))
    img = ax.imshow(wave_arr[0], extent=[0,1,0,1], cmap=plt.get_cmap("viridis"), norm=norm)
    anim = FuncAnimation(fig, animation, interval=1, frames=np.arange(0,len(wave_arr)), repeat=True, blit=0)
    # plt.show()

    anim.save("./animation.mp4", writer="ffmpeg", bitrate=10000, fps=15)