import streamlit as st
import joblib
import numpy as np
import pandas as pd

# ------------------------------
# 1. Load Model & Scaler
# ------------------------------
MODEL_PATH = "src/voting_model.joblib"
SCALER_PATH = "src/scaler.joblib"

@st.cache_resource
def load_artifacts():
    try:
        model = joblib.load(MODEL_PATH)
        scaler = joblib.load(SCALER_PATH)
        return model, scaler
    except Exception as e:
        st.error(f"Error loading model: {e}")
        return None, None

model, scaler = load_artifacts()

# ------------------------------
# 2. Feature Engineering Function
# ------------------------------
def engineer_features(df):
    data = df.copy()
    epsilon = 1e-6  
    
    data['stiffness_low'] = data['Columns 1-3 I mm4*10^6'] / (data['Floor height m']**3 + epsilon)
    data['stiffness_high'] = data['Columns 4-6 I mm4*10^6'] / (data['Floor height m']**3 + epsilon)
    data['stiffness_diff'] = data['stiffness_low'] - data['stiffness_high']

    data['strength_ratio'] = data['Column fy Mpa'] / (data['Beam fy Mpa'] + epsilon)

    data['total_height'] = data['Number of floors'] * data['Floor height m']
    data['total_width'] = data['Spans'] * data['Span width m']
    data['slenderness'] = data['total_height'] / (data['total_width'] + epsilon)
    data['total_area_low'] = data['Columns 1-3 A mm2'] * (data['Spans'] + 1)

    data['seismic_power'] = data['PGA g'] * data['Magnitude']
    data['fault_attenuation'] = data['Magnitude'] / np.log1p(data['Distance to fault km'])

    data['total_mass'] = data['Floor mass kg'] * data['Number of floors']
    data['base_shear'] = data['total_mass'] * data['PGA g']
    
    data = data.replace([np.inf, -np.inf], 0)
    return data

# ------------------------------
# 3. User Input Section
# ------------------------------
st.title("🏢 Seismic Building Safety Prediction")
st.write("""
This AI model predicts the **Maximum Interstorey Drift (mm)** a building might experience during an earthquake.
Lower drift values generally indicate safer buildings.
""")

RAW_INPUTS = [
    "Column fy Mpa", "Beam fy Mpa", 
    "Columns 1-3 I mm4*10^6", "Columns 4-6 I mm4*10^6",
    "Columns 1-3 A mm2", "Columns 4-6 A mm2", "Beam I mm4*10^6",
    "Spans", "Number of floors", "Floor height m", "Span width m",
    "LLRS tributary width m", "Floor mass kg", "Facade Load kN/m",
    "PGA g", "Magnitude", "Distance to fault km", "Period s",
    "Final Dead Load", "Final Live Load"
]

input_data = {}

with st.form("input_form"):
    c1, c2 = st.columns(2)
    
    for i, col_name in enumerate(RAW_INPUTS):
        if i % 2 == 0:
            with c1:
                input_data[col_name] = st.number_input(col_name, value=0.0, format="%.2f")
        else:
            with c2:
                input_data[col_name] = st.number_input(col_name, value=0.0, format="%.2f")

    soil_class = st.selectbox("Soil Class", ["A (Other)", "B", "C"])
    input_data['soil_class__B'] = 1.0 if soil_class == "B" else 0.0
    input_data['soil_class__C'] = 1.0 if soil_class == "C" else 0.0

    submitted = st.form_submit_button("Predict Max Drift")

# ------------------------------
# 4. Prediction Logic
# ------------------------------
if submitted:
    if model is None or scaler is None:
        st.error("Model could not be loaded!")
    else:
        try:
            df_raw = pd.DataFrame([input_data])

            df_engineered = engineer_features(df_raw)

            X_scaled = scaler.transform(df_engineered)

            log_pred = model.predict(X_scaled)

            real_pred = np.expm1(log_pred)[0]
            
            st.divider()
            st.success(f"### 🎯 Predicted Drift: **{real_pred:.2f} mm**")
            
            if real_pred < 10:
                st.info("Risk: 🟢 Low (Safe)")
            elif real_pred < 50:
                st.warning("Risk: 🟡 Medium (Potential Damage)")
            else:
                st.error("Risk: 🔴 High (Collapse Hazard)")

        except Exception as e:
            st.error(f"Calculation error: {e}")
            st.write("Please ensure all values are entered correctly.")