import matplotlib.pyplot as plt
import numpy as np
from astropy.modeling.models import Gaussian2D
from photutils.datasets import make_noise_image
from photutils.isophote import (Ellipse, EllipseGeometry,
                                build_ellipse_model)

g = Gaussian2D(100.0, 75, 75, 20, 12, theta=np.deg2rad(40.0))
ny = nx = 150
y, x = np.mgrid[0:ny, 0:nx]
noise = make_noise_image((ny, nx), distribution='gaussian', mean=0.0,
                         stddev=2.0, seed=1234)
data = g(x, y) + noise
geometry = EllipseGeometry(x0=75, y0=75, sma=20, eps=0.5,
                           pa=np.deg2rad(20.0))
ellipse = Ellipse(data, geometry=geometry)
isolist = ellipse.fit_image()

model_image = build_ellipse_model(data.shape, isolist)
residual = data - model_image

fig, (ax1, ax2, ax3) = plt.subplots(ncols=3, figsize=(14, 5))
fig.subplots_adjust(left=0.04, right=0.98, bottom=0.02, top=0.98)
ax1.imshow(data, origin='lower')
ax1.set_title('Data')

smas = np.linspace(10, 50, 5)
for sma in smas:
    iso = isolist.get_closest(sma)
    x, y, = iso.sampled_coordinates()
    ax1.plot(x, y, color='white')

ax2.imshow(model_image, origin='lower')
ax2.set_title('Ellipse Model')

ax3.imshow(residual, origin='lower')
ax3.set_title('Residual')