from pyspark.sql import functions as F
from pyspark.ml.feature import HashingTF, IDF
from pyspark.ml import Pipeline
from pyspark.sql.types import *
from pyspark.sql.functions import udf, split
import collections

def trim_array_elements(array_column):
    return [e.strip() for e in array_column]

trim_array_elements_udf = F.udf(trim_array_elements, ArrayType(StringType()))

class CategoriaKeyword:
    def __init__(self, categories_df, text_df, category_col, keyword_col, text_col):
        self.categories_df = categories_df
        self.text_df = text_df
        self.category_col = category_col
        self.keyword_col = keyword_col
        self.text_col = text_col

    def buscar_palavras_chave(self):
        # Prepare keywords for efficient lookup
        keywords = self.categories_df.select(self.keyword_col)
        keywords = keywords.withColumn('trimmed_keywords', trim_array_elements_udf(self.keyword_col))

        # Create a map for fast keyword-to-category lookup
        keyword_to_category_map = keywords.select('trimmed_keywords').distinct().rdd.flatMap(lambda x: x).collectAsMap()

        # Create a HashingTF tokenizer
        tokenizer = HashingTF(inputCol='trimmed_keywords', outputCol='tfFeatures')

        # Create an IDF model
        idf = IDF(inputCol='tfFeatures', outputCol='idfFeatures')

        # Create the feature transform pipeline
        feature_pipeline = Pipeline(stages=[tokenizer, idf])

        # Apply the pipeline to the keywords
        result = feature_pipeline.fit(keywords).transform(keywords)

        # Function to process individual texts and identify matching categories
        def processar_texto(text_features):
            text_vec = text_features.idfFeatures
            matching_categories = []
            for keyword in keyword_to_category_map.keys():
                if text_vec.indexOf(keyword_to_category_map[keyword]) >= 0:
                    matching_categories.append(keyword)
            return matching_categories

        # Apply the processing function to each text
        result_df = self.text_df.withColumn('text_features', result.select('idfFeatures').where(col('id') == col('id')).collect()[0]['idfFeatures'])

        result_df = result_df.withColumn("Categorias Encontradas", udf(processar_texto)(col('text_features')))

        return result_df[['id', 'tx_trsc_segm', 'Categorias Encontradas']]

# spark = SparkSession.builder.appName('CategoriaKeyword').getOrCreate()

# # DataFrames de exemplo (categories_df e text_df)
# categories_df = spark.createDataFrame([
#     (1, 'Cartões - Anuidade - Negociação', ['anuidade, negociação']), 
#     (2, 'Cartões - Anuidade - Estorno', ['anuidade, estorno']), 
#     (3, 'Cartões - Bloqueio', ['bloqueio'])
# ], ['id', 'Categoria', 'Palavras-Chave'])

# text_df = spark.createDataFrame([
#     (1, 'Estou com problemas com a anuidade do meu cartão.'), 
#     (2, 'Preciso de um bloqueio imediato no meu cartão.'), 
#     (3, 'Gostaria de negociar a anuidade.')
# ], ['id', 'text'])

# Criar o objeto CategoriaKeyword
categorizer = CategoriaKeyword(
    categories_df, 
    text_df, 
    category_col='Categoria',  
    keyword_col='Palavras-Chave',    
    text_col='tx_trsc_segm'    
)

# Executar a busca de palavras-chave
result_df = categorizer.buscar_palavras_chave()

# Exibir o resultado
result_df.show(truncate=False) py4j.security.Py4JSecurityException: Method public org.apache.spark.rdd.RDD org.apache.spark.api.java.JavaRDD.rdd() is not whitelisted on class class org.apache.spark.api.java.JavaRDD