#!/usr/bin/env python3
"""Offline synthetic data generation (T1.5).

Produces JSONL datasets by composing :mod:`tsmm.data` primitives in parallel:

* ``data/align.jsonl``  — alignment-style samples (default: many)
* ``data/sft.jsonl``    — Evol-Instruct rephrased (subset)
* ``data/eval_synth.jsonl`` — held-out eval set (fixed seed)

Usage::

    python scripts/gen_synthetic.py --n 1000 --out data/align.jsonl --seed 0 --workers 8

Each line is one JSON sample with keys: series, attributes, timestamps,
events, instruction, answer, labels, category, segments.
"""

from __future__ import annotations

import argparse
import json
import os
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path

from tqdm import tqdm

from tsmm.data.gen_pipeline import gen_one_sample, write_jsonl
from tsmm.data.instruct import evolve_instruction


def _one(args):
    """Worker: generate a single sample (top-level for pickling)."""
    T, C, seed, anomaly_types, missing_rate, event_prob, do_evolve = args
    s = gen_one_sample(
        T=T, C=C, seed=seed,
        anomaly_types=anomaly_types,
        missing_rate=missing_rate,
        event_prob=event_prob,
    )
    if do_evolve:
        s["instruction"] = evolve_instruction(s["instruction"], seed=seed)
    return s


def run(
    n: int,
    out: str,
    *,
    seed: int = 0,
    workers: int = max(1, (os.cpu_count() or 2) - 1),
    T: int = 512,
    C: int = 5,
    anomaly_types=("spike", "level_shift", "variance_change", "drift"),
    missing_rate: float = 0.05,
    event_prob: float = 0.5,
    do_evolve: bool = False,
) -> int:
    Path(out).parent.mkdir(parents=True, exist_ok=True)
    work = [
        (T, C, seed + i, anomaly_types, missing_rate, event_prob, do_evolve)
        for i in range(n)
    ]
    samples = [None] * n
    with ProcessPoolExecutor(max_workers=workers) as ex:
        futures = {ex.submit(_one, w): i for i, w in enumerate(work)}
        for fut in tqdm(as_completed(futures), total=n, desc=os.path.basename(out)):
            i = futures[fut]
            samples[i] = fut.result()
    return write_jsonl(samples, out)


def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(description="Generate synthetic TS-QA JSONL data.")
    p.add_argument("--n", type=int, default=1000, help="number of samples")
    p.add_argument("--out", type=str, default="data/align.jsonl", help="output JSONL path")
    p.add_argument("--seed", type=int, default=0, help="base RNG seed")
    p.add_argument("--workers", type=int, default=max(1, (os.cpu_count() or 2) - 1))
    p.add_argument("--T", type=int, default=512, help="series length")
    p.add_argument("--C", type=int, default=5, help="number of channels")
    p.add_argument("--missing-rate", type=float, default=0.05)
    p.add_argument("--event-prob", type=float, default=0.5)
    p.add_argument("--evolve", action="store_true", help="apply Evol-Instruct rephrasing (for sft set)")
    args = p.parse_args(argv)

    n = run(
        n=args.n,
        out=args.out,
        seed=args.seed,
        workers=args.workers,
        T=args.T,
        C=args.C,
        missing_rate=args.missing_rate,
        event_prob=args.event_prob,
        do_evolve=args.evolve,
    )
    print(f"wrote {n} samples -> {args.out}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())