module dataframe

import encoding.csv
import math
import strconv

// CsvConfig configures CSV loading for a DataFrame.
@[params]
pub struct CsvConfig {
pub:
    has_header   bool = true
    separator    u8   = `,`
    comment      u8   = `#`
    default_cell string
    empty_cell   string
    end_line_len int = csv.endline_cr_len
    quote        u8  = `"`
    quote_remove bool
}

// SortOrder controls the direction used by sorting helpers.
pub enum SortOrder {
    asc
    desc
}

// Summary contains basic numeric statistics for a Series.
pub struct Summary {
pub:
    count  int
    sum    f64
    mean   f64
    min    f64
    max    f64
    median f64
    stddev f64
}

// Series is a single named string column with numeric helpers.
pub struct Series {
pub:
    name   string
    values []string
}

// Row is a named view of a DataFrame row.
pub struct Row {
pub:
    values map[string]string
}

// DataFrame stores rectangular tabular data as string cells.
pub struct DataFrame {
    index map[string]int
pub:
    columns []string
    rows    [][]string
}

// new creates a DataFrame from column names and rows.
pub fn new(columns []string, rows [][]string) !DataFrame {
    index := build_index(columns)!
    mut copied_rows := [][]string{cap: rows.len}
    for row_index, row in rows {
        if row.len != columns.len {
            return error('row ${row_index} has ${row.len} values, expected ${columns.len}')
        }
        copied_rows << row.clone()
    }
    return DataFrame{
        index:   index
        columns: columns.clone()
        rows:    copied_rows
    }
}

// empty creates an empty DataFrame with the given columns.
pub fn empty(columns []string) !DataFrame {
    return new(columns, [][]string{})!
}

// from_series creates a DataFrame from named columns.
pub fn from_series(series []Series) !DataFrame {
    if series.len == 0 {
        return error('at least one series is required')
    }
    row_count := series[0].values.len
    mut columns := []string{cap: series.len}
    for s in series {
        if s.values.len != row_count {
            return error('series `${s.name}` has ${s.values.len} values, expected ${row_count}')
        }
        columns << s.name
    }
    mut rows := [][]string{cap: row_count}
    for row_index in 0 .. row_count {
        mut row := []string{cap: series.len}
        for s in series {
            row << s.values[row_index]
        }
        rows << row
    }
    return new(columns, rows)!
}

// from_columns creates a DataFrame from a map of columns.
pub fn from_columns(columns map[string][]string) !DataFrame {
    if columns.len == 0 {
        return error('at least one column is required')
    }
    mut names := columns.keys()
    names.sort()
    mut series := []Series{cap: names.len}
    for name in names {
        series << Series{
            name:   name
            values: columns[name].clone()
        }
    }
    return from_series(series)!
}

// from_csv creates a DataFrame from CSV text.
pub fn from_csv(text string, cfg CsvConfig) !DataFrame {
    if text.len == 0 {
        return error('csv input is empty')
    }
    mut reader := csv.csv_reader(
        scr_buf:      text.str
        scr_buf_len:  text.len
        separator:    cfg.separator
        comment:      cfg.comment
        default_cell: cfg.default_cell
        empty_cell:   cfg.empty_cell
        end_line_len: cfg.end_line_len
        quote:        cfg.quote
        quote_remove: cfg.quote_remove
    )!
    defer {
        reader.dispose_csv_reader()
    }
    return from_csv_reader(mut reader, cfg)!
}

// read_csv creates a DataFrame from a CSV file.
pub fn read_csv(path string, cfg CsvConfig) !DataFrame {
    mut reader := csv.csv_reader(
        file_path:    path
        separator:    cfg.separator
        comment:      cfg.comment
        default_cell: cfg.default_cell
        empty_cell:   cfg.empty_cell
        end_line_len: cfg.end_line_len
        quote:        cfg.quote
        quote_remove: cfg.quote_remove
    )!
    defer {
        reader.dispose_csv_reader()
    }
    return from_csv_reader(mut reader, cfg)!
}

fn from_csv_reader(mut reader csv.RandomAccessReader, cfg CsvConfig) !DataFrame {
    if reader.csv_map.len == 0 {
        return error('csv input has no rows')
    }
    first_row := reader.get_row(0)!
    if first_row.len == 0 {
        return error('csv input has no columns')
    }
    columns := if cfg.has_header {
        normalize_columns(first_row)
    } else {
        default_columns(first_row.len)
    }
    start_row := if cfg.has_header { 1 } else { 0 }
    mut rows := [][]string{cap: math.max(0, reader.csv_map.len - start_row)}
    for row_index in start_row .. reader.csv_map.len {
        row := reader.get_row(row_index)!
        rows << normalize_row(row, columns.len, cfg.empty_cell)
    }
    return new(columns, rows)!
}

// height returns the number of rows.
pub fn (df DataFrame) height() int {
    return df.rows.len
}

// width returns the number of columns.
pub fn (df DataFrame) width() int {
    return df.columns.len
}

// shape returns the row and column count.
pub fn (df DataFrame) shape() (int, int) {
    return df.rows.len, df.columns.len
}

// cell returns a single cell by row index and column name.
pub fn (df DataFrame) cell(row_index int, column string) !string {
    if row_index < 0 || row_index >= df.rows.len {
        return error('row index ${row_index} is out of range')
    }
    column_index := df.column_index(column)!
    return df.rows[row_index][column_index]
}

// row returns a named row by index.
pub fn (df DataFrame) row(row_index int) !Row {
    if row_index < 0 || row_index >= df.rows.len {
        return error('row index ${row_index} is out of range')
    }
    mut values := map[string]string{}
    row := df.rows[row_index]
    for column_index, column in df.columns {
        values[column] = row[column_index]
    }
    return Row{
        values: values
    }
}

// column returns a Series by column name.
pub fn (df DataFrame) column(name string) !Series {
    column_index := df.column_index(name)!
    mut values := []string{cap: df.rows.len}
    for row in df.rows {
        values << row[column_index]
    }
    return Series{
        name:   name
        values: values
    }
}

// select returns a DataFrame with only the requested columns.
pub fn (df DataFrame) select(names []string) !DataFrame {
    mut column_indices := []int{cap: names.len}
    for name in names {
        column_indices << df.column_index(name)!
    }
    mut rows := [][]string{cap: df.rows.len}
    for row in df.rows {
        mut selected := []string{cap: names.len}
        for column_index in column_indices {
            selected << row[column_index]
        }
        rows << selected
    }
    return new(names, rows)!
}

// head returns the first n rows.
pub fn (df DataFrame) head(n int) DataFrame {
    if n <= 0 {
        return df.with_rows([][]string{})
    }
    end := math.min(n, df.rows.len)
    return df.with_rows(df.rows[..end])
}

// tail returns the last n rows.
pub fn (df DataFrame) tail(n int) DataFrame {
    if n <= 0 {
        return df.with_rows([][]string{})
    }
    start := math.max(0, df.rows.len - n)
    return df.with_rows(df.rows[start..])
}

// filter returns rows accepted by the predicate.
pub fn (df DataFrame) filter(predicate fn (Row) bool) DataFrame {
    mut rows := [][]string{cap: df.rows.len}
    for row_index in 0 .. df.rows.len {
        row := df.row(row_index) or { continue }
        if predicate(row) {
            rows << df.rows[row_index].clone()
        }
    }
    return df.with_rows(rows)
}

// sort_by returns rows sorted lexicographically by column.
pub fn (df DataFrame) sort_by(name string, order SortOrder) !DataFrame {
    column_index := df.column_index(name)!
    return df.sorted(column_index, order, false)!
}

// sort_by_f64 returns rows sorted numerically by column.
pub fn (df DataFrame) sort_by_f64(name string, order SortOrder) !DataFrame {
    column_index := df.column_index(name)!
    return df.sorted(column_index, order, true)!
}

// value_counts counts unique values in a column.
pub fn (df DataFrame) value_counts(name string) !map[string]int {
    column_index := df.column_index(name)!
    mut counts := map[string]int{}
    for row in df.rows {
        value := row[column_index]
        counts[value]++
    }
    return counts
}

// describe returns numeric statistics for a named column.
pub fn (df DataFrame) describe(name string) !Summary {
    series := df.column(name)!
    return series.describe()!
}

fn (df DataFrame) sorted(column_index int, order SortOrder, numeric bool) !DataFrame {
    mut rows := clone_rows(df.rows)
    for i in 1 .. rows.len {
        mut j := i
        for j > 0 {
            if !should_swap(rows[j - 1][column_index], rows[j][column_index], order, numeric)! {
                break
            }
            rows[j - 1], rows[j] = rows[j], rows[j - 1]
            j--
        }
    }
    return df.with_rows(rows)
}

fn (df DataFrame) with_rows(rows [][]string) DataFrame {
    return DataFrame{
        index:   df.index.clone()
        columns: df.columns.clone()
        rows:    clone_rows(rows)
    }
}

fn (df DataFrame) column_index(name string) !int {
    if name !in df.index {
        return error('unknown column `${name}`')
    }
    return df.index[name]
}

// get returns a value from the Row by column name.
pub fn (row Row) get(name string) !string {
    if name !in row.values {
        return error('unknown column `${name}`')
    }
    return row.values[name]
}

// len returns the number of values in the Series.
pub fn (s Series) len() int {
    return s.values.len
}

// get returns a value from the Series by row index.
pub fn (s Series) get(index int) !string {
    if index < 0 || index >= s.values.len {
        return error('series index ${index} is out of range')
    }
    return s.values[index]
}

// f64s converts every value in the Series to f64.
pub fn (s Series) f64s() ![]f64 {
    mut values := []f64{cap: s.values.len}
    for index, raw_value in s.values {
        values << parse_f64(raw_value, '${s.name}[${index}]')!
    }
    return values
}

// sum returns the numeric sum of the Series.
pub fn (s Series) sum() !f64 {
    values := s.f64s()!
    mut total := 0.0
    for value in values {
        total += value
    }
    return total
}

// mean returns the numeric mean of the Series.
pub fn (s Series) mean() !f64 {
    values := s.f64s()!
    if values.len == 0 {
        return error('series `${s.name}` is empty')
    }
    return sum_f64(values) / f64(values.len)
}

// min returns the smallest numeric value in the Series.
pub fn (s Series) min() !f64 {
    values := s.f64s()!
    if values.len == 0 {
        return error('series `${s.name}` is empty')
    }
    mut min_value := values[0]
    for value in values[1..] {
        if value < min_value {
            min_value = value
        }
    }
    return min_value
}

// max returns the largest numeric value in the Series.
pub fn (s Series) max() !f64 {
    values := s.f64s()!
    if values.len == 0 {
        return error('series `${s.name}` is empty')
    }
    mut max_value := values[0]
    for value in values[1..] {
        if value > max_value {
            max_value = value
        }
    }
    return max_value
}

// median returns the numeric median of the Series.
pub fn (s Series) median() !f64 {
    mut values := s.f64s()!
    if values.len == 0 {
        return error('series `${s.name}` is empty')
    }
    values.sort()
    return median_f64(values)
}

// stddev returns the population standard deviation of the Series.
pub fn (s Series) stddev() !f64 {
    values := s.f64s()!
    if values.len == 0 {
        return error('series `${s.name}` is empty')
    }
    mean := sum_f64(values) / f64(values.len)
    mut variance_sum := 0.0
    for value in values {
        diff := value - mean
        variance_sum += diff * diff
    }
    return math.sqrt(variance_sum / f64(values.len))
}

// describe returns basic numeric statistics for the Series.
pub fn (s Series) describe() !Summary {
    mut values := s.f64s()!
    if values.len == 0 {
        return error('series `${s.name}` is empty')
    }
    values.sort()
    sum := sum_f64(values)
    mean := sum / f64(values.len)
    mut variance_sum := 0.0
    for value in values {
        diff := value - mean
        variance_sum += diff * diff
    }
    return Summary{
        count:  values.len
        sum:    sum
        mean:   mean
        min:    values.first()
        max:    values.last()
        median: median_f64(values)
        stddev: math.sqrt(variance_sum / f64(values.len))
    }
}

fn build_index(columns []string) !map[string]int {
    if columns.len == 0 {
        return error('at least one column is required')
    }
    mut index := map[string]int{}
    for column_index, column in columns {
        name := column.trim_space()
        if name.len == 0 {
            return error('column ${column_index} is empty')
        }
        if name in index {
            return error('duplicate column `${name}`')
        }
        index[name] = column_index
    }
    return index
}

fn clone_rows(rows [][]string) [][]string {
    mut copied_rows := [][]string{cap: rows.len}
    for row in rows {
        copied_rows << row.clone()
    }
    return copied_rows
}

fn default_columns(count int) []string {
    mut columns := []string{cap: count}
    for index in 0 .. count {
        columns << 'column_${index}'
    }
    return columns
}

fn normalize_columns(columns []string) []string {
    mut normalized := []string{cap: columns.len}
    for column in columns {
        normalized << column.trim_space()
    }
    return normalized
}

fn normalize_row(row []string, column_count int, fill string) []string {
    mut normalized := []string{cap: column_count}
    for column_index in 0 .. column_count {
        if column_index < row.len {
            normalized << row[column_index]
        } else {
            normalized << fill
        }
    }
    return normalized
}

fn should_swap(left string, right string, order SortOrder, numeric bool) !bool {
    comparison := if numeric {
        compare_f64(left, right)!
    } else {
        compare_string(left, right)
    }
    return match order {
        .asc { comparison > 0 }
        .desc { comparison < 0 }
    }
}

fn compare_string(left string, right string) int {
    if left < right {
        return -1
    }
    if left > right {
        return 1
    }
    return 0
}

fn compare_f64(left string, right string) !int {
    left_number := parse_f64(left, 'left value')!
    right_number := parse_f64(right, 'right value')!
    if left_number < right_number {
        return -1
    }
    if left_number > right_number {
        return 1
    }
    return 0
}

fn parse_f64(value string, label string) !f64 {
    trimmed_value := value.trim_space()
    $if js {
        return strconv.atof64(trimmed_value) or { error('${label} is not a number: `${value}`') }
    } $else {
        return strconv.atof64(trimmed_value, strconv.AtoF64Param{}) or {
            error('${label} is not a number: `${value}`')
        }
    }
}

fn sum_f64(values []f64) f64 {
    mut total := 0.0
    for value in values {
        total += value
    }
    return total
}

fn median_f64(sorted_values []f64) f64 {
    mid := sorted_values.len / 2
    if sorted_values.len % 2 == 1 {
        return sorted_values[mid]
    }
    return (sorted_values[mid - 1] + sorted_values[mid]) / 2.0
}