/*
 * markerAssociationDP.cpp
 *
 *  Created on: Feb 6, 2013
 *      Author: koerner
 */

#include <boost/config.hpp>

#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/dijkstra_shortest_paths.hpp>
#include <boost/graph/graph_traits.hpp>
#include <boost/graph/graphviz.hpp>
#include <boost/graph/iteration_macros.hpp>
#include <boost/graph/properties.hpp>

#include <boost/property_map/property_map.hpp>

#include <boost/program_options.hpp>
#include <boost/format.hpp>

#include <opencv2/opencv.hpp>

#include <limits>
#include <iostream>
#include <fstream>
#include <utility>
#include <vector>

/**
 * command line parameters
 */

// command line parameters
struct Params {
    // filename of 3d point filen
    std::string strTracklets;

    // output filename
    std::string outputFilename;

    // number of tracklets to be extracted
    size_t trackCount;

    // maximum temporal distance (in frames) for two tracklets to get linked
    size_t maxTemporalDist;

    // maximum spatial distance (in mm) for two tracklets to get linked
    double maxSpatialDist;

    // minimum length of tracks to be selected
    double minTrackLength;
};

// get command line parameters
bool evalCmdLine(int argc, char **argv, Params &p) {
    // define parameters
    boost::program_options::options_description desc("Allowed options");
    desc.add_options()
            ("help", "produce help message")
            ("trackletFile,t", boost::program_options::value< std::string >(), "tracklet input file (as generated by \"graphBasedTracklets\")")
            ("outputFile,o", boost::program_options::value< std::string >(), "output file")
            ("trackCount,c", boost::program_options::value< size_t >()->default_value(1), "number of tracks to be extracted")
            ("maxSpatialDist,d",
             boost::program_options::value< double >()->default_value(1.5),
             "maximum spatial distance (in mm) for two tracklets to get linked")
            ("maxTemporalDist,s",
             boost::program_options::value< size_t >()->default_value(150),
             "maximum temporal distance (in frames) for two tracklets to get linked")
            ("minTrackLength,l", boost::program_options::value< double >()->default_value(0.85), "minimum length of tracks to be selected (if value is <= 1, it is interpreted as fraction of the sequence length)")
            ;

    // parse parameters
    boost::program_options::variables_map vm;
    boost::program_options::store(boost::program_options::parse_command_line(argc, argv, desc), vm);
    boost::program_options::notify(vm);

    if (vm.count("help") || vm.empty()) {
        std::cout << desc << std::endl;
        exit(EXIT_SUCCESS);
    }

    // default parameters
    if (vm.count("trackletFile")) {
        p.strTracklets = vm["trackletFile"].as< std::string >();
    } else {
        std::cerr << "No tracklet input file passed!\n";
        std::cout << desc << std::endl;
        exit(EXIT_SUCCESS);
    }
    if (vm.count("outputFile")) {
        p.outputFilename = vm["outputFile"].as< std::string >();
    } else {
        std::cerr << "No output file passed!" << std::endl;
        std::cout << desc << std::endl;
        exit(EXIT_SUCCESS);
    }

    p.trackCount = vm["trackCount"].as< size_t >();
    p.maxTemporalDist = vm["maxTemporalDist"].as< size_t >();
    p.maxSpatialDist = vm["maxSpatialDist"].as< double >();
    p.minTrackLength = vm["minTrackLength"].as< double >();

    return true;
}

/**
 * helper functions
 */

double mean(const std::vector< double > &values) {
    if (values.size() == 0) {
        return 0.0;
    }

    double sum = 0.0;
    for (size_t k = 0; k < values.size(); k++) {
        sum += values[k];
    }

    return sum / double(values.size());
}

/**
 * tracklet handling
 */

// tracklet class
class Tracklet {
public:
    // constructors
    Tracklet() {
        this->rawPoints.clear();
        this->ids.clear();
    }

    Tracklet(const Tracklet &tracklet) {
        this->rawPoints = tracklet.rawPoints;
        this->ids = tracklet.ids;
    }

    Tracklet& operator=(const Tracklet &tracklet) {
        this->rawPoints = tracklet.rawPoints;
        this->ids = tracklet.ids;
        return *this;
    }

    // getters/setters
    bool isVirtual() const {
        return this->rawPoints.empty();
    }

    std::string getIdString() const {
        std::stringstream ss;
        for (size_t k = 0; k < this->ids.size(); k++) {
            if (k > 0) {
                ss << ",";
            }
            ss << ids[k];
        }
        return ss.str();
    }

    std::vector< size_t > &getIds() {
        return this->ids;
    }

    const std::vector< size_t > &getIds() const {
        return this->ids;
    }

    void appendId(size_t _id) {
        this->ids.push_back(_id);
    }

    void appendIds(std::vector< size_t > _ids) {
        this->ids.insert(this->ids.end(), _ids.begin(), _ids.end());
    }

    size_t getFirstFrame() const {
        //return this->firstFrame;
        return this->rawPoints.begin()->first;
    }

    size_t getLastFrame() const {
        return this->rawPoints.rbegin()->first;
    }

    size_t getRange() const {
        return this->getLastFrame() - this->getFirstFrame() + 1;
    }

    const cv::Point3d &getFirstPoint() const {
        return this->rawPoints.begin()->second;
    }

    const cv::Point3d &getLastPoint() const {
        return this->rawPoints.rbegin()->second;
    }

    void printDebugInfo() const {
        std::cout << "=== DEBUG ===" << std::endl;
        std::cout << "IDs: " << this->getIdString() << std::endl;
        std::cout << "frames: (" << this->getFirstFrame() << "," << this->getLastFrame() << ")" << std::endl;
        std::cout << "raw point count: " << this->pRawPoints()->size() << std::endl;
        std::cout << "first point: " << this->rawPoints.begin()->second << std::endl;
    }

    void interpolatePoints(void) {
        for (std::map< size_t, cv::Point3d >::iterator it = this->rawPoints.begin();
             it != this->rawPoints.end(); ++it) {
            cv::Point3d p = it->second;

//        std::cout << it->first << p;

            std::map< size_t, cv::Point3d >::iterator itNext = it;
            std::advance(itNext, 1);

            if (itNext != this->rawPoints.end()) {
                size_t dT = itNext->first - it->first;

                cv::Point3d q = itNext->second;
                cv::Point3d v = (q - p) * (1.0 / (double)(dT));

                if (dT > 0) {
                    for (size_t i = 1; i < dT; ++i) {
//              std::cout << " -> " << it->first + i << (p + v * (double)i);
                        this->rawPoints.insert(std::make_pair(it->first + i, p + v * (double)i));
                    }
                }
            }
            //std::cout << std::endl;
        }
    }

    const double getMatchingScore(const Tracklet &tracklet) {
        size_t tStart = std::max(this->getFirstFrame(), tracklet.getFirstFrame());
        size_t tEnd = std::min(this->getLastFrame(), tracklet.getLastFrame());

        int tDiff = tEnd - tStart;
        double tRatio = 2.0 * (double)tDiff / (double)(this->pRawPoints()->size() + tracklet.pRawPoints()->size() - 2);

        if (tRatio > 0.0) {
            const std::map< size_t, cv::Point3d >* p = this->pRawPoints();
            const std::map< size_t, cv::Point3d >* q = tracklet.pRawPoints();
            double dist = 0.0f;
            for (size_t i = tStart; i <= tEnd; ++i) {
                dist += cv::norm(p->at(i) - q->at(i));
            }
            dist /= (double)tDiff;

            return dist;
        } else {
            return std::numeric_limits< double >::max();
        }
    }

    const std::map< size_t, cv::Point3d >* pRawPoints() const {
        return &this->rawPoints;
    }

    void addPoint(size_t _frame, cv::Point3d &_point) {
        this->rawPoints[_frame] = _point;
    }

    void append(Tracklet &tracklet) {
        this->rawPoints.insert(tracklet.pRawPoints()->begin(), tracklet.pRawPoints()->end());
        this->appendIds(tracklet.getIds());
    }

    void getVelocities(std::vector< double > &velocities) {
        // clear object to be returned
        velocities = std::vector< double >(0);

        // need at least two points
        if (this->rawPoints.size() <= 1) {
            return;
        }

        // for consecutive points (in time), calculate magnitude of difference (a.k.a. velocity)
        for (std::map< size_t, cv::Point3d >::iterator it = this->rawPoints.begin(); it != this->rawPoints.end(); it++) {
            std::map< size_t, cv::Point3d >::iterator it2 = it;
            it2++;
            if (it2 != this->rawPoints.end()) {
                double dx = it->second.x - it2->second.x;
                double dy = it->second.y - it2->second.y;
                double v = sqrt(dx * dx + dy * dy);
                velocities.push_back(v);
            }
        }
    }

    // linking cost to another tracklet (assuming "this" comes before "target" in time)
    double linkingCost(Tracklet &target, const Params &params) {
        // temporal distance (if "target" does not start after "this" ends, linking is not possible)
        int dTemporal = int(target.getFirstFrame()) - this->getLastFrame();
        if ((dTemporal <= 0) || (dTemporal > int(params.maxTemporalDist))) {
            return std::numeric_limits< double >::infinity();
        }

        // distance between nearest end-points
        double dx = target.getFirstPoint().x - this->getLastPoint().x;
        double dy = target.getFirstPoint().y - this->getLastPoint().y;
        double dSpatial = sqrt(dx * dx + dy * dy);
        if (dSpatial > params.maxSpatialDist) {
            return std::numeric_limits< double >::infinity();
        }

        // angular distance between nearest end-points
        double dz = target.getFirstPoint().z - this->getLastPoint().z;
        double dAngular = std::max(0.0, abs(dz) - 5.0);
        if (dAngular > 25) {
            return std::numeric_limits< double >::infinity();
        }

        // difference between mean velocities
        std::vector< double > velocitiesThis, velocitiesTarget;
        this->getVelocities(velocitiesThis);
        target.getVelocities(velocitiesTarget);
        double dMeanVelocities = std::abs(mean(velocitiesThis) - mean(velocitiesTarget));

        //

        // final cost
        return 0.1 * double(dTemporal) * (dAngular + dSpatial + dMeanVelocities);
    }

protected:
    std::map< size_t, cv::Point3d > rawPoints;
    std::vector< size_t > ids;
};

void mergeDuplicates(std::vector< Tracklet > &tracklets, const double maxScore = 2.0) {
    // iterate over all tracklets
    for (std::vector< Tracklet >::iterator it1 = tracklets.begin(); it1 != tracklets.end(); ++it1) {

        // get maximum time span
        size_t tMin = it1->getFirstFrame();
        size_t tMax = it1->getLastFrame();

        // collect tracklets to be merged
        std::vector< std::vector< Tracklet >::iterator > toBeMerged;

        // first one by default
        toBeMerged.push_back(it1);

        // iterate over all succeeding tracklets
        for (std::vector< Tracklet >::iterator it2 = it1 + 1; it2 != tracklets.end(); ++it2) {
            double match = it1->getMatchingScore(*it2);
            if (match < maxScore) {
                /*
                 std::cout
                 << std::distance(tracklets.begin(), it1) << " " << it1->getIdString() << " -> "
                 << std::distance(tracklets.begin(), it2) << " " << it2->getIdString() << ": "
                 << match << std::endl;
                 */

                // update time span
                tMin = std::min(tMin, it2->getFirstFrame());
                tMax = std::max(tMax, it2->getLastFrame());

                // add to merge candidates
                toBeMerged.push_back(it2);
            }
        }

        if (toBeMerged.size() > 1) {
            std::cout << "Merging " << toBeMerged.size() << " tracks!" << std::endl;
            //std::cout << "tMin: " << tMin << ", tMax: " << tMax << std::endl;

            // our new tracklet
            Tracklet mergedTracklet;

            // iterate over complete time span
            for (size_t t = tMin; t <= tMax; ++t) {
                cv::Point3d p;
                size_t nP = 0;

                // get point for current time step of each merging candidate
                for (size_t i = 0; i < toBeMerged.size(); ++i) {
                    std::map< size_t, cv::Point3d >::const_iterator itT = toBeMerged[i]->pRawPoints()->find(t);
                    if (itT != toBeMerged[i]->pRawPoints()->end()) {
                        p = p + itT->second;
                        ++nP;
                    }
                }

                // result is the average of all candidate points
                p *= 1.0 / (double)nP;

                // add to new tracklet
                mergedTracklet.addPoint(t, p);
            }

            // delete merged tracks
            for (int i = toBeMerged.size() - 1; i >= 0; --i) {
                mergedTracklet.appendIds(toBeMerged[i]->getIds());
                if (i > 0) {
                    tracklets.erase(toBeMerged[i]);
                }
            }

            // replace first tracklet by the merged one
            *it1 = mergedTracklet;
        }
    }
}

void readTracklets(const std::string &filename,
                   std::vector< Tracklet > &tracklets,
                   size_t &nFrames) {
    // open file
    std::cout << "Reading tracklets from file " << filename << std::endl;
    std::ifstream pointFile(filename.c_str(), std::ios_base::in);
    if (!pointFile.is_open()) {
        std::cerr << "Could not open file " << std::endl;
        return;
    }

    // result object
    tracklets.clear();

    // read in all 3d points
    std::string line;

    // for each line...
    nFrames = 0;
    size_t id = 0;
    while (std::getline(pointFile, line)) {
        // print current frame number
        //std::cerr << "\rReading points... " << std::flush;

        // points of the current frame
        std::vector< cv::Point3d > tracklet;

        // read 3d points for each camera
        Tracklet tmpTracklet;
        tmpTracklet.appendId(id++);

        std::stringstream split(line);
        cv::Point3d tmpPoint;
        size_t frame = 0;

        // for each point...
        while (split.good()) {
            // read all points of this line
            if ((split >> tmpPoint.x >> tmpPoint.y >> tmpPoint.z >> frame).good()) {
                tmpTracklet.addPoint(frame, tmpPoint);

                // update global frame count
                nFrames = std::max(nFrames, tmpTracklet.getLastFrame() + 1);
            }
        }

        // add tracklet to list
        tmpTracklet.interpolatePoints();
        tracklets.push_back(tmpTracklet);
    }
    uint nTracklets = tracklets.size();
    std::cerr << std::endl << "Read " << nTracklets << " tracklets" << std::endl;

    // merge duplicate tracklets
    mergeDuplicates(tracklets);
    nTracklets = tracklets.size();
    std::cerr << std::endl << "Merged to " << nTracklets << " tracklets" << std::endl;
}

/**
 * graph-based typedefs
 */

typedef double Weight;
typedef boost::property< boost::edge_weight_t, Weight > WeightProperty;
typedef boost::property< boost::vertex_name_t, Tracklet > NameProperty;

typedef boost::adjacency_list<
        boost::listS,
        boost::vecS,
        boost::directedS,
        NameProperty,
        WeightProperty > Graph;

typedef boost::graph_traits< Graph >::vertex_descriptor Vertex;

typedef boost::property_map< Graph, boost::vertex_index_t >::type IndexMap;
typedef boost::property_map< Graph, boost::vertex_name_t >::type NameMap;

typedef boost::iterator_property_map< Vertex*, IndexMap, Vertex, Vertex& > PredecessorMap;
typedef boost::iterator_property_map< Weight*, IndexMap, Weight, Weight& > DistanceMap;

//
template< class G >
class VertexWriter {
public:
    VertexWriter(G _g) :
            g(_g) {
        this->nameMap = boost::get(boost::vertex_name, g);
    }

    template< class VertexOrEdge >
    void operator()(std::ostream& out, const VertexOrEdge& v) const {
        if (nameMap[v].isVirtual()) {
            out << "[label=\"" << nameMap[v].getIdString() << "\",shape=\"circle\"]";
        } else {
            out << "[label=\"" << nameMap[v].getIdString() << ":\\n(" << nameMap[v].getFirstFrame() << ","
            << nameMap[v].getLastFrame() << ")\",shape=\"box\"]";
        }
        //
    }
private:
    G g;
    NameMap nameMap;
};

//
template< class G >
class EdgeWriter {
public:
    EdgeWriter(G _g) :
            g(_g) {
    }
    template< class VertexOrEdge >
    void operator()(std::ostream& out, const VertexOrEdge& v) const {
        out << "[label=\"" << boost::get(boost::edge_weight, g, v) << "\"]";
    }
private:
    G g;
};

// tracklet graph to dot file
void trackletsToDot(std::string filename, const Graph &g) {
    std::cout << "Writing graph to file \"" << filename << "\"" << std::endl;
    std::ofstream graphFile(filename.c_str());
    if (!graphFile.is_open()) {
        std::cerr << "Could not open file \"" << filename << "\"" << std::endl;
        return;
    }

    VertexWriter< Graph > vertexWriter(g);
    EdgeWriter< Graph > edgeWriter(g);

    boost::write_graphviz(graphFile, g, vertexWriter, edgeWriter);

    graphFile.close();
}

/**
 * main tracking functions
 */

// nFrames: global frame count (needed to calculated weights to vSink)
// return tracks
void createTracksFromTracklets(const std::vector< Tracklet > &tracklets,
                               const Params &params,
                               const size_t nFrames,
                               std::vector< Tracklet > &tracks) {
    //
    // build graph
    //

    Graph g;
    std::vector< Vertex > vertices;

    size_t trackletStart = 0;
    size_t trackletMax = tracklets.size();

    // adding vertices
    for (size_t iTracklet = trackletStart; iTracklet < trackletMax; ++iTracklet) {
        std::cout << "\rAdding node for tracklet " << iTracklet + 1 << std::flush;
        Vertex v = boost::add_vertex(tracklets[iTracklet], g);
        vertices.push_back(v);
    }

    // virtual source and sink
    Vertex vSource = boost::add_vertex(Tracklet(), g);
    Vertex vSink = boost::add_vertex(Tracklet(), g);

    // debug
    ///boost::add_edge(vSource, vSink, 0, g);

    NameMap nameMap = boost::get(boost::vertex_name, g);

    std::cout << std::endl;

    // adding edges (layer eq tracklets)
    for (size_t iTracklet = 0; iTracklet < vertices.size(); ++iTracklet) {
        // all points of frame "iLayers"
        Vertex tracklet1 = vertices[iTracklet];

        // create edges to source and sink
        std::cout << "\rAdding edges vSource -> " << iTracklet + 1 << " -> vSink" << std::flush;
        boost::add_edge(vSource, tracklet1, (nameMap[tracklet1].getFirstFrame() + 1) * 1, g);
        boost::add_edge(tracklet1, vSink, (nFrames - nameMap[tracklet1].getLastFrame()) * 1, g);

        for (size_t jTracklet = 0; jTracklet < vertices.size(); ++jTracklet) {
            // get second tracklet for linking
            Vertex tracklet2 = vertices[jTracklet];

            // this is where the magic happens... the cost function!
            double ijCost = nameMap[tracklet1].linkingCost(nameMap[tracklet2], params);

            // create edge between both nodes
            if (ijCost < std::numeric_limits< double >::infinity()) {
                boost::add_edge(tracklet1, tracklet2, ijCost, g);
            }
        }
    }
    std::cout << std::endl;

    //
    // find tracks (shortest paths in tracklet graph)
    //

    // Create things for Dijkstra
    std::vector< Vertex > predecessors(boost::num_vertices(g));  // To store parents
    std::vector< Weight > distances(boost::num_vertices(g));  // To store distances

    IndexMap indexMap = boost::get(boost::vertex_index, g);
    PredecessorMap predecessorMap(&predecessors[0], indexMap);
    DistanceMap distanceMap(&distances[0], indexMap);

    // export graph to graphviz dot file
    //trackletsToDot("tmp2.dot", g);

    // each track consists of merged tracklets (which in turn is again a tracklet)
    tracks.clear();

    // Compute shortest paths from input layer vertices to the sink
    for (size_t tracklet = 0; tracklet < params.trackCount; ++tracklet) {
        // preparations for Dijkstra
        boost::dijkstra_shortest_paths(g,         // the graph
                                       vSource,   // the source vertex
                                       boost::distance_map(distanceMap).predecessor_map(predecessorMap));

        // stop if no more tracks can be found
        if (distanceMap[vSink] == std::numeric_limits<double>::max()) {
            std::cout << "No more tracks can be found" << std::endl;
            break;
        }

        // Extract the shortest path
        typedef std::vector< Graph::edge_descriptor > PathType;

        PathType path;

        Vertex v = vSink;  // We want to start at the sink and work our way back to the source
        for (Vertex u = predecessorMap[v]; u != v; v = u, u = predecessorMap[v]) {
            std::pair< Graph::edge_descriptor, bool > edgePair = boost::edge(u, v, g);
            Graph::edge_descriptor edge = edgePair.first;
            path.push_back(edge);
        }

        // traverse shortest path (actually: traverse reverse path in a reverse direction)
        Tracklet currentTrack;
        bool vFirst = true;
        for (PathType::reverse_iterator pathIterator = path.rbegin(); pathIterator != path.rend(); ++pathIterator) {
            // append non-virtual tracklets to this path
            if (!nameMap[boost::source(*pathIterator, g)].isVirtual()) {
                currentTrack.append(nameMap[boost::source(*pathIterator, g)]);
            }

            // for each node of the path (but not the first, i.e. vSource), set the weights of all outgoing edges to Inf (for Dijkstra equivalent to deletion of the node)
            if (!vFirst) {
                std::pair< Graph::out_edge_iterator, Graph::out_edge_iterator > edgeIts = boost::out_edges(boost::source(*pathIterator,
                                                                                                                         g),
                                                                                                           g);
                for (Graph::out_edge_iterator edgeIt = edgeIts.first; edgeIt != edgeIts.second; edgeIt++) {
                    boost::get(boost::edge_weight, g, *edgeIt) = std::numeric_limits< Weight >::infinity();
                }
            }
            vFirst = false;
        }

        // interpolate
        currentTrack.interpolatePoints();

        tracks.push_back(currentTrack);

        std::cout << "Found track #" << tracklet + 1 << "/" << params.trackCount << ": " << path.size() - 1 << " tracklet(s), cost: "
        << distanceMap[vSink] << std::endl;
    }

    // merge duplicate tracks
    mergeDuplicates(tracks);
    std::cerr << std::endl << "Merged to " << tracks.size() << " tracks" << std::endl;
}

void writeTracks(std::string filename, std::vector< Tracklet > &tracks, size_t nFrames, double minTrackLength) {
    // open file
    std::ofstream f(filename.c_str());

    // get minimum track length
    size_t minLength;
    if (minTrackLength <= 1.0) {
        minLength = std::min(size_t(minTrackLength * double(nFrames)), nFrames);
    } else {
        minLength = std::min(size_t(minTrackLength), nFrames);
    }

    // print each track
    size_t finalCount = 0;
    for (size_t track = 0; track < tracks.size(); ++track) {
        if (tracks[track].getRange() >= minLength) {
            const std::map< size_t, cv::Point3d > *p = tracks[track].pRawPoints();
            for (std::map< size_t, cv::Point3d >::const_iterator it = p->begin(); it != p->end(); it++) {
                f << it->second.x << " " << it->second.y << " " << it->second.z << " " << it->first << " ";
            }
            //f << tracks[track].getIdString()
            f << std::endl;
            finalCount++;
        }
    }

    f.close();

    std::cout << finalCount << " tracks with length >= " << minLength << " were found" << std::endl;
}

int main(int argc, char** argv) {
    // get command line parameters
    Params p;
    if (!evalCmdLine(argc, argv, p)) {
        std::cerr << "Error while parsing arguments!" << std::endl;
        return 1;
    }

    size_t nFrames;
    std::vector< Tracklet > tracklets, tracks;
    readTracklets(p.strTracklets, tracklets, nFrames);
    createTracksFromTracklets(tracklets, p, nFrames, tracks);
    writeTracks(p.outputFilename, tracks, nFrames, p.minTrackLength);

    return EXIT_SUCCESS;
}