train_utils.py

"""
Shared utilities for DISMIS training and prediction scripts.

Contains feature definitions, classifier configurations, data loading,
balancing strategies, custom training objectives, and fit helpers.
"""

import numpy as np
import os
import polars as pl

from sklearn.linear_model import SGDClassifier, LogisticRegression
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import make_pipeline
from sklearn.metrics import recall_score, precision_score
from sklearn.neural_network import MLPClassifier
from sklearn.ensemble import RandomForestClassifier, AdaBoostClassifier
from sklearn.tree import DecisionTreeClassifier
from sklearn.naive_bayes import GaussianNB
from sklearn.svm import SVC, LinearSVC
from sklearn.discriminant_analysis import QuadraticDiscriminantAnalysis
from sklearn.cluster import MiniBatchKMeans
import xgboost as xgb


# ============================================================================
# Constants
# ============================================================================

TYPE_MAPPING = {'numeric': 0, 'date': 1, 'categorical': 2, 'text': 3}
TARGET_TYPE_STR_MAP = {v: k for k, v in TYPE_MAPPING.items()}


# ============================================================================
# Feature set definitions
# ============================================================================

def get_featuresets(size):
    if size == 'big':
        featureset_num = ['key_distance_outlier_distribution', 'length_outlier_distribution', 'repeated_substring1_outlier_distribution', 'repeated_substring2_outlier_distribution', 'repeated_substring3_outlier_distribution', 'sign_outlier_feature', 'non_alphanumerical_outlier_distribution', 'capital_letter_outlier_distribution', 'frequent_values_1', 'frequent_values_10', 'approximate_similar_samples_25', 'approximate_similar_samples_corr_25', 'no_duplicate_similar_samples_25', 'BucketPDFGoF', 'bucket_knn_square', 'pyod_mad', 'nan_outlier', 'type_feature', 'type', 'syntactic_outlier']
        featureset_date = ['key_distance_outlier_distribution', 'length_outlier_distribution', 'repeated_substring1_outlier_distribution', 'repeated_substring2_outlier_distribution', 'repeated_substring3_outlier_distribution', 'sign_outlier_feature', 'non_alphanumerical_outlier_distribution', 'capital_letter_outlier_distribution', 'frequent_values_1', 'frequent_values_10', 'approximate_similar_samples_25', 'approximate_similar_samples_corr_25', 'no_duplicate_similar_samples_25', 'BucketPDFGoF', 'bucket_knn_square', 'pyod_mad', 'nan_outlier', 'type_feature', 'type', 'syntactic_outlier']
        featureset_cat = ['semantic_comments', 'semantic_placeholder', 'semantic_unsure', 'semantic_valid', 'semantic_outlier_3_new_dub', 'semantic_outlier_10_new_dub', 'semantic_outlier_25_new_dub', 'semantic_outlier_100_new_dub', 'nan_outlier',  'type_feature', 'type', 'syntactic_outlier', 'key_distance_outlier_distribution', 'length_outlier_distribution', 'repeated_substring1_outlier_distribution', 'repeated_substring2_outlier_distribution', 'repeated_substring3_outlier_distribution', 'sign_outlier_feature', 'capital_letter_outlier_distribution', 'frequency_outlier', 'approximate_similar_samples_25', 'approximate_similar_samples_corr_25', 'no_duplicate_similar_samples_25']
        featureset_text = ['semantic_comments', 'semantic_placeholder', 'semantic_unsure', 'semantic_valid', 'semantic_outlier_3_new_dub', 'semantic_outlier_10_new_dub', 'semantic_outlier_25_new_dub', 'semantic_outlier_100_new_dub', 'nan_outlier',  'type_feature', 'type', 'syntactic_outlier', 'key_distance_outlier_distribution', 'length_outlier_distribution', 'repeated_substring1_outlier_distribution', 'repeated_substring2_outlier_distribution', 'repeated_substring3_outlier_distribution', 'sign_outlier_feature', 'capital_letter_outlier_distribution', 'frequency_outlier', 'approximate_similar_samples_25', 'approximate_similar_samples_corr_25', 'no_duplicate_similar_samples_25']

    else:
        featureset_num = ['length_outlier_distribution', 'repeated_substring1_outlier_distribution', 'repeated_substring2_outlier_distribution', 'repeated_substring3_outlier_distribution', 'sign_outlier_feature', 'non_alphanumerical_outlier_distribution', 'frequency_outlier', 'frequent_values_1', 'approximate_similar_samples_25', 'approximate_similar_samples_corr_25', 'no_duplicate_similar_samples_25', 'BucketPDFGoF', 'bucket_knn_square', 'pyod_mad', 'nan_outlier', 'type_feature', 'type_feature_2', 'type', 'syntactic_outlier']
        featureset_date = ['key_distance_outlier_distribution', 'repeated_substring1_outlier_distribution', 'repeated_substring2_outlier_distribution', 'repeated_substring3_outlier_distribution', 'sign_outlier_feature', 'capital_letter_outlier_distribution', 'frequency_outlier', 'frequent_values_1', 'frequent_values_10', 'approximate_similar_samples_25', 'approximate_similar_samples_corr_25', 'no_duplicate_similar_samples_25', 'BucketPDFGoF', 'bucket_knn_square', 'pyod_mad', 'nan_outlier', 'quantile', 'type_feature', 'type_feature_2', 'type']
        featureset_cat = ['semantic_comments', 'semantic_placeholder', 'semantic_unsure', 'semantic_valid', 'semantic_outlier_3_new_dub', 'semantic_outlier_10_new_dub', 'semantic_outlier_25_new_dub', 'semantic_outlier_100_new_dub', 'nan_outlier', 'type_feature', 'type', 'syntactic_outlier', 'key_distance_outlier_distribution', 'length_outlier_distribution', 'repeated_substring1_outlier_distribution', 'repeated_substring2_outlier_distribution', 'repeated_substring3_outlier_distribution', 'frequency_outlier']
        featureset_text = ['semantic_comments', 'semantic_placeholder', 'semantic_unsure', 'semantic_valid', 'nan_outlier', 'type_feature', 'repeated_substring1_outlier_distribution', 'repeated_substring3_outlier_distribution', 'repeated_substring2_outlier_distribution']

    return {
        'num': featureset_num,
        'date': featureset_date,
        'cat': featureset_cat,
        'text': featureset_text,
    }

def get_featureset(target_type, featuresets):
    if target_type == 0:
        return featuresets['num']
    elif target_type == 1:
        return featuresets['date']
    elif target_type == 2:
        return featuresets['cat']
    elif target_type == 3:
        return featuresets['text']
    else:
        raise ValueError(f"Unknown target type: {target_type}")


def get_classifier(classifier_name, pos_weight=None):
    """
    Return a classifier instance based on the name.
    
    Includes all classifiers used by both predict.py and train_dismis_robust_old.py.
    """
    # ------------------------------------------------------------------
    # Standard sklearn classifiers
    # ------------------------------------------------------------------
    if classifier_name == 'mlp':
        return MLPClassifier(max_iter=1000, random_state=42, early_stopping=True)
    elif classifier_name == 'random_forest':
        return RandomForestClassifier(n_estimators=100, random_state=42, n_jobs=1)
    elif classifier_name == 'mlp5':
        return MLPClassifier(max_iter=1000, random_state=42, early_stopping=True, hidden_layer_sizes=(64, 256, 64))
    elif classifier_name == 'random_forest_5':
        return RandomForestClassifier(n_estimators=100, max_depth=5, random_state=42, n_jobs=1)
    elif classifier_name == 'logistic_sgd':
        return SGDClassifier(loss='log_loss', max_iter=1000, n_jobs=1, random_state=42, early_stopping=True)
    elif classifier_name == 'logistic_regression':
        return LogisticRegression(max_iter=1000, random_state=42, n_jobs=1)
    elif classifier_name == 'decision_tree_5':
        return DecisionTreeClassifier(max_depth=5, random_state=42)
    elif classifier_name == 'decision_tree':
        return DecisionTreeClassifier(random_state=42)
    elif classifier_name == 'gaussian_nb':
        return GaussianNB()
    elif classifier_name == 'adaboost':
        return AdaBoostClassifier(n_estimators=100, random_state=42)
    elif classifier_name == 'rbf_svm':
        return SVC(kernel='rbf', random_state=42)
    elif classifier_name == 'linear_svm':
        return LinearSVC(max_iter=1000, random_state=42)
    elif classifier_name == 'qda':
        return QuadraticDiscriminantAnalysis()
    elif classifier_name == 'xgboost_robust':
        return xgb.XGBClassifier(
            random_state=42, n_jobs=1, use_label_encoder=False,
            eval_metric='logloss',
            scale_pos_weight=40 if pos_weight is None else pos_weight,
            max_depth=5, min_child_weight=8,
            subsample=0.75, colsample_bytree=0.8,
            reg_alpha=0.5, reg_lambda=0.5,
            n_estimators=150, learning_rate=0.07, gamma=0
        )
    elif classifier_name == 'xgboost_stratified_categorical':
        """Stratified sampling for categorical. Use with --balance-type stratified."""
        return xgb.XGBClassifier(
            random_state=42, n_jobs=1, use_label_encoder=False,
            eval_metric='logloss',
            scale_pos_weight=50, max_depth=8, min_child_weight=3,
            subsample=0.85, colsample_bytree=0.85,
            reg_alpha=0.2, reg_lambda=0.2,
            n_estimators=150, learning_rate=0.08
        )
    else:
        raise ValueError(f"Unknown classifier: {classifier_name}")


def fit_classifier(classifier_name, train_features, train_labels, pos_weight=None):
    base_clf = get_classifier(classifier_name, pos_weight)
    clf = make_pipeline(StandardScaler(), base_clf)
    clf.fit(train_features, train_labels)
    return clf

def stratified_undersample(X, y, target_ratio=5.0, n_clusters=10, random_state=42):
    pos_mask = y == 1
    neg_mask = y == 0

    X_pos = X[pos_mask]
    X_neg = X[neg_mask]

    n_pos = len(X_pos)
    n_neg_target = int(n_pos * target_ratio)

    if n_neg_target >= len(X_neg):
        return X, y

    print(f"    Stratified undersampling: {len(X_neg):,} negatives -> {n_neg_target:,} (preserving {n_clusters} subgroups)")

    kmeans = MiniBatchKMeans(
        n_clusters=min(n_clusters, len(X_neg)),
        random_state=random_state,
        batch_size=min(1000, len(X_neg))
    )
    neg_clusters = kmeans.fit_predict(X_neg)

    samples_per_cluster = n_neg_target // n_clusters
    sampled_indices = []

    for cluster_id in range(n_clusters):
        cluster_indices = np.where(neg_clusters == cluster_id)[0]
        if len(cluster_indices) > 0:
            n_to_sample = min(samples_per_cluster, len(cluster_indices))
            sampled = np.random.choice(cluster_indices, n_to_sample, replace=False)
            sampled_indices.extend(sampled)

    if len(sampled_indices) < n_neg_target:
        remaining = n_neg_target - len(sampled_indices)
        available = np.setdiff1d(np.arange(len(X_neg)), sampled_indices)
        if len(available) > 0:
            extra = np.random.choice(available, min(remaining, len(available)), replace=False)
            sampled_indices.extend(extra)

    X_neg_sampled = X_neg[sampled_indices]

    X_balanced = np.vstack([X_pos, X_neg_sampled])
    y_balanced = np.concatenate([np.ones(len(X_pos)), np.zeros(len(X_neg_sampled))])

    shuffle_idx = np.random.permutation(len(X_balanced))
    return X_balanced[shuffle_idx], y_balanced[shuffle_idx]


def balance_data(train_features, train_labels, balance_type,
                 balance_fraction=0.5, random_state=42):
    if balance_type == 'weight':
        return train_features, train_labels

    rng = np.random.RandomState(random_state)
    num_positive = int(train_labels.sum())
    num_negative = len(train_labels) - num_positive

    if balance_type == 'stratified':
        return stratified_undersample(
            train_features, train_labels,
            target_ratio=5.0, n_clusters=10, random_state=random_state
        )

    elif balance_type == 'drop':
        if balance_fraction is None:
            balance_fraction = 0.1
        desired_negative = int(num_positive * (1 - balance_fraction) / balance_fraction)

        neg_indices = np.where(train_labels == 0)[0]
        pos_indices = np.where(train_labels == 1)[0]

        sampled_neg = rng.choice(
            neg_indices,
            min(desired_negative, len(neg_indices)),
            replace=False
        )
        keep_indices = np.concatenate([pos_indices, sampled_neg])
        return train_features[keep_indices], train_labels[keep_indices]

    elif balance_type == 'oversample':
        if balance_fraction is None:
            balance_fraction = 0.5
        desired_positive = int(num_negative * balance_fraction / (1 - balance_fraction))

        pos_indices = np.where(train_labels == 1)[0]
        neg_indices = np.where(train_labels == 0)[0]

        sampled_pos = rng.choice(pos_indices, desired_positive, replace=True)
        features = np.vstack([train_features[neg_indices], train_features[sampled_pos]])
        labels = np.concatenate([train_labels[neg_indices], train_labels[sampled_pos]])
        return features, labels

    # Unknown balance_type — return as-is
    return train_features, train_labels

def load_columnar_data(data_dir, csv_files):
    """
    Load columnar CSV files and merge them horizontally.
    
    Args:
        data_dir: Directory containing the CSV files.
        csv_files: List of column names (without .csv extension).
    
    Returns:
        pandas DataFrame with TYPE_MAPPING applied to 'type' column.
    """
    column_dfs = []
    for csv_file in csv_files:
        col_path = os.path.join(data_dir, csv_file + ".csv")
        col_df = pl.read_csv(col_path)
        col_df = col_df.rename({col: col.replace("_inference", "") for col in col_df.columns})
        column_dfs.append(col_df)

    eval_data_pl = pl.concat(column_dfs, how="horizontal")
    print(f"Loaded data with {len(eval_data_pl)} rows and {len(eval_data_pl.columns)} columns")

    eval_data = eval_data_pl.to_pandas()
    eval_data['type'] = eval_data['type'].map(TYPE_MAPPING)
    return eval_data