from agi_web import python_ui as st
import numpy as np
import pandas as pd
import altair as alt
import os
import sys

# Ensure project root is in path for imports
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from forecast_core import make_fixture, predict, run_analysis

st.title("Promotion forecast lab")
st.caption("Chronos-2 Small on synthetic data. This is a scenario forecast, not a causal estimate.")

# Input Form
with st.form("forecast_form"):
    col1, col2 = st.columns(2)
    with col1:
        horizon = st.selectbox("Horizon (days)", [14, 28, 56], index=1)
        seed = st.number_input("Seed", min_value=0, max_value=10000, value=42)
    with col2:
        promotion_start = st.number_input("Promotion Start (day)", min_value=0, max_value=100, value=7)
        promotion_days = st.number_input("Promotion Duration (days)", min_value=1, max_value=28, value=7)
    
    submitted = st.form_submit_button("Run analysis")

# Show context immediately if not submitted yet, or if we have session state
if 'fixture' not in st.session_state:
    # Generate a default fixture for immediate display
    st.session_state.fixture = make_fixture(seed=42, horizon=28, promotion_start=7, promotion_days=7)

fixture = st.session_state.fixture

# Display Context Data
st.subheader("Synthetic Context Data")
context_df = pd.DataFrame({
    "t": np.arange(len(fixture["context"])),
    "sales": fixture["context"],
    "promotion": fixture["past_promotion"],
    "weekday_sin": fixture["past_weekday"]
})

chart_context = alt.Chart(context_df).mark_line().encode(
    x=alt.X("t:Q", title="Time Step"),
    y=alt.Y("sales:Q", title="Sales"),
    color=alt.value("blue")
).properties(
    title="Historical Sales (Context)",
    width=600,
    height=300
)

st.altair_chart(chart_context, use_container_width=True)

if submitted:
    with st.spinner("Running Chronos-2 inference..."):
        try:
            # Run the full analysis
            results = run_analysis(
                seed=int(seed),
                horizon=int(horizon),
                promotion_start=int(promotion_start),
                promotion_days=int(promotion_days)
            )
            st.session_state.results = results
            st.session_state.fixture = results["fixture"]
        except Exception as e:
            st.error(f"Analysis failed: {str(e)}")
            st.session_state.results = None

if 'results' in st.session_state and st.session_state.results is not None:
    results = st.session_state.results
    fixture = results["fixture"]
    predictions = results["predictions"]
    metrics = results["metrics"]
    
    st.subheader("Forecast Results")
    
    # Prepare data for plotting
    horizon_len = len(predictions["forecast"])
    context_len = len(fixture["context"])
    
    # Historical part
    hist_df = pd.DataFrame({
        "t": np.arange(context_len),
        "value": fixture["context"],
        "type": "Historical"
    })
    
    # Future part
    future_t = np.arange(context_len, context_len + horizon_len)
    future_df = pd.DataFrame({
        "t": future_t,
        "actual": fixture["actual"],
        "forecast": predictions["forecast"],
        "lower": predictions["lower"],
        "upper": predictions["upper"],
        "baseline": results["baseline"],
        "type": "Future"
    })
    
    # Plot Actual vs Forecast
    chart_forecast = alt.Chart(future_df).mark_line().encode(
        x=alt.X("t:Q", title="Time Step"),
        y=alt.Y("value:Q", title="Sales"),
        color=alt.Color("series:N", legend=alt.Legend(title="Series"))
    ).transform_fold(
        ["actual", "forecast", "baseline"],
        as_=["series", "value"]
    ).properties(
        title="Forecast vs Actual vs Baseline",
        width=600,
        height=300
    )
    
    # Add uncertainty band
    band = alt.Chart(future_df).mark_area(opacity=0.3).encode(
        x="t:Q",
        y="lower:Q",
        y2="upper:Q",
        color=alt.value("lightblue")
    )
    
    st.altair_chart(band + chart_forecast, use_container_width=True)
    
    # Metrics
    st.subheader("Metrics")
    col1, col2, col3 = st.columns(3)
    col1.metric("Chronos MAE", f"{metrics['mae']:.2f}")
    col2.metric("Baseline MAE", f"{metrics['baseline_mae']:.2f}")
    col3.metric("Observed Coverage (p10-p90)", f"{metrics['coverage']:.2%}")
    
    st.info("Note: The nominal 80% interval is not calibrated on this synthetic fixture. Observed coverage reflects the specific realization.")