#include "network_trainer.hh"
#include "../../TradeCore/exceptions.hh"
#include <fstream>

using namespace OpenNN;

namespace {

/*std::shared_ptr<OpenNN::GradientDescent::GradientDescentResults>
trainNetwork(PerformanceFunctional& performanceFunctional)
{
    TrainingStrategy strategy(&performanceFunctional);

    strategy.set_initialization_type(TrainingStrategy::EVOLUTIONARY_ALGORITHM);
    strategy.get_evolutionary_algorithm_pointer()->set_population_size(50);
    strategy.get_evolutionary_algorithm_pointer()->set_maximum_generations_number(150);
    strategy.get_evolutionary_algorithm_pointer()->set_display_period(100);
    strategy.get_evolutionary_algorithm_pointer()->set_display(true);

    strategy.set_main_type(TrainingStrategy::GRADIENT_DESCENT);
    strategy.get_gradient_descent_pointer()->set_maximum_iterations_number(7500);
    strategy.get_gradient_descent_pointer()->set_display_period(100);
    strategy.get_gradient_descent_pointer()->set_performance_goal(0.05);
    strategy.get_gradient_descent_pointer()->set_maximum_time(1800);

    strategy.get_gradient_descent_pointer()
        ->set_reserve_all_training_history(true);

    // TODO: possible memory leak
    return
        std::shared_ptr<OpenNN::GradientDescent::GradientDescentResults>(
        strategy.perform_training().gradient_descent_results_pointer);
}*/

std::shared_ptr<OpenNN::GradientDescent::GradientDescentResults>
    trainNetwork(PerformanceFunctional& performanceFunctional)
{
    GradientDescent algorithm(&performanceFunctional);

    algorithm.set_display(true);
    algorithm.set_display_period(10);
    algorithm.set_maximum_iterations_number(8000);
    algorithm.set_minimum_performance_increase(0.0);
    algorithm.set_reserve_all_training_history(true);
    algorithm.set_maximum_time(3600);
    algorithm.set_performance_goal(0.04);

    return
        std::shared_ptr<OpenNN::GradientDescent::GradientDescentResults>(
        algorithm.perform_training());
}

void saveTrainingResult(
    std::shared_ptr<OpenNN::GradientDescent::GradientDescentResults> results,
    const std::string& resultFileName)
{
    std::ofstream stream;
    stream.open(resultFileName);

    stream << "Iteration, " << "MSE, " << "Training rate, ";
    stream << "Time" << std::endl;

    for (unsigned int i = 0; i < results->performance_history.size(); i++) {
        stream << i << ", ";
        stream << results->performance_history[i] << ", ";
        stream << results->training_rate_history[i] << ", ";
        stream << results->elapsed_time_history[i] << std::endl;
    }
}

} // namespace

void NetworkTrainer::networkFile(const std::string& networkFileName)
{
    networkFile_ = networkFileName;
}

void NetworkTrainer::resultFile(const std::string& resultFileName)
{
    resultFile_ = resultFileName;
}

void NetworkTrainer::trainingDataset(
std::shared_ptr<OpenNN::DataSet> trainingDataset)
{
    trainingSet_ = trainingDataset;
}

std::shared_ptr<OpenNN::NeuralNetwork> NetworkTrainer::trainNetwork()
{
    if (networkFile_.empty())
        throw LogicException("No neural network file given");

    if (!trainingSet_)
        throw LogicException("No training dataset given");

    if (resultFile_.empty())
        throw LogicException("No training result file given");

    Vector<size_t> topology;
    topology.push_back(11);
    topology.push_back(15);
    topology.push_back(7);
    topology.push_back(1);

    Vector<Perceptron::ActivationFunction> func;
    func.push_back(Perceptron::HyperbolicTangent);
    func.push_back(Perceptron::HyperbolicTangent);
    func.push_back(Perceptron::HyperbolicTangent);

    MultilayerPerceptron perceptron(topology);
    perceptron.set_layers_activation_function(func);

    auto network = std::make_shared<OpenNN::NeuralNetwork>(perceptron);
    network->randomize_parameters_uniform(-0.3, 0.3);

    PerformanceFunctional performanceFunctional(
        network.get(), trainingSet_.get());

    auto results = ::trainNetwork(performanceFunctional);
    saveTrainingResult(results, resultFile_);

    network->save(networkFile_);
    return network;
}
