import gradio as gr
import firebase_admin
from firebase_admin import credentials, db
import pandas as pd
from datetime import datetime
import joblib
import logging
import plotly.express as px
import io

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Load model and scaler
rf_model = joblib.load('random_forest_air_quality.pkl')
scaler = joblib.load('scaler_air_quality.pkl')
logger.info("Model and scaler loaded successfully.")

# Calibration R0 values
mq3_r0 = 17200
mq9_r0 = 21940

# Firebase setup
cred = credentials.Certificate('smell-42bf5-firebase-adminsdk-fbsvc-06651f2cef.json')
firebase_admin.initialize_app(cred, {
    'databaseURL': 'https://smell-42bf5-default-rtdb.asia-southeast1.firebasedatabase.app/'
})
logger.info("Firebase initialized successfully.")

# Rs/R0 calculation
def rs_r0(adc, r0, vcc=5.0, rl=5000, adc_vmax=3.3):
    vout = (adc / 4095) * adc_vmax
    if vout == 0:
        return float('inf')
    rs = rl * ((vcc / vout) - 1)
    return rs / r0

# Global state
seen_keys = set()
logs_df = pd.DataFrame(columns=["timestamp", "mq3_value", "mq9_value", "prediction"])

# Main inference logic
def fetch_and_infer():
    global seen_keys, logs_df
    ref = db.reference('/sensor_data')
    all_data = ref.get()
    logger.info("Fetched data from Firebase.")

    if not all_data:
        return "No data available", logs_df

    latest_prediction = "No new prediction"
    for key, value in sorted(all_data.items()):
        if key in seen_keys:
            continue

        try:
            mq3_value = int(value['mq3_value'])
            mq9_value = int(value['mq9_value'])
            timestamp = value['timestamp']
        except (KeyError, ValueError):
            continue

        try:
            timestamp_obj = datetime.strptime(timestamp, "%Y-%m-%d %H:%M:%S")
            hour = timestamp_obj.hour
        except ValueError:
            timestamp_obj = datetime.now()
            hour = timestamp_obj.hour
            timestamp = timestamp_obj.strftime("%Y-%m-%d %H:%M:%S")

        mq3_rs = rs_r0(mq3_value, mq3_r0)
        mq9_rs = rs_r0(mq9_value, mq9_r0)

        new_data = pd.DataFrame([{
            'mq3_rs_r0': mq3_rs,
            'mq9_rs_r0': mq9_rs,
            'hour': hour
        }])

        new_data_scaled = scaler.transform(new_data)
        prediction = rf_model.predict(new_data_scaled)[0]
        latest_prediction = prediction

        new_row = pd.DataFrame([{
            "timestamp": timestamp,
            "mq3_value": mq3_value,
            "mq9_value": mq9_value,
            "prediction": prediction
        }])

        logs_df = pd.concat([logs_df, new_row]).tail(50).reset_index(drop=True)
        seen_keys.add(key)

    return latest_prediction, logs_df

def update_ui():
    prediction, log_data = fetch_and_infer()
    alert = ""
    if prediction in ["foul", "very foul"]:
        alert = f"⚠️ Immediate action required! Please clean or sanitize. Detected condition: {prediction.upper()}"
    return prediction, alert, log_data

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# 🌫️ Real-Time Air Quality Monitoring Dashboard")
    gr.Markdown("Device ID: **SMELL-42BF5** | Location: **Public Restroom**")

    alert_box = gr.Markdown("")
    prediction_text = gr.Textbox(label="Latest Prediction", interactive=False)
    logs_table = gr.Dataframe(headers=["timestamp", "mq3_value", "mq9_value", "prediction"], datatype="str")

    interface = gr.Interface(
        fn=update_ui,
        inputs=[],
        outputs=[prediction_text, alert_box, logs_table],
        live=True
    )
    interface.render()

# Launch Gradio App
demo.launch(share=True)