import json
import os
import sys
import traceback
from multiprocessing import Pool, cpu_count

import SimpleITK as sitk
from pydicom import dicomio
import lmdb
import numpy as np
from datetime import datetime

from resample import ItkResample

BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))  # 添加路径
sys.path.append(BASE_DIR)

class MyEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        elif isinstance(obj, datetime):
            return obj.__str__()
        else:
            return super(MyEncoder, self).default(obj)
        
def load_ct_from_dicom(dcm_path, sort_by_distance=True):
    class DcmInfo(object):
        def __init__(self, dcm_path, series_instance_uid, acquisition_number, sop_instance_uid, instance_number,
                     image_orientation_patient, image_position_patient):
            super(DcmInfo, self).__init__()

            self.dcm_path = dcm_path
            self.series_instance_uid = series_instance_uid
            self.acquisition_number = acquisition_number
            self.sop_instance_uid = sop_instance_uid
            self.instance_number = instance_number
            self.image_orientation_patient = image_orientation_patient
            self.image_position_patient = image_position_patient

            self.slice_distance = self._cal_distance()

        def _cal_distance(self):
            normal = [self.image_orientation_patient[1] * self.image_orientation_patient[5] -
                      self.image_orientation_patient[2] * self.image_orientation_patient[4],
                      self.image_orientation_patient[2] * self.image_orientation_patient[3] -
                      self.image_orientation_patient[0] * self.image_orientation_patient[5],
                      self.image_orientation_patient[0] * self.image_orientation_patient[4] -
                      self.image_orientation_patient[1] * self.image_orientation_patient[3]]

            distance = 0
            for i in range(3):
                distance += normal[i] * self.image_position_patient[i]
            return distance

    def is_sop_instance_uid_exist(dcm_info, dcm_infos):
        for item in dcm_infos:
            if dcm_info.sop_instance_uid == item.sop_instance_uid:
                return True
        return False

    def get_dcm_path(dcm_info):
        return dcm_info.dcm_path

    reader = sitk.ImageSeriesReader()
    if sort_by_distance:
        dcm_infos = []

        files = os.listdir(dcm_path)
        for file in files:
            file_path = os.path.join(dcm_path, file)

            dcm = dicomio.read_file(file_path, force=True)
            _series_instance_uid = dcm.SeriesInstanceUID
            _sop_instance_uid = dcm.SOPInstanceUID
            _instance_number = dcm.InstanceNumber
            _acquisition_number = dcm.AcquisitionNumber
            _image_orientation_patient = dcm.ImageOrientationPatient
            _image_position_patient = dcm.ImagePositionPatient

            dcm_info = DcmInfo(file_path, _series_instance_uid, _acquisition_number, _sop_instance_uid,
                               _instance_number, _image_orientation_patient, _image_position_patient)

            if is_sop_instance_uid_exist(dcm_info, dcm_infos):
                continue

            dcm_infos.append(dcm_info)

        dcm_infos.sort(key=lambda x: x.slice_distance)
        dcm_series = list(map(get_dcm_path, dcm_infos))
    else:
        dcm_series = reader.GetGDCMSeriesFileNames(dcm_path)

    reader.SetFileNames(dcm_series)
    sitk_image = reader.Execute()
    return sitk_image

class Trans_Dicom_2_NII(object):
    def __init__(self, in_path, out_path, raw_dict_dir):
        '''
        :param in_path: path of dicom
        :param out_path: path of nii.gz
        '''
        self.in_path = in_path
        self.out_path = out_path
        self.raw_dict_dir = raw_dict_dir

        self.target_spacing = 0.5
        self.itk_resample = ItkResample()

        if not os.path.exists(self.out_path):
            os.makedirs(self.out_path)

    def __call__(self):
        # raw_dicts = {}
        uids = os.listdir(self.in_path)
        ## cpu_count() 112
        # pool = Pool(int(cpu_count() * 0.7))

        for uid in uids:
            if not os.path.exists(os.path.join(self.out_path, uid + '.nii.gz')):
                try:
                    self._single_transform(uid, self.raw_dict_dir)
                    # pool.apply_async(self._single_transform, (uid, self.raw_dict_dir))
                    # raw_info = raw_info.get()
                    # raw_dicts[uid] = raw_info
                    
                except Exception as err:
                    traceback.print_exc()
                    print('Transform dicom to nii.gz throws exception %s, with series uid %s!' % (err, uid))

        # pool.close()
        # pool.join()
        
        # if not os.path.exists(os.path.join(os.path.dirname(self.out_path), 'raw_dicts.json')):
        #     with open(os.path.join(os.path.dirname(self.out_path), 'raw_dicts.json'), 'w') as file:
        #         json.dump(raw_dicts, file, indent=4)
        # else:
        #     raw_dicts = json.load(open(os.path.join(os.path.dirname(self.out_path), 'raw_dicts.json'), 'r'))

        # return raw_dicts

    def _single_transform(self, uid, raw_dict_dir):
        # dcm -> nii.gz，并斤进行排序
        print('Processing series uid %s' % uid)

        nii_file = os.path.join(self.out_path, uid + '.nii.gz')

        dcm_folder = os.path.join(self.in_path, uid)

        # convert dicom to nii
        try:
            itk_image = load_ct_from_dicom(dcm_folder)
        except Exception as err:
            print('!!!!! Read %s throws exception %s.' % (uid, err))
            return
        
        # resample spacing
        print('Resample uid %s to spacing %.1f.' % (uid, self.target_spacing))
        itk_image_resample = self.itk_resample.resample_to_spacing(
            itk_image=itk_image,
            target_spacing=(tuple([self.target_spacing] * 3)),
            interpolator=sitk.sitkBSpline)
        itk_image_resample.SetOrigin(itk_image.GetOrigin())
        itk_image_resample.SetSpacing([self.target_spacing] * 3)

        try:
            sitk.WriteImage(itk_image_resample, nii_file)
        except Exception as err:
            print('!!!!! Write %s throws exception %s.' % (uid, err))
            return
        
        info_dict = dict()
        info_dict['raw_origin'] = list(np.array(itk_image_resample.GetOrigin()))
        info_dict['raw_spacing'] = list(np.array(itk_image_resample.GetSpacing()))
        info_dict['raw_size'] = list(np.array(itk_image_resample.GetSize()).astype(float))
        raw_dict = {uid: info_dict}

        with open(os.path.join(raw_dict_dir, uid+'.json'), 'w+') as file:
            json.dump(raw_dict, file, indent=4)

        # info_dict = dict()
        # info_dict['raw_origin'] = list(np.array(itk_image_resample.GetOrigin()))
        # info_dict['raw_spacing'] = list(np.array(itk_image_resample.GetSpacing()))
        # info_dict['raw_size'] = list(np.array(itk_image_resample.GetSize()).astype(float))

        # return info_dict

