import streamlit as st
import pandas as pd
import pickle
import os
from pathlib import Path

# Set page title and configuration
st.set_page_config(page_title="ISA414: Energy Consumption Estimator", layout="wide")
st.title("ISA 414: Energy Consumption Estimator")

# Create a function to load models
@st.cache_resource
def load_models():
    # Check if model files exist in the current directory
    model_path = Path("model.pickle")
    encoder_path = Path("encoder.pickle")
    
    if not model_path.exists() or not encoder_path.exists():
        st.error("Model files not found. Please make sure 'model.pickle' and 'encoder.pickle' are in the same directory as this script.")
        st.stop()
    
    # Load model and encoding from files
    model = pickle.load(open("model.pickle", 'rb'))
    enc = pickle.load(open("encoder.pickle", 'rb'))
    return model, enc

# Try to load models
try:
    model, enc = load_models()
except Exception as e:
    st.error(f"Error loading models: {str(e)}")
    st.info("For demonstration purposes, you can continue using the app without the models.")
    model, enc = None, None

# Create sidebar with form inputs
with st.form("consumption_form"):
    st.subheader("Enter Client Information")
    
    # Create two columns for form layout
    col1, col2 = st.columns(2)
    
    with col1:
        age = st.number_input("Age", min_value=0, max_value=120, value=30)
        
        marital_status_options = ["Single", "Married", "Separated", "Divorced", "Widowed", "Other"]
        marital_status = st.selectbox("Marital Status", options=marital_status_options, index=0)
        marital_status_value = str(marital_status_options.index(marital_status) + 1)
        
        consumption = st.number_input("Day/Night Consumption", min_value=0.0, value=1.0, format="%.2f")
        
        income_options = ["Less than 10,000", "Between 10,000 and 30,000", 
                          "Between 40,000 and 60,000", "Between 60,000 and 100,000", "More than 100,000"]
        income = st.selectbox("Income Level", options=income_options, index=0)
        income_value = str(income_options.index(income) + 1)
        
        area_options = ["City centre", "City outskirts", "Rural", "Remote"]
        area = st.selectbox("Dwelling Area", options=area_options, index=0)
        area_value = str(area_options.index(area) + 1)
    
    with col2:
        children_options = ["No", "Yes"]
        children = st.selectbox("Children", options=children_options, index=0)
        children_value = False if children_options.index(children) == 0 else True
        
        solar_options = ["No", "Yes"]
        solar = st.selectbox("Solar Panel", options=solar_options, index=0)
        solar_value = False if solar_options.index(solar) == 0 else True
        
        sustainability_options = ["Favorable", "Undecided", "Unfavorable"]
        sustainability = st.selectbox("Attitude Towards Sustainability", options=sustainability_options, index=0)
        sustainability_value = str(sustainability_options.index(sustainability) + 1)
        
        tariff_options = ["Flat", "Time-of-Use", "Dynamic"]
        tariff = st.selectbox("Tariff", options=tariff_options, index=0)
        tariff_value = "Tariff " + str(tariff_options.index(tariff) + 1)
    
    # Submit button
    submitted = st.form_submit_button("Estimate Consumption")

# Function to make prediction
def forecast(age, marital_status, consumption, income, area, children, solar, sustainability, tariff):
    # Creating a DataFrame having a new client
    new_client = [age, marital_status, consumption, income, area, children, solar, sustainability, tariff]
    
    new_client = pd.DataFrame([new_client], 
                    columns=("Age", "MaritalStatus", "DayNightConsumption", "IncomeLevel", 
                             "DwellingArea", "HasChildren", "SolarRoof", "AttitudeSustainability", "Tariff"))
    
    # Categorical columns
    cat_columns = ["MaritalStatus", "IncomeLevel", "DwellingArea", "AttitudeSustainability", "Tariff"]
    
    # Using the previously created encoding and list of categorical variables to generate dummies
    dummies = enc.transform(new_client[cat_columns])
    
    # Creating a DataFrame of the dummies
    dummies_df = pd.DataFrame(dummies, columns=enc.get_feature_names_out(input_features=cat_columns))
    
    # Concatenating with the original data
    new_client = pd.concat([new_client, dummies_df], axis=1)
    
    # Dropping old columns
    new_client = new_client.drop(columns=cat_columns, axis=1)
    
    # Making prediction
    prediction = model.predict(new_client)[0]
    
    return prediction

# Display results if form is submitted
if submitted:
    try:
        if model is not None and enc is not None:
            # Make prediction
            prediction = forecast(
                age, 
                marital_status_value, 
                consumption, 
                income_value, 
                area_value, 
                children_value, 
                solar_value, 
                sustainability_value, 
                tariff_value
            )
            
            # Display result
            st.success(f"Estimated annual consumption: {prediction:.2f} kWh")
        else:
            st.warning("This is a demo mode. Models are not loaded.")
            st.info("In a real deployment, the prediction would be calculated here.")
    except Exception as e:
        st.error(f"Error making prediction: {str(e)}")

# Add information about the app
with st.expander("About this app"):
    st.write("""
    This app estimates energy consumption based on various client attributes.
    It uses a machine learning model to make predictions.
    
    To use the app:
    1. Enter the client information in the form
    2. Click 'Estimate Consumption' to see the prediction
    
    Note: This app requires 'model.pickle' and 'encoder.pickle' files to make actual predictions.
    """)