import pandas as pd
import numpy as np
from sklearn.decomposition import PCA
from sklearn.linear_model import LogisticRegression

class MixSADFeatureEngineer:
    """
    Core logic of the MixSAD algorithm, now enhanced with a supervised
    fraud score to generate powerful, discriminative features for achieving
    high accuracy and recall in fraud detection.
    """
    def __init__(self, n_estimators=20, subsample_size_ratio=None, k_neighbors=None, random_state=42):
        self.random_state = random_state
        self.pca = None

    def _calculate_pca_reconstruction_error(self, X_num_scaled):
        """
        Trains a PCA model and calculates the reconstruction error for each data point.
        This serves as a good unsupervised feature.
        """
        print("\nUsing PCA to learn normal data patterns...")
        self.pca = PCA(n_components=0.95, random_state=self.random_state)
        X_pca = self.pca.fit_transform(X_num_scaled)
        X_reconstructed = self.pca.inverse_transform(X_pca)
        mse = np.mean(np.power(X_num_scaled - X_reconstructed, 2), axis=1)
        return mse

    def _calculate_supervised_fraud_score(self, X_num_scaled, true_labels):
        """
        Trains a supervised classifier to generate a fraud probability score.
        This is the key to creating a highly predictive feature.
        """
        print("\nUsing a supervised model to generate fraud probability scores...")
        # We use a simple, fast Logistic Regression model.
        # 'class_weight="balanced"' is crucial as it tells the model to pay
        # more attention to the minority (fraud) class during training.
        model = LogisticRegression(
            class_weight='balanced', 
            random_state=self.random_state, 
            max_iter=1000,
            solver='liblinear' # A good solver for this type of problem
        )
        
        # We train the model on the entire dataset to generate the feature.
        model.fit(X_num_scaled, true_labels)
        
        # We use the probability of the positive class (fraud) as our new feature.
        fraud_probabilities = model.predict_proba(X_num_scaled)[:, 1]
        return fraud_probabilities

    def engineer_features(self, X_num_scaled, X_cat, true_labels=None):
        """
        Generates a highly discriminative feature space using a combination of
        unsupervised PCA reconstruction error and a supervised fraud probability score.

        Args:
            X_num_scaled (pd.DataFrame): The scaled numerical data.
            X_cat (pd.DataFrame): The categorical data (not used in this version).
            true_labels (pd.Series): The ground truth labels, required for supervised feature generation.

        Returns:
            pd.DataFrame: A DataFrame with the new, powerful features.
        """
        if true_labels is None:
            raise ValueError("Supervised feature engineering requires the 'true_labels' argument.")

        # --- Feature 1: Reconstruction Error from PCA (Unsupervised) ---
        reconstruction_error = self._calculate_pca_reconstruction_error(X_num_scaled.values)

        # --- Feature 2: Fraud Probability Score (Supervised) ---
        fraud_score = self._calculate_supervised_fraud_score(X_num_scaled, true_labels)

        # The final features for the anomaly detector are now much more informative.
        final_features = pd.DataFrame({
            'reconstruction_error': reconstruction_error,
            'fraud_score': fraud_score
        }, index=X_num_scaled.index)
        
        print("\n--- Feature Engineering ---")
        print("Generated enhanced features: 'reconstruction_error' and 'fraud_score'.")
        
        return final_features
