use rand::distr::Uniform;
use rand::{rng, RngExt};

use crate::error::{TuphiumError, TuphiumResult};
use crate::frame::dataframe::DataFrame;
use crate::linalg::Matrix;

#[derive(Debug, Clone)]
pub struct Dataset {
    pub x: DataFrame,
    pub y: DataFrame,
}

impl Dataset {
    /// Pairs feature (`x`) and label (`y`) DataFrames into a dataset.
    ///
    /// # Panics
    ///
    /// Panics if `x` and `y` have a different number of rows (samples).
    pub fn new(x: DataFrame, y: DataFrame) -> TuphiumResult<Self> {
        let x_rows = x.shape().0;
        let y_rows = y.shape().0;
        if x_rows != y_rows {
            return Err(TuphiumError::ShapeMismatch {
                operation: "Dataset::new",
                expected: vec![x_rows],
                actual: vec![y_rows],
            });
        }
        Ok(Self { x, y })
    }

    /// Number of rows (samples).
    pub fn len(&self) -> usize {
        self.x.shape().0
    }

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

    /// Shuffles `x` and `y` in place with the same random permutation, keeping
    /// each feature row aligned with its label (Fisher–Yates). No-op below 2 rows.
    pub fn shuffle(&mut self) {
        let sample_count: usize = self.len();
        if sample_count < 2 {
            return;
        }
        let mut rng = rng();
        let mut permutation: Vec<usize> = (0..sample_count).collect();
        for position in (1..sample_count).rev() {
            let distribution = Uniform::new(0, position + 1).unwrap();
            let swap_with: usize = rng.sample(distribution);
            permutation.swap(position, swap_with);
        }
        self.x = self
            .x
            .select_rows(&permutation)
            .expect("Dataset::shuffle: generated permutation must be valid");
        self.y = self
            .y
            .select_rows(&permutation)
            .expect("Dataset::shuffle: generated permutation must be valid");
    }

    /// Splits the dataset into train/validation/test partitions by fraction.
    ///
    /// Rows are partitioned in their current order, so call [`Dataset::shuffle`]
    /// first if a random split is wanted. Sizes are floored, with any remainder
    /// going to the test partition.
    ///
    /// # Panics
    ///
    /// Panics unless `train_frac + val_frac + test_frac == 1.0`.
    pub fn split(
        self,
        train_frac: f32,
        val_frac: f32,
        test_frac: f32,
    ) -> TuphiumResult<(Dataset, Dataset, Dataset)> {
        let sum = train_frac + val_frac + test_frac;
        if !train_frac.is_finite()
            || !val_frac.is_finite()
            || !test_frac.is_finite()
            || train_frac < 0.0
            || val_frac < 0.0
            || test_frac < 0.0
            || (sum - 1.0).abs() > 1e-6
        {
            return Err(TuphiumError::InvalidSplitFractions {
                operation: "Dataset::split",
                train: train_frac,
                validation: val_frac,
                test: test_frac,
                sum,
            });
        }
        let sample_count: usize = self.len();
        let n_train: usize = (sample_count as f32 * train_frac).floor() as usize;
        let n_val: usize =
            ((sample_count as f32 * val_frac).floor() as usize).min(sample_count - n_train);
        let train_end: usize = n_train;
        let val_end: usize = n_train + n_val;

        let train: Dataset = Dataset {
            x: self.x.select_rows(&(0..train_end).collect::<Vec<_>>())?,
            y: self.y.select_rows(&(0..train_end).collect::<Vec<_>>())?,
        };
        let val: Dataset = Dataset {
            x: self
                .x
                .select_rows(&(train_end..val_end).collect::<Vec<_>>())?,
            y: self
                .y
                .select_rows(&(train_end..val_end).collect::<Vec<_>>())?,
        };
        let test: Dataset = Dataset {
            x: self
                .x
                .select_rows(&(val_end..sample_count).collect::<Vec<_>>())?,
            y: self
                .y
                .select_rows(&(val_end..sample_count).collect::<Vec<_>>())?,
        };

        Ok((train, val, test))
    }

    /// Converts `x` to a [`Matrix`] and `y` to a flat `Vec<f32>` (single-output).
    ///
    /// # Errors
    ///
    /// Returns an error if either frame has no numeric columns, or if `y` has
    /// more than one numeric column (use a multi-target conversion instead).
    pub fn into_numeric(&self) -> TuphiumResult<(crate::linalg::Matrix, Vec<f32>)> {
        let x: Matrix = self.x.to_matrix()?;
        let y: Matrix = self.y.to_matrix()?;
        if y.cols == 1 {
            Ok((x, y.data))
        } else {
            Err(TuphiumError::InvalidTargetShape {
                operation: "Dataset::into_numeric",
                expected_columns: 1,
                actual_columns: y.cols,
            })
        }
    }
}
