write.go

package dicom

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"

	"github.com/ginuerzh/dicom/pkg/dicomio"
	"github.com/ginuerzh/dicom/pkg/tag"
	"github.com/ginuerzh/dicom/pkg/uid"
)

var (
	// ErrorUnimplemented is for not yet finished things.
	ErrorUnimplemented = errors.New("this functionality is not yet implemented")
	// ErrorMismatchValueTypeAndVR is for when there's a discrepency betweeen the ValueType and what the VR specifies.
	ErrorMismatchValueTypeAndVR = errors.New("ValueType does not match the VR required")
	// ErrorUnexpectedValueType indicates an unexpected value type was seen.
	ErrorUnexpectedValueType = errors.New("Unexpected ValueType")
	// ErrorUnsupportedBitsPerSample indicates that the BitsPerSample in this
	// Dataset is not supported when unpacking native PixelData.
	ErrorUnsupportedBitsPerSample = errors.New("unsupported BitsPerSample value")
)

// TODO(suyashkumar): consider adding an element-by-element write API.

// Write will write the input DICOM dataset to the provided io.Writer as a complete DICOM (including any header
// information if available).
func Write(out io.Writer, ds Dataset, opts ...WriteOption) error {
	optSet := toOptSet(opts...)
	w := dicomio.NewWriter(out, nil, false)
	var metaElems []*Element
	for _, elem := range ds.Elements {
		if elem.Tag.Group == tag.MetadataGroup {
			metaElems = append(metaElems, elem)
		}
	}

	err := writeFileHeader(w, &ds, metaElems, *optSet)
	if err != nil {
		return err
	}

	endian, implicit, err := ds.transferSyntax()
	if (err != nil && err != ErrorElementNotFound) || (err == ErrorElementNotFound && !optSet.defaultMissingTransferSyntax) {
		return err
	}

	if err == ErrorElementNotFound && optSet.defaultMissingTransferSyntax {
		w.SetTransferSyntax(binary.LittleEndian, true)
	} else {
		w.SetTransferSyntax(endian, implicit)
	}

	for _, elem := range ds.Elements {
		if elem.Tag.Group != tag.MetadataGroup {
			err = writeElement(w, elem, *optSet)
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// WriteOption represents an option that can be passed to WriteDataset. Later options will override previous options if
// applicable.
type WriteOption func(*writeOptSet)

// SkipVRVerification returns a WriteOption that skips VR verification.
func SkipVRVerification() WriteOption {
	return func(set *writeOptSet) {
		set.skipVRVerification = true
	}
}

// SkipValueTypeVerification returns WriteOption function that skips checking ValueType
// for concurrency with VR and casting
func SkipValueTypeVerification() WriteOption {
	return func(set *writeOptSet) {
		set.skipValueTypeVerification = true
	}
}

// DefaultMissingTransferSyntax returns a WriteOption indicating that a missing
// transferSyntax should not raise an error, and instead the default
// LittleEndian Implicit transfer syntax should be used and written out as a
// Metadata element in the Dataset.
func DefaultMissingTransferSyntax() WriteOption {
	return func(set *writeOptSet) {
		set.defaultMissingTransferSyntax = true
	}
}

// writeOptSet represents the flattened option set after all WriteOptions have been applied.
type writeOptSet struct {
	skipVRVerification           bool
	skipValueTypeVerification    bool
	defaultMissingTransferSyntax bool
}

func toOptSet(opts ...WriteOption) *writeOptSet {
	optSet := &writeOptSet{}
	for _, opt := range opts {
		opt(optSet)
	}
	return optSet
}

func writeFileHeader(w dicomio.Writer, ds *Dataset, metaElems []*Element, opts writeOptSet) error {
	// File headers are always written in littleEndian explicit
	w.SetTransferSyntax(binary.LittleEndian, false)

	metaBytes := &bytes.Buffer{}
	subWriter := dicomio.NewWriter(metaBytes, binary.LittleEndian, false)
	tagsUsed := make(map[tag.Tag]bool)
	tagsUsed[tag.FileMetaInformationGroupLength] = true

	err := writeMetaElem(subWriter, tag.FileMetaInformationVersion, ds, &tagsUsed, opts)
	if err != nil && err != ErrorElementNotFound {
		return err
	}
	err = writeMetaElem(subWriter, tag.MediaStorageSOPClassUID, ds, &tagsUsed, opts)
	if err != nil && err != ErrorElementNotFound {
		return err
	}
	err = writeMetaElem(subWriter, tag.MediaStorageSOPInstanceUID, ds, &tagsUsed, opts)
	if err != nil && err != ErrorElementNotFound {
		return err
	}
	err = writeMetaElem(subWriter, tag.TransferSyntaxUID, ds, &tagsUsed, opts)
	if err != nil && err != ErrorElementNotFound || err == ErrorElementNotFound && !opts.defaultMissingTransferSyntax {
		return err
	}
	if err == ErrorElementNotFound && opts.defaultMissingTransferSyntax {
		// Write the default transfer syntax
		if err = writeElement(subWriter, mustNewElement(tag.TransferSyntaxUID, []string{uid.ImplicitVRLittleEndian}), opts); err != nil {
			return err
		}
	}

	for _, elem := range metaElems {
		if elem.Tag.Group == tag.MetadataGroup {
			if _, ok := tagsUsed[elem.Tag]; !ok {
				err = writeElement(subWriter, elem, opts)
				if err != nil {
					return err
				}
			}
		}
	}

	if len(ds.Preamble) == 128 {
		if err := w.WriteBytes(ds.Preamble); err != nil {
			return err
		}
	} else {
		if err := w.WriteZeros(128); err != nil {
			return err
		}
	}

	if err := w.WriteString(magicWord); err != nil {
		return err
	}
	lengthElem, err := NewElement(tag.FileMetaInformationGroupLength, []uint64{uint64(len(metaBytes.Bytes()))})
	if err != nil {
		return err
	}

	err = writeElement(w, lengthElem, opts)
	if err != nil {
		return err
	}
	w.WriteBytes(metaBytes.Bytes())

	return nil
}

func writeElement(w dicomio.Writer, elem *Element, opts writeOptSet) error {
	vr := elem.RawValueRepresentation
	if !opts.skipVRVerification {
		var err error
		vr, err = verifyVROrDefault(elem.Tag, elem.RawValueRepresentation)
		if err != nil {
			return err
		}
	}
	if !opts.skipValueTypeVerification && elem.Value != nil {
		err := verifyValueType(elem.Tag, elem.Value, vr)
		if err != nil {
			return err
		}
	}

	length := elem.ValueLength
	var valueData = &bytes.Buffer{}
	if elem.Value != nil {
		bo, implicit := w.GetTransferSyntax()
		subWriter := dicomio.NewWriter(valueData, bo, implicit)
		err := writeValue(subWriter, elem.Tag, elem.Value, elem.Value.ValueType(), vr, elem.ValueLength, opts)
		if err != nil {
			return err
		}

		length = uint32(len(valueData.Bytes()))
		if elem.ValueLength == tag.VLUndefinedLength {
			length = tag.VLUndefinedLength
		}
	}

	err := encodeElementHeader(w, elem.Tag, vr, length)
	if err != nil {
		return err
	}

	if elem.Value != nil {
		// Write the bytes to the original writer
		w.WriteBytes(valueData.Bytes())
	}
	return nil
}

func writeMetaElem(w dicomio.Writer, t tag.Tag, ds *Dataset, tagsUsed *map[tag.Tag]bool, optSet writeOptSet) error {
	elem, err := ds.FindElementByTag(t)
	if err != nil {
		return err
	}
	err = writeElement(w, elem, optSet)
	if err != nil {
		return err
	}
	(*tagsUsed)[t] = true
	return nil
}

func verifyVROrDefault(t tag.Tag, vr string) (string, error) {
	tagInfo, err := tag.Find(t)
	if err != nil {
		return "UN", nil
	}
	if vr == "" {
		return tagInfo.VR, nil
	}
	if tagInfo.VR != vr {
		return "", fmt.Errorf("ERROR dicomio.veryifyElement: VR mismatch for tag %v. Element.VR=%v, but DICOM standard defines VR to be %v",
			tag.DebugString(t), vr, tagInfo.VR)
	}

	return vr, nil
}

func verifyValueType(t tag.Tag, value Value, vr string) error {
	valueType := value.ValueType()
	var ok bool
	switch vr {
	case "SL", "SS", "SV":
		ok = valueType == Ints
	case "AT", "OL", "OV", "UL", "US", "UV":
		ok = valueType == UInts
	case "SQ":
		ok = valueType == Sequences
	case "NA":
		ok = valueType == SequenceItem
	case "OW", "OB", "UN":
		if t == tag.PixelData {
			ok = valueType == PixelData
		} else {
			ok = valueType == Bytes
		}
	case "FL", "FD", "OD", "OF":
		ok = valueType == Floats
	default:
		ok = valueType == Strings
	}

	if !ok {
		return fmt.Errorf("%v: ValueType %s does not match the specified type in the VR %s", t, valueType, vr)
	}
	return nil
}

func writeTag(w dicomio.Writer, t tag.Tag, vl uint32) error {
	if vl%2 != 0 && vl != tag.VLUndefinedLength {
		return fmt.Errorf("ERROR dicomio.writeTag: Value Length must be even, but for Tag=%v, ValueLength=%v",
			tag.DebugString(t), vl)
	}
	if err := w.WriteUInt16(t.Group); err != nil {
		return err
	}
	return w.WriteUInt16(t.Element)
}

func writeVRVL(w dicomio.Writer, t tag.Tag, vr string, vl uint32) error {
	// Rectify Undefined Length VL
	if vl == 0xffff {
		vl = tag.VLUndefinedLength
	}

	if vr == "SQ" || t == tag.Item {
		// We are going to write these out with undefined length always.
		vl = tag.VLUndefinedLength
	}

	if len(vr) != 2 && vl != tag.VLUndefinedLength &&
		t != tag.SequenceDelimitationItem && t != tag.ItemDelimitationItem {
		return fmt.Errorf("ERROR dicomio.writeVRVL: Value Representation must be of length 2, e.g. 'UN'. For tag=%v, it was RawValueRepresentation=%v",
			tag.DebugString(t), vr)
	}

	// Write VR then VL
	_, implicit := w.GetTransferSyntax()
	if t.Group == tag.GROUP_ItemSeq {
		implicit = true
	}
	if !implicit { // Explicit
		if err := w.WriteString(vr); err != nil {
			return err
		}
		switch vr {
		case "NA", "OB", "OD", "OF", "OL", "OV", "OW", "SQ", "SV", "UN", "UC", "UR", "UV", "UT":
			if err := w.WriteZeros(2); err != nil {
				return err
			}
			if err := w.WriteUInt32(vl); err != nil {
				return err
			}
		default:
			if err := w.WriteUInt16(uint16(vl)); err != nil {
				return err
			}
		}
	} else {
		if err := w.WriteUInt32(vl); err != nil {
			return err
		}
	}
	return nil
}

func writeRawItem(w dicomio.Writer, data []byte) error {
	length := uint32(len(data))
	if err := writeTag(w, tag.Item, length); err != nil {
		return err
	}
	if err := writeVRVL(w, tag.Item, "NA", length); err != nil {
		return err
	}
	w.WriteBytes(data)
	return nil
}

func writeBasicOffsetTable(w dicomio.Writer, offsets []uint32) error {
	byteOrder, implicit := w.GetTransferSyntax()
	data := &bytes.Buffer{}
	subWriter := dicomio.NewWriter(data, byteOrder, implicit)
	for _, offset := range offsets {
		if err := subWriter.WriteUInt32(offset); err != nil {
			return err
		}
	}
	return writeRawItem(w, data.Bytes())
}

func encodeElementHeader(w dicomio.Writer, t tag.Tag, vr string, vl uint32) error {
	err := writeTag(w, t, vl)
	if err != nil {
		return err
	}
	err = writeVRVL(w, t, vr, vl)
	if err != nil {
		return err
	}
	return nil
}

func writeValue(w dicomio.Writer, t tag.Tag, value Value, valueType ValueType, vr string, vl uint32, opts writeOptSet) error {
	if vl == tag.VLUndefinedLength && valueType <= 3 { // strings, bytes, ints or uints
		return fmt.Errorf("encoding undefined-length element not yet supported: %v", t)
	}

	v := value.GetValue()
	switch valueType {
	case Strings:
		return writeStrings(w, v.([]string), vr)
	case Bytes:
		return writeBytes(w, v.([]byte), vr)
	case Ints:
		return writeInts(w, v.([]int64), vr)
	case UInts:
		return writeUInts(w, v.([]uint64), vr)
	case PixelData:
		return writePixelData(w, t, value, vr, vl)
	case SequenceItem:
		return writeSequenceItem(w, t, v.([]*Element), vr, vl, opts)
	case Sequences:
		return writeSequence(w, t, v.([]*SequenceItemValue), vr, vl, opts)
	case Floats:
		return writeFloats(w, value, vr)
	default:
		return fmt.Errorf("ValueType not supported")
	}
}

func writeStrings(w dicomio.Writer, values []string, vr string) error {
	s := ""
	for i, substr := range values {
		if i > 0 {
			s += "\\"
		}
		s += substr
	}
	if err := w.WriteString(s); err != nil {
		return err
	}
	if len(s)%2 == 1 {
		switch vr {
		case "AE", "CS", "DS", "DT", "IS", "LO", "LT", "PN", "SH", "ST", "TM", "UC", "UN", "UR", "UT":
			if err := w.WriteString(" "); err != nil { // http://dicom.nema.org/medical/dicom/current/output/html/part05.html#sect_6.2
				return err
			}
		default:
			if err := w.WriteByte(0); err != nil {
				return err
			}
		}
	}
	return nil
}

func writeBytes(w dicomio.Writer, values []byte, vr string) error {
	var err error
	switch vr {
	case "OW":
		err = writeOtherWordString(w, values)
	case "OB", "UN":
		err = writeOtherByteString(w, values)
	default:
		return ErrorMismatchValueTypeAndVR
	}
	if err != nil {
		return err
	}
	return nil
}

func writeInts(w dicomio.Writer, values []int64, vr string) error {
	for _, value := range values {
		switch vr {
		case "SS":
			if err := w.WriteUInt16(uint16(value)); err != nil {
				return err
			}
		case "SL":
			if err := w.WriteUInt32(uint32(value)); err != nil {
				return err
			}
		case "SV":
			if err := w.WriteUInt64(uint64(value)); err != nil {
				return err
			}
		default:
			return ErrorMismatchValueTypeAndVR
		}
	}
	return nil
}

func writeUInts(w dicomio.Writer, values []uint64, vr string) error {
	for _, value := range values {
		switch vr {
		case "US":
			if err := w.WriteUInt16(uint16(value)); err != nil {
				return err
			}
		case "AT", "UL", "OL":
			if err := w.WriteUInt32(uint32(value)); err != nil {
				return err
			}
		case "UV", "OV":
			if err := w.WriteUInt64(uint64(value)); err != nil {
				return err
			}
		default:
			return ErrorMismatchValueTypeAndVR
		}
	}
	return nil
}

func writeFloats(w dicomio.Writer, v Value, vr string) error {
	if v.ValueType() != Floats {
		return ErrorUnexpectedValueType
	}
	floats := MustGetFloats(v)
	for _, fl := range floats {
		switch vr {
		case "FL", "OF":
			// NOTE: this is a conversion from float64 -> float32 which may lead to a loss in precision. The assumption
			// is that the value sitting in the float64 was originally at float32 precision if the VR is FL for this
			// element. We will need to revisit this. Maybe we can detect if there will be a loss of precision and if so
			// indicate an error or warning.
			err := w.WriteFloat32(float32(fl))
			if err != nil {
				return err
			}
		case "FD", "OD":
			err := w.WriteFloat64(fl)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func writePixelData(w dicomio.Writer, t tag.Tag, value Value, vr string, vl uint32) error {
	image := MustGetPixelDataInfo(value)
	if vl == tag.VLUndefinedLength {
		if err := writeBasicOffsetTable(w, image.Offsets); err != nil {
			return err
		}
		for _, frame := range image.Frames {
			if err := writeRawItem(w, frame.EncapsulatedData.Data); err != nil {
				return err
			}
		}
		err := encodeElementHeader(w, tag.SequenceDelimitationItem, "", 0)
		if err != nil {
			return err
		}
	} else {
		for _, frame := range image.Frames {
			if err := w.WriteBytes(frame.NativeData.Data); err != nil {
				return err
			}
		}
	}
	return nil
}

var sequenceDelimitationItem = &Element{
	Tag:         tag.SequenceDelimitationItem,
	ValueLength: 0, // This should be 00000000H in base32
}

func writeSequence(w dicomio.Writer, t tag.Tag, values []*SequenceItemValue, vr string, vl uint32, opts writeOptSet) error {
	// We always write out sequences using the undefined length encoding.
	// Note: we currently don't validate that the length of the sequence matches
	// the VL if it's not undefined VL.
	// More details about the sequence structure can be found at:
	// http://dicom.nema.org/medical/dicom/current/output/chtml/part05/sect_7.5.html

	// Write out the items.
	for _, seqItem := range values {
		if err := writeSequenceItem(w, t, seqItem.elements, vr, vl, opts); err != nil {
			return err
		}
	}

	// Write Sequence Delimitation Item as implicit VR
	oldBO, oldImplicit := w.GetTransferSyntax()
	w.SetTransferSyntax(oldBO, true)
	if err := writeElement(w, sequenceDelimitationItem, opts); err != nil {
		return err
	}
	w.SetTransferSyntax(oldBO, oldImplicit) // Return TS to what it was before.

	return nil
}

var sequenceItemDelimitationItem = &Element{
	Tag:         tag.ItemDelimitationItem,
	ValueLength: 0, // This should be 00000000H in base32
}

var item = &Element{
	Tag:         tag.Item,
	ValueLength: tag.VLUndefinedLength,
}

func writeSequenceItem(w dicomio.Writer, t tag.Tag, values []*Element, vr string, vl uint32, opts writeOptSet) error {
	// Write out item header.
	if err := writeElement(w, item, opts); err != nil {
		return err
	}

	// Write out nested Dataset elements.
	for _, elem := range values {
		if err := writeElement(w, elem, opts); err != nil {
			return err
		}
	}

	// Write ItemDelimitationItem.
	return writeElement(w, sequenceItemDelimitationItem, opts)
}

func writeOtherWordString(w dicomio.Writer, data []byte) error {
	if len(data)%2 != 0 {
		return ErrorOWRequiresEvenVL
	}
	bo, _ := w.GetTransferSyntax()
	r, err := dicomio.NewReader(bufio.NewReader(bytes.NewBuffer(data)), bo, int64(len(data)))
	if err != nil {
		return err
	}
	for i := 0; i < int(len(data)/2); i++ {
		v, err := r.ReadUInt16()
		if err != nil {
			return err
		}
		if err := w.WriteUInt16(v); err != nil {
			return err
		}
	}
	return nil
}

func writeOtherByteString(w dicomio.Writer, data []byte) error {
	if err := w.WriteBytes(data); err != nil {
		return err
	}
	if len(data)%2 == 1 {
		if err := w.WriteByte(0); err != nil {
			return err
		}
	}
	return nil
}