// This file is part of the AliceVision project.
// Copyright (c) 2018 AliceVision contributors.
// This Source Code Form is subject to the terms of the Mozilla Public License,
// v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at https://mozilla.org/MPL/2.0/.

#include <aliceVision/system/Logger.hpp>
#include <aliceVision/cmdline/cmdline.hpp>
#include <aliceVision/system/main.hpp>
#include <aliceVision/system/ProgressDisplay.hpp>
#include <aliceVision/system/Timer.hpp>
#include <aliceVision/sfmDataIO/sfmDataIO.hpp>
#include <aliceVision/sfmDataIO/AlembicExporter.hpp>
#include <aliceVision/sfmDataIO/viewIO.hpp>
#include <aliceVision/image/Image.hpp>
#include <aliceVision/image/io.hpp>
#include <aliceVision/camera/cameraUndistortImage.hpp>

#include <aliceVision/utils/filesIO.hpp>
#include <aliceVision/utils/regexFilter.hpp>

#include <boost/program_options.hpp>

#include <filesystem>
#include <cstdlib>
#include <limits>
#include <string>
#include <regex>

// These constants define the current software version.
// They must be updated when the command line is changed.
#define ALICEVISION_SOFTWARE_VERSION_MAJOR 2
#define ALICEVISION_SOFTWARE_VERSION_MINOR 0

using namespace aliceVision;

namespace po = boost::program_options;
namespace fs = std::filesystem;

oiio::ROI computeRod(const camera::IntrinsicBase* intrinsic, bool correctPrincipalPoint)
{
    std::vector<Vec2> pointToBeChecked;
    pointToBeChecked.push_back(Vec2(0, 0));
    pointToBeChecked.push_back(Vec2(intrinsic->w() - 1, 0));
    pointToBeChecked.push_back(Vec2(0, intrinsic->h() - 1));
    pointToBeChecked.push_back(Vec2(intrinsic->w() - 1, intrinsic->h() - 1));
    const Vec2 center(intrinsic->w() * 0.5, intrinsic->h() * 0.5);
    Vec2 ppCorrection(0, 0);
    if (camera::EINTRINSIC::PINHOLE_CAMERA & intrinsic->getType())
    {
        const camera::Pinhole* pinholePtr = dynamic_cast<const camera::Pinhole*>(intrinsic);
        ppCorrection = pinholePtr->getPrincipalPoint() - center;
    }
    const Vec2 opticalCenter = center + ppCorrection;
    pointToBeChecked.push_back(Vec2(opticalCenter[0], 0));
    pointToBeChecked.push_back(Vec2(opticalCenter[0], intrinsic->h() - 1));
    pointToBeChecked.push_back(Vec2(0, opticalCenter[1]));
    pointToBeChecked.push_back(Vec2(intrinsic->w() - 1, opticalCenter[1]));

    std::vector<Vec2> maxDistortionVector;
    for (const Vec2& n : pointToBeChecked)
    {
        // Undistort pixel without principal point correction
        const Vec2 n_undist = intrinsic->getUndistortedPixel(n);
        maxDistortionVector.push_back(n_undist);
    }

    std::sort(std::begin(maxDistortionVector), std::end(maxDistortionVector), [](Vec2 a, Vec2 b) { return a[0] > b[0]; });
    const int xRoiMax = std::round(maxDistortionVector.front()[0]);
    const int xRoiMin = std::round(maxDistortionVector.back()[0]);
    std::sort(std::begin(maxDistortionVector), std::end(maxDistortionVector), [](Vec2 a, Vec2 b) { return a[1] > b[1]; });
    const int yRoiMax = std::round(maxDistortionVector.front()[1]);
    const int yRoiMin = std::round(maxDistortionVector.back()[1]);

    oiio::ROI rod(xRoiMin, xRoiMax + 1, yRoiMin, yRoiMax + 1);

    if (correctPrincipalPoint)
    {
        rod.xbegin -= ppCorrection(0);
        rod.xend -= ppCorrection(0);
        rod.ybegin -= ppCorrection(1);
        rod.yend -= ppCorrection(1);
    }
    return rod;
}

oiio::ROI convertRodToRoi(const camera::IntrinsicBase* intrinsic, const oiio::ROI& rod)
{
    const int xOffset = rod.xbegin;
    const int yOffset = rod.ybegin;  // (intrinsic->h() - rod.yend);
    const oiio::ROI roi(-xOffset, intrinsic->w() - xOffset, -yOffset, intrinsic->h() - yOffset);

    ALICEVISION_LOG_DEBUG("roi:" << roi.xbegin << ";" << roi.xend << ";" << roi.ybegin << ";" << roi.yend);
    return roi;
}

int aliceVision_main(int argc, char** argv)
{
    // Command-line parameters
    std::string sfmDataFilename;
    std::string outFolder;

    // User optional parameters
    bool undistortedImages = false;
    bool exportSTMaps = false;
    bool exportFullROD = false;
    bool correctPrincipalPoint = true;
    std::map<IndexT, oiio::ROI> roiForIntrinsic;
    std::string viewFilter;
    std::string sfmDataFilterFilepath;
    std::string outImageFileTypeName = image::EImageFileType_enumToString(image::EImageFileType::JPEG);
    std::string outMapFileTypeName = image::EImageFileType_enumToString(image::EImageFileType::EXR);
    double frameRate = 24.0;

    // clang-format off
    po::options_description requiredParams("Required parameters");
    requiredParams.add_options()
        ("input,i", po::value<std::string>(&sfmDataFilename)->required(),
         "SfMData file containing a complete SfM.")
        ("output,o", po::value<std::string>(&outFolder)->required(),
         "Output folder.");

    po::options_description optionalParams("Optional parameters");
    optionalParams.add_options()
        ("exportUndistortedImages", po::value<bool>(&undistortedImages)->default_value(undistortedImages),
         "Export undistorted images for the animated camera(s).\n"
         "If false, animated camera(s) exported with original frame paths.")
        ("exportFullROD", po::value<bool>(&exportFullROD)->default_value(exportFullROD),
         "Export undistorted images with the full Region of Definition (RoD). Only supported by the EXR image file format.")
        ("exportSTMaps", po::value<bool>(&exportSTMaps)->default_value(exportSTMaps),
         "Export ST Maps in exr format to apply distort/undistort transformations in a compositing software.")
        ("correctPrincipalPoint", po::value<bool>(&correctPrincipalPoint)->default_value(correctPrincipalPoint),
         "Apply an offset to correct the position of the principal point.")
        ("viewFilter", po::value<std::string>(&viewFilter)->default_value(viewFilter),
         "Select the cameras to export using an expression based on the image filepath. Export all cameras if empty.")
        ("sfmDataFilter", po::value<std::string>(&sfmDataFilterFilepath)->default_value(sfmDataFilterFilepath),
         "Filter out cameras from the export if they are part of this SfMData. Export all cameras if empty.")
        ("undistortedImageType", po::value<std::string>(&outImageFileTypeName)->default_value(outImageFileTypeName),
         image::EImageFileType_informations().c_str())
         ("frameRate", po::value<double>(&frameRate)->default_value(frameRate),
         "Define the camera's Frames per seconds.");
    // clang-format on

    CmdLine cmdline("AliceVision exportAnimatedCamera");
    cmdline.add(requiredParams);
    cmdline.add(optionalParams);
    if (!cmdline.execute(argc, argv))
    {
        return EXIT_FAILURE;
    }

    // Set output file type
    const image::EImageFileType outputFileType = image::EImageFileType_stringToEnum(outImageFileTypeName);
    const image::EImageFileType outputMapFileType = image::EImageFileType_stringToEnum(outMapFileTypeName);

    if (exportFullROD && outputFileType != image::EImageFileType::EXR)
    {
        ALICEVISION_LOG_ERROR("Export full RoD (Region Of Definition) is only possible in EXR file format and not in '" << outputFileType << "'.");
        return EXIT_FAILURE;
    }

    // Load SfMData files
    sfmData::SfMData sfmData;
    if (!sfmDataIO::load(sfmData, sfmDataFilename, sfmDataIO::ESfMData::ALL))
    {
        ALICEVISION_LOG_ERROR("The input SfMData file '" << sfmDataFilename << "' cannot be read.");
        return EXIT_FAILURE;
    }

    if (sfmData.getViews().empty())
    {
        ALICEVISION_LOG_ERROR("The input SfMData file '" << sfmDataFilename << "' is empty.");
        return EXIT_FAILURE;
    }

    sfmData::SfMData sfmDataFilter;
    if (!sfmDataFilterFilepath.empty())
    {
        if (!sfmDataIO::load(sfmDataFilter, sfmDataFilterFilepath, sfmDataIO::ESfMData::VIEWS))
        {
            ALICEVISION_LOG_ERROR("The input filter SfMData file '" << sfmDataFilterFilepath << "' cannot be read.");
            return EXIT_FAILURE;
        }
    }
    system::Timer timer;

    // Decide the views and intrinsics to export
    sfmData::SfMData sfmDataExport;
    for (auto& viewPair : sfmData.getViews())
    {
        sfmData::View& view = *(viewPair.second);

        // Regex filter
        if (!viewFilter.empty())
        {
            // Skip the view if it does not match the expression filter
            const std::regex regexFilter = utils::filterToRegex(viewFilter);
            if (!std::regex_match(view.getImage().getImagePath(), regexFilter))
                continue;
        }

        // SfMData filter
        if (!sfmDataFilterFilepath.empty())
        {
            // Skip the view if it exists in the sfmDataFilter
            if (sfmDataFilter.getViews().find(view.getViewId()) != sfmDataFilter.getViews().end())
                continue;
        }

        sfmDataExport.getViews().emplace(view.getViewId(), viewPair.second);

        // Export intrinsics with at least one view with a valid pose
        if (sfmData.isPoseAndIntrinsicDefined(view))
        {
            // std::map::emplace does nothing if the key already exist
            sfmDataExport.getIntrinsics().emplace(view.getIntrinsicId(), sfmData.getIntrinsics().at(view.getIntrinsicId()));
        }
    }

    const fs::path undistortedImagesFolderPath = fs::path(outFolder) / "undistort";
    const bool writeUndistordedResult = undistortedImages || exportSTMaps;

    if (writeUndistordedResult && !utils::exists(undistortedImagesFolderPath))
        fs::create_directory(undistortedImagesFolderPath);

    std::map<std::string, std::map<std::size_t, IndexT>> videoViewPerFrame;
    std::map<std::string, std::vector<std::pair<std::size_t, IndexT>>> dslrViewPerKey;

    // Export distortion map / one image per intrinsic
    if (exportSTMaps)
    {
        for (const auto& intrinsicPair : sfmDataExport.getIntrinsics())
        {
            const camera::IntrinsicBase& intrinsic = *(intrinsicPair.second);
            image::Image<image::RGBfColor> image_dist;
            // Init image as black (no distortion)
            image_dist.resize(int(intrinsic.w()), int(intrinsic.h()), true, image::FBLACK);

            // Compute ST vertors for distortion
            const Vec2 center(intrinsic.w() * 0.5, intrinsic.h() * 0.5);
            Vec2 ppCorrection(0.0, 0.0);

            if ((camera::EINTRINSIC::PINHOLE_CAMERA & intrinsic.getType()) && correctPrincipalPoint)  // correct principal point
            {
                const camera::Pinhole* pinholePtr = dynamic_cast<const camera::Pinhole*>(intrinsicPair.second.get());
                ppCorrection = pinholePtr->getPrincipalPoint() - center;
            }
            ALICEVISION_LOG_DEBUG("ppCorrection:" + std::to_string(ppCorrection[0]) + ";" + std::to_string(ppCorrection[1]));

            // ST Map: Undistort
            {
// Flip and normalize as STMap
#pragma omp parallel for
                for (int y = 0; y < int(intrinsic.h()); ++y)
                {
                    for (int x = 0; x < int(intrinsic.w()); ++x)
                    {
                        const Vec2 undisto_pix(x, y);
                        // Compute coordinates with distortion
                        const Vec2 disto_pix = intrinsic.getDistortedPixel(undisto_pix) + ppCorrection;

                        image_dist(y, x).r() = float((disto_pix[0]) / (intrinsic.w() - 1));
                        image_dist(y, x).g() = float((intrinsic.h() - 1 - disto_pix[1]) / (intrinsic.h() - 1));
                    }
                }

                const std::string dstImage = (undistortedImagesFolderPath / (std::to_string(intrinsicPair.first) + "_STMap_Undistort." +
                                                                             image::EImageFileType_enumToString(outputMapFileType)))
                                               .string();
                image::writeImage(dstImage, image_dist, image::ImageWriteOptions().storageDataType(image::EStorageDataType::Float));
            }

            // ST Map: Distort
            {
// Flip and normalize as STMap
#pragma omp parallel for
                for (int y = 0; y < int(intrinsic.h()); ++y)
                {
                    for (int x = 0; x < int(intrinsic.w()); ++x)
                    {
                        const Vec2 disto_pix(x, y);
                        // Compute coordinates without distortion
                        const Vec2 undisto_pix = intrinsic.getUndistortedPixel(disto_pix) - ppCorrection;

                        image_dist(y, x).r() = float((undisto_pix[0]) / (intrinsic.w() - 1));
                        image_dist(y, x).g() = float((intrinsic.h() - 1 - undisto_pix[1]) / (intrinsic.h() - 1));
                    }
                }

                const std::string dstImage = (undistortedImagesFolderPath / (std::to_string(intrinsicPair.first) + "_STMap_Distort." +
                                                                             image::EImageFileType_enumToString(outputMapFileType)))
                                               .string();
                image::writeImage(dstImage, image_dist, image::ImageWriteOptions().storageDataType(image::EStorageDataType::Float));
            }
        }
    }

    ALICEVISION_LOG_INFO("Build animated camera(s)...");

    image::Image<image::RGBfColor> image, image_ud;
    auto progressDisplay = system::createConsoleProgressDisplay(sfmDataExport.getViews().size(), std::cout);

    for (const auto& viewPair : sfmDataExport.getViews())
    {
        const sfmData::View& view = *(viewPair.second);

        ++progressDisplay;

        const std::string imagePathStem = fs::path(viewPair.second->getImage().getImagePath()).stem().string();

        // Undistort camera images
        if (undistortedImages)
        {
            sfmData::Intrinsics::const_iterator iterIntrinsic = sfmDataExport.getIntrinsics().find(view.getIntrinsicId());
            const std::string dstImage = (undistortedImagesFolderPath / (std::to_string(view.getIntrinsicId()) + "_" + imagePathStem + "." +
                                                                         image::EImageFileType_enumToString(outputFileType)))
                                           .string();
            const camera::IntrinsicBase* cam = iterIntrinsic->second.get();

            image::readImage(view.getImage().getImagePath(), image, image::EImageColorSpace::LINEAR);
            oiio::ParamValueList metadata = image::readImageMetadata(view.getImage().getImagePath());

            if (cam != nullptr && cam->isValid() && cam->hasDistortion())
            {
                // Undistort the image and save it
                if (exportFullROD)
                {
                    // Build a ROI
                    const IndexT key = view.getIntrinsicId();
                    oiio::ROI rod;
                    const camera::IntrinsicBase& intrinsic = (*cam);
                    if (roiForIntrinsic.find(key) == roiForIntrinsic.end())
                    {
                        rod = computeRod(cam, correctPrincipalPoint);
                        roiForIntrinsic[key] = rod;
                    }
                    else
                    {
                        rod = roiForIntrinsic[key];
                    }

                    ALICEVISION_LOG_DEBUG("rod:" + std::to_string(rod.xbegin) + ";" + std::to_string(rod.xend) + ";" + std::to_string(rod.ybegin) +
                                          ";" + std::to_string(rod.yend));
                    camera::UndistortImage(image, cam, image_ud, image::FBLACK, correctPrincipalPoint, rod);
                    const oiio::ROI roi = convertRodToRoi(cam, rod);
                    writeImage(dstImage, image_ud, image::ImageWriteOptions(), oiio::ParamValueList(), roi);
                }
                else
                {
                    camera::UndistortImage(image, cam, image_ud, image::FBLACK, correctPrincipalPoint);
                    image::writeImage(dstImage, image_ud, image::ImageWriteOptions(), metadata);
                }
            }
            else  // No distortion
            {
                if(cam == nullptr)
                {
                    ALICEVISION_LOG_ERROR("One camera intrinsic is undefined (intrinsic Id: " << view.getIntrinsicId() << ", view Id: " << view.getViewId() << ")");
                }
                else if(!cam->isValid())
                {
                    ALICEVISION_LOG_ERROR("One camera intrinsic is not valid (intrinsic Id: " << view.getIntrinsicId() << ", view Id: " << view.getViewId() << ")");
                }
                // Copy the image since there is no distortion
                image::writeImage(dstImage, image, image::ImageWriteOptions(), metadata);
            }
        }

        // Pose and intrinsic defined
        // Note: we use "sfmData" and not "sfmDataExport" to have access to poses
        if (!sfmData.isPoseAndIntrinsicDefined(view))
            continue;

        std::string cameraName = view.getImage().getMetadataMake() + "_" + view.getImage().getMetadataModel();
        IndexT frameN = 0;
        bool isSequence = false;

        if (view.isPartOfRig())
            cameraName += std::string("_") + std::to_string(view.getSubPoseId());

        {
            std::string prefix;
            std::string suffix;

            if (sfmDataIO::extractNumberFromFileStem(imagePathStem, frameN, prefix, suffix))
            {
                if (prefix.empty() && suffix.empty())
                    cameraName = std::string("Undefined") + "_" + cameraName;
                else
                    cameraName = prefix + "frame" + suffix + "_" + cameraName;

                isSequence = true;
            }
        }

        ALICEVISION_LOG_TRACE("imagePathStem: " << imagePathStem << ", frameN: " << frameN << ", isSequence: " << isSequence
                                                << ", cameraName: " << cameraName);

        if (isSequence)  // Video
        {
            const std::size_t frame = frameN;
            videoViewPerFrame[cameraName][frame] = view.getViewId();
        }
        else if (view.getImage().hasMetadataDateTimeOriginal())  // Picture
        {
            const std::size_t key = view.getImage().getMetadataDateTimestamp();

            dslrViewPerKey[cameraName].push_back({key, view.getViewId()});
        }
        else  // No time or sequence information
        {
            dslrViewPerKey[cameraName].push_back({0, view.getViewId()});
        }
    }

    // Print results
    {
        std::stringstream ss;

        ss << "Camera(s) found:" << std::endl << "\t- # video camera(s): " << videoViewPerFrame.size() << std::endl;

        for (const auto& camera : videoViewPerFrame)
            ss << "\t    - " << camera.first << " | " << camera.second.size() << " frame(s)" << std::endl;

        ss << "\t- # DSLR camera(s): " << dslrViewPerKey.size() << std::endl;

        for (const auto& camera : dslrViewPerKey)
            ss << "\t    - " << camera.first << " | " << camera.second.size() << " image(s)" << std::endl;

        ss << "\t- # Used camera intrinsics: " << sfmDataExport.getIntrinsics().size() << std::endl;

        for (const auto& intrinsicIt : sfmDataExport.getIntrinsics())
        {
            const auto intrinsic = intrinsicIt.second;
            ss << "\t    - " << intrinsicIt.first << " | " << intrinsic->w() << "x" << intrinsic->h() << " " << intrinsic->serialNumber()
               << std::endl;
        }

        ALICEVISION_LOG_INFO(ss.str());
    }

    ALICEVISION_LOG_INFO("Export animated camera(s)...");

    sfmDataIO::AlembicExporter exporter((fs::path(outFolder) / "camera.abc").string());

    for (const auto& cameraViews : videoViewPerFrame)
    {
        const std::map<std::size_t, IndexT>& frameToView = cameraViews.second;
        const std::size_t firstFrame = cameraViews.second.begin()->first;

        exporter.initAnimatedCamera(cameraViews.first, firstFrame, frameRate);

        for (std::size_t frame = firstFrame; frame <= frameToView.rbegin()->first; ++frame)
        {
            const auto findFrameIt = frameToView.find(frame);

            if (findFrameIt != frameToView.end())
            {
                const IndexT viewId = findFrameIt->second;

                // Note: we use "sfmData" and not "sfmDataExport" to have access to poses

                const auto findViewIt = sfmData.getViews().find(viewId);
                assert(findViewIt != sfmData.getViews().end());

                ALICEVISION_LOG_DEBUG("[" + cameraViews.first + "][video] Keyframe added");
                const IndexT intrinsicId = findViewIt->second->getIntrinsicId();
                const camera::Pinhole* cam = dynamic_cast<camera::Pinhole*>(sfmData.getIntrinsicPtr(intrinsicId));
                const sfmData::CameraPose pose = sfmData.getPose(*findViewIt->second);
                const std::string& imagePath = findViewIt->second->getImage().getImagePath();
                const std::string undistortedImagePath =
                  (undistortedImagesFolderPath / (std::to_string(intrinsicId) + "_" + fs::path(imagePath).stem().string() + "." +
                                                  image::EImageFileType_enumToString(outputFileType)))
                    .string();

                exporter.addCameraKeyframe(pose.getTransform(), cam, (undistortedImages) ? undistortedImagePath : imagePath, viewId, intrinsicId);
            }
            else
            {
                exporter.jumpKeyframe(std::to_string(frame));
            }
        }
    }

    for (auto& cameraViews : dslrViewPerKey)
    {
        exporter.initAnimatedCamera(cameraViews.first);
        std::sort(cameraViews.second.begin(), cameraViews.second.end());

        for (const auto& cameraView : cameraViews.second)
        {
            ALICEVISION_LOG_DEBUG("[" + cameraViews.first + "][dslr] Keyframe added");
            const sfmData::View& view = *(sfmData.getViews().at(cameraView.second));
            const camera::Pinhole* cam = dynamic_cast<camera::Pinhole*>(sfmData.getIntrinsicPtr(view.getIntrinsicId()));
            const sfmData::CameraPose pose = sfmData.getPose(view);
            const std::string& imagePath = view.getImage().getImagePath();
            const std::string undistortedImagePath =
              (undistortedImagesFolderPath / (std::to_string(view.getIntrinsicId()) + "_" + fs::path(imagePath).stem().string() + "." +
                                              image::EImageFileType_enumToString(outputFileType)))
                .string();

            exporter.addCameraKeyframe(
              pose.getTransform(), cam, (undistortedImages) ? undistortedImagePath : imagePath, view.getViewId(), view.getIntrinsicId());
        }
    }

    ALICEVISION_LOG_INFO("Task done in (s): " + std::to_string(timer.elapsed()));
    return EXIT_SUCCESS;
}
