#include "sfm_positioning_evaluation.h"

#include <maps/wikimap/mapspro/services/mrc/libs/carsegm/include/carsegm.h>
#include <maps/wikimap/mapspro/services/mrc/libs/common/include/exif.h>
#include <maps/wikimap/mapspro/services/mrc/signs_detector/lib/camera.h>
#include <maps/wikimap/mapspro/services/mrc/tools/experiment_sfm_signs_positioning/lib/sfm_positioning.h>
#include <maps/wikimap/mapspro/services/mrc/libs/common/include/threadpool_wrapper.h>
#include <maps/libs/geolib/include/conversion.h>
#include <maps/libs/geolib/include/direction.h>
#include <maps/libs/geolib/include/distance.h>
#include <maps/libs/log8/include/log8.h>

#include <iostream>
#include <iterator>
#include <mutex>

namespace maps::mrc::experiment_sfm_positioning {

namespace tws = maps::mrc::tracks_with_sensors;
namespace sfp = maps::mrc::sensors_feature_positioner;
namespace sfm = maps::mrc::experiment_sfm_positioning;

namespace {

using pos_improvment::ImprovedGpsEvent;

enum class CameraMode { Portrait, Landscape };

CameraMode inferCameraMode(const ImprovedGpsEvent& improvedPosition)
{
    if (abs(improvedPosition.cameraUpDirection().z())
        > abs(improvedPosition.cameraRightDirection().z())) {
        return CameraMode::Portrait;
    }
    return CameraMode::Landscape;
}

// Translate camera orientation from the Android phone coordinate system into
// the image coordinate system SfM positioning expects it to be. Notice that
// SfM expects that a vertical axis is to be directed down.
cv::Mat inferImageGlobalOrientation(const ImprovedGpsEvent& improvedPosition)
{
    // These are default values if camera is in normal orientation while
    // taking pictures
    geolib3::Vector3 X = improvedPosition.cameraRightDirection();     // Right
    geolib3::Vector3 Y = improvedPosition.cameraUpDirection() * -1.0; // Down
    geolib3::Vector3 Z = improvedPosition.cameraFrontDirection(); // Forward

    switch (inferCameraMode(improvedPosition)) {
    default:
        REQUIRE(false, "Undefined camera mode");
    case CameraMode::Portrait:
        if (improvedPosition.cameraUpDirection().z() < 0) {
            // The camera is upside down
            X *= -1.0;
            Y *= -1.0;
        }
        break;
    case CameraMode::Landscape:
        std::swap(X, Y);

        if (improvedPosition.cameraRightDirection().z() < 0) {
            // The camera is turned to the right side
            X *= -1.0;
        } else {
            // The camera is turned to the left side
            Y *= -1.0;
        }
        break;
    }

    // clang-format off
    // The camera orientation vectors are of unit length and ortogonal to each
    // over so just make a rotation matrix out of them
    double values[9] = {X.x(), Y.x(), Z.x(),
                        X.y(), Y.y(), Z.y(),
                        X.z(), Y.z(), Z.z()};
    // clang-format on
    cv::Mat orientation{3, 3, CV_64F, values};
    orientation.convertTo(orientation, CV_32F);
    return orientation;
}

std::map<db::TId, cv::Mat> makeNonCarMasksMap(const tws::Photos& photos)
{
    std::map<db::TId, cv::Mat> featureNonCarMaskMap;
    carsegm::CarSegmentator carSegmentator;
    for (const auto& photo : photos) {
        featureNonCarMaskMap[photo.featureId]
            = (carSegmentator.segment(photo.image) == 0);
        // TODO: remove the next two statements
        //(void)carSegmentator;
        //featureNonCarMaskMap[photo.featureId]
        //    = cv::Mat::ones(photo.image.size(), CV_8UC1);
    }
    return featureNonCarMaskMap;
}

PositionedObjects averagePosedObjects(const PositionedObjects& objects) {
    std::map<const ObjectInfo*, PositionedObjects*> objectToClusterMap;
    std::list<PositionedObjects> clusters; // std::list do not reallocate its
                                           // elements which is crucial here.
    for (const auto& object : objects) {
        PositionedObjects* cluster = nullptr;
        if (objectToClusterMap.count(&object.object1)) {
            cluster = objectToClusterMap[&object.object1];
        } else if (objectToClusterMap.count(&object.object2)) {
            cluster = objectToClusterMap[&object.object2];
        }

        if (!cluster) {
            cluster = &clusters.emplace_back();
        }

        cluster->push_back(object);
        objectToClusterMap[&object.object1] = cluster;
        objectToClusterMap[&object.object2] = cluster;
    }

    PositionedObjects averagedObjects;
    for (const auto& cluster : clusters) {
        if (cluster.empty()) {
            WARN() << "Empty objects cluster detected!";
            continue;
        }

        // Compute a geomean position out of positions of an object images.
        double x = 1.0;
        double y = 1.0;
        const double exp = 1.0 / cluster.size();
        for (const auto& object : cluster) {
            // Note: taking a root first accumulates less error due to
            // multiplication.
            x *= std::pow(object.mercatorPos.x(), exp);
            y *= std::pow(object.mercatorPos.y(), exp);
        }
        // Just take the first couple of objects to reference necessarey info
        // later on.
        averagedObjects.emplace_back(PositionedObject{
            cluster.front().object1, cluster.front().object2, {x, y}});
    }

    return averagedObjects;
}

void debugDump(db::TId featureId,
    const cv::Mat& img,
    const ObjectsInfo& objects,
    const CameraPose& imgPose,
    const PositionedObjects& posedObjects)
{
    const std::string filename = std::to_string(featureId) + ".jpeg";
    cv::Mat imgToDraw = img.clone();
    for (const auto& object : objects) {
        cv::rectangle(imgToDraw, object.bbox, {0, 0, 255});
    }

    float zData[3] = {0.0f, 0.0f, 1.0f};
    cv::Mat Z{3, 1, CV_32F, zData};
    cv::Mat imgDir = imgPose.rotation * Z;

    const geolib3::Heading imgHeading = geolib3::Direction2{
        geolib3::Vector2{imgDir.at<float>(0),
            imgDir.at<float>(1)}}.heading();

    std::cout.precision(12);
    cv::imwrite(filename, imgToDraw);
    const auto geoPos
        = geolib3::convertMercatorToGeodetic(imgPose.mercatorPos);
    std::cout << "Saved an image to " << filename << '\n';
    std::cout << "Image position " << geoPos.y() << ", " << geoPos.x()
              << '\n';
    std::cout << "Image heading " << imgHeading << '\n';
    for (const auto& posedObj : posedObjects) {
        const auto objVec = posedObj.mercatorPos - imgPose.mercatorPos;
        const MercatorUnits mercatorDist{
            geolib3::distance(posedObj.mercatorPos, imgPose.mercatorPos)};
        const auto distance = toMeters(imgPose.mercatorPos, mercatorDist);

        std::cout << "==========================\n";
        std::cout << "Object distance " << distance.value() << " m.\n";
        std::cout << "Direction to object "
                  << geolib3::Direction2{objVec}.heading() << '\n';
    }
    std::cout << "==========================\n";
    std::cout << "press any key\n\n";

    char fake;
    std::cin.get(fake);
    (void)fake;
}
}

tracks_with_sensors::Signs locateSignsByTwoPhotos(
    const sfp::SensorsFeaturePositioners& featurePositioners,
    const tws::Photos& photos,
    const tws::SignPhotos& signPhotos,
    std::size_t featuresLimit)
{
    constexpr Meters DEFAULT_OBJECT_RADIUS{0.5f};
    tracks_with_sensors::Signs results;

    if (photos.size() < 2) {
        return results;
    }

    std::map<db::TId, CameraPose> featurePoseMap;
    std::map<db::TId, ObjectsInfo> featureObjectsMap;
    std::map<const ObjectInfo*, const tws::SignPhoto&> objectToSignPhotoMap;
    auto featureNonCarMaskMap = makeNonCarMasksMap(photos);

    auto featurePositionerIt = featurePositioners.begin();
    const auto getImprovedPosition = [&](chrono::TimePoint timestamp)
        -> std::optional<pos_improvment::ImprovedGpsEvent> {
        while (featurePositionerIt != featurePositioners.end()) {
            if (timestamp < featurePositionerIt->trackStartTime()) {
                if (featurePositionerIt == featurePositioners.begin()) {
                    return std::nullopt;
                }
                --featurePositionerIt;
                continue;
            } else if (featurePositionerIt->trackEndTime() < timestamp) {
                ++featurePositionerIt;
                continue;
            }
            return featurePositionerIt->getPositionByTime(timestamp);
        }
        return std::nullopt;
    };

    for (auto& photo : photos) {
        const auto improvedPosition = getImprovedPosition(photo.timestamp);
        if (improvedPosition) {
            auto& camPose = featurePoseMap[photo.featureId];
            camPose.mercatorPos = improvedPosition->mercatorPosition();
            camPose.rotation = inferImageGlobalOrientation(*improvedPosition);
        }
    }

    for (const auto& signPhoto : signPhotos) {
        auto& featureObjects = featureObjectsMap[signPhoto.featureId];
        const ObjectInfo* objectKey
            = &featureObjects.emplace_back(
                ObjectInfo{DEFAULT_OBJECT_RADIUS, signPhoto.bbox});
        // Remember the backward mapping from an object info to its original
        // sign photo
        objectToSignPhotoMap.emplace(objectKey, signPhoto);
    }

    common::ThreadpoolWrapper threadPool{
        std::thread::hardware_concurrency(),
        std::thread::hardware_concurrency() * 2};

    // Position objects which are visible on a pair of photos
    PositionedObjects posedObjects;
    std::mutex posedObjectsMutex;
    const auto end = std::prev(photos.end());
    for (auto photoOneIt = photos.begin(); photoOneIt != end; ++photoOneIt) {
        if (!featurePoseMap.count(photoOneIt->featureId)) {
            continue;
        }

        // FIXME: get an actual camera matrix instead of a default one
        const auto camParams = signs_detector::getCameraParameters(
            "", mrc::common::ImageOrientation{}, photoOneIt->image.size());

        // SfM positioning requires greyscale images
        cv::Mat imgOne;
        cv::cvtColor(photoOneIt->image, imgOne, CV_BGR2GRAY);

        const cv::Mat& maskOne = featureNonCarMaskMap[photoOneIt->featureId];
        const CameraPose& imgOnePose = featurePoseMap[photoOneIt->featureId];
        const ObjectsInfo& objectsOne
            = featureObjectsMap[photoOneIt->featureId];

        for (auto photoTwoIt = std::next(photoOneIt);
             photoTwoIt != photos.end(); ++photoTwoIt) {
            if (!featurePoseMap.count(photoTwoIt->featureId)) {
                continue;
            }

            cv::Mat imgTwo;
            cv::cvtColor(photoTwoIt->image, imgTwo, CV_BGR2GRAY);

            const cv::Mat& maskTwo
                = featureNonCarMaskMap[photoTwoIt->featureId];
            const CameraPose& imgTwoPose
                = featurePoseMap[photoTwoIt->featureId];
            const ObjectsInfo& objectsTwo
                = featureObjectsMap[photoTwoIt->featureId];

            if (!sfm::isCameraMoveInAllowedRange(imgOnePose, imgTwoPose)) {
                continue;
            }

            const auto computePositions =
                [&, photoOneIt, photoTwoIt, imgOne, imgTwo] {
                INFO() << "Trying to position objects on a couple "
                       << photoOneIt->featureId << " -> "
                       << photoTwoIt->featureId;

                const auto objects = computeObjectPositionsWith2Images(
                    camParams.cameraMatrix, imgOne, maskOne, imgOnePose,
                    objectsOne, imgTwo, maskTwo, imgTwoPose, objectsTwo,
                    featuresLimit);
                if (objects.empty()) {
                    return;
                }

                std::lock_guard<std::mutex> lock(posedObjectsMutex);
                for (const auto& object : objects) {
                    posedObjects.push_back(object);
                }
                INFO() << objects.size() << " were positioned on a couple "
                       << photoOneIt->featureId << " -> "
                       << photoTwoIt->featureId;

                if (std::getenv("DEBUG_DUMP")) {
                    debugDump(photoOneIt->featureId, imgOne, objectsOne,
                        imgOnePose, objects);
                }
            };

            if (std::getenv("MRC_PARALLEL_SFM_EVALUATION")) {
                threadPool->add(computePositions);
            } else {
                computePositions();
            }
        }
    }

    if (std::getenv("MRC_PARALLEL_SFM_EVALUATION")) {
        threadPool->join();
    }

    const auto averagedPosedObjects = averagePosedObjects(posedObjects);

    for (const auto& object : averagedPosedObjects) {
        const auto signPhoto = objectToSignPhotoMap.at(&object.object1);
        results.push_back(tws::Sign{object.mercatorPos,
            signPhoto.signType, signPhoto.signGtId});
    }
    return results;
}

} // namespace maps::mrc::experiment_sfm_positioning
