dataset/scripts/export_training_data.py

"""
Export validated dataset to train/val/test splits.

Produces two task datasets from the same source samples:

1. SQL generation  (prompt = question + candidates CSV, completion = SQL)
2. Place extraction (prompt = question only, completion = PlacesResult JSON)

Place extraction pairs are derived automatically: for each SQL sample the
selected_candidates give us the correct place names that the extractor should
return.

Output layout (all paths relative to dataset/):
    output/runs/{run_name}/sql/train.jsonl
    output/runs/{run_name}/sql/val.jsonl
    output/runs/{run_name}/sql/test.jsonl
    output/runs/{run_name}/places/train.jsonl
    output/runs/{run_name}/places/val.jsonl
    output/runs/{run_name}/places/test.jsonl
    output/runs/{run_name}/stats.json
"""

import copy
import json
import random
import sys
from collections import defaultdict
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import sqlparse
import yaml

from gazet.config import DIVISIONS_AREA_PATH, NATURAL_EARTH_PATH


# ---------------------------------------------------------------------------
# Loading
# ---------------------------------------------------------------------------

def load_samples(jsonl_path: Path) -> List[Dict[str, Any]]:
    samples = []
    with open(jsonl_path) as f:
        for line in f:
            line = line.strip()
            if line:
                samples.append(json.loads(line))
    return samples


def load_run_name(config_path: Optional[Path]) -> str:
    if config_path and config_path.exists():
        with open(config_path) as f:
            cfg = yaml.safe_load(f)
        return cfg.get("run_name", "default")
    return "default"


# ---------------------------------------------------------------------------
# Splitting
# ---------------------------------------------------------------------------

def stratified_split(
    samples: List[Dict[str, Any]],
    train_ratio: float = 0.8,
    val_ratio: float = 0.1,
    seed: int = 42,
) -> Tuple[List[Dict], List[Dict], List[Dict]]:
    """Split stratified by template_id so every template is represented in each split.

    Stratifying by task_family let rare template variants (e.g. partial_05,
    diff_02) land entirely in train and never appear in val/test.
    """
    random.seed(seed)
    by_tpl: Dict[str, List] = defaultdict(list)
    for s in samples:
        key = s["metadata"].get("template_id") or s["metadata"].get("task_family", "unknown")
        by_tpl[key].append(s)

    train, val, test = [], [], []
    for tpl_samples in by_tpl.values():
        random.shuffle(tpl_samples)
        n = len(tpl_samples)
        n_train = int(n * train_ratio)
        n_val = int(n * val_ratio)
        train.extend(tpl_samples[:n_train])
        val.extend(tpl_samples[n_train : n_train + n_val])
        test.extend(tpl_samples[n_train + n_val :])

    random.shuffle(train)
    random.shuffle(val)
    random.shuffle(test)
    return train, val, test


# ---------------------------------------------------------------------------
# SQL generation format
# Conversational prompt-completion: model sees system + user, generates SQL.
# ---------------------------------------------------------------------------

_SQL_SYSTEM = """You are a text to SQL query translator that helps in natural language geocoding.

You have access to two DuckDB parquet tables. Given a set of candidate entities and a user query, generate the SQL to retrieve the desired geometry.

<SCHEMA>
1. divisions_area  -- Overture polygon/multipolygon admin boundaries
   query: read_parquet('divisions_area')
   columns:
     id VARCHAR              -- unique feature id
     names STRUCT("primary" VARCHAR, ...)
     country VARCHAR         -- ISO 3166-1 alpha-2
     subtype VARCHAR         -- country | region | county
     class VARCHAR
     region VARCHAR
     admin_level INTEGER
     division_id VARCHAR
     is_land BOOLEAN
     is_territorial BOOLEAN
     geometry GEOMETRY       -- WGS-84 polygon/multipolygon (spatial ext loaded)

2. natural_earth  -- Natural Earth geography polygons (oceans, seas, rivers, terrain)
   query: read_parquet('natural_earth')
   columns:
     id VARCHAR              -- unique feature id prefixed 'ne_'
     names STRUCT("primary" VARCHAR, ...)
     country VARCHAR
     subtype VARCHAR         -- e.g. 'ocean', 'sea', 'bay', 'range/mtn', 'island group'
     class VARCHAR
     region VARCHAR
     admin_level INTEGER
     is_land BOOLEAN
     is_territorial BOOLEAN
     geometry GEOMETRY       -- WGS-84 polygon/multipolygon (spatial ext loaded)
</SCHEMA>

The candidates table has a 'source' column: 'divisions_area' or 'natural_earth'.
Use read_parquet('divisions_area') or read_parquet('natural_earth') accordingly.
Use ST_AsGeoJSON(geometry) for all geometry outputs."""

_CANDIDATES_COLS = [
    "source", "id", "name", "subtype", "country", "region",
    "admin_level",
]


def _candidates_csv(candidates: List[Dict]) -> str:
    import io
    import csv
    rows = []
    for c in candidates:
        row = {col: c.get(col, "") for col in _CANDIDATES_COLS if col in c}
        rows.append(row)
    if not rows:
        return ""
    buf = io.StringIO()
    writer = csv.DictWriter(buf, fieldnames=[k for k in _CANDIDATES_COLS if k in rows[0]])
    writer.writeheader()
    writer.writerows(rows)
    return buf.getvalue().strip()


def _to_symbolic_sql(sql: str) -> str:
    """Normalize any hardcoded or runtime paths back to symbolic names."""
    sql = sql.replace(DIVISIONS_AREA_PATH, "divisions_area")
    sql = sql.replace(NATURAL_EARTH_PATH, "natural_earth")
    sql = sql.replace("/data/overture/division_area/*.parquet",          "divisions_area")
    sql = sql.replace("/data/overture/divisions_area/*.parquet",         "divisions_area")
    sql = sql.replace("/data/natural_earth_geoparquet/ne_geography.parquet", "natural_earth")
    return sql


def _format_sql(sql: str) -> str:
    """Pretty-print SQL so the model learns clean, readable style."""
    return sqlparse.format(
        sql,
        reindent=True,
        keyword_case="upper",
        indent_width=4,
    ).strip()


def _shuffle_candidates_for_export(sample: Dict[str, Any]) -> Dict[str, Any]:
    """Return a copy of sample with candidate row order shuffled deterministically.

    Candidate IDs remain unchanged; only row order changes. This removes the
    positional shortcut where the true anchor often appears first in the raw
    synthetic samples, while keeping selected_candidates valid.
    """
    shuffled = copy.deepcopy(sample)
    candidates = shuffled.get("candidates", [])
    if len(candidates) <= 1:
        return shuffled

    rng = random.Random(shuffled.get("id", "sample"))
    rng.shuffle(candidates)
    return shuffled


def sample_to_sql_pair(sample: Dict[str, Any]) -> Optional[Dict]:
    """Convert a raw sample to a conversational prompt-completion pair for SQL generation."""
    sql = sample.get("target", {}).get("sql", "").strip()
    if not sql:
        return None
    sql = _format_sql(_to_symbolic_sql(sql))

    user_content = (
        f"<CANDIDATES>\n{_candidates_csv(sample.get('candidates', []))}\n</CANDIDATES>\n\n"
        f"<USER_QUERY>\n{sample['question']}\n</USER_QUERY>"
    )

    return {
        "messages": [
            {"role": "system",    "content": _SQL_SYSTEM},
            {"role": "user",      "content": user_content},
            {"role": "assistant", "content": sql},
        ],
        "metadata": sample.get("metadata", {}),
    }


# ---------------------------------------------------------------------------
# Place extraction format
# Derived from the same SQL samples: selected_candidates → PlacesResult JSON.
# ---------------------------------------------------------------------------

_PLACE_SYSTEM = """You are a geographic entity extractor. Extract the place names the user is asking about and return valid JSON only.

OUTPUT FORMAT:
{"places": [{"place": "<name>"}]}

RULES:
- Extract the place or places that are the actual anchors of the query.
- Physical features are valid places: oceans, seas, gulfs, bays, straits, rivers, lakes, basins, mountain ranges, peninsulas, island groups, deserts, and terrain regions.
- When a place is followed by its containing region, state, or country as disambiguation context ("Puri, Odisha", "Lisboa, Portugal", "Goa, India", "Manchester in US"), extract ONLY the specific place. Do not return the container as a separate place.
- When a query names two or more distinct anchors joined by words like "and", "both", "between", or mixes an admin area with a physical feature as separate anchors, extract every anchor in the order they appear.
- Do not infer or expand category nouns like "regions", "districts", "counties", "rivers", or "mountains" when they refer to a type rather than a specific named place ("regions of India" -> extract "India" only).
- Only extract places explicitly mentioned.
- No duplicate place names.

EXAMPLES:
Query: "Puri, Odisha"
-> {"places": [{"place": "Puri"}]}

Query: "Lisboa, Portugal"
-> {"places": [{"place": "Lisboa"}]}

Query: "Goa, India"
-> {"places": [{"place": "Goa"}]}

Query: "Manchester in US"
-> {"places": [{"place": "Manchester"}]}

Query: "Springfield, Illinois"
-> {"places": [{"place": "Springfield"}]}

Query: "coastal districts of Brazil"
-> {"places": [{"place": "Brazil"}]}

Query: "northern half of India"
-> {"places": [{"place": "India"}]}

Query: "what's within 50 km of Paris?"
-> {"places": [{"place": "Paris"}]}

Query: "countries the Nile crosses"
-> {"places": [{"place": "Nile"}]}

Query: "which countries touch the Gulf of Maine"
-> {"places": [{"place": "Gulf of Maine"}]}

Query: "10 km buffer around Odisha"
-> {"places": [{"place": "Odisha"}]}

Query: "part of Ecuador in the Amazon basin"
-> {"places": [{"place": "Ecuador"}, {"place": "Amazon basin"}]}

Query: "Amazon basin inside Ecuador"
-> {"places": [{"place": "Amazon basin"}, {"place": "Ecuador"}]}

Query: "the part of Chad in Lake Chad"
-> {"places": [{"place": "Chad"}, {"place": "Lake Chad"}]}

Query: "which regions border both France and Germany?"
-> {"places": [{"place": "France"}, {"place": "Germany"}]}

Query: "merge Nairobi and Mombasa"
-> {"places": [{"place": "Nairobi"}, {"place": "Mombasa"}]}"""


def _candidate_to_place(c: Dict) -> Optional[Dict]:
    """Convert a selected candidate to a minimal Place dict for PlacesResult."""
    name = c.get("name", "").strip()
    if not name:
        return None

    return {"place": name}


def sample_to_place_pair(sample: Dict[str, Any]) -> Optional[Dict]:
    """Convert a raw sample to a conversational prompt-completion pair for place extraction.

    Uses selected_candidates to determine the correct PlacesResult output.
    Skips samples where no valid places can be derived.
    """
    selected_ids = sample.get("target", {}).get("selected_candidates", [])
    if not selected_ids:
        return None

    id_to_candidate = {c["candidate_id"]: c for c in sample.get("candidates", [])}
    places = []
    seen_names: set = set()

    for cid in selected_ids:
        c = id_to_candidate.get(cid)
        if not c:
            continue
        place = _candidate_to_place(c)
        if place and place["place"].lower() not in seen_names:
            places.append(place)
            seen_names.add(place["place"].lower())

    if not places:
        return None

    completion_json = json.dumps({"places": places}, ensure_ascii=False)

    return {
        "messages": [
            {"role": "system",    "content": _PLACE_SYSTEM},
            {"role": "user",      "content": sample["question"]},
            {"role": "assistant", "content": completion_json},
        ],
        "metadata": sample.get("metadata", {}),
    }


# ---------------------------------------------------------------------------
# I/O helpers
# ---------------------------------------------------------------------------

def save_jsonl(records: List[Dict], path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w") as f:
        for r in records:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")


def split_stats(samples: List[Dict]) -> Dict[str, int]:
    counts: Dict[str, int] = defaultdict(int)
    for s in samples:
        counts[s.get("metadata", {}).get("task_family", "unknown")] += 1
    return dict(sorted(counts.items()))


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main(config_path: Optional[Path] = None) -> None:
    script_dir = Path(__file__).parent
    dataset_dir = script_dir.parent
    output_dir = dataset_dir / "output"

    run_name = load_run_name(config_path or dataset_dir / "config.yaml")

    validated_file = output_dir / "dataset_validated.jsonl"
    if not validated_file.exists():
        print(f"Error: {validated_file} not found. Run validate first.")
        sys.exit(1)

    run_dir = output_dir / "runs" / run_name
    sql_dir = run_dir / "sql"
    places_dir = run_dir / "places"

    print(f"Run name   : {run_name}")
    print(f"Output dir : {run_dir}")

    # Load
    print("\nLoading validated samples...")
    samples = load_samples(validated_file)
    samples = [_shuffle_candidates_for_export(s) for s in samples]
    print(f"  {len(samples):,} samples loaded")
    print("  Candidate row order shuffled deterministically for export")

    # Split once, reuse for both tasks
    print("\nSplitting 80 / 10 / 10 (stratified by task family)...")
    train_raw, val_raw, test_raw = stratified_split(samples)
    print(f"  train={len(train_raw):,}  val={len(val_raw):,}  test={len(test_raw):,}")

    # --- SQL generation ---
    print("\nBuilding SQL generation splits...")
    sql_stats: Dict = {}
    for split_name, raw in [("train", train_raw), ("val", val_raw), ("test", test_raw)]:
        pairs = [p for s in raw if (p := sample_to_sql_pair(s)) is not None]
        save_jsonl(pairs, sql_dir / f"{split_name}.jsonl")
        sql_stats[split_name] = {"total": len(pairs), "by_family": split_stats(pairs)}
        print(f"  sql/{split_name}.jsonl  — {len(pairs):,} pairs")

    # --- Place extraction ---
    print("\nBuilding place extraction splits...")
    place_stats: Dict = {}
    for split_name, raw in [("train", train_raw), ("val", val_raw), ("test", test_raw)]:
        pairs = [p for s in raw if (p := sample_to_place_pair(s)) is not None]
        save_jsonl(pairs, places_dir / f"{split_name}.jsonl")
        place_stats[split_name] = {"total": len(pairs), "by_family": split_stats(pairs)}
        print(f"  places/{split_name}.jsonl  — {len(pairs):,} pairs")

    # --- Stats ---
    stats = {
        "run_name": run_name,
        "total_samples": len(samples),
        "sql_generation": sql_stats,
        "place_extraction": place_stats,
    }
    stats_path = run_dir / "stats.json"
    with open(stats_path, "w") as f:
        json.dump(stats, f, indent=2)

    print(f"\nStats written to {stats_path}")
    print("\nDone. Training-ready files:")
    print(f"  SQL generation  : {sql_dir}/{{train,val,test}}.jsonl")
    print(f"  Place extraction: {places_dir}/{{train,val,test}}.jsonl")


if __name__ == "__main__":
    main()