#include <random>
#include <bitset>
#include <magic_enum/magic_enum.hpp>
#include <stdio.h>
#include <type_traits>
#include <iostream>
#include <format>
#include <vector>

namespace s2ga
{
    template<class... T>
    struct always_false: std::false_type
    {
    };

    class lehmer64
    {
    public:
        lehmer64(__uint128_t seed = 0)
        {
            this->seed(seed);
        }

        void seed(__uint128_t z)
        {
            if(z == 0)
            {
                z = 0xC0FFEE;
            }

            s = z;
        }

        uint64_t operator()() noexcept
        {
            s *= 0xda942042e4dd58b5;
            return s >> 64;
        }

        static constexpr uint64_t min()
        {
            return 0;
        }

        static constexpr uint64_t max()
        {
            return UINT64_MAX;
        }

        template<typename T>
        T random(T min, T max)
        {
            if constexpr(std::is_integral_v<T>)
            {
                return std::uniform_int_distribution<T>(min, max)(*this);
            }
            else if constexpr(std::is_floating_point_v<T>)
            {
                return std::uniform_real_distribution<T>(min, max)(*this);
            }
            else
            {
                static_assert(always_false<T>::value,
                              "unsupported distribution type");
            }
        }

    private:
        __uint128_t s;
    };

    static_assert(std::uniform_random_bit_generator<lehmer64>);

    template<typename E, typename... Args>
    requires std::is_enum_v<E> && ((std::is_same_v<E, Args>), ...)
    std::bitset<magic_enum::enum_count<E>()> bm(const Args&... v) //aka "bitmask"
    {
        constexpr size_t N = magic_enum::enum_count<E>();
        std::bitset<N> bitmask;

        ((bitmask.set(magic_enum::enum_integer(v))), ...);
        return bitmask;
    }

    template<typename E>
    requires std::is_enum_v<E>
    struct action 
    {
        static constexpr size_t N = magic_enum::enum_count<E>();

        std::string name = "ACTION";
        float cost = 1.0f;
        
        std::bitset<N> positive_effects{};
        std::bitset<N> negative_effects{};
        std::bitset<N> positive_preconds{};
        std::bitset<N> negative_preconds{};

        bool preconds_met(const std::bitset<N>& state) const
        {
            return ((state & positive_preconds) == positive_preconds) && 
                (state & negative_preconds).none();
        }

        std::bitset<N> apply(const std::bitset<N>& state) const
        {
            return (state | positive_effects) & ~negative_effects;
        }

        void print()
        {
            std::cout << std::format("[{}, {}]", name, cost) << std::endl;

            if(positive_effects.any() || negative_effects.any())
            {
                std::cout << "Effects:" << std::endl;
                for(size_t i = 0; i < N; ++i)
                {
                    if(positive_effects.test(i) || negative_effects.test(i))
                    {
                        auto enum_value = magic_enum::enum_cast<E>(i).value();
                        auto value_name = magic_enum::enum_name(enum_value);

                        std::cout << std::format(" {}{}", 
                                negative_effects.test(i) ? "-" : "+",
                                value_name) << std::endl;
                    }
                }
            }
        
            if(positive_preconds.any() || negative_preconds.any())
            {
                std::cout << "Preconditions:" << std::endl;
                for(size_t i = 0; i < N; ++i)
                {
                    if(positive_preconds.test(i) || negative_preconds.test(i))
                    {
                        auto enum_value = magic_enum::enum_cast<E>(i).value();
                        auto value_name = magic_enum::enum_name(enum_value);
                        
                        std::cout << std::format(" {}{}", 
                                negative_preconds.test(i) ? "-" : "+",
                                value_name) << std::endl;
                    }
                }
            }
        }
    };

    template<typename E>
    requires std::is_enum_v<E>
    using action_list = std::vector<action<E>>;

    template<typename E>
    requires std::is_enum_v<E>
    struct goal 
    {
        static constexpr size_t N = magic_enum::enum_count<E>();
        std::bitset<N> positive_goals;
        std::bitset<N> negative_goals;

        int distance(const std::bitset<N>& state)
        {
            /*
             *V000111
             *P111100
             *&000100
             *^111000 <- distance 3
             */
            /*V000111
             *P000000
             *&000000
             *^000000 <- distance 0
             */
            /*V000111
             *P111111
             *&000111
             *^111000 <- distance 3
             */

            auto positive_bm = (state & positive_goals) ^ positive_goals;
            
            /*
             *V000111 
             *N111100
             *&000100 <- distance 1 
             */
            /*
             *V000111 
             *N000000
             *&000000 <- distance 0 
             */
            auto negative_bm = state & negative_goals;

            return positive_bm.count() + negative_bm.count(); 
        }
    };

    template<typename E, int MaxDepth = 20>
    requires std::is_enum_v<E>
    class genetic_goap
    {
        static constexpr size_t N = magic_enum::enum_count<E>();
    
        struct genotype 
        {
            int distance(const genotype& g)
            {
                int d = 0;

                for(int i = 0; i < MaxDepth; ++i)
                {
                    if(action_ids[i] != g.action_ids[i])
                    {
                        d++;
                    }
                }

                return d;
            }

            std::array<int, MaxDepth> action_ids;
            float fitness = 9999.0f;
        };

        void evaluate(genotype& g)
        {
            float cost = 0.0f;

            auto state = initial_state;

            int valid_actions = 0;

            int d = 99999;

            for(int i = 0; i < MaxDepth; ++i)
            {
                auto aid = g.action_ids[i];

                if(aid < 0)
                {
                    continue;
                }

                if(!actions[aid].preconds_met(state))
                {
                    continue;
                }

                state = actions[aid].apply(state);
                d = final_goal.distance(state);
                cost += actions[aid].cost;
                valid_actions++;

                if(d == 0)
                {
                    break;
                }
            }

            if(valid_actions == 0)
            {
                g.fitness = 9999.0f;
                return;
            }

            g.fitness = 9999.0f * d * float(valid_actions) + cost;
            //std::cout << std::format("{}({})", g.fitness, cost) << std::endl;
        }

        void evaluate(std::vector<genotype>& g)
        {
            for(auto& i: g)
            {
                evaluate(i);
            }
        }

        std::vector<genotype> generate_population(int size)
        {
            std::vector<genotype> population;
            population.reserve(size);

            for(int i = 0; i < size; ++i)
            {
                population.emplace_back();
                
                auto& d = population.back();
                
                std::generate(std::begin(d.action_ids), std::end(d.action_ids), 
                [&]()
                {
                    return -1;
                });

                const int ATTEMPTS = 4;
               
                auto last_state = initial_state;

                for(int j = 0; j < MaxDepth; ++j)
                {
                    bool found = false;
                    for(int attempt = 0; attempt < ATTEMPTS; attempt++)
                    {
                        int act_id = rng.random(0, int(actions.size()) - 1);
                        auto& act = actions[act_id];
                        if(act.preconds_met(last_state))
                        {
                            last_state = act.apply(last_state);
                            d.action_ids[j] = act_id;
                            found = true;
                            break;
                        }
                    }

                    if(!found)
                    {
                        d.action_ids[j] = rng.random(-1, int(actions.size()) - 1);
                    }
                }

            }

            return population;
        }

        action_list<E> to_action_list(const genotype& g)
        {
            action_list<E> result;
            auto state = initial_state;

            int d = 9999;

            for(int i = 0; i < MaxDepth; ++i)
            {
                auto aid = g.action_ids[i];

                if(aid < 0)
                {
                    continue;
                }

                if(!actions[aid].preconds_met(state))
                {
                    continue;
                }

                state = actions[aid].apply(state);
                result.emplace_back(actions[aid]);
                
                d = final_goal.distance(state);

                if(d == 0)
                {
                    break;
                }
            }

            return result;
        }

        std::vector<genotype> elitism(int N, const std::vector<genotype>& p)
        {
            std::vector<genotype> result;
            result.reserve(N);

            for(int i = 0; i < N; ++i)
            {
                result.emplace_back(p[i]);
            }

            return result;
        }

        const genotype& tornament_selection(const std::vector<genotype>& p, int k = 8)
        {
            int id = rng.random(0, int(p.size()) - 1);
                
            for(int i = 1; i < k; ++i)
            {
                int nid = rng.random(0, int(p.size()) - 1);

                if(p[nid].fitness < p[id].fitness)
                {
                    id = nid;
                }
            }

            return p[id];
        }

        genotype crossover(const genotype& a, const genotype& b)
        {
            genotype result;


            int split_index = rng.random(0, MaxDepth - 1);
            for(int i = 0; i < MaxDepth; ++i)
            {
                bool from_a = i < split_index;
                result.action_ids[i] = from_a ? a.action_ids[i] : b.action_ids[i];
            }

            return result;
        }

        void mutate(genotype& g)
        {
            int index = rng.random(0, MaxDepth - 1);
            g.action_ids[index] = rng.random(-1, int(actions.size() - 1));
        }
    public:
        action_list<E> plan(int population_size, int max_attempts)
        {
            auto population = generate_population(population_size);

            int attempts = 0;


            float last_f = 9999999.0f;

            while(attempts < max_attempts)
            {
                evaluate(population);
                std::sort(population.begin(), population.end(), 
                        [](const genotype& a, const genotype& b)
                        {
                            return a.fitness < b.fitness;
                        });

                if(population[0].fitness < last_f)
                {
                    last_f = population[0].fitness;
                    attempts = 0;
                }
                else 
                {
                    attempts++;
                }

                auto new_pop = elitism(3, population);

                while(new_pop.size() < population.size())
                {
                    auto solution = crossover(tornament_selection(population), 
                            tornament_selection(population));

                    if(rng.random(0, 100) <= 5)
                    {
                        mutate(solution);
                    }

                    new_pop.emplace_back(solution);
                }

                population = new_pop;
            }

            return to_action_list(population[0]);
        }

        void set_possible_actions(const action_list<E>& list)
        {
            actions = list;
        }

        void set_initial_state(const std::bitset<N>& state)
        {
            initial_state = state;
        }

        void set_goal(const goal<E>& goal)
        {
            final_goal = goal;
        }
    private:
        action_list<E> actions;
        std::bitset<N> initial_state;
        goal<E> final_goal;
        lehmer64 rng;
    };
}; // namespace s2ga