extern crate smyl;

use std::fs::File;
use std::io::{BufWriter, Write};
use std::time::Instant;

use rand::seq::SliceRandom;
use rand::thread_rng;
use smyl::prelude::*;

const IMAGE_HEIGHT: usize = 28;
const IMAGE_WIDTH: usize = 28;

#[derive(Debug, Copy, Clone)]
struct Label(u8);

#[derive(Debug, Copy, Clone)]
struct Image([u8; IMAGE_HEIGHT * IMAGE_WIDTH]);

impl Read for Image {
    fn read<T>(reader: &mut Reader<T>) -> std::io::Result<Self> {
        let mut pixels = [0; IMAGE_HEIGHT * IMAGE_WIDTH];
        reader.read_buf(&mut pixels)?;

        Ok(Self(pixels))
    }
}

impl Idx for Image {
    #[allow(clippy::cast_possible_truncation)]
    const EXPECTED_DIMENSIONS: &'static [u32] = &[IMAGE_HEIGHT as u32, IMAGE_WIDTH as u32];
    const MAGIC_WORD: u32 = 0x0000_0803;
}

impl Read for Label {
    fn read<T>(reader: &mut Reader<T>) -> std::io::Result<Self> { Ok(Self(reader.read_other()?)) }
}

impl Idx for Label {
    const EXPECTED_DIMENSIONS: &'static [u32] = &[];
    const MAGIC_WORD: u32 = 0x0000_0801;
}

const LEARNING_RATE: f64 = 0.01;
const BATCH: usize = 1;

struct MNISTDataset {
    train_labels: Vec<Matrix<f64>>,
    train_inputs: Vec<Matrix<f64>>,
    test_labels:  Vec<Matrix<f64>>,
    test_inputs:  Vec<Matrix<f64>>,
}

impl MNISTDataset {
    const TESTING_SAMPLE_COUNT: usize = 10_000;
    const TRAINING_SAMPLE_COUNT: usize = 60_000;

    pub fn read() -> std::io::Result<Self> {
        let mut train_labels = Vec::with_capacity(Self::TRAINING_SAMPLE_COUNT);
        let mut train_inputs = Vec::with_capacity(Self::TRAINING_SAMPLE_COUNT);
        let mut test_labels = Vec::with_capacity(Self::TESTING_SAMPLE_COUNT);
        let mut test_inputs = Vec::with_capacity(Self::TESTING_SAMPLE_COUNT);

        let mut train_labels_reader: Reader<Label> = Reader::try_new(include_bytes!(
            "../../datasets/mnist/train-labels-idx1-ubyte.gz"
        ))?;
        let mut train_images_reader: Reader<Image> = Reader::try_new(include_bytes!(
            "../../datasets/mnist/train-images-idx3-ubyte.gz"
        ))?;
        let mut test_labels_reader: Reader<Label> = Reader::try_new(include_bytes!(
            "../../datasets/mnist/t10k-labels-idx1-ubyte.gz"
        ))?;
        let mut test_images_reader: Reader<Image> = Reader::try_new(include_bytes!(
            "../../datasets/mnist/t10k-images-idx3-ubyte.gz"
        ))?;

        for i in 0..Self::TRAINING_SAMPLE_COUNT {
            let mut label = Matrix::zero(1, 10);
            label.set(0, train_labels_reader.read()?.0 as usize, 1.0);

            let image = train_images_reader.read()?.0;
            let input = Matrix::from([image.map(|p| f64::from(p) / 255.0)]);

            train_labels.insert(i, label);
            train_inputs.insert(i, input);
        }

        for i in 0..Self::TESTING_SAMPLE_COUNT {
            let mut label = Matrix::zero(1, 10);
            label.set(0, test_labels_reader.read()?.0 as usize, 1.0);

            let image = test_images_reader.read()?.0;
            let input = Matrix::from([image.map(|p| f64::from(p) / 255.0)]);

            test_labels.insert(i, label);
            test_inputs.insert(i, input);
        }

        Ok(Self {
            train_labels,
            train_inputs,
            test_labels,
            test_inputs,
        })
    }
}

#[allow(clippy::large_stack_frames)]
fn main() -> std::io::Result<()> {
    let mut profiling = false;
    for arg in std::env::args() {
        if arg == "--profiling" {
            profiling = true;
        }
    }
    if profiling {
        println!("Profiling mode is enabled.");
    }

    println!("Retrieving mnist dataset");
    let dataset = MNISTDataset::read()?;

    let error_margin = f64::EPSILON;

    let mut ann = ann!(784, 60, 60, 10);
    let mut rng = thread_rng();

    let training_sample_count = if profiling {
        MNISTDataset::TRAINING_SAMPLE_COUNT / 6
    } else {
        MNISTDataset::TRAINING_SAMPLE_COUNT
    };

    let mut j = 1;
    loop {
        let mut sequence: Vec<usize> = (0..training_sample_count).collect();
        sequence.shuffle(&mut rng);
        let start = Instant::now();

        for (i, index) in sequence.iter().enumerate() {
            let label = dataset.train_labels[*index].clone();
            let input = dataset.train_inputs[*index].clone();

            ann.forward(input);
            ann.backward(label);

            #[allow(clippy::modulo_one)]
            if i % BATCH == 0 {
                ann.apply_chunk_gradient(LEARNING_RATE);
            }

            #[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)]
            if i % 100 == 0 && i > 0 {
                let elapsed = start.elapsed();
                let speed = i as f64 / elapsed.as_secs_f64();
                let eta = ((training_sample_count as f64 - i as f64) / speed) as i64;
                println!(
                    "Training {i} / {} | Speed: {} /s | ETA: {}{} s ",
                    training_sample_count,
                    speed.round(),
                    if eta >= 60 {
                        format!("{} min ", eta / 60)
                    } else {
                        String::new()
                    },
                    eta % 60
                );
            }
        }
        ann.apply_chunk_gradient(LEARNING_RATE);

        let mut error_sum = 0;
        for i in 0..MNISTDataset::TESTING_SAMPLE_COUNT {
            let label = dataset.test_labels[i].clone();
            let input = dataset.test_inputs[i].clone();

            let output_signal = ann.forward(input);
            let mut max = 0.0;
            let mut output = 0;
            let mut expected = 0;
            for i in 0..10 {
                if (label.get(0, i) - 1.0).abs() < error_margin {
                    expected = i;
                }

                if output_signal.get(0, i) > max {
                    max = output_signal.get(0, i);
                    output = i;
                }
            }

            if output != expected {
                error_sum += 1;
            }

            if i % 100 == 0 {
                println!("Testing {i} / {}", MNISTDataset::TESTING_SAMPLE_COUNT);
            }
        }

        #[allow(clippy::cast_precision_loss)]
        let average_error = f64::from(error_sum) / MNISTDataset::TESTING_SAMPLE_COUNT as f64;
        println!(
            "Passage n°{j}: {}",
            100.0 - (average_error * 10000.0).round() / 100.0
        );

        let file = File::options()
            .create(true)
            .write(true)
            .truncate(true)
            .open("ann.json")?;

        let mut writer = BufWriter::new(file);
        serde_json::to_writer(&mut writer, &ann)?;
        writer.flush()?;

        j += 1;

        if profiling {
            break;
        }
    }

    Ok(())
}