#!/bin/bash

usage() {
	echo 
	echo "Calculates viscosity using time decomposition method"
	echo "J. Chem. Theory Comput. 2015, 11, 3537−3546"
	echo
	echo "USAGE"
	echo "	-b Number of blocks"
	echo "	-t Time length of each block (ps)"
	echo "	-x Total time length of full trajectory (ps))"
	echo "	-g Displays graphs of running integral and its standard deviation"
	echo 
	exit 1
}

# setting working directory and GROMACS stuff
WORKDIR="$(dirname "$(readlink -f "$0")")"
PWD="$(pwd)"
source /usr/local/gromacs/bin/GMXRC

# set variables from command line
GRAPH=false
while getopts ":b:t:x:gh" options; do
	case ${options} in 
		b)
			BLOCKNUM=$OPTARG ;;
		t)
			BLOCKTIME=$OPTARG ;;
		x)
			FULLTIME=$OPTARG ;;
		g)
			GRAPH=true ;;
		h)
			usage
			exit 1 ;;
		\?)
			echo "Unknown flag. Please fix input and check the help with -h"
			exit 1 ;;
		:)
			echo "Invalid input. You must input an argument for -${OPTARG}. Check the help with -h"
			exit 1 ;;
	esac
done

# check for input on command line
if [ -z $BLOCKNUM ] || [ -z $BLOCKNUM ] | [ -z $BLOCKNUM ]; then
	echo "Missing one of the required flags -b, -t, or -x for command line input. Stopping script"
	exit 1
fi


# Check to make sure the trajectory can be broken down evenly
# if blocks can split the trajectory evenly, then mention perfect even split
# else if total blocking time is less than full trajectory time, then still possible to do blocking
# else consists of instances where total blocking time exceeds full trajectory time, which cannot be done
BLOCKTOTALTIME=$(( $BLOCKNUM * $BLOCKTIME ))
if [ $BLOCKTOTALTIME -eq $FULLTIME ]; then
	echo "Number of blocks $BLOCKNUM of length $BLOCKTIME ps split the entire trajectory of time $FULLTIME ps evenly."
elif [ $BLOCKTOTALTIME -lt $FULLTIME ]; then
	echo "Number of blocks $BLOCKNUM of length $BLOCKTIME ps can be taken from part of the trajectory of time $FULLTIME ps."
	echo "Will only use first $BLOCKTOTALTIME ps of full trajectory $FULLTIME ps for blocking."
	echo "Number of blocks $BLOCKNUM of length $BLOCKTIME ps taken from first $BLOCKTOTALTIME ps of trajectory."
	echo "If this is not what you want, then please stop the script and re-enter the correct blocking values."
else 
	echo "Number of blocks $BLOCKNUM of length $BLOCKTIME ps cannot be taken from trajectory of time $FULLTIME ps."
	echo "Please fix the number of blocks and their lengths so it does fit in the trajectory."
	exit 1
fi
	
# Running green-kubo calculations
mkdir "$WORKDIR"/viscosity
VISCODIR="$WORKDIR"/viscosity
VISCONUM=1
BEGINTIME=0
ENDTIME=$BLOCKTIME
NVTBOXLENGTH=$(tail -1 "$WORKDIR"/*.gro | rev | cut -d " " -f 1 | rev ) # box length in nm
NVTVOLUME=$(echo "${NVTBOXLENGTH}^3" | bc -l )  # box volume in nm^3
echo "Box volume is $NVTVOLUME nm^3"
while [ $VISCONUM -lt $(( $BLOCKNUM + 1 )) ]; do
	echo $NVTVOLUME | gmx energy \
		-f "$WORKDIR"/out/*.edr \
		-o "$VISCODIR"/energy_visco${VISCONUM}.xvg \
		-vis "$VISCODIR"/visco${VISCONUM}.xvg \
		-evisco "$VISCODIR"/evisco${VISCONUM}.xvg \
		-eviscoi "$VISCODIR"/eviscoi${VISCONUM}.xvg \
		-b $BEGINTIME \
		-e $ENDTIME \
		-quiet
	mv "$PWD"/enecorr.xvg "$VISCODIR"/enecorrbulkvisc${VISCONUM}.xvg 2>/dev/null; true
	mv "$PWD"/"#enecorr.xvg.1#" "$VISCODIR"/enecorrshearvisc${VISCONUM}.xvg 2>/dev/null; true
	let BEGINTIME=$BEGINTIME+$BLOCKTIME
	let ENDTIME=$ENDTIME+$BLOCKTIME
	let VISCONUM++
done

if [ "$GRAPH" = true ]; then  # if we want a graph (i.e. GRAPH = TRUE)
	python3 /home/mmh/BrockportMDFiles/python_scripts/viscosity.py "$VISCODIR"/visco*.xvg -g
else
	python3 /home/mmh/BrockportMDFiles/python_scripts/viscosity.py "$VISCODIR"/visco*.xvg
fi


