import matplotlib
matplotlib.use('Agg')

import yt
from yt import derived_field
import numpy as np
import matplotlib.pyplot as plt
import os
import glob
import math





runNumber = '176_g' # enter the number of the run (last part of the name of the flashFolder directory)
createArchive = True # False if you DON'T WANT a .tar.gz archive containing the entire all_Slices_Run_<runNumber>
basename = 'super3d_' # sets the prefix of the plotfiles and checkpoint files. Usually 'super3d_' for MHD, 'relax_' for purehydro.


def main(filenames):
	 for n in range (len(filenames)):
                ds = yt.load(filenames[n])
		grad = ds.add_gradient_fields(("gas","gravitational_potential"))
                time[n] =ds.current_time
                #dd= ds.all_data()

                
####...x-component of negative of gradient of gravitational potential...
def _GradientX(field, data):
     Gx = -1.0*data["gravitational_potential_gradient_x"]
     return Gx
yt.add_field("Gx", function=_GradientX, take_log=False, units="code length/ code time**2")

####...y-component of negative of gradient of gravitational potential...
def _GradientY(field, data):
     Gy = -1.0*data["gravitational_potential_gradient_y"]
     return Gy
yt.add_field("Gy", function=_GradientY, take_log=False, units="code length/ code time**2")

####...z-component of negative of gradient of gravitational potential...
def _GradientZ(field, data):
     Gz = -1.0*data["gravitational_potential_gradient_z"]
     return Gz
yt.add_field("Gz", function=_GradientZ, take_log=False, units="code length/ code time**2")







varlist = ['Gx', 'Gy', 'Gz']
axislist = ['z', 'z', 'z']

# ---------------------------------

if len(varlist) != len(axislist):
        print('Unequal number of entries in varlist and axislist, please check!')
# ===========================================

# Gets the location of the plotfiles and creates a list of the filenames

#scriptFolder = os.getcwd()
flashFolder = '/work/03858/thaque56/runs/run_176/'

# uncomment the following line if your plotfiles are located at a different location, and specify the path:
#flashFolder = '/barycenter/maengenheyster/Stampede_Runs/super3D_Run_'+runNumber

if os.path.isdir(flashFolder)==False:
        print('flashFolder directory does not exist, enter correct one!')

# filenames
plotFilenames = glob.glob(flashFolder + basename + 'hdf5_plt_cnt_[0-9][0-9][0-9][0-9]') # list of plotfiles
plotFilenames.sort()

if plotFilenames == []:
        print('No plotfiles found at '+flashFolder+', check path and filenames!')

# ============================================
# checks if main directory sliceFolder and subdirectories for the plots plotFolders exist, if not, creates them.
# main directory
sliceFolder = '/work/03858/thaque56/runs/run_176/'+runNumber

if os.path.isdir(sliceFolder)==False:
        os.system('mkdir '+sliceFolder+'')

os.chdir(sliceFolder)

# subdiretories for slice plots
plotFolders = [] # list of subfolders for the different slice plots
for n in range(len(varlist)):
        plotFolderName = 'Slice_'+axislist[n]+'_'+varlist[n]
        if os.path.isdir(plotFolderName)==False:
                os.system('mkdir '+plotFolderName)
        plotFolders.append(plotFolderName)

#=========================

for filename in plotFilenames:

  ds=yt.load(filename)
  for i in range(len(varlist)):
          slc =yt.SlicePlot(ds,axislist[i],varlist[i],origin="center-domain")
          slc.zoom(8)
          slc.annotate_title('t= '+str(ds.current_time))
          slc.annotate_grids()
          if(varlist[i]=='temp'):slc.set_cmap(varlist[i],'gist_ncar')
          slc.save(plotFolders[i]+'/'+ds.basename)

#os.chdir(scriptFolder)
os.chdir(sliceFolder)
if createArchive == True:
        zip_command = 'tar -zcvpf '+sliceFolder+'.tar.gz '+sliceFolder+'/'
        os.system(zip_command)