from typing import Any, Dict, List
import importlib.util
from itertools import combinations
import multiprocessing as mp
import numpy as np
import pandas as pd
from numba import njit
from parallel_pandas import ParallelPandas
from concurrent.futures import ThreadPoolExecutor

from app.main.engine.action.ActionInterface import ActionInterface

# RELACION DE IDENTIFICADOR DE ACCION Y NOMBRE DE CLASE
relation_classname_identifier = {
    "match-and-exclude-records-actions": "MatchAndExcludeRecordsAction"
}

# CONFIGURACION DE SESION DE SPARK
MASTER = "local[*]"
DRIVER_MEMORY = "8g"
EXECUTOR_MEMORY = "8g"
MYSQL_JAR_PATH = "jars/mysql-connector-java-8.0.30.jar"

# EXCLUDE VALIDATION FIELD
EXCLUDE_ROWS_FIELD = "EXCLUDE_VALID"
# REDONDEO DE DECIMALES
ROUND_DECIMAL = 2


class MatchAndExcludeRecordsAction(ActionInterface):

    library_required = "pyspark"
    version_required = "3.4.0"

    def __init__(self, app) -> None:
        super().__init__(app)
        self.max_combinations = None
        self.comb_per_group = None
        self.exclude_pivot = None
        self.pivot_params = None
        self.ctp_params = None
        self.output = None
        self.config_params = ["max-records-per-combinations", "max-combinations-per-group", "exclude-entity-pivot"]

    def parser(self, descriptor: Dict[str, Any]):
        # Validar si pyspark y su versión está instalada
        pyspark_lib = importlib.util.find_spec(self.library_required)
        if pyspark_lib is None:
            raise ImportError(f"Librería requerida *{self.library_required}* no instalada")
        import pyspark
        version = pyspark.__version__
        if version != self.version_required:
            raise ImportError(f"Versión requerida no instalada. Requerida: {self.version_required}. Instalada: {version}")

        # Validación de parámetros de entrada
        entity_config_params = ["tablename", "id-column", "amount-column", "columns-group", "columns-transaction"]

        if "config-params" not in descriptor.keys():
            raise ReferenceError(f"Parámetro *params-input* no definido en descriptor")

        configs = descriptor["config-params"]
        for param in self.config_params:
            if param not in configs.keys():
                raise ReferenceError(f"Parámetro *{param}* no definido en configuracion")

        if "params-input" not in descriptor.keys():
            raise ReferenceError(f"Parámetro *params-input* no definido en descriptor")

        params_input = descriptor["params-input"]
        if "pivot-config" not in params_input.keys() or "counterpart-config" not in params_input.keys():
            raise ReferenceError(f"Pivot o Contraparte no definidos en descriptor")

        pivot_params, ctp_params = params_input["pivot-config"], params_input["counterpart-config"]
        for param in entity_config_params:
            if param not in pivot_params.keys() or param not in ctp_params.keys():
                raise ReferenceError(f"Parámetro *{param}* no encontrado en pivot o contraparte")

        self.max_combinations = configs["max-records-per-combinations"]
        self.comb_per_group = configs["max-combinations-per-group"]
        self.exclude_pivot = configs["exclude-entity-pivot"]
        self.pivot_params = pivot_params
        self.ctp_params = ctp_params

    def process(self, source_obj):

        # Inicializar la sesion de Spark
        session = self.createSession()

        # Traer la data desde BD tanto pivot como contraparte
        pivot_table, ctp_table = self.pivot_params["tablename"], self.ctp_params["tablename"]
        jdbc_conn = source_obj.create_spark_connection()
        jdbc_url = jdbc_conn["url"]
        jdbc_properties = jdbc_conn["properties"]
        pivot_df = session.read.jdbc(url=jdbc_url, table=pivot_table, properties=jdbc_properties)
        ctp_df = session.read.jdbc(url=jdbc_url, table=ctp_table, properties=jdbc_properties)

        # Agregar un prefijo a cada columna, tanto del pivot como contraparte. Actualizar campos del input
        # pivot: 'PIVOT_', contraparte: 'COUNTERPART_'
        for column in pivot_df.columns:
            if column == EXCLUDE_ROWS_FIELD:
                continue
            pivot_df = pivot_df.withColumnRenamed(column, "PIVOT_"+column)

        for column in ctp_df.columns:
            if column == EXCLUDE_ROWS_FIELD:
                continue
            ctp_df = ctp_df.withColumnRenamed(column, "COUNTERPART_"+column)

        for key_p, key_c in zip(self.pivot_params.keys(), self.ctp_params.keys()):
            if isinstance(self.pivot_params[key_p], str):
                self.pivot_params[key_p] = "PIVOT_"+self.pivot_params[key_p]
                self.ctp_params[key_c] = "COUNTERPART_"+self.ctp_params[key_c]
            else:
                self.pivot_params[key_p] = ["PIVOT_"+column for column in self.pivot_params[key_p]]
                self.ctp_params[key_c] = ["COUNTERPART_" + column for column in self.ctp_params[key_c]]

        from pyspark.sql.functions import sum, collect_list, round, when, col, lit

        pivot_cols = self.pivot_params["columns-transaction"].copy()
        if self.pivot_params["amount-column"] in pivot_cols:
            pivot_cols.remove(self.pivot_params["amount-column"])

        ctp_cols = self.ctp_params["columns-transaction"].copy()
        if self.ctp_params["amount-column"] in ctp_cols:
            ctp_cols.remove(self.ctp_params["amount-column"])

        comb_per_group = self.comb_per_group
        max_combinations = self.max_combinations

        # Ejecutamos lógica de excluir registros
        if len(self.pivot_params["columns-group"]) == 0 and len(self.ctp_params["columns-group"]) == 0:
            raise RuntimeError(f"Debe haber al menos pivot o contraparte agrupado")

        # Caso: 1 - Muchos
        elif len(self.pivot_params["columns-group"]) == 0 and len(self.ctp_params["columns-group"]) > 0:
            ctp_df2 = ctp_df.groupby(self.ctp_params["columns-group"]). \
                agg(round(sum(self.ctp_params["amount-column"]), ROUND_DECIMAL).alias(self.ctp_params["amount-column"]),
                    collect_list(self.ctp_params["id-column"]).alias(self.ctp_params["id-column"]))

            pivot_df2 = pivot_df

        # Caso: Muchos - 1
        elif len(self.pivot_params["columns-group"]) > 0 and len(self.ctp_params["columns-group"]) == 0:
            pivot_df2 = pivot_df.groupby(self.pivot_params["columns-group"]).\
                agg(round(sum(self.pivot_params["amount-column"]), ROUND_DECIMAL).alias(self.pivot_params["amount-column"]),
                    collect_list(self.pivot_params["id-column"]).alias(self.pivot_params["id-column"]))

            ctp_df2 = ctp_df

        # Caso: Muchos - Muchos
        elif len(self.pivot_params["columns-group"]) > 0 and len(self.ctp_params["columns-group"]) > 0:
            pivot_df2 = pivot_df.groupby(self.pivot_params["columns-group"]). \
                agg(round(sum(self.pivot_params["amount-column"]), ROUND_DECIMAL).alias(self.pivot_params["amount-column"]),
                    collect_list(self.pivot_params["id-column"]).alias(self.pivot_params["id-column"]))
            ctp_df2 = ctp_df.groupby(self.ctp_params["columns-group"]). \
                agg(round(sum(self.ctp_params["amount-column"]), ROUND_DECIMAL).alias(self.ctp_params["amount-column"]),
                    collect_list(self.ctp_params["id-column"]).alias(self.ctp_params["id-column"]))

        condition = [pivot_df2[col1] == ctp_df2[col2] for col1, col2 in zip(self.pivot_params["columns-transaction"],
                                                                            self.ctp_params["columns-transaction"])]
        total_merged = pivot_df2.join(ctp_df2, condition, 'left')

        total_merged = total_merged.withColumn("DIFF", when(col(self.ctp_params["columns-transaction"][0]).isNotNull(),
                                                            lit(0)).otherwise(lit(None)))
        total_merged = total_merged.select(*pivot_df2.columns, "DIFF")

        condition = [total_merged[col1] == ctp_df2[col2] for col1, col2 in zip(pivot_cols, ctp_cols)]
        merged = total_merged.join(ctp_df2, condition)

        merged = merged.withColumn("DIFF", when(col("DIFF").isNull(),
                                                total_merged[self.pivot_params["amount-column"]] - ctp_df2[self.ctp_params["amount-column"]]).otherwise(col("DIFF")))

        merged_df = merged.withColumn("DIFF", round(merged["DIFF"], ROUND_DECIMAL))
        if self.exclude_pivot:
            df = pivot_df
            group_cols = self.pivot_params["columns-group"]
            amount_col = self.pivot_params["amount-column"]
            id_col = self.pivot_params["id-column"]
        else:
            df = ctp_df
            group_cols = self.ctp_params["columns-group"]
            amount_col = self.ctp_params["amount-column"]
            id_col = self.ctp_params["id-column"]

        total_tmp_cols = group_cols + ["DIFF"]
        df3 = df.join(merged_df.select(*total_tmp_cols), group_cols)

        df3 = df3.toPandas()
        total_cols = group_cols + [amount_col, id_col, EXCLUDE_ROWS_FIELD,  "DIFF"]

        # ParallelPandas.initialize(n_cpu=mp.cpu_count(), split_factor=8, disable_pr_bar=True)
        df3 = df3.sort_values(group_cols + [amount_col])

        resultado = df3[total_cols].groupby(group_cols).apply(lambda x: custom_func(x, amount_col, id_col, max_combinations))
        resultado = resultado.reset_index()
        if len(resultado.columns) == 1:
            resultado = pd.DataFrame([], columns=group_cols + ["LISTA_DIFF"])
        else:
            resultado.columns = group_cols + ["LISTA_DIFF"]
        # print(resultado["LISTA_DIFF"].apply(lambda x: x if pd.notna(x) and x[0]!=-1 else x))
        meged2 = resultado.merge(merged_df.toPandas(), 'left', group_cols)
        print(meged2)
        meged2["LISTA_DIFF"] = meged2["LISTA_DIFF"].apply(self.handle_array)
        meged2 = meged2[(meged2['DIFF'] == 0) | ((meged2['DIFF'] != 0) & (meged2['LISTA_DIFF'].apply(len) > 0))]
        if meged2.empty:
            pass
        elif self.exclude_pivot:
            meged2['INTER_PIVOT_ID'] = meged2.apply(lambda row: self.array_except(row[self.pivot_params["id-column"]], row['LISTA_DIFF']), axis=1)
            meged2 = meged2.rename(columns={self.ctp_params["id-column"]: "INTER_CTP_ID"})
            if meged2['INTER_CTP_ID'].dtype == 'int64':
                merged_df['INTER_CTP_ID'] = merged_df['INTER_CTP_ID'].apply(lambda x: [x]).astype('object')
        else:
            meged2['INTER_CTP_ID'] = meged2.apply(lambda row: self.array_except(row[self.ctp_params["id-column"]], row['LISTA_DIFF']), axis=1)
            meged2 = meged2.rename(columns={self.pivot_params["id-column"]: "INTER_PIVOT_ID"})
            if meged2['INTER_PIVOT_ID'].dtype == 'int64':
                merged_df['INTER_PIVOT_ID'] = merged_df['INTER_PIVOT_ID'].apply(lambda x: [x]).astype('object')

        self.output = meged2

    def response(self):
        return self.output

    def handle_array(self, x):
        # print(type(x))
        if isinstance(x, np.ndarray):
            return x
        else:
            return []

    def array_except(self, arr1, arr2):
        if arr2 is None:
            return arr1
        else:
            return [item for item in arr1 if item not in arr2]

    def createSession(self, name: str = "app_engine_spark"):
        try:
            from pyspark.sql import SparkSession
            session = SparkSession.builder.master(MASTER) \
                .appName(name) \
                .config("spark.jars", MYSQL_JAR_PATH) \
                .config("spark.executor.extraClassPath", MYSQL_JAR_PATH) \
                .config("spark.driver.extraClassPath", MYSQL_JAR_PATH) \
                .config("spark.driver.memory", DRIVER_MEMORY) \
                .config("spark.executor.memory", EXECUTOR_MEMORY) \
                .getOrCreate()
            self.app.logger.info(f"Sesión creada exitosamente")
            return session
        except Exception as e:
            raise Exception(f"Error creando sesion Spark. {e}")


def custom_func(group, amount_field, id_field, max_combinations):
    diff = group["DIFF"].values[0]
    if pd.isna(diff) or diff == 0:
        return None
    group = group[group[EXCLUDE_ROWS_FIELD] == 'S']
    group[amount_field] = group[amount_field].astype(float)
    group = group.reset_index(drop=True)
    values = group[amount_field].values
    ids = group[id_field].values

    n = len(values)
    valores1 = encontrar_comb_1(values, diff)
    if valores1[0] != -1:
        indices = ids[valores1]
        return indices
    valores2 = encontrar_comb_2(values, diff, n)
    if valores2[0] != -1:
        indices = ids[valores2]
        return indices
    # Iterar sobre todos los índices posibles
    # valores4 = encontrar_comb_4(values, diff, n)
    # if valores4[0] != -1:
    #     indices = ids[valores4]
    #     return indices
    valores5 = encontrar_comb_5(values, diff, n)
    if valores5[0] != -1:
        indices = ids[valores5]
        return indices


@njit
def encontrar_comb_1(valores, target):
    indice = [-1]
    for idx, value in enumerate(valores):
        suma = value
        if round(suma, ROUND_DECIMAL) == target:
            indice = [idx for idx, val in enumerate(valores) if val in [value]]
            return indice

    return indice

@njit
def encontrar_comb_2(valores, target, n):
    indice = [-1]
    for i in range(n):
        array_except = np.delete(valores, i)
        for idx, value in enumerate(array_except):
            suma = value + valores[i]
            if round(suma, ROUND_DECIMAL) == target:
                indice = [idx for idx, val in enumerate(valores) if val in [value, valores[i]]]
                return indice

    return indice

@njit
def encontrar_comb_4(valores, target, n):
    indice = [-1]
    for i in range(n):
        a1 = np.delete(valores, i)
        for j in range(len(a1)):
            a2 = np.delete(a1, j)
            for k in range(len(a2)):
                array_except = np.delete(a2, k)
                for idx, value in enumerate(array_except):
                    suma = value + valores[i] + a1[j] + a2[k]
                    if round(suma, ROUND_DECIMAL) == target:
                        indice = [idx for idx, val in enumerate(valores) if val in [value, valores[i], a1[j], a2[k]]]
                        return indice

    return indice

@njit
def encontrar_comb_5(valores, target, n):
    indice = [-1]
    for i in range(n):
        a1 = np.delete(valores, i)
        for j in range(len(a1)):
            a2 = np.delete(a1, j)
            for k in range(len(a2)):
                a3 = np.delete(a2, k)
                for l in range(len(a3)):
                    array_except = np.delete(a2, l)
                    for idx, value in enumerate(array_except):
                        suma = value + valores[i] + a1[j] + a2[k] + a3[l]
                        if round(suma, ROUND_DECIMAL) == target:
                            indice = [idx for idx, val in enumerate(valores) if val in [value, valores[i], a1[j], a2[k], a3[l]]]
                            return indice
    return indice

