import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.impute import SimpleImputer
from sklearn.utils import resample

RANDOM_STATE = 42
TEST_SIZE    = 0.20

# Manual oversampling

def oversample_minority(X, y, random_state=RANDOM_STATE):
    """Random oversample minority class to match majority class count."""
    classes, counts = np.unique(y, return_counts=True)
    majority_class  = classes[np.argmax(counts)]
    minority_class  = classes[np.argmin(counts)]
    majority_count  = counts.max()

    X_maj = X[y == majority_class]
    y_maj = y[y == majority_class]
    X_min = X[y == minority_class]
    y_min = y[y == minority_class]

    X_min_up, y_min_up = resample(
        X_min, y_min,
        replace=True,
        n_samples=majority_count,
        random_state=random_state
    )
    return np.vstack([X_maj, X_min_up]), np.hstack([y_maj, y_min_up])


# Shared pipeline

def _preprocess(X, y, oversample=True):
    X = X.values.astype(float) if hasattr(X, "values") else np.array(X, dtype=float)
    y = y.values.astype(int)   if hasattr(y, "values") else np.array(y, dtype=int)

    # Impute missing values with column median
    imputer = SimpleImputer(strategy="median")
    X = imputer.fit_transform(X)

    # Stratified split
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=TEST_SIZE, random_state=RANDOM_STATE, stratify=y
    )

    # Standard scaling (fit on train only)
    scaler = StandardScaler()
    X_train = scaler.fit_transform(X_train)
    X_test  = scaler.transform(X_test)

    # Oversample on training set only to avoid data leakage
    if oversample:
        X_train, y_train = oversample_minority(X_train, y_train)

    return X_train, X_test, y_train, y_test

# Disease-specific loaders

def load_diabetes(path="data/diabetes.csv"):
    df = pd.read_csv(path)
    # Replace physiologically impossible zeros with NaN
    for col in ["Glucose", "BloodPressure", "SkinThickness", "Insulin", "BMI"]:
        df[col] = df[col].replace(0, np.nan)
    X, y = df.drop("Outcome", axis=1), df["Outcome"]
    return _preprocess(X, y) + (X.columns.tolist(),)


def load_heart(path="data/heart_disease.csv"):
    df = pd.read_csv(path, na_values="?")
    df.dropna(thresh=len(df.columns) - 2, inplace=True)
    X, y = df.drop("target", axis=1), df["target"]
    return _preprocess(X, y) + (X.columns.tolist(),)


def load_sepsis(path="data/sepsis.csv"):
    df = pd.read_csv(path)
    # Drop columns with >60% missing
    df = df.dropna(thresh=int(0.4 * len(df)), axis=1)
    X, y = df.drop("SepsisLabel", axis=1), df["SepsisLabel"]
    return _preprocess(X, y) + (X.columns.tolist(),)


if __name__ == "__main__":
    print("Testing preprocessing…")
    for name, fn, path in [
        ("Diabetes",      load_diabetes, "data/diabetes.csv"),
        ("Heart Disease", load_heart,    "data/heart_disease.csv"),
        ("Sepsis",        load_sepsis,   "data/sepsis.csv"),
    ]:
        X_tr, X_te, y_tr, y_te, feats = fn(path)
        print(f"  {name:14s} | train={X_tr.shape} test={X_te.shape} "
              f"| features={len(feats)} "
              f"| class balance train: 0={int((y_tr==0).sum())} 1={int((y_tr==1).sum())}")
    print("Preprocessing OK.")