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.


varlist = ['density']
axislist = ['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/'

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)