import streamlit as st
import pandas as pd
import joblib
import numpy as np
from tensorflow.keras.models import load_model 

MODEL_PATH = 'src/lstm_energy_prediction.h5' 
SCALER_PATH = 'src/scaler_energy.joblib'

FEATURES = ['Global_active_power', 'Global_reactive_power', 'Voltage',
            'Global_intensity', 'Sub_metering_2', 'Sub_metering_1', 'Sub_metering_3']

@st.cache_resource
def load_assets():
    try:
        # Load the LSTM model
        model = load_model(MODEL_PATH) 
        # Load the scaler
        scaler = joblib.load(SCALER_PATH)
        return model, scaler
    except Exception as e:
        st.error(f"Error loading assets. Ensure '{MODEL_PATH}' (H5 model) and '{SCALER_PATH}' (Scaler) are uploaded. Error: {e}")
        return None, None

def prepare_input_and_predict(model, scaler, input_data):
    # 1. Convert input data to a numpy array for scaling
    raw_input = np.array([input_data[f] for f in FEATURES]).reshape(1, -1)
    
    # 2. Scale the input data
    scaled_input = scaler.transform(raw_input) 
    
    # 3. Reshape for LSTM: [Samples, Time Steps, Features] = [1, 1, 7]
    lstm_input = scaled_input[:, :len(FEATURES)].reshape(1, 1, len(FEATURES))
    
    # 4. Predict the scaled output (var1(t) = Global_active_power at time t)
    scaled_prediction = model.predict(lstm_input)
    
    # 5. Inverse Transform the prediction
    dummy_inverse = np.zeros(raw_input.shape)
    dummy_inverse[0, 0] = scaled_prediction[0, 0]
    
    actual_prediction = scaler.inverse_transform(dummy_inverse)[0, 0]
    
    return actual_prediction

# --- Streamlit Interface ---
st.set_page_config(page_title="Energy Prediction App", layout="centered")
st.title("⚡ Hourly Energy Consumption Prediction (LSTM)")
st.markdown("Enter the consumption values from the last hour to predict the next hour's **Global Active Power**.")


model, scaler = load_assets()

if model is not None and scaler is not None:
    st.sidebar.header("Last Hour's Consumption Data")

    # --- INPUT WIDGETS (Matching the 7 features) ---
    # These input keys MUST match the FEATURES list.
    global_active_power = st.sidebar.number_input("Global Active Power:", value=0.5, step=0.01)
    global_reactive_power = st.sidebar.number_input("Global Reactive Power:", value=0.08, step=0.01)
    voltage = st.sidebar.number_input("Voltage:", value=240.0, step=0.1)
    global_intensity = st.sidebar.number_input("Global Intensity:", value=2.5, step=0.1)
    sub_metering_2 = st.sidebar.number_input("Sub Metering 2:", value=0.0, step=0.1)
    sub_metering_1 = st.sidebar.number_input("Sub Metering 1:", value=0.0, step=0.1)
    sub_metering_3 = st.sidebar.number_input("Sub Metering 3:", value=17.0, step=0.1)

    # Collect inputs
    input_data = {
        'Global_active_power': global_active_power,
        'Global_reactive_power': global_reactive_power,
        'Voltage': voltage,
        'Global_intensity': global_intensity,
        'Sub_metering_2': sub_metering_2,
        'Sub_metering_1': sub_metering_1,
        'Sub_metering_3': sub_metering_3
    }

    if st.button("Predict Next Hour's Power"):
        with st.spinner('Predicting...'):
            predicted_power = prepare_input_and_predict(model, scaler, input_data)
            
            if predicted_power is not None:
                st.success("Prediction Successful!")
                st.markdown(f"### Predicted Next Hour's Global Active Power:")
                st.markdown(f"**{predicted_power:,.2f} kW**")
                
                st.info("Note: Prediction is based on the complex feature transformation (lag features) used during model training.")