import json

import re

import matplotlib.pyplot as plt

import collections

import matplotlib

import seaborn as sns

from collections import defaultdict

import os

matplotlib.use("Agg")

matplotlib.rcParams.update({'font.family': 'Times New Roman'})

matplotlib.rcParams['pdf.fonttype'] = 42

matplotlib.rcParams['ps.fonttype'] = 42

sns.set_theme(style="ticks", font="Times New Roman", font_scale=2.1, rc={'grid.linestyle': ':', 'axes.grid': True})

for dataset in ["mimic_iii", "eicu"]:

    with open(f"<YOUR_DATASET_PATH>", 'r') as f_in, \

            open(f"<YOUR_SAVED_DATASET_PATH>", "w") as f_out:

        list_num_q_tag_var = []

        list_num_tables = []

        list_num_columns = []

        num_q_tag_var_dict = defaultdict(list)

        num_tables_dict = defaultdict(list)

        num_columns_dict = defaultdict(list)

        for lines in f_in:

            x = json.loads(lines)

            # number of variables in q_tag

            num_q_tag_var = x["q_tag"].count("{") + x["q_tag"].count("[")

            # number of different tables used in the sql

            # (we only count the tables in the original datasets, but not the new ones created in the sql)

            tables = re.findall(r'\bfrom\s+(\w+)\b', x["query"])

            num_tables = len(set(tables))

            # number of different columns used in the sql

            # (we only count the columns in the original datasets, but not the new ones created in the sql)

            columns = re.findall(r'\b\w*\.\w*\b', x["query"])

            columns = [item for item in columns if not re.match(r'^t\d', item)]

            num_columns = len(set(columns))

            x["num_q_tag_var"] = num_q_tag_var

            x["num_tables"] = num_tables

            x["num_columns"] = num_columns

            f_out.write(json.dumps(x) + "\n")

            list_num_q_tag_var.append(num_q_tag_var)

            list_num_tables.append(num_tables)

            list_num_columns.append(num_columns)

            num_q_tag_var_dict[num_q_tag_var].append(x["id"])

            num_tables_dict[num_tables].append(x["id"])

            num_columns_dict[num_columns].append(x["id"])

    os.makedirs(f"{dataset}/num_q_tag_var/", exist_ok=True)

    for k, id_list in num_q_tag_var_dict.items():

        with open(f"{dataset}/num_q_tag_var/{k}.jsonl", "w") as f:

            json.dump(id_list, f)

    os.makedirs(f"{dataset}/num_tables/", exist_ok=True)

    for k, id_list in num_tables_dict.items():

        with open(f"{dataset}/num_tables/{k}.jsonl", "w") as f:

            json.dump(id_list, f)

    os.makedirs(f"{dataset}/num_columns/", exist_ok=True)

    for k, id_list in num_columns_dict.items():

        with open(f"{dataset}/num_columns/{k}.jsonl", "w") as f:

            json.dump(id_list, f)

    # plot the distribution

    all_list = [list_num_q_tag_var, list_num_tables, list_num_columns]

    xlabel = ["# q_tag variables", "# tables", "# columns"]

    file_name = ["num_q_tag_ver_distri.pdf", "num_tables_distri.pdf", "num_columns_distri.pdf"]

    os.makedirs(f"{dataset}/figures/", exist_ok=True)

    for i in range(len(all_list)):

        c = collections.Counter(all_list[i])

        c = sorted(c.items())

        plt.figure(figsize=(6, 5.5), dpi=120)

        ax = sns.barplot(x=[i[0] for i in c], y=[i[1] for i in c], color=sns.color_palette()[0])

        plt.ylabel("Frequency", size=30)

        plt.xlabel(xlabel[i], size=30)

        plt.tight_layout(rect=[-0.05, -0.05, 1.05, 1.05])

        plt.savefig(f"{dataset}/figures/{file_name[i]}")