/*********************************************************************
 *  Author  : Tino Weinkauf and Wiebke K�pp
 *  Init    : Monday, October 09, 2017 - 16:06:40
 *
 *  Project : KTH Inviwo Modules
 *
 *  License : Follows the Inviwo BSD license model
 *********************************************************************
 */

#include "treewriter.h"
#include "treeorder.h"

TemporalTreeWriter::TemporalTreeWriter() {
    this->propOverwrite = true;
    this->propPrettyPrint = true;
    this->propNTGAddOrder = true;
    this->propNTGAddWeights = true;
    this->propNTGScaleWeights = 1.0;
}

namespace {

void GetLayerOrder(const TemporalTree& Tree, const std::map<size_t, size_t>& LeafOrderMap,
                   const std::vector<size_t>& LevelIndices, std::vector<size_t>& LayerOrder) {
    // Prepare memory
    const size_t NumLeaves = LeafOrderMap.size();
    if (NumLeaves == 0) return;

    // For actually sorting the parents
    std::vector<std::pair<size_t, size_t>> ParentFirstLeaf;
    ParentFirstLeaf.reserve(LevelIndices.size());

    // For each parent, get all its leaves and find the first one.
    std::set<size_t> LeavesOfParent;
    for (const size_t idxParent : LevelIndices) {
        // Get this parent's leaves
        LeavesOfParent.clear();

        const uint64_t tMinParent = Tree.nodes[idxParent].startTime();
        const uint64_t tMaxParent = Tree.nodes[idxParent].endTime();

        Tree.getLeaves(idxParent, tMinParent, tMaxParent, tMinParent, tMaxParent, LeavesOfParent);

        if (LeavesOfParent.size()) {
            // Find first leaf according to the given order
            size_t MinOrderRow(NumLeaves);
            for (const size_t idxLeaf : LeavesOfParent) {
                if (LeafOrderMap.count(idxLeaf) == 0) {
                    clog<<"Leaf " << idxLeaf << " not found! Size = "
                                                        << LeafOrderMap.size();
                    continue;
                }
                const size_t OrderRow = LeafOrderMap.at(idxLeaf);
                if (OrderRow < MinOrderRow) MinOrderRow = OrderRow;
            }

            // Add it to the sorting vector
            ParentFirstLeaf.push_back(std::make_pair(idxParent, MinOrderRow));
        } else {
            // We are a child
            ParentFirstLeaf.push_back(std::make_pair(idxParent, LeafOrderMap.at(idxParent)));
        }
    }

    // Sort parents according to their first child
    std::sort(ParentFirstLeaf.begin(), ParentFirstLeaf.end(),
              [&](const auto& a, const auto& b) { return a.second < b.second; });

    // Condense to a layer order
    const size_t NumElems = ParentFirstLeaf.size();
    LayerOrder.resize(NumElems);
    for (size_t i(0); i < NumElems; i++) {
        LayerOrder[i] = ParentFirstLeaf[i].first;
    }
}

}  // namespace

json TemporalTreeWriter::createJSON(std::shared_ptr<TemporalTree> tree, const bool bNTG,
                                    const bool bNTGAddWeights, const bool bNTGAddOrder,
                                    double nTGScaleWeights) {
    // Create a JSON from the tree
    json jAll;

    // Our version or Nested Tracking Graph?
    if (!bNTG) {
        json jNodes = json::array();
        for (auto& node : tree->nodes) {
            json jSingleNode;
            jSingleNode["name"] = node.name;
            jSingleNode["values"] = node.values;

            jNodes.push_back(jSingleNode);
        }

        jAll["nodes"] = jNodes;
        jAll["edgesHierarchy"] = tree->edgesHierarchy;
        jAll["edgesTime"] = tree->edgesTime;
        jAll["order"] = tree->order;
    } else {
        // Nested Tracking Graphs are stored and processed in a non-aggregated way
        // Node property w is responsible for the width of the band

        // Options like bNTGAddWeights, bNTGAddOrder, nTGScaleWeights are given as arguments to the
        // function

        // The three big arrays of the data structure
        json jNodes;
        json jTemporalEdges;
        json jHierarchicalEdges;

        // De-aggregate
        TemporalTree NTGTree;
        NTGTree = tree->getSelfDeaggregated();

        auto times = NTGTree.getTimes();

        for (auto time : times) {
            jHierarchicalEdges[std::to_string(time)] = json::object();
        }

        // Get a proper order
        TemporalTree::TTreeOrder Order(NTGTree.order);
        const size_t NumLeaves = Order.size();
        const bool bSaveWithOrder = (NumLeaves != 0) && bNTGAddOrder;
        TemporalTree::TTreeOrderMap OrderMap;
        treeorder::toOrderMap(OrderMap, Order);

        std::vector<size_t> LayerOrder;
        size_t Level(1);

        std::vector<size_t> LevelIndices;
        LevelIndices = NTGTree.getLevel(Level, LevelIndices, 0);
        while (!LevelIndices.empty()) {
            // In which order will we run through the nodes of this hierarchy layer?
            // The default is the order in which they initially come
            std::vector<size_t>* pNodeIndexArray = &LevelIndices;

            // Create an order of nodes in this layer, but depending on the leaves layer
            if (bSaveWithOrder) {
                GetLayerOrder(NTGTree, OrderMap, LevelIndices, LayerOrder);
                // Run through nodes in our sorted order
                pNodeIndexArray = &LayerOrder;
            }

            // We count how many nodes we already written per time step. Needed for topological
            // order.
            std::map<uint64_t, int> Stacks;

            // Order in which we run through the nodes
            const std::vector<size_t>& NodeIndexArray = *pNodeIndexArray;

            // Add to JSON nodes
            for (size_t i(0); i < NodeIndexArray.size(); i++) {
                // Shorthand
                const size_t& idxNode = NodeIndexArray[i];
                const TemporalTree::TNode& ThisNode = NTGTree.nodes[idxNode];
                const uint64_t NTGTime = ThisNode.values.cbegin()->first;

                // Prepare stack
                if (Stacks.count(NTGTime) == 0) Stacks[NTGTime] = 0;

                // Node
                json jSingleNode;
                jSingleNode["l"] = Level - 1;
                jSingleNode["t"] = NTGTime;
                if (bNTGAddWeights)
                    jSingleNode["w"] = ThisNode.values.cbegin()->second * nTGScaleWeights;

                // Node Layout based on order
                // - a layout array property with x, y, w variables.
                if (bSaveWithOrder) {
                    json jLayout;
                    jLayout["x"] = NTGTime;
                    jLayout["y"] = Stacks[NTGTime];
                    Stacks[NTGTime]++;

                    // Get number of siblings
                    // - get parent
                    std::vector<size_t> Parents = NTGTree.getHierarchicalParents(idxNode);
                    // - get number of its children
                    const size_t NumSiblings =
                        Parents.empty() ? 0 : NTGTree.getHierarchicalChildren(Parents[0]).size();
                    const float SiblingsFactor =
                        (Level - 1 > 0) ? (1.0f / float(NumSiblings)) : 1.0f;

                    // Topological size
                    // jLayout["w"] = 0.3 * SiblingsFactor * float((MaxLevel+1) - Level) /
                    // float(MaxLevel);

                    // Data size
                    jLayout["w"] = ThisNode.values.cbegin()->second * nTGScaleWeights;

                    jSingleNode["layout"] = jLayout;
                }

                // Add this node to the list of all nodes
                jNodes[std::to_string(idxNode)] = jSingleNode;

                // Add temporal edges
                const auto& Succ = NTGTree.getTemporalSuccessors(idxNode);
                if (!Succ.empty()) {
                    jTemporalEdges[std::to_string(Level - 1)][std::to_string(idxNode)] = Succ;
                }

                // Add hierarchical edges
                const auto& Children = NTGTree.getHierarchicalChildren(idxNode);
                if (!Children.empty()) {
                    jHierarchicalEdges[std::to_string(NTGTime)][std::to_string(idxNode)] = Children;
                }
            }

            // Forward to the next level
            LevelIndices = NTGTree.getLevel(Level + 1, LevelIndices, Level);
            Level++;
        }

        // Join all arrays in one big json
        jAll["N"] = jNodes;
        jAll["ET"] = jTemporalEdges;
        jAll["EN"] = jHierarchicalEdges;
    }

    return jAll;
}

void TemporalTreeWriter::process() {
    // Get filename and open file
    const std::string& Filename = propFilename;

    // Get the tree
    std::shared_ptr<TemporalTree> InTree = portInTree;

    std::ofstream outfile;
    outfile.exceptions(std::ofstream::failbit | std::ofstream::badbit);
    try {
        outfile.open(Filename);
    } catch (const std::ofstream::failure& e) {
        clog<<"File could not be opened: " << Filename;
        cerr<<"  Error Code: " << e.code() << "    . " << e.what();
        return;
    }

    // Filetype
    const bool bASCII = propFilename.rfind(".tree") != -1;
    const bool bCBOR = propFilename.rfind(".cbortree") != -1;
    const bool bNTG = propFilename.rfind(".ntg") != -1;
    // const bool bMsgPack = !(bASCII || bCBOR);

    json jAll = createJSON(InTree, bNTG, propNTGAddWeights, propNTGAddOrder,
                           propNTGScaleWeights);

    // Stream it out as ASCII or Binary
    try {
        if (bASCII || bNTG) {
            // ASCII
            if (propPrettyPrint) outfile << std::setw(4);
            outfile << jAll << std::endl;
        } else if (bCBOR) {
            // Binary CBOR
            std::vector<std::uint8_t> v_cbor = json::to_cbor(jAll);
            outfile.write(reinterpret_cast<char*>(&v_cbor[0]), v_cbor.size());
        } else {
            // Binary MessagePack
            std::vector<std::uint8_t> v_pack = json::to_msgpack(jAll);
            outfile.write(reinterpret_cast<char*>(&v_pack[0]), v_pack.size());
        }
    } catch (const std::ofstream::failure& e) {
        cerr<<"Error during save: " << Filename;
        cerr<<"  Error Code: " << e.code() << "    . " << e.what();
        return;
    }

    outfile.close();
}
