#!/bin/bash

# ==============================================================================
# Configuration
# ==============================================================================
LODIS_EXE="/home/tmlkyza/Documents/Universita/Master/nanopants/lcm/LoDiS/LODIS_all"
POT_FILE="/home/tmlkyza/Documents/Universita/Master/nanopants/lcm/POT_FILE/Ag_Ag.pot"
END_NUM=$2
SCRIPT_DIR=$(pwd)
CONFIG_FILE=$SCRIPT_DIR"/in_g.in"
CONFIG_MELT_FILE=$SCRIPT_DIR"/in_melt.in"
SEED_FILE=$SCRIPT_DIR"/"$1

if [ -z "$1" ] || [ -z "$2" ]; then
    echo "Usage: $0 [SEED_FILE] [END_NUM]"
    exit 1
fi

# ==============================================================================
# Helper Functions
# ==============================================================================

checkStart(){
    # Extract starting atom count from filename
    FileNum=$(echo "$1" | grep -o -E '[0-9]+' | tail -1)
    
    # Update master config with quotes for Fortran strings
    sed -i "s|natom *= *.*|natom = ${FileNum}|" "$CONFIG_FILE"
    sed -i "s|filepos *= *.*|filepos = '$SEED_FILE'|" "$CONFIG_FILE"
    sed -i "s|filepot *= *.*|filepot = '$POT_FILE'|" "$CONFIG_FILE"
    
    START_NUM=$FileNum
}

getLastTemp(){
    ENERGYOUT=$1
    COLUMN=$2
    temp=$(tail -n 1 "$ENERGYOUT" | awk -v col="$COLUMN" '{print int($col) ".d0"}')
    echo $temp
}

runGrowth(){
    local STAGE_DIR="$SCRIPT_DIR/growthSims"
    mkdir -p "$STAGE_DIR"
    cp "$SEED_FILE" "$STAGE_DIR/seed.xyz"
    cp "$POT_FILE" "$STAGE_DIR/pot.pot"
    
    cd "$STAGE_DIR" || exit
    
    ITERATIONS=$(($END_NUM - $START_NUM))

    for (( i=1; i<=$ITERATIONS; i++ )); do
        CURRENT_NATOMS=$(($START_NUM + i))
        mkdir -p "run_$CURRENT_NATOMS"
        cd "run_$CURRENT_NATOMS" || exit

        cp "$CONFIG_FILE" "growth.in"

        sed -i "s|filepos *= *.*|filepos = '../seed.xyz'|" "growth.in"
        sed -i "s|filepot *= *.*|filepot = '../pot.pot'|" "growth.in"
        sed -i "s|ndepmax *= *.*|ndepmax = $i|" "growth.in"
        
        "$LODIS_EXE" < "growth.in" > "growth.log" &
        
        cd ..
    done

    echo "Waiting for all Growth processes to finish..."
    wait 
    cd "$SCRIPT_DIR"
}

runMelting(){
    local INPUT_DIR=$1
    local STAGE_DIR="$SCRIPT_DIR/meltingSims"

    
    mkdir -p "$STAGE_DIR"
    cp "$POT_FILE" "$STAGE_DIR/pot.pot"
    
    cd "$STAGE_DIR" || exit

    for dir in "$INPUT_DIR"/run_*; do
        [ -e "$dir" ] || continue
        natoms=$(echo "$dir" | grep -o -E '[0-9]+' | tail -1)
        
        mkdir -p "run_$natoms"
        cd "run_$natoms" || exit

        # Extract last frame from the PROVIDED input directory
        tail -n $(($natoms + 2)) "$dir/movie.xyz" > "toMelt_$natoms.xyz"

        cp "$CONFIG_MELT_FILE" "melt.in"

        initTemp=$(getLastTemp "$INPUT_DIR/run_$natoms/energy.out" "8")
        
        sed -i "s|tinit *= *.*|tinit = $initTemp|" "melt.in"
        sed -i "s|natom *= *.*|natom = $natoms|" "melt.in"
        sed -i "s|filepos *= *.*|filepos = 'toMelt_$natoms.xyz'|" "melt.in"
        sed -i "s|filepot *= *.*|filepot = '../pot.pot'|" "melt.in"

        "$LODIS_EXE" < "melt.in" > "melt.log" &
        
        cd ..
    done

    echo "Waiting for all Melting processes to finish..."
    wait 
    cd "$SCRIPT_DIR"
}

runFreezing(){
    local INPUT_DIR=$1
    local STAGE_DIR="$SCRIPT_DIR/freezeSims"
    
    mkdir -p "$STAGE_DIR"
    cp "$POT_FILE" "$STAGE_DIR/pot.pot"
    
    cd "$STAGE_DIR" || exit

    for dir in "$INPUT_DIR"/run_*; do
        [ -e "$dir" ] || continue
        natoms=$(echo "$dir" | grep -o -E '[0-9]+' | tail -1)
        
        mkdir -p "run_$natoms"
        cd "run_$natoms" || exit

        tail -n $(($natoms + 2)) "$dir/movie.xyz" > "toFreeze_$natoms.xyz"

        cp "$CONFIG_FILE" "freeze.in"

        initTemp=$(getLastTemp "$INPUT_DIR/run_$natoms/energy.out" "7")

        sed -i "s|tinit *= *.*|tinit = $initTemp|" "freeze.in"
        sed -i "s|natom *= *.*|natom = $natoms|" "freeze.in"
        sed -i "s|type_process *= *.*|type_process = 'freezing'|" "freeze.in"
        sed -i "s|filepos *= *.*|filepos = 'toFreeze_$natoms.xyz'|" "freeze.in"
        sed -i "s|filepot *= *.*|filepot = '../pot.pot'|" "freeze.in"

        "$LODIS_EXE" < "freeze.in" > "freeze.log" &
        
        cd ..
    done

    echo "Waiting for all Freezing processes to finish..."
    wait 
    cd "$SCRIPT_DIR"
}

runQuenching(){
    local INPUT_DIR=$1
    local STAGE_DIR="$SCRIPT_DIR/quenchSims"
    
    mkdir -p "$STAGE_DIR"
    cp "$POT_FILE" "$STAGE_DIR/pot.pot"
    
    cd "$STAGE_DIR" || exit

    for dir in "$INPUT_DIR"/run_*; do
        [ -e "$dir" ] || continue
        natoms=$(echo "$dir" | grep -o -E '[0-9]+' | tail -1)
        
        mkdir -p "run_$natoms"
        cd "run_$natoms" || exit

        tail -n $(($natoms + 2)) "$dir/movie.xyz" > "toQuench_$natoms.xyz"

        cp "$CONFIG_FILE" "quench.in"

        initTemp=$(getLastTemp "$INPUT_DIR/run_$natoms/energy.out" "7")
        
        sed -i "s|tinit *= *.*|tinit = $initTemp|" "quench.in"
        sed -i "s|natom *= *.*|natom = $natoms|" "quench.in"
        sed -i "s|type_process *= *.*|type_process = 'quenching'|" "quench.in"
        sed -i "s|filepos *= *.*|filepos = 'toQuench_$natoms.xyz'|" "quench.in"
        sed -i "s|filepot *= *.*|filepot = '../pot.pot'|" "quench.in"

        "$LODIS_EXE" < "quench.in" > "quench.log" &
        
        cd ..
    done

    echo "Waiting for all Quenching processes to finish..."
    wait 
    cd "$SCRIPT_DIR"
}

# ==============================================================================
# Logic Flow (Chain the stages here)
# ==============================================================================

# 1. Initialize variables
checkStart "$1"

# 2. Start with Growth
runGrowth
# Resulting folder: $SCRIPT_DIR/growthSims

# 3. Perform Melting (using Growth results)
runMelting "$SCRIPT_DIR/growthSims"
# Resulting folder: $SCRIPT_DIR/meltingSims

# 4. Perform Freezing (using Growth results)
runFreezing "$SCRIPT_DIR/meltingSims"
# Resulting folder: $SCRIPT_DIR/freezeSims

# 5. Perform Quenching (using Freezing results)
runQuenching "$SCRIPT_DIR/freezeSims"
# Resulting folder: $SCRIPT_DIR/quenchSims

echo "All simulations completed successfully."