# -*- coding: utf-8 -*-
# Author: allen
# Copyright (c) 2022 RuiCare All rights reserved.
import cv2
import random
import numpy as np
from keras.utils import Sequence
from util import fileutil, imgutil


class Generator(Sequence):

    def __init__(self,
                 train_dir,
                 batch_size=2,
                 shape=(16, 512, 512, 2),
                 ):
        self.train_dir = train_dir
        self.batch_size = batch_size
        self.shape = shape
        self.pat_dirs = None
        self.steps_per_epoch = None
        self.initial()

    def initial(self):
        self.pat_dirs = fileutil.list_dir(self.train_dir)
        assert len(self.pat_dirs) > 0, 'not found patient'
        nii_path = f'{self.pat_dirs[0]}/ct.nii.gz'
        s = imgutil.load_nii(nii_path, zfirst=True).shape
        self.steps_per_epoch = int((s[0] * s[1] * s[2]) / (
                self.shape[0] * self.shape[1] * self.shape[2]) * self.pat_dirs.__len__() / self.batch_size)

    def __len__(self):
        return self.steps_per_epoch

    def __getitem__(self, i):
        s = self.shape
        imgs = np.zeros(shape=(self.batch_size, s[0], s[1], s[2], 2), dtype=np.float16)
        masks = np.zeros(shape=(self.batch_size, s[0], s[1], s[2], 1), dtype=np.float16)
        pat_dirs = random.sample(self.pat_dirs, self.batch_size)

        for b, pat_dir in enumerate(pat_dirs):
            ct = imgutil.load_nii(f'{pat_dir}/ct.nii.gz', dtype=np.float16, zfirst=True)
            dose_plan = imgutil.load_nii(f'{pat_dir}/dose_plan.nii.gz', dtype=np.float16, zfirst=True)
            dose_qa = imgutil.load_nii(f'{pat_dir}/dose_qa.nii.gz', dtype=np.float16, zfirst=True)

            begin_z = 0 if ct.shape[0] == s[0] else random.randint(0, ct.shape[0] - s[0])

            imgs[b, :, :, :, 0] = ct[begin_z:begin_z + s[0], :, :]
            imgs[b, :, :, :, 1] = dose_plan[begin_z:begin_z + s[0], :, :]
            masks[b, :, :, :, 0] = dose_qa[begin_z:begin_z + s[0], :, :]

        return imgs, masks


if __name__ == '__main__':
    ct_gt_4d = Generator(train_dir='/home/allen/Temp/PlanQA/train',
                         batch_size=1,
                         shape=(16, 512, 512, 2))

    imgs, masks = ct_gt_4d.__getitem__(0)
    print(imgs.shape, masks.shape)

    ct = imgutil.normalize(imgs[0, 0, :, :, 0])
    plan = imgutil.normalize(imgs[0, 0, :, :, 1])
    qa = imgutil.normalize(masks[0, 0, :, :, 0])
    cv2.imwrite('../resource/generator/ct.png', ct)
    cv2.imwrite('../resource/generator/plan.png', plan)
    cv2.imwrite('../resource/generator/qa.png', qa)
