/*
 * 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;
using CartographerSharp.Models.Mapping;
using CartographerSharp.Models.Transform;
using CartographerSharp.Sensor;
using  RobotNet10.Shared.Numbers;
using PoseGraphProto = CartographerSharp.Models.Mapping.PoseGraph;

namespace CartographerSharp.IO;

/// <summary>
/// The current serialization format version.
/// </summary>
internal static class MappingStateSerialization
{
    private const uint FormatVersion = 2;
    private const uint FormatVersionWithoutSubmapHistograms = 1;

    /// <summary>
    /// Serializes mapping state to a pbstream.
    /// </summary>
    public static void WritePbStream(
        IPoseGraph poseGraph,
        List<TrajectoryBuilderOptionsWithSensorIds> trajectoryBuilderOptions,
        IProtoStreamWriter writer,
        bool includeUnfinishedSubmaps)
    {
        ArgumentNullException.ThrowIfNull(poseGraph);
        ArgumentNullException.ThrowIfNull(writer);

        // Write header
        var header = new SerializationHeader(FormatVersion);
        var headerData = new SerializedData { SerializationHeader = header };
        writer.WriteProto(headerData);

        // Serialize pose graph
        var poseGraphProto = poseGraph.ToProto(includeUnfinishedSubmaps);
        var poseGraphData = new SerializedData { PoseGraph = poseGraphProto };
        writer.WriteProto(poseGraphData);

        // Get valid trajectory IDs (not deleted)
        var trajectoryStates = poseGraph.GetTrajectoryStates();
        var validTrajectoryIds = GetValidTrajectoryIds(trajectoryStates);

        // Serialize trajectory builder options
        var allOptions = CreateAllTrajectoryBuilderOptionsProto(trajectoryBuilderOptions, validTrajectoryIds);
        var optionsData = new SerializedData { AllTrajectoryBuilderOptions = allOptions };
        writer.WriteProto(optionsData);

        // Serialize submaps
        var submapData = poseGraph.GetAllSubmapData();
        SerializeSubmaps(submapData, includeUnfinishedSubmaps, writer);

        // Serialize trajectory nodes
        var trajectoryNodes = poseGraph.GetTrajectoryNodes();
        SerializeTrajectoryNodes(trajectoryNodes, writer);

        // Serialize trajectory data
        var allTrajectoryData = poseGraph.GetTrajectoryData();
        SerializeTrajectoryData(allTrajectoryData, writer);

        // Serialize IMU data
        var imuData = poseGraph.GetImuData();
        SerializeImuData(imuData, writer);

        // Serialize odometry data
        var odometryData = poseGraph.GetOdometryData();
        SerializeOdometryData(odometryData, writer);

        // Serialize fixed frame pose data
        var fixedFramePoseData = poseGraph.GetFixedFramePoseData();
        SerializeFixedFramePoseData(fixedFramePoseData, writer);

        // Serialize landmark data
        var landmarkNodes = poseGraph.GetLandmarkNodes();
        SerializeLandmarkData(landmarkNodes, writer);
    }

    private static List<int> GetValidTrajectoryIds(Dictionary<int, IPoseGraph.TrajectoryState> trajectoryStates)
    {
        var validTrajectories = new List<int>();
        foreach (var kvp in trajectoryStates)
        {
            if (kvp.Value != IPoseGraph.TrajectoryState.Deleted)
            {
                validTrajectories.Add(kvp.Key);
            }
        }
        return validTrajectories;
    }

    private static AllTrajectoryBuilderOptions CreateAllTrajectoryBuilderOptionsProto(
        List<TrajectoryBuilderOptionsWithSensorIds> allOptionsWithSensorIds,
        List<int> trajectoryIdsToSerialize)
    {
        var optionsList = new List<TrajectoryBuilderOptionsWithSensorIds>();
        foreach (var id in trajectoryIdsToSerialize)
        {
            if (id >= 0 && id < allOptionsWithSensorIds.Count)
            {
                optionsList.Add(allOptionsWithSensorIds[id]);
            }
        }
        return new AllTrajectoryBuilderOptions(optionsList);
    }

    private static void SerializeSubmaps(
        MapById<SubmapId, IPoseGraph.SubmapData> submapData,
        bool includeUnfinishedSubmaps,
        IProtoStreamWriter writer)
    {
        foreach (var kvp in submapData)
        {
            if (!includeUnfinishedSubmaps)
            {
                if (kvp.Data.Submap != null && !kvp.Data.Submap.InsertionFinished)
                {
                    continue; // Skip unfinished submaps
                }
            }

            if (kvp.Data.Submap == null)
                continue;

            var submapProto = kvp.Data.Submap.ToProto(includeGridData: true);
            
            var submapWithId = new Models.Mapping.Submap(
                new PoseGraphProto.SubmapId(kvp.Id.TrajectoryId, kvp.Id.SubmapIndex),
                submapProto.Submap2D,
                submapProto.Submap3D
            );

            var submapDataProto = new SerializedData { Submap = submapWithId };
            writer.WriteProto(submapDataProto);
        }
    }

    private static void SerializeTrajectoryNodes(
        MapById<NodeId, TrajectoryNode> trajectoryNodes,
        IProtoStreamWriter writer)
    {
        foreach (var kvp in trajectoryNodes)
        {
            if (kvp.Data.ConstantData == null)
                continue;

            var nodeData = TrajectoryNodeOperations.ToProto(kvp.Data.ConstantData);
            var nodeWithId = new Models.Mapping.Node(
                new PoseGraphProto.NodeId(kvp.Id.TrajectoryId, kvp.Id.NodeIndex),
                nodeData
            );

            var nodeDataProto = new SerializedData { Node = nodeWithId };
            writer.WriteProto(nodeDataProto);
        }
    }

    private static void SerializeTrajectoryData(
        Dictionary<int, IPoseGraph.TrajectoryData> allTrajectoryData,
        IProtoStreamWriter writer)
    {
        foreach (var kvp in allTrajectoryData)
        {
            var trajectoryData = kvp.Value;
            var imuCalibration = trajectoryData.ImuCalibration != Quaternion.Identity
                ? (Quaterniond?)new Quaterniond(
                    trajectoryData.ImuCalibration.X,
                    trajectoryData.ImuCalibration.Y,
                    trajectoryData.ImuCalibration.Z,
                    trajectoryData.ImuCalibration.W
                )
                : null;

            var fixedFrameOriginInMap = trajectoryData.FixedFrameOriginInMap.HasValue
                ? (Rigid3dProto?)(Rigid3dProto)trajectoryData.FixedFrameOriginInMap.Value
                : null;

            var serializedTrajectoryData = new SerializedTrajectoryData(
                kvp.Key,
                trajectoryData.GravityConstant,
                imuCalibration,
                fixedFrameOriginInMap
            );

            var trajectoryDataProto = new SerializedData { SerializedTrajectoryData = serializedTrajectoryData };
            writer.WriteProto(trajectoryDataProto);
        }
    }

    private static void SerializeImuData(
        Dictionary<int, List<ImuData>> imuData,
        IProtoStreamWriter writer)
    {
        foreach (var kvp in imuData)
        {
            var trajectoryId = kvp.Key;
            foreach (var imu in kvp.Value)
            {
                var imuProto = Sensor.ImuDataOperations.ToProto(imu);
                var serializedImuData = new SerializedImuData(trajectoryId, imuProto);
                var serializedData = new SerializedData { ImuData = serializedImuData };
                writer.WriteProto(serializedData);
            }
        }
    }

    private static void SerializeOdometryData(
        Dictionary<int, List<OdometryData>> odometryData,
        IProtoStreamWriter writer)
    {
        foreach (var kvp in odometryData)
        {
            var trajectoryId = kvp.Key;
            foreach (var odometry in kvp.Value)
            {
                var odometryProto = Sensor.OdometryDataOperations.ToProto(odometry);
                var serializedOdometryData = new SerializedOdometryData(trajectoryId, odometryProto);
                var serializedData = new SerializedData { OdometryData = serializedOdometryData };
                writer.WriteProto(serializedData);
            }
        }
    }

    private static void SerializeFixedFramePoseData(
        Dictionary<int, List<FixedFramePoseData>> fixedFramePoseData,
        IProtoStreamWriter writer)
    {
        foreach (var kvp in fixedFramePoseData)
        {
            var trajectoryId = kvp.Key;
            foreach (var fixedFramePose in kvp.Value)
            {
                var fixedFramePoseProto = Sensor.FixedFramePoseDataOperations.ToProto(fixedFramePose);
                var serializedFixedFramePoseData = new SerializedFixedFramePoseData(trajectoryId, fixedFramePoseProto);
                var serializedData = new SerializedData { FixedFramePoseData = serializedFixedFramePoseData };
                writer.WriteProto(serializedData);
            }
        }
    }

    private static void SerializeLandmarkData(
        Dictionary<string, IPoseGraph.LandmarkNode> landmarkNodes,
        IProtoStreamWriter writer)
    {
        foreach (var kvp in landmarkNodes)
        {
            var landmarkId = kvp.Key;
            var landmarkNode = kvp.Value;

            // Serialize each landmark observation
            foreach (var observation in landmarkNode.LandmarkObservations)
            {
                // Create landmark data from observation
                var landmarkData = new Sensor.LandmarkData(
                    observation.Time,
                    [
                        new(
                            landmarkId,
                            observation.LandmarkToTrackingTransform,
                            observation.TranslationWeight,
                            observation.RotationWeight
                        )
                    ]
                );

                // Convert to proto
                var landmarkDataProto = Sensor.LandmarkDataOperations.ToProto(landmarkData);

                var serializedLandmarkData = new SerializedLandmarkData(
                    observation.TrajectoryId,
                    landmarkDataProto
                );

                var serializedData = new SerializedData { LandmarkData = serializedLandmarkData };
                writer.WriteProto(serializedData);
            }
        }
    }
}