from pylinac.core.image_generator import AS1000Image, FilteredFieldLayer, GaussianFilterLayer
from pylinac.core.image import DicomImage
from pylinac.metrics.image import MetricBase

class ImageCenterMetric(MetricBase):
    name = "Image Center"

    def calculate(self):
        return self.image.center

    def plot(self, axis: plt.Axes):
        axis.plot(self.image.center.x, self.image.center.y, 'rx', markersize=10)

# now we create an image to compute over
as1000 = AS1000Image(sid=1000)  # this will set the pixel size and shape automatically
as1000.add_layer(
    FilteredFieldLayer(field_size_mm=(100, 100))
)  # create a 100x100mm square field
as1000.add_layer(
    GaussianFilterLayer(sigma_mm=2)
)  # add an image-wide gaussian to simulate penumbra/scatter
ds = as1000.as_dicom()

# now we can compute the metric on the image
img = DicomImage.from_dataset(ds)
center = img.compute(metrics=ImageCenterMetric())
print(center)
img.plot()