"""Galaxy ellipse + inner-ring measurement, ported from TamanoGalaxia.ipynb.

Given a Lupton RGB image (uint8, HxWx3) this reproduces the notebook's geometry
measurement: it segments the main galaxy, fits an ellipse to the mask by PCA of
the pixel coordinates, and detects an inner-ring candidate from the elliptical
radial profile. The returned parameters (center, orientation, semi-axes, ring
radius) are in pixel coordinates of the input image, ready to be drawn with
canvas ctx.ellipse(cx, cy, rx, ry, theta, 0, 2*pi).

The algorithm matches the notebook function-for-function. The only difference is
the base image: here we measure the luminance of the already-built Lupton RGB
from the cache, instead of rebuilding an asinh RGB from raw FITS bands. The
geometry is driven by luminance, so the fitted ellipse aligns with the galaxy as
shown on screen.
"""

from __future__ import annotations

import numpy as np
from scipy.ndimage import binary_fill_holes, gaussian_filter, gaussian_filter1d
from scipy.signal import find_peaks
from skimage import measure, morphology
from skimage.filters import threshold_otsu


def limpiar_imagen(img: np.ndarray) -> np.ndarray:
    """Cast to float32 and replace NaN/inf with the median."""
    arr = np.asarray(img, dtype=np.float32).squeeze()
    if arr.size == 0:
        return arr
    finitos = np.isfinite(arr)
    if not finitos.any():
        return np.zeros_like(arr, dtype=np.float32)
    mediana = np.nanmedian(arr[finitos])
    return np.where(np.isfinite(arr), arr, mediana).astype(np.float32)


def escalar_percentiles(img: np.ndarray, p_low: float = 1, p_high: float = 99, eps: float = 1e-8) -> np.ndarray:
    """Normalize to [0, 1] using percentile clipping (robust to bright stars)."""
    x = limpiar_imagen(img)
    lo, hi = np.nanpercentile(x, [p_low, p_high])
    if not np.isfinite(lo) or not np.isfinite(hi) or hi <= lo:
        return np.zeros_like(x, dtype=np.float32)
    return np.clip((x - lo) / (hi - lo + eps), 0, 1).astype(np.float32)


def calcular_luminancia(rgb: np.ndarray) -> np.ndarray:
    """RGB (float [0,1]) -> grayscale brightness via Rec.709 luminance."""
    return (0.2126 * rgb[:, :, 0] + 0.7152 * rgb[:, :, 1] + 0.0722 * rgb[:, :, 2]).astype(np.float32)


def centro_luminoso(lum: np.ndarray) -> tuple[float, float]:
    """Brightness-weighted center estimate (cy, cx)."""
    x = escalar_percentiles(lum, 1, 99.7)
    h, w = x.shape
    yy, xx = np.indices(x.shape)
    pesos = np.clip(x, 0, None) ** 1.5
    total = float(np.nansum(pesos))
    if not np.isfinite(total) or total <= 1e-8:
        return h / 2.0, w / 2.0
    cy = float(np.nansum(yy * pesos) / total)
    cx = float(np.nansum(xx * pesos) / total)
    return cy, cx


def mascara_galaxia(lum: np.ndarray, centro: tuple[float, float]) -> np.ndarray:
    """Binary mask of the main galaxy (mixed sky+Otsu+percentile threshold)."""
    x = escalar_percentiles(lum, 0.5, 99.7)
    suave = gaussian_filter(x, sigma=2.0)
    h, w = suave.shape

    borde = int(max(2, 0.12 * min(h, w)))
    pixeles_borde = np.concatenate([
        suave[:borde, :].ravel(),
        suave[-borde:, :].ravel(),
        suave[:, :borde].ravel(),
        suave[:, -borde:].ravel(),
    ])
    sky = float(np.nanmedian(pixeles_borde))
    sky_sigma = 1.4826 * float(np.nanmedian(np.abs(pixeles_borde - sky)))
    if not np.isfinite(sky_sigma) or sky_sigma <= 0:
        sky_sigma = float(np.nanstd(pixeles_borde))
    if not np.isfinite(sky_sigma) or sky_sigma <= 0:
        sky_sigma = 1e-6

    try:
        otsu = threshold_otsu(suave)
    except Exception:
        otsu = np.nanpercentile(suave, 70)

    umbral = max(sky + 2.0 * sky_sigma, 0.55 * otsu, np.nanpercentile(suave, 58))
    mask = suave > umbral

    min_size = max(25, int(0.0015 * h * w))
    mask = morphology.remove_small_objects(mask, min_size=min_size)
    mask = morphology.closing(mask, morphology.disk(2))
    mask = binary_fill_holes(mask)

    labeled = measure.label(mask)
    props = measure.regionprops(labeled, intensity_image=suave)

    if len(props) == 0:
        cy, cx = centro
        yy, xx = np.indices(suave.shape)
        rr = np.sqrt((yy - cy) ** 2 + (xx - cx) ** 2)
        return rr <= 0.25 * min(h, w)

    cy, cx = centro
    mejor_region = None
    mejor_score = np.inf
    for region in props:
        ry, rx = region.centroid
        distancia = np.sqrt((ry - cy) ** 2 + (rx - cx) ** 2)
        score = distancia - 0.20 * np.sqrt(max(region.area, 1)) - 5.0 * max(region.intensity_mean, 0)
        if score < mejor_score:
            mejor_score = score
            mejor_region = region

    mask_principal = labeled == mejor_region.label
    mask_principal = morphology.dilation(mask_principal, morphology.disk(2))
    mask_principal = binary_fill_holes(mask_principal)
    return mask_principal.astype(bool)


def parametros_elipse(mask: np.ndarray, lum: np.ndarray):
    """Ellipse center, semi-axes, q and orientation from the mask via PCA.

    Returns (cy, cx, semi_major, semi_minor, q, theta, area).
    """
    mask = np.asarray(mask, dtype=bool)
    h, w = mask.shape
    labeled = measure.label(mask)
    props = measure.regionprops(labeled)
    if len(props) == 0:
        return h / 2.0, w / 2.0, np.nan, np.nan, 1.0, 0.0, 0

    region = max(props, key=lambda p: p.area)
    coords = region.coords.astype(float)
    if coords.shape[0] < 10:
        cy, cx = region.centroid
        return float(cy), float(cx), np.nan, np.nan, 1.0, 0.0, int(region.area)

    yy = coords[:, 0]
    xx = coords[:, 1]
    cy = float(np.mean(yy))
    cx = float(np.mean(xx))
    x = xx - cx
    y = yy - cy

    cov = np.cov(np.vstack([x, y]), bias=True)
    try:
        eigenvalues, eigenvectors = np.linalg.eigh(cov)
    except Exception:
        return cy, cx, np.nan, np.nan, 1.0, 0.0, int(region.area)

    order = np.argsort(eigenvalues)[::-1]
    eigenvalues = eigenvalues[order]
    eigenvectors = eigenvectors[:, order]
    v_major = eigenvectors[:, 0]
    theta = float(np.arctan2(v_major[1], v_major[0]))
    if theta > np.pi / 2:
        theta -= np.pi
    if theta < -np.pi / 2:
        theta += np.pi

    semi_major = float(2.0 * np.sqrt(max(eigenvalues[0], 0.0)))
    semi_minor = float(2.0 * np.sqrt(max(eigenvalues[1], 0.0)))
    if not np.isfinite(semi_major) or semi_major <= 0:
        semi_major = np.nan
    if not np.isfinite(semi_minor) or semi_minor <= 0:
        semi_minor = np.nan

    if np.isfinite(semi_major) and semi_major > 0 and np.isfinite(semi_minor):
        q = float(np.clip(semi_minor / semi_major, 0.15, 1.0))
    else:
        q = 1.0
    return cy, cx, semi_major, semi_minor, q, theta, int(region.area)


def mapa_radio_eliptico(shape, centro, q, theta) -> np.ndarray:
    """Elliptical radius of each pixel, aligned with the galaxy axes."""
    h, w = shape
    cy, cx = centro
    yy, xx = np.indices((h, w))
    x = xx - cx
    y = yy - cy
    xr = x * np.cos(theta) + y * np.sin(theta)
    yr = -x * np.sin(theta) + y * np.cos(theta)
    q = max(float(q), 0.15)
    rr = np.sqrt(xr ** 2 + (yr / q) ** 2)
    return rr.astype(np.float32)


def perfil_radial(lum, centro, q, theta, max_radius=None):
    """Mean intensity in elliptical annuli."""
    rr_float = mapa_radio_eliptico(lum.shape, centro, q, theta)
    if max_radius is None:
        max_radius = int(np.nanmax(rr_float))
    rr = np.clip(rr_float.astype(int), 0, max_radius)
    suma = np.bincount(rr.ravel(), weights=lum.ravel(), minlength=max_radius + 1)
    conteo = np.bincount(rr.ravel(), minlength=max_radius + 1)
    perfil = suma / np.maximum(conteo, 1)
    radios = np.arange(len(perfil), dtype=np.float32)
    return radios, perfil.astype(np.float32), conteo.astype(np.float32)


def radio_por_fraccion_flujo(lum, centro, q, theta, fraccion, max_radius) -> float:
    """Radius enclosing a given fraction of the flux (e.g. R90)."""
    x = escalar_percentiles(lum, 0.5, 99.7)
    rr_float = mapa_radio_eliptico(x.shape, centro, q, theta)
    rr = np.clip(rr_float.astype(int), 0, max_radius)
    flujo_radial = np.bincount(rr.ravel(), weights=np.clip(x, 0, None).ravel(), minlength=max_radius + 1)
    flujo_acumulado = np.cumsum(flujo_radial)
    total = float(flujo_acumulado[-1])
    if not np.isfinite(total) or total <= 0:
        return np.nan
    idx = int(np.searchsorted(flujo_acumulado, fraccion * total))
    return float(np.clip(idx, 0, max_radius))


def measure_ellipse(rgb_uint8: np.ndarray) -> dict:
    """Measure the galaxy ellipse and inner-ring candidate from a Lupton RGB.

    Args:
        rgb_uint8: HxWx3 uint8 image (the cached Lupton composite).

    Returns:
        dict with pixel-coordinate geometry for drawing:
          cx, cy            ellipse center
          theta             major-axis orientation (radians)
          q                 axis ratio (minor / major)
          radius_major      drawn galaxy semi-major radius (px)
          radius_minor      radius_major * q (px)
          ring_radius       inner-ring candidate semi-major radius (px) or None
          r50, r90          flux radii (px)
          status            "ok" or "error"
    """
    try:
        rgb = np.asarray(rgb_uint8, dtype=np.float32) / 255.0
        lum = calcular_luminancia(rgb)
        h, w = lum.shape

        centro_inicial = centro_luminoso(lum)
        mask = mascara_galaxia(lum, centro_inicial)
        cy, cx, semi_major, semi_minor, q, theta, area_mask = parametros_elipse(mask, lum)

        rr_map = mapa_radio_eliptico(lum.shape, (cy, cx), q, theta)
        radios_mask = rr_map[mask]
        r_mask95 = float(np.nanpercentile(radios_mask, 95)) if len(radios_mask) > 0 else np.nan

        max_radius = int(min(np.nanmax(rr_map), 0.95 * max(h, w)))
        r50 = radio_por_fraccion_flujo(lum, (cy, cx), q, theta, 0.50, max_radius)
        r90 = radio_por_fraccion_flujo(lum, (cy, cx), q, theta, 0.90, max_radius)

        candidatos = [v for v in [r_mask95, semi_major, r90] if np.isfinite(v) and v > 0]
        radio_mayor = float(np.nanmedian(candidatos)) if candidatos else np.nan

        radios, perfil, _ = perfil_radial(lum, (cy, cx), q, theta, max_radius=max_radius)
        perfil_suave = gaussian_filter1d(perfil, sigma=2)
        base_suave = gaussian_filter1d(perfil_suave, sigma=9)
        residual = perfil_suave - base_suave
        residual[:max(5, int(0.06 * min(h, w)))] = 0
        prominencia = max(float(np.nanstd(residual)) * 0.65, 1e-5)
        peaks, _ = find_peaks(residual, prominence=prominencia, distance=5)
        ring = float(peaks[0]) if len(peaks) > 0 else None

        if not np.isfinite(radio_mayor) or radio_mayor <= 0:
            return {"status": "error"}

        return {
            "cx": round(cx, 2),
            "cy": round(cy, 2),
            "theta": round(float(theta), 5),
            "q": round(float(q), 4),
            "radius_major": round(radio_mayor, 2),
            "radius_minor": round(radio_mayor * q, 2),
            "ring_radius": round(ring, 2) if ring is not None and np.isfinite(ring) else None,
            "r50": round(r50, 2) if np.isfinite(r50) else None,
            "r90": round(r90, 2) if np.isfinite(r90) else None,
            "status": "ok",
        }
    except Exception as exc:  # never let one bad image stop a batch
        return {"status": "error", "error": str(exc)}