risk_model.py

import os
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from joblib import dump, load

MODEL_PATH = "heating_model.pkl"
DATA_PATH = "mantle_training.csv"
HISTORY = []

# Train the model and save
def train_and_save_model():
    data = pd.read_csv(DATA_PATH)
    X = data[["temperature", "duration"]]
    y = data["risk_level"]
    model = RandomForestClassifier()
    model.fit(X, y)
    dump(model, MODEL_PATH)
    return model

# Load model safely
def load_model():
    if not os.path.exists(MODEL_PATH):
        return train_and_save_model()
    return load(MODEL_PATH)

# Load once at startup
model = load_model()

def predict_risk(temp, duration):
    global model
    pred = model.predict([[temp, duration]])[0]
    score = max(model.predict_proba([[temp, duration]])[0]) * 100
    HISTORY.append({"Temperature": temp, "Duration": duration, "Risk": pred, "Confidence": round(score, 2)})
    return pred, round(score, 2)

def retrain_model():
    try:
        data = pd.read_csv(DATA_PATH)
        X = data[["temperature", "duration"]]
        y = data["risk_level"]
        clf = RandomForestClassifier().fit(X, y)
        dump(clf, MODEL_PATH)
        global model
        model = clf
        return "✅ Model retrained successfully"
    except Exception as e:
        return f"❌ Error: {str(e)}"

def get_history_df():
    return pd.DataFrame(HISTORY)