/*
 * Copyright 2018 The Cartographer Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

using CartographerSharp.Mapping.Internal.D3D.Optimization;
using CeresSharp;

namespace CartographerSharp.Mapping.Internal.Optimization;

/// <summary>
/// Cost function measuring the weighted error between the observed pose given by
/// the landmark measurement and the linearly interpolated pose.
/// </summary>
public class LandmarkCostFunction3D
{
    private readonly IPoseGraph.LandmarkNode.LandmarkObservation _observation;
    private readonly NodeSpec3D _prevNode;
    private readonly NodeSpec3D _nextNode;
    private readonly double _interpolationParameter;

    /// <summary>
    /// Creates an AutoDiff cost function for landmark constraints in 3D.
    /// </summary>
    public static AutoDiffCostFunction CreateAutoDiffCostFunction(
        IPoseGraph.LandmarkNode.LandmarkObservation observation,
        NodeSpec3D prevNode,
        NodeSpec3D nextNode)
    {
        var costFunction = new LandmarkCostFunction3D(observation, prevNode, nextNode);
        return new AutoDiffCostFunction(
            costFunction.Evaluate,
            numResiduals: 6, // [dx, dy, dz, dqx, dqy, dqz]
            parameterBlockSizes: [4, 3, 4, 3, 4, 3] // [prev_rotation[4], prev_translation[3], next_rotation[4], next_translation[3], landmark_rotation[4], landmark_translation[3]]
        );
    }

    private LandmarkCostFunction3D(
        IPoseGraph.LandmarkNode.LandmarkObservation observation,
        NodeSpec3D prevNode,
        NodeSpec3D nextNode)
    {
        _observation = observation;
        _prevNode = prevNode;
        _nextNode = nextNode;

        // Compute interpolation parameter
        _interpolationParameter = OptimizationHelpers.ComputeInterpolationParameter(
            _observation.Time,
            _prevNode.Time,
            _nextNode.Time
        );
    }

    /// <summary>
    /// Evaluates the cost function.
    /// </summary>
    private bool Evaluate(double[][] parameters, double[] residuals)
    {
        if (parameters == null || parameters.Length < 6)
            return false;
        if (parameters[0].Length < 4 || parameters[1].Length < 3 ||
            parameters[2].Length < 4 || parameters[3].Length < 3 ||
            parameters[4].Length < 4 || parameters[5].Length < 3)
            return false;
        if (residuals == null || residuals.Length < 6)
            return false;

        var prevNodeRotation = parameters[0]; // [w, x, y, z]
        var prevNodeTranslation = parameters[1]; // [x, y, z]
        var nextNodeRotation = parameters[2]; // [w, x, y, z]
        var nextNodeTranslation = parameters[3]; // [x, y, z]
        var landmarkRotation = parameters[4]; // [w, x, y, z]
        var landmarkTranslation = parameters[5]; // [x, y, z]

        // Interpolate node poses
        var (interpolatedRotationQuat, interpolatedTranslationVec) = CostHelpers.InterpolateNodes3D(
            prevNodeRotation,
            prevNodeTranslation,
            nextNodeRotation,
            nextNodeTranslation,
            _interpolationParameter
        );

        var landmarkRotationQuat = OptimizationHelpers.ParametersToQuaternion(landmarkRotation);
        var landmarkTranslationVec = OptimizationHelpers.ParametersToVector3(landmarkTranslation);

        // Compute error
        var unscaledError = CostHelpers.ComputeUnscaledError3D(
            _observation.LandmarkToTrackingTransform,
            interpolatedRotationQuat,
            interpolatedTranslationVec,
            landmarkRotationQuat,
            landmarkTranslationVec
        );

        // Scale error
        var scaledError = CostHelpers.ScaleError3D(
            unscaledError,
            _observation.TranslationWeight,
            _observation.RotationWeight
        );

        for (int i = 0; i < 6; i++)
        {
            residuals[i] = scaledError[i];
        }

        return true;
    }
}