app.py

"""
EDAgent — a lightweight LLM agent that writes, executes, repairs and reports
its own exploratory data analysis.

Pipeline
    1. load + normalize        deterministic
    2. context card            deterministic
    3. statistics              deterministic  (roles, integrity, shapes,
                                               mixed-type associations, IQR outliers)
    4. narrative               LLM, one sentence per pre-computed fact
    5. code generation         LLM
    6. sandboxed execution     deterministic  (+ plot-quality guards)
    7. one repair pass         LLM, given a cleaned traceback

Design principle: the model writes analysis code and phrases facts. Every table,
statistic, association and "which is highest" lookup is computed in pandas/scipy.

Caches are keyed by dataset id: a second run on the same dataset reuses the
download, the statistics and the narrative, and only regenerates the plots.
"""

import base64
import concurrent.futures
import contextlib
import glob
import io
import json
import os
import re
import shutil
import time
import traceback

import gradio as gr
import matplotlib
import numpy as np
import pandas as pd
import spaces
import torch

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import seaborn as sns  # noqa: E402
from matplotlib.ticker import (FuncFormatter, LogLocator,  # noqa: E402
                               NullFormatter, ScalarFormatter)
from scipy import stats  # noqa: E402
from huggingface_hub import HfApi, hf_hub_download  # noqa: E402
from transformers import (AutoModelForCausalLM, AutoTokenizer,  # noqa: E402
                          StoppingCriteria, StoppingCriteriaList)

# --------------------------------------------------------------------------- #
# Model — must be placed on cuda at module level for ZeroGPU
# --------------------------------------------------------------------------- #
MODEL_ID = "Qwen/Qwen2.5-Coder-1.5B-Instruct"

# On ZeroGPU, `spaces` enables CUDA emulation at import, so placing the model on
# "cuda" at module level is required. On CPU Spaces (or locally) there is no CUDA,
# so fall back to float32 on CPU — the app boots, but generation is far slower.
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.bfloat16 if DEVICE == "cuda" else torch.float32

tok = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=DTYPE).to(DEVICE)
model.eval()
if tok.pad_token_id is None:
    tok.pad_token = tok.eos_token
print(f"[startup] device={DEVICE} dtype={DTYPE}")

PLOTS_DIR = "plots"
os.makedirs(PLOTS_DIR, exist_ok=True)


class StopOnClosingFence(StoppingCriteria):
    """Stop once a fenced block has been opened and closed.

    Decoding the tail is more reliable than matching token ids: the fence
    tokenizes differently depending on what precedes it. Decoding on every step
    would be wasteful, so we only check every `check_every` tokens.
    """

    def __init__(self, tokenizer, prompt_len, fences=2, check_every=8):
        self.tok, self.prompt_len = tokenizer, prompt_len
        self.fences, self.check_every = fences, check_every
        self._step = 0

    def __call__(self, input_ids, scores, **kwargs):
        self._step += 1
        done = False
        if self._step % self.check_every == 0:
            text = self.tok.decode(input_ids[0][self.prompt_len:],
                                   skip_special_tokens=True)
            done = text.count("```") >= self.fences
        # one flag per batch row: StoppingCriteriaList ORs the results together
        return torch.full((input_ids.shape[0],), done,
                          dtype=torch.bool, device=input_ids.device)


@torch.inference_mode()
def generate(system: str, user: str, max_new_tokens: int = 900,
             stop_fences: int = 0) -> str:
    """Single entry point to the LLM. Returns only the newly generated text.

    stop_fences=2 halts as soon as the closing fence of a code block is emitted,
    instead of running on to max_new_tokens with trailing prose.
    """
    msgs = [{"role": "system", "content": system}, {"role": "user", "content": user}]
    text = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
    inputs = tok(text, return_tensors="pt").to(model.device)
    prompt_len = inputs["input_ids"].shape[1]

    criteria = None
    if stop_fences:
        criteria = StoppingCriteriaList(
            [StopOnClosingFence(tok, prompt_len, fences=stop_fences)])

    out = model.generate(
        **inputs,
        max_new_tokens=max_new_tokens,
        do_sample=False,                      # greedy -> reproducible
        pad_token_id=tok.pad_token_id,
        stopping_criteria=criteria,
    )
    return tok.decode(out[0][prompt_len:], skip_special_tokens=True)


# --------------------------------------------------------------------------- #
# 1. Data layer
# --------------------------------------------------------------------------- #
TABULAR_EXT = (".csv", ".tsv", ".parquet", ".json", ".jsonl")


def _read_any(path: str) -> pd.DataFrame:
    if path.endswith(".parquet"):
        return pd.read_parquet(path)
    if path.endswith(".tsv"):
        return pd.read_csv(path, sep="\t")
    if path.endswith((".json", ".jsonl")):
        return pd.read_json(path, lines=path.endswith(".jsonl"))
    return pd.read_csv(path)


def load_hf_dataframe(hf_dataset: str) -> pd.DataFrame:
    """Hub dataset id -> DataFrame. Tries load_dataset, falls back to raw files.

    The fallback matters: datasets>=3.0 dropped script-based repos, so datasets
    such as mstz/titanic can no longer be loaded the standard way.
    """
    try:
        from datasets import load_dataset

        ds = load_dataset(hf_dataset)
        split = "train" if "train" in ds else list(ds.keys())[0]
        return ds[split].to_pandas()
    except Exception:
        pass

    files = HfApi().list_repo_files(hf_dataset, repo_type="dataset")
    cands = [f for f in files if f.lower().endswith(TABULAR_EXT)]
    if not cands:
        raise RuntimeError(f"No tabular file found in '{hf_dataset}'.")
    cands.sort(key=lambda f: (("train" not in f.lower()), len(f)))
    return _read_any(hf_hub_download(hf_dataset, cands[0], repo_type="dataset"))


def normalize(df: pd.DataFrame) -> pd.DataFrame:
    """Strip literal quote chars, turn empty strings into real NaN, re-infer numerics."""
    df = df.copy()
    for c in df.select_dtypes(include="object").columns:
        s = df[c].astype(str).str.strip()
        s = s.str.replace(r"^'(.*)'$", r"\1", regex=True)
        df[c] = s.str.strip().replace(
            {"": np.nan, "nan": np.nan, "None": np.nan, "?": np.nan})
    for c in df.select_dtypes(include="object").columns:
        conv = pd.to_numeric(df[c], errors="coerce")
        if conv.isna().sum() == df[c].isna().sum():      # lossless only
            df[c] = conv
    return df


def drop_index_cols(df: pd.DataFrame):
    """Remove monotonic unique integer columns - row ids, not variables."""
    drop = [c for c in df.columns
            if pd.api.types.is_integer_dtype(df[c])
            and df[c].is_monotonic_increasing
            and df[c].nunique() == len(df)]
    return df.drop(columns=drop), drop


# --------------------------------------------------------------------------- #
# 2. Context card
# --------------------------------------------------------------------------- #
def build_context(df: pd.DataFrame, name: str) -> str:
    lines = [f"Dataset: {name}",
             f"Shape: {df.shape[0]} rows x {df.shape[1]} columns", "", "Columns:"]
    for c in df.columns:
        s, nulls = df[c], df[c].isna().sum()
        if pd.api.types.is_numeric_dtype(s):
            detail = f"min={s.min():.4g} max={s.max():.4g} mean={s.mean():.4g}"
        else:
            detail = "examples=" + ", ".join(map(str, s.dropna().unique()[:4]))[:70]
        lines.append(
            f"- {c} ({s.dtype}) nulls={nulls} ({100 * nulls / len(df):.1f}%) "
            f"unique={s.nunique(dropna=True)} {detail}")
    lines += ["", "First 3 rows:", df.head(3).to_string(max_colwidth=22)]
    return "\n".join(lines)


def schema_only(df: pd.DataFrame, name: str) -> str:
    """Context minus sample rows, so stray values cannot be copied into prose."""
    return (f"Dataset: {name}\nShape: {df.shape[0]} rows x {df.shape[1]} columns\n"
            "Columns: " + ", ".join(f"{c} ({df[c].dtype})" for c in df.columns))


# --------------------------------------------------------------------------- #
# 3. Code generation
# --------------------------------------------------------------------------- #
CODE_SYSTEM = """You are a data analyst. You write Python code for exploratory data analysis.

Rules:
- A pandas DataFrame named `df` is already loaded. Never load, download or create data.
- Use pandas, matplotlib and seaborn.
- print() every result you want reported.
- Save each figure with plt.savefig("plots/<name>.png", bbox_inches="tight") then plt.close(). Never call plt.show().
- Reply with exactly one ```python code block and no other text."""

FIX_SYSTEM = """You fix broken Python data-analysis code.

Rules:
- A pandas DataFrame named `df` is already loaded. Never load or create data.
- Return the complete corrected script, not just the changed lines.
- Reply with exactly one ```python code block and no other text."""

# Chart names follow the "Main Plots" / "Basic Plots" slides from Lecture 02.
PLOT_MENU = {
    "Histogram":           "histograms of the main numeric columns",
    "Bar Graph":           ("bar charts of value counts for the low-cardinality "
                            "categorical columns"),
    "Scatter Plot":        ("scatter plots of a few pairs of numeric columns that look "
                            "related"),
    "Area Plot":           ("an area plot of one numeric column across the sorted row "
                            "order"),
    "Pie Plot":            ("a pie chart of the share of each level in one "
                            "low-cardinality categorical column"),
    "Correlation Heatmap": "one correlation heatmap of the numeric columns",
}

ALL_PLOTS = list(PLOT_MENU)
DEFAULT_PLOTS = ["Histogram", "Bar Graph", "Scatter Plot", "Correlation Heatmap"]
CODE_MAX_TOKENS = 700     # observed scripts are ~500 tokens; more only buys rambling


def build_code_prompt(context: str, instruction: str, selected_plots=None) -> str:
    plot_req = ""
    if selected_plots:
        wants = "; ".join(PLOT_MENU[p] for p in selected_plots if p in PLOT_MENU)
        if wants:
            plot_req = (f"\nCreate only these plot types: {wants}. "
                        "Do not create any other plots.")
    return f"{context}\n\nTask: {instruction}{plot_req}"


def extract_code(text: str) -> str:
    """Pull python out of markdown fences; concatenate if the model emitted several."""
    blocks = re.findall(r"```(?:python|py)?\s*\n(.*?)```", text, flags=re.DOTALL)
    return "\n\n".join(b.strip() for b in blocks) if blocks else text.strip()


# --------------------------------------------------------------------------- #
# 4. Execution sandbox (+ plot-quality guards)
# --------------------------------------------------------------------------- #
BANNED = re.compile(r"\b(pd\.read_\w+|load_dataset|sns\.load_dataset)\s*\(")
CORR_FIX = re.compile(r"\.corr\(\s*\)")          # bare df.corr() -> numeric_only

# Speed bump, not a sandbox. This Space runs LLM-generated code publicly;
# never store secrets in its settings.
UNSAFE = re.compile(
    r"(__import__|\bsubprocess\b|\bsocket\b|\brequests\b|\bshutil\b|"
    r"\bimport\s+(os|sys)\b|\bopen\s*\(|\beval\s*\(|\bexec\s*\(|\bgetattr\s*\()"
)

MAX_ANNOT_COLS    = 12
MAX_TICKS         = 20
MAX_CAT_LEVELS    = 20
MAX_PAIRPLOT_COLS = 5      # a pair grid is quadratic in columns
MAX_PAIRPLOT_ROWS = 2000   # ... and linear in rows, per panel
PLOT_SAMPLE_ROWS  = 5000   # plotting sees a sample; the report's stats do not
CODE_TIMEOUT_S    = 120    # generated code is not allowed to run forever
SKIP_ATTR         = "_edagent_skip"
TITLE_ATTR        = "_edagent_title"


def sanitize(code: str) -> str:
    """Neutralize data reloads; patch the one error the model cannot reliably fix.

    df.corr() on mixed-type data failed on every mixed-dtype dataset tested. The
    repair loop fixed it correctly once and incorrectly once, so it moves to the
    deterministic side. Everything else still goes through the retry.
    """
    out = []
    for ln in code.splitlines():
        if BANNED.search(ln):
            out.append("# [removed by agent] " + ln.strip())
        else:
            out.append(CORR_FIX.sub(".corr(numeric_only=True)", ln))
    return "\n".join(out)


def _short(s, n=18):
    s = str(s)
    return s if len(s) <= n else f"{s[:n // 2 - 1]}...{s[-(n // 2 - 1):]}"


def _tick(v, _=None):
    """Plain numbers with thousands separators - never 1e7."""
    return f"{v:,.0f}" if abs(v) >= 1000 else f"{v:g}"


def _tidy(fig):
    """Make any figure legible: thin dense ticks, rotate, kill scientific notation."""
    for ax in fig.get_axes():
        if len(ax.get_xticklabels()) > MAX_TICKS:
            step = max(1, len(ax.get_xticks()) // MAX_TICKS)
            ax.set_xticks(ax.get_xticks()[::step])
        for axis, scale in ((ax.xaxis, ax.get_xscale()), (ax.yaxis, ax.get_yscale())):
            # only touch numeric axes: a heatmap's labels use a FixedFormatter
            if scale == "linear" and isinstance(axis.get_major_formatter(),
                                                ScalarFormatter):
                axis.set_major_formatter(FuncFormatter(_tick))
        plt.setp(ax.get_xticklabels(), rotation=45, ha="right", fontsize=8)
        plt.setp(ax.get_yticklabels(), fontsize=8)


def _install_plot_guards():
    """Patch the real modules, so `import matplotlib.pyplot as plt` cannot bypass us."""
    real = {"savefig": plt.savefig, "heatmap": sns.heatmap, "histplot": sns.histplot,
            "boxplot": sns.boxplot, "violinplot": sns.violinplot,
            "countplot": sns.countplot, "pairplot": sns.pairplot, "pie": plt.pie}

    def _vars(a, kw):
        """Recover the (name, Series) pairs a seaborn call is plotting."""
        d = kw.get("data", a[0] if a else None)
        out = []
        for key in ("x", "y"):
            v = kw.get(key)
            if isinstance(v, str) and hasattr(d, "columns"):
                out.append((v, d[v]))
            elif isinstance(v, pd.Series):
                out.append((v.name, v))
        if not out and isinstance(d, pd.Series):
            out.append((d.name, d))
        return out

    def _counts_instead(name, s):
        """A box plot of labels is meaningless; show how often each level occurs."""
        ax = plt.gca()
        s.value_counts().head(MAX_CAT_LEVELS).sort_values().plot(kind="barh", ax=ax)
        ax.set_xlabel("count")
        ax.set_ylabel(str(name or ""))
        setattr(plt.gcf(), TITLE_ATTR, f"Counts of {name}" if name else "Counts")
        return ax

    def _categorical_guard(a, kw):
        """Return an axes if we handled it, else None to fall through."""
        pairs = _vars(a, kw)
        if not pairs or any(pd.api.types.is_numeric_dtype(s) for _, s in pairs):
            return None                    # a numeric variable is present: legitimate
        name, s = pairs[0]
        if s.nunique(dropna=True) > MAX_CAT_LEVELS:
            setattr(plt.gcf(), SKIP_ATTR, True)     # ids, free-text descriptions
            return plt.gca()
        return _counts_instead(name, s)

    def savefig(fname, *a, **kw):
        fig = plt.gcf()
        if getattr(fig, SKIP_ATTR, False):
            return None                    # nothing worth writing to disk
        title = getattr(fig, TITLE_ATTR, None)
        if title and fig.get_axes():
            fig.get_axes()[0].set_title(title)   # the model's title would lie
        w, h = fig.get_size_inches()
        fig.set_size_inches(max(w, 8), max(h, 5))
        _tidy(fig)
        kw.setdefault("bbox_inches", "tight")
        kw.setdefault("dpi", 110)
        return real["savefig"](fname, *a, **kw)

    def heatmap(data, *a, **kw):
        n = getattr(data, "shape", (0, 0))[1]
        if n > MAX_ANNOT_COLS:
            kw["annot"] = False
        kw.setdefault("cmap", "coolwarm")
        side = min(max(7, 0.5 * n + 4), 20)
        plt.gcf().set_size_inches(side, side * 0.85)
        ax = real["heatmap"](data, *a, **kw)
        if hasattr(data, "columns"):
            fs = 8 if n <= 15 else 6
            ax.set_xticklabels([_short(c, 24) for c in data.columns],
                               rotation=45, ha="right", fontsize=fs)
            ax.set_yticklabels([_short(c, 24) for c in data.index],
                               rotation=0, fontsize=fs)
        return ax

    def _numeric_series(a, kw):
        for _, s in _vars(a, kw):
            s = pd.to_numeric(s, errors="coerce").dropna()
            if len(s):
                return s
        return None

    def histplot(*a, **kw):
        s = _numeric_series(a, kw)
        logged = False
        # heavy right skew (prices, fares, incomes) -> log x, else one tall bar
        if (s is not None and len(s) > 20 and s.min() > 0
                and s.max() / max(s.median(), 1e-9) > 50 and "log_scale" not in kw):
            kw["log_scale"] = (True, False)
            logged = True
        ax = real["histplot"](*a, **kw)
        if logged:
            ax.xaxis.set_major_locator(LogLocator(base=10, subs=(1.0, 2.0, 5.0),
                                                  numticks=12))
            ax.xaxis.set_major_formatter(FuncFormatter(_tick))
            ax.xaxis.set_minor_formatter(NullFormatter())
            ax.set_xlabel(f"{ax.get_xlabel()} (log scale)")
        return ax

    def boxplot(*a, **kw):
        handled = _categorical_guard(a, kw)
        return handled if handled is not None else real["boxplot"](*a, **kw)

    def violinplot(*a, **kw):
        handled = _categorical_guard(a, kw)
        return handled if handled is not None else real["violinplot"](*a, **kw)

    def countplot(*a, **kw):
        pairs = _vars(a, kw)
        if pairs and pairs[0][1].nunique(dropna=True) > MAX_CAT_LEVELS:
            setattr(plt.gcf(), SKIP_ATTR, True)
            return plt.gca()
        return real["countplot"](*a, **kw)

    def pairplot(*a, **kw):
        """A pair grid is O(cols^2) panels x rows points. Unbounded, it hangs."""
        d = kw.get("data", a[0] if a else None)
        if isinstance(d, pd.DataFrame):
            num = d.select_dtypes(include="number")
            keep = list(num.columns)[:MAX_PAIRPLOT_COLS]
            d2 = num[keep] if keep else num
            if len(d2) > MAX_PAIRPLOT_ROWS:
                d2 = d2.sample(MAX_PAIRPLOT_ROWS, random_state=0)
            if "data" in kw:
                kw["data"] = d2
                return real["pairplot"](*a, **kw)
            return real["pairplot"](d2, *a[1:], **kw)
        return real["pairplot"](*a, **kw)

    def pie(x, *a, **kw):
        """Hundreds of labelled wedges is neither fast nor readable."""
        try:
            n = len(x)
        except Exception:
            n = 0
        if n > MAX_CAT_LEVELS:
            setattr(plt.gcf(), SKIP_ATTR, True)
            return ([], [], []) if kw.get("autopct") else ([], [])
        return real["pie"](x, *a, **kw)

    plt.savefig, plt.pie = savefig, pie
    sns.heatmap, sns.histplot = heatmap, histplot
    sns.boxplot, sns.violinplot, sns.countplot = boxplot, violinplot, countplot
    sns.pairplot = pairplot
    return real


def _restore_plot_guards(real):
    plt.savefig, plt.pie = real["savefig"], real["pie"]
    sns.heatmap, sns.histplot = real["heatmap"], real["histplot"]
    sns.boxplot, sns.violinplot = real["boxplot"], real["violinplot"]
    sns.countplot, sns.pairplot = real["countplot"], real["pairplot"]


def _exec_with_timeout(code, ns, seconds):
    """Run generated code with a wall clock.

    Python cannot kill a running thread, so on timeout the worker is abandoned
    rather than stopped - this unblocks the UI, it does not free the CPU. The
    real defence is the sampling and plot guards above; this is the backstop.
    """
    pool = concurrent.futures.ThreadPoolExecutor(max_workers=1)
    try:
        fut = pool.submit(exec, code, ns)
        try:
            fut.result(timeout=seconds)
            return None
        except concurrent.futures.TimeoutError:
            return (f"TimeoutError: the generated code was still running after "
                    f"{seconds}s and was abandoned.\n")
        except Exception:
            return traceback.format_exc(limit=3)
    finally:
        pool.shutdown(wait=False)


def run_code(code: str, df: pd.DataFrame) -> dict:
    """Execute code with df in scope. Return stdout, error traceback, new plot paths.

    The sandbox gets a SAMPLE of large frames. Plots of a 5,000-row sample are
    visually identical and finish in seconds; every number in the report is
    computed separately, on the full frame, by the synthesis layer.
    """
    if UNSAFE.search(code):
        return {"ok": False, "stdout": "", "plots": [],
                "error": "Blocked: generated code contained a disallowed operation."}

    before = set(glob.glob(f"{PLOTS_DIR}/*.png"))

    plot_df, note = df, ""
    if len(df) > PLOT_SAMPLE_ROWS:
        plot_df = df.sample(PLOT_SAMPLE_ROWS, random_state=0)
        note = (f"[agent] plotting on a random sample of {PLOT_SAMPLE_ROWS:,} of "
                f"{len(df):,} rows; report statistics use every row.\n")

    ns = {"df": plot_df.copy(), "pd": pd, "np": np, "plt": plt, "sns": sns}

    real = _install_plot_guards()
    buf = io.StringIO()
    try:
        with contextlib.redirect_stdout(buf):
            err = _exec_with_timeout(code, ns, CODE_TIMEOUT_S)
    finally:
        for i, num in enumerate(plt.get_fignums(), start=1):
            fig = plt.figure(num)
            if fig.get_axes() and not getattr(fig, SKIP_ATTR, False):
                fig.savefig(f"{PLOTS_DIR}/plot_{i}.png")
        plt.close("all")
        _restore_plot_guards(real)

    return {"ok": err is None, "stdout": note + buf.getvalue(), "error": err,
            "plots": sorted(set(glob.glob(f"{PLOTS_DIR}/*.png")) - before)}


def format_error(err: str, code: str) -> str:
    """Turn a harness traceback into a signal about the model's own code.

    exec'd code has no source file, so Python prints the failing line number
    with no source text. We reconstruct it from the code we sent.
    """
    last = err.strip().splitlines()[-1]
    hits = re.findall(r'File "<string>", line (\d+)', err)
    if not hits:
        return last
    ln, src = int(hits[-1]), code.splitlines()
    line = src[ln - 1].strip() if 0 < ln <= len(src) else ""
    return f"Line {ln}: {line}\n{last}"


def generate_and_run(context: str, instruction: str, df: pd.DataFrame,
                     retries: int = 1, selected_plots=None) -> dict:
    """Generate, execute, repair once. Phases are timed so a slow run can be
    attributed to generation or to execution rather than guessed at."""
    shutil.rmtree(PLOTS_DIR, ignore_errors=True)
    os.makedirs(PLOTS_DIR, exist_ok=True)

    t0 = time.time()
    raw = generate(CODE_SYSTEM, build_code_prompt(context, instruction, selected_plots),
                   max_new_tokens=CODE_MAX_TOKENS, stop_fences=2)
    t_gen = time.time() - t0
    code = sanitize(extract_code(raw))

    t0 = time.time()
    res = run_code(code, df)
    t_exec = time.time() - t0
    print(f"[attempt 1] ok={res['ok']}  generate={t_gen:.1f}s  execute={t_exec:.1f}s")

    for attempt in range(retries):
        if res["ok"]:
            break
        fix_prompt = (f"{context}\n\nThis code failed:\n```python\n{code}\n```\n\n"
                      f"Error:\n{format_error(res['error'], code)}\n\n"
                      f"Return the corrected script.")
        t0 = time.time()
        code = sanitize(extract_code(
            generate(FIX_SYSTEM, fix_prompt, max_new_tokens=CODE_MAX_TOKENS,
                     stop_fences=2)))
        t_gen = time.time() - t0
        t0 = time.time()
        res = run_code(code, df)
        t_exec = time.time() - t0
        print(f"[attempt {attempt + 2}] ok={res['ok']}  generate={t_gen:.1f}s  "
              f"execute={t_exec:.1f}s")

    res["code"] = code
    res["plots"] = sorted(glob.glob(f"{PLOTS_DIR}/*.png"))
    return res


# --------------------------------------------------------------------------- #
# 5. Deterministic analysis
# --------------------------------------------------------------------------- #
REDUNDANT_R         = 0.99
MAX_GROUP_BLOCKS    = 10
IQR_K               = 1.5
OUTLIER_FLAG_PCT    = 5.0
DROP_MISSING_PCT    = 70.0
ONEHOT_LEVELS       = 2
MAX_ASSOC_COLS      = 25
MAX_ASSOC_LEVELS    = 20
FREETEXT_UNIQ_RATIO = 0.9
FREETEXT_MIN_LEN    = 15
FREETEXT_MIN_TOKENS = 1.0

ACTION_ICON = {"DROP": "\U0001F5D1", "IMPUTE": "\U0001F504", "ENCODE": "\U0001F522",
               "OUTLIERS": "⚠️", "FEATURE": "✍️"}


def fmt_num(v):
    """Plain, comma-separated numbers. Never 1.02011e+06."""
    if pd.isna(v):
        return ""
    v = float(v)
    if v.is_integer():
        return f"{int(v):,}"
    if abs(v) >= 1000:
        return f"{v:,.0f}"
    return f"{v:,.2f}"


def detect_free_text(df: pd.DataFrame) -> list:
    """A nearly-unique text column is usually free text, not an identifier."""
    out = []
    for c in df.select_dtypes(include=["object", "category"]).columns:
        s = df[c].dropna().astype(str)
        if len(s) < 5:
            continue
        uniq_ratio = s.nunique() / len(s)
        mean_len = s.str.len().mean()
        mean_tok = s.str.split().str.len().mean()
        if uniq_ratio > FREETEXT_UNIQ_RATIO and (mean_len > FREETEXT_MIN_LEN
                                                 or mean_tok > FREETEXT_MIN_TOKENS):
            out.append(c)
    return out


def cramers_v(a, b):
    """Categorical <-> categorical. Uncorrected; fine for ranking, not inference."""
    ct = pd.crosstab(a, b)
    if ct.shape[0] < 2 or ct.shape[1] < 2:
        return np.nan
    chi2 = stats.chi2_contingency(ct, correction=False)[0]
    n = ct.to_numpy().sum()
    if n <= 0:
        return np.nan
    phi2 = chi2 / n
    r, k = ct.shape
    return float(np.sqrt(phi2 / max(min(k - 1, r - 1), 1)))


def correlation_ratio(cats, vals):
    """Categorical <-> numeric: eta = sqrt(between-group var / total var)."""
    d = pd.DataFrame({"c": cats, "v": pd.to_numeric(vals, errors="coerce")}).dropna()
    if d["c"].nunique() < 2 or len(d) < 3:
        return np.nan
    grand = d["v"].mean()
    num = sum(len(g) * (g["v"].mean() - grand) ** 2 for _, g in d.groupby("c"))
    den = float(((d["v"] - grand) ** 2).sum())
    return float(np.sqrt(num / den)) if den > 0 else np.nan


def association_pairs(df: pd.DataFrame, free_text=(), id_like=()):
    """Every pair of usable columns, ranked by absolute effect size."""
    skip = set(free_text) | set(id_like)
    num = [c for c in df.select_dtypes(include="number").columns
           if c not in skip and df[c].nunique(dropna=True) > 1][:MAX_ASSOC_COLS]
    cat = [c for c in df.columns
           if c not in skip and c not in num
           and 2 <= df[c].nunique(dropna=True) <= MAX_ASSOC_LEVELS][:MAX_ASSOC_COLS]

    rows = []
    for i, a in enumerate(num):
        for b in num[i + 1:]:
            d = df[[a, b]].dropna()
            if len(d) < 3:
                continue
            r = d[a].corr(d[b])
            if pd.isna(r):
                continue
            rows.append({"pair": f"{a} ~ {b}", "type": "num-num",
                         "measure": "Pearson r", "strength": round(abs(float(r)), 3),
                         "direction": "positive" if r > 0 else "negative"})
    for i, a in enumerate(cat):
        for b in cat[i + 1:]:
            v = cramers_v(df[a], df[b])
            if pd.isna(v):
                continue
            rows.append({"pair": f"{a} ~ {b}", "type": "cat-cat",
                         "measure": "Cramer's V", "strength": round(float(v), 3),
                         "direction": "n/a"})
    for a in cat:
        for b in num:
            e = correlation_ratio(df[a], df[b])
            if pd.isna(e):
                continue
            rows.append({"pair": f"{a} ~ {b}", "type": "cat-num",
                         "measure": "corr. ratio", "strength": round(float(e), 3),
                         "direction": "n/a"})

    rows.sort(key=lambda r: r["strength"], reverse=True)
    return ([r for r in rows if r["strength"] < REDUNDANT_R],
            [r for r in rows if r["strength"] >= REDUNDANT_R])


def shape_label(skew, exkurt):
    if pd.isna(skew):
        return "undetermined"
    if abs(skew) < 0.5:
        lab = "symmetric"
    elif abs(skew) < 1:
        lab = "moderately right-skewed" if skew > 0 else "moderately left-skewed"
    else:
        lab = "strongly right-skewed" if skew > 0 else "strongly left-skewed"
    if not pd.isna(exkurt) and exkurt > 3:
        lab += ", heavy-tailed"
    return lab


def distribution_profile(df: pd.DataFrame) -> list:
    """Per numeric column: location, spread, shape. pandas.kurtosis is excess."""
    rows = []
    for c in df.select_dtypes(include="number").columns:
        s = df[c].dropna()
        if len(s) < 3 or s.nunique() <= 1:
            continue
        sk = float(s.skew())
        ku = float(s.kurtosis()) if len(s) > 3 else np.nan
        med = float(s.median())
        ratio = float(s.mean() / med) if med != 0 else np.nan
        rows.append({"column": c, "count": int(s.count()),
                     "mean": fmt_num(s.mean()), "median": fmt_num(med),
                     "std": fmt_num(s.std()), "min": fmt_num(s.min()),
                     "max": fmt_num(s.max()),
                     "skew": round(sk, 2),
                     "kurtosis": ("" if pd.isna(ku) else round(ku, 2)),
                     "mean/median": ("" if pd.isna(ratio) else round(ratio, 2)),
                     "shape": shape_label(sk, ku)})
    return rows


def run_advanced_eda_checks(df: pd.DataFrame) -> dict:
    """Column roles, integrity findings, shapes, IQR outliers and the checklist."""
    n_rows, n_cols = df.shape
    checks = {}

    free_text = detect_free_text(df)
    dup_rows = int(df.duplicated().sum())
    const_cols = [c for c in df.columns if df[c].nunique(dropna=False) <= 1]
    id_like = [c for c in df.columns
               if n_rows > 0 and df[c].nunique(dropna=True) == n_rows
               and not pd.api.types.is_float_dtype(df[c])
               and c not in free_text]
    miss_pct = (df.isna().sum() / max(n_rows, 1) * 100)
    overall_missing = float(df.isna().sum().sum()) / max(n_rows * n_cols, 1) * 100

    checks["free_text"] = free_text
    checks["dup_rows"] = dup_rows
    checks["dup_pct"] = 100 * dup_rows / max(n_rows, 1)
    checks["const_cols"] = const_cols
    checks["id_like"] = id_like
    checks["overall_missing_pct"] = overall_missing
    checks["cols_with_missing"] = int((miss_pct > 0).sum())
    checks["distribution"] = distribution_profile(df)

    outliers = []
    for c in df.select_dtypes(include="number").columns:
        s = df[c].dropna()
        if len(s) < 10 or s.nunique() <= 2:
            continue
        q1, q3 = s.quantile(0.25), s.quantile(0.75)
        iqr = q3 - q1
        if iqr == 0:
            continue
        lo, hi = q1 - IQR_K * iqr, q3 + IQR_K * iqr
        ex = s[(s < lo) | (s > hi)]
        cnt = int(len(ex))
        if not cnt:
            continue
        pct = round(100 * cnt / len(s), 1)
        far = ex.reindex((ex - s.median()).abs().sort_values(ascending=False).index)
        examples = ", ".join(fmt_num(v) for v in far.head(3))
        if pct > 10:
            rec = "large share - check whether these are a distinct sub-population"
        elif pct > OUTLIER_FLAG_PCT:
            rec = "inspect domain validity; consider winsorizing or filtering"
        else:
            rec = "few extremes - usually safe to keep"
        outliers.append({"column": c, "pct": pct, "count": cnt,
                         "examples": examples, "recommendation": rec})
    outliers.sort(key=lambda r: r["pct"], reverse=True)
    checks["outliers"] = outliers

    cats = []
    for c in df.select_dtypes(include=["object", "category"]).columns:
        k = df[c].nunique(dropna=True)
        if c in free_text:
            rec = "free text - extract features (length, tokens, title) rather than dropping"
        elif k == ONEHOT_LEVELS:
            rec = "binary - one-hot encode with pd.get_dummies"
        elif k < ONEHOT_LEVELS:
            rec = "single level - carries no information, drop"
        elif k <= 10:
            rec = "more than 2 levels - keep as is or group rare levels"
        else:
            rec = "high cardinality - group rare levels or drop"
        cats.append({"column": c, "levels": k, "recommendation": rec})
    checks["categoricals"] = cats

    cl = []
    for c in df.columns:
        if c in free_text:
            continue
        if miss_pct.get(c, 0) > DROP_MISSING_PCT:
            cl.append(("DROP", c,
                       f"{miss_pct[c]:.1f}% missing (over {DROP_MISSING_PCT:.0f}%)"))
    for c in const_cols:
        cl.append(("DROP", c, "constant column - zero variance"))
    for c in id_like:
        cl.append(("DROP", c, "unique per row - an identifier, not a feature"))

    dropped_set = {c for _, c, _ in cl}
    for c in free_text:
        if c not in dropped_set:
            cl.append(("FEATURE", c,
                       "free text - extract features (length, token count, title) "
                       "rather than dropping"))
    for c in df.columns:
        if c in dropped_set or c in free_text or miss_pct.get(c, 0) == 0:
            continue
        if pd.api.types.is_numeric_dtype(df[c]):
            cl.append(("IMPUTE", c,
                       f"{miss_pct[c]:.1f}% missing - fill with the median "
                       f"({fmt_num(df[c].median())})"))
        else:
            mode = df[c].mode(dropna=True)
            mv = str(mode.iloc[0]) if len(mode) else "?"
            cl.append(("IMPUTE", c,
                       f"{miss_pct[c]:.1f}% missing - fill with the mode ('{mv}')"))
    for row in cats:
        if (row["column"] not in dropped_set and row["column"] not in free_text
                and row["levels"] == ONEHOT_LEVELS):
            cl.append(("ENCODE", row["column"],
                       "2 levels - one-hot encode with pd.get_dummies"))
    for row in outliers:
        if row["pct"] > OUTLIER_FLAG_PCT and row["column"] not in dropped_set:
            cl.append(("OUTLIERS", row["column"],
                       f"{row['pct']}% extreme values (e.g. {row['examples']}) - "
                       f"{row['recommendation']}"))

    checks["checklist"] = cl
    return checks


def integrity_findings(checks: dict) -> list:
    """Section 2 as findings, not a score.

    A single number invited the wrong reading - a synthetic dataset with uniform
    columns and zero correlations scored 100/100.
    """
    out = []
    if checks["dup_rows"]:
        out.append(f"WARNING **{fmt_num(checks['dup_rows'])} duplicate rows** "
                   f"({checks['dup_pct']:.1f}% of the table) - consider "
                   f"`df.drop_duplicates()`.")
    else:
        out.append("OK No duplicate rows.")

    if checks["const_cols"]:
        out.append(f"WARNING **{len(checks['const_cols'])} constant column(s)**: "
                   f"`{'`, `'.join(checks['const_cols'])}` - zero variance, drop them.")
    else:
        out.append("OK No constant columns.")

    if checks["id_like"]:
        out.append(f"WARNING **{len(checks['id_like'])} identifier column(s)**: "
                   f"`{'`, `'.join(checks['id_like'])}` - unique per row, not features.")
    else:
        out.append("OK No identifier columns.")

    if checks["free_text"]:
        out.append(f"TEXT **{len(checks['free_text'])} free-text column(s)**: "
                   f"`{'`, `'.join(checks['free_text'])}` - extract features "
                   f"(length, tokens, title) rather than dropping.")

    if checks["overall_missing_pct"] > 0:
        out.append(f"WARNING **{checks['overall_missing_pct']:.1f}% of all cells are "
                   f"missing**, across {checks['cols_with_missing']} column(s) - "
                   f"see section 3.")
    else:
        out.append("OK No missing values anywhere.")

    icons = {"OK ": "✅ ", "WARNING ": "⚠️ ", "TEXT ": "✍️ "}
    return [next((v + line[len(k):] for k, v in icons.items()
                  if line.startswith(k)), line) for line in out]


def format_checklist(cl: list) -> str:
    """A grouped, scannable bullet list. A dry table nobody reads is a wasted section."""
    if not cl:
        return "_Nothing to do - the dataset is ML-ready as loaded._"
    out = []
    for act in ("DROP", "IMPUTE", "ENCODE", "OUTLIERS", "FEATURE"):
        rows = [r for r in cl if r[0] == act]
        if not rows:
            continue
        out += [f"**{ACTION_ICON[act]} {act} - {len(rows)} column(s)**", ""]
        out += [f"- {ACTION_ICON[act]} **[{act}]** `{col}`: {why}"
                for _, col, why in rows]
        out += [""]
    return "\n".join(out).strip()


def make_tables(df: pd.DataFrame, checks: dict, max_card: int = 6) -> dict:
    """Every table rendered by pandas. The model never transcribes one."""
    t, n_rows = {}, len(df)

    t["health"] = "\n".join(f"- {line}" for line in integrity_findings(checks))

    n = df.isna().sum()
    n = n[n > 0].sort_values(ascending=False)
    if len(n):
        head = n.head(15)
        strategy = ["median" if pd.api.types.is_numeric_dtype(df[c]) else "mode"
                    for c in head.index]
        tbl = pd.DataFrame({"missing": head.map(fmt_num),
                            "pct": (100 * head / n_rows).round(1),
                            "imputation": strategy}).to_markdown(disable_numparse=True)
        if len(n) > 15:
            tbl += f"\n\n_+{len(n) - 15} more columns with missing values._"
        t["missing"] = tbl
    else:
        t["missing"] = "_No missing values._"

    if checks["distribution"]:
        t["describe"] = pd.DataFrame(checks["distribution"]).to_markdown(
            index=False, disable_numparse=True)
    else:
        t["describe"] = "_No numeric columns with enough variation to profile._"

    if checks["categoricals"]:
        t["encoding"] = pd.DataFrame(checks["categoricals"]).rename(
            columns={"levels": "unique levels"}).to_markdown(index=False)
    else:
        t["encoding"] = "_No categorical columns._"

    num = df.select_dtypes(include="number")
    bins = [c for c in num.columns if df[c].nunique() == 2]
    blocks = [f"**mean `{b}` by `{c}`**\n\n"
              + df.groupby(c)[b].mean().round(3).to_frame().to_markdown()
              for c in df.columns if df[c].nunique(dropna=True) <= max_card
              for b in bins if b != c]
    if blocks:
        t["groups"] = "\n\n".join(blocks[:MAX_GROUP_BLOCKS])
        if len(blocks) > MAX_GROUP_BLOCKS:
            t["groups"] += (f"\n\n_+{len(blocks) - MAX_GROUP_BLOCKS} more group "
                            "comparisons omitted._")
    else:
        t["groups"] = "_No low-cardinality groupings._"

    real, dup = association_pairs(df, checks["free_text"], checks["id_like"])
    note = ("_Ranked by absolute effect size across mixed types: Pearson r for "
            "numeric pairs, Cramer's V for categorical pairs, and the correlation "
            "ratio for categorical-numeric pairs._\n\n")
    if real and real[0]["strength"] < 0.1:
        note += (f"_All associations are near zero (largest = "
                 f"{real[0]['strength']:.3f}) - the columns appear mutually "
                 f"independent, often a sign of synthetic data._\n\n")
    t["corr"] = note + (
        pd.DataFrame(real[:10]).to_markdown(index=False, disable_numparse=True)
        if real else "_Too few usable columns._")
    if dup:
        t["redundant"] = (
            f"_{len(dup)} column pairs associate at >= {REDUNDANT_R} - likely "
            "duplicate encodings rather than findings._\n\n"
            + pd.DataFrame([{k: r[k] for k in ("pair", "type", "strength")}
                            for r in dup[:8]]).to_markdown(index=False,
                                                           disable_numparse=True))

    if checks["outliers"]:
        t["outliers"] = pd.DataFrame(
            [{"column": r["column"], "outliers %": r["pct"],
              "example extreme values": r["examples"],
              "recommendation": r["recommendation"]}
             for r in checks["outliers"]]).to_markdown(index=False,
                                                       disable_numparse=True)
    else:
        t["outliers"] = f"_No values fall outside the {IQR_K}x IQR fences._"

    t["checklist"] = format_checklist(checks["checklist"])
    return t


def key_facts(df: pd.DataFrame, checks: dict, max_card: int = 6) -> dict:
    """Every superlative computed in pandas, never looked up by the model."""
    issues = []
    if checks["dup_rows"]:
        issues.append(f"{fmt_num(checks['dup_rows'])} duplicate rows")
    if checks["const_cols"]:
        issues.append(f"{len(checks['const_cols'])} constant columns")
    if checks["id_like"]:
        issues.append(f"{len(checks['id_like'])} identifier columns")
    if checks["overall_missing_pct"] > 0:
        issues.append(f"{checks['overall_missing_pct']:.1f}% of cells missing")
    f = {"health": ("the integrity checks found " + ", ".join(issues)
                    + f", and {len(checks['checklist'])} preprocessing actions "
                      "are recommended")
                   if issues else
                   ("the integrity checks found no duplicate rows, no constant "
                    "columns, no identifier columns and no missing values")}

    n = df.isna().sum()
    n = n[n > 0]
    if len(n):
        c = n.idxmax()
        f["missing"] = (f"{c} has the most missing values: "
                        f"{fmt_num(n.max())} ({100 * n.max() / len(df):.1f}%)")

    if checks["distribution"]:
        worst = max(checks["distribution"], key=lambda r: abs(r["skew"]))
        f["shape"] = (f"{worst['column']} is the most skewed numeric column "
                      f"(skew {worst['skew']}), described as {worst['shape']}")

    if checks["outliers"]:
        r = checks["outliers"][0]
        f["outliers"] = (f"{r['column']} has the highest share of extreme values: "
                         f"{r['pct']}% of its rows")

    real, _ = association_pairs(df, checks["free_text"], checks["id_like"])
    if real:
        r = real[0]
        if r["strength"] < 0.1:
            f["corr"] = (f"no meaningful association exists between the columns: "
                         f"the largest is {r['pair']} at {r['strength']}")
        else:
            direction = "" if r["direction"] == "n/a" else f", {r['direction']}"
            f["corr"] = (f"the strongest association is {r['pair']} "
                         f"({r['measure']} = {r['strength']}{direction})")
    return f


# --------------------------------------------------------------------------- #
# 6. Narrative — the model only phrases facts Python already computed
# --------------------------------------------------------------------------- #
def first_sentence(text: str) -> str:
    lines = [l.strip().lstrip("-*# ").strip() for l in text.splitlines() if l.strip()]
    lines = [l for l in lines if len(l.split()) >= 4] or lines
    s = lines[0] if lines else ""
    m = re.match(r"(.+?[.!?])(\s|$)", s)
    return m.group(1) if m else s[:220]


def first_sentences(text: str, n: int = 3) -> str:
    lines = [l.strip().lstrip("-*# ").strip() for l in text.splitlines() if l.strip()]
    lines = [l for l in lines if len(l.split()) >= 4] or lines
    blob = " ".join(lines)
    parts = re.findall(r"[^.!?]+[.!?]", blob)
    return " ".join(p.strip() for p in parts[:n]) or blob[:400]


def one_liner(fact: str) -> str:
    sys_p = ("Rewrite the given fact as one short sentence. Add nothing. "
             "No lists, no headings.")
    return first_sentence(generate(sys_p, fact, max_new_tokens=90))


SEMANTIC_SYSTEM = (
    "You explain what a dataset is about, using only its name and column names.\n"
    "Rules:\n"
    "- Two or three short sentences of plain prose.\n"
    "- Never state a count, percentage or any other statistic.\n"
    "- No lists, no headings, no code."
)

HYPOTHESIS_SYSTEM = (
    "You write short analytical conclusions for a data report.\n"
    "Rules:\n"
    "- Exactly three markdown bullets, one sentence each.\n"
    "- Use only the findings given. Never invent a number.\n"
    "- Each bullet is either a hypothesis worth testing or a concrete next step.\n"
    "- No preamble, no headings."
)


def semantic_overview(df: pd.DataFrame, name: str) -> str:
    """The one place the model draws on world knowledge rather than the data."""
    cols = ", ".join(map(str, df.columns))
    txt = generate(SEMANTIC_SYSTEM,
                   f"Dataset name: {name}\nColumns: {cols}\n\n"
                   "What kind of data is this, and what do the main columns appear "
                   "to measure?",
                   max_new_tokens=200)
    return first_sentences(txt, 3)


def hypotheses(facts: dict) -> str:
    """Section 10. The model reasons over facts Python already computed."""
    body = "\n".join(f"- {v}" for v in facts.values())
    txt = generate(HYPOTHESIS_SYSTEM, f"Findings:\n{body}", max_new_tokens=260)
    bullets = [l.strip() for l in txt.splitlines()
               if l.strip().startswith(("-", "*"))][:3]
    return "\n".join(bullets) or "_No hypotheses generated._"


def audit_numbers(md: str, allowed_text: str) -> list:
    """Flag numerals in the prose absent from the evidence.

    Table rows and headings are skipped: heading numbers ("## 7.") are structure,
    not claims. Detects fabrication, not misinterpretation.
    """
    allowed = set(re.findall(r"\d+\.?\d*", allowed_text.replace(",", "")))
    prose = "\n".join(l for l in md.splitlines()
                      if not l.strip().startswith(("|", "#")))
    return sorted({x for x in re.findall(r"\d+\.?\d*", prose.replace(",", ""))
                   if x not in allowed})


def embed_images(md: str) -> str:
    """Inline PNGs as data URIs so they render inside the Gradio markdown component."""
    def repl(m):
        path = m.group(1)
        if not os.path.exists(path):
            return m.group(0)
        b64 = base64.b64encode(open(path, "rb").read()).decode()
        return f"](data:image/png;base64,{b64})"
    return re.sub(r"\]\(([^)]+\.png)\)", repl, md)


def report_to_pdf(report_md: str, plots: list, path: str = "eda_report.pdf") -> str:
    """Render the markdown report plus every plot to a single PDF."""
    from fpdf import FPDF
    from fpdf.enums import XPos, YPos

    def clean(s):
        for a, b in [("—", "-"), ("…", "..."), ("≥", ">="),
                     ("×", "x"), ("−", "-"), ("’", "'"),
                     ("“", '"'), ("”", '"')]:
            s = s.replace(a, b)
        # fpdf2's core fonts are latin-1: drop emoji rather than render them as "?"
        s = re.sub(r"[^\x00-\xff]", "", s)
        return s.encode("latin-1", "replace").decode("latin-1")

    pdf = FPDF(format="A4")
    pdf.set_auto_page_break(True, margin=15)
    pdf.add_page()
    usable = pdf.w - pdf.l_margin - pdf.r_margin
    max_table_chars = int(usable / (6.5 * 0.6 * 0.3528))

    def write(text, style="", size=10, h=5):
        """Always start at the left margin - otherwise multi_cell width goes to zero."""
        pdf.set_font("Courier" if style == "mono" else "Helvetica",
                     "B" if style == "B" else "", size)
        pdf.set_x(pdf.l_margin)
        pdf.multi_cell(0, h, text, new_x=XPos.LMARGIN, new_y=YPos.NEXT)

    for raw in report_md.splitlines():
        line = clean(raw.rstrip())
        if line.startswith("!["):
            continue
        is_table = line.startswith("|")
        if not is_table:
            line = re.sub(r"[*_`]", "", line)
            line = " ".join(w if len(w) <= 50 else
                            " ".join(w[i:i + 50] for i in range(0, len(w), 50))
                            for w in line.split())
        elif len(line) > max_table_chars:
            line = line[:max_table_chars - 3] + "..."

        if line.startswith("# "):
            write(line[2:], "B", 15, 8)
            pdf.ln(1)
        elif line.startswith("## "):
            pdf.ln(2)
            write(line[3:], "B", 12, 7)
        elif is_table:
            write(line, "mono", 6.5, 3.4)
        elif line.strip():
            write(line, "", 10, 5)
        else:
            pdf.ln(2)

    for p in plots:
        pdf.add_page()
        write(clean(os.path.basename(p)[:-4].replace("_", " ")), "B", 11, 7)
        pdf.image(p, w=usable)

    pdf.output(path)
    return path


# --------------------------------------------------------------------------- #
# 7. Caches and report assembly
# --------------------------------------------------------------------------- #
_DATA_CACHE: dict = {}    # name -> (df, dropped, context)
_STATS_CACHE: dict = {}   # name -> (checks, tables, facts)
_NARR_CACHE: dict = {}    # name -> {section: sentence}


def get_dataframe(name: str):
    if name not in _DATA_CACHE:
        df = normalize(load_hf_dataframe(name))
        df, dropped = drop_index_cols(df)
        _DATA_CACHE[name] = (df, dropped, build_context(df, name))
    return _DATA_CACHE[name]


def get_stats(name: str, df: pd.DataFrame):
    if name not in _STATS_CACHE:
        checks = run_advanced_eda_checks(df)
        _STATS_CACHE[name] = (checks, make_tables(df, checks), key_facts(df, checks))
    return _STATS_CACHE[name]


def get_narrative(name: str, df: pd.DataFrame, facts: dict):
    if name not in _NARR_CACHE:
        nar = {"overview": semantic_overview(df, name)}
        for k in ("health", "missing", "shape", "corr", "outliers"):
            if k in facts:
                nar[k] = one_liner(facts[k])
        nar["hypotheses"] = hypotheses(facts)
        _NARR_CACHE[name] = nar
    return _NARR_CACHE[name]


def clear_caches(name=None):
    for c in (_DATA_CACHE, _STATS_CACHE, _NARR_CACHE):
        c.pop(name, None) if name else c.clear()


def assemble_report(df, name, instruction, dropped, tables, facts, narrative,
                    plots) -> tuple:
    """Pure string assembly - no computation, no model calls."""
    t, nar = tables, narrative
    P = [f"# EDA Report - {name}", "", f"*{instruction}*", ""]

    P += ["## 1. Dataset Overview & Schema",
          f"**{df.shape[0]:,} rows x {df.shape[1]} columns.** " + nar["overview"], ""]
    if dropped:
        P += [f"_Index-like columns excluded from analysis: {', '.join(dropped)}._", ""]

    P += ["## 2. Data Quality & Integrity Inspection", t["health"], ""]
    if "health" in nar:
        P += [nar["health"], ""]

    P += ["## 3. Missing Values & Imputation Strategy", t["missing"], ""]
    if "missing" in nar:
        P += [nar["missing"], ""]

    P += ["## 4. Distribution & Summary Characteristics", t["describe"], ""]
    if "shape" in nar:
        P += [nar["shape"], ""]

    P += ["## 5. Required Categorical Transformations", t["encoding"], ""]

    P += ["## 6. Correlations & Group Differences", t["corr"], ""]
    if "corr" in nar:
        P += [nar["corr"], ""]
    P += [t["groups"], ""]
    if "redundant" in t:
        P += ["**Redundant columns**", "", t["redundant"], ""]

    P += [f"## 7. Outliers Detection ({IQR_K}x IQR Method)", t["outliers"], ""]
    if "outliers" in nar:
        P += [nar["outliers"], ""]

    P += ["## 8. Machine Learning Readiness & Preprocessing Checklist",
          t["checklist"], ""]

    P += ["## 9. Visualizations", ""]
    for p in plots:
        P += [f"**{os.path.basename(p)[:-4].replace('_', ' ')}**", "", f"![]({p})", ""]

    P += ["## 10. Hypotheses & Next Steps",
          nar.get("hypotheses", "_No hypotheses generated._"), ""]

    report = "\n".join(P)
    allowed = "\n".join([*t.values(), *facts.values(), schema_only(df, name)])
    return report, audit_numbers(report, allowed)


# --------------------------------------------------------------------------- #
# 8. Pipeline + Gradio plumbing
# --------------------------------------------------------------------------- #
DATASET_RE = re.compile(r"^[\w.-]+(/[\w.-]+)?$")
BLANK = (None, "", "", [], "", "")


def _pipeline(instruction: str, dataset: str):
    """Streams: status, pdf, code, stdout, gallery, report, flags."""
    instruction = (instruction or "").strip()
    dataset = (dataset or "").strip()
    if not DATASET_RE.match(dataset):
        yield ("**Invalid dataset id.** Use the form `owner/name`.", *BLANK)
        return

    warm_data = dataset in _DATA_CACHE
    yield ((f"Reusing cached data for `{dataset}` ..." if warm_data
            else f"Downloading `{dataset}` ..."), *BLANK)
    try:
        df, dropped, context = get_dataframe(dataset)
    except Exception as e:
        yield (f"**Could not load `{dataset}`.** {type(e).__name__}: {e}", *BLANK)
        return

    warm_stats = dataset in _STATS_CACHE
    yield ((f"Loaded **{df.shape[0]:,} x {df.shape[1]}**. "
            + ("Statistics already computed ..." if warm_stats
               else "Computing statistics - integrity, missingness, shapes, "
                    "associations, outliers ...")), *BLANK)
    checks, tables, facts = get_stats(dataset, df)

    warm_narr = dataset in _NARR_CACHE
    yield (("Narrative already written ..." if warm_narr
            else "Writing the narrative - overview, section summaries and "
                 "hypotheses ..."), *BLANK)
    narrative = get_narrative(dataset, df, facts)

    yield ("Generating and running the plot code ...", *BLANK)
    res = generate_and_run(context, instruction, df, selected_plots=DEFAULT_PLOTS)
    status = ("Code ran successfully" if res["ok"]
              else "Code still failing after one repair - report built from data only")

    yield (f"{status}. Assembling the report ...",
           None, res["code"], res["stdout"], res["plots"], "", "")
    report, flags = assemble_report(df, dataset, instruction, dropped,
                                    tables, facts, narrative, res["plots"])

    safe = re.sub(r"[^\w.-]", "_", dataset)
    try:
        pdf = report_to_pdf(report, res["plots"], path=f"eda_report_{safe}.pdf")
        pdf_note = ""
    except Exception as e:
        pdf, pdf_note = None, f" (PDF unavailable: {type(e).__name__})"

    reused = " Statistics and narrative reused from cache." if warm_stats else ""
    yield (f"Done - {status.lower()}, {len(res['plots'])} plots.{pdf_note}{reused}",
           pdf, res["code"], res["stdout"], res["plots"],
           embed_images(report), ", ".join(flags) if flags else "none")


@spaces.GPU(duration=150)
def _run_gpu(instruction: str, dataset: str):
    yield from _pipeline(instruction, dataset)


_CACHE: dict = {}


def run_agent(instruction: str, dataset: str):
    """Cache check happens outside @spaces.GPU, so a repeat run costs no quota.

    Free ZeroGPU is 5 min/day and the requested duration is checked upfront,
    so re-demoing the same dataset would otherwise exhaust the allowance.
    """
    key = ((instruction or "").strip(), (dataset or "").strip())
    if key in _CACHE:
        status, *rest = _CACHE[key]
        yield (status + "  _(cached - no GPU used)_", *rest)
        return

    last = None
    for out in _run_gpu(*key):
        last = out
        yield out
    if last and last[5]:                 # only cache runs that produced a report
        _CACHE[key] = last


# --------------------------------------------------------------------------- #
# 9. Quick Starters — pre-generated, instant, no model call
# --------------------------------------------------------------------------- #
DEFAULT_INSTRUCTION = ("Run a comprehensive exploratory data analysis, highlighting "
                       "missing values, distributions, and key feature correlations.")
DEFAULT_DATASET = "Kogann/stockmatch-synthetic"

QS_DIR = "quickstarts"
QUICKSTARTS = [
    {"slug": "titanic",
     "label": "Titanic - hidden missing values",
     "blurb": "77% of `cabin` is missing, but the raw file hides it as empty strings.",
     "dataset": "mstz/titanic",
     "instruction": DEFAULT_INSTRUCTION},
    {"slug": "stockmatch",
     "label": "StockMatch - synthetic stocks",
     "blurb": "12,500 generated stocks - my own synthetic dataset from the final project.",
     "dataset": "Kogann/stockmatch-synthetic",
     "instruction": "Explore this dataset: missing values, distributions, correlations."},
    {"slug": "housing",
     "label": "Canada housing - skewed prices",
     "blurb": "35,768 listings with a long price tail the agent replots on a log axis.",
     "dataset": "imanmalhi/canada_realestate_listings",
     "instruction": DEFAULT_INSTRUCTION},
]


def load_quickstart(spec: dict):
    """Read a pre-generated run from disk. Returns None if it was never built."""
    d = os.path.join(QS_DIR, spec["slug"])
    meta_path = os.path.join(d, "meta.json")
    if not os.path.exists(meta_path):
        return None
    meta = json.load(open(meta_path))
    report = open(os.path.join(d, "report.md")).read()
    code = open(os.path.join(d, "code.py")).read()
    stdout = open(os.path.join(d, "stdout.txt")).read()
    plots = sorted(glob.glob(os.path.join(d, "plots", "*.png")))
    pdf = os.path.join(d, "report.pdf")
    return (spec["instruction"], spec["dataset"],
            f"**{spec['label']}** - pre-generated example, loaded instantly "
            f"(no GPU used).",
            pdf if os.path.exists(pdf) else None,
            code, stdout, plots, embed_images(report),
            ", ".join(meta.get("flags") or []) or "none")


def quickstart_handler(spec: dict):
    """Serve from cache; fall back to a live run if the cache was not uploaded."""
    def handler():
        cached = load_quickstart(spec)
        if cached is not None:
            return cached
        last = None
        for out in run_agent(spec["instruction"], spec["dataset"]):
            last = out
        status, pdf, code, stdout, plots, report, flags = last
        return (spec["instruction"], spec["dataset"], status, pdf,
                code, stdout, plots, report, flags)
    return handler


def reset_form():
    return DEFAULT_INSTRUCTION, DEFAULT_DATASET


def clear_all_caches():
    _CACHE.clear()
    clear_caches()
    return "Caches cleared - the next run will recompute from scratch."


# --------------------------------------------------------------------------- #
# 10. UI
# --------------------------------------------------------------------------- #
CSS = """
.hero {text-align:center;}
.qs-card {border:1px solid var(--border-color-primary);
          border-radius:12px; padding:10px 12px;}
"""

with gr.Blocks(title="EDAgent", theme=gr.themes.Soft(), css=CSS) as demo:
    gr.Markdown(
        "# EDAgent\n"
        "**Point it at any Hugging Face dataset. It writes its own analysis code, "
        "runs it, fixes it when it crashes, and hands back a report.**",
        elem_classes="hero",
    )
    gr.Markdown(
        "Powered by `Qwen2.5-Coder-1.5B-Instruct`. The report follows the EDA steps "
        "from the course and ends with an **ML preprocessing checklist** and "
        "**hypotheses for next steps**. Every table, statistic and association is "
        "computed in pandas and scipy - the model writes plotting code and short "
        "interpretations only, so the numbers cannot be hallucinated. A built-in "
        "check flags any figure in the text that is missing from the data."
    )

    gr.Markdown("### Quick Starters - one click, instant, no GPU used")
    with gr.Row():
        qs_buttons = []
        for spec in QUICKSTARTS:
            with gr.Column(elem_classes="qs-card"):
                btn = gr.Button(spec["label"], variant="secondary", size="lg")
                gr.Markdown(f"<small>{spec['blurb']}</small>")
                qs_buttons.append((btn, spec))

    gr.Markdown("### Or run it on any dataset")
    with gr.Row():
        instruction = gr.Textbox(label="Prompt instruction", value=DEFAULT_INSTRUCTION,
                                 lines=3, scale=3)
        dataset = gr.Textbox(label="Hugging Face dataset id", value=DEFAULT_DATASET,
                             placeholder="owner/name", lines=1, scale=1)

    with gr.Row():
        run_btn = gr.Button("Run EDA Agent", variant="primary", size="lg", scale=4)
        reset_btn = gr.Button("Reset", variant="secondary", size="lg", scale=1)
        clear_btn = gr.Button("Clear cache", variant="secondary", size="lg", scale=1)

    gr.Markdown(
        "<small>A fresh run takes one to two minutes on ZeroGPU, most of it the model "
        "writing the plotting code. Free daily GPU quota is per visitor - if you hit "
        "the limit, the Quick Starters above always work.</small>"
    )

    status = gr.Markdown()
    pdf_file = gr.File(label="Full report as PDF (text, tables and every plot)")

    with gr.Tabs():
        with gr.Tab("Report"):
            report_md = gr.Markdown()
            flags_box = gr.Textbox(label="Unverified numbers (fabrication check)",
                                   interactive=False)
        # No fixed heights: a nested scroll region with nothing to scroll makes
        # Safari swallow the wheel event instead of passing it to the page.
        with gr.Tab("Plots"):
            gallery = gr.Gallery(label="Plots", columns=2)
        with gr.Tab("Generated code"):
            code_box = gr.Code(language="python", label="Model-written analysis code")
        with gr.Tab("Execution output"):
            stdout_box = gr.Textbox(label="stdout / traceback", lines=8,
                                    max_lines=30, interactive=False)

    live_outputs = [status, pdf_file, code_box, stdout_box, gallery, report_md,
                    flags_box]
    qs_outputs = [instruction, dataset] + live_outputs

    run_btn.click(run_agent, [instruction, dataset], live_outputs)
    reset_btn.click(reset_form, None, [instruction, dataset])
    clear_btn.click(clear_all_caches, None, status)
    for btn, spec in qs_buttons:
        btn.click(quickstart_handler(spec), None, qs_outputs)

if __name__ == "__main__":
    demo.queue().launch()