import json
import os
import argparse
import pandas as pd

def summarize(state_file):
    if not os.path.exists(state_file):
        print(f"Error: File {state_file} not found.")
        return

    with open(state_file, 'r') as f:
        data = json.load(f)

    results = data.get('results', [])
    if not results:
        print("No results found in the state file.")
        return

    df = pd.DataFrame(results)
    
    # Format the table
    print("\n" + "="*50)
    print(f" LOPO TRAINING SUMMARY: {os.path.basename(state_file)}")
    print("="*50)
    
    # If best_epoch exists, show it, else show N/A
    columns = ['participant', 'best_mae']
    if 'best_epoch' in df.columns:
        columns.append('best_epoch')
    
    # Print the table
    print(df[columns].to_string(index=False, justify='center'))
    
    # Calculate and print Mean
    mean_mae = df['best_mae'].mean()
    print("-" * 50)
    print(f"MEAN MAE: {mean_mae:.4f} degrees")
    print(f"Total Participants Completed: {len(df)}/15")
    print("="*50 + "\n")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--file', type=str, default='report/training_state_fusion.json', 
                        help='Path to the training state JSON file')
    args = parser.parse_args()
    
    summarize(args.file)