import json
import os
import pandas as pd
import argparse
from rdflib import Graph
from collections import Counter

from mimicsql.evaluation.utils import query
from sql2sparql import SQL2SPARQL, split_entity
from evaluation_sparql import isequal
from build_mimicsparql_kg.build_complex_kg_from_mimicsqlstar_db import clean_text


def sparql_tokenize(sparql):
    sparql = split_entity(sparql)
    sparql = ' ^^'.join(sparql.split('^^'))
    sparql_tok = ' '.join(sparql.split(' '))
    return sparql_tok.split()


def convert_sql2sparql(complex=True, filename='train.json', dataset_type='natural', execution=True):
    if complex:
        savedir = f'./dataset/mimic_sparqlstar/{dataset_type}/'
        datadir = f'./dataset/mimicsqlstar/{dataset_type}/'

        sql2sparql = SQL2SPARQL(complex=complex, root='subject_id')

        if execution:
            print('LOAD  ... mimicqlstar.db')
            db_file = './build_mimicsqlstar_db/mimicsqlstar.db'
            model = query(db_file)
            print('DONE')

            print('LOAD KG ... mimic_kg')
            kg = Graph()
            kg.parse('./build_mimicsparql_kg/mimic_sparqlstar_kg.xml', format='xml', publicID='/')
            print('DONE')

    else:
        print(f'This dataset is Simple')
        savedir = f'./dataset/mimic_sparql/{dataset_type}/'
        datadir = f'./dataset/mimicsql/mimicsql_{dataset_type}/'

        sql2sparql = SQL2SPARQL(complex=complex, root='hadm_id')

        if execution:
            print('LOAD ... mimic.db')
            db_file = './mimicsql/evaluation/mimic_db/mimic.db'
            model = query(db_file)
            print('DONE')

            print('LOAD KG ... mimic_sparql_kg')
            kg = Graph()
            kg.parse('./build_mimicsparql_kg/mimic_sparql_kg.xml', format='xml', publicID='/')
            print('DONE')

    data = []
    with open(os.path.join(datadir, filename)) as json_file:
        for line in json_file:
            data.append(json.loads(line))

    df = pd.DataFrame(data)

    correct = 0
    sparqls = []
    for i, sql in enumerate(df['sql']):
        sql = sql.lower()
        sql_answer = []
        sparql_answer = []

        print("-" * 50)
        print(i, sql)

        if execution:
            sql_res = model.execute_sql(sql).fetchall()
            for res in sql_res:
                val = '|'
                temp = []
                for t in res:
                    val += str(t) + '|\t\t|'
                    temp.append(str(t))
                print(val[:-1])
                sql_answer.append(tuple(temp))
        print()

        sparql = sql2sparql.convert(sql)
        sparql = clean_text(sparql)

        print(i, sparql)
        if execution:
            sparql_res = kg.query(sparql)
            for res in sparql_res:
                val = '|'
                temp = []
                for t in res:
                    val += str(t.toPython()) + '|\t\t|'
                    temp.append(str(t.toPython()))
                print(val[:-1])
                sparql_answer.append(tuple(temp))

            print(sql_answer, sparql_answer, isequal(sql_answer, sparql_answer))
            if isequal(sql_answer, sparql_answer):
                correct += 1
            else:
                print("[incorrect]")
        print()

        sparql = sparql.lower()
        sparql_tok = sparql_tokenize(sparql)
        sparqls.append({'sql': sparql, 'sql_tok': sparql_tok})

    if execution:
        print(f'[SQL2SPARQL] filenmae: {filename}, Answer Accuracy: {correct/len(df):.4f}')

    sparql_data = []
    for d, sparql_d in zip(data, sparqls):
        d['sql'] = sparql_d['sql']
        d['sql_tok'] = sparql_d['sql_tok']
        sparql_data.append(d)

    save_filename = os.path.join(savedir, filename)
    with open(save_filename, 'w') as json_file:
        for dic in sparql_data:
            json.dump(dic, json_file)
            json_file.write('\n')

    print(f"Write to {save_filename}")


def build_vocab(complex=True, dataset_type='natural'):
    if complex:
        datadir = f'./dataset/mimic_sparqlstar/{dataset_type}'
    else:
        datadir = f'./dataset/mimic_sparql/{dataset_type}'

    filenames = ['train.json']
    counter = Counter()
    for filename in filenames:
        with open(os.path.join(datadir, filename)) as json_file:
            for line in json_file:
                dic = json.loads(line)
                counter.update(dic['question_refine_tok'])
                counter.update(dic['sql_tok'])

    with open(os.path.join(datadir, 'vocab'), 'w') as f:
        for k, v in counter.most_common():

            if len(k.split()) == 0:
                continue

            if k == ' ':
                continue
            f.write(f'{k} {v}\n')

    print(f'vocab builded: {len(counter)}')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='mimicsql to mimic-sparql')
    parser.add_argument('--complex', default=False, type=lambda x: (str(x).lower() == 'true'))
    parser.add_argument('--dataset_type', type=str, default='natural', choices=['natural','template'])
    parser.add_argument('--execution', default=False, type=lambda x: (str(x).lower() == 'true'))
    args = parser.parse_args()

    execution = args.execution
    dataset_type = args.dataset_type
    complex = args.complex

    filenames = ['train.json', 'dev.json', 'test.json']
    for filename in filenames:
        convert_sql2sparql(complex=complex, filename=filename, dataset_type=dataset_type, execution=execution)
    build_vocab(complex=complex, dataset_type=dataset_type)