"""Shared constants and helpers for the NGC 5128 surface-photometry analysis."""
import numpy as np, os, warnings
warnings.filterwarnings('ignore')
from astropy.io import fits
from astropy.wcs import WCS

import layout

DATA = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
PIXSCALE = 0.5376                 # arcsec/px (verified from WCS)
PIXAREA  = PIXSCALE**2            # arcsec^2 per pixel
ZP5      = 27.942                 # sep.sum_circle r=5px zero point (verified)
# curve of growth on 13 isolated field stars: F(r=25)/F(r=5) = 1.2158
APCOR    = +2.5*np.log10(1.2158)  # = +0.2121 mag, r=5 -> r=25 ("total")
ZPTOT    = ZP5 + APCOR            # 28.154, zero point for total flux
MU0      = ZPTOT + 2.5*np.log10(PIXAREA)   # 26.807; mu = MU0 - 2.5*log10(I_ADU_per_px)
X0, Y0   = 2397.88, 1592.09       # nucleus, from WCS + Gaia-verified astrometry
DIST_MPC = 3.8
KPC_PER_ARCSEC = DIST_MPC*1e3*np.pi/180/3600   # 0.01842 kpc/arcsec

def path(*p): return layout.path(*p)

def load(ch):
    """Load one master as a contiguous native-endian float32 array."""
    d = fits.getdata(path('master-%s.fit' % ch))
    return np.ascontiguousarray(d.astype(np.float32))

def wcs():
    return WCS(fits.getheader(path('master-Luminance.fit')))

def mu(I):
    """Surface brightness (mag/arcsec^2) from intensity in ADU/pixel."""
    I = np.asarray(I, float)
    out = np.full(I.shape, np.nan)
    m = I > 0
    out[m] = MU0 - 2.5*np.log10(I[m])
    return out

def ell_radius(shape, x0, y0, eps, pa_rad):
    """Semi-major-axis-equivalent radius map for a fixed ellipse geometry.
    pa_rad measured counter-clockwise from the +x axis (photutils convention)."""
    ny, nx = shape
    y, x = np.mgrid[0:ny, 0:nx].astype(np.float32)
    x -= np.float32(x0); y -= np.float32(y0)
    c, s = np.float32(np.cos(pa_rad)), np.float32(np.sin(pa_rad))
    xp = x*c + y*s
    yp = -x*s + y*c
    del x, y
    return np.sqrt(xp*xp + (yp/np.float32(1.0-eps))**2)


def sky_pa(pa_deg):
    """Convert a photutils isophote PA (deg CCW from +x) to sky PA (deg E of N).

    For this frame north lies 0.96 deg CCW of the +x axis and east lies along
    -y, so the two conventions differ by very nearly 90 deg with a flip.
    """
    cd = wcs().pixel_scale_matrix
    t = np.radians(np.asarray(pa_deg, float))
    xi = cd[0, 0]*np.cos(t) + cd[0, 1]*np.sin(t)      # +east
    eta = cd[1, 0]*np.cos(t) + cd[1, 1]*np.sin(t)     # +north
    return np.degrees(np.arctan2(xi, eta)) % 180.


def north_east_pixel():
    """Unit vectors (dx, dy) pointing north and east in pixel coordinates."""
    cd = wcs().pixel_scale_matrix
    det = cd[0, 0]*cd[1, 1] - cd[0, 1]*cd[1, 0]
    n = np.array([-cd[0, 1], cd[0, 0]])/det
    e = np.array([cd[1, 1], -cd[1, 0]])/det
    return n/np.hypot(*n), e/np.hypot(*e)


ISO_HDR = ('sma_px,sma_arcsec,sma_arcmin,sma_kpc,intens_adu_px,intens_err,'
           'rms_adu,mu_mag_arcsec2,mu_err,ellipticity,ellipticity_err,'
           'pa_deg_ccw_from_x,pa_err_deg,pa_deg_east_of_north,ndata,nflag,'
           'stop_code')


def write_isophote_csv(tab, fn):
    """Write an isophote table to CSV.

    Rows with stop_code != 0 had too little unmasked azimuth for the geometry
    to converge: photutils still measures a valid intensity along the held
    ellipse, but its eps/PA are carried over from the previous isophote and its
    formal errors are meaningless (they come back as values like 1169 and
    24006).  Those five geometry columns are therefore blanked to NaN, so the
    file cannot be read as if the geometry had been measured there.  The
    intensity columns are kept, because they are real.
    """
    import os
    a = tab['sma']*PIXSCALE
    conv = tab['stop'] == 0
    blank = lambda v: np.where(conv, v, np.nan)
    with np.errstate(all='ignore'):
        me = 2.5/np.log(10)*tab['int_err']/np.where(tab['intens'] > 0,
                                                    tab['intens'], np.nan)
    out = np.column_stack([tab['sma'], a, a/60., a*KPC_PER_ARCSEC,
                           tab['intens'], tab['int_err'], tab['rms'],
                           mu(tab['intens']), me,
                           blank(tab['eps']), blank(tab['eps_err']),
                           blank(tab['pa']), blank(tab['pa_err']),
                           blank(sky_pa(tab['pa'])),
                           tab['ndata'], tab['nflag'], tab['stop']])
    np.savetxt(path(fn), out, delimiter=',', header=ISO_HDR, comments='',
               fmt='%.5f')
    print('wrote %s (%d rows, %d with converged geometry)'
          % (fn, len(out), conv.sum()))