# mean_inference_app.py
# Streamlit ≥1.32 — Accessible, minimal color design

import streamlit as st
import pandas as pd
import numpy as np
from scipy.stats import norm, t
import io

# ---------- Page Config ----------
st.set_page_config(
    page_title="Inference for Means",
    page_icon="📈",
    layout="centered",
    initial_sidebar_state="collapsed"
)

# ---------- Accessible CSS - Compact for embedding ----------
st.markdown("""
<style>
    /* Maximize vertical space usage */
    .block-container { 
        padding-top: 0.5rem !important; 
        padding-bottom: 0.5rem !important;
        max-width: 100% !important;
    }
    
    /* Tighter element spacing */
    .element-container { margin-bottom: 0.3rem !important; }
    .stRadio > div { margin-bottom: 0 !important; }
    .stSelectSlider { padding-top: 0 !important; padding-bottom: 0 !important; }
    h2 { margin-top: 0 !important; margin-bottom: 0.3rem !important; font-size: 1.4rem !important; }
    h3 { margin-top: 0.3rem !important; margin-bottom: 0.2rem !important; font-size: 1.1rem !important; }
    p { margin-bottom: 0.3rem !important; }
    .stNumberInput { margin-bottom: 0 !important; }
    
    /* Compact info boxes */
    .stAlert { padding: 0.4rem 0.7rem !important; margin: 0.2rem 0 !important; }
    
    /* Clean result box - compact */
    .result-box {
        background: #f8f9fa;
        border: 2px solid #dee2e6;
        padding: 0.6rem;
        border-radius: 8px;
        text-align: center;
        margin: 0.2rem 0;
    }
    .result-box .label {
        font-size: 0.75rem;
        color: #6c757d;
        margin-bottom: 0.1rem;
    }
    .result-box .value {
        font-size: 1.2rem;
        font-weight: 600;
        color: #212529;
    }
    
    /* Decision boxes - compact */
    .decision-box {
        padding: 0.6rem;
        border-radius: 8px;
        text-align: center;
        margin: 0.3rem 0;
        font-weight: 600;
        font-size: 1rem;
    }
    .decision-reject {
        background: #fff;
        border: 3px solid #212529;
    }
    .decision-accept {
        background: #fff;
        border: 3px dashed #6c757d;
    }
    
    /* Compact dividers */
    hr { margin: 0.5rem 0 !important; border: none; border-top: 1px solid #dee2e6; }
    
    /* Compact metrics */
    [data-testid="stMetricValue"] { font-size: 1.1rem !important; }
    [data-testid="stMetricLabel"] { font-size: 0.75rem !important; }
    
    /* Compact expander */
    .streamlit-expanderHeader { padding: 0.3rem 0 !important; font-size: 0.9rem !important; }
    
    /* Hide Streamlit branding for cleaner embed */
    #MainMenu {visibility: hidden;}
    footer {visibility: hidden;}
    header {visibility: hidden;}
</style>
""", unsafe_allow_html=True)

# ---------- Title ----------
st.markdown("## 📈 Inference for Means")

# ---------- INPUTS ----------
col1, col2 = st.columns(2)
with col1:
    inf_type = st.radio("Inference type:", ["One-Sample Mean", "Two-Sample Mean (independent)"],
                       key="inf_type")
with col2:
    analysis_type = st.radio("Analysis:", ["Confidence Interval", "Hypothesis Test"],
                            key="analysis_type")

# Distribution choice
dist_choice = st.radio(
    "Distribution:",
    ["z (large sample)", "t (small sample, σ unknown)"],
    key="dist_choice",
    horizontal=True
)
dist_is_z = dist_choice.startswith("z")

st.markdown("---")

# ---------- Sample Data Inputs ----------
if inf_type == "One-Sample Mean":
    col1, col2, col3 = st.columns(3)
    with col1:
        n = st.number_input("Sample size (n)", min_value=2, step=1, value=30, key="n")
    with col2:
        xbar = st.number_input("Sample mean (x̄)", format="%.4f", value=0.0, key="xbar")
    with col3:
        s = st.number_input("Sample std dev (s)", min_value=0.0001, format="%.4f", value=1.0, key="s")
    
    se = s / np.sqrt(n)
    st.info(f"Standard Error: **SE = s/√n = {se:.4f}**")

else:  # Two-Sample
    col1, col2 = st.columns(2)
    
    with col1:
        st.subheader("Group 1")
        n1 = st.number_input("Sample size n₁", min_value=2, step=1, value=30, key="n1")
        xbar1 = st.number_input("Sample mean x̄₁", format="%.4f", value=0.0, key="xbar1")
        s1 = st.number_input("Sample std dev s₁", min_value=0.0001, format="%.4f", value=1.0, key="s1")
    
    with col2:
        st.subheader("Group 2")
        n2 = st.number_input("Sample size n₂", min_value=2, step=1, value=30, key="n2")
        xbar2 = st.number_input("Sample mean x̄₂", format="%.4f", value=0.0, key="xbar2")
        s2 = st.number_input("Sample std dev s₂", min_value=0.0001, format="%.4f", value=1.0, key="s2")

st.markdown("---")

# ---------- CI / HT specific controls ----------
if analysis_type == "Confidence Interval":
    conf_level = st.select_slider("Confidence level:",
                                  options=[0.90, 0.95, 0.99], value=0.95,
                                  format_func=lambda x: f"{x*100:.0f}%", key="conf_level")
else:
    col1, col2 = st.columns(2)
    with col1:
        alpha = st.select_slider("Significance level (α):", 
                                options=[0.01, 0.05, 0.10], value=0.05, key="alpha")
    with col2:
        if inf_type == "One-Sample Mean":
            mu0 = st.number_input("Null mean (μ₀):", format="%.4f", value=0.0, key="mu0")
    
    if inf_type == "One-Sample Mean":
        alt = st.radio("Alternative hypothesis (H₁):",
                       ["μ ≠ μ₀ (Two-sided)", "μ > μ₀ (Right-sided)", "μ < μ₀ (Left-sided)"],
                       horizontal=True, key="alt_one")
    else:
        alt = st.radio("Alternative hypothesis (H₁):",
                       ["μ₁ ≠ μ₂ (Two-sided)", "μ₁ > μ₂ (Right-sided)", "μ₁ < μ₂ (Left-sided)"],
                       horizontal=True, key="alt_two")

# ---------- RUN ----------
run = st.button("▶ Run Analysis", type="primary", use_container_width=True)

# ---------- Helper function ----------
def result_box(label, value):
    return f'<div class="result-box"><div class="label">{label}</div><div class="value">{value}</div></div>'

# ===== RESULTS =====
if run:
    st.markdown("---")
    
    # ===== ONE-SAMPLE =====
    if inf_type == "One-Sample Mean" and n >= 2:
        se = s / np.sqrt(n)
        df = n - 1
        
        if analysis_type == "Confidence Interval":
            crit = norm.ppf(1 - (1 - conf_level)/2) if dist_is_z else t.ppf(1 - (1 - conf_level)/2, df)
            margin = crit * se
            lower, upper = xbar - margin, xbar + margin
            
            cols = st.columns(2)
            with cols[0]:
                st.markdown(result_box("Sample Mean (x̄)", f"{xbar:.4f}"), unsafe_allow_html=True)
            with cols[1]:
                st.markdown(result_box(f"{conf_level*100:.0f}% Confidence Interval", 
                                       f"[{lower:.4f}, {upper:.4f}]"), unsafe_allow_html=True)
            
            st.markdown("---")
            st.subheader("Calculation Details")
            col1, col2, col3 = st.columns(3)
            with col1:
                st.metric("Standard Error", f"{se:.4f}")
            with col2:
                label = "Critical Value (z)" if dist_is_z else f"Critical Value (t, df={df})"
                st.metric(label, f"±{crit:.4f}")
            with col3:
                st.metric("Margin of Error", f"{margin:.4f}")
            
            results = {"Analysis": ["CI"], "n": [n], "x̄": [xbar], "s": [s],
                      "SE": [se], "Lower": [lower], "Upper": [upper]}

        else:  # Hypothesis Test
            stat = (xbar - mu0) / se
            
            if "≠" in alt:
                p_val = 2 * (1 - (norm.cdf(abs(stat)) if dist_is_z else t.cdf(abs(stat), df)))
                crit = norm.ppf(1 - alpha/2) if dist_is_z else t.ppf(1 - alpha/2, df)
                crit_display = f"±{crit:.4f}"
            elif ">" in alt:
                p_val = 1 - (norm.cdf(stat) if dist_is_z else t.cdf(stat, df))
                crit = norm.ppf(1 - alpha) if dist_is_z else t.ppf(1 - alpha, df)
                crit_display = f"{crit:.4f}"
            else:
                p_val = norm.cdf(stat) if dist_is_z else t.cdf(stat, df)
                crit = -(norm.ppf(1 - alpha) if dist_is_z else t.ppf(1 - alpha, df))
                crit_display = f"{crit:.4f}"

            reject = p_val <= alpha
            stat_name = "z" if dist_is_z else "t"
            
            cols = st.columns(2)
            with cols[0]:
                st.markdown(result_box(f"{stat_name}-statistic", f"{stat:.4f}"), unsafe_allow_html=True)
            with cols[1]:
                st.markdown(result_box("p-value", f"{p_val:.4g}"), unsafe_allow_html=True)
            
            if reject:
                st.markdown(f'<div class="decision-box decision-reject">✗ REJECT H₀ at α = {alpha}</div>', 
                           unsafe_allow_html=True)
            else:
                st.markdown(f'<div class="decision-box decision-accept">— Fail to reject H₀ at α = {alpha}</div>', 
                           unsafe_allow_html=True)
            
            st.markdown("---")
            st.subheader("Calculation Details")
            st.write(f"**Hypotheses:** H₀: μ = {mu0:.4f} vs H₁: {alt}")
            st.write(f"**Sample mean:** x̄ = {xbar:.4f}")
            st.write(f"**Standard Error:** {se:.4f}")
            if not dist_is_z:
                st.write(f"**Degrees of freedom:** df = {df}")
            st.write(f"**Critical value(s):** {crit_display}")
            
            results = {"Analysis": ["HT"], "n": [n], "x̄": [xbar], "μ₀": [mu0],
                      "stat": [stat], "p-value": [p_val], "Reject": [reject]}

    # ===== TWO-SAMPLE =====
    elif inf_type.startswith("Two-Sample") and n1 >= 2 and n2 >= 2:
        se = np.sqrt(s1**2/n1 + s2**2/n2)
        diff = xbar1 - xbar2
        
        # Welch df
        df_num = (s1**2/n1 + s2**2/n2)**2
        df_den = (s1**2/n1)**2/(n1-1) + (s2**2/n2)**2/(n2-1)
        df = df_num / df_den
        
        if analysis_type == "Confidence Interval":
            crit = norm.ppf(1 - (1 - conf_level)/2) if dist_is_z else t.ppf(1 - (1 - conf_level)/2, df)
            margin = crit * se
            lower, upper = diff - margin, diff + margin
            
            cols = st.columns(2)
            with cols[0]:
                st.markdown(result_box("Difference (x̄₁ − x̄₂)", f"{diff:.4f}"), unsafe_allow_html=True)
            with cols[1]:
                st.markdown(result_box(f"{conf_level*100:.0f}% Confidence Interval", 
                                       f"[{lower:.4f}, {upper:.4f}]"), unsafe_allow_html=True)
            
            st.markdown("---")
            st.subheader("Calculation Details")
            st.write(f"**x̄₁ = {xbar1:.4f}**, **x̄₂ = {xbar2:.4f}**")
            col1, col2, col3 = st.columns(3)
            with col1:
                st.metric("SE (Welch)", f"{se:.4f}")
            with col2:
                st.metric("df (Welch)", f"{df:.1f}")
            with col3:
                st.metric("Margin of Error", f"{margin:.4f}")
            
            results = {"Analysis": ["CI-2"], "x̄₁": [xbar1], "x̄₂": [xbar2],
                      "Diff": [diff], "Lower": [lower], "Upper": [upper]}

        else:  # Hypothesis Test
            stat = diff / se
            
            if "≠" in alt:
                p_val = 2 * (1 - (norm.cdf(abs(stat)) if dist_is_z else t.cdf(abs(stat), df)))
                crit = norm.ppf(1 - alpha/2) if dist_is_z else t.ppf(1 - alpha/2, df)
                crit_display = f"±{crit:.4f}"
            elif ">" in alt:
                p_val = 1 - (norm.cdf(stat) if dist_is_z else t.cdf(stat, df))
                crit = norm.ppf(1 - alpha) if dist_is_z else t.ppf(1 - alpha, df)
                crit_display = f"{crit:.4f}"
            else:
                p_val = norm.cdf(stat) if dist_is_z else t.cdf(stat, df)
                crit = -(norm.ppf(1 - alpha) if dist_is_z else t.ppf(1 - alpha, df))
                crit_display = f"{crit:.4f}"

            reject = p_val <= alpha
            stat_name = "z" if dist_is_z else "t"
            
            cols = st.columns(2)
            with cols[0]:
                st.markdown(result_box(f"{stat_name}-statistic", f"{stat:.4f}"), unsafe_allow_html=True)
            with cols[1]:
                st.markdown(result_box("p-value", f"{p_val:.4g}"), unsafe_allow_html=True)
            
            if reject:
                st.markdown(f'<div class="decision-box decision-reject">✗ REJECT H₀ at α = {alpha}</div>', 
                           unsafe_allow_html=True)
            else:
                st.markdown(f'<div class="decision-box decision-accept">— Fail to reject H₀ at α = {alpha}</div>', 
                           unsafe_allow_html=True)
            
            st.markdown("---")
            st.subheader("Calculation Details")
            st.write(f"**Hypotheses:** H₀: μ₁ = μ₂ vs H₁: {alt}")
            st.write(f"**x̄₁ = {xbar1:.4f}**, **x̄₂ = {xbar2:.4f}**, **Difference = {diff:.4f}**")
            st.write(f"**Standard Error (Welch):** {se:.4f}")
            st.write(f"**Degrees of freedom (Welch):** {df:.1f}")
            st.write(f"**Critical value(s):** {crit_display}")
            
            results = {"Analysis": ["HT-2"], "x̄₁": [xbar1], "x̄₂": [xbar2],
                      "stat": [stat], "p-value": [p_val], "Reject": [reject]}

    # Download
    if 'results' in locals():
        df_out = pd.DataFrame(results)
        buff = io.BytesIO()
        with pd.ExcelWriter(buff, engine="xlsxwriter") as writer:
            df_out.to_excel(writer, index=False)
        st.download_button("📥 Download Results", data=buff.getvalue(),
                          file_name="mean_inference.xlsx",
                          mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet")

# ---------- Formulas (collapsed) ----------
with st.expander("📚 Formulas & Theory"):
    st.markdown(r"""
**One-Sample Mean**
- SE: $s/\sqrt{n}$
- CI: $\bar x \pm t_{\alpha/2, df} \cdot SE$
- Test stat: $t = (\bar x - \mu_0)/SE$, $df = n-1$

**Two Independent Means (Welch)**
- SE: $\sqrt{s_1^2/n_1 + s_2^2/n_2}$
- Welch df: $\dfrac{(s_1^2/n_1 + s_2^2/n_2)^2}{(s_1^2/n_1)^2/(n_1-1) + (s_2^2/n_2)^2/(n_2-1)}$
- CI: $(\bar x_1 - \bar x_2) \pm t_{\alpha/2, df} \cdot SE$
- Test stat: $t = (\bar x_1 - \bar x_2)/SE$

**When to use z vs t:**
- **z**: Large sample (n ≥ 30) or σ known
- **t**: Small sample with σ unknown (assumes normal population)
""")