# Plot script for LoCoOp evaluation results
import numpy as np
import seaborn as sns
from matplotlib import pyplot as plt
import os
import pandas as pd
import torch
import torch.nn.functional as F
import argparse


def plot_distribution(in_dataset, out_dataset, id_scores, ood_scores, output_dir):
    sns.set(style="white", palette="muted")
    palette = ['#A8BAE3', '#55AB83']

    data = {
        f"ID ({in_dataset})": [-1 * id_score for id_score in id_scores],
        f"OOD ({out_dataset})": [-1 * ood_score for ood_score in ood_scores]
    }

    sns.displot(data, label="id", kind="kde", palette=palette, fill=True, alpha=0.8)

    plt.suptitle(f"Score Distribution\nID: ImageNet, OOD: {out_dataset}", fontsize=14, y=1.1)

    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    plt.savefig(os.path.join(output_dir, f"{out_dataset}.png"), bbox_inches='tight')


def main():
    parser = argparse.ArgumentParser(description="Plot evaluation results")
    parser.add_argument("--input-file", type=str, required=True, help="Path to scores.npz file")
    parser.add_argument("--output-dir", type=str, required=True, help="Output directory for plots")
    
    args = parser.parse_args()
    
    # Load scores
    scores = np.load(args.input_file, allow_pickle=True)
    print(scores)

    in_score_mcm = scores["ImageNet"]


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

    for dataset in scores.keys():
        if dataset == "ImageNet":
            continue
        out_score_mcm = scores[dataset]
        # plot
        plot_distribution("ImageNet", dataset, in_score_mcm, out_score_mcm, args.output_dir)

if __name__ == "__main__":
    main()