import pandas as pd
import numpy as np
import re
from datetime import datetime
from statsmodels.tsa.stattools import acf


def classify(val,lower,upper):
    if val < lower:
        return 'low'
    elif val > upper:
        return 'high'
    else:
        return 'none'


def create_full_calendar_and_interpolate(
    master_data, unique_ids, variable, date_column, freq, max_records, imputation_method
):
    # --- FIX: Cast variable to float to handle decimal results from mean/median ---
    master_data[variable] = master_data[variable].astype(float)
    # 1. Ensure datetime format
    master_data[date_column] = pd.to_datetime(master_data[date_column])
    
    # 2. Optimized Calendar Generation
    def get_group_range(group):
        start, end = group[date_column].min(), group[date_column].max()
        dr = pd.date_range(start=start, end=end, freq=freq)
        if max_records and len(dr) > max_records:
            dr = dr[-max_records:]
        return pd.DataFrame({date_column: dr})

    # Use include_groups=False to handle the FutureWarning
    # and reset_index() to ensure unique_ids are actual columns
    full_calendar = (
        master_data.groupby(unique_ids, group_keys=True)
        .apply(get_group_range, include_groups=False)
        .reset_index()
    )
    
    # Remove the 'level_X' column that apply sometimes creates
    if 'level_len' in full_calendar.columns or f'level_{len(unique_ids)}' in full_calendar.columns:
         full_calendar = full_calendar.drop(columns=[c for c in full_calendar.columns if 'level_' in str(c)])

    # 3. Merge raw data into the full calendar
    merged = pd.merge(full_calendar, master_data, on=unique_ids + [date_column], how="left")
    merged["is_missing_record"] = merged[variable].isna()

    # 4. Vectorized Imputation
    if imputation_method == 'zero':
        merged[variable] = merged[variable].fillna(0)
        
    elif imputation_method in ['mean', 'median']:
        fill_values = merged.groupby(unique_ids)[variable].transform(imputation_method)
        merged[variable] = merged[variable].fillna(fill_values)
        
    elif imputation_method == 'mode':
        def fill_mode(s):
            m = s.mode()
            return s.fillna(m[0] if not m.empty else np.nan)
        merged[variable] = merged.groupby(unique_ids)[variable].transform(fill_mode)
        
    else:
        # Standard interpolation (linear, time, etc.)
        merged[variable] = merged.groupby(unique_ids)[variable].transform(
            lambda x: x.interpolate(method=imputation_method, limit_direction="both")
        )

    # 5. Success Metrics
    # Note: Using .groupby(unique_ids) here is safe because merged definitely has these columns now
    success_report = merged.groupby(unique_ids).agg(
        initial_records=("is_missing_record", lambda x: (~x).sum()),
        interpolated_records=("is_missing_record", "sum"),
        final_records=("is_missing_record", "size")
    ).reset_index()
    
    success_report["imputation_pct"] = (
        (success_report["interpolated_records"] / success_report["final_records"]) * 100
    ).round(2)

    return merged, success_report

import re

class Colors:

    BLUE = '\033[94m'

    RED = '\033[91m'

    END = '\033[0m'

    BOLD = '\033[1m'

    BOLD_ITALIC = "\033[1;3m"

    RESET = "\033[0m"
    
    BG_WHITE = '\033[47m'
    
    BLACK = '\033[30m'
    
    GREEN = '\033[92m'
    
    YELLOW = '\033[93m'
    
    CYAN = '\033[36m'
    
    BOX = '\033[48;5;243m\033[38;5;255m'
    
    
#  PRINTING UTILITIES ---

def print_header(title, color=Colors.BLUE):
    """
    Prints a centered, bold header. 
    Borders are always Blue, but the Title color is configurable (default Red).
    """
    width = 70
    border = "=" * width
    
    # Border is blue
    print(f"\n{Colors.BLUE}{Colors.BOLD}{border}")
    
    # Title uses the passed 'color' parameter (e.g., Red)
    print(f"{color}{title:^{width}}")
    
    # Border is blue
    print(f"{Colors.BLUE}{border}{Colors.END}")


def print_anomaly_stats(df, final_results, success_report, unique_ids, imputation_method, is_assigned):
    # --- 1. SUCCESS NOTIFICATION & AUTO-SAVE WARNING ---
    total_groups = len(success_report)
    
    # Calculate unique groups actually evaluated in final_results
    evaluated_groups = final_results.groupby(unique_ids).ngroups
    
    # Success Message
    print("-" * 70)
    print(f"{Colors.BOLD}{Colors.GREEN}✔ Anomaly detection pipeline successfully ran for {evaluated_groups:,} unique_ids.{Colors.END}")
    
    # Logic for auto-saving notification
    if not is_assigned:
        print(f"{Colors.YELLOW}📀 NOTE: Since variables are not assigned from the user,{Colors.END}")
        print(f"{Colors.YELLOW}the following dataframes were generated as part of the output:{Colors.END}")
        print(f" {Colors.BOX} final_results {Colors.END}     : dataframe containing the evaluation of each record")
        print(f" {Colors.BOX} evaluation_report {Colors.END} : contains aggregate metrics for each unique_id")
    
    print("-" * 70)

    # --- 2. DATA PREPARATION & CALCULATIONS ---
    Evaluated_records = len(final_results)
    total_anomalies = final_results['is_Anomaly'].fillna(False).astype(bool).sum()
    anomaly_rate = (total_anomalies / Evaluated_records) * 100 if Evaluated_records > 0 else 0
    
    # Helper formatters
    def print_stats_rows(stats):
        label_width = 45
        value_width = 22
        for label, val in stats:
            # Strip ANSI colors for padding calculation
            visible_val = re.sub(r'\033\[[0-9;]*m', '', str(val))
            padding = " " * (value_width - len(visible_val))
            print(f"{label:<{label_width}} : {padding}{val}")

    # --- 3. EXECUTION SUMMARY ---
    print_header("ANOMALY DETECTION EXECUTIVE SUMMARY", Colors.GREEN)

    exec_stats = [
        ["Evaluated Unique IDs", f"{evaluated_groups:,}"],
        ["Evaluated Records", f"{Evaluated_records:,}"],
        ["Detected Anomalies", f"{Colors.RED}{total_anomalies:,}{Colors.END}"],
        ["Anomaly Rate", f"{Colors.RED}{anomaly_rate:.1f}%{Colors.END}"]
    ]

    print_stats_rows(exec_stats)
    
    # Using Cyan for the footer line as discussed earlier
    print(f"{Colors.BLUE}{'='*70}{Colors.END}\n")


def top_anomaly_stats(final_results, unique_ids):
    # Check if data exists to avoid printing empty headers
    if final_results.empty:
        print(f"{Colors.YELLOW}No data available for anomaly stats.{Colors.END}")
        return

    title = f"TOP 5 unique_ids BY ANOMALY RATE ({' | '.join(unique_ids)}):"
    print_header(title) # Using your updated print_header

    group_stats = final_results.groupby(unique_ids)['is_Anomaly'].agg(['mean', 'sum']).sort_values(by='mean', ascending=False).head(5)

    for label, row in group_stats.iterrows():
        # Robust handling for single vs multi-column labels
        if isinstance(label, tuple):
            group_label = " | ".join(map(str, label))
        else:
            group_label = str(label)

        # Formatting values
        # Note: Added 30-char padding to group_label for alignment
        rate_val = row['mean'] * 100
        rate_str = f"{Colors.RED}{rate_val:.1f}%{Colors.END}"
        count_str = f"({Colors.RED}{int(row['sum']):}{Colors.END} anomalies)"
        
        print(f" - {group_label:} : {rate_str} {count_str}")

    print(f"{Colors.BLUE}{'='*70}{Colors.END}\n")
    

def calculate_ensemble_scores(df, variable, unique_ids, date_column, enabled_models):
    """
    Calculates the normalized consensus score across all anomaly models.
    """
    
    # Identify all columns that are model flags (is_..._anomaly)
    anomaly_flags = [col for col in df.columns if col.startswith('is_') and col.endswith('_anomaly') and col != 'is_Anomaly']
    
    # 1. Total Votes (Count of True)
    df['Anomaly_Votes'] = df[anomaly_flags].sum(axis=1).astype(int)

    # 2. Total Models active for that row (Count of non-NaN values)
    df['Vote_Cnt'] = df[anomaly_flags].notna().sum(axis=1).astype(int)
    
    # 3. Anomaly Votes Score Display (x out of N)
    df['Anomaly_Votes_Display'] = df['Anomaly_Votes'].astype(int).astype(str) + " out of " + df['Vote_Cnt'].astype(int).astype(str)
    
    # 5. Final Boolean Consensus (e.g., majority rule)
    df['is_Anomaly'] = df['Anomaly_Votes'] / df['Vote_Cnt'] >= 0.5
    
    def cleanse_anomalies_directional(df, variable, enabled_models):
        df = df.copy()

        model_map = {
            'PERCENTILE': ('is_Percentile_anomaly', 'Percentile_anomaly', 'Percentile_low', 'Percentile_high'),
            'SD':         ('is_SD_anomaly',         'SD_anomaly',         'SD_low',         'SD_high'),
            'MAD':        ('is_MAD_anomaly',        'MAD_anomaly',        'MAD_low',        'MAD_high'),
            'IQR':        ('is_IQR_anomaly',        'IQR_anomaly',        'IQR_low',        'IQR_high'),
            'EWMA':       ('is_EWMA_anomaly',       'EWMA_anomaly',       'EWMA_low',       'EWMA_high'),
            'FB':         ('is_FB_anomaly',         'FB_anomaly',         'FB_low',         'FB_high')
        }

        # --- 1. Determine Consensus & Conflict ---
        dir_cols = []
        for m in enabled_models:
            if m in model_map and model_map[m][1] in df.columns:
                tmp_col = f"tmp_dir_{m}"
                df[tmp_col] = df[model_map[m][1]].map({'high': 1, 'low': -1}).fillna(0)
                dir_cols.append(tmp_col)

        # Logic to find conflicts:
        # Does the row contain both a 1 and a -1?
        has_high = (df[dir_cols] == 1).any(axis=1)
        has_low = (df[dir_cols] == -1).any(axis=1)
        df['is_conflict'] = has_high & has_low

        # OVERWRITE: If there is a conflict, it's not a reliable anomaly
        df['is_Anomaly'] = np.where(df['is_conflict'], False, df['is_Anomaly'])

        # Determine direction for non-conflicting rows
        df['net_direction'] = df[dir_cols].sum(axis=1)
        df['consensus_type'] = np.where(df['net_direction'] > 0, 'high', 
                                        np.where(df['net_direction'] < 0, 'low', 'none'))

        # --- 2. Calculate Lenient Bounds ---
        high_candidates, low_candidates = [], []
        for m in enabled_models:
            if m in model_map:
                _, label_col, low_col, high_col = model_map[m]
                if all(c in df.columns for c in [label_col, low_col, high_col]):
                    df[f'm_hi_{m}'] = np.where((df['consensus_type'] == 'high') & (df[label_col] == 'high'), df[high_col], np.nan)
                    df[f'm_lo_{m}'] = np.where((df['consensus_type'] == 'low') & (df[label_col] == 'low'), df[low_col], np.nan)
                    high_candidates.append(f'm_hi_{m}')
                    low_candidates.append(f'm_lo_{m}')

        # --- 3. Project and Cleanse ---
        raw_lenient_high = df[high_candidates].max(axis=1).fillna(df['base_signal'])
        raw_lenient_low = df[low_candidates].min(axis=1).fillna(df['base_signal'])


        df['max_threshold'] = np.where(df['is_Anomaly'], raw_lenient_high, df[variable])
        df['min_threshold'] = np.where(df['is_Anomaly'], raw_lenient_low, df[variable])

        df[f'{variable}_cleansed'] = df[variable].clip(lower=df['min_threshold'], upper=df['max_threshold'])

        # Cleanup
        cols_to_drop = dir_cols + high_candidates + low_candidates + ['net_direction']
        df = df.drop(columns=[c for c in cols_to_drop if c in df.columns])

        return df
    
    df = cleanse_anomalies_directional(
            df,
            variable, 
            enabled_models
        )
        
    groups = list(df.groupby(unique_ids))
    
    
    # 6. Scale all the model scores to be between -1 and 1
    def get_scaled_score(group, score_scaled_col):
        group[score_scaled_col] = 100 * np.where(group[score_scaled_col] < 0,
                                                 -group[score_scaled_col]/group[score_scaled_col].min(),
                                                 group[score_scaled_col]/group[score_scaled_col].max())
        group[score_scaled_col] = np.where(group[score_scaled_col] < 0, group[score_scaled_col] - 1, group[score_scaled_col] + 1)
        group[score_scaled_col] = np.where(group[score_scaled_col] < 0,
                                           -np.log(-group[score_scaled_col]),
                                           np.log(group[score_scaled_col]))
        if len(group[group[score_scaled_col] > 0]) >= 1:
            is_min = group[group[score_scaled_col] > 0][score_scaled_col].min()
            is_max = group[group[score_scaled_col] > 0][score_scaled_col].max()
        else:
            is_min = 1
            is_max = 2
        if len(group[group[score_scaled_col] < 0]) >= 1:
            not_min = group[group[score_scaled_col] < 0][score_scaled_col].min()
            not_max = group[group[score_scaled_col] < 0][score_scaled_col].max()
        else:
            not_min = 1
            not_max = 2
        group[score_scaled_col] = 100 * np.where(group[score_scaled_col] > 0,
                                              0.51 + 0.49 * (group[score_scaled_col] - is_min)/(is_max - is_min),
                                              0.49 * (group[score_scaled_col] - not_min)/(not_max - not_min))
        return group
    
    score_scaled_cols = ['Percentile_score_scaled', 'SD_score_scaled', 'MAD_score_scaled', 'IQR_score_scaled',
                         'EWMA_score_scaled', 'FB_score_scaled', 'IsolationForest_score_scaled', 'DBSCAN_score_scaled']

    df_scores = []
    for group in groups:
        group = group[1]
        try:
            group['Percentile_score_scaled'] = np.where(group['is_Percentile_anomaly'].isna()==False,
                                                     abs(group[variable] - (group['Percentile_high'] + group['Percentile_low'])/2)/((group['Percentile_high'] - group['Percentile_low'])/2) - 1,
                                                     np.nan)
            group = get_scaled_score(group, 'Percentile_score_scaled')
        except:
            group['Percentile_score_scaled'] = np.nan

        try:
            group['SD_score_scaled'] = np.where(group['is_SD_anomaly'].isna()==False,
                                             abs(group[variable] - (group['SD2_high'] + group['SD2_low'])/2)/((group['SD2_high'] - group['SD2_low'])/2) - 1,
                                             np.nan)
            group = get_scaled_score(group, 'SD_score_scaled')
        except:
            group['SD_score_scaled'] = np.nan

        try:
            group['MAD_score_scaled'] = np.where(group['is_MAD_anomaly'].isna()==False,
                                              abs(group[variable] - (group['MAD_high'] + group['MAD_low'])/2)/((group['MAD_high'] - group['MAD_low'])/2) - 1,
                                              np.nan)
            group = get_scaled_score(group, 'MAD_score_scaled')
        except:
            group['MAD_score_scaled'] = np.nan

        try:
            group['IQR_score_scaled'] = np.where(group['is_IQR_anomaly'].isna()==False,
                                              abs(group[variable] - (group['IQR_high'] + group['IQR_low'])/2)/((group['IQR_high'] - group['IQR_low'])/2) - 1,
                                              np.nan)
            group = get_scaled_score(group, 'IQR_score_scaled')
        except:
            group['IQR_score_scaled'] = np.nan

        try:
            group['EWMA_score_scaled'] = np.where(group['is_EWMA_anomaly'].isna()==False,
                                               abs(group[variable] - (group['EWMA_high'] + group['EWMA_low'])/2)/((group['EWMA_high'] - group['EWMA_low'])/2) - 1,
                                               np.nan)
            group = get_scaled_score(group, 'EWMA_score_scaled')
        except:
            group['EWMA_score_scaled'] = np.nan

        try:
            group['FB_score_scaled'] = np.where(group['is_FB_anomaly'].isna()==False,
                                             abs(group[variable] - (group['FB_high'] + group['FB_low'])/2)/((group['FB_high'] - group['FB_low'])/2) - 1,
                                             np.nan)
            group = get_scaled_score(group, 'FB_score_scaled')
        except:
            group['FB_score_scaled'] = np.nan

        try:
            group['IsolationForest_score_scaled'] = np.where(group['is_IsolationForest_anomaly'] == True,
                                                          group['IsolationForest_score'] - group[group['is_IsolationForest_anomaly'] == True]['IsolationForest_score'].min(),
                                                          group['IsolationForest_score'] - group[group['is_IsolationForest_anomaly'] == False]['IsolationForest_score'].min())

            group['IsolationForest_score_scaled'] = 100 * np.where(group['is_IsolationForest_anomaly'] == True,
                                                                0.49*group['IsolationForest_score_scaled']/group[group['is_IsolationForest_anomaly'] == True]['IsolationForest_score_scaled'].max() + 0.51,
                                                                0.49*group['IsolationForest_score_scaled']/group[group['is_IsolationForest_anomaly'] == False]['IsolationForest_score_scaled'].max())
        except:
            group['IsolationForest_score_scaled'] = np.nan

        try:
            group['DBSCAN_score_scaled'] = np.where(group['is_DBSCAN_anomaly'] == True,
                                                 group['DBSCAN_score'] - group[group['is_DBSCAN_anomaly'] == True]['DBSCAN_score'].min(),
                                                 group['DBSCAN_score'] - group[group['is_DBSCAN_anomaly'] == False]['DBSCAN_score'].min())
            group['DBSCAN_score_scaled'] = 100 * np.where(group['is_DBSCAN_anomaly'] == True,
                                                       0.51 + 0.49*group['DBSCAN_score_scaled']/group[group['is_DBSCAN_anomaly'] == True]['DBSCAN_score_scaled'].max(),
                                                       0.49*group['DBSCAN_score_scaled']/group[group['is_DBSCAN_anomaly'] == False]['DBSCAN_score_scaled'].max())
        except:
            group['DBSCAN_score_scaled'] = np.nan
        
        try:
            group['Anomaly_Score'] = group[score_scaled_cols].mean(axis=1)
            
            # Rescale all non anomalies between 0 and 0.5 and anomalies between 0.5 and 1.0
            if len(group[group['is_Anomaly'] == True]) >= 1:
                is_min = group[group['is_Anomaly'] == True]['Anomaly_Score'].min()
                is_max = group[group['is_Anomaly'] == True]['Anomaly_Score'].max()
            else:
                is_min = 0
                is_max = 1
            if len(group[group['is_Anomaly'] == False]) >= 1:
                not_min = group[group['is_Anomaly'] == False]['Anomaly_Score'].min()
                not_max = group[group['is_Anomaly'] == False]['Anomaly_Score'].max()
            else:
                not_min = 0
                not_max = 1

            group['Anomaly_Score'] = (100 * np.where(group['is_Anomaly'] == True,
                                                     0.51 + 0.49 * (group['Anomaly_Score'] - is_min)/(is_max - is_min),
                                                     0.49 * (group['Anomaly_Score'] - not_min)/(not_max - not_min))).astype(int)
            
        except:
            group['Anomaly_Score'] = np.nan

        df_scores.append(group[unique_ids + [date_column] + score_scaled_cols + ['Anomaly_Score']])
    
    df_scores = pd.concat(df_scores)
    df = df.merge(df_scores, on=unique_ids + [date_column], how='left')
    
    
    # 7. Reposition is_Anomaly column to the end
    df['is_Anomaly'] = df.pop('is_Anomaly')

    return df


def min_records_extraction(freq,eval_period):
    freq_upper = freq.upper()
    
    if freq_upper.startswith('W'):
        annual_count = 52
    elif freq_upper.startswith('D') or freq_upper.startswith('B'):
        annual_count = 365
    elif freq_upper.startswith('M'):
        annual_count = 12
    else:
        # Fallback to weekly if custom/unknown
        annual_count = 52

    # Logic: 1 year for min, 2 years for max
    min_records = annual_count + eval_period
    #max_records = (2 * annual_count) + eval_period
    
    return min_records

def get_dynamic_lags(series: pd.Series) -> list:
    n = len(series)
    nlags = min(int(n * 0.5), 60)
    if nlags < 5:
        return [1, 2, 3]

    autocorrelations, confint = acf(series.dropna(), nlags=nlags, alpha=0.25, fft=True)
    autocorr_values = autocorrelations[1:]
    conf_limit = confint[1:, 1] - autocorr_values
    is_significant = np.abs(autocorr_values) > conf_limit
    significant_lags_indices = np.where(is_significant)[0] + 1
    
    ranked_indices = np.argsort(np.abs(autocorr_values[is_significant]))[::-1]
    top_lags = significant_lags_indices[ranked_indices[:10]].tolist()
    
    return sorted(list(set([1, 2, 3] + top_lags)))[:10]


def get_optimal_window(group, date_column):
    """
    Dynamically determines window size based on data cadence.
    """
    # Ensure sorted by date
    temp_df = group.sort_values(date_column)
    
    # Calculate the median delta between rows
    delta = temp_df[date_column].diff().median()
    
    if delta <= pd.Timedelta(days=1):      # Daily (or sub-daily)
        return 90  # 3 months of context
    elif delta <= pd.Timedelta(days=7):    # Weekly
        return 52  # 1 year of context
    elif delta <= pd.Timedelta(days=31):   # Monthly
        return 24  # 2 years of context
    else:
        return 12  # Default fallback