mod env;
mod float;
mod linear_algebra;
mod neural_network;
mod preprocessing;
mod transforms;

use preprocessing::{dataset_from_file, export_distributions_to, export_to};
use std::io;

use crate::neural_network::{NNPoint, NeuralNetworkInTraining, NeuralNetworkParameters};

fn main() -> Result<(), io::Error> {
    let size_of_full_training_dataset = 60000;
    let size_of_training_dataset = 50000;
    let size_of_validation_dataset = 10000;
    let size_of_testing_dataset = 10000;

    println!("Importing datasets...");
    let full_training_dataset: Vec<NNPoint> =
        dataset_from_file(env::TRAIN_CSVS, size_of_full_training_dataset)?;
    let mut training_dataset: Vec<NNPoint> = {
        let mut training_dataset = Vec::with_capacity(size_of_training_dataset);
        training_dataset.extend_from_slice(&full_training_dataset[0..size_of_training_dataset]);
        training_dataset
    };
    let validation_dataset: Vec<NNPoint> = {
        let mut validation_dataset = Vec::with_capacity(size_of_validation_dataset);
        validation_dataset.extend_from_slice(
            &full_training_dataset
                [size_of_training_dataset..size_of_training_dataset + size_of_validation_dataset],
        );
        validation_dataset
    };
    let test_dataset: Vec<NNPoint> = dataset_from_file(env::TEST_CSVS, size_of_testing_dataset)?;

    // let mut nn = NeuralNetworkInTraining::new(vec![env::NUMBER_OF_PIXELS_PER_IMAGE, 70, 10]);
    // let mut nn = NeuralNetworkInTraining::new(vec![env::NUMBER_OF_PIXELS_PER_IMAGE, 150, 10]);
    // let mut nn = NeuralNetworkInTraining::new(vec![env::NUMBER_OF_PIXELS_PER_IMAGE, 60, 40, 10]);
    // let mut nn = NeuralNetworkInTraining::new(vec![env::NUMBER_OF_PIXELS_PER_IMAGE, 150, 10]);
    // let mut nn = NeuralNetworkInTraining::new(vec![env::NUMBER_OF_PIXELS_PER_IMAGE, 60, 40, 10]); // batch=20, rate=2 is pretty good. Seems to reliably reach 89%
    // let mut nn = NeuralNetworkInTraining::new(vec![env::NUMBER_OF_PIXELS_PER_IMAGE, 80, 60, 40, 10]); // batch=20, rate=2, pretty good like 88 % then suddenly drops to 10%
    // let mut nn = NeuralNetworkInTraining::new(vec![env::NUMBER_OF_PIXELS_PER_IMAGE, 150, 150, 10]); // this gets me over 90%, nice
    let mut nn = NeuralNetworkInTraining::new(vec![env::NUMBER_OF_PIXELS_PER_IMAGE, 150, 150, 10]);

    println!("Begin training");

    let params = NeuralNetworkParameters {
        epochs: 30,
        batch_size: 30,
        learning_rate: 2.00,
    };
    nn.train(params, &mut training_dataset, Some(&validation_dataset));
    // nn.train(params, &mut training_dataset, None);

    nn.show_accuracy_on(&full_training_dataset, "training");
    nn.show_accuracy_on(&test_dataset, "test");

    let predictions_on_full_training_set = nn.test(&full_training_dataset);
    let predictions_on_test_set = nn.test(&test_dataset);
    println!("Exporting to {}", env::TRAIN_PREDICTIONS_CSV);
    export_to(
        &predictions_on_full_training_set,
        env::TRAIN_PREDICTIONS_CSV,
    )?;
    println!("Exporting to {}", env::TEST_PREDICTIONS_CSV);
    export_to(&predictions_on_test_set, env::TEST_PREDICTIONS_CSV)?;

    // TODO: Comment this out
    // let distributions_on_full_training_set = nn.test_distributions(&full_training_dataset);
    // let distributions_on_test_set = nn.test_distributions(&test_dataset);
    // export_distributions_to(&distributions_on_full_training_set, env::TRAIN_DISTRIBUTIONS_CSV)?;
    // export_distributions_to(&distributions_on_test_set, env::TEST_DISTRIBUTIONS_CSV)?;

    Ok(())
}
