import os
import nibabel as nib
import numpy as np
import argparse

parser = argparse.ArgumentParser()
parser.add_argument(
    '--dir',
    type=str,
    default=r'D:\DATA\Mri_temp_c\ALS\659141/test\t2/',
    help='feature matrix directory')
parser.add_argument('--random_state',
                    type=int,
                    default=42,
                    help='seed for random initialisation')
args = parser.parse_args()
seg_label_path = [0]*4
for c in [3, 4, 5, 6]:
    seg_label_path[c-3] = args.dir + \
        't2s_seg_c' + str(c) + '.nii.gz'

for j in range(4):
    if 't2s_seg_c' + str(j+3) not in os.listdir(args.dir):
        os.mkdir(args.dir + 't2s_seg_c' + str(j+3))
    n = 0
    nii = nib.load(seg_label_path[j])
    img_seg_label = np.transpose(nii.get_fdata(), [2, 1, 0])
    a = np.zeros(img_seg_label.shape)
    for k in range(img_seg_label.shape[0]):
        if img_seg_label[k, :, :].max() > 10e-3:
            n = n+1
            a[k, :, :] = img_seg_label[k, :, :]
            b=np.transpose(a, [2, 1, 0])
            niiNew = nib.Nifti1Image(b, nii.affine)
            nib.save(niiNew, args.dir + 't2s_seg_c' + str(j+3)+'/'
                     + str(n) + '.nii.gz')
