/**
 * 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 "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());
    this->domain_bounds = domain_bounds;
    for(unsigned int i = 0; i < this->coordinate_dim; ++i){
        std::uniform_real_distribution<double> dist_coord(-1, 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 = 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() {

    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));

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

        this->c1 = c1;

        this->c2 = c2;

        this->c3 = (c1 + c2) / 2.0;

        this->w = w;

        this->gamma = gamma;

        this->epsilon = epsilon;

        this->delta = delta;

        this->n_particles = n_particles;

        this->iter_max = iter_max;

        this->particle_swarm = new std::vector<Particle *>(this->n_particles);

        this->domain_bounds = domain_bounds;

        std::fill(this->particle_swarm->begin(), this->particle_swarm->end(), nullptr);

    }

    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->p_min_glob) {
            delete this->p_min_glob;
            this->p_min_glob = nullptr;
        }

    }

/**
 *
 * @param gamma
 * @param epsilon
 * @param delta
 */
    void ParticleSwarm::optimize(lib4neuro::ErrorFunction &ef) {
        if (epsilon < 0 || gamma < 0 || delta < 0) {
            throw std::invalid_argument(
                    "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, this->domain_bounds);
        }

        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;

        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();
//    }
        printf("Initial best value: %10.8f\n", optimal_value);

        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]);
//            }
//        }

            /* 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++;
//        this->w *= 0.99;
        }

        this->determine_optimal_coordinate_and_value(*this->p_min_glob, optimal_value);
        if (outer_it < this->iter_max) {
            /* Convergence reached */
            printf("\nFound optimum in %d iterations. Objective function value: %10.8f\n", (int) outer_it,
                   optimal_value);
        } else {
            /* Maximal number of iterations reached */
            printf("\nMax number of iterations reached (%d)!  Objective function value: %10.8f\n", (int) outer_it,
                   optimal_value);
        }
//    for (size_t i = 0; i <= this->func_dim - 1; ++i) {
//        printf("%10.8f \n", (*this->p_min_glob)[i]);
//    }
//
//    this->f->eval( this->get_solution() );


        delete centroid;
    }

    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::vector<double> *ParticleSwarm::get_parameters() {
        return this->p_min_glob;
    }

}