"""Step 4: smooth elliptical model, model-subtracted residual, and renders.

The model is the Pass B isophote table (stars and dust lane masked) turned into
a 2-D image with sb_model.build, with a Sersic extrapolation inside a = 62 px
where no dust-free azimuth exists.  Subtracting it leaves everything that is
not a smooth ellipse: the dust lane, foreground stars, and any shell, tidal
feature or halo asymmetry.

Outputs
  sb-model.fits          the smooth model
  sb-residual.fits       luminance minus model
  NGC5128-sb-model-residual.png  4-panel: data / model / residual / binned deep residual
  NGC5128-sb-residual-deep.png   heavily binned residual alone, for shell hunting
"""
import numpy as np
from astropy.io import fits
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy.ndimage import gaussian_filter
from sb_common import *
import sb_model

XC, YC = np.load(path('_geom.npy'))
P = np.load(path('_profile.npz'))
tB = dict(np.load(path('_isoB.npz')))
star = fits.getdata(path('sb-mask-stars.fits')).astype(bool)
dust = fits.getdata(path('sb-mask-dust.fits')).astype(bool)
A_IN, A_OUT = float(P['a_in']), float(P['a_out'])
mue, re_as, nser = P['sersic']


def sersic_mu(a_as):
    bn = 2 * nser - 1 / 3. + 0.009876 / nser
    return mue + 2.5 * bn / np.log(10) * ((a_as / re_as) ** (1. / nser) - 1.)


# ---- build a profile that is defined at every radius -----------------------
ac = P['a']
Lp = P['Luminance'].copy()
inner = ac < A_IN
Lp[inner] = 10 ** ((MU0 - sersic_mu(ac[inner] * PIXSCALE)) / 2.5)
okp = np.isfinite(Lp) & (ac < 1500)
tab = dict(sma=ac[okp], intens=Lp[okp],
           eps=np.interp(ac[okp], tB['sma'], np.where(np.isfinite(tB['eps']),
                                                      tB['eps'], 0.15)),
           pa=np.interp(ac[okp], tB['sma'], np.where(np.isfinite(tB['pa']),
                                                     tB['pa'], 150.)))
model, a_map = sb_model.build((3194, 4788), XC, YC, tab, block=2)
fits.PrimaryHDU(model).writeto(path('sb-model.fits'), overwrite=True)

PEDL = float(P['ped'])
L = load('Luminance')
res = (L - model).astype(np.float32)
fits.PrimaryHDU(res).writeto(path('sb-residual.fits'), overwrite=True)
print('model built; residual rms inside a<600 px: %.2f ADU/px'
      % res[(a_map < 600) & ~star & ~dust].std())


def binned(img, mask, B):
    """Masked block mean, returning NaN where a block is mostly masked."""
    H, W = img.shape
    h, w = H // B, W // B
    a = img[:h * B, :w * B].reshape(h, B, w, B)
    m = (~mask)[:h * B, :w * B].reshape(h, B, w, B)
    n = m.sum(axis=(1, 3))
    s = np.where(m, a, 0).sum(axis=(1, 3))
    return np.where(n > 0.35 * B * B, s / np.maximum(n, 1), np.nan)




# ------------------------------------------------------------------ renders
import sb_render_tail
sb_render_tail.run(dict(res=res, model=model, a_map=a_map, L=L, star=star,
                        dust=dust, XC=XC, YC=YC, PEDL=PEDL, A_OUT=A_OUT,
                        tab=tab, binned=binned))