#   ev-to-dicom
#   Copyright © 2023–2024  Lee Yingtong Li
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU Affero General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU Affero General Public License for more details.
#
#   You should have received a copy of the GNU Affero General Public License
#   along with this program.  If not, see <https://www.gnu.org/licenses/>.

import msgpack
import pydicom
from pydicom.dataset import FileDataset, FileMetaDataset
from pydicom.encaps import encapsulate

from datetime import datetime

from .ev_tlv import read_ev_tlv

def ev_files_to_dcm_file(study_uid, series_uid, series_number, image_number, input_filenames, output_filename):
	# Read Enhanced Viewer input data
	
	pixel_datas = []
	
	for filename in input_filenames:
		with open(filename, 'rb') as f:
			data = f.read()
		
		metadata, pixel_data = read_ev_tlv(data)
		
		pixel_datas.append(pixel_data)
	
	# Prepare DICOM file
	
	file_meta = FileMetaDataset()
	
	if metadata['metadata']['modality'] == 'US':
		if len(pixel_datas) == 1:
			file_meta.MediaStorageSOPClassUID = '1.2.840.10008.5.1.4.1.1.6.1'  # Ultrasound Image Storage
		else:
			file_meta.MediaStorageSOPClassUID = '1.2.840.10008.5.1.4.1.1.3.1'  # Ultrasound Multi-frame Image Storage
	elif metadata['metadata']['modality'] == 'PX':
		file_meta.MediaStorageSOPClassUID = '1.2.840.10008.5.1.4.1.1.1.1'  # Digital X-Ray Image Storage - For Presentation
	else:
		raise Exception('Unknown modality {}'.format(metadata['metadata']['modality']))
	
	file_meta.MediaStorageSOPInstanceUID = metadata['sopInstanceUid']
	
	if metadata['imageFormat'] == 0:
		file_meta.TransferSyntaxUID = '1.2.840.10008.1.2.4.91'  # JPEG 2000
	elif metadata['imageFormat'] == 1:
		file_meta.TransferSyntaxUID = '1.2.840.10008.1.2.4.50'  # JPEG
	else:
		raise Exception('Unknown imageFormat {}'.format(metadata['metadata']['imageFormat']))
	
	ds = FileDataset(output_filename, {}, file_meta=file_meta, preamble=b'\0' * 128)
	ds.SOPClassUID = file_meta.MediaStorageSOPClassUID
	ds.SOPInstanceUID = metadata['sopInstanceUid']
	
	ds.is_little_endian = True
	ds.is_implicit_VR = False
	
	ds.StudyInstanceUID = study_uid
	ds.SeriesInstanceUID = series_uid
	ds.SeriesNumber = series_number
	ds.InstanceNumber = image_number
	
	ds.PatientName = metadata['metadata']['patientName'].replace(', ', '^')
	ds.PatientSex = metadata['metadata']['patientSex']
	ds.SeriesDescription = metadata['metadata']['seriesDescription']
	ds.Modality = metadata['metadata']['modality']
	ds.PatientID = metadata['metadata']['patientId']
	ds.StudyDate = datetime.strptime(metadata['metadata']['seriesDate'], '%Y-%m-%dT%H:%M:%S.%f').strftime('%Y%m%d')
	ds.StudyTime = datetime.strptime(metadata['metadata']['seriesDate'], '%Y-%m-%dT%H:%M:%S.%f').strftime('%H%M%S.%f')
	ds.SeriesDate = datetime.strptime(metadata['metadata']['seriesDate'], '%Y-%m-%dT%H:%M:%S.%f').strftime('%Y%m%d')
	ds.SeriesTime = datetime.strptime(metadata['metadata']['seriesDate'], '%Y-%m-%dT%H:%M:%S.%f').strftime('%H%M%S.%f')
	ds.StudyDescription = metadata['metadata']['studyDescription']
	ds.AccessionNumber = metadata['metadata']['accessionNumber']
	ds.PatientBirthDate = datetime.strptime(metadata['metadata']['patientBirthDate'], '%Y-%m-%d').strftime('%Y%m%d')
	ds.InstitutionName = metadata['metadata']['institutionName']
	ds.PatientAge = metadata['metadata']['patientAge']
	ds.SamplesPerPixel = 3 if metadata['color'] else 1
	ds.PixelRepresentation = 1 if metadata['isSigned'] else 0
	#ds.BitsAllocated = metadata['bytesPerSample'] * 8  # Incorrect value!
	#ds.BitsStored = metadata['bytesPerSample'] * 8
	ds.Columns = metadata['columns']
	ds.WindowCenter = metadata['windowCenter']
	ds.LossyImageCompression = '01' if metadata['isLossyImage'] else '00'
	ds.PixelSpacing = [metadata['rowPixelSpacing'], metadata['columnPixelSpacing']]
	ds.WindowWidth = metadata['windowWidth']
	ds.PlanarConfiguration = 1 if metadata['planarConfiguration'] else 0
	ds.Rows = metadata['rows']
	
	if metadata['color']:
		# FIXME: Any proper way to detect this?
		
		if file_meta.TransferSyntaxUID == '1.2.840.10008.1.2.4.50':  # JPEG
			ds.PhotometricInterpretation = 'YBR_FULL_422'  # Specify YBR as that is how JPEG compression works, even though the JSON says RGB
		else:
			ds.PhotometricInterpretation = 'RGB'
	else:
		ds.PhotometricInterpretation = 'MONOCHROME2'
	
	if len(pixel_datas) > 1:
		ds.RecommendedDisplayFrameRate = round(metadata['recommendedDisplayFrameRate'])
		ds.FrameTime = 1000 / metadata['recommendedDisplayFrameRate']
		ds.FrameIncrementPointer = (0x0018, 0x1063)  # Frame Time
		ds.NumberOfFrames = len(pixel_datas)
	
	# Store copy of raw metadata
	block = ds.private_block(0x0075, 'RunasSudo ev-to-dicom', create=True)
	block.add_new(0x01, 'UL', 0)  # Data format version
	del metadata['totalAvailableBytes']
	block.add_new(0x02, 'OB', msgpack.packb(metadata))
	
	if metadata['imageFormat'] == 0:
		# JPEG 2000 - For some reason, the end of codestream marker is (always?) missing
		ds.PixelData = encapsulate([pixel_data if pixel_data.endswith(b'\xff\xd9') else pixel_data + b'\xff\xd9' for pixel_data in pixel_datas])
	else:
		ds.PixelData = encapsulate(pixel_datas)
	
	# Save DICOM data
	
	ds.save_as(output_filename)