import json

import os

import re

import numpy as np

import pandas as pd

import networkx as nx

from build_mimicsqlstar_db.schema_mimic import SCHEMA, MAP_WITH_MIMICSQL

from collections import OrderedDict

from itertools import repeat

class SQL2SQL:

    def __init__(self):

        self.schema_graph = nx.Graph()

        for k, vs in SCHEMA.items():

            for v in vs:

                self.schema_graph.add_edge(k, v[0])

        self.inner_join_template = 'INNER JOIN {} ON {} = {}'

    def find_table(self, block):

        cols = re.findall('[A-Z_]+[.]"[\dA-Z_]+"', block)

        return [c.split('.')[0] for c in cols]

    def from_caluse(self, new_select):

        cols = re.findall('[A-Z_]+[.]"[\dA-Z_]+"', new_select)

        for col in cols:

            t, c = col.split('.')

            return t

    def cols_clause(self, block):

        cols = re.findall('[A-Z_]+[.]"[\dA-Z_]+"', block)

        #print(cols)

        for col in cols:

            t, c = col.split('.')

            new_tc = MAP_WITH_MIMICSQL[t][c.replace('"', '')]

            #print(new_tc)

            block = block.replace(col, new_tc)

            # print(col, new_tc)

        return block

    def translate(self, sql):

        remain, where_block = sql.split('WHERE')

        select_block, remain = remain.split('FROM')

        new_select = self.cols_clause(select_block)

        select_tables = self.find_table(new_select)

        from_table = self.from_caluse(new_select)

        new_where = self.cols_clause(where_block)

        where_tables = self.find_table(new_where)

        # print(new_select)

        # print(from_table)

        # print(new_where)

        # print(where_tables)

        inner_join_blocks = list()

        tables = list(set(where_tables + select_tables))

        for wt in tables:

            path = nx.shortest_path(self.schema_graph, source=from_table.replace(' ', ''), target=wt.replace(' ', ''))

            for i in range(len(path) - 1):

                a = path[i]

                b = path[i+1]

                key = [con[1] for con in SCHEMA[a] if con[0] == b][0]

                inner_join_block = self.inner_join_template.format(b, f'{a}."{key}"', f'{b}."{key}"')

                #print(inner_join_block)

                inner_join_blocks.append(inner_join_block)

        inner_join_blocks = list(OrderedDict(zip(inner_join_blocks, repeat(None))))

        new_inner_join = ' '.join(list(inner_join_blocks))

        #print(new_inner_join)

        temp = 'FROM '.join([new_select, from_table])

        temp = ' '.join([temp, new_inner_join])

        temp = ' '.join(temp.split(' ')) + ' '

        new_sql = 'WHERE'.join([temp, new_where])

        return new_sql

if __name__ == '__main__':

    convertor = SQL2SQL()

    sql = 'SELECT DEMOGRAPHIC."GENDER",DEMOGRAPHIC."INSURANCE" FROM DEMOGRAPHIC WHERE DEMOGRAPHIC."SUBJECT_ID" = "81923"'

    new_sql = convertor.translate(sql)

    print(sql)

    print(new_sql)