# Import necessary libraries
import numpy as np
import joblib  # For loading the serialized model
import pandas as pd  # For data manipulation
from flask import Flask, request, jsonify  # For creating the Flask API
import os # For accessing environment variables if needed
import logging # For logging

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# Initialize the Flask application
predictive_maintenance_api = Flask("Predictive Maintenance API")

# Define the model path
# Use the correct path as defined in the notebook
#MODEL_PATH = "backend_files/Predictive_Maintenance_XGB_Tuned_model_v1_0.joblib"
MODEL_PATH = "Predictive_Maintenance_XGB_Tuned_model_v1_0.joblib" # Corrected path

# Load the trained machine learning model
model = None
try:
    model = joblib.load(MODEL_PATH)
    logging.info(f"Model loaded successfully from {MODEL_PATH}")
except Exception as e:
    logging.error(f"Error loading model: {e}")

# Define a route for the home page (GET request)
@predictive_maintenance_api.get('/')
def home():
    """
    This function handles GET requests to the root URL ('/') of the API.
    It returns a simple welcome message.
    """
    logging.info("Home page accessed.")
    return "Welcome to the Predictive Maintenance API for Engine Health!"

# Define an endpoint for engine condition prediction (POST request)
@predictive_maintenance_api.post('/v1/engine_condition_prediction')
def predict_engine_condition():
    """
    This function handles POST requests to the '/v1/engine_condition_prediction' endpoint.
    It expects a JSON payload containing engine sensor data and returns
    the predicted engine condition (0 for Normal, 1 for Faulty) as a JSON response.
    """
    logging.info("Prediction request received.")
    if model is None:
        logging.error("Model not loaded when prediction request came. Returning 500.")
        return jsonify({"error": "Model not loaded. Please check server logs."}), 500

    # Get the JSON data from the request body
    engine_data = request.get_json()
    logging.info(f"Received engine data: {engine_data}")

    # Extract relevant features from the JSON data
    # Features: 'Engine_RPM', 'Lub_Oil_Pressure', 'Fuel_Pressure', 'Coolant_Pressure', 'Lub_Oil_Temperature', 'Coolant_Temperature'
    try:
        sample_data = {
            'Engine_RPM': [engine_data['Engine_RPM']],
            'Lub_Oil_Pressure': [engine_data['Lub_Oil_Pressure']],
            'Fuel_Pressure': [engine_data['Fuel_Pressure']],
            'Coolant_Pressure': [engine_data['Coolant_Pressure']],
            'Lub_Oil_Temperature': [engine_data['Lub_Oil_Temperature']],
            'Coolant_Temperature': [engine_data['Coolant_Temperature']]
        }
        logging.debug(f"Extracted sample data: {sample_data}")
    except KeyError as e:
        logging.error(f"Missing data for feature: {e}. Returning 400.")
        return jsonify({"error": f"Missing data for feature: {e}. Please provide all required sensor readings."}), 400
    except Exception as e:
        logging.error(f"Unexpected error during data extraction: {e}. Returning 400.")
        return jsonify({"error": f"An unexpected error occurred during data processing: {e}"}), 400


    # Convert the extracted data into a Pandas DataFrame
    # Ensure the order of columns matches the training data
    input_df = pd.DataFrame(sample_data)

    # Make prediction
    try:
        prediction_proba = model.predict_proba(input_df)
        predicted_class = model.predict(input_df)[0]
        logging.info(f"Prediction made: class={predicted_class}, probabilities={prediction_proba}")
    except Exception as e:
        logging.error(f"Error during model prediction: {e}. Returning 500.")
        return jsonify({"error": f"Error during model prediction: {e}"}), 500

    # Map the predicted class to a readable string
    condition_map = {0: "Normal", 1: "Faulty"}
    predicted_condition_str = condition_map.get(predicted_class, "Unknown")

    # Return the prediction as a JSON response
    return jsonify({
        "predicted_engine_condition_class": int(predicted_class), # Ensure it's a standard Python int
        "predicted_engine_condition_label": predicted_condition_str,
        "probability_normal": round(float(prediction_proba[0][0]), 4),
        "probability_faulty": round(float(prediction_proba[0][1]), 4)
    })

# To run the Flask app (for local testing)
if __name__ == '__main__':
    # You might want to get the port from an environment variable for deployment
    port = int(os.environ.get('PORT', 5000))
    predictive_maintenance_api.run(host='0.0.0.0', port=port, debug=True)