server.py

"""
server.py — SeasonShift Flask API
Serves data to the UI: forecast chart data, agent activity log,
and time-cursor advance endpoint.
"""

import json
import pathlib
import numpy as np
import pandas as pd
from flask import Flask, jsonify, request
from flask_cors import CORS

import model as m
import arize_logging as al
import agent as ag

app = Flask(__name__)
CORS(app)

DATA_DIR = pathlib.Path("data")
_df: pd.DataFrame | None = None
_gemini_client = None

# Current simulation cursor index
_cursor_idx: int = 0
_cursor_step: int = 30  # advance by 30 days each step


def _load_data() -> pd.DataFrame:
    global _df
    if _df is None:
        csv = DATA_DIR / "season_shift_dataset.csv"
        _df = pd.read_csv(csv, index_col=0, parse_dates=True)
        try:
            m.load_model()
        except FileNotFoundError:
            pass  # predictions will fail gracefully
    return _df


@app.route("/api/health")
def health():
    return jsonify({"status": "ok", "active_pack": m.get_active_pack()})


# ---------------------------------------------------------------------------
# Chart data: actuals + predictions for a date window
# ---------------------------------------------------------------------------

@app.route("/api/forecast")
def forecast():
    df = _load_data()
    # Query params: start_date, end_date (ISO strings)
    start = request.args.get("start")
    end = request.args.get("end")

    if start and end:
        mask = (df.index >= start) & (df.index <= end)
        window = df[mask]
    else:
        # Default: show last 120 days from cursor
        global _cursor_idx
        end_idx = min(_cursor_idx + _cursor_step, len(df))
        window = df.iloc[max(0, end_idx - 120):end_idx]

    if len(window) == 0:
        return jsonify({"dates": [], "actuals": [], "predicted_default": [], "predicted_holiday": []})

    # Compute trust weight at end of window (cursor position)
    date_str = window.index[-1].strftime("%Y-%m-%d")
    state = build_simulation_state(date_str)
    trust_weight = state["trust_weight"]

    try:
        preds_default = m.predict(window, pack="pack_default")
        preds_holiday = m.predict(window, pack="pack_holiday")
        # Active = trust-blended when holiday pack is live
        if state["active_pack"] == "pack_holiday":
            preds_active = m.predict(window, pack="pack_holiday", trust_weight=trust_weight)
        else:
            preds_active = m.predict(window, pack="pack_default")
    except Exception:
        preds_default = pd.Series([None] * len(window), index=window.index)
        preds_holiday = pd.Series([None] * len(window), index=window.index)
        preds_active = pd.Series([None] * len(window), index=window.index)

    def _clean(s):
        return [None if (v is None or (isinstance(v, float) and np.isnan(v))) else round(float(v), 1) for v in s]

    return jsonify({
        "dates": [d.strftime("%Y-%m-%d") for d in window.index],
        "actuals": _clean(window["sales"].tolist()),
        "predicted_default": _clean(preds_default.tolist()),
        "predicted_holiday": _clean(preds_holiday.tolist()),
        "predicted_active": _clean(preds_active.tolist()),
        "active_pack": m.get_active_pack(),
        "trust_weight": round(trust_weight, 3),
        "cursor_date": df.index[min(_cursor_idx, len(df) - 1)].strftime("%Y-%m-%d"),
    })


# ---------------------------------------------------------------------------
# Time cursor
# ---------------------------------------------------------------------------

@app.route("/api/cursor", methods=["GET"])
def get_cursor():
    df = _load_data()
    return jsonify({
        "cursor_idx": _cursor_idx,
        "cursor_date": df.index[min(_cursor_idx, len(df) - 1)].strftime("%Y-%m-%d"),
        "total_days": len(df),
        "active_pack": m.get_active_pack(),
    })


@app.route("/api/cursor/advance", methods=["POST"])
def advance_cursor():
    global _cursor_idx
    df = _load_data()
    body = request.get_json(silent=True) or {}
    step = int(body.get("step", _cursor_step))
    _cursor_idx = min(_cursor_idx + step, len(df) - 1)
    current_date = df.index[_cursor_idx]

    date_str = current_date.strftime("%Y-%m-%d")

    # Run agent step at new cursor position
    activity = ag.run_agent_step(df.iloc[:_cursor_idx + 1], current_date, gemini_client=_gemini_client)

    # Re-derive model pack from replay — no mutable tracking variables
    state = build_simulation_state(date_str)
    m._active_pack = state["active_pack"]

    return jsonify({
        "cursor_idx":    _cursor_idx,
        "cursor_date":   date_str,
        "total_days":    len(df),
        "active_pack":   m.get_active_pack(),
        "agent_activity": activity,
    })


@app.route("/api/cursor/reset", methods=["POST"])
def reset_cursor():
    global _cursor_idx
    df = _load_data()
    _cursor_idx = 0
    m._active_pack = "pack_default"
    # Clear agent log so replay starts fresh
    if ag.AGENT_LOG_PATH.exists():
        ag.AGENT_LOG_PATH.unlink()
    return jsonify({
        "cursor_idx": 0,
        "cursor_date": df.index[0].strftime("%Y-%m-%d"),
        "total_days": len(df),
        "active_pack": m.get_active_pack(),
    })


@app.route("/api/cursor/set", methods=["POST"])
def set_cursor():
    global _cursor_idx
    df = _load_data()
    body = request.get_json(silent=True) or {}
    idx = int(body.get("idx", _cursor_idx))
    _cursor_idx = max(0, min(idx, len(df) - 1))
    date_str = df.index[_cursor_idx].strftime("%Y-%m-%d")

    # Replay log to re-derive pack — this is what makes rewind work
    state = build_simulation_state(date_str)
    m._active_pack = state["active_pack"]

    return jsonify({
        "cursor_idx":  _cursor_idx,
        "cursor_date": date_str,
        "total_days":  len(df),
        "active_pack": m.get_active_pack(),
    })


# ---------------------------------------------------------------------------
# Inference view — deterministic replay, no mutable global state
# ---------------------------------------------------------------------------

BASELINE_FEATURES = ["day_of_week", "month", "is_school_holiday", "is_promo",
                     "sales_lag_7", "sales_lag_365"]


def build_simulation_state(cursor_date_str: str) -> dict:
    """
    Replay the agent log up to cursor_date_str and reconstruct what the
    system would have known at that point in simulated time.

    This is the single source of truth for promoted features, active pack,
    and the primary signal currently being watched.  No mutable globals.
    """
    all_entries = ag.get_agent_log()
    visible = [e for e in all_entries if e.get("simulation_date", "") <= cursor_date_str]

    promoted: list[str] = []
    rejected_set: set[str] = set()
    for e in visible:
        for ev in e.get("promotion_events", []):
            f = ev.get("feature")
            if f and f not in promoted:
                promoted.append(f)
        for ev in e.get("rejection_events", []):
            f = ev.get("feature")
            # Only surface non-baseline features as "rejected" — baseline features
            # (is_promo, month, sales_lag_7 etc.) are expected candidates that are
            # simply "not yet promoted", not confirmed decoys.
            if f and f not in BASELINE_FEATURES:
                rejected_set.add(f)

    # A feature is "rejected" only if it was never subsequently promoted
    # AND it stopped appearing in recent decisions (signal died out — a true decoy).
    # Features still being actively evaluated (keeps appearing in decisions) are
    # shown in the "watching" panel, not the rejected chips.
    recently_active: set[str] = set()
    for e in visible[-4:]:  # last 4 log entries
        for d in e.get("decisions", []):
            recently_active.add(d["feature"])

    rejected = [
        f for f in rejected_set
        if f not in promoted and f not in recently_active
    ]
    active   = BASELINE_FEATURES + [f for f in promoted if f not in BASELINE_FEATURES]
    active_pack = "pack_holiday" if promoted else "pack_default"

    # Primary candidate: top non-promoted, non-baseline decision in most recent step.
    # Skip baseline features — they're always active and not candidates for promotion.
    # Include features that were previously rejected (they may still be accumulating
    # evidence — e.g. days_until_christmas in early months before first Dec cycle).
    primary_candidate = None
    for e in reversed(visible):
        for d in e.get("decisions", []):
            if d["feature"] not in promoted and d["feature"] not in BASELINE_FEATURES:
                primary_candidate = {
                    "feature":          d["feature"],
                    "verdict":          d["verdict"],
                    "confidence":       d["confidence"],
                    "recurrence_count": d["recurrence_count"],
                    "recent_signal":    d["recent_signal"],
                }
                break
        if primary_candidate:
            break

    # Trust weight: derived from days_until_christmas current value (not from calendar dates).
    # After days_until_christmas is promoted, trust scales 0→1 as DUC approaches 0 (Christmas).
    # Drops to 0 immediately after Christmas when DUC resets to ~365.
    # DORMANT_THRESHOLD = 0.15 → feature shows DORMANT when DUC > ~77 days from Christmas.
    DORMANT_THRESHOLD = 0.15
    trust_weight = 0.0
    if promoted:
        df = _load_data()
        trust_weight = m.compute_trust_weight(cursor_date_str, df)

    dormant_features = [
        f for f in promoted
        if f not in BASELINE_FEATURES and trust_weight < DORMANT_THRESHOLD
    ]
    # A promoted feature is "currently active" when trust is above threshold
    currently_active_features = [f for f in promoted if f not in dormant_features]

    return {
        "active_features":             active,
        "promoted_features":           promoted,
        "currently_active_features":   currently_active_features,
        "dormant_features":            dormant_features,
        "rejected_features":           rejected,
        "baseline_features":           BASELINE_FEATURES,
        "active_pack":                 active_pack,
        "primary_candidate":           primary_candidate,
        "trust_weight":                trust_weight,
        "visible_entries":             visible,
    }


@app.route("/api/inference-view")
def inference_view():
    df = _load_data()
    date_str = df.index[min(_cursor_idx, len(df) - 1)].strftime("%Y-%m-%d")
    state = build_simulation_state(date_str)
    return jsonify({
        "active_features":           state["active_features"],
        "promoted_features":         state["promoted_features"],
        "currently_active_features": state["currently_active_features"],
        "dormant_features":          state["dormant_features"],
        "rejected_features":         state["rejected_features"],
        "baseline_features":         state["baseline_features"],
        "primary_candidate":         state["primary_candidate"],
        "trust_weight":              state["trust_weight"],
    })


# ---------------------------------------------------------------------------
# Agent log — filtered to cursor date so rewind hides future decisions
# ---------------------------------------------------------------------------

@app.route("/api/agent/log")
def agent_log():
    df = _load_data()
    date_str = df.index[min(_cursor_idx, len(df) - 1)].strftime("%Y-%m-%d")
    all_entries = ag.get_agent_log()
    visible = [e for e in all_entries if e.get("simulation_date", "") <= date_str]
    return jsonify({"entries": visible[-50:]})


# ---------------------------------------------------------------------------
# Feature performance (for debugging / Arize read-path)
# ---------------------------------------------------------------------------

@app.route("/api/feature-performance")
def feature_performance():
    feature = request.args.get("feature", "days_until_christmas")
    date = request.args.get("date")
    if not date:
        df = _load_data()
        date = df.index[min(_cursor_idx, len(df) - 1)].strftime("%Y-%m-%d")
    result = al.get_feature_performance(feature, date)
    return jsonify(result)


# ---------------------------------------------------------------------------
# Model info
# ---------------------------------------------------------------------------

@app.route("/api/model/packs")
def model_packs():
    return jsonify({
        "active_pack": m.get_active_pack(),
        "available_packs": list(m.FEATURE_PACKS.keys()),
        "pack_features": m.FEATURE_PACKS,
    })


@app.route("/api/model/switch", methods=["POST"])
def manual_switch():
    body = request.get_json(silent=True) or {}
    pack = body.get("pack", "pack_holiday")
    result = m.switch_feature_pack(pack)
    return jsonify(result)


if __name__ == "__main__":
    _load_data()
    _gemini_client = ag.get_gemini_client()
    print("[server] SeasonShift API running on http://localhost:5050")
    app.run(host="0.0.0.0", port=5050, debug=False)