"""Smooth elliptical model builder shared by the later steps.

Given an isophote table (sma, intens, eps, pa) and a fixed centre, assign every
pixel the semi-major axis a of the isophote passing through it.  Because eps(a)
and pa(a) vary slowly this is solved by fixed-point iteration starting from the
circular radius, which converges in a handful of passes.  The model intensity is
then a log-log interpolation of intens(a).

This is used instead of photutils.isophote.build_ellipse_model because it is
much faster on a 4788x3194 frame and because it guarantees a strictly smooth,
monotonic-in-a model with no interpolation artefacts to confuse the residual.
"""
import numpy as np
from scipy.interpolate import interp1d
from scipy.ndimage import gaussian_filter1d, zoom


def smooth_geometry(sma, eps, pa_deg, sig=2.0):
    """Return (sma, eps, pa_rad) with eps and PA lightly smoothed along a."""
    ok = np.isfinite(eps) & np.isfinite(pa_deg) & np.isfinite(sma)
    s = sma[ok]
    e = gaussian_filter1d(np.clip(eps[ok], 0.0, 0.7), sig, mode='nearest')
    p = np.unwrap(np.radians(pa_deg[ok]) * 2.0) / 2.0     # PA is defined mod 180
    p = gaussian_filter1d(p, sig, mode='nearest')
    return s, e, p


def radius_map(shape, xc, yc, sma, eps, pa_rad, nit=15, dtype=np.float32):
    """Semi-major axis of the isophote through each pixel."""
    fe = interp1d(sma, eps, bounds_error=False, fill_value=(eps[0], eps[-1]))
    fp = interp1d(sma, pa_rad, bounds_error=False, fill_value=(pa_rad[0], pa_rad[-1]))
    ny, nx = shape
    Y, X = np.mgrid[0:ny, 0:nx].astype(np.float64)
    X -= xc
    Y -= yc
    a = np.maximum(np.hypot(X, Y), 0.3)
    for _ in range(nit):
        q = 1.0 - fe(a)
        th = fp(a)
        c, s = np.cos(th), np.sin(th)
        xp = X * c + Y * s
        yp = -X * s + Y * c
        a = np.maximum(np.sqrt(xp * xp + (yp / q) ** 2), 0.3)
    return a.astype(dtype)


def build(shape, xc, yc, tab, nit=15, block=1):
    """Return (model, a_map) at full resolution.

    block > 1 computes the radius map on a coarser grid and bilinearly
    upsamples it; the model is smooth on scales far larger than block so this
    costs nothing in accuracy and a lot less in time and memory.
    """
    s, e, p = smooth_geometry(tab['sma'], tab['eps'], tab['pa'])
    if block > 1:
        sh = (shape[0] // block, shape[1] // block)
        a = radius_map(sh, (xc - (block - 1) / 2.) / block,
                       (yc - (block - 1) / 2.) / block, s / block, e, p, nit)
        # radius_map worked in block units: convert back to full-resolution pixels
        a = zoom(a.astype(np.float32) * block,
                 (shape[0] / sh[0], shape[1] / sh[1]), order=1)
        if a.shape != tuple(shape):
            b = np.zeros(shape, np.float32)
            n0, n1 = min(a.shape[0], shape[0]), min(a.shape[1], shape[1])
            b[:n0, :n1] = a[:n0, :n1]
            if n0 < shape[0]:
                b[n0:, :] = b[n0 - 1, :]
            if n1 < shape[1]:
                b[:, n1:] = b[:, [n1 - 1]]
            a = b
    else:
        a = radius_map(shape, xc, yc, s, e, p, nit)

    ok = np.isfinite(tab['intens']) & (tab['intens'] > 1e-3)
    ls, li = np.log10(tab['sma'][ok]), np.log10(tab['intens'][ok])
    o = np.argsort(ls)
    ls, li = ls[o], li[o]
    mod = (10 ** np.interp(np.log10(np.maximum(a, 0.3)), ls, li,
                           left=li[0], right=-3.0)).astype(np.float32)
    return mod, a