import pandas as pd
import os

def main():
    # Read the CSV file containing tree analysis results
    try:
        df = pd.read_csv('results/tree_analysis_results.csv')
    except FileNotFoundError:
        print("Error: tree_analysis_results.csv not found.")
        return
    except pd.errors.EmptyDataError:
        print("Error: tree_analysis_results.csv is empty.")
        return
    
    # Create dictionaries to store unpruned and pruned tree sizes
    unpruned_sizes = {}
    pruned_sizes = {}
    pruned_errors = {}
    
    # Iterate through the rows of the dataframe
    for index, row in df.iterrows():
        tree_name = row['Tree']
        tree_size = row['Tree Size']
        errors = row['Classification Errors']
        
        # Extract dataset name and tree type from the tree name
        if '_data-' in tree_name:
            parts = tree_name.split('_data-')
            dataset = parts[0]
            tree_type = parts[1]
            
            # Store the tree size based on the tree type
            if tree_type == 'unpru':
                unpruned_sizes[dataset] = tree_size
            elif tree_type == 'repl':
                pruned_sizes[dataset] = tree_size
                pruned_errors[dataset] = errors
    
    # Create a list to store the results
    results = []
    
    # Calculate the difference between unpruned and pruned tree sizes
    for dataset in sorted(unpruned_sizes.keys()):
        if dataset in pruned_sizes:
            unpruned_size = unpruned_sizes[dataset]
            pruned_size = pruned_sizes[dataset]
            difference = unpruned_size - pruned_size
            errors = pruned_errors[dataset]
            results.append({'Dataset': dataset, 'Heuristic pruned': difference, 'Heuristic errors': errors})
    
    # Convert the results to a dataframe
    results_df = pd.DataFrame(results)
    
    # Create the results directory if it doesn't exist
    os.makedirs('results', exist_ok=True)
    
    # Write the results to the specified CSV file
    results_df.to_csv('results/heuristic-pruned.csv', index=False)
    
    print(f"Results written to results/heuristic-pruned.csv")

if __name__ == "__main__":
    main()
