"""Step 1: masks and basic calibration checks for the NGC 5128 analysis.

Produces
  sb-mask-stars.fits   foreground stars / saturated cores / compact objects
  sb-mask-dust.fits    the dust lane, defined from the B-R colour excess
  sb-colour-excess.npy E(B-R) instrumental colour excess map (float32)

The dust mask is built from COLOUR, not from a model residual: the lane is the
only thing in the frame that is strongly red relative to the smooth stellar
body, so a colour cut is far more specific than a brightness-residual cut and
does not eat the galaxy itself.
"""
import numpy as np, sep
from astropy.io import fits
from scipy import ndimage
from sb_common import *

# ------------------------------------------------------------------ saturation
d = load('Luminance')
H, W = d.shape
print('-- saturation census (luminance master) --')
for thr in [30000, 50000, 60000, 62000, 63000, 64000, 65000, 66000]:
    print('   > %6d ADU : %6d px' % (thr, (d > thr).sum()))
SAT = 63000.0
core = ndimage.median_filter(d[int(Y0)-400:int(Y0)+400, int(X0)-400:int(X0)+400], 25)
print('   star-free (median-25) peak of galaxy light, inner 800 px: %.1f ADU' % core.max())
print('   => the galaxy core is a factor %.0f below the clip level: NOT saturated'
      % (SAT/core.max()))
del core

bkg = sep.Background(d, bw=128, bh=128)
rms = float(bkg.globalrms)
print('\nglobal sky rms %.3f ADU/px  ->  1-sigma = %.2f mag/arcsec^2' % (rms, mu(rms)))
np.save(path('_rms.npy'), np.array([rms]))

# ------------------------------------------------------------------ star mask
mask = np.zeros((H, W), bool)
z = np.load(path('_gaia_deep.npz'))
gx, gy = wcs().world_to_pixel_values(z['ra'], z['dec']); gg = z['g']
rad = np.clip(4.0 + 3.6*(16.5 - gg), 4, 120)
sel = (gx > -150) & (gx < W+150) & (gy > -150) & (gy < H+150) & (gg < 19.0)
gx, gy, rad = gx[sel], gy[sel], rad[sel]
print('\nGaia stars masked: %d  (radii %.0f-%.0f px)' % (gx.size, rad.min(), rad.max()))
for x, y, r in zip(gx, gy, rad):
    i0, i1 = max(0, int(y-r)), min(H, int(y+r)+1)
    j0, j1 = max(0, int(x-r)), min(W, int(x+r)+1)
    if i1 <= i0 or j1 <= j0: continue
    sy = np.arange(i0, i1)[:, None] - y; sx = np.arange(j0, j1)[None, :] - x
    mask[i0:i1, j0:j1] |= (sx*sx + sy*sy) < r*r

rr = np.hypot(np.arange(W)[None, :]-X0, np.arange(H)[:, None]-Y0)

# compact non-Gaia objects, but leave the crowded inner 150 px to the
# isophote fitter's own sigma clipping (masking there kills the fit)
sub = d - bkg.back()
obj, seg = sep.extract(sub, 8.0, err=rms, minarea=6, deblend_cont=0.005,
                       segmentation_map=True)
compact = obj['npix'] < 20000
segmask = ndimage.binary_dilation(np.isin(seg, np.nonzero(compact)[0]+1), np.ones((5, 5)))
mask |= segmask & (rr > 150)
print('sep compact objects: %d (applied outside r=150 px)' % compact.sum())
del sub, seg, segmask

sat = ndimage.binary_dilation(d > SAT, np.ones((5, 5)), iterations=6)
mask |= sat
print('saturated-core mask (grown 12 px): %d px' % sat.sum())
del sat, d

print('star mask: %.2f%% of frame' % (100*mask.mean()))
for a, b in [(0,25),(25,50),(50,100),(100,200),(200,400),(400,800),(800,1600)]:
    s = (rr >= a) & (rr < b)
    print('   r %4d-%4d px: %.1f%%' % (a, b, 100*mask[s].mean()))
fits.PrimaryHDU(mask.astype(np.uint8)).writeto(path('sb-mask-stars.fits'), overwrite=True)

# ------------------------------------------------------------------ dust mask
# instrumental B-R from smoothed R and B masters
R = ndimage.gaussian_filter(load('Red'), 3.0)
B = ndimage.gaussian_filter(load('Blue'), 3.0)
good = (R > 15) & (B > 4)
col = np.full((H, W), np.nan, np.float32)
col[good] = -2.5*np.log10(B[good]/R[good])
del R, B

# Unobscured baseline colour vs elliptical radius.  Dust only ever reddens, so
# the blue tail of the colour distribution in each annulus is the dust-free
# stellar colour.  The 15th percentile is measured over 150 < a < 900 px (where
# both the colour SNR is high and the lane does not fill the annulus) and fitted
# with a quadratic in log a, which is then extrapolated inwards and outwards.
a_map = ell_radius((H, W), X0, Y0, 0.17, np.radians(150.0))
bins = np.geomspace(5, 2200, 60)
ib = np.digitize(a_map, bins)
base_r, base_v = [], []
valid = good & ~mask
for k in range(1, len(bins)):
    s = (ib == k) & valid
    if s.sum() < 400: continue
    base_r.append(0.5*(bins[k-1]+bins[k])); base_v.append(np.nanpercentile(col[s], 15))
base_r, base_v = np.array(base_r), np.array(base_v)
fitr = (base_r > 150) & (base_r < 900)
pcoef = np.polyfit(np.log10(base_r[fitr]), base_v[fitr], 2)
print()
print('colour baseline: quadratic in log10(a), coeffs', np.round(pcoef, 4))
print('   baseline B-R at a = 20/50/150/400/900 px:',
      np.round(np.polyval(pcoef, np.log10([20, 50, 150, 400, 900])), 3))
base = np.polyval(pcoef, np.log10(np.clip(a_map, 5, 3000))).astype(np.float32)
exc = (col - base).astype(np.float32)               # E(B-R), instrumental
np.save(path('sb-colour-excess.npy'), exc)
np.save(path('_colbase.npy'), np.c_[base_r, base_v])
np.save(path('_colbasefit.npy'), pcoef)

# Hysteresis threshold: seed on a firm colour excess, grow into the fainter
# wings of the same connected structure.  A flat low threshold alone picks up a
# spurious ring at the edge of the colour-SNR region, so it is not used.
e0 = np.nan_to_num(exc, nan=-9.0)
seed = (e0 > 0.30) & (a_map < 650)
grow = (e0 > 0.20) & (a_map < 650)
seed = ndimage.binary_opening(seed, np.ones((5, 5)))
lab, n = ndimage.label(ndimage.binary_closing(grow, np.ones((7, 7))))
keep = np.unique(lab[seed & (lab > 0)])
dust = np.isin(lab, keep[keep > 0])
dust = ndimage.binary_closing(dust, np.ones((15, 15)))
lab, n = ndimage.label(dust); sz = np.bincount(lab.ravel()); sz[0] = 0
dust = np.isin(lab, np.nonzero(sz > 3000)[0])
print('dust mask: %d px = %.1f arcmin^2 (%.2f%% of frame)'
      % (dust.sum(), dust.sum()*PIXAREA/3600., 100*dust.mean()))
for a, b in [(0,25),(25,50),(50,100),(100,200),(200,400),(400,800)]:
    s = (a_map >= a) & (a_map < b)
    print('   dust a %4d-%4d px: %5.1f%%   star+dust %5.1f%%'
          % (a, b, 100*dust[s].mean(), 100*(dust | mask)[s].mean()))
fits.PrimaryHDU(dust.astype(np.uint8)).writeto(path('sb-mask-dust.fits'), overwrite=True)
print()
print('wrote sb-mask-stars.fits, sb-mask-dust.fits, sb-colour-excess.npy')