/**
 * DESCRIPTION OF THE FILE
 *
 * @author Michal Kravčenko
 * @date 2.7.18 -
 */

#include <cstdlib>
#include <ctime>
#include <cmath>
#include <set>
#include <stdexcept>
#include <random>
#include <iterator>
#include <algorithm>
#include <iostream>
#include <format.hpp>

#include "message.h"
#include "../Network/NeuralNetwork.h"
#include "../DataSet/DataSet.h"
#include "exceptions.h"

#include "ParticleSwarm.h"

/**
 * TODO
 * domain_bound out_of_range check
 * @param f_dim
 * @param domain_bounds
 * @param F
 */


void Particle::randomize_coordinates() {

    std::random_device seeder;
    std::mt19937 gen(seeder());
    for(unsigned int i = 0; i < this->coordinate_dim; ++i){
        std::uniform_real_distribution<double> dist_coord(this->domain_bounds->at(2 * i), this->domain_bounds->at(2 * i + 1));
        (*this->coordinate)[i] = dist_coord(gen);
    }
}

void Particle::randomize_parameters() {

    std::random_device seeder;
    std::mt19937 gen(seeder());
    std::uniform_real_distribution<double> dist_vel(0.5, 1.0);
    this->r1 = dist_vel(gen);
    this->r2 = dist_vel(gen);
    this->r3 = dist_vel(gen);
}

void Particle::randomize_velocity() {
    std::random_device seeder;
    std::mt19937 gen(seeder());
    std::uniform_real_distribution<double> dist_vel(0.5, 1.0);
    for(unsigned int i = 0; i < this->coordinate_dim; ++i){
        (*this->velocity)[i] = dist_vel(gen);
    }
}

Particle::Particle(lib4neuro::ErrorFunction *ef, std::vector<double> *domain_bounds) {

    this->ef = ef;
    this->domain_bounds = new std::vector<double>(*domain_bounds);
    this->coordinate_dim = ef->get_dimension();
    this->ef = ef;

    this->coordinate = new std::vector<double>(this->coordinate_dim);
    this->velocity = new std::vector<double>(this->coordinate_dim);
    this->optimal_coordinate = new std::vector<double>(this->coordinate_dim);


    this->randomize_velocity();
    this->randomize_parameters();
    this->randomize_coordinates();

    for(unsigned int i = 0; i < this->coordinate_dim; ++i){
        (*this->optimal_coordinate)[i] = (*this->coordinate)[i];
    }

    this->optimal_value = this->ef->eval(this->coordinate);

}

Particle::Particle(lib4neuro::ErrorFunction *ef, std::vector<double> *central_system, double dispersion_coeff) {

    this->ef = ef;

    if( this->domain_bounds ){
        delete this->domain_bounds;
    }

    this->domain_bounds = new std::vector<double>(2 * central_system->size());
//    return;


    for( size_t i = 0; i < central_system->size(); ++i ){
        this->domain_bounds->at(2 * i) = central_system->at( i ) - dispersion_coeff;
        this->domain_bounds->at(2 * i + 1) = central_system->at( i ) + dispersion_coeff;
    }

    this->coordinate_dim = ef->get_dimension();
    this->ef = ef;

    this->coordinate = new std::vector<double>(this->coordinate_dim);
    this->velocity = new std::vector<double>(this->coordinate_dim);
    this->optimal_coordinate = new std::vector<double>(this->coordinate_dim);


    this->randomize_velocity();
    this->randomize_parameters();
    this->randomize_coordinates();

    for(unsigned int i = 0; i < this->coordinate_dim; ++i){
        (*this->optimal_coordinate)[i] = (*this->coordinate)[i];
    }

    this->optimal_value = this->ef->eval(this->coordinate);

}

Particle::~Particle() {

    if( this->optimal_coordinate ){
        delete this->optimal_coordinate;
    }

    if( this->coordinate ){
        delete this->coordinate;
    }

    if( this->velocity ){
        delete this->velocity;
    }

}

std::vector<double>* Particle::get_coordinate() {
    return this->coordinate;
}

double Particle::get_current_value() {
    return this->current_val;
}

double Particle::get_optimal_value() {
    return this->optimal_value;
}

void Particle::get_optimal_coordinate(std::vector<double> &ref_coordinate) {
    for( unsigned int i = 0; i < this->coordinate_dim; ++i ){
        ref_coordinate[i] = (*this->optimal_coordinate)[i];
    }
}

double Particle::change_coordinate(double w, double c1, double c2, std::vector<double> &glob_min_coord, std::vector<std::vector<double>> &global_min_vec, double penalty_coef) {

    /**
     * v = w * v + c1r1(p_min_loc - x) + c2r2(p_min_glob - x) + c3r3(random_global_min - x)
     * x = x + v
     */

    double vel_mem;
    double output = 0.0;

    /* Choose random global minima */
    std::vector<double> *random_global_best;
    std::random_device rand_dev;
    std::mt19937 engine{rand_dev()};
    std::uniform_int_distribution<size_t> dist(0, global_min_vec.size() - 1);
    random_global_best = &global_min_vec[dist(engine)];

    // TODO use std::sample to choose random vector
    //std::sample(global_min_vec.begin(), global_min_vec.end(), std::back_inserter(random_global_best), 1, std::mt19937{std::random_device{}()});

    for(unsigned int i = 0; i < this->coordinate_dim; ++i) {
        vel_mem = w * (*this->velocity)[i]
                  + c1 * this->r1 * ((*this->optimal_coordinate)[i] - (*this->coordinate)[i])
                  + c2 * this->r2 * (glob_min_coord[i] - (*this->coordinate)[i])
                  + (c1+c2)/2 * this->r3 * ((*random_global_best)[i] - (*this->coordinate)[i]);

        if ((*this->coordinate)[i] + vel_mem > this->domain_bounds->at(2 * i + 1)) {
            this->randomize_velocity();
            this->randomize_parameters();
            this->randomize_coordinates();
            break;
        } else if ((*this->coordinate)[i] + vel_mem < this->domain_bounds->at(2 * i)) {
            this->randomize_velocity();
            this->randomize_parameters();
            this->randomize_coordinates();
            break;
        }
    }

    for(unsigned int i = 0; i < this->coordinate_dim; ++i){
        vel_mem = w * (*this->velocity)[i]
                + c1 * this->r1 * ((*this->optimal_coordinate)[i] - (*this->coordinate)[i])
                + c2 * this->r2 * (glob_min_coord[i] - (*this->coordinate)[i])
                + (c1+c2)/2 * this->r3 * ((*random_global_best)[i] - (*this->coordinate)[i]);


        (*this->velocity)[i] = vel_mem;
        (*this->coordinate)[i] += vel_mem;

        output += std::abs(vel_mem);
    }

    vel_mem = this->ef->eval(this->coordinate);
    this->current_val = vel_mem;

    if(vel_mem < this->optimal_value){
        this->optimal_value = vel_mem;
        for(unsigned int i = 0; i < this->coordinate_dim; ++i){
            (*this->optimal_coordinate)[i] = (*this->coordinate)[i];
        }
    }

    return output;
}

void Particle::print_coordinate() {
    for(unsigned int i = 0; i < this->coordinate_dim - 1; ++i){
        std::cout << (*this->coordinate)[i] << " ";
    }
    std::cout << (*this->coordinate)[this->coordinate_dim - 1] << std::endl;
}

namespace lib4neuro {
    ParticleSwarm::ParticleSwarm(std::vector<double> *domain_bounds,
                                 double c1,
                                 double c2,
                                 double w,
                                 double gamma,
                                 double epsilon,
                                 double delta,
                                 size_t n_particles,
                                 size_t iter_max) {
        srand(time(NULL));  //TODO rewrite using boost.random

        if (epsilon < 0 || gamma < 0 || delta < 0) {
            THROW_INVALID_ARGUMENT_ERROR(
                    "Parameters 'gamma', 'epsilon' and 'delta' must be greater than or equal to zero!");
        }

        this->gamma = gamma;
        this->epsilon = epsilon;
        this->delta = delta;
        this->pst = PARTICLE_SWARM_TYPE::GENERAL;

        this->init_constructor(domain_bounds, c1, c2, w, n_particles, iter_max);
    }

    ParticleSwarm::~ParticleSwarm() {

//        if (this->particle_swarm) {
//            for (size_t i = 0; i < this->n_particles; ++i) {
//                delete this->particle_swarm.at(i);
//            }
//
//            delete this->particle_swarm;
//            this->particle_swarm = nullptr;
//        }
//
//        if( this->domain_bounds ){
//            delete this->domain_bounds;
//        }
//
//        if (this->p_min_glob) {
//            delete this->p_min_glob;
//            this->p_min_glob = nullptr;
//        }

    }

    /**
     *
     *
     */
    void ParticleSwarm::optimize(lib4neuro::ErrorFunction &ef, std::ofstream* ofs) {
        //TODO add output to the 'ofs'

        COUT_INFO("Finding optima via Globalized Particle Swarm method..." << std::endl);

        /* Copy data set max and min values, if it's normalized */
        if(ef.get_dataset()->is_normalized()) {
            ef.get_network_instance()->set_normalization_strategy_instance(
                    ef.get_dataset()->get_normalization_strategy());
        }

        if (this->epsilon < 0 || this->gamma < 0 || this->delta < 0) {
            THROW_INVALID_ARGUMENT_ERROR(
                    "Parameters 'gamma', 'epsilon' and 'delta' must be greater than or equal to zero!");
        }

        this->func_dim = ef.get_dimension();


        /* initialize the particles */
        for (size_t pi = 0; pi < this->n_particles; ++pi) {
            if (this->particle_swarm.at(pi)) {
                delete this->particle_swarm.at(pi);
            }
            this->particle_swarm.at(pi) = new Particle(&ef, ef.get_parameters().get(), this->radius_factor);
        }
        this->radius_factor *= 1.25;

//        if (!this->p_min_glob) {
//            this->p_min_glob = new std::vector<double>(this->func_dim);
//        } else {
            this->p_min_glob.resize(this->func_dim);
//        }

        size_t outer_it = 0;
        Particle *particle;

        std::vector<std::vector<double>> global_best_vec;
        double optimal_value = 0.0;

        std::set<Particle *> cluster; //!< Particles in a cluster
        std::vector<double> *centroid = new std::vector<double>(this->func_dim);//<! Centroid coordinates

        double tmp_velocity;
        double prev_max_velocity = 0;
        double max_velocity;
        double max_vel_step = 0;
        double prev_max_vel_step;
        double euclidean_dist;
        double current_err = -1;

        this->determine_optimal_coordinate_and_value(this->p_min_glob, optimal_value);
//    for(unsigned int i = 0; i < this->n_particles; ++i){
//        this->particle_swarm[i]->print_coordinate();
//    }
        COUT_INFO("Initial best value: " << optimal_value << std::endl);

        while (outer_it < this->iter_max) {
            max_velocity = 0;
            euclidean_dist = 0;

            //////////////////////////////////////////////////
            // Clustering algorithm - termination condition //
            //////////////////////////////////////////////////
            particle = this->determine_optimal_coordinate_and_value(this->p_min_glob, optimal_value);

            if (std::find(global_best_vec.begin(), global_best_vec.end(), this->p_min_glob) == global_best_vec.end()) {
                global_best_vec.emplace_back(this->p_min_glob); // TODO rewrite as std::set
            }

            cluster.insert(particle);

            //for(unsigned int i=0; i < 5; i++) {
            /* Zero AVG coordinates */
            std::fill(centroid->begin(), centroid->end(), 0);
            std::vector<double> *c_ptr;

            /* Looking for a centroid */
            for (auto it : cluster) {
                c_ptr = it->get_coordinate();
                for (size_t di = 0; di < this->func_dim; di++) {
                    (*centroid)[di] += (*c_ptr)[di];
                }
            }

            for (size_t di = 0; di < this->func_dim; di++) {
                (*centroid)[di] /= cluster.size();
            }

            for (size_t pi = 0; pi < this->n_particles; pi++) {
                particle = this->particle_swarm.at(pi);
                tmp_velocity = particle->change_coordinate(this->w, this->c1, this->c2, this->p_min_glob,
                                                           global_best_vec);
//                particle->print_coordinate();

                if (tmp_velocity > max_velocity) {
                    prev_max_velocity = max_velocity;
                    max_velocity = tmp_velocity;
                }

                /* Looking for nearby particles */
                //printf("iter: %d, pi: %d, euclidean dist: %f\n", outer_it, pi, this->get_euclidean_distance(particle->get_coordinate(), coords, this->func_dim));

                // TODO - only in verbose mode
                // only for info purposes
                euclidean_dist += this->get_euclidean_distance(particle->get_coordinate(), centroid);

                if (this->get_euclidean_distance(particle->get_coordinate(), centroid) < epsilon) {
                    cluster.insert(particle);
                }
            }
            //}

            prev_max_vel_step = max_vel_step;
            max_vel_step = max_velocity - prev_max_velocity;

            //TODO only in verbose mode
            euclidean_dist /= this->n_particles;
            if (outer_it % 10 == 0) {
                //printf("Iteration %d, avg euclidean distance: %f, cluster percent: %f, f-value: %f\r", (int)outer_it, euclidean_dist,
                //       double(cluster.size())/this->n_particles, optimal_value);
                //std::cout.flush();
            }

//        for(unsigned int i=0; i < this->n_particles; i++) {
//            printf("Particle %d (%f): \n", i, this->particle_swarm[i]->get_current_value());
//            for(unsigned int j=0; j < this->func_dim; j++) {
//                printf("\t%f\n", this->particle_swarm[i]->get_coordinate()[j]);
//            }
//        }

            current_err = ef.eval(&this->p_min_glob);

            COUT_DEBUG(std::string("Iteration: ") << (unsigned int)(outer_it)
                                                  << ". Total error: " << current_err
                                                  << ". Objective function value: " << optimal_value
                                                  << ".\r");

            if(this->err_thresh) {

                /* If the error threshold is given, then check the current error */
                if(current_err <= this->err_thresh) {
                    break;
                }
            } else {

                /* Check if the particles are near to each other AND the maximal velocity is less than 'gamma' */
                if (cluster.size() >= this->delta * this->n_particles &&
                    prev_max_vel_step <= this->gamma * max_vel_step) {
                    break;
                }
            }

            outer_it++;

            //TODO parameter for inertia weight decrease?
//        this->w *= 0.99;

        }
        COUT_DEBUG(std::string("Iteration: ") << (unsigned int)(outer_it)
                                              << ". Total error: " << current_err
                                              << ". Objective function value: " << optimal_value
                                              << "." << std::endl );

        this->determine_optimal_coordinate_and_value(this->p_min_glob, optimal_value);
        //TODO rewrite following output using COUT_INFO
        if (outer_it < this->iter_max) {
            /* Convergence reached */
            COUT_INFO( std::endl << "Found optimum in "  <<  outer_it << " iterations. Objective function value: " << optimal_value << std::endl);
        } else {
            /* Maximal number of iterations reached */
            COUT_INFO( std::endl << "Max number of iterations reached ("  <<  outer_it << ")!  Objective function value: " << optimal_value <<std:: endl);
        }

        ef.get_network_instance()->copy_parameter_space(&this->p_min_glob);

        delete centroid;
    }

    ParticleSwarm::ParticleSwarm(std::vector<double>* domain_bounds,
                                 double err_thresh,
                                 PARTICLE_SWARM_TYPE pst,
                                 double c1,
                                 double c2,
                                 double w,
                                 size_t n_particles,
                                 size_t iter_max) {

        if(err_thresh <= 0 ) {
            THROW_INVALID_ARGUMENT_ERROR("Error threshold has to be greater then 0!");
        }

        this->err_thresh = err_thresh;
        this->pst = pst;

        this->init_constructor(domain_bounds, c1, c2, w, n_particles, iter_max);
    }

    Particle *ParticleSwarm::determine_optimal_coordinate_and_value(std::vector<double> &coord, double &val) {

        Particle *p;

        val = this->particle_swarm.at(0)->get_optimal_value();
        this->particle_swarm.at(0)->get_optimal_coordinate(coord);
        p = this->particle_swarm.at(0);

        for (size_t i = 1; i < this->n_particles; ++i) {

            double val_m = this->particle_swarm.at(i)->get_optimal_value();

            if (val_m < val) {
                val = val_m;
                this->particle_swarm.at(i)->get_optimal_coordinate(coord);
                p = this->particle_swarm.at(i);
            }
        }

        return p;
    }

    std::vector<double> *ParticleSwarm::get_centroid_coordinates() {
        std::vector<double> *coords = new std::vector<double>(this->func_dim);
        std::vector<double> *tmp;

        for (size_t pi = 0; pi < this->n_particles; pi++) {
            tmp = this->particle_swarm.at(pi)->get_coordinate();

            for (size_t di = 0; di < this->func_dim; di++) {
                (*coords)[di] += (*tmp)[di];
            }
        }

        for (size_t di = 0; di < this->func_dim; di++) {
            (*coords)[di] /= this->n_particles;
        }

        return coords;
    }

    double ParticleSwarm::get_euclidean_distance(std::vector<double> *a, std::vector<double> *b) {
        double dist = 0, m;
        for (size_t i = 0; i < a->size(); i++) {
            m = (*a)[i] - (*b)[i];
            m *= m;
            dist += m;
        }
        return std::sqrt(dist);
    }

    std::shared_ptr<std::vector<double>> ParticleSwarm::get_parameters() {
        std::shared_ptr<std::vector<double>> ret;
        ret.reset(&this->p_min_glob);
        return ret;
    }

    void ParticleSwarm::init_constructor(std::vector<double>* domain_bounds,
                                         double c1,
                                         double c2,
                                         double w,
                                         size_t n_particles,
                                         size_t iter_max) {
        this->c1 = c1;
        this->c2 = c2;
        this->c3 = (c1 + c2) / 2.0;
        this->w = w;
        this->n_particles = n_particles;
        this->iter_max = iter_max;
//        this->particle_swarm = new std::vector<Particle *>(this->n_particles);
//        this->domain_bounds = new std::vector<double>(*domain_bounds);
        std::fill(this->particle_swarm.begin(), this->particle_swarm.end(), nullptr);
    }

}