#! /usr/bin/env python

"""
Plot the time-dependent density expectation value in Lieb-Liniger (n-hole wavepacket).

Setup:
- ensure matplotlib and numpy are installed in your python environment
- make this file executable (chmod +x plot_animated.py)

Usage: ./plot_animated.py [rho file name]
"""


import math
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import numpy as np
import sys

filename = str(sys.argv[1])

data = filename.rpartition('/')[2]

data = data.partition('c_')[2]
c = float(data.partition('_')[0])
data = data.partition('N_')[2]
N = int(data.partition('_')[0])
L = float(N)
data = data.partition('nholes_')[2]
nholes = int(data.partition('_')[0])
data = data.partition('width_')[2]
width = int(data.partition('_')[0])
data = data.partition('offset_')[2]
offset = int(data.partition('_')[0])
data = data.partition('tmax_')[2].partition('.rhoxt')[0]
tmax = float(data.partition('_')[0])


x = np.loadtxt(filename.replace('.rhoxt','.x'), unpack=True)

Nx = x.size

rho = np.loadtxt(filename, unpack=True)

Nt = int(rho.size/rho[0].size)
dt = tmax/(Nt-1)

fig, ax = plt.subplots()

artists = []

for i in range(Nt):
    container = ax.plot(x, rho[i], '.', markersize=4, color='blue')
    artists.append(container)

ani = animation.ArtistAnimation(fig=fig, artists=artists, interval=100)

plt.xlabel('x')
plt.ylabel(r'$\rho(x, t)$')
plt.title(f"Density expectation value\n{c=}, {N=}, nh={nholes}, w={width}, o={offset}")

ani.save(filename.replace("rhoxt", "mp4"))

plt.show()