src/tune.py

import argparse
import json
import warnings
from pathlib import Path

import numpy as np
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.ensemble import (
    HistGradientBoostingClassifier,
    HistGradientBoostingRegressor,
    RandomForestClassifier,
    RandomForestRegressor,
)
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LogisticRegression, Ridge
from sklearn.metrics import confusion_matrix, f1_score, mean_squared_error, roc_auc_score
from sklearn.model_selection import ParameterGrid
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler
from sklearn.svm import SVC, SVR
from tqdm import tqdm

from nn.torch_ft_transformer import (
    TorchFTTransformerClassifier,
    TorchFTTransformerRegressor,
)
from nn.torch_mlp import TorchMLPClassifier, TorchMLPRegressor
from preprocessing import (
    CommaSeparatedMultiLabelBinarizer,
    UnixTimestampTransformer,
    infer_task_type,
    to_bool_if_binary,
)
from utils.logger import logger

warnings.filterwarnings("ignore", category=UserWarning)


def load_json(path: str) -> dict:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def validate_required_columns(df: pd.DataFrame, required_columns: list[str], context: str):
    missing_columns = [col for col in required_columns if col not in df.columns]
    if not missing_columns:
        return

    available_columns = ", ".join(df.columns.astype(str).tolist())
    missing_list = ", ".join(missing_columns)
    raise ValueError(
        f"Missing required columns for {context}: {missing_list}\n"
        f"Available dataset columns: {available_columns}"
    )


def build_preprocessor(input_features, cols_string, cols_date, cols_multi, use_scaler=True):
    transformers = []

    numeric_cols = [
        c for c in input_features if c not in cols_string + cols_date + cols_multi
    ]

    if numeric_cols:
        steps = [("imputer", SimpleImputer(strategy="median"))]
        if use_scaler:
            steps.append(("scaler", StandardScaler()))
        transformers.append(("numeric", Pipeline(steps), numeric_cols))

    if cols_string:
        steps = [
            ("imputer", SimpleImputer(strategy="most_frequent")),
            ("onehot", OneHotEncoder(handle_unknown="ignore", sparse_output=False)),
        ]
        transformers.append(("categorical", Pipeline(steps), cols_string))

    if cols_date:
        steps = [
            ("unix_ts", UnixTimestampTransformer()),
            ("imputer", SimpleImputer(strategy="median")),
        ]
        if use_scaler:
            steps.append(("scaler", StandardScaler()))
        transformers.append(("date", Pipeline(steps), cols_date))

    if cols_multi:
        steps = [
            ("imputer", SimpleImputer(strategy="constant", fill_value="")),
            ("multilabel", CommaSeparatedMultiLabelBinarizer()),
        ]
        transformers.append(("multi", Pipeline(steps), cols_multi))

    return ColumnTransformer(transformers=transformers, remainder="drop")


def get_registry():
    return {
        "binary": {
            "hgb": HistGradientBoostingClassifier,
            "rf": RandomForestClassifier,
            "lr": LogisticRegression,
            "svc": SVC,
            "torch_mlp": TorchMLPClassifier,
            "torch_ft_transformer": TorchFTTransformerClassifier,
        },
        "categorical": {
            "hgb": HistGradientBoostingClassifier,
            "rf": RandomForestClassifier,
            "lr": LogisticRegression,
            "svc": SVC,
            "torch_mlp": TorchMLPClassifier,
            "torch_ft_transformer": TorchFTTransformerClassifier,
        },
        "continuous": {
            "hgb": HistGradientBoostingRegressor,
            "rf": RandomForestRegressor,
            "ridge": Ridge,
            "svc": SVR,
            "torch_mlp": TorchMLPRegressor,
            "torch_ft_transformer": TorchFTTransformerRegressor,
        },
    }


def get_model(model_name: str, task_type: str, params: dict):
    registry = get_registry()
    if model_name not in registry[task_type]:
        raise ValueError(f"Model '{model_name}' not valid for task '{task_type}'.")
    return registry[task_type][model_name](**params)


def balanced_weights_from_y(y_series):
    y_np = np.asarray(y_series)
    classes, counts = np.unique(y_np, return_counts=True)
    n_samples = len(y_np)
    n_classes = len(classes)
    weights = {
        cls: float(n_samples / (n_classes * count))
        for cls, count in zip(classes, counts, strict=False)
    }
    return weights


def sample_weight_from_y(y_series):
    class_weights = balanced_weights_from_y(y_series)
    y_np = np.asarray(y_series)
    return np.asarray([class_weights[v] for v in y_np], dtype=float)


def normalize_class_weight_keys(class_weight: dict, task_type: str):
    normalized = {}
    for key, value in class_weight.items():
        new_key = key
        if task_type == "binary":
            if isinstance(key, str):
                lk = key.strip().lower()
                if lk in {"true", "1"}:
                    new_key = True
                elif lk in {"false", "0"}:
                    new_key = False
        normalized[new_key] = value
    return normalized


def apply_imbalance_strategy(model_name: str, task_type: str, params: dict, y_train):
    new_params = dict(params)
    if "class_weight" in new_params and isinstance(new_params["class_weight"], dict):
        new_params["class_weight"] = normalize_class_weight_keys(
            new_params["class_weight"], task_type
        )

    if task_type not in {"binary", "categorical"}:
        return new_params

    weights = balanced_weights_from_y(y_train)
    if not weights:
        return new_params

    if model_name in {"rf", "lr", "svc"} and "class_weight" not in new_params:
        new_params["class_weight"] = {k.item() if hasattr(k, "item") else k: v for k, v in weights.items()}

    if model_name in {"torch_mlp", "torch_ft_transformer"}:
        classes_sorted = sorted(weights.keys())
        if task_type == "binary" and "pos_weight" not in new_params:
            neg_label, pos_label = classes_sorted[0], classes_sorted[-1]
            neg_w = weights[neg_label]
            pos_w = weights[pos_label]
            if neg_w > 0:
                new_params["pos_weight"] = float(pos_w / neg_w)
        elif task_type == "categorical" and "class_weights" not in new_params:
            new_params["class_weights"] = [float(weights[c]) for c in classes_sorted]

    return new_params


def threshold_metrics(y_true_bin, y_prob_pos, threshold, beta, fn_cost, fp_cost):
    y_pred = (y_prob_pos >= threshold).astype(int)
    tn, fp, fn, tp = confusion_matrix(y_true_bin, y_pred, labels=[0, 1]).ravel()

    recall = tp / (tp + fn) if (tp + fn) else 0.0
    precision = tp / (tp + fp) if (tp + fp) else 0.0
    specificity = tn / (tn + fp) if (tn + fp) else 0.0
    npv = tn / (tn + fn) if (tn + fn) else 0.0

    beta2 = beta * beta
    f_beta = (
        (1 + beta2) * precision * recall / (beta2 * precision + recall)
        if (precision + recall)
        else 0.0
    )
    cost = fn_cost * fn + fp_cost * fp
    return {
        "threshold": float(threshold),
        "tn": int(tn),
        "fp": int(fp),
        "fn": int(fn),
        "tp": int(tp),
        "recall": float(recall),
        "precision": float(precision),
        "specificity": float(specificity),
        "npv": float(npv),
        "f_beta": float(f_beta),
        "cost": float(cost),
    }


def select_threshold(y_true_bin, y_prob_pos, min_recall, beta, fn_cost, fp_cost):
    thresholds = np.linspace(0.01, 0.99, 199)
    metrics = [
        threshold_metrics(y_true_bin, y_prob_pos, t, beta, fn_cost, fp_cost)
        for t in thresholds
    ]
    feasible = [m for m in metrics if m["recall"] >= min_recall]

    if feasible:
        best = max(
            feasible,
            key=lambda m: (m["precision"], m["f_beta"], -m["cost"], m["specificity"]),
        )
        best["meets_recall_constraint"] = True
        return best

    best = max(
        metrics,
        key=lambda m: (m["recall"], m["precision"], m["f_beta"], -m["cost"]),
    )
    best["meets_recall_constraint"] = False
    return best


def evaluate_val(
    pipeline,
    X_val,
    y_val,
    task_type,
    min_recall,
    beta,
    fn_cost,
    fp_cost,
):
    y_pred = pipeline.predict(X_val)

    if task_type == "continuous":
        rmse = float(np.sqrt(mean_squared_error(y_val, y_pred)))
        return {
            "sort_key": rmse,
            "metric_name": "rmse",
            "display_score": rmse,
            "details": {"rmse": rmse},
        }

    y_prob = (
        pipeline.predict_proba(X_val) if hasattr(pipeline, "predict_proba") else None
    )

    if task_type == "binary" and y_prob is not None:
        pos_idx = 1 if y_prob.shape[1] > 1 else 0
        classes = getattr(pipeline.named_steps["model"], "classes_", [False, True])
        pos_label = classes[pos_idx]
        y_true_bin = (y_val == pos_label).astype(int).to_numpy()
        y_prob_pos = y_prob[:, pos_idx]

        try:
            auc = float(roc_auc_score(y_true_bin, y_prob_pos))
        except ValueError:
            auc = float("nan")

        op = select_threshold(y_true_bin, y_prob_pos, min_recall, beta, fn_cost, fp_cost)
        op["auc_roc"] = auc

        if op["meets_recall_constraint"]:
            sort_key = (2, op["precision"], op["f_beta"], auc if np.isfinite(auc) else -1, -op["cost"])
            return {
                "sort_key": sort_key,
                "metric_name": f"precision@recall>={min_recall:.2f}",
                "display_score": op["precision"],
                "details": op,
            }

        sort_key = (1, op["recall"], op["precision"], auc if np.isfinite(auc) else -1, -op["cost"])
        return {
            "sort_key": sort_key,
            "metric_name": f"max_recall_if_<{min_recall:.2f}",
            "display_score": op["recall"],
            "details": op,
        }

    f1m = float(f1_score(y_val, y_pred, average="macro"))
    return {
        "sort_key": f1m,
        "metric_name": "f1_macro",
        "display_score": f1m,
        "details": {"f1_macro": f1m},
    }


def is_better_eval(task_type, new_eval, best_eval):
    if best_eval is None:
        return True
    if task_type == "continuous":
        return new_eval["sort_key"] < best_eval["sort_key"]
    return new_eval["sort_key"] > best_eval["sort_key"]


def main():
    parser = argparse.ArgumentParser(description="Temporal Grid Search for MedModel")
    parser.add_argument("--target", required=True, help="Target column.")
    parser.add_argument("--data_config", default="data_config.json", help="Data configuration.")
    parser.add_argument("--search_space", default="search_space.json", help="Grid search parameters.")
    parser.add_argument("--output_file", default="best_parameters.json", help="Where to save the best configs.")
    parser.add_argument("--date_column", default="Date of surgery", help="Column used for temporal sorting.")
    parser.add_argument("--test_size", type=float, default=0.15, help="Held-out test set size (ignored during tuning).")
    parser.add_argument("--val_size", type=float, default=0.15, help="Validation set size (used to evaluate params).")
    parser.add_argument("--min_recall", type=float, default=0.90, help="Binary tuning constraint: minimum recall target.")
    parser.add_argument("--f_beta", type=float, default=2.0, help="Beta for F-beta during threshold optimization.")
    parser.add_argument("--fn_cost", type=float, default=5.0, help="Relative cost assigned to each false negative.")
    parser.add_argument("--fp_cost", type=float, default=1.0, help="Relative cost assigned to each false positive.")
    args = parser.parse_args()

    data_config = load_json(args.data_config)
    search_space = load_json(args.search_space)

    logger.info(f"Loading data from {data_config['input_file']}...")
    if str(data_config["input_file"]).endswith((".xlsx", ".xls")):
        df = pd.read_excel(data_config["input_file"])
    else:
        try:
            df = pd.read_csv(data_config["input_file"], encoding="utf-8")
        except UnicodeDecodeError:
            df = pd.read_csv(data_config["input_file"], encoding="latin1")

    col_output = args.target
    validate_required_columns(df, [col_output], "target")
    validate_required_columns(df, [args.date_column], "temporal split")

    configured_columns = [
        *data_config.get("input_features", []),
        *data_config.get("cols_string", []),
        *data_config.get("cols_date", []),
        *data_config.get("cols_multi", []),
    ]
    validate_required_columns(df, list(dict.fromkeys(configured_columns)), "data_config")

    df = df.dropna(subset=[col_output, args.date_column]).copy()

    task_type = infer_task_type(df[col_output])
    logger.info(f"Task: {task_type.upper()} | Target: {col_output}")

    if task_type == "binary":
        df[col_output] = to_bool_if_binary(df[col_output])

    df["_temp_date"] = pd.to_datetime(df[args.date_column], errors="coerce")
    df = (
        df.dropna(subset=["_temp_date"])
        .sort_values(by="_temp_date")
        .drop(columns=["_temp_date"])
    )

    n_total = len(df)
    n_test = int(n_total * args.test_size)
    n_val = int(n_total * args.val_size)
    n_train = n_total - n_val - n_test

    df_train = df.iloc[:n_train]
    df_val = df.iloc[n_train : n_train + n_val]

    X_train, y_train = df_train, df_train[col_output]
    X_val, y_val = df_val, df_val[col_output]

    logger.info(f"Temporal Split -> Train: {n_train}, Val: {n_val}, Test (held out): {n_test}")

    registry = get_registry()
    valid_models = set(registry[task_type].keys())

    best_overall_params = {}
    best_selection_details = {}

    for model_name, param_grid in search_space.items():
        logger.info(f"\n--- Tuning {model_name.upper()} ---")

        if model_name not in valid_models:
            logger.info(f"Skipping {model_name}: not valid for task '{task_type}'.")
            best_overall_params[model_name] = None
            best_selection_details[model_name] = {"status": "invalid_for_task"}
            continue

        use_scaler = model_name in {
            "lr",
            "ridge",
            "svc",
            "torch_mlp",
            "torch_ft_transformer",
        }
        grid = list(ParameterGrid(param_grid))

        best_params = None
        best_eval = None
        best_metric_name = ""
        fail_count = 0

        pbar = tqdm(grid, desc=f"Grid Search ({model_name})")
        for raw_params in pbar:
            try:
                params = apply_imbalance_strategy(model_name, task_type, raw_params, y_train)

                preprocessor = build_preprocessor(
                    data_config["input_features"],
                    data_config["cols_string"],
                    data_config["cols_date"],
                    data_config["cols_multi"],
                    use_scaler,
                )
                model = get_model(model_name, task_type, params)
                pipeline = Pipeline([("preprocess", preprocessor), ("model", model)])

                if model_name in {"torch_mlp", "torch_ft_transformer"}:
                    X_train_t = pipeline.named_steps["preprocess"].fit_transform(X_train, y_train)
                    X_val_t = pipeline.named_steps["preprocess"].transform(X_val)
                    pipeline.named_steps["model"].fit(X_train_t, y_train, eval_set=(X_val_t, y_val))
                else:
                    fit_kwargs = {}
                    if model_name == "hgb" and task_type in {"binary", "categorical"}:
                        fit_kwargs["model__sample_weight"] = sample_weight_from_y(y_train)
                    pipeline.fit(X_train, y_train, **fit_kwargs)

                eval_result = evaluate_val(
                    pipeline,
                    X_val,
                    y_val,
                    task_type,
                    min_recall=args.min_recall,
                    beta=args.f_beta,
                    fn_cost=args.fn_cost,
                    fp_cost=args.fp_cost,
                )

                if is_better_eval(task_type, eval_result, best_eval):
                    best_eval = eval_result
                    best_metric_name = eval_result["metric_name"]
                    best_params = params

                pbar.set_postfix({"Best": f"{best_eval['display_score']:.4f}" if best_eval else "n/a"})
            except Exception as e:
                fail_count += 1
                logger.warning(f"Failed with params {raw_params}: {e}")

        detail = {
            "status": "ok" if best_params is not None else "failed",
            "metric_name": best_metric_name,
            "metric_value": best_eval["display_score"] if best_eval else None,
            "failed_trials": fail_count,
        }
        if best_eval and "details" in best_eval:
            detail.update(best_eval["details"])

        logger.info(
            f"Best {model_name} Params: {best_params} | "
            f"Best Validation {best_metric_name.upper() if best_metric_name else 'N/A'}: "
            f"{best_eval['display_score']:.4f}" if best_eval else f"Best {model_name} Params: None"
        )

        best_overall_params[model_name] = best_params
        best_selection_details[model_name] = detail

    out_path = Path(args.output_file)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    with open(out_path, "w", encoding="utf-8") as f:
        json.dump(best_overall_params, f, indent=4)

    details_path = out_path.with_name(f"{out_path.stem}_selection.json")
    with open(details_path, "w", encoding="utf-8") as f:
        json.dump(best_selection_details, f, indent=4)

    logger.info(f"\nGrid Search Complete! Best parameters saved to '{out_path}'.")
    logger.info(f"Selection details saved to '{details_path}'.")


if __name__ == "__main__":
    main()