/*
 * 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 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 of embedded in 3D
/// space node poses.
/// </summary>
public class LandmarkCostFunction2D
{
    private readonly IPoseGraph.LandmarkNode.LandmarkObservation _observation;
    private readonly NodeSpec2D _prevNode;
    private readonly NodeSpec2D _nextNode;
    private readonly double _interpolationParameter;

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

    private LandmarkCostFunction2D(
        IPoseGraph.LandmarkNode.LandmarkObservation observation,
        NodeSpec2D prevNode,
        NodeSpec2D 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 < 4)
            return false;
        if (parameters[0].Length < 3 || parameters[1].Length < 3 ||
            parameters[2].Length < 4 || parameters[3].Length < 3)
            return false;
        if (residuals == null || residuals.Length < 6)
            return false;

        var prevNodePose = parameters[0]; // [x, y, theta]
        var nextNodePose = parameters[1]; // [x, y, theta]
        var landmarkRotation = parameters[2]; // [w, x, y, z]
        var landmarkTranslation = parameters[3]; // [x, y, z]

        // Interpolate node poses
        var (interpolatedRotation, interpolatedTranslation) = CostHelpers.InterpolateNodes2D(
            prevNodePose,
            _prevNode.GravityAlignment,
            nextNodePose,
            _nextNode.GravityAlignment,
            _interpolationParameter
        );

        // Landmark pose parameters
        var landmarkRotationQuat = OptimizationHelpers.ParametersToQuaternion(landmarkRotation);
        var landmarkTranslationVec = OptimizationHelpers.ParametersToVector3(landmarkTranslation);

        // The landmark cost function computes error between:
        // - observed: landmark_to_tracking_transform (from observation)
        // - computed: (interpolated_tracking_pose^-1 * landmark_pose)
        // Error = observed - computed
        // This is equivalent to: landmark_to_tracking_transform - (interpolated_pose^-1 * landmark_pose)
        var unscaledError = CostHelpers.ComputeUnscaledError3D(
            _observation.LandmarkToTrackingTransform,
            interpolatedRotation,
            interpolatedTranslation,
            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;
    }
}