import matplotlib.pyplot as plt

from pylinac.core.profile import FWXMProfile, ProfileMetric
from pylinac.core.image_generator import AS1000Image, FilteredFieldLayer, GaussianFilterLayer
from pylinac.core.array_utils import normalize


# same as above; included so we can plot
class CenterMetric(ProfileMetric):
    name = 'Center Index'  # human-readable string

    def calculate(self) -> float:
        """Return the index of the center of the profile."""
        return self.profile.center_idx

    def plot(self, axis: plt.Axes) -> None:
        """Plot the center index."""
        axis.plot(self.profile.center_idx, self.profile.y_at_x(self.profile.center_idx), 'o', color='red',
                  markersize=10, label=self.name)

# this is our set up to get a nice profile
as1000 = AS1000Image()
as1000.add_layer(
    FilteredFieldLayer(field_size_mm=(100, 100))
)
as1000.add_layer(
    GaussianFilterLayer(sigma_mm=2)
)  # add an image-wide gaussian to simulate penumbra/scatter

# pull out the profile array
array = normalize(as1000.image[:, as1000.shape[1] // 2])

# create the profile
profile = FWXMProfile(array)

# compute the metric with our plugin
profile.compute(metrics=CenterMetric())

# plot the profile
profile.plot()