// This file is part of the AliceVision project.
// Copyright (c) 2021 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/sfmData/SfMData.hpp>
#include <aliceVision/sfmDataIO/sfmDataIO.hpp>
#include <aliceVision/system/Logger.hpp>
#include <aliceVision/cmdline/cmdline.hpp>
#include <aliceVision/system/main.hpp>
#include <aliceVision/config.hpp>
#include <aliceVision/utils/regexFilter.hpp>

#include <boost/program_options.hpp>
#include <boost/system/error_code.hpp>

#include <algorithm>
#include <filesystem>
#include <string>
#include <regex>
#include <iostream>
#include <fstream>
#include <list>

#include <boost/property_tree/ptree.hpp>
#include <boost/property_tree/xml_parser.hpp>
#include <boost/property_tree/json_parser.hpp>

// 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 json = boost::property_tree;
namespace fs = std::filesystem;

struct XMPData
{
    std::vector<double> rotation;
    std::vector<double> position;
    std::vector<double> distortionCoefficients;
    std::string distortionModel;
    double focalLength35mm = 0.0;
    int skew = 0;
    double aspectRatio = 1.0;
    double principalPointU = 0.0;
    double principalPointV = 0.0;
    std::string calibrationPrior = "DEFAULT";
    int calibrationGroup = 0;
    int distortionGroup = 0;
    int inTexturing = 0;
    int inMeshing = 0;
};

XMPData read_xmp(const std::string& xmpFilepath, const std::string& knownPosesFilePath, const std::string& stem, fs::directory_entry pathIt)
{
    XMPData xmp;
    const fs::path path = pathIt.path();
    if (!is_regular_file(path))
        ALICEVISION_THROW_ERROR("Path isn't a regulat file: " << path);
    std::string extension = path.extension().string();
    boost::to_lower(extension);
    if (extension != ".xmp")
    {
        ALICEVISION_THROW_ERROR("Unknown extension: " << extension);
    }

    json::ptree tree;
    read_xml(knownPosesFilePath + '/' + stem + ".xmp", tree);

    xmp.distortionModel = tree.get<std::string>("x:xmpmeta.rdf:RDF.rdf:Description.<xmlattr>.xcr:DistortionModel", "DEFAULT");
    xmp.focalLength35mm = tree.get<double>("x:xmpmeta.rdf:RDF.rdf:Description.<xmlattr>.xcr:FocalLength35mm", 0.0);
    xmp.skew = tree.get<int>("x:xmpmeta.rdf:RDF.rdf:Description.<xmlattr>.xcr:Skew", 0);
    xmp.aspectRatio = tree.get<double>("x:xmpmeta.rdf:RDF.rdf:Description.<xmlattr>.xcr:AspectRatio", 1.0);
    xmp.principalPointU = tree.get<float>("x:xmpmeta.rdf:RDF.rdf:Description.<xmlattr>.xcr:PrincipalPointU", 0);
    xmp.principalPointV = tree.get<float>("x:xmpmeta.rdf:RDF.rdf:Description.<xmlattr>.xcr:PrincipalPointV", 0);
    xmp.calibrationPrior = tree.get<std::string>("x:xmpmeta.rdf:RDF.rdf:Description.<xmlattr>.xcr:CalibrationPrior", "DEFAULT");
    xmp.calibrationGroup = tree.get<int>("x:xmpmeta.rdf:RDF.rdf:Description.<xmlattr>.xcr:CalibrationGroup", 0);
    xmp.distortionGroup = tree.get<int>("x:xmpmeta.rdf:RDF.rdf:Description.<xmlattr>.xcr:DistortionGroup", 0);
    xmp.inTexturing = tree.get<int>("x:xmpmeta.rdf:RDF.rdf:Description.<xmlattr>.xcr:InTexturing", 0);
    xmp.inMeshing = tree.get<int>("x:xmpmeta.rdf:RDF.rdf:Description.<xmlattr>.xcr:InMeshing", 0);

    std::string rotationStr = tree.get<std::string>("x:xmpmeta.rdf:RDF.rdf:Description.xcr:Rotation", "");
    std::vector<std::string> rotationStrings;
    boost::split(rotationStrings, rotationStr, boost::is_any_of(" \t"), boost::token_compress_on);
    ALICEVISION_LOG_TRACE("stem: " << stem);
    ALICEVISION_LOG_TRACE("rotation: " << rotationStrings);

    for (const auto& rot_val : rotationStrings)
    {
        xmp.rotation.push_back(std::stod(rot_val));
    }

    std::string positionStr = tree.get<std::string>("x:xmpmeta.rdf:RDF.rdf:Description.xcr:Position", "");
    std::vector<std::string> positionStrings;
    if (!positionStr.empty())
    {
        boost::split(positionStrings, positionStr, boost::is_any_of(" \t"), boost::token_compress_on);
        ALICEVISION_LOG_TRACE("position: " << positionStrings);
    }
    else
    {
        positionStr = tree.get<std::string>("x:xmpmeta.rdf:RDF.rdf:Description.<xmlattr>.xcr:Position", "");
        boost::split(positionStrings, positionStr, boost::is_any_of(" \t"), boost::token_compress_on);
    }

    for (const auto& pos_val : positionStrings)
    {
        xmp.position.push_back(std::stod(pos_val));
    }
    std::string distortionStr = tree.get<std::string>("x:xmpmeta.rdf:RDF.rdf:Description.xcr:DistortionCoeficients", "");
    std::vector<std::string> distortionStrings;
    boost::split(distortionStrings, distortionStr, boost::is_any_of(" \t"), boost::token_compress_on);
    ALICEVISION_LOG_TRACE("distortion: " << distortionStrings);

    for (const auto& disto_val : distortionStrings)
    {
        xmp.distortionCoefficients.push_back(std::stod(disto_val));
    }
    return xmp;
}

// import from a SfMData format to another
int aliceVision_main(int argc, char** argv)
{
    // command-line parameters
    std::string knownPosesFilePath;
    std::string sfmDataFilePath;
    std::string outputFilename;

    sfmData::SfMData sfmData;

    // clang-format off
    po::options_description requiredParams("Required parameters");
    requiredParams.add_options()
        ("knownPosesData", po::value<std::string>(&knownPosesFilePath)->required(),
         "Input path to a json file or a folder containing an XMP file per image.")
        ("sfmData", po::value<std::string>(&sfmDataFilePath)->required(),
         "SfmData filepath.")
        ("output,o", po::value<std::string>(&outputFilename)->required(),
         "Output sfmData filepath.");
    // clang-format on

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

    // Loading the sfmData to modify it
    if (!sfmDataIO::load(sfmData, sfmDataFilePath, sfmDataIO::ESfMData::ALL))
    {
        ALICEVISION_LOG_ERROR("The input SfMData file '" << sfmDataFilePath << "' cannot be read.");
        return EXIT_FAILURE;
    }
    if (sfmData.getViews().empty())
    {
        ALICEVISION_LOG_ERROR("The input SfMData file '" << sfmDataFilePath << "' is empty.");
        return EXIT_FAILURE;
    }

    std::map<std::string, IndexT> viewIdPerStem;
    for (const auto viewIt : sfmData.getViews())
    {
        const std::string stem = fs::path(viewIt.second->getImage().getImagePath()).stem().string();
        viewIdPerStem[stem] = viewIt.first;
    }
    fs::path knownPosesPath(knownPosesFilePath);
    if (fs::is_directory(knownPosesPath))
    {
        try
        {
            for (const auto& pathIt : fs::directory_iterator(knownPosesPath))
            {
                const std::string stem = pathIt.path().stem().string();
                if (viewIdPerStem.count(stem) == 0)
                {
                    continue;
                }

                const XMPData xmp = read_xmp(pathIt.path().string(), knownPosesFilePath, stem, pathIt);

                const IndexT viewId = viewIdPerStem[stem];
                sfmData::View& view = sfmData.getView(viewId);
                sfmData::CameraPose& pose = sfmData.getPoses()[view.getPoseId()];

                std::shared_ptr<camera::IntrinsicBase> intrinsicBase = sfmData.getIntrinsicSharedPtr(view.getIntrinsicId());

                Eigen::Map<const Eigen::Matrix<double, 3, 3, Eigen::RowMajor>> rot(xmp.rotation.data());

                Vec3 pos_vec(xmp.position[0], xmp.position[1], xmp.position[2]);
                Vec3 translation = -rot * pos_vec;

                Eigen::Matrix4d T;
                T.setIdentity();
                T.block<3, 3>(0, 0) = rot;
                T.block<3, 1>(0, 3) = translation;

                Eigen::Matrix4d av_T_cr = Eigen::Matrix4d::Zero();
                av_T_cr(0, 0) = 1.0;
                av_T_cr(1, 2) = -1.0;
                av_T_cr(2, 1) = 1.0;
                av_T_cr(3, 3) = 1.0;

                T = T * av_T_cr;
                Eigen::Matrix3d R = T.block<3, 3>(0, 0);
                translation = T.block<3, 1>(0, 3);
                pos_vec = -R.transpose() * translation;

                geometry::Pose3 pos3(R, pos_vec);
                pose.setTransform(pos3);

                std::shared_ptr<camera::IntrinsicScaleOffsetDisto> intrinsic =
                  std::dynamic_pointer_cast<camera::IntrinsicScaleOffsetDisto>(intrinsicBase);
                if (intrinsic == nullptr)
                {
                    ALICEVISION_THROW_ERROR("Invalid intrinsic");
                }

                const double imageRatio = static_cast<double>(view.getImage().getWidth()) / static_cast<double>(view.getImage().getHeight());
                const double sensorWidth = intrinsic->sensorWidth();
                const double maxSize = std::max(view.getImage().getWidth(), view.getImage().getHeight());
                const double focalLengthmm = (sensorWidth * xmp.focalLength35mm) / 36.0;
                const double focalLengthPix = maxSize * focalLengthmm / sensorWidth;
                const double offsetX = (double(view.getImage().getWidth()) * 0.5) + (xmp.principalPointU * maxSize);
                const double offsetY = (double(view.getImage().getHeight()) * 0.5) + (xmp.principalPointV * maxSize);

                intrinsic->setScale({focalLengthPix, focalLengthPix});
                intrinsic->setOffset({offsetX, offsetY});

                std::cout << focalLengthPix << std::endl;

                if (xmp.distortionModel == "brown3t2")
                {
                    std::shared_ptr<camera::Pinhole> camera = std::dynamic_pointer_cast<camera::Pinhole>(intrinsic);
                    if (camera == nullptr)
                    {
                        camera = camera::createPinhole(camera::EDISTORTION::DISTORTION_NONE, camera::EUNDISTORTION::UNDISTORTION_NONE);
                        camera->copyFrom(*intrinsic);
                        sfmData.getIntrinsics().at(view.getIntrinsicId()) = camera;
                    }

                    if (xmp.distortionCoefficients.size() == 6)
                    {
                        std::vector<double> distortionCoefficients;

                        distortionCoefficients.push_back(xmp.distortionCoefficients[0]);
                        distortionCoefficients.push_back(xmp.distortionCoefficients[1]);
                        distortionCoefficients.push_back(xmp.distortionCoefficients[2]);
                        // Skip element at index 3 as it is empty
                        distortionCoefficients.push_back(xmp.distortionCoefficients[5]);
                        distortionCoefficients.push_back(xmp.distortionCoefficients[4]);
                        camera->setDistortionParams(distortionCoefficients);  // vector of 5 elements (r1, r2, r3, t1, t2)
                    }
                    else
                    {
                        ALICEVISION_THROW_ERROR("Error in xmp file: " << stem << " the distortion coefficient doesn't have the right size.");
                    }
                }
                else if (xmp.distortionModel == "brown3")
                {
                    std::shared_ptr<camera::Pinhole> camera = std::dynamic_pointer_cast<camera::Pinhole>(intrinsic);
                    if (camera == nullptr)
                    {
                        camera = camera::createPinhole(camera::EDISTORTION::DISTORTION_NONE, camera::EUNDISTORTION::UNDISTORTION_NONE);
                        camera->copyFrom(*intrinsic);
                        sfmData.getIntrinsics().at(view.getIntrinsicId()) = camera;
                    }

                    if (xmp.distortionCoefficients.size() == 3)
                    {
                        std::vector<double> distortionCoefficients = xmp.distortionCoefficients;
                        camera->setDistortionParams(distortionCoefficients);  // vector of 5 elements (r1, r2, r3)
                    }
                    else
                    {
                        ALICEVISION_THROW_ERROR("Error in xmp file: " << stem << " the distortion coefficient doesn't have the right size.");
                    }
                }
                else
                {
                    ALICEVISION_THROW_ERROR("Unsupported distortion model: " << xmp.distortionModel);
                }
            }
        }
        catch (boost::program_options::error& e)
        {
            ALICEVISION_CERR("ERROR: " << e.what() << std::endl);
            return EXIT_FAILURE;
        }
    }
    else if (is_regular_file(knownPosesPath))
    {
        std::string extension = knownPosesPath.extension().string();
        boost::to_lower(extension);
        if (extension == ".json")
        {
            std::ifstream jsonFile(knownPosesFilePath);
            if (!jsonFile)
            {
                ALICEVISION_LOG_ERROR("Error opening file: " << knownPosesFilePath);
                return EXIT_FAILURE;
            }

            std::string line;
            size_t count = 0;
            std::vector<std::pair<size_t, int>> records;
            std::vector<std::pair<IndexT, IndexT>> frameIdToPoseId;

            // Here we are making a vector that associate a frameId to a PoseId so we can access each easier
            for (const auto& view : sfmData.getViews())
            {
                frameIdToPoseId.emplace_back(view.second->getFrameId(), view.second->getPoseId());
            }
            std::sort(frameIdToPoseId.begin(), frameIdToPoseId.end());

            // ensure there is no duplicated frameId
            auto it =
              std::adjacent_find(frameIdToPoseId.begin(), frameIdToPoseId.end(), [](const auto& a, const auto& b) { return a.first == b.first; });
            if (it != frameIdToPoseId.end())
            {
                ALICEVISION_THROW_ERROR("Duplicated frameId in sfmData: " << sfmDataFilePath << ", frameID: " << it->first);
            }
            // This is where we start to read our json line by line
            while (getline(jsonFile, line))
            {
                std::stringstream linestream(line);
                std::vector<double> up;
                std::vector<double> forward;
                std::vector<double> pose;
                json::ptree pt;

                // We put each line in a stringstream because that is what boost's parser needs.
                // The parser turns the json into a property tree.
                json::json_parser::read_json(linestream, pt);
                const int sensor = pt.get<int>("sensorwidth", 0);
                const double fov = pt.get<double>("xFovDegrees", 0);
                const long timestamp = pt.get<long>("tstamp", 0);
                const int frmcnt = pt.get<int>("frmcnt", 0);
                const double expoff = pt.get<double>("exposureOff", 0);
                const double expdur = pt.get<double>("exposureDur", 0);
                // These arguments are lists so we need to loop to store them properly
                for (json::ptree::value_type& up_val : pt.get_child("up"))
                {
                    std::string value = up_val.second.data();
                    up.push_back(std::stof(value));
                }
                for (json::ptree::value_type& for_val : pt.get_child("forward"))
                {
                    std::string value = for_val.second.data();
                    forward.push_back(std::stof(value));
                }
                for (json::ptree::value_type& pose_val : pt.get_child("pose"))
                {
                    std::string value = pose_val.second.data();
                    pose.push_back(std::stof(value));
                }
                // We use records to indexify our frame count this way we know which frame started the list, this will be our offset
                records.emplace_back(count, frmcnt);

                // Without surprise we store our vector pose to get our position
                Vec3 pos_vec(pose[0], pose[1], pose[2]);
                Mat3 rot;
                // And we need those two vectors to calculate the rotation matrix
                Vec3 up_vec(up[0], up[1], up[2]);
                Vec3 forward_vec(forward[0], forward[1], forward[2]);

                rot.row(0) = up_vec.cross(forward_vec);
                rot.row(1) = -up_vec;
                rot.row(2) = forward_vec;
                // we store this new information into a pose3
                geometry::Pose3 pos3(rot, pos_vec);

                // And we set this pose into the sfmData using our frameId (which corresponds to the count) to set a new pos3 transform
                IndexT sfmPoseId = frameIdToPoseId[count].second;
                sfmData.getPoses()[sfmPoseId].setTransform(pos3);
                count++;
                // We repeat this to each line of the file which contains a json
            }
        }
        else if (extension == ".ma")
        {
            std::ifstream file(knownPosesPath.string());

            std::string line;
            std::string name;
            bool hasName = false;
            bool hasPosition = false;
            bool hasRotation = false;
            double pos[3] = {0.0, 0.0, 0.0};
            double rot[3] = {0.0, 0.0, 0.0};

            while (std::getline(file, line))
            {
                std::regex regex("[^\\s\\t;]+");
                std::vector<std::string> words;

                for (auto it = std::sregex_iterator(line.begin(), line.end(), regex); it != std::sregex_iterator(); it++)
                {
                    std::string tok = it->str();
                    tok.erase(std::remove(tok.begin(), tok.end(), '\"'), tok.end());
                    words.push_back(tok);
                }

                if (words.empty())
                    continue;

                if (words[0] == "createNode")
                {
                    if (words.size() == 4)
                    {
                        name = words[3];
                        hasName = true;
                        hasPosition = false;
                        hasRotation = false;
                    }
                }

                if (words[0] == "setAttr")
                {
                    if (words[1] == ".translate")
                    {
                        if (hasName && (!hasPosition))
                        {
                            hasPosition = true;
                            pos[0] = std::stod(words[4]);
                            pos[1] = std::stod(words[5]);
                            pos[2] = std::stod(words[6]);
                        }
                    }

                    if (words[1] == ".rotate")
                    {
                        if (hasName && (!hasRotation))
                        {
                            hasRotation = true;
                            rot[0] = std::stod(words[4]);
                            rot[1] = std::stod(words[5]);
                            rot[2] = std::stod(words[6]);
                        }
                    }
                }

                if (hasName && hasRotation && hasPosition)
                {
                    if (viewIdPerStem.count(name) == 0)
                    {
                        continue;
                    }

                    const IndexT viewId = viewIdPerStem[name];
                    sfmData::View& view = sfmData.getView(viewId);
                    sfmData::CameraPose& pose = sfmData.getPoses()[view.getPoseId()];

                    Eigen::Matrix3d R = Eigen::Matrix3d::Identity();
                    const Eigen::AngleAxis<double> MX(degreeToRadian(rot[0]), Eigen::Vector3d::UnitX());
                    const Eigen::AngleAxis<double> MY(degreeToRadian(rot[1]), Eigen::Vector3d::UnitY());
                    const Eigen::AngleAxis<double> MZ(degreeToRadian(rot[2]), Eigen::Vector3d::UnitZ());
                    R = MZ * MY * MX;

                    Eigen::Vector3d position;
                    position(0) = pos[0];
                    position(1) = pos[1];
                    position(2) = pos[2];

                    Vec3 translation = -R * position;

                    Eigen::Matrix3d alice_R_maya = Eigen::Matrix3d::Identity();

                    alice_R_maya(0, 0) = 1.0;
                    alice_R_maya(1, 1) = -1.0;
                    alice_R_maya(2, 2) = -1.0;
                    position = position;

                    R = R * alice_R_maya;

                    geometry::Pose3 pose3(R.transpose(), position);
                    pose.setTransform(pose3);
                    ALICEVISION_LOG_TRACE("Read maya: " << name);

                    hasName = false;
                    hasRotation = false;
                    hasPosition = false;
                }
            }
        }
    }

    // export the SfMData scene in the expected format
    if (!sfmDataIO::save(sfmData, outputFilename, sfmDataIO::ESfMData::ALL))
    {
        ALICEVISION_LOG_ERROR("An error occurred while trying to save '" << outputFilename << "'");
        return EXIT_FAILURE;
    }
    return EXIT_SUCCESS;
}
