from typing import Any, Dict, List
import importlib.util
import numpy as np
import pandas as pd
import json
from dask import dataframe as dd
from numba import jit, types, typed
from wrapt_timeout_decorator import timeout

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"
}

# 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.timeout = None
        self.exclude_pivot = None
        self.pivot_params = None
        self.ctp_params = None
        self.output = None
        self.config_params = ["max-records-per-combinations", "max-timeout-per-combinations", "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.timeout = configs["max-timeout-per-combinations"]
        self.exclude_pivot = configs["exclude-entity-pivot"]
        self.pivot_params = pivot_params
        self.ctp_params = ctp_params

    def process(self, source_obs):
        try:
            @timeout(self.timeout)
            def __process(source_obj):
                # Traer la data desde BD tanto pivot como contraparte
                pivot_table, ctp_table = self.pivot_params["tablename"], self.ctp_params["tablename"]
                dialect = source_obj.get_dialect()

                pivot_df = dd.read_sql_table(pivot_table, dialect, index_col=self.pivot_params["id-column"],
                                             npartitions=4)
                pivot_df = pivot_df.reset_index()
                ctp_df = dd.read_sql_table(ctp_table, dialect, index_col=self.ctp_params["id-column"], npartitions=4)
                ctp_df = ctp_df.reset_index()

                # Agregar un prefijo a cada columna, tanto del pivot como contraparte. Actualizar campos del input
                # pivot: 'PIVOT_', contraparte: 'COUNTERPART_'
                # Iterar sobre las columnas del DataFrame
                for column in pivot_df.columns:
                    if column == EXCLUDE_ROWS_FIELD:
                        continue
                    new_column_name = "PIVOT_" + column
                    pivot_df = pivot_df.rename(columns={column: new_column_name})

                for column in ctp_df.columns:
                    if column == EXCLUDE_ROWS_FIELD:
                        continue
                    new_column_name = "COUNTERPART_" + column
                    ctp_df = ctp_df.rename(columns={column: new_column_name})

                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"])

                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({
                        self.ctp_params["amount-column"]: 'sum',  # Sumar la columna de cantidades
                        self.ctp_params["id-column"]: list
                    })
                    ctp_df2 = ctp_df2.reset_index()

                    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({
                        self.pivot_params["amount-column"]: 'sum',
                        self.pivot_params["id-column"]: list
                    })
                    pivot_df2 = pivot_df2.reset_index()

                    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({
                        self.pivot_params["amount-column"]: 'sum',
                        self.pivot_params["id-column"]: list
                    })
                    pivot_df2 = pivot_df2.reset_index()

                    ctp_df2 = ctp_df.groupby(self.ctp_params["columns-group"]).agg({
                        self.ctp_params["amount-column"]: 'sum',  # Sumar la columna de cantidades
                        self.ctp_params["id-column"]: list
                    })
                    ctp_df2 = ctp_df2.reset_index()

                pivot_df2[self.pivot_params["amount-column"]] = pivot_df2[self.pivot_params["amount-column"]].round(
                    ROUND_DECIMAL)
                ctp_df2[self.ctp_params["amount-column"]] = ctp_df2[self.ctp_params["amount-column"]].round(
                    ROUND_DECIMAL)
                total_merged = pivot_df2.merge(ctp_df2, 'left', left_on=self.pivot_params["columns-transaction"],
                                               right_on=self.ctp_params["columns-transaction"])
                total_merged = total_merged.map_partitions(self.add_diff_column)
                selected_columns = list(pivot_df2.columns) + ['DIFF']
                total_merged = total_merged[selected_columns]

                merged = total_merged.merge(ctp_df2, 'inner', left_on=pivot_cols, right_on=ctp_cols)

                merged['DIFF'] = merged['DIFF'].where(merged['DIFF'].notnull(),
                                                      merged[self.pivot_params["amount-column"]] - merged[
                                                          self.ctp_params["amount-column"]])

                if len(self.pivot_params["columns-group"]) == 0 and len(self.ctp_params["columns-group"]) > 0:
                    merged = merged.drop_duplicates(subset=pivot_cols)
                elif len(self.pivot_params["columns-group"]) > 0 and len(self.ctp_params["columns-group"]) == 0:
                    merged = merged.drop_duplicates(subset=ctp_cols)

                merged_df = merged.assign(DIFF=lambda partition: partition["DIFF"].round(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.merge(merged_df[total_tmp_cols], 'inner', on=group_cols)
                df3 = df3.compute()
                total_cols = group_cols + [amount_col, id_col, EXCLUDE_ROWS_FIELD,  "DIFF"]
                resultado = df3.groupby(group_cols)[total_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"]
                resultado = dd.from_pandas(resultado, npartitions=4)

                meged2 = resultado.merge(merged_df, 'left', group_cols)

                meged2 = meged2.map_partitions(lambda partition: partition.assign(
                    LISTA_DIFF=partition['LISTA_DIFF'].apply(lambda x: [] if pd.isna(x) else x)), meta=meged2.dtypes.to_dict())

                meged2 = meged2[
                    (meged2['DIFF'] == 0) |
                    ((meged2['DIFF'] != 0) & meged2['LISTA_DIFF'].apply(
                        lambda x: True if not pd.isna(x) and ((isinstance(x, List) and len(x) > 0) or (isinstance(x, str) and len(x) > 2)) else False))
                    ]
                meged2 = meged2.compute()

                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':
                        meged2['INTER_CTP_ID'] = meged2['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':
                        meged2['INTER_PIVOT_ID'] = meged2['INTER_PIVOT_ID'].apply(lambda x: [x]).astype('object')

                return meged2

        except TimeoutError as e:
            raise TimeoutError(f"Tiempo límite superado. {e}")

        self.output = __process(source_obs)

    def response(self):
        return self.output

    def add_diff_column(self, partition):
        partition['DIFF'] = np.where(partition[self.ctp_params["columns-transaction"][0]].notnull(), 0, np.nan)
        return partition

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

    def array_except(self, arr1, arr2):
        # print(arr2)
        if arr2 is None:
            return arr1
        elif not isinstance(arr2, List):
            cadena_sin_corchetes = arr2.strip('[]')
            partes = cadena_sin_corchetes.split()
            # print(partes)
            arr2 = [int(numero) for numero in partes]
        arr1 = json.loads(arr1.replace(" ", ""))
        return [item for item in arr1 if item not in arr2]


def custom_func(group, amount_field, id_field, max_combinations):
    diff_value = group["DIFF"].values[0]
    if np.isnan(diff_value):
        return None
    diff = int(diff_value*(10**ROUND_DECIMAL))
    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
    values *= (10**ROUND_DECIMAL)
    values = values.astype(np.int64)
    ids = group[id_field].values
    tam = len(values)
    tam = tam if tam <= max_combinations else max_combinations

    result = subset_sum_iter(values, diff, tam)
    indices = ids[np.isin(values, result)]
    return indices


@jit(nopython=False)
def subset_sum_iter(numbers, target, num_elements):

    # Initialize solutions list
    final = typed.List.empty_list(types.int64)
    for step in range(1, num_elements+1):
        # Build first index by taking the first num_elements from the numbers
        indices = list(range(step))

        while True:
            for i in range(step):
                if indices[i] != i + len(numbers) - step:
                    break
            else:
                # No combinations left
                break

            # Increase current index and all its following ones
            indices[i] += 1
            for j in range(i + 1, step):
                indices[j] = indices[j - 1] + 1

            # Check current solution
            solution = typed.List.empty_list(types.int64)
            for i in indices:
                solution.append(numbers[i])
            if round(sum(solution), ROUND_DECIMAL) == target:
                final = solution
                break
        if len(final) > 0:
            break

    return final

