#!/usr/bin/env python
# -*- coding: utf-8 -*-

############################ WARNING ############################
# For running this script with Linux or MacOS, open a terminal in
# the subdirectory EXAMPLES, then type the command
#
#    python movie_rtin.py
#
#################################################################

import sys
sys.path.insert(0, '../PYTHON')

# Import the module for plotting the contents of HDF5 files
from plots import *
# Import the module for encoding movies
from make_movie import *



# The input files are RESULTS/EXAMPLES/test_rtin/diags_i.h5 with i=0,...,34
dir_inputh5 = '../RESULTS/EXAMPLES/test_rtin'
istart = 0
iend = 34

# We aim to encode the dynamics of
# - rho with contours
# - p with pseudocolor
# - u with vectors


# Warning: the size of output PNGs should be prescribed
figsize = (19.2,12.)

# Build the list of input files
inputh5set = []
for i in range(istart, iend+1):
	inputh5set.append(dir_inputh5+'/diags_'+str(i)+'.h5')


# Plot the contour of ux (Default: 10 isovalues that do not change with time)
plot_contour_set(inputh5set, 'Density', aliasfield='rho', tuple_inches_figsize=figsize, \
                 outputdir='./tmp', forcesave=True)

# Plot the pseudocolor of p (Default: bounds of colormap do not change with time)
plot_contour_set(inputh5set, 'Pressure', aliasfield='p', tuple_inches_figsize=figsize, \
                 outputdir='./tmp', forcesave=True)

# Plot the vector field u (Default: 200 arrows)
plot_vectorfield_set(inputh5set, 'Velocity_x', 'Velocity_y', aliasfield='u', \
                     tuple_inches_figsize=figsize, outputdir='./tmp', forcesave=True)

# For each set {diags_i_rho.png, diags_i_p.png, diags_i_u.png}, build a 2x2 mosaic named diags_i.png and prepare a list of PNGs for video encoding
# Warning: i = 000, 001, 002, ... 034
list_encoding = []
ndigits = int(numpy.ceil(numpy.log10(iend)+1))
for i in range(istart, iend+1):
	inputpngs = ['./tmp/diags_'+str(i).rjust(ndigits, '0')+'_rho.png', \
                 './tmp/diags_'+str(i).rjust(ndigits, '0')+'_p.png', \
                 './tmp/diags_'+str(i).rjust(ndigits, '0')+'_u.png']
	outputpng = './tmp/diags_'+str(i).rjust(ndigits, '0')+'.png'
	build_mosaic(inputpngs, 2, 2, outputpng)
	list_encoding.append(outputpng)

# Encode the video
encode_video(list_encoding, './tmp', 'rtin.mp4', forcesave=True)

# Remove the directory ./tmp and its contents
shutil.rmtree('./tmp')
print('Directory ./tmp deleted')