app.py

"""
app.py
======
Causal Productivity Intelligence System
Streamlit Interactive Dashboard

LAYOUT
------
Page 1 — Profile & Prediction   : sliders + instant prediction + uncertainty
Page 2 — Counterfactual Planner : simulate interventions, see causal lifts
Page 3 — Statistical Deep Dive  : regression, causal, Bayesian outputs
Page 4 — LLM Research Insights  : AI-generated, research-grounded narration

Run: streamlit run app.py
"""

import os
import sys
import json
import warnings
import numpy as np
import pandas as pd
import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
import streamlit as st
from copy import deepcopy

warnings.filterwarnings("ignore")

# Add src to path
SRC_DIR  = os.path.join(os.path.dirname(__file__), "src")
BASE_DIR = os.path.dirname(__file__)
sys.path.insert(0, SRC_DIR)

from simulation import (
    BehavioralProfile, simulate_with_uncertainty,
    counterfactual_comparison, variable_sweep
)
from llm_insights import (
    generate_profile_insight, generate_counterfactual_insight,
    generate_causal_insight, generate_weekly_plan
)

# ─────────────────────────────────────────────────────────────────────────────
# PAGE CONFIG
# ─────────────────────────────────────────────────────────────────────────────

st.set_page_config(
    page_title="Causal Productivity Intelligence",
    page_icon="🧠",
    layout="wide",
    initial_sidebar_state="expanded",
)

# ── Custom CSS ────────────────────────────────────────────────────────────────
st.markdown("""
<style>
    .main { background-color: #F8FAFC; }
    .metric-card {
        background: white;
        border-radius: 12px;
        padding: 1.2rem 1.5rem;
        box-shadow: 0 1px 4px rgba(0,0,0,0.08);
        border-left: 4px solid #2563EB;
        margin-bottom: 0.8rem;
    }
    .metric-card.warn  { border-left-color: #F59E0B; }
    .metric-card.good  { border-left-color: #10B981; }
    .metric-card.bad   { border-left-color: #EF4444; }
    .section-header {
        font-size: 1.1rem; font-weight: 700;
        color: #1E293B; margin: 1.2rem 0 0.6rem;
        border-bottom: 2px solid #E2E8F0; padding-bottom: 0.3rem;
    }
    .research-note {
        background: #EFF6FF; border: 1px solid #BFDBFE;
        border-radius: 8px; padding: 0.7rem 1rem;
        font-size: 0.82rem; color: #1D4ED8; margin: 0.5rem 0;
    }
    .insight-box {
        background: #F0FDF4; border: 1px solid #BBF7D0;
        border-radius: 10px; padding: 1rem 1.2rem;
        font-size: 0.88rem; color: #14532D; margin: 0.6rem 0;
    }
    .stSlider label { font-weight: 600; color: #374151; }
    div[data-testid="stMetricValue"] { font-size: 2rem; font-weight: 800; }
</style>
""", unsafe_allow_html=True)


# ─────────────────────────────────────────────────────────────────────────────
# HELPERS
# ─────────────────────────────────────────────────────────────────────────────

def score_color(score: float) -> str:
    if score >= 72:  return "good"
    if score >= 60:  return "warn"
    return "bad"


def score_label(score: float) -> str:
    if score >= 75:  return "🟢 High"
    if score >= 63:  return "🟡 Moderate"
    return "🔴 Low"


def research_note(text: str):
    st.markdown(f'<div class="research-note">📚 {text}</div>', unsafe_allow_html=True)


def insight_box(text: str):
    st.markdown(f'<div class="insight-box">{text}</div>', unsafe_allow_html=True)


@st.cache_data(ttl=600)
def load_data():
    path = os.path.join(BASE_DIR, "data", "productivity_data.csv")
    if os.path.exists(path):
        return pd.read_csv(path, parse_dates=["date"])
    return None


@st.cache_data(ttl=600)
def load_outputs():
    """Load pre-computed model outputs if available."""
    out = {}
    out_dir = os.path.join(BASE_DIR, "data", "outputs")
    for fname in ["bayesian_summary.csv"]:
        path = os.path.join(out_dir, fname)
        if os.path.exists(path):
            out[fname.replace(".csv","")] = pd.read_csv(path)
    for fname in ["llm_insights.json"]:
        path = os.path.join(out_dir, fname)
        if os.path.exists(path):
            with open(path) as f:
                out["llm_insights"] = json.load(f)
    return out


def build_profile_from_sidebar() -> BehavioralProfile:
    """Build a BehavioralProfile from sidebar slider values."""
    sleep   = st.session_state.get("sleep_hours", 7.0)
    start   = st.session_state.get("start_hour", 9.0)
    hours   = st.session_state.get("hours_worked", 7.0)
    breaks  = st.session_state.get("break_frequency", 0.7)
    caff    = st.session_state.get("caffeine_intake", 2.0)
    caff_h  = st.session_state.get("caffeine_hour", 9.0)
    screen  = st.session_state.get("leisure_screen_time", 3.0)
    task    = st.session_state.get("task_type", "deep")
    stress  = st.session_state.get("stress_level", None)
    dow     = st.session_state.get("day_of_week", 2)
    return BehavioralProfile(
        sleep_hours=sleep, start_hour=start, hours_worked=hours,
        break_frequency=breaks, caffeine_intake=caff, caffeine_hour=caff_h,
        leisure_screen_time=screen, task_type=task,
        stress_level=stress if stress != 0 else None,
        day_of_week=dow,
    )


# ─────────────────────────────────────────────────────────────────────────────
# SIDEBAR — BEHAVIORAL PROFILE INPUT
# ─────────────────────────────────────────────────────────────────────────────

def render_sidebar():
    st.sidebar.image("https://img.icons8.com/fluency/96/brain.png", width=56)
    st.sidebar.title("🧠 Your Behavioral Profile")
    st.sidebar.markdown("*Adjust sliders to reflect your typical work day.*")
    st.sidebar.markdown("---")

    st.sidebar.markdown("**💤 Sleep**")
    st.sidebar.slider("Sleep Hours", 3.0, 11.0, 7.0, 0.25,
                       key="sleep_hours",
                       help="Previous night's sleep. Optimal: 7–9h (Van Dongen, 2003)")

    st.sidebar.markdown("**⏰ Timing**")
    st.sidebar.slider("Work Start Hour", 6.0, 21.0, 9.0, 0.5,
                       key="start_hour",
                       help="Hour you start working (24h). Peak alertness ~10AM (Anderson 2014)")
    day_names = ["Monday","Tuesday","Wednesday","Thursday","Friday","Saturday","Sunday"]
    st.sidebar.selectbox("Day of Week", day_names, index=2, key="day_name_sel")
    st.session_state["day_of_week"] = day_names.index(
        st.session_state.get("day_name_sel", "Wednesday"))

    st.sidebar.markdown("**☕ Caffeine**")
    st.sidebar.slider("Caffeine (cups)", 0.0, 5.0, 2.0, 0.25,
                       key="caffeine_intake",
                       help="Daily coffee/tea cups. Optimal: 2–3 (Lieberman, 2002)")
    st.sidebar.slider("First Coffee Hour", 6.0, 20.0, 9.0, 0.5,
                       key="caffeine_hour",
                       help="Late caffeine (>14:00) disrupts sleep (Drake, 2013)")

    st.sidebar.markdown("**🏢 Work Session**")
    st.sidebar.slider("Hours Worked", 1.0, 12.0, 7.0, 0.5,
                       key="hours_worked",
                       help="Session length. Performance degrades beyond ~6h")
    st.sidebar.slider("Breaks per Hour", 0.0, 2.0, 0.7, 0.1,
                       key="break_frequency",
                       help="Break frequency. ~1/90min optimal (Kleitman BRAC)")

    st.sidebar.markdown("**📱 Lifestyle**")
    st.sidebar.slider("Evening Screen Time (hrs)", 0.0, 8.0, 3.0, 0.25,
                       key="leisure_screen_time",
                       help="Non-work screen use. Causes stress & disrupts sleep")

    st.sidebar.markdown("**🎯 Task Type**")
    st.sidebar.radio("Primary Task", ["deep", "shallow"], index=0,
                     key="task_type",
                     help="Deep: cognitively demanding. Shallow: routine/admin (Newport, 2016)")

    st.sidebar.markdown("---")
    st.sidebar.markdown("**😰 Stress Override** *(0 = auto-compute)*")
    st.sidebar.slider("Stress Level (1–10)", 0.0, 10.0, 0.0, 0.5,
                       key="stress_level",
                       help="0 = computed from screen time + day. Override for manual control")


# ─────────────────────────────────────────────────────────────────────────────
# TAB 1 — PROFILE & PREDICTION
# ─────────────────────────────────────────────────────────────────────────────

def tab_profile(profile: BehavioralProfile, sim: dict):
    st.markdown("## 📊 Your Productivity Prediction")
    st.markdown("*Predictions based on the causal structural model, with full uncertainty quantification.*")

    # ── Top metrics ────────────────────────────────────────────────────────────
    col1, col2, col3, col4 = st.columns(4)
    score = sim["mean"]
    cls   = score_color(score)

    with col1:
        st.markdown(f"""<div class="metric-card {cls}">
            <div style="font-size:0.8rem;color:#64748B;font-weight:600">PREDICTED PRODUCTIVITY</div>
            <div style="font-size:2.4rem;font-weight:800;color:#1E293B">{score:.1f}<span style="font-size:1rem;color:#64748B">/100</span></div>
            <div style="font-size:0.8rem;color:#64748B">{score_label(score)}</div>
        </div>""", unsafe_allow_html=True)

    with col2:
        st.markdown(f"""<div class="metric-card warn">
            <div style="font-size:0.8rem;color:#64748B;font-weight:600">90% INTERVAL</div>
            <div style="font-size:1.5rem;font-weight:700;color:#1E293B">{sim['ci_5']:.0f} – {sim['ci_95']:.0f}</div>
            <div style="font-size:0.8rem;color:#64748B">Uncertainty range</div>
        </div>""", unsafe_allow_html=True)

    with col3:
        stress_cls = "bad" if sim["stress_level"] > 6 else ("warn" if sim["stress_level"] > 4 else "good")
        st.markdown(f"""<div class="metric-card {stress_cls}">
            <div style="font-size:0.8rem;color:#64748B;font-weight:600">STRESS LEVEL</div>
            <div style="font-size:1.8rem;font-weight:700;color:#1E293B">{sim['stress_level']:.1f}<span style="font-size:1rem;color:#64748B">/10</span></div>
            <div style="font-size:0.8rem;color:#64748B">Cohen PSS estimate</div>
        </div>""", unsafe_allow_html=True)

    with col4:
        mot_cls = "good" if sim["motivation_score"] > 6 else ("warn" if sim["motivation_score"] > 4 else "bad")
        st.markdown(f"""<div class="metric-card {mot_cls}">
            <div style="font-size:0.8rem;color:#64748B;font-weight:600">MOTIVATION</div>
            <div style="font-size:1.8rem;font-weight:700;color:#1E293B">{sim['motivation_score']:.1f}<span style="font-size:1rem;color:#64748B">/10</span></div>
            <div style="font-size:0.8rem;color:#64748B">SDT estimate (Deci & Ryan)</div>
        </div>""", unsafe_allow_html=True)

    st.markdown("")

    # ── Predictive distribution + component bars ───────────────────────────────
    col_a, col_b = st.columns([1.3, 1])

    with col_a:
        st.markdown('<div class="section-header">Predictive Distribution</div>', unsafe_allow_html=True)
        research_note("Uncertainty bands include coefficient estimation error + irreducible noise (σ≈4.5 pts)")

        fig = go.Figure()
        # KDE-like histogram
        fig.add_trace(go.Histogram(
            x=sim["draws"], nbinsx=60,
            marker_color="#2563EB", opacity=0.75,
            histnorm="probability density",
            name="Simulated draws"
        ))
        fig.add_vline(x=sim["mean"],  line_color="#1E293B", line_width=2.5,
                      annotation_text=f"Mean: {sim['mean']:.1f}")
        fig.add_vline(x=sim["ci_5"],  line_color="#DC2626", line_width=1.5,
                      line_dash="dash", annotation_text=f"5%: {sim['ci_5']:.0f}")
        fig.add_vline(x=sim["ci_95"], line_color="#DC2626", line_width=1.5,
                      line_dash="dash", annotation_text=f"95%: {sim['ci_95']:.0f}")
        fig.update_layout(
            height=280, margin=dict(l=20,r=20,t=30,b=30),
            xaxis_title="Productivity Score",
            showlegend=False,
            plot_bgcolor="white", paper_bgcolor="white",
        )
        st.plotly_chart(fig, use_container_width=True)

    with col_b:
        st.markdown('<div class="section-header">Component Contributions</div>', unsafe_allow_html=True)
        research_note("Each component maps to a specific research-grounded mechanism")

        comps = sim["components"]
        comp_labels = ["Sleep\nEffectiveness", "Circadian\nAlerting",
                       "Fatigue\nResistance", "Caffeine\nBoost"]
        comp_vals   = [
            comps["sleep_effect"]   * 100,
            comps["circadian"]      * 100,
            comps["fatigue"]        * 100,
            comps["caffeine_boost"] * 100 * 3  # scaled for visibility
        ]
        colors = ["#2563EB","#7C3AED","#10B981","#F59E0B"]
        fig2 = go.Figure(go.Bar(
            x=comp_labels, y=comp_vals,
            marker_color=colors, opacity=0.85,
            text=[f"{v:.0f}%" for v in comp_vals],
            textposition="outside",
        ))
        fig2.update_layout(
            height=280, margin=dict(l=10,r=10,t=30,b=30),
            yaxis_title="Effectiveness (%)",
            yaxis_range=[0, 115],
            plot_bgcolor="white", paper_bgcolor="white",
        )
        st.plotly_chart(fig2, use_container_width=True)

    # ── Dose-response curves ──────────────────────────────────────────────────
    st.markdown('<div class="section-header">Dose-Response: How Each Variable Affects You</div>',
                unsafe_allow_html=True)

    c1, c2, c3 = st.columns(3)

    def dose_fig(var, lo, hi, steps, label, color, profile=profile):
        vals  = np.linspace(lo, hi, steps)
        sweep = variable_sweep(profile, var, vals, n_boot=200)
        fig   = go.Figure()
        fig.add_trace(go.Scatter(
            x=sweep[var], y=sweep["mean"],
            line=dict(color=color, width=2.5), name="Mean"
        ))
        fig.add_trace(go.Scatter(
            x=pd.concat([sweep[var], sweep[var][::-1]]),
            y=pd.concat([sweep["ci_95"], sweep["ci_5"][::-1]]),
            fill="toself", fillcolor=color,
            opacity=0.15, line=dict(width=0), name="90% CI"
        ))
        cur_val = getattr(profile, var)
        cur_sim = variable_sweep(profile, var, [cur_val], n_boot=50)
        fig.add_trace(go.Scatter(
            x=[cur_val], y=[cur_sim["mean"].iloc[0]],
            mode="markers", marker=dict(size=12, color="#DC2626", symbol="circle"),
            name="Your value"
        ))
        fig.update_layout(
            height=230, margin=dict(l=20,r=10,t=10,b=30),
            xaxis_title=label, yaxis_title="Productivity",
            plot_bgcolor="white", paper_bgcolor="white",
            showlegend=False,
        )
        return fig

    with c1:
        st.markdown("**💤 Sleep Hours**")
        st.plotly_chart(dose_fig("sleep_hours", 4, 10, 25, "Sleep (hrs)", "#2563EB"),
                        use_container_width=True)
        research_note("Van Dongen (2003): non-linear drop below 6h")

    with c2:
        st.markdown("**☕ Caffeine Cups**")
        st.plotly_chart(dose_fig("caffeine_intake", 0, 5, 20, "Caffeine (cups)", "#F59E0B"),
                        use_container_width=True)
        research_note("Lieberman (2002): inverted-U, peak at 2-3 cups")

    with c3:
        st.markdown("**📱 Screen Time**")
        st.plotly_chart(dose_fig("leisure_screen_time", 0, 7, 20, "Screen (hrs)", "#EF4444"),
                        use_container_width=True)
        research_note("Screen time → stress → less sleep → lower productivity")


# ─────────────────────────────────────────────────────────────────────────────
# TAB 2 — COUNTERFACTUAL PLANNER
# ─────────────────────────────────────────────────────────────────────────────

def tab_counterfactual(profile: BehavioralProfile, sim: dict):
    st.markdown("## 🔮 Counterfactual Planner")
    st.markdown("*Simulate: 'What would my productivity be if I changed X?' — using Pearl's do-operator, not correlation.*")

    research_note(
        "do(sleep=8h) ≠ 'people who sleep 8h'. The do-operator INTERVENES and propagates "
        "effects through the full causal graph — accounting for downstream changes to "
        "stress, motivation, and task selection."
    )

    st.markdown("---")

    # ── Custom intervention builder ───────────────────────────────────────────
    st.markdown('<div class="section-header">Build Your Intervention</div>', unsafe_allow_html=True)
    col1, col2 = st.columns(2)

    with col1:
        new_sleep  = st.slider("🛌 New sleep hours",       3.0, 11.0, profile.sleep_hours, 0.25)
        new_caff   = st.slider("☕ New caffeine cups",      0.0, 5.0,  profile.caffeine_intake, 0.25)
        new_caff_h = st.slider("⏰ Coffee timing (hour)",  6.0, 20.0, profile.caffeine_hour, 0.5)

    with col2:
        new_screen = st.slider("📱 New screen time (hrs)", 0.0, 8.0,  profile.leisure_screen_time, 0.25)
        new_breaks = st.slider("☕ Breaks per hour",        0.0, 2.0,  profile.break_frequency, 0.1)
        new_task   = st.radio("🎯 Task type", ["deep", "shallow"],
                               index=0 if profile.task_type == "deep" else 1)

    # ── Predefined + custom scenarios ─────────────────────────────────────────
    interventions = {
        "Sleep +1h":             {"sleep_hours": min(profile.sleep_hours + 1, 11)},
        "Sleep 8h (optimal)":    {"sleep_hours": 8.0},
        "Morning coffee (9AM)":  {"caffeine_hour": 9.0},
        "Fewer coffees (2)":     {"caffeine_intake": 2.0},
        "No evening screen":     {"leisure_screen_time": 0.5},
        "Switch to deep work":   {"task_type": "deep"},
        "More breaks (1/hr)":    {"break_frequency": 1.0},
        "Custom scenario": {
            "sleep_hours":         new_sleep,
            "caffeine_intake":     new_caff,
            "caffeine_hour":       new_caff_h,
            "leisure_screen_time": new_screen,
            "break_frequency":     new_breaks,
            "task_type":           new_task,
        },
    }

    with st.spinner("⚙️  Running counterfactual simulations..."):
        cf_df = counterfactual_comparison(profile, interventions, n_boot=600)

    # ── Results table ─────────────────────────────────────────────────────────
    col_res, col_chart = st.columns([1, 1.6])

    with col_res:
        st.markdown('<div class="section-header">Simulation Results</div>', unsafe_allow_html=True)
        display_df = cf_df[["Scenario","Mean","Delta_mean","Delta_CI_5","Delta_CI_95"]].copy()
        display_df.columns = ["Scenario", "Score", "Δ Mean", "Δ CI 5%", "Δ CI 95%"]

        def color_delta(val):
            if isinstance(val, float):
                if val > 1:   return "color: #059669; font-weight:bold"
                if val < -1:  return "color: #DC2626; font-weight:bold"
            return ""

        styled = display_df.style.applymap(color_delta, subset=["Δ Mean"])
        st.dataframe(styled, height=310, use_container_width=True)

    with col_chart:
        st.markdown('<div class="section-header">Causal Lift vs Baseline</div>', unsafe_allow_html=True)
        plot_df = cf_df[cf_df["Scenario"] != "Current Baseline"].copy()
        plot_df = plot_df.sort_values("Delta_mean", ascending=True)

        colors = ["#10B981" if d > 0 else "#EF4444" for d in plot_df["Delta_mean"]]

        fig = go.Figure()
        fig.add_trace(go.Bar(
            x=plot_df["Delta_mean"],
            y=plot_df["Scenario"],
            orientation="h",
            marker_color=colors, opacity=0.85,
            error_x=dict(
                type="data",
                symmetric=False,
                array=plot_df["Delta_CI_95"] - plot_df["Delta_mean"],
                arrayminus=plot_df["Delta_mean"] - plot_df["Delta_CI_5"],
                color="#1E293B", thickness=1.5, width=5,
            ),
            text=[f"{v:+.1f}" for v in plot_df["Delta_mean"]],
            textposition="outside",
        ))
        fig.add_vline(x=0, line_color="#374151", line_dash="dash", line_width=1)
        fig.update_layout(
            height=310, margin=dict(l=20,r=40,t=20,b=30),
            xaxis_title="Δ Productivity Score (causal lift)",
            plot_bgcolor="white", paper_bgcolor="white",
            showlegend=False,
        )
        st.plotly_chart(fig, use_container_width=True)

    # ── Best intervention callout ──────────────────────────────────────────────
    best = cf_df[cf_df["Scenario"] != "Current Baseline"].nlargest(1, "Delta_mean").iloc[0]
    if best["Delta_mean"] > 0:
        insight_box(
            f"🏆 <b>Highest impact intervention:</b> <i>{best['Scenario']}</i> — "
            f"predicted gain of <b>+{best['Delta_mean']:.1f} pts</b> "
            f"[90% CI: +{best['Delta_CI_5']:.1f} to +{best['Delta_CI_95']:.1f}]"
        )


# ─────────────────────────────────────────────────────────────────────────────
# TAB 3 — STATISTICAL DEEP DIVE
# ─────────────────────────────────────────────────────────────────────────────

def tab_statistics(df):
    st.markdown("## 📈 Statistical Deep Dive")
    st.markdown("*Explore the raw data and model outputs underlying your predictions.*")

    if df is None:
        st.warning("Dataset not found. Run `python src/data_generation.py` first.")
        return

    sub1, sub2, sub3 = st.tabs(["📊 Data Explorer", "🔗 Causal DAG", "🎯 Bayesian Posteriors"])

    with sub1:
        st.markdown('<div class="section-header">Dataset Overview</div>', unsafe_allow_html=True)
        col1, col2, col3 = st.columns(3)
        col1.metric("Observations", f"{len(df):,}")
        col2.metric("Deep Work Sessions", f"{(df['task_type']=='deep').sum():,}")
        col3.metric("Mean Productivity", f"{df['productivity_score'].mean():.1f}")

        # Scatter: sleep vs productivity
        c1, c2 = st.columns(2)
        with c1:
            fig = px.scatter(df, x="sleep_hours", y="productivity_score",
                             color="task_type",
                             color_discrete_map={"deep":"#2563EB","shallow":"#EF4444"},
                             opacity=0.5, trendline="lowess",
                             labels={"sleep_hours":"Sleep Hours",
                                     "productivity_score":"Productivity Score"},
                             title="Sleep vs Productivity (by task type)")
            fig.update_layout(height=320, plot_bgcolor="white", paper_bgcolor="white",
                              margin=dict(l=10,r=10,t=40,b=10))
            st.plotly_chart(fig, use_container_width=True)
            research_note("Van Dongen (2003): non-linear sleep-performance curve")

        with c2:
            fig2 = px.scatter(df, x="stress_level", y="productivity_score",
                              color="caffeine_intake",
                              color_continuous_scale="YlOrRd",
                              opacity=0.5,
                              labels={"stress_level":"Stress Level",
                                      "productivity_score":"Productivity",
                                      "caffeine_intake":"Caffeine (cups)"},
                              title="Stress vs Productivity (caffeine coloured)")
            fig2.update_layout(height=320, plot_bgcolor="white", paper_bgcolor="white",
                               margin=dict(l=10,r=10,t=40,b=10))
            st.plotly_chart(fig2, use_container_width=True)
            research_note("Cohen (1983): stress impairs working memory via cortisol")

        # Correlation heatmap
        st.markdown('<div class="section-header">Correlation Matrix</div>', unsafe_allow_html=True)
        num_cols = ["sleep_hours","stress_level","caffeine_intake","motivation_score",
                    "start_hour","hours_worked","break_frequency",
                    "leisure_screen_time","productivity_score"]
        corr = df[num_cols].corr()
        fig3 = px.imshow(corr, text_auto=".2f", aspect="auto",
                         color_continuous_scale="RdBu_r", zmin=-1, zmax=1,
                         title="Variable Correlation Matrix")
        fig3.update_layout(height=420, margin=dict(l=10,r=10,t=40,b=10))
        st.plotly_chart(fig3, use_container_width=True)
        research_note("Correlations ≠ causation. See Causal DAG tab for directed effects.")

    with sub2:
        st.markdown('<div class="section-header">Causal DAG — v2 Enriched</div>',
                    unsafe_allow_html=True)
        dag_path = os.path.join(BASE_DIR, "data", "outputs", "causal_dag_v2.png")
        if os.path.exists(dag_path):
            st.image(dag_path, use_column_width=True)
        else:
            st.info("Run `python src/causal_analysis.py` to generate the DAG.")

        st.markdown('<div class="section-header">Effect Decomposition</div>',
                    unsafe_allow_html=True)
        decomp_path = os.path.join(BASE_DIR, "data", "outputs", "effect_decomposition.png")
        if os.path.exists(decomp_path):
            st.image(decomp_path, use_column_width=True)
            research_note(
                "26.6% of sleep's benefit is mediated via motivation (Deci & Ryan 2000 SDT). "
                "The direct path (74%) acts through raw cognitive capacity (Van Dongen 2003)."
            )

        st.markdown('<div class="section-header">Sensitivity Analysis</div>',
                    unsafe_allow_html=True)
        sens_path = os.path.join(BASE_DIR, "data", "outputs", "sensitivity_analysis.png")
        if os.path.exists(sens_path):
            st.image(sens_path, use_column_width=True)
            research_note(
                "Rosenbaum (2002) Γ-bounds: at Γ=1.25, unmeasured chronotype confounding "
                "could explain the result. This quantifies our sensitivity to unmeasured variables."
            )

    with sub3:
        bayes_path = os.path.join(BASE_DIR, "data", "outputs", "posterior_distributions.png")
        cri_path   = os.path.join(BASE_DIR, "data", "outputs", "credible_intervals.png")
        ppc_path   = os.path.join(BASE_DIR, "data", "outputs", "posterior_predictive_check.png")

        st.markdown('<div class="section-header">Posterior Distributions</div>',
                    unsafe_allow_html=True)
        research_note(
            "Bayesian 95% Credible Interval = '95% probability the true value lies here' "
            "(Gelman et al. 2013). Richer than frequentist confidence intervals for decision-making."
        )
        if os.path.exists(bayes_path):
            st.image(bayes_path, use_column_width=True)

        c1, c2 = st.columns(2)
        with c1:
            if os.path.exists(cri_path):
                st.markdown('<div class="section-header">Credible Intervals</div>',
                            unsafe_allow_html=True)
                st.image(cri_path, use_column_width=True)
        with c2:
            if os.path.exists(ppc_path):
                st.markdown('<div class="section-header">Posterior Predictive Check</div>',
                            unsafe_allow_html=True)
                st.image(ppc_path, use_column_width=True)
                research_note("Gelman et al. (2013): PPC validates that the model generates "
                              "data resembling what was actually observed.")


# ─────────────────────────────────────────────────────────────────────────────
# TAB 4 — LLM INSIGHTS
# ─────────────────────────────────────────────────────────────────────────────

def tab_llm(profile: BehavioralProfile, sim: dict, cf_df):
    st.markdown("## 🤖 Research-Backed AI Insights")
    st.markdown(
        "*Each insight is grounded in the peer-reviewed research that motivated the model. "
        "The AI cites only concepts verified in the statistical analysis.*"
    )

    causal_results = {
        "total":  {"theta": 1.89, "ci_lo": 1.52, "ci_hi": 2.27, "p": 0.000001},
        "direct": {"theta": 1.66, "ci_lo": 1.32, "ci_hi": 2.00, "p": 0.000001},
        "mediation": {"prop_mediated": 0.266},
    }

    insight_type = st.selectbox(
        "Choose insight type:",
        ["🔍 Profile Analysis",
         "🔮 Counterfactual Explanation",
         "📐 Causal Effect Interpretation",
         "📅 Weekly Action Plan"],
        index=0,
    )

    if st.button("✨ Generate Insight", type="primary"):
        with st.spinner("Consulting the research literature..."):
            if insight_type == "🔍 Profile Analysis":
                text = generate_profile_insight(profile.to_dict(), sim)
            elif insight_type == "🔮 Counterfactual Explanation":
                text = generate_counterfactual_insight(cf_df, profile.to_dict())
            elif insight_type == "📐 Causal Effect Interpretation":
                text = generate_causal_insight(causal_results)
            else:
                text = generate_weekly_plan(profile.to_dict(), cf_df)

        st.markdown("---")
        st.markdown(text)

    st.markdown("---")

    # Pre-generated insights (if available)
    outputs = load_outputs()
    if "llm_insights" in outputs:
        with st.expander("📚 View pre-generated insights (from last full run)"):
            for key, text in outputs["llm_insights"].items():
                st.markdown(f"### {key.replace('_',' ').title()}")
                st.markdown(text)
                st.markdown("---")

    # Research citations
    with st.expander("📖 Research references used in this analysis"):
        st.markdown("""
| Finding | Reference |
|---|---|
| Sleep deprivation & cognitive decline | Van Dongen et al. (2003, SLEEP) |
| Morning alertness peak | Anderson et al. (2014, Psych Science) |
| Sleep-circadian interaction | Dijk & Czeisler (1995, J Neuroscience) |
| Deep work & deliberate practice | Newport (2016); Ericsson et al. (1993) |
| Motivation as mediator | Deci & Ryan (2000, SDT) |
| Stress & cortisol pathway | Cohen et al. (1983); Salleh (2008) |
| Caffeine alertness & sleep | Lieberman (2002); Drake et al. (2013) |
| Break timing | Kleitman BRAC; Lorist et al. (2005) |
| Causal identification | Pearl (2009); Chernozhukov et al. (2018) |
| Sensitivity analysis | Rosenbaum (2002) |
| Bayesian methods | Gelman et al. (2013, BDA3) |
""")


# ─────────────────────────────────────────────────────────────────────────────
# MAIN APP
# ─────────────────────────────────────────────────────────────────────────────

def main():
    render_sidebar()

    # Build profile from sidebar
    profile = build_profile_from_sidebar()

    # Run simulation
    with st.spinner(""):
        sim = simulate_with_uncertainty(profile, n_boot=800, seed=42)

    # Run quick counterfactual for LLM tab
    quick_interventions = {
        "Sleep +1h":          {"sleep_hours": min(profile.sleep_hours + 1, 11)},
        "Morning coffee":     {"caffeine_hour": 9.0},
        "Less screen (1h)":   {"leisure_screen_time": 1.0},
        "Full package":       {"sleep_hours": 8.0, "caffeine_hour": 9.0,
                               "leisure_screen_time": 1.0, "task_type": "deep"},
    }
    cf_df = counterfactual_comparison(profile, quick_interventions, n_boot=400)

    # Header
    st.markdown("""
    <div style="background: linear-gradient(135deg, #1E293B, #2563EB);
                border-radius: 14px; padding: 1.5rem 2rem; margin-bottom: 1.5rem;">
        <h1 style="color:white;margin:0;font-size:1.8rem">
            🧠 Causal Productivity Intelligence System
        </h1>
        <p style="color:#93C5FD;margin:0.4rem 0 0;font-size:0.95rem">
            Statistical Modelling • Bayesian Methods • Causal Inference • LLM Insights
        </p>
    </div>
    """, unsafe_allow_html=True)

    # Tabs
    tab1, tab2, tab3, tab4 = st.tabs([
        "📊 Profile & Prediction",
        "🔮 Counterfactual Planner",
        "📈 Statistical Deep Dive",
        "🤖 AI Research Insights",
    ])

    df = load_data()

    with tab1: tab_profile(profile, sim)
    with tab2: tab_counterfactual(profile, sim)
    with tab3: tab_statistics(df)
    with tab4: tab_llm(profile, sim, cf_df)


if __name__ == "__main__":
    main()