import pandas as pd
import numpy as np
from sklearn.ensemble import IsolationForest
from statsmodels.tsa.stattools import acf
from .Preprocessing import get_dynamic_lags

def detect_time_series_anomalies_isoforest(
    group,
    variable,
    date_column,
    eval_period,
    ):
    
    # ... (Docstring omitted for brevity) ...
    group = group.reset_index(drop=True)
    
    # Get copy of the original group df
    group_orig = group.copy()
    
    # Keep only relevant columns for analysis, which are variable and date_column
    # group_orig may have columns that have NaNs or NoneTypes and those rows will be removed in a .dropna() function below
    group = group[[variable, date_column]]
    group[date_column] = pd.to_datetime(group[date_column])
    group = group.copy().sort_values(date_column).reset_index(drop=True)
    group['set'] = np.where(np.arange(len(group)) >= len(group) - eval_period, 'TEST', 'TRAIN')
    
    
    try:
        all_results = []

        # ===================================================================
        # STEP 1: Evaluate all points in the initial TRAIN period
        # ===================================================================
        
        # Get the cutoff date for initial train period
        initial_cutoff_date = group[group['set'] == 'TRAIN'][date_column].max()
        
        # Prepare the full group with features
        model_group_initial = group.copy()
        
        # Get train set to determine lags
        train_initial = model_group_initial[model_group_initial['set'] == 'TRAIN'].copy()
        lags = get_dynamic_lags(train_initial[variable])

        # Create lag features on the entire model_group DF
        for lag in lags:
            model_group_initial[f'lag{lag}'] = model_group_initial[variable].shift(lag)

        # Get rolling stats features for the entire model_group DF
        rolling_stats_features = []    
        for w in [int(np.ceil(max(lags)/4)), int(np.ceil(max(lags)/2)), int(max(lags))]:
            if w >= 3:
                rolling_stats_features.append('roll_mean' + str(w))
                rolling_stats_features.append('roll_std' + str(w))
                model_group_initial['roll_mean' + str(w)] = model_group_initial[variable].shift(1).rolling(w).mean()
                model_group_initial['roll_std' + str(w)] = model_group_initial[variable].shift(1).rolling(w).std()

        # Get trend feature
        model_group_initial['trend'] = model_group_initial.index
        
        # Drop records with NAs
        model_group_initial = model_group_initial.copy().dropna()
        
        # Get just the initial train set
        train_initial = model_group_initial[model_group_initial['set'] == 'TRAIN'].copy()

        # Identify all model features (lags, rolling stats, trend, and the variable itself)
        features = [f'lag{i}' for i in lags] + rolling_stats_features + ['trend'] + [variable]

        # Create and fit the model on initial training data
        iso_forest_model = IsolationForest(
            n_estimators=200,
            contamination=0.01,
            random_state=42
        )
        iso_forest_model.fit(train_initial[features])

        # Score all training points
        train_initial['IsolationForest_score_timeseries'] = iso_forest_model.decision_function(train_initial[features])
        
        # Calculate anomaly threshold
        positive_scores = train_initial[train_initial['IsolationForest_score_timeseries'] > 0]['IsolationForest_score_timeseries']
        if len(positive_scores) > 0:
            anomaly_threshold = min(0, positive_scores.mean() - 3 * positive_scores.std())
        else:
            anomaly_threshold = 0
        
        # Predict anomalies for training points
        train_initial['contamination_anomaly'] = iso_forest_model.predict(train_initial[features])  # -1 = anomaly, 1 = normal
        train_initial['IsolationForest_score_low_timeseries'] = anomaly_threshold
        train_initial['threshold_anomaly'] = np.where(
            train_initial['IsolationForest_score_timeseries'] < anomaly_threshold, -1, 1
        )
        
        # Dual threshold: both contamination and statistical threshold must agree
        train_initial['is_IsolationForest_anomaly_timeseries'] = np.where(
            (train_initial['contamination_anomaly'] == -1) & (train_initial['threshold_anomaly'] == -1), 
            True, 
            False
        )
        
        train_initial['IsolationForest_anomaly_timeseries'] = np.where(train_initial['is_IsolationForest_anomaly_timeseries'] == True, 'low', 'none')
        
        # Select relevant columns
        train_initial_result = train_initial[[
            variable, 
            date_column, 
            'IsolationForest_score_timeseries', 
            'IsolationForest_score_low_timeseries', 
            'is_IsolationForest_anomaly_timeseries',
            'IsolationForest_anomaly_timeseries'
        ]]
        all_results.append(train_initial_result)

        # ===================================================================
        # STEP 2: Walk-forward evaluation for TEST period (one-step-ahead)
        # ===================================================================
        
        for t in list(range(eval_period - 1, -1, -1)):

            try:

                # Boundary between rolling train and rolling forecast region
                cutoff_date = group[date_column].max() - pd.Timedelta(weeks=t)

                # Get train set to determine lags
                model_group = group.copy()
                train = model_group[model_group[date_column] <= cutoff_date].copy()
                lags = get_dynamic_lags(train[variable])

                # Create lag features on the entire model_group DF
                for lag in lags:
                    model_group[f'lag{lag}'] = model_group[variable].shift(lag)

                # Get rolling stats features for the entire model_group DF
                rolling_stats_features = []    
                for w in [int(np.ceil(max(lags)/4)), int(np.ceil(max(lags)/2)), int(max(lags))]:
                    if w >= 3:
                        rolling_stats_features.append('roll_mean' + str(w))
                        rolling_stats_features.append('roll_std' + str(w))
                        model_group['roll_mean' + str(w)] = model_group[variable].shift(1).rolling(w).mean()
                        model_group['roll_std' + str(w)] = model_group[variable].shift(1).rolling(w).std()

                # Get trend feature
                model_group['trend'] = group.index

                # Drop records with NAs
                model_group = model_group.copy().dropna()

                # Split into train and test (train and test now both have all the features
                train = model_group[model_group[date_column] <= cutoff_date].copy()
                test = model_group[model_group[date_column] == cutoff_date].copy()

                # Identify all model features (lags, rolling stats, trend, and the variable itself)
                features = [f'lag{i}' for i in lags] + rolling_stats_features +  ['trend'] + [variable]

                # Create and fit the model
                iso_forest_model = IsolationForest(
                    n_estimators=200,
                    contamination=0.01,
                    random_state=42
                    )
                iso_forest_model.fit(train[features])

                train['IsolationForest_score_timeseries'] = iso_forest_model.decision_function(train[features])
                
                # Calculate anomaly threshold
                positive_scores = train[train['IsolationForest_score_timeseries'] > 0]['IsolationForest_score_timeseries']
                if len(positive_scores) > 0:
                    anomaly_threshold = min(0, positive_scores.mean() - 3 * positive_scores.std())
                else:
                    anomaly_threshold = 0
                
                test['IsolationForest_score_timeseries'] = iso_forest_model.decision_function(test[features])
                test['contamination_anomaly'] = iso_forest_model.predict(test[features])  # -1 = anomaly, 1 = normal
                test['IsolationForest_score_low_timeseries'] = anomaly_threshold
                test['threshold_anomaly'] = np.where(test['IsolationForest_score_timeseries'] < anomaly_threshold, -1, 1)
                
                test['is_IsolationForest_anomaly_timeseries'] = np.where((test['contamination_anomaly'] == -1) & (test['threshold_anomaly'] == -1), True, False)
                test['IsolationForest_anomaly_timeseries'] = np.where(test['is_IsolationForest_anomaly_timeseries'] == True, 'low', 'none')
                
                test = test[[
                    variable, 
                    date_column, 
                    'IsolationForest_score_timeseries', 
                    'IsolationForest_score_low_timeseries', 
                    'is_IsolationForest_anomaly_timeseries',
                    'IsolationForest_anomaly_timeseries'
                ]]
                all_results.append(test)
                
            except Exception as e:
                print(f"Error in iteration {t}: {e}")
                pass
        
        # ===================================================================
        # STEP 3: Combine all results and merge back to original group
        # ===================================================================
        
        try:
            all_results_df = pd.concat(all_results, ignore_index=True)
            
            # Merge back to original group
            group = group.merge(
                all_results_df[[
                    variable,
                    date_column,
                    'IsolationForest_score_timeseries',
                    'IsolationForest_score_low_timeseries',
                    'is_IsolationForest_anomaly_timeseries',
                    'IsolationForest_anomaly_timeseries'
                ]],
                on=[variable, date_column],
                how='left'
            )
            
            # ===================================================================
            # STEP 4: SCALE SCORES [0 to 1]
            # ===================================================================
            
            # Ensure flags are boolean (handle NaNs from left merge if any)
            is_anomaly = (group['is_IsolationForest_anomaly_timeseries'] == True)
            is_normal = (group['is_IsolationForest_anomaly_timeseries'] == False)

            # --- A. Scale Anomalies to range [0.51, 1.0] ---
            # Standard IsoForest: Lower (more negative) raw score = More Anomalous
            # Target: Lower raw score -> Higher scaled score (closer to 1.0)
            if is_anomaly.sum() > 0:
                raw_anom = group.loc[is_anomaly, 'IsolationForest_score_timeseries']
                min_a, max_a = raw_anom.min(), raw_anom.max()
                
                if min_a != max_a:
                    # Inverse Min-Max Scaling
                    # Formula: 0.51 + 0.49 * ((Max - x) / (Max - Min))
                    group.loc[is_anomaly, 'IsolationForest_score_timeseries'] = \
                        0.51 + 0.49 * ((max_a - raw_anom) / (max_a - min_a))
                else:
                    # If there is only 1 anomaly (or all have same score), assign max anomaly score
                    group.loc[is_anomaly, 'IsolationForest_score_timeseries'] = 1.0

            # --- B. Scale Normals to range [0.0, 0.49] ---
            # Standard IsoForest: Higher (positive) raw score = More Normal
            # Target: Higher raw score -> Lower scaled score (closer to 0.0)
            if is_normal.sum() > 0:
                raw_norm = group.loc[is_normal, 'IsolationForest_score_timeseries']
                min_n, max_n = raw_norm.min(), raw_norm.max()
                
                if min_n != max_n:
                    # Inverse Min-Max Scaling
                    # Formula: 0.49 * ((Max - x) / (Max - Min))
                    group.loc[is_normal, 'IsolationForest_score_timeseries'] = \
                        0.49 * ((max_n - raw_norm) / (max_n - min_n))
                else:
                    # If there is only 1 normal point, assign perfectly normal score
                    group.loc[is_normal, 'IsolationForest_score_timeseries'] = 0.0

        except Exception as e:
            print(f"Error in concatenating/scaling results: {e}")
            group["IsolationForest_score_timeseries"] = np.nan
            group["IsolationForest_score_low_timeseries"] = np.nan
            group["is_IsolationForest_anomaly_timeseries"] = np.nan
            group['IsolationForest_anomaly_timeseries'] = np.nan
    
    except Exception as e:
        # Fallback error handling
        try:
            group_id_cols = group.select_dtypes(include=['object', 'string']).columns.tolist()
            group_id = " ".join(group[group_id_cols].reset_index(drop=True).iloc[0].astype(str).to_list())
        except:
            group_id = "Unknown Group ID"
        print(f'Isolation Forest Anomaly Detection failed for {group_id}. Error: {e}')
        group["IsolationForest_score_timeseries"] = np.nan
        group["IsolationForest_score_low_timeseries"] = np.nan
        group["is_IsolationForest_anomaly_timeseries"] = np.nan
        group['IsolationForest_anomaly_timeseries'] = np.nan
    
    group = pd.concat([group_orig, group.drop(columns=[variable, date_column])], axis=1)
    
    return group