/*
 * 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.IO;
using CartographerSharp.Mapping.Internal;
using CartographerSharp.Mapping.Internal.D2D;
using CartographerSharp.Mapping.Internal.D3D;
using CartographerSharp.Models.Mapping;
using CartographerSharp.Transform;
using ThreadPool = CartographerSharp.Common.Threading.ThreadPool;

namespace CartographerSharp.Mapping;

using CartographerSharp.Common.Math;
using CartographerSharp.Mapping.D2D;
using CartographerSharp.Models.Transform;
using System;
using System.IO;
using System.IO.Compression;
using  RobotNet10.Shared.Numbers;


/// <summary>
/// Wires up the complete SLAM stack with TrajectoryBuilders (for local submaps)
/// and a PoseGraph for loop closure.
/// </summary>
public class MapBuilder : IMapBuilder
{
    private readonly MapBuilderOptions _options;
    private readonly ThreadPool _threadPool;
    private readonly IPoseGraph _poseGraph;
    private readonly List<ITrajectoryBuilder> _trajectoryBuilders = [];
    private readonly List<TrajectoryBuilderOptionsWithSensorIds> _allTrajectoryBuilderOptions = [];
    // CRITICAL FIX: Map trajectoryId to builder because trajectoryId may not match index in _trajectoryBuilders
    // when trajectories are loaded from map (they don't have builders)
    private readonly Dictionary<int, ITrajectoryBuilder> _trajectoryIdToBuilder = [];
    private bool _disposed = false;

    public MapBuilder(MapBuilderOptions options)
    {
        _options = options;
        
        if (options.UseTrajectoryBuilder2D == options.UseTrajectoryBuilder3D)
        {
            throw new ArgumentException("Exactly one of UseTrajectoryBuilder2D or UseTrajectoryBuilder3D must be true");
        }

        _threadPool = new ThreadPool(options.NumBackgroundThreads);

        if (options.UseTrajectoryBuilder2D)
        {
            // Match C++: Pass thread_pool to PoseGraph2D (map_builder.cc line 86-90)
            _poseGraph = new PoseGraph2D(options.PoseGraphOptions, _threadPool);
        }
        else if (options.UseTrajectoryBuilder3D)
        {
            // Match C++: Pass thread_pool to PoseGraph3D (map_builder.cc line 92-97)
            _poseGraph = new PoseGraph3D(options.PoseGraphOptions, optimizationProblem: null, threadPool: _threadPool);
        }
        else
        {
            throw new ArgumentException("Exactly one of UseTrajectoryBuilder2D or UseTrajectoryBuilder3D must be true");
        }
    }

    public int AddTrajectoryBuilder(
        HashSet<ITrajectoryBuilder.SensorId> expectedSensorIds,
        TrajectoryBuilderOptions trajectoryOptions)
    {
        // CRITICAL FIX: Use _allTrajectoryBuilderOptions.Count instead of _trajectoryBuilders.Count
        // because AddTrajectoryForDeserialization() adds to _allTrajectoryBuilderOptions but not to _trajectoryBuilders.
        // This ensures new trajectory IDs don't conflict with trajectories loaded from map.
        var trajectoryId = _allTrajectoryBuilderOptions.Count;
        
        // Select range sensor IDs
        var rangeSensorIds = expectedSensorIds
            .Where(s => s.Type == ITrajectoryBuilder.SensorId.SensorType.Range)
            .Select(s => s.Id)
            .ToList();

        ITrajectoryBuilder trajectoryBuilder;
        
        if (_options.UseTrajectoryBuilder2D)
        {
            // Create LocalTrajectoryBuilder2D with options
            // Note: TrajectoryBuilder2DOptions is JsonElement?, need to deserialize if present
            LocalTrajectoryBuilderOptions2D? localOptions = null;
            if (trajectoryOptions.TrajectoryBuilder2DOptions.HasValue)
            {
                localOptions = System.Text.Json.JsonSerializer.Deserialize<LocalTrajectoryBuilderOptions2D>(
                    trajectoryOptions.TrajectoryBuilder2DOptions.Value.GetRawText());
            }
            
            var localTrajectoryBuilder = new LocalTrajectoryBuilder2D(
                localOptions ?? new LocalTrajectoryBuilderOptions2D(),
                rangeSensorIds
            );

            // Set initial pose if provided in trajectory options
            if (trajectoryOptions.InitialTrajectoryPose.HasValue)
            {
                var initialTrajectoryPose = trajectoryOptions.InitialTrajectoryPose.Value;
                var initialPose = (Rigid3d)initialTrajectoryPose.RelativePose;
            }

            // Create motion filter for odometry if configured
            MotionFilter? motionFilter = null;
            if (trajectoryOptions.PoseGraphOdometryMotionFilter.HasValue)
            {
                motionFilter = new MotionFilter(trajectoryOptions.PoseGraphOdometryMotionFilter.Value);
            }

            // Create GlobalTrajectoryBuilder2D wrapping the local builder
            trajectoryBuilder = new GlobalTrajectoryBuilder2D(
                localTrajectoryBuilder,
                trajectoryId,
                (PoseGraph2D)_poseGraph,
                motionFilter
            );
        }
        else if (_options.UseTrajectoryBuilder3D)
        {
            // Create LocalTrajectoryBuilder3D with options
            // Note: TrajectoryBuilder3DOptions is JsonElement?, need to deserialize if present
            LocalTrajectoryBuilderOptions3D? localOptions = null;
            if (trajectoryOptions.TrajectoryBuilder3DOptions.HasValue)
            {
                localOptions = System.Text.Json.JsonSerializer.Deserialize<LocalTrajectoryBuilderOptions3D>(
                    trajectoryOptions.TrajectoryBuilder3DOptions.Value.GetRawText());
            }
            
            var localTrajectoryBuilder = new LocalTrajectoryBuilder3D(
                localOptions ?? new LocalTrajectoryBuilderOptions3D(),
                rangeSensorIds
            );

            // Create motion filter for odometry if configured
            MotionFilter? motionFilter = null;
            if (trajectoryOptions.PoseGraphOdometryMotionFilter.HasValue)
            {
                motionFilter = new MotionFilter(trajectoryOptions.PoseGraphOdometryMotionFilter.Value);
            }

            // Create GlobalTrajectoryBuilder3D wrapping the local builder
            trajectoryBuilder = new GlobalTrajectoryBuilder3D(
                localTrajectoryBuilder,
                trajectoryId,
                (PoseGraph3D)_poseGraph,
                motionFilter
            );
        }
        else
        {
            throw new ArgumentException("Exactly one of UseTrajectoryBuilder2D or UseTrajectoryBuilder3D must be true");
        }

        _trajectoryBuilders.Add(trajectoryBuilder);
        
        // CRITICAL FIX: Map trajectoryId to builder for GetTrajectoryBuilder() lookup
        _trajectoryIdToBuilder[trajectoryId] = trajectoryBuilder;

        // Store options
        var optionsWithSensorIds = new TrajectoryBuilderOptionsWithSensorIds
        {
            TrajectoryBuilderOptions = trajectoryOptions,
            SensorIds = [.. expectedSensorIds.Select(s => SensorIdOperations.ToProto(s))]
        };
        _allTrajectoryBuilderOptions.Add(optionsWithSensorIds);

        // Match C++ MaybeAddPureLocalizationTrimmer (map_builder.cc 151-152)
        if (trajectoryOptions.PureLocalizationTrimmer.HasValue)
        {
            var trimmerOpts = trajectoryOptions.PureLocalizationTrimmer.Value;
            _poseGraph.AddTrimmer(new PureLocalizationTrimmer(trajectoryId, trimmerOpts.MaxSubmapsToKeep));
        }

        // Set initial trajectory pose if provided (localization: loadFrozenState = true)
        if (trajectoryOptions.InitialTrajectoryPose.HasValue)
        {
            var initialPose = trajectoryOptions.InitialTrajectoryPose.Value;
            var relativePose = (Rigid3d)initialPose.RelativePose;
            _poseGraph.SetInitialTrajectoryPose(
                trajectoryId,
                initialPose.ToTrajectoryId,
                relativePose,
                initialPose.Timestamp
            );
            // Match C++ (map_builder.cc 164-170): when localization_mode, set pose graph mode and enable_matching_score
            if (_poseGraph is PoseGraph2D poseGraph2D)
            {
                poseGraph2D.SetLocalizationMode(true);
                // C++: if (trajectory_options.enable_matching_score()) pose_graph_->set_enable_matching_score(true);
                // C#: SetEnableMatchingScore can be added to PoseGraph2D when TrajectoryBuilderOptions has it.
            }
            // Original Cartographer: no SetInitialPose on GlobalTrajectoryBuilder. Initial pose is in pose graph (SetInitialTrajectoryPose); relocalization constraints via SetLocalizationInitialPoses. MCL refines pose before adding trajectory.
        }

        return trajectoryId;
    }

    /// <summary>
    /// Starts relocalization (match C++ MapBuilder::StartRelocalization).
    /// Sets pose graph localization callback and initial poses so constraint builder
    /// uses them when finding constraints (MaybeAddLocalizationConstraint).
    /// </summary>
    public void StartRelocalization(IReadOnlyList<Rigid3d> initialPoses, Action<int, long, Rigid3d>? callback)
    {
        if (_poseGraph is PoseGraph2D poseGraph2D)
        {
            poseGraph2D.SetLocalizationCallback(
                callback != null ? (int tid, long t, Rigid3d p) => callback(tid, t, p) : null);
            poseGraph2D.SetLocalizationInitialPoses(initialPoses ?? []);
        }
        // PoseGraph3D: add SetLocalizationCallback/SetLocalizationInitialPoses if needed for 3D
    }

    public int AddTrajectoryForDeserialization(
        TrajectoryBuilderOptionsWithSensorIds optionsWithSensorIdsProto)
    {
        // CRITICAL FIX: Use _allTrajectoryBuilderOptions.Count to ensure trajectory IDs are sequential
        // and don't conflict with trajectories created via AddTrajectoryBuilder()
        var trajectoryId = _allTrajectoryBuilderOptions.Count;
        _allTrajectoryBuilderOptions.Add(optionsWithSensorIdsProto);
        // No trajectory builder is created for deserialization
        return trajectoryId;
    }

    public ITrajectoryBuilder? GetTrajectoryBuilder(int trajectoryId)
    {
        // CRITICAL FIX: Use dictionary lookup instead of index because trajectoryId may not match
        // index in _trajectoryBuilders when trajectories are loaded from map (they don't have builders)
        if (_trajectoryIdToBuilder.TryGetValue(trajectoryId, out var builder))
        {
            return builder;
        }
        return null;
    }

    public void FinishTrajectory(int trajectoryId)
    {
        // CRITICAL FIX: Don't check _trajectoryBuilders.Count because trajectoryId may not match index
        // when trajectories are loaded from map. Just call FinishTrajectory on pose graph,
        // which will handle the check internally.
        if (trajectoryId >= 0)
        {
            _poseGraph.FinishTrajectory(trajectoryId);
        }
    }

    public string SubmapToProto(SubmapId submapId, out SubmapQuery.Response response)
    {
        // Convert submap to proto format for visualization and serialization
        var submapData = _poseGraph.GetSubmapData(submapId);
        
        if (submapData.Submap == null)
        {
            response = new SubmapQuery.Response(0, []);
            return "Submap not found";
        }

        if (submapData.Submap is Mapping.D2D.Submap2D submap2D)
        {
            // Get submap version (number of range data inserted)
            int version = submap2D.NumRangeData;

            var textures = new List<SubmapQuery.Texture>();
            
            // Extract grid data if available
            var grid = submap2D.Grid;
            if (grid != null)
            {
                // Compute cropped grid to minimal size
                var croppedGrid = grid.ComputeCroppedGrid();
                var limits = croppedGrid.Limits;
                
                var width = limits.CellLimits.NumXCells;
                var height = limits.CellLimits.NumYCells;
                var resolution = limits.Resolution;
                
                // Compute slice pose (origin of the texture)
                // C++: local_pose.inverse() * transform::Rigid3d::Translation(Eigen::Vector3d(max_x, max_y, 0.));
                // This means the slice pose is relative to the submap frame.
                var max = limits.Max;
                var sliceTranslation = new Vector3(max.X, max.Y, 0.0);
                var slicePose = submap2D.LocalPose.Inverse() * new Rigid3d(sliceTranslation, Quaternion.Identity);
                
                // Texture format is 2 bytes per pixel: Value, Alpha
                var pixels = new List<byte>(width * height * 2);
                var probabilityGrid = (ProbabilityGrid)croppedGrid;

                // Iterate cells and convert to pixels
                for (int y = 0; y < height; y++)
                {
                    for (int x = 0; x < width; x++)
                    {
                        var cellIndex = new Array2i(x, y);
                        if (probabilityGrid.IsKnown(cellIndex))
                        {
                            double probability = probabilityGrid.GetProbability(cellIndex);
                            // C++: const int delta = 128 - ProbabilityToLogOddsInteger(probability);
                            byte logOddsInteger = SubmapProbabilityUtils.ProbabilityToLogOddsInteger(probability);
                            int delta = 128 - logOddsInteger;
                            
                            byte value = (byte)(delta > 0 ? delta : 0);
                            byte alpha = (byte)(delta > 0 ? 0 : -delta);
                            
                            pixels.Add(value);
                            pixels.Add(alpha);
                        }
                        else
                        {
                            pixels.Add(0); // Unknown
                            pixels.Add(0);
                        }
                    }
                }
                
                // Compress data
                byte[] compressedBytes;
                using (var memoryStream = new MemoryStream())
                {
                    using (var gzipStream = new GZipStream(memoryStream, CompressionMode.Compress))
                    {
                        var pixelArray = pixels.ToArray();
                        gzipStream.Write(pixelArray, 0, pixelArray.Length);
                    }
                    compressedBytes = memoryStream.ToArray();
                }
                
                textures.Add(new SubmapQuery.Texture(
                    [.. compressedBytes],
                    width,
                    height,
                    resolution,
                    (Rigid3dProto)slicePose
                ));
            }

            response = new SubmapQuery.Response(version, textures);
            return string.Empty; // Empty string indicates success
        }
        else if (submapData.Submap is Mapping.D3D.Submap3D)
        {
            // 3D submaps would require voxel grid conversion
            // This is more complex and typically involves:
            // 1. Extracting hybrid grid data
            // 2. Converting TSDF or probability values
            // 3. Generating 3D texture or point cloud representation
            
            response = new SubmapQuery.Response(0, []);
            return "3D submap conversion not yet fully implemented";
        }
        
        response = new SubmapQuery.Response(0, []);
        return "Unknown submap type";
    }

    public void SerializeState(bool includeUnfinishedSubmaps, IO.IProtoStreamWriter writer)
    {
        ArgumentNullException.ThrowIfNull(writer);

        var trajectoryBuilderOptions = GetAllTrajectoryBuilderOptions();
        MappingStateSerialization.WritePbStream(
            _poseGraph,
            trajectoryBuilderOptions,
            writer,
            includeUnfinishedSubmaps
        );
    }

    public bool SerializeStateToFile(bool includeUnfinishedSubmaps, string filename)
    {
        try
        {
            using var writer = new ProtoStreamWriter(filename);
            SerializeState(includeUnfinishedSubmaps, writer);
            return writer.Close();
        }
        catch (Exception)
        {
            return false;
        }
    }

    public Dictionary<int, int> LoadState(IO.IProtoStreamReader reader, bool loadFrozenState)
    {
        ArgumentNullException.ThrowIfNull(reader);

        var deserializer = new ProtoStreamDeserializer(reader);

        // Create a copy of the pose_graph_proto, such that we can re-write the trajectory ids.
        var poseGraphProto = deserializer.PoseGraph;
        var allBuilderOptionsProto = deserializer.AllTrajectoryBuilderOptions;

        var trajectoryRemapping = new Dictionary<int, int>();
        
        // Create trajectories and build remapping
        for (int i = 0; i < poseGraphProto.Trajectories.Count; ++i)
        {
            var trajectoryProto = poseGraphProto.Trajectories[i];
            var optionsWithSensorIdsProto = allBuilderOptionsProto.OptionsWithSensorIds[i];
            
            var newTrajectoryId = AddTrajectoryForDeserialization(optionsWithSensorIdsProto);
            
            if (trajectoryRemapping.ContainsKey(trajectoryProto.TrajectoryId))
            {
                throw new InvalidOperationException(
                    $"Duplicate trajectory ID: {trajectoryProto.TrajectoryId}");
            }
            
            trajectoryRemapping[trajectoryProto.TrajectoryId] = newTrajectoryId;
            
            // Update trajectory ID in the proto (we need to modify the list element)
            var updatedTrajectory = trajectoryProto;
            updatedTrajectory.TrajectoryId = newTrajectoryId;
            poseGraphProto.Trajectories[i] = updatedTrajectory;
            
            if (loadFrozenState)
            {
                _poseGraph.FreezeTrajectory(newTrajectoryId);
            }
        }

        // Apply the calculated remapping to constraints in the pose graph proto.
        for (int i = 0; i < poseGraphProto.Constraints.Count; ++i)
        {
            var constraintProto = poseGraphProto.Constraints[i];
            
            var updatedSubmapId = constraintProto.SubmapId;
            updatedSubmapId.TrajectoryId = trajectoryRemapping[constraintProto.SubmapId.TrajectoryId];
            
            var updatedNodeId = constraintProto.NodeId;
            updatedNodeId.TrajectoryId = trajectoryRemapping[constraintProto.NodeId.TrajectoryId];
            
            var updatedConstraint = constraintProto;
            updatedConstraint.SubmapId = updatedSubmapId;
            updatedConstraint.NodeId = updatedNodeId;
            poseGraphProto.Constraints[i] = updatedConstraint;
        }

        // Build submap poses map
        var submapPoses = new MapById<SubmapId, Rigid3d>();
        foreach (var trajectoryProto in poseGraphProto.Trajectories)
        {
            foreach (var submapProto in trajectoryProto.Submaps)
            {
                var submapId = new SubmapId(trajectoryProto.TrajectoryId, submapProto.SubmapIndex);
                var pose = (Rigid3d)submapProto.Pose;
                submapPoses.Insert(submapId, pose);
            }
        }

        // Build node poses map
        var nodePoses = new MapById<NodeId, Rigid3d>();
        foreach (var trajectoryProto in poseGraphProto.Trajectories)
        {
            foreach (var nodeProto in trajectoryProto.Nodes)
            {
                var nodeId = new NodeId(trajectoryProto.TrajectoryId, nodeProto.NodeIndex);
                var pose = (Rigid3d)nodeProto.Pose;
                nodePoses.Insert(nodeId, pose);
            }
        }

        // Set global poses of landmarks.
        if (poseGraphProto.LandmarkPoses != null)
        {
            foreach (var landmark in poseGraphProto.LandmarkPoses)
            {
                _poseGraph.SetLandmarkPose(
                    landmark.LandmarkId,
                    (Rigid3d)landmark.GlobalPose,
                    true);
            }
        }

        // Check format version for 3D
        if (_options.UseTrajectoryBuilder3D)
        {
            const uint FormatVersionWithoutSubmapHistograms = 1;
            if (deserializer.Header.FormatVersion == FormatVersionWithoutSubmapHistograms)
            {
                throw new NotSupportedException(
                    "The pbstream file contains submaps without rotational histograms. " +
                    "This can be converted with the 'pbstream migrate' tool, see the " +
                    "Cartographer documentation for details.");
            }
        }

        // Read and process serialized data
        while (deserializer.ReadNextSerializedData(out var protoNullable))
        {
            if (!protoNullable.HasValue)
                break;
                
            var proto = protoNullable.Value;

            // Handle different data types
            if (proto.PoseGraph.HasValue)
            {
                // Found multiple serialized `PoseGraph`. Serialized stream likely corrupt!
                // Log error but continue
                continue;
            }
            
            if (proto.AllTrajectoryBuilderOptions.HasValue)
            {
                // Found multiple serialized `AllTrajectoryBuilderOptions`. Serialized stream likely corrupt!
                // Log error but continue
                continue;
            }
            
            if (proto.Submap.HasValue)
            {
                var submapProto = proto.Submap.Value;
                var oldTrajectoryId = submapProto.SubmapId.TrajectoryId;
                var newTrajectoryId = trajectoryRemapping[oldTrajectoryId];
                
                var submapId = new SubmapId(newTrajectoryId, submapProto.SubmapId.SubmapIndex);
                var globalPose = submapPoses[submapId];
                
                var updatedSubmap = submapProto;
                var updatedSubmapId = submapProto.SubmapId;
                updatedSubmapId.TrajectoryId = newTrajectoryId;
                updatedSubmap.SubmapId = updatedSubmapId;
                
                _poseGraph.AddSubmapFromProto(globalPose, updatedSubmap);
            }
            else if (proto.Node.HasValue)
            {
                var nodeProto = proto.Node.Value;
                var oldTrajectoryId = nodeProto.NodeId.TrajectoryId;
                var newTrajectoryId = trajectoryRemapping[oldTrajectoryId];
                
                var nodeId = new NodeId(newTrajectoryId, nodeProto.NodeId.NodeIndex);
                var nodePose = nodePoses[nodeId];
                
                var updatedNode = nodeProto;
                var updatedNodeId = nodeProto.NodeId;
                updatedNodeId.TrajectoryId = newTrajectoryId;
                updatedNode.NodeId = updatedNodeId;
                
                _poseGraph.AddNodeFromProto(nodePose, updatedNode);
            }
            else if (proto.SerializedTrajectoryData.HasValue)
            {
                var trajectoryDataProto = proto.SerializedTrajectoryData.Value;
                var oldTrajectoryId = trajectoryDataProto.TrajectoryId;
                var newTrajectoryId = trajectoryRemapping[oldTrajectoryId];
                
                var updatedTrajectoryData = new TrajectoryData(
                    newTrajectoryId,
                    trajectoryDataProto.GravityConstant,
                    trajectoryDataProto.ImuCalibration,
                    trajectoryDataProto.FixedFrameOriginInMap
                );
                
                _poseGraph.SetTrajectoryDataFromProto(updatedTrajectoryData);
            }
            else if (proto.ImuData.HasValue)
            {
                if (!loadFrozenState)
                {
                    var imuDataProto = proto.ImuData.Value;
                    var oldTrajectoryId = imuDataProto.TrajectoryId;
                    var newTrajectoryId = trajectoryRemapping[oldTrajectoryId];
                    
                    var imuData = Sensor.ImuDataOperations.FromProto(imuDataProto.ImuDataValue);
                    _poseGraph.AddImuData(newTrajectoryId, imuData);
                }
            }
            else if (proto.OdometryData.HasValue)
            {
                if (!loadFrozenState)
                {
                    var odometryDataProto = proto.OdometryData.Value;
                    var oldTrajectoryId = odometryDataProto.TrajectoryId;
                    var newTrajectoryId = trajectoryRemapping[oldTrajectoryId];
                    
                    var odometryData = Sensor.OdometryDataOperations.FromProto(odometryDataProto.OdometryDataValue);
                    _poseGraph.AddOdometryData(newTrajectoryId, odometryData);
                }
            }
            else if (proto.FixedFramePoseData.HasValue)
            {
                if (!loadFrozenState)
                {
                    var fixedFramePoseDataProto = proto.FixedFramePoseData.Value;
                    var oldTrajectoryId = fixedFramePoseDataProto.TrajectoryId;
                    var newTrajectoryId = trajectoryRemapping[oldTrajectoryId];
                    
                    var fixedFramePoseData = Sensor.FixedFramePoseDataOperations.FromProto(fixedFramePoseDataProto.FixedFramePoseDataValue);
                    _poseGraph.AddFixedFramePoseData(newTrajectoryId, fixedFramePoseData);
                }
            }
            else if (proto.LandmarkData.HasValue)
            {
                if (!loadFrozenState)
                {
                    var landmarkDataProto = proto.LandmarkData.Value;
                    var oldTrajectoryId = landmarkDataProto.TrajectoryId;
                    var newTrajectoryId = trajectoryRemapping[oldTrajectoryId];
                    
                    var landmarkData = Sensor.LandmarkDataOperations.FromProto(landmarkDataProto.LandmarkDataValue);
                    _poseGraph.AddLandmarkData(newTrajectoryId, landmarkData);
                }
            }
        }

        if (loadFrozenState)
        {
            // Add information about which nodes belong to which submap.
            // This is required, even without constraints.
            foreach (var constraintProto in poseGraphProto.Constraints)
            {
                if (constraintProto.ConstraintTag != Models.Mapping.PoseGraph.Constraint.Tag.IntraSubmap)
                {
                    continue;
                }
                
                var nodeId = new NodeId(
                    constraintProto.NodeId.TrajectoryId,
                    constraintProto.NodeId.NodeIndex);
                var submapId = new SubmapId(
                    constraintProto.SubmapId.TrajectoryId,
                    constraintProto.SubmapId.SubmapIndex);
                
                _poseGraph.AddNodeToSubmap(nodeId, submapId);
            }
        }
        else
        {
            // When loading unfrozen trajectories, 'AddSerializedConstraints' will
            // take care of adding information about which nodes belong to which submap.
            // Use the static method from ConstraintOperations class
            var constraints = ConstraintOperations.FromProto(poseGraphProto.Constraints);
            _poseGraph.AddSerializedConstraints(constraints);
        }

        // Match C++: apply transform_to_map when loading state (map_builder.cc LoadState)
        if (poseGraphProto.TransformToMap != null)
        {
            _poseGraph.SetTransformToMap((Rigid3d)poseGraphProto.TransformToMap);
        }

        if (!reader.Eof)
        {
            throw new InvalidOperationException("Reader is not at end of file after deserialization");
        }

        return trajectoryRemapping;
    }

    public Dictionary<int, int> LoadStateFromFile(string filename, bool loadFrozenState)
    {
        const string suffix = ".pbstream";
        if (filename.Length >= suffix.Length &&
            filename[^suffix.Length..] != suffix)
        {
            // Log warning: The file containing the state should be a .pbstream file.
        }

        using var reader = new ProtoStreamReader(filename);
        return LoadState(reader, loadFrozenState);
    }

    public int NumTrajectoryBuilders => _trajectoryBuilders.Count;

    public IPoseGraph PoseGraph => _poseGraph;

    public List<TrajectoryBuilderOptionsWithSensorIds> GetAllTrajectoryBuilderOptions()
    {
        return [.. _allTrajectoryBuilderOptions];
    }
    
    /// <summary>
    /// Disposes resources, including pose graph and thread pool.
    /// </summary>
    public void Dispose()
    {
        if (_disposed)
        {
            return;
        }
        
        // Dispose pose graph first (waits for optimization threads)
        if (_poseGraph is IDisposable disposablePoseGraph)
        {
            disposablePoseGraph.Dispose();
        }

        // Dispose trajectory builders (they hold CeresScanMatcher instances with native resources)
        foreach (var builder in _trajectoryBuilders)
        {
            if (builder is IDisposable disposableBuilder)
            {
                disposableBuilder.Dispose();
            }
        }

        // Dispose thread pool
        _threadPool?.Dispose();
        
        _disposed = true;
        
        GC.SuppressFinalize(this);
    }
}