import nibabel as nib
import numpy as np

# Load NIfTI viral tracing data
viral_tracing_img = nib.load(r"E:\2024_Ruthe_SND\output\Viral_tracing_flipped\Mop_projection_density_VT_mask.nii.gz"")
viral_tracing_data = viral_tracing_img.get_fdata()

# Load NIfTI mask data
mask_img = nib.load(r"C:\Users\aswen\Desktop\Code\2024_Ruthe_SND\output\Tract_Mask_registered\CC_Mop_dilated_registered.nii")
mask_data = mask_img.get_fdata()

# Load template
template_img = nib.load(r"C:\Users\aswen\Desktop\Code\2024_Ruthe_SND\output\C57BL6_mouse.iso_registerd.nii.gz")
template_data = template_img.get_fdata()

# Check dimensions and shapes to ensure they align if needed
print("Viral Tracing Data Shape:", viral_tracing_data.shape)
print("Mask Data Shape:", mask_data.shape)
print("Template Data Shape:", template_data.shape)

# Now you can put them together based on your requirements
# For example, if you want to overlay viral tracing data on the template using the mask
# Assuming all data have the same dimensions

# Apply mask to viral tracing data
viral_tracing_masked = np.where(mask_data > 0, viral_tracing_data, 0)

# Overlay masked viral tracing data on the template
final_data = template_data + viral_tracing_masked

# Save or further process the final_data as needed