#pragma once

#include "vector"
#include "map"

#include "../CFGraph/CFGraph.h"
#include <unordered_map>

using namespace std;

namespace SAPFOR {
class DominatorFinder {
private:
    BasicBlock* entry;
    std::vector<BasicBlock*> vertices;
    std::unordered_map<BasicBlock*, int> dfs_num;
    std::vector<int> parent, semi, vertex, ancestor, label;
    std::vector<std::vector<int>> bucket;
    int n;

    void DFS(BasicBlock* v, int parent_num) {
        dfs_num[v] = n;
        vertex[n] = n;
        semi[n] = n;
        label[n] = n;
        ancestor[n] = -1;
        parent[n] = parent_num;
        vertices[n++] = v;

        for (const auto& w : v->getNext()) {
            if (dfs_num[w] == -1) {
                DFS(w, dfs_num[v]);
            }
        }
    }

    void Compress(int v) {
        if (ancestor[ancestor[v]] != -1) {
            Compress(ancestor[v]);
            if (semi[label[ancestor[v]]] < semi[label[v]])
                label[v] = label[ancestor[v]];
            ancestor[v] = ancestor[ancestor[v]];
        }
    }

    int Eval(int v) {
        if (ancestor[v] == -1) return v;
        Compress(v);
        return label[v];
    }

    void Link(int v, int w) {
        ancestor[w] = v;
    }

public:
    DominatorFinder(std::vector<BasicBlock*>& blocks) {
        if (blocks.empty()) return;
        entry = blocks[0];
        n = 0;

        for (auto block : blocks) dfs_num[block] = -1;

        int max_size = blocks.size();
        vertices.resize(max_size);
        parent.assign(max_size, -1);
        semi.assign(max_size, -1);
        vertex.assign(max_size, -1);
        ancestor.assign(max_size, -1);
        label.assign(max_size, -1);
        bucket.resize(max_size);

        DFS(entry, -1);

        for (int i = n - 1; i > 0; --i) {
            int w = vertex[i];

            for (BasicBlock* v : vertices[w]->getPrev()) {
                int u = Eval(dfs_num[v]);
                if (semi[u] < semi[w])
                    semi[w] = semi[u];
            }
            bucket[vertex[semi[w]]].push_back(w);
            Link(parent[w], w);

            for (int v : bucket[parent[w]])
            {
                int u = Eval(v);
                if (semi[u] < semi[v])
                    vertices[v]->setIdom(vertices[u]);
                else
                    vertices[v]->setIdom(vertices[parent[w]]);
            }
            bucket[parent[w]].clear();
        }

        for (int i = 1; i < n; ++i) {
            int w = vertex[i];
            if (vertices[w]->getIdom() != vertices[vertex[semi[w]]])
                vertices[w]->setIdom(vertices[w]->getIdom()->getIdom());
        }

        entry->setIdom(nullptr);
    }
};

void buildDominatorTreeLT(std::vector<BasicBlock*>& blocks) {
    DominatorFinder finder(blocks);
}

}