use serde::{Serialize, Deserialize};

use crate::error::{TuphiumError, TuphiumResult};

#[derive(Serialize, Deserialize, Clone, Debug, Default)]
pub struct Matrix {
    pub rows: usize,
    pub cols: usize,
    pub data: Vec<f32>,
}

impl Matrix {
    /// Creates a `rows × cols` matrix filled with zeros.
    pub fn zeros(rows: usize, cols: usize) -> TuphiumResult<Self> {
        let len = rows
            .checked_mul(cols)
            .ok_or(TuphiumError::DimensionOverflow {
                operation: "Matrix::zeros",
                rows,
                cols,
            })?;
        Ok(Self {
            rows,
            cols,
            data: vec![0.0; len],
        })
    }

    /// Builds a `rows × cols` matrix from a flat, row-major `data` buffer.
    ///
    /// Element `(r, c)` is read from `data[r * cols + c]`.
    ///
    /// # Panics
    ///
    /// Panics if `data.len()` is not exactly `rows * cols`.
    pub fn from_data(rows: usize, cols: usize, data: Vec<f32>) -> TuphiumResult<Self> {
        let expected = rows
            .checked_mul(cols)
            .ok_or(TuphiumError::DimensionOverflow {
                operation: "Matrix::from_data",
                rows,
                cols,
            })?;
        if data.len() != expected {
            return Err(TuphiumError::LengthMismatch {
                operation: "Matrix::from_data",
                expected,
                actual: data.len(),
            });
        }
        Ok(Self { rows, cols, data })
    }

    /// Returns the element at row `r`, column `c`.
    ///
    /// # Panics
    ///
    /// Panics if the index falls outside the backing buffer.
    #[inline]
    pub fn get(&self, r: usize, c: usize) -> TuphiumResult<f32> {
        let index = self.checked_index(r, c, "Matrix::get")?;
        Ok(self.data[index])
    }

    /// Writes `v` into row `r`, column `c`.
    ///
    /// # Panics
    ///
    /// Panics if the index falls outside the backing buffer.
    #[inline]
    pub fn set(&mut self, r: usize, c: usize, v: f32) -> TuphiumResult<()> {
        let index = self.checked_index(r, c, "Matrix::set")?;
        self.data[index] = v;
        Ok(())
    }

    /// Returns row `i` as a contiguous slice of length `cols`.
    ///
    /// # Panics
    ///
    /// Panics if `i >= rows`.
    #[inline]
    pub fn row(&self, i: usize) -> TuphiumResult<&[f32]> {
        if i >= self.rows {
            return Err(TuphiumError::RowIndexOutOfBounds {
                operation: "Matrix::row",
                row: i,
                rows: self.rows,
            });
        }
        let start = i * self.cols;
        Ok(&self.data[start..start + self.cols])
    }

    /// Matrix–vector product `y = self · x`.
    ///
    /// `x` must have length `cols`; the result has length `rows`.
    ///
    /// # Panics
    ///
    /// Panics if `x.len() != cols`.
    pub fn matvec(&self, x: &[f32]) -> TuphiumResult<Vec<f32>> {
        if x.len() != self.cols {
            return Err(TuphiumError::ShapeMismatch {
                operation: "Matrix::matvec",
                expected: vec![self.cols],
                actual: vec![x.len()],
            });
        }
        let mut result = vec![0.0; self.rows];
        for row in 0..self.rows {
            let row_start = row * self.cols;
            let mut sum = 0.0;
            for col in 0..self.cols {
                sum += self.data[row_start + col] * x[col];
            }
            result[row] = sum;
        }
        Ok(result)
    }

    /// Transposed matrix–vector product `y = selfᵀ · v`.
    ///
    /// `v` must have length `rows`; the result has length `cols`. Used in the
    /// backward pass to propagate gradients to a layer's inputs.
    ///
    /// # Panics
    ///
    /// Panics if `v.len() != rows`.
    pub fn t_matvec(&self, v: &[f32]) -> TuphiumResult<Vec<f32>> {
        if v.len() != self.rows {
            return Err(TuphiumError::ShapeMismatch {
                operation: "Matrix::t_matvec",
                expected: vec![self.rows],
                actual: vec![v.len()],
            });
        }
        let mut result = vec![0.0; self.cols];
        for row in 0..self.rows {
            let row_start = row * self.cols;
            let v_row = v[row];
            for col in 0..self.cols {
                result[col] += self.data[row_start + col] * v_row;
            }
        }
        Ok(result)
    }

    /// Accumulates the outer product into the matrix: `self += a ⊗ b`.
    ///
    /// `a` indexes rows (length `rows`) and `b` indexes columns (length `cols`).
    /// Used in the backward pass to accumulate weight gradients (`dz ⊗ x`).
    ///
    /// # Panics
    ///
    /// Panics if `a.len() != rows` or `b.len() != cols`.
    pub fn add_outer(&mut self, a: &[f32], b: &[f32]) -> TuphiumResult<()> {
        if a.len() != self.rows || b.len() != self.cols {
            return Err(TuphiumError::ShapeMismatch {
                operation: "Matrix::add_outer",
                expected: vec![self.rows, self.cols],
                actual: vec![a.len(), b.len()],
            });
        }
        for row in 0..self.rows {
            let row_start = row * self.cols;
            let a_row = a[row];
            for col in 0..self.cols {
                self.data[row_start + col] += a_row * b[col];
            }
        }
        Ok(())
    }

    /// Overwrites every element with `val`.
    pub fn fill(&mut self, val: f32) {
        for element in self.data.iter_mut() {
            *element = val;
        }
    }

    /// Uploads the matrix to the compute device. No-op on the CPU backend.
    ///
    /// Exists so code can be written backend-agnostically; the GPU backend
    /// performs a real host-to-device transfer here.
    pub fn sync_to_device(&mut self) -> TuphiumResult<()> {
        Ok(())
    }

    /// Downloads the matrix from the compute device. No-op on the CPU backend.
    ///
    /// The GPU backend copies device memory back into `data` here.
    pub fn sync_from_device(&mut self) -> TuphiumResult<()> {
        Ok(())
    }

    /// Marks the host buffer as modified so it is re-uploaded on next use.
    /// No-op on the CPU backend.
    pub fn mark_dirty(&mut self) {}

    fn checked_index(
        &self,
        row: usize,
        col: usize,
        operation: &'static str,
    ) -> TuphiumResult<usize> {
        if row >= self.rows || col >= self.cols {
            return Err(TuphiumError::IndexOutOfBounds {
                operation,
                row,
                col,
                rows: self.rows,
                cols: self.cols,
            });
        }
        row.checked_mul(self.cols)
            .and_then(|base| base.checked_add(col))
            .filter(|&index| index < self.data.len())
            .ok_or(TuphiumError::IndexOutOfBounds {
                operation,
                row,
                col,
                rows: self.rows,
                cols: self.cols,
            })
    }
}
