#include <array>
#include <algorithm>
#include <vector>
#include <cstdio>
#include <random>
#include <set>
#include <future>

constexpr std::size_t num_strips = 200;
constexpr std::size_t strip_size = 16;
constexpr std::size_t num_steps = 100;
constexpr std::size_t num_threads = 4;
constexpr double num_ant_factor = 0.1 + 1.0 / num_threads;
constexpr double evap_rate = 0.8;

using strip_t = std::array<unsigned char, strip_size>;
using node_int_t = unsigned short;
using path_t = std::vector<node_int_t>;
using rng_t = std::minstd_rand;

double trail_preference(double x)
{
    return x;
}

double greedy_preference(double x)
{
    x *= x; // x^2
    x *= x; // x^4
    x *= x; // x^8
    x *= x; // x^16
    x *= x; // x^32
    x *= x; // x^64
    return x;
}

struct ant_t
{
    std::size_t total_cost;
    std::vector<unsigned char> visited;
    path_t path;

    node_int_t node() const { return path.back(); }
    void visit(node_int_t node, std::size_t cost)
    {
        total_cost += cost;
        path.push_back(node);
        visited[node] = true;
    }
};

class ants_t
{
public:
    ants_t(rng_t& rng,
           std::vector<strip_t> const& strips,
           std::vector<unsigned char> const& strip_costs)
    : global_rng(rng)
    , strips(strips)
    , strip_costs(strip_costs)
    {}

    path_t optimize()
    {
        trails.assign(strips.size()*strips.size(), 1.0);
        lowest_cost = SIZE_MAX;

        for(std::size_t steps = 0; steps != num_steps; ++steps)
        {
            step();
            std::printf("%lu i=%lu\n", lowest_cost, steps);
        }
        return lowest_path;
    }

private:
    rng_t& global_rng;
    std::vector<strip_t> const& strips;
    std::vector<unsigned char> const& strip_costs;
    std::vector<double> trails;
    std::size_t lowest_cost;
    path_t lowest_path;
    std::array<std::vector<ant_t>, num_threads> ants;

    unsigned char cost(node_int_t i, node_int_t j) const
    {
        return strip_costs[i*strips.size() + j];
    }

    double& trail(node_int_t i, node_int_t j)
    {
        return trails[i*strips.size() + j];
    }

    double trail(node_int_t i, node_int_t j) const
    {
        return trails[i*strips.size() + j];
    }

    node_int_t select_next_node(rng_t& rng, ant_t ant) const
    {
        std::vector<double> probs(strips.size(), 0.0);
        node_int_t node = ant.node();

        // Calculate probabilities for each node.
        double denom = 0.0;
        for(std::size_t i = 0; i != strips.size(); ++i)
        {
            if(!ant.visited[i])
            {
                denom += (trail_preference(trail(node, i))
                          * greedy_preference(1.0 / cost(node, i)));
            }
        }

        for(std::size_t i = 0; i != strips.size(); ++i)
        {
            if(!ant.visited[i])
            {
                double numer = (trail_preference(trail(node, i))
                                * greedy_preference(1.0 / cost(node, i)));
                probs[i] = numer / denom;
            }
        }

        // Randomly select according to probs.
        std::uniform_real_distribution<> d;
        double const r = d(rng);
        double t = 0.0;
        for(std::size_t i = 0; i != strips.size(); ++i)
        {
            t += probs[i];
            if(t >= r)
                return i;
        }

        throw 0;
    }

    void step_ants(rng_t& rng, std::vector<ant_t>& ants)
    {
        // Create ants.
        ants.clear();
        std::uniform_int_distribution<> nd(0, strips.size()-1);
        std::size_t num_ants = strips.size() * num_ant_factor;
        for(std::size_t i = 0; i != num_ants; ++i)
        {
            ant_t new_ant = {0};
            new_ant.visited.assign(strips.size(), 0);
            new_ant.visit(nd(rng), 0);
            ants.push_back(new_ant);
        }

        // Move ants.
        for(std::size_t i = 1; i < strips.size(); ++i)
        {
            for(ant_t& ant : ants)
            {
                node_int_t const next_node = select_next_node(rng, ant);
                ant.visit(next_node, cost(ant.node(), next_node));
            }
        }
        for(ant_t& ant : ants)
            ant.total_cost += cost(ant.node(), ant.path[0]);
    }

    void step()
    {
        using lowest_t = std::pair<std::size_t, std::size_t>;
        std::vector<std::future<lowest_t> > futures;

        for(std::size_t thread = 0; thread != num_threads; ++thread)
        {
            futures.push_back(std::async(std::launch::async,
            [&](rng_t rng, std::vector<ant_t>* ants) -> lowest_t
            {
                step_ants(rng, *ants);

                std::size_t thread_lowest_cost = SIZE_MAX;
                std::size_t thread_lowest_ant;
                for(std::size_t i = 0; i != ants->size(); ++i)
                {
                    if((*ants)[i].total_cost < thread_lowest_cost)
                    {
                        thread_lowest_cost = (*ants)[i].total_cost;
                        thread_lowest_ant = i;
                    }
                }
                return std::make_pair(thread_lowest_cost, thread_lowest_ant);
            },
            rng_t(global_rng()), &ants[thread]));
        }

        // Update lowest.
        for(std::size_t thread = 0; thread != num_threads; ++thread)
        {
            lowest_t pair = futures[thread].get();
            if(pair.first < lowest_cost)
            {
                lowest_cost = pair.first;
                lowest_path = ants[thread][pair.second].path;
            }
        }

        // Evaporate trails.
        for(double& f : trails)
            f *= evap_rate;

        // Add pheramones.
        for(std::size_t thread = 0; thread != num_threads; ++thread)
        for(ant_t const& ant : ants[thread])
        {
            double contribution = 1.0 / ant.total_cost;
            for(int i = 0; i < strips.size()-1; ++i)
                trail(ant.path[i], ant.path[i+1]) += contribution;
            trail(ant.path[strips.size()-1], ant.path[0]) += contribution;
        }
    }

};

std::size_t best_overlap(strip_t l, strip_t r)
{
    for(std::size_t i = 0; i != strip_size; ++i)
    {
        for(std::size_t j = 0; j != strip_size - i; ++j)
            if(l[i+j] != r[j])
                goto no_match;
        return i;
    no_match:;
    }
    return strip_size;
}

std::vector<unsigned char> get_strip_costs(std::vector<strip_t> const& strips)
{
    std::vector<unsigned char> strip_costs(strips.size() * strips.size());
    for(std::size_t i = 0; i != strips.size(); ++i)
    for(std::size_t j = 0; j != strips.size(); ++j)
        strip_costs[i*strips.size()+j] = best_overlap(strips[i], strips[j]);
    return strip_costs;
}

int main()
{
    std::random_device rd;
    rng_t rng;

    std::uniform_int_distribution<> d(0, 2);
    std::vector<strip_t> strips;
    std::set<strip_t> unique_strips;
    for(std::size_t i = 0; i != num_strips; ++i)
    {
        strip_t strip;
        for(auto& c : strip)
            c = '0' + d(rng);
        if(unique_strips.count(strip) == 0)
        {
            strips.push_back(strip);
            unique_strips.insert(strip);
        }
    }

    auto strip_costs = get_strip_costs(strips);

    ants_t ants(rng, strips, strip_costs);
    ants.optimize();
    std::printf("uncompressed: %lu\n", num_strips * strip_size);
    std::printf("uncompressed (no duplicates): %lu\n", strips.size() * strip_size);
}