import re
import networkx as nx
from build_mimicsparql_kg.kg_complex_schema import SCHEMA_DTYPE, KG_SCHEMA
from build_mimicsparql_kg.kg_simple_schema import SIMPLE_KG_SCHEMA, SIMPLE_SCHEMA_DTYPE
from collections import OrderedDict
from itertools import repeat


def cond_syntax_fix(sparql):
    match = re.findall('"?[^>]*"?\^\^', sparql)
    for m in match:
        sparql = sparql.replace(m, ' "' + m.split('^^')[0].strip().replace('"', '') + '"^^')
    return sparql


def split_entity(sparql):
    match = re.findall('</[a-z_]+/[\d.]+>',sparql)
    for m in match:
        val = re.findall('[\d.]+',m)[0]

        f, a = re.sub(val,' ', m).split()
        repla = ' '.join([f, val, a])
        sparql = re.sub(m, repla, sparql)
    return sparql


def join_entity(sparql):
    match = re.findall('</[a-z_]+/ [\d\.]+ >[\. ]', sparql)
    for m in match:
        sparql = re.sub(m, ''.join(m.split() + [' ']), sparql)
    return sparql


def clean_text(val):
    if type(val) == str:
        val = val.replace("\\", ' ')
    return val


def value2entity(sparql):
    where = sparql.split('where')[-1].split('filter')[0]
    ms = re.findall('</[a-z_\d]+> [a-z\d.]+', where)
    for m in ms:
        rel, val = m.split()
        val = val.replace('.','')
        entity_val = f'{rel} {rel[:-1]}/{val}>.'
        sparql = sparql.replace(m,entity_val)
    return sparql


def sparql_postprocessing(sparql):
    sparql = clean_text(sparql.lower())
    sparql = sparql.replace(' <stop>', '')
    sparql = sparql.replace('/xmlschema#', '/XMLSchema#')
    sparql = sparql.replace(' ^^<http://www', '^^<http://www')
    sparql = cond_syntax_fix(sparql)
    sparql = value2entity(sparql)
    return sparql


class SQL2SPARQL:
    def __init__(self, complex, root='subject_id'):
        self.schema = KG_SCHEMA if complex else SIMPLE_KG_SCHEMA
        self.schema_type = SCHEMA_DTYPE if complex else SIMPLE_SCHEMA_DTYPE
        self.schema_type = {k.lower(): v for k, v in self.schema_type.items()}
        self.schema_graph = nx.DiGraph()
        for k, vs in self.schema.items():
            for v in vs:
                self.schema_graph.add_edge(k.lower(), v.lower())
        self.agg_func = ['count', 'max', 'min', 'avg']
        self.sel_p = re.compile('"[^"]*"')
        self.cond_p = re.compile('"[^"]*"|[=><]+')
        self.sparql_agg_template = "select ( {AGG} ( {DISTINCT} ?{COL} ) as ?agg ) "
        self.sparql_select_template = "select"
        self.root = root
        self.duplicate_columns = ['short_title', 'long_title', 'icd9_code']

    def _replace_dulicate_column(self, sql):
        for col in self.duplicate_columns:
            tokens = [token.split('.') for token in re.findall(f'[a-z]+."{col}"', sql)]
            for table, _ in tokens:
                sql = re.sub(f'{table}."{col}"', f'"{table}_{col}"', sql)
        return sql

    def get_max_hop(self, sql):
        sql = self._replace_dulicate_column(sql)
        remain, where_part = sql.split(' where ')
        select_part, remain = remain.split(' from ')

        distinct_term = 'distinct' if 'distinct' in select_part else ''
        agg_f = self._get_agg_func(select_part)
        select_cols = self._get_select_col(select_part)
        select_term = self._make_select_term(select_cols, agg_f, distinct_term)

        conds = self._get_conds(where_part)
        triples, max_length = self._get_sparql_where_triples(select_cols, conds)
        return max_length

    def convert(self, sql):
        sql = self._replace_dulicate_column(sql)
        sparql = self._parse_sql(sql)
        return sparql

    def _parse_sql(self, sql):
        remain, where_part = sql.split(' where ')
        select_part, remain = remain.split(' from ')

        distinct_term = 'distinct' if 'distinct' in select_part else ''
        agg_f = self._get_agg_func(select_part)
        select_cols = self._get_select_col(select_part)
        select_term = self._make_select_term(select_cols, agg_f, distinct_term)

        conds = self._get_conds(where_part)
        where_term = self._make_where_term(select_cols, conds)
        sql = f'{select_term} {where_term}'
        return sql

    def _get_agg_func(self, select_part):
        for f in self.agg_func:
            if f in select_part.split():
                return f
        return None

    def _get_select_col(self, select_part):
        return [col.replace('"', '') for col in re.findall(self.sel_p, select_part)]

    def _get_conds(self, where_part):
        tokens = re.findall(self.cond_p, where_part)
        assert len(tokens) % 3 == 0
        conds = [[tokens[i].replace('"', '')] + tokens[i+1:i+3] for i in range(0, len(tokens), 3)]
        return conds

    def _make_select_term(self, select_cols, agg_f, distinct):
        term = self.sparql_select_template
        if agg_f:
            term = self.sparql_agg_template.format(AGG=agg_f, DISTINCT=distinct, COL=select_cols[0].lower())
            return term

        for col in select_cols:
            term += f' ?{col.lower()}'

        return term

    def _make_where_term(self, sel_cols, conds):
        triples, max_length = self._get_sparql_where_triples(sel_cols, conds)
        term = f'where {{ {". ".join(triples)}. }}'
        return term

    def _path2triples(self, path):
        path.reverse()
        t = [f"?{path.pop().lower()}"]
        triples = []
        while len(path) > 0:
            if len(t) == 0:
                t.append(f"?{path[-1].lower()}")
                path.pop()
            elif len(t) == 1:
                t.append(f"</{path[-1]}>")
            elif len(t) == 2:
                t.append(f"?{path[-1].lower()}")
            elif len(t) == 3:
                triples.append(t)
                t = []
            else:
                print('error')
        return triples

    def _get_sparql_where_triples(self, sel_cols, conds):
        max_length = 0
        triples = []
        filters = []
        for sel_col in sel_cols:
            for cond_col, op, val in conds:
                root2sel = nx.shortest_path(self.schema_graph, source=self.root, target=sel_col)
                sel_triples = self._path2triples(root2sel)

                root2con = nx.shortest_path(self.schema_graph, source=self.root, target=cond_col)
                if max_length < len(root2con):
                    max_length = len(root2con)

                if max_length < len(root2sel):
                    max_length = len(root2sel)

                if len(root2con) > 1:
                    con_triples = self._path2triples(root2con)
                    if op == '=':
                        con_triples = [self._fill_cond_value(t, cond_col, val) for t in con_triples]
                    else:
                        filters.append(f"""filter( ?{cond_col.lower()} {op} {val.replace('"', '')} )""")
                    a = [' '.join(t) for t in sel_triples]
                    b = [' '.join(t) for t in con_triples]

                else:
                    if op == '=':
                        sel_triples = [self._fill_cond_value(t, cond_col, val) for t in sel_triples]
                    else:
                        raise Exception()
                    a = [' '.join(t) for t in sel_triples]
                    b = []

                b = [t for t in b if b not in a]

                triples += list(OrderedDict(zip(a + b, repeat(None))))

        return list(OrderedDict(zip(triples, repeat(None)))) + filters, max_length

    def _fill_cond_value(self, t, cond_col, val):
        sub = t[0]
        rel = t[1]
        cond = t[2]

        if cond.replace('?', '') == cond_col:
            cond_name = cond.replace('?', '')
            if self.schema_type[cond_name] == 'entity':
                cond = f"""</{cond_name}/{val.replace('"', '')}>"""
            else:
                cond = f'{val}^^<{self.schema_type[cond_name]}>'

        elif sub.replace('?', '') == cond_col:
            sub_name = sub.replace('?', '')
            if self.schema_type[sub_name] == 'entity':
                sub = f"""</{sub_name}/{val.replace('"', '')}>"""
            else:
                raise Exception()

        return sub, rel, cond


if __name__ == '__main__':
    sql2sparql = SQL2SPARQL(True)

    sql = """SELECT MIN ( ADMISSIONS."AGE" ) 
                FROM ADMISSIONS
                WHERE ADMISSIONS."DIAGNOSIS" = "S/P FALL" AND ADMISSIONS."ADMITYEAR" >= "2119"
    """

    sql = sql2sparql._replace_dulicate_column(sql)
    print(sql)
    result_sparql = sql2sparql._parse_sql(sql)
    print(result_sparql)
