import pydicom
import glob
import sys, os
import numpy as np
import cv2
import matplotlib.pyplot as plt
from tqdm import tqdm


class DicomProcessor:

    def __init__(self, dicom_dir, args):
        self.args = args
        self.dicom_dir = dicom_dir
        self.dicom_file_list = glob.glob(os.path.join(dicom_dir, '*'))
        series_set = set([pydicom.read_file(x).SeriesDescription for x in self.dicom_file_list])
        show_dict = {(i + 1): series_name for i, series_name in enumerate(series_set)}
        print("対象シリーズ名を選択してください")
        for key, item in show_dict.items():
            print("{}: {}".format(key, item))
        target = show_dict[int(input())]
        self.dicom_file_list = [x for x in self.dicom_file_list if pydicom.read_file(x).SeriesDescription == target]
        print("対象シリーズを抽出しました．({})".format(target))
        self.dicom_file_list = sorted(self.dicom_file_list, key=lambda x: float(pydicom.read_file(x).SliceLocation), reverse=False)
        self.dicom_size = len(self.dicom_file_list)

        if self.dicom_size == 0:
            print("dicom file has no files")
            raise

        sample_dfile = pydicom.read_file(self.dicom_file_list[0])
        sample_dfile2 = pydicom.read_file(self.dicom_file_list[1])
        self.dcm_wc = sample_dfile.WindowCenter
        self.dcm_ww = sample_dfile.WindowWidth

        # 一時的に変更
        try:
            self.dcm_wc = self.dcm_wc[0]
            self.dcm_ww = self.dcm_ww[0]
        except:
            pass

        print("original wc: {}".format(self.dcm_wc))
        print("original ww: {}".format(self.dcm_ww))
        self.dicom_row_num = sample_dfile.Rows
        self.dicom_col_num = sample_dfile.Columns

        self.pixel_spacing_dx, self.pixel_spacing_dy = sample_dfile.PixelSpacing
        self.slice_spacing = abs(sample_dfile.ImagePositionPatient[2] - sample_dfile2.ImagePositionPatient[2])
        self.scaling_f = np.array([self.pixel_spacing_dx, self.pixel_spacing_dy, self.slice_spacing])
        self.model_center_vec = np.array([self.dicom_row_num / 2, self.dicom_col_num / 2, self.dicom_size / 2])
        self.image_position_patient = sample_dfile.ImagePositionPatient
        #なぜか3DSlicer上ではSが-0.25されている
        # self.image_position_patient[2] -= 2.5
        self.CT_for_imshow = self.load_initial_CT()

    def calc_ijk2LPS_mat(self, dst=None):
        ijk_samples = self.calc_ijk_samples()
        LPS_samples = self.calc_LPS_samples()

        Ls = LPS_samples[:, 0]
        Ps = LPS_samples[:, 1]
        Ss = LPS_samples[:, 2]

        row1 = np.linalg.solve(ijk_samples, Ls)
        row2 = np.linalg.solve(ijk_samples, Ps)
        row3 = np.linalg.solve(ijk_samples, Ss)
        row4 = np.array([0, 0, 0, 1])

        IJKtoLPS_mat = np.array([row1, row2, row3, row4])

        if dst != None:
            dst = IJKtoLPS_mat

        return IJKtoLPS_mat

    def calc_ijk_samples(self):
        origin_k = len(self.dicom_file_list) - 1
        ijk = np.array([[0, 0, origin_k, 1],
                        [1, 0, origin_k, 1],
                        [0, 1, origin_k, 1],
                        [0, 0, origin_k - 1, 1]])

        return ijk

    def calc_LPS_samples(self):
        origin_L, origin_P, origin_S = self.image_position_patient

        LPS = np.array([[origin_L, origin_P, origin_S],
                        [origin_L + self.pixel_spacing_dx, origin_P, origin_S],
                        [origin_L, origin_P + self.pixel_spacing_dy, origin_S],
                        [origin_L, origin_P, origin_S - self.slice_spacing]])

        return LPS

    def load_initial_CT(self):
        CT_data = pydicom.read_file(self.dicom_file_list[0])
        CT_row = CT_data.pixel_array + CT_data.RescaleIntercept

        return self.row2uint8(CT_row)

    def row2uint8(self, CT_row, delete_0s=False):
        if delete_0s:
            CT_row = np.where(CT_row == 0, 0, CT_row - np.min(CT_row[CT_row != 0]))

        window_max = self.dcm_wc + self.dcm_ww / 2
        window_min = self.dcm_wc - self.dcm_ww / 2

        CT_img = CT_row.astype(np.float)
        CT_img[CT_img < window_min] = window_min
        CT_img[window_max < CT_img] = window_max

        CT_img -= np.mean(CT_img)
        CT_img = CT_img / (np.max(np.abs(CT_img)) + 1e-5) * 256.0
        CT_img -= np.mean(CT_img)
        CT_img = np.clip(CT_img, 0, 255).astype(np.uint8)
        CT_img = cv2.cvtColor(CT_img, cv2.COLOR_GRAY2BGR)

        return CT_img

    def ijk2relative_vec(self, ijk_coords):
        relative_vec = ijk_coords - self.model_center_vec
        relative_vec[1], relative_vec[2] = -relative_vec[1], -relative_vec[2]
        return relative_vec * self.scaling_f

    def get_CT_by_index(self, dicom_index):
        # print(self.dicom_file_list[dicom_index])
        CT_data = pydicom.read_file(self.dicom_file_list[dicom_index])
        CT_row = CT_data.pixel_array + CT_data.RescaleIntercept

        return self.row2uint8(CT_row)

    def calc_k_on_ijk_coordinates(self, dicom_index):
        return self.dicom_size - dicom_index - 1

    def linear_pad_list(self, ijk_eso_centers):
        pad_center_list = []
        Is_first = True
        prev_index = 0
        for index, eso_center in enumerate(ijk_eso_centers):
            if (eso_center is not None) and Is_first:
                prev_index = index
                Is_first = False
            elif (eso_center is not None) and (Is_first is not True):
                item_distance = index - prev_index
                if item_distance != 1:
                    average_weight = 1 / item_distance
                    for new_index in range((prev_index + 1), index, 1):
                        pad_center = (ijk_eso_centers[prev_index] * (average_weight * (index - new_index)) +
                                      ijk_eso_centers[index] * (average_weight * (new_index - prev_index)))
                        pad_center_list[new_index] = pad_center
                prev_index = index
            pad_center_list.append(eso_center)
        return pad_center_list