use std::fmt;

use crate::error::{TuphiumError, TuphiumResult};
use crate::frame::csv::{CsvColumnType, CsvData, CsvValue};
use crate::linalg::Matrix;

/// A single homogeneous column of a DataFrame.
#[derive(Debug, Clone, PartialEq)]
pub enum Column {
    Integer(Vec<i64>),
    Float(Vec<f64>),
    Boolean(Vec<bool>),
    Text(Vec<String>),
}

impl Column {
    /// Number of values in the column.
    pub fn len(&self) -> usize {
        match self {
            Column::Integer(integers) => integers.len(),
            Column::Float(floats) => floats.len(),
            Column::Boolean(booleans) => booleans.len(),
            Column::Text(texts) => texts.len(),
        }
    }

    /// True when the column has no values.
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Type tag for this column.
    pub fn column_type(&self) -> CsvColumnType {
        match self {
            Column::Integer(_) => CsvColumnType::Integer,
            Column::Float(_) => CsvColumnType::Float,
            Column::Boolean(_) => CsvColumnType::Boolean,
            Column::Text(_) => CsvColumnType::Text,
        }
    }

    /// True for Integer/Float columns.
    pub fn is_numeric(&self) -> bool {
        matches!(self, Column::Integer(_) | Column::Float(_))
    }

    /// Numeric view as f64 (Integer is promoted). None for non-numeric columns.
    pub fn as_f64(&self) -> Option<Vec<f64>> {
        match self {
            Column::Integer(integers) => {
                Some(integers.iter().map(|&int_val| int_val as f64).collect())
            }
            Column::Float(floats) => Some(floats.clone()),
            _ => None,
        }
    }

    /// Keep only the values whose mask entry is true.
    fn filter_mask(&self, mask: &[bool]) -> Column {
        fn keep<T: Clone>(slice: &[T], mask: &[bool]) -> Vec<T> {
            slice
                .iter()
                .zip(mask)
                .filter(|&(_, &should_keep)| should_keep)
                .map(|(item, _)| item.clone())
                .collect()
        }
        match self {
            Column::Integer(integers) => Column::Integer(keep(integers, mask)),
            Column::Float(floats) => Column::Float(keep(floats, mask)),
            Column::Boolean(booleans) => Column::Boolean(keep(booleans, mask)),
            Column::Text(texts) => Column::Text(keep(texts, mask)),
        }
    }

    /// New column with the values at `indices`, in the given order (may repeat).
    fn gather(&self, indices: &[usize]) -> Column {
        fn take<T: Clone>(slice: &[T], indices: &[usize]) -> Vec<T> {
            indices.iter().map(|&idx| slice[idx].clone()).collect()
        }
        match self {
            Column::Integer(integers) => Column::Integer(take(integers, indices)),
            Column::Float(floats) => Column::Float(take(floats, indices)),
            Column::Boolean(booleans) => Column::Boolean(take(booleans, indices)),
            Column::Text(texts) => Column::Text(take(texts, indices)),
        }
    }

    /// String representation of the value at `row` (for display).
    fn cell_string(&self, row: usize) -> String {
        match self {
            Column::Integer(integers) => integers[row].to_string(),
            Column::Float(floats) => floats[row].to_string(),
            Column::Boolean(booleans) => booleans[row].to_string(),
            Column::Text(texts) => texts[row].clone(),
        }
    }
}

/// A column-oriented table with named, typed columns.
#[derive(Debug, Clone)]
pub struct DataFrame {
    names: Vec<String>,
    columns: Vec<Column>,
    n_rows: usize,
}

impl DataFrame {
    /// Builds a DataFrame from a CsvData (read with a header).
    ///
    /// Column types are inferred via [`CsvData::column_types`] and every cell is
    /// coerced to its column type (Integer promotes to Float when needed).
    ///
    /// # Errors
    ///
    /// Returns an error if a row's field count differs from the header width, or
    /// if a cell cannot be coerced to its inferred column type.
    pub fn from_csv_data(data: CsvData) -> TuphiumResult<DataFrame> {
        let names: Vec<String> = data.header.clone();
        let n_cols: usize = names.len();
        let n_rows: usize = data.rows.len();

        // All rows must match the header width.
        for (row_index, row) in data.rows.iter().enumerate() {
            if row.len() != n_cols {
                return Err(TuphiumError::MalformedCsv {
                    operation: "DataFrame::from_csv_data",
                    row: row_index,
                    expected_fields: n_cols,
                    actual_fields: row.len(),
                });
            }
        }

        let types: Vec<CsvColumnType> = data.column_types();
        let mut columns: Vec<Column> = Vec::with_capacity(n_cols);

        for (col_index, col_type) in types.iter().enumerate() {
            let column: Column = match col_type {
                CsvColumnType::Integer => {
                    let mut column_values = Vec::with_capacity(n_rows);
                    for row in &data.rows {
                        match &row[col_index] {
                            CsvValue::Integer(int_val) => column_values.push(*int_val),
                            other => {
                                return Err(TuphiumError::ColumnTypeMismatch {
                                    operation: "DataFrame::from_csv_data",
                                    column: names[col_index].clone(),
                                    expected: "Integer".to_string(),
                                    actual: csv_value_type_name(other).to_string(),
                                });
                            }
                        }
                    }
                    Column::Integer(column_values)
                }
                CsvColumnType::Float => {
                    let mut column_values: Vec<f64> = Vec::with_capacity(n_rows);
                    for row in &data.rows {
                        match &row[col_index] {
                            CsvValue::Float(float_val) => column_values.push(*float_val),
                            CsvValue::Integer(int_val) => column_values.push(*int_val as f64), // promote
                            other => {
                                return Err(TuphiumError::ColumnTypeMismatch {
                                    operation: "DataFrame::from_csv_data",
                                    column: names[col_index].clone(),
                                    expected: "Float".to_string(),
                                    actual: csv_value_type_name(other).to_string(),
                                });
                            }
                        }
                    }
                    Column::Float(column_values)
                }
                CsvColumnType::Boolean => {
                    let mut column_values: Vec<bool> = Vec::with_capacity(n_rows);
                    for row in &data.rows {
                        match &row[col_index] {
                            CsvValue::Boolean(bool_val) => column_values.push(*bool_val),
                            other => {
                                return Err(TuphiumError::ColumnTypeMismatch {
                                    operation: "DataFrame::from_csv_data",
                                    column: names[col_index].clone(),
                                    expected: "Boolean".to_string(),
                                    actual: csv_value_type_name(other).to_string(),
                                });
                            }
                        }
                    }
                    Column::Boolean(column_values)
                }
                CsvColumnType::Text => {
                    // Text is the widest type: render any value as its string form.
                    let column_values: Vec<String> = data
                        .rows
                        .iter()
                        .map(|row| match &row[col_index] {
                            CsvValue::Text(text_val) => text_val.clone(),
                            CsvValue::Integer(int_val) => int_val.to_string(),
                            CsvValue::Float(float_val) => float_val.to_string(),
                            CsvValue::Boolean(bool_val) => bool_val.to_string(),
                        })
                        .collect();
                    Column::Text(column_values)
                }
            };
            columns.push(column);
        }

        Ok(DataFrame {
            names,
            columns,
            n_rows,
        })
    }

    /// Build a DataFrame from named columns.
    pub fn from_columns(names: Vec<String>, columns: Vec<Column>) -> TuphiumResult<Self> {
        if names.len() != columns.len() {
            return Err(TuphiumError::LengthMismatch {
                operation: "DataFrame::from_columns",
                expected: columns.len(),
                actual: names.len(),
            });
        }
        let n_rows: usize = columns.first().map(|column| column.len()).unwrap_or(0);
        for column in &columns {
            if column.len() != n_rows {
                return Err(TuphiumError::ShapeMismatch {
                    operation: "DataFrame::from_columns",
                    expected: vec![n_rows],
                    actual: vec![column.len()],
                });
            }
        }
        Ok(DataFrame {
            names,
            columns,
            n_rows,
        })
    }

    // ─── Inspection ─────────────────────────────────────────────────────────

    /// Returns the dimensions as `(n_rows, n_cols)`.
    pub fn shape(&self) -> (usize, usize) {
        (self.n_rows, self.columns.len())
    }

    /// Returns the column names in order.
    pub fn column_names(&self) -> &[String] {
        &self.names
    }

    fn index_of(&self, name: &str) -> Option<usize> {
        self.names.iter().position(|col_name| col_name == name)
    }

    /// Returns the named column, or `None` if it does not exist.
    pub fn column(&self, name: &str) -> Option<&Column> {
        self.index_of(name).map(|col_idx| &self.columns[col_idx])
    }

    /// Returns the inferred type of the named column, or `None` if it is missing.
    pub fn column_type(&self, name: &str) -> Option<CsvColumnType> {
        self.column(name).map(|column| column.column_type())
    }

    /// Fallible numeric column access as f64.
    pub fn as_f64(&self, name: &str) -> TuphiumResult<Vec<f64>> {
        let column = self
            .column(name)
            .ok_or_else(|| TuphiumError::ColumnNotFound {
                operation: "DataFrame::as_f64",
                column: name.to_string(),
            })?;
        column
            .as_f64()
            .ok_or_else(|| TuphiumError::ColumnTypeMismatch {
                operation: "DataFrame::as_f64",
                column: name.to_string(),
                expected: "numeric".to_string(),
                actual: csv_column_type_name(&column.column_type()).to_string(),
            })
    }

    // ─── Column selection ───────────────────────────────────────────────────

    /// New DataFrame with only the named columns, in the given order.
    ///
    /// # Errors
    ///
    /// Returns an error if any requested column name is not present.
    pub fn select(&self, names: &[&str]) -> TuphiumResult<DataFrame> {
        let mut selected_names = Vec::with_capacity(names.len());
        let mut selected_columns = Vec::with_capacity(names.len());
        for &name in names {
            let column_index = self
                .index_of(name)
                .ok_or_else(|| TuphiumError::ColumnNotFound {
                    operation: "DataFrame::select",
                    column: name.to_string(),
                })?;
            selected_names.push(self.names[column_index].clone());
            selected_columns.push(self.columns[column_index].clone());
        }
        Ok(DataFrame {
            names: selected_names,
            columns: selected_columns,
            n_rows: self.n_rows,
        })
    }

    /// New DataFrame without the named column.
    ///
    /// # Errors
    ///
    /// Returns an error if the column name is not present.
    pub fn drop(&self, name: &str) -> TuphiumResult<DataFrame> {
        let column_index = self
            .index_of(name)
            .ok_or_else(|| TuphiumError::ColumnNotFound {
                operation: "DataFrame::drop",
                column: name.to_string(),
            })?;
        let mut remaining_names = self.names.clone();
        let mut remaining_columns = self.columns.clone();
        remaining_names.remove(column_index);
        remaining_columns.remove(column_index);
        Ok(DataFrame {
            names: remaining_names,
            columns: remaining_columns,
            n_rows: self.n_rows,
        })
    }

    // ─── Row filtering ──────────────────────────────────────────────────────

    /// New DataFrame keeping rows where `mask[i]` is true.
    pub fn filter_mask(&self, mask: &[bool]) -> TuphiumResult<DataFrame> {
        if mask.len() != self.n_rows {
            return Err(TuphiumError::LengthMismatch {
                operation: "DataFrame::filter_mask",
                expected: self.n_rows,
                actual: mask.len(),
            });
        }
        let columns: Vec<Column> = self
            .columns
            .iter()
            .map(|column| column.filter_mask(mask))
            .collect();
        let n_rows: usize = mask.iter().filter(|&&should_keep| should_keep).count();
        Ok(DataFrame {
            names: self.names.clone(),
            columns,
            n_rows,
        })
    }

    /// New DataFrame with the rows at `indices`, in the given order. Indices may
    /// repeat or be reordered, which makes this the building block for shuffling
    /// and partitioning. Panics if any index is out of bounds.
    pub fn select_rows(&self, indices: &[usize]) -> TuphiumResult<DataFrame> {
        for &row in indices {
            if row >= self.n_rows {
                return Err(TuphiumError::RowIndexOutOfBounds {
                    operation: "DataFrame::select_rows",
                    row,
                    rows: self.n_rows,
                });
            }
        }
        let columns: Vec<Column> = self
            .columns
            .iter()
            .map(|column| column.gather(indices))
            .collect();
        Ok(DataFrame {
            names: self.names.clone(),
            columns,
            n_rows: indices.len(),
        })
    }

    // ─── Statistics (numeric columns only) ──────────────────────────────────

    /// Sum of a numeric column. `None` if the column is missing or non-numeric.
    pub fn sum(&self, name: &str) -> Option<f64> {
        self.as_f64(name).ok().map(|values| values.iter().sum())
    }

    /// Arithmetic mean of a numeric column.
    ///
    /// `None` if the column is missing, non-numeric, or empty.
    pub fn mean(&self, name: &str) -> Option<f64> {
        let values: Vec<f64> = self.as_f64(name).ok()?;
        if values.is_empty() {
            return None;
        }
        Some(values.iter().sum::<f64>() / values.len() as f64)
    }

    /// Minimum value of a numeric column. `None` if missing or non-numeric.
    pub fn min(&self, name: &str) -> Option<f64> {
        let values: Vec<f64> = self.as_f64(name).ok()?;
        values
            .into_iter()
            .fold(None, |minimum, current| match minimum {
                Some(min_so_far) if min_so_far <= current => Some(min_so_far),
                _ => Some(current),
            })
    }

    /// Maximum value of a numeric column. `None` if missing or non-numeric.
    pub fn max(&self, name: &str) -> Option<f64> {
        let values: Vec<f64> = self.as_f64(name).ok()?;
        values
            .into_iter()
            .fold(None, |maximum, current| match maximum {
                Some(max_so_far) if max_so_far >= current => Some(max_so_far),
                _ => Some(current),
            })
    }

    /// Sample standard deviation (ddof=1). None if non-numeric or n < 2.
    pub fn std(&self, name: &str) -> Option<f64> {
        let values: Vec<f64> = self.as_f64(name).ok()?;
        let count: usize = values.len();
        if count < 2 {
            return None;
        }
        let mean: f64 = values.iter().sum::<f64>() / count as f64;
        let variance: f64 = values
            .iter()
            .map(|&value| (value - mean).powi(2))
            .sum::<f64>()
            / (count as f64 - 1.0);
        Some(variance.sqrt())
    }

    // ─── ML bridge ──────────────────────────────────────────────────────────

    /// Builds a Matrix (n_rows × n_numeric_cols, f32) from every numeric column,
    /// preserving column order.
    ///
    /// # Errors
    ///
    /// Returns an error if the DataFrame has no numeric columns to convert.
    pub fn to_matrix(&self) -> TuphiumResult<Matrix> {
        let numeric_columns: Vec<&Column> = self
            .columns
            .iter()
            .filter(|column| column.is_numeric())
            .collect();
        if numeric_columns.is_empty() {
            return Err(TuphiumError::NoNumericColumns {
                operation: "DataFrame::to_matrix",
            });
        }
        let n_cols: usize = numeric_columns.len();
        let len = self
            .n_rows
            .checked_mul(n_cols)
            .ok_or(TuphiumError::DimensionOverflow {
                operation: "DataFrame::to_matrix",
                rows: self.n_rows,
                cols: n_cols,
            })?;
        let mut data: Vec<f32> = Vec::with_capacity(len);
        for row in 0..self.n_rows {
            for numeric_col in &numeric_columns {
                let cell_value = match numeric_col {
                    Column::Integer(integers) => integers[row] as f32,
                    Column::Float(floats) => floats[row] as f32,
                    _ => unreachable!("filtered to numeric columns"),
                };
                data.push(cell_value);
            }
        }
        Matrix::from_data(self.n_rows, n_cols, data)
    }

    /// Prints the header and the first `n` rows.
    pub fn head(&self, n: usize) {
        let rows: usize = n.min(self.n_rows);
        println!("{}", self.names.join("\t"));
        for row in 0..rows {
            let cells: Vec<String> = self
                .columns
                .iter()
                .map(|column| column.cell_string(row))
                .collect();
            println!("{}", cells.join("\t"));
        }
    }
}

fn csv_value_type_name(value: &CsvValue) -> &'static str {
    match value {
        CsvValue::Integer(_) => "Integer",
        CsvValue::Float(_) => "Float",
        CsvValue::Boolean(_) => "Boolean",
        CsvValue::Text(_) => "Text",
    }
}

fn csv_column_type_name(column_type: &CsvColumnType) -> &'static str {
    match column_type {
        CsvColumnType::Integer => "Integer",
        CsvColumnType::Float => "Float",
        CsvColumnType::Boolean => "Boolean",
        CsvColumnType::Text => "Text",
    }
}

impl fmt::Display for DataFrame {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        writeln!(f, "{}", self.names.join("\t"))?;
        for row in 0..self.n_rows {
            let cells: Vec<String> = self
                .columns
                .iter()
                .map(|column| column.cell_string(row))
                .collect();
            writeln!(f, "{}", cells.join("\t"))?;
        }
        Ok(())
    }
}
