import argparse
import numpy as np
import nibabel as nib
import torch


parser = argparse.ArgumentParser()
parser.add_argument(
    '--dir',
    type=str,
    default=r'D:\DATA\MRI_Datasets\CSM\PA1\test\t2/',
    help='feature matrix directory')
parser.add_argument(
    '--c',
    type=str,
    default='t2s',
    help='feature matrix directory')
parser.add_argument('--random_state',
                    type=int,
                    default=42,
                    help='seed for random initialisation')
args = parser.parse_args()

patient_path = args.dir +'/'
img_path = patient_path + 't2s.nii.gz'
seg_label_path = [0]*5
for c in [3, 4, 5, 6]:
    seg_label_path[c-3] = patient_path + \
        't2andt2s_t2s_labeled_c' + str(c) + '.nii.gz'
seg_path = patient_path + 't2s_seg.nii.gz'
# nii2tensor
img_nii = nib.load(img_path)
img = torch.Tensor(np.transpose(img_nii.get_fdata(), [2, 1, 0]))

img_shape = np.transpose(img_nii.get_fdata(), [2, 1, 0]).shape
seg_nii = nib.load(seg_path)
img_seg = torch.Tensor(np.transpose(seg_nii.get_fdata(), [2, 1, 0]))
# segmenta
for j in range(4):
    nii = nib.load(seg_label_path[j])
    img_seg_label = np.transpose(nii.get_fdata(), [2, 1, 0])
    label_n = img_shape[1]
    seg_label = np.zeros((img_shape[0], label_n, label_n))
    for k in range(img_seg_label.shape[0]):
        if img_seg_label[k, :, :].max() > 10e-3:
            seg_label[k, :, :] = np.eye(label_n)
    seg_label = torch.Tensor(seg_label)
    img_ = torch.bmm(seg_label, img)
    seg_ = img_ * img_seg
    seg_ = np.transpose(seg_.numpy(), [2, 1, 0])
    Niinew = nib.Nifti1Image(seg_, seg_nii.affine)
    nib.save(Niinew, patient_path + 't2s_seg_c'+str(j+3)+'.nii.gz')
