echo "------------------------------------------------------------------------------------------------"
echo "----------------------------------------Operator File Generator---------------------------------"
echo "------------------------------------------------------------------------------------------------"
echo ""
echo ""
echo ""

echo "---------This program generates the complete operator file---------"
echo "Ensure the following files exist in current working directory: "
echo ""
echo "a) frequency.dat in following format"
echo "       35.2816"
echo "      46.7087"
echo "Frequencies are in cm^-1 units......."
echo ""
echo "b) diabat-filename.log"
echo "Diabatisation output file of the optimized/reference geometry"
echo ""
echo "c) The matrix files consisting of couplings obtained from distortions"
echo "The format of matrix file must look like : "
echo "11   -0.00001040   -0.00380270   0.00479339   0.00"
echo "12   0.05437256   -0.01129996   0.02069676   0.00"
echo ""
echo "d) Operator.f90.back"
echo "A fortran code to write the coupling terms as per operator requirement"
echo ""
echo ""

read -p "Number of electronic states: " ns
echo ""
read -p "Enter the ground state diabatisation log file: " logfile
echo ""


#-----------------ENERGY-FILE-GENERATION----------------------------------------------------
> energy.dat

for state in $(seq 1 $ns)
do
    grep -H "STATE #  $state'S MCQDPT-LEVEL DIABATIC ENERGY=" "$logfile" >> diag.dummy
done

    sed -i  "s/.*ENERGY= *//"  diag.dummy

E0=$(head -n 1 diag.dummy)

i=1
while read Ei
do
    if [ $i -eq 1 ]; then
        echo "0.00000000" >> energy.dat
    else
        awk -v e="$Ei" -v e0="$E0" 'BEGIN{printf "%.8f\n", 27.2*(e-e0)}' >> energy.dat
    fi
    i=$((i+1))
done < diag.dummy

for state1 in $(seq 1 $ns)
do
for state2 in $(seq 1 $ns)
do
    if [ $state1 -lt $state2 ]; then

        grep -H "STATE #  $state1 &  $state2'S MCQDPT-LEVEL COUPLING" "$logfile" > tmp.dummy
         sed -i "s/.*= *//"  tmp.dummy

        if [ -s tmp.dummy ]; then
            val=$(head -n 1 tmp.dummy)
            awk -v v="$val" 'BEGIN{printf "%.8f\n", 27.2*v}' >> energy.dat
        else
            echo "0.00000000" >> energy.dat
        fi
    fi
done
done

rm -rf *.dummy
echo "energy.dat generated"
echo ""
#-----------------------------------------------------------------------------------------------------
echo ""
echo "Your matrix files must be named such as modenumber+suffix"
echo " e.g.    05-matrix-all.dat /11-shift.dat /81-unitary.dat  " 
echo "suffix  => -matrix-all.dat /  -shift.dat /  -unitary.dat  " 
read -p "Enter the suffix for your files : " suffix
echo ""
echo ""



# ---------------- FILE CHECKS -----------------------------------------------------------------------
if [ ! -f frequency.dat ]; then
    echo "ERROR: frequency.dat not found"
    exit 1
fi

if [ ! -f $logfile ]; then
    echo "ERROR: $logfile not found"
    exit 1
fi

if [ ! -f energy.dat ]; then
    echo "ERROR: energy.dat not found"
    exit 1
fi

# check matrix files
for m in $modes
do
    file=$(printf "%02d$suffix" $m)
    if [ ! -f $file ]; then
        echo "ERROR: Missing $file"
        exit 1
    fi
done


echo "................................Done with the file check "
echo "Everthing intact: Moving to the operator file generation"
echo ""
# ---------------- INPUT -----------------------------------------------------------------------------
read -p "Operator file name (without .op): " opname
echo ""
read -p "Number of modes: " nv
echo ""
read -p "Enter mode numbers (space separated): " modes
echo ""
read -p "Title: " title
echo ""

outfile="$opname.op"
# ----------------START-OP----------------------------------------------------------------------------
echo "#Operator file for $title" > $outfile
echo "OP_DEFINE-SECTION" >> $outfile
echo "title" >> $outfile
echo "$title with $ns electronic states and $nv modes" >> $outfile
echo "end-title" >> $outfile
echo "end-op_define-section" >> $outfile
echo "" >> $outfile
echo "PARAMETER-SECTION" >> $outfile
echo "# vib. freq."      >> $outfile
echo "# modes included are mode number" >> $outfile

#------------------------ADD-MODES--------------------------------------------------------------------
line="# "
count=0
i=1

for m in $modes
do
freq=$(awk -v mode=$m 'NR==mode {printf "%.1f", $1}' frequency.dat)

entry="v$i=$m (${freq} cm^-1)"

# add comma if not first entry
if [ $count -ne 0 ]; then
    line="${line}, "
fi

line="${line}${entry}"

count=$((count+1))
i=$((i+1))

# if 5 entries reached, print line and reset
if [ $count -eq 5 ]; then
    echo "$line" >> $outfile
    line="# "
    count=0
fi
done

# print remaining (if any)
if [ $count -ne 0 ]; then
    echo "$line" >> $outfile
fi

# -------- Frequencies -------------------------------------------------------------------------------
i=1
for m in $modes
do
freq=$(awk -v mode=$m 'NR==mode {print $1}' frequency.dat)

if [ -z "$freq" ]; then
    echo "ERROR: Mode $m not found in frequency.dat"
    exit 1
fi

wev=$(awk -v f=$freq 'BEGIN{printf "%.11f", f*1.23981e-4}')
echo "w$i = $wev, ev" >> $outfile
i=$((i+1))
done

# -------- Energies ----------------------------------------------------------------------------------
echo "#intra state energies" >> $outfile

readarray -t energy < energy.dat

if [ ${#energy[@]} -lt $((ns*(ns+1)/2)) ]; then
    echo "ERROR: energy.dat does not have enough entries"
    exit 1
fi

index=0

# diagonal
for ((i=1;i<=ns;i++))
do
echo "e$i$i =  ${energy[index]}, ev" >> $outfile
index=$((index+1))
done

# off-diagonal
echo "#inter state energies" >> $outfile
for ((i=1;i<=ns;i++))
do
for ((j=i+1;j<=ns;j++))
do
echo "e$i$j = ${energy[index]}, ev" >> $outfile
index=$((index+1))
done
done

# -------- Couplings ---------------------------------------------------------------------------------
echo "#first and second order elements" >> $outfile

mode_index=1

for m in $modes
do
file=$(printf "%02d$suffix" $m)

cp operator.f90.back operator.f90
sed -i "s/input_file/$file/g" operator.f90
sed -i "s/no_of_states/$ns/g" operator.f90
sed -i "s/mode_number/$mode_index/g" operator.f90
 
gfortran operator.f90 -o operator.exe
./operator.exe

more operator-elements.out >> $outfile

mode_index=$((mode_index+1))
rm -f operator-elements.out operator.f90 operator.exe
done

echo "end-parameter-section" >> $outfile
echo "" >> $outfile

# ---------------- HAMILTONIAN -------------------------------------------------------------------------
echo "HAMILTONIAN-SECTION" >> $outfile
echo "------------------------------------------------------------------" >> $outfile


count=0
line="modes |"

for ((i=1;i<=nv;i++))
do
if [ $i -lt 10 ]; then
    entry=" v$i  |"
else
    entry=" v$i |"
fi

line="${line}${entry}"
count=$((count+1))

if [ $count -eq 10 ]; then
    line=${line%|}   # remove last pipe
    echo "$line" >> $outfile
    line="modes |"
    count=0
fi
done

# last line
if [ $count -ne 0 ]; then
    line=${line%|}
    line="${line}| el"
    echo "$line" >> $outfile
fi

echo "------------------------------------------------------------------" >> $outfile



#--------- KINETIC+POTENTIAL--------------------------------------------------------------------------
for ((i=1;i<=nv;i++))
do
echo "w$i            |$i KE" >> $outfile
done

for ((i=1;i<=nv;i++))
do
echo "0.5*w$i          |$i   q^2" >> $outfile
done
echo "" >> $outfile

# ----------ELECTRONIC ENERGIES-----------------------------------------------------------------------
for ((i=1;i<=ns;i++))
do
for ((j=i;j<=ns;j++))
do
echo "e$i$j            |$((nv+1)) S$i&$j" >> $outfile
done
done

echo " " >> $outfile

# -------------FIRST ORDER COUPLING-------------------------------------------------------------------
for ((v=1;v<=nv;v++))
do
for ((i=1;i<=ns;i++))
do
for ((j=i;j<=ns;j++))
do
echo "h$i$j""_k$v    |$v q  |$((nv+1)) S$i&$j" >> $outfile
done
done
done

echo " " >> $outfile
#--------------- SECOND ORDER COUPLING----------------------------------------------------------------
for ((v=1;v<=nv;v++))
do
for ((i=1;i<=ns;i++))
do
for ((j=i;j<=ns;j++))
do
echo "0.5*h$i$j""_g$v    |$v q^2  |$((nv+1)) S$i&$j" >> $outfile
done
done
done

echo "end-hamiltonian-section" >> $outfile

#---------------DIAGONALISATION-HAMILTONIAN-----------------------------------------------------------
for ((v=1; v<=nv; v++))
do
echo "" >> $outfile
echo "hamiltonian-section_h$v" >> $outfile
echo "modes | v$v" >> $outfile

printf "w%-4d |1 KE\n" $v >> $outfile
printf "0.5*w%d |1 q^2\n" $v >> $outfile
printf "h11_k%-2d |1 q\n" $v >> $outfile
printf "0.5*h11_g%d |1 q^2\n" $v >> $outfile

echo "end-hamiltonian-section" >> $outfile
done

echo "" >> $outfile
echo "end-operator" >> $outfile


echo "The operator file is generated: $outfile"
