import os
import subprocess
import time


def run_predict_property(input_file, output_dir):
    """
    Run Predict_Property on the input file and save the results.
    """
    base_filename = os.path.basename(input_file)
    sequence_name = os.path.splitext(base_filename)[0]
    output_subdir = os.path.join(output_dir, f"{sequence_name}_PROP")

    # Path to the Predict_Property script
    predict_property_path = "/mnt/e/Predict_Property/Predict_Property.sh"

    # Construct the Predict_Property command
    command = f"{predict_property_path} -i {input_file} -o {output_subdir}"

    print(f"Running Predict_Property on: {base_filename}")
    try:
        process = subprocess.Popen(
            command,
            shell=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        stdout, stderr = process.communicate()

        if process.returncode != 0:
            error_message = stderr.decode()
            print(f"Error processing {base_filename}: {error_message}")
            with open(os.path.join(output_dir, "error_log.txt"), "a") as log_file:
                log_file.write(f"{base_filename}: {error_message}\n")
            return False

        # Check if output files were created
        if not os.path.exists(os.path.join(output_subdir, f"{sequence_name}.all")):
            print(f"Error: Output files not found for {base_filename}")
            return False

        print(f"Results saved to: {output_subdir}")
        return True
    except Exception as e:
        print(f"Error processing {base_filename}: {str(e)}")
        with open(os.path.join(output_dir, "error_log.txt"), "a") as log_file:
            log_file.write(f"{base_filename}: {str(e)}\n")
        return False


# Define input and output directories for each type
directories = {
    "iORFs_proteins": {
        "input": "/mnt/e/Pasteur/TM/Actinobacteria/sORFs_proteins",
        "output": "/mnt/e/Pasteur/TM/Actinobacteria/predict_property_results/sORFs_proteins",
    },
    "iORFs_shuffled_proteins": {
        "input": "/mnt/e/Pasteur/TM/Actinobacteria/sORFs_proteins_shuffled",
        "output": "/mnt/e/Pasteur/TM/Actinobacteria/predict_property_results/sORFs_proteins_shuffled",
    },
}

# Process each type
total_successful = 0
total_failed = 0

for key, dirs in directories.items():
    input_dir = dirs["input"]
    output_dir = dirs["output"]

    print(f"\nProcessing: {key}")
    print(f"Input directory: {input_dir}")
    print(f"Output directory: {output_dir}")

    # Create the output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)

    successful = 0
    failed = 0

    # Process each FASTA file in the input directory
    fasta_files = [
        f for f in os.listdir(input_dir) if f.endswith((".fa", ".fasta", ".faa"))
    ]
    total_files = len(fasta_files)

    for idx, filename in enumerate(fasta_files, 1):
        input_file = os.path.join(input_dir, filename)

        # Run prediction
        if run_predict_property(input_file, output_dir):
            successful += 1
        else:
            failed += 1

        # Print progress
        print(
            f"Progress: {idx}/{total_files} files processed ({successful} successful, {failed} failed)"
        )

        # Add a small delay between files
        time.sleep(1)

    total_successful += successful
    total_failed += failed

    print(f"\nResults for {key}:")
    print(f"Successfully processed: {successful}")
    print(f"Failed: {failed}")

print("\nFinal Summary:")
print(f"Total successfully processed: {total_successful}")
print(f"Total failed: {total_failed}")

# Print output file summary
print("\nResults Summary:")
for key, dirs in directories.items():
    output_dir = dirs["output"]
    if os.path.exists(output_dir):
        results_count = len([d for d in os.listdir(output_dir) if d.endswith("_PROP")])
        print(f"{key}: {results_count} sequences processed")