analysis.py

import os
import glob
import json
import shutil
import logging
from pathlib import Path

import polars as pl
import pandas as pd
import numpy as np
from typing import List, Dict
from sklearn.feature_selection import mutual_info_regression
from data_analysis.config import GOOGLE_APPLICATION_CREDENTIALS, OWNER, REPO, BUCKET_NAME, INGESTION_JOB_ID

from utils.gcs_utils import GCSClient

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def load_data_from_gcs(bucket_name: str, prefix: str, local_dir: str, schema: dict) -> pl.DataFrame:
    """
    Download JSON files from a GCS bucket and load them into a single Polars DataFrame.
    """
    with GCSClient(gcp_credentials_path=GOOGLE_APPLICATION_CREDENTIALS) as gcs:
        gcs.download_folder(bucket_name=bucket_name, prefix=prefix, local_dir=local_dir)
    
    json_files = glob.glob(os.path.join(local_dir, "batch*.json"))
    if not json_files:
        logger.warning(f"No JSON files starting with 'batch' found in {local_dir}")
        return pl.DataFrame()
    
    dfs = []
    all_columns = {}
    ordered_columns = []
    
    # Read each JSON into a Polars DataFrame and track schema
    for file in json_files:
        df = pl.read_json(file, schema=schema).with_columns(
            pl.col(["is_weekend", "is_us_holiday"]).cast(pl.Int8)
        )
        # if "created_at" in df.columns:
        #     df = df.with_columns(pl.col("created_at").cast(pl.Datetime("ms")))
        dfs.append(df)
        
        for col in df.columns:
            if col not in all_columns:
                all_columns[col] = df.schema[col]
    
    # Determine column order based on first appearance
    for df in dfs:
        for col in df.columns:
            if col not in ordered_columns:
                ordered_columns.append(col)
    for col in all_columns:
        if col not in ordered_columns:
            ordered_columns.append(col)
    
    # Ensure each DataFrame has all columns and consistent order
    dfs_fixed = []
    for df in dfs:
        missing_cols = set(all_columns.keys()) - set(df.columns)
        for col in missing_cols:
            df = df.with_columns(pl.lit(None, dtype=all_columns[col]).alias(col))
        df = df.select(ordered_columns)
        dfs_fixed.append(df)
    
    final_df = pl.concat(dfs_fixed, how='vertical')
    logger.info(f"Loaded {len(dfs_fixed)} files into DataFrame with {final_df.height} rows.")
    return final_df


def split_dataset(df: pl.DataFrame, date_col: str, train_ratio: float = 0.7, val_ratio: float = 0.15, seed: int = 42) -> tuple:
    """
    Split the DataFrame into train, validation, and test sets based on chronological order for test set.
    """
    df_sorted = df.sort(date_col)
    
    total_len = df_sorted.height
    n_train = int(total_len * train_ratio)
    n_val = int(total_len * val_ratio)
    n_train_val = n_train + n_val
    
    test_df = df_sorted.slice(n_train_val, total_len - n_train_val)
    train_val_pool = df_sorted.slice(0, n_train_val).sample(fraction=1.0, shuffle=True, seed=seed)
    train_df = train_val_pool.slice(0, n_train)
    val_df = train_val_pool.slice(n_train, n_val)
    
    logger.info(f"Split data into Train: {train_df.height}, Val: {val_df.height}, Test: {test_df.height}")
    return train_df, val_df, test_df


def compute_summary_statistics(df: pl.DataFrame, columns: list) -> pl.DataFrame:
    """
    Compute summary statistics (mean, std, min, max, skew, median) for specified columns.
    Returns a DataFrame with 'feature' and a dict of stats.
    """
    agg_exprs = []
    for col in columns:
        agg_exprs.extend([
            pl.col(col).mean().alias(f"{col}_mean"),
            pl.col(col).std().alias(f"{col}_std"),
            pl.col(col).min().alias(f"{col}_min"),
            pl.col(col).max().alias(f"{col}_max"),
            pl.col(col).median().alias(f"{col}_median"),
            pl.col(col).skew().alias(f"{col}_skew"),
        ])
    summary_df = df.select(agg_exprs)
    
    summary_json = {}
    for col in columns:
        summary_json[col] = {
            "mean": float(summary_df[0, f"{col}_mean"]) if summary_df[0, f"{col}_mean"] is not None else None,
            "std": float(summary_df[0, f"{col}_std"]) if summary_df[0, f"{col}_std"] is not None else None,
            "min": float(summary_df[0, f"{col}_min"]) if summary_df[0, f"{col}_min"] is not None else None,
            "max": float(summary_df[0, f"{col}_max"]) if summary_df[0, f"{col}_max"] is not None else None,
            "median": float(summary_df[0, f"{col}_median"]) if summary_df[0, f"{col}_median"] is not None else None,
            "skew": float(summary_df[0, f"{col}_skew"]) if summary_df[0, f"{col}_skew"] is not None else None,
        }

    rows = []
    for col in columns:
        rows.append({
            "feature": col,
            "mean": summary_json[col]["mean"],
            "std": summary_json[col]["std"],
            "min": summary_json[col]["min"],
            "max": summary_json[col]["max"],
            "median": summary_json[col]["median"],
            "skew": summary_json[col]["skew"],
        })
    # Convert to a Polars DataFrame with feature and stats
    return pl.DataFrame(rows)


def compute_correlations(df: pl.DataFrame, target_col: str, numeric_cols: list) -> pd.Series:
    """
    Compute Pearson correlation of each numeric column with the target column.
    Returns a pandas Series of correlations indexed by feature name (excluding target).
    """
    cols = [target_col] + [col for col in numeric_cols if col != target_col]
    df_clean = df.select(cols).drop_nulls()
    corr_df = df_clean.to_pandas().corr(method='pearson')
    target_corr = corr_df[target_col].drop(labels=[target_col])
    target_corr = target_corr.reindex(target_corr.abs().sort_values(ascending=False).index)
    return target_corr


def compute_mutual_info(df: pl.DataFrame, target_col: str, feature_cols: list) -> pd.Series:
    """
    Compute mutual information regression between target and each feature.
    Returns a pandas Series of MI scores indexed by feature name.
    """
    df_clean = df.select(feature_cols + [target_col]).drop_nulls()
    X = df_clean.select(feature_cols).to_pandas()
    y = df_clean.select(target_col).to_pandas().values.ravel()
    mi = mutual_info_regression(X, y, random_state=42)
    mi_series = pd.Series(mi, index=feature_cols, name="mi_score")
    return mi_series.sort_values(ascending=False)


def detect_outliers_iqr(df: pl.DataFrame, col: str, multiplier: float = 1.5) -> tuple:
    """
    Calculate IQR-based outlier bounds for a column.
    Returns (lower_bound, upper_bound).
    """
    q1 = df.select(pl.col(col).quantile(0.25)).to_numpy()[0][0]
    q3 = df.select(pl.col(col).quantile(0.75)).to_numpy()[0][0]
    iqr = q3 - q1
    lower_bound = q1 - multiplier * iqr
    upper_bound = q3 + multiplier * iqr
    return lower_bound, upper_bound


def filter_outliers(df: pl.DataFrame, col: str, lower: float = None, upper: float = None) -> pl.DataFrame:
    """
    Filter the DataFrame by given column bounds (retain only rows within bounds).
    """
    mask = None
    if lower is not None:
        mask = (pl.col(col) >= lower)
    if upper is not None:
        if mask is None:
            mask = (pl.col(col) <= upper)
        else:
            mask = mask & (pl.col(col) <= upper)
    if mask is not None:
        return df.filter(mask)
    else:
        return df


def remove_group_outliers(
    df: pl.DataFrame,
    feature_cols: list[str],
    target_col: str,
    bin_sizes: dict[str, float],
    multiple: float = 3.0,
    min_group_size: int = 5
) -> pl.DataFrame:

    group_bin_cols = []
    df_binned = df

    # Create bins
    for col in feature_cols:
        bin_size = bin_sizes.get(col)
        if not bin_size or bin_size <= 0:
            continue

        bin_col = f"{col}_bin"
        df_binned = df_binned.with_columns(
            ((pl.col(col) / bin_size).floor() * bin_size).alias(bin_col)
        )
        group_bin_cols.append(bin_col)

    if not group_bin_cols:
        return df

    # Compute group size + median
    group_stats = (
        df_binned
        .group_by(group_bin_cols)
        .agg([
            pl.len().alias("group_size"),
            pl.col(target_col).median().alias("median")
        ])
    )

    df_merged = df_binned.join(group_stats, on=group_bin_cols, how="left")

    # Apply filtering ONLY if group is large enough
    clean_df = df_merged.filter(
        (pl.col("group_size") < min_group_size) |
        (pl.col(target_col) <= pl.col("median") * multiple)
    ).drop(group_bin_cols + ["median", "group_size"])

    logger.info(
        f"Removed {df.height - clean_df.height} rows."
    )

    return clean_df


def suggest_bin_sizes(
    df: pl.DataFrame,
    feature_cols: List[str],
    min_bins: int = 5,
    max_bins: int = 50
) -> Dict[str, float]:
    """
    Suggest bin size for each feature using Freedman-Diaconis rule.
    Ensures number of bins is within reasonable limits.
    """

    bin_sizes = {}
    n = df.height

    for col in feature_cols:
        if col not in df.columns:
            continue

        series = df[col].drop_nulls()

        if series.len() < 10:
            continue

        q1 = series.quantile(0.25)
        q3 = series.quantile(0.75)
        iqr = q3 - q1

        if iqr == 0:
            # fallback if constant column
            bin_sizes[col] = 1.0
            continue

        # Freedman–Diaconis bin width
        bin_width = (2 * iqr) / (n ** (1/3))

        # Clamp number of bins to avoid extreme fragmentation
        value_range = series.max() - series.min()
        estimated_bins = value_range / bin_width if bin_width > 0 else 1

        estimated_bins = max(min_bins, min(max_bins, estimated_bins))

        # Recalculate bin width based on clamped bin count
        final_bin_width = value_range / estimated_bins if estimated_bins > 0 else 1.0

        bin_sizes[col] = float(final_bin_width)

    return bin_sizes


def apply_outlier_thresholds(df: pl.DataFrame, thresholds: dict) -> pl.DataFrame:
    """
    Apply outlier thresholds to remove rows outside specified bounds.
    `thresholds` is a dict mapping column -> {'lower': val, 'upper': val}.
    """
    clean_df = df
    for col, bounds in thresholds.items():
        lower = bounds.get('lower')
        upper = bounds.get('upper')
        clean_df = filter_outliers(clean_df, col, lower=lower, upper=upper)
    return clean_df


def save_dataframe_to_csv(df: pl.DataFrame, path: str):
    """
    Save a Polars DataFrame to a CSV file (creates directories if needed).
    """
    Path(os.path.dirname(path)).mkdir(parents=True, exist_ok=True)
    df.write_csv(path)
    logger.info(f"Saved DataFrame to {path}")


def upload_to_gcs(local_path: str, bucket_name: str, dest_blob: str):
    """
    Upload a local file to a GCS bucket.
    """
    with GCSClient(gcp_credentials_path=GOOGLE_APPLICATION_CREDENTIALS) as gcs:
        gcs.upload_file(bucket_name, local_path, dest_blob)

    # # Remove file only if it exists
    if os.path.exists(local_path):
        os.remove(local_path)
        logger.info(f"Deleted local file: {local_path}")


def generate_analysis_payload(df: pl.DataFrame, target_col: str, numeric_cols: list, ordinal_cols: list, output_path: str):
    """
    Generate a JSON payload summarizing the dataset analysis (for LLM input).
    Includes features, target, summary statistics, correlations, and mutual information.
    """
    features = [col for col in numeric_cols + ordinal_cols if col != target_col]
    summary_stats_df = compute_summary_statistics(df, numeric_cols + [target_col])
    summary_stats = {row[0]: row[1] for row in summary_stats_df.rows()}
    correlations = compute_correlations(df, target_col, numeric_cols + [target_col]).to_dict()
    logger.info("Correlation Values :")
    for key, value in correlations.items():
        print(f"{key:<30}\t{value:.6f}")
    mutual_info = compute_mutual_info(df, target_col, numeric_cols).to_dict()
    logger.info("Mutual Info Values :")
    for key, value in mutual_info.items():
        print(f"{key:<30}\t{value:.6f}")
    
    payload = {
        "features": features,
        "target": target_col,
        "summary_statistics": summary_stats,
        "correlations": correlations,
        "mutual_information": mutual_info,
    }

    Path(os.path.dirname(output_path)).mkdir(parents=True, exist_ok=True)
    with open(output_path, 'w') as f:
        json.dump(payload, f, indent=4)
    logger.info(f"Saved analysis payload to {output_path}")

    return summary_stats, correlations, mutual_info


# Example usage (without invoking an LLM):
if __name__ == "__main__":

    ingestion_id = INGESTION_JOB_ID
    owner = OWNER
    repo = REPO
    bucket = BUCKET_NAME

    prefix = f"etl_pipeline/{owner}/{repo}/{ingestion_id}"
    
    target_col = "time_to_merge_hours"
    numeric_cols = [
        "avg_time_since_last_mod_days", "commit_count", "unique_commit_authors", "assignee_count",
        "file_count", "files_added", "files_modified", "files_deleted", "unique_file_types",
        "line_additions", "line_deletions", "code_churn", "title_length", "body_length",
        "label_count" 
    ]
    ordinal_cols = ["is_site_admin", "author_association"]

    schema = {
        "time_to_merge_hours": pl.Float64,
        "avg_time_since_last_mod_days": pl.Float64,
        "created_at": pl.String,
        "commit_count": pl.Int64,
        "unique_commit_authors": pl.Int64,
        "assignee_count": pl.Int64,
        "file_count": pl.Int64,
        "files_added": pl.Int64,
        "files_modified": pl.Int64,
        "files_deleted": pl.Int64,
        "unique_file_types": pl.Int64,
        "line_additions": pl.Int64,
        "line_deletions": pl.Int64,
        "code_churn": pl.Int64,
        "title_length": pl.Int64,
        "body_length": pl.Int64,
        "label_count": pl.Int64,
        "hour": pl.Int64,
        "weekday":pl.Int64,
        "is_weekend":pl.Boolean,
        "is_us_holiday":pl.Boolean,
        "is_site_admin": pl.Int64,
        "author_association": pl.Utf8
    }

    output_prefix = f"analysis_preprocessing/{owner}/{repo}/{ingestion_id}"
    local_dir = f"etl_pipeline/{owner}/{repo}/{ingestion_id}"
    os.makedirs(local_dir, exist_ok=True)

    output_dir = os.path.join(output_prefix,"preprocessed")

    df_all = load_data_from_gcs(bucket_name=bucket, prefix=prefix, local_dir=local_dir, schema=schema)

    train_df, val_df, test_df = split_dataset(df_all, date_col="created_at", train_ratio=0.70, val_ratio=0.15)

    analysis_output_path = os.path.join(output_prefix,"analysis","analysis_output.json")
    _, correlations, mutual_info = generate_analysis_payload(train_df, target_col, numeric_cols, ordinal_cols, analysis_output_path)

    upload_to_gcs(
        local_path=analysis_output_path,
        bucket_name=bucket,
        dest_blob=analysis_output_path,
    )

    train_df_raw_path = os.path.join(output_dir, "train_raw.csv")
    save_dataframe_to_csv(train_df, train_df_raw_path)

    upload_to_gcs(
        local_path=train_df_raw_path,
        bucket_name=bucket,
        dest_blob=train_df_raw_path
    )

    # LLM can take the inputs for the available feature columsn and generated analysis file data to generate feature subsets and suggested thresholds for outlier

    high_mutual_info_features = []
    suggested_features_subsets = []
    suggested_thresholds = {}

    # Manual Calculation :
    corr_vals = {k: abs(v) for k, v in correlations.items() if not np.isnan(v)}
 
    max_corr = max(corr_vals.values())
    norm_corr = {k: v / max_corr for k, v in corr_vals.items()}

    max_mi = max(mutual_info.values())
    norm_mi = {k: v / max_mi if max_mi > 0 else 0 for k, v in mutual_info.items()}

    combined_score = {}
    for k in norm_corr:
        combined_score[k] = 0.6 * norm_corr.get(k, 0) + 0.4 * norm_mi.get(k, 0)

    logger.info("Final Weightage :")
    for key, value in mutual_info.items():
        logger.info(f"{key:<30}\t{value:.6f}")

    top_4_features_aligned_with_target = sorted(combined_score, key=combined_score.get, reverse=True)[:4]

    bin_sizes = suggest_bin_sizes(train_df, top_4_features_aligned_with_target)
    for key, value in bin_sizes.items():
        logger.info(f"{key:<30}\t{value:.6f}")

    cleaned_train = pl.DataFrame()

    cleaned_train = remove_group_outliers(
        df=train_df, 
        feature_cols=top_4_features_aligned_with_target,
        target_col="time_to_merge_hours",
        bin_sizes=bin_sizes,
        min_group_size=5
    )

    if cleaned_train.is_empty():
        cleaned_train = train_df

    logger.info(f'Final Rows in Training Data : {len(cleaned_train)}/{len(train_df)}')
    
    train_cleaned_path = os.path.join(output_dir, "train_cleaned.csv")
    save_dataframe_to_csv(cleaned_train, train_cleaned_path)

    upload_to_gcs(
        local_path=train_cleaned_path,
        bucket_name=bucket,
        dest_blob=train_cleaned_path
    )

    val_path = os.path.join(output_dir, "val.csv")
    save_dataframe_to_csv(val_df, val_path)

    upload_to_gcs(
        local_path=val_path,
        bucket_name=bucket,
        dest_blob=val_path
    )

    test_path = os.path.join(output_dir, "test.csv")
    save_dataframe_to_csv(test_df, test_path)

    upload_to_gcs(
        local_path=test_path,
        bucket_name=bucket,
        dest_blob=test_path
    )
    
    # Save LLM outputs (feature list and thresholds) to JSON
    llm_output = {
        "high_mutual_info_features": high_mutual_info_features,
        "selected_feature_subsets": suggested_features_subsets,
        "outlier_thresholds": suggested_thresholds
    }

    llm_output_path = os.path.join(output_dir, "llm_output.json")
    with open(llm_output_path, "w") as f:
        json.dump(llm_output, f, indent=4)

    upload_to_gcs(
        local_path=llm_output_path,
        bucket_name=bucket,
        dest_blob=llm_output_path
    )

    logger.info("Saved LLM output JSON.")

    if os.path.exists("etl_pipeline") and os.path.isdir("etl_pipeline"):
        shutil.rmtree("etl_pipeline")
    
    if os.path.exists("analysis_preprocessing") and os.path.isdir("analysis_preprocessing"):
        shutil.rmtree("analysis_preprocessing")