import pandas as pd
from itertools import combinations
from collections import defaultdict
import json
import os

class AprioriModel:
    def __init__(self, min_support=0.01):
        self.min_support = min_support
        self.transactions = []

    def load_data(self, folder_path):
        """Load CSV files from a folder and combine them into a single DataFrame."""
        combined_df = pd.DataFrame()
        
        for file in os.listdir(folder_path):
            if file.endswith('.csv'):
                file_path = os.path.join(folder_path, file)
                df = pd.read_csv(file_path)
                combined_df = pd.concat([combined_df, df], ignore_index=True)
        
        return combined_df

    def preprocess_data(self, data):
        """Prepare basket data and transaction sets from raw data."""
        basket = data.groupby(['Transaction ID', 'Product'])['Quantity'].sum().unstack().fillna(0)
        basket = basket.applymap(lambda x: 1 if x > 0 else 0)

        for _, row in basket.iterrows():
            self.transactions.append(set(row[row == 1].index))

        return basket

    def calculate_pair_support(self):
        """Calculate support for all pairs in the transactions."""
        pair_count = defaultdict(int)
        
        for transaction in self.transactions:
            for pair in combinations(transaction, 2):
                pair_count[frozenset(pair)] += 1
        
        return {
            pair: count for pair, count in pair_count.items()
            if count / len(self.transactions) >= self.min_support
        }

    def generate_rules(self, frequent_pairs):
        """Generate association rules based on frequent pairs."""
        rules = []

        for pair, support_count in frequent_pairs.items():
            for item in pair:
                antecedent = frozenset([item])
                consequent = pair - antecedent

                antecedent_support = sum(1 for t in self.transactions if antecedent.issubset(t)) / len(self.transactions)
                confidence = (support_count / len(self.transactions)) / antecedent_support
                consequent_support = sum(1 for t in self.transactions if consequent.issubset(t)) / len(self.transactions)
                lift = confidence / consequent_support

                rules.append({
                    "Driver": list(antecedent)[0],
                    "Attach": list(consequent)[0],
                    "Support": support_count / len(self.transactions),
                    "Confidence": confidence,
                    "Lift": lift
                })

        return rules

    def save_rules(self, rules, output_path):
        """Save generated rules to a CSV file."""
        rules_df = pd.DataFrame(rules)
        rules_df.to_csv(output_path, index=False)
        print("Rules saved successfully to", output_path)

# This function is optional for Hugging Face Inference API (if using API)
def predict(data):
    # Example of how to process input for inference
    model = AprioriModel(min_support=0.01)
    processed_data = model.load_data(data['folder_path'])
    model.preprocess_data(processed_data)
    frequent_pairs = model.calculate_pair_support()
    rules = model.generate_rules(frequent_pairs)
    return rules