import matplotlib.pyplot as plt
import pandas as pd
import os
import argparse
import json

split = os.getenv('split', '')
suf_split = f'-{split}' if split else ''

RESULT_DIR = f'./results{suf_split}'
cat_apps_file = f'./cat_apps{suf_split}.json'


def plot_box(random_data, guided_data, datatype, min_apps = 0):
    DATATYPE_MAP = {
        'effective_interacts_cnt': 'Effective Interacts Count',
        'effective_interacts_rate': 'Effective Interacts Rate',
        'coverage_rate': 'IGE Coverage Rate'
    }

    # categories = guided_data.columns[1:]  # 获取所有category
    with open(cat_apps_file, 'r') as f:
        category_apps_map = json.load(f)
    categories = [cat for cat, apps in category_apps_map.items() if len(apps) >= min_apps and cat != 'All']
    categories = ['All'] + categories
    num_categories = len(categories)
    num_cols = 6
    num_rows = (num_categories + num_cols - 1) // num_cols  # 计算行数

    fig, axes = plt.subplots(num_rows, num_cols, figsize=(15.5, num_rows * 2.5))
    axes = axes.flatten()

    for i, category in enumerate(categories):
        data = [random_data[category], guided_data[category]]
        box = axes[i].boxplot(data, labels=['Base', 'Ori.'], patch_artist=True, widths=0.6)
        colors = ['purple', 'green']
        for patch, color in zip(box['boxes'], colors):
            patch.set_facecolor(color)
        for median in box['medians']:
            median.set_color('blue')
            median.set_linewidth(2)
        axes[i].set_title(category, fontsize=20)
        axes[i].tick_params(axis='both', which='major', labelsize=20)
        axes[i].grid(True)

    for i in range(num_categories, len(axes)):
        fig.delaxes(axes[i])  # 删除多余的子图

    fig.text(0.5, 0.01, 'App Category', ha='center', fontsize=20)
    fig.text(0.01, 0.5, DATATYPE_MAP[datatype], va='center', rotation='vertical', fontsize=20)

    plt.tight_layout(rect=[0.02, 0.02, 1, 1])
    plt.savefig(f'{RESULT_DIR}/category_eval_{datatype}.png', format='png')
    plt.savefig(f'{RESULT_DIR}/category_eval_{datatype}.pdf', format='pdf')



def main(args):
    random_data = pd.read_csv(f'{args.random}/{args.type}.csv')
    guided_data = pd.read_csv(f'{args.guided}/{args.type}.csv')

    plot_box(random_data, guided_data, args.type, min_apps=args.min_apps)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Plot the category evaluation results')
    parser.add_argument('-r', '--random', type=str, help='The random interact evaluation result')
    parser.add_argument('-g', '--guided', type=str, help='The our interact evaluation result')
    parser.add_argument('-t', '--type', type=str, help='The type of the evaluation result')
    parser.add_argument('-m', '--min_apps', type=int, default=1, help='Minimum number of apps per category')
    args = parser.parse_args()
    main(args)