extern crate rust_nn;

use std::error::Error;
use std::f64::consts::PI;

use ndarray::Array1;
use ndarray_rand::rand_distr::Uniform;
use ndarray_rand::RandomExt;
use plotters::prelude::*;
use rust_nn::functions::{activation_functions, loss_functions};
use rust_nn::layers::activation_layer::ActivationLayer;
use rust_nn::layers::fc_layer::{FCLayer, Initializer};
use rust_nn::Network;

fn main() {
    // training data
    let training_interval = (0.0f64, 2.0f64 * PI);
    let steps = 100000;
    let training_values = Array1::random(
        steps,
        Uniform::new(training_interval.0, training_interval.1),
    )
    .to_vec();
    let mut x_train = Vec::with_capacity(steps);
    let mut y_train = Vec::with_capacity(steps);
    for x in training_values {
        x_train.push(Array1::from_elem(1usize, x));
        y_train.push(Array1::from_elem(1usize, x.sin()));
    }
    // test data
    let test_steps = 1000;
    let interval_length = training_interval.1 - training_interval.0;
    let step_size = interval_length / test_steps as f64;
    let testing_values = Array1::range(training_interval.0, training_interval.1, step_size);
    let mut x_test = Vec::with_capacity(test_steps);
    let mut y_test_true = Vec::with_capacity(test_steps);
    for x in testing_values {
        x_test.push(Array1::from_elem(1usize, x));
        y_test_true.push(Array1::from_elem(1usize, x.sin()));
    }

    // initialize neural network
    let mut network = Network::new(loss_functions::Type::MSE);

    // add layers
    network.add_layer(Box::new(FCLayer::new(
        8,
        Initializer::GaussianWFactor(0.0, 1.0, 0.1),
        Initializer::GaussianWFactor(0.0, 1.0, 0.1),
    )));
    network.add_layer(Box::new(ActivationLayer::new(
        activation_functions::Type::Gelu,
    )));
    network.add_layer(Box::new(FCLayer::new(
        8,
        Initializer::GaussianWFactor(0.0, 1.0, 0.1),
        Initializer::GaussianWFactor(0.0, 1.0, 0.1),
    )));
    network.add_layer(Box::new(ActivationLayer::new(
        activation_functions::Type::Gelu,
    )));
    network.add_layer(Box::new(FCLayer::new(
        1,
        Initializer::GaussianWFactor(0.0, 1.0, 0.1),
        Initializer::GaussianWFactor(0.0, 1.0, 0.1),
    )));

    // train network on training data
    network.fit(&x_train, &y_train, 100, 0.05, true);

    // predict test dataset
    let y_test_pred = network.predict(&x_test);

    // show results
    if let Ok(()) = draw_results(&training_interval, &x_test, &y_test_true, &y_test_pred) {
        println!("results can be seen in ./examples/sine.png");
    } else {
        println!("failed to draw results");
    }
}

fn draw_results(
    training_interval: &(f64, f64),
    x_test: &[Array1<f64>],
    y_test_true: &[Array1<f64>],
    y_test_pred: &[Array1<f64>],
) -> Result<(), Box<dyn Error>> {
    // create the chart
    let buf = BitMapBackend::new("./examples/sine.png", (800, 600)).into_drawing_area();
    buf.fill(&WHITE)?;
    let mut chart = ChartBuilder::on(&buf)
        //.caption("sin(x)", ("sans-serif", 30))
        .x_label_area_size(30)
        .y_label_area_size(30)
        .build_cartesian_2d(training_interval.0..training_interval.1, -1.0f64..1.0f64)?;

    chart
        .configure_mesh()
        .disable_x_mesh()
        .disable_y_mesh()
        .draw()?;

    // add the first plot
    let data1: Vec<(f64, f64)> = x_test
        .iter()
        .zip(y_test_true.iter())
        .map(|(x, y)| (x[0], y[0]))
        .collect();
    chart
        .draw_series(LineSeries::new(data1, &RED))
        .unwrap()
        .label("true values")
        .legend(|(x, y)| PathElement::new(vec![(x, y), (x + 1, y)], RED));

    // add the second plot
    let data2: Vec<(f64, f64)> = x_test
        .iter()
        .zip(y_test_pred.iter())
        .map(|(x, y)| (x[0], y[0]))
        .collect();
    chart
        .draw_series(LineSeries::new(data2, &BLUE))
        .unwrap()
        .label("predicted values")
        .legend(|(x, y)| PathElement::new(vec![(x, y), (x + 1, y)], BLUE));

    Ok(())
}