import pandas as pd
import numpy as np
from datetime import date
import joblib
from joblib import Parallel, delayed
import contextlib
from tqdm.auto import tqdm
from statsmodels.tsa.seasonal import STL
from .helpers.percentile import detect_outliers_percentile
from .helpers.STD import detect_outliers_sd
from .helpers.MAD import detect_outliers_mad
from .helpers.IQR import detect_outliers_iqr
from .helpers.iso_forest_general import detect_outliers_isf_general
from .helpers.ewma import ewma_with_anomalies_rolling_group
from .helpers.fb_prophet import detect_time_series_anomalies_fb_walkforward
from .helpers.iso_forest_timeseries import detect_time_series_anomalies_isoforest
from .helpers.DB_scan import detect_time_series_anomalies_dbscan
from .helpers.Preprocessing import (create_full_calendar_and_interpolate, 
                                    print_anomaly_stats,
                                    top_anomaly_stats,
                                    calculate_ensemble_scores,
                                    min_records_extraction)

from .helpers.evaluation_plots import (summary_pie_plot,
                                       anomaly_stacked_bar_plot, 
                                       anomaly_variable_score_plot, 
                                       anomaly_overview_plot,
                                       model_performance_chart)

from IPython.display import display, Markdown

# --- Progress Bar Context Manager ---
@contextlib.contextmanager
def tqdm_joblib(tqdm_object):
    """Context manager to patch joblib to report into tqdm progress bar"""
    class TqdmBatchCompletionCallback(joblib.parallel.BatchCompletionCallBack):
        def __call__(self, *args, **kwargs):
            tqdm_object.update(self.batch_size)
            return super().__call__(*args, **kwargs)

    old_callback = joblib.parallel.BatchCompletionCallBack
    joblib.parallel.BatchCompletionCallBack = TqdmBatchCompletionCallback
    try:
        yield tqdm_object
    finally:
        joblib.parallel.BatchCompletionCallBack = old_callback
        # We DO NOT close the pbar here so it can be reused for the next model
        
def get_stl_period(freq):
    """Determines the seasonal period for STL based on frequency string."""
    if not isinstance(freq, str): return 7
    f = freq.upper()
    if f.startswith('D'): return 7
    if f.startswith('W'): return 4
    if f.startswith('M'): return 12
    return 7

def process_stats_group(group, variable, date_column, eval_period, freq, percentile_threshold,
                        SD_threshold, mad_threshold, mad_scale_factor,IQR_threshold):
    """
    Parallel worker: Decomposes group once, runs 4 models on residuals, 
    and reconstructs original scale.
    """
    group = group.copy().sort_values(date_column)
    
    # 1. STL Decomposition
    period = get_stl_period(freq)
    stl = STL(group[variable], period=period, robust=True).fit()
    group['residual'] = stl.resid
    group['base_signal'] = stl.trend + stl.seasonal

    # 2. Run Models on Residuals (Note: Ensure internal helpers don't force max(0) yet)
    res_p = detect_outliers_percentile(group, 'residual', date_column, eval_period, percentile_threshold)
    res_s = detect_outliers_sd(group, 'residual', date_column, eval_period,SD_threshold)
    res_m = detect_outliers_mad(group, 'residual', date_column, mad_threshold, mad_scale_factor, eval_period)
    res_i = detect_outliers_iqr(group, 'residual', date_column, eval_period, IQR_threshold)

    # 3. Scale Reconstruction
    for res_df in [res_p, res_s,res_m,res_i]:
        # Add the decomposition columns to each result df so they are available for merging
        res_df['base_signal'] = group['base_signal'].values
        res_df['residual'] = group['residual'].values
        cols = [c for c in res_df.columns if any(w in c for w in ['low', 'high', 'Mean', 'Median', 'Q1', 'Q3'])]
        # Shift back to original scale
        res_df[cols] = res_df[cols].add(res_df['base_signal'], axis=0)
        # Apply zero-floor safety only on final reconstructed bounds
        low_cols = [c for c in cols if 'low' in c]
        res_df[low_cols] = res_df[low_cols].clip(lower=0)
        high_cols = [c for c in cols if 'high' in c]
        res_df[high_cols] = res_df[high_cols].clip(lower=0)

    return (res_p, res_s,res_m, res_i)


def process_group(model, name, group, unique_ids, variable,
                  date_column, alpha, sigma, eval_period, prophet_CI, contamination, random_state):

    if model == "ISF_general":
        return detect_outliers_isf_general(group, variable, contamination, random_state, eval_period)

    if model == "EWMA":
        return ewma_with_anomalies_rolling_group(
            group, unique_ids, variable, date_column, alpha, sigma, eval_period
        )

    if model == "FB":
        return detect_time_series_anomalies_fb_walkforward(
            group, variable, date_column, eval_period, prophet_CI
        )
    
    if model == 'ISF_timeseries':
        return detect_time_series_anomalies_isoforest(
            group, variable, date_column, eval_period
        )
    
    if model == 'DBSCAN':
        return detect_time_series_anomalies_dbscan(
            group, variable, date_column, eval_period
        )

MODEL_CONFIG = {
    'DECOMPOSITION': ['base_signal', 'residual'],
    'PERCENTILE': ['Percentile_low', 'Percentile_high','Percentile_anomaly', 'is_Percentile_anomaly'],
    'SD': ['Mean', 'SD', 'SD_low', 'SD_high', 'SD_anomaly','is_SD_anomaly'],
    'MAD': ['Median', 'MAD', 'MAD_low', 'MAD_high', 'MAD_anomaly','is_MAD_anomaly'],
    'IQR': ['Q1', 'Q3', 'IQR', 'IQR_low', 'IQR_high', 'IQR_anomaly', 'is_IQR_anomaly'],
    'EWMA': ['EWMA_forecast', 'EWMA_high', 'EWMA_low', 'EWMA_residual', 'EWMA_anomaly', 'is_EWMA_anomaly'],
    'FB': ['FB_forecast', 'FB_low', 'FB_high', 'FB_residual', 'FB_anomaly', 'is_FB_anomaly'],
    'ISF': ['IsolationForest_score', 'IsolationForest_score_low', 'is_IsolationForest_anomaly', 'IsolationForest_anomaly'],
    'DBSCAN': ['DBSCAN_score', 'DBSCAN_score_high', 'is_DBSCAN_anomaly', 'DBSCAN_anomaly']
}

def run_pipeline(master_data, unique_ids, variable, date_column, freq,
                 eval_period, max_records, imputation_method,
                 percentile_threshold, SD_threshold, mad_threshold, 
                 mad_scale_factor, IQR_threshold, alpha, sigma, 
                 prophet_CI, contamination, random_state, is_assigned,
                 enabled_models):
    
    # 1. Initialization & Preprocessing
    enabled_models = [m.upper() for m in enabled_models]
    if max_records is not None:
        max_records = max_records + eval_period

    final_data, success_report = create_full_calendar_and_interpolate(
        master_data, unique_ids, variable, date_column, freq, max_records, imputation_method
    )
    
    # Filter groups with at least 5 valid data points
    groups = [(n, g) for n, g in final_data.groupby(unique_ids) if g[variable].notna().sum() >= 5]
    
    # Calculate Total Operations for Progress Bar
    stats_keys = {'PERCENTILE', 'SD', 'MAD', 'IQR'}
    run_stats = any(m in enabled_models for m in stats_keys)
    other_ts_models = [m for m in ['EWMA', 'FB', 'DBSCAN'] if m in enabled_models]
    isf_op = 2 if 'ISF' in enabled_models else 0 # ISF runs two parallel passes
    
    total_operations = len(groups) * ((1 if run_stats else 0) + len(other_ts_models) + isf_op)
    results_map = {}

    # 2. Execution Loop
    with tqdm(total=total_operations, desc="Running Anomaly Detection Pipeline") as pbar:
        
        # --- BATCH 1: STATS ON RESIDUALS (STL) ---
        if run_stats:
            with tqdm_joblib(pbar):
                stats_raw = Parallel(n_jobs=-1)(delayed(process_stats_group)(
                    g, variable, date_column, eval_period, freq, percentile_threshold,
                    SD_threshold, mad_threshold, mad_scale_factor, IQR_threshold
                ) for n, g in groups)
            
            # Unpack results into map if they were requested
            if 'PERCENTILE' in enabled_models: results_map['PERCENTILE'] = pd.concat([r[0] for r in stats_raw])
            if 'SD' in enabled_models: results_map['SD'] = pd.concat([r[1] for r in stats_raw])
            if 'MAD' in enabled_models: results_map['MAD'] = pd.concat([r[2] for r in stats_raw])
            if 'IQR' in enabled_models: results_map['IQR'] = pd.concat([r[3] for r in stats_raw])
            
            first_available = next(iter(stats_keys.intersection(results_map.keys())))
            results_map['DECOMPOSITION'] = results_map[first_available][unique_ids + [date_column, 'base_signal', 'residual']]
            
            
        # --- BATCH 2: STANDARD TIME SERIES MODELS ---
        for model_key in ['EWMA', 'FB', 'DBSCAN']:
            if model_key in enabled_models:
                with tqdm_joblib(pbar):
                    raw = Parallel(n_jobs=-1)(delayed(process_group)(
                        model_key, n, g, unique_ids, variable, date_column, 
                        alpha, sigma, eval_period, prophet_CI, contamination, random_state
                    ) for n, g in groups)
                    results_map[model_key] = pd.concat(raw)

        # --- BATCH 3: ISOLATION FOREST (Ensemble Logic) ---
        if 'ISF' in enabled_models:
            with tqdm_joblib(pbar):
                res_gen = Parallel(n_jobs=-1)(delayed(process_group)(
                    'ISF_general', n, g, unique_ids, variable, date_column, 
                    alpha, sigma, eval_period, prophet_CI, contamination, random_state
                ) for n, g in groups)
                res_ts = Parallel(n_jobs=-1)(delayed(process_group)(
                    'ISF_timeseries', n, g, unique_ids, variable, date_column, 
                    alpha, sigma, eval_period, prophet_CI, contamination, random_state
                ) for n, g in groups)
            
            isf_gen_df = pd.concat(res_gen)
            isf_ts_df = pd.concat(res_ts)
            
            # Combine General and Timeseries results based on Train/Test set
            isf_combined = isf_gen_df.merge(
                isf_ts_df[unique_ids + [date_column, 'IsolationForest_score_timeseries', 
                                        'IsolationForest_score_low_timeseries', 
                                        'is_IsolationForest_anomaly_timeseries', 
                                        'IsolationForest_anomaly_timeseries']], 
                on=unique_ids + [date_column], how='inner'
            )
            
           
            for suffix in ['', '_low']:
                # For boolean/categorical 'anomaly' and numeric scores
                main_col = f'IsolationForest_score{suffix}'
                ts_col = f'IsolationForest_score{suffix}_timeseries'
                gen_col = f'IsolationForest_score{suffix}_general'
                
                # Coalesce: Use TS if not null, else use General
                isf_combined[main_col] = isf_combined[ts_col].fillna(isf_combined[gen_col])

            # Specifically handle the 'is_Anomaly' boolean flag
            isf_combined['IsolationForest_anomaly'] = (
            isf_combined['IsolationForest_anomaly_timeseries']
                .fillna(isf_combined['IsolationForest_anomaly_general'])
            )
            isf_combined['is_IsolationForest_anomaly'] = (
                isf_combined['is_IsolationForest_anomaly_timeseries']
                .fillna(isf_combined['is_IsolationForest_anomaly_general'])
            )
                                                                              
            #print(isf_combined.head(50))
            #print(isf_combined.tail(50))                                                                 
                                                                              
            results_map['ISF'] = isf_combined[unique_ids + [date_column] + MODEL_CONFIG['ISF']]

    # 3. Dynamic Merging
    anomaly = final_data.copy()
    

        # Add DECOMPOSITION to the keys we iterate over
    models_to_merge = (['DECOMPOSITION'] if run_stats else [])+enabled_models 

    for model_key in models_to_merge:
        if model_key in results_map:
            cols_to_extract = unique_ids + [date_column] + MODEL_CONFIG[model_key]
            # Ensure we only pick columns that actually exist in the result dataframe
            valid_cols = [c for c in cols_to_extract if c in results_map[model_key].columns]
            
            anomaly = anomaly.merge(
                results_map[model_key][valid_cols], 
                on=unique_ids + [date_column], 
                how='inner'
            )
    
    
    # --- Final Ensemble & Scoring ---
    anomaly_final = calculate_ensemble_scores(anomaly, variable, unique_ids, date_column, enabled_models)
        
    # # --- Cleansing Mechanism ---
    # anomaly_final = cleanse_anomalies_directional(
    #         anomaly_final, 
    #         variable, 
    #         enabled_models
    #     )
    
    globals()['anomaly_df'] = anomaly_final
    #print(anomaly_final.head())
    #print(f"Successfully processed {len(success_report)} groups.")
    
    # --- UPDATE SUCCESS REPORT WITH ANOMALY STATS ---
    #  Calculate stats per group from the final results
    group_stats = anomaly_final.groupby(unique_ids)['is_Anomaly'].agg(
        anomalies='sum',
        # Multiply mean by 100 and round to 1 decimal
        anomaly_rate=lambda x: round(x.mean() * 100, 1)
    ).reset_index()

    # 2. Merge these stats into the success_report
    evaluation_report = success_report.merge(group_stats, on=unique_ids, how='left')
    
    
    print_anomaly_stats(master_data, anomaly_final, success_report, unique_ids, imputation_method, is_assigned)
    
    print("")

    # Plot summary charts
    # ------------------------------
    
    # Get data for pie chart
    pie_chart_df = anomaly_final['is_Anomaly'].value_counts().reset_index()
    pie_chart_df['is_Anomaly'] = np.where(pie_chart_df['is_Anomaly'] == True, 'Anomalous Records', 'Normal Records')
    pie_chart_df = pie_chart_df.rename(columns={'is_Anomaly': 'Records'})
   
    summary_pie_plot(pie_chart_df, title=f"Anomaly Detection Summary for {len(anomaly_final[unique_ids].drop_duplicates())} Unique IDs")
    anomaly_stacked_bar_plot(anomaly_final, unique_ids, variable, date_column, secondary_line=variable)
    anomaly_variable_score_plot(anomaly_final, unique_ids, variable, date_column)
    

    model_performance_chart(anomaly_final)

    # print_model_performance_chart(anomaly_final)
    
    top_5_anomaly_groups = anomaly_final.groupby(unique_ids)['is_Anomaly'].agg(['mean', 'sum', 'count']).reset_index()\
        .sort_values('mean', ascending=False).reset_index(drop=True).head(5)

    eval_plots_msg = f"""
---
### Overall Evaluation Plots of the {len(top_5_anomaly_groups)} Unique IDs with the Highest Anomaly Rates
---
"""
    top_anomaly_stats(anomaly_final, unique_ids)
    
    display(Markdown(eval_plots_msg))
    
    uid_example = ""
    id_values = []
    for i, col in enumerate(unique_ids):
        uid_val = top_5_anomaly_groups[unique_ids].head(1).iloc[0 , i]
        id_values.append(uid_val)
        uid_str = "`" + col + "` = '" + str(uid_val) + "'"
        if i == 0:
            uid_example = uid_example + uid_str
        else:
            uid_example = uid_example + " and " + uid_str

    code_content = f"""id_values = {id_values}
mask = final_results[unique_ids].eq(id_values).all(axis=1)
group_df = final_results[mask]

evaluation_info(
    group_df,
    unique_ids,
    variable,
    date_column,
    eval_period)"""

    code_block_html = f"""
<div style="background-color: ghostwhite; padding: 10px; border: 1px solid #e0e0e0; border-radius: 5px; width: fit-content;">
<pre style="background-color: ghostwhite; margin: 0; font-family: monospace; white-space: pre;">{code_content}</pre></div>
"""

    help_info = f"""
---
You can generate detailed plots of individual anomaly detection models per unique_id.<br>
If your anomaly detection table is `final_results` (which is the default table name),<br>
and you want to investigate the unique_id where {uid_example},<br>
then copy and paste this block of code into another cell and run it:<br><br>

{code_block_html}
---
"""

    group_nbr = 1
    for group_key, group in top_5_anomaly_groups.groupby(unique_ids, sort=False):
        anomaly_rate = group['mean'].values[0]
        group_df = anomaly_final.merge(group[unique_ids], on=unique_ids, how='inner')
        group_id = group_df[unique_ids].drop_duplicates().astype(str).apply(lambda x: ' -- '.join(x), axis=1).values[0]
        group_msg = f"""#### {group_id} Anomaly Rate: {anomaly_rate:.1%}"""
        display(Markdown(group_msg))
        anomaly_overview_plot(group_df, unique_ids, variable, date_column, eval_period, show_anomaly_scores_on_main_plot=False)
        group_nbr += 1
    
    display(Markdown(help_info))  
    anomaly_final = anomaly_final.drop(["normal_val","anomaly_val"], axis =1)
    
    return anomaly_final, evaluation_report
