
import os
import ants

from glob import glob

from dlaais.dlaais.atlas import get_atlas_nii_file
from dlaais.dlaais.ants.registration_antspy import rigid_reg_4d_nifti
from dlaais.dlaais.ctp.brain_mask import generate_brain_mask

#t1w_nii_file = get_atlas_nii_file('mni_icbm152_t1_tal_nlin_sym_09c')
#ref_image_ants = ants.image_read(t1w_nii_file)

import numpy as np
import nibabel as nib

#--------------------------------------------------------------------------------------------------------
#
def generate_brain_mask(ctp_nii_file, mask_nii_file):
    """
    """
    nib_ctp = nib.load(ctp_nii_file)
    np_ctp = nib_ctp.get_fdata()
    np_max = np.max(np_ctp, axis=-1)
    np_mean = np.mean(np_ctp, axis=-1)
    np_working = np_max - np_mean
    np_mask = np.zeros(np_working.shape, dtype=np.int8)
    np_mask[np_working>10] = 1
    nib_mask = nib.Nifti1Image(np_mask, affine=nib_ctp.affine)
    nib.save(nib_mask, mask_nii_file)

"""
working_root = '/data/public/data/003_STROKE_CTP/002_NIIX_SITE005'
CTP_nii_files = glob(os.path.join(working_root, 'PID*PNAME*', '*', 'CTP_1mm.nii.gz'))
failed_cases = []
for CTP_nii_file in CTP_nii_files:
    series_root = os.path.dirname(CTP_nii_file)
    date = os.path.basename(series_root)
    pname = os.path.basename(os.path.dirname(series_root))
    print('working on %s - %s'%(pname, date))
    try:
        rigid_reg_4d_nifti(os.path.dirname(CTP_nii_file), 'CTP_1mm', ref_nii_file='First')
    except:
        failed_cases.append(pname)
"""
#ctp_nii_file = '/data/public/data/003_STROKE_CTP/002_NIIX_SITE005/PID_3052429_PNAME_ZHANG_BING_LIANG/20210813/CTP_1mm_mc.nii.gz'
#mask_nii_file = '/data/public/data/003_STROKE_CTP/002_NIIX_SITE005/PID_3052429_PNAME_ZHANG_BING_LIANG/20210813/CTP_1mm_mc_mask.nii.gz'
#generate_brain_mask(ctp_nii_file, mask_nii_file)