import matplotlib.pyplot as plt
import numpy as np
from astropy.modeling.models import Gaussian2D
from photutils.centroids import centroid_2dg
from photutils.datasets import make_noise_image
from photutils.profiles import EnsquaredCurveOfGrowth

# Create an artificial single source
gmodel = Gaussian2D(42.1, 47.8, 52.4, 4.7, 4.7, 0)
yy, xx = np.mgrid[0:100, 0:100]
data = gmodel(xx, yy)
bkg_sig = 2.1
noise = make_noise_image(data.shape, mean=0., stddev=bkg_sig, seed=0)
data += noise
error = np.zeros_like(data) + bkg_sig

# Find the source centroid
xycen = centroid_2dg(data)

# Create the ensquared curve of growth
half_sizes = np.arange(1, 26)
ecog = EnsquaredCurveOfGrowth(data, xycen, half_sizes, error=error)
ecog.normalize(method='max')
ee_half_sizes = np.array([3, 6, 9])
ee_vals = ecog.calc_ee_at_half_size(ee_half_sizes)

# Plot the ensquared curve of growth
fig, ax = plt.subplots(figsize=(8, 6))
ecog.plot(ax=ax, label='Ensquared Curve of Growth')
ecog.plot_error(ax=ax)
ax.legend()

xmin, xmax = ax.get_xlim()
ymin, ymax = ax.get_ylim()
ax.vlines(ee_half_sizes, ymin, ee_vals, colors='C1', linestyles='dashed')
ax.hlines(ee_vals, xmin, ee_half_sizes, colors='C1', linestyles='dashed')
ax.set_xlim(xmin, xmax)
ax.set_ylim(ymin, ymax)

for ee_hs, ee_val in zip(ee_half_sizes, ee_vals):
    ax.text(ee_hs/2, ee_val, f'{ee_val:.2f}', ha='center',
            va='bottom')