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

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

def get_ist_time():
    ist = pytz.timezone('Asia/Kolkata')
    return datetime.now(ist).strftime("%Y-%m-%d %H:%M:%S %Z")

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

def load_model():
    if not os.path.exists(MODEL_PATH):
        return train_and_save_model()
    return load(MODEL_PATH)

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
    timestamp = get_ist_time()
    HISTORY.append({
        "Temperature": temp,
        "Duration": duration,
        "Risk": pred,
        "Confidence": round(score, 2),
        "Timestamp": timestamp
    })
    return pred, round(score, 2), timestamp

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)