import os
import pandas as pd
from matplotlib import pyplot as plt, ticker

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

NAMES = {
    # 'Orienter (GPT-4v)': f'{RESULT_DIR}/aaa_icse_test_set_merged_gpt4v_gpt4v_ape_d_object_bbox_r1i_bf.csv',
    # 'Orienter (Claude)': f'{RESULT_DIR}/realfse_union3_claude35sonnet_ape_d_object_bbox_gpu0123_i_bf.csv',
    'Orienter (Gemini)': f'{RESULT_DIR}/realfse_union3_gemini15pro_ape_d_object_bbox_gpu0123_i_bf.csv',
    # 'Orienter (GPT-4o)': f'{RESULT_DIR}/realfse_union3_gpt4v_ape_d_object_bbox_gpu0123_i_bf.csv',
    'YOLOv8': f'{RESULT_DIR}/YOLO.csv',
    'Claude-4.5-sonnet': f'{RESULT_DIR}/Claude4_5-sonnet-E2E.csv',
    'InternVL 3.5': f'{RESULT_DIR}/internVL-E2E.csv',
    # 'GPT-4V-E2E': f'{RESULT_DIR}/GPT4V-E2E.csv',
    'Random-based Fuzzing': f'{RESULT_DIR}/random.csv',
}

MARKERS = {
    # 'Orienter (GPT-4v)': 'o',
    # 'Orienter (Claude)': 's',
    'Orienter (Gemini)': 'D',
    # 'Orienter (GPT-4o)': '^',
    'YOLOv8': '*',
    'Claude-4.5-sonnet': 'P',
    'InternVL 3.5': 'v',
    # 'GPT-4V-E2E': 'v',
    'Random-based Fuzzing': 'x',
}


def format_func(x, pos):
    return format(int(x), ',').replace(',', ' ')


def plot_effective_interacts_cnt(data):
    plt.figure()
    ax = plt.gca()
    plt.gcf().set_size_inches(8, 6)
    ax.set_xlabel('Time/min')
    ax.xaxis.set_major_locator(ticker.MultipleLocator(10))
    ax.set_ylabel('Effective Interactions')
    ax.yaxis.set_major_locator(ticker.MultipleLocator(1000))
    ax.yaxis.set_major_formatter(ticker.FuncFormatter(format_func))
    for name, df in data.items():
        plt.plot(df.index, df['effective_interacts_cnt'], label=name, marker=MARKERS[name], markevery=10, linestyle='-' if name.startswith('Orienter') else '-', markersize=12)
    # Add legend
    # plt.legend(loc='upper center', bbox_to_anchor=(0.5, -0.15), ncol=2, frameon=False)
    # set the size of the fig
    plt.savefig(f'{RESULT_DIR}/test-effective-interacts-cnt.png', dpi=300, bbox_inches='tight')
    plt.savefig(f'{RESULT_DIR}/test-effective-interacts-cnt.pdf', dpi=300, bbox_inches='tight', format='pdf')


def plot_coverage_rate(data):
    plt.figure()
    ax = plt.gca()
    plt.gcf().set_size_inches(8, 6)
    ax.set_xlabel('Time/min')
    ax.xaxis.set_major_locator(ticker.MultipleLocator(10))
    ax.set_ylabel('IGE Coverage')
    ax.yaxis.set_major_locator(ticker.MultipleLocator(0.1))
    for name, df in data.items():
        plt.plot(df.index, df['coverage_rate'], label=name, marker=MARKERS[name], markevery=10, linestyle='-' if name.startswith('Orienter') else '-', markersize=12)
    # Add legend inside the plot
    plt.legend(loc='lower right', ncol=1, frameon=False, fontsize=19)
    plt.savefig(f'{RESULT_DIR}/test-coverage-rate.png', dpi=300, bbox_inches='tight')
    plt.savefig(f'{RESULT_DIR}/test-coverage-rate.pdf', dpi=300, bbox_inches='tight', format='pdf')


def main():

    plt.rcParams['font.size'] = 28

    # Load the data
    data = {}
    for name, path in NAMES.items():
        data[name] = pd.read_csv(path)
        data[name] = data[name][1:]

    plot_effective_interacts_cnt(data)
    plot_coverage_rate(data)


if __name__ == '__main__':
    main()