#! /usr/bin/env python

"""
Plot DSF.

Usage: python plot_dsf_k_range.py [omega file] [dsf file]
"""

import glob
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import numpy as np
import os
import sys

N = str(sys.argv[1])
width = str(sys.argv[2])


# Get the list of interactions which have been computed
dirlist = os.listdir('../data/N_%s' % N)
clist = sorted([float(c.lstrip('c_')) for c in dirlist])

# Get the K file
kfile = glob.glob('../data/N_%s/c_%s*/K_*' % (N, str(clist[0]).rstrip('.0')[:12]))[0]
k = np.loadtxt(kfile)

# Get the Omega file
omegafile = glob.glob('../data/N_%s/c_%s*/Omega*' % (N, str(clist[0]).rstrip('.0')[:12]))[0]
omega = np.loadtxt(omegafile)

# Load all the available dsfs from the data store
dsfs = {}
for c in clist:
    # Do some black magic here: for matching the interaction,
    # first try for exact match, stripping '.0' to treat integer values, e.g. 4.0 into 4
    # and then (if it doesn't work) use only the first 12 characters, plus wildcard, to cover rounding errors
    try:
        dsffile = glob.glob('../data/N_%s/c_%s/*_w_%s.dsfs' % (N, str(c).rstrip('.0'), width))[0]
    except IndexError:
        dsffile = glob.glob('../data/N_%s/c_%s*/*_w_%s.dsfs' % (N, str(c).rstrip('.0')[:12], width))[0]
    dsfs[str(c)] = np.loadtxt(dsffile)

# Read some useful parameters from (last) file name:
elements = dsffile.rpartition('/')[2].split('_')
L = elements[5]
#N = elements[7]
iKmax = elements[14]
width = elements[22].rpartition('.')[0]
# which leads to
rho = int(N)/int(L)
#kokF = int(iK)*0.5/int(L)


fig, ax = plt.subplots()

# To determine the y axis limits, we look at the lowest value of c (sharpest peak)
dsfsmax = 0.5

ymax = omega[-1]
xtext = 0.05 * k[-1]
ytext = 0.9 * ymax

ax.text(xtext, ytext, f'c = {clist[0]:10.6f}', color='white', fontsize='large')
pcm = ax.pcolormesh(k, omega, dsfs[str(clist[0])], vmax=dsfsmax)
fig.colorbar(pcm, ax=ax)
def animate(i):
    ax.clear()
    ax.set_title(f'c scan, rho={rho} (N={N}), w={width}', fontsize='x-large')
    ax.text(xtext, ytext, f'c = {clist[i]:10.6f}', color='white', fontsize='x-large')
    ax.pcolormesh(k, omega, dsfs[str(clist[i])], vmax=dsfsmax)

ani = animation.FuncAnimation(fig, animate, frames=len(clist), interval=100, repeat=False)

outfilename = (dsffile.partition('_c_')[0].rpartition('/')[2] + '_c_scan_' +
               dsffile.partition('_c_')[2].partition('_')[2].partition('.dsfs')[0])

ani.save(outfilename + '.mp4')

with open(outfilename + '.html', 'w') as file:
    file.write(ani.to_html5_video())

#plt.show()