"""The canonical synthetic corpus: recipe, shard builder and verifier.

TinyCast's pretraining mix is GIFT-Eval-Pretrain and Chronos KernelSynth, taken
from their publishers, plus four synthetic shards. This module is the record of
how those four shards were made: four shards of 62,500 series at length 4,096,
families mixed 70/15/15 over the generators in :mod:`tinycast.synth`.

Everything that changes the generated bytes is a default in
:data:`CANONICAL_RECIPE` rather than an argument, so ``build_shard()`` with no
arguments reproduces published shard 0 and there is nothing to tune. Treat the
values in that dictionary as versioned data: changing one produces a different
corpus, not a differently configured build of the same one.

WHAT REPRODUCIBILITY MEANS HERE. The shards are published as files, so the
corpus can be downloaded and read without regenerating anything. Regeneration
reproduces the published float16 payload exactly on the stack it was generated
on: 0 of 327,680 values differ across 4 shards by 20 rows on an RTX 3090 with
torch 2.10.0+cu128 and driver 580.159.03. Reproduction on other GPUs, CUDA
versions or torch builds is untested, and the GP family runs through cuSOLVER,
so it is not safe to assume it carries over. :func:`verify_shard` is how you
find out for a given stack: it regenerates a prefix and reports how many
float16 values differ.

Verify against ``series.f16``, which is the payload, and not against
``series_mean.f32``. That sidecar is a float32 sum with heavy cancellation on
near-zero-mean rows, so its last bit moves with the numpy build while the f16
payload it feeds is unaffected. ``series_stdev.f32`` does not cancel.

CUDA IS REQUIRED, and is enforced rather than detected. The GP family samples
by dense Cholesky, and off CUDA that factorization and the normal draws feeding
it come from different generators, so the same seed gives different series: the
CPU path moves up to 182 of the 4,096 float32 values in a row. The published
shards came from the CUDA branch. A build with no GPU therefore fails here
instead of falling back and quietly producing a different corpus.
"""
from __future__ import annotations

import io
import json
import math
import os
import time
from pathlib import Path
from typing import Iterator, Optional, Tuple, Union

import numpy as np

from .synth import generate_gp, generate_spikes, generate_tsi

_F16_BYTES = 2

# ---------------------------------------------------------------------------
# The canonical recipe.
#
# Every entry is a default. A reviewer runs the build with no flags and gets the
# published corpus.
# ---------------------------------------------------------------------------
CANONICAL_RECIPE = {
    "version": "tinycast-synth-4096/v1",
    "n_shards": 4,
    "series_per_shard": 62_500,
    "length": 4096,
    "mix": (0.70, 0.15, 0.15),          # gp, spikes, tsi
    "chunk": 2048,                       # series generated per call
    "device": "cuda",
    # Part of the data, not a memory knob. cuSOLVER factorizes a batch of one on
    # a different code path from a real batch, and these covariances are close
    # to singular (periodic and linear kernels with only 1e-6 jitter), so the
    # factors differ by around 1e-8 and the published rows do not reproduce at
    # batch=1. Batches of 2, 4, 8 and 32 all agree with each other and with the
    # published data. Lowering this to fit a smaller GPU changes the corpus.
    "gp_batch": 32,
    "scale_factor_range": (0.05, 6.0),   # log-uniform, per series
    "verified_stack": "RTX 3090 / torch 2.10.0+cu128 / driver 580.159.03",
    # Per-shard seed and family salts. The seed is 50,000 x shard, confirmed by
    # reproducing all 62,500 scale factors of each shard bit for bit.
    #
    # The salts cannot be computed, which is why they are carried as data. They
    # were a per-family hash of the family name taken under CPython's randomized
    # string hashing, so every build process drew its own set and no two shards
    # share a triple. Replacing them with a stable hash would make future builds
    # reproducible while making the published shards unreproducible. They were
    # recovered by brute force against the published bytes (997 candidates per
    # family) and frozen here.
    "shards": {
        0: {"seed": 0,       "salt": {"gp": 724, "spikes": 351, "tsi": 900}},
        1: {"seed": 50_000,  "salt": {"gp":  23, "spikes": 431, "tsi": 840}},
        2: {"seed": 100_000, "salt": {"gp": 762, "spikes": 505, "tsi": 641}},
        3: {"seed": 150_000, "salt": {"gp": 691, "spikes": 491, "tsi": 745}},
    },
}

FAMILIES = ("gp", "spikes", "tsi")

# Files a shard directory holds. The first six are the payload and its indices;
# the last two attribute each row to a family.
CACHE_FILES = (
    "series.f16", "offsets.npy", "lengths.npy", "scale_factors.f32",
    "series_mean.f32", "series_stdev.f32",
)
FAMILY_FILES = ("dataset_id.u16", "dataset_names.json")

DEFAULT_OUT_DIR = "synth4096_{shard}"

# Rows compared by verify_shard. Twenty is what the published fidelity statement
# was measured over, and at 4,096 samples a row it is already 81,920 values.
DEFAULT_VERIFY_ROWS = 20
DEFAULT_VERIFY_SCALE_FACTORS = 4096


# ---------------------------------------------------------------------------
# Recipe lookups
# ---------------------------------------------------------------------------

def _shard_record(shard: int) -> dict:
    try:
        return CANONICAL_RECIPE["shards"][int(shard)]
    except KeyError:
        raise KeyError(
            f"shard {shard} is not in the published corpus, which covers shards "
            f"{sorted(CANONICAL_RECIPE['shards'])}."
        ) from None


def shard_seed(shard: int) -> int:
    """Base seed for a shard. The scale-factor stream is drawn from it."""
    return int(_shard_record(shard)["seed"])


def family_salt(family: str, shard: int) -> int:
    """Per-shard, per-family seed offset, looked up rather than computed.

    See :data:`CANONICAL_RECIPE` for why these are data. A missing entry means
    the salt has not been recovered for that shard, and generation refuses
    rather than substituting a computed one: a computed salt would produce a
    different corpus while appearing to succeed.
    """
    salt = _shard_record(shard)["salt"].get(family)
    if salt is None:
        raise RuntimeError(
            f"the salt for family {family!r} on shard {shard} is not recorded, so "
            f"that shard cannot be reproduced. It is a single integer in [0, 997) "
            f"and is recoverable by brute force against the published bytes. "
            f"Refusing to substitute a computed value."
        )
    return int(salt)


def plan_counts(
    n_series: Optional[int] = None, mix: Optional[tuple] = None,
) -> Tuple[int, int, int]:
    """Series per family, in emission order (gp, spikes, tsi).

    This is also how a published shard's family boundaries are read. The
    per-row family sidecars of the published shards were overwritten by a
    post-processing step that collapsed them to a single label with all ids
    zero, so the attribution cannot be recovered from the files themselves; the
    payload, indices and scale factors were not touched. A shard that ends up
    one row short lost a GP row to a covariance that failed to factorize (see
    :func:`iter_shard_series`), which shifts the boundaries down by that many.
    """
    n_series = CANONICAL_RECIPE["series_per_shard"] if n_series is None else int(n_series)
    mix = CANONICAL_RECIPE["mix"] if mix is None else tuple(mix)
    n_gp = int(round(n_series * mix[0]))
    n_spikes = int(round(n_series * mix[1]))
    return n_gp, n_spikes, n_series - n_gp - n_spikes


def _sample_scale_factor(rng: np.random.Generator) -> float:
    """Log-uniform scale factor, one draw per emitted series.

    Synthetic series carry no sampling frequency, so each gets a scale factor
    drawn from the range that the real corpora's frequency-derived factors span.
    """
    low, high = CANONICAL_RECIPE["scale_factor_range"]
    return float(np.exp(rng.uniform(math.log(low), math.log(high))))


def _require_cuda(device: Optional[str]) -> str:
    """Resolve and check the device, refusing anything but the canonical one."""
    canonical = CANONICAL_RECIPE["device"]
    device = canonical if device is None else str(device)
    if device != canonical:
        raise ValueError(
            f"device={device!r} is not the device this corpus was generated on "
            f"({canonical!r}). The GP family draws from a different generator on "
            f"each branch, so the same seed gives different series and overriding "
            f"the device does not reproduce the published shards."
        )
    try:
        import torch
    except ImportError as exc:
        raise RuntimeError(
            "torch is required: this corpus was generated on the CUDA branch."
        ) from exc
    if not torch.cuda.is_available():
        raise RuntimeError(
            "no CUDA device is available. This corpus was generated on the CUDA "
            "branch, and a CPU run would produce different data rather than "
            "reproducing it, so this fails instead of falling back."
        )
    return device


# ---------------------------------------------------------------------------
# Generation
# ---------------------------------------------------------------------------

def iter_shard_series(
    shard: int = 0,
    *,
    max_series: Optional[int] = None,
    device: Optional[str] = None,
) -> Iterator[Tuple[np.ndarray, float, str]]:
    """Yield ``(series, scale_factor, family_label)`` for one published shard.

    Families are emitted in order, GP first, each generated in chunks of
    ``CANONICAL_RECIPE["chunk"]`` series under a seed of
    ``shard_seed + 1000 * chunk_index + family_salt``. GP rows whose covariance
    failed to factorize come back non-finite and are dropped here, which is the
    only way a shard ends up holding fewer rows than planned.

    ``max_series`` truncates the emitted stream and leaves the plan alone, so
    what it yields is the published prefix. Generation still runs in canonical
    chunks, so the first chunk is produced in full however early the truncation
    falls; :func:`verify_shard` is the cheap way to check a stack.
    ``device`` exists only so that asking for a non-canonical one fails loudly.

    Validation runs before the first row is generated. The refusals above are
    worth nothing if they wait for the caller to start iterating, which is what
    happens when a generator function does its own argument checking.
    """
    device = _require_cuda(device)
    seed = shard_seed(shard)
    salts = {f: family_salt(f, shard) for f in FAMILIES}
    if max_series is not None and int(max_series) < 0:
        raise ValueError("max_series must be non-negative")
    return _iter_shard_series(shard, seed, salts, max_series, device)


def _iter_shard_series(
    shard: int,
    seed: int,
    salts: dict,
    max_series: Optional[int],
    device: str,
) -> Iterator[Tuple[np.ndarray, float, str]]:
    """Generator body for :func:`iter_shard_series`, validation already done."""
    length = CANONICAL_RECIPE["length"]
    chunk = CANONICAL_RECIPE["chunk"]
    gp_batch = CANONICAL_RECIPE["gp_batch"]
    counts = dict(zip(FAMILIES, plan_counts()))
    generators = {
        "gp": lambda n, s: generate_gp(n, length, seed=s, device=device,
                                       batch=gp_batch),
        "spikes": lambda n, s: generate_spikes(n, length, seed=s),
        "tsi": lambda n, s: generate_tsi(n, length, seed=s),
    }

    rng = np.random.default_rng(seed)
    emitted_total = 0
    dropped = 0
    for family in FAMILIES:
        total = counts[family]
        label = f"{family}_{length}"
        generated = 0
        chunk_index = 0
        while generated < total:
            cur = min(chunk, total - generated)
            block = generators[family](cur, seed + 1000 * chunk_index + salts[family])
            chunk_index += 1
            generated += cur
            for row in block:
                if not np.isfinite(row).all():
                    dropped += 1
                    continue
                yield row, _sample_scale_factor(rng), label
                emitted_total += 1
                if max_series is not None and emitted_total >= int(max_series):
                    return
        print(f"  {family}: {total} generated, {dropped} dropped so far", flush=True)


# ---------------------------------------------------------------------------
# Writing a shard
# ---------------------------------------------------------------------------

def normalize_for_f16(series: np.ndarray) -> Tuple[np.ndarray, float, float]:
    """Per-series z-normalization ahead of the float16 cast.

    float16 stores the payload at half the size of float32 but caps magnitudes
    at 65,504, which raw series exceed. Storing a z-normalized payload with the
    mean and standard deviation alongside it keeps any realistic magnitude in
    range, and the read path is ``x = payload * stdev + mean`` with no
    branching: a constant series is stored with ``stdev = 1`` so its payload is
    identically zero and denormalizes exactly.
    """
    series = np.asarray(series, dtype=np.float32)
    clean = np.where(np.isfinite(series), series, np.nan)
    finite = clean[~np.isnan(clean)]
    if finite.size == 0:
        return clean, 0.0, 1.0
    mu = float(finite.mean())
    sd = float(finite.std())
    sd_safe = sd if sd > 0.0 else 1.0
    return ((clean - mu) / sd_safe).astype(np.float32, copy=False), mu, sd_safe


def build_shard(
    out_dir: Union[str, Path, None] = None,
    shard: int = 0,
    *,
    max_series: Optional[int] = None,
    overwrite: bool = False,
) -> int:
    """Generate one published shard and write it as a cache directory.

    Called with no arguments this reproduces shard 0 into ``synth4096_0``. The
    written files are the six payload and index files, the two family sidecars,
    and ``recipe.json``, which records the recipe and the stack that produced
    the directory. Returns the number of series written.

    The payload is streamed to a temporary file and renamed on success, so an
    interrupted build leaves no directory that looks complete.
    """
    _require_cuda(None)
    for family in FAMILIES:                # refuse an unreproducible shard first
        family_salt(family, shard)
    out_dir = Path(DEFAULT_OUT_DIR.format(shard=shard) if out_dir is None else out_dir)
    if all((out_dir / f).exists() for f in CACHE_FILES) and not overwrite:
        raise RuntimeError(
            f"{out_dir} already holds a shard. Pass overwrite=True to rebuild it, "
            f"or build into a new directory."
        )
    out_dir.mkdir(parents=True, exist_ok=True)

    length = CANONICAL_RECIPE["length"]
    n_gp, n_spikes, n_tsi = plan_counts()
    print(f"building shard {shard} into {out_dir}: gp={n_gp} spikes={n_spikes} "
          f"tsi={n_tsi} length={length} device={CANONICAL_RECIPE['device']}",
          flush=True)

    offsets: list[int] = []
    lengths: list[int] = []
    scale_factors: list[float] = []
    means: list[float] = []
    stdevs: list[float] = []
    labels: list[str] = []
    label_ids: dict[str, int] = {}
    ids: list[int] = []

    tmp_payload = out_dir / "series.f16.tmp"
    cursor = 0
    n_written = 0
    t0 = time.time()
    try:
        with open(tmp_payload, "wb") as fp:
            for row, sf, label in iter_shard_series(shard, max_series=max_series):
                normalized, mu, sd = normalize_for_f16(row)
                payload = normalized.astype(np.float16, copy=False)
                fp.write(payload.tobytes())
                offsets.append(cursor)
                lengths.append(int(payload.size))
                scale_factors.append(float(sf))
                means.append(mu)
                stdevs.append(sd)
                if label not in label_ids:
                    label_ids[label] = len(labels)
                    labels.append(label)
                ids.append(label_ids[label])
                cursor += int(payload.size) * _F16_BYTES
                n_written += 1
                if n_written % 10_000 == 0:
                    rate = n_written / max(time.time() - t0, 1e-6)
                    print(f"  {n_written:,} written, {rate:.0f} series/s, "
                          f"{cursor / 1e9:.2f} GB", flush=True)
            fp.flush()
            os.fsync(fp.fileno())
    except BaseException:
        tmp_payload.unlink(missing_ok=True)
        raise

    if n_written == 0:
        tmp_payload.unlink(missing_ok=True)
        raise RuntimeError("no series were generated, so nothing was written")

    os.replace(tmp_payload, out_dir / "series.f16")
    np.save(out_dir / "offsets.npy", np.asarray(offsets, dtype=np.int64))
    np.save(out_dir / "lengths.npy", np.asarray(lengths, dtype=np.int32))
    _write_bytes(out_dir / "scale_factors.f32",
                 np.asarray(scale_factors, dtype=np.float32).tobytes())
    _write_bytes(out_dir / "series_mean.f32",
                 np.asarray(means, dtype=np.float32).tobytes())
    _write_bytes(out_dir / "series_stdev.f32",
                 np.asarray(stdevs, dtype=np.float32).tobytes())
    _write_bytes(out_dir / "dataset_id.u16",
                 np.asarray(ids, dtype=np.uint16).tobytes())
    _write_bytes(out_dir / "dataset_names.json",
                 json.dumps({"version": 2, "datasets": labels}).encode("utf-8"))
    _write_bytes(out_dir / "recipe.json",
                 json.dumps(_provenance(shard, n_written), indent=2).encode("utf-8"))

    print(f"wrote {n_written:,} series in {time.time() - t0:.0f}s "
          f"({cursor / 1e9:.2f} GB) to {out_dir}", flush=True)
    return n_written


def _write_bytes(path: Path, data: bytes) -> None:
    tmp = path.with_suffix(path.suffix + ".tmp")
    with open(tmp, "wb") as fp:
        fp.write(data)
        fp.flush()
        os.fsync(fp.fileno())
    os.replace(tmp, path)


def _provenance(shard: int, n_written: int) -> dict:
    record = {
        "version": CANONICAL_RECIPE["version"],
        "shard": int(shard),
        "seed": shard_seed(shard),
        "salt": dict(_shard_record(shard)["salt"]),
        "series_written": int(n_written),
        "series_planned": CANONICAL_RECIPE["series_per_shard"],
        "length": CANONICAL_RECIPE["length"],
        "mix": list(CANONICAL_RECIPE["mix"]),
        "chunk": CANONICAL_RECIPE["chunk"],
        "gp_batch": CANONICAL_RECIPE["gp_batch"],
        "device": CANONICAL_RECIPE["device"],
        "verified_stack": CANONICAL_RECIPE["verified_stack"],
    }
    try:
        import torch
        record["torch"] = torch.__version__
        if torch.cuda.is_available():
            record["gpu"] = torch.cuda.get_device_name(0)
    except Exception:                      # provenance is best effort
        pass
    record["numpy"] = np.__version__
    return record


# ---------------------------------------------------------------------------
# Verifying a shard against the published one
# ---------------------------------------------------------------------------

def _resolve_source(path_or_url: Union[str, Path]) -> Union[Path, str]:
    """A local directory, or an https prefix ending in a slash.

    ``hf://<owner>/<repo>[/<subdir>]`` is rewritten to the dataset's resolve
    URL, which serves byte ranges. That matters: checking twenty rows pulls
    about 160 KB rather than the half gigabyte a shard payload occupies.
    """
    text = str(path_or_url)
    if text.startswith("hf://"):
        parts = [p for p in text[len("hf://"):].split("/") if p]
        if len(parts) < 2:
            raise ValueError("hf:// source must be hf://<owner>/<repo>[/<subdir>]")
        repo = "/".join(parts[:2])
        subdir = "/".join(parts[2:])
        url = f"https://huggingface.co/datasets/{repo}/resolve/main/"
        return url + (subdir + "/" if subdir else "")
    if text.startswith("http://") or text.startswith("https://"):
        return text if text.endswith("/") else text + "/"
    return Path(text)


def _read_file(source: Union[Path, str], name: str) -> bytes:
    if isinstance(source, Path):
        return (source / name).read_bytes()
    import urllib.request
    with urllib.request.urlopen(source + name) as response:
        return response.read()


def _read_range(source: Union[Path, str], name: str, start: int, count: int) -> bytes:
    if isinstance(source, Path):
        with open(source / name, "rb") as fp:
            fp.seek(start)
            return fp.read(count)
    import urllib.request
    request = urllib.request.Request(
        source + name, headers={"Range": f"bytes={start}-{start + count - 1}"})
    with urllib.request.urlopen(request) as response:
        data = response.read()
    if len(data) != count:                 # server ignored the range header
        data = data[start:start + count]
    return data


def verify_shard(
    path_or_url: Union[str, Path],
    shard: int = 0,
    *,
    n_rows: int = DEFAULT_VERIFY_ROWS,
    n_scale_factors: int = DEFAULT_VERIFY_SCALE_FACTORS,
    verbose: bool = True,
) -> dict:
    """Regenerate a prefix of a published shard and report how far it agrees.

    ``path_or_url`` is a shard directory, an https prefix, or
    ``hf://<owner>/<repo>[/<subdir>]``. The comparison is against
    ``series.f16``, the payload, cast the same way the builder casts it. The
    returned report gives ``values_differing`` out of ``values_compared``, which
    is the number this module's fidelity claim is stated in.

    ``n_rows`` rows are regenerated from the GP family, which is what a shard
    opens with, rounded up to a whole number of GP batches so the batch shape
    matches the published run. It must not exceed one chunk. Set it to 0 to skip
    the payload check, which is the only part that needs a GPU; the scale-factor
    check runs anywhere and still exercises the seed scheme.

    A shard that holds fewer rows than planned lost GP rows to covariances that
    failed to factorize. Those rows are dropped here too, so the compared prefix
    stays aligned, but the scale-factor stream shifts by one draw per drop, and
    a mismatch after the drop index is expected rather than a fidelity failure.
    """
    source = _resolve_source(path_or_url)
    length = CANONICAL_RECIPE["length"]
    chunk = CANONICAL_RECIPE["chunk"]
    gp_batch = CANONICAL_RECIPE["gp_batch"]
    n_rows = int(n_rows)
    if n_rows > chunk:
        raise ValueError(
            f"n_rows={n_rows} exceeds one chunk ({chunk}); rows beyond the first "
            f"chunk are generated under a different seed and would need the whole "
            f"shard regenerated."
        )

    published_lengths = np.load(io.BytesIO(_read_file(source, "lengths.npy")))
    published_offsets = np.load(io.BytesIO(_read_file(source, "offsets.npy")))
    published_rows = int(published_lengths.size)
    planned_rows = CANONICAL_RECIPE["series_per_shard"]

    report = {
        "source": str(path_or_url),
        "shard": int(shard),
        "published_rows": published_rows,
        "planned_rows": planned_rows,
        "rows_short_of_plan": max(planned_rows - published_rows, 0),
        "rows_compared": 0,
        "values_compared": 0,
        "values_differing": None,
        "scale_factors_compared": 0,
        "scale_factors_differing": None,
        "payload_checked": False,
        "verified_stack": CANONICAL_RECIPE["verified_stack"],
    }

    if n_scale_factors:
        published_sf = np.frombuffer(
            _read_file(source, "scale_factors.f32"), dtype=np.float32)
        n_sf = min(int(n_scale_factors), published_sf.size)
        rng = np.random.default_rng(shard_seed(shard))
        ours = np.asarray([_sample_scale_factor(rng) for _ in range(n_sf)],
                          dtype=np.float32)
        differing = int(np.count_nonzero(ours != published_sf[:n_sf]))
        report["scale_factors_compared"] = n_sf
        report["scale_factors_differing"] = differing

    if n_rows > 0:
        _require_cuda(None)
        n_generate = int(math.ceil(n_rows / gp_batch) * gp_batch)
        block = generate_gp(
            n_generate, length,
            seed=shard_seed(shard) + family_salt("gp", shard),
            device=CANONICAL_RECIPE["device"], batch=gp_batch,
        )
        kept = [row for row in block if np.isfinite(row).all()]
        if len(kept) < n_rows:
            raise RuntimeError(
                f"regeneration produced only {len(kept)} usable rows of the "
                f"{n_rows} requested: {n_generate - len(kept)} covariances failed "
                f"to factorize. Investigate that before reading anything into a "
                f"comparison, since it is far above the published failure rate."
            )
        differing = 0
        compared = 0
        max_abs_diff = 0.0
        for i in range(n_rows):
            n = int(published_lengths[i])
            if n != length:
                raise RuntimeError(
                    f"published row {i} has length {n}, not {length}: this is not "
                    f"a shard of this corpus."
                )
            raw = _read_range(source, "series.f16",
                              int(published_offsets[i]), n * _F16_BYTES)
            theirs = np.frombuffer(raw, dtype=np.float16, count=n)
            normalized, _, _ = normalize_for_f16(kept[i])
            ours = normalized.astype(np.float16, copy=False)
            delta = ours.astype(np.float32) - theirs.astype(np.float32)
            differing += int(np.count_nonzero(delta))
            max_abs_diff = max(max_abs_diff, float(np.abs(delta).max()))
            compared += n
        report.update({
            "rows_compared": n_rows,
            "values_compared": compared,
            "values_differing": differing,
            "max_abs_diff": max_abs_diff,
            "payload_checked": True,
        })

    if verbose:
        print(f"shard {shard} at {path_or_url}")
        print(f"  published rows: {published_rows:,} of {planned_rows:,} planned")
        if report["scale_factors_compared"]:
            print(f"  scale factors: {report['scale_factors_differing']} of "
                  f"{report['scale_factors_compared']:,} differ")
        if report["payload_checked"]:
            print(f"  payload: {report['values_differing']} of "
                  f"{report['values_compared']:,} float16 values differ "
                  f"across {report['rows_compared']} rows "
                  f"(max abs difference {report['max_abs_diff']:g})")
        else:
            print("  payload: not checked (n_rows=0)")
        print(f"  the published fidelity statement was measured on "
              f"{CANONICAL_RECIPE['verified_stack']}")
    return report


# ---------------------------------------------------------------------------
# Command line
# ---------------------------------------------------------------------------

def main(argv: Optional[list] = None) -> None:
    import argparse

    parser = argparse.ArgumentParser(
        prog="python -m tinycast.corpus",
        description="Build or verify a shard of TinyCast's synthetic corpus.")
    sub = parser.add_subparsers(dest="command", required=True)

    build = sub.add_parser("build", help="generate a shard (needs a CUDA GPU)")
    build.add_argument("--shard", type=int, default=0)
    build.add_argument("--out", default=None,
                       help=f"output directory (default {DEFAULT_OUT_DIR})")
    build.add_argument("--max-series", type=int, default=None,
                       help="stop after this many series; the prefix is unchanged")
    build.add_argument("--overwrite", action="store_true")

    verify = sub.add_parser("verify", help="compare a published shard to a rebuild")
    verify.add_argument("source",
                        help="shard directory, https prefix, or hf://<owner>/<repo>")
    verify.add_argument("--shard", type=int, default=0)
    verify.add_argument("--rows", type=int, default=DEFAULT_VERIFY_ROWS,
                        help="rows to regenerate and compare; 0 skips the payload")
    verify.add_argument("--scale-factors", type=int,
                        default=DEFAULT_VERIFY_SCALE_FACTORS)

    args = parser.parse_args(argv)
    if args.command == "build":
        build_shard(args.out, args.shard, max_series=args.max_series,
                    overwrite=args.overwrite)
    else:
        verify_shard(args.source, args.shard, n_rows=args.rows,
                     n_scale_factors=args.scale_factors)


if __name__ == "__main__":
    main()