import json
import numpy as np
import pandas as pd
import plotly.graph_objects as go
import plotly.express as px
import streamlit as st
import os

st.set_page_config(page_title="Hybrid Attendance Monte Carlo", layout="wide")
st.title("Hybrid Office Attendance: Monte Carlo PoC")

here=os.path.dirname(os.path.abspath(__file__))
preset_path = os.path.join(here,"presets.json")

def load_presets():
    with open(preset_path, "r") as f:
        return json.load(f)

presets=load_presets()        
preset_names = list(presets.keys())

with st.sidebar:
    st.header("Scenario Presets")
    chosen = st.selectbox("Pick a preset", preset_names, index=0)
    seed = st.number_input("Random seed", 0, 10000, 42, step=1)
    st.caption("Start with a preset, then tweak below as needed.")

# ---- Inputs (main panel) ----
cfg = presets[chosen].copy()

colA, colB, colC, colD = st.columns(4)
n_employees = colA.number_input("Employees", 10, 20000, int(cfg.get("n_employees", 300)), step=10)
n_days      = colB.number_input("Sim days", 5, 365, int(cfg.get("n_days", 60)), step=5)
n_runs      = colC.number_input("Monte Carlo runs", 100, 50000, int(cfg.get("n_runs", 2000)), step=100)
capacity    = colD.number_input("Capacity (seats/desk count)", 0, 100000, int(cfg.get("capacity", 220)), step=10)

st.subheader("Day-of-Week Base Patterns")
with st.expander("Weekly attendance patterns", expanded=True):
    st.write("Set base attendance probability for each day of the week:")
    dow_cols = st.columns(7)
    days = ["Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday", "Sunday"]
    day_abbrev = ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"]
    
    # Default pattern: Thu highest, Tue second, Fri lowest
    default_dow_rates = cfg.get("dow_base_rates", [0.35, 0.55, 0.45, 0.65, 0.25, 0.10, 0.05])
    dow_base_rates = []
    
    for i, (day, abbr) in enumerate(zip(days, day_abbrev)):
        rate = dow_cols[i].slider(f"{abbr}", 0.0, 1.0, float(default_dow_rates[i]), 0.05, key=f"dow_{i}")
        dow_base_rates.append(rate)

st.subheader("Model Parameters")
with st.expander("Simulation settings", expanded=True):
    c1, c2, c3 = st.columns(3)
    noise_sd    = c1.slider("Day-to-day volatility (σ)", 0.0, 0.5, float(cfg.get("noise_sd", 0.15)), 0.01)
    team_corr   = c2.slider("Team synchronicity (0=independent, 1=lockstep)", 0.0, 0.95, float(cfg.get("team_corr", 0.25)), 0.05)
    dow_strength = c3.slider("Day-of-week pattern strength", 0.0, 1.0, float(cfg.get("dow_strength", 0.8)), 0.05)
    st.caption("Higher volatility = more day-to-day variation. Lower day-of-week strength = more random attendance patterns.")

st.subheader("Employee Groups")
with st.expander("Define employee attendance groups", expanded=True):
    st.write("Create different employee groups with distinct attendance patterns:")
    
    # Default groups
    default_groups = cfg.get("employee_groups", {
        "Regular (3-4 days/week)": {"share": 0.4, "target_days": 3.5},
        "Light (1-2 days/week)": {"share": 0.35, "target_days": 1.5}, 
        "Heavy (4-5 days/week)": {"share": 0.15, "target_days": 4.2},
        "Rare (< 1 day/week)": {"share": 0.1, "target_days": 0.4}
    })
    
    employee_groups = {}
    group_names = list(default_groups.keys())
    
    # Group configuration
    group_cols = st.columns(len(group_names))
    for i, group_name in enumerate(group_names):
        with group_cols[i]:
            st.write(f"**{group_name}**")
            share = st.slider("Share of employees", 0.0, 1.0, 
                            float(default_groups[group_name]["share"]), 0.05, 
                            key=f"share_{i}")
            target_days = st.slider("Target days/week", 0.0, 5.0, 
                                  float(default_groups[group_name]["target_days"]), 0.1,
                                  key=f"target_{i}")
            employee_groups[group_name] = {"share": share, "target_days": target_days}
    
    # Normalize shares
    total_share = sum(g["share"] for g in employee_groups.values())
    if total_share > 0:
        for group_name in employee_groups:
            employee_groups[group_name]["share"] /= total_share

st.subheader("Advanced: Day-of-Week Modifiers by Group")
with st.expander("Customize day-of-week patterns for each employee group", expanded=False):
    st.write("**Day-of-Week Modifiers by Group**")
    st.caption("Multiply base day-of-week rates by these factors for each group")
    
    # Day-of-week modifiers by group
    group_dow_modifiers = {}
    for group_name in employee_groups.keys():
        st.write(f"*{group_name}*")
        modifier_cols = st.columns(7)
        modifiers = []
        default_mods = cfg.get("group_dow_modifiers", {}).get(group_name, [1.0] * 7)
        for j, abbr in enumerate(day_abbrev):
            mod = modifier_cols[j].slider(f"{abbr}", 0.0, 3.0, 
                                        float(default_mods[j] if j < len(default_mods) else 1.0), 
                                        0.1, key=f"mod_{group_name}_{j}")
            modifiers.append(mod)
        group_dow_modifiers[group_name] = modifiers

rng = np.random.default_rng(int(seed))

# ---- Build per-employee probabilities from groups ----
# Assign employees to groups
group_counts = []
group_assignments = []
for group_name, group_info in employee_groups.items():
    count = int(round(group_info["share"] * n_employees))
    group_counts.append(count)
    group_assignments.extend([group_name] * count)

# Adjust for rounding errors
diff = n_employees - len(group_assignments)
if diff > 0:
    # Add employees to first group
    first_group = list(employee_groups.keys())[0]
    group_assignments.extend([first_group] * diff)
elif diff < 0:
    # Remove from end
    group_assignments = group_assignments[:n_employees]

# Shuffle assignments
rng.shuffle(group_assignments)

# Create per-employee day-of-week probability matrix
emp_dow_probs = np.zeros((n_employees, 7))
for i, group_name in enumerate(group_assignments):
    group_info = employee_groups[group_name]
    group_modifiers = group_dow_modifiers[group_name]
    
    # Calculate group's day-of-week probabilities
    group_dow_probs = np.array(dow_base_rates) * np.array(group_modifiers)
    
    # Scale to match target days per week
    current_avg = np.mean(group_dow_probs)
    target_avg = group_info["target_days"] / 7.0
    if current_avg > 0:
        group_dow_probs = group_dow_probs * (target_avg / current_avg)
    
    # Clip to valid probability range
    group_dow_probs = np.clip(group_dow_probs, 0.0, 1.0)
    
    # Add some individual variation
    variation = rng.normal(0, 0.05, 7)  # Small random variation per person
    emp_dow_probs[i] = np.clip(group_dow_probs + variation, 0.0, 1.0)

# Create day-of-week pattern for simulation period
weekday = np.arange(int(n_days)) % 7

def simulate_run():
    # Generate base probabilities for each employee-day combination
    base_probs = np.zeros((int(n_employees), int(n_days)))
    
    # Calculate average probability across all days for each employee (for non-dow component)
    avg_emp_probs = np.mean(emp_dow_probs, axis=1)
    
    for day in range(int(n_days)):
        dow = weekday[day]  # day of week (0=Monday, 6=Sunday)
        dow_probs = emp_dow_probs[:, dow]
        
        # Blend day-of-week pattern with average (controlled by dow_strength)
        base_probs[:, day] = dow_strength * dow_probs + (1 - dow_strength) * avg_emp_probs
    
    # MUCH LARGER noise for visible day-to-day variation
    # Common shocks affect everyone (weather, company events, etc.)
    common_shock = rng.normal(0, noise_sd * 2.0, size=int(n_days))
    
    # Individual variation - each person has their own random daily variation
    indiv_noise = rng.normal(0, noise_sd * 1.5, size=(int(n_employees), int(n_days)))
    
    # Apply shocks with much stronger impact
    shock = team_corr * common_shock + np.sqrt(max(1 - team_corr**2, 0)) * indiv_noise
    
    # Final probabilities with noise - allow broader range before clipping
    final_probs = np.clip(base_probs + shock, 0.0, 1.0)
    
    # Generate attendance (Bernoulli draws)
    draws = rng.uniform(size=(int(n_employees), int(n_days))) < final_probs
    daily_counts = draws.sum(axis=0)
    return daily_counts

st.divider()
run_btn = st.button("Run Simulation", type="primary")
if run_btn:
    daily_matrix = np.vstack([simulate_run() for _ in range(int(n_runs))])  # [runs x days]
    summary = pd.DataFrame({
        "day": np.arange(1, int(n_days) + 1),
        "p10": np.percentile(daily_matrix, 10, axis=0),
        "p25": np.percentile(daily_matrix, 25, axis=0),
        "p50": np.percentile(daily_matrix, 50, axis=0),
        "p75": np.percentile(daily_matrix, 75, axis=0),
        "p90": np.percentile(daily_matrix, 90, axis=0),
    }).astype({"day": int})

    # ----- Visualization: percentile fan chart + capacity line -----
    fig = go.Figure()
    
    # Add uncertainty bands (filled areas)
    fig.add_trace(go.Scatter(
        x=list(summary["day"]) + list(summary["day"][::-1]),
        y=list(summary["p90"]) + list(summary["p10"][::-1]),
        fill='toself', fillcolor='rgba(0,100,80,0.1)',
        line=dict(color='rgba(255,255,255,0)'),
        name='p10-p90 band'
    ))
    
    fig.add_trace(go.Scatter(
        x=list(summary["day"]) + list(summary["day"][::-1]),
        y=list(summary["p75"]) + list(summary["p25"][::-1]),
        fill='toself', fillcolor='rgba(0,100,80,0.2)',
        line=dict(color='rgba(255,255,255,0)'),
        name='p25-p75 band'
    ))
    
    # Add percentile lines
    fig.add_trace(go.Scatter(x=summary["day"], y=summary["p90"], mode="lines", name="p90", line=dict(color='red', width=1)))
    fig.add_trace(go.Scatter(x=summary["day"], y=summary["p75"], mode="lines", name="p75", line=dict(color='orange', width=1)))
    fig.add_trace(go.Scatter(x=summary["day"], y=summary["p50"], mode="lines", name="p50 (median)", line=dict(color='blue', width=2)))
    fig.add_trace(go.Scatter(x=summary["day"], y=summary["p25"], mode="lines", name="p25", line=dict(color='orange', width=1)))
    fig.add_trace(go.Scatter(x=summary["day"], y=summary["p10"], mode="lines", name="p10", line=dict(color='red', width=1)))
    
    if capacity and capacity > 0:
        fig.add_trace(go.Scatter(x=summary["day"], y=[capacity]*len(summary), mode="lines", name="capacity", line=dict(dash="dash", color="black", width=2)))
    
    fig.update_layout(
        title="Daily In-Office Headcount (percentiles across Monte Carlo runs)",
        xaxis_title="Day", 
        yaxis_title="Headcount", 
        legend_title="Series",
        hovermode='x'
    )
    st.plotly_chart(fig, use_container_width=True)

    # Individual simulation runs visualization
    st.subheader("Individual Simulation Runs")
    with st.expander("Sample of individual runs (shows actual variability)", expanded=True):
        # Sample runs to show
        n_sample_runs = min(100, n_runs)
        
        # Sample evenly spaced runs
        sample_indices = np.linspace(0, n_runs-1, n_sample_runs, dtype=int)
        
        # Create line chart of individual runs
        fig_runs = go.Figure()
        
        # Add individual run lines (lighter, thinner)
        for i, run_idx in enumerate(sample_indices):
            if i < 50:  # Limit lines for performance
                fig_runs.add_trace(go.Scatter(
                    x=summary["day"], 
                    y=daily_matrix[run_idx], 
                    mode="lines",
                    name=f"Run {run_idx+1}",
                    line=dict(width=0.8, color=f'rgba(100,150,200,0.3)'),
                    showlegend=False,
                    hovertemplate=f"Run {run_idx+1}<br>Day: %{{x}}<br>Attendance: %{{y}}<extra></extra>"
                ))
        
        # Add median line on top
        fig_runs.add_trace(go.Scatter(
            x=summary["day"], 
            y=summary["p50"], 
            mode="lines", 
            name="Median (p50)",
            line=dict(color='blue', width=3)
        ))
        
        # Add capacity line
        if capacity and capacity > 0:
            fig_runs.add_trace(go.Scatter(
                x=summary["day"], 
                y=[capacity]*len(summary), 
                mode="lines", 
                name="Capacity", 
                line=dict(dash="dash", color="black", width=2)
            ))
        
        fig_runs.update_layout(
            title=f"Individual Monte Carlo Runs (showing {len(sample_indices)} of {n_runs} runs)",
            xaxis_title="Day", 
            yaxis_title="Daily Attendance", 
            hovermode='x unified'
        )
        
        st.plotly_chart(fig_runs, use_container_width=True)
        
        # Show variability statistics for the sample
        sample_data = daily_matrix[sample_indices]
        
        col1, col2, col3, col4 = st.columns(4)
        col1.metric("Runs Shown", len(sample_indices))
        col2.metric("Avg Daily Range", f"{np.mean(sample_data.max(axis=0) - sample_data.min(axis=0)):.1f}")
        col3.metric("Max Single Day", f"{sample_data.max()}")
        col4.metric("Min Single Day", f"{sample_data.min()}")
        
        st.caption("Each thin line is one complete simulation run. You can see the actual day-to-day variation!")

    util = (summary["p50"] / n_employees).mean()
    over_days_p50 = int((summary["p50"] > capacity).sum()) if capacity else 0
    over_days_p90 = int((summary["p90"] > capacity).sum()) if capacity else 0
    c1, c2, c3 = st.columns(3)
    c1.metric("Median utilization (avg across days)", f"{util:.0%}")
    c2.metric("Days over capacity @ p50", f"{over_days_p50}")
    c3.metric("Days over capacity @ p90", f"{over_days_p90}")

    # Show group summary
    col1, col2 = st.columns(2)
    with col1:
        st.write("**Employee Group Summary**")
        group_summary = []
        for group_name, group_info in employee_groups.items():
            count = sum(1 for g in group_assignments if g == group_name)
            avg_prob = np.mean([emp_dow_probs[i] for i, g in enumerate(group_assignments) if g == group_name])
            group_summary.append({
                "Group": group_name,
                "Employees": count,
                "Target Days/Week": f"{group_info['target_days']:.1f}",
                "Avg Daily Prob": f"{avg_prob:.1%}"
            })
        st.dataframe(pd.DataFrame(group_summary), hide_index=True)
    
    with col2:
        st.write("**Day-of-Week Pattern & Volatility**")
        # Calculate actual volatility by day of week
        dow_volatility = []
        dow_p50_values = []
        for i in range(7):
            dow_days = summary[weekday == i]["p50"]
            if len(dow_days) > 1:
                volatility = dow_days.std()
                avg_p50 = dow_days.mean()
            else:
                volatility = 0
                avg_p50 = dow_days.iloc[0] if len(dow_days) == 1 else 0
            dow_volatility.append(volatility)
            dow_p50_values.append(avg_p50)
        
        dow_summary = pd.DataFrame({
            "Day": ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"],
            "Base Rate": [f"{r:.1%}" for r in dow_base_rates],
            "Avg P50": [f"{v:.0f}" for v in dow_p50_values],
            "StdDev": [f"{v:.1f}" for v in dow_volatility]
        })
        st.dataframe(dow_summary, hide_index=True)
        st.caption("StdDev shows day-to-day variation within each day of week")

    st.dataframe(summary, use_container_width=True, height=400)
    st.download_button("Download percentile table (CSV)", summary.to_csv(index=False), "attendance_percentiles.csv")
    
    # Raw simulation data output
    st.subheader("Raw Simulation Data")
    with st.expander("All Monte Carlo runs (full data)", expanded=False):
        st.write(f"**Complete dataset**: {n_runs:,} simulation runs × {n_days} days = {n_runs * n_days:,} data points")
        
        # Convert daily_matrix to DataFrame for better viewing
        raw_data = pd.DataFrame(daily_matrix.T)  # Transpose so days are rows, runs are columns
        raw_data.index = range(1, n_days + 1)  # Day numbers as index
        raw_data.index.name = "Day"
        raw_data.columns = [f"Run_{i+1}" for i in range(n_runs)]
        
        # Show statistics
        col1, col2, col3, col4 = st.columns(4)
        col1.metric("Min Daily Attendance", f"{raw_data.values.min()}")
        col2.metric("Max Daily Attendance", f"{raw_data.values.max()}")
        col3.metric("Overall Std Dev", f"{raw_data.values.std():.1f}")
        col4.metric("Mean Daily Attendance", f"{raw_data.values.mean():.1f}")
        
        # Show sample of the data
        st.write("**Sample of raw data** (first 20 days, first 10 runs):")
        sample_data = raw_data.iloc[:20, :10] if n_runs >= 10 else raw_data.iloc[:20, :]
        st.dataframe(sample_data, use_container_width=True)
        
        # Full data download
        csv_data = raw_data.to_csv()
        st.download_button(
            "Download ALL simulation data (CSV)", 
            csv_data, 
            f"attendance_raw_data_{n_runs}runs_{n_days}days.csv",
            help=f"Downloads {n_runs:,} runs × {n_days} days = {len(csv_data):,} characters of raw data"
        )
        
        # Show variability by day of week in raw data
        st.write("**Variability by Day of Week** (from raw data):")
        dow_analysis = []
        days_of_week = ["Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday", "Sunday"]
        
        for dow in range(7):
            dow_mask = (raw_data.index - 1) % 7 == dow  # Day of week mask
            if dow_mask.sum() > 0:
                dow_data = raw_data[dow_mask].values.flatten()
                dow_analysis.append({
                    "Day": days_of_week[dow],
                    "Count": len(dow_data),
                    "Mean": f"{dow_data.mean():.1f}",
                    "Std": f"{dow_data.std():.1f}",
                    "Min": f"{dow_data.min()}",
                    "Max": f"{dow_data.max()}",
                    "Range": f"{dow_data.max() - dow_data.min()}"
                })
        
        st.dataframe(pd.DataFrame(dow_analysis), hide_index=True)
        st.caption("This shows the actual variability in your simulation data")

    st.markdown("### Notes on interpretation")
    st.markdown(
        "- The **uncertainty band** between p10 and p90 tells you how spiky or stable demand is.\n"
        "- Increase **team synchronicity** to model all-hands/meeting cadence—watch for capacity collisions.\n"
        "- Tune **role base propensities** to reflect lab staff, field staff, design review rituals, etc."
    )
else:
    st.info("Pick a preset and tweak parameters, then click **Run Simulation**.")