/*
 * 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.Models.GroundTruth;
using CartographerSharp.Models.Mapping;
using CartographerSharp.Models.Transform;
using CartographerSharp.Transform;
using GroundTruthProto = CartographerSharp.Models.GroundTruth.GroundTruth;
using PoseGraphProto = CartographerSharp.Models.Mapping.PoseGraph;

namespace CartographerSharp.GroundTruth;

/// <summary>
/// Generates GroundTruth proto from the given pose graph using the specified
/// criteria parameters. See
/// 'https://google-cartographer.readthedocs.io/en/latest/evaluation.html' for
/// more details.
/// </summary>
public static class AutogenerateGroundTruth
{
    /// <summary>
    /// Generates ground truth from pose graph.
    /// </summary>
    /// <param name="poseGraph">Pose graph proto.</param>
    /// <param name="minCoveredDistance">Minimum covered distance between nodes.</param>
    /// <param name="outlierThresholdMeters">Outlier threshold in meters.</param>
    /// <param name="outlierThresholdRadians">Outlier threshold in radians.</param>
    /// <returns>GroundTruth proto with relations.</returns>
    public static GroundTruthProto GenerateGroundTruth(
        PoseGraph poseGraph,
        double minCoveredDistance,
        double outlierThresholdMeters,
        double outlierThresholdRadians)
    {
        if (poseGraph.Trajectories == null || poseGraph.Trajectories.Count == 0)
        {
            return new GroundTruthProto { Relations = [] };
        }

        var trajectory = poseGraph.Trajectories[0];
        if (trajectory.Nodes == null || trajectory.Nodes.Count == 0)
        {
            return new GroundTruthProto { Relations = [] };
        }

        var coveredDistance = ComputeCoveredDistance(trajectory);
        var submapToNodeIndex = ComputeSubmapRepresentativeNode(poseGraph);

        int numOutliers = 0;
        var groundTruth = new GroundTruthProto
        {
            Relations = []
        };

        if (poseGraph.Constraints == null)
        {
            return groundTruth;
        }

        foreach (var constraint in poseGraph.Constraints)
        {
            // We're only interested in loop closure constraints.
            if (constraint.ConstraintTag == PoseGraphProto.Constraint.Tag.IntraSubmap)
            {
                continue;
            }

            // For some submaps at the very end, we have not chosen a representative
            // node, but those should not be part of loop closure anyway.
            if (constraint.SubmapId.TrajectoryId != 0 ||
                constraint.NodeId.TrajectoryId != 0)
            {
                continue;
            }

            if (constraint.SubmapId.SubmapIndex >= submapToNodeIndex.Count)
            {
                continue;
            }

            var matchedNode = constraint.NodeId.NodeIndex;
            var representativeNode = submapToNodeIndex[constraint.SubmapId.SubmapIndex];

            // Covered distance between the two should not be too small.
            var coveredDistanceInConstraint = Math.Abs(
                coveredDistance[matchedNode] - coveredDistance[representativeNode]);
            if (coveredDistanceInConstraint < minCoveredDistance)
            {
                continue;
            }

            // Compute the transform between the nodes according to the solution and
            // the constraint.
            var solutionPose1 = (Rigid3d)trajectory.Nodes[representativeNode].Pose;
            var solutionPose2 = (Rigid3d)trajectory.Nodes[matchedNode].Pose;
            var solution = solutionPose1.Inverse() * solutionPose2;

            var submapSolution = (Rigid3d)trajectory.Submaps[constraint.SubmapId.SubmapIndex].Pose;
            var submapSolutionToNodeSolution = solutionPose1.Inverse() * submapSolution;
            var nodeToSubmapConstraint = (Rigid3d)constraint.RelativePose;
            var expected = submapSolutionToNodeSolution * nodeToSubmapConstraint;

            var error = solution * expected.Inverse();

            if (error.Translation.Length() > outlierThresholdMeters ||
                TransformOperations.GetAngle(error) > outlierThresholdRadians)
            {
                numOutliers++;
                continue;
            }

            var relation = new Relation
            {
                Timestamp1 = trajectory.Nodes[representativeNode].Timestamp,
                Timestamp2 = trajectory.Nodes[matchedNode].Timestamp,
                Expected = (Rigid3dProto)expected,
                CoveredDistance = coveredDistanceInConstraint
            };

            groundTruth.Relations.Add(relation);
        }

        // Log number of relations and outliers for debugging and analysis
        return groundTruth;
    }

    /// <summary>
    /// Computes covered distance for each node in the trajectory.
    /// </summary>
    private static List<double> ComputeCoveredDistance(Trajectory trajectory)
    {
        var coveredDistance = new List<double> { 0.0 };

        if (trajectory.Nodes == null || trajectory.Nodes.Count == 0)
        {
            return coveredDistance;
        }

        for (int i = 1; i < trajectory.Nodes.Count; i++)
        {
            var lastPose = (Rigid3d)trajectory.Nodes[i - 1].Pose;
            var thisPose = (Rigid3d)trajectory.Nodes[i].Pose;
            var relativeTransform = lastPose.Inverse() * thisPose;
            coveredDistance.Add(coveredDistance[^1] + relativeTransform.Translation.Length());
        }

        return coveredDistance;
    }

    /// <summary>
    /// We pick the representative node in the middle of the submap.
    /// </summary>
    private static List<int> ComputeSubmapRepresentativeNode(PoseGraphProto poseGraph)
    {
        var submapToNodeIndex = new List<int>();

        if (poseGraph.Constraints == null)
        {
            return submapToNodeIndex;
        }

        foreach (var constraint in poseGraph.Constraints)
        {
            if (constraint.ConstraintTag != PoseGraphProto.Constraint.Tag.IntraSubmap)
            {
                continue;
            }

            if (constraint.SubmapId.TrajectoryId != 0 ||
                constraint.NodeId.TrajectoryId != 0)
            {
                continue;
            }

            var nextSubmapIndex = submapToNodeIndex.Count;
            var submapIndex = constraint.SubmapId.SubmapIndex;
            if (submapIndex <= nextSubmapIndex)
            {
                continue;
            }

            if (submapIndex != nextSubmapIndex + 1)
            {
                continue;
            }

            submapToNodeIndex.Add(constraint.NodeId.NodeIndex);
        }

        return submapToNodeIndex;
    }
}