Blame for datainventor/visualize.py

"""Visualize: ask a question about a dataset, get a chart and the rows behind every mark.

Built for small datasets first, the 25 to 200 rows a hosted run writes. Every row stays in
memory, charts are exact counts over all rows (nothing is sampled), and when a question needs
something no column has (tone, whether a reply gives steps), a model reads the rows in a few
batched calls and the answer becomes a new column that later questions can reuse.
"""
from __future__ import annotations

import json
import math
import re
import statistics
import time
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
from typing import Callable, Optional

ID_NAMES = {"id", "_id", "idx", "index", "row_idx", "uuid", "row_id"}
CLASS_NAMES = {"label", "labels", "class", "category", "target", "sentiment", "rating", "stars", "grade"}
MAX_COLUMNS = 30
MAX_CELL = 4000           # characters kept per cell
LABEL_CHARS = 600         # characters of each row a labelling call reads
LABEL_BATCH = 8
MAX_GROUPS = 12           # categories on one axis before the rest fold into "Other"
MAX_COLORS = 6
KINDS = ("label", "score", "pattern", "count")
HEDGE_AFTER = 8.0         # seconds before a slow call gets a twin; the first good answer wins
_HEDGE = ThreadPoolExecutor(max_workers=24)
CHARTS = ("bar", "stack", "box", "hist", "scatter")


class PlanError(ValueError):
    pass


# ------------------------------------------------------------------ rows and columns
def _turns(v) -> Optional[list]:
    """Chat-style cells: [{"role", "content"}] or ShareGPT [{"from", "value"}]."""
    if isinstance(v, list) and v and all(isinstance(x, dict) for x in v):
        if all("content" in x for x in v):
            return [(str(x.get("role") or "turn"), x.get("content")) for x in v]
        if all("value" in x for x in v):
            return [(str(x.get("from") or "turn"), x.get("value")) for x in v]
    return None


def _cell(v):
    if v is None:
        return ""
    if isinstance(v, bool):
        return "yes" if v else "no"
    if isinstance(v, (int, float)):
        return v if math.isfinite(v) else ""
    if isinstance(v, str):
        return v[:MAX_CELL]
    t = _turns(v)
    if t:
        return "\n".join(f"{r}: {c if isinstance(c, str) else json.dumps(c, ensure_ascii=False)}" for r, c in t)[:MAX_CELL]
    return json.dumps(v, ensure_ascii=False)[:MAX_CELL]


def _num(v) -> Optional[float]:
    if isinstance(v, bool):
        return None
    if isinstance(v, (int, float)):
        return float(v) if math.isfinite(v) else None
    if isinstance(v, str) and v.strip():
        try:
            f = float(v.strip().replace(",", ""))
            return f if math.isfinite(f) else None
        except ValueError:
            return None
    return None


def words(text) -> int:
    return len(str(text or "").split())


def prepare(raw: list[dict]) -> tuple[list[dict], dict]:
    """Normalize raw rows and type every column: text, category or number.

    Returns (rows, columns). Each row gets `_id` (1-based). Long text columns get a derived
    "<name> words" number column, and chat columns a "<name> turns" one.
    """
    names: list[str] = []
    for r in raw:
        for k in r:
            if k not in names and str(k).lower() not in ID_NAMES and not str(k).startswith("_"):
                names.append(k)
    names = names[:MAX_COLUMNS]
    rows, turns = [], {}
    for i, r in enumerate(raw, 1):
        row = {"_id": i}
        for k in names:
            v = r.get(k)
            t = _turns(v)
            if t:
                turns.setdefault(k, {})[i] = len(t)
            row[k] = _cell(v)
        rows.append(row)
    cols: dict[str, dict] = {}
    n = len(rows)
    for k in names:
        vals = [row[k] for row in rows if row[k] != ""]
        if not vals:
            continue
        nums = [_num(v) for v in vals]
        few = all(x is not None for x in nums) and all(float(x).is_integer() for x in nums) and \
            (len(set(nums)) <= 2 or (str(k).lower() in CLASS_NAMES and len(set(nums)) <= 12))
        if few:  # 0/1 flags and class ids are categories, not amounts
            for row in rows:
                row[k] = str(int(_num(row[k]))) if row[k] != "" else ""
            cols[k] = {"type": "cat"}
            continue
        if all(x is not None for x in nums) and not all(isinstance(v, str) and len(v) > 12 for v in vals):
            for row in rows:
                row[k] = _num(row[k]) if row[k] != "" else None
            cols[k] = {"type": "num"}
            continue
        strs = [str(v) for v in vals]
        avg = sum(len(s) for s in strs) / len(strs)
        distinct = len(set(strs))
        if distinct == len(strs) and len(strs) > 1 and not any(" " in x for x in strs) and k not in turns:
            continue  # hashes, ids, urls: unique and never a phrase
        if avg > 60 or k in turns:
            cols[k] = {"type": "text"}
        elif distinct <= 30 and (distinct <= max(6, 0.6 * len(strs)) if len(strs) > 6 else distinct < len(strs)):
            cols[k] = {"type": "cat"}
        elif avg >= 12:
            cols[k] = {"type": "text"}
        else:
            continue  # short and unique: names, ids, codes
        for row in rows:
            if not isinstance(row[k], str):
                row[k] = str(row[k])
    for k, c in list(cols.items()):
        if c["type"] != "text":
            continue
        if k in turns:
            name = f"{k} turns"
            for row in rows:
                row[name] = float(turns[k].get(row["_id"], 0))
            cols[name] = {"type": "num", "derived": "turns", "from": k}
        avg_words = sum(words(row[k]) for row in rows) / max(1, n)
        if avg_words >= 8:
            name = f"{k} words"
            for row in rows:
                row[name] = float(words(row[k]))
            cols[name] = {"type": "num", "derived": "words", "from": k}
    order = {"text": 0, "cat": 1, "num": 2}
    cols = dict(sorted(cols.items(), key=lambda kv: (order[kv[1]["type"]], bool(kv[1].get("derived")))))
    return rows, cols


def describe(rows: list[dict], cols: dict) -> list[dict]:
    """Per-column facts for the UI and the planner."""
    out = []
    for k, c in cols.items():
        d = {"name": k, "type": c["type"], "derived": c.get("derived", ""), "from": c.get("from", "")}
        vals = [r.get(k) for r in rows]
        if c["type"] == "cat":
            counts: dict[str, int] = {}
            for v in vals:
                counts[v or "(blank)"] = counts.get(v or "(blank)", 0) + 1
            d["values"] = sorted(counts.items(), key=lambda kv: -kv[1])[:12]
            d["distinct"] = len(counts)
        elif c["type"] == "num":
            xs = [v for v in vals if isinstance(v, (int, float))]
            if xs:
                d.update(min=min(xs), max=max(xs), median=statistics.median(xs))
        else:
            xs = [words(v) for v in vals if v]
            d["avg_words"] = round(sum(xs) / max(1, len(xs)))
            d["example"] = next((str(v)[:140] for v in vals if v), "")
        if c.get("labels"):
            d["labels"] = c["labels"]
        if c.get("guide"):
            d["guide"] = c["guide"]
        out.append(d)
    return out


# ------------------------------------------------------------------ charts
def _q(xs: list[float], p: float) -> float:
    s = sorted(xs)
    if len(s) == 1:
        return s[0]
    pos = (len(s) - 1) * p
    lo = int(pos)
    return s[lo] + (s[min(lo + 1, len(s) - 1)] - s[lo]) * (pos - lo)


def _nice(x: float) -> float:
    if x <= 0:
        return 1.0
    e = 10 ** math.floor(math.log10(x))
    for m in (1, 2, 2.5, 5, 10):
        if x <= m * e:
            return m * e
    return 10 * e


def _key(v) -> str:
    return str(v) if v not in ("", None) else "(blank)"


def _fold(groups: dict[str, list[int]], limit: int) -> list[tuple[str, list[int]]]:
    items = sorted(groups.items(), key=lambda kv: -len(kv[1]))
    if len(items) <= limit:
        return items
    keep, rest = items[:limit - 1], items[limit - 1:]
    ids = [i for _, g in rest for i in g]
    return keep + [(f"Other ({len(rest)})", ids)]


def chart(rows: list[dict], cols: dict, spec: dict) -> dict:
    """Exact aggregation of `rows` for one chart spec. Every mark carries the ids of its rows."""
    t, x, y, color, agg = spec.get("type"), spec.get("x"), spec.get("y"), spec.get("color"), spec.get("agg") or "count"
    by = {r["_id"]: r for r in rows}
    out = {"type": t, "n": len(rows), "x": x, "y": y, "color": color, "agg": agg}

    def groups_of(col):
        g: dict[str, list[int]] = {}
        for r in rows:
            g.setdefault(_key(r.get(col)), []).append(r["_id"])
        return g

    def nums(ids, col):
        return [(i, by[i][col]) for i in ids if isinstance(by[i].get(col), (int, float))]

    if t == "bar":
        marks = []
        for k, ids in _fold(groups_of(x), MAX_GROUPS):
            if y and agg in ("mean", "median"):
                pts = nums(ids, y)
                if not pts:
                    continue
                vs = [v for _, v in pts]
                val = statistics.fmean(vs) if agg == "mean" else statistics.median(vs)
                marks.append({"key": k, "value": round(val, 1 if abs(val) >= 10 else 2), "n": len(pts), "ids": [i for i, _ in pts]})
            else:
                marks.append({"key": k, "value": len(ids), "n": len(ids), "ids": ids})
        marks.sort(key=lambda m: -m["value"])
        out["marks"] = marks
    elif t == "stack":
        colors = [k for k, _ in _fold(groups_of(color), MAX_COLORS)]
        color_ids = dict(_fold(groups_of(color), MAX_COLORS))
        rev = {i: k for k, ids in color_ids.items() for i in ids}
        marks = []
        for k, ids in _fold(groups_of(x), MAX_GROUPS):
            parts = []
            for c in colors:
                pids = [i for i in ids if rev.get(i) == c]
                if pids:
                    parts.append({"key": c, "value": len(pids), "ids": pids})
            marks.append({"key": k, "n": len(ids), "ids": ids, "parts": parts})
        marks.sort(key=lambda m: -m["n"])
        out.update(marks=marks, colors=colors)
    elif t == "box":
        gs = _fold(groups_of(x), MAX_GROUPS) if x else [("All rows", [r["_id"] for r in rows])]
        marks = []
        for k, ids in gs:
            pts = nums(ids, y)
            if not pts:
                continue
            vs = [v for _, v in pts]
            q1, med, q3 = _q(vs, .25), _q(vs, .5), _q(vs, .75)
            iqr = q3 - q1
            lo = min(v for v in vs if v >= q1 - 1.5 * iqr)
            hi = max(v for v in vs if v <= q3 + 1.5 * iqr)
            marks.append({"key": k, "n": len(pts), "q1": q1, "median": med, "q3": q3, "lo": lo, "hi": hi,
                          "ids": [i for i, _ in pts], "points": [[i, v] for i, v in pts]})
        marks.sort(key=lambda m: -m["median"])
        out["marks"] = marks
    elif t == "hist":
        pts = nums([r["_id"] for r in rows], x)
        vs = [v for _, v in pts]
        if vs and all(float(v).is_integer() for v in vs) and max(vs) - min(vs) <= 12:
            # small whole numbers (scores, counts): one bar per value
            out["discrete"] = True
            out["marks"] = [{"x0": v, "x1": v + 1, "value": len(ids), "ids": ids}
                            for v in range(int(min(vs)), int(max(vs)) + 1) for ids in [[i for i, x in pts if x == v]]]
        elif vs:
            lo, hi = min(vs), max(vs)
            k = 6 if len(vs) <= 30 else min(14, int(math.sqrt(len(vs))) + 2)
            step = _nice((hi - lo) / k) if hi > lo else 1.0
            start = math.floor(lo / step) * step
            nb = max(1, min(40, math.ceil((hi - start) / step + 1e-9)))
            if start + nb * step <= hi:
                nb += 1
            bins = [{"x0": round(start + b * step, 4), "x1": round(start + (b + 1) * step, 4), "value": 0, "ids": []} for b in range(nb)]
            for i, v in pts:
                b = min(nb - 1, int((v - start) / step))
                bins[b]["ids"].append(i)
                bins[b]["value"] += 1
            out["marks"] = bins
        else:
            out["marks"] = []
    elif t == "scatter":
        rev = {}
        if color:
            for k, ids in _fold(groups_of(color), MAX_COLORS):
                for i in ids:
                    rev[i] = k
            out["colors"] = [k for k, _ in _fold(groups_of(color), MAX_COLORS)]
        marks = [{"id": r["_id"], "x": r[x], "y": r[y], "c": rev.get(r["_id"], "")} for r in rows
                 if isinstance(r.get(x), (int, float)) and isinstance(r.get(y), (int, float))]
        out["marks"] = marks
        if len(marks) >= 3:
            xs, ys = [m["x"] for m in marks], [m["y"] for m in marks]
            try:
                out["r"] = round(statistics.correlation(xs, ys), 2)
            except statistics.StatisticsError:
                pass
    else:
        raise PlanError(f"Unknown chart type {t!r}.")
    return out


def check_spec(spec: dict, cols: dict) -> dict:
    """Validate a chart spec against the columns, fixing the mistakes that have one obvious fix."""
    s = {k: spec.get(k) for k in ("type", "x", "y", "color", "agg")}
    for k in ("x", "y", "color"):
        if s[k] in ("", "null", "none", "None"):
            s[k] = None
    typ = lambda c: (cols.get(c) or {}).get("type")
    for k in ("x", "y", "color"):
        if s[k] is not None and s[k] not in cols:
            raise PlanError(f"There is no column named {s[k]!r}. Columns: {', '.join(cols)}.")
    t = s["type"]
    if t not in CHARTS:
        raise PlanError(f"chart.type must be one of {', '.join(CHARTS)}.")
    if t == "bar" and typ(s["x"]) == "num":
        t, s["y"] = "hist", None
    if t == "box" and s["y"] is None and typ(s["x"]) == "num":
        s["x"], s["y"] = None, s["x"]
    if t == "stack" and not s["color"]:
        t = "bar"
    if t == "bar":
        if typ(s["x"]) not in ("cat",):
            raise PlanError("A bar chart needs x to be a category column.")
        if s["y"] and typ(s["y"]) != "num":
            raise PlanError("A bar chart's y must be a number column, or null to count rows.")
        s["agg"] = (s["agg"] if s["agg"] in ("mean", "median") else "mean") if s["y"] else "count"
    elif t == "stack":
        if typ(s["x"]) != "cat" or typ(s["color"]) != "cat":
            raise PlanError("A stacked chart needs x and color to be category columns.")
        s["y"], s["agg"] = None, s["agg"] if s["agg"] in ("share", "count") else "share"
    elif t == "box":
        if typ(s["y"]) != "num":
            raise PlanError("A box plot needs y to be a number column.")
        if s["x"] and typ(s["x"]) != "cat":
            raise PlanError("A box plot's x must be a category column or null.")
    elif t == "hist":
        if typ(s["x"]) != "num":
            raise PlanError("A histogram needs x to be a number column.")
    elif t == "scatter":
        if typ(s["x"]) != "num" or typ(s["y"]) != "num":
            raise PlanError("A scatter plot needs x and y to be number columns.")
        if s["color"] and typ(s["color"]) != "cat":
            s["color"] = None
    s["type"] = t
    return s


# ------------------------------------------------------------------ instant views (no model)
def _main_cat(cols: dict, rows: list[dict]) -> Optional[str]:
    prefer = ("area", "category", "label", "topic", "subtopic", "department", "language")
    cats = [k for k, c in cols.items() if c["type"] == "cat" and len({r.get(k) for r in rows}) > 1]
    for p in prefer:
        if p in cats:
            return p
    return min(cats, key=lambda k: abs(len({r.get(k) for r in rows}) - 5)) if cats else None


def _main_len(cols: dict) -> Optional[str]:
    prefer = ("completion words", "chosen words", "response words", "output words", "answer words", "text words")
    nums = [k for k, c in cols.items() if c.get("derived") == "words"]
    for p in prefer:
        if p in nums:
            return p
    return nums[-1] if nums else None


def glance(rows: list[dict], cols: dict) -> list[dict]:
    """Two or three charts that need no model, so a dataset shows something the moment it loads."""
    out = []
    cat, ln = _main_cat(cols, rows), _main_len(cols)
    if cat:
        out.append({"title": f"Rows by {cat}", **chart(rows, cols, {"type": "bar", "x": cat})})
    if ln:
        spec = {"type": "box", "x": cat, "y": ln} if cat else {"type": "hist", "x": ln}
        out.append({"title": f"{ln.capitalize()}" + (f" by {cat}" if cat else ""), **chart(rows, cols, spec)})
    if not cat:
        more = [k for k, c in cols.items() if c.get("derived") in ("words", "turns") and k != ln]
        if more:
            out.append({"title": more[0].capitalize(), **chart(rows, cols, {"type": "hist", "x": more[0]})})
    second = [k for k, c in cols.items() if c["type"] == "cat" and k != cat and 1 < len({r.get(k) for r in rows}) <= 8]
    if second:
        k = "language" if "language" in second else "difficulty" if "difficulty" in second else second[0]
        out.append({"title": f"Rows by {k}", **chart(rows, cols, {"type": "bar", "x": k})})
    return out[:3]


def _pl(name: str) -> str:
    return name if name.endswith("s") else name + "s"


def presets(rows: list[dict], cols: dict) -> dict[str, dict]:
    """Starter questions with their plans already written, so clicking one skips the planning call."""
    cat, ln = _main_cat(cols, rows), _main_len(cols)
    texts = [k for k, c in cols.items() if c["type"] == "text"]
    reply = next((k for k in ("completion", "chosen", "response", "output", "answer", "messages", "conversations") if k in texts), None)
    ask = next((k for k in ("prompt", "instruction", "question", "input", "query") if k in texts), None)
    other = next((k for k in texts if k not in (reply, ask)), None)
    out: dict[str, dict] = {}

    def add(q, title, spec, derive=(), ups=()):
        known = dict(cols)
        for d in derive:
            known[d["name"]] = {"type": "num" if d["kind"] in ("score", "count") else "cat"}
        try:
            out[q] = {"title": title, "derive": [dict(d) for d in derive], "chart": check_spec(spec, known), "follow_ups": list(ups)}
        except PlanError:
            pass

    if reply:
        tone = {"name": "tone", "kind": "label", "from": reply, "guide": f"the overall tone of the {reply}",
                "labels": ["friendly", "formal", "empathetic", "technical", "blunt"]}
        add(f"What tone do the {_pl(reply)} take" + (f" in each {cat}?" if cat else "?"), f"Tone of the {_pl(reply)}" + (f" by {cat}" if cat else ""),
            {"type": "stack", "x": cat, "color": "tone", "agg": "share"} if cat else {"type": "bar", "x": "tone"}, [tone],
            [f"How long are the {_pl(reply)} for each tone?"] if ln else [])
    if ask:
        add(f"What do the {_pl(ask)} ask for? Group them into a few themes.", f"What the {_pl(ask)} ask for",
            {"type": "bar", "x": "request"}, [{"name": "request", "kind": "label", "from": ask, "guide": f"what the {ask} asks for", "labels": []}])
    if not ask and not reply and other:
        add(f"What themes come up in the {_pl(other)}?", f"Themes in the {_pl(other)}", {"type": "bar", "x": "theme"},
            [{"name": "theme", "kind": "label", "from": other, "guide": f"the main theme of the {other}", "labels": []}])
        add(f"How positive or negative is each {other}?", f"Sentiment of the {_pl(other)}", {"type": "hist", "x": "sentiment"},
            [{"name": "sentiment", "kind": "score", "from": other, "guide": "1 = very negative, 3 = neutral or mixed, 5 = very positive"}])
    if ln and cat:
        add(f"Which {cat} gets the longest {_pl(ln.replace(' words', ''))}?", f"{ln.replace(' words', '').capitalize()} length by {cat}",
            {"type": "box", "x": cat, "y": ln})
    elif ln:
        add(f"How long are the {_pl(ln.replace(' words', ''))}?", f"{ln.capitalize()}", {"type": "hist", "x": ln})
    if reply:
        add(f"How many {_pl(reply)} give step-by-step instructions?", f"{_pl(reply).capitalize()} with step-by-step instructions",
            {"type": "bar", "x": "gives steps"},
            [{"name": "gives steps", "kind": "label", "from": reply, "guide": "whether it walks the reader through numbered or ordered steps",
              "labels": ["yes", "no"]}])
    seconds = [k for k, c in cols.items() if c["type"] == "cat" and k != cat and 1 < len({r.get(k) for r in rows}) <= 8]
    if cat and seconds:
        add(f"How does {seconds[0]} vary across {cat}?", f"{seconds[0].capitalize()} within each {cat}",
            {"type": "stack", "x": cat, "color": seconds[0], "agg": "share"})
    return dict(list(out.items())[:4])


def _pl(name: str) -> str:
    return name if name.endswith("s") else name + "s"


def suggestions(rows: list[dict], cols: dict) -> list[str]:
    return list(presets(rows, cols))


# ------------------------------------------------------------------ the model's part
def first_json(text: str):
    """The first JSON object or list in a model reply (fences and chatter around it are ignored)."""
    dec = json.JSONDecoder()
    for i, ch in enumerate(text):
        if ch in "{[":
            try:
                return dec.raw_decode(text[i:])[0]
            except ValueError:
                continue
    raise ValueError("no JSON in the reply")


def _snip(v, n=LABEL_CHARS) -> str:
    return re.sub(r"\s+", " ", str(v or "")).strip()[:n]


def plan_prompt(question: str, name: str, rows: list[dict], cols: dict, scope: str = "", error: str = "") -> str:
    lines = []
    for d in describe(rows, cols):
        if d["type"] == "cat":
            vals = ", ".join(f"{k} {v}" for k, v in d["values"][:8])
            more = f" (+{d['distinct'] - 8} more)" if d["distinct"] > 8 else ""
            lines.append(f'- "{d["name"]}" (category, {d["distinct"]} values: {vals}{more})')
        elif d["type"] == "num":
            if "min" in d:
                lines.append(f'- "{d["name"]}" (number, {d["min"]:g} to {d["max"]:g}, median {d["median"]:g})')
        else:
            lines.append(f'- "{d["name"]}" (text, about {d["avg_words"]} words) e.g. "{_snip(d["example"], 100)}"')
        if d.get("derived") in ("label", "score", "pattern", "count"):
            lines[-1] += f' [made earlier by reading "{d["from"]}"]'
    texts = [k for k, c in cols.items() if c["type"] in ("text", "cat")][:5]
    ex = []
    for r in rows[:5]:
        ex.append(f"#{r['_id']} " + " | ".join(f"{k}: {_snip(r.get(k), 110)}" for k in texts))
    return f"""You turn a question about a dataset into a chart plan.

Dataset: {name}, {len(rows)} rows{scope}.
Columns:
{chr(10).join(lines)}

A few rows:
{chr(10).join(ex)}

Question: "{question}"

Reply with only a JSON object:
{{"title": "...", "derive": [], "chart": {{"type": "...", "x": "...", "y": null, "color": null, "agg": "count"}}, "follow_ups": ["...", "...", "..."]}}

Rules:
- Use the columns above when they can answer the question. They are instant.
- Questions about what the text says, asks for, means or sounds like (themes, intent, tone, quality) need a derived column read from that text, unless a category column already names exactly that.
- Chart a length or other number column only when the question is about length or amount.
- Only when no column answers it, add one or two entries to "derive". Each becomes a new column a model fills by reading one column of every row:
  {{"name": "tone", "kind": "label", "from": "completion", "labels": ["formal", "friendly", "curt"], "guide": "the overall tone of the reply"}}
  kind "label": one of 2 to 7 short labels per row. Give "labels" yourself when the question or the rows above suggest them (faster); use [] only for open-ended themes, to have them discovered from the rows.
  kind "score": an integer from 1 to 5; "guide" must say what 1 and 5 mean.
  kind "pattern": yes or no, whether the case-insensitive regex in "pattern" matches. kind "count": how many times it matches.
  Use "pattern" or "count" only for literal marks (links, numbered lists, tables, code blocks, digits); judge meaning with "label" or "score".
  "name" is 1 to 3 lowercase words, "from" is a text or category column.
- chart.type:
  "bar": x is a category; y null counts rows, or y is a number column with agg "mean" or "median".
  "stack": x and color are categories; agg "share" or "count". Use it to compare make-up across groups.
  "box": y is a number column; x is a category or null. Use it for how a number spreads in each group.
  "hist": x is a number column.
  "scatter": x and y are number columns; color is a category or null.
  A derived column can be used in the chart by its name.
- "title": what the chart shows, under 60 characters, sentence case, no claims about the result.
- "follow_ups": three short questions a person exploring this data might ask next.{chr(10) + chr(10) + "Your last plan was rejected: " + error + " Fix it." if error else ""}"""


def _clean_name(raw: str) -> str:
    return re.sub(r"\s+", " ", re.sub(r"[^a-z0-9 ]", "", re.sub(r"[_\-]+", " ", raw.lower()))).strip()[:24]


def parse_plan(text: str, cols: dict) -> dict:
    try:
        p = first_json(text)
    except ValueError:
        raise PlanError("The reply was not JSON.")
    if not isinstance(p, dict) or not isinstance(p.get("chart"), dict):
        raise PlanError('The reply needs a "chart" object.')
    derive = []
    known = dict(cols)
    for d in (p.get("derive") or [])[:2]:
        if not isinstance(d, dict):
            continue
        kind = d.get("kind")
        raw_name = str(d.get("name") or "")
        name = _clean_name(raw_name)
        for k in ("x", "y", "color"):  # the chart may refer to the name as the model wrote it
            if p["chart"].get(k) is not None and p["chart"].get(k) not in cols and _clean_name(str(p["chart"][k])) == name:
                p["chart"][k] = name
        src = d.get("from")
        if kind not in KINDS or not name:
            raise PlanError(f"derive entries need a name and a kind from {', '.join(KINDS)}.")
        if src not in cols or cols[src]["type"] not in ("text", "cat"):
            raise PlanError(f'derive "from" must be a text or category column, not {src!r}.')
        item = {"name": name, "kind": kind, "from": src, "guide": str(d.get("guide") or "")[:200]}
        if kind == "label":
            labels = [str(x).strip()[:40] for x in (d.get("labels") or []) if str(x).strip()]
            item["labels"] = list(dict.fromkeys(labels))[:7]
        elif kind in ("pattern", "count"):
            pat = re.sub(r"^\(\?i\)", "", str(d.get("pattern") or ""))
            if not pat or len(pat) > 160 or re.search(r"\([^)]*[+*]\)[+*{]", pat):
                raise PlanError("pattern must be a short regex without nested repeats.")
            try:
                re.compile(pat, re.I)
            except re.error:
                raise PlanError(f"{pat!r} is not a valid regex.")
            item["pattern"] = pat
        if name in cols and cols[name].get("derived") == kind and cols[name].get("from") == src \
                and (kind != "label" or not item["labels"] or set(item["labels"]) == set(cols[name].get("labels") or [])):
            continue  # already made by an earlier question: reuse it
        if name in cols and cols[name].get("derived") not in KINDS:  # an original column: never overwrite it
            new = f"{name} (model)"
            for k in ("x", "y", "color"):
                if p["chart"].get(k) == name:
                    p["chart"][k] = new
            item["name"] = name = new
        derive.append(item)
        known[name] = {"type": "num" if kind in ("score", "count") else "cat"}
    spec = check_spec(p["chart"], known)
    title = re.sub(r"\s+", " ", str(p.get("title") or "")).strip()[:80] or "What the data shows"
    ups = [re.sub(r"\s+", " ", str(q)).strip()[:140] for q in (p.get("follow_ups") or []) if str(q).strip()][:3]
    return {"title": title, "derive": derive, "chart": spec, "follow_ups": ups}


def fallback_plan(question: str, rows: list[dict], cols: dict) -> dict:
    """Used when the model cannot produce a valid plan: answer with the closest instant chart."""
    q = question.lower()
    cat = next((k for k, c in cols.items() if c["type"] == "cat" and k.lower() in q), None) or _main_cat(cols, rows)
    num = next((k for k, c in cols.items() if c["type"] == "num" and k.lower().replace(" words", "") in q), None)
    if not num and re.search(r"\blong|length|short|words?\b", q):
        num = _main_len(cols)
    if num:
        spec = {"type": "box", "x": cat, "y": num} if cat else {"type": "hist", "x": num}
        title = f"{num.capitalize()}" + (f" by {cat}" if cat else "")
    elif cat:
        spec, title = {"type": "bar", "x": cat}, f"Rows by {cat}"
    else:
        raise PlanError("This data has no column a chart can use.")
    return {"title": title, "derive": [], "chart": check_spec(spec, cols), "follow_ups": [], "fallback": True}


def discover_prompt(question: str, d: dict, rows: list[dict]) -> str:
    snips = "\n".join(f"- {_snip(r.get(d['from']), 220)}" for r in rows[:40])
    return f"""Here are {min(40, len(rows))} entries from the "{d['from']}" column of a dataset. Someone asked: "{question}"

We need a label for every row: "{d['name']}"{f" ({d['guide']})" if d['guide'] else ""}.
Propose 3 to 7 short labels (1 to 4 words each, lowercase) that together cover these entries, most common first. Use labels a reader would understand without seeing the data.

{snips}

Reply with only a JSON list of strings."""


def label_prompt(d: dict, batch: list[dict]) -> str:
    items = "\n".join(f"[{k}] {_snip(r.get(d['from']))}" for k, r in enumerate(batch, 1))
    if d["kind"] == "label":
        ask = f'Pick exactly one label for each item from: {", ".join(json.dumps(x) for x in d["labels"])}.'
        eg = ", ".join(f'"{k}": {json.dumps(d["labels"][k % len(d["labels"])])}' for k in (1, 2))
    else:
        ask = "Give each item an integer score from 1 to 5."
        eg = '"1": 4, "2": 2'
    return f"""You are labelling rows of a dataset.
{d['name'].capitalize()}: {d['guide'] or d['name']}
{ask}

{items}

Reply with only a JSON object mapping each item number to its {"label" if d["kind"] == "label" else "score"}, like {{{eg}}}."""


def _read_labels(text: str, d: dict, k: int) -> Optional[dict]:
    try:
        obj = first_json(text)
    except ValueError:
        return None
    if isinstance(obj, list):
        obj = {str(i): v for i, v in enumerate(obj, 1)}
    if not isinstance(obj, dict):
        return None
    out = {}
    canon = {x.lower(): x for x in d.get("labels", [])}
    for key, v in obj.items():
        m = re.search(r"\d+", str(key))
        if not m or not 1 <= int(m.group()) <= k:
            continue
        i = int(m.group())
        if d["kind"] == "label":
            s = str(v).strip().lower()
            hit = canon.get(s) or next((canon[c] for c in canon if s.startswith(c) or c.startswith(s)), None)
            if hit:
                out[i] = hit
        else:
            n = _num(v)
            if n is not None and 1 <= n <= 5:
                out[i] = float(round(n))
    return out if len(out) >= max(1, int(0.8 * k)) else None


def hedged(fn: Callable[[float], object], after: float = HEDGE_AFTER, tries: int = 2):
    """Call fn(temperature) and return its first truthy result. A call still running after `after`
    seconds gets a twin, so one slow provider does not hold up the chart. None if every try fails."""
    for attempt in range(tries):
        futs = {_HEDGE.submit(fn, 0.1 if attempt == 0 else 0.4)}
        start, twin = time.time(), False
        while futs:
            done, futs = wait(futs, timeout=None if twin else max(0.05, after - (time.time() - start)), return_when=FIRST_COMPLETED)
            for f in done:
                try:
                    got = f.result()
                except Exception:
                    got = None
                if got:
                    return got
            if not twin and (time.time() - start >= after or not futs):
                futs = set(futs) | {_HEDGE.submit(fn, 0.5)}
                twin = True
    return None


def fill(d: dict, rows: list[dict], chat: Callable[..., str], question: str = "",
         progress: Optional[Callable[[int], None]] = None, workers: int = 6) -> dict:
    """Compute one derived column for `rows`. Returns {row id: value} and may set d["labels"]."""
    if d["kind"] in ("pattern", "count"):
        rx = re.compile(d["pattern"], re.I)
        vals = {}
        for r in rows:
            hits = len(rx.findall(str(r.get(d["from"]) or "")[:MAX_CELL]))
            vals[r["_id"]] = ("yes" if hits else "no") if d["kind"] == "pattern" else float(hits)
        if progress:
            progress(len(rows))
        return vals
    if d["kind"] == "label" and len(d.get("labels") or []) < 2:
        def discover(temp):
            try:
                obj = first_json(chat(discover_prompt(question, d, rows), max_tokens=200, temperature=temp + 0.2))
            except ValueError:
                return None
            if isinstance(obj, dict):  # {"labels": [...]} and similar
                obj = next((v for v in obj.values() if isinstance(v, list)), [])
            labels = [str(x).strip().lower()[:40] for x in obj if isinstance(x, (str, int, float)) and str(x).strip()]
            labels = list(dict.fromkeys(labels))[:7]
            return labels if len(labels) >= 2 else None
        labels = hedged(discover)
        if not labels:
            raise PlanError(f'Could not find labels for "{d["name"]}". Try asking again, or name the groups you expect.')
        d["labels"] = labels
    batches = [rows[i:i + LABEL_BATCH] for i in range(0, len(rows), LABEL_BATCH)]
    vals: dict[int, object] = {}
    done = [0]

    def run(batch):
        got = hedged(lambda temp: _read_labels(chat(label_prompt(d, batch), max_tokens=12 * len(batch) + 40, temperature=temp), d, len(batch)))
        res = {}
        for k, r in enumerate(batch, 1):
            res[r["_id"]] = (got or {}).get(k, "unclear" if d["kind"] == "label" else None)
        done[0] += len(batch)
        if progress:
            progress(done[0])
        return res

    with ThreadPoolExecutor(max_workers=max(1, min(workers, len(batches)))) as pool:
        for res in pool.map(run, batches):
            vals.update(res)
    return vals


def _rows(k) -> str:
    return f"{k} row" if k == 1 else f"{k} rows"


def _unit(col) -> str:
    col = str(col or "")
    return " words" if col.endswith(" words") else " turns" if col.endswith(" turns") else ""


def facts(c: dict, limit: int = 14) -> list[str]:
    """The numbers behind a chart in plain lines, for the write-up (and nothing else)."""
    t, marks, n = c["type"], c.get("marks") or [], max(1, c["n"])
    u = _unit(c.get("y") if t in ("bar", "box") else c.get("x"))
    out = []
    if t == "bar":
        for m in marks[:limit]:
            if c["agg"] == "count":
                out.append(f"{m['key']}: {_rows(m['value'])} ({round(100 * m['value'] / n)}%)")
            else:
                out.append(f"{m['key']}: {c['agg']} {c['y']} {m['value']:g}{u} ({_rows(m['n'])})")
    elif t == "stack":
        for m in marks[:limit]:
            out.append(f"{m['key']} ({_rows(m['n'])}): " + ", ".join(f"{p['key']} {p['value']} ({round(100 * p['value'] / m['n'])}%)" for p in m["parts"]))
    elif t == "box":
        for m in marks[:limit]:
            out.append(f"{m['key']} ({_rows(m['n'])}): median {m['median']:g}{u}, middle half {m['q1']:g} to {m['q3']:g}{u}, range {m['lo']:g} to {m['hi']:g}{u}")
    elif t == "hist":
        for m in marks:
            if m["value"]:
                out.append(f"{c['x']} {m['x0']:g}{u}: {_rows(m['value'])}" if c.get("discrete") else f"{m['x0']:g} to {m['x1']:g}{u}: {_rows(m['value'])}")
    elif t == "scatter":
        out.append(f"{len(marks)} rows plotted")
        if "r" in c:
            out.append(f"correlation between {c['x']} and {c['y']}: r = {c['r']}")
    return out


def summary_prompt(question: str, title: str, c: dict, derived: list[dict]) -> str:
    made = "".join(f'\n"{d["name"]}" was assigned by a model reading each row\'s "{d["from"]}" text.' for d in derived)
    return f"""Someone asked about a dataset: "{question}"
The chart "{title}" covers {c['n']} rows. Its exact numbers:
{chr(10).join(facts(c))}{made}

Answer the question in 2 to 4 short sentences of plain prose. Use only these numbers and keep them exact, but write them into sentences; do not paste the lines above. Lead with the clearest difference. Mention groups with fewer than 5 rows at most once, as too small to read much into. No quotation marks, markdown, lists, headings or advice."""


def notes(c: dict, derived: list[dict], fallback: bool = False) -> list[str]:
    out = []
    if fallback:
        out.append("The model could not plan a chart for that question, so this is the closest chart the columns give directly.")
    for d in derived:
        out.append(f'"{d["name"].capitalize()}" was assigned by a model reading each row\'s {d["from"]}. Open a few rows before relying on it.')
    if c["n"] < 20:
        out.append(f"Only {c['n']} rows, so a difference of one or two rows can change the picture.")
    thin = [m["key"] for m in (c.get("marks") or []) if isinstance(m, dict) and "key" in m and 0 < m.get("n", 99) < 3]
    if thin and c["type"] in ("bar", "box", "stack") and c.get("agg") != "count":
        out.append(f"{', '.join(thin[:3])} {'has' if len(thin) == 1 else 'have'} fewer than three rows.")
    return out