//! Timestamp alignment for multi-asset backtests.
//!
//! When multiple symbols are backtested together, each may have a different
//! set of timestamps (gaps, different listing dates). This module builds a
//! union time index and forward-fills each symbol's bars to match.
//!
//! **Optimisations** (vs. the naïve BTreeSet + HashMap approach):
//! 1. *Identical-timestamps fast path* — when all symbols share the exact same
//!    timestamp array we skip alignment entirely (common for same-exchange data).
//! 2. *Sorted k-way merge* — symbol bars arrive sorted, so the union is built
//!    in O(N) via merge instead of O(N log N) BTreeSet insertion.
//! 3. *Two-pointer mapping* — forward-fill indices are resolved with a single
//!    linear scan instead of per-row HashMap lookups.

use std::collections::HashMap;
use std::sync::Arc;

use arrow::array::{
    Array, ArrayRef, BooleanArray, Float64Array, TimestampNanosecondArray, UInt32Array, UInt8Array,
};
use arrow::datatypes::{DataType, TimeUnit};
use arrow::record_batch::RecordBatch;
use bt_kernel::{BtError, Result, SymbolId};

/// Result of aligning multiple symbol bars to a common timestamp grid.
pub struct AlignedBars {
    /// Sorted union of all symbol timestamps (nanoseconds).
    pub master_timestamps: Vec<i64>,
    /// Per-symbol bars forward-filled to match `master_timestamps`.
    /// Every batch has exactly `master_timestamps.len()` rows.
    pub symbol_bars: HashMap<SymbolId, RecordBatch>,
}

/// Align bars from multiple symbols to a common timestamp grid.
///
/// 1. Builds the union of all timestamps across symbols.
/// 2. For each symbol, forward-fills missing bars from the last known row.
/// 3. Before a symbol's first bar, values are filled with NaN (float), 0 (int), false (bool).
///
/// For a single symbol, this is essentially a no-op (returns the same data).
pub fn align_symbol_bars(raw_bars: HashMap<SymbolId, RecordBatch>) -> Result<AlignedBars> {
    if raw_bars.is_empty() {
        return Err(BtError::Data("no symbol bars to align".to_string()));
    }

    // Fast path: single symbol — no alignment needed.
    if raw_bars.len() == 1 {
        let (symbol, batch) = raw_bars.into_iter().next().unwrap();
        let ts = extract_timestamps(&batch)?;
        return Ok(AlignedBars {
            master_timestamps: ts,
            symbol_bars: HashMap::from([(symbol, batch)]),
        });
    }

    // Fast path: all symbols share identical timestamp arrays.
    if timestamps_identical(&raw_bars)? {
        tracing::debug!("alignment fast-path: all symbols share identical timestamps");
        let ts = {
            let first_batch = raw_bars.values().next().unwrap();
            extract_timestamps(first_batch)?
        };
        let symbol_bars: HashMap<SymbolId, RecordBatch> = raw_bars;
        return Ok(AlignedBars {
            master_timestamps: ts,
            symbol_bars,
        });
    }

    // 1. Build union via sorted merge (timestamps are pre-sorted from parquet).
    let ts_vecs: Vec<Vec<i64>> = raw_bars
        .values()
        .map(extract_timestamps)
        .collect::<Result<_>>()?;
    let master_timestamps = sorted_merge_dedup(&ts_vecs);
    let master_len = master_timestamps.len();

    // 2. For each symbol, build forward-fill mapping via two-pointer scan.
    let mut aligned = HashMap::new();
    for (symbol, batch) in &raw_bars {
        let sym_ts = extract_timestamps(batch)?;
        let mapping = two_pointer_forward_fill(&master_timestamps, &sym_ts);
        let aligned_batch = forward_fill_batch(batch, &master_timestamps, &mapping)?;
        aligned.insert(*symbol, aligned_batch);
    }

    tracing::debug!(
        master_len,
        symbols = raw_bars.len(),
        "alignment: sorted merge + two-pointer fill"
    );

    Ok(AlignedBars {
        master_timestamps,
        symbol_bars: aligned,
    })
}

/// Slice aligned bars to a sub-range of master timestamps.
///
/// Uses binary search to find the inclusive range `[start, end)` within
/// `master_timestamps`, then slices all symbol batches accordingly.
/// This avoids reloading data from disk when splitting into train/test folds.
pub fn slice_aligned_bars(
    aligned: &AlignedBars,
    range: bt_kernel::TimeRange,
) -> Result<AlignedBars> {
    let master = &aligned.master_timestamps;
    if master.is_empty() {
        return Err(BtError::Data("cannot slice empty aligned bars".to_string()));
    }

    // Binary search for [start, end) within master_timestamps.
    let start_idx = master.partition_point(|&t| t < range.start.as_nanos());
    let end_idx = master.partition_point(|&t| t < range.end.as_nanos());

    if start_idx >= end_idx {
        return Err(BtError::Data(format!(
            "slice_aligned_bars: no bars in range [{}, {})",
            range.start.as_nanos(),
            range.end.as_nanos()
        )));
    }

    let new_master = master[start_idx..end_idx].to_vec();
    let mut new_symbol_bars = HashMap::new();

    for (&symbol, batch) in &aligned.symbol_bars {
        let sliced = batch.slice(start_idx, end_idx - start_idx);
        new_symbol_bars.insert(symbol, sliced);
    }

    Ok(AlignedBars {
        master_timestamps: new_master,
        symbol_bars: new_symbol_bars,
    })
}

/// Check if all symbol batches have identical timestamp arrays.
fn timestamps_identical(raw_bars: &HashMap<SymbolId, RecordBatch>) -> Result<bool> {
    let mut iter = raw_bars.values();
    let first = iter.next().unwrap();
    let first_ts = timestamp_array(first)?;
    let first_len = first_ts.len();

    for batch in iter {
        let ts = timestamp_array(batch)?;
        if ts.len() != first_len {
            return Ok(false);
        }
        // Compare raw i64 slices directly — O(n) memcmp-friendly.
        let a = first_ts.values();
        let b = ts.values();
        if a.as_ref() != b.as_ref() {
            return Ok(false);
        }
    }
    Ok(true)
}

/// Merge k sorted, deduplicated timestamp vectors into one sorted, deduplicated vector.
fn sorted_merge_dedup(vecs: &[Vec<i64>]) -> Vec<i64> {
    if vecs.len() == 1 {
        return vecs[0].clone();
    }

    // Two-way merge (handles 2 symbols directly; generalises by folding).
    let mut result = vecs[0].clone();
    for other in &vecs[1..] {
        result = merge_two_sorted(&result, other);
    }
    result
}

/// Merge two sorted, deduplicated slices into one sorted, deduplicated Vec.
fn merge_two_sorted(a: &[i64], b: &[i64]) -> Vec<i64> {
    let mut out = Vec::with_capacity(a.len().max(b.len()));
    let (mut i, mut j) = (0, 0);
    while i < a.len() && j < b.len() {
        match a[i].cmp(&b[j]) {
            std::cmp::Ordering::Less => {
                out.push(a[i]);
                i += 1;
            }
            std::cmp::Ordering::Greater => {
                out.push(b[j]);
                j += 1;
            }
            std::cmp::Ordering::Equal => {
                out.push(a[i]);
                i += 1;
                j += 1;
            }
        }
    }
    if i < a.len() {
        out.extend_from_slice(&a[i..]);
    }
    if j < b.len() {
        out.extend_from_slice(&b[j..]);
    }
    out
}

/// Build forward-fill mapping using a two-pointer scan.
///
/// Both `master` and `symbol_ts` must be sorted. Returns a Vec of length
/// `master.len()` where each entry is `Some(source_row)` if the symbol has
/// data at or before that master timestamp, or `None` if the symbol is not
/// yet active.
fn two_pointer_forward_fill(master: &[i64], symbol_ts: &[i64]) -> Vec<Option<usize>> {
    let mut mapping = Vec::with_capacity(master.len());
    let mut sym_idx: usize = 0;
    let mut last_row: Option<usize> = None;

    for &t in master {
        // Advance sym_idx while the symbol's next timestamp <= t.
        while sym_idx < symbol_ts.len() && symbol_ts[sym_idx] <= t {
            last_row = Some(sym_idx);
            sym_idx += 1;
        }
        mapping.push(last_row);
    }
    mapping
}

/// Create a new RecordBatch with `master_timestamps` as the timestamp column,
/// and all other columns forward-filled according to `mapping`.
fn forward_fill_batch(
    source: &RecordBatch,
    master_timestamps: &[i64],
    mapping: &[Option<usize>],
) -> Result<RecordBatch> {
    let schema = source.schema();
    let master_len = master_timestamps.len();
    let mut columns: Vec<ArrayRef> = Vec::with_capacity(schema.fields().len());

    for (col_idx, field) in schema.fields().iter().enumerate() {
        let source_col = source.column(col_idx);

        if field.name() == "timestamp" {
            // Replace with master timestamps.
            let ts =
                TimestampNanosecondArray::from(master_timestamps.to_vec()).with_timezone("UTC");
            columns.push(Arc::new(ts));
            continue;
        }

        let filled = forward_fill_column(
            source_col,
            mapping,
            master_len,
            field.data_type(),
            field.is_nullable(),
        )?;
        columns.push(filled);
    }

    RecordBatch::try_new(schema, columns).map_err(Into::into)
}

/// Forward-fill a single column according to the mapping.
///
/// When `nullable` is false, unmapped rows use a default value (0 / 0.0 / false)
/// instead of null to satisfy the Arrow schema constraint.
fn forward_fill_column(
    source: &ArrayRef,
    mapping: &[Option<usize>],
    len: usize,
    data_type: &DataType,
    nullable: bool,
) -> Result<ArrayRef> {
    match data_type {
        DataType::Float64 => {
            let src = source
                .as_any()
                .downcast_ref::<Float64Array>()
                .ok_or_else(|| BtError::Data("expected Float64Array".to_string()))?;
            if nullable {
                let values: Vec<Option<f64>> = mapping
                    .iter()
                    .map(|opt| opt.map(|row| src.value(row)))
                    .collect();
                Ok(Arc::new(Float64Array::from(values)))
            } else {
                let values: Vec<f64> = mapping
                    .iter()
                    .map(|opt| opt.map(|row| src.value(row)).unwrap_or(f64::NAN))
                    .collect();
                Ok(Arc::new(Float64Array::from(values)))
            }
        }
        DataType::UInt32 => {
            let src = source
                .as_any()
                .downcast_ref::<UInt32Array>()
                .ok_or_else(|| BtError::Data("expected UInt32Array".to_string()))?;
            if nullable {
                let values: Vec<Option<u32>> = mapping
                    .iter()
                    .map(|opt| opt.map(|row| src.value(row)))
                    .collect();
                Ok(Arc::new(UInt32Array::from(values)))
            } else {
                let values: Vec<u32> = mapping
                    .iter()
                    .map(|opt| opt.map(|row| src.value(row)).unwrap_or(0))
                    .collect();
                Ok(Arc::new(UInt32Array::from(values)))
            }
        }
        DataType::UInt8 => {
            let src = source
                .as_any()
                .downcast_ref::<UInt8Array>()
                .ok_or_else(|| BtError::Data("expected UInt8Array".to_string()))?;
            if nullable {
                let values: Vec<Option<u8>> = mapping
                    .iter()
                    .map(|opt| opt.map(|row| src.value(row)))
                    .collect();
                Ok(Arc::new(UInt8Array::from(values)))
            } else {
                let values: Vec<u8> = mapping
                    .iter()
                    .map(|opt| opt.map(|row| src.value(row)).unwrap_or(0))
                    .collect();
                Ok(Arc::new(UInt8Array::from(values)))
            }
        }
        DataType::Boolean => {
            let src = source
                .as_any()
                .downcast_ref::<BooleanArray>()
                .ok_or_else(|| BtError::Data("expected BooleanArray".to_string()))?;
            if nullable {
                let values: Vec<Option<bool>> = mapping
                    .iter()
                    .map(|opt| opt.map(|row| src.value(row)))
                    .collect();
                Ok(Arc::new(BooleanArray::from(values)))
            } else {
                let values: Vec<bool> = mapping
                    .iter()
                    .map(|opt| opt.map(|row| src.value(row)).unwrap_or(false))
                    .collect();
                Ok(Arc::new(BooleanArray::from(values)))
            }
        }
        DataType::Timestamp(TimeUnit::Nanosecond, tz) => {
            let src = source
                .as_any()
                .downcast_ref::<TimestampNanosecondArray>()
                .ok_or_else(|| BtError::Data("expected TimestampNanosecondArray".to_string()))?;
            if nullable {
                let values: Vec<Option<i64>> = mapping
                    .iter()
                    .map(|opt| opt.map(|row| src.value(row)))
                    .collect();
                let mut arr = TimestampNanosecondArray::from(values);
                if let Some(tz) = tz {
                    arr = arr.with_timezone(tz.as_ref());
                }
                Ok(Arc::new(arr))
            } else {
                let values: Vec<i64> = mapping
                    .iter()
                    .map(|opt| opt.map(|row| src.value(row)).unwrap_or(0))
                    .collect();
                let mut arr = TimestampNanosecondArray::from(values);
                if let Some(tz) = tz {
                    arr = arr.with_timezone(tz.as_ref());
                }
                Ok(Arc::new(arr))
            }
        }
        _ => {
            // Generic fallback: fill with nulls for unsupported types.
            let nulls = arrow::array::new_null_array(data_type, len);
            Ok(nulls)
        }
    }
}

fn timestamp_array(batch: &RecordBatch) -> Result<&TimestampNanosecondArray> {
    batch
        .column_by_name("timestamp")
        .ok_or_else(|| BtError::Data("missing timestamp column".to_string()))?
        .as_any()
        .downcast_ref::<TimestampNanosecondArray>()
        .ok_or_else(|| BtError::Data("timestamp must be TimestampNanosecondArray".to_string()))
}

fn extract_timestamps(batch: &RecordBatch) -> Result<Vec<i64>> {
    let ts = timestamp_array(batch)?;
    Ok((0..ts.len()).map(|i| ts.value(i)).collect())
}

// ---------------------------------------------------------------------------
// OHLCV Resampling
// ---------------------------------------------------------------------------

/// Resample aligned bars to a coarser time resolution.
///
/// Groups rows into time buckets of `target_interval_ns` nanoseconds and
/// aggregates OHLCV columns appropriately (open=first, high=max, low=min,
/// close=last, volume=sum, vwap=weighted, etc.).
///
/// Returns the original `AlignedBars` unchanged if the native resolution is
/// already >= the target, or if the data has fewer than 2 rows.
pub fn resample_aligned_bars(aligned: AlignedBars, target_interval_ns: i64) -> Result<AlignedBars> {
    let master = &aligned.master_timestamps;
    if master.len() < 2 {
        return Ok(aligned);
    }

    // Detect native resolution from median diff of first few timestamps.
    let native_ns = master[1] - master[0];
    if target_interval_ns <= native_ns {
        return Ok(aligned);
    }

    // Compute resample groups: (start, end) inclusive for each time bucket.
    let groups = compute_resample_groups(master, target_interval_ns);
    let new_len = groups.len();

    // Build new master timestamps (last timestamp of each bucket).
    let new_master: Vec<i64> = groups.iter().map(|&(_, end)| master[end]).collect();

    // Resample each symbol's bars.
    let mut new_symbol_bars = HashMap::new();
    for (symbol, batch) in aligned.symbol_bars {
        let resampled = resample_batch(&batch, &groups, new_len)?;
        new_symbol_bars.insert(symbol, resampled);
    }

    Ok(AlignedBars {
        master_timestamps: new_master,
        symbol_bars: new_symbol_bars,
    })
}

/// Compute (start, end) index pairs for time buckets.
pub fn compute_resample_groups(timestamps: &[i64], interval_ns: i64) -> Vec<(usize, usize)> {
    let mut groups = Vec::new();
    if timestamps.is_empty() {
        return groups;
    }
    let mut start = 0;
    let mut current_bucket = timestamps[0] / interval_ns;
    for i in 1..timestamps.len() {
        let bucket = timestamps[i] / interval_ns;
        if bucket != current_bucket {
            groups.push((start, i - 1));
            start = i;
            current_bucket = bucket;
        }
    }
    groups.push((start, timestamps.len() - 1));
    groups
}

/// Resample a single RecordBatch using pre-computed groups.
pub fn resample_batch(
    batch: &RecordBatch,
    groups: &[(usize, usize)],
    new_len: usize,
) -> Result<RecordBatch> {
    let schema = batch.schema();
    let mut columns: Vec<ArrayRef> = Vec::with_capacity(schema.fields().len());

    // Pre-extract volume column for vwap weighting.
    let volume_col = batch
        .column_by_name("volume")
        .and_then(|c| c.as_any().downcast_ref::<Float64Array>());

    for (col_idx, field) in schema.fields().iter().enumerate() {
        let col = batch.column(col_idx);
        let resampled = resample_column(
            col,
            field.name(),
            field.data_type(),
            field.is_nullable(),
            groups,
            new_len,
            volume_col,
        )?;
        columns.push(resampled);
    }

    RecordBatch::try_new(schema, columns).map_err(Into::into)
}

/// Resample a single column according to its name and type.
fn resample_column(
    col: &ArrayRef,
    name: &str,
    _data_type: &DataType,
    nullable: bool,
    groups: &[(usize, usize)],
    new_len: usize,
    volume_col: Option<&Float64Array>,
) -> Result<ArrayRef> {
    match name {
        // Timestamp: last of bucket
        "timestamp" => {
            let src = col
                .as_any()
                .downcast_ref::<TimestampNanosecondArray>()
                .ok_or_else(|| BtError::Data("expected TimestampNanosecondArray".to_string()))?;
            let values: Vec<i64> = groups.iter().map(|&(_, end)| src.value(end)).collect();
            Ok(Arc::new(
                TimestampNanosecondArray::from(values).with_timezone("UTC"),
            ))
        }

        // Symbol ID: first of bucket (constant within symbol)
        "symbol_id" => {
            let src = col
                .as_any()
                .downcast_ref::<UInt32Array>()
                .ok_or_else(|| BtError::Data("expected UInt32Array for symbol_id".to_string()))?;
            let values: Vec<u32> = groups.iter().map(|&(start, _)| src.value(start)).collect();
            Ok(Arc::new(UInt32Array::from(values)))
        }

        // OHLC
        "open" => agg_f64_first(col, groups),
        "high" => agg_f64_max(col, groups),
        "low" => agg_f64_min(col, groups),
        "close" | "bid" | "ask" | "spread" => agg_f64_last(col, groups, nullable),

        // VWAP: volume-weighted average
        "vwap" => {
            let src = col
                .as_any()
                .downcast_ref::<Float64Array>()
                .ok_or_else(|| BtError::Data("expected Float64Array for vwap".to_string()))?;
            let vol = volume_col.ok_or_else(|| {
                BtError::Data("volume column needed for vwap resampling".to_string())
            })?;
            let mut values = Vec::with_capacity(new_len);
            for &(start, end) in groups {
                let mut wsum = 0.0_f64;
                let mut vsum = 0.0_f64;
                for i in start..=end {
                    let v = vol.value(i);
                    wsum += src.value(i) * v;
                    vsum += v;
                }
                values.push(if vsum > 0.0 {
                    wsum / vsum
                } else {
                    src.value(end)
                });
            }
            Ok(Arc::new(Float64Array::from(values)))
        }

        // Summable Float64 columns
        "volume" | "buy_volume" | "sell_volume" => agg_f64_sum(col, groups, nullable),

        // Trade count: sum
        "trade_count" => {
            let src = col
                .as_any()
                .downcast_ref::<UInt32Array>()
                .ok_or_else(|| BtError::Data("expected UInt32Array for trade_count".to_string()))?;
            let values: Vec<u32> = groups
                .iter()
                .map(|&(start, end)| {
                    let mut sum = 0u32;
                    for i in start..=end {
                        sum = sum.saturating_add(src.value(i));
                    }
                    sum
                })
                .collect();
            Ok(Arc::new(UInt32Array::from(values)))
        }

        // is_gap: true only if ALL bars in bucket are gaps
        "is_gap" => {
            let src = col
                .as_any()
                .downcast_ref::<BooleanArray>()
                .ok_or_else(|| BtError::Data("expected BooleanArray for is_gap".to_string()))?;
            let values: Vec<bool> = groups
                .iter()
                .map(|&(start, end)| (start..=end).all(|i| src.value(i)))
                .collect();
            Ok(Arc::new(BooleanArray::from(values)))
        }

        // gap_fill_method: first of bucket
        "gap_fill_method" => {
            let src = col.as_any().downcast_ref::<UInt8Array>().ok_or_else(|| {
                BtError::Data("expected UInt8Array for gap_fill_method".to_string())
            })?;
            let values: Vec<u8> = groups.iter().map(|&(start, _)| src.value(start)).collect();
            Ok(Arc::new(UInt8Array::from(values)))
        }

        // Fallback: take last value
        _ => agg_f64_last(col, groups, nullable),
    }
}

// --- Float64 aggregation helpers ---

fn agg_f64_first(col: &ArrayRef, groups: &[(usize, usize)]) -> Result<ArrayRef> {
    let src = col
        .as_any()
        .downcast_ref::<Float64Array>()
        .ok_or_else(|| BtError::Data("expected Float64Array".to_string()))?;
    let values: Vec<f64> = groups.iter().map(|&(start, _)| src.value(start)).collect();
    Ok(Arc::new(Float64Array::from(values)))
}

fn agg_f64_last(col: &ArrayRef, groups: &[(usize, usize)], nullable: bool) -> Result<ArrayRef> {
    let src = col
        .as_any()
        .downcast_ref::<Float64Array>()
        .ok_or_else(|| BtError::Data("expected Float64Array".to_string()))?;
    if nullable {
        let values: Vec<Option<f64>> = groups
            .iter()
            .map(|&(_, end)| {
                if src.is_null(end) {
                    None
                } else {
                    Some(src.value(end))
                }
            })
            .collect();
        Ok(Arc::new(Float64Array::from(values)))
    } else {
        let values: Vec<f64> = groups.iter().map(|&(_, end)| src.value(end)).collect();
        Ok(Arc::new(Float64Array::from(values)))
    }
}

fn agg_f64_max(col: &ArrayRef, groups: &[(usize, usize)]) -> Result<ArrayRef> {
    let src = col
        .as_any()
        .downcast_ref::<Float64Array>()
        .ok_or_else(|| BtError::Data("expected Float64Array".to_string()))?;
    let values: Vec<f64> = groups
        .iter()
        .map(|&(start, end)| {
            let mut max = f64::NEG_INFINITY;
            for i in start..=end {
                let v = src.value(i);
                if v > max {
                    max = v;
                }
            }
            max
        })
        .collect();
    Ok(Arc::new(Float64Array::from(values)))
}

fn agg_f64_min(col: &ArrayRef, groups: &[(usize, usize)]) -> Result<ArrayRef> {
    let src = col
        .as_any()
        .downcast_ref::<Float64Array>()
        .ok_or_else(|| BtError::Data("expected Float64Array".to_string()))?;
    let values: Vec<f64> = groups
        .iter()
        .map(|&(start, end)| {
            let mut min = f64::INFINITY;
            for i in start..=end {
                let v = src.value(i);
                if v < min {
                    min = v;
                }
            }
            min
        })
        .collect();
    Ok(Arc::new(Float64Array::from(values)))
}

fn agg_f64_sum(col: &ArrayRef, groups: &[(usize, usize)], nullable: bool) -> Result<ArrayRef> {
    let src = col
        .as_any()
        .downcast_ref::<Float64Array>()
        .ok_or_else(|| BtError::Data("expected Float64Array".to_string()))?;
    if nullable {
        let values: Vec<Option<f64>> = groups
            .iter()
            .map(|&(start, end)| {
                let mut sum = 0.0_f64;
                let mut all_null = true;
                for i in start..=end {
                    if !src.is_null(i) {
                        sum += src.value(i);
                        all_null = false;
                    }
                }
                if all_null {
                    None
                } else {
                    Some(sum)
                }
            })
            .collect();
        Ok(Arc::new(Float64Array::from(values)))
    } else {
        let values: Vec<f64> = groups
            .iter()
            .map(|&(start, end)| {
                let mut sum = 0.0_f64;
                for i in start..=end {
                    sum += src.value(i);
                }
                sum
            })
            .collect();
        Ok(Arc::new(Float64Array::from(values)))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use arrow::datatypes::{Field, Schema};

    fn simple_batch(symbol_id: u32, timestamps: &[i64], closes: &[f64]) -> RecordBatch {
        let schema = Arc::new(Schema::new(vec![
            Field::new(
                "timestamp",
                DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())),
                false,
            ),
            Field::new("symbol_id", DataType::UInt32, false),
            Field::new("close", DataType::Float64, false),
        ]));

        RecordBatch::try_new(
            schema,
            vec![
                Arc::new(TimestampNanosecondArray::from(timestamps.to_vec()).with_timezone("UTC")),
                Arc::new(UInt32Array::from(vec![symbol_id; timestamps.len()])),
                Arc::new(Float64Array::from(closes.to_vec())),
            ],
        )
        .unwrap()
    }

    #[test]
    fn single_symbol_passthrough() {
        let bars = HashMap::from([(
            SymbolId(1),
            simple_batch(1, &[1000, 2000, 3000], &[100.0, 101.0, 102.0]),
        )]);

        let aligned = align_symbol_bars(bars).unwrap();
        assert_eq!(aligned.master_timestamps, vec![1000, 2000, 3000]);
        assert_eq!(aligned.symbol_bars.len(), 1);
        assert_eq!(aligned.symbol_bars[&SymbolId(1)].num_rows(), 3);
    }

    #[test]
    fn two_symbols_same_timestamps() {
        let bars = HashMap::from([
            (
                SymbolId(1),
                simple_batch(1, &[1000, 2000, 3000], &[100.0, 101.0, 102.0]),
            ),
            (
                SymbolId(2),
                simple_batch(2, &[1000, 2000, 3000], &[200.0, 198.0, 202.0]),
            ),
        ]);

        let aligned = align_symbol_bars(bars).unwrap();
        assert_eq!(aligned.master_timestamps, vec![1000, 2000, 3000]);
        assert_eq!(aligned.symbol_bars[&SymbolId(1)].num_rows(), 3);
        assert_eq!(aligned.symbol_bars[&SymbolId(2)].num_rows(), 3);
    }

    #[test]
    fn two_symbols_different_timestamps_forward_fill() {
        // Symbol 1 has bars at [1000, 2000, 3000]
        // Symbol 2 has bars at [1000, 3000] (missing 2000)
        let bars = HashMap::from([
            (
                SymbolId(1),
                simple_batch(1, &[1000, 2000, 3000], &[100.0, 101.0, 102.0]),
            ),
            (SymbolId(2), simple_batch(2, &[1000, 3000], &[200.0, 202.0])),
        ]);

        let aligned = align_symbol_bars(bars).unwrap();
        assert_eq!(aligned.master_timestamps, vec![1000, 2000, 3000]);

        // Symbol 2 should be forward-filled at timestamp 2000
        let s2_close = aligned.symbol_bars[&SymbolId(2)]
            .column_by_name("close")
            .unwrap()
            .as_any()
            .downcast_ref::<Float64Array>()
            .unwrap();
        assert_eq!(s2_close.value(0), 200.0); // t=1000
        assert_eq!(s2_close.value(1), 200.0); // t=2000 (forward-filled from t=1000)
        assert_eq!(s2_close.value(2), 202.0); // t=3000
    }

    #[test]
    fn symbol_not_yet_active_filled_with_default() {
        // Symbol 1 starts at 1000, Symbol 2 starts at 2000
        let bars = HashMap::from([
            (
                SymbolId(1),
                simple_batch(1, &[1000, 2000, 3000], &[100.0, 101.0, 102.0]),
            ),
            (SymbolId(2), simple_batch(2, &[2000, 3000], &[200.0, 202.0])),
        ]);

        let aligned = align_symbol_bars(bars).unwrap();
        assert_eq!(aligned.master_timestamps, vec![1000, 2000, 3000]);

        let s2_close = aligned.symbol_bars[&SymbolId(2)]
            .column_by_name("close")
            .unwrap()
            .as_any()
            .downcast_ref::<Float64Array>()
            .unwrap();
        // At t=1000, symbol 2 is not yet active → NaN (non-nullable Float64)
        assert!(s2_close.value(0).is_nan());
        assert_eq!(s2_close.value(1), 200.0);
        assert_eq!(s2_close.value(2), 202.0);

        // Non-nullable UInt32 symbol_id → filled with 0 before first bar
        let s2_sid = aligned.symbol_bars[&SymbolId(2)]
            .column_by_name("symbol_id")
            .unwrap()
            .as_any()
            .downcast_ref::<UInt32Array>()
            .unwrap();
        assert_eq!(s2_sid.value(0), 0); // before first bar → default 0
        assert_eq!(s2_sid.value(1), 2);
        assert_eq!(s2_sid.value(2), 2);
    }
}
