//
// Created by wrede on 02.06.16.
//

#include "Berclaz.h"
#include "../util/Parser.h"
#include "../util/Logger.h"
#include "KShortestPaths.h"

namespace algo
{
    Berclaz::Berclaz(int h_res, int v_res, int vicinity_size)
    {
        h_res_ = h_res;
        v_res_ = v_res;
        vicinity_size_ = vicinity_size;
    }

    void Berclaz::CreateGraph(DirectedGraph& graph, Vertex& source, Vertex& sink, util::Grid& grid)
    {
        util::Logger::LogDebug("add vertices");

        // Add grid vertices
        for (int z = 0; z < grid.GetDepthCount(); ++z)
        {
            for (int y = 0; y < grid.GetHeightCount(); ++y)
            {
                for (int x = 0; x < grid.GetWidthCount(); ++x)
                {
                    boost::add_vertex(grid.GetValue(x, y, z), graph);
                }
            }
        }

        // Add source and sink vertex
        source = boost::add_vertex(core::ObjectDataPtr(new core::ObjectData()), graph);
        sink = boost::add_vertex(core::ObjectDataPtr(new core::ObjectData()), graph);

        util::Logger::LogDebug("add edges");

        double min_weight = std::numeric_limits<double>::max();
        double max_weight = std::numeric_limits<double>::lowest();
        double min_score = std::numeric_limits<double>::max();
        double max_score = std::numeric_limits<double>::lowest();

        // Iterate all vertices but source and sink
        VertexIndexMap vertices = boost::get(boost::vertex_index, graph);
        VertexValueMap values = boost::get(boost::vertex_name, graph);
        int layer_size = grid.GetWidthCount() * grid.GetHeightCount();
        for (int z = 0; z < grid.GetDepthCount(); ++z)
        {
            for (int y = 0; y < grid.GetHeightCount(); ++y)
            {
                for (int x = 0; x < grid.GetWidthCount(); ++x)
                {
                    // First vertex index
                    int vi = x + y * grid.GetWidthCount() + z * layer_size;

                    // Get the score, clamp it, prevent division by zero and
                    // logarithm of zero
                    double score = values[vi]->GetDetectionScore();

                    if (score < min_score)
                        min_score = score;

                    if (score > max_score)
                        max_score = score;

                    if (score > MAX_SCORE_VALUE)
                    {
                        score = MAX_SCORE_VALUE;
                    }
                    else if (score < MIN_SCORE_VALUE)
                    {
                        score = MIN_SCORE_VALUE;
                    }

                    // Calculate the edge weight
                    double weight = -std::log(score / (1 - score));

                    if (weight < min_weight)
                        min_weight = weight;

                    if (weight > max_weight)
                        max_weight = weight;

                    // Connect with the next frame only if there is a next frame
                    if (z < grid.GetDepthCount() - 1)
                    {
                        // Iterate all nearby cells in the next frame
                        for (int ny = std::max(0, y - vicinity_size_);
                             ny < std::min(grid.GetHeightCount(), y + vicinity_size_ + 1);
                             ++ny)
                        {
                            for (int nx = std::max(0, x - vicinity_size_);
                                 nx < std::min(grid.GetWidthCount(), x + vicinity_size_ + 1);
                                 ++nx)
                            {
                                // Second vertex index
                                int vj = nx + ny * grid.GetWidthCount() + (z + 1) * layer_size;

                                // Connect to nearby cells
                                boost::add_edge(vertices[vi], vertices[vj], weight, graph);
                            }
                        }

                        boost::add_edge(vertices[vi], sink, VIRTUAL_EDGE_WEIGHT, graph);
                    }
                    else
                    {
                        boost::add_edge(vertices[vi], sink, weight, graph);
                    }

                    // Connect with source and sink
                    boost::add_edge(source, vertices[vi], VIRTUAL_EDGE_WEIGHT, graph);
                }
            }
        }

        util::Logger::LogDebug("weight " + std::to_string(min_weight) + " to " + std::to_string(max_weight));
        util::Logger::LogDebug("score " + std::to_string(min_score) + " to " + std::to_string(max_score));
        util::Logger::LogDebug("vertex count " + std::to_string(boost::num_vertices(graph)));
        util::Logger::LogDebug("edge count " + std::to_string(boost::num_edges(graph)));
    }

    void Berclaz::ExtractTracks(DirectedGraph& graph, MultiPredecessorMap& map, Vertex origin,
                                std::vector<core::TrackletPtr>& tracks)
    {
        VertexValueMap values = boost::get(boost::vertex_name, graph);

        // Move along all paths in reverse, starting at the origin
        for (Vertex first : map[origin])
        {
            core::TrackletPtr tracklet(new core::Tracklet());

            // The paths are node disjoint, so there should always be only one
            // node to proceed to
            for (Vertex u = first, v = (*map[u].begin());
                 u != v; u = v, v = (*map[v].begin()))
            {
                tracklet->AddPathObject(values[u]);
            }

            tracks.push_back(tracklet);
        }
    }

    void Berclaz::Run(core::DetectionSequence& sequence,
                      size_t batch_size, size_t max_track_count,
                      std::vector<core::TrackletPtr> & tracks, util::Filter2D & filter)
    {
        size_t batch_count = 0;
        for (size_t i = sequence.GetFrameOffset(); i < sequence.GetFrameCount(); i += batch_size)
        {
            util::Logger::LogDebug("batch offset: " + std::to_string(i));

            util::Grid grid = util::Parser::ParseGrid(sequence, i, i + batch_size,
                                                      0.0, 1.0, h_res_, 0.0, 1.0, v_res_);

            grid.Convolve2D(filter);

            util::Logger::LogDebug("create graph");
            DirectedGraph graph;
            Vertex source, sink;
            CreateGraph(graph, source, sink, grid);

            util::Logger::LogDebug("run ksp");
            KShortestPaths ksp(graph, source, sink);
            ksp.Run(max_track_count);

            util::Logger::LogDebug("get paths");
            std::vector<std::vector<Vertex>> paths;
            ksp.GetPaths(paths);

            util::Logger::LogDebug("extract tracks");
            VertexValueMap values = boost::get(boost::vertex_name, graph);

            util::Logger::LogDebug("");
            for (auto p : paths)
            {
                std::string path_string = "";
                core::TrackletPtr tlt(new core::Tracklet());
                for (auto v : p)
                {
                    path_string += std::to_string(v) + " ";
                    tlt->AddPathObject(values[v]);
                }
                tracks.push_back(tlt);
                util::Logger::LogDebug(path_string);
            }
            util::Logger::LogDebug("");

            batch_count++;
        }

        // Only connect tracks if the sequence was split
        if (batch_count > 1)
        {
            //TODO find a better way to connect tracks (n-stage)
            util::Logger::LogDebug("connect tracks");
            ConnectTracks(tracks);
        }
    }

    void Berclaz::ConnectTracks(std::vector<core::TrackletPtr>& tracks)
    {
        for (size_t i = 0; i < tracks.size(); ++i)
        {
            // find the best matching tracklet
            double best_value = std::numeric_limits<double>::max();
            size_t best_index = 0;
            for (size_t k = i + 1; k < tracks.size(); ++k)
            {
                if (tracks[i]->GetLastFrameIndex() < tracks[k]->GetFirstFrameIndex())
                {
                    double value = tracks[i]->CompareTo(tracks[k]);
                    if (value < best_value)
                    {
                        best_index = k;
                    }
                }
            }

            // if a match was found
            if (best_index != 0)
            {
                // merge the two tracks
                tracks[i]->Combine(tracks[best_index]);
                tracks.erase(tracks.begin() + best_index);
            }
        }
    }
}