"""
AsteroidNET Source Extractor (source_extractor.detector)

Two-pass source detection: bright pass at 5σ to build PSF model,
faint pass at 3σ for asteroid candidates.
"""
from __future__ import annotations
import logging
from typing import Optional

import numpy as np
from astropy.table import Table
from astropy.wcs import WCS
import astropy.units as u

logger = logging.getLogger(__name__)

_CATALOG_COLS = [
    "x_pixel", "y_pixel", "ra_deg", "dec_deg",
    "flux_adu", "flux_err", "mag", "snr",
    "fwhm", "roundness", "sharpness",
]


def extract_sources(
    data_sub: np.ndarray,
    bkg_rms: np.ndarray,
    wcs: WCS,
    config: Optional[dict] = None,
) -> Table:
    """
    Detect sources in a background-subtracted image and return a catalog.

    Parameters
    ----------
    data_sub : ndarray
        Background-subtracted image (float32, NaN where masked).
    bkg_rms : ndarray
        Per-pixel background RMS (for threshold computation).
    wcs : WCS
        Frame WCS for pixel→sky coordinate conversion.
    config : dict, optional
        Pipeline configuration.

    Returns
    -------
    Table
        Source catalog with standardized columns.
    """
    cfg  = (config or {}).get("detection", {})
    thresh_hi = float(cfg.get("bright_threshold_sigma", 5.0))
    thresh_lo = float(cfg.get("threshold_sigma", 3.0))
    fwhm_lo, fwhm_hi = cfg.get("fwhm_range", [2.0, 8.0])

    rms_median = float(np.nanmedian(bkg_rms))
    if rms_median <= 0:
        rms_median = float(np.nanstd(data_sub[np.isfinite(data_sub)])) or 1.0

    try:
        from photutils.detection import DAOStarFinder
        from photutils.aperture import CircularAperture, aperture_photometry

        # Estimate typical FWHM from bright sources
        fwhm_est = _estimate_fwhm(data_sub, bkg_rms, thresh_hi, fwhm_lo, fwhm_hi)

        # Low-threshold pass for asteroid candidates
        finder = DAOStarFinder(
            fwhm=fwhm_est,
            threshold=thresh_lo * rms_median,
            sharplo=0.2, sharphi=1.0,
            roundlo=-1.0, roundhi=1.0,
            exclude_border=True,
        )
        # Replace NaN with 0 for detection only
        data_clean = np.nan_to_num(data_sub, nan=0.0)
        sources = finder(data_clean)

        if sources is None or len(sources) == 0:
            logger.info("No sources detected (threshold=%.1fσ)", thresh_lo)
            return _empty_catalog()

        # Aperture photometry
        positions = np.column_stack([sources["xcentroid"], sources["ycentroid"]])
        ap = CircularAperture(positions, r=fwhm_est)
        phot = aperture_photometry(data_sub, ap, error=bkg_rms)

        flux = np.array(phot["aperture_sum"], dtype=float)
        flux_err = np.array(phot["aperture_sum_err"], dtype=float) if "aperture_sum_err" in phot.colnames else np.full(len(flux), rms_median * np.sqrt(np.pi * fwhm_est**2))
        flux = np.maximum(flux, 0.0)

        snr = np.where(flux_err > 0, flux / flux_err, 0.0)

        # Sky coordinates via WCS
        sky = wcs.pixel_to_world(sources["xcentroid"], sources["ycentroid"])
        ra  = np.atleast_1d(sky.icrs.ra.deg)
        dec = np.atleast_1d(sky.icrs.dec.deg)

        # Magnitude (relative, no ZP needed for detection)
        with np.errstate(invalid="ignore", divide="ignore"):
            mag = np.where(flux > 0, -2.5 * np.log10(flux), 99.0)

        catalog = Table({
            "x_pixel":   np.array(sources["xcentroid"]),
            "y_pixel":   np.array(sources["ycentroid"]),
            "ra_deg":    ra,
            "dec_deg":   dec,
            "flux_adu":  flux,
            "flux_err":  flux_err,
            "mag":       mag,
            "snr":       snr,
            "fwhm":      np.full(len(flux), fwhm_est),
            "roundness": np.array(sources["roundness1"]),
            "sharpness": np.array(sources["sharpness"]),
        })

        logger.info("Extracted %d sources (fwhm=%.2fpx, threshold=%.1fσ)",
                    len(catalog), fwhm_est, thresh_lo)
        return catalog

    except ImportError:
        logger.warning("photutils not available — falling back to sigma-clip peak finder")
        return _fallback_extract(data_sub, bkg_rms, wcs, thresh_lo)


def _estimate_fwhm(data, bkg_rms, thresh, lo, hi):
    """Estimate FWHM from bright sources, clamped to [lo, hi]."""
    try:
        from photutils.detection import DAOStarFinder
        rms = float(np.nanmedian(bkg_rms))
        finder = DAOStarFinder(fwhm=3.5, threshold=thresh * rms,
                               sharplo=0.3, sharphi=0.9,
                               roundlo=-0.5, roundhi=0.5,
                               exclude_border=True)
        data_clean = np.nan_to_num(data, nan=0.0)
        sources = finder(data_clean)
        if sources and len(sources) > 5:
            fwhm = float(np.median(sources["fwhm"]))
            return float(np.clip(fwhm, lo, hi))
    except Exception:
        pass
    return 3.5  # default


def _empty_catalog() -> Table:
    return Table({c: [] for c in _CATALOG_COLS})


def _fallback_extract(data, bkg_rms, wcs, thresh):
    """Simple connected-component fallback when photutils absent."""
    try:
        from scipy.ndimage import label, center_of_mass
        rms = float(np.nanmedian(bkg_rms))
        binary = np.nan_to_num(data) > thresh * rms
        labeled, n = label(binary)
        if n == 0:
            return _empty_catalog()
        indices = list(range(1, n + 1))
        coms = center_of_mass(data, labeled, indices)
        xs = np.array([c[1] for c in coms])
        ys = np.array([c[0] for c in coms])
        sky = wcs.pixel_to_world(xs, ys)
        ra  = np.atleast_1d(sky.icrs.ra.deg)
        dec = np.atleast_1d(sky.icrs.dec.deg)
        flux = np.array([float(np.nansum(data[labeled == i])) for i in indices])
        flux = np.maximum(flux, 0.0)
        return Table({
            "x_pixel": xs, "y_pixel": ys,
            "ra_deg": ra, "dec_deg": dec,
            "flux_adu": flux, "flux_err": np.full(n, rms),
            "mag": np.where(flux > 0, -2.5 * np.log10(np.maximum(flux, 1e-10)), 99.0),
            "snr": flux / rms,
            "fwhm": np.full(n, 3.5),
            "roundness": np.zeros(n),
            "sharpness": np.zeros(n),
        })
    except Exception as exc:
        logger.error("Fallback extraction failed: %s", exc)
        return _empty_catalog()