/*==============================================================================

  Copyright (c) Laboratory for Percutaneous Surgery (PerkLab)
  Queen's University, Kingston, ON, Canada. All Rights Reserved.

  See COPYRIGHT.txt
  or http://www.slicer.org/copyright/copyright.txt for details.

  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.

  This file was originally developed by Adam Rankin and Csaba Pinter, PerkLab, Queen's
  University and was supported through the Applied Cancer Research Unit program of Cancer
  Care Ontario with funds provided by the Ontario Ministry of Health and Long-Term Care

==============================================================================*/

// Segmentations includes
#include "vtkMRMLSegmentationStorageNode.h"

#include "vtkSegmentation.h"
#include "vtkOrientedImageData.h"
#include "vtkOrientedImageDataResample.h"

// MRML includes
#include "vtkMRMLI18N.h"
#include "vtkMRMLMessageCollection.h"
#include <vtkMRMLScalarVolumeNode.h>
#include <vtkMRMLScene.h>
#include "vtkMRMLSegmentationNode.h"
#include "vtkMRMLSegmentationDisplayNode.h"

// VTK includes
#include <vtkDataObject.h>
#include <vtkDoubleArray.h>
#include <vtkErrorCode.h>
#include <vtkFieldData.h>
#include <vtkImageAccumulate.h>
#include <vtkImageAppendComponents.h>
#include <vtkImageCast.h>
#include <vtkImageConstantPad.h>
#include <vtkImageExtractComponents.h>
#include <vtkInformation.h>
#include <vtkInformationIntegerVectorKey.h>
#include <vtkInformationStringKey.h>
#include <vtkITKArchetypeImageSeriesVectorReaderFile.h>
#include <vtkMultiBlockDataSet.h>
#include <vtkNew.h>
#include <vtkTeemNRRDWriter.h>
#include <vtkObjectFactory.h>
#include <vtkPointData.h>
#include <vtkPolyData.h>
#include <vtkStringArray.h>
#include <vtkTransform.h>
#include <vtkXMLMultiBlockDataWriter.h>
#include <vtkXMLMultiBlockDataReader.h>
#include <vtksys/SystemTools.hxx>

#ifdef SUPPORT_4D_SPATIAL_NRRD
// ITK includes
# include <itkMacro.h>
# include <itkImageFileWriter.h>
# include <itkImageFileReader.h>
# include <itkMetaDataDictionary.h>
# include <itkMetaDataObject.h>
# include <itkNrrdImageIO.h>
#endif

#include "teem/nrrd.h"

// STL & C++ includes
#include <iterator>
#include <sstream>

//----------------------------------------------------------------------------
static const char LIST_SEPARATOR = '|';
static const std::string KEY_SEGMENT_ID = "ID";
static const std::string KEY_SEGMENT_NAME = "Name";
static const std::string KEY_SEGMENT_COLOR = "Color";
static const std::string KEY_SEGMENT_TAGS = "Tags";
static const std::string KEY_SEGMENT_EXTENT = "Extent";
static const std::string KEY_SEGMENT_LABEL_VALUE = "LabelValue";
static const std::string KEY_SEGMENT_LAYER = "Layer";
static const std::string KEY_SEGMENT_NAME_AUTO_GENERATED = "NameAutoGenerated";
static const std::string KEY_SEGMENT_COLOR_AUTO_GENERATED = "ColorAutoGenerated";
static const std::string KEY_SEGMENTATION_SOURCE_REPRESENTATION = "MasterRepresentation";
static const std::string KEY_SEGMENTATION_CONVERSION_PARAMETERS = "ConversionParameters";
static const std::string KEY_SEGMENTATION_EXTENT = "Extent"; // Deprecated, kept only for being able to read legacy files.
static const std::string KEY_SEGMENTATION_REFERENCE_IMAGE_EXTENT_OFFSET = "ReferenceImageExtentOffset";
static const std::string KEY_SEGMENTATION_CONTAINED_REPRESENTATION_NAMES = "ContainedRepresentationNames";

static const int SINGLE_SEGMENT_INDEX = -1; // used as segment index when there is only a single segment
//----------------------------------------------------------------------------
vtkMRMLNodeNewMacro(vtkMRMLSegmentationStorageNode);

//----------------------------------------------------------------------------
vtkMRMLSegmentationStorageNode::vtkMRMLSegmentationStorageNode() = default;

//----------------------------------------------------------------------------
vtkMRMLSegmentationStorageNode::~vtkMRMLSegmentationStorageNode() = default;

//----------------------------------------------------------------------------
void vtkMRMLSegmentationStorageNode::PrintSelf(ostream& os, vtkIndent indent)
{
  this->TypeDisplayName = vtkMRMLTr("vtkMRMLSegmentationStorageNode", "Segmentation Storage");

  Superclass::PrintSelf(os, indent);
  vtkMRMLPrintBeginMacro(os, indent);
  vtkMRMLPrintBooleanMacro(CropToMinimumExtent);
  vtkMRMLPrintEndMacro();
}

//----------------------------------------------------------------------------
void vtkMRMLSegmentationStorageNode::ReadXMLAttributes(const char** atts)
{
  MRMLNodeModifyBlocker blocker(this);
  Superclass::ReadXMLAttributes(atts);
  vtkMRMLReadXMLBeginMacro(atts);
  vtkMRMLReadXMLBooleanMacro(CropToMinimumExtent, CropToMinimumExtent);
  vtkMRMLReadXMLEndMacro();
}

//----------------------------------------------------------------------------
void vtkMRMLSegmentationStorageNode::WriteXML(ostream& of, int nIndent)
{
  Superclass::WriteXML(of, nIndent);
  vtkMRMLWriteXMLBeginMacro(of);
  vtkMRMLWriteXMLBooleanMacro(CropToMinimumExtent, CropToMinimumExtent);
  vtkMRMLWriteXMLEndMacro();
}

//----------------------------------------------------------------------------
// Copy the node's attributes to this object.
// Does NOT copy: ID, FilePrefix, Name, StorageID
void vtkMRMLSegmentationStorageNode::Copy(vtkMRMLNode* anode)
{
  MRMLNodeModifyBlocker blocker(anode);
  Superclass::Copy(anode);
  vtkMRMLCopyBeginMacro(anode);
  vtkMRMLCopyBooleanMacro(CropToMinimumExtent);
  vtkMRMLCopyEndMacro();
}

//----------------------------------------------------------------------------
void vtkMRMLSegmentationStorageNode::InitializeSupportedReadFileTypes()
{
  //: File format name
  std::string fileType = vtkMRMLTr("vtkMRMLSegmentationStorageNode", "Segmentation");
  this->SupportedReadFileTypes->InsertNextValue(fileType + " (.seg.nrrd)");
  this->SupportedReadFileTypes->InsertNextValue(fileType + " (.seg.nhdr)");
  this->SupportedReadFileTypes->InsertNextValue(fileType + " (.seg.vtm)");
  this->SupportedReadFileTypes->InsertNextValue(fileType + " (.nrrd)");
  this->SupportedReadFileTypes->InsertNextValue(fileType + " (.vtm)");
  this->SupportedReadFileTypes->InsertNextValue(fileType + " (.nii.gz)");
  this->SupportedReadFileTypes->InsertNextValue(fileType + " (.hdr)");
  this->SupportedReadFileTypes->InsertNextValue(fileType + " (.nii)");
}

//----------------------------------------------------------------------------
void vtkMRMLSegmentationStorageNode::InitializeSupportedWriteFileTypes()
{
  Superclass::InitializeSupportedWriteFileTypes();
  vtkMRMLSegmentationNode* segmentationNode = this->GetAssociatedDataNode();
  bool masterIsImage = true;
  bool masterIsPolyData = true;
  if (segmentationNode)
  {
    // restrict write file types to those that are suitable for current source representation
    masterIsImage = segmentationNode->GetSegmentation()->IsSourceRepresentationImageData();
    masterIsPolyData = segmentationNode->GetSegmentation()->IsSourceRepresentationPolyData();
    if (!masterIsImage && !masterIsPolyData)
    {
      // if contains unknown representation then enable all formats
      masterIsImage = true;
      masterIsPolyData = true;
    }
  }
  //: File format name
  std::string fileType = vtkMRMLTr("vtkMRMLSegmentationStorageNode", "Segmentation");
  if (masterIsImage)
  {
    this->SupportedWriteFileTypes->InsertNextValue(fileType + " (.seg.nrrd)");
    this->SupportedWriteFileTypes->InsertNextValue(fileType + " (.seg.nhdr)");
    this->SupportedWriteFileTypes->InsertNextValue(fileType + " (.nrrd)");
    this->SupportedWriteFileTypes->InsertNextValue(fileType + " (.nhdr)");
  }
  if (masterIsPolyData)
  {
    this->SupportedWriteFileTypes->InsertNextValue(fileType + " (.seg.vtm)");
    this->SupportedWriteFileTypes->InsertNextValue(fileType + " (.vtm)");
  }
}

//----------------------------------------------------------------------------
vtkMRMLSegmentationNode* vtkMRMLSegmentationStorageNode::GetAssociatedDataNode()
{
  if (!this->GetScene())
  {
    return nullptr;
  }

  std::vector<vtkMRMLNode*> segmentationNodes;
  unsigned int numberOfNodes = this->GetScene()->GetNodesByClass("vtkMRMLSegmentationNode", segmentationNodes);
  for (unsigned int nodeIndex = 0; nodeIndex < numberOfNodes; nodeIndex++)
  {
    vtkMRMLSegmentationNode* node = vtkMRMLSegmentationNode::SafeDownCast(segmentationNodes[nodeIndex]);
    if (node)
    {
      const char* storageNodeID = node->GetStorageNodeID();
      if (storageNodeID && !strcmp(storageNodeID, this->ID))
      {
        return vtkMRMLSegmentationNode::SafeDownCast(node);
      }
    }
  }

  return nullptr;
}

//----------------------------------------------------------------------------
const char* vtkMRMLSegmentationStorageNode::GetDefaultWriteFileExtension()
{
  vtkMRMLSegmentationNode* segmentationNode = this->GetAssociatedDataNode();
  if (!segmentationNode)
  {
    return nullptr;
  }
  if (segmentationNode->GetSegmentation()->IsSourceRepresentationImageData())
  {
    return "seg.nrrd";
  }
  else if (segmentationNode->GetSegmentation()->IsSourceRepresentationPolyData())
  {
    return "seg.vtm";
  }
  // Source representation is not supported for writing to file
  return nullptr;
}

//----------------------------------------------------------------------------
void vtkMRMLSegmentationStorageNode::ResetSupportedWriteFileTypes()
{
  this->InitializeSupportedWriteFileTypes();
}

//----------------------------------------------------------------------------
bool vtkMRMLSegmentationStorageNode::CanReadInReferenceNode(vtkMRMLNode* refNode)
{
  return refNode->IsA("vtkMRMLSegmentationNode");
}

//----------------------------------------------------------------------------
int vtkMRMLSegmentationStorageNode::ReadDataInternal(vtkMRMLNode* refNode)
{
  vtkMRMLSegmentationNode* segmentationNode = vtkMRMLSegmentationNode::SafeDownCast(refNode);
  if (!segmentationNode)
  {
    vtkErrorToMessageCollectionMacro(this->GetUserMessages(),
                                     "vtkMRMLSegmentationStorageNode::ReadDataInternal",
                                     "Node for storing reading result (" << (this->ID ? this->ID : "(unknown)") << ") is not a valid segmentation node.");
    return 0;
  }

  MRMLNodeModifyBlocker blocker(segmentationNode);

  std::string fullName = this->GetFullNameFromFileName();
  if (fullName.empty())
  {
    vtkErrorToMessageCollectionMacro(
      this->GetUserMessages(), "vtkMRMLSegmentationStorageNode::ReadDataInternal", "Filename is not specified (" << (this->ID ? this->ID : "(unknown)") << ").");
    return 0;
  }

  // Check that the file exists
  if (vtksys::SystemTools::FileExists(fullName.c_str()) == false)
  {
    vtkErrorToMessageCollectionMacro(this->GetUserMessages(),
                                     "vtkMRMLSegmentationStorageNode::ReadDataInternal",
                                     "Segmentation file '" << fullName.c_str() << "' is not found while trying to read node (" << (this->ID ? this->ID : "(unknown)") << ").");
    return 0;
  }

  bool success = false;
  // Try to read as labelmap first then as poly data
  if (this->ReadBinaryLabelmapRepresentation(segmentationNode, fullName))
  {
    success = true;
  }
#ifdef SUPPORT_4D_SPATIAL_NRRD
  else if (this->ReadBinaryLabelmapRepresentation4DSpatial(segmentationNode, fullName))
  {
    success = true;
  }
#endif
  else if (this->ReadPolyDataRepresentation(segmentationNode, fullName))
  {
    success = true;
  }

  // Create display node if segmentation there is none
  if (success && !segmentationNode->GetDisplayNode())
  {
    segmentationNode->CreateDefaultDisplayNodes();
  }

  if (!success)
  {
    // Failed to read
    vtkErrorToMessageCollectionMacro(this->GetUserMessages(),
                                     "vtkMRMLSegmentationStorageNode::ReadDataInternal",
                                     "File '" << fullName << "' could not be read neither as labelmap nor polydata"
                                              << " while trying to read node (" << (this->ID ? this->ID : "(unknown)") << ").");
    return 0;
  }

  return 1;
}

#ifdef SUPPORT_4D_SPATIAL_NRRD
//----------------------------------------------------------------------------
int vtkMRMLSegmentationStorageNode::ReadBinaryLabelmapRepresentation4DSpatial(vtkMRMLSegmentationNode* segmentationNode, std::string path)
{
  if (!vtksys::SystemTools::FileExists(path.c_str()))
  {
    vtkErrorToMessageCollectionMacro(
      this->GetUserMessages(), "vtkMRMLSegmentationStorageNode::ReadBinaryLabelmapRepresentation4DSpatial", "Input file " << path << " does not exist.");
    return 0;
  }

  // Set up output segmentation
  if (!segmentationNode)
  {
    vtkErrorToMessageCollectionMacro(this->GetUserMessages(), "vtkMRMLSegmentationStorageNode::ReadBinaryLabelmapRepresentation4DSpatial", "Output segmentation must exist.");
    return 0;
  }

  // Read 4D NRRD image file
  typedef itk::ImageFileReader<BinaryLabelmap4DImageType> FileReaderType;
  FileReaderType::Pointer reader = FileReaderType::New();
  reader->SetFileName(path);
  try
  {
    reader->Update();
  }
  catch (itk::ImageFileReaderException& error)
  {
    // Do not report error as the file might contain poly data in which case ReadPolyDataRepresentation will read it alright
    vtkDebugMacro("ReadBinaryLabelmapRepresentation: Failed to load file " << path << " as segmentation. Exception:\n" << error);
    return 0;
  }
  BinaryLabelmap4DImageType::Pointer allSegmentLabelmapsImage = reader->GetOutput();

  // Read succeeded

  // Make sure there is a valid segmentation object in the node
  vtkSegmentation* segmentation = segmentationNode->GetSegmentation();
  if (!segmentation)
  {
    vtkNew<vtkSegmentation> newSegmentation;
    segmentation = newSegmentation;
    segmentationNode->SetAndObserveSegmentation(newSegmentation);
  }

  MRMLNodeModifyBlocker blocker(segmentationNode);

  // Clean out the segmentation before adding the new segments
  if (segmentation->GetNumberOfSegments() > 0)
  {
    segmentation->RemoveAllSegments();
  }

  // Set source representation
  segmentation->SetSourceRepresentationName(vtkSegmentationConverter::GetSegmentationBinaryLabelmapRepresentationName());

  // Get metadata dictionary from image
  itk::MetaDataDictionary metadata = allSegmentLabelmapsImage->GetMetaDataDictionary();

  // Read common geometry extent
  std::string commonExtent;
  itk::ExposeMetaData<std::string>(metadata, GetSegmentationMetaDataKey(KEY_SEGMENTATION_EXTENT).c_str(), commonExtent);
  int commonGeometryExtent[6] = { 0, -1, 0, -1, 0, -1 };
  GetImageExtentFromString(commonGeometryExtent, commonExtent);
  // Read conversion parameters
  std::string conversionParameters;
  itk::ExposeMetaData<std::string>(metadata, GetSegmentationMetaDataKey(KEY_SEGMENTATION_CONVERSION_PARAMETERS).c_str(), conversionParameters);
  segmentation->DeserializeConversionParameters(conversionParameters);
  // Read contained representation names
  std::string containedRepresentationNames;
  itk::ExposeMetaData<std::string>(metadata, GetSegmentationMetaDataKey(KEY_SEGMENTATION_CONTAINED_REPRESENTATION_NAMES).c_str(), containedRepresentationNames);

  // Get image properties
  BinaryLabelmap4DImageType::RegionType itkRegion = allSegmentLabelmapsImage->GetLargestPossibleRegion();
  BinaryLabelmap4DImageType::PointType itkOrigin = allSegmentLabelmapsImage->GetOrigin();
  BinaryLabelmap4DImageType::SpacingType itkSpacing = allSegmentLabelmapsImage->GetSpacing();
  BinaryLabelmap4DImageType::DirectionType itkDirections = allSegmentLabelmapsImage->GetDirection();
  // Make image properties accessible for VTK
  double origin[3] = { itkOrigin[0], itkOrigin[1], itkOrigin[2] };
  double spacing[3] = { itkSpacing[0], itkSpacing[1], itkSpacing[2] };
  double directions[3][3] = { { 1.0, 0.0, 0.0 }, { 0.0, 1.0, 0.0 }, { 0.0, 0.0, 1.0 } };
  for (unsigned int col = 0; col < 3; col++)
  {
    for (unsigned int row = 0; row < 3; row++)
    {
      directions[row][col] = itkDirections[row][col];
    }
  }

  // Read segment binary labelmaps
  for (unsigned int segmentIndex = itkRegion.GetIndex()[3]; segmentIndex < itkRegion.GetIndex()[3] + itkRegion.GetSize()[3]; ++segmentIndex)
  {
    // Create segment
    vtkSmartPointer<vtkSegment> currentSegment = vtkSmartPointer<vtkSegment>::New();

    // Get metadata for current segment

    // ID
    std::string currentSegmentID;
    itk::ExposeMetaData<std::string>(metadata, GetSegmentMetaDataKey(segmentIndex, KEY_SEGMENT_ID), currentSegmentID);

    // Name
    std::string currentSegmentName;
    itk::ExposeMetaData<std::string>(metadata, GetSegmentMetaDataKey(segmentIndex, KEY_SEGMENT_NAME), currentSegmentName);
    currentSegment->SetName(currentSegmentName.c_str());

    // Color
    std::string colorString;
    itk::ExposeMetaData<std::string>(metadata, GetSegmentMetaDataKey(segmentIndex, KEY_SEGMENT_COLOR), colorString);
    if (colorString.empty())
    {
      // Backwards compatibility
      itk::ExposeMetaData<std::string>(metadata, GetSegmentMetaDataKey(segmentIndex, "DefaultColor"), colorString);
    }
    double currentSegmentColor[3] = { 0.0, 0.0, 0.0 };
    GetSegmentColorFromString(currentSegmentColor, colorString);

    // Extent
    std::string extentValue;
    itk::ExposeMetaData<std::string>(metadata, GetSegmentMetaDataKey(segmentIndex, KEY_SEGMENT_EXTENT), extentValue);
    int currentSegmentExtent[6] = { 0, -1, 0, -1, 0, -1 };
    GetImageExtentFromString(currentSegmentExtent, extentValue);

    // Tags
    std::string tagsValue;
    itk::ExposeMetaData<std::string>(metadata, GetSegmentMetaDataKey(segmentIndex, KEY_SEGMENT_TAGS), tagsValue);
    SetSegmentTagsFromString(currentSegment, tagsValue);

    // NameAutoGenerated
    std::string currentSegmentNameAutoGenerated;
    itk::ExposeMetaData<std::string>(metadata, GetSegmentMetaDataKey(segmentIndex, KEY_SEGMENT_NAME_AUTO_GENERATED), currentSegmentNameAutoGenerated);
    currentSegment->SetNameAutoGenerated(!currentSegmentNameAutoGenerated.compare("1"));

    // ColorAutoGenerated
    std::string currentSegmentColorAutoGenerated;
    itk::ExposeMetaData<std::string>(metadata, GetSegmentMetaDataKey(segmentIndex, KEY_SEGMENT_COLOR_AUTO_GENERATED), currentSegmentColorAutoGenerated);
    currentSegment->SetColorAutoGenerated(!currentSegmentColorAutoGenerated.compare("1"));

    // Create binary labelmap volume
    vtkSmartPointer<vtkOrientedImageData> currentBinaryLabelmap = vtkSmartPointer<vtkOrientedImageData>::New();
    currentBinaryLabelmap->SetOrigin(origin);
    currentBinaryLabelmap->SetSpacing(spacing);
    currentBinaryLabelmap->SetDirections(directions);
    currentBinaryLabelmap->SetExtent(currentSegmentExtent);
    currentBinaryLabelmap->AllocateScalars(VTK_UNSIGNED_CHAR, 1);
    unsigned char* labelmapPtr = (unsigned char*)currentBinaryLabelmap->GetScalarPointerForExtent(currentSegmentExtent);

    // Define ITK region for current segment
    BinaryLabelmap4DImageType::RegionType segmentRegion;
    BinaryLabelmap4DImageType::SizeType segmentRegionSize;
    BinaryLabelmap4DImageType::IndexType segmentRegionIndex;
    segmentRegionIndex[0] = segmentRegionIndex[1] = segmentRegionIndex[2] = 0;
    segmentRegionIndex[3] = segmentIndex;
    segmentRegionSize = itkRegion.GetSize();
    segmentRegionSize[3] = 1;
    segmentRegion.SetIndex(segmentRegionIndex);
    segmentRegion.SetSize(segmentRegionSize);

    // Iterate through current segment's region and read voxel values into segment labelmap
    BinaryLabelmap4DIteratorType segmentLabelmapIterator(allSegmentLabelmapsImage, segmentRegion);
    for (segmentLabelmapIterator.GoToBegin(); !segmentLabelmapIterator.IsAtEnd(); ++segmentLabelmapIterator)
    {
      // Skip region outside extent of current segment (consider common extent boundaries)
      BinaryLabelmap4DImageType::IndexType segmentIndex = segmentLabelmapIterator.GetIndex();
      if (segmentIndex[0] + commonGeometryExtent[0] < currentSegmentExtent[0]    //
          || segmentIndex[0] + commonGeometryExtent[0] > currentSegmentExtent[1] //
          || segmentIndex[1] + commonGeometryExtent[2] < currentSegmentExtent[2] //
          || segmentIndex[1] + commonGeometryExtent[2] > currentSegmentExtent[3] //
          || segmentIndex[2] + commonGeometryExtent[4] < currentSegmentExtent[4] //
          || segmentIndex[2] + commonGeometryExtent[4] > currentSegmentExtent[5])
      {
        continue;
      }

      // Get voxel value
      unsigned char voxelValue = segmentLabelmapIterator.Get();

      // Set voxel value in current segment labelmap
      (*labelmapPtr) = voxelValue;
      ++labelmapPtr;
    }

    // Set loaded binary labelmap to segment
    currentSegment->AddRepresentation(vtkSegmentationConverter::GetSegmentationBinaryLabelmapRepresentationName(), currentBinaryLabelmap);

    // Add segment to segmentation
    segmentation->AddSegment(currentSegment, currentSegmentID);
  }

  // Create contained representations now that all the data is loaded
  this->CreateRepresentationsBySerializedNames(segmentation, containedRepresentationNames);

  return 1;
}
#endif // SUPPORT_4D_SPATIAL_NRRD

//----------------------------------------------------------------------------
int vtkMRMLSegmentationStorageNode::ReadBinaryLabelmapRepresentation(vtkMRMLSegmentationNode* segmentationNode, std::string path)
{
  if (!vtksys::SystemTools::FileExists(path.c_str()))
  {
    vtkErrorToMessageCollectionMacro(this->GetUserMessages(), "vtkMRMLSegmentationStorageNode::ReadBinaryLabelmapRepresentation", "Input file " << path << " does not exist.");
    return 0;
  }
  if (!segmentationNode)
  {
    vtkErrorToMessageCollectionMacro(this->GetUserMessages(), "vtkMRMLSegmentationStorageNode::ReadBinaryLabelmapRepresentation", "Output segmentation must exist.");
    return 0;
  }

  // Make sure there is a valid segmentation object in the node
  vtkSegmentation* segmentation = segmentationNode->GetSegmentation();
  if (!segmentation)
  {
    vtkNew<vtkSegmentation> newSegmentation;
    segmentation = newSegmentation;
    segmentationNode->SetAndObserveSegmentation(newSegmentation);
  }

  vtkSmartPointer<vtkImageData> imageData = nullptr;

  vtkNew<vtkITKArchetypeImageSeriesVectorReaderFile> archetypeImageReader;
  archetypeImageReader->SetSingleFile(1);
  archetypeImageReader->SetUseOrientationFromFile(1);
  archetypeImageReader->ResetFileNames();
  archetypeImageReader->SetArchetype(path.c_str());
  archetypeImageReader->SetOutputScalarTypeToNative();
  archetypeImageReader->SetDesiredCoordinateOrientationToNative();
  archetypeImageReader->SetUseNativeOriginOn();

  int numberOfSegments = 0;
  std::map<int, std::vector<int>> segmentIndexInLayer;
  std::string containedRepresentationNames;
  vtkMatrix4x4* rasToFileIjk = nullptr;
  int imageExtentInFile[6] = { 0, -1, 0, -1, 0, -1 };
  int commonGeometryExtent[6] = { 0, -1, 0, -1, 0, -1 };
  bool isExtentValid = false;
  int referenceImageExtentOffset[3] = { 0, 0, 0 };

  if (archetypeImageReader->CanReadFile(path.c_str()))
  {
    // Read the volume
    this->GetUserMessages()->SetObservedObject(archetypeImageReader);
    archetypeImageReader->Update();
    this->GetUserMessages()->SetObservedObject(nullptr);
    if (archetypeImageReader->GetErrorCode() != vtkErrorCode::NoError)
    {
      vtkErrorToMessageCollectionMacro(this->GetUserMessages(), "vtkMRMLSegmentationStorageNode::ReadBinaryLabelmapRepresentation", "Error reading image.");
      return 0;
    }

    // Copy image data to sequence of volume nodes
    imageData = archetypeImageReader->GetOutput();
    rasToFileIjk = archetypeImageReader->GetRasToIjkMatrix();
    imageData->GetExtent(imageExtentInFile);
    imageData->GetExtent(commonGeometryExtent);

    // Get metadata dictionary from image
    itk::MetaDataDictionary dictionary = archetypeImageReader->GetMetaDataDictionary();

    std::string segmentationExtentString;
    if (this->GetSegmentationMetaDataFromDicitionary(segmentationExtentString, dictionary, KEY_SEGMENTATION_EXTENT))
    {
      // Legacy format. Return and read using ReadBinaryLabelmapRepresentation4DSpatial if available.
      return 0;
    }

    // Read common geometry
    std::string referenceImageExtentOffsetStr;
    if (this->GetSegmentationMetaDataFromDicitionary(referenceImageExtentOffsetStr, dictionary, KEY_SEGMENTATION_REFERENCE_IMAGE_EXTENT_OFFSET))
    {
      // Common geometry extent is specified by an offset (extent[0], extent[2], extent[4]) and the size of the image
      // NRRD file cannot store start extent, so we store that in KEY_SEGMENTATION_REFERENCE_IMAGE_EXTENT_OFFSET and from imageExtentInFile we
      // only use the extent size.
      std::stringstream ssExtentValue(referenceImageExtentOffsetStr);
      ssExtentValue >> referenceImageExtentOffset[0] >> referenceImageExtentOffset[1] >> referenceImageExtentOffset[2];
      commonGeometryExtent[0] = referenceImageExtentOffset[0];
      commonGeometryExtent[1] = referenceImageExtentOffset[0] + imageExtentInFile[1] - imageExtentInFile[0];
      commonGeometryExtent[2] = referenceImageExtentOffset[1];
      commonGeometryExtent[3] = referenceImageExtentOffset[1] + imageExtentInFile[3] - imageExtentInFile[2];
      commonGeometryExtent[4] = referenceImageExtentOffset[2];
      commonGeometryExtent[5] = referenceImageExtentOffset[2] + imageExtentInFile[5] - imageExtentInFile[4];
    }
    else
    {
      // KEY_SEGMENTATION_REFERENCE_IMAGE_EXTENT_OFFSET is not specified,
      // which means that this is probably a regular NRRD file that should be imported as a segmentation.
      // Use the image extent as common geometry extent.
      vtkInfoMacro(<< KEY_SEGMENTATION_REFERENCE_IMAGE_EXTENT_OFFSET << " attribute was not found in NRRD segmentation file. Assume no offset.");
      imageData->GetExtent(commonGeometryExtent);
    }

    // Special case: extent = [0, 0, 0, 0, 0, 0] means there is no image data
    isExtentValid = true;
    if (imageExtentInFile[0] == 0    //
        && imageExtentInFile[1] == 0 //
        && imageExtentInFile[2] == 0 //
        && imageExtentInFile[3] == 0 //
        && imageExtentInFile[4] == 0 //
        && imageExtentInFile[5] == 0)
    {
      imageExtentInFile[1] = -1;
      imageExtentInFile[3] = -1;
      imageExtentInFile[5] = -1;

      commonGeometryExtent[0] = 0;
      commonGeometryExtent[1] = -1;
      commonGeometryExtent[2] = 0;
      commonGeometryExtent[3] = -1;
      commonGeometryExtent[4] = 0;
      commonGeometryExtent[5] = -1;

      isExtentValid = false;
    }

    // Read conversion parameters
    std::string conversionParameters;
    if (this->GetSegmentationMetaDataFromDicitionary(conversionParameters, dictionary, KEY_SEGMENTATION_CONVERSION_PARAMETERS))
    {
      segmentation->DeserializeConversionParameters(conversionParameters);
    }

    // Read contained representation names
    this->GetSegmentationMetaDataFromDicitionary(containedRepresentationNames, dictionary, KEY_SEGMENTATION_CONTAINED_REPRESENTATION_NAMES);

    // Read contained segment layer numbers
    while (dictionary.HasKey(GetSegmentMetaDataKey(numberOfSegments, KEY_SEGMENT_ID)))
    {
      int segmentIndex = numberOfSegments;
      std::string layerValue;
      if (this->GetSegmentMetaDataFromDicitionary(layerValue, dictionary, segmentIndex, KEY_SEGMENT_LAYER))
      {
        segmentIndexInLayer[vtkVariant(layerValue).ToInt()].push_back(segmentIndex);
      }
      else
      {
        segmentIndexInLayer[segmentIndex].push_back(segmentIndex);
      }
      ++numberOfSegments;
    }
  }
  else
  {
    vtkDebugMacro("ReadBinaryLabelmapRepresentation: File is not using a supported format");
    return 0;
  }

  if (imageData == nullptr)
  {
    vtkErrorToMessageCollectionMacro(this->GetUserMessages(), "vtkMRMLSegmentationStorageNode::ReadBinaryLabelmapRepresentation", "Error reading image: invalid image data.");
    return 0;
  }

  // Read succeeded
  int ret = vtkOrientedImageDataResample::IsImageScalarTypeValid(imageData);
  switch (ret)
  {
    case vtkOrientedImageDataResample::TYPE_CONVERSION_TRUNCATION_NEEDED:
      vtkWarningToMessageCollectionMacro(this->GetUserMessages(),
                                         "vtkMRMLSegmentationStorageNode::ReadBinaryLabelmapRepresentation",
                                         "Segmentation is a floating point scalar type and will be cast to an integer type by truncation (rounding towards 0).");
      break;
    case vtkOrientedImageDataResample::TYPE_CONVERSION_CLAMPING_NEEDED:
      vtkWarningToMessageCollectionMacro(
        this->GetUserMessages(),
        "vtkMRMLSegmentationStorageNode::ReadBinaryLabelmapRepresentation",
        "Segmentation is a floating point scalar type and are outside of the possible range of integer. Voxel values will be clamped to integer range.");
      break;
    case vtkOrientedImageDataResample::TYPE_ERROR:
      vtkErrorToMessageCollectionMacro(this->GetUserMessages(), "vtkMRMLSegmentationStorageNode::ReadBinaryLabelmapRepresentation", "Failed to cast image to integer type.");
      return 0;
    case vtkOrientedImageDataResample::TYPE_OK:
    default: break;
  }
  if (ret != vtkOrientedImageDataResample::TYPE_OK //
      && !vtkOrientedImageDataResample::CastSegmentationToSmallestIntegerType(imageData))
  {
    vtkErrorToMessageCollectionMacro(this->GetUserMessages(), "vtkMRMLSegmentationStorageNode::ReadBinaryLabelmapRepresentation", "Failed to cast image to integer type.");
    return 0;
  }

  int numberOfFrames = imageData->GetNumberOfScalarComponents();

  MRMLNodeModifyBlocker blocker(segmentationNode);

  // Clean out the segmentation before adding the new segments
  if (segmentation->GetNumberOfSegments() > 0)
  {
    segmentation->RemoveAllSegments();
  }

  // Set source representation
  segmentation->SetSourceRepresentationName(vtkSegmentationConverter::GetSegmentationBinaryLabelmapRepresentationName());

  // Compensate for the extent shift in the image origin.
  // We change the origin so that if a reader ignores private fields, such as
  // referenceImageExtentOffset, the image is placed correctly in physical coordinate system.
  vtkNew<vtkMatrix4x4> fileIjkToIjk;
  fileIjkToIjk->SetElement(0, 3, referenceImageExtentOffset[0]);
  fileIjkToIjk->SetElement(1, 3, referenceImageExtentOffset[1]);
  fileIjkToIjk->SetElement(2, 3, referenceImageExtentOffset[2]);
  vtkNew<vtkMatrix4x4> rasToIjk;
  vtkMatrix4x4::Multiply4x4(fileIjkToIjk.GetPointer(), rasToFileIjk, rasToIjk.GetPointer());
  vtkNew<vtkMatrix4x4> imageToWorldMatrix; // = ijkToRas;
  vtkMatrix4x4::Invert(rasToIjk.GetPointer(), imageToWorldMatrix.GetPointer());

  imageData->SetExtent(commonGeometryExtent);
  vtkNew<vtkImageExtractComponents> extractComponents;
  extractComponents->SetInputData(imageData);

  vtkNew<vtkImageConstantPad> padder;
  padder->SetInputConnection(extractComponents->GetOutputPort());

  // Get metadata for current segment
  itk::MetaDataDictionary dictionary = archetypeImageReader->GetMetaDataDictionary();

  std::vector<vtkSmartPointer<vtkSegment>> segments(numberOfSegments);
  std::map<int, vtkSmartPointer<vtkOrientedImageData>> layerToImage;
  for (int frameIndex = 0; frameIndex < numberOfFrames; ++frameIndex)
  {
    // Create binary labelmap volume
    vtkSmartPointer<vtkOrientedImageData> currentBinaryLabelmap = nullptr;

    if (numberOfSegments == 0 && isExtentValid)
    {
      // No segment metadata. We are loading from a plain volume (not seg.nrrd).

      currentBinaryLabelmap = vtkSmartPointer<vtkOrientedImageData>::New();
      extractComponents->SetComponents(frameIndex);
      padder->SetOutputWholeExtent(imageExtentInFile);
      padder->Update();
      currentBinaryLabelmap->ShallowCopy(padder->GetOutput());
      currentBinaryLabelmap->SetImageToWorldMatrix(imageToWorldMatrix.GetPointer());

      double scalarRange[2] = { 0 };
      currentBinaryLabelmap->GetScalarRange(scalarRange);
      vtkIdType lowLabel = (vtkIdType)(floor(scalarRange[0]));
      vtkIdType highLabel = (vtkIdType)(ceil(scalarRange[1]));

      vtkSmartPointer<vtkImageAccumulate> imageAccumulate = vtkSmartPointer<vtkImageAccumulate>::New();
      imageAccumulate->SetInputData(currentBinaryLabelmap);
      imageAccumulate->IgnoreZeroOn(); // Ignore background value
      imageAccumulate->SetComponentExtent(lowLabel, highLabel, 0, 0, 0, 0);
      imageAccumulate->SetComponentOrigin(0, 0, 0);
      imageAccumulate->SetComponentSpacing(1, 1, 1);
      imageAccumulate->Update();

      std::vector<vtkIdType> labelValues;
      for (vtkIdType labelValue = lowLabel; labelValue <= highLabel; ++labelValue)
      {
        if (labelValue == 0)
        {
          continue;
        }
        double frequency = imageAccumulate->GetOutput()->GetPointData()->GetScalars()->GetTuple1(labelValue - lowLabel);
        if (frequency == 0.0)
        {
          continue;
        }
        vtkSmartPointer<vtkSegment> currentSegment = vtkSmartPointer<vtkSegment>::New();
        currentSegment->SetLabelValue(labelValue);
        currentSegment->AddRepresentation(vtkSegmentationConverter::GetBinaryLabelmapRepresentationName(), currentBinaryLabelmap);
        segments.push_back(currentSegment);
      }

      // Set segmentation geometry from loaded image
      std::string imageGeometryString = vtkSegmentationConverter::SerializeImageGeometry(currentBinaryLabelmap);
      segmentation->SetConversionParameter(vtkSegmentationConverter::GetReferenceImageGeometryParameterName(), imageGeometryString);
    }
    else
    {
      for (int segmentIndex : segmentIndexInLayer[frameIndex])
      {
        // Create segment
        vtkSmartPointer<vtkSegment> currentSegment = vtkSmartPointer<vtkSegment>::New();
        std::string segmentColor;
        if (this->GetSegmentMetaDataFromDicitionary(segmentColor, dictionary, segmentIndex, KEY_SEGMENT_COLOR))
        {
          double currentSegmentColor[3] = { 0.0, 0.0, 0.0 };
          this->GetSegmentColorFromString(currentSegmentColor, segmentColor);
          currentSegment->SetColor(currentSegmentColor);
        }
        else if (this->GetSegmentMetaDataFromDicitionary(segmentColor, dictionary, segmentIndex, "DefaultColor"))
        {
          double defaultSegmentColor[3] = { 0.0, 0.0, 0.0 };
          this->GetSegmentColorFromString(defaultSegmentColor, segmentColor);
          currentSegment->SetColor(defaultSegmentColor);
        }

        // Tags
        std::string segmentTags;
        if (this->GetSegmentMetaDataFromDicitionary(segmentTags, dictionary, segmentIndex, KEY_SEGMENT_TAGS))
        {
          this->SetSegmentTagsFromString(currentSegment, segmentTags);
        }

        // NameAutoGenerated
        std::string nameAutoGenerated;
        if (this->GetSegmentMetaDataFromDicitionary(nameAutoGenerated, dictionary, segmentIndex, KEY_SEGMENT_NAME_AUTO_GENERATED))
        {
          currentSegment->SetNameAutoGenerated(!strcmp(nameAutoGenerated.c_str(), "1"));
        }

        // ColorAutoGenerated
        std::string colorAutoGenerated;
        if (this->GetSegmentMetaDataFromDicitionary(colorAutoGenerated, dictionary, segmentIndex, KEY_SEGMENT_COLOR_AUTO_GENERATED))
        {
          currentSegment->SetColorAutoGenerated(!strcmp(colorAutoGenerated.c_str(), "1"));
        }

        // Label value
        std::string labelValue;
        if (this->GetSegmentMetaDataFromDicitionary(labelValue, dictionary, segmentIndex, KEY_SEGMENT_LABEL_VALUE))
        {
          currentSegment->SetLabelValue(vtkVariant(labelValue).ToInt());
        }

        if (currentBinaryLabelmap == nullptr)
        {
          currentBinaryLabelmap = vtkSmartPointer<vtkOrientedImageData>::New();

          // Get layer extent: union of all segment extents in this layer
          int layerExtent[6] = { 0, -1, 0, -1, 0, -1 };
          for (int segmentIndexForExtent : segmentIndexInLayer[frameIndex])
          {
            int currentSegmentExtent[6] = { 0, -1, 0, -1, 0, -1 };
            std::string currentExtentString;
            if (this->GetSegmentMetaDataFromDicitionary(currentExtentString, dictionary, segmentIndexForExtent, KEY_SEGMENT_EXTENT))
            {
              GetImageExtentFromString(currentSegmentExtent, currentExtentString);
              vtkMRMLSegmentationStorageNode::AddToExtent(layerExtent, currentSegmentExtent);
            }
            else
            {
              // Extent is not specified, we don't know the extent of this layer, so use the entire region
              vtkWarningToMessageCollectionMacro(
                this->GetUserMessages(), "vtkMRMLSegmentationStorageNode::ReadBinaryLabelmapRepresentation", "Segment extent is missing for segment " << segmentIndexForExtent);
              for (int i = 0; i < 6; i++)
              {
                layerExtent[i] = imageExtentInFile[i];
              }
              break;
            }
          }
          for (int i = 0; i < 3; i++)
          {
            layerExtent[i * 2] += referenceImageExtentOffset[i];
            layerExtent[i * 2 + 1] += referenceImageExtentOffset[i];
          }
          // Copy with clipping to specified extent
          if (layerExtent[0] <= layerExtent[1]    //
              && layerExtent[2] <= layerExtent[3] //
              && layerExtent[4] <= layerExtent[5])
          {
            // non-empty segment
            extractComponents->SetComponents(frameIndex);
            padder->SetOutputWholeExtent(layerExtent);
            padder->Update();
            currentBinaryLabelmap->DeepCopy(padder->GetOutput());
          }
          else
          {
            // empty segment
            for (int i = 0; i < 3; ++i)
            {
              layerExtent[2 * i] = 0;
              layerExtent[2 * i + 1] = -1;
            }
            currentBinaryLabelmap->SetExtent(layerExtent);
            currentBinaryLabelmap->AllocateScalars(VTK_UNSIGNED_CHAR, 1);
          }
          currentBinaryLabelmap->SetImageToWorldMatrix(imageToWorldMatrix.GetPointer());
        }

        // Set loaded binary labelmap to segment
        currentSegment->AddRepresentation(vtkSegmentationConverter::GetSegmentationBinaryLabelmapRepresentationName(), currentBinaryLabelmap);
        segments[segmentIndex] = currentSegment;
      }
    }
  }

  vtkMRMLColorTableNode* exportColorTableNode = segmentationNode->GetLabelmapConversionColorTableNode();

  // Add the created segments to the segmentation
  for (int segmentIndex = 0; segmentIndex < static_cast<int>(segments.size()); ++segmentIndex)
  {
    vtkSegment* currentSegment = segments[segmentIndex];
    if (!currentSegment)
    {
      vtkErrorToMessageCollectionMacro(this->GetUserMessages(), "vtkMRMLSegmentationStorageNode::ReadBinaryLabelmapRepresentation", "Could not find segment" << segmentIndex);
      continue;
    }

    std::string currentSegmentID;
    if (numberOfSegments != 0)
    {
      // ID
      if (!this->GetSegmentMetaDataFromDicitionary(currentSegmentID, dictionary, segmentIndex, KEY_SEGMENT_ID))
      {
        // No segment ID is specified, which may mean that it is an empty segmentation.
        // We consider a segmentation empty if it has only one scalar component that is empty.
        if (numberOfFrames == 1)
        {
          extractComponents->SetComponents(segmentIndex);
          extractComponents->Update();
          vtkImageData* labelmap = extractComponents->GetOutput();
          double* scalarRange = labelmap->GetScalarRange();
          if (scalarRange[0] >= scalarRange[1])
          {
            // Segmentation contains a single blank segment without segment ID,
            // which means that it is an empty segmentation.
            // It may still contain valuable metadata, but we don't create any segments.
            break;
          }
        }

        currentSegmentID = segmentation->GenerateUniqueSegmentID("Segment");
        vtkWarningToMessageCollectionMacro(this->GetUserMessages(),
                                           "vtkMRMLSegmentationStorageNode::ReadBinaryLabelmapRepresentation",
                                           "Segment ID is missing for segment " << segmentIndex << " adding segment with ID: " << currentSegmentID);
      }

      // Name
      std::string segmentName;
      if (this->GetSegmentMetaDataFromDicitionary(segmentName, dictionary, segmentIndex, KEY_SEGMENT_NAME))
      {
        currentSegment->SetName(segmentName.c_str());
      }
      else
      {
        vtkWarningToMessageCollectionMacro(
          this->GetUserMessages(), "vtkMRMLSegmentationStorageNode::ReadBinaryLabelmapRepresentation", "Segment name is missing for segment " << segmentIndex);
        currentSegment->SetName(currentSegmentID.c_str());
      }
    }
    else
    {
      std::stringstream segmentNameSS;
      segmentNameSS << "Segment_" << currentSegment->GetLabelValue();
      std::string segmentName = segmentNameSS.str();
      if (exportColorTableNode)
      {
        const char* colorName = exportColorTableNode->GetColorName(currentSegment->GetLabelValue());
        if (colorName)
        {
          segmentName = colorName;
        }

        double color[4] = { 0.0, 0.0, 0.0, 1.0 };
        exportColorTableNode->GetColor(currentSegment->GetLabelValue(), color);
        currentSegment->SetColor(color);
      }
      else
      {
        vtkMRMLSegmentationDisplayNode* displayNode = vtkMRMLSegmentationDisplayNode::SafeDownCast(segmentationNode->GetDisplayNode());
        if (displayNode)
        {
          double color[4] = { 0.0, 0.0, 0.0, 1.0 };
          displayNode->GenerateSegmentColor(color);
          currentSegment->SetColor(color);
        }
      }
      currentSegmentID = segmentation->GenerateUniqueSegmentID(segmentName);
      currentSegment->SetName(currentSegmentID.c_str());
    }
    segmentation->AddSegment(currentSegment, currentSegmentID);
  }

  // Create contained representations now that all the data is loaded
  this->CreateRepresentationsBySerializedNames(segmentation, containedRepresentationNames);

  return 1;
}

//----------------------------------------------------------------------------
void vtkMRMLSegmentationStorageNode::AddToExtent(int extent[6], int extentToAdd[6])
{
  if (extentToAdd[0] > extentToAdd[1]    //
      || extentToAdd[2] > extentToAdd[3] //
      || extentToAdd[4] > extentToAdd[5])
  {
    // extentToAdd is empty, no need to change the extent
    return;
  }
  if (extent[0] > extent[1]    //
      || extent[2] > extent[3] //
      || extent[4] > extent[5])
  {
    // initial extent is empty, use the added extent
    for (int i = 0; i < 6; i++)
    {
      extent[i] = extentToAdd[i];
    }
  }
  else
  {
    // add extentToAdd to extent
    for (int i = 0; i < 3; i++)
    {
      if (extentToAdd[i * 2] < extent[i * 2])
      {
        extent[i * 2] = extentToAdd[i * 2];
      }
      if (extentToAdd[i * 2 + 1] > extent[i * 2 + 1])
      {
        extent[i * 2 + 1] = extentToAdd[i * 2 + 1];
      }
    }
  }
}

//----------------------------------------------------------------------------
int vtkMRMLSegmentationStorageNode::ReadPolyDataRepresentation(vtkMRMLSegmentationNode* segmentationNode, std::string path)
{
  if (!vtksys::SystemTools::FileExists(path.c_str()))
  {
    vtkErrorToMessageCollectionMacro(
      this->GetUserMessages(), "vtkMRMLSegmentationStorageNode::ReadPolyDataRepresentation", "ReadPolyDataRepresentation: Input file " << path << " does not exist.");
    return 0;
  }

  // Set up output segmentation
  if (!segmentationNode)
  {
    vtkErrorToMessageCollectionMacro(this->GetUserMessages(), "vtkMRMLSegmentationStorageNode::ReadPolyDataRepresentation", "Output segmentation must exist.");
    return 0;
  }

  // Read multiblock dataset from disk
  vtkSmartPointer<vtkXMLMultiBlockDataReader> reader = vtkSmartPointer<vtkXMLMultiBlockDataReader>::New();
  reader->SetFileName(path.c_str());
  reader->Update();
  vtkMultiBlockDataSet* multiBlockDataset = vtkMultiBlockDataSet::SafeDownCast(reader->GetOutput());
  if (!multiBlockDataset || multiBlockDataset->GetNumberOfBlocks() == 0)
  {
    vtkErrorToMessageCollectionMacro(
      this->GetUserMessages(), "vtkMRMLSegmentationStorageNode::ReadPolyDataRepresentation", "ReadPolyDataRepresentation: Failed to read file " << path);
    return 0;
  }

  MRMLNodeModifyBlocker blocker(segmentationNode);

  // Make sure there is a valid segmentation object in the node
  vtkSegmentation* segmentation = segmentationNode->GetSegmentation();
  if (!segmentation)
  {
    vtkNew<vtkSegmentation> newSegmentation;
    segmentation = newSegmentation;
    segmentationNode->SetAndObserveSegmentation(newSegmentation);
  }

  // Clean out the segmentation before adding the new segments
  std::set<std::string> originalRepresentationNames;
  if (segmentation->GetNumberOfSegments() > 0)
  {
    segmentation->GetAvailableRepresentationNames(originalRepresentationNames);
    segmentation->RemoveAllSegments();
  }

  // Add all files to storage node (multiblock dataset writes segments to individual files in a separate folder)
  this->AddPolyDataFileNames(path, segmentation);

  // Read segment poly data
  std::string sourceRepresentationName;
  std::string containedRepresentationNames;
  std::string conversionParameters;
  for (unsigned int blockIndex = 0; blockIndex < multiBlockDataset->GetNumberOfBlocks(); ++blockIndex)
  {
    // Get poly data representation
    vtkPolyData* currentPolyData = vtkPolyData::SafeDownCast(multiBlockDataset->GetBlock(blockIndex));
    if (!currentPolyData)
    {
      vtkErrorToMessageCollectionMacro(this->GetUserMessages(), "vtkMRMLSegmentationStorageNode::ReadPolyDataRepresentation", "Could not read block " << blockIndex);
      continue;
    }

    // Set source representation if it has not been set yet
    // (there is no global place to store it, but every segment field data contains a copy of it)
    if (sourceRepresentationName.empty())
    {
      vtkStringArray* sourceRepresentationArray =
        vtkStringArray::SafeDownCast(currentPolyData->GetFieldData()->GetAbstractArray(GetSegmentationMetaDataKey(KEY_SEGMENTATION_SOURCE_REPRESENTATION).c_str()));
      if (!sourceRepresentationArray)
      {
        vtkErrorToMessageCollectionMacro(
          this->GetUserMessages(), "vtkMRMLSegmentationStorageNode::ReadPolyDataRepresentation", "Unable to find source representation for segmentation in file " << path);
        return 0;
      }
      sourceRepresentationName = sourceRepresentationArray->GetValue(0);
      segmentation->SetSourceRepresentationName(sourceRepresentationName.c_str());
    }
    // Read conversion parameters (stored in each segment file, but need to set only once)
    if (conversionParameters.empty() //
        && currentPolyData->GetFieldData()->GetAbstractArray(GetSegmentationMetaDataKey(KEY_SEGMENTATION_CONVERSION_PARAMETERS).c_str()))
    {
      vtkStringArray* conversionParametersArray =
        vtkStringArray::SafeDownCast(currentPolyData->GetFieldData()->GetAbstractArray(GetSegmentationMetaDataKey(KEY_SEGMENTATION_CONVERSION_PARAMETERS).c_str()));
      conversionParameters = conversionParametersArray->GetValue(0);
      segmentation->DeserializeConversionParameters(conversionParameters);
    }
    // Read contained representation names
    if (containedRepresentationNames.empty() //
        && currentPolyData->GetFieldData()->GetAbstractArray(GetSegmentationMetaDataKey(KEY_SEGMENTATION_CONTAINED_REPRESENTATION_NAMES).c_str()))
    {
      containedRepresentationNames =
        vtkStringArray::SafeDownCast(currentPolyData->GetFieldData()->GetAbstractArray(GetSegmentationMetaDataKey(KEY_SEGMENTATION_CONTAINED_REPRESENTATION_NAMES).c_str()))
          ->GetValue(0);
    }

    // Create segment
    vtkSmartPointer<vtkSegment> currentSegment = vtkSmartPointer<vtkSegment>::New();
    currentSegment->AddRepresentation(segmentation->GetSourceRepresentationName(), currentPolyData);

    // Set segment properties

    std::string currentSegmentID;
    vtkStringArray* idArray = vtkStringArray::SafeDownCast(currentPolyData->GetFieldData()->GetAbstractArray(GetSegmentMetaDataKey(SINGLE_SEGMENT_INDEX, KEY_SEGMENT_ID).c_str()));
    if (idArray && idArray->GetNumberOfValues() > 0)
    {
      currentSegmentID = idArray->GetValue(0);
    }
    else
    {
      vtkWarningToMessageCollectionMacro(this->GetUserMessages(),
                                         "vtkMRMLSegmentationStorageNode::ReadPolyDataRepresentation",
                                         "Segment ID property not found when reading segment " << blockIndex << " from file " << path);
    }

    std::string currentSegmentName;
    vtkStringArray* nameArray =
      vtkStringArray::SafeDownCast(currentPolyData->GetFieldData()->GetAbstractArray(GetSegmentMetaDataKey(SINGLE_SEGMENT_INDEX, KEY_SEGMENT_NAME).c_str()));
    if (nameArray && nameArray->GetNumberOfValues() > 0)
    {
      currentSegmentName = nameArray->GetValue(0);
    }
    else
    {
      vtkWarningToMessageCollectionMacro(this->GetUserMessages(),
                                         "vtkMRMLSegmentationStorageNode::ReadPolyDataRepresentation",
                                         "Segment Name property not found when reading segment " << blockIndex << " from file " << path);
      std::stringstream ssCurrentSegmentName;
      ssCurrentSegmentName << "Segment " << blockIndex;
      currentSegmentName = ssCurrentSegmentName.str();
    }
    currentSegment->SetName(currentSegmentName.c_str());

    double color[3] = { 1.0, 0.0, 0.0 };
    vtkDoubleArray* colorArray = vtkDoubleArray::SafeDownCast(currentPolyData->GetFieldData()->GetArray(GetSegmentMetaDataKey(SINGLE_SEGMENT_INDEX, KEY_SEGMENT_COLOR).c_str()));
    if (!colorArray || colorArray->GetNumberOfTuples() == 0)
    {
      // Backwards compatibility
      colorArray = vtkDoubleArray::SafeDownCast(currentPolyData->GetFieldData()->GetArray(GetSegmentMetaDataKey(SINGLE_SEGMENT_INDEX, "DefaultColor").c_str()));
    }
    if (colorArray && colorArray->GetNumberOfTuples() > 0 && colorArray->GetNumberOfComponents() == 3)
    {
      colorArray->GetTuple(0, color);
    }
    else
    {
      vtkWarningToMessageCollectionMacro(this->GetUserMessages(),
                                         "vtkMRMLSegmentationStorageNode::ReadPolyDataRepresentation",
                                         "Segment color property not found when reading segment " << blockIndex << " from file " << path);
    }
    currentSegment->SetColor(color);

    // Tags
    vtkStringArray* tagsArray =
      vtkStringArray::SafeDownCast(currentPolyData->GetFieldData()->GetAbstractArray(GetSegmentMetaDataKey(SINGLE_SEGMENT_INDEX, KEY_SEGMENT_TAGS).c_str()));
    if (tagsArray)
    {
      std::stringstream tagsValueStr(tagsArray->GetValue(0));
      std::string tagValue;
      while (getline(tagsValueStr, tagValue, LIST_SEPARATOR))
      {
        size_t colonPosition = tagValue.find(":");
        if (colonPosition == std::string::npos)
        {
          continue;
        }
        currentSegment->SetTag(tagValue.substr(0, colonPosition), tagValue.substr(colonPosition + 1));
      }
    }

    // NameAutoGenerated
    vtkStringArray* nameAutoGeneratedArray =
      vtkStringArray::SafeDownCast(currentPolyData->GetFieldData()->GetAbstractArray(GetSegmentMetaDataKey(SINGLE_SEGMENT_INDEX, KEY_SEGMENT_NAME_AUTO_GENERATED).c_str()));
    if (nameAutoGeneratedArray && nameAutoGeneratedArray->GetNumberOfValues() > 0)
    {
      currentSegment->SetNameAutoGenerated(!nameAutoGeneratedArray->GetValue(0).compare("1"));
    }

    // ColorAutoGenerated
    vtkStringArray* colorAutoGeneratedArray =
      vtkStringArray::SafeDownCast(currentPolyData->GetFieldData()->GetAbstractArray(GetSegmentMetaDataKey(SINGLE_SEGMENT_INDEX, KEY_SEGMENT_COLOR_AUTO_GENERATED).c_str()));
    if (colorAutoGeneratedArray && colorAutoGeneratedArray->GetNumberOfValues() > 0)
    {
      currentSegment->SetColorAutoGenerated(!colorAutoGeneratedArray->GetValue(0).compare("1"));
    }

    // Add segment to segmentation
    segmentation->AddSegment(currentSegment, currentSegmentID);
  }

  // Create contained representations now that all the data is loaded
  this->CreateRepresentationsBySerializedNames(segmentation, containedRepresentationNames);

  return 1;
}

//----------------------------------------------------------------------------
int vtkMRMLSegmentationStorageNode::WriteDataInternal(vtkMRMLNode* refNode)
{
  std::string fullName = this->GetFullNameFromFileName();
  if (fullName.empty())
  {
    vtkErrorToMessageCollectionMacro(this->GetUserMessages(), "vtkMRMLSegmentationStorageNode::WriteDataInternal", "File name not specified");
    return 0;
  }

  vtkMRMLSegmentationNode* segmentationNode = vtkMRMLSegmentationNode::SafeDownCast(refNode);
  if (segmentationNode == nullptr)
  {
    vtkErrorToMessageCollectionMacro(this->GetUserMessages(), "vtkMRMLSegmentationStorageNode::WriteDataInternal", "Segmentation node expected. Unable to write node to file.");
    return 0;
  }

  if (segmentationNode->GetSegmentation() == nullptr)
  {
    vtkErrorToMessageCollectionMacro(
      this->GetUserMessages(), "vtkMRMLSegmentationStorageNode::WriteDataInternal", "Segmentation node does not contain segmentation object. Unable to write node to file.");
    return 0;
  }

  // Write only source representation
  if (segmentationNode->GetSegmentation()->IsSourceRepresentationImageData())
  {
    return this->WriteBinaryLabelmapRepresentation(segmentationNode, fullName);
  }
  else if (segmentationNode->GetSegmentation()->IsSourceRepresentationPolyData())
  {
    return this->WritePolyDataRepresentation(segmentationNode, fullName);
  }

  vtkErrorToMessageCollectionMacro(this->GetUserMessages(),
                                   "vtkMRMLSegmentationStorageNode::WriteDataInternal",
                                   "Segmentation source representation " << segmentationNode->GetSegmentation()->GetSourceRepresentationName() << " cannot be written to file");
  return 0;
}

//----------------------------------------------------------------------------
int vtkMRMLSegmentationStorageNode::WriteBinaryLabelmapRepresentation(vtkMRMLSegmentationNode* segmentationNode, std::string fullName)
{
  if (!segmentationNode)
  {
    vtkErrorToMessageCollectionMacro(this->GetUserMessages(), "vtkMRMLSegmentationStorageNode::WriteBinaryLabelmapRepresentation", "Invalid segmentation to write to disk");
    return 0;
  }
  vtkSegmentation* segmentation = segmentationNode->GetSegmentation();
  segmentation->CollapseBinaryLabelmaps(false);

  // Get and check source representation
  if (!segmentationNode->GetSegmentation()->IsSourceRepresentationImageData())
  {
    vtkErrorToMessageCollectionMacro(
      this->GetUserMessages(), "vtkMRMLSegmentationStorageNode::WriteBinaryLabelmapRepresentation", "Invalid source representation to write as image data");
    return 0;
  }

  vtkIdType scalarType = VTK_UNSIGNED_CHAR;
  vtkIdType scalarSize = 0;
  std::vector<std::string> segmentIDs;
  segmentation->GetSegmentIDs(segmentIDs);
  for (std::vector<std::string>::const_iterator segmentIdIt = segmentIDs.begin(); segmentIdIt != segmentIDs.end(); ++segmentIdIt)
  {
    std::string currentSegmentID = *segmentIdIt;
    vtkSegment* currentSegment = segmentation->GetSegment(*segmentIdIt);
    vtkSmartPointer<vtkOrientedImageData> currentBinaryLabelmap =
      vtkOrientedImageData::SafeDownCast(currentSegment->GetRepresentation(segmentationNode->GetSegmentation()->GetSourceRepresentationName()));
    if (currentBinaryLabelmap->GetScalarSize() > scalarSize)
    {
      scalarSize = currentBinaryLabelmap->GetScalarSize();
      scalarType = currentBinaryLabelmap->GetScalarType();
    }
  }

  // Determine shared labelmap dimensions and properties
  vtkSmartPointer<vtkOrientedImageData> commonGeometryImage = vtkSmartPointer<vtkOrientedImageData>::New();
  int commonGeometryExtent[6] = { 0, -1, 0, -1, 0, -1 };
  if (segmentation->GetNumberOfSegments() > 0)
  {
    std::string commonGeometryString = segmentation->DetermineCommonLabelmapGeometry(
      this->CropToMinimumExtent ? vtkSegmentation::EXTENT_UNION_OF_EFFECTIVE_SEGMENTS : vtkSegmentation::EXTENT_UNION_OF_EFFECTIVE_SEGMENTS_AND_REFERENCE_GEOMETRY);
    if (!commonGeometryString.empty())
    {
      vtkSegmentationConverter::DeserializeImageGeometry(commonGeometryString, commonGeometryImage, true, scalarType, 1);
      commonGeometryImage->GetExtent(commonGeometryExtent);
    }
  }
  if (commonGeometryExtent[0] > commonGeometryExtent[1]    //
      || commonGeometryExtent[2] > commonGeometryExtent[3] //
      || commonGeometryExtent[4] > commonGeometryExtent[5])
  {
    // common image is empty, which cannot be written to image file
    // change it to a 1x1x1 image instead
    commonGeometryExtent[0] = 0;
    commonGeometryExtent[1] = 0;
    commonGeometryExtent[2] = 0;
    commonGeometryExtent[3] = 0;
    commonGeometryExtent[4] = 0;
    commonGeometryExtent[5] = 0;
    commonGeometryImage->SetExtent(commonGeometryExtent);
    commonGeometryImage->AllocateScalars(scalarType, 1);
  }
  vtkOrientedImageDataResample::FillImage(commonGeometryImage, 0);

  vtkNew<vtkTeemNRRDWriter> writer;
  writer->SetFileName(fullName.c_str());
  writer->SetUseCompression(this->GetUseCompression());
  writer->SetSpace(nrrdSpaceLeftPosteriorSuperior);
  writer->SetMeasurementFrameMatrix(nullptr);

  // Create metadata dictionary

  // Save extent start of common geometry image so that we can restore original extents when reading from file
  // writer->SetAttribute(GetSegmentationMetaDataKey(KEY_SEGMENTATION_EXTENT).c_str(), GetImageExtentAsString(commonGeometryImage));
  int referenceImageExtentOffset[3] = { commonGeometryExtent[0], commonGeometryExtent[2], commonGeometryExtent[4] };
  std::stringstream ssReferenceImageExtentOffset;
  ssReferenceImageExtentOffset << referenceImageExtentOffset[0] << " " << referenceImageExtentOffset[1] << " " << referenceImageExtentOffset[2];
  writer->SetAttribute(GetSegmentationMetaDataKey(KEY_SEGMENTATION_REFERENCE_IMAGE_EXTENT_OFFSET).c_str(), ssReferenceImageExtentOffset.str());

  vtkNew<vtkMatrix4x4> rasToIjk;
  commonGeometryImage->GetWorldToImageMatrix(rasToIjk.GetPointer());
  // Compensate for the extent shift in the image origin:
  vtkNew<vtkMatrix4x4> ijkToFileIjk;
  ijkToFileIjk->SetElement(0, 3, -referenceImageExtentOffset[0]);
  ijkToFileIjk->SetElement(1, 3, -referenceImageExtentOffset[1]);
  ijkToFileIjk->SetElement(2, 3, -referenceImageExtentOffset[2]);
  vtkNew<vtkMatrix4x4> rasToFileIjk;
  vtkMatrix4x4::Multiply4x4(ijkToFileIjk.GetPointer(), rasToIjk.GetPointer(), rasToFileIjk.GetPointer());
  vtkNew<vtkMatrix4x4> fileIjkToRas;
  vtkMatrix4x4::Invert(rasToFileIjk.GetPointer(), fileIjkToRas.GetPointer());
  writer->SetIJKToRASMatrix(fileIjkToRas.GetPointer());

  // Save source representation name
  writer->SetAttribute(GetSegmentationMetaDataKey(KEY_SEGMENTATION_SOURCE_REPRESENTATION).c_str(), segmentationNode->GetSegmentation()->GetSourceRepresentationName());
  // Save conversion parameters
  std::string conversionParameters = segmentation->SerializeAllConversionParameters();
  writer->SetAttribute(GetSegmentationMetaDataKey(KEY_SEGMENTATION_CONVERSION_PARAMETERS).c_str(), conversionParameters);
  // Save created representation names so that they are re-created when loading
  std::string containedRepresentationNames = this->SerializeContainedRepresentationNames(segmentation);
  writer->SetAttribute(GetSegmentationMetaDataKey(KEY_SEGMENTATION_CONTAINED_REPRESENTATION_NAMES).c_str(), containedRepresentationNames);

  vtkNew<vtkImageAppendComponents> appender;

  unsigned int layerIndex = 0;
  std::map<vtkDataObject*, int> labelmapLayers;

  // Dimensions of the output 4D NRRD file: (i, j, k, segment)
  unsigned int segmentIndex = 0;
  for (std::vector<std::string>::const_iterator segmentIdIt = segmentIDs.begin(); segmentIdIt != segmentIDs.end(); ++segmentIdIt, ++segmentIndex)
  {
    std::string currentSegmentID = *segmentIdIt;
    vtkSegment* currentSegment = segmentation->GetSegment(*segmentIdIt);

    // Get source representation from segment
    vtkSmartPointer<vtkOrientedImageData> currentBinaryLabelmap =
      vtkOrientedImageData::SafeDownCast(currentSegment->GetRepresentation(segmentationNode->GetSegmentation()->GetSourceRepresentationName()));
    if (!currentBinaryLabelmap)
    {
      vtkErrorToMessageCollectionMacro(
        this->GetUserMessages(), "vtkMRMLSegmentationStorageNode::WriteBinaryLabelmapRepresentation", "Failed to retrieve source representation from segment " << currentSegmentID);
      continue;
    }

    int currentBinaryLabelmapExtent[6] = { 0, -1, 0, -1, 0, -1 };
    currentBinaryLabelmap->GetExtent(currentBinaryLabelmapExtent);
    if (currentBinaryLabelmapExtent[0] <= currentBinaryLabelmapExtent[1]    //
        && currentBinaryLabelmapExtent[2] <= currentBinaryLabelmapExtent[3] //
        && currentBinaryLabelmapExtent[4] <= currentBinaryLabelmapExtent[5])
    {
      // There is a valid labelmap

      // Get transformed extents of the segment in the common labelmap geometry
      vtkNew<vtkTransform> currentBinaryLabelmapToCommonGeometryImageTransform;
      vtkOrientedImageDataResample::GetTransformBetweenOrientedImages(currentBinaryLabelmap, commonGeometryImage, currentBinaryLabelmapToCommonGeometryImageTransform.GetPointer());
      int currentBinaryLabelmapExtentInCommonGeometryImageFrame[6] = { 0, -1, 0, -1, 0, -1 };
      vtkOrientedImageDataResample::TransformExtent(
        currentBinaryLabelmapExtent, currentBinaryLabelmapToCommonGeometryImageTransform.GetPointer(), currentBinaryLabelmapExtentInCommonGeometryImageFrame);
      for (int i = 0; i < 3; i++)
      {
        currentBinaryLabelmapExtent[i * 2] = std::max(currentBinaryLabelmapExtentInCommonGeometryImageFrame[i * 2], commonGeometryExtent[i * 2]);
        currentBinaryLabelmapExtent[i * 2 + 1] = std::min(currentBinaryLabelmapExtentInCommonGeometryImageFrame[i * 2 + 1], commonGeometryExtent[i * 2 + 1]);
      }
      // TODO: maybe calculate effective extent to make sure the data is as compact as possible? (saving may be a good time to make segments more compact)

      // Pad/resample current binary labelmap representation to common geometry
      vtkSmartPointer<vtkOrientedImageData> resampledCurrentBinaryLabelmap = vtkSmartPointer<vtkOrientedImageData>::New();
      bool success = vtkOrientedImageDataResample::ResampleOrientedImageToReferenceOrientedImage(currentBinaryLabelmap, commonGeometryImage, resampledCurrentBinaryLabelmap);
      if (!success)
      {
        vtkWarningToMessageCollectionMacro(this->GetUserMessages(),
                                           "vtkMRMLSegmentationStorageNode::WriteBinaryLabelmapRepresentation",
                                           "Segment " << currentSegmentID << " cannot be resampled to common geometry");
        continue;
      }

      // currentBinaryLabelmap smart pointer will keep the temporary labelmap valid until it is needed
      currentBinaryLabelmap = resampledCurrentBinaryLabelmap;
      if (currentBinaryLabelmap->GetScalarType() != scalarType)
      {
        vtkNew<vtkImageCast> castFilter;
        castFilter->SetInputData(resampledCurrentBinaryLabelmap);
        castFilter->SetOutputScalarType(scalarType);
        castFilter->Update();
        currentBinaryLabelmap->ShallowCopy(castFilter->GetOutput());
      }
    }
    else
    {
      // empty segment, use the commonGeometryImage (filled with 0)
      currentBinaryLabelmap = commonGeometryImage;
    }

    // Set metadata for current segment
    writer->SetAttribute(GetSegmentMetaDataKey(segmentIndex, KEY_SEGMENT_ID).c_str(), currentSegmentID);
    writer->SetAttribute(GetSegmentMetaDataKey(segmentIndex, KEY_SEGMENT_NAME).c_str(), currentSegment->GetName());
    writer->SetAttribute(GetSegmentMetaDataKey(segmentIndex, KEY_SEGMENT_COLOR).c_str(), GetSegmentColorAsString(segmentationNode, currentSegmentID));
    writer->SetAttribute(GetSegmentMetaDataKey(segmentIndex, KEY_SEGMENT_NAME_AUTO_GENERATED).c_str(), (currentSegment->GetNameAutoGenerated() ? "1" : "0"));
    writer->SetAttribute(GetSegmentMetaDataKey(segmentIndex, KEY_SEGMENT_COLOR_AUTO_GENERATED).c_str(), (currentSegment->GetColorAutoGenerated() ? "1" : "0"));
    // Save the geometry relative to the current image (so that the extent in the file describe the extent of the segment in the
    // saved image buffer)
    for (int i = 0; i < 3; i++)
    {
      currentBinaryLabelmapExtent[i * 2] -= referenceImageExtentOffset[i];
      currentBinaryLabelmapExtent[i * 2 + 1] -= referenceImageExtentOffset[i];
    }
    writer->SetAttribute(GetSegmentMetaDataKey(segmentIndex, KEY_SEGMENT_EXTENT).c_str(), GetImageExtentAsString(currentBinaryLabelmapExtent));
    writer->SetAttribute(GetSegmentMetaDataKey(segmentIndex, KEY_SEGMENT_TAGS).c_str(), GetSegmentTagsAsString(currentSegment));
    std::stringstream labelValueSS;
    labelValueSS << currentSegment->GetLabelValue();
    writer->SetAttribute(GetSegmentMetaDataKey(segmentIndex, KEY_SEGMENT_LABEL_VALUE).c_str(), labelValueSS.str());

    vtkDataObject* originalRepresentation = currentSegment->GetRepresentation(segmentationNode->GetSegmentation()->GetSourceRepresentationName());
    if (labelmapLayers.find(originalRepresentation) == labelmapLayers.end())
    {
      labelmapLayers[originalRepresentation] = layerIndex;
      appender->AddInputData(currentBinaryLabelmap);
      ++layerIndex;
    }
    unsigned int layer = labelmapLayers[originalRepresentation];
    std::stringstream layerIndexSS;
    layerIndexSS << layer;
    writer->SetAttribute(GetSegmentMetaDataKey(segmentIndex, KEY_SEGMENT_LAYER).c_str(), layerIndexSS.str());

  } // For each segment

  this->GetUserMessages()->SetObservedObject(writer);
  if (segmentationNode->GetSegmentation()->GetNumberOfSegments() > 0)
  {
    appender->Update();
    writer->SetInputConnection(appender->GetOutputPort());
    writer->SetVectorAxisKind(nrrdKindList);
  }
  else
  {
    // If there are no segments, we still write the data so that we can store
    // various metadata fields.
    writer->SetInputData(commonGeometryImage);
  }

  writer->Write();
  this->GetUserMessages()->SetObservedObject(nullptr);
  int writeSuccess = true;
  if (writer->GetWriteError())
  {
    vtkErrorToMessageCollectionMacro(this->GetUserMessages(),
                                     "vtkMRMLSegmentationStorageNode::WriteBinaryLabelmapRepresentation",
                                     "Error writing NRRD file " << (writer->GetFileName() == nullptr ? "null" : writer->GetFileName()));
    writeSuccess = false;
  }

  return (writeSuccess ? 1 : 0);
}

//----------------------------------------------------------------------------
int vtkMRMLSegmentationStorageNode::WritePolyDataRepresentation(vtkMRMLSegmentationNode* segmentationNode, std::string path)
{
  if (!segmentationNode || segmentationNode->GetSegmentation()->GetNumberOfSegments() == 0)
  {
    vtkErrorToMessageCollectionMacro(this->GetUserMessages(), "vtkMRMLSegmentationStorageNode::WritePolyDataRepresentation", "Invalid segmentation to write to disk");
    return 0;
  }
  vtkSegmentation* segmentation = segmentationNode->GetSegmentation();

  // Get and check source representation
  if (!segmentationNode->GetSegmentation()->IsSourceRepresentationPolyData())
  {
    vtkErrorToMessageCollectionMacro(this->GetUserMessages(), "vtkMRMLSegmentationStorageNode::WritePolyDataRepresentation", "Invalid source representation to write as poly data");
    return 0;
  }

  // Initialize dataset to write
  vtkSmartPointer<vtkMultiBlockDataSet> multiBlockDataset = vtkSmartPointer<vtkMultiBlockDataSet>::New();
  multiBlockDataset->SetNumberOfBlocks(segmentation->GetNumberOfSegments());

  // Add segment polydata objects to dataset
  std::vector<std::string> segmentIDs;
  segmentation->GetSegmentIDs(segmentIDs);
  unsigned int segmentIndex = 0;
  for (std::vector<std::string>::const_iterator segmentIdIt = segmentIDs.begin(); segmentIdIt != segmentIDs.end(); ++segmentIdIt, ++segmentIndex)
  {
    std::string currentSegmentID = *segmentIdIt;
    vtkSegment* currentSegment = segmentation->GetSegment(*segmentIdIt);

    // Get source representation from segment
    vtkPolyData* currentPolyData = vtkPolyData::SafeDownCast(currentSegment->GetRepresentation(segmentationNode->GetSegmentation()->GetSourceRepresentationName()));
    if (!currentPolyData)
    {
      vtkErrorToMessageCollectionMacro(
        this->GetUserMessages(), "vtkMRMLSegmentationStorageNode::WritePolyDataRepresentation", "Failed to retrieve source representation from segment " << currentSegmentID);
      continue;
    }
    // Make temporary duplicate of the poly data so that adding the metadata does not cause invalidating the other
    // representations (which is done when the source representation is modified)
    vtkSmartPointer<vtkPolyData> currentPolyDataCopy = vtkSmartPointer<vtkPolyData>::New();
    currentPolyDataCopy->ShallowCopy(currentPolyData);

    // Set metadata for current segment

    // SourceRepresentation
    vtkSmartPointer<vtkStringArray> sourceRepresentationArray = vtkSmartPointer<vtkStringArray>::New();
    sourceRepresentationArray->SetNumberOfValues(1);
    sourceRepresentationArray->SetValue(0, segmentationNode->GetSegmentation()->GetSourceRepresentationName());
    sourceRepresentationArray->SetName(GetSegmentationMetaDataKey(KEY_SEGMENTATION_SOURCE_REPRESENTATION).c_str());
    currentPolyDataCopy->GetFieldData()->AddArray(sourceRepresentationArray);

    // ID
    vtkSmartPointer<vtkStringArray> idArray = vtkSmartPointer<vtkStringArray>::New();
    idArray->SetNumberOfValues(1);
    idArray->SetValue(0, currentSegmentID.c_str());
    idArray->SetName(GetSegmentMetaDataKey(SINGLE_SEGMENT_INDEX, KEY_SEGMENT_ID).c_str());
    currentPolyDataCopy->GetFieldData()->AddArray(idArray);

    // Name
    vtkSmartPointer<vtkStringArray> nameArray = vtkSmartPointer<vtkStringArray>::New();
    nameArray->SetNumberOfValues(1);
    nameArray->SetValue(0, currentSegment->GetName());
    nameArray->SetName(GetSegmentMetaDataKey(SINGLE_SEGMENT_INDEX, KEY_SEGMENT_NAME).c_str());
    currentPolyDataCopy->GetFieldData()->AddArray(nameArray);

    // Color
    vtkSmartPointer<vtkDoubleArray> colorArray = vtkSmartPointer<vtkDoubleArray>::New();
    colorArray->SetNumberOfComponents(3);
    colorArray->SetNumberOfTuples(1);
    colorArray->SetTuple(0, currentSegment->GetColor());
    colorArray->SetName(GetSegmentMetaDataKey(SINGLE_SEGMENT_INDEX, KEY_SEGMENT_COLOR).c_str());
    currentPolyDataCopy->GetFieldData()->AddArray(colorArray);

    // Tags
    std::map<std::string, std::string> tags;
    currentSegment->GetTags(tags);
    std::stringstream ssTags;
    std::map<std::string, std::string>::iterator tagIt;
    for (tagIt = tags.begin(); tagIt != tags.end(); ++tagIt)
    {
      ssTags << tagIt->first << ":" << tagIt->second << "|";
    }
    vtkSmartPointer<vtkStringArray> tagsArray = vtkSmartPointer<vtkStringArray>::New();
    tagsArray->SetNumberOfValues(1);
    tagsArray->SetValue(0, ssTags.str().c_str());
    tagsArray->SetName(GetSegmentMetaDataKey(SINGLE_SEGMENT_INDEX, KEY_SEGMENT_TAGS).c_str());
    currentPolyDataCopy->GetFieldData()->AddArray(tagsArray);

    // NameAutoGenerated
    vtkSmartPointer<vtkStringArray> nameAutoGeneratedArray = vtkSmartPointer<vtkStringArray>::New();
    nameAutoGeneratedArray->SetNumberOfValues(1);
    nameAutoGeneratedArray->SetValue(0, (currentSegment->GetNameAutoGenerated() ? "1" : "0"));
    nameAutoGeneratedArray->SetName(GetSegmentMetaDataKey(SINGLE_SEGMENT_INDEX, KEY_SEGMENT_NAME_AUTO_GENERATED).c_str());
    currentPolyDataCopy->GetFieldData()->AddArray(nameAutoGeneratedArray);

    // ColorAutoGenerated
    vtkSmartPointer<vtkStringArray> colorAutoGeneratedArray = vtkSmartPointer<vtkStringArray>::New();
    colorAutoGeneratedArray->SetNumberOfValues(1);
    colorAutoGeneratedArray->SetValue(0, (currentSegment->GetColorAutoGenerated() ? "1" : "0"));
    colorAutoGeneratedArray->SetName(GetSegmentMetaDataKey(SINGLE_SEGMENT_INDEX, KEY_SEGMENT_COLOR_AUTO_GENERATED).c_str());
    currentPolyDataCopy->GetFieldData()->AddArray(colorAutoGeneratedArray);

    // Save conversion parameters as metadata (save in each segment file)
    std::string conversionParameters = segmentation->SerializeAllConversionParameters();
    vtkSmartPointer<vtkStringArray> conversionParametersArray = vtkSmartPointer<vtkStringArray>::New();
    conversionParametersArray->SetNumberOfValues(1);
    conversionParametersArray->SetValue(0, conversionParameters);
    conversionParametersArray->SetName(GetSegmentationMetaDataKey(KEY_SEGMENTATION_CONVERSION_PARAMETERS).c_str());
    currentPolyDataCopy->GetFieldData()->AddArray(conversionParametersArray);

    // Save contained representation names as metadata (save in each segment file)
    std::string containedRepresentationNames = this->SerializeContainedRepresentationNames(segmentation);
    vtkSmartPointer<vtkStringArray> containedRepresentationNamesArray = vtkSmartPointer<vtkStringArray>::New();
    containedRepresentationNamesArray->SetNumberOfValues(1);
    containedRepresentationNamesArray->SetValue(0, containedRepresentationNames);
    containedRepresentationNamesArray->SetName(GetSegmentationMetaDataKey(KEY_SEGMENTATION_CONTAINED_REPRESENTATION_NAMES).c_str());
    currentPolyDataCopy->GetFieldData()->AddArray(containedRepresentationNamesArray);

    // Set segment poly data to dataset
    multiBlockDataset->SetBlock(segmentIndex, currentPolyDataCopy);
  }

  // Write multiblock dataset to disk
  vtkSmartPointer<vtkXMLMultiBlockDataWriter> writer = vtkSmartPointer<vtkXMLMultiBlockDataWriter>::New();
  writer->SetInputData(multiBlockDataset);
  writer->SetFileName(path.c_str());
  if (this->UseCompression)
  {
    writer->SetDataModeToBinary();
    writer->SetCompressorTypeToZLib();
  }
  else
  {
    writer->SetDataModeToAscii();
    writer->SetCompressorTypeToNone();
  }

  this->GetUserMessages()->SetObservedObject(writer);
  writer->Write();
  this->GetUserMessages()->SetObservedObject(nullptr);

  // Add all files to storage node (multiblock dataset writes segments to individual files in a separate folder)
  this->AddPolyDataFileNames(path, segmentation);

  return 1;
}

//----------------------------------------------------------------------------
void vtkMRMLSegmentationStorageNode::AddPolyDataFileNames(std::string path, vtkSegmentation* segmentation)
{
  if (!segmentation)
  {
    vtkErrorMacro("AddPolyDataFileNames: Invalid segmentation");
    return;
  }

  this->AddFileName(path.c_str());

  std::string fileNameWithoutExtension = vtksys::SystemTools::GetFilenameWithoutLastExtension(path);
  std::string parentDirectory = vtksys::SystemTools::GetParentDirectory(path);
  std::string multiBlockDirectory = parentDirectory + "/" + fileNameWithoutExtension;
  for (int segmentIndex = 0; segmentIndex < segmentation->GetNumberOfSegments(); ++segmentIndex)
  {
    std::stringstream ssSegmentFilePath;
    ssSegmentFilePath << multiBlockDirectory << "/" << fileNameWithoutExtension << "_" << segmentIndex << ".vtp";
    std::string segmentFilePath = ssSegmentFilePath.str();
    this->AddFileName(segmentFilePath.c_str());
  }
}

//----------------------------------------------------------------------------
std::string vtkMRMLSegmentationStorageNode::SerializeContainedRepresentationNames(vtkSegmentation* segmentation)
{
  if (!segmentation)
  {
    vtkErrorMacro("SerializeContainedRepresentationNames: Invalid segmentation");
    return "";
  }

  std::stringstream ssRepresentationNames;
  std::vector<std::string> containedRepresentationNames;
  segmentation->GetContainedRepresentationNames(containedRepresentationNames);
  for (std::vector<std::string>::iterator reprIt = containedRepresentationNames.begin(); reprIt != containedRepresentationNames.end(); ++reprIt)
  {
    ssRepresentationNames << (*reprIt) << LIST_SEPARATOR;
  }

  return ssRepresentationNames.str();
}

//----------------------------------------------------------------------------
void vtkMRMLSegmentationStorageNode::CreateRepresentationsBySerializedNames(vtkSegmentation* segmentation, std::string representationNames)
{
  if (!segmentation)
  {
    vtkErrorMacro("CreateRepresentationsBySerializedNames: Invalid segmentation");
    return;
  }
  if (segmentation->GetNumberOfSegments() == 0)
  {
    vtkDebugMacro("CreateRepresentationsBySerializedNames: Segmentation is empty, nothing to do");
    return;
  }
  if (representationNames.empty())
  {
    vtkDebugMacro("CreateRepresentationsBySerializedNames: Empty representation names list, nothing to create");
    return;
  }

  std::string sourceRepresentation(segmentation->GetSourceRepresentationName());
  std::string representationName;
  std::stringstream representationNamesStr(representationNames);
  while (getline(representationNamesStr, representationName, LIST_SEPARATOR))
  {
    // Only create non-source representations
    if (representationName.compare(sourceRepresentation))
    {
      segmentation->CreateRepresentation(representationName);
    }
  }
}

//----------------------------------------------------------------------------
bool vtkMRMLSegmentationStorageNode::GetSegmentMetaDataFromDicitionary(std::string& headerValue, itk::MetaDataDictionary dictionary, int segmentIndex, std::string keyName)
{
  if (!dictionary.HasKey(GetSegmentMetaDataKey(segmentIndex, keyName)))
  {
    return false;
  }
  itk::ExposeMetaData<std::string>(dictionary, GetSegmentMetaDataKey(segmentIndex, keyName), headerValue);
  return true;
}

//----------------------------------------------------------------------------
bool vtkMRMLSegmentationStorageNode::GetSegmentationMetaDataFromDicitionary(std::string& headerValue, itk::MetaDataDictionary dictionary, std::string keyName)
{
  if (!dictionary.HasKey(GetSegmentationMetaDataKey(keyName)))
  {
    return false;
  }
  itk::ExposeMetaData<std::string>(dictionary, GetSegmentationMetaDataKey(keyName), headerValue);
  return true;
}

//----------------------------------------------------------------------------
std::string vtkMRMLSegmentationStorageNode::GetSegmentMetaDataKey(int segmentIndex, const std::string& keyName)
{
  std::stringstream key;
  key << "Segment";
  if (segmentIndex != SINGLE_SEGMENT_INDEX)
  {
    key << segmentIndex;
  }
  key << "_" << keyName;
  return key.str();
}

//----------------------------------------------------------------------------
std::string vtkMRMLSegmentationStorageNode::GetSegmentationMetaDataKey(const std::string& keyName)
{
  std::stringstream key;
  key << "Segmentation" << "_" << keyName;
  return key.str();
}

//----------------------------------------------------------------------------
std::string vtkMRMLSegmentationStorageNode::GetSegmentTagsAsString(vtkSegment* segment)
{
  std::map<std::string, std::string> tags;
  if (segment)
  {
    segment->GetTags(tags);
  }
  std::stringstream ssTagsValue;
  std::map<std::string, std::string>::iterator tagIt;
  for (tagIt = tags.begin(); tagIt != tags.end(); ++tagIt)
  {
    ssTagsValue << tagIt->first << ":" << tagIt->second << "|";
  }
  std::string tagsValue = ssTagsValue.str();
  return tagsValue;
}

//----------------------------------------------------------------------------
void vtkMRMLSegmentationStorageNode::SetSegmentTagsFromString(vtkSegment* segment, std::string tagsValue)
{
  std::stringstream tagsValueStr(tagsValue);
  std::string tagValue;
  while (getline(tagsValueStr, tagValue, LIST_SEPARATOR))
  {
    size_t colonPosition = tagValue.find(":");
    if (colonPosition == std::string::npos)
    {
      continue;
    }
    segment->SetTag(tagValue.substr(0, colonPosition), tagValue.substr(colonPosition + 1));
  }
}

//----------------------------------------------------------------------------
std::string vtkMRMLSegmentationStorageNode::GetImageExtentAsString(vtkOrientedImageData* image)
{
  int extent[6] = { 0, -1, 0, -1, 0, -1 };
  if (image)
  {
    image->GetExtent(extent);
  }
  return GetImageExtentAsString(extent);
}

//----------------------------------------------------------------------------
std::string vtkMRMLSegmentationStorageNode::GetImageExtentAsString(int extent[6])
{
  std::stringstream ssExtentValue;
  ssExtentValue << extent[0] << " " << extent[1] << " " << extent[2] << " " << extent[3] << " " << extent[4] << " " << extent[5];
  std::string extentValue = ssExtentValue.str();
  return extentValue;
}

//----------------------------------------------------------------------------
void vtkMRMLSegmentationStorageNode::GetImageExtentFromString(int extent[6], std::string extentValue)
{
  std::stringstream ssExtentValue(extentValue);
  extent[0] = 0;
  extent[1] = -1;
  extent[2] = 0;
  extent[3] = -1;
  extent[4] = 0;
  extent[5] = -1;
  ssExtentValue >> extent[0] >> extent[1] >> extent[2] >> extent[3] >> extent[4] >> extent[5];
}

//----------------------------------------------------------------------------
std::string vtkMRMLSegmentationStorageNode::GetSegmentColorAsString(vtkMRMLSegmentationNode* segmentationNode, const std::string& segmentId)
{
  double color[3] = { vtkSegment::SEGMENT_COLOR_INVALID[0], vtkSegment::SEGMENT_COLOR_INVALID[1], vtkSegment::SEGMENT_COLOR_INVALID[2] };
  vtkSegment* segment = segmentationNode->GetSegmentation()->GetSegment(segmentId);
  if (segment)
  {
    segment->GetColor(color);
  }
  std::stringstream colorStream;
  colorStream << color[0] << " " << color[1] << " " << color[2];
  return colorStream.str();
}

//----------------------------------------------------------------------------
void vtkMRMLSegmentationStorageNode::GetSegmentColorFromString(double color[3], std::string colorString)
{
  std::stringstream colorStream(colorString);
  color[0] = 0.5;
  color[1] = 0.5;
  color[2] = 0.5;
  colorStream >> color[0] >> color[1] >> color[2];
}
