"""
Simple SMOTE Data Balancer for CVE training Dataset
Creates balanced training dataset using SMOTE oversampling technique
"""

import pandas as pd
from imblearn.over_sampling import SMOTE


def balance_with_smote(df, random_state=42):
    """Balance dataset using SMOTE oversampling"""
    print("Original dataset shape:", df.shape)
    print("\nOriginal class distribution:")
    class_counts = df["severity_label"].value_counts()
    print(class_counts)
    print("\nClass distribution percentages:")
    print(df["severity_label"].value_counts(normalize=True) * 100)

    # Separate features and target
    X = df.drop(["cve_id", "severity_label", "severity_numeric"], axis=1)
    X = X.astype(int)
    y = df["severity_label"]

    # Apply SMOTE
    smote = SMOTE(random_state=random_state)
    X_balanced, y_balanced = smote.fit_resample(X, y)

    # Create balanced dataframe
    balanced_df = pd.DataFrame(X_balanced, columns=X.columns)
    balanced_df["severity_label"] = y_balanced
    label2num = (
        df[["severity_label", "severity_numeric"]]
        .drop_duplicates()
        .set_index("severity_label")["severity_numeric"]
        .to_dict()
    )
    balanced_df["severity_numeric"] = balanced_df["severity_label"].map(label2num)

    # Generate synthetic CVE IDs for new samples
    original_ids = df["cve_id"].tolist()
    synthetic_count = len(balanced_df) - len(df)
    synthetic_ids = [f"SYNTHETIC-{i:05d}" for i in range(synthetic_count)]

    # Combine original and synthetic IDs
    all_ids = original_ids + synthetic_ids
    balanced_df["cve_id"] = all_ids[: len(balanced_df)]

    # Reorder columns to match original
    balanced_df = balanced_df[df.columns]

    print("\nAfter SMOTE balancing:")
    print("Balanced dataset shape:", balanced_df.shape)
    print("\nBalanced class distribution:")
    print(balanced_df["severity_label"].value_counts())
    print("\nBalanced class distribution percentages:")
    print(balanced_df["severity_label"].value_counts(normalize=True) * 100)

    return balanced_df


def main():
    # Load the dataset
    input_file = "data/cve_training_dataset.csv"
    output_file = "data/balanced-training-dataset-through-smote.csv"

    try:
        df = pd.read_csv(input_file)
        print(f"Loaded dataset from: {input_file}")
    except FileNotFoundError:
        print(f"Error: Could not find file {input_file}")
        return

    # Create SMOTE balanced dataset
    balanced_df = balance_with_smote(df)

    # Save the balanced dataset
    balanced_df.to_csv(output_file, index=False)
    print(f"\nSMOTE balanced dataset saved as: {output_file}")
    print("✅ Done!")


if __name__ == "__main__":
    main()
