#!/usr/bin/env python3

import os
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

def main():
    # Set style
    sns.set_style("whitegrid")
    plt.rcParams.update({
        'font.size': 14,
        'axes.labelsize': 14,
        'axes.titlesize': 16,
        'xtick.labelsize': 14,
        'ytick.labelsize': 14,
        'legend.fontsize': 14
    })
    
    # Path to the combined metrics file
    metrics_file = "experiments/results/files/generative-metrics/all-models-comparison.csv"
    
    if not os.path.exists(metrics_file):
        print(f"Error: Metrics file not found at {metrics_file}")
        print("Please run generative-metrics-experiment.py first to generate the metrics.")
        return
    
    # Read the metrics data
    df = pd.read_csv(metrics_file)
    
    # Select only the columns we want to display
    metrics_to_display = ['Model', 'Diversity', 'Uniqueness', 'FitnessScore']
    
    if 'Diversity' not in df.columns:
        print("Warning: 'Diversity' column not found in the metrics file.")
        metrics_to_display.remove('Diversity')
    
    # Filter the dataframe to include only the desired metrics
    table_df = df[metrics_to_display].copy()
    
    # Rename the columns
    if 'Diversity' in table_df.columns:
        table_df = table_df.rename(columns={'Diversity': 'Diversity'})
    
    # Round numeric values to 3 decimal places for better readability
    for col in table_df.columns:
        if col != 'Model':
            table_df[col] = table_df[col].round(3)
    
    # Create a figure and axis for the table - make it wider
    fig, ax = plt.subplots(figsize=(20, len(table_df) + 2))
    
    # Hide axes
    ax.axis('off')
    ax.axis('tight')
    
    # Create the table
    table = ax.table(
        cellText=table_df.values,
        colLabels=table_df.columns,
        loc='center',
        cellLoc='center',
        colColours=['#f2f2f2'] * len(table_df.columns)
    )
    
    # Style the table
    table.auto_set_font_size(False)
    table.set_fontsize(14)
    table.scale(1.5, 1.5)  # Increased width scaling
    
    # Adjust column widths to make the table wider
    for (i, j), cell in table.get_celld().items():
        if j == 0:  # Model column
            cell.set_width(0.4)
        else:
            cell.set_width(0.2)
    
    # Adjust layout
    plt.tight_layout()
    
    # Create output directory if it doesn't exist
    output_dir = "experiments/results/plots"
    os.makedirs(output_dir, exist_ok=True)
    
    # Save the table as an SVG only
    output_path = os.path.join(output_dir, "generative_metrics_table.svg")
    plt.savefig(output_path, format="svg", bbox_inches="tight")
    
    print(f"Metrics table saved to {output_path}")
    
    # Also print the table to console
    print("\nGenerative Model Metrics:")
    print(table_df.to_string(index=False))

if __name__ == "__main__":
    main()
