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

# ---------------------------------
# 1. Setup & Model Loading
# ---------------------------------
st.set_page_config(page_title="DistractIQ", page_icon="🧠", layout="wide")

@st.cache_resource
def load_trained_model():
    # Checking common paths for the .pkl file
    paths = ['src/distraction_model.pkl', 'distraction_model.pkl']
    for path in paths:
        try:
            return joblib.load(path)
        except:
            continue
    return None

trained_model = load_trained_model()

# ---------------------------------
# 2. Sidebar: Manual Data Entry
# ---------------------------------
st.sidebar.title("📥 Input Data")
st.sidebar.markdown("Enter your daily habits below to analyze your focus.")

st.sidebar.header("Daily Habits")
# All ranges set to 0-24
study = st.sidebar.slider("Study Hours", 0, 24, 5)
phone = st.sidebar.slider("Total Phone Usage (Hours)", 0, 24, 4)
social = st.sidebar.slider("Social Media Usage (Hours)", 0, 24, 2)
sleep = st.sidebar.slider("Sleep Hours", 0, 24, 7)
notifications = st.sidebar.slider("Daily Notifications", 0, 1000, 80)

# ---------------------------------
# 3. Math Validation (The "Mathing" Fixes)
# ---------------------------------
total_hours = study + phone + sleep
math_error = False

# Paradox 1: More than 24 hours in a day
if total_hours > 24:
    st.sidebar.error(f"⚠️ Paradox! Total hours ({total_hours}h) exceed 24h limit.")
    math_error = True

# Paradox 2: Social media cannot exceed total phone time
if social > phone:
    st.sidebar.warning("⚠️ Social media hours adjusted (cannot exceed total phone).")
    social = phone

# ---------------------------------
# 4. Main Dashboard
# ---------------------------------
st.title("🧠 DistractIQ")
st.subheader("AI Digital Distraction Risk Predictor + Focus Optimizer")

if st.button("🚀 Analyze My Focus"):
    if math_error:
        st.error("Please fix the 24-hour limit error in the sidebar first.")
    else:
        # Use Trained Model if loaded, otherwise fallback logic
        if trained_model:
            features = pd.DataFrame([[study, phone, social, sleep, 5]], 
                                    columns=['study_hours_per_day', 'phone_usage_hours', 'social_media_hours', 'sleep_hours', 'stress_level'])
            prediction = trained_model.predict(features)[0]
        else:
            prediction = 0 if phone > 7 else 1 # Fallback: 0=Distracted, 1=Focused

        # Calculate Scores
        # Formula: Weights phone and social heavily
        score = int(((phone * 6) + (social * 8) + (notifications * 0.08)) / 2)
        score = min(100, max(0, score))
        
        risk = min(100, score + 10) if prediction == 0 else max(0, score - 10)
        productivity = max(0, min(100, 100 - score + (study * 2)))

        # --- Display Results ---
        st.divider()
        c1, c2, c3 = st.columns(3)
        c1.metric("Distraction Score", f"{score}/100")
        c2.metric("Risk Forecast", f"{risk}%")
        c3.metric("Productivity", f"{int(productivity)}%")

        st.subheader("Focus Analysis")
        st.progress(score/100)

        # Persona
        st.subheader("Focus Persona")
        if productivity > 80:
            st.success("🏆 Deep Work Ninja")
        elif phone > 8:
            st.error("📱 Digital Drifter")
        elif notifications > 150:
            st.warning("⚡ Chronic Multitasker")
        else:
            st.info("⚖️ Balanced Performer")

        # Drivers
        st.subheader("Top Distraction Drivers")
        drivers = []
        if social > 4: drivers.append(f"High social media usage ({social}h)")
        if notifications > 120: drivers.append(f"Notification Overload ({notifications} pings)")
        if phone > 7: drivers.append(f"Excessive phone usage ({phone}h)")
        if sleep < 6: drivers.append("Sleep deficit")
        
        for d in drivers if drivers else ["No major distraction drivers detected"]:
            st.write("•", d)

        # What-if Simulator
        st.divider()
        st.subheader("Focus Optimizer")
        st.info(f"By reducing social media by 1 hour, your potential productivity could rise by **{int(social * 4)}%**.")

        # Focus XP
        st.subheader("Focus XP")
        st.progress(productivity/100)
        if productivity > 75: st.success("LEVEL: Focus Master")
        elif productivity > 50: st.info("LEVEL: Focus Builder")
        else: st.warning("LEVEL: Focus Beginner")

st.markdown("---")
st.caption("Developed by Team Ai Projecttt")