#include "treeordercomputationiterative.h"

list<size_t> getInitialOrder(TemporalTree &tree, size_t v, std::mt19937 &rng)
{
    list<size_t> l;
    if (tree.isLeaf(v))
    {
        l.push_back(v);
    }
    else
    {
        auto ws = tree.edgesHierarchy[v];
        shuffle(ws.begin(), ws.end(), rng);
        for (auto w : tree.edgesHierarchy[v])
        {
            auto lchild = getInitialOrder(tree, w, rng);
            l.splice(l.end(), lchild);
        }
    }
    return l;
}

void setSizeWeight(size_t v, TemporalTree &tree, unordered_map<size_t, double> &index_order, unordered_map<size_t, double> &weight, unordered_map<size_t, int> &size, bool rev, bool leafedges)
{
    double sum = 0;
    vector<size_t> adj;
    if (rev)
    {
        if (!leafedges)
            adj = tree.edgesTime[v];
        else
            adj = tree.leafEdges[v];
    }
    else
    {
        if (!leafedges)
            adj = tree.reverseEdgesTime[v];
        else
            adj = tree.reverseLeafEdges[v];
    }
    for (auto w : adj)
    {
        sum += index_order[w];
    }
    weight[v] = sum;
    size[v] = adj.size();
}

list<size_t> getBarycenterOrder(TemporalTree &tree, size_t v, unordered_map<size_t, double> &index_order, unordered_map<size_t, double> &weight, unordered_map<size_t, int> &size, bool rev, bool leafedges)
{
    list<size_t> l;
    TemporalTree::TAdjacency *adjacency;
    if (rev)
    {
        if (leafedges)
            adjacency = &tree.leafEdges;
        else
            adjacency = &tree.edgesTime;
    }
    else
    {
        if (leafedges)
            adjacency = &tree.reverseLeafEdges;
        else
            adjacency = &tree.reverseEdgesTime;
    }
    if ((*adjacency).find(v) != (*adjacency).end() && (*adjacency)[v].size() > 0)
    {
        setSizeWeight(v, tree, index_order, weight, size, rev, leafedges);
    }
    else
    {
        weight[v] = size[v] = 0;
    }
    if (tree.isLeaf(v))
    {
        l.push_back(v);
    }
    else
    {
        vector<list<size_t>> orders;
        auto ws = tree.edgesHierarchy[v];
        vector<int> indices;
        vector<int> indicesnoprev;
        int i = 0;
        for (auto w : ws)
        {
            orders.push_back(getBarycenterOrder(tree, w, index_order, weight, size, rev, leafedges));
            weight[v] += weight[w];
            size[v] += size[w];
            if (size[w] == 0)
            {
                indicesnoprev.push_back(i);
            }
            else
            {
                indices.push_back(i);
            }
            i++;
        }

        sort(indices.begin(), indices.end(),
             [&](int A, int B) -> bool
             {
                 if ((size[ws[A]] == 0 || size[ws[B]] == 0) && index_order.find(tree.someLeafChild(ws[A])) != index_order.end() && index_order.find(tree.someLeafChild(ws[B])) != index_order.end())
                 {
                     return index_order[ws[A]] < index_order[ws[B]];
                 }
                 double wA = size[ws[A]] > 0 ? (weight[ws[A]] / (double)size[ws[A]]) : 0;
                 double wB = size[ws[B]] > 0 ? (weight[ws[B]] / (double)size[ws[B]]) : 0;
                 return wA < wB;
             });
        for (auto x : indicesnoprev)
            indices.push_back(x);
        sort(indices.begin(), indices.end(),
             [&](int A, int B) -> bool
             {
                 if ((size[ws[A]] == 0 || size[ws[B]] == 0) && index_order.find(tree.someLeafChild(ws[A])) != index_order.end() && index_order.find(tree.someLeafChild(ws[B])) != index_order.end())
                 {
                     return index_order[ws[A]] < index_order[ws[B]];
                 }
                 double wA = size[ws[A]] > 0 ? (weight[ws[A]] / (double)size[ws[A]]) : 0;
                 double wB = size[ws[B]] > 0 ? (weight[ws[B]] / (double)size[ws[B]]) : 0;
                 return wA < wB;
             });

        for (auto i : indices)
        {
            l.splice(l.end(), orders[i]);
        }
    }
    return l;
}

void setInitialPrevVec(size_t v, TemporalTree &tree, unordered_map<size_t, double> &index_order, unordered_map<size_t, multiset<double>> &prev_indices, bool rev, bool leafedges)
{
    prev_indices[v].clear();
    vector<size_t> adj;
    if (rev)
    {
        if (!leafedges)
            adj = tree.edgesTime[v];
        else
            adj = tree.leafEdges[v];
    }
    else
    {
        if (!leafedges)
            adj = tree.reverseEdgesTime[v];
        else
            adj = tree.reverseLeafEdges[v];
    }
    for (auto w : adj)
    {
        prev_indices[v].insert(index_order[w]);
    }
}

template <class T>
double GetMedian(const std::multiset<T> &data)
{
    if (data.empty())
        return 0;

    const size_t n = data.size();
    double median = 0;

    auto iter = data.cbegin();
    std::advance(iter, n / 2);

    // Middle or average of two middle values
    if (n % 2 == 0)
    {
        const auto iter2 = iter--;
        median = double(*iter + *iter2) / 2; // data[n/2 - 1] AND data[n/2]
    }
    else
    {
        median = *iter;
    }

    return median;
}

list<size_t> getMedianOrder(TemporalTree &tree, size_t v, unordered_map<size_t, double> &index_order, unordered_map<size_t, multiset<double>> &prev_indices, bool rev, bool leafedges)
{
    list<size_t> l;
    TemporalTree::TAdjacency *adjacency;
    if (rev)
    {
        if (leafedges)
            adjacency = &tree.leafEdges;
        else
            adjacency = &tree.edgesTime;
    }
    else
    {
        if (leafedges)
            adjacency = &tree.reverseLeafEdges;
        else
            adjacency = &tree.reverseEdgesTime;
    }
    if ((*adjacency).find(v) != (*adjacency).end() && (*adjacency)[v].size() > 0)
    {
        setInitialPrevVec(v, tree, index_order, prev_indices, rev, leafedges);
    }
    else
    {
        prev_indices[v].clear();
    }
    if (tree.isLeaf(v))
    {
        l.push_back(v);
    }
    else
    {
        vector<list<size_t>> orders;
        auto ws = tree.edgesHierarchy[v];
        vector<double> medians;
        vector<int> indices;
        vector<int> indicesnoprev;
        int j = 0;
        for (auto w : ws)
        {
            orders.push_back(getMedianOrder(tree, w, index_order, prev_indices, rev, leafedges));
            prev_indices[v].insert(prev_indices[w].begin(), prev_indices[w].end());
            medians.push_back(GetMedian(prev_indices[w]));
            if (prev_indices[w].size() == 0)
            {
                indicesnoprev.push_back(j);
            }
            else
            {
                indices.push_back(j);
            }
            j++;
        }
        sort(indices.begin(), indices.end(),
             [&](int A, int B) -> bool
             {
                 if ((prev_indices[ws[A]].size() == 0 || prev_indices[ws[B]].size() == 0) && index_order.find(ws[A]) != index_order.end() && index_order.find(ws[B]) != index_order.end())
                 {
                     return index_order[ws[A]] < index_order[ws[B]];
                 }
                 double wA = prev_indices[ws[A]].size() > 0 ? medians[A] : 0;
                 double wB = prev_indices[ws[B]].size() > 0 ? medians[B] : 0;
                 return wA < wB;
             });
        for (auto x : indicesnoprev)
            indices.push_back(x);
        sort(indices.begin(), indices.end(),
             [&](int A, int B) -> bool
             {
                 if ((prev_indices[ws[A]].size() == 0 || prev_indices[ws[B]].size() == 0) && index_order.find(ws[A]) != index_order.end() && index_order.find(ws[B]) != index_order.end())
                 {
                     return index_order[ws[A]] < index_order[ws[B]];
                 }
                 double wA = prev_indices[ws[A]].size() > 0 ? medians[A] : 0;
                 double wB = prev_indices[ws[B]].size() > 0 ? medians[B] : 0;
                 return wA < wB;
             });
        for (auto w : ws)
        {
            prev_indices[w].clear();
        }

        for (auto i : indices)
        {
            l.splice(l.end(), orders[i]);
        }
    }
    return l;
}

void fillIndex(TemporalTree &tree, size_t v, unordered_map<size_t, double> &index_order)
{
    if (tree.isLeaf(v))
        return;
    else
    {
        int s = 0;
        vector<double> indices;
        for (auto w : tree.edgesHierarchy[v])
        {
            fillIndex(tree, w, index_order);
            indices.push_back(index_order[w]);
            s += index_order[w];
        }
        assert(tree.edgesHierarchy[v].size() > 0);
        // index_order[v] = s / (double)tree.edgesHierarchy[v].size();
        index_order[v] = *min_element(indices.begin(), indices.end());
    }
}

TemporalTree computeTreeOrderBarycenter(TemporalTree tree, int sweeps, bool leafedges)
{
    TemporalTree outTree = tree.getSelfDeaggregated();
    outTree.computeReverseEdges();
    if (leafedges)
    {
        outTree.computeLeafEdges();
    }
    std::random_device rd;
    std::mt19937 rng(rd());
    map<int, vector<size_t>> t_nodes;
    map<int, vector<size_t>> t_roots;
    map<int, vector<size_t>> t_leaves;

    unordered_map<int, vector<size_t>> t_order;
    unordered_map<size_t, double> index_order;
    unordered_map<size_t, double> weight;
    unordered_map<size_t, int> size;

    set<int> stimes;
    for (size_t i = 1; i < outTree.nodes.size(); i++)
    {
        stimes.insert(outTree.nodes[i].startTime());
        t_nodes[outTree.nodes[i].startTime()].push_back(i);
        if (outTree.isLeaf(i))
            t_leaves[outTree.nodes[i].startTime()].push_back(i);
        if (outTree.reverseEdgesHierachy.find(i) != outTree.reverseEdgesHierachy.end() && outTree.reverseEdgesHierachy[i][0] == 0)
        {
            t_roots[outTree.nodes[i].startTime()].push_back(i);
        }
    }
    vector<int> times(stimes.begin(), stimes.end());
    int lasttime = -1;
    for (int RUNS = 0; RUNS < sweeps; RUNS++)
    {
        for (auto const &[time, nodesattime] : t_nodes)
        {
            list<size_t> leaforder;
            if (lasttime == -1)
            {
                shuffle(t_roots[time].begin(), t_roots[time].end(), rng);
                for (auto root : t_roots[time])
                {
                    auto order = getInitialOrder(outTree, root, rng);
                    leaforder.splice(leaforder.end(), order);
                }
                /* leaforder = list<size_t>{53, 24, 77, 16, 108, 134, 131, 128, 114, 101, 85, 62, 34, 1};
                for (auto leaf: leaforder) cout<<leaf<<" ";
                cout<<endl;
                for (auto leaf: t_leaves[time]) cout<<leaf<<" ";
                cout<<endl; */
            }
            else
            {
                if (time == times[0])
                    continue;
                vector<list<size_t>> orders;
                vector<double> barycenters;
                vector<size_t> ws;
                vector<int> indices;
                vector<int> indicesnoprev;
                int j = 0;
                for (auto root : t_roots[time])
                {
                    auto order = getBarycenterOrder(outTree, root, index_order, weight, size, false, leafedges);
                    orders.push_back(order);
                    ws.push_back(root);
                    if (size[root] == 0)
                    {
                        indicesnoprev.push_back(j);
                    }
                    else
                    {
                        indicesnoprev.push_back(j);
                    }
                    j++;
                }
                sort(indices.begin(), indices.end(),
                     [&](int A, int B) -> bool
                     {
                         if ((size[ws[A]] == 0 || size[ws[B]] == 0) && index_order.find(tree.someLeafChild(ws[A])) != index_order.end() && index_order.find(tree.someLeafChild(ws[B])) != index_order.end())
                         {
                             return index_order[ws[A]] < index_order[ws[B]];
                         }
                         double wA = size[ws[A]] > 0 ? (weight[ws[A]] / (double)size[ws[A]]) : 0;
                         double wB = size[ws[B]] > 0 ? (weight[ws[B]] / (double)size[ws[B]]) : 0;
                         return wA < wB;
                     });
                for (auto x : indicesnoprev)
                    indices.push_back(x);
                sort(indices.begin(), indices.end(),
                     [&](int A, int B) -> bool
                     {
                         if ((size[ws[A]] == 0 || size[ws[B]] == 0) && index_order.find(tree.someLeafChild(ws[A])) != index_order.end() && index_order.find(tree.someLeafChild(ws[B])) != index_order.end())
                         {
                             return index_order[ws[A]] < index_order[ws[B]];
                         }
                         double wA = size[ws[A]] > 0 ? (weight[ws[A]] / (double)size[ws[A]]) : 0;
                         double wB = size[ws[B]] > 0 ? (weight[ws[B]] / (double)size[ws[B]]) : 0;
                         return wA < wB;
                     });

                for (auto i : indices)
                {
                    leaforder.splice(leaforder.end(), orders[i]);
                }
            }
            t_order[time].clear();
            int i = 0;
            for (auto v : leaforder)
            {
                t_order[time].push_back(v);
                index_order[v] = i;
                i += 1;
            }
            for (auto root : t_roots[time])
            {
                fillIndex(outTree, root, index_order);
            }
            lasttime = time;
        }

        for (int ti = times.size() - 2; ti >= 0; ti--)
        {
            int time = times[ti];
            vector<list<size_t>> orders;
            list<size_t> leaforder;
            vector<double> barycenters;
            vector<size_t> ws;
            vector<int> indices;
            vector<int> indicesnoprev;
            int j = 0;
            for (auto root : t_roots[time])
            {
                auto order = getBarycenterOrder(outTree, root, index_order, weight, size, true, leafedges);
                orders.push_back(order);
                ws.push_back(root);
                if (size[root] == 0)
                {
                    indicesnoprev.push_back(j);
                }
                else
                {
                    indices.push_back(j);
                }
                j++;
            }
            sort(indices.begin(), indices.end(),
                 [&](int A, int B) -> bool
                 {
                     if ((size[ws[A]] == 0 || size[ws[B]] == 0) && index_order.find(tree.someLeafChild(ws[A])) != index_order.end() && index_order.find(tree.someLeafChild(ws[B])) != index_order.end())
                     {
                         return index_order[ws[A]] < index_order[ws[B]];
                     }
                     double wA = size[ws[A]] > 0 ? (weight[ws[A]] / (double)size[ws[A]]) : 0;
                     double wB = size[ws[B]] > 0 ? (weight[ws[B]] / (double)size[ws[B]]) : 0;
                     return wA < wB;
                 });
            for (auto x : indicesnoprev)
                indices.push_back(x);
            sort(indices.begin(), indices.end(),
                 [&](int A, int B) -> bool
                 {
                     if ((size[ws[A]] == 0 || size[ws[B]] == 0) && index_order.find(tree.someLeafChild(ws[A])) != index_order.end() && index_order.find(tree.someLeafChild(ws[B])) != index_order.end())
                     {
                         return index_order[ws[A]] < index_order[ws[B]];
                     }
                     double wA = size[ws[A]] > 0 ? (weight[ws[A]] / (double)size[ws[A]]) : 0;
                     double wB = size[ws[B]] > 0 ? (weight[ws[B]] / (double)size[ws[B]]) : 0;
                     return wA < wB;
                 });

            for (auto i : indices)
            {
                leaforder.splice(leaforder.end(), orders[i]);
            }
            t_order[time].clear();
            int i = 0;
            for (auto v : leaforder)
            {
                t_order[time].push_back(v);
                index_order[v] = i;
                i += 1;
            }
            for (auto root : t_roots[time])
            {
                fillIndex(outTree, root, index_order);
            }
            lasttime = time;
        }
    }

    outTree.order.clear();
    for (auto time : times)
    {
        outTree.order.insert(outTree.order.end(), t_order[time].begin(), t_order[time].end());
    }

    return outTree;
}

TemporalTree computeTreeOrderMedian(TemporalTree tree, int sweeps, bool leafedges)
{
    TemporalTree outTree = tree.getSelfDeaggregated();
    outTree.computeReverseEdges();
    std::random_device rd;
    std::mt19937 rng(rd());
    map<int, vector<size_t>> t_nodes;
    map<int, vector<size_t>> t_roots;
    map<int, vector<size_t>> t_leaves;

    unordered_map<int, vector<size_t>> t_order;
    unordered_map<size_t, double> index_order;
    unordered_map<size_t, multiset<double>> prev_indices;

    if (leafedges)
    {
        outTree.computeLeafEdges();
    }

    set<int> stimes;
    for (size_t i = 1; i < outTree.nodes.size(); i++)
    {
        stimes.insert(outTree.nodes[i].startTime());
        t_nodes[outTree.nodes[i].startTime()].push_back(i);
        if (outTree.isLeaf(i))
            t_leaves[outTree.nodes[i].startTime()].push_back(i);
        if (outTree.reverseEdgesHierachy.find(i) != outTree.reverseEdgesHierachy.end() && outTree.reverseEdgesHierachy[i][0] == 0)
        {
            t_roots[outTree.nodes[i].startTime()].push_back(i);
        }
    }
    vector<int> times(stimes.begin(), stimes.end());
    int lasttime = -1;
    for (int RUNS = 0; RUNS < sweeps; RUNS++)
    {
        for (auto const &[time, nodesattime] : t_nodes)
        {
            list<size_t> leaforder;
            if (lasttime == -1)
            {
                shuffle(t_roots[time].begin(), t_roots[time].end(), rng);
                for (auto root : t_roots[time])
                {
                    auto order = getInitialOrder(outTree, root, rng);
                    leaforder.splice(leaforder.end(), order);
                }
                /* leaforder = list<size_t>{53, 24, 77, 16, 108, 134, 131, 128, 114, 101, 85, 62, 34, 1};
                for (auto leaf: leaforder) cout<<leaf<<" ";
                cout<<endl;
                for (auto leaf: t_leaves[time]) cout<<leaf<<" ";
                cout<<endl; */
            }
            else
            {
                if (time == times[0])
                    continue;
                vector<list<size_t>> orders;
                vector<double> medians;
                vector<size_t> ws;
                vector<int> indices;
                vector<int> indicesnoprev;
                int j = 0;
                for (auto root : t_roots[time])
                {
                    auto order = getMedianOrder(outTree, root, index_order, prev_indices, false, leafedges);
                    orders.push_back(order);
                    ws.push_back(root);
                    medians.push_back(GetMedian(prev_indices[root]));
                    
                    if (prev_indices[root].size() == 0)
                    {
                        indicesnoprev.push_back(j);
                    }
                    else
                    {
                        indices.push_back(j);
                    }
                    j++;
                }
                
                sort(indices.begin(), indices.end(),
                     [&](int A, int B) -> bool
                     {
                         if ((prev_indices[ws[A]].size() == 0 || prev_indices[ws[B]].size() == 0) && index_order.find(ws[A]) != index_order.end() && index_order.find(ws[B]) != index_order.end())
                         {
                             return index_order[ws[A]] < index_order[ws[B]];
                         }
                         double wA = prev_indices[ws[A]].size() > 0 ? medians[A] : 0;
                         double wB = prev_indices[ws[B]].size() > 0 ? medians[B] : 0;
                         return wA < wB;
                     });
                for (auto x : indicesnoprev)
                    indices.push_back(x);
                sort(indices.begin(), indices.end(),
                     [&](int A, int B) -> bool
                     {
                         if ((prev_indices[ws[A]].size() == 0 || prev_indices[ws[B]].size() == 0) && index_order.find(ws[A]) != index_order.end() && index_order.find(ws[B]) != index_order.end())
                        {
                            return index_order[ws[A]] < index_order[ws[B]];
                        }
                         double wA = prev_indices[ws[A]].size() > 0 ? medians[A] : 0;
                         double wB = prev_indices[ws[B]].size() > 0 ? medians[B] : 0;
                         return wA < wB;
                     });

                for (auto i : indices)
                {
                    leaforder.splice(leaforder.end(), orders[i]);
                }
            }
            int i = 0;
            t_order[time].clear();
            for (auto v : leaforder)
            {
                t_order[time].push_back(v);
                index_order[v] = i;
                i += 1;
            }
            for (auto root : t_roots[time])
            {
                fillIndex(outTree, root, index_order);
            }
            lasttime = time;
        }

        for (int ti = times.size() - 2; ti >= 0; ti--)
        {
            int time = times[ti];
            vector<list<size_t>> orders;
            list<size_t> leaforder;
            vector<double> medians;
            vector<size_t> ws;
            vector<int> indices;
            vector<int> indicesnoprev;
            int j = 0;
            for (auto root : t_roots[time])
            {
                auto order = getMedianOrder(outTree, root, index_order, prev_indices, true, leafedges);
                medians.push_back(GetMedian(prev_indices[root]));
                orders.push_back(order);
                ws.push_back(root);
                if (prev_indices[root].size() == 0)
                {
                    indicesnoprev.push_back(j);
                }
                else
                {
                    indices.push_back(j);
                }
                j++;
            }
            sort(indices.begin(), indices.end(),
                 [&](int A, int B) -> bool
                 {
                     if ((prev_indices[ws[A]].size() == 0 || prev_indices[ws[B]].size() == 0) && index_order.find(ws[A]) != index_order.end() && index_order.find(ws[B]) != index_order.end())
                    {
                        return index_order[ws[A]] < index_order[ws[B]];
                    }
                     double wA = prev_indices[ws[A]].size() > 0 ? medians[A] : 0;
                     double wB = prev_indices[ws[B]].size() > 0 ? medians[B] : 0;
                     return wA < wB;
                 });
            for (auto x : indicesnoprev)
                indices.push_back(x);
            sort(indices.begin(), indices.end(),
                 [&](int A, int B) -> bool
                 {
                     if ((prev_indices[ws[A]].size() == 0 || prev_indices[ws[B]].size() == 0) && index_order.find(ws[A]) != index_order.end() && index_order.find(ws[B]) != index_order.end())
                    {
                        return index_order[ws[A]] < index_order[ws[B]];
                    }
                     double wA = prev_indices[ws[A]].size() > 0 ? medians[A] : 0;
                     double wB = prev_indices[ws[B]].size() > 0 ? medians[B] : 0;
                     return wA < wB;
                 });

            for (auto i : indices)
            {
                leaforder.splice(leaforder.end(), orders[i]);
            }
            t_order[time].clear();
            int i = 0;
            for (auto v : leaforder)
            {
                t_order[time].push_back(v);
                index_order[v] = i;
                i += 1;
            }
            for (auto root : t_roots[time])
            {
                fillIndex(outTree, root, index_order);
            }
            lasttime = time;
        }
    }

    outTree.order.clear();
    for (auto time : times)
    {
        outTree.order.insert(outTree.order.end(), t_order[time].begin(), t_order[time].end());
    }

    return outTree;
}

TemporalTree computeTreeOrderBarycenterLeafEdges(TemporalTree tree, int sweeps)
{
    return computeTreeOrderBarycenter(tree, sweeps, true);
}

TemporalTree computeTreeOrderMedianLeafEdges(TemporalTree tree, int sweeps)
{
    return computeTreeOrderMedian(tree, sweeps, true);
}