/*
 * 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.Transform;
using CeresSharp;
using  RobotNet10.Shared.Numbers;

namespace CartographerSharp.Mapping.Internal.Optimization;

/// <summary>
/// Sparse Pose Adjustment (SPA) cost function for 2D pose graph optimization.
/// Computes the error between observed relative pose and computed relative pose.
/// </summary>
public class SpaCostFunction2D
{
    private readonly IPoseGraph.Constraint.Pose _observedRelativePose;
    private readonly Rigid2d _observedRelativePose2D;

    /// <summary>
    /// Creates an AutoDiff cost function for SPA.
    /// </summary>
    /// <param name="observedRelativePose">The observed relative pose constraint.</param>
    /// <returns>AutoDiff cost function.</returns>
    public static AutoDiffCostFunction CreateAutoDiffCostFunction(
        IPoseGraph.Constraint.Pose observedRelativePose)
    {
        var costFunction = new SpaCostFunction2D(observedRelativePose);
        return new AutoDiffCostFunction(
            costFunction.Evaluate,
            numResiduals: 3, // [dx, dy, dtheta]
            parameterBlockSizes: [3, 3] // [start_pose[3], end_pose[3]]
        );
    }

    private SpaCostFunction2D(IPoseGraph.Constraint.Pose observedRelativePose)
    {
        _observedRelativePose = observedRelativePose;
        // Project 3D pose to 2D
        _observedRelativePose2D = TransformOperations.Project2D(observedRelativePose.ZbarIj);
    }

    /// <summary>
    /// Evaluates the cost function.
    /// Match C++ spa_cost_function_2d.h operator() implementation.
    /// </summary>
    /// <param name="parameters">Parameter blocks [start_pose[3], end_pose[3]].</param>
    /// <param name="residuals">Output residuals [dx, dy, dtheta].</param>
    /// <returns>True on success.</returns>
    private bool Evaluate(double[][] parameters, double[] residuals)
    {
        if (parameters == null || parameters.Length < 2)
        {
            return false;
        }
        if (parameters[0].Length < 3 || parameters[1].Length < 3)
        {
            return false;
        }
        if (residuals == null || residuals.Length < 3)
        {
            return false;
        }

        var startPose = parameters[0];
        var endPose = parameters[1];
        
        // Validate parameters for NaN/Infinity
        for (int i = 0; i < 3; i++)
        {
            if (double.IsNaN(startPose[i]) || double.IsInfinity(startPose[i]))
            {
                return false;
            }
            if (double.IsNaN(endPose[i]) || double.IsInfinity(endPose[i]))
            {
                return false;
            }
        }
        
        // NOTE: Weight validation removed to match C++ behavior.
        // C++ does not validate weights - Ceres handles invalid weights internally.
        // Validation was causing constraints to be incorrectly rejected.

        // NOTE: Pose explosion handling REMOVED to match C++ behavior.
        // The original C++ spa_cost_function_2d.h does NOT have any pose distance checks.
        // Returning zero residuals was causing optimization to skip constraints incorrectly,
        // leading to optimization failures and incorrect pose graph results.
        // If poses diverge, Ceres will handle it through its own convergence criteria.

        // Compute unscaled error (match C++ cost_helpers_impl.h ComputeUnscaledError)
        var unscaledError = ComputeUnscaledError(
            _observedRelativePose2D,
            startPose,
            endPose
        );

        // Scale error with weights (match C++ ScaleError)
        var translationWeight = _observedRelativePose.TranslationWeight;
        var rotationWeight = _observedRelativePose.RotationWeight;

        var scaledError = ScaleError(
            unscaledError,
            translationWeight,
            rotationWeight
        );

        residuals[0] = scaledError[0];
        residuals[1] = scaledError[1];
        residuals[2] = scaledError[2];

        return true;
    }

    /// <summary>
    /// Computes unscaled error between observed and computed relative pose.
    /// Match C++: Uses direct formula for numerical stability with Ceres autodiff.
    /// </summary>
    private static double[] ComputeUnscaledError(
        Rigid2d observedRelativePose,
        double[] startPose,
        double[] endPose)
    {
        // Match C++ implementation in cost_helpers_impl.h
        // startPose = [x1, y1, theta1]
        // endPose = [x2, y2, theta2]
        // observedRelativePose = relative pose from start to end (in start frame)

        var cosThetaI = Math.Cos(startPose[2]);
        var sinThetaI = Math.Sin(startPose[2]);
        var deltaX = endPose[0] - startPose[0];
        var deltaY = endPose[1] - startPose[1];
        
        // Compute h = relative pose from start to end (in start frame)
        // h[0] = cos_theta_i * delta_x + sin_theta_i * delta_y
        // h[1] = -sin_theta_i * delta_x + cos_theta_i * delta_y
        // h[2] = end[2] - start[2]
        var h0 = cosThetaI * deltaX + sinThetaI * deltaY;
        var h1 = -sinThetaI * deltaX + cosThetaI * deltaY;
        var h2 = endPose[2] - startPose[2];

        // Error = observed - computed
        var translationErrorX = observedRelativePose.Translation.X - h0;
        var translationErrorY = observedRelativePose.Translation.Y - h1;
        
        // Rotation error (normalize angle difference)
        var rotationError = OptimizationHelpers.NormalizeAngleDifference(
            observedRelativePose.Rotation - h2
        );

        return
        [
            translationErrorX,
            translationErrorY,
            rotationError
        ];
    }

    /// <summary>
    /// Scales error with translation and rotation weights.
    /// </summary>
    private static double[] ScaleError(
        double[] unscaledError,
        double translationWeight,
        double rotationWeight)
    {
        return
        [
            translationWeight * unscaledError[0],
            translationWeight * unscaledError[1],
            rotationWeight * unscaledError[2]
        ];
    }

}