mod raw_operation {
    use crate::float::Float;

    pub fn zero(out: &mut [Float]) {
        for y in out {
            *y = 0.0
        }
    }

    pub fn inner_product(xs: &[Float], ys: &[Float]) -> Float {
        let mut result: Float = 0.0;
        for (x, y) in xs.iter().zip(ys) {
            result += x * y;
        }
        result
    }

    pub fn scale(xs: &[Float], k: Float, out: &mut [Float]) {
        for (x, y) in xs.iter().zip(out) {
            *y = *x * k
        }
    }

    pub fn add(xs: &[Float], ys: &[Float], out: &mut [Float]) {
        for ((x, y), z) in xs.iter().zip(ys).zip(out) {
            *z = *x + *y
        }
    }

    pub fn sub(xs: &[Float], ys: &[Float], out: &mut [Float]) {
        for ((x, y), z) in xs.iter().zip(ys).zip(out) {
            *z = *x - *y
        }
    }
}

use crate::float::Float;
use rand_distr::{Distribution, Normal};
use std::iter::FromIterator;
use std::ops::{Index, IndexMut};
use std::slice::{Iter, IterMut, SliceIndex};

#[derive(Clone, Debug)]
pub struct Vector(Vec<Float>);

impl<Idx> Index<Idx> for Vector
where
    Idx: SliceIndex<[Float]>,
{
    type Output = Idx::Output;

    fn index(&self, index: Idx) -> &Self::Output {
        &self.0[index]
    }
}

impl<Idx> IndexMut<Idx> for Vector
where
    Idx: SliceIndex<[Float]>,
{
    fn index_mut(&mut self, index: Idx) -> &mut Self::Output {
        &mut self.0[index]
    }
}

impl FromIterator<Float> for Vector {
    fn from_iter<I: IntoIterator<Item = Float>>(iter: I) -> Self {
        let mut v = vec![];
        for x in iter {
            v.push(x)
        }
        Vector(v)
    }
}

impl Vector {
    pub fn iter(&self) -> Iter<'_, Float> {
        self.0.iter()
    }

    pub fn iter_mut(&mut self) -> IterMut<'_, Float> {
        self.0.iter_mut()
    }

    pub fn as_slice(&self) -> &[Float] {
        self.0.as_slice()
    }

    pub fn as_mut_slice(&mut self) -> &mut [Float] {
        self.0.as_mut_slice()
    }

    pub fn copy_from_slice(&mut self, src: &[Float]) {
        self.0.copy_from_slice(src)
    }

    pub fn to_vec(self) -> Vec<Float> {
        self.0
    }
}

impl Vector {
    pub fn new(vector: Vec<Float>) -> Self {
        Self(vector)
    }

    pub fn zero(size: usize) -> Self {
        Self(vec![0.0; size])
    }

    pub fn add_mut(&self, w: &[Float], out: &mut [Float]) {
        raw_operation::add(&self[..], w, out)
    }

    pub fn sub_mut(&self, w: &[Float], out: &mut [Float]) {
        raw_operation::add(&self[..], w, out)
    }

    pub fn scale_mut(&self, k: Float, out: &mut [Float]) {
        raw_operation::scale(&self[..], k, out)
    }

    pub fn inner_product(&self, w: &[Float]) -> Float {
        raw_operation::inner_product(&self[..], w)
    }
}

#[derive(Clone, Debug)]
pub struct ColumnEfficientMatrix {
    pub input_dimension: usize,
    pub output_dimension: usize,
    pub components: Vec<Float>,
}

impl Index<(usize, usize)> for ColumnEfficientMatrix {
    type Output = Float;

    fn index(&self, (column, row): (usize, usize)) -> &Self::Output {
        &self.components[column * self.output_dimension + row]
    }
}

impl IndexMut<(usize, usize)> for ColumnEfficientMatrix {
    fn index_mut(&mut self, (column, row): (usize, usize)) -> &mut Self::Output {
        &mut self.components[column * self.output_dimension + row]
    }
}

impl ColumnEfficientMatrix {
    pub fn from_rows(rows: Vec<Vec<Float>>) -> Self {
        let output_dimension = rows.len();
        if output_dimension == 0 {
            Self {
                input_dimension: 0,
                output_dimension: 0,
                components: vec![],
            }
        } else {
            let input_dimension = rows[0].len();
            let mut components = Vec::with_capacity(input_dimension * output_dimension);
            for i in 0..input_dimension {
                for j in 0..output_dimension {
                    components.push(rows[j][i])
                }
            }
            Self {
                input_dimension,
                output_dimension,
                components,
            }
        }
    }

    pub fn zero(input_dimension: usize, output_dimension: usize) -> ColumnEfficientMatrix {
        Self {
            input_dimension,
            output_dimension,
            components: vec![0.0; input_dimension * output_dimension],
        }
    }

    pub fn random_with_normal_distribution(
        input_dimension: usize,
        output_dimension: usize,
        distribution: Normal<Float>,
    ) -> ColumnEfficientMatrix {
        let mut components = Vec::with_capacity(input_dimension * output_dimension);
        for _ in 0..input_dimension {
            for _ in 0..output_dimension {
                components.push(distribution.sample(&mut rand::thread_rng()))
            }
        }

        Self {
            input_dimension,
            output_dimension,
            components,
        }
    }

    pub fn zero_mut(&mut self) {
        for a in &mut self.components {
            *a = 0.0
        }
    }

    pub fn add_scaled_mut(&self, k: Float, b: &Self, c: &mut Self) {
        for ((a, b), c) in self
            .components
            .iter()
            .zip(&b.components)
            .zip(&mut c.components)
        {
            *c = *a + *b * k
        }
    }

    pub fn add_to_self_scaled_mut(&mut self, k: Float, b: &Self) {
        for (a, b) in self.components.iter_mut().zip(&b.components) {
            *a += *b * k
        }
    }

    pub fn apply_mut(&self, v: &[Float], out: &mut [Float]) {
        for j in 0..self.output_dimension {
            let mut result: Float = 0.0;
            for i in 0..self.input_dimension {
                result += self[(i, j)] * v[i]
            }
            out[j] = result;
        }
    }

    // Apply the transpose
    pub fn coapply_mut(&self, w: &[Float], out: &mut [Float]) {
        for i in 0..self.input_dimension {
            let start_index = i * self.output_dimension;
            out[i] = raw_operation::inner_product(
                &self.components[start_index..start_index + self.output_dimension],
                w,
            );
        }
    }

    pub fn drop_first_column_coapply_mut(&self, w: &[Float], out: &mut [Float]) {
        for i in 1..self.input_dimension {
            let start_index = i * self.output_dimension;
            out[i - 1] = raw_operation::inner_product(
                &self.components[start_index..start_index + self.output_dimension],
                w,
            );
        }
    }
}

#[derive(Clone, Debug)]
pub struct DiagonalMatrix<'a> {
    pub diagonal: &'a [Float],
}

impl<'a> DiagonalMatrix<'a> {
    pub fn new(diagonal: &'a [Float]) -> DiagonalMatrix<'a> {
        Self { diagonal }
    }

    pub fn apply_mut(&self, v: &[Float], out: &mut [Float]) {
        for (i, (d, x)) in self.diagonal.iter().zip(v.iter()).enumerate() {
            out[i] = d * x;
        }
    }
}

// Given vector `w : W` and a covector `f : V -> Float`,
// the linear map `w tensor f : V -> W` computes `v : V ~> f(v) * w`
#[derive(Clone, Debug)]
pub struct VectorTensorCovectorMatrix<'a> {
    pub output_vector: &'a [Float],  // w : W
    pub input_covector: &'a [Float], // f : V -> Float
}

impl<'a> VectorTensorCovectorMatrix<'a> {
    pub fn new(
        output_vector: &'a [Float],
        input_covector: &'a [Float],
    ) -> VectorTensorCovectorMatrix<'a> {
        Self {
            output_vector,
            input_covector,
        }
    }

    pub fn apply_mut(&self, v: &[Float], out: &mut [Float]) {
        let scalar = raw_operation::inner_product(self.input_covector, v);
        raw_operation::scale(self.output_vector, scalar, out)
    }

    pub fn add_to_mut(&self, matrix: &mut ColumnEfficientMatrix) {
        // TODO: Surely this can be optimized by iterating over the columns of the matrix directly
        for (i, x) in self.input_covector.iter().enumerate() {
            for (j, y) in self.output_vector.iter().enumerate() {
                matrix[(i, j)] += x * y
            }
        }
    }
}
