import os
import json
import numpy as np
import argparse
import SimpleITK as sitk
import cv2
from PIL import Image, ImageDraw, ImageFont
import imageio

class_tag_1 = {1: '磨玻璃结节',
               2: '实性结节',
               3: '混合结节',
               4: '钙化结节'}

class_tag_2 = {1: '良性',
               2: '恶性'}

def norm(image, hu_min=-1000.0, hu_max=600.0):
    image = (np.clip(image.astype(np.float32), hu_min, hu_max) - hu_min) / float(hu_max - hu_min)
    return image * 255.

def vis(args):
    # 按病人id和时间整理匹配结果
    followup_root = os.path.join(args.job_data_root, 'output/elastix')
    files = os.listdir(followup_root)
    nii_path = os.path.join(args.job_data_root, 'output/preprocess/preprocess_file_details/nii')
    nii_files = os.listdir(nii_path)

    working_dicts = {}
    for name in files:
        patientid = name.split('_')[0].split('-')[-1]
        if patientid not in working_dicts.keys():
            working_dicts[patientid] = []
            working_dicts[patientid].append(os.path.join(followup_root, name))
        else:
            working_dicts[patientid].append(os.path.join(followup_root, name))
    
    for key in list(working_dicts.keys()):
        working_dicts[key] = sorted(working_dicts[key])

    # 计算每一个病人的信息，（size、类别）
    working_infos = {}
    for key in list(working_dicts.keys()):
        files = working_dicts[key]
        for file in files:
            working_infos[file] = []
            temp= {'info': [],
                   'matched': []}
            match_results = json.load(open(os.path.join(file, 'match_results.json'), 'r'))
            temp['info'] = match_results['info']
            if len(match_results['matched']) != 0:
                temp['matched'] = match_results['matched']
        
            working_infos[file].append(temp)
        

    for file in list(working_infos.keys()):
        fix_uid = working_infos[file][0]['info'][0]
        fix_time = working_infos[file][0]['info'][1]
        matched_pairs = working_infos[file][0]['matched']
        moving_uid = working_infos[file][0]['info'][2]
        moving_time = working_infos[file][0]['info'][3]
        fix_nii = [os.path.join(nii_path, file) for file in nii_files if fix_uid in file][0]
        moving_nii = [os.path.join(nii_path, file) for file in nii_files if moving_uid in file][0]

        fix_img = sitk.ReadImage(fix_nii)
        fix_array = sitk.GetArrayFromImage(fix_img)
        fix_origin = fix_img.GetOrigin()
        fix_spacing = fix_img.GetSpacing()

        moving_img = sitk.ReadImage(moving_nii)
        moving_array = sitk.GetArrayFromImage(moving_img)
        moving_origin = moving_img.GetOrigin()
        moving_spacing = moving_img.GetSpacing()

        for idx, match in enumerate(matched_pairs):
            p1 = match[0][:-2]
            p2 = match[1][:-2]
            p1[:3] = (np.array(p1[:3]) - np.array(list(fix_origin))) / np.array(list(fix_spacing))
            p2[:3] = (np.array(p2[:3]) - np.array(list(moving_origin))) / np.array(list(moving_spacing))

            fix_x1 = int(max(0, p1[0] - p1[3] / 2))
            fix_y1 = int(max(0, p1[1] - p1[4] / 2))
            fix_z1 = int(max(0, p1[2] - p1[5] / 2))
            fix_x2 = int(min(fix_array.shape[2] - 1, p1[0] + p1[3] / 2))
            fix_y2 = int(min(fix_array.shape[1] - 1, p1[1] + p1[4] / 2))
            fix_z2 = int(min(fix_array.shape[0] - 1, p1[2] + p1[5] / 2))
            
            fix_images = []
            z_min = max(0, fix_z1-4)
            z_max = min(fix_array.shape[0]-1, fix_z2+4)
            for z in range(z_min, z_max, 1):
                img_z = fix_array.copy()[z, :, :]
                img_z = norm(img_z)
                img_z = cv2.cvtColor(img_z, cv2.COLOR_GRAY2BGR)
                # cv2.rectangle(img_z, (int(y1), int(x1)), (int(y2), int(x2)), (255,0,0), 2)
                if z>=fix_z1 and z<=fix_z2:
                    cv2.rectangle(img_z, (int(fix_x1), int(fix_y1)), (int(fix_x2), int(fix_y2)), (255,0,0), 2)

                cut_x_min = max(0, fix_x1-128)
                cut_y_min = max(0, fix_y1-128)
                cut_x_max = min(fix_array.shape[2]-1, fix_x2+128)
                cut_y_max = min(fix_array.shape[1]-1, fix_y2+128)

                img_z = img_z[cut_y_min: cut_y_max, cut_x_min: cut_x_max, :]
                # img_z = img_z[cut_x_min: cut_x_max, cut_y_min: cut_y_max, :]
                img_pil = Image.fromarray(cv2.cvtColor(img_z, cv2.COLOR_BGR2RGB).astype(np.uint8))
                draw = ImageDraw.Draw(img_pil)
                font = cv2.FONT_HERSHEY_SIMPLEX
                font_scale = 0.2
                font_color = (255,0,0)
                font_thickness = 1
                font_size = 10
                text_1 = fix_uid
                (_, text_h_1), _ = cv2.getTextSize(text_1, font, font_scale, font_thickness)
                adjusted_position_1 = (0, 0+text_h_1)
                text_2 = fix_time
                (_, text_h_2), _ = cv2.getTextSize(text_2, font, font_scale, font_thickness)
                adjusted_position_2 = (0, 0+text_h_1+text_h_2+5)
                text_3 = str(0)
                (_, text_h_3), _ = cv2.getTextSize(text_3, font, font_scale, font_thickness)
                adjusted_position_3 = (0, 0+text_h_1+text_h_2+text_h_3+10)
                img_z = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
                cv2.putText(img_z, text_1, adjusted_position_1, font, font_scale, font_color, font_thickness)
                cv2.putText(img_z, text_2, adjusted_position_2, font, font_scale, font_color, font_thickness)
                cv2.putText(img_z, text_3, adjusted_position_3, font, font_scale, font_color, font_thickness)
                fix_images.append(img_z)

            fps = 24.0
            imageio.mimsave(os.path.join(file, fix_uid + ':' + fix_time + '-' + str(idx) + '.gif'), fix_images, 'GIF', duration=1/fps)

            moving_x1 = int(max(0, p2[0] - p2[3] / 2))
            moving_y1 = int(max(0, p2[1] - p2[4] / 2))
            moving_z1 = int(max(0, p2[2] - p2[5] / 2))
            moving_x2 = int(min(moving_array.shape[2] - 1, p2[0] + p2[3] / 2))
            moving_y2 = int(min(moving_array.shape[1] - 1, p2[1] + p2[4] / 2))
            moving_z2 = int(min(moving_array.shape[0] - 1, p2[2] + p2[5] / 2))
            
            moving_images = []
            z_min = max(0, moving_z1-4)
            z_max = min(moving_array.shape[0]-1, moving_z2+4)
            for z in range(z_min, z_max, 1):
                img_z = moving_array.copy()[z, :, :]
                img_z = norm(img_z)
                img_z = cv2.cvtColor(img_z, cv2.COLOR_GRAY2BGR)
                # cv2.rectangle(img_z, (int(y1), int(x1)), (int(y2), int(x2)), (0,0,255), 2)
                if z>=moving_z1 and z<=moving_z2:
                    cv2.rectangle(img_z, (int(moving_x1), int(moving_y1)), (int(moving_x2), int(moving_y2)), (255,0,0), 2)

                cut_x_min = max(0, moving_x1-128)
                cut_y_min = max(0, moving_y1-128)
                cut_x_max = min(moving_array.shape[2]-1, moving_x2+128)
                cut_y_max = min(moving_array.shape[1]-1, moving_y2+128)

                img_z = img_z[cut_y_min: cut_y_max, cut_x_min: cut_x_max, :]
                # img_z = img_z[cut_x_min: cut_x_max, cut_y_min: cut_y_max, :]
                img_pil = Image.fromarray(cv2.cvtColor(img_z, cv2.COLOR_BGR2RGB).astype(np.uint8))
                draw = ImageDraw.Draw(img_pil)
                text_1 = moving_uid
                (_, text_h_1), _ = cv2.getTextSize(text_1, font, font_scale, font_thickness)
                adjusted_position_1 = (0, 0+text_h_1)
                text_2 = moving_time
                (_, text_h_2), _ = cv2.getTextSize(text_2, font, font_scale, font_thickness)
                adjusted_position_2 = (0, 0+text_h_1+text_h_2+5)
                text_3 = str(idx)
                (_, text_h_3), _ = cv2.getTextSize(text_3, font, font_scale, font_thickness)
                adjusted_position_3 = (0, 0+text_h_1+text_h_2+text_h_3+10)
                font = cv2.FONT_HERSHEY_SIMPLEX
                font_scale = 0.2
                font_color = (255,0,0)
                font_thickness = 1
                font_size = 10
                img_z = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
                cv2.putText(img_z, text_1, adjusted_position_1, font, font_scale, font_color, font_thickness)
                cv2.putText(img_z, text_2, adjusted_position_2, font, font_scale, font_color, font_thickness)
                cv2.putText(img_z, text_3, adjusted_position_3, font, font_scale, font_color, font_thickness)
                moving_images.append(img_z)

            fps = 24.0
            imageio.mimsave(os.path.join(file, moving_uid + ':' + moving_time + '-' + str(idx) + '.gif'), moving_images, 'GIF', duration=1/fps)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='followup processing')
    parser.add_argument('--job_data_root', default='/data/job_715/job_data_preprocess', type=str, help='model')
    args = parser.parse_args()

    vis(args)