use tuphium::error::TuphiumError;
use tuphium::frame::csv::{CsvData, CsvValue};
use tuphium::frame::dataframe::{Column, DataFrame};
use tuphium::frame::dataset::Dataset;

fn ids_dataset(n: usize) -> Dataset {
    fn make_rows(n: usize) -> Vec<Vec<CsvValue>> {
        (0..n).map(|i| vec![CsvValue::Integer(i as i64)]).collect()
    }

    let x = DataFrame::from_csv_data(CsvData {
        header: vec!["x".into()],
        rows: make_rows(n),
    })
    .unwrap();

    let y = DataFrame::from_csv_data(CsvData {
        header: vec!["y".into()],
        rows: make_rows(n),
    })
    .unwrap();

    Dataset::new(x, y).unwrap()
}

fn ids(df: &DataFrame) -> Vec<i64> {
    match df.column("x").unwrap() {
        Column::Integer(v) => v.clone(),
        other => panic!("expected Integer column, got {:?}", other),
    }
}

#[test]
fn split_sizes_sum_to_total_and_lose_no_rows() {
    let ds = ids_dataset(10);
    let (train, val, test) = ds.split(0.6, 0.2, 0.2).unwrap();

    assert_eq!(train.len(), 6);
    assert_eq!(val.len(), 2);
    assert_eq!(test.len(), 2);

    let mut all = ids(&train.x);
    all.extend(ids(&val.x));
    all.extend(ids(&test.x));
    all.sort();
    assert_eq!(all, (0..10).collect::<Vec<i64>>());
}

#[test]
fn split_remainder_goes_to_test() {
    let ds = ids_dataset(7);
    let (train, val, test) = ds.split(0.5, 0.2, 0.3).unwrap();
    assert_eq!(train.len(), 3);
    assert_eq!(val.len(), 1);
    assert_eq!(test.len(), 3);
}

#[test]
fn split_with_zero_val_is_train_test_only() {
    let ds = ids_dataset(10);
    let (train, val, test) = ds.split(0.8, 0.0, 0.2).unwrap();
    assert_eq!(train.len(), 8);
    assert_eq!(val.len(), 0);
    assert_eq!(test.len(), 2);
}

#[test]
fn split_returns_invalid_split_fractions() {
    let ds = ids_dataset(10);
    assert_eq!(
        ds.split(0.5, 0.5, 0.5).unwrap_err(),
        TuphiumError::InvalidSplitFractions {
            operation: "Dataset::split",
            train: 0.5,
            validation: 0.5,
            test: 0.5,
            sum: 1.5,
        }
    );
}

#[test]
fn split_rejects_negative_fractions() {
    let ds = ids_dataset(10);
    assert!(matches!(
        ds.split(1.0, -0.1, 0.1).unwrap_err(),
        TuphiumError::InvalidSplitFractions {
            operation: "Dataset::split",
            ..
        }
    ));
}

#[test]
fn shuffle_preserves_all_rows_and_reorders() {
    let mut ds = ids_dataset(50);
    let original_ids = ids(&ds.x);
    ds.shuffle();
    let shuffled_ids = ids(&ds.x);
    assert_ne!(shuffled_ids, original_ids);
    let mut sorted = shuffled_ids.clone();
    sorted.sort();
    assert_eq!(sorted, original_ids);
}

#[test]
fn new_returns_shape_mismatch_for_different_sample_counts() {
    let x =
        DataFrame::from_columns(vec!["x".to_string()], vec![Column::Integer(vec![1, 2])]).unwrap();
    let y = DataFrame::from_columns(vec!["y".to_string()], vec![Column::Integer(vec![1])]).unwrap();
    assert_eq!(
        Dataset::new(x, y).unwrap_err(),
        TuphiumError::ShapeMismatch {
            operation: "Dataset::new",
            expected: vec![2],
            actual: vec![1],
        }
    );
}

#[test]
fn into_numeric_returns_invalid_target_shape() {
    let x =
        DataFrame::from_columns(vec!["x".to_string()], vec![Column::Integer(vec![1, 2])]).unwrap();
    let y = DataFrame::from_columns(
        vec!["y1".to_string(), "y2".to_string()],
        vec![Column::Integer(vec![1, 2]), Column::Integer(vec![0, 1])],
    )
    .unwrap();
    let ds = Dataset::new(x, y).unwrap();
    assert_eq!(
        ds.into_numeric().unwrap_err(),
        TuphiumError::InvalidTargetShape {
            operation: "Dataset::into_numeric",
            expected_columns: 1,
            actual_columns: 2,
        }
    );
}
