models.py

import pandas as pd
from sklearn.linear_model import LogisticRegression as SklearnLogisticRegression
from sklearn.naive_bayes import BernoulliNB
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import AdaBoostClassifier
from logistic_regression import LogisticRegressionManual
from naivebayes import NaiveBayes


def get_model(algorithm, use_custom=True, **model_params):
    algorithm = algorithm.lower()

    if algorithm == "logistic_regression":
        if use_custom:
            return LogisticRegressionManual(
                learning_rate=1e-4,
                max_iter=5000,
                reg_lambda=0.01,
                use_regularization=True
            )
        else:
            return SklearnLogisticRegression(max_iter=1000, **model_params)

    if algorithm == "naive_bayes":
        if use_custom:
            return NaiveBayes()
        return BernoulliNB(**model_params)
    if algorithm == "id3":
        return DecisionTreeClassifier(criterion="entropy", random_state=42, **model_params)

    if algorithm == "adaboost":
        base = DecisionTreeClassifier(max_depth=1, random_state=42)
        return AdaBoostClassifier(estimator=base, n_estimators=200, random_state=42, **model_params)

    return None


def train_model(train_df, target_product, model, exclude_columns=None):
    exclude_columns = set(exclude_columns or [])
    exclude_columns.add(target_product)

    y = train_df[target_product].astype(int)
    X = train_df.drop(columns=list(exclude_columns), errors="ignore")

    model.fit(X, y)
    return model, X.columns.tolist()


def train_all_models(train_df, products, algorithms, exclude_columns=None):
    model_store = {}

    for algo in algorithms:
        for prod in products:
            model = get_model(algo, use_custom=False)
            model, feature_cols = train_model(train_df, prod, model, exclude_columns)

            model_store[(algo, prod)] = {
                "model": model,
                "feature_cols": feature_cols,
            }

    return model_store


def predict_probability(model, feature_columns, partial_cart, numeric_features=None):
    x = pd.DataFrame(0.0, index=[0], columns=feature_columns)

    for prod in partial_cart:
        if prod in x.columns:
            x.at[0, prod] = 1.0

    if numeric_features:
        for k, v in numeric_features.items():
            if k in x.columns:
                x.at[0, k] = float(v)

    return float(model.predict_proba(x)[0, 1])


def rank_products(model_store, algorithm, candidate_products, partial_cart, product_prices, numeric_features=None):
    rankings = []

    for prod in candidate_products:
        entry = model_store[(algorithm, prod)]
        model = entry["model"]
        feature_cols = entry["feature_cols"]

        prob = predict_probability(model, feature_cols, partial_cart, numeric_features)
        rankings.append((prod, prob * product_prices.get(prod, 0.0)))

    return sorted(rankings, key=lambda x: x[1], reverse=True)