#include "forecasting_data.hh"
#include "../../Indicators/Source/sma.hh"
#include "../../3rdParty/opennn/data_set.h"
#include "../../TradeCore/chart_candle.hh"
#include "../../TradeCore/exceptions.hh"
#include <QDebug>
#include <random>
#include <deque>
#include <chrono>

namespace {

bool checkProbability(double probability, int& seed)
{
    std::mt19937 engine;
    std::uniform_real_distribution<> dist(0, 1);
    engine.seed(++seed);
    return probability >= dist(engine);
}

std::vector<double> calculateAverages(
    const std::deque<ChartCandle>& buffer)
{
    std::vector<double> averages;

    for (int i = 0; i < 12; i++) {
        double sum = 0.0;

        for (int j = 0; j < 4; j++)
            sum += buffer[i*4 + j].close();

        averages.push_back(sum / 4.0);
    }

    return averages;
}

} // namespace

ForecastingData::ForecastingData()
    : validation_(false)
{}

void ForecastingData::data(std::shared_ptr<std::vector<ChartCandle>> data)
{
    data_ = data;
}

void ForecastingData::trainingFile(const std::string& trainingFileName)
{
    trainingFile_ = trainingFileName;
}

void ForecastingData::testingFile(const std::string& testingFileName)
{
    testingFile_ = testingFileName;
}

void ForecastingData::priceBounds(
    const std::vector<StatisticalClass>& bounds)
{
    bounds_ = bounds;
}

void ForecastingData::dateTimeFilter(DateTimeFilter filter)
{
    filter_ = filter;
}

void ForecastingData::validation(bool validation)
{
    validation_ = validation;
}

std::pair<OpenNN::DataSet, OpenNN::DataSet>
ForecastingData::makeDataset() const
{
    if (!data_)
        throw LogicException("No price data given");

    if (trainingFile_.empty())
        throw LogicException("No training dataset file given");

    if (bounds_.empty())
        throw LogicException("No price boundaries given");

    if (!filter_)
        throw LogicException("No date time filter given");

    auto datasets = generateDataset();

    datasets.first.set_data_file_name(trainingFile_);
    datasets.first.save_data();

    if (!testingFile_.empty()) {
        datasets.second.set_data_file_name(testingFile_);
        datasets.second.save_data();
    }

    return datasets;
}

std::pair<OpenNN::DataSet, OpenNN::DataSet>
ForecastingData::generateDataset() const
{
    std::vector<OpenNN::Vector<double>> patterns;

    AI::SMA currentAverage(9);
    AI::SMA futureAverage(9);
    std::deque<ChartCandle> buffer;

    auto duration =
        std::chrono::steady_clock::now().time_since_epoch();

    static int seed = std::chrono::duration_cast<
        std::chrono::milliseconds>(duration).count();

    double weakBearishProbability =
        ((bounds_[0].count() + bounds_[2].count() +
         bounds_[4].count()) / 3.0) / bounds_[1].count();

    double weakBullishProbability =
        ((bounds_[0].count() + bounds_[2].count() +
         bounds_[4].count()) / 3.0) / bounds_[3].count();

    if (validation_) {
        weakBearishProbability = 1.0;
        weakBullishProbability = 1.0;
    }

    for (unsigned int i = 0; i < 48; i++)
        buffer.push_front((*data_)[i]);

    for (unsigned int i = 48; i < data_->size()-12; i++) {
        currentAverage((*data_)[i].close());
        futureAverage((*data_)[i+12].close());

        buffer.push_front((*data_)[i]);
        buffer.pop_back();
        assert(buffer.size() == 48);

        if (!filter_((*data_)[i].dateTime()))
            continue;

        auto averages = calculateAverages(buffer);
        OpenNN::Vector<double> instance;

        for (int j = averages.size() - 1; j > 0; j--)
            instance.push_back((averages[j] - averages[j-1]) * 1000);

        double diff = futureAverage() - currentAverage();       

        if (bounds_[0].isMember(diff)) {
            instance.push_back(-1.0);
            patterns.push_back(instance);
        }
        else if (bounds_[4].isMember(diff)) {
            instance.push_back(1.0);
            patterns.push_back(instance);
        }
        else if (bounds_[2].isMember(diff)) {
            instance.push_back(0.0);
            patterns.push_back(instance);
        }
        else if (
            bounds_[1].isMember(diff) &&
            checkProbability(weakBearishProbability, seed))
        {
            instance.push_back(bounds_[1].interpolate(diff, -1, 0));
            patterns.push_back(instance);
        }
        else if (
            bounds_[3].isMember(diff) &&
            checkProbability(weakBullishProbability, seed))
        {
            instance.push_back(bounds_[3].interpolate(diff, 0, 1));
            patterns.push_back(instance);
        }
    }

    std::vector<OpenNN::Vector<double>> trainingPatterns;
    std::vector<OpenNN::Vector<double>> testingPatterns;

    if (!testingFile_.empty()) {
        for (auto& pattern : patterns) {
            if (checkProbability(0.85, seed))
                trainingPatterns.push_back(pattern);
            else
                testingPatterns.push_back(pattern);
        }

        OpenNN::DataSet trainingDataset(11, 1, trainingPatterns.size());
        OpenNN::DataSet testingDataset(11, 1, testingPatterns.size());

        for (unsigned int i = 0; i < trainingPatterns.size(); i++)
            trainingDataset.set_instance(i, trainingPatterns[i]);

        for (unsigned int i = 0; i < testingPatterns.size(); i++)
            testingDataset.set_instance(i, testingPatterns[i]);

        return std::make_pair(trainingDataset, testingDataset);
    }
    else {
        OpenNN::DataSet trainingDataset(11, 1, patterns.size());

        for (unsigned int i = 0; i < patterns.size(); i++)
            trainingDataset.set_instance(i, patterns[i]);

        return std::make_pair(trainingDataset, OpenNN::DataSet(11, 1, 1));
    }
}
