import pandas as pd

from openlog.datamodel.connection.interfaces.categories_interface import (
    CategoriesInterface,
)
from openlog.datamodel.connection.spatialite.pivot_spatialite_assay_interface import (
    PivotSpatialiteAssayInterface,
)


class GeoticRegistry:
    def __init__(self):
        # key(spatialite table) : value(geotic table)
        self.assay = {}
        # Nested dictionnary. key : spatialite table value : key(spatialite column) : value(geotic column)
        self.columns = {}
        self.col_types = {}

    def add_assay(self, assay, table):
        self.assay[assay] = table

    def add_column(self, assay, col_name, col_table):
        if self.columns.get(assay) is None:
            self.columns[assay] = {col_name: col_table}
        else:
            self.columns[assay][col_name] = col_table

    def add_column_type(self, assay, col_name, col_type):
        if self.col_types.get(assay) is None:
            self.col_types[assay] = {col_name: col_type}
        else:
            self.col_types[assay][col_name] = col_type

    def get_geotic_table(self, assay):
        return self.assay[assay]

    def get_geotic_column(self, assay, col_name):
        if col_name in ["Nom", "Profondeur", "ProfA", "ProfDe"]:
            return col_name
        else:
            return self.columns[assay][col_name]

    def get_column_type(self, assay, col_name):
        return self.col_types[assay][col_name]


class GeoticPivotSpatialiteAssayInterface(PivotSpatialiteAssayInterface):
    def __init__(
        self, engine, session, spatialite_session, categories_iface: CategoriesInterface
    ):
        # Geotic registry to map column/table names
        self.registry = GeoticRegistry()
        super().__init__(engine, session, spatialite_session, categories_iface)

    def check_if_column_exist(self, table, col):
        """
        Check if column really exist in origin database table.
        """

        q = f"SELECT COUNT(*) FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_NAME = '{table}' AND COLUMN_NAME = '{col}'"
        res = self.origin_session.execute(q).fetchall()[0][0]
        return res > 0

    def initialize_tblperso(self):
        """
        Extract informations of user-defined geotic tables and fill definition tables.
        Informations are stored in tblPerso and tblColonnePerso tables.
        """
        # geotic tables definition.
        query = """
        SELECT NomTable, NomAffiche, TypeTable FROM dbo.tblTablePerso;
        """
        assays = self.origin_session.execute(query).fetchall()
        assays = [assay for assay in assays if assay[0] is not None]
        # assays and assay_definition
        domain = "depth"
        for assay in assays:
            variable, display, data_extent = assay
            display = self.clean_name(display)
            self.registry.add_assay(display, variable)

            if data_extent > 1:
                continue

            data_extent = "discrete" if data_extent == 0 else "extended"

            q = f'INSERT INTO assays VALUES("{display}", "{domain}", "{data_extent}")'
            self.session.execute(q)

            x_col = "Profondeur" if data_extent == "discrete" else "ProfDe"
            x_col_end = "NULL" if data_extent == "discrete" else '"ProfA"'
            q = f'INSERT INTO assay_definition VALUES("{display}", "{display}", "Nom", "{x_col}", "dataset", {x_col_end}, "")'
            self.session.execute(q)

        self.session.commit()

        # column definitions
        query = """
        SELECT NomTable, NomAffiche, NomColonne, TypeColonne, NomDisplay FROM dbo.tblColonnePerso AS c LEFT JOIN dbo.tblTablePerso AS t ON t.idKey = c.idTablePerso;
        """
        columns = self.origin_session.execute(query).fetchall()
        columns = [column for column in columns if column[0] is not None]
        # assay_column and assay_column_definition
        keys = []
        for column in columns:
            table, table_name, col, col_type, col_name = column
            col_name = self.clean_name(col_name)
            table_name = self.clean_name(table_name)
            key = (col_name, table)
            if key not in keys:
                if col_type in [0, 1, 15, 16, 21]:
                    serie_type = "numerical"
                elif col_type == 2:
                    serie_type = "nominal"
                else:
                    continue

                self.registry.add_column_type(table_name, col_name, serie_type)
                self.registry.add_column(table_name, col_name, col)
                # check if column exist in Geotic table
                if not self.check_if_column_exist(
                    table=self.registry.get_geotic_table(table_name), col=col
                ):
                    continue

                q = f'INSERT INTO assay_column_definition VALUES("{col_name}", "{table_name}", "{col_name}");'
                self.session.execute(q)

                q = f'INSERT INTO assay_column VALUES("{col_name}", "{table_name}", "{serie_type}", "", "", "", "", "", "", "");'
                self.session.execute(q)
                keys.append(key)
        self.session.commit()

    def initialize_other(self):

        columns = self.origin_session.execute(
            "SELECT TABLE_NAME, COLUMN_NAME, DATA_TYPE FROM INFORMATION_SCHEMA.COLUMNS;"
        ).fetchall()
        df = pd.DataFrame.from_records(
            columns, columns=["TABLE_NAME", "COLUMN_NAME", "DATA_TYPE"]
        )

        # extended data
        extended = (
            df.loc[df["COLUMN_NAME"].isin(["idNuSondage", "ProfA", "ProfDe"]), :]
            .pivot(index="TABLE_NAME", columns="COLUMN_NAME", values="DATA_TYPE")
            .reset_index()
        )
        extended = extended[extended.isna().sum(axis=1) == 0]
        extended = extended[~extended["TABLE_NAME"].str.contains("tblPerso")]

        all_cols = df[df["TABLE_NAME"].isin(extended["TABLE_NAME"])]
        all_cols = all_cols[
            ~all_cols["COLUMN_NAME"].isin(["idKey", "idNuSondage", "ProfDe", "ProfA"])
        ]

        def add_row_number(x):
            return list(range(x.shape[0]))

        all_cols["row"] = 0
        all_cols["row"] = all_cols.groupby("TABLE_NAME")["row"].transform(
            add_row_number
        )
        all_cols_name = all_cols.pivot(
            index="TABLE_NAME", columns="row", values="COLUMN_NAME"
        )
        all_cols_type = all_cols.pivot(
            index="TABLE_NAME", columns="row", values="DATA_TYPE"
        )

        domain = "depth"
        for names, types in zip(all_cols_name.iterrows(), all_cols_type.iterrows()):
            table_name, col_names = names
            _, col_types = types

            # assay definitions
            self.registry.add_assay(table_name, table_name)
            q = f'INSERT INTO assays VALUES("{table_name}", "{domain}", "extended")'
            self.session.execute(q)
            q = f'INSERT INTO assay_definition VALUES("{table_name}", "{table_name}", "Nom", "ProfDe", "dataset", "ProfA", "")'
            self.session.execute(q)

            # column definitions
            keys = []
            for col_name, col_type in zip(col_names, col_types):
                key = (col_name, table_name)

                if key not in keys:
                    if col_type in ["numeric", "float", "int"]:
                        serie_type = "numerical"
                    elif col_type == "nvarchar":
                        serie_type = "nominal"
                    else:
                        continue

                    self.registry.add_column_type(table_name, col_name, serie_type)
                    self.registry.add_column(table_name, col_name, col_name)

                    q = f'INSERT INTO assay_column_definition VALUES("{col_name}", "{table_name}", "{col_name}");'
                    self.session.execute(q)

                    q = f'INSERT INTO assay_column VALUES("{col_name}", "{table_name}", "{serie_type}", "", "", "", "", "", "", "");'
                    self.session.execute(q)
                    keys.append(key)

        self.session.commit()

        # discrete data
        discrete = (
            df.loc[df["COLUMN_NAME"].isin(["idNuSondage", "Profondeur"]), :]
            .pivot(index="TABLE_NAME", columns="COLUMN_NAME", values="DATA_TYPE")
            .reset_index()
        )
        discrete = discrete[discrete.isna().sum(axis=1) == 0]
        discrete = discrete[~discrete["TABLE_NAME"].str.contains("tblPerso")]

        all_cols = df[df["TABLE_NAME"].isin(discrete["TABLE_NAME"])]
        all_cols = all_cols[
            ~all_cols["COLUMN_NAME"].isin(["idKey", "idNuSondage", "Profondeur"])
        ]
        all_cols["row"] = 0
        all_cols["row"] = all_cols.groupby("TABLE_NAME")["row"].transform(
            add_row_number
        )
        all_cols_name = all_cols.pivot(
            index="TABLE_NAME", columns="row", values="COLUMN_NAME"
        )
        all_cols_type = all_cols.pivot(
            index="TABLE_NAME", columns="row", values="DATA_TYPE"
        )

        for names, types in zip(all_cols_name.iterrows(), all_cols_type.iterrows()):
            table_name, col_names = names
            _, col_types = types

            # assay definitions
            self.registry.add_assay(table_name, table_name)
            q = f'INSERT INTO assays VALUES("{table_name}", "{domain}", "discrete")'
            self.session.execute(q)
            q = f'INSERT INTO assay_definition VALUES("{table_name}", "{table_name}", "Nom", "Profondeur", "dataset", "", "")'
            self.session.execute(q)

            # column definitions
            keys = []
            for col_name, col_type in zip(col_names, col_types):
                key = (col_name, table_name)

                if key not in keys:
                    if col_type in ["numeric", "float", "int"]:
                        serie_type = "numerical"
                    elif col_type == "nvarchar":
                        serie_type = "nominal"
                    else:
                        continue

                    self.registry.add_column_type(table_name, col_name, serie_type)
                    self.registry.add_column(table_name, col_name, col_name)

                    q = f'INSERT INTO assay_column_definition VALUES("{col_name}", "{table_name}", "{col_name}");'
                    self.session.execute(q)

                    q = f'INSERT INTO assay_column VALUES("{col_name}", "{table_name}", "{serie_type}", "", "", "", "", "", "", "");'
                    self.session.execute(q)
                    keys.append(key)

        self.session.commit()

    def set_categorical(self, assays: [str]):
        for assay in assays:
            q = f'UPDATE assay_column SET series_type = "categorical" WHERE assay = "{assay}"'
            self.session.execute(q)
        self.session.commit()

    def initialize_definitions(self):
        """
        Extract geotic tables informations and fill assay definitions in xplordb schema.
        """

        self.initialize_tblperso()
        self.initialize_other()
        self.set_categorical(assays=["tblLithos1", "tblLithos2"])

    def copy_assay_from_origin(self, variable: str):

        # check if table already exists
        if self.check_if_table_exist(variable):
            return

        db_def = self.get_assay_table_definition(variable)

        # extract col types : key is (col_name, col_table)
        y_types = {
            key: self.registry.get_column_type(variable, key)
            for key in db_def.y_col.keys()
        }
        for key, value in y_types.items():
            if value == "numerical":
                y_types[key] = "FLOAT"
            else:
                y_types[key] = "VARCHAR"
        col_types = {}
        col_types["Nom"] = "VARCHAR"
        col_types[db_def.x_col] = "FLOAT"
        if db_def.x_end_col:
            col_types[db_def.x_end_col] = "FLOAT"

        col_types.update(y_types)
        # define an empty table in spatialite
        body = ",".join(
            [f"{col_name} {col_type}" for col_name, col_type in col_types.items()]
        )

        q = f"""
            CREATE TABLE IF NOT EXISTS {variable} (
            {body}
            );
            """
        self.session.execute(q)
        self.session.commit()

        # extract geotic data
        q = "SELECT "
        q += ",".join(
            [
                f"{self.registry.get_geotic_column(variable, col_name)}"
                for col_name in col_types.keys()
            ]
        )
        q += f" FROM dbo.{self.registry.get_geotic_table(db_def.table_name)} as t"
        q += " JOIN dbo.tblNuSondage as s ON t.idNuSondage = s.idKey;"

        data = self.origin_session.execute(q).fetchall()

        # insert into spatialite table
        for row in data:
            values = []
            for elt, col_type in zip(row, col_types.values()):
                if elt is None:
                    values.append("NULL")
                    continue
                if col_type == "INT":
                    values.append(f"{int(str(elt))}")
                elif col_type == "FLOAT":
                    values.append(f"{float(str(elt))}")
                elif col_type == "VARCHAR":
                    values.append(f'"{str(elt)}"')

            q = f"INSERT INTO {variable} VALUES({','.join(values)});"
            self.session.execute(q)
        self.session.commit()
