use crate::env;
use crate::float::Float;
use crate::neural_network::NNPoint;
use std::fs::File;
use std::io;
use std::io::{BufRead, BufReader, BufWriter, Write};

#[derive(Debug)]
enum ParsingError {
    CouldNotParseLabel,
    CouldNotParseVector,
    ImageHasWrongSize,
}

fn parse_dataset(
    labels_buffer: impl BufRead,
    vectors_buffer: impl BufRead,
    number_of_points: usize,
) -> Result<Result<Vec<NNPoint>, ParsingError>, io::Error> {
    let mut output = vec![];
    for (label_result, vector_result) in labels_buffer
        .lines()
        .zip(vectors_buffer.lines())
        .take(number_of_points)
    {
        let label_str = label_result?;
        let image_str = vector_result?;

        let label: u8 = match label_str.parse() {
            Ok(label) => label,
            Err(_) => return Ok(Err(ParsingError::CouldNotParseLabel)),
        };

        let mut image: Vec<u8> = vec![];
        for str in image_str.split(',') {
            match str.parse() {
                Ok(pixel_value) => {
                    image.push(pixel_value);
                }
                Err(_) => return Ok(Err(ParsingError::CouldNotParseVector)),
            }
        }

        if image.len() != env::NUMBER_OF_PIXELS_PER_IMAGE {
            return Ok(Err(ParsingError::ImageHasWrongSize));
        }

        output.push(NNPoint::new(label, normalize_input(&image)));
    }
    Ok(Ok(output))
}

fn average(vec: &[u8]) -> Float {
    let mut result: u32 = 0; // the result is bounded by NUMBER_OF_PIXELS_PER_IMAGE * 255 ~~ 200 k
    for x in vec {
        result += *x as u32;
    }
    (result as Float) / (vec.len() as Float)
}

fn variance(vec: &[Float]) -> Float {
    // assumes that the average of vec is 0
    let mut result: Float = 0.0;
    for x in vec {
        result += x * x;
    }
    result
}

fn normalize_input(vec: &[u8]) -> Vec<Float> {
    let average = average(vec);
    let mut result: Vec<Float> = vec.iter().map(|x| (*x as Float) - average).collect();
    let stddev = variance(&result).sqrt();
    for i in &mut result {
        *i /= stddev;
    }
    result
}

pub fn dataset_from_file(
    (labels_file_path, vectors_file_path): (&str, &str),
    number_of_points: usize,
) -> Result<Vec<NNPoint>, io::Error> {
    let (labels_file, vectors_file): (File, File) = (
        File::open(labels_file_path)?,
        File::open(vectors_file_path)?,
    );

    match parse_dataset(
        BufReader::new(labels_file),
        BufReader::new(vectors_file),
        number_of_points,
    )? {
        Ok(points) => Ok(points),
        Err(parsing_error) => {
            println!("ERROR: {:?}", parsing_error);
            todo!()
        }
    }
}

pub fn export_to(outputs: &[usize], file_path: &str) -> Result<(), io::Error> {
    let mut file = BufWriter::new(File::create(file_path)?);
    for x in outputs {
        writeln!(file, "{}", x)?
    }
    Ok(())
}

pub fn export_distributions_to(outputs: &[Vec<Float>], file_path: &str) -> Result<(), io::Error> {
    let mut file = BufWriter::new(File::create(file_path)?);
    for ps in outputs {
        let mut s = "".to_string();
        for p in ps {
            let p_str = format!("{}, ", p);
            s += &p_str
        }
        writeln!(file, "[{}]", s)?
    }
    Ok(())
}
