# -*- coding: utf-8 -*-
import os
import nibabel as nib
import numpy as np

# Set up the directory paths
code_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(code_dir)
file_path = os.path.join(parent_dir, "input", "C57BL6_mouse_qa.nii.gz")
output_path = os.path.join(parent_dir, "output", "C57BL6_mouse_qa.nii.gz")


# Load the image
img = nib.load(file_path)
data = img.get_fdata()

# Thresholding
threshold_value = 0.3 * np.max(data)
mask_data = np.where(data >= threshold_value, 1, 0)

# Save the mask
output_path = output_path.replace("C57BL6_mouse_qa.nii.gz", "White_matter_mask.nii.gz")
mask_img = nib.Nifti1Image(mask_data, img.affine)
nib.save(mask_img, output_path)