import os
import time
import pandas as pd
import numpy as np
import joblib
import requests
import streamlit as st
from streamlit_autorefresh import st_autorefresh

# Auto-refresh every 5 seconds
st_autorefresh(interval=5000, key="refresh")

# Load model
@st.cache_resource
def load_model():
    return joblib.load("rf_model.pkl")

model = load_model()

# Supabase config
SUPABASE_URL = os.environ["SUPABASE_URL"]
SUPABASE_KEY = os.environ["SUPABASE_KEY"]
TABLE = "smart_meter_readings_1year"

# Initialize session state
if "row_index" not in st.session_state:
    st.session_state.row_index = 0
if "history" not in st.session_state:
    st.session_state.history = pd.DataFrame()

# Fetch all data
@st.cache_data
def fetch_all_data():
    url = f"{SUPABASE_URL}/rest/v1/{TABLE}?select=*&order=timestamp.asc"
    headers = {
        "apikey": SUPABASE_KEY,
        "Authorization": f"Bearer {SUPABASE_KEY}"
    }
    r = requests.get(url, headers=headers)
    if r.ok:
        return pd.DataFrame(r.json())
    else:
        st.error(f"❌ Error fetching data: {r.status_code}")
        return pd.DataFrame()

df_all = fetch_all_data()

# Debug sidebar
st.sidebar.title("🛠 Debug Info")
st.sidebar.write("Row index:", st.session_state.row_index)
st.sidebar.write("Total rows:", len(df_all))
if not df_all.empty and st.session_state.row_index < len(df_all):
    st.sidebar.write("Next row:", df_all.iloc[st.session_state.row_index].to_dict())

# Get next row
def get_next_row():
    if st.session_state.row_index < len(df_all):
        row = df_all.iloc[[st.session_state.row_index]]
        st.session_state.row_index += 1
        return row
    return pd.DataFrame()

# Feature engineering
def engineer(df):
    # Handle timestamp
    if pd.api.types.is_numeric_dtype(df["timestamp"]):
        df["datetime"] = pd.to_datetime(df["timestamp"], unit="s")
    else:
        df["datetime"] = pd.to_datetime(df["timestamp"])

    df["hour_of_day"] = df["datetime"].dt.hour
    df["lag_30min"] = df["power_consumption_kwh"].shift(1)
    df["lag_1h"] = df["power_consumption_kwh"].shift(2)
    df["rolling_avg_1h"] = df["power_consumption_kwh"].rolling(2).mean()
    df["rolling_avg_2h"] = df["power_consumption_kwh"].rolling(4).mean()
    df["is_weekend"] = df["datetime"].dt.weekday >= 5
    df["hour_sin"] = np.sin(2 * np.pi * df["hour_of_day"] / 24)
    df["hour_cos"] = np.cos(2 * np.pi * df["hour_of_day"] / 24)

    # One-hot encode property_type and region
    df = pd.get_dummies(df, columns=["property_type", "region"], drop_first=False)

    # Ensure all expected features exist
    expected_features = [
        'lag_30min', 'lag_1h',
        'rolling_avg_1h', 'rolling_avg_2h',
        'hour_of_day', 'is_weekend',
        'hour_sin', 'hour_cos',
        'temperature_c', 'ev_owner', 'solar_installed',
        'property_type_commercial', 'property_type_residential',
        'region_north', 'region_south', 'region_east', 'region_west'
    ]

    for col in expected_features:
        if col not in df.columns:
            df[col] = 0

    return df

# UI layout
st.title("⚡ Gridflux: Live Smart-Meter Forecast")

placeholder_chart = st.empty()
placeholder_metric = st.empty()

new_row = get_next_row()

if not new_row.empty:
    st.session_state.history = pd.concat([st.session_state.history, new_row], ignore_index=True)
    df_feat = engineer(st.session_state.history).dropna()

    if not df_feat.empty:
        latest_input = df_feat.iloc[[-1]][[
            'lag_30min', 'lag_1h',
            'rolling_avg_1h', 'rolling_avg_2h',
            'hour_of_day', 'is_weekend',
            'hour_sin', 'hour_cos',
            'temperature_c', 'ev_owner', 'solar_installed',
            'property_type_commercial', 'property_type_residential',
            'region_north', 'region_south', 'region_east', 'region_west'
        ]]

        prediction = model.predict(latest_input)[0]

        # Show outputs
        chart_df = st.session_state.history.copy()
        chart_df["datetime"] = pd.to_datetime(chart_df["timestamp"])
        chart_df.set_index("datetime", inplace=True)

        placeholder_chart.line_chart(chart_df["power_consumption_kwh"])
        placeholder_metric.metric("🔮 Predicted Power Usage (kWh)", f"{prediction:.3f}")
else:
    st.success("✅ All data processed.")