
import subprocess
import os
import shutil
import pymatgen
import mp_api
from lattice import lattice
from pressure_eV import pressure_eV
from position_force import position_force
from energy import energy
from atom_symbol import atom_symbol
from combine import combine
from npz_to_extxyz import npz_to_extxyz


def process_and_run_script(input_folder):                       # Define a wrapper function that orchestrates the entire data processing workflow
    
    # Step 1: 
    input_folder_path = os.path.abspath(input_folder)           # Get the absolute path of the input folder
    script_dir = os.path.dirname(os.path.abspath(__file__))     # Set the path to the directory where the Bash script resides  
    bash_script_path = os.path.join(script_dir, '..', 'scripts', 'hydrostatic_strain_post_processing.sh')  # Find the path to the Bash script
    #bash_script_path = os.path.join(script_dir, '..', 'scripts', 'AIMD_data_extract.sh')  # Find the path to the Bash script

    # Check if the Bash script exists
    if not os.path.exists(bash_script_path):
        print("Error: Bash script not found.")
        return

    # Ask if the user wants to run Step 1
    run_step1 = input("Do you want to run the first step (execute hydrostatic_strain_post_processing.sh)? (yes/no): ").strip().lower()
    
    if run_step1 == 'yes':       
        try:                                            # Copy the Bash script to the input folder
            shutil.copy(bash_script_path, input_folder_path)
            print("Bash script copied successfully.")
        except Exception as e:
            print(f"Error copying Bash script: {e}")
            return

        os.chdir(input_folder_path)                       # Change directory to the input folder
        try:                                              # Execute the Bash script
            subprocess.run(['bash', 'hydrostatic_strain_post_processing.sh'], capture_output=True, text=True, check=True)
            #subprocess.run(['bash', 'AIMD_data_extract.sh'], capture_output=True, text=True, check=True)
            print("Bash script executed successfully.")
        except subprocess.CalledProcessError as e:
            print(f"Error executing Bash script: {e}")
            return
    else:
        print("Skipping the first step and proceeding to step 2.")
    
    print("Input folder path:", input_folder_path)   # Debugging line: Print the function name and input file path


    # Step 2: Process the file generated by the Bash script using *.py
    os.chdir(input_folder)
    print("Starting Step 2: Processing files with Python scripts.")
    lattice_output_file = lattice(input_folder)
    pressure_eV_output_file = pressure_eV(input_folder)
    position_force_output_file = position_force(input_folder)
    energy_output_file = energy(input_folder)
    atom_symbol_output_file = atom_symbol(input_folder)
    combined_output_file = combine(input_folder)
    npz_to_extxyz_output_file = npz_to_extxyz(combined_output_file)

    
    print(f"Final output file directory from the workflow : {npz_to_extxyz_output_file}") # Debugging line: Print the final output file path

if __name__ == "__main__":
    # Replace 'path/to/input_file.txt' with the actual path to your input file
    #process_and_run_script(r'C:\Users\Galib\Downloads\hydrostatic_strain_-0.1_0.0001_0.1\alumina')
    #process_and_run_script(r'F:\Research\ML\package\AtomProNet\example_dataset\AIMD\alumina')
    #pass

    # Prompt the user to enter the folder path
    input_folder = input("Please enter the full path to the folder where the operations should be performed: ").strip()
    process_and_run_script(input_folder)