import os
import joblib
import pandas as pd
import numpy as np

# Path to model artifacts
ARTIFACTS_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "model_artifacts"))

class FraudPredictor:
    def __init__(self):
        self.artifacts_dir = ARTIFACTS_DIR
        self.models = {}
        self.scaler = None
        self.selector = None
        self.train_median = None
        self.ohe_columns = None
        self.cat_cols = None
        self.constant_cols = None
        self.heavy_missing_cols = None
        
        self.load_artifacts()

    def load_artifacts(self):
        print(f"Loading model artifacts from {self.artifacts_dir}...")
        try:
            self.cat_cols = joblib.load(os.path.join(self.artifacts_dir, 'cat_cols.pkl'))
            self.constant_cols = joblib.load(os.path.join(self.artifacts_dir, 'constant_cols.pkl'))
            self.heavy_missing_cols = joblib.load(os.path.join(self.artifacts_dir, 'heavy_missing_cols.pkl'))
            self.ohe_columns = joblib.load(os.path.join(self.artifacts_dir, 'model_features.pkl'))
            self.train_median = joblib.load(os.path.join(self.artifacts_dir, 'train_median.pkl'))
            self.scaler = joblib.load(os.path.join(self.artifacts_dir, 'robust_scaler.pkl'))
            self.selector = joblib.load(os.path.join(self.artifacts_dir, 'variance_selector.pkl'))
            
            # Load models on demand or pre-load them
            self.models['voting'] = joblib.load(os.path.join(self.artifacts_dir, 'voting_classifier_model.pkl'))
            self.models['xgboost'] = joblib.load(os.path.join(self.artifacts_dir, 'xgboost_model.pkl'))
            self.models['random_forest'] = joblib.load(os.path.join(self.artifacts_dir, 'random_forest_model.pkl'))
            print("Successfully loaded all artifacts and models.")
        except Exception as e:
            print(f"Error loading model artifacts: {e}")
            raise e

    def predict_dataframe(self, df: pd.DataFrame, model_name: str = 'voting', threshold: float = 0.30) -> pd.DataFrame:
        """
        Processes a whole DataFrame and returns prediction results (probability, flag, risk_level).
        Uses a highly optimized vectorized approach.
        """
        if model_name not in self.models:
            model_name = 'voting'
        model = self.models[model_name]

        processed_df = df.copy()

        # 1. Convert dtypes for consistency
        for col_dtype in processed_df.select_dtypes(include='float64').columns:
            processed_df[col_dtype] = processed_df[col_dtype].astype("float32")

        # 2. Drop constant and heavy missing columns from training
        cols_to_drop = [col for col in self.constant_cols if col in processed_df.columns] + \
                       [col for col in self.heavy_missing_cols if col in processed_df.columns]
        if cols_to_drop:
            processed_df.drop(columns=cols_to_drop, inplace=True, errors='ignore')

        # 3. Feature Engineering (Date-Time)
        if 'F3888' in processed_df.columns:
            # Parse datetime safely
            dates = pd.to_datetime(processed_df['F3888'], errors='coerce')
            
            # Combine all new time features in a single concat to avoid fragmentation
            time_df = pd.DataFrame({
                'transaction_month': dates.dt.month,
                'year': dates.dt.year,
                'day': dates.dt.day,
                'day_of_week': dates.dt.dayofweek,
                'transaction_quarter': dates.dt.quarter,
                'week_of_year': dates.dt.isocalendar().week.astype(float),
                'day_of_year': dates.dt.dayofyear
            }, index=processed_df.index)
            
            time_df['is_weekend'] = ((time_df['day_of_week'] == 5) | (time_df['day_of_week'] == 6)).astype(int)
            
            processed_df = pd.concat([processed_df, time_df], axis=1)
            processed_df.drop(columns=['F3888'], inplace=True, errors='ignore')

        # 4. Drop indices & target
        cols_to_drop_pre_ohe = ['Unnamed: 0', 'F3923']
        processed_df.drop(columns=[col for col in cols_to_drop_pre_ohe if col in processed_df.columns], inplace=True, errors='ignore')

        # 5. One-Hot Encoding for categorical columns
        cols_to_ohe_exist = [col for col in self.cat_cols if col in processed_df.columns]
        if cols_to_ohe_exist:
            processed_df = pd.get_dummies(processed_df, columns=cols_to_ohe_exist, drop_first=True)

        # 6. Align columns with training OHE columns
        processed_df = processed_df.reindex(columns=self.ohe_columns, fill_value=0)

        # 7. Impute missing values
        imputed_data = processed_df.fillna(self.train_median)

        # 8. Scale
        scaled_data = self.scaler.transform(imputed_data)

        # 9. Feature Selection
        final_processed_data = self.selector.transform(scaled_data)

        # 10. Predict probabilities
        probs = model.predict_proba(final_processed_data)[:, 1]

        # 11. Create results dataframe
        results = pd.DataFrame({
            'probability': np.round(probs, 4),
            'is_suspicious': probs >= threshold
        })
        
        # Risk levels
        results['risk_level'] = np.where(
            probs > 0.7, 'High',
            np.where(probs >= threshold, 'Medium', 'Low')
        )
        
        return results

    def predict_transaction(self, data_row: dict, model_name: str = 'voting', threshold: float = 0.30) -> dict:
        """
        Accepts a single raw row (as dict) and returns risk assessment.
        """
        df = pd.DataFrame([data_row])
        results_df = self.predict_dataframe(df, model_name, threshold)
        row_res = results_df.iloc[0]
        return {
            'probability': float(row_res['probability']),
            'is_suspicious': bool(row_res['is_suspicious']),
            'risk_level': str(row_res['risk_level'])
        }

# Global predictor instance
predictor = None

def get_predictor():
    global predictor
    if predictor is None:
        predictor = FraudPredictor()
    return predictor