# Author: Ian May
# File: gpplot.py
# Purpose: Make plots and movies from solution files generated
#          by my GP-WENO code

import numpy as np
import h5py
import glob
import matplotlib.pyplot as plt
import sys

# Return hdf5 dictionary corresponding to a single file
def dataFromFile(fStr):
    print('Opening ', fStr)
    f = h5py.File(fStr,'r')
    data = {}
    for k in f.keys():
        data[k] = f[k][()]
    data['ext'] = tuple(data['ext'])
    if 'time' in data:
        data['time'] = data['time'][0]
    data['filename'] = f.filename
    f.close()
    return data

# Return a list of hdf5 dictionaries for all file matching a glob
def dataFromGlob(gStr):
    dset = []
    for fStr in glob.iglob(gStr):
        dset.append(dataFromFile(fStr))
    return dset

# Plot a single field from one dataset
# Optionally supply fname to override showing the plot window
#   saving directly to a file
def plotField(d,fld,nCont=0,cmap='jet',fname='',vr=[]):
    plt.figure(figsize=(9.6,5.4),dpi=200)
    if len(vr)==2:
        plt.imshow(d[fld],origin='lower',extent=d['ext'],interpolation='gaussian',cmap=cmap,vmin=vr[0],vmax=vr[1])
    else:
        plt.imshow(d[fld],origin='lower',extent=d['ext'],interpolation='gaussian',cmap=cmap)
    plt.colorbar()
    if nCont != 0:
        if len(vr)==2:
            plt.contour(d['x'],d['y'],d[fld],nCont,colors='black',linewidths=0.5,vmin=vr[0],vmax=vr[1])
        else:
            plt.contour(d['x'],d['y'],d[fld],nCont,colors='black',linewidths=0.5)
    if 'time' in d:
        tstr = ('Time: %1.3e' % d['time'])
        plt.title(tstr,loc='left')
    plt.tight_layout()
    if fname != '':
        plt.savefig(fname)
        plt.close()
    else:
        plt.show()

# Save many plots for a given field split into frames through time
def saveFrames(gStr,fld):
    fldMin = sys.float_info.max
    fldMax = -sys.float_info.max
    for fStr in glob.iglob(gStr):
        d = dataFromFile(fStr)
        fldMin = min(np.min(d[fld]),fldMin)
        fldMax = max(np.max(d[fld]),fldMax)
    for fStr in glob.iglob(gStr):
        d = dataFromFile(fStr)
        ln = d['filename'].split('.')
        del ln[-1]
        fname = '.'.join(ln)+'_'+fld+'.png'
        plotField(d,fld,fname=fname,vr=[fldMin,fldMax])

def numericError(dc,df,fld):
    fx = df[fld].shape[1]
    fy = df[fld].shape[0]
    cx = dc[fld].shape[1]
    cy = dc[fld].shape[0]
    dx = (dc['ext'][1]-dc['ext'][0])/cx
    dy = (dc['ext'][3]-dc['ext'][2])/cy
    rx = int(fx/cx)
    ry = int(fy/cy)
    errfld = fld+'_err'
    dc[errfld] = np.zeros_like(dc[fld])
    # Verify layouts are compatible
    if rx!=ry or fx%rx+fy%ry!=0:
        print('Invalid refinement ratio\n')
        dc[errfld] = np.Nan*err
    # Compute error by averaging fine solution
    for j in range(0,cx):
        for i in range(0,cy):
            dc[errfld][i,j] = dc[fld][i,j]-np.average(df[fld][rx*i:rx*i+rx,rx*j:rx*j+rx])
    # Summarize
    return (dx*dy*np.sum(np.abs(dc[errfld])),np.max(np.abs(dc[errfld])))
    
# Calling convention: python3 gpplot.py <path/to/data.hdf5> <field name>
if __name__ == "__main__":
    # Open dataset
    print(sys.argv)
    d = dataFromFile(sys.argv[1])
    if len(sys.argv)>3:
        plotField(d,sys.argv[2],nCont=int(sys.argv[3]))
    else:
        plotField(d,sys.argv[2])