import streamlit as st
import pandas as pd
import numpy as np
import joblib
from huggingface_hub import hf_hub_download
import json
import os
import plotly.graph_objects as go

# =============================================================================
# STREAMLIT CONFIG
# =============================================================================

st.set_page_config(
    page_title="Predictive Maintenance",
    layout="wide",
    page_icon="🛠️"
)

# =============================================================================
# LOAD MODEL
# =============================================================================

@st.cache_resource
def load_artifacts():
    try:
        repo_id = "SharleyK/predictive-maintenance-model"
        token = os.getenv("HF_TOKEN")

        model = joblib.load(hf_hub_download(repo_id, "best_model.pkl", token=token))
        scaler = joblib.load(hf_hub_download(repo_id, "scaler.pkl", token=token))

        metadata_path = hf_hub_download(repo_id, "metadata.json", token=token)
        metadata = json.load(open(metadata_path))

        return model, scaler, metadata, True

    except Exception as e:
        st.error(f"❌ Error loading model: {e}")
        return None, None, {}, False

model, scaler, metadata, MODEL_LOADED = load_artifacts()
MODEL_NAME = metadata.get("model_name", "Predictive Maintenance Model")

# =============================================================================
# FEATURE ENGINEERING
# =============================================================================

def engineer_features(df):
    df = df.copy()

    df["temp_pressure_ratio"] = df["coolant_temp"] / (df["coolant_pressure"] + 1e-6)
    df["oil_temp_pressure_ratio"] = df["lub_oil_temp"] / (df["lub_oil_pressure"] + 1e-6)

    df["pressure_diff_oil_coolant"] = df["lub_oil_pressure"] - df["coolant_pressure"]
    df["pressure_diff_fuel_oil"] = df["fuel_pressure"] - df["lub_oil_pressure"]

    df["temp_diff_oil_coolant"] = df["lub_oil_temp"] - df["coolant_temp"]

    df["avg_temp"] = (df["lub_oil_temp"] + df["coolant_temp"]) / 2
    df["avg_pressure"] = (
        df["lub_oil_pressure"] +
        df["fuel_pressure"] +
        df["coolant_pressure"]
    ) / 3

    df["high_rpm"] = (df["engine_rpm"] > 934).astype(int)
    df["high_temp"] = (df["coolant_temp"] > 82.9).astype(int)
    df["operating_stress"] = df["high_rpm"] * df["high_temp"]
    df["low_oil_pressure"] = (df["lub_oil_pressure"] < 2.52).astype(int)

    return df

# =============================================================================
# PREDICTION
# =============================================================================

def predict(df):

    df_eng = engineer_features(df)
    X_scaled = scaler.transform(df_eng)

    pred = model.predict(X_scaled)[0]

    if hasattr(model, "predict_proba"):
        confidence = model.predict_proba(X_scaled)[0][pred]
    else:
        confidence = 0.85

    return pred, confidence

# =============================================================================
# VISUALS
# =============================================================================

def gauge_chart(confidence):

    fig = go.Figure(go.Indicator(
        mode="gauge+number",
        value=confidence * 100,
        title={'text': "Confidence %"},
        gauge={
            'axis': {'range': [0, 100]},
            'steps': [
                {'range': [0, 50], 'color': 'lightgreen'},
                {'range': [50, 70], 'color': 'yellow'},
                {'range': [70, 85], 'color': 'orange'},
                {'range': [85, 100], 'color': 'red'}
            ],
        }
    ))

    fig.update_layout(height=300)
    return fig

def radar_chart(params):

    categories = [
        'RPM',
        'Oil Pressure',
        'Fuel Pressure',
        'Coolant Pressure',
        'Oil Temp',
        'Coolant Temp'
    ]

    values = [
        params["engine_rpm"]/25,
        params["lub_oil_pressure"]/0.08,
        params["fuel_pressure"]/0.22,
        params["coolant_pressure"]/0.08,
        params["lub_oil_temp"],
        params["coolant_temp"]/2
    ]

    fig = go.Figure()

    fig.add_trace(go.Scatterpolar(
        r=values,
        theta=categories,
        fill='toself'
    ))

    return fig

# =============================================================================
# RECOMMENDATIONS
# =============================================================================

def recommendations(pred, rpm, oil_p, fuel_p, coolant_p, oil_t, coolant_t):

    rec = []

    if pred == 1:
        rec.append("### 🚨 Maintenance Recommended")

        if oil_p < 2.5:
            rec.append("• Check oil pump / leaks")

        if coolant_t > 85:
            rec.append("• Cooling system inspection needed")

        if rpm < 600:
            rec.append("• Engine load or fuel intake check")

        if fuel_p > 8:
            rec.append("• Possible fuel regulator issue")

    else:
        rec.append("### ✅ Engine Healthy")
        rec.append("• Continue routine maintenance")
        rec.append("• Monitor coolant weekly")

    return "\n".join(rec)

# =============================================================================
# UI
# =============================================================================

st.title("🛠️ Predictive Maintenance Dashboard")
st.caption("Real-time AI Engine Health Monitoring")

st.info(f"Model: {MODEL_NAME}")

st.divider()

col1, col2, col3 = st.columns(3)

with col1:
    engine_rpm = st.slider("Engine RPM", 50, 2500, 800)

with col2:
    lub_oil_pressure = st.slider("Lub Oil Pressure", 0.0, 8.0, 3.3)

with col3:
    fuel_pressure = st.slider("Fuel Pressure", 0.0, 22.0, 6.5)

col4, col5, col6 = st.columns(3)

with col4:
    coolant_pressure = st.slider("Coolant Pressure", 0.0, 8.0, 2.3)

with col5:
    lub_oil_temp = st.slider("Lub Oil Temp", 70.0, 95.0, 77.0)

with col6:
    coolant_temp = st.slider("Coolant Temp", 60.0, 200.0, 78.0)

# =============================================================================
# PREDICT BUTTON
# =============================================================================

if st.button("🔮 Predict Engine Condition"):

    df = 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
    }])

    pred, conf = predict(df)

    colA, colB = st.columns(2)

    with colA:
        if pred == 1:
            st.error("🚨 Faulty Engine Detected")
        else:
            st.success("✅ Engine Healthy")

        st.metric("Confidence", f"{conf*100:.2f}%")

    with colB:
        st.plotly_chart(gauge_chart(conf), use_container_width=True)

    st.divider()

    st.subheader("📊 Engine Parameter Profile")

    st.plotly_chart(
        radar_chart(df.iloc[0]),
        use_container_width=True
    )

    st.divider()

    st.subheader("💡 Recommendations")

    st.markdown(
        recommendations(
            pred,
            engine_rpm,
            lub_oil_pressure,
            fuel_pressure,
            coolant_pressure,
            lub_oil_temp,
            coolant_temp
        )
    )