"""Step 2: isophote fitting of the NGC 5128 luminance master.

Pass A  free-centre fits on the dust-free outer body            -> adopted centre
Pass B  fixed-centre fit, stars AND the dust lane masked        -> PRIMARY table
Pass C  the same fit on a de-reddened image, where the
        extinction is estimated from the B-R colour excess and
        calibrated against the Pass B model over 150<a<600 px   -> inner extension

Outputs: sb-isophotes.csv, sb-isophotes-dereddened.csv,
         _isoB.npz, _isoC.npz, _geom.npy, _extcal.npy
"""
import numpy as np
import time
from astropy.io import fits
from photutils.isophote import Ellipse, EllipseGeometry
from sb_common import *
import sb_model

SMA_MIN, SMA_MAX, STEP = 8.0, 1750.0, 0.11

L = load('Luminance')
starmask = fits.getdata(path('sb-mask-stars.fits')).astype(bool)
dustmask = fits.getdata(path('sb-mask-dust.fits')).astype(bool)
exc = np.load(path('sb-colour-excess.npy'))


def table(iso):
    k = [i for i in iso if i.sma > 0 and np.isfinite(i.intens)]

    def g(f):
        return np.array([(f(i) if f(i) is not None else np.nan) for i in k], float)

    return dict(sma=g(lambda i: i.sma), intens=g(lambda i: i.intens),
                int_err=g(lambda i: i.int_err), rms=g(lambda i: i.rms),
                eps=g(lambda i: i.eps), eps_err=g(lambda i: i.ellip_err),
                pa=np.degrees(g(lambda i: i.pa)) % 180.,
                pa_err=np.degrees(g(lambda i: i.pa_err)),
                ndata=g(lambda i: i.ndata), nflag=g(lambda i: i.nflag),
                stop=g(lambda i: i.stop_code))


def fit(img, mask, x0, y0, label):
    arr = np.ma.masked_array(img, mask=mask)
    g = EllipseGeometry(x0=x0, y0=y0, sma=300., eps=0.15, pa=np.radians(150.))
    g.fix_center = True
    t = time.time()
    iso = Ellipse(arr, geometry=g).fit_image(
        sma0=300., minsma=SMA_MIN, maxsma=SMA_MAX, step=STEP, linear=False,
        nclip=3, sclip=3.0, fix_center=True)
    print('%s: %d isophotes in %.0f s' % (label, len(iso), time.time() - t))
    return table(iso)


# ------------------------------------------------------------------- Pass A
arr = np.ma.masked_array(L, mask=starmask | dustmask)
cen = []
for s in [350., 500., 650., 800., 1000.]:
    try:
        it = Ellipse(arr, geometry=EllipseGeometry(
            x0=X0, y0=Y0, sma=s, eps=0.18, pa=np.radians(148.))
        ).fit_image(sma0=s, minsma=s * 0.98, maxsma=s * 1.02, step=0.1,
                    nclip=3, sclip=3.)
        for i in it:
            if np.isfinite(i.x0):
                cen.append((i.x0, i.y0))
    except Exception as e:
        print('  passA sma=%.0f: %s' % (s, e))
cen = np.array(cen)
XC, YC = float(np.median(cen[:, 0])), float(np.median(cen[:, 1]))
off = np.hypot(XC - X0, YC - Y0)
print('Pass A: outer-isophote centre %.2f, %.2f  (scatter %.1f, %.1f px, n=%d)'
      % (XC, YC, cen[:, 0].std(), cen[:, 1].std(), len(cen)))
print('        WCS/Gaia nucleus      %.2f, %.2f  -> offset %.1f px = %.1f arcsec'
      % (X0, Y0, off, off * PIXSCALE))
np.save(path('_geom.npy'), np.array([XC, YC]))
del arr

# ------------------------------------------------------------------- Pass B
tB = fit(L, starmask | dustmask, XC, YC, 'Pass B (stars+dust masked)')
np.savez(path('_isoB.npz'), **tB)

# ---------------------------------------------------------------- extinction
# A_L from the Pass B model, used only where that model is directly constrained
good = np.isfinite(tB['intens']) & (tB['ndata'] > 150) & (tB['sma'] > 90)
tBg = {k: v[good] for k, v in tB.items()}
modB, aB = sb_model.build(L.shape, XC, YC, tBg, block=4)
with np.errstate(all='ignore'):
    A_L = -2.5 * np.log10(np.clip(L, 1e-3, None) / np.clip(modB, 1e-3, None))
cal = (dustmask & ~starmask & (aB > 150) & (aB < 600) & np.isfinite(exc)
       & (exc > 0.05) & np.isfinite(A_L) & (A_L > -0.5) & (A_L < 4.0))
cal &= exc > 0.15                       # restrict to a well-measured colour excess
x, y = exc[cal].astype(float), A_L[cal].astype(float)
# robust slope through the origin: median of the per-pixel ratios
k_ratio = float(np.median(y / x))
scatter = float(np.median(np.abs(y - k_ratio * x)) * 1.4826)
print('extinction calibration on %d px: A_L = %.3f * E(B-R),  scatter %.3f mag'
      % (cal.sum(), k_ratio, scatter))
print('   (least-squares through origin for comparison: %.3f)'
      % (np.sum(x * y) / np.sum(x * x)))
np.save(path('_extcal.npy'), np.array([k_ratio, scatter, cal.sum()]))

E = np.clip(np.nan_to_num(exc, nan=0.0), 0.0, None)
A_est = np.clip(k_ratio * E, 0.0, 2.5)                  # >2.5 mag is unreliable
heavy = (k_ratio * E) > 2.5
Lc = (L * 10 ** (0.4 * A_est)).astype(np.float32)
print('de-reddening: median A_L inside the lane mask %.2f mag; %d px above the '
      '2.5 mag cap (masked in Pass C)' % (np.median(A_est[dustmask]), heavy.sum()))
del modB, aB, A_L, A_est, E

# ------------------------------------------------------------------- Pass C
tC = fit(Lc, starmask | heavy, XC, YC, 'Pass C (de-reddened)')
np.savez(path('_isoC.npz'), **tC)
del Lc, L

# ------------------------------------------------------------------- CSVs
write_isophote_csv(tB, 'sb-isophotes.csv')
write_isophote_csv(tC, 'sb-isophotes-dereddened.csv')

for name, t in [('Pass B (primary, dust masked)', tB), ('Pass C (de-reddened)', tC)]:
    print('')
    print(name)
    print(' sma_px  arcsec     mu    eps     PA  ndata nflag stop')
    for i in range(len(t['sma'])):
        if t['sma'][i] > 30 and i % 3:
            continue
        print('%7.1f %7.1f %6.2f %6.3f %6.1f %6d %5d %4d'
              % (t['sma'][i], t['sma'][i] * PIXSCALE, mu(t['intens'][i]),
                 t['eps'][i], t['pa'][i], t['ndata'][i], t['nflag'][i], t['stop'][i]))