virtual-frame 0.1.1

//! virtual-frame — Deterministic data pipeline toolkit for LLM training.
//!
//! Bitmask-filtered virtual views, NFA regex, Kahan summation, NLP primitives,
//! CSV ingestion, and a deterministic RNG. Python bindings via PyO3.

pub mod bitmask;
pub mod column;
pub mod csv;
pub mod dataframe;
pub mod expr;
pub mod kahan;
pub mod nlp;
pub mod regex_engine;
pub mod rng;
pub mod tidyview;

// ── PyO3 Python bindings ──────────────────────────────────────────────────

use pyo3::prelude::*;
use pyo3::exceptions::{PyRuntimeError, PyValueError};
use pyo3::types::{PyDict, PyList};

// ── Python wrappers ───────────────────────────────────────────────────────

/// Python-visible DataFrame wrapper.
#[pyclass(name = "DataFrame")]
#[derive(Clone)]
struct PyDataFrame {
    inner: dataframe::DataFrame,
}

#[pymethods]
impl PyDataFrame {
    /// Create a DataFrame from a dict of column_name → list.
    #[new]
    fn new(columns: &Bound<'_, PyDict>) -> PyResult<Self> {
        let mut cols: Vec<(String, column::Column)> = Vec::new();
        for (key, value) in columns.iter() {
            let name: String = key.extract()?;
            let list = value.downcast::<PyList>()?;
            let col = py_list_to_column(list)?;
            cols.push((name, col));
        }
        let df = dataframe::DataFrame::from_columns(cols)
            .map_err(|e| PyValueError::new_err(format!("{}", e)))?;
        Ok(PyDataFrame { inner: df })
    }

    /// Number of rows.
    fn nrows(&self) -> usize {
        self.inner.nrows()
    }

    /// Number of columns.
    fn ncols(&self) -> usize {
        self.inner.ncols()
    }

    /// Column names.
    fn column_names(&self) -> Vec<String> {
        self.inner.column_names().into_iter().map(|s| s.to_string()).collect()
    }

    /// Get a column as a Python list.
    fn get_column(&self, name: &str) -> PyResult<PyObject> {
        let col = self.inner.get_column(name)
            .ok_or_else(|| PyValueError::new_err(format!("column `{}` not found", name)))?;
        Python::with_gil(|py| column_to_py(py, col))
    }

    fn __repr__(&self) -> String {
        format!("DataFrame(nrows={}, ncols={}, columns={:?})",
            self.inner.nrows(), self.inner.ncols(), self.inner.column_names())
    }
}

/// Python-visible TidyView wrapper.
#[pyclass(name = "TidyView", unsendable)]
#[derive(Clone)]
struct PyTidyView {
    inner: tidyview::TidyView,
}

#[pymethods]
impl PyTidyView {
    /// Create a TidyView from a DataFrame.
    #[new]
    fn new(df: &PyDataFrame) -> Self {
        let tv = tidyview::TidyView::new(df.inner.clone());
        PyTidyView { inner: tv }
    }

    /// Number of visible rows.
    fn nrows(&self) -> usize {
        self.inner.nrows()
    }

    /// Number of visible columns.
    fn ncols(&self) -> usize {
        self.inner.ncols()
    }

    /// Visible column names.
    fn column_names(&self) -> Vec<String> {
        self.inner.column_names().into_iter().map(|s| s.to_string()).collect()
    }

    /// Filter rows where column > value (integer).
    fn filter_gt_int(&self, col_name: &str, value: i64) -> PyResult<Self> {
        let pred = expr::binop(
            expr::BinOp::Gt,
            expr::col(col_name),
            expr::DExpr::LitInt(value),
        );
        let inner = self.inner.filter(&pred)
            .map_err(|e| PyRuntimeError::new_err(format!("{}", e)))?;
        Ok(PyTidyView { inner })
    }

    /// Filter rows where column < value (integer).
    fn filter_lt_int(&self, col_name: &str, value: i64) -> PyResult<Self> {
        let pred = expr::binop(
            expr::BinOp::Lt,
            expr::col(col_name),
            expr::DExpr::LitInt(value),
        );
        let inner = self.inner.filter(&pred)
            .map_err(|e| PyRuntimeError::new_err(format!("{}", e)))?;
        Ok(PyTidyView { inner })
    }

    /// Filter rows where column == value (integer).
    fn filter_eq_int(&self, col_name: &str, value: i64) -> PyResult<Self> {
        let pred = expr::binop(
            expr::BinOp::Eq,
            expr::col(col_name),
            expr::DExpr::LitInt(value),
        );
        let inner = self.inner.filter(&pred)
            .map_err(|e| PyRuntimeError::new_err(format!("{}", e)))?;
        Ok(PyTidyView { inner })
    }

    /// Filter rows where column == value (string).
    fn filter_eq_str(&self, col_name: &str, value: &str) -> PyResult<Self> {
        let pred = expr::binop(
            expr::BinOp::Eq,
            expr::col(col_name),
            expr::DExpr::LitStr(value.to_string()),
        );
        let inner = self.inner.filter(&pred)
            .map_err(|e| PyRuntimeError::new_err(format!("{}", e)))?;
        Ok(PyTidyView { inner })
    }

    /// Filter rows where column > value (float).
    fn filter_gt_float(&self, col_name: &str, value: f64) -> PyResult<Self> {
        let pred = expr::binop(
            expr::BinOp::Gt,
            expr::col(col_name),
            expr::DExpr::LitFloat(value),
        );
        let inner = self.inner.filter(&pred)
            .map_err(|e| PyRuntimeError::new_err(format!("{}", e)))?;
        Ok(PyTidyView { inner })
    }

    /// Select specific columns by name.
    fn select(&self, columns: Vec<String>) -> PyResult<Self> {
        let refs: Vec<&str> = columns.iter().map(|s| s.as_str()).collect();
        let inner = self.inner.select(&refs)
            .map_err(|e| PyRuntimeError::new_err(format!("{}", e)))?;
        Ok(PyTidyView { inner })
    }

    /// Sort by a column (ascending).
    fn arrange(&self, col_name: &str) -> PyResult<Self> {
        let keys = vec![tidyview::ArrangeKey {
            col_name: col_name.to_string(),
            descending: false,
        }];
        let inner = self.inner.arrange(&keys)
            .map_err(|e| PyRuntimeError::new_err(format!("{}", e)))?;
        Ok(PyTidyView { inner })
    }

    /// Sort by a column (descending).
    fn arrange_desc(&self, col_name: &str) -> PyResult<Self> {
        let keys = vec![tidyview::ArrangeKey {
            col_name: col_name.to_string(),
            descending: true,
        }];
        let inner = self.inner.arrange(&keys)
            .map_err(|e| PyRuntimeError::new_err(format!("{}", e)))?;
        Ok(PyTidyView { inner })
    }

    /// Take the first n rows.
    fn slice_head(&self, n: usize) -> Self {
        PyTidyView { inner: self.inner.slice_head(n) }
    }

    /// Take the last n rows.
    fn slice_tail(&self, n: usize) -> Self {
        PyTidyView { inner: self.inner.slice_tail(n) }
    }

    /// Deterministic random sample of n rows.
    fn slice_sample(&self, n: usize, seed: u64) -> Self {
        PyTidyView { inner: self.inner.slice_sample(n, seed) }
    }

    /// Distinct rows by specified columns.
    fn distinct(&self, columns: Vec<String>) -> PyResult<Self> {
        let refs: Vec<&str> = columns.iter().map(|s| s.as_str()).collect();
        let inner = self.inner.distinct(&refs)
            .map_err(|e| PyRuntimeError::new_err(format!("{}", e)))?;
        Ok(PyTidyView { inner })
    }

    /// Group by columns and summarise with an aggregation.
    ///
    /// `agg_fn` is one of: "count", "sum", "mean", "min", "max", "sd", "var",
    /// "first", "last", "n_distinct".
    /// `agg_col` is the source column to aggregate.
    /// `output_name` is the name for the output column (e.g., "mean_score").
    fn group_summarise(&self, group_cols: Vec<String>, agg_col: &str, agg_fn: &str, output_name: &str) -> PyResult<PyDataFrame> {
        let refs: Vec<&str> = group_cols.iter().map(|s| s.as_str()).collect();
        let agg = parse_agg(agg_fn, agg_col)?;
        let grouped = self.inner.group_by(&refs)
            .map_err(|e| PyRuntimeError::new_err(format!("{}", e)))?;
        let result_df = grouped.summarise(&[(output_name, agg)])
            .map_err(|e| PyRuntimeError::new_err(format!("{}", e)))?;
        Ok(PyDataFrame { inner: result_df })
    }

    /// Inner join with another TidyView on specified columns.
    ///
    /// `by` is a list of column names to join on (same name in both views).
    fn inner_join(&self, other: &PyTidyView, by: Vec<String>) -> PyResult<PyDataFrame> {
        let pairs: Vec<(&str, &str)> = by.iter().map(|s| (s.as_str(), s.as_str())).collect();
        let result_df = self.inner.inner_join(&other.inner, &pairs)
            .map_err(|e| PyRuntimeError::new_err(format!("{}", e)))?;
        Ok(PyDataFrame { inner: result_df })
    }

    /// Left join with another TidyView on specified columns.
    fn left_join(&self, other: &PyTidyView, by: Vec<String>) -> PyResult<PyDataFrame> {
        let pairs: Vec<(&str, &str)> = by.iter().map(|s| (s.as_str(), s.as_str())).collect();
        let result_df = self.inner.left_join(&other.inner, &pairs)
            .map_err(|e| PyRuntimeError::new_err(format!("{}", e)))?;
        Ok(PyDataFrame { inner: result_df })
    }

    /// Materialize the view into a concrete DataFrame.
    fn materialize(&self) -> PyResult<PyDataFrame> {
        let df = self.inner.materialize()
            .map_err(|e| PyRuntimeError::new_err(format!("{}", e)))?;
        Ok(PyDataFrame { inner: df })
    }

    fn __repr__(&self) -> String {
        format!("TidyView(nrows={}, ncols={}, columns={:?})",
            self.nrows(), self.ncols(), self.column_names())
    }
}

/// Python-visible Kahan accumulator.
#[pyclass(name = "KahanAccumulator")]
struct PyKahanAccumulator {
    inner: kahan::KahanAccumulator,
}

#[pymethods]
impl PyKahanAccumulator {
    #[new]
    fn new() -> Self {
        PyKahanAccumulator { inner: kahan::KahanAccumulator::new() }
    }

    fn add(&mut self, value: f64) {
        self.inner.add(value);
    }

    fn add_slice(&mut self, values: Vec<f64>) {
        self.inner.add_slice(&values);
    }

    fn finalize(&self) -> f64 {
        self.inner.finalize()
    }

    fn count(&self) -> usize {
        self.inner.count()
    }
}

/// Python-visible deterministic RNG.
#[pyclass(name = "Rng")]
struct PyRng {
    inner: rng::Rng,
}

#[pymethods]
impl PyRng {
    #[new]
    fn new(seed: u64) -> Self {
        PyRng { inner: rng::Rng::seeded(seed) }
    }

    fn next_u64(&mut self) -> u64 {
        self.inner.next_u64()
    }

    fn next_f64(&mut self) -> f64 {
        self.inner.next_f64()
    }

    fn next_normal(&mut self) -> f64 {
        self.inner.next_normal()
    }

    fn fork(&mut self) -> Self {
        PyRng { inner: self.inner.fork() }
    }
}

// ── CSV functions ─────────────────────────────────────────────────────────

/// Parse CSV text into a DataFrame.
#[pyfunction]
fn read_csv(text: &str) -> PyResult<PyDataFrame> {
    let reader = csv::CsvReader::new(csv::CsvConfig::default());
    let df = reader.parse(text.as_bytes())
        .map_err(|e| PyValueError::new_err(format!("{}", e)))?;
    Ok(PyDataFrame { inner: df })
}

/// Parse CSV with a custom delimiter.
#[pyfunction]
fn read_csv_delim(text: &str, delimiter: &str) -> PyResult<PyDataFrame> {
    let delim = delimiter.as_bytes().first().copied().unwrap_or(b',');
    let config = csv::CsvConfig {
        delimiter: delim,
        ..Default::default()
    };
    let reader = csv::CsvReader::new(config);
    let df = reader.parse(text.as_bytes())
        .map_err(|e| PyValueError::new_err(format!("{}", e)))?;
    Ok(PyDataFrame { inner: df })
}

// ── Regex functions ───────────────────────────────────────────────────────

/// Test if a regex pattern matches anywhere in the text.
#[pyfunction]
#[pyo3(signature = (pattern, text, flags=None))]
fn regex_is_match(pattern: &str, text: &str, flags: Option<&str>) -> bool {
    regex_engine::is_match(pattern, flags.unwrap_or(""), text.as_bytes())
}

/// Find the first match span (start, end) or None.
#[pyfunction]
#[pyo3(signature = (pattern, text, flags=None))]
fn regex_find(pattern: &str, text: &str, flags: Option<&str>) -> Option<(usize, usize)> {
    regex_engine::find(pattern, flags.unwrap_or(""), text.as_bytes())
}

/// Find all non-overlapping match spans.
#[pyfunction]
#[pyo3(signature = (pattern, text, flags=None))]
fn regex_find_all(pattern: &str, text: &str, flags: Option<&str>) -> Vec<(usize, usize)> {
    regex_engine::find_all(pattern, flags.unwrap_or(""), text.as_bytes())
}

/// Split text by a regex pattern, returning segment spans.
#[pyfunction]
#[pyo3(signature = (pattern, text, flags=None))]
fn regex_split(pattern: &str, text: &str, flags: Option<&str>) -> Vec<(usize, usize)> {
    regex_engine::split(pattern, flags.unwrap_or(""), text.as_bytes())
}

// ── NLP functions ─────────────────────────────────────────────────────────

/// Levenshtein edit distance between two strings.
#[pyfunction]
fn levenshtein(a: &str, b: &str) -> usize {
    nlp::levenshtein(a, b)
}

/// Normalized Levenshtein similarity in [0.0, 1.0].
#[pyfunction]
fn levenshtein_similarity(a: &str, b: &str) -> f64 {
    nlp::levenshtein_similarity(a, b)
}

/// Jaccard similarity between character n-gram sets.
#[pyfunction]
fn jaccard_ngram_similarity(a: &str, b: &str, n: usize) -> f64 {
    nlp::jaccard_ngram_similarity(a, b, n)
}

/// Extract character n-grams with frequency counts.
#[pyfunction]
fn char_ngrams(text: &str, n: usize) -> std::collections::BTreeMap<String, usize> {
    nlp::char_ngrams(text, n)
}

/// Extract word n-grams with frequency counts.
#[pyfunction]
fn word_ngrams(text: &str, n: usize) -> std::collections::BTreeMap<String, usize> {
    nlp::word_ngrams(text, n)
}

/// Tokenize by whitespace, returning (start, end) byte spans.
#[pyfunction]
fn tokenize_whitespace(text: &str) -> Vec<(usize, usize)> {
    nlp::tokenize_whitespace(text)
}

/// Tokenize into words and punctuation.
#[pyfunction]
fn tokenize_words(text: &str) -> Vec<String> {
    nlp::tokenize_words(text)
}

/// Term frequency (TF) for each word.
#[pyfunction]
fn term_frequency(text: &str) -> std::collections::BTreeMap<String, f64> {
    nlp::term_frequency(text)
}

/// One-shot Kahan-compensated sum of a list of floats.
#[pyfunction]
fn kahan_sum(values: Vec<f64>) -> f64 {
    kahan::kahan_sum(&values)
}

// ── Helper functions ──────────────────────────────────────────────────────

fn py_list_to_column(list: &Bound<'_, PyList>) -> PyResult<column::Column> {
    if list.is_empty() {
        return Ok(column::Column::Str(Vec::new()));
    }

    // Detect type from the first element
    let first = list.get_item(0)?;
    if first.extract::<bool>().is_ok() {
        let vals: Vec<bool> = list.extract()?;
        Ok(column::Column::Bool(vals))
    } else if first.extract::<i64>().is_ok() {
        let vals: Vec<i64> = list.extract()?;
        Ok(column::Column::Int(vals))
    } else if first.extract::<f64>().is_ok() {
        let vals: Vec<f64> = list.extract()?;
        Ok(column::Column::Float(vals))
    } else {
        let vals: Vec<String> = list.extract()?;
        Ok(column::Column::Str(vals))
    }
}

fn column_to_py(py: Python<'_>, col: &column::Column) -> PyResult<PyObject> {
    match col {
        column::Column::Int(v) => Ok(v.to_object(py)),
        column::Column::Float(v) => Ok(v.to_object(py)),
        column::Column::Str(v) => Ok(v.to_object(py)),
        column::Column::Bool(v) => Ok(v.to_object(py)),
    }
}

fn parse_agg(name: &str, col: &str) -> PyResult<tidyview::TidyAgg> {
    let c = col.to_string();
    match name.to_lowercase().as_str() {
        "count" => Ok(tidyview::TidyAgg::Count),
        "sum" => Ok(tidyview::TidyAgg::Sum(c)),
        "mean" => Ok(tidyview::TidyAgg::Mean(c)),
        "min" => Ok(tidyview::TidyAgg::Min(c)),
        "max" => Ok(tidyview::TidyAgg::Max(c)),
        "sd" => Ok(tidyview::TidyAgg::Sd(c)),
        "var" => Ok(tidyview::TidyAgg::Var(c)),
        "first" => Ok(tidyview::TidyAgg::First(c)),
        "last" => Ok(tidyview::TidyAgg::Last(c)),
        "n_distinct" => Ok(tidyview::TidyAgg::NDistinct(c)),
        _ => Err(PyValueError::new_err(format!("unknown aggregation: {}", name))),
    }
}

// ── Python module ─────────────────────────────────────────────────────────

/// virtual-frame Python module.
#[pymodule]
fn virtual_frame(m: &Bound<'_, PyModule>) -> PyResult<()> {
    // Classes
    m.add_class::<PyDataFrame>()?;
    m.add_class::<PyTidyView>()?;
    m.add_class::<PyKahanAccumulator>()?;
    m.add_class::<PyRng>()?;

    // CSV functions
    m.add_function(wrap_pyfunction!(read_csv, m)?)?;
    m.add_function(wrap_pyfunction!(read_csv_delim, m)?)?;

    // Regex functions
    m.add_function(wrap_pyfunction!(regex_is_match, m)?)?;
    m.add_function(wrap_pyfunction!(regex_find, m)?)?;
    m.add_function(wrap_pyfunction!(regex_find_all, m)?)?;
    m.add_function(wrap_pyfunction!(regex_split, m)?)?;

    // NLP functions
    m.add_function(wrap_pyfunction!(levenshtein, m)?)?;
    m.add_function(wrap_pyfunction!(levenshtein_similarity, m)?)?;
    m.add_function(wrap_pyfunction!(jaccard_ngram_similarity, m)?)?;
    m.add_function(wrap_pyfunction!(char_ngrams, m)?)?;
    m.add_function(wrap_pyfunction!(word_ngrams, m)?)?;
    m.add_function(wrap_pyfunction!(tokenize_whitespace, m)?)?;
    m.add_function(wrap_pyfunction!(tokenize_words, m)?)?;
    m.add_function(wrap_pyfunction!(term_frequency, m)?)?;

    // Math
    m.add_function(wrap_pyfunction!(kahan_sum, m)?)?;

    Ok(())
}