from flask import Flask, request, jsonify
import pandas as pd
import joblib
from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import StandardScaler
import numpy as np

app = Flask(__name__)

# Load the initial model and scaler
MODEL_PATH = 'garbage_model.pkl'
SCALER_PATH = 'scaler.pkl'
DATA_PATH = 'real_world_garbage_data.csv'

def retrain_model():
    """Internal function to refresh the model when data updates."""
    df = pd.read_csv(DATA_PATH)
    X = df[['Hour', 'Weight_kg', 'Distance_cm', 'Is_Weekend']]
    y = df['Status']
    
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    
    model = RandomForestClassifier(n_estimators=100, max_depth=5)
    model.fit(X_scaled, y)
    
    # Save the updated versions
    joblib.dump(model, MODEL_PATH)
    joblib.dump(scaler, SCALER_PATH)
    return model, scaler

# Load current versions on startup
try:
    model = joblib.load(MODEL_PATH)
    scaler = joblib.load(SCALER_PATH)
except:
    print("No model found. Training initial model...")
    model, scaler = retrain_model()

@app.route('/predict', methods=['POST'])
def predict():
    """Endpoint for real-time sensor predictions."""
    data = request.get_json()
    # Required format: [Hour, Weight_kg, Distance_cm, Is_Weekend]
    features = np.array([[data['hour'], data['weight'], data['distance'], data['weekend']]])
    
    # Scale and Predict
    features_scaled = scaler.transform(features)
    prediction = int(model.predict(features_scaled)[0])
    
    return jsonify({'status': prediction, 'message': 'Prediction successful'})

@app.route('/update', methods=['POST'])
def update_and_retrain():
    """Endpoint to add new data and trigger retraining."""
    global model, scaler
    new_data = request.get_json() # List of data points
    
    df_new = pd.DataFrame(new_data)
    df_new.to_csv(DATA_PATH, mode='a', header=False, index=False)
    
    # Trigger the retraining loop
    model, scaler = retrain_model()
    
    return jsonify({'message': 'Model retrained successfully with new data'})

if __name__ == '__main__':
    app.run(port=5000, debug=True)