Update scenario_engine.py

scenario_engine.py

@@ -1,37 +1,277 @@
-import pandas as pd
+from __future__ import annotations
+from typing import Dict, List, Any, Tuple, Union, Optional
+import re, math, json, ast
 import numpy as np
-import math, json, re, ast
-from schemas import ScenarioPlan, TaskSpec
+import pandas as pd
+from schema import ScenarioPlan, TaskPlan
+from column_resolver import resolve_cols
+
+_ALLOWED_FUNCS = {
+    "abs": abs, "round": round, "sqrt": math.sqrt, "log": math.log, "exp": math.exp,
+    "min": np.minimum, "max": np.maximum,
+    "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),
+}
+
+class _SafeExpr(ast.NodeTransformer):
+    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: {node.id}")
+        return node
+    def visit_Call(self, node):
+        if not isinstance(node.func, ast.Name): raise ValueError("Only simple calls allowed")
+        if node.func.id not in _ALLOWED_FUNCS: raise ValueError(f"Function not allowed: {node.func.id}")
+        return self.generic_visit(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.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:
+    names = set(df.columns) | {"True","False","None"}
+    tree = ast.parse(expr, mode="eval")
+    _SafeExpr(names).visit(tree)
+    code = compile(tree, "<expr>", "eval")
+    env = {**{k: df[k] for k in df.columns}, **_ALLOWED_FUNCS}
+    return eval(code, {"__builtins__": {}}, env)
 
 class ScenarioEngine:
     @staticmethod
-    def render_plan(plan: ScenarioPlan, results: dict) -> str:
-        sections = ["# Scenario Output\n"]
+    def _as_df(v: Any) -> Optional[pd.DataFrame]:
+        if isinstance(v, list):
+            if not v: return pd.DataFrame()
+            return pd.DataFrame(v) if isinstance(v[0], dict) else pd.DataFrame({"value": v})
+        if isinstance(v, dict):
+            if any(isinstance(val, (int, float, str, bool, type(None))) for val in v.values()):
+                return pd.DataFrame([v])
+            rows = []
+            for k, val in v.items():
+                if isinstance(val, dict):
+                    rec = {"item": k}; rec.update(val); rows.append(rec)
+            if rows: return pd.DataFrame(rows)
+        if isinstance(v, pd.DataFrame): return v
+        return None
+
+    # ---------- Plan-first API ----------
+    @staticmethod
+    def execute_plan(plan: ScenarioPlan, datasets: Dict[str, Any]) -> str:
+        sections: List[str] = ["# Scenario Output\n"]
         for t in plan.tasks:
-            sections.append(ScenarioEngine._render_task(t, results))
-        return "\n".join(sections)
+            sections.append(ScenarioEngine._exec_task(t, datasets))
+        return "\n".join(sections).strip()
 
     @staticmethod
-    def _df(v):
-        if isinstance(v, pd.DataFrame): return v
-        if isinstance(v, list) and v and isinstance(v[0], dict): return pd.DataFrame(v)
-        if isinstance(v, dict): return pd.DataFrame([v])
-        return None
+    def _get_df(datasets: Dict[str, Any], key: Optional[str]) -> Optional[pd.DataFrame]:
+        if key and key in datasets:
+            v = datasets[key]
+        else:
+            v = next((vv for vv in datasets.values() if isinstance(vv, (list, dict, pd.DataFrame))), None)
+        return ScenarioEngine._as_df(v) if v is not None else None
 
     @staticmethod
-    def _render_task(t: TaskSpec, results: dict) -> str:
-        out=[f"## {t.title}\n"]
-        df=ScenarioEngine._df(results.get(t.data_key)) if t.data_key else None
-        if df is None: return "\n".join(out+["_No data_"])
-        if t.filter: df=df.query(t.filter)
-        if t.group_by or t.agg: df=df.groupby(t.group_by).agg("first").reset_index()
-        if t.sort_by: df=df.sort_values(by=t.sort_by, ascending=(t.sort_dir=="asc"))
-        if t.top: df=df.head(t.top)
-        if t.fields: df=df[t.fields]
+    def _apply_filter(df: pd.DataFrame, expr: str) -> pd.DataFrame:
+        m = _eval_series_expr(expr, df)
+        return df.loc[m.astype(bool)].copy()
 
-        if t.format=="list":
-            out += [f"- {row.to_dict()}" for _,row in df.iterrows()]
+    @staticmethod
+    def _apply_derive(df: pd.DataFrame, spec: str) -> pd.DataFrame:
+        parts = re.split(r'[;,]\s*', spec)
+        for p in parts:
+            if not p.strip(): continue
+            if "=" not in p: raise ValueError(f"derive requires col=expr: '{p}'")
+            col, expr = p.split("=", 1); df[col.strip()] = _eval_series_expr(expr.strip(), df)
+        return df
+
+    @staticmethod
+    def _parse_aggs(spec: Optional[str]) -> List[Tuple[str, str]]:
+        if not spec: return []
+        items = [x.strip() for x in spec.split(",") if x.strip()]
+        out = []
+        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: {it}")
+            func, arg = m.group(1).lower(), m.group(2).strip()
+            out.append((f"{func}_{arg}", f"{func}({arg})"))
+        return out
+
+    @staticmethod
+    def _apply_agg_call(df: pd.DataFrame, call: str):
+        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)
+        func, arg = m.group(1).lower(), m.group(2).strip()
+        if arg not in df.columns: raise ValueError(f"Unknown column: {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"): return float(getattr(np, func)(col)) if len(col) else float("nan")
+        if func.startswith("p") and func[1:].isdigit(): return float(np.percentile(col, int(func[1:]))) if len(col) else float("nan")
+        raise ValueError(f"Unsupported agg: {func}")
+
+    @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:
+            return pd.DataFrame([{k: ScenarioEngine._apply_agg_call(df, call) for k, call in aggs}])
+        rows = []
+        gb = df.groupby(group_by, dropna=False)
+        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))}
+            if aggs:
+                for out_col, call in aggs: rec[out_col] = ScenarioEngine._apply_agg_call(g, call)
+            else:
+                rec["count"] = len(g)
+            rows.append(rec)
+        return pd.DataFrame(rows)
+
+    @staticmethod
+    def _pivot(df: pd.DataFrame, spec: str) -> pd.DataFrame:
+        parts = dict(re.findall(r'(\w+)\s*=\s*([^\s,]+)', spec))
+        idx = [x.strip() for x in parts.get("index","").split(",") if x.strip()]
+        cols = parts.get("columns"); vals = parts.get("values")
+        if not (idx and cols and vals): raise ValueError("pivot requires index=.. columns=.. values=..")
+        pv = df.pivot_table(index=idx, columns=cols, values=vals, aggfunc="first").reset_index()
+        if isinstance(pv.columns, pd.MultiIndex):
+            pv.columns = ["_".join([str(c) for c in tup if c!=""]) for tup in pv.columns]
+        return pv
+
+    @staticmethod
+    def _render_table(df: pd.DataFrame) -> str:
+        if df.empty: return "_No rows._"
+        dff = df.copy()
+        for c in dff.columns:
+            dff[c] = dff[c].apply(lambda v: "NaN" if (isinstance(v,float) and math.isnan(v)) else f"{v:,.4g}" if isinstance(v,float) else 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._"
+        primary = df.columns[0]
+        lines = []
+        for i, row in enumerate(df.itertuples(index=False), 1):
+            extras = [f"{c}: {getattr(row,c)}" for c in df.columns if c != primary]
+            lines.append(f"{i}. {getattr(row, primary)}" + (f" ({', '.join(extras)})" if extras else ""))
+        return "\n".join(lines)
+
+    @staticmethod
+    def _render_comparison(df: pd.DataFrame) -> str:
+        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 requires 'current' and 'previous' columns._"
+        header = "| Item | Current | Previous | Change |"; sep="|---|---:|---:|---:|"; body=[]
+        for _, r in df.iterrows():
+            c, p = r[cur], r[prev]
+            ch = (c - p) if isinstance(c,(int,float)) and isinstance(p,(int,float)) else "N/A"
+            body.append(f"| {r[name]} | {c} | {p} | {ch} |")
+        return "\n".join([header, sep, *body])
+
+    @staticmethod
+    def _render_map(df: pd.DataFrame) -> str:
+        col = {c.lower(): c for c in df.columns}
+        name = col.get("facility") or col.get("name") or df.columns[0]
+        lat = col.get("latitude") or col.get("lat"); lon = col.get("longitude") or col.get("lon")
+        zone = col.get("zone"); city = col.get("city")
+        show = [x for x in [name, city, zone, lat, lon] if x]
+        if not show: return "_No geographic fields._"
+        tmp = df[show].copy()
+        if lat and lon:
+            tmp["coordinates"] = tmp[lat].astype(str) + ", " + tmp[lon].astype(str)
+            show = [name, city or "city", zone or "zone", "coordinates"]
+        return ScenarioEngine._render_table(tmp[show])
+
+    @staticmethod
+    def _render_chart(df: pd.DataFrame, d: Dict[str, Any]) -> str:
+        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```"
+
+    @staticmethod
+    def _exec_task(t: TaskPlan, datasets: Dict[str, Any]) -> str:
+        section = [f"## {t.title_override or t.title}\n"]
+        df = ScenarioEngine._get_df(datasets, t.data_key)
+        if df is None or df.empty:
+            section += ["_No matching data for this task._", "\n**Provenance**", f"- Data key: `{t.data_key or 'auto'}`"]
+            return "\n".join(section)
+
+        if t.filter: df = ScenarioEngine._apply_filter(df, t.filter)
+        if t.derive:
+            for d in t.derive: df = ScenarioEngine._apply_derive(df, d)
+        if t.joins:
+            for j in t.joins:
+                rk, lo, ro, how = j["right_key"], j["left_on"], j["right_on"], j.get("how","left").lower()
+                r = ScenarioEngine._as_df(datasets.get(rk))
+                if r is not None:
+                    df = df.merge(r, left_on=lo, right_on=ro, how=how)
+        if t.group_by or t.agg:
+            df = ScenarioEngine._group_agg(df, t.group_by, ", ".join(t.agg or []))
+        if t.pivot:
+            spec = t.pivot
+            df = ScenarioEngine._pivot(df, f"index={','.join(spec.get('index', []))} columns={spec['columns']} values={spec['values']}")
+        if t.sort_by and t.sort_by in df.columns:
+            df = df.sort_values(by=t.sort_by, ascending=(t.sort_dir or "desc").lower()=="asc")
+        if t.top and t.top>0: df = df.head(t.top)
+        if t.fields:
+            cols = resolve_cols(t.fields, df.columns.tolist())
+            cols = [c for c in cols if c in df.columns]
+            if cols: df = df[cols]
+
+        if t.number_format:
+            for col, fmt in t.number_format.items():
+                if col in df.columns:
+                    if fmt.endswith("%"):
+                        decimals = len(fmt.split(".")[-1].rstrip("%")) if "." in fmt else 0
+                        df[col] = (df[col].astype(float) * 100).round(decimals).astype(str) + "%"
+                    else:
+                        try:
+                            decimals = int(fmt.split(".")[-1]) if "." in fmt else 0
+                            df[col] = df[col].astype(float).round(decimals)
+                        except Exception:
+                            pass
+
+        fmt = (t.format or "table").lower()
+        if fmt == "list": body = ScenarioEngine._render_list(df)
+        elif fmt == "comparison": body = ScenarioEngine._render_comparison(df)
+        elif fmt == "map": body = ScenarioEngine._render_map(df)
+        elif fmt == "chart":
+            enc = t.encodings or {}
+            d = {"chart": t.chart or "bar", **enc}
+            body = ScenarioEngine._render_chart(df, d)
+        elif fmt == "narrative":
+            lines = []
+            for i, rec in enumerate(df.to_dict(orient="records"), 1):
+                parts = [f"**{k}**: {v}" for k, v in rec.items()]
+                lines.append(f"{i}. " + "; ".join(parts))
+            body = "\n".join(lines) if lines else "_No content._"
         else:
-            out.append(df.to_markdown(index=False))
-        return "\n".join(out)
+            body = ScenarioEngine._render_table(df)
+
+        section.append(body)
+        section.append("\n**Provenance**")
+        section.append(f"- Data key: `{t.data_key or 'auto'}`")
+        return "\n".join(section)