#include <iostream>
#include <blt/std/allocator.h>
#include <blt/gfx/window.h>
#include <blt/gfx/state.h>
#include <blt/gfx/stb/stb_image.h>
#include <functions.h>
#include "blt/gfx/renderer/resource_manager.h"
#include "blt/gfx/renderer/batch_2d_renderer.h"
#include "blt/std/assert.h"
#include "blt/std/system.h"
#include "imgui.h"
#include "extern.h"
#include "blt/std/format.h"
#include "blt/profiling/profiler_v2.h"
#include <variant>
#include <random>
#include <queue>
#include <stack>
#include <config.h>
#include <gp.h>

blt::gfx::matrix_state_manager global_matrices;
blt::gfx::resource_manager resources;
blt::gfx::batch_renderer_2d renderer_2d(resources);

class tree
{
    public:
        struct search_result_t
        {
            node* child;
            node* parent;
            blt::size_t index;
        };
        struct crossover_result_t
        {
            std::unique_ptr<tree> c1;
            std::unique_ptr<tree> c2;
        };
        struct mutation_result_t
        {
            std::unique_ptr<tree> c;
        };
    private:
        std::unique_ptr<node, node_deleter> root = nullptr;
        
        explicit tree(node* root): root(root)
        {}
        
        void normalize_image()
        {
            float mx = 0, my = 0, mz = 0;
            float sx = 0, sy = 0, sz = 0;
            for (auto& v : root->getImage().getData())
            {
                for (int i = 0; i < 3; i++)
                {
                    if (std::isnan(v[i]))
                    {
                        v[i] = 0.0f;
                    }
                }
                mx = std::max(v.x(), mx);
                my = std::max(v.y(), my);
                mz = std::max(v.z(), mz);
                sx = std::min(v.x(), sx);
                sy = std::min(v.y(), sy);
                sz = std::min(v.z(), sz);
            }
            for (auto& v : root->getImage().getData())
            {
                if (mx - sx != 0)
                    v[0] = (v.x() - sx) / (mx - sx);
                if (my - sy != 0)
                    v[1] = (v.y() - sy) / (my - sy);
                if (mz - sz != 0)
                    v[2] = (v.z() - sz) / (mz - sz);
            }
        }
        
        search_result_t select_random_child()
        {
            static std::random_device dev;
            static std::mt19937_64 engine{dev()};
            std::uniform_int_distribution select(0, 3);
            
            node* current = root.get();
            node* parent = nullptr;
            blt::size_t index = 0;
            
            while (true)
            {
                std::uniform_int_distribution<size_t> children(0, current->argc - 1);
                if (select(engine) == 0 || current->argc == 0)
                    break;
                index = children(engine);
                auto* next = current->sub_nodes[index];
                if (next == nullptr)
                    break;
                parent = current;
                current = next;
            }
            return {current, parent, index};
        }
    
    public:
        std::unique_ptr<tree> clone()
        {
            return std::make_unique<tree>(tree{root->clone()});
        }
        
        static std::unique_ptr<tree> construct_random_tree()
        {
            return std::make_unique<tree>(tree{node::construct_random_tree()});
        }
        
        static blt::size_t depth(node* root_node)
        {
            // depth -> node
            std::stack<std::pair<blt::size_t, node*>> stack;
            blt::size_t max_depth = 0;
            
            stack.emplace(0, root_node);
            while (!stack.empty())
            {
                auto top = stack.top();
                auto* node = top.second;
                auto depth = top.first;
                stack.pop();
                max_depth = std::max(max_depth, depth);
                for (blt::size_t i = 0; i < node->argc; i++)
                {
                    if (node->sub_nodes[i] != nullptr)
                    {
                        stack.emplace(depth + 1, node->sub_nodes[i]);
                    }
                }
            }
            
            return max_depth;
        }
        
        static std::optional<crossover_result_t> crossover(tree* p1, tree* p2)
        {
            if (p1 == nullptr || p2 == nullptr)
                return {};
            auto c1 = p1->clone();
            auto c2 = p2->clone();
            
            auto n1 = c1->select_random_child();
            auto n2 = c2->select_random_child();
            
            if (n1.parent == nullptr || n2.parent == nullptr)
                return {};
            
            const auto& p1_allowed = function_arg_allowed_set_map[to_underlying(n1.parent->type)][n1.index];
            const auto& p2_allowed = function_arg_allowed_set_map[to_underlying(n2.parent->type)][n2.index];
            
            if (!p1_allowed.contains(n2.child->type))
                return {};
            if (!p2_allowed.contains(n1.child->type))
                return {};
            
            n1.parent->sub_nodes[n1.index] = n2.child;
            n2.parent->sub_nodes[n2.index] = n1.child;
            
            return crossover_result_t{std::move(c1), std::move(c2)};
        }
        
        static std::optional<mutation_result_t> mutate(tree* p)
        {
            if (p == nullptr)
                return {};
            static std::random_device dev;
            static std::mt19937_64 engine{dev()};
            std::uniform_int_distribution choice(0, 1);
            
            auto c = p->clone();
            
            auto n = c->select_random_child();
            
            if (n.parent == nullptr)
                return {};
            
            auto d = depth(n.child);
            
            node* new_subtree = node::construct_random_tree(d + 1);
            n.parent->sub_nodes[n.index] = new_subtree;
            destroyNode(n.child);
            return mutation_result_t{std::move(c)};
        }
        
        void evaluate()
        {
            root->evaluate_tree();
            normalize_image();
        }
        
        image& getImage()
        {
            return root->getImage();
        }
        
        bool hasImage()
        {
            return root->hasImage();
        }
        
        bool hasTree()
        {
            return root != nullptr;
        }
        
        float fitness()
        {
            auto& img = root->getImage();
            return eval_DNF_SW(img) * eval_DNF_SW_1(img) * static_cast<float>(std::min(depth(root.get()), 5ul));
        }
        
        void printTree()
        {
            root->printTree();
        }
};

class gp_population
{
    public:
        struct gp_i
        {
            std::unique_ptr<tree> t = nullptr;
            float fitness = 0;
        };
        blt::size_t best = -1;
    private:
        std::array<gp_i, POPULATION_SIZE> pop;
    public:
        gp_population() = default;
        
        static blt::size_t choice()
        {
            static std::random_device dev;
            static std::mt19937_64 engine{dev()};
            static std::uniform_int_distribution c(0, POPULATION_SIZE - 1);
            
            return c(engine);
        }
        
        void init()
        {
            for (auto& v : pop)
                v = {tree::construct_random_tree()};
            evaluate_population();
        }
        
        tree* select()
        {
            tree* n = nullptr;
            float fitness = -2 * 8192;
            //BLT_TRACE("With inital %f", fitness);
            for (int i = 0; i < TOURNAMENT_SIZE; i++)
            {
                blt::size_t index = 0;
                do
                {
                    index = choice();
                    auto& v = pop[index];
                    if (v.fitness >= fitness)
                    {
                        n = v.t.get();
                        fitness = v.fitness;
                    }
                    //BLT_TRACE("%d: %p -> %f ? %p -> %f || %d %d", index, v.t.get(), v.fitness, n, fitness, n != v.t.get(), v.fitness > fitness);
                } while (n == pop[index].t.get());
            }
            //BLT_DEBUG("%p -> %f", n, fitness);
            BLT_ASSERT(n != nullptr);
            return n;
        }
        
        void run_step()
        {
            std::array<gp_i, POPULATION_SIZE> new_pop;
            static std::random_device dev;
            static std::mt19937_64 engine{dev()};
            static std::uniform_real_distribution rand(0.0, 1.0);
            
            auto b = get_best();
            new_pop[0] = {pop[b.first].t->clone(), b.second};
            blt::size_t insert_pos = 1;
            
            blt::size_t crossover_count = 0;
            blt::size_t mutation_count = 0;
            BLT_TRACE("Running Crossover");
            for (blt::size_t i = 0; i < POPULATION_SIZE; i++)
            {
                if (rand(engine) < CROSSOVER_RATE)
                {
                    auto* p1 = select();
                    auto* p2 = select();
                    
                    if (auto r = tree::crossover(p1, p2))
                    {
                        new_pop[insert_pos++] = {std::move(r->c1)};
                        new_pop[insert_pos++] = {std::move(r->c2)};
                        crossover_count++;
                    }
                }
            }
            
            BLT_TRACE("Running Mutation");
            for (blt::size_t i = 0; i < POPULATION_SIZE; i++)
            {
                if (rand(engine) < MUTATION_RATE)
                {
                    auto* p1 = select();
                    
                    if (auto r = tree::mutate(p1))
                    {
                        new_pop[insert_pos++] = gp_i{std::move(r->c), 0};
                        mutation_count++;
                    }
                }
            }
            
            while (insert_pos < POPULATION_SIZE)
            {
                new_pop[insert_pos++] = {select()->clone()};
            }
            pop = std::move(new_pop);
            BLT_TRACE("ran %d crossovers and %d mutations", crossover_count, mutation_count);
        }
        
        void evaluate_population()
        {
            BLT_TRACE("Running Eval");
            for (auto& v : pop)
            {
                v.t->evaluate();
                v.fitness = v.t->fitness();
            }
            BLT_TRACE("Complete");
        }
        
        image& display(blt::size_t i)
        {
            return pop[i].t->getImage();
        }
        
        std::pair<blt::size_t, float> get_best()
        {
            blt::size_t i = 0;
            float fitness = -2 * 8192;
            for (blt::size_t j = 0; j < POPULATION_SIZE; j++)
            {
                if (pop[j].fitness > fitness)
                {
                    i = j;
                    fitness = pop[j].fitness;
                }
            }
            best = i;
            return {i, fitness};
        }
        
        image& display_best()
        {
            return display(get_best().first);
        }
        
        float best_fitness()
        {
            return get_best().second;
        }
        
        void print_best()
        {
            pop[get_best().first].t->printTree();
        }
};

gp_population pop;
blt::gfx::texture_gl2D* texture;

void print_bits(float f)
{
    auto u = blt::mem::type_cast<blt::u32>(f);
    for (size_t i = 31; i > 0; i--)
    {
        std::cout << ((u >> i) & 0x1);
        if (i % 8 == 0)
            std::cout << ", ";
        else if (i % 4 == 0)
            std::cout << " ";
    }
    std::cout << std::endl;
}

void init()
{
    global_matrices.create_internals();
    renderer_2d.create();
    
    texture = new blt::gfx::texture_gl2D(width, height);
    resources.set("img", texture);
    resources.enqueue("../libraries/BLT-With-Graphics-Template/resources/textures/parkerfemBOY.png", "cum");
    
    resources.load_resources();
}

float best = 0;

void update(std::int32_t w, std::int32_t h)
{
    global_matrices.update_perspectives(w, h, 90, 0.1, 2000);
    glDisable(GL_BLEND);
    
    if (ImGui::Button("Regenerate"))
    {
        BLT_INFO("Regen tree");
        pop = {};
        pop.init();
        best = pop.best_fitness();
    }
    
    if (ImGui::Button("Run Step"))
    {
        BLT_INFO("Running Step");
        pop.run_step();
        BLT_INFO("Evaluating Population");
        pop.evaluate_population();
        best = pop.best_fitness();
    }
    
    if (ImGui::Button("Display"))
    {
        BLT_INFO("Uploading");
        texture->upload((void*) pop.display_best().getData().data(), GL_RGB, 0, 0, 0, -1, -1, GL_FLOAT);
    }
    
    if (ImGui::Button("Print"))
    {
        pop.print_best();
    }
    ImGui::Text("Best Fitness: %f", best);
    auto data = blt::system::get_memory_process();
    ImGui::Text("Physical Memory Usage: %s", blt::string::fromBytes(data.resident).c_str());
    ImGui::Text("Shared Memory Usage: %s", blt::string::fromBytes(data.shared).c_str());
    ImGui::Text("Total Memory Usage: %s", blt::string::fromBytes(data.size).c_str());
    
    auto lw = 512.0f;
    auto lh = 512.0f;
    //renderer_2d.drawRectangle(blt::vec4{0.5, 0.0, 1.0, 1.0},
    //                          {static_cast<float>(w) / 2.0f, static_cast<float>(h) / 2.0f, static_cast<float>(w), static_cast<float>(h)});
    renderer_2d.drawRectangle("img", {static_cast<float>(w) / 2.0f, static_cast<float>(h) / 2.0f, lw, lh});
    
    global_matrices.update();
    renderer_2d.render();
}

int main()
{
//    auto& funcs = function_arg_allowed_map[to_underlying(function_t::IF)];
//    for (auto v : blt::enumerate(funcs))
//        for (auto f : v.second)
//            std::cout << "arg " << v.first << ": " << function_name_map[to_underlying(f)] << std::endl;
    //shapiro_test_run();
    blt::gfx::init(blt::gfx::window_data{"Window of GP test", init, update}.setSyncInterval(1));
    global_matrices.cleanup();
    resources.cleanup();
    renderer_2d.cleanup();
    blt::gfx::cleanup();
    return 0;
}