sft_generator.py

"""
Ultra-Fast SFT Data Generator — No sklearn, No Environment Execution.

Instead of running the environment live, we generate realistic prompt/response
pairs directly from templates using known dataset states.

This is correct because:
 - We know exactly what inspect_dataset returns (from dataset_generator)
 - We know what query_cleaner returns (from specialist_agents)
 - We know the reward trajectory
 - The actual RL training will run the real environment — SFT just warms up
   the LLM's action distribution (command grammar + strategy)

Output: ~1000+ diverse examples in under 10 seconds.
"""

import json
import os
import random
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from server.dataset_generator import TASK_CONFIGS

rng = random.Random(42)

TASKS = list(TASK_CONFIGS.keys())

# ── Prompt templates ──────────────────────────────────────────────────────────

def make_prompt(
    task: str,
    step: int,
    max_steps: int,
    current_acc: float,
    target_acc: float,
    baseline_acc: float,
    dataset_shape: str,
    rows_pct: float,
    quality: float,
    budget: int,
    session: str,
    validate_left: int,
    last_obs: str,
) -> str:
    gap = max(0.0, target_acc - current_acc)
    return (
        f"You are a Data-Centric AI agent improving an ML dataset.\n\n"
        f"Task: {task}\n"
        f"Step: {step}/{max_steps}\n"
        f"Current accuracy: {current_acc:.4f}  "
        f"Target: {target_acc:.4f}  Gap: {gap:.4f}\n"
        f"Baseline accuracy: {baseline_acc:.4f}\n"
        f"Dataset: {dataset_shape} | "
        f"Rows preserved: {rows_pct*100:.1f}%\n"
        f"Quality score: {quality:.4f} | "
        f"Budget remaining: {budget}\n"
        f"Active session: {session} | "
        f"Validate calls left: {validate_left}\n\n"
        f"Last observation:\n{last_obs}\n\n"
        f"What is your next command?"
    )


# ── Observation text snippets ────────────────────────────────────────────────

INSPECT_OBS_TEMPLATES = [
    "=== Dataset Inspection ===\nShape: {rows} rows × {cols} features\nOriginal rows: {rows} | Preserved: 100.0%\nDuplicates: {dups}\nMissing values:\n  {col}: {missing}\nClass distribution: {dist}\nDtypes: {{'age': 'float64', 'score': 'float64', 'target': 'int64'}}",
    "=== Dataset Inspection ===\nShape: {rows} rows × {cols} features\nDuplicates: {dups}\nMissing values:\n  {col}: {missing}\nClass distribution: {dist}",
]

INSPECT_MODEL_TEMPLATES = [
    "=== Model Inspection ===\nAccuracy: {acc:.4f}\n  Class 0: precision={p0:.3f} recall={r0:.3f} f1={f0:.3f}\n  Class 1: precision={p1:.3f} recall={r1:.3f} f1={f1:.3f}\nTarget: {target:.4f} | Not yet",
    "=== Model Inspection (cached) ===\nAccuracy: {acc:.4f}\nTarget: {target:.4f} | Not yet",
]

CLEANER_OBS_TEMPLATES = [
    "=== Cleaner Recommendations ===\n[1] Fill {n} missing values in '{col}' using mean ({mean:.2f})\n    type=fill_missing impact=+0.075 confidence=0.90\n[2] Remove {dups} duplicate rows\n    type=remove_duplicates impact=+0.020 confidence=0.95",
    "=== Cleaner Recommendations ===\n[1] Fill {n} missing values in '{col}' using mean ({mean:.2f})\n    type=fill_missing impact=+0.075 confidence=0.90\n[2] Fix {typos} type errors in 'income'\n    type=fix_type_errors impact=+0.040 confidence=0.75",
    "=== Cleaner Recommendations ===\n[1] Fill {n} missing values in '{col}' using mean ({mean:.2f})\n    type=fill_missing impact=+0.075 confidence=0.90",
]

BALANCER_OBS_TEMPLATES = [
    "=== Balancer Recommendations ===\n[1] Upsample minority class 1 from {min_c} to {maj_c} rows via random oversampling (imbalance ratio: {ratio:.2f})\n    type=oversample impact=+0.053 confidence=0.80",
    "=== Balancer Recommendations ===\n[1] Downsample majority class 0 from {maj_c} to {min_c} rows\n    type=undersample impact=+0.030 confidence=0.70",
]

APPLY_OBS_TEMPLATES = [
    "Applied: fill_missing [Fill {n} missing values in '{col}' using mean ({mean:.2f})]\n\nDataset health check:\n  Missing values: {remaining} remaining (was {was})\n  Duplicates: ✓ (was 0)\n  Row count: {rows}/{orig} (100.0% preserved)\n\nEstimated quality score: {quality:.4f}\nBudget remaining: {budget}",
    "Applied: remove_duplicates [Remove {dups} duplicate rows]\n\nDataset health check:\n  Missing values: {remaining} remaining (was {was})\n  Duplicates: ✓ (was {dups})\n  Row count: {rows}/{orig} ({pct:.1f}% preserved)\n\nEstimated quality score: {quality:.4f}\nBudget remaining: {budget}",
    "Applied: oversample [Upsample minority class 1 via random oversampling]\n\nDataset health check:\n  Missing values: 0 remaining (was 0)\n  Duplicates: 2 remaining (was 0)\n  Row count: {rows}/{orig} (102.0% preserved)\n\nEstimated quality score: {quality:.4f}\nBudget remaining: {budget}",
]

VALIDATE_OBS_TEMPLATES = [
    "=== Validate ===\nRF Accuracy: {acc:.4f}  (primary)\nLR Accuracy: {lr_acc:.4f}  (secondary)\nAgreement:   BOTH_AGREE_IMPROVE -- fix is robust and generalises\n  Class 0: p={p:.3f} r={r:.3f} f1={f:.3f}\n  Class 1: p={p:.3f} r={r:.3f} f1={f:.3f}\nTarget: {target:.4f} | {status}",
    "=== Validate ===\nRF Accuracy: {acc:.4f}  (primary)\nLR Accuracy: {lr_acc:.4f}  (secondary)\nAgreement:   BOTH_AGREE_IMPROVE -- fix is robust and generalises\nTarget: {target:.4f} | {status}",
    "=== Validate (cached) ===\nRF Accuracy: {acc:.4f}  (primary)\nLR Accuracy: {lr_acc:.4f}  (secondary)\nTarget: {target:.4f} | {status}",
]

ERROR_OBS_TEMPLATES = [
    "Error: Recommendation 1 has already been applied this session. Duplicate apply not allowed.",
    "Validate on cooldown. Take 1 more action(s) before validating again.",
    "Error: stale recommendation ID 99. Please re-query for fresh recommendations.",
]

RESET_OBS = (
    "Episode started: {task}\n"
    "Baseline accuracy: {baseline:.4f} | Target: {target:.4f}\n"
    "Dataset: {rows} rows \u00d7 {cols} features\n"
    "Budget: {budget} steps\n\n"
    "Available commands:\n"
    "  inspect_dataset          \u2014 shape, dtypes, missing, class distribution\n"
    "  inspect_model            \u2014 accuracy (RF + LR), F1, feature importance\n"
    "  query_analyst            \u2014 holistic diagnosis + prioritised action plan (costs 2 budget total)\n"
    "  query_cleaner            \u2014 get cleaning recommendations\n"
    "  query_augmenter [class]  \u2014 get augmentation suggestions\n"
    "  query_balancer           \u2014 get resampling recommendations\n"
    "  query_validator          \u2014 check rule violations (costs 2 budget total)\n"
    "  apply [id]               \u2014 apply recommendation by ID\n"
    "  reject [id]              \u2014 reject a recommendation\n"
    "  validate                 \u2014 retrain and score (cooldown applies)\n"
    "  submit                   \u2014 finalize episode"
)

ANALYST_OBS_TEMPLATES = [
    "=== Analyst Report (costs 1 budget) ===\nDIAGNOSIS:\n  - Class Imbalance: severity={imb:.2f} [HIGH] -> use query_balancer\n  - Missing Values: severity={miss:.2f} [MEDIUM] -> use query_cleaner\n  - Type Errors: severity=0.00 [NONE]\n  - Accuracy gap: {gap:.4f} (significant gap)\n\nRECOMMENDED PLAN (budget remaining: {budget}):\n  1. query_balancer -> apply best recommendation\n  2. query_cleaner -> apply best recommendation\n  3. validate (check accuracy improvement)\n  4. submit if accuracy >= target\n\nPRIORITY NOTE: Class imbalance is the dominant issue -- fix this first.",
    "=== Analyst Report (costs 1 budget) ===\nDIAGNOSIS:\n  - Missing Values: severity={miss:.2f} [HIGH] -> use query_cleaner\n  - Class Imbalance: severity={imb:.2f} [LOW] -> use query_balancer\n  - Type Errors: severity=0.00 [NONE]\n  - Accuracy gap: {gap:.4f} (significant gap)\n\nRECOMMENDED PLAN (budget remaining: {budget}):\n  1. query_cleaner -> apply best recommendation\n  2. query_balancer -> apply best recommendation\n  3. validate\n  4. submit",
]


# ── Episode builders ─────────────────────────────────────────────────────────

def sample_dataset_params(task: str, seed: int):
    """Sample realistic dataset params for a given task."""
    cfg = TASK_CONFIGS[task]
    rng2 = random.Random(seed)
    rows_map = {"task_0_tutorial": 100, "task_1_easy": 200,
                "task_2_medium": 500, "task_3_hard": 900}
    cols_map = {"task_0_tutorial": 4, "task_1_easy": 5,
                "task_2_medium": 7, "task_3_hard": 10}
    rows = rows_map[task]
    cols = cols_map[task]
    missing_cols = ["age", "income", "score"][:rng2.randint(1, 3)]
    missing_pct = rng2.uniform(0.10, 0.30)
    n_missing = int(rows * missing_pct)
    mean_val = rng2.uniform(30.0, 60.0)
    dups = rng2.randint(0, int(rows * 0.05))
    maj_class = int(rows * rng2.uniform(0.52, 0.65))
    min_class = rows - maj_class
    return {
        "task": task, "rows": rows, "cols": cols,
        "missing_col": missing_cols[0], "n_missing": n_missing,
        "mean_val": round(mean_val, 2), "dups": dups,
        "maj_class": maj_class, "min_class": min_class,
        "baseline": cfg["baseline_accuracy"],
        "target": cfg["target_accuracy"],
        "budget": cfg["budget"],
    }


def build_episode(task: str, seed: int, strategy: list) -> list:
    """
    Build a synthetic SFT episode using template obs + fixed action sequence.
    Returns list of {prompt, response} dicts.
    """
    p = sample_dataset_params(task, seed)
    cfg = TASK_CONFIGS[task]
    examples = []

    acc = p["baseline"]
    quality = round(rng.uniform(0.45, 0.65), 4)
    rows = p["rows"]
    missing_remaining = p["n_missing"]
    budget = p["budget"]
    session = "none"
    validate_left = 3
    prev_obs = RESET_OBS.format(
        task=task, baseline=p["baseline"], target=p["target"],
        rows=rows, cols=p["cols"], budget=budget
    )

    for step, action in enumerate(strategy):
        prompt = make_prompt(
            task=task, step=step, max_steps=p["budget"],
            current_acc=acc, target_acc=p["target"], baseline_acc=p["baseline"],
            dataset_shape=f"{rows} rows × {p['cols']} columns",
            rows_pct=rows / p["rows"], quality=quality, budget=budget,
            session=session, validate_left=validate_left, last_obs=prev_obs,
        )
        examples.append({"prompt": prompt, "response": action})

        # Simulate observation update
        budget -= 1
        cmd = action.split()[0].lower()

        if cmd == "inspect_dataset":
            t = rng.choice(INSPECT_OBS_TEMPLATES)
            dist = f"class 0: {p['maj_class']}, class 1: {p['min_class']}"
            prev_obs = t.format(
                rows=rows, cols=p["cols"], dups=p["dups"],
                col=p["missing_col"], missing=missing_remaining, dist=dist,
            )
        elif cmd == "inspect_model":
            t = rng.choice(INSPECT_MODEL_TEMPLATES)
            p0 = round(rng.uniform(0.55, 0.75), 3)
            r0 = round(rng.uniform(0.55, 0.75), 3)
            prev_obs = t.format(
                acc=acc, target=p["target"],
                p0=p0, r0=r0, f0=round(2*p0*r0/(p0+r0+1e-9), 3),
                p1=p0, r1=r0, f1=round(2*p0*r0/(p0+r0+1e-9), 3),
            )
        elif cmd == "query_cleaner":
            t = rng.choice(CLEANER_OBS_TEMPLATES)
            session = f"cleaner:{seed:08x}"
            prev_obs = t.format(
                n=missing_remaining, col=p["missing_col"],
                mean=p["mean_val"], dups=p["dups"], typos=rng.randint(2, 8),
            )
        elif cmd == "query_balancer":
            t = rng.choice(BALANCER_OBS_TEMPLATES)
            session = f"balancer:{seed:08x}"
            ratio = round(p["min_class"] / max(p["maj_class"], 1), 2)
            prev_obs = t.format(
                min_c=p["min_class"], maj_c=p["maj_class"], ratio=ratio
            )
        elif cmd == "query_augmenter":
            session = f"augmenter:{seed:08x}"
            cls = action.split()[1] if len(action.split()) > 1 else "0"
            n_synth = rng.randint(5, 25)
            prev_obs = (
                f"=== Augmenter Recommendations ===\n"
                f"[1] Synthesize {n_synth} samples for class {cls} via SMOTE\n"
                f"    type=augment_class impact=+0.040 confidence=0.72"
            )
        elif cmd == "query_analyst":
            budget -= 1  # costs 1 extra
            t = rng.choice(ANALYST_OBS_TEMPLATES)
            imb = round(rng.uniform(0.3, 0.8), 2)
            miss = round(rng.uniform(0.1, 0.5), 2)
            gap = round(p["target"] - acc, 4)
            prev_obs = t.format(imb=imb, miss=miss, gap=gap, budget=budget)
        elif cmd == "query_validator":
            budget -= 1  # costs 2
            prev_obs = (
                "=== Validator Report (costs 2 budget) ===\n"
                f"  [WARNING] [{p['missing_col']}] rule=no_missing "
                f"count={missing_remaining}\n"
                f"    Column '{p['missing_col']}' has {missing_remaining} missing values."
            )
        elif cmd == "apply":
            rec_id = int(action.split()[1]) if len(action.split()) > 1 else 1
            t = rng.choice(APPLY_OBS_TEMPLATES)
            was_missing = missing_remaining
            missing_remaining = max(0, missing_remaining - p["n_missing"])
            quality = min(1.0, quality + rng.uniform(0.10, 0.35))
            quality = round(quality, 4)
            prev_obs = t.format(
                n=p["n_missing"], col=p["missing_col"], mean=p["mean_val"],
                remaining=missing_remaining, was=was_missing,
                rows=rows, orig=p["rows"], pct=rows/p["rows"]*100,
                dups=p["dups"], quality=quality, budget=budget,
            )
        elif cmd == "reject":
            prev_obs = f"Recommendation {action.split()[1] if len(action.split())>1 else 1} rejected."
        elif cmd == "validate":
            if validate_left > 0:
                acc = min(1.0, acc + rng.uniform(0.05, 0.35))
                acc = round(acc, 4)
                lr_acc = round(min(1.0, acc + rng.uniform(-0.03, 0.03)), 4)
                validate_left -= 1
                t = rng.choice(VALIDATE_OBS_TEMPLATES)
                status = "HIT \u2713" if acc >= p["target"] else "Not yet"
                pv = round(rng.uniform(0.75, 0.98), 3)
                rv = round(rng.uniform(0.75, 0.98), 3)
                prev_obs = t.format(
                    acc=acc, lr_acc=lr_acc, target=p["target"], status=status,
                    p=pv, r=rv, f=round(2*pv*rv/(pv+rv+1e-9), 3),
                )
            else:
                prev_obs = "Validate on cooldown. Take 2 more action(s) before validating again."
        elif cmd == "submit":
            break

    return examples


# ── Strategy sequences ────────────────────────────────────────────────────────

STRATEGIES = {
    "minimal_clean":             ["inspect_dataset", "query_cleaner", "apply 1", "apply 2", "inspect_dataset", "validate", "submit"],
    "inspect_model_first":       ["inspect_dataset", "inspect_model", "query_cleaner", "apply 1", "inspect_dataset", "validate", "submit"],
    "clean_then_balance":        ["inspect_dataset", "query_cleaner", "apply 1", "apply 2", "query_balancer", "apply 1", "inspect_dataset", "validate", "submit"],
    "reject_then_apply":         ["inspect_dataset", "query_cleaner", "reject 1", "apply 2", "inspect_dataset", "validate", "submit"],
    "baseline_validate_first":   ["inspect_dataset", "validate", "query_cleaner", "apply 1", "inspect_dataset", "validate", "submit"],
    "augment_path":              ["inspect_dataset", "query_cleaner", "apply 1", "query_augmenter 0", "apply 1", "inspect_dataset", "validate", "submit"],
    "with_validator":            ["inspect_dataset", "query_validator", "query_cleaner", "apply 1", "inspect_dataset", "validate", "submit"],
    "deep_clean_requery":        ["inspect_dataset", "query_cleaner", "apply 1", "apply 2", "query_cleaner", "apply 1", "inspect_dataset", "validate", "submit"],
    "fast_submit":               ["query_cleaner", "apply 1", "apply 2", "inspect_dataset", "submit"],
    "balance_heavy":             ["inspect_dataset", "query_balancer", "apply 1", "query_cleaner", "apply 1", "inspect_dataset", "validate", "submit"],
    "reject_requery":            ["inspect_dataset", "query_cleaner", "reject 1", "reject 2", "query_cleaner", "apply 1", "inspect_dataset", "validate", "submit"],
    "multi_augment":             ["inspect_dataset", "query_cleaner", "apply 1", "query_augmenter 1", "apply 1", "inspect_dataset", "validate", "submit"],
    "model_then_balance":        ["inspect_model", "query_balancer", "apply 1", "inspect_dataset", "validate", "submit"],
    "full_pipeline":             ["inspect_dataset", "inspect_model", "query_cleaner", "apply 1", "query_balancer", "apply 1", "query_augmenter 0", "apply 1", "inspect_dataset", "validate", "submit"],
    "suboptimal_no_validate":    ["inspect_dataset", "query_cleaner", "apply 1", "submit"],
    "inspect_only_submit":       ["inspect_dataset", "inspect_model", "submit"],
    "reject_all_then_requery":   ["inspect_dataset", "query_cleaner", "reject 1", "reject 2", "query_balancer", "apply 1", "inspect_dataset", "validate", "submit"],
    "apply3_then_validate":      ["inspect_dataset", "query_cleaner", "apply 1", "apply 2", "query_balancer", "apply 1", "query_augmenter 0", "apply 1", "inspect_dataset", "validate", "submit"],
    # NEW: analyst-led strategies
    "analyst_led_clean":         ["query_analyst", "inspect_dataset", "query_cleaner", "apply 1", "apply 2", "validate", "submit"],
    "analyst_led_balance":       ["query_analyst", "query_balancer", "apply 1", "query_cleaner", "apply 1", "validate", "submit"],
    "analyst_full_pipeline":     ["query_analyst", "inspect_dataset", "inspect_model", "query_cleaner", "apply 1", "query_balancer", "apply 1", "validate", "submit"],
}


def generate_sft_data(output_file: str = "sft_data.jsonl", seeds_per_combo: int = 15):
    sft_examples = []

    print(f"Generating SFT data: {len(STRATEGIES)} strategies × {len(TASKS)} tasks × {seeds_per_combo} seeds")

    for strategy_name, sequence in STRATEGIES.items():
        strategy_examples = []
        for task in TASKS:
            for seed in range(seeds_per_combo):
                episode = build_episode(task, seed, sequence)
                strategy_examples.extend(episode)
        sft_examples.extend(strategy_examples)
        print(f"  {strategy_name:<30} +{len(strategy_examples)} examples")

    rng.shuffle(sft_examples)

    out_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), output_file)
    with open(out_path, "w", encoding="utf-8") as f:
        for ex in sft_examples:
            f.write(json.dumps(ex) + "\n")

    # Diversity report
    from collections import Counter
    responses = [ex["response"] for ex in sft_examples]
    unique_cmds = set(responses)

    print(f"\n{'='*55}")
    print(f"Total examples:   {len(sft_examples)}")
    print(f"Unique commands:  {len(unique_cmds)}")
    print(f"Unique prompts:   {len(set(ex['prompt'] for ex in sft_examples))}")
    print(f"\nResponse distribution:")
    for cmd, cnt in Counter(responses).most_common():
        pct = cnt / len(responses) * 100
        bar = "#" * int(pct / 2)
        flag = " ← DOMINANT" if pct > 25 else ""
        print(f"  {cmd:<32} {cnt:>5}  ({pct:5.1f}%)  {bar}{flag}")

    print(f"\nOutput: {out_path}")
    print("✓ SFT generation complete (no sklearn, instant).")
    return sft_examples


if __name__ == "__main__":
    generate_sft_data()