scenario_engine.py

# scenario_engine.py
from __future__ import annotations
from typing import Dict, List, Any, Tuple, Union, Optional
import re
import math
import statistics
import json
import ast

import pandas as pd
import numpy as np

# ----------------------------
# Safe expression evaluation
# ----------------------------
_ALLOWED_FUNCS = {
    "abs": abs,
    "round": round,
    "sqrt": math.sqrt,
    "log": math.log,
    "exp": math.exp,
    "min": np.minimum,   # vectorized
    "max": np.maximum,   # vectorized
    "mean": np.mean,
    "avg": np.mean,
    "median": np.median,
    "sum": np.sum,
    "count": lambda x: np.size(x),
    "p50": lambda x: np.percentile(x, 50),
    "p75": lambda x: np.percentile(x, 75),
    "p90": lambda x: np.percentile(x, 90),
    "p95": lambda x: np.percentile(x, 95),
    "p99": lambda x: np.percentile(x, 99),
    "ceil": np.ceil,
    "floor": np.floor,
}

class _SafeExpr(ast.NodeTransformer):
    """
    Restrict expressions to:
    - Names (columns), numbers, strings, booleans
    - Arithmetic: + - * / // % **, comparisons, and/or/not
    - Calls to allowed functions (above)
    """
    def __init__(self, allowed_names: set):
        self.allowed_names = allowed_names

    def visit_Name(self, node):
        if node.id not in self.allowed_names and node.id not in ("True","False","None"):
            raise ValueError(f"Unknown name in expression: {node.id}")
        return node

    def visit_Call(self, node):
        if not isinstance(node.func, ast.Name):
            raise ValueError("Only simple function calls are allowed")
        func = node.func.id
        if func not in _ALLOWED_FUNCS:
            raise ValueError(f"Function not allowed: {func}")
        self.generic_visit(node)
        return node

    def generic_visit(self, node):
        allowed = (
            ast.Expression, ast.BoolOp, ast.BinOp, ast.UnaryOp,
            ast.Compare, ast.Call, ast.Name, ast.Load, ast.Constant,
            ast.And, ast.Or, ast.Not,
            ast.Add, ast.Sub, ast.Mult, ast.Div, ast.Mod, ast.Pow, ast.FloorDiv,
            ast.Eq, ast.NotEq, ast.Lt, ast.LtE, ast.Gt, ast.GtE, ast.In, ast.NotIn,
            ast.USub, ast.UAdd
        )
        if not isinstance(node, allowed):
            raise ValueError(f"Unsupported syntax: {type(node).__name__}")
        return super().generic_visit(node)

def _eval_series_expr(expr: str, df: pd.DataFrame) -> pd.Series:
    allowed_names = set(df.columns) | {"True", "False", "None"}
    tree = ast.parse(expr, mode="eval")
    _SafeExpr(allowed_names).visit(tree)
    code = compile(tree, "<expr>", "eval")
    env = {**{k: df[k] for k in df.columns}, **_ALLOWED_FUNCS}
    return eval(code, {"__builtins__": {}}, env)

# ----------------------------
# Engine
# ----------------------------

class ScenarioEngine:
    """
    Scenario-first engine:
    - Parse tasks + inline directives from scenario text
    - For each task, execute a pipeline over analysis_results:
        load -> filter -> derive -> groupby/agg -> pivot -> sort/top -> select fields -> render
    - Render formats: table | list | comparison | map | narrative | chart (Vega-Lite spec)
    - Strict: only what is asked is emitted.
    """

    @staticmethod
    def render(scenario_text: str, analysis_results: Dict[str, Any]) -> str:
        scen = ScenarioEngine._parse_scenario(scenario_text)
        out: List[str] = ["# Scenario Output\n"]
        for task in scen["tasks"]:
            out.append(ScenarioEngine._render_task(task, analysis_results))
        return "\n".join(out).strip()

    # ------------- Parsing -------------
    @staticmethod
    def _parse_scenario(s: str) -> Dict[str, Any]:
        """
        Detect a 'Tasks/Deliverables/Requirements/Your Tasks' block; fallback to any bullet/numbered lines.
        Each task may include inline directives: key: value
        Supported directives (per task):
            format: table|list|comparison|map|narrative|chart
            data_key: <key in analysis_results>
            filter: <expr using columns>   e.g., zone == "North" and wait_time > 5
            derive: <col>=<expr>[, <col2>=<expr2> ...]
            group_by: col1[, col2 ...]
            agg: avg(x), median(y), sum(z), p90(wait), count(*)
            pivot: index=a[,b] columns=c values=v  (values must be an aggregated column)
            sort_by: col  sort_dir: asc|desc
            top: N
            fields: col1 col2 col3   (space or comma separated)
            title: Custom name
            chart: bar|line|area|point   (Vega-Lite spec emitted)
            x: <field>  y: <field>  color: <field>  column: <facet>
        """
        lines = [ln.rstrip() for ln in s.splitlines()]
        task_hdr = re.compile(r'^\s*(tasks?|deliverables|requirements|your tasks?)\s*$', re.I)
        bullet = re.compile(r'^\s*(?:\d+\.\s+|[-*•]\s+)')
        in_tasks = False
        raw_tasks: List[str] = []

        for ln in lines:
            if task_hdr.match(ln):
                in_tasks = True
                continue
            if in_tasks:
                if bullet.match(ln.strip()):
                    raw_tasks.append(ln.strip())
                elif ln.strip() == "":
                    continue
                else:
                    # stop when we hit a non-task looking line after capturing some tasks
                    if raw_tasks:
                        in_tasks = False

        if not raw_tasks:
            # fallback: grab any bullet/numbered lines
            raw_tasks = [ln.strip() for ln in lines if bullet.match(ln.strip())]

        tasks: List[Dict[str, Any]] = []
        for raw in raw_tasks:
            directives = ScenarioEngine._extract_directives(raw)
            title = directives.get("title") or ScenarioEngine._strip_bullet(raw)
            tasks.append({"title": title, "raw": raw, "d": directives})
        return {"tasks": tasks}

    @staticmethod
    def _strip_bullet(line: str) -> str:
        return re.sub(r'^\s*(?:\d+\.\s+|[-*•]\s+)', '', line).strip()

    @staticmethod
    def _extract_directives(text: str) -> Dict[str, Any]:
        d: Dict[str, Any] = {}
        # key: value pairs (value extends until ; or end or two spaces before next key:)
        for m in re.finditer(r'([a-z_]+)\s*:\s*([^|,\n;]+)', text, re.I):
            k = m.group(1).strip().lower()
            v = m.group(2).strip()
            d[k] = v

        def _split_csv(val: str) -> List[str]:
            return [x.strip() for x in re.split(r'[,\s]+', val) if x.strip()]

        if "fields" in d:
            d["fields"] = _split_csv(d["fields"])
        if "group_by" in d:
            d["group_by"] = _split_csv(d["group_by"])
        if "top" in d:
            try:
                d["top"] = int(re.findall(r'\d+', d["top"])[0])
            except Exception:
                d["top"] = None
        if "sort_dir" in d:
            d["sort_dir"] = "desc" if d["sort_dir"].lower().startswith("d") else "asc"
        if "format" in d:
            d["format"] = d["format"].lower()
        if "chart" in d:
            d["chart"] = d["chart"].lower()
        return d

    # ------------- Rendering -------------
    @staticmethod
    def _render_task(task: Dict[str, Any], analysis_results: Dict[str, Any]) -> str:
        title, d = task["title"], task["d"]
        section: List[str] = [f"## {title}\n"]

        # 1) Resolve data
        df, key_used, why = ScenarioEngine._resolve_df(d, analysis_results)
        if df is None:
            section.append("_No matching data for this task._")
            section.append(f"\n> Resolver note: {why}")
            return "\n".join(section)

        # 2) Filter
        if "filter" in d:
            mask = ScenarioEngine._safe_filter(df, d["filter"])
            df = df.loc[mask].copy()

        # 3) Derive columns
        if "derive" in d:
            df = ScenarioEngine._apply_derive(df, d["derive"])

        # 4) Group & aggregate
        if "group_by" in d or "agg" in d:
            df = ScenarioEngine._group_agg(df, d.get("group_by"), d.get("agg"))

        # 5) Pivot
        if "pivot" in d:
            df = ScenarioEngine._pivot(df, d["pivot"])

        # 6) Sort + Top
        if "sort_by" in d:
            asc = (d.get("sort_dir", "desc") == "asc")
            df = df.sort_values(by=d["sort_by"], ascending=asc)
        if isinstance(d.get("top"), int) and d["top"] > 0:
            df = df.head(d["top"])

        # 7) Fields selection
        if "fields" in d:
            cols = [c for c in d["fields"] if c in df.columns]
            if cols:
                df = df[cols]

        # 8) Render by format
        fmt = d.get("format", "table")
        if fmt == "list":
            section.append(ScenarioEngine._render_list(df))
        elif fmt == "comparison":
            section.append(ScenarioEngine._render_comparison(df))
        elif fmt == "map":
            section.append(ScenarioEngine._render_map(df))
        elif fmt == "narrative":
            section.append(ScenarioEngine._render_narrative(df))
        elif fmt == "chart":
            section.append(ScenarioEngine._render_chart_spec(df, d))
        else:
            section.append(ScenarioEngine._render_table(df))

        # 9) Per-task provenance (kept minimal)
        section.append("\n**Provenance**")
        section.append(f"- Data key: `{key_used}`")
        section.append(f"- Match note: {why}")

        return "\n".join(section)

    # ------------- Data resolution -------------
    @staticmethod
    def _resolve_df(d: Dict[str, Any], analysis_results: Dict[str, Any]) -> Tuple[Optional[pd.DataFrame], Optional[str], str]:
        # explicit key
        if "data_key" in d and d["data_key"] in analysis_results:
            return ScenarioEngine._as_df(analysis_results[d["data_key"]]), d["data_key"], "explicit data_key"

        # jaccard match on keys using hinted fields + any words in title/sort/agg
        hints = set()
        for k in ("fields", "sort_by"):
            v = d.get(k)
            if isinstance(v, list):
                hints |= set(v)
            elif isinstance(v, str):
                hints |= set(re.findall(r'[A-Za-z0-9_]+', v.lower()))
        best_key, best_score = None, 0.0
        for k in analysis_results:
            words = set(re.findall(r'[A-Za-z0-9_]+', k.lower()))
            if not words:
                continue
            inter = len(hints & words)
            union = len(hints | words) or 1
            score = inter / union
            if score > best_score:
                best_key, best_score = k, score

        if best_key:
            return ScenarioEngine._as_df(analysis_results[best_key]), best_key, f"keyword match (score={best_score:.2f})"

        # fallback: first list-of-dicts or dict-like
        for k, v in analysis_results.items():
            df = ScenarioEngine._as_df(v)
            if df is not None and not df.empty:
                return df, k, "fallback first structured"

        return None, None, "no suitable dataset found"

    @staticmethod
    def _as_df(v: Any) -> Optional[pd.DataFrame]:
        if isinstance(v, list):
            if not v:
                return pd.DataFrame()
            if isinstance(v[0], dict):
                return pd.DataFrame(v)
            return pd.DataFrame({"value": v})
        if isinstance(v, dict):
            # expand nested dicts into columns where sensible
            flat = {}
            any_scalar = False
            for k, val in v.items():
                if isinstance(val, (int, float, str, bool, type(None))):
                    flat[k] = [val]
                    any_scalar = True
            if any_scalar:
                return pd.DataFrame(flat)
            # complex dict -> try records
            recs = []
            for k, val in v.items():
                if isinstance(val, dict):
                    rec = {"item": k}
                    rec.update({kk: valv for kk, valv in val.items()})
                    recs.append(rec)
            if recs:
                return pd.DataFrame(recs)
        return None

    # ------------- Pipeline ops -------------
    @staticmethod
    def _safe_filter(df: pd.DataFrame, expr: str) -> pd.Series:
        try:
            s = _eval_series_expr(expr, df)
            if not isinstance(s, (pd.Series, np.ndarray)):
                raise ValueError("filter must evaluate to a boolean Series/array")
            return pd.Series(s).astype(bool).reindex(df.index, fill_value=False)
        except Exception as e:
            raise ValueError(f"Invalid filter expression: {e}")

    @staticmethod
    def _apply_derive(df: pd.DataFrame, spec: str) -> pd.DataFrame:
        # e.g., "load = patients / capacity, rate = 100*admits/pop"
        parts = re.split(r'[;,]\s*', spec)
        for p in parts:
            if not p.strip():
                continue
            if "=" not in p:
                raise ValueError(f"derive requires assignments: '{p}'")
            col, expr = p.split("=", 1)
            col = col.strip()
            expr = expr.strip()
            df[col] = _eval_series_expr(expr, df)
        return df

    @staticmethod
    def _parse_aggs(spec: Optional[str]) -> List[Tuple[str, str]]:
        """
        Returns list of (out_col, func_call_string), e.g. [("avg_wait_time","avg(wait_time)")]
        """
        if not spec:
            return []
        items = [x.strip() for x in spec.split(",") if x.strip()]
        out: List[Tuple[str, str]] = []
        for it in items:
            m = re.match(r'([a-zA-Z_][a-zA-Z0-9_]*)\s*\(([^)]+)\)', it)
            if not m:
                if it.lower() in ("count", "count(*)"):
                    out.append(("count", "count(*)"))
                    continue
                raise ValueError(f"Bad agg item: '{it}' (use avg(x), median(y), p90(z), sum(a), count(*))")
            func = m.group(1)
            arg = m.group(2).strip()
            out_col = f"{func.lower()}_{arg}"
            out.append((out_col, f"{func}({arg})"))
        return out

    @staticmethod
    def _group_agg(df: pd.DataFrame, group_by: Optional[List[str]], agg_spec: Optional[str]) -> pd.DataFrame:
        aggs = ScenarioEngine._parse_aggs(agg_spec)
        if not aggs and not group_by:
            return df
        if not group_by:
            # reduce to single row with requested aggs
            res = {}
            for out_col, call in aggs:
                val = ScenarioEngine._apply_agg_call(df, call)
                res[out_col] = val
            return pd.DataFrame([res])
        # grouped
        gb = df.groupby(group_by, dropna=False)
        rows = []
        for keys, g in gb:
            if not isinstance(keys, tuple):
                keys = (keys,)
            rec = {group_by[i]: keys[i] for i in range(len(group_by))}
            for out_col, call in aggs:
                rec[out_col] = ScenarioEngine._apply_agg_call(g, call)
            if not aggs:
                # no aggs? carry counts by default
                rec["count"] = len(g)
            rows.append(rec)
        return pd.DataFrame(rows)

    @staticmethod
    def _apply_agg_call(df: pd.DataFrame, call: str) -> Any:
        call = call.strip()
        if call.lower() in ("count", "count(*)"):
            return int(len(df))
        m = re.match(r'([a-zA-Z_][a-zA-Z0-9_]*)\s*\(([^)]+)\)', call)
        if not m:
            raise ValueError(f"Bad agg call: {call}")
        func, arg = m.group(1).lower(), m.group(2).strip()
        if arg not in df.columns:
            raise ValueError(f"Unknown column in agg: {arg}")
        col = df[arg].dropna()
        if func in ("avg", "mean"):
            return float(np.mean(col)) if len(col) else float("nan")
        if func == "median":
            return float(np.median(col)) if len(col) else float("nan")
        if func == "sum":
            return float(np.sum(col)) if len(col) else 0.0
        if func in ("min", "max"):
            f = getattr(np, func)
            return float(f(col)) if len(col) else float("nan")
        if func.startswith("p") and func[1:].isdigit():
            q = int(func[1:])
            return float(np.percentile(col, q)) if len(col) else float("nan")
        raise ValueError(f"Unsupported agg function: {func}")

    @staticmethod
    def _pivot(df: pd.DataFrame, spec: str) -> pd.DataFrame:
        # spec: index=a[,b] columns=c values=v
        parts = dict(re.findall(r'(\w+)\s*=\s*([^\s,]+)', spec))
        idx = parts.get("index")
        cols = parts.get("columns")
        vals = parts.get("values")
        if not (idx and cols and vals):
            raise ValueError("pivot requires 'index=.. columns=.. values=..'")
        idx = [x.strip() for x in idx.split(",")]
        pv = df.pivot_table(index=idx, columns=cols, values=vals, aggfunc="first")
        pv = pv.reset_index()
        # flatten columns if needed
        if isinstance(pv.columns, pd.MultiIndex):
            pv.columns = ["_".join([str(c) for c in tup if c != ""]) for tup in pv.columns]
        return pv

    # ------------- Output renderers -------------
    @staticmethod
    def _render_table(df: pd.DataFrame) -> str:
        if df.empty:
            return "_No rows to display._"
        # convert all to string-friendly
        dff = df.copy()
        for c in dff.columns:
            dff[c] = dff[c].apply(lambda v: ScenarioEngine._fmt_val(v))
        header = "| " + " | ".join(dff.columns) + " |"
        sep = "|" + "|".join(["---"] * len(dff.columns)) + "|"
        rows = ["| " + " | ".join(map(str, r)) + " |" for r in dff.to_numpy().tolist()]
        return "\n".join([header, sep, *rows])

    @staticmethod
    def _render_list(df: pd.DataFrame) -> str:
        if df.empty:
            return "_No items._"
        # pick first column as primary
        primary = df.columns[0]
        lines = []
        for i, row in enumerate(df.itertuples(index=False), 1):
            parts = []
            for c, v in zip(df.columns, row):
                if c == primary:
                    continue
                parts.append(f"{c}: {ScenarioEngine._fmt_val(v)}")
            extra = f" ({', '.join(parts)})" if parts else ""
            lines.append(f"{i}. {ScenarioEngine._fmt_val(getattr(row, primary))}{extra}")
        return "\n".join(lines)

    @staticmethod
    def _render_comparison(df: pd.DataFrame) -> str:
        # look for columns named like current/previous
        cols = {c.lower(): c for c in df.columns}
        cur = cols.get("current") or cols.get("now") or cols.get("value")
        prev = cols.get("previous") or cols.get("prior") or cols.get("past")
        name = cols.get("name") or cols.get("metric") or cols.get("item") or df.columns[0]
        if not (cur and prev):
            return "_Comparison format requires columns 'current' and 'previous' (or aliases)._"
        header = "| Item | Current | Previous | Change |"
        sep = "|---|---:|---:|---:|"
        body = []
        for _, r in df.iterrows():
            c, p = r[cur], r[prev]
            change = (c - p) if isinstance(c, (int, float)) and isinstance(p, (int, float)) else "N/A"
            body.append(f"| {ScenarioEngine._fmt_val(r[name])} | {ScenarioEngine._fmt_val(c)} | {ScenarioEngine._fmt_val(p)} | {ScenarioEngine._fmt_val(change)} |")
        return "\n".join([header, sep, *body])

    @staticmethod
    def _render_map(df: pd.DataFrame) -> str:
        # simple location table
        colmap = {c.lower(): c for c in df.columns}
        name = colmap.get("name") or colmap.get("facility") or colmap.get("title") or df.columns[0]
        zone = colmap.get("zone")
        city = colmap.get("city")
        region = colmap.get("region")
        lat = colmap.get("latitude") or colmap.get("lat")
        lon = colmap.get("longitude") or colmap.get("lon")
        cols = [x for x in [name, city, region, zone, lat, lon] if x]
        if not cols:
            return "_No geographic fields to show._"
        dff = df[cols].copy()
        dff["coordinates"] = np.where((lat is not None) & (lon is not None) & dff[lat].notna() & dff[lon].notna(),
                                      dff[lat].astype(str) + ", " + dff[lon].astype(str), "N/A")
        show = [name, city or "city", region or "region", zone or "zone", "coordinates"]
        # ensure all exist
        for c in show:
            if c not in dff.columns:
                dff[c] = ""
        dff = dff[show]
        return ScenarioEngine._render_table(dff)

    @staticmethod
    def _render_narrative(df: pd.DataFrame) -> str:
        if df.empty:
            return "_No content._"
        paras = []
        for i, row in enumerate(df.to_dict(orient="records"), 1):
            parts = [f"**{k}**: {ScenarioEngine._fmt_val(v)}" for k, v in row.items()]
            paras.append(f"{i}. " + "; ".join(parts))
        return "\n".join(paras)

    @staticmethod
    def _render_chart_spec(df: pd.DataFrame, d: Dict[str, Any]) -> str:
        """
        Emits a Vega-Lite spec in a fenced code block that downstream renderers can plot exactly.
        Accepts: chart (bar|line|area|point), x, y, color, column (facet)
        """
        mark = d.get("chart", "bar")
        spec = {
            "$schema": "https://vega.github.io/schema/vega-lite/v5.json",
            "description": d.get("title") or "Chart",
            "data": {"values": df.to_dict(orient="records")},
            "mark": mark,
            "encoding": {}
        }
        for enc in ("x", "y", "color", "column"):
            if enc in d and d[enc] in df.columns:
                spec["encoding"][enc] = {"field": d[enc], "type": "quantitative" if pd.api.types.is_numeric_dtype(df[d[enc]]) else "nominal"}
        return "```vega-lite\n" + json.dumps(spec, ensure_ascii=False, indent=2) + "\n```"

    # ------------- Helpers -------------
    @staticmethod
    def _fmt_val(v: Any) -> str:
        if isinstance(v, float):
            if math.isnan(v):
                return "NaN"
            return f"{v:,.4g}"
        if isinstance(v, (int, np.integer)):
            return f"{int(v):,}"
        return str(v)