"""Step 6: how deep does the shell search actually go, and is anything there?

Two questions:
  1. On what surface-brightness level would a shell have had to sit to be seen?
     Binning the residual to ever coarser scales shows whether the noise
     integrates down like photon noise (it does not: it is dominated by
     correlated large-scale systematics), which sets the real limit.
  2. Is there any significant azimuthal structure?  The m = 1..4 Fourier
     amplitudes of the residual in each elliptical annulus are compared with
     the amplitude expected from the noise alone.
"""
import numpy as np
from astropy.io import fits
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from sb_common import *

XC, YC = np.load(path('_geom.npy'))
star = fits.getdata(path('sb-mask-stars.fits')).astype(bool)
dust = fits.getdata(path('sb-mask-dust.fits')).astype(bool)
a_map = np.load(path('_amap.npy'))
res = fits.getdata(path('sb-residual-flat.fits')).astype(np.float32)
rms = float(np.load(path('_rms.npy'))[0])


def blockstat(img, mask, B, sel):
    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)
    s = sel[:h * B, :w * B].reshape(h, B, w, B)
    n = m.sum(axis=(1, 3))
    v = np.where(n > 0.35 * B * B, np.where(m, a, 0).sum(axis=(1, 3)) /
                 np.maximum(n, 1), np.nan)
    keep = np.isfinite(v) & (s.mean(axis=(1, 3)) > 0.8)
    return v[keep]


print('depth of the shell search, measured on the model-subtracted residual')
print('(outer field, 1200 < a < 2400 px, stars and the dust lane excluded)')
print('')
print(' bin    bin size    rms        3 sigma limit   ideal if noise were white')
print(' [px]   [arcsec]    [ADU/px]   [mag/arcsec2]   [mag/arcsec2]')
sel = (a_map > 1200) & (a_map < 2400)  # outer field, beyond the measured profile
base = None
rows = []
for B in [8, 16, 32, 64, 128]:
    v = blockstat(res, star | dust, B, sel)
    if v.size < 30:
        continue
    sd = float(np.std(v))
    if base is None:
        base, base_b = sd, B
    ideal = base * (base_b / B)
    print('  %4d   %7.1f    %8.3f   %13.2f   %13.2f'
          % (B, B * PIXSCALE, sd, mu(3 * sd), mu(3 * ideal)))
    rows.append((B, sd))
print('')
print('The rms barely falls as the bins grow (%.2f -> %.2f ADU/px from 8 to 128 px'
      % (rows[0][1], rows[-1][1]))
print('bins, against a factor 16 if it were white), so the floor is correlated')
print('large-scale structure -- flat-field residual plus the sky-plane')
print('systematic -- not photon noise.  Pure photon noise would reach')
print('%.2f mag/arcsec2 at 128 px bins; the real limit is %.1f mag SHALLOWER.'
      % (mu(3 * rms / 128), mu(3 * rms / 128) - mu(3 * rows[-1][1])))

# ------------------------------------------------------- Fourier amplitudes
print('')
print('azimuthal Fourier amplitudes of the residual, normalised to the model')
resd = np.load(path('_resd16.npy'))
a16 = np.load(path('_a16.npy'))
H, W = resd.shape
Y, X = np.mgrid[0:H, 0:W]
phi = np.arctan2(Y * 16 + 8 - YC, X * 16 + 8 - XC)
P = np.load(path('_profile.npz'))
prof, ac = P['Luminance'], P['a']
okp = np.isfinite(prof) & (prof > 0)

# stop where the measured profile itself runs out: beyond that I(a) is a fill
# value and the fractional amplitudes would be meaningless
A_OUT = float(P['a_out'])
edges = np.geomspace(150, A_OUT, 11)
out = []
print('   a range [px]     I(a)     m=1     m=2     m=3     m=4   noise   n')
for lo, hi in zip(edges[:-1], edges[1:]):
    m = np.isfinite(resd) & (a16 >= lo) & (a16 < hi)
    n = m.sum()
    if n < 60:
        continue
    r, ph = resd[m], phi[m]
    amp = [2 * np.abs(np.mean(r * np.exp(-1j * k * ph))) for k in (1, 2, 3, 4)]
    noise = np.std(r) * np.sqrt(2. / n) * 2
    Im = np.interp(np.sqrt(lo * hi), ac[okp], prof[okp])
    out.append((np.sqrt(lo * hi), Im, amp, noise, n))
    print('  %5.0f-%5.0f  %8.2f  ' % (lo, hi, Im) +
          '  '.join('%6.3f' % (a / max(Im, 1e-3)) for a in amp) +
          '  %6.3f  %5d' % (noise / max(Im, 1e-3), n))
print('')
print('(amplitudes are fractional: A_m / I(a).  A value is only meaningful if it')
print(' exceeds the "noise" column, which is the amplitude a pure-noise annulus')
print(' would produce.)')

fig, ax = plt.subplots(figsize=(9.5, 6.4))
aa = np.array([o[0] for o in out]) * PIXSCALE
for k in range(4):
    ax.plot(aa, [o[2][k] / max(o[1], 1e-3) for o in out], 'o-', ms=4,
            label='m = %d' % (k + 1))
ax.plot(aa, [o[3] / max(o[1], 1e-3) for o in out], 'k--', lw=1.6,
        label='formal noise expectation (lower bound: it assumes' + chr(10) +
              'independent bins and ignores correlated systematics)')
ax.set_xscale('log')
ax.set_yscale('log')
ax.set_xlabel('semi-major axis  a  [arcsec]')
ax.set_ylabel('fractional Fourier amplitude  $A_m / I(a)$')
ax.set_title('NGC 5128: azimuthal structure in the model-subtracted residual' +
             chr(10) + 'amplitudes are 4-10% of the local surface brightness at '
             'every radius' + chr(10) + 'inside ~350 arcsec this is demonstrably '
             'the dust lane; outside it, correlated systematics')
ax.grid(alpha=.3, which='both')
ax.legend(fontsize=8.5, loc='lower left')
fig.tight_layout()
fig.savefig(path('NGC5128-sb-residual-fourier.png'), dpi=140)
plt.close(fig)

with open(path('sb-derived-quantities.txt'), 'a') as f:
    wr = lambda t: (f.write(t + chr(10)), print(t))
    wr('')
    wr('Shell / faint-structure search depth (outer field 1200 < a < 2400 px)')
    for B, sd in rows:
        wr('  %4d px bins (%5.1f arcsec): rms %.3f ADU/px, 3 sigma = %.2f mag/arcsec^2'
           % (B, B * PIXSCALE, sd, mu(3 * sd)))
    wr('  the rms does not integrate down like photon noise: the floor is')
    wr('  correlated large-scale structure, not shot noise')
    wr('  NO shells, arcs or tidal features were detected')
print('')
print('wrote NGC5128-sb-residual-fourier.png')