import matplotlib.pyplot as plt
import pandas as pd
import os
import argparse
import json
import numpy as np
from matplotlib.patches import Patch

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)

    base, ori = {}, {}
    for i, category in enumerate(categories):
        base[category] = random_data[category].dropna().tolist()
        ori[category] = guided_data[category].dropna().tolist()

    positions = np.arange(len(categories))
    width = 0.35

    dataA = [base[c] for c in categories]
    dataB = [ori[c] for c in categories]

    # plt.figure(figsize=(15, 5))
    fig, ax = plt.subplots(figsize=(15, 5))
    box_base = ax.boxplot(dataA, positions=positions - width/2, widths=0.3, vert=False, patch_artist=True)
    for box in box_base['boxes']:
        box.set(facecolor="purple", alpha=0.7)

    box_ori = ax.boxplot(dataB, positions=positions + width/2, widths=0.3, vert=False, patch_artist=True)
    for box in box_ori['boxes']:
        box.set(facecolor="blue", alpha=0.7)

    for y in positions[:-1]:
        ax.axhline(
            y + 0.5,
            linestyle="--",
            linewidth=0.8,
            alpha=0.5,
            color="gray"
    )

    # plt.xticks(positions, categories, rotation=90)
    ax.set_yticks(positions)
    ax.set_yticklabels(categories, fontsize=16)
    ax.set_xlabel("Value")
    ax.set_ylabel("App Category")
    # ax.legend(["Orienter", "Baseline"])
    ax.legend(
        handles=[
            Patch(facecolor='blue', alpha=0.7, label='Orienter'),
            Patch(facecolor='purple', alpha=0.7, label='Baseline')
        ]
    )

    fig.tight_layout(rect=[0.02, 0.02, 1, 1])
    fig.savefig(f'{RESULT_DIR}/category_eval_{datatype}.png', format='png')
    fig.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)