/*
 * Copyright 2016 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.Common.Math;
using CartographerSharp.Transform;
using System.Text.Json;
using GroundTruthProto = CartographerSharp.Models.GroundTruth.GroundTruth;
using PoseGraphProto = CartographerSharp.Models.Mapping.PoseGraph;

namespace CartographerSharp.GroundTruth;

/// <summary>
/// Error structure for computing metrics.
/// </summary>
internal struct Error
{
    public double TranslationalSquared { get; set; }
    public double RotationalSquared { get; set; }
}

/// <summary>
/// Computes relations metrics from pose graph and ground truth.
/// </summary>
public static class ComputeRelationsMetrics
{
    private static readonly JsonSerializerOptions JsonOptions = new()
    {
        PropertyNameCaseInsensitive = true
    };

    /// <summary>
    /// Computes error between two poses and expected transform.
    /// </summary>
    private static Error ComputeError(
        Rigid3d pose1,
        Rigid3d pose2,
        Rigid3d expected)
    {
        var error = (pose1.Inverse() * pose2) * expected.Inverse();
        return new Error
        {
            TranslationalSquared = error.Translation.Length() * error.Translation.Length(),
            RotationalSquared = MathUtils.Pow2(TransformOperations.GetAngle(error))
        };
    }

    /// <summary>
    /// Computes mean and standard deviation string from values.
    /// </summary>
    private static string MeanAndStdDevString(List<double> values)
    {
        if (values.Count < 2)
            return "N/A";

        var mean = values.Average();
        var sumOfSquaredDifferences = values.Sum(v => MathUtils.Pow2(v - mean));
        var standardDeviation = Math.Sqrt(sumOfSquaredDifferences / (values.Count - 1));
        return $"{mean:F5} +/- {standardDeviation:F5}";
    }

    /// <summary>
    /// Computes statistics string from errors.
    /// </summary>
    private static string StatisticsString(List<Error> errors)
    {
        var translationalErrors = errors.Select(e => Math.Sqrt(e.TranslationalSquared)).ToList();
        var squaredTranslationalErrors = errors.Select(e => e.TranslationalSquared).ToList();
        var rotationalErrorsDegrees = errors.Select(e => MathUtils.RadToDeg(Math.Sqrt(e.RotationalSquared))).ToList();
        var squaredRotationalErrorsDegrees = errors.Select(e => MathUtils.Pow2(MathUtils.RadToDeg(Math.Sqrt(e.RotationalSquared)))).ToList();

        return $"Translational error (m): {MeanAndStdDevString(translationalErrors)}\n" +
               $"Translational error squared (m²): {MeanAndStdDevString(squaredTranslationalErrors)}\n" +
               $"Rotational error (deg): {MeanAndStdDevString(rotationalErrorsDegrees)}\n" +
               $"Rotational error squared (deg²): {MeanAndStdDevString(squaredRotationalErrorsDegrees)}";
    }

    /// <summary>
    /// Computes relations metrics from pose graph and ground truth.
    /// </summary>
    /// <param name="poseGraph">Pose graph proto.</param>
    /// <param name="groundTruth">Ground truth relations.</param>
    /// <returns>Metrics string.</returns>
    public static string ComputeMetrics(
        PoseGraphProto poseGraph,
        GroundTruthProto groundTruth)
    {
        if (poseGraph.Trajectories == null || poseGraph.Trajectories.Count == 0)
        {
            return "No trajectories in pose graph.";
        }

        var trajectory = poseGraph.Trajectories[0];
        if (trajectory.Nodes == null || trajectory.Nodes.Count == 0)
        {
            return "No nodes in trajectory.";
        }

        var errors = new List<Error>();

        if (groundTruth.Relations == null)
        {
            return "No relations in ground truth.";
        }

        foreach (var relation in groundTruth.Relations)
        {
            // Find nodes with matching timestamps
            var node1 = trajectory.Nodes.FirstOrDefault(n => n.Timestamp == relation.Timestamp1);
            var node2 = trajectory.Nodes.FirstOrDefault(n => n.Timestamp == relation.Timestamp2);

            // Check if nodes were found (Node is a struct, so FirstOrDefault returns default(Node) if not found)
            // We check if any node with matching timestamp exists
            var found1 = trajectory.Nodes.Any(n => n.Timestamp == relation.Timestamp1);
            var found2 = trajectory.Nodes.Any(n => n.Timestamp == relation.Timestamp2);
            if (!found1 || !found2)
            {
                continue; // Nodes not found
            }

            var pose1 = (Rigid3d)node1.Pose;
            var pose2 = (Rigid3d)node2.Pose;
            var expected = (Rigid3d)relation.Expected;

            var error = ComputeError(pose1, pose2, expected);
            errors.Add(error);
        }

        if (errors.Count == 0)
        {
            return "No matching relations found.";
        }

        return $"Number of relations: {errors.Count}\n" +
               StatisticsString(errors);
    }

    /// <summary>
    /// Computes relations metrics from pose graph and ground truth file.
    /// Note: Pose graph should be obtained from MapBuilder.ToProto() method.
    /// </summary>
    /// <param name="poseGraph">Pose graph proto (obtained from MapBuilder.ToProto()).</param>
    /// <param name="relationsFilename">Path to relations text file or ground truth proto file.</param>
    /// <param name="readTextFileWithUnixTimestamps">Whether to read text file with Unix timestamps.</param>
    /// <param name="writeRelationMetrics">Whether to write relation metrics to CSV file.</param>
    /// <param name="outputCsvFilename">Output CSV filename (optional).</param>
    /// <returns>Metrics string.</returns>
    public static string ComputeMetricsFromFiles(
        PoseGraphProto poseGraph,
        string relationsFilename,
        bool readTextFileWithUnixTimestamps = false,
        bool writeRelationMetrics = false,
        string? outputCsvFilename = null)
    {
        if (string.IsNullOrEmpty(relationsFilename))
            throw new ArgumentException("Relations filename cannot be null or empty", nameof(relationsFilename));

        // Load ground truth
        GroundTruthProto groundTruth;
        if (readTextFileWithUnixTimestamps)
        {
            groundTruth = RelationsTextFile.ReadRelationsTextFile(relationsFilename);
        }
        else
        {
            // Try to read as proto file first, fall back to text file if it fails
            try
            {
                groundTruth = ReadGroundTruthProto(relationsFilename);
            }
            catch
            {
                // Fall back to text file if proto reading fails
                groundTruth = RelationsTextFile.ReadRelationsTextFile(relationsFilename);
            }
        }

        var metrics = ComputeMetrics(poseGraph, groundTruth);

        // Write metrics to CSV if requested
        if (writeRelationMetrics)
        {
            var csvFilename = outputCsvFilename ?? "relation_metrics.csv";
            WriteMetricsToCsv(csvFilename, poseGraph, groundTruth);
        }

        return metrics;
    }

    /// <summary>
    /// Reads ground truth from proto file.
    /// </summary>
    /// <param name="filename">Path to ground truth proto file.</param>
    /// <returns>GroundTruth proto.</returns>
    private static GroundTruthProto ReadGroundTruthProto(string filename)
    {
        if (!File.Exists(filename))
        {
            throw new FileNotFoundException($"Ground truth file not found: {filename}", filename);
        }

        // Try reading as proto stream format first (pbstream or compressed format)
        try
        {
            using var reader = new IO.ProtoStreamReader(filename);
            if (reader.ReadProto<GroundTruthProto>(out var proto))
            {
                return proto;
            }
        }
        catch
        {
            // If proto stream reading fails, try JSON deserialization
        }

        // Try reading as JSON file (if not proto stream format)
        try
        {
            var jsonContent = File.ReadAllText(filename);
            var proto = JsonSerializer.Deserialize<GroundTruthProto>(jsonContent, JsonOptions);

            // Check if deserialization was successful (struct has default relations list if failed)
            if (proto.Relations != null && proto.Relations.Count > 0)
            {
                return proto;
            }
            // Also check if empty relations list is valid (could be empty ground truth)
            if (proto.Relations != null)
            {
                return proto;
            }
        }
        catch (Exception ex)
        {
            throw new InvalidOperationException($"Failed to read ground truth from file '{filename}'. " +
                "Expected either proto stream format (.pbstream) or JSON format. " +
                $"Error: {ex.Message}", ex);
        }

        throw new InvalidOperationException($"Failed to read ground truth from file '{filename}'. " +
            "File format not recognized or file is empty.");
    }

    /// <summary>
    /// Writes relation metrics to CSV file.
    /// </summary>
    private static void WriteMetricsToCsv(
        string csvFilename,
        PoseGraphProto poseGraph,
        GroundTruthProto groundTruth)
    {
        using var writer = new StreamWriter(csvFilename);
        writer.WriteLine("timestamp1,timestamp2,translational_error_m,rotational_error_deg,covered_distance");

        if (poseGraph.Trajectories == null || poseGraph.Trajectories.Count == 0 ||
            groundTruth.Relations == null)
        {
            return;
        }

        var trajectory = poseGraph.Trajectories[0];
        if (trajectory.Nodes == null)
        {
            return;
        }

        foreach (var relation in groundTruth.Relations)
        {
            var node1 = trajectory.Nodes.FirstOrDefault(n => n.Timestamp == relation.Timestamp1);
            var node2 = trajectory.Nodes.FirstOrDefault(n => n.Timestamp == relation.Timestamp2);

            // Check if nodes were found
            var found1 = trajectory.Nodes.Any(n => n.Timestamp == relation.Timestamp1);
            var found2 = trajectory.Nodes.Any(n => n.Timestamp == relation.Timestamp2);
            if (!found1 || !found2)
            {
                continue;
            }

            var pose1 = (Rigid3d)node1.Pose;
            var pose2 = (Rigid3d)node2.Pose;
            var expected = (Rigid3d)relation.Expected;

            var error = ComputeError(pose1, pose2, expected);
            var translationalError = Math.Sqrt(error.TranslationalSquared);
            var rotationalErrorDeg = MathUtils.RadToDeg(Math.Sqrt(error.RotationalSquared));

            writer.WriteLine($"{relation.Timestamp1},{relation.Timestamp2}," +
                            $"{translationalError:F6},{rotationalErrorDeg:F6}," +
                            $"{relation.CoveredDistance:F6}");
        }
    }
}