import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

if __name__ == "__main__":
    # read in the csv
    df = pd.read_csv("./results/compress/compression_results.csv")

    for model_type in df["model_type"].unique():
        model_df = df[df["model_type"] == model_type]

        # execution time
        plt.figure()
        grouped = model_df.groupby("context_length")["compression_time"].mean() / 1e9
        labels = grouped.index.astype(str)  # "128", "256"
        x = np.arange(len(labels))  # [0, 1]

        plt.bar(x, grouped.values, width=0.6)
        plt.title(f"{model_type} mean compression time")
        plt.xticks(x, labels)
        plt.xlabel("Context length")
        plt.ylabel("Mean compression time [s]")
        plt.tight_layout()
        plt.savefig(f"./graphs/{model_type}_{}_compression_time.png")

        plt.figure()
        grouped = model_df.groupby("context_length")["decompression_time"].mean() / 1e9
        labels = grouped.index.astype(str)  # "128", "256"
        x = np.arange(len(labels))  # [0, 1]

        plt.bar(x, grouped.values, width=0.6)
        plt.title(f"{model_type} mean decompression time")
        plt.xticks(x, labels)
        plt.xlabel("Context length")
        plt.ylabel("Mean decompression time [s]")
        plt.tight_layout()
        plt.savefig(f"./graphs/{model_type}_{}_decompression_time.png")

        # accuracy
        plt.figure()
        bar_height = 0.25
        files = model_df["input_file_name"].unique()
        y = np.arange(len(files))
        c256 = model_df[model_df["context_length"] == 256]
        c128 = model_df[model_df["context_length"] == 128]

        plt.barh(
            y - bar_height / 2,
            c256["match_percentage"] * 100,
            height=bar_height,
            label="256"
        )

        plt.barh(
            y + bar_height / 2,
            c128["match_percentage"] * 100,
            height=bar_height,
            label="128"
        )
        plt.yticks(y, files)
        plt.title(f"{model_type} time for different context lengths")
        plt.xlabel("accuracy")
        plt.ylabel("Filename")
        plt.legend()
        plt.savefig(f"./graphs/{model_type}_{}_accuracy.png")

        # compression ratio
        plt.figure()
        c256 = model_df[model_df["context_length"] == 256]
        c128 = model_df[model_df["context_length"] == 128]

        plt.plot(c256["original_file_size"] / 1_000_000, c256["compressed_file_size"] / 1_000_000, label="256")
        plt.plot(c128["original_file_size"] / 1_000_000, c128["compressed_file_size"] / 1_000_000, label="128")
        plt.title(f"{model_type} compressed file evolution")
        plt.xlabel("Original file size [MB]")
        plt.ylabel("Compressed file size [MB]")
        plt.legend()
        plt.savefig(f"./graphs/{model_type}_{}_compression_ratio.png")