import nibabel as nib
from os import path
from os.path import isdir
import os
import numpy as np

def convert_nii(path_to_nii=None):
  nii=nib.load(path_to_nii)
  data=nii.get_fdata()
  shape=data.shape
  print(shape)

file_path = path.dirname(path.abspath(__file__))
root_dir=path.join(file_path, '../..')

path_to_input=path.join(root_dir, 'input')
dirs=os.listdir(path_to_input)

for dir in dirs:
  print(f'processing {dir}')

  os.makedirs(path.join(root_dir, f'output/{dir}'), exist_ok=True)
  path_to_map=path.join(path_to_input, dir, 'maps.nii.gz')

  num_chan=int(dir)
  
  nii=nib.load(path_to_map)
  nii_hdr=nii.header
  nii_hdr.set_xyzt_units('mm', 'sec')
  nii_affine=nii.affine

  chan_idx=0
  while chan_idx < num_chan:
    data=nii.dataobj[..., chan_idx:chan_idx+1]

    img=nib.Nifti1Image(data,header=nii_hdr,affine=nii_affine)
    nib.save(img, path.join(root_dir, f'output/{dir}/{str(chan_idx + 1)}.nii.gz'))
    if chan_idx % 10 == 0:
      print(f'finished processing {chan_idx}')
    chan_idx += 1
