#!/usr/bin/python3

import os
import platipy # https://pyplati.github.io/platipy/index.html

from matplotlib import pyplot as plt

import SimpleITK as sitk

from platipy.imaging.projects.cardiac.run import run_hybrid_segmentation
from platipy.imaging import ImageVisualiser
from platipy.imaging.label.utils import get_com

### Change this ###
filename = f'volume_{1}'
### Change this ###

os.makedirs(f"./output/{filename}", exist_ok=True)
test_image = sitk.ReadImage(f"./input/volumes/{filename}.nii.gz")
auto_structures, _ = run_hybrid_segmentation(test_image)

output_directory = './output'

for struct_name in list(auto_structures.keys()):
    sitk.WriteImage(auto_structures[struct_name], str(f"{output_directory}/{filename}/{filename}_{struct_name}.nii.gz"))

print(f"Segmentations saved to: {output_directory}")

vis = ImageVisualiser(test_image, cut=get_com(auto_structures["Heart"]))
vis.add_contour({struct: auto_structures[struct] for struct in auto_structures.keys()})
fig = vis.show()

plt.savefig(f"{output_directory}/{filename}/{filename}_heart_visualization.png")
plt.close(fig)
print(f"Visualization saved to: {output_directory}/{filename}/{filename}_heart_visualization.png")