#include <algorithm>
#include <cmath>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <stdexcept>
#include <vector>
#include <boost_tree.h>
#include <genetic_search.h>

int main(int argc, char** argv) {
    try {
        if (argc != 2) {
            throw std::invalid_argument("Usage: boosted-tree-genetic-search output.json");
        }
        const std::vector<std::vector<double>> training{
            {0, 0}, {0, 1}, {1, 0}, {1, 1},
            {0.1, 0.1}, {0.1, 0.9}, {0.9, 0.1}, {0.9, 0.9},
        };
        const std::vector<int> labels{0, 1, 1, 0, 0, 1, 1, 0};
        const std::vector<std::vector<double>> validation{
            {0.25, 0.25}, {0.25, 0.75}, {0.75, 0.25}, {0.75, 0.75},
            {0.35, 0.35}, {0.35, 0.65}, {0.65, 0.35}, {0.65, 0.65},
        };
        const std::vector<std::vector<double>> test{
            {0.05, 0.8}, {0.8, 0.05}, {0.05, 0.05}, {0.8, 0.8},
        };
        const std::vector<int> test_labels{1, 1, 0, 0};
        auto decode = [](const std::vector<double>& genes) {
            BoostTreeParameters params;
            params.num_estimators = 20 + static_cast<unsigned>(80 * genes[0] + 0.5);
            params.learning_rate = 0.03 + 0.17 * genes[1];
            params.max_depth = 1 + static_cast<unsigned>(3 * genes[2] + 0.5);
            return params;
        };
        auto loss = [&](const BoostTreeClassifier& model) {
            double total = 0;
            for (std::size_t i = 0; i < validation.size(); ++i) {
                const auto& classes = model.classes();
                const auto index = std::find(classes.begin(), classes.end(), labels[i])
                                 - classes.begin();
                total -= std::log(std::max(1e-15, model.predict_proba(validation[i])[index]));
            }
            return total / validation.size();
        };
        auto accuracy = [&](const BoostTreeClassifier& model) {
            unsigned correct = 0;
            for (std::size_t i = 0; i < test.size(); ++i) {
                correct += model.predict(test[i]) == test_labels[i];
            }
            return static_cast<double>(correct) / test.size();
        };
        const std::vector<double> initial{0.5, (0.1 - 0.03) / 0.17, 0.0};
        const auto baseline_params = decode(initial);
        BoostTreeClassifier baseline(baseline_params);
        baseline.fit(training, labels);
        opt::GeneticSearch::Config config;
        config.dimensions = 3;
        config.population_size = 12;
        config.generations = 8;
        config.mutation_rate = 0.2;
        config.mutation_sigma = 0.15;
        config.lower_bound = 0;
        config.upper_bound = 1;
        config.seed = 42;
        opt::GeneticSearch search(config);
        const auto winner = search.optimize([&](const std::vector<double>& genes) {
            BoostTreeClassifier candidate(decode(genes));
            candidate.fit(training, labels);
            return loss(candidate);
        }, initial);
        const auto winner_params = decode(winner);
        BoostTreeClassifier candidate(winner_params);
        candidate.fit(training, labels);
        const bool improved = loss(candidate) < loss(baseline);
        const auto& selected = improved ? candidate : baseline;
        const auto selected_params = improved ? winner_params : baseline_params;
        std::ofstream output(argv[1]);
        output.exceptions(std::ios::failbit | std::ios::badbit);
        output << std::setprecision(17);
        output << "{\"grid_size\":60,\"models\":[";
        auto write_model = [&](const BoostTreeClassifier& model,
                               const BoostTreeParameters& params) {
            output << "{\"num_estimators\":" << params.num_estimators
                   << ",\"learning_rate\":" << params.learning_rate
                   << ",\"max_depth\":" << params.max_depth
                   << ",\"validation_loss\":" << loss(model)
                   << ",\"test_accuracy\":" << accuracy(model) << ",\"cells\":[";
            const auto& classes = model.classes();
            const auto class_one = std::find(classes.begin(), classes.end(), 1) - classes.begin();
            for (int row = 0; row < 60; ++row) {
                for (int column = 0; column < 60; ++column) {
                    if (row || column) output << ',';
                    const std::vector<double> point{(column + 0.5) / 60, 1 - (row + 0.5) / 60};
                    output << '[' << model.predict(point) << ','
                           << model.predict_proba(point)[class_one] << ']';
                }
            }
            output << "]}";
        };
        write_model(baseline, baseline_params);
        output << ',';
        write_model(selected, selected_params);
        output << "]}\n";
        output.close();
        std::cout << "Baseline: depth " << baseline_params.max_depth
                  << ", validation loss " << loss(baseline)
                  << ", test accuracy " << accuracy(baseline) << '\n';
        std::cout << "Selected: depth " << selected_params.max_depth
                  << ", trees " << selected_params.num_estimators
                  << ", rate " << selected_params.learning_rate
                  << ", validation loss " << loss(selected)
                  << ", test accuracy " << accuracy(selected) << '\n';
        return 0;
    } catch (const std::exception& error) {
        std::cerr << error.what() << '\n';
        return 1;
    }
}