static USAGE: &str = r#"
Randomly samples CSV data.

It supports ten sampling methods:
* RESERVOIR: the default sampling method when NO INDEX is present and no sampling method
  is specified. Visits every CSV record exactly once, using MEMORY PROPORTIONAL to the
  sample size (k) - O(k).
  https://en.wikipedia.org/wiki/Reservoir_sampling

* INDEXED: the default sampling method when an INDEX is present and no sampling method
  is specified. Uses random I/O to sample efficiently, as it only visits records selected
  by random indexing, using MEMORY PROPORTIONAL to the sample size (k) - O(k).
  https://en.wikipedia.org/wiki/Random_access

* BERNOULLI: the sampling method when the --bernoulli option is specified.
  Each record has an independent probability p of being selected, where p is
  specified by the <sample-size> argument. For example, if p=0.1, then each record
  has a 10% chance of being selected, regardless of the other records. The final
  sample size is random and follows a binomial distribution. Uses CONSTANT MEMORY - O(1).
  When sampling from a remote URL, processes the file in chunks without downloading it
  entirely, making it especially efficient for sampling large remote files.
  https://en.wikipedia.org/wiki/Bernoulli_sampling

* SYSTEMATIC: the sampling method when the --systematic option is specified.
  Selects every nth record from the input, where n is the integer part of <sample-size>
  and the fraction part is the percentage of the population to sample.
  For example, if <sample-size> is 10.5, it will select every 10th record and 50% of the
  population. If <sample-size> is a whole number (no fractional part), it will select
  every nth record for the whole population. Uses CONSTANT memory - O(1). The starting
  point can be specified as "random" or "first". Useful for time series data or when you
  want evenly spaced samples.
  https://en.wikipedia.org/wiki/Systematic_sampling

* STRATIFIED: the sampling method when the --stratified option is specified.
  Stratifies the population by the specified column and then samples from each stratum.
  Particularly useful when a population has distinct subgroups (strata) that are
  heterogeneous within but homogeneous between in terms of the variable of interest. 
  For example, if you want to sample 1,000 records from a population of 100,000 across the US,
  you can stratify the population by US state and then sample 20 records from each stratum.
  This will ensure that you have a representative sample from each of the 50 states.
  The sample size must be a whole number. Uses MEMORY PROPORTIONAL to the
  number of strata (s) and samples per stratum (k) as specified by <sample-size> - O(s*k).
  https://en.wikipedia.org/wiki/Stratified_sampling

* WEIGHTED: the sampling method when the --weighted option is specified.
  Samples records with probabilities proportional to values in a specified weight column.
  Records with higher weights are more likely to be selected. For example, if you have
  sales data and want to sample transactions weighted by revenue, high-value transactions
  will have a higher chance of being included. Non-numeric weights are treated as zero.
  The weights are automatically normalized using the maximum weight in the dataset.
  Specify the desired sample size with <sample-size>. Uses MEMORY PROPORTIONAL to the
  sample size (k) - O(k).
  "Weighted random sampling with a reservoir" https://doi.org/10.1016/j.ipl.2005.11.003

* VAROPT: the sampling method when the --varopt option is specified.
  Variance-bounded weighted reservoir sampling using the A-ExpJ keying scheme of
  Efraimidis and Spirakis (2006). For each record, computes a key u^(1/w) and
  retains the <sample-size> items with the largest keys. Unlike the --weighted
  method, it does NOT require a stats cache, runs in a single pass, and supports
  merge across partitions through the --sketch-out and --sketch-in options.
  Suitable for heavy-tailed weight distributions where bounded-variance
  estimators are needed. Uses MEMORY PROPORTIONAL to the sample size (k) - O(k).
  This is a native Rust implementation written from the original paper; the
  analogous VarOpt sketches in the Apache DataSketches library use the same
  family of algorithms but are NOT used here.
  Algorithm: "Weighted random sampling with a reservoir"
  doi 10.1016/j.ipl.2005.11.003

* MERGEABLE-RESERVOIR: the sampling method when the --mergeable-reservoir flag is set.
  Uniform reservoir sample using Vitter's Algorithm R. Same statistical
  distribution as the default RESERVOIR method, but the sampler state is
  mergeable: a sketch written by one run can be combined with sketches from
  other runs via the --sketch-out and --sketch-in options, producing a uniform
  sample of the combined stream WITHOUT re-reading the input files. Useful
  for sharded or incremental sampling pipelines. Uses MEMORY PROPORTIONAL to
  the sample size (k) - O(k). Native Rust implementation; the analogous
  ReservoirItemsSketch in the Apache DataSketches library implements the same
  algorithm but is NOT used here.
  See en.wikipedia.org/wiki/Reservoir_sampling

* CLUSTER: the sampling method when the --cluster option is specified.
  Samples entire groups of records together based on a cluster identifier column.
  The number of clusters is specified by the <sample-size> argument.
  Useful when records are naturally grouped (e.g., by household, neighborhood, etc.).
  For example, if you have records grouped by neighborhood and specify a sample size of 10,
  it will randomly select 10 neighborhoods and include ALL records from those neighborhoods
  in the output. This ensures that natural groupings in the data are preserved.
  Uses MEMORY PROPORTIONAL to the number of clusters (c) - O(c).
  https://en.wikipedia.org/wiki/Cluster_sampling

* TIMESERIES: the sampling method when the --timeseries option is specified.
  Samples records based on time intervals from a time-series dataset. Groups records by
  time windows (e.g., hourly, daily, weekly) and selects one record per interval.
  Supports adaptive sampling (e.g., prefer business hours or weekends) and aggregation
  (e.g., mean, sum, min, max) within each interval. The starting point can be "first"
  (earliest), "last" (most recent), or "random". Particularly useful for time-series data
  where simple row-based sampling would always return the same records due to sorting.
  Uses MEMORY PROPORTIONAL to the number of records - O(n).

Supports sampling from CSVs on remote URLs. Note that the entire file is downloaded first
to a temporary file before sampling begins for all sampling methods except Bernoulli, which
streams the file as it samples it, stopping when the desired sample size is reached or the
end of the file is reached.

Sampling from stdin is also supported for all sampling methods, copying stdin to a in-memory
buffer first before sampling begins.

If a stats cache is available, it will be used to do extra checks on systematic,
weighted and cluster sampling, and to speed up sampling in general.

This command is intended to provide a means to sample from a CSV data set that
is too big to fit into memory (for example, for use with commands like
'qsv stats' with the '--everything' option). 

Examples:

  # Take a sample of 1000 records from data.csv using RESERVOIR or INDEXED sampling
  # depending on whether an INDEX is present. 
  qsv sample 1000 data.csv

  # Take a sample of approximately 10% of the records from data.csv using RESERVOIR
  # or INDEXED sampling depending on whether an INDEX is present.
  qsv sample 0.1 data.csv

  # Take a sample using BERNOULLI sampling where each record has a 5% chance of being selected
  qsv sample --bernoulli 0.05 data.csv

  # Take a sample using SYSTEMATIC sampling where every 10th record is selected
  # and approximately 50% of the population is sampled, starting from a random point.
  qsv sample --systematic random 10.5 data.csv

  # Take a sample using STRATIFIED sampling where 20 records are sampled from each
  # stratum defined by the 'State' column.
  qsv sample --stratified State 20 data.csv

  # Take a sample using WEIGHTED sampling where records are sampled with probabilities
  # proportional to the 'Revenue' column, for a total sample size of 1000 records.
  qsv sample --weighted Revenue 1000 data.csv

  # Take a sample using CLUSTER sampling where 10 clusters defined by the
  # 'Neighborhood' column are randomly selected and all records from those clusters
  # are included in the sample.
  qsv sample --cluster Neighborhood 10 data.csv

  # Take a sample using VAROPT (A-ExpJ weighted reservoir) sampling, weighted by
  # the 'Revenue' column, for a sample size of 1000 records. Unlike --weighted,
  # this does NOT require a stats cache.
  qsv sample --varopt Revenue 1000 data.csv

  # Sample two shards and merge their sketches into a single uniform sample
  # without re-reading the inputs.
  qsv sample --mergeable-reservoir --sketch-out a.sk 1000 shard_a.csv
  qsv sample --mergeable-reservoir --sketch-out b.sk 1000 shard_b.csv
  qsv sample --sketch-in a.sk,b.sk 1000 -o merged.csv

For more examples, see https://github.com/dathere/qsv/blob/master/tests/test_sample.rs.
See also https://github.com/dathere/qsv/wiki/Selection-and-Inspection#sample

Usage:
    qsv sample [options] <sample-size> [<input>]
    qsv sample --help

sample arguments:
    <input>                The CSV file to sample. This can be a local file,
                           stdin, or a URL (http and https schemes supported).

    <sample-size>          When using INDEXED, RESERVOIR or WEIGHTED sampling, the sample size.
                             Can either be a whole number or a value between value between 0 and 1.
                             If a fraction, specifies the sample size as a percentage of the population. 
                             (e.g. 0.15 - 15 percent of the CSV)
                           When using BERNOULLI sampling, the probability of selecting each record
                             (between 0 and 1).
                           When using SYSTEMATIC sampling, the integer part is the interval between
                             records to sample & the fractional part is the percentage of the
                             population to sample. When there is no fractional part, it will
                             select every nth record for the entire population.
                           When using STRATIFIED sampling, the stratum sample size.
                           When using CLUSTER sampling, the number of clusters.
                           When using TIMESERIES sampling, the interval number (treated as hours
                             by default, e.g., 1 = 1 hour). Use --ts-interval for custom intervals
                             like "1d" (daily), "1w" (weekly), "1m" (monthly), "1y" (yearly), etc.                       

sample options:
    --seed <number>        Random Number Generator (RNG) seed.
    --rng <kind>           The Random Number Generator (RNG) algorithm to use.
                           Three RNGs are supported:
                            * standard: Use the standard RNG.
                              1.5 GB/s throughput.
                            * faster: Use faster RNG using the Xoshiro256Plus algorithm.
                              8 GB/s throughput.
                            * cryptosecure: Use cryptographically secure HC128 algorithm.
                              Recommended by eSTREAM (https://www.ecrypt.eu.org/stream/).
                              2.1 GB/s throughput though slow initialization.
                           [default: standard]

                           SAMPLING METHODS:
    --bernoulli            Use Bernoulli sampling instead of indexed or reservoir sampling.
                           When this flag is set, <sample-size> must be between
                           0 and 1 and represents the probability of selecting each record.
    --systematic <arg>     Use systematic sampling (every nth record as specified by <sample-size>).
                           If <arg> is "random", the starting point is randomly chosen between 0 & n.
                           If <arg> is "first", the starting point is the first record.
                           The sample size must be a whole number. Uses CONSTANT memory - O(1).
    --stratified <col>     Use stratified sampling. The strata column is specified by <col>.
                           Can be either a column name or 0-based column index.
                           The sample size must be a whole number. Uses MEMORY PROPORTIONAL to the
                           number of strata (s) and samples per stratum (k) - O(s*k).
    --weighted <col>       Use weighted sampling. The weight column is specified by <col>.
                           Can be either a column name or 0-based column index.
                           The column will be parsed as a number. Records with non-number weights
                           will be skipped.
                           Uses MEMORY PROPORTIONAL to the sample size (k) - O(k).
    --varopt <col>         Use VAROPT weighted reservoir sampling (A-ExpJ keying).
                           The weight column is specified by <col> (column name or 0-based index).
                           Variance-bounded, single-pass, no stats-cache required, mergeable
                           via --sketch-out / --sketch-in. Records with non-positive or
                           non-numeric weights are silently skipped.
                           Uses MEMORY PROPORTIONAL to the sample size (k) - O(k).
    --mergeable-reservoir  Use a mergeable Algorithm-R reservoir sampler. Distribution is
                           identical to the default RESERVOIR method, but the resulting sketch
                           is mergeable via --sketch-out / --sketch-in. Cannot be combined
                           with another sampling-method flag.
    --cluster <col>        Use cluster sampling. The cluster column is specified by <col>.
                           Can be either a column name or 0-based column index.
                           Uses MEMORY PROPORTIONAL to the number of clusters (c) - O(c).
    --timeseries <col>     Use time-series sampling. The time column is specified by <col>.
                           Can be either a column name or 0-based column index.
                           Sorts records by the specified time column and then groups by time intervals
                           and selects one record per interval.
                           Supports various date formats (19 formats recognized by qsv-dateparser).
                           Uses MEMORY PROPORTIONAL to the number of records - O(n).

                           TIME-SERIES SAMPLING OPTIONS:
    --ts-interval <intvl>  Time interval for grouping records. Format: <number><unit>
                           where unit is h (hour), d (day), w (week), m (month), y (year).
                           Examples: "1h", "1d", "1w", "2d" (every 2 days).
                           If not specified, <sample-size> is treated as hours.
    --ts-start <mode>      Starting point for time-series sampling.
                           Options: "first" (earliest timestamp, default), "last" (most recent timestamp),
                           "random" (random starting point).
                           [default: first]
    --ts-adaptive <mode>   Adaptive sampling mode for time-series data.
                           Options: "business-hours" (prefer 9am-5pm Mon-Fri),
                           "weekends" (prefer weekends), "business-days" (prefer weekdays),
                           "both" (combine business-hours and weekends).
    --ts-aggregate <func>  Aggregation function to apply within each time interval.
                           Options: "first", "last", "mean", "sum", "count", "min", "max", "median".
                           When specified, aggregates all records in each interval instead of selecting a single record.
    --ts-input-tz <tz>     Timezone for parsing input timestamps. Can be an IANA timezone name or "local" for the local timezone.
                           [default: UTC]
    --ts-prefer-dmy        Prefer to parse dates in dmy format. Otherwise, use mdy format.

                           SKETCH OPTIONS:
    --sketch-out <file>    After sampling, also write a binary sketch describing the internal
                           sampler state to <file>. The blob can later be merged into another
                           run via --sketch-in. Only valid with --varopt or --mergeable-reservoir.
                           The format is qsv-specific and is not interoperable with serialized
                           sketches from other tools.
    --sketch-in <files>    Comma-separated list of sketch files produced by --sketch-out.
                           CSV input is NOT read; the listed sketches (which must all be of
                           the same sampler kind) are merged and the resulting sample is
                           emitted as CSV. <sample-size> may be used to cap the merged sample
                           below the sketches' own k.

                           REMOTE FILE OPTIONS:
    --user-agent <agent>   Specify custom user agent to use when the input is a URL.
                           It supports the following variables -
                           $QSV_VERSION, $QSV_TARGET, $QSV_BIN_NAME, $QSV_KIND and $QSV_COMMAND.
                           Try to follow the syntax here -
                           https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/User-Agent
    --timeout <secs>       Inactivity timeout for downloading URLs in seconds. Aborts only if
                           no data is received from the server for this long. If 0, no timeout.
                           [default: 30]
    --max-size <mb>        Maximum size of the file to download in MB before sampling.
                           Will download the entire file if not specified.
                           If the CSV is partially downloaded, the sample will be taken
                           only from the downloaded portion.
    --force                Do not use stats cache, even if its available.

Common options:
    -h, --help             Display this message
    -o, --output <file>    Write output to <file> instead of stdout.
    -n, --no-headers       When set, the first row will be considered as part of
                           the population to sample from. (When not set, the
                           first row is the header row and will always appear
                           in the output.)
    -d, --delimiter <arg>  The field delimiter for reading/writing CSV data.
                           Must be a single character. (default: ,)
"#;

use std::{io, str::FromStr};

use chrono::{DateTime, Datelike, Duration, TimeZone, Timelike, Utc, Weekday};
use chrono_tz::Tz;
use foldhash::{HashMap, HashMapExt, HashSet, HashSetExt};
use futures_util::StreamExt;
use qsv_dateparser::parse_with_preference_and_timezone;
use rand::{
    Rng, RngExt, SeedableRng,
    distr::{Bernoulli, Distribution},
    prelude::IndexedRandom,
    rngs::StdRng,
};
use rand_hc::Hc128Rng;
use rand_xoshiro::Xoshiro256Plus;
use rayon::prelude::ParallelSliceMut;
use serde::Deserialize;
use strum_macros::EnumString;
use tempfile::NamedTempFile;
use url::Url;

use crate::{
    CliResult,
    config::{Config, Delimiter},
    select::SelectColumns,
    util,
    util::{SchemaArgs, StatsMode, get_stats_records},
};

mod sketches;

use sketches::{ReservoirItemsSketch, SketchFamily, VarOptItemsSketch, peek_family};

#[derive(Deserialize)]
struct Args {
    arg_input:                Option<String>,
    arg_sample_size:          f64,
    flag_output:              Option<String>,
    flag_no_headers:          bool,
    flag_delimiter:           Option<Delimiter>,
    flag_seed:                Option<u64>,
    flag_rng:                 String,
    flag_user_agent:          Option<String>,
    flag_timeout:             Option<u16>,
    flag_max_size:            Option<u64>,
    flag_bernoulli:           bool,
    flag_systematic:          Option<String>,
    flag_stratified:          Option<String>,
    flag_weighted:            Option<String>,
    flag_varopt:              Option<String>,
    flag_mergeable_reservoir: bool,
    flag_sketch_out:          Option<String>,
    flag_sketch_in:           Option<String>,
    flag_cluster:             Option<String>,
    flag_timeseries:          Option<String>,
    flag_ts_interval:         Option<String>,
    flag_ts_start:            Option<String>,
    flag_ts_adaptive:         Option<String>,
    flag_ts_aggregate:        Option<String>,
    flag_ts_input_tz:         Option<String>,
    flag_ts_prefer_dmy:       bool,
    flag_force:               bool,
}

impl Args {
    fn get_column_index(
        header: &csv::ByteRecord,
        column_spec: &str,
        purpose: &str,
    ) -> CliResult<usize> {
        // Try parsing as number first
        if let Ok(idx) = column_spec.parse::<usize>() {
            if idx < header.len() {
                return Ok(idx);
            }
            return fail_incorrectusage_clierror!(
                "{} column index {} is out of bounds (max: {})",
                purpose,
                idx,
                header.len() - 1
            );
        }

        // If not a number, try to find column by name (compare raw bytes —
        // header fields are CSV bytes, the spec is a UTF-8 string).
        let needle = column_spec.as_bytes();
        for (i, field) in header.iter().enumerate() {
            if field == needle {
                return Ok(i);
            }
        }

        fail_incorrectusage_clierror!("Could not find {} column named '{}'", purpose, column_spec)
    }

    fn get_strata_column(&self, header: &csv::ByteRecord) -> CliResult<usize> {
        match &self.flag_stratified {
            Some(col) => Self::get_column_index(header, col, "strata"),
            None => {
                fail_incorrectusage_clierror!(
                    "--stratified <col> is required for stratified sampling"
                )
            },
        }
    }

    fn get_weight_column(&self, header: &csv::ByteRecord) -> CliResult<usize> {
        match &self.flag_weighted {
            Some(col) => Self::get_column_index(header, col, "weight"),
            None => {
                fail_incorrectusage_clierror!("--weighted <col> is required for weighted sampling")
            },
        }
    }

    fn get_varopt_column(&self, header: &csv::ByteRecord) -> CliResult<usize> {
        match &self.flag_varopt {
            Some(col) => Self::get_column_index(header, col, "varopt weight"),
            None => {
                fail_incorrectusage_clierror!("--varopt <col> is required for varopt sampling")
            },
        }
    }

    fn get_cluster_column(&self, header: &csv::ByteRecord) -> CliResult<usize> {
        match &self.flag_cluster {
            Some(col) => Self::get_column_index(header, col, "cluster"),
            None => {
                fail_incorrectusage_clierror!("--cluster <col> is required for cluster sampling")
            },
        }
    }

    fn get_timeseries_column(&self, header: &csv::ByteRecord) -> CliResult<usize> {
        match &self.flag_timeseries {
            Some(col) => Self::get_column_index(header, col, "timeseries"),
            None => {
                fail_incorrectusage_clierror!(
                    "--timeseries <col> is required for timeseries sampling"
                )
            },
        }
    }
}

#[derive(Debug, EnumString, PartialEq)]
#[strum(ascii_case_insensitive)]
enum RngKind {
    Standard,
    Faster,
    Cryptosecure,
}

// Dispatches `$body` against the chosen RNG kind, binding `$rng` to a freshly
// created RNG of the matching concrete type. Replaces the recurring three-arm
// `match rng_kind { Standard => …, Faster => …, Cryptosecure => … }` blocks.
macro_rules! with_rng {
    ($rng_kind:expr, $seed:expr, |$rng:ident| $body:block) => {
        match $rng_kind {
            RngKind::Standard => {
                let mut $rng = StandardRng::create($seed);
                $body
            },
            RngKind::Faster => {
                let mut $rng = FasterRng::create($seed);
                $body
            },
            RngKind::Cryptosecure => {
                let mut $rng = CryptoRng::create($seed);
                $body
            },
        }
    };
}

#[derive(PartialEq)]
enum SamplingMethod {
    Bernoulli,
    Systematic,
    Stratified,
    Weighted,
    VarOpt,
    MergeableReservoir,
    SketchMerge,
    Cluster,
    Timeseries,
    Default,
}

// trait to handle different RNG types
trait RngProvider: Sized {
    type RngType: Rng + SeedableRng;

    fn get_name() -> &'static str;

    fn create(seed: Option<u64>) -> Self::RngType {
        if let Some(seed) = seed {
            Self::RngType::seed_from_u64(seed) // DevSkim: ignore DS148264
        } else {
            rand::make_rng::<Self::RngType>()
        }
    }
}

// Implement for each RNG type
struct StandardRng;
impl RngProvider for StandardRng {
    type RngType = StdRng;

    fn get_name() -> &'static str {
        "standard"
    }
}

struct FasterRng;
impl RngProvider for FasterRng {
    type RngType = Xoshiro256Plus;

    fn get_name() -> &'static str {
        "faster"
    }
}

struct CryptoRng;
impl RngProvider for CryptoRng {
    type RngType = Hc128Rng;

    fn get_name() -> &'static str {
        "cryptosecure"
    }
}

// Time-series start mode
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum TSStartMode {
    First,
    Last,
    Random,
}

impl std::str::FromStr for TSStartMode {
    type Err = &'static str;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_lowercase().as_str() {
            "first" => Ok(TSStartMode::First),
            "last" => Ok(TSStartMode::Last),
            "random" => Ok(TSStartMode::Random),
            _ => Err("Time-series start mode must be 'first', 'last' or 'random'"),
        }
    }
}

// Time-series sampling helper functions

#[derive(Debug, Clone, Copy, PartialEq)]
enum AggregationFunction {
    First,
    Last,
    Mean,
    Sum,
    Count,
    Min,
    Max,
    Median,
}

impl FromStr for AggregationFunction {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_lowercase().as_str() {
            "first" => Ok(AggregationFunction::First),
            "last" => Ok(AggregationFunction::Last),
            "mean" => Ok(AggregationFunction::Mean),
            "sum" => Ok(AggregationFunction::Sum),
            "count" => Ok(AggregationFunction::Count),
            "min" => Ok(AggregationFunction::Min),
            "max" => Ok(AggregationFunction::Max),
            "median" => Ok(AggregationFunction::Median),
            _ => Err(format!(
                "Invalid aggregation function: {s}. Supported: first, last, mean, sum, count, \
                 min, max, median"
            )),
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
enum AdaptiveMode {
    BusinessHours,
    Weekends,
    BusinessDays,
    Both,
}

impl FromStr for AdaptiveMode {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_lowercase().as_str() {
            "business-hours" | "businesshours" => Ok(AdaptiveMode::BusinessHours),
            "weekends" => Ok(AdaptiveMode::Weekends),
            "business-days" | "businessdays" => Ok(AdaptiveMode::BusinessDays),
            "both" => Ok(AdaptiveMode::Both),
            _ => Err(format!(
                "Invalid adaptive mode: {s}. Supported: business-hours, weekends, business-days, \
                 both"
            )),
        }
    }
}

fn parse_time_interval(interval_str: &str) -> CliResult<Duration> {
    let s = interval_str.trim().to_lowercase();

    // Try to parse as number + unit (e.g., "1h", "2d", "3w")
    if s.len() < 2 {
        return fail_incorrectusage_clierror!(
            "Invalid time interval format: {interval_str}. Expected format: <number><unit> (e.g., \
             1h, 1d, 1w, 1m, 1y)"
        );
    }

    let (num_str, unit) = s.split_at(s.len() - 1);
    let num: i64 = num_str.parse().map_err(|_| {
        format!(
            "Invalid time interval number: {num_str}. Expected format: <number><unit> (e.g., 1h, \
             1d, 1w, 1m, 1y)"
        )
    })?;

    if num <= 0 {
        return fail_incorrectusage_clierror!("Time interval must be positive");
    }

    let duration = match unit {
        "h" => Duration::hours(num),
        "d" => Duration::days(num),
        "w" => Duration::weeks(num),
        "m" => Duration::days(num * 30), // Approximate month as 30 days
        "y" => Duration::days(num * 365), // Approximate year as 365 days
        _ => {
            return fail_incorrectusage_clierror!(
                "Invalid time interval unit: {unit}. Supported units: h (hour), d (day), w \
                 (week), m (month), y (year)"
            );
        },
    };

    Ok(duration)
}

fn parse_timestamp(
    value: &[u8],
    prefer_dmy: bool,
    input_tz: Option<&str>,
) -> CliResult<DateTime<Utc>> {
    // Try to parse as UTF-8 string first

    let Ok(value_str) = simdutf8::basic::from_utf8(value) else {
        return fail_incorrectusage_clierror!("Time column value is not valid UTF-8");
    };

    // Try parsing as Unix timestamp first (simple integer check).
    // Disambiguate seconds vs milliseconds by magnitude: chrono accepts a huge
    // year range, so a millisecond value like 1_704_067_200_000 would silently
    // parse as seconds (year ~55,899) if we tried seconds first. We pick a
    // 10^11 cutoff (covers any reasonable date in seconds well past year 5000)
    // and treat anything larger as milliseconds.
    if let Ok(ts_val) = atoi_simd::parse::<i64, false, false>(value) {
        const SEC_LIMIT: i64 = 100_000_000_000; // 10^11
        // Range form (not `.abs()`) so i64::MIN doesn't panic in debug.
        if (-SEC_LIMIT..SEC_LIMIT).contains(&ts_val) {
            if let Some(dt) = Utc.timestamp_opt(ts_val, 0).single() {
                return Ok(dt);
            }
        } else if let Some(dt) = Utc.timestamp_millis_opt(ts_val).single() {
            return Ok(dt);
        }
    }

    // Parse timezone
    let tz: Tz = if let Some(tz_str) = input_tz {
        if tz_str.eq_ignore_ascii_case("local") {
            if let Ok(tz_name) = iana_time_zone::get_timezone() {
                tz_name.parse::<Tz>().unwrap_or(chrono_tz::UTC)
            } else {
                chrono_tz::UTC
            }
        } else {
            tz_str.parse::<Tz>().unwrap_or(chrono_tz::UTC)
        }
    } else {
        chrono_tz::UTC
    };

    // Parse using qsv_dateparser
    parse_with_preference_and_timezone(value_str, prefer_dmy, &tz)
        .map_err(|e| format!("Failed to parse timestamp '{value_str}': {e}").into())
}

fn is_business_hours(dt: &DateTime<Utc>) -> bool {
    let hour = dt.hour();
    (9..=17).contains(&hour)
}

fn is_weekend(dt: &DateTime<Utc>) -> bool {
    matches!(dt.weekday(), Weekday::Sat | Weekday::Sun)
}

fn is_business_day(dt: &DateTime<Utc>) -> bool {
    !is_weekend(dt)
}

fn check_stats_cache(
    args: &Args,
    method: &SamplingMethod,
) -> CliResult<(Option<f64>, Option<u64>)> {
    if args.flag_force {
        return Ok((None, None));
    }

    // We returned early above when flag_force was set, so this is always
    // false here. Asserting (rather than hard-coding the literal) keeps the
    // two callsites in sync if the early-return guard ever moves.
    debug_assert!(!args.flag_force);
    let schema_args = SchemaArgs {
        arg_input:            args.arg_input.clone(),
        flag_no_headers:      args.flag_no_headers,
        flag_delimiter:       args.flag_delimiter,
        flag_jobs:            None,
        flag_polars:          false,
        flag_memcheck:        false,
        flag_force:           args.flag_force,
        flag_prefer_dmy:      false,
        flag_dates_whitelist: String::new(),
        flag_enum_threshold:  0,
        flag_ignore_case:     false,
        flag_strict_dates:    false,
        flag_strict_formats:  false,
        flag_pattern_columns: SelectColumns::parse("")?,
        flag_stdout:          false,
        flag_output:          None,
    };

    // Get stats records
    match get_stats_records(&schema_args, StatsMode::Frequency) {
        Ok((csv_fields, stats)) => {
            // Extract relevant stats based on sampling method
            let mut max_weight = None;
            let mut cardinality = None;
            match method {
                SamplingMethod::Weighted => {
                    // For weighted sampling, get max weight
                    if let Some(weight_col) = &args.flag_weighted {
                        let idx = weight_col.parse::<usize>().ok().or_else(|| {
                            csv_fields
                                .iter()
                                .position(|field| field == weight_col.as_bytes())
                        });

                        if let Some(idx) = idx
                            && let Some(col_stats) = stats.get(idx)
                        {
                            let min_weight = col_stats
                                .min
                                .clone()
                                .unwrap_or_default()
                                .parse::<f64>()
                                .unwrap_or_default();
                            if min_weight < 0.0 {
                                return fail_incorrectusage_clierror!(
                                    "Weights must be non-negative. Lowest weight: {min_weight}"
                                );
                            }

                            max_weight = col_stats.max.clone().unwrap().parse::<f64>().ok();
                        }
                    }
                },
                SamplingMethod::Cluster => {
                    // For cluster sampling, get cardinality
                    if let Some(cluster_col) = &args.flag_cluster {
                        let idx = cluster_col.parse::<usize>().ok().or_else(|| {
                            csv_fields
                                .iter()
                                .position(|field| field == cluster_col.as_bytes())
                        });

                        if let Some(idx) = idx {
                            cardinality = stats.get(idx).map(|col_stats| col_stats.cardinality);
                        }
                    }
                },
                _ => {},
            }

            Ok((max_weight, cardinality))
        },
        _ => Ok((None, None)),
    }
}

// "streaming" bernoulli sampling
//
// Boundary detection: instead of scanning the byte buffer for raw `\n` (which
// would incorrectly split CSVs whose fields contain quoted newlines), we drive
// the csv parser itself and use `Reader::position().byte()` to learn where each
// record actually ends. We only commit records whose terminator we know lies
// WITHIN the current buffer — if the parser consumes all of it, the trailing
// record might be partial, so we hold it back until either more data arrives
// or the stream closes naturally.
//
// `--max-size` truncation is treated as NOT-EOF: a capped buffer may have cut
// the final record in half, so we never let the parser's "treat trailing
// bytes as a record" behavior fire at the cap boundary.
#[allow(clippy::future_not_send)]
async fn stream_bernoulli_sampling(uri: &str, args: &Args, rng_kind: &RngKind) -> CliResult<()> {
    // Resolve the delimiter ONCE: --delimiter wins, then QSV_DEFAULT_DELIMITER,
    // then comma. The same byte is reused for every csv::ReaderBuilder probe
    // below so header and data parses agree.
    let delim_byte = if let Some(d) = args.flag_delimiter {
        d.as_byte()
    } else if let Ok(delim) = std::env::var("QSV_DEFAULT_DELIMITER") {
        Delimiter::decode_delimiter(&delim).map_or(b',', super::super::config::Delimiter::as_byte)
    } else {
        b','
    };

    let mut wtr = Config::new(args.flag_output.as_ref())
        .delimiter(args.flag_delimiter)
        .writer()?;

    let client = util::create_reqwest_async_client(
        args.flag_user_agent.clone(),
        util::timeout_secs(args.flag_timeout.unwrap_or(30)).map(|t| t as u16)?,
        Some(uri.to_string()),
    )?;

    // Fail fast on non-2xx — reqwest's `.send()` does NOT error on HTTP error
    // status, so without this a 404/500 HTML body would be streamed straight
    // into the csv parser and produce confusing record errors.
    let response = client.get(uri).send().await?.error_for_status()?;
    let mut stream = response.bytes_stream();

    let max_bytes = args.flag_max_size.map(|mb| mb * 1024 * 1024);

    // Create only the RNG we'll actually use — Cryptosecure init is slow.
    let mut std_rng = match rng_kind {
        RngKind::Standard => Some(StandardRng::create(args.flag_seed)),
        _ => None,
    };
    let mut faster_rng = match rng_kind {
        RngKind::Faster => Some(FasterRng::create(args.flag_seed)),
        _ => None,
    };
    let mut crypto_rng = match rng_kind {
        RngKind::Cryptosecure => Some(CryptoRng::create(args.flag_seed)),
        _ => None,
    };
    let probability = args.arg_sample_size;

    let mut buffer: Vec<u8> = Vec::new();
    let mut bytes_read: u64 = 0;
    let mut header_handled = args.flag_no_headers;
    let mut record = csv::ByteRecord::new();
    let mut size_capped = false;
    let mut stream_done = false;

    while !stream_done {
        match stream.next().await {
            Some(chunk) => {
                let chunk = chunk?;
                if let Some(cap) = max_bytes {
                    let remaining = cap.saturating_sub(bytes_read);
                    if remaining == 0 {
                        size_capped = true;
                        stream_done = true;
                    } else {
                        let take = (chunk.len() as u64).min(remaining) as usize;
                        buffer.extend_from_slice(&chunk[..take]);
                        bytes_read += take as u64;
                        if take < chunk.len() {
                            size_capped = true;
                        }
                    }
                } else {
                    buffer.extend_from_slice(&chunk);
                    bytes_read += chunk.len() as u64;
                }
            },
            None => stream_done = true,
        }

        // EOF flag for the csv parser: only true on a *natural* stream end.
        // A size-capped buffer may have truncated the final record mid-way,
        // so we must NOT trust the parser's trailing-record fallback there.
        let parser_eof = stream_done && !size_capped;

        // Header: read exactly one complete record from the buffer.
        if !header_handled {
            let mut probe = csv::ReaderBuilder::new()
                .has_headers(false)
                .delimiter(delim_byte)
                .from_reader(&buffer[..]);
            let mut hdr = csv::ByteRecord::new();
            if matches!(probe.read_byte_record(&mut hdr), Ok(true)) {
                let pos = probe.position().byte() as usize;
                // A record terminator was definitely inside the buffer iff the
                // parser stopped before consuming all of it. (Or the stream is
                // genuinely over and what we have IS the whole header.)
                if parser_eof || pos < buffer.len() {
                    wtr.write_byte_record(&hdr)?;
                    buffer.drain(..pos);
                    header_handled = true;
                }
                // else: the parser may have treated EOF-of-slice as a record
                // terminator; wait for more data before committing.
            }
        }

        // Data: read every record whose terminator is INSIDE the buffer (or
        // every remaining record once we know the stream is naturally done).
        if header_handled && !buffer.is_empty() {
            let mut probe = csv::ReaderBuilder::new()
                .has_headers(false)
                .delimiter(delim_byte)
                .from_reader(&buffer[..]);
            let mut last_consumed = 0usize;
            while matches!(probe.read_byte_record(&mut record), Ok(true)) {
                let pos = probe.position().byte() as usize;
                if !parser_eof && pos == buffer.len() {
                    // Parser ate the rest of the buffer — this last "record"
                    // might be a partial one. Hold it back.
                    break;
                }
                let pick = match rng_kind {
                    RngKind::Standard => std_rng.as_mut().unwrap().random_bool(probability),
                    RngKind::Faster => faster_rng.as_mut().unwrap().random_bool(probability),
                    RngKind::Cryptosecure => crypto_rng.as_mut().unwrap().random_bool(probability),
                };
                if pick {
                    wtr.write_byte_record(&record)?;
                }
                last_consumed = pos;
            }
            if last_consumed > 0 {
                buffer.drain(..last_consumed);
            }
        }

        if size_capped {
            break;
        }
    }

    Ok(wtr.flush()?)
}

pub fn run(argv: &[&str]) -> CliResult<()> {
    let mut args: Args = util::get_args(USAGE, argv)?;

    if !args.arg_sample_size.is_finite() {
        return fail_incorrectusage_clierror!("Sample size must be a finite number.");
    }
    if args.arg_sample_size.is_sign_negative() {
        return fail_incorrectusage_clierror!("Sample size cannot be negative.");
    }

    // --sketch-in is its own input mode (CSV is NOT read); it cannot be
    // combined with any sampling-method flag. Detect this first so the user
    // gets a targeted error rather than the generic "Only one sampling
    // method" message.
    if args.flag_sketch_in.is_some() {
        let conflicting = args.flag_bernoulli
            || args.flag_systematic.is_some()
            || args.flag_stratified.is_some()
            || args.flag_weighted.is_some()
            || args.flag_varopt.is_some()
            || args.flag_mergeable_reservoir
            || args.flag_cluster.is_some()
            || args.flag_timeseries.is_some();
        if conflicting {
            return fail_incorrectusage_clierror!(
                "--sketch-in does not read CSV input; do not combine it with another \
                 sampling-method flag (--bernoulli, --systematic, --stratified, --weighted, \
                 --varopt, --mergeable-reservoir, --cluster, --timeseries)."
            );
        }
    }

    // Validate that only one sampling method is selected
    let methods = [
        args.flag_bernoulli,
        args.flag_systematic.is_some(),
        args.flag_stratified.is_some(),
        args.flag_weighted.is_some(),
        args.flag_varopt.is_some(),
        args.flag_mergeable_reservoir,
        args.flag_cluster.is_some(),
        args.flag_timeseries.is_some(),
    ];
    if methods.iter().filter(|&&x| x).count() > 1 {
        return fail_incorrectusage_clierror!("Only one sampling method can be specified");
    }

    // --sketch-out is only meaningful with --varopt, --mergeable-reservoir, or
    // --sketch-in (re-emitting a merged sketch). Reject other combinations.
    if args.flag_sketch_out.is_some()
        && args.flag_varopt.is_none()
        && !args.flag_mergeable_reservoir
        && args.flag_sketch_in.is_none()
    {
        return fail_incorrectusage_clierror!(
            "--sketch-out is only valid with --varopt, --mergeable-reservoir, or --sketch-in"
        );
    }

    // For Bernoulli, sample-size IS the probability — validate up front so the
    // streaming-URL path (which dispatches before the per-method checks below)
    // can't reach random_bool() with an out-of-range value and panic.
    if args.flag_bernoulli && (args.arg_sample_size >= 1.0 || args.arg_sample_size <= 0.0) {
        return fail_incorrectusage_clierror!(
            "Bernoulli sampling requires a probability between 0 and 1"
        );
    }

    let Ok(rng_kind) = RngKind::from_str(&args.flag_rng) else {
        return fail_incorrectusage_clierror!(
            "Invalid RNG algorithm `{}`. Supported RNGs are: standard, faster, cryptosecure.",
            args.flag_rng
        );
    };

    let sampling_method = if args.flag_bernoulli {
        SamplingMethod::Bernoulli
    } else if args.flag_systematic.is_some() {
        SamplingMethod::Systematic
    } else if args.flag_stratified.is_some() {
        SamplingMethod::Stratified
    } else if args.flag_weighted.is_some() {
        SamplingMethod::Weighted
    } else if args.flag_varopt.is_some() {
        SamplingMethod::VarOpt
    } else if args.flag_sketch_in.is_some() {
        SamplingMethod::SketchMerge
    } else if args.flag_mergeable_reservoir {
        SamplingMethod::MergeableReservoir
    } else if args.flag_cluster.is_some() {
        SamplingMethod::Cluster
    } else if args.flag_timeseries.is_some() {
        SamplingMethod::Timeseries
    } else {
        SamplingMethod::Default
    };

    // SketchMerge does not read CSV input — handle it entirely here.
    if sampling_method == SamplingMethod::SketchMerge {
        let mut wtr = Config::new(args.flag_output.as_ref())
            .delimiter(args.flag_delimiter)
            .writer()?;
        sample_sketch_merge(&args, &mut wtr, &rng_kind)?;
        return Ok(wtr.flush()?);
    }

    let temp_download = NamedTempFile::new()?;

    args.arg_input = match args.arg_input {
        Some(ref uri) if Url::parse(uri).is_ok() && uri.starts_with("http") => {
            // For bernoulli sampling with remote file, handle specially
            if sampling_method == SamplingMethod::Bernoulli {
                log::info!("Streaming Bernoulli sampling remote file");

                let rt = tokio::runtime::Runtime::new()?;
                rt.block_on(stream_bernoulli_sampling(uri, &args, &rng_kind))?;
                return Ok(());
            }

            // For other cases, download entire file
            let max_size_bytes = args.flag_max_size.map(|mb| mb * 1024 * 1024);
            let future = util::download_file(
                uri,
                temp_download.path().to_path_buf(),
                false,
                Some(util::set_user_agent(args.flag_user_agent.clone())?),
                args.flag_timeout,
                max_size_bytes,
            );
            tokio::runtime::Runtime::new()?.block_on(future)?;
            // safety: temp_download is a NamedTempFile, so we know can unwrap.to_string
            Some(temp_download.path().to_str().unwrap().to_string())
        },
        Some(uri) => Some(uri), // local file
        None => None,
    };

    let rconfig = Config::new(args.arg_input.as_ref())
        .delimiter(args.flag_delimiter)
        .no_headers_flag(args.flag_no_headers)
        .flexible(true)
        .skip_format_check(true);

    let mut rdr = rconfig.reader()?;
    let mut wtr = Config::new(args.flag_output.as_ref())
        .delimiter(args.flag_delimiter)
        .writer()?;

    // Write headers unless --no-headers is specified
    rconfig.write_headers(&mut rdr, &mut wtr)?;

    let mut sample_size = args.arg_sample_size;

    match sampling_method {
        SamplingMethod::Bernoulli => {
            // probability range was validated up front (see run() prelude)
            sample_bernoulli(
                &mut rdr,
                &mut wtr,
                args.arg_sample_size,
                args.flag_seed,
                &rng_kind,
            )?;
        },
        SamplingMethod::Systematic => {
            let starting_point = match args.flag_systematic.as_deref().map(str::to_lowercase) {
                Some(arg) if arg == "random" || arg == "first" => arg,
                Some(_) => {
                    return fail_incorrectusage_clierror!(
                        "Systematic sampling starting point must be either 'random' or 'first'"
                    );
                },
                None => String::from("random"),
            };

            let row_count: u64 = if let Ok(rc) = util::count_rows(&rconfig) {
                rc
            } else {
                return fail!("Systematic sampling requires rowcount.");
            };

            sample_systematic(
                &mut rdr,
                &mut wtr,
                args.arg_sample_size,
                row_count,
                &starting_point,
                args.flag_seed,
                &rng_kind,
            )?;
        },
        SamplingMethod::Stratified => {
            let strata_column = args.get_strata_column(&rdr.byte_headers()?.clone())?;
            sample_stratified(
                &mut rdr,
                &mut wtr,
                strata_column,
                args.arg_sample_size as usize,
                args.flag_seed,
                &rng_kind,
            )?;
        },
        SamplingMethod::Weighted => {
            let weight_column = args.get_weight_column(&rdr.byte_headers()?.clone())?;

            // Get max_weight from cache if available
            let (max_weight, _) = check_stats_cache(&args, &SamplingMethod::Weighted)?;

            // determine sample size
            #[allow(clippy::cast_precision_loss)]
            let sample_size = if args.arg_sample_size < 1.0 {
                let row_count: u64 = if let Ok(rc) = util::count_rows(&rconfig) {
                    rc
                } else {
                    return fail!("Weighted fractional sampling requires rowcount.");
                };
                (row_count as f64 * args.arg_sample_size).round() as usize
            } else {
                args.arg_sample_size as usize
            };

            sample_weighted(
                &rconfig,
                &mut rdr,
                &mut wtr,
                weight_column,
                max_weight,
                sample_size,
                args.flag_seed,
                &rng_kind,
            )?;
        },
        SamplingMethod::VarOpt => {
            let headers = rdr.byte_headers()?.clone();
            let weight_column = args.get_varopt_column(&headers)?;

            // Embed the CSV header into the sketch so a future --sketch-in
            // can re-emit a schema-bearing CSV. Skip when --no-headers is set
            // (the "header" then is actually the first data row).
            let embedded_header = if args.flag_no_headers {
                None
            } else {
                Some(headers)
            };

            #[allow(clippy::cast_precision_loss)]
            let sample_size = if args.arg_sample_size < 1.0 {
                let row_count: u64 = if let Ok(rc) = util::count_rows(&rconfig) {
                    rc
                } else {
                    return fail!("VarOpt fractional sampling requires rowcount.");
                };
                (row_count as f64 * args.arg_sample_size).round() as usize
            } else {
                args.arg_sample_size as usize
            };

            sample_varopt(
                &mut rdr,
                &mut wtr,
                weight_column,
                sample_size,
                args.flag_seed,
                &rng_kind,
                args.flag_sketch_out.as_deref(),
                embedded_header,
            )?;
        },
        SamplingMethod::MergeableReservoir => {
            // Embed the CSV header into the sketch so a future --sketch-in
            // can re-emit a schema-bearing CSV. Skip when --no-headers is set
            // (the "header" then is actually the first data row).
            let embedded_header = if args.flag_no_headers {
                None
            } else {
                Some(rdr.byte_headers()?.clone())
            };

            #[allow(clippy::cast_precision_loss)]
            let sample_size = if args.arg_sample_size < 1.0 {
                let row_count: u64 = if let Ok(rc) = util::count_rows(&rconfig) {
                    rc
                } else {
                    return fail!("Mergeable-reservoir fractional sampling requires rowcount.");
                };
                (row_count as f64 * args.arg_sample_size).round() as usize
            } else {
                args.arg_sample_size as usize
            };

            sample_mergeable_reservoir(
                &mut rdr,
                &mut wtr,
                sample_size,
                args.flag_seed,
                &rng_kind,
                args.flag_sketch_out.as_deref(),
                embedded_header,
            )?;
        },
        SamplingMethod::SketchMerge => {
            // handled above before the CSV reader is opened
            unreachable!("SketchMerge handled before CSV pipeline");
        },
        SamplingMethod::Cluster => {
            let cluster_column = args.get_cluster_column(&rdr.byte_headers()?.clone())?;

            // Get cardinality from cache if available
            let (_, cardinality) = check_stats_cache(&args, &SamplingMethod::Cluster)?;

            sample_cluster(
                &rconfig,
                &mut rdr,
                &mut wtr,
                cluster_column,
                cardinality,
                args.arg_sample_size as usize,
                args.flag_seed,
                &rng_kind,
            )?;
        },
        SamplingMethod::Timeseries => {
            let time_column = args.get_timeseries_column(&rdr.byte_headers()?.clone())?;

            // Parse interval - prefer --ts-interval flag, otherwise use sample_size as hours
            let interval_str = if let Some(interval) = &args.flag_ts_interval {
                interval.clone()
            } else if args.arg_sample_size.fract() == 0.0 && args.arg_sample_size > 0.0 {
                // If it's a whole number, treat as hours
                format!("{}h", args.arg_sample_size as i64)
            } else {
                return fail_incorrectusage_clierror!(
                    "Time-series sampling requires either --ts-interval (e.g., '1h', '1d', '1w', \
                     '1m', '1y') or a positive whole number for <sample-size> (treated as hours)"
                );
            };

            let start_mode = match args
                .flag_ts_start
                .as_deref()
                .unwrap_or("first")
                .parse::<TSStartMode>()
            {
                Ok(mode) => mode,
                Err(msg) => return fail_incorrectusage_clierror!("{msg}"),
            };

            // Parse adaptive mode
            let adaptive_mode = if let Some(adaptive_str) = &args.flag_ts_adaptive {
                Some(
                    AdaptiveMode::from_str(adaptive_str)
                        .map_err(|e| format!("Invalid adaptive mode: {e}"))?,
                )
            } else {
                None
            };

            // Parse aggregation function
            let aggregate_func = if let Some(agg_str) = &args.flag_ts_aggregate {
                Some(
                    AggregationFunction::from_str(agg_str)
                        .map_err(|e| format!("Invalid aggregation function: {e}"))?,
                )
            } else {
                None
            };

            // Get timezone and prefer_dmy settings
            let prefer_dmy = args.flag_ts_prefer_dmy || rconfig.get_dmy_preference();
            let input_tz = match args.flag_ts_input_tz.as_deref() {
                Some(tz_str) => {
                    if tz_str.eq_ignore_ascii_case("local") {
                        if let Ok(tz_name) = iana_time_zone::get_timezone() {
                            if tz_name.parse::<chrono_tz::Tz>().is_ok() {
                                Some(tz_str)
                            } else {
                                wwarn!(
                                    "Invalid local timezone from iana_time_zone, falling back to \
                                     UTC."
                                );
                                None
                            }
                        } else {
                            wwarn!("Could not determine local timezone, falling back to UTC.");
                            None
                        }
                    } else if tz_str.parse::<chrono_tz::Tz>().is_ok() {
                        Some(tz_str)
                    } else {
                        wwarn!("Invalid timezone '{tz_str}', falling back to UTC.");
                        None
                    }
                },
                None => None,
            };

            sample_timeseries(
                &rconfig,
                &mut rdr,
                &mut wtr,
                time_column,
                &interval_str,
                start_mode,
                adaptive_mode,
                aggregate_func,
                prefer_dmy,
                input_tz,
                args.flag_seed,
                &rng_kind,
            )?;
        },
        SamplingMethod::Default => {
            // no sampling method is specified, so we do indexed sampling
            // if an index is present
            if let Some(mut idx) = rconfig.indexed()? {
                #[allow(clippy::cast_precision_loss)]
                if sample_size < 1.0 {
                    sample_size *= idx.count() as f64;
                }

                let sample_count = sample_size as usize;
                let total_count = idx.count().try_into().unwrap();

                log::info!("doing {rng_kind:?} INDEXED sampling...");
                with_rng!(&rng_kind, args.flag_seed, |rng| {
                    sample_indices(&mut rng, total_count, sample_count, |i| {
                        idx.seek(i as u64)?;
                        Ok(wtr.write_byte_record(&idx.byte_records().next().unwrap()?)?)
                    })?;
                });
            } else {
                // No sampling method is specified and no index is present
                // do reservoir sampling

                #[allow(clippy::cast_precision_loss)]
                let sample_size = if args.arg_sample_size < 1.0 {
                    let row_count: u64 = if let Ok(rc) = util::count_rows(&rconfig) {
                        rc
                    } else {
                        return fail!("Fractional sampling requires rowcount.");
                    };
                    (row_count as f64 * args.arg_sample_size).round() as u64
                } else {
                    args.arg_sample_size as u64
                };

                sample_reservoir(&mut rdr, &mut wtr, sample_size, args.flag_seed, &rng_kind)?;
            }
        },
    }

    Ok(wtr.flush()?)
}

fn sample_reservoir<R: io::Read, W: io::Write>(
    rdr: &mut csv::Reader<R>,
    wtr: &mut csv::Writer<W>,
    sample_size: u64,
    seed: Option<u64>,
    rng_kind: &RngKind,
) -> CliResult<()> {
    let mut reservoir = Vec::with_capacity(sample_size as usize);
    let mut records = rdr.byte_records().enumerate();

    // Pre-fill reservoir
    // Note that we use by_ref() to avoid consuming the iterator
    // and we only take the first sample_size records
    for (_, row) in records.by_ref().take(sample_size as usize) {
        reservoir.push(row?);
    }

    match rng_kind {
        RngKind::Standard => {
            do_reservoir_sampling::<StandardRng>(&mut records, &mut reservoir, sample_size, seed)
        },
        RngKind::Faster => {
            do_reservoir_sampling::<FasterRng>(&mut records, &mut reservoir, sample_size, seed)
        },
        RngKind::Cryptosecure => {
            do_reservoir_sampling::<CryptoRng>(&mut records, &mut reservoir, sample_size, seed)
        },
    }?;

    // Write the reservoir to output
    for record in reservoir {
        wtr.write_byte_record(&record)?;
    }

    Ok(())
}

// Generic reservoir sampling implementation. Memory: O(k) for the reservoir
// passed in by the caller (the algorithm itself uses only constant extra
// state).
fn do_reservoir_sampling<T: RngProvider>(
    records: &mut impl Iterator<Item = (usize, Result<csv::ByteRecord, csv::Error>)>,
    reservoir: &mut [csv::ByteRecord],
    sample_size: u64,
    seed: Option<u64>,
) -> CliResult<()> {
    log::info!("doing {} RESERVOIR sampling...", T::get_name());
    let mut rng = T::create(seed);

    // Bound writes against the reservoir's actual length, not sample_size.
    // sample_reservoir() pre-fills via take(sample_size), which yields fewer
    // when the input is shorter than k — in that case we have nothing to
    // sample, but checking reservoir.len() keeps this defensible if a future
    // caller passes an undersized slice.
    let reservoir_cap = reservoir.len().min(sample_size as usize);

    // Process remaining records using Algorithm R (Robert Floyd)
    for (i, row) in records {
        let random_idx = rng.random_range(0..=i);
        if random_idx < reservoir_cap {
            reservoir[random_idx] = row?;
        }
    }
    Ok(())
}

fn sample_bernoulli<R: io::Read, W: io::Write>(
    rdr: &mut csv::Reader<R>,
    wtr: &mut csv::Writer<W>,
    probability: f64,
    seed: Option<u64>,
    rng_kind: &RngKind,
) -> CliResult<()> {
    let mut records = rdr.byte_records();

    match rng_kind {
        RngKind::Standard => {
            do_bernoulli_sampling::<StandardRng>(&mut records, wtr, probability, seed)
        },
        RngKind::Faster => do_bernoulli_sampling::<FasterRng>(&mut records, wtr, probability, seed),
        RngKind::Cryptosecure => {
            do_bernoulli_sampling::<CryptoRng>(&mut records, wtr, probability, seed)
        },
    }
}

// Generic bernoulli sampling implementation using constant memory
fn do_bernoulli_sampling<T: RngProvider>(
    records: &mut impl Iterator<Item = Result<csv::ByteRecord, csv::Error>>,
    wtr: &mut csv::Writer<impl io::Write>,
    probability: f64,
    seed: Option<u64>,
) -> CliResult<()> {
    log::info!("doing {} BERNOULLI sampling...", T::get_name());
    let mut rng = T::create(seed);

    let dist =
        Bernoulli::new(probability).map_err(|_| "probability must be between 0.0 and 1.0")?;

    for row in records {
        if dist.sample(&mut rng) {
            wtr.write_byte_record(&row?)?;
        }
    }
    Ok(())
}

// Helper function to sample indices. Memory: O(k) for the selected-index
// buffer (the algorithm itself uses constant extra state).
fn sample_indices<F>(
    rng: &mut impl Rng,
    total_count: usize,
    sample_count: usize,
    mut process_index: F,
) -> CliResult<()>
where
    F: FnMut(usize) -> CliResult<()>,
{
    if sample_count > total_count {
        return fail!("Sample size cannot be larger than population size");
    }

    // Fill first k positions, then reservoir-sample (Algorithm R) the rest.
    let mut selected: Vec<usize> = (0..sample_count).collect();

    for i in sample_count..total_count {
        let j = rng.random_range(0..=i);
        if j < sample_count {
            selected[j] = i;
        }
    }

    // Process indices in order to avoid seeking back and forth
    selected.par_sort_unstable();
    for idx in selected {
        process_index(idx)?;
    }

    Ok(())
}

// Systematic sampling implementation
fn sample_systematic<R: io::Read, W: io::Write>(
    rdr: &mut csv::Reader<R>,
    wtr: &mut csv::Writer<W>,
    sample_size: f64,
    row_count: u64,
    starting_point: &str,
    seed: Option<u64>,
    rng_kind: &RngKind,
) -> CliResult<()> {
    if sample_size <= 0.0 {
        return fail_incorrectusage_clierror!("Sample size must be positive");
    }

    // Split sample_size into integer and fractional parts
    let interval = sample_size.trunc() as usize;
    let percentage = sample_size.fract();

    if interval == 0 {
        return fail_incorrectusage_clierror!("Interval must be at least 1");
    }

    // Calculate target sample size based on percentage
    #[allow(clippy::cast_precision_loss)]
    let target_count = if percentage > 0.0 {
        ((row_count as f64) * percentage).round() as u64
    } else {
        row_count
    };

    // Select starting point
    let start = if starting_point == "random" {
        with_rng!(rng_kind, seed, |rng| { rng.random_range(0..interval) })
    } else {
        0 // starting point is the first record
    };

    // Select records at regular intervals
    let mut selected_count = 0;
    for (i, record) in rdr.byte_records().enumerate().skip(start) {
        if i.is_multiple_of(interval) && selected_count < target_count {
            wtr.write_byte_record(&record?)?;
            selected_count += 1;
        }
    }

    Ok(())
}

// Stratified sampling implementation
fn sample_stratified<R: io::Read, W: io::Write>(
    rdr: &mut csv::Reader<R>,
    wtr: &mut csv::Writer<W>,
    strata_column: usize,
    samples_per_stratum: usize,
    seed: Option<u64>,
    rng_kind: &RngKind,
) -> CliResult<()> {
    const ESTIMATED_STRATA_COUNT: usize = 100;

    // Single-pass: collect records, and discover strata into a per-stratum
    // state map as we go. Each entry holds (reservoir, records_seen) — packing
    // both into one HashMap halves per-record lookups in do_stratified_sampling
    // and removes the awkward "lookup-or-insert" dance over a parallel map.
    let mut strata: HashMap<Vec<u8>, (Vec<csv::ByteRecord>, usize)> =
        HashMap::with_capacity(ESTIMATED_STRATA_COUNT);
    let mut records = Vec::with_capacity(ESTIMATED_STRATA_COUNT * samples_per_stratum);

    for record in rdr.byte_records() {
        let curr_record = record?;
        let stratum_bytes = curr_record
            .get(strata_column)
            .ok_or_else(|| format!("Strata column index {strata_column} out of bounds"))?;
        if !strata.contains_key(stratum_bytes) {
            strata.insert(
                stratum_bytes.to_vec(),
                (Vec::with_capacity(samples_per_stratum), 0),
            );
        }
        records.push(curr_record);
    }

    if strata.is_empty() {
        return fail_incorrectusage_clierror!("No valid strata found in the data");
    }

    // Create RNG and perform sampling
    with_rng!(rng_kind, seed, |rng| {
        do_stratified_sampling(
            records.into_iter(),
            &mut strata,
            strata_column,
            samples_per_stratum,
            &mut rng,
        )?;
    });

    // Write results in deterministic order
    let mut keys: Vec<_> = strata.keys().collect();
    keys.par_sort_unstable();
    for k in keys {
        if let Some((reservoir, _)) = strata.get(k) {
            for record in reservoir {
                wtr.write_byte_record(record)?;
            }
        }
    }

    Ok(())
}

fn do_stratified_sampling<T: Rng + ?Sized>(
    records: impl Iterator<Item = csv::ByteRecord>,
    strata: &mut HashMap<Vec<u8>, (Vec<csv::ByteRecord>, usize)>,
    strata_column: usize,
    samples_per_stratum: usize,
    rng: &mut T,
) -> CliResult<()> {
    for record in records {
        let stratum_bytes = record
            .get(strata_column)
            .ok_or_else(|| format!("Strata column index {strata_column} out of bounds"))?;

        // One Borrow<[u8]>-based lookup per record, no allocation. The strata
        // map was pre-populated by the caller in its first pass.
        let Some((reservoir, seen)) = strata.get_mut(stratum_bytes) else {
            continue;
        };

        if reservoir.len() < samples_per_stratum {
            reservoir.push(record);
        } else {
            let j = rng.random_range(0..=*seen);
            if j < samples_per_stratum {
                reservoir[j] = record;
            }
        }
        *seen += 1;
    }
    Ok(())
}

// Weighted sampling implementation
fn sample_weighted<R: io::Read, W: io::Write>(
    rconfig: &Config,
    rdr: &mut csv::Reader<R>,
    wtr: &mut csv::Writer<W>,
    weight_column: usize,
    max_weight_stats: Option<f64>,
    sample_size: usize,
    seed: Option<u64>,
    rng_kind: &RngKind,
) -> CliResult<()> {
    let max_weight = if let Some(wt) = max_weight_stats {
        wt
    } else {
        // We don't have a stats cache, do a first pass to find maximum weight
        let mut max_weight_scan = 0.0f64;
        let mut curr_record;
        for record in rdr.byte_records() {
            curr_record = record?;

            let weight: f64 = fast_float2::parse(
                curr_record
                    .get(weight_column)
                    .ok_or_else(|| format!("Weight column index {weight_column} out of bounds"))?,
            )
            .unwrap_or(0.0);

            if weight < 0.0 {
                return fail_incorrectusage_clierror!("Weights must be non-negative: ({weight})");
            }
            max_weight_scan = max_weight_scan.max(weight);
        }
        max_weight_scan
    };

    if max_weight == 0.0 {
        return fail_incorrectusage_clierror!("All weights are zero");
    }

    // Second pass: acceptance-rejection sampling
    let mut rdr2 = rconfig.reader()?;

    log::info!("doing {rng_kind:?} WEIGHTED sampling...");
    with_rng!(rng_kind, seed, |rng| {
        do_weighted_sampling(
            &mut rdr2.byte_records(),
            wtr,
            weight_column,
            sample_size,
            max_weight,
            &mut rng,
        )?;
    });

    Ok(())
}

// Single-pass acceptance-rejection sampling. The caller already determined
// `max_weight` (either from the stats cache or via a prior pass), and the
// records iterator is consumed exactly once — so there is no way to "retry"
// rejected records. If acceptance ends up under-filling the reservoir we
// just warn (the previous outer-loop "retry" was dead code: the iterator
// is unbounded only at the source level, not here).
fn do_weighted_sampling<T: Rng + ?Sized>(
    records: &mut impl Iterator<Item = Result<csv::ByteRecord, csv::Error>>,
    wtr: &mut csv::Writer<impl io::Write>,
    weight_column: usize,
    sample_size: usize,
    max_weight: f64,
    rng: &mut T,
) -> CliResult<()> {
    let mut selected_len = 0;

    for record in records {
        if selected_len >= sample_size {
            break;
        }

        let curr_record = record?;
        let weight: f64 = fast_float2::parse(
            curr_record
                .get(weight_column)
                .ok_or_else(|| format!("Weight column index {weight_column} out of bounds"))?,
        )
        .unwrap_or(0.0);

        if weight < 0.0 {
            return fail_incorrectusage_clierror!("Weights must be non-negative: ({weight})");
        }

        // Skip zero-weight records; otherwise accept with probability w/max_weight.
        if weight == 0.0 {
            continue;
        }
        if rng.random::<f64>() <= (weight / max_weight) {
            wtr.write_byte_record(&curr_record)?;
            selected_len += 1;
        }
    }

    if selected_len < sample_size {
        wwarn!("Could only sample {selected_len} records out of requested {sample_size}");
    }

    Ok(())
}

// Aggregation helper functions for time-series sampling
fn aggregate_numeric_values(values: &[f64], func: AggregationFunction) -> f64 {
    if values.is_empty() {
        return 0.0;
    }

    #[allow(clippy::cast_precision_loss)]
    match func {
        AggregationFunction::First => *values.first().unwrap_or(&0.0),
        AggregationFunction::Last => *values.last().unwrap_or(&0.0),
        AggregationFunction::Mean => {
            let sum: f64 = values.iter().sum();
            sum / values.len() as f64
        },
        AggregationFunction::Sum => values.iter().sum(),
        AggregationFunction::Count => values.len() as f64,
        AggregationFunction::Min => values.iter().copied().fold(f64::INFINITY, f64::min),
        AggregationFunction::Max => values.iter().copied().fold(f64::NEG_INFINITY, f64::max),
        AggregationFunction::Median => {
            let mut sorted = values.to_vec();
            sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
            let mid = sorted.len() / 2;
            if sorted.len().is_multiple_of(2) {
                f64::midpoint(sorted[mid - 1], sorted[mid])
            } else {
                sorted[mid]
            }
        },
    }
}

// Aggregates `records` into a single ByteRecord using `func`.
// `numeric_cols[i]` must be true iff column i is numeric across the WHOLE
// dataset (computed once by the caller). For numeric columns we apply the
// aggregation function over the parsed values; for non-numeric ones we pick
// the first or last raw value depending on the function.
fn aggregate_records(
    records: &[&csv::ByteRecord],
    headers: &csv::ByteRecord,
    func: AggregationFunction,
    numeric_cols: &[bool],
) -> CliResult<csv::ByteRecord> {
    if records.is_empty() {
        return fail_incorrectusage_clierror!("Cannot aggregate empty record set");
    }

    let mut result_fields = Vec::with_capacity(headers.len());

    for col_idx in 0..headers.len() {
        let is_numeric = numeric_cols.get(col_idx).copied().unwrap_or(false);

        if is_numeric {
            // Parse and aggregate. The global numeric scan guarantees every
            // value parses, but we still keep a defensive fallback to 0.0.
            let numeric_values: Vec<f64> = records
                .iter()
                .map(|r| {
                    r.get(col_idx)
                        .and_then(|f| fast_float2::parse::<f64, &[u8]>(f).ok())
                        .unwrap_or(0.0)
                })
                .collect();
            let aggregated = aggregate_numeric_values(&numeric_values, func);
            result_fields.push(aggregated.to_string().into_bytes());
        } else {
            // For non-numeric columns, use first or last based on function
            let value = match func {
                AggregationFunction::First
                | AggregationFunction::Min
                | AggregationFunction::Mean
                | AggregationFunction::Sum
                | AggregationFunction::Count => records[0].get(col_idx).unwrap_or(b"").to_vec(),
                AggregationFunction::Last
                | AggregationFunction::Max
                | AggregationFunction::Median => records[records.len() - 1]
                    .get(col_idx)
                    .unwrap_or(b"")
                    .to_vec(),
            };
            result_fields.push(value);
        }
    }

    Ok(csv::ByteRecord::from(result_fields))
}

// Time-series sampling implementation
fn sample_timeseries<R: io::Read, W: io::Write>(
    _rconfig: &Config,
    rdr: &mut csv::Reader<R>,
    wtr: &mut csv::Writer<W>,
    time_column: usize,
    interval_str: &str,
    start_mode: TSStartMode,
    adaptive_mode: Option<AdaptiveMode>,
    aggregate_func: Option<AggregationFunction>,
    prefer_dmy: bool,
    input_tz: Option<&str>,
    seed: Option<u64>,
    rng_kind: &RngKind,
) -> CliResult<()> {
    let interval = parse_time_interval(interval_str)?;
    let headers = rdr.byte_headers()?.clone();

    // First pass: collect all records with their timestamps
    let mut records_with_times: Vec<(DateTime<Utc>, csv::ByteRecord)> = Vec::new();

    for record_result in rdr.byte_records() {
        let record = record_result?;
        if let Some(time_field) = record.get(time_column) {
            match parse_timestamp(time_field, prefer_dmy, input_tz) {
                Ok(dt) => {
                    records_with_times.push((dt, record));
                },
                Err(e) => {
                    log::warn!("Skipping record with invalid timestamp: {e}");
                },
            }
        } else {
            log::warn!("Skipping record with missing time column");
        }
    }

    if records_with_times.is_empty() {
        return fail_incorrectusage_clierror!("No valid timestamps found in time column");
    }

    // Sort by timestamp - parallel unstable sort for maximum performance
    records_with_times.par_sort_unstable_by(|a, b| a.0.cmp(&b.0));

    // If we'll be aggregating, pre-compute which columns are numeric across
    // the WHOLE dataset (one O(records × cols) scan, instead of repeating
    // the per-column "are all values numeric?" check inside every interval's
    // aggregate_records call).
    let numeric_cols: Vec<bool> = if aggregate_func.is_some() {
        let mut mask = vec![true; headers.len()];
        for (_, record) in &records_with_times {
            for (col_idx, m) in mask.iter_mut().enumerate() {
                if !*m {
                    continue;
                }
                match record.get(col_idx) {
                    Some(field) if fast_float2::parse::<f64, &[u8]>(field).is_ok() => {},
                    _ => *m = false,
                }
            }
            if mask.iter().all(|m| !*m) {
                break;
            }
        }
        mask
    } else {
        Vec::new()
    };

    // Determine starting point
    let start_time = match start_mode {
        TSStartMode::Last => {
            // Start from the last (most recent)
            // safety: we know there are records because we checked above
            records_with_times.last().unwrap().0
        },
        TSStartMode::Random => {
            // Random starting point
            // safety: we know there are records because we checked above
            let earliest = records_with_times.first().unwrap().0;
            let latest = records_with_times.last().unwrap().0;
            let range_secs = (latest - earliest).num_seconds();
            if range_secs > 0 {
                let random_offset =
                    with_rng!(rng_kind, seed, |rng| { rng.random_range(0..=range_secs) });
                earliest + Duration::seconds(random_offset)
            } else {
                earliest
            }
        },
        TSStartMode::First => {
            // Start from earliest
            // safety: we know there are records because we checked above
            records_with_times.first().unwrap().0
        },
    };

    // Group records by time intervals
    let mut interval_groups: HashMap<i64, Vec<(DateTime<Utc>, csv::ByteRecord)>> = HashMap::new();

    for (dt, record) in records_with_times {
        // Calculate which interval this timestamp belongs to
        let elapsed = dt - start_time;
        let interval_num = if elapsed.num_seconds() >= 0 {
            elapsed.num_seconds() / interval.num_seconds()
        } else {
            // Handle negative elapsed time (before start_time)
            (elapsed.num_seconds() - interval.num_seconds() + 1) / interval.num_seconds()
        };

        // Add all records to groups (adaptive filtering happens during selection)
        interval_groups
            .entry(interval_num)
            .or_default()
            .push((dt, record));
    }

    // Sort intervals and process them
    let mut interval_keys: Vec<i64> = interval_groups.keys().copied().collect();
    interval_keys.sort_unstable();

    for interval_key in interval_keys {
        // safety: interval_key is from interval_groups so it exists
        let group = interval_groups.get(&interval_key).unwrap();

        if let Some(agg_func) = aggregate_func {
            // Aggregate records in this interval — pass refs so we don't clone
            // each record into a temporary Vec<ByteRecord>.
            let records_only: Vec<&csv::ByteRecord> = group.iter().map(|(_, r)| r).collect();
            let aggregated = aggregate_records(&records_only, &headers, agg_func, &numeric_cols)?;
            wtr.write_byte_record(&aggregated)?;
        } else {
            // Select one record per interval
            // If adaptive mode is set, prefer records matching criteria
            // Otherwise, select the first record in the interval
            let selected = match adaptive_mode {
                Some(AdaptiveMode::BusinessHours) => {
                    // Prefer business hours records
                    group
                        .iter()
                        .find(|(dt, _)| is_business_hours(dt) && is_business_day(dt))
                        .or_else(|| group.first())
                },
                Some(AdaptiveMode::Weekends) => {
                    // Prefer weekend records
                    group
                        .iter()
                        .find(|(dt, _)| is_weekend(dt))
                        .or_else(|| group.first())
                },
                Some(AdaptiveMode::BusinessDays) => {
                    // Prefer business day records
                    group
                        .iter()
                        .find(|(dt, _)| is_business_day(dt))
                        .or_else(|| group.first())
                },
                Some(AdaptiveMode::Both) => {
                    // Prefer business hours or weekends
                    group
                        .iter()
                        .find(|(dt, _)| {
                            (is_business_hours(dt) && is_business_day(dt)) || is_weekend(dt)
                        })
                        .or_else(|| group.first())
                },
                None => {
                    // Default: first record in interval
                    group.first()
                },
            };

            if let Some((_, record)) = selected {
                wtr.write_byte_record(record)?;
            }
        }
    }

    Ok(())
}

// Cluster sampling implementation
fn sample_cluster<R: io::Read, W: io::Write>(
    rconfig: &Config,
    rdr: &mut csv::Reader<R>,
    wtr: &mut csv::Writer<W>,
    cluster_column: usize,
    cluster_cardinality: Option<u64>,
    requested_clusters: usize,
    seed: Option<u64>,
    rng_kind: &RngKind,
) -> CliResult<()> {
    const ESTIMATED_CLUSTER_COUNT: usize = 100;

    // Pre-allocate for the *number of unique clusters in the file* (cardinality),
    // not for the sample size. requested_clusters is the user-asked sample —
    // sizing the unique-cluster collections to it caused excessive rehashing
    // any time cardinality > requested_clusters.
    let prealloc = if let Some(cardinality) = cluster_cardinality {
        if requested_clusters > cardinality as usize {
            return fail_incorrectusage_clierror!(
                "Requested sample size ({requested_clusters}) exceeds number of clusters \
                 ({cardinality})",
            );
        }
        cardinality as usize
    } else {
        ESTIMATED_CLUSTER_COUNT
    };

    let mut unique_clusters: HashSet<Vec<u8>> = HashSet::with_capacity(prealloc);
    let mut all_clusters: Vec<Vec<u8>> = Vec::with_capacity(prealloc);

    // First pass: collect unique clusters
    for record in rdr.byte_records() {
        let curr_record = record?;
        let cluster = curr_record
            .get(cluster_column)
            .ok_or_else(|| format!("Cluster column index {cluster_column} out of bounds"))?
            .to_vec();

        if unique_clusters.insert(cluster.clone()) {
            all_clusters.push(cluster);
        }
    }

    if unique_clusters.is_empty() {
        return fail_incorrectusage_clierror!("No valid clusters found in the data");
    }

    let take = requested_clusters.min(all_clusters.len());
    let selected_clusters: HashSet<Vec<u8>> = with_rng!(rng_kind, seed, |rng| {
        all_clusters.sample(&mut rng, take).cloned().collect()
    });

    // Second pass: output records from selected clusters
    let mut rdr2 = rconfig.reader()?;
    for record in rdr2.byte_records() {
        let curr_record = record?;
        let cluster = curr_record
            .get(cluster_column)
            .ok_or_else(|| format!("Cluster column index {cluster_column} out of bounds"))?;

        if selected_clusters.contains(cluster) {
            wtr.write_byte_record(&curr_record)?;
        }
    }

    Ok(())
}

// --- Sketch-backed sampling helpers (sample_varopt,
// sample_mergeable_reservoir, sample_sketch_merge) ----------------------------
//
// These drive the input stream through native Rust sampler-state objects
// (`ReservoirItemsSketch`, `VarOptItemsSketch`) defined in `sketches.rs`.
// They are not bindings to Apache DataSketches.

// Drive the input through a `VarOptItemsSketch` (A-ExpJ keying, native Rust)
// keyed by `weight_column`, then emit retained items as CSV. Optionally
// serialize the sketch to `sketch_out`.
fn sample_varopt(
    rdr: &mut csv::Reader<impl io::Read>,
    wtr: &mut csv::Writer<impl io::Write>,
    weight_column: usize,
    sample_size: usize,
    seed: Option<u64>,
    rng_kind: &RngKind,
    sketch_out: Option<&str>,
    header: Option<csv::ByteRecord>,
) -> CliResult<()> {
    if sample_size == 0 {
        return Ok(());
    }
    log::info!(
        "doing {rng_kind:?} VAROPT (A-ExpJ) weighted reservoir sampling with k={sample_size}..."
    );

    let mut sketch: VarOptItemsSketch<csv::ByteRecord> = VarOptItemsSketch::new(sample_size);
    if let Some(h) = header {
        sketch.set_header(h);
    }
    let mut skipped = 0u64;

    with_rng!(rng_kind, seed, |rng| {
        for record in rdr.byte_records() {
            let curr_record = record?;
            let raw = curr_record
                .get(weight_column)
                .ok_or_else(|| format!("Weight column index {weight_column} out of bounds"))?;
            let weight: f64 = fast_float2::parse(raw).unwrap_or(0.0);
            if !weight.is_finite() || weight <= 0.0 {
                skipped += 1;
                continue;
            }
            sketch.update(curr_record, weight, &mut rng);
        }
    });

    if skipped > 0 {
        wwarn!("VAROPT: skipped {skipped} record(s) with non-positive or non-numeric weight");
    }

    let kept = sketch.samples_with_weights();
    let actual = kept.len();
    if actual < sample_size {
        wwarn!("VAROPT: produced {actual} record(s) (requested {sample_size})");
    }
    for (record, _w) in kept {
        wtr.write_byte_record(&record)?;
    }

    if let Some(path) = sketch_out {
        let bytes = sketch
            .to_bytes()
            .map_err(|e| format!("Failed to serialize VarOpt sketch: {e}"))?;
        std::fs::write(path, &bytes)?;
        log::info!("VAROPT: wrote sketch ({} bytes) to {path}", bytes.len());
    }

    Ok(())
}

// Drive the input through a `ReservoirItemsSketch` (Vitter Algorithm R, native
// Rust; same distribution as the default RESERVOIR but with a serializable,
// mergeable state), then emit retained items as CSV. Optionally serialize the
// sketch to `sketch_out`.
fn sample_mergeable_reservoir(
    rdr: &mut csv::Reader<impl io::Read>,
    wtr: &mut csv::Writer<impl io::Write>,
    sample_size: usize,
    seed: Option<u64>,
    rng_kind: &RngKind,
    sketch_out: Option<&str>,
    header: Option<csv::ByteRecord>,
) -> CliResult<()> {
    if sample_size == 0 {
        return Ok(());
    }
    log::info!(
        "doing {rng_kind:?} MERGEABLE-RESERVOIR (Algorithm R) sampling with k={sample_size}..."
    );

    let mut sketch: ReservoirItemsSketch<csv::ByteRecord> = ReservoirItemsSketch::new(sample_size);
    if let Some(h) = header {
        sketch.set_header(h);
    }

    with_rng!(rng_kind, seed, |rng| {
        for record in rdr.byte_records() {
            let curr_record = record?;
            sketch.update(curr_record, &mut rng);
        }
    });

    for record in sketch.samples() {
        wtr.write_byte_record(record)?;
    }

    if let Some(path) = sketch_out {
        let bytes = sketch
            .to_bytes()
            .map_err(|e| format!("Failed to serialize Reservoir sketch: {e}"))?;
        std::fs::write(path, &bytes)?;
        log::info!(
            "MERGEABLE-RESERVOIR: wrote sketch ({} bytes) to {path}",
            bytes.len()
        );
    }

    Ok(())
}

// Resolve `<sample-size>` for sketch-merge mode against an already-retained
// item count (the merged sketch's sample size). Semantics mirror the other
// sample paths: a value in [0, 1) is a *fraction* of the retained items; a
// value >= 1 is an absolute count. Non-finite or non-positive values mean
// "no cap" — emit everything in the sketch.
fn resolve_sketch_cap(sample_size: f64, retained: usize) -> usize {
    if !sample_size.is_finite() || sample_size <= 0.0 {
        return retained;
    }
    #[allow(
        clippy::cast_precision_loss,
        clippy::cast_possible_truncation,
        clippy::cast_sign_loss
    )]
    if sample_size < 1.0 {
        // Fraction of retained items, rounded to nearest. min ensures we
        // never request more rows than the sketch actually holds.
        let n = (retained as f64 * sample_size).round() as usize;
        n.min(retained)
    } else {
        (sample_size as usize).min(retained)
    }
}

// Load and merge a comma-separated list of sketch files (produced by
// --sketch-out), emit the merged sample as CSV, and optionally re-emit the
// merged sketch via --sketch-out. All input sketches must share the same
// family (Reservoir or VarOpt).
fn sample_sketch_merge(
    args: &Args,
    wtr: &mut csv::Writer<impl io::Write>,
    rng_kind: &RngKind,
) -> CliResult<()> {
    let sketch_in = args
        .flag_sketch_in
        .as_deref()
        .ok_or("--sketch-in is required for sketch-merge mode")?;
    let paths: Vec<&str> = sketch_in
        .split(',')
        .map(str::trim)
        .filter(|s| !s.is_empty())
        .collect();
    if paths.is_empty() {
        return fail_incorrectusage_clierror!("--sketch-in requires at least one sketch path");
    }

    // Read first sketch to determine family; then merge the rest.
    let first_bytes = std::fs::read(paths[0])
        .map_err(|e| format!("Failed to read sketch file {}: {e}", paths[0]))?;
    let family = peek_family(&first_bytes)
        .map_err(|e| format!("Failed to identify sketch family in {}: {e}", paths[0]))?;
    log::info!(
        "SKETCH-MERGE: merging {} {:?} sketch(es) with {rng_kind:?} RNG",
        paths.len(),
        family
    );

    match family {
        SketchFamily::Reservoir => {
            let mut merged: ReservoirItemsSketch<csv::ByteRecord> =
                ReservoirItemsSketch::from_bytes(&first_bytes).map_err(|e| {
                    format!("Failed to load Reservoir sketch from {}: {e}", paths[0])
                })?;
            with_rng!(rng_kind, args.flag_seed, |rng| {
                for p in &paths[1..] {
                    let bytes = std::fs::read(p)
                        .map_err(|e| format!("Failed to read sketch file {p}: {e}"))?;
                    let other_family = peek_family(&bytes)
                        .map_err(|e| format!("Failed to identify sketch family in {p}: {e}"))?;
                    if other_family != SketchFamily::Reservoir {
                        return fail_incorrectusage_clierror!(
                            "Sketch family mismatch: {} is Reservoir but {p} is {:?}",
                            paths[0],
                            other_family
                        );
                    }
                    let other = ReservoirItemsSketch::<csv::ByteRecord>::from_bytes(&bytes)
                        .map_err(|e| format!("Failed to load Reservoir sketch from {p}: {e}"))?;
                    merged.merge(other, &mut rng);
                }
                Ok::<(), crate::CliError>(())
            })?;

            // Emit the embedded header first (if present and --no-headers not set).
            if !args.flag_no_headers
                && let Some(h) = merged.header()
            {
                wtr.write_byte_record(h)?;
            }

            // Optionally cap the emitted sample below the sketch's k.
            // 0 < size < 1 means "fraction of the merged sketch's retained
            // items"; size >= 1 means "this many rows".
            let cap = resolve_sketch_cap(args.arg_sample_size, merged.samples().len());
            for (i, record) in merged.samples().iter().enumerate() {
                if i >= cap {
                    break;
                }
                wtr.write_byte_record(record)?;
            }

            if let Some(path) = &args.flag_sketch_out {
                let bytes = merged
                    .to_bytes()
                    .map_err(|e| format!("Failed to serialize merged Reservoir sketch: {e}"))?;
                std::fs::write(path, &bytes)?;
                log::info!(
                    "SKETCH-MERGE: wrote merged sketch ({} bytes) to {path}",
                    bytes.len()
                );
            }
        },
        SketchFamily::VarOpt => {
            let mut merged: VarOptItemsSketch<csv::ByteRecord> =
                VarOptItemsSketch::from_bytes(&first_bytes)
                    .map_err(|e| format!("Failed to load VarOpt sketch from {}: {e}", paths[0]))?;
            with_rng!(rng_kind, args.flag_seed, |rng| {
                for p in &paths[1..] {
                    let bytes = std::fs::read(p)
                        .map_err(|e| format!("Failed to read sketch file {p}: {e}"))?;
                    let other_family = peek_family(&bytes)
                        .map_err(|e| format!("Failed to identify sketch family in {p}: {e}"))?;
                    if other_family != SketchFamily::VarOpt {
                        return fail_incorrectusage_clierror!(
                            "Sketch family mismatch: {} is VarOpt but {p} is {:?}",
                            paths[0],
                            other_family
                        );
                    }
                    let other = VarOptItemsSketch::<csv::ByteRecord>::from_bytes(&bytes)
                        .map_err(|e| format!("Failed to load VarOpt sketch from {p}: {e}"))?;
                    merged.merge(other, &mut rng);
                }
                Ok::<(), crate::CliError>(())
            })?;

            // Emit the embedded header first (if present and --no-headers not set).
            if !args.flag_no_headers
                && let Some(h) = merged.header()
            {
                wtr.write_byte_record(h)?;
            }

            // Same fractional / count semantics as the Reservoir branch above.
            let cap = resolve_sketch_cap(args.arg_sample_size, merged.samples_with_weights().len());
            for (i, (record, _w)) in merged.samples_with_weights().into_iter().enumerate() {
                if i >= cap {
                    break;
                }
                wtr.write_byte_record(&record)?;
            }

            if let Some(path) = &args.flag_sketch_out {
                let bytes = merged
                    .to_bytes()
                    .map_err(|e| format!("Failed to serialize merged VarOpt sketch: {e}"))?;
                std::fs::write(path, &bytes)?;
                log::info!(
                    "SKETCH-MERGE: wrote merged sketch ({} bytes) to {path}",
                    bytes.len()
                );
            }
        },
    }

    Ok(())
}
