#!/usr/bin/env python3
"""Synthetic data generation pipeline: seed -> teacher -> filter -> JSONL.

Demonstrates each stage of the Self-Instruct-inspired pipeline:
1. Small seed set of hand-written examples
2. Template-based expansion (simulates teacher LLM)
3. Quality filtering (dedup, length, validation)
4. JSONL output for fine-tuning (TRL/Unsloth format)
"""
from __future__ import annotations

import hashlib
import json
import random
from dataclasses import dataclass, asdict


@dataclass
class Instruction:
    instruction: str
    input: str = ""
    output: str = ""


# --- Stage 1: Small, high-quality seed set ---

SEED_SET: list[Instruction] = [
    Instruction("Summarize the key points", "The report shows Q3 revenue grew 12% to $4.2B.",
                "Q3 revenue grew 12% YoY to $4.2 billion."),
    Instruction("Classify sentiment",
                "This product is amazing! Works perfectly.",
                "positive"),
    Instruction("Extract the date",
                "The meeting is scheduled for March 15, 2024 at 2pm.",
                "2024-03-15"),
    Instruction("Fix the grammar",
                "He go to store yesterday.",
                "He went to the store yesterday."),
]


# --- Stage 2: Template-based expansion (simulating a teacher LLM) ---

TEMPLATES: list[str] = [
    "Explain {concept} in simple terms.",
    "What are the benefits of {concept}?",
    "Compare {concept} with {concept2}.",
    "List {n} {things} for {topic}.",
]

CONCEPTS = ["vector search", "tokenization", "attention", "RAG", "fine-tuning"]
CONCEPT_PAIRS = [("HNSW", "IVF"), ("SQL", "NoSQL"), ("RAG", "fine-tuning")]
THINGS = ["best practices", "common mistakes", "tools", "metrics"]
TOPICS = ["NLP", "databases", "ML pipelines", "API design"]


def expand_instructions(seeds: list[Instruction], n: int) -> list[Instruction]:
    """Simulate a teacher expanding the seed set via templates."""
    expanded = list(seeds)
    rng = random.Random(42)

    for _ in range(n - len(seeds)):
        template = rng.choice(TEMPLATES)
        instruction = template.format(
            concept=rng.choice(CONCEPTS),
            concept2=rng.choice(CONCEPTS),
            n=rng.randint(3, 7),
            things=rng.choice(THINGS),
            topic=rng.choice(TOPICS),
        )
        expanded.append(Instruction(
            instruction=instruction,
            input="",
            output=f"[SYNTHETIC] Response for: {instruction[:40]}...",
        ))
    return expanded


# --- Stage 3: Quality filtering ---

def deduplicate(samples: list[Instruction]) -> list[Instruction]:
    """Remove near-duplicate instructions by hashing."""
    seen: set[str] = set()
    unique = []
    for s in samples:
        h = hashlib.md5(s.instruction.encode()).hexdigest()
        if h not in seen:
            seen.add(h)
            unique.append(s)
    return unique


def filter_by_length(samples: list[Instruction], min_out: int = 5, max_out: int = 500) -> list[Instruction]:
    """Drop samples with too-short or too-long outputs."""
    return [s for s in samples if min_out <= len(s.output) <= max_out]


def filter_pipeline(samples: list[Instruction]) -> list[Instruction]:
    before = len(samples)
    samples = deduplicate(samples)
    samples = filter_by_length(samples)
    after = len(samples)
    print(f"Filtered: {before} -> {after} samples (dropped {before - after})")
    return samples


# --- Stage 4: Output JSONL ---

def write_jsonl(samples: list[Instruction], path: str) -> None:
    with open(path, "w") as f:
        for s in samples:
            f.write(json.dumps(asdict(s)) + "\n")
    print(f"Wrote {len(samples)} samples to {path}")


def main() -> None:
    # Expand
    expanded = expand_instructions(SEED_SET, n=100)
    print(f"Expanded: {len(expanded)} samples")

    # Filter
    filtered = filter_pipeline(expanded)

    # Report diversity
    instruction_lengths = [len(s.instruction) for s in filtered]
    print(f"Instruction length: min={min(instruction_lengths)}, "
          f"max={max(instruction_lengths)}, mean={sum(instruction_lengths)/len(instruction_lengths):.0f}")

    # Split train/eval (90/10)
    random.Random(42).shuffle(filtered)
    split = int(len(filtered) * 0.9)
    train, eval_set = filtered[:split], filtered[split:]

    write_jsonl(train, "train.jsonl")
    write_jsonl(eval_set, "eval.jsonl")

    print(f"\nTrain: {len(train)}, Eval: {len(eval_set)}")
    print("Ready for fine-tuning with TRL/Unsloth.")
    # Clean up
    import os
    os.unlink("train.jsonl")
    os.unlink("eval.jsonl")


if __name__ == "__main__":
    main()