import sys
sys.path.append('.')
sys.path.append('..')
import json
import os
import re
import pandas as pd
from rdflib import Graph
from sql2sparql import SQL2SPARQL, sparql_postprocessing, join_entity
from mimicsql.evaluation.utils import query
from build_mimicsparql_kg.build_complex_kg_from_mimicsqlstar_db import clean_text


def split_triples(sparql):
    try:
        select_part, where_part = sparql.split(' where ')
    except:
        print(sparql)
        select_part, where_part = sparql.split(' where ')[0], sparql.split(' where ')[-1]

    where_part = where_part.replace('{','').replace('}','')
    triple = [t.strip() for t in where_part.split('. ')]
    return select_part, [t for t in triple if len(t) != 0]


def none2zero(answer):
    if answer is None:
        return 0.0

    if type(answer) != str:
        return answer

    if answer.lower() == 'none':
        return 0.0

    try:
        answer = float(answer)
    except:
        pass

    return answer


def answer_normalization(answers):
    if len(answers) == 0:
        answers = [(0.0, )]
    return [tuple([none2zero(val) for val in answer]) for answer in answers]


def entity2value(entity):
    match = re.findall('/[a-z_\d]+/[a-z\d]+', entity)
    if len(match) > 0:
        return re.sub('/[a-z_\d]+/', '', entity)
    else:
        return entity


def replace_cond_val(sparql):
    try:
        where_part = re.findall('{[^{^}].*}', sparql)[0]
    except Exception as e:
        print(e)
        return sparql

    where_part = where_part.replace('{', '').replace('}', '').strip()

    try:
        ent_rel_cond = re.findall('\?[a-z_\d]+ </[a-z_\d]+> [^?][^.^]+', where_part)
        for m in ent_rel_cond:
            token = m.split()
            ent, rel = token[0], token[1]
            re_m = ' '.join([ent, rel, '<COND>'])
            sparql = sparql.replace(m, re_m)

        cond_rel_ent = re.findall('<[^?^ ]+> </[a-z_\d]+> \?[a-z_\d]+', where_part)
        for m in cond_rel_ent:
            cond, rel, ent = m.split()
            re_m = ' '.join(['<COND>', rel, ent])
            sparql = sparql.replace(m, re_m)

        filter_cond = re.findall('filter\( \?[a-z_\d]+ [<=>]+ [^?]+ \)', where_part)
        for m in filter_cond:
            ft, var, op, cond, _ = m.split()
            re_m = ' '.join([ft, var, op, '<COND>', _])
            sparql = sparql.replace(m, re_m)

    except Exception as e:
        print(e)

    return sparql


def isequal(sql_answer, sparql_answer): # list of tuple
    sql_answer = [row for row in sql_answer if 'None' not in row]
    sql_answer = [tuple([clean_text(a.lower()) if type(a) == str else a for a in row]) for row in sql_answer]

    sparql_answer = [tuple([entity2value(a) if type(a) == str else a for a in row]) for row in sparql_answer]

    if set(sql_answer) == set(sparql_answer):
        return True

    sparql_answer = answer_normalization(sparql_answer)
    sql_answer = answer_normalization(sql_answer)

    if set(sql_answer) == set(sparql_answer):
        return True

    return False


def check_no_cond_val(sparql):
    cond = []
    cond += re.findall('\^\^<http://', sparql) # value
    cond += re.findall('</[a-z_\d]+/[\d]+>', sparql) # entity
    cond += re.findall('"[a-z\d ]+"', sparql) # value
    cond += re.findall('filter', sparql) # fiter
    if len(cond) == 0:
        return True
    else:
        return False


def n_inner_join(x):
    return len(re.findall('inner join', x))


def compare_sql_and_spqral_pred():
    datadir = '../TREQS/mimicsql_data/mimicsql_natural/'
    filename = 'test.json'
    outputdir = '../TREQS/evaluation/generated_sql/'
    output_filename = 'output.json'

    covertor = SQL2SPARQL()
    print('LOAD output.json')
    sparql_preds = []
    sparql_golds = []
    with open(os.path.join(outputdir, output_filename)) as json_file:
        for line in json_file:
            dic = json.loads(line)
            sparql_preds.append(dic['sql_pred'])
            sparql_golds.append(dic['sql_gold'])
    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)

    print('LOAD DB ...')
    db_file = './evaluation/mimic_db/mimic.db'
    model = query(db_file)
    print('DONE')

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

    lf_permu_correct = 0
    lf_permu_cond_correct = 0
    cond_lf_correct = 0
    lf_correct = 0
    gold_correct = 0
    pred_correct = 0
    ablation_results = []
    for i, sql in enumerate(df['sql']):
        ablation_dic = {}
        sql = sql.lower()

        sql_answer = []
        sparql_pred_answer = []
        sparql_gold_answer = []

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

        ablation_dic['n_inner'] = n_inner_join(sql)
        ablation_dic['n_hop'] = covertor.get_max_hop(sql)

        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_pred = sparql_preds[i]
        sparql_gold = sparql_golds[i]

        sparql_pred = sparql_postprocessing(sparql_pred)
        sparql_pred = join_entity(sparql_pred)
        sparql_gold = sparql_postprocessing(sparql_gold)
        sparql_gold = join_entity(sparql_gold)


        if sparql_pred.split() == sparql_gold.split():
            lf_correct += 1
            ablation_dic['lf_correct'] = 1

        print(sparql_gold)
        print(sparql_pred)

        cond_sp = replace_cond_val(sparql_pred)
        cond_sg = replace_cond_val(sparql_gold)

        if cond_sp.split() == cond_sg.split():
            cond_lf_correct += 1
            ablation_dic['cond_lf_correct'] = 1
            print(cond_sg)
            print(cond_sp)

        cond_sps, cond_spw = split_triples(cond_sp)
        cond_sgs, cond_sgw = split_triples(cond_sg)

        sps, spw = split_triples(sparql_pred)
        sgs, sgw = split_triples(sparql_gold)

        if cond_sps.split() == cond_sgs.split() and set(cond_spw) == set(cond_sgw):
            lf_permu_cond_correct += 1

        if sps.split() == sgs.split() and set(spw) == set(sgw):
            lf_permu_correct += 1

        print(i, sparql_gold)
        sparql_res = kg.query(sparql_gold)
        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_gold_answer.append(tuple(temp))
        print(sql_answer, sparql_gold_answer, isequal(sql_answer, sparql_gold_answer))

        if isequal(sql_answer, sparql_gold_answer):
            gold_correct += 1
        else:
            print('sql gold false')

        print(i, sparql_pred)

        if check_no_cond_val(sparql_pred):
            print(f'[NO COND]: {sparql_pred}')
            print()
            ablation_results.append(ablation_dic)
            continue

        try:
            sparql_res = kg.query(sparql_pred)
            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_pred_answer.append(tuple(temp))
            print(sql_answer, sparql_pred_answer, isequal(sql_answer, sparql_pred_answer))

            if isequal(sql_answer, sparql_pred_answer):
                ablation_dic['ex_correct'] = 1
                pred_correct += 1

        except:
            print(sparql_pred)
            print("syntax error")

        ablation_results.append(ablation_dic)

        print()

    print(f'[SQL2SPARQL] filenmae: {filename}, Answer Accuracy: {gold_correct / len(data):.4f}')
    print(f'[SQL2SPARQL] filenmae: {output_filename}, Answer Accuracy: {pred_correct / len(data):.4f}')
    print(f'[SQL2SPARQL] filenmae: {output_filename}, Logic Form Accuracy: {lf_correct / len(data):.4f}')
    print(f'[SQL2SPARQL] filenmae: {output_filename}, Cond Invariant Logic Form Accuracy: {cond_lf_correct / len(data):.4f}')
    print(f'[SQL2SPARQL] filenmae: {output_filename}, Logic Form Accuracy*: {lf_permu_correct / len(data):.4f}')
    print(f'[SQL2SPARQL] filenmae: {output_filename}, Cond Invariant Logic Form Accuracy*: {lf_permu_cond_correct / len(data):.4f}')

    df = pd.DataFrame(ablation_results)
    df.fillna(0, inplace=True)
    df.info()
    print(df['n_inner'].value_counts())
    print(df[df['n_inner'] == 0]['ex_correct'].sum())
    print(df[df['n_inner'] == 1]['ex_correct'].sum())
    print(df[df['n_inner'] == 2]['ex_correct'].sum())

    print('*'*50)
    print(df['n_hop'].value_counts())
    print(df[df['n_hop'] == 1]['ex_correct'].sum())
    print(df[df['n_hop'] == 2]['ex_correct'].sum())
    print(df[df['n_hop'] == 3]['ex_correct'].sum())
    print(df[df['n_hop'] == 4]['ex_correct'].sum())

    df.to_csv(f'./ablation_results_{output_filename}.csv')


if __name__ == '__main__':
    #compare_sql_and_spqral()
    compare_sql_and_spqral_pred()