import joblib
import pandas as pd
from huggingface_hub import hf_hub_download
import sys
import streamlit as st
import datetime # Added import for datetime

# Initialize model to None
model = None

# Define repository ID and model path in repo
repo_id = "Srinivas1969/dl-capstone-dataset"
path_in_repo = "model/tuned_random_forest.joblib"

# Download the model file
try:
    model_path_local = hf_hub_download(repo_id=repo_id, repo_type='dataset', filename=path_in_repo)
    # st.write(f"Model downloaded to: {model_path_local}") # For debugging in Streamlit, can be removed
except Exception as e:
    st.error(f"Error downloading model from Hugging Face Hub: {e}")
    st.error("Please ensure 'HF_TOKEN' is correctly set as an environment variable if the repo is private.")
    sys.exit(1) # Exit with an error code

# Load the downloaded model
try:
    model = joblib.load(model_path_local)
    # st.write("Model loaded successfully.") # For debugging in Streamlit, can be removed
except Exception as e:
    st.error(f"Error loading the model: {e}")
    sys.exit(1) # Exit with an error code

def predict_engine_condition(engine_rpm, lub_oil_pressure, fuel_pressure, coolant_pressure, lub_oil_temp, coolant_temp):
    """
    Predicts the engine condition (0 = Normal, 1 = Faulty) based on input sensor readings.

    Args:
        engine_rpm (int): The number of revolutions per minute (RPM) of the engine.
        lub_oil_pressure (float): The pressure of the lubricating oil in the engine (bar/kPa).
        fuel_pressure (float): The pressure at which fuel is supplied to the engine (bar/kPa).
        coolant_pressure (float): The pressure of the engine coolant (bar/kPa).
        lub_oil_temp (float): The temperature of the lubricating oil (°C).
        coolant_temp (float): The temperature of the engine coolant (°C).

    Returns:
        int: Predicted engine condition (0 for Normal, 1 for Faulty).
    """
    # Defensive check: Ensure the model is loaded before making predictions.
    if model is None:
        raise RuntimeError("Model is not loaded. Cannot make predictions.")

    # Create a DataFrame from the input parameters
    input_data = pd.DataFrame([[engine_rpm, lub_oil_pressure, fuel_pressure, coolant_pressure, lub_oil_temp, coolant_temp]],
                              columns=['Engine rpm', 'Lub oil pressure', 'Fuel pressure', 'Coolant pressure', 'lub oil temp', 'Coolant temp'])

    # Make prediction
    prediction = model.predict(input_data)

    return int(prediction[0])

# Streamlit Interface
st.set_page_config(layout="wide")
st.title('Engine Condition Predictor')
st.write(datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')) # Display current date and time
st.write('Enter sensor readings to predict engine condition (0=Normal, 1=Faulty).')

# Input widgets for sensor readings (using ranges from df.describe() and common sense)
with st.sidebar:
    st.header("Sensor Readings")
    engine_rpm = st.number_input('Engine RPM', min_value=60, max_value=2300, value=791, step=10, help="Revolutions per minute")
    lub_oil_pressure = st.number_input('Lub Oil Pressure (bar/kPa)', min_value=0.0, max_value=8.0, value=3.3, step=0.1, format="%.2f")
    fuel_pressure = st.number_input('Fuel Pressure (bar/kPa)', min_value=0.0, max_value=22.0, value=6.65, step=0.1, format="%.2f")
    coolant_pressure = st.number_input('Coolant Pressure (bar/kPa)', min_value=0.0, max_value=8.0, value=2.33, step=0.1, format="%.2f")
    lub_oil_temp = st.number_input('Lub Oil Temperature (°C)', min_value=70.0, max_value=90.0, value=77.64, step=0.1, format="%.2f")
    coolant_temp = st.number_input('Coolant Temperature (°C)', min_value=60.0, max_value=200.0, value=78.42, step=0.1, format="%.2f")

if st.button('Predict Engine Condition'):
    try:
        prediction = predict_engine_condition(
            engine_rpm,
            lub_oil_pressure,
            fuel_pressure,
            coolant_pressure,
            lub_oil_temp,
            coolant_temp
        )

        if prediction == 0:
            st.success('Predicted Engine Condition: Normal (0)')
        else:
            st.error('Predicted Engine Condition: Faulty (1)')

        st.write("### Input Data:")
        st.write(pd.DataFrame({
            'Engine rpm': [engine_rpm],
            'Lub oil pressure': [lub_oil_pressure],
            'Fuel pressure': [fuel_pressure],
            'Coolant pressure': [coolant_pressure],
            'lub oil temp': [lub_oil_temp],
            'Coolant temp': [coolant_temp]
        }))

    except Exception as e:
        st.error(f"An error occurred during prediction: {e}")