app/models/dosing/train_supervised.py

"""Dosing supervised training placeholder."""

from __future__ import annotations

import pickle
from pathlib import Path

import numpy as np
from sklearn.ensemble import RandomForestRegressor
from sklearn.multioutput import MultiOutputRegressor

from app.common.enums import Difficulty
from app.models.dosing.dose_policy_features import build_dose_features
from app.simulator.patient_generator import generate_patient_profile

def train_dosing_surrogate(dataset_size: int) -> dict[str, float | str]:
    feature_rows: list[list[float]] = []
    target_rows: list[list[float]] = []
    for i in range(dataset_size):
        difficulty = Difficulty.HARD if i % 2 == 0 else Difficulty.MEDIUM
        patient = generate_patient_profile(seed=5000 + i, difficulty=difficulty)
        drug = patient.medications[0].drug if patient.medications else "warfarin_like"
        feats = build_dose_features(patient, drug)
        organ = feats.get("organ_stress", 0.0)
        interaction = feats.get("interaction_load", 0.0)
        adherence = feats.get("adherence", 0.7)
        target_attainment = max(0.0, min(1.0, 0.72 + adherence * 0.15 - interaction * 0.2))
        toxicity = max(0.0, min(1.0, 0.15 + organ * 0.5 + interaction * 0.25))
        underdose = max(0.0, min(1.0, 0.25 + (1.0 - adherence) * 0.35 + max(0.0, 0.4 - interaction) * 0.1))
        measurement_need = max(toxicity, underdose)
        feature_rows.append(list(feats.values()))
        target_rows.append([target_attainment, toxicity, underdose, measurement_need])

    x = np.array(feature_rows, dtype=float)
    y = np.array(target_rows, dtype=float)
    model = MultiOutputRegressor(RandomForestRegressor(n_estimators=80, random_state=42))
    model.fit(x, y)
    preds = model.predict(x)
    mae = float(np.mean(np.abs(preds - y)))

    artifact = {
        "model": model,
        "feature_keys": list(build_dose_features(generate_patient_profile(seed=1, difficulty=Difficulty.EASY), "warfarin_like").keys()),
        "target_keys": ["target_attainment", "toxicity_proxy", "underdose_proxy", "measurement_need"],
    }
    path = Path("outputs/models/dose_model.pkl")
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("wb") as f:
        pickle.dump(artifact, f)
    return {
        "dataset_size": float(dataset_size),
        "status": "trained",
        "train_mae": round(mae, 4),
        "model_path": str(path),
    }