import numpy as np
import matplotlib.pyplot as plt
from scipy import ndimage as ndi
import os
import re




def atoi(text):
    return int(text) if text.lstrip('-+').isdigit() else text

def natural_keys(text):
    '''
    alist.sort(key=natural_keys) sorts in human order
    http://nedbatchelder.com/blog/200712/human_sorting.html
    (See Toothy's implementation in the comments)
    '''
    return [ atoi(c) for c in re.split(r'(-?\d+)', text) ]


fsize=[10,6]
fig = plt.figure(1, figsize=fsize)

gs1 = plt.GridSpec(1, 1)
ax1 = plt.subplot(gs1[0])


what = 'HF_occ.dat'
state = 28

prefix = os.path.dirname(os.path.abspath(__file__))+'/'

folds = os.listdir(prefix)
folds = sorted([x+'/' for x in folds if 'delay' in x], key=natural_keys)

ref = np.genfromtxt(prefix+'noNIR/'+what)[:,state]

color_list = plt.cm.Spectral(np.linspace(0,1, len(folds)))[::-1]

for ff,fold in enumerate(folds):
    tmp = np.genfromtxt(prefix+fold+what)
    x = tmp[:, 0]*0.024
    y = tmp[:, state]-ref
    
    lab = int(fold.strip('delay').strip('/'))*0.024
    
    ax1.plot(x, y, label=r"$\tau=%.2f$~fs"%lab, color=color_list[ff])


ax1.ticklabel_format(style='sci', axis='y', scilimits=(0,0))
ax1.set_xlabel(r"$t (\rm{fs})$")
ax1.set_ylabel(r"$\delta n_{%d}(t)$"%state)


ax1.legend(ncol=2)

plt.show()