"""
White-stack bright unmatched-source masking for tract7dt.

Runs SEP on the in-memory white stack, matches detections against active
catalog sources, and builds ellipse-based masks for unmatched bright
contaminants.  The resulting mask is stored in ``state["bright_mask"]``
and later composed with saturation and non-finite masks into per-band
``bad`` arrays by the pipeline's ``_compose_final_masks`` step, so that
downstream Tractor fitting sees ``invvar = 0`` for those pixels.
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd
import sep
from astropy.io import fits
from scipy.ndimage import binary_dilation

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from astropy.visualization import ZScaleInterval

logger = logging.getLogger("tract7dt.bright_mask")


def apply_bright_mask(cfg: dict[str, Any], state: dict[str, Any]) -> np.ndarray:
    """Detect bright unmatched sources on the white stack, build ellipse masks,
    store in ``state["bright_mask"]``, and save diagnostics.

    Returns the 2-D boolean mask (True = masked).
    """
    bm_cfg = cfg.get("bright_mask", {})
    if not bool(bm_cfg.get("enabled", True)):
        logger.info("bright_mask disabled, skipping.")
        state["input_catalog"]["affected_by_bright_mask"] = False
        h, w = state["white"].shape
        state["bright_mask"] = np.zeros((h, w), dtype=bool)
        return state["bright_mask"]

    white = np.asarray(state["white"], dtype=np.float32)
    sigma_white = np.asarray(state["sigma_white"], dtype=np.float32)
    input_catalog = state["input_catalog"]
    image_dict = state["image_dict"]

    h, w = white.shape

    # --- SEP detection on white stack ---
    white_sep = white.copy()
    bad_pix = ~np.isfinite(white_sep)
    white_sep[bad_pix] = 0.0

    sig_sep = sigma_white.copy()
    sig_bad = ~np.isfinite(sig_sep) | (sig_sep <= 0)
    sig_sep[sig_bad] = 1e30

    thresh = float(bm_cfg.get("thresh_sigma", 7.0))
    objs = sep.extract(
        white_sep,
        thresh=thresh,
        err=sig_sep,
        mask=bad_pix.astype(np.uint8),
        minarea=int(bm_cfg.get("minarea", 20)),
        deblend_nthresh=int(bm_cfg.get("deblend_nthresh", 32)),
        deblend_cont=float(bm_cfg.get("deblend_cont", 0.005)),
        clean=bool(bm_cfg.get("clean", True)),
        clean_param=float(bm_cfg.get("clean_param", 1.0)),
    )
    logger.info("SEP detected %d objects on white stack (thresh=%.1f sigma)", len(objs), thresh)

    if len(objs) == 0:
        mask = np.zeros((h, w), dtype=bool)
        input_catalog["affected_by_bright_mask"] = False
        state["bright_mask"] = mask
        _save_diagnostics(cfg, white, mask, objs, np.array([], dtype=bool),
                          input_catalog, image_dict)
        return mask

    det_x = objs["x"]
    det_y = objs["y"]

    # --- Match detections to active catalog sources ---
    match_r = float(bm_cfg.get("match_radius_pix", 3.0))
    matched = _match_to_catalog(det_x, det_y, input_catalog, image_dict, match_r)
    n_matched = int(np.sum(matched))
    n_unmatched = len(objs) - n_matched
    logger.info("Matched %d / %d SEP detections to catalog; %d unmatched to mask",
                n_matched, len(objs), n_unmatched)

    if n_unmatched == 0:
        mask = np.zeros((h, w), dtype=bool)
        input_catalog["affected_by_bright_mask"] = False
        state["bright_mask"] = mask
        _save_diagnostics(cfg, white, mask, objs, matched,
                          input_catalog, image_dict)
        return mask

    # --- Build ellipse masks for unmatched detections ---
    ellipse_scale = float(bm_cfg.get("ellipse_scale", 1.5))
    min_radius = float(bm_cfg.get("min_radius_pix", 6.0))
    dilate_pix = int(bm_cfg.get("dilate_pix", 2))

    mask = _build_ellipse_mask(
        h, w, objs, ~matched, ellipse_scale, min_radius,
    )

    if dilate_pix > 0:
        struct = np.ones((2 * dilate_pix + 1, 2 * dilate_pix + 1), dtype=bool)
        mask = binary_dilation(mask, structure=struct).astype(bool)

    n_masked_pix = int(np.sum(mask))
    logger.info("Bright-source mask: %d pixels (%.2f%% of frame)",
                n_masked_pix, 100.0 * n_masked_pix / (h * w))

    # --- Store mask in state (composed into bad later by _compose_final_masks) ---
    state["bright_mask"] = mask
    logger.info("Stored bright mask in state (will be composed into bad later)")

    # --- Flag catalog sources whose positions fall inside the mask ---
    _flag_affected_sources(input_catalog, image_dict, mask)

    # --- Save diagnostics ---
    _save_diagnostics(cfg, white, mask, objs, matched,
                      input_catalog, image_dict)

    return mask


def _match_to_catalog(
    det_x: np.ndarray,
    det_y: np.ndarray,
    input_catalog: pd.DataFrame,
    image_dict: dict[str, Any],
    match_radius_pix: float,
) -> np.ndarray:
    """Return a boolean array of length len(det_x), True where a detection
    matches an active catalog source within match_radius_pix."""

    is_active = ~input_catalog.get("excluded_any", pd.Series(False, index=input_catalog.index)).fillna(False).astype(bool)
    active_cat = input_catalog.loc[is_active]

    if len(active_cat) == 0:
        return np.zeros(len(det_x), dtype=bool)

    wcs_white = next(iter(image_dict.values()))["wcs"]
    col_map = {str(c).strip().lower(): str(c) for c in active_cat.columns}
    ra_col = col_map.get("ra", "RA")
    dec_col = col_map.get("dec", "DEC")

    ra = pd.to_numeric(active_cat[ra_col], errors="coerce").to_numpy(dtype=float)
    dec = pd.to_numeric(active_cat[dec_col], errors="coerce").to_numpy(dtype=float)
    ok = np.isfinite(ra) & np.isfinite(dec)

    if not np.any(ok):
        return np.zeros(len(det_x), dtype=bool)

    cat_x = np.full(len(ra), np.nan, dtype=float)
    cat_y = np.full(len(dec), np.nan, dtype=float)
    xw, yw = wcs_white.all_world2pix(ra[ok], dec[ok], 0)
    cat_x[ok] = xw
    cat_y[ok] = yw

    cat_x = cat_x[ok]
    cat_y = cat_y[ok]

    matched = np.zeros(len(det_x), dtype=bool)
    r2 = match_radius_pix ** 2
    for i in range(len(det_x)):
        dx = det_x[i] - cat_x
        dy = det_y[i] - cat_y
        if np.any(dx * dx + dy * dy <= r2):
            matched[i] = True

    return matched


def _flag_affected_sources(
    input_catalog: pd.DataFrame,
    image_dict: dict[str, Any],
    mask: np.ndarray,
) -> None:
    """Flag catalog sources whose pixel positions fall inside the bright mask."""
    input_catalog["affected_by_bright_mask"] = False

    wcs_white = next(iter(image_dict.values()))["wcs"]
    col_map = {str(c).strip().lower(): str(c) for c in input_catalog.columns}
    ra_col = col_map.get("ra", "RA")
    dec_col = col_map.get("dec", "DEC")

    ra = pd.to_numeric(input_catalog[ra_col], errors="coerce").to_numpy(dtype=float)
    dec = pd.to_numeric(input_catalog[dec_col], errors="coerce").to_numpy(dtype=float)
    ok = np.isfinite(ra) & np.isfinite(dec)

    if not np.any(ok):
        return

    cat_x = np.full(len(ra), np.nan, dtype=float)
    cat_y = np.full(len(dec), np.nan, dtype=float)
    xw, yw = wcs_white.all_world2pix(ra[ok], dec[ok], 0)
    cat_x[ok] = xw
    cat_y[ok] = yw

    h, w = mask.shape
    for idx in np.where(ok)[0]:
        ix = int(np.round(cat_x[idx]))
        iy = int(np.round(cat_y[idx]))
        if 0 <= ix < w and 0 <= iy < h and mask[iy, ix]:
            input_catalog.iat[idx, input_catalog.columns.get_loc("affected_by_bright_mask")] = True

    n_affected = int(input_catalog["affected_by_bright_mask"].sum())
    if n_affected > 0:
        logger.info("Flagged %d catalog sources as affected_by_bright_mask", n_affected)


def _build_ellipse_mask(
    h: int,
    w: int,
    objs: np.ndarray,
    unmatched: np.ndarray,
    ellipse_scale: float,
    min_radius: float,
) -> np.ndarray:
    """Build a boolean mask from scaled SEP ellipses of unmatched detections."""
    mask = np.zeros((h, w), dtype=np.bool_)

    idx = np.where(unmatched)[0]
    if len(idx) == 0:
        return mask

    xs = objs["x"][idx].astype(np.float64)
    ys = objs["y"][idx].astype(np.float64)
    as_ = np.maximum(objs["a"][idx].astype(np.float64) * ellipse_scale, min_radius)
    bs = np.maximum(objs["b"][idx].astype(np.float64) * ellipse_scale, min_radius)
    thetas = objs["theta"][idx].astype(np.float64)

    sep.mask_ellipse(mask, xs, ys, as_, bs, thetas)
    return mask


def _save_diagnostics(
    cfg: dict[str, Any],
    white: np.ndarray,
    mask: np.ndarray,
    objs: np.ndarray,
    matched: np.ndarray,
    input_catalog: pd.DataFrame,
    image_dict: dict[str, Any],
) -> None:
    """Save bright-mask overlay PNG and raw FITS mask."""
    outputs = cfg.get("outputs", {})
    work_dir = outputs.get("work_dir")
    if work_dir is None:
        logger.warning("Cannot save bright-mask diagnostics: no work_dir")
        return

    bm_dir = Path(work_dir) / "bright_mask"
    bm_dir.mkdir(parents=True, exist_ok=True)

    # --- FITS mask (bright ellipses | any-band saturation) ---
    combined_mask = mask.copy()
    for d in image_dict.values():
        sm = d.get("satur_mask")
        if sm is not None:
            combined_mask = combined_mask | np.asarray(sm, dtype=bool)
    fits_path = bm_dir / "white_bright_mask.fits"
    try:
        hdu = fits.PrimaryHDU(data=combined_mask.astype(np.uint8))
        hdu.header["EXTNAME"] = "BRIGHT_MASK"
        hdu.header["COMMENT"] = "1 = masked (bright unmatched source OR saturated pixel)"
        hdu.writeto(str(fits_path), overwrite=True)
        logger.info("Saved bright-mask FITS (incl. saturation): %s", fits_path)
    except Exception as e:
        logger.warning("Failed to write bright-mask FITS: %s", e)

    # --- Overlay PNG ---
    png_path = bm_dir / "white_bright_mask_overlay.png"
    try:
        dpi = int(cfg.get("plotting", {}).get("dpi", 150))
        bm_cfg = cfg.get("bright_mask", {})
        match_r = float(bm_cfg.get("match_radius_pix", 3.0))
        h, w = white.shape

        fig, ax = plt.subplots(figsize=(12, 8), constrained_layout=True)

        disp = white.copy()
        m = np.isfinite(disp)
        if np.any(m):
            vmin, vmax = ZScaleInterval().get_limits(disp[m])
        else:
            vmin, vmax = 0, 1
        disp[~m] = 0

        ax.imshow(disp, origin="lower", cmap="gray", vmin=vmin, vmax=vmax,
                  interpolation="nearest", rasterized=True)

        if np.any(mask):
            mask_rgba = np.zeros((h, w, 4), dtype=np.float32)
            mask_rgba[mask, 0] = 1.0
            mask_rgba[mask, 3] = 0.35
            ax.imshow(mask_rgba, origin="lower", interpolation="nearest")

        # -- SEP detections --
        if len(objs) > 0:
            sep_matched_x = [objs["x"][i] for i in range(len(objs)) if matched[i]]
            sep_matched_y = [objs["y"][i] for i in range(len(objs)) if matched[i]]
            sep_unmatched_x = [objs["x"][i] for i in range(len(objs)) if not matched[i]]
            sep_unmatched_y = [objs["y"][i] for i in range(len(objs)) if not matched[i]]
            if sep_matched_x:
                ax.scatter(sep_matched_x, sep_matched_y, marker="o", s=20,
                           facecolors="none", edgecolors="lime", linewidths=0.8,
                           label=f"SEP matched ({len(sep_matched_x)})", zorder=4)
            if sep_unmatched_x:
                ax.scatter(sep_unmatched_x, sep_unmatched_y, marker="x", s=20,
                           c="red", linewidths=0.8,
                           label=f"SEP unmatched ({len(sep_unmatched_x)})", zorder=4)

        # -- Input catalog sources --
        if input_catalog is not None and len(input_catalog) > 0:
            wcs_white = next(iter(image_dict.values()))["wcs"]
            col_map = {str(c).strip().lower(): str(c) for c in input_catalog.columns}
            ra_col = col_map.get("ra", "RA")
            dec_col = col_map.get("dec", "DEC")

            ra_arr = pd.to_numeric(input_catalog[ra_col], errors="coerce").to_numpy(dtype=float)
            dec_arr = pd.to_numeric(input_catalog[dec_col], errors="coerce").to_numpy(dtype=float)
            coord_ok = np.isfinite(ra_arr) & np.isfinite(dec_arr)

            if np.any(coord_ok):
                cx_all = np.full(len(ra_arr), np.nan)
                cy_all = np.full(len(dec_arr), np.nan)
                xw, yw = wcs_white.all_world2pix(ra_arr[coord_ok], dec_arr[coord_ok], 0)
                cx_all[coord_ok] = xw
                cy_all[coord_ok] = yw

                r2 = match_r ** 2
                cat_has_sep = np.zeros(len(input_catalog), dtype=bool)
                if len(objs) > 0:
                    det_x = objs["x"]
                    det_y = objs["y"]
                    for ci in np.where(coord_ok)[0]:
                        dx = cx_all[ci] - det_x
                        dy = cy_all[ci] - det_y
                        if np.any(dx * dx + dy * dy <= r2):
                            cat_has_sep[ci] = True

                affected = input_catalog.get("affected_by_bright_mask",
                                             pd.Series(False, index=input_catalog.index)
                                             ).fillna(False).astype(bool).to_numpy()

                cat_matched_m = coord_ok & cat_has_sep
                cat_masked_m = coord_ok & ~cat_has_sep & affected

                if np.any(cat_matched_m):
                    ax.scatter(cx_all[cat_matched_m], cy_all[cat_matched_m],
                               marker="D", s=18, facecolors="none",
                               edgecolors="deepskyblue", linewidths=0.7,
                               label=f"Cat matched ({int(cat_matched_m.sum())})", zorder=3)
                if np.any(cat_masked_m):
                    ax.scatter(cx_all[cat_masked_m], cy_all[cat_masked_m],
                               marker="^", s=22, facecolors="none",
                               edgecolors="orange", linewidths=0.8,
                               label=f"Cat unmatched masked ({int(cat_masked_m.sum())})", zorder=3)

        n_unmatched = int(np.sum(~matched)) if len(matched) > 0 else 0
        n_total = len(objs)
        ax.set_title(f"Bright mask: {n_unmatched} unmatched / {n_total} SEP detections")
        ax.legend(loc="upper right", fontsize=7, framealpha=0.8)
        ax.set_xticks([])
        ax.set_yticks([])

        plt.savefig(str(png_path), dpi=dpi)
        plt.close(fig)
        logger.info("Saved bright-mask overlay: %s", png_path)
    except Exception as e:
        logger.warning("Failed to save bright-mask overlay: %s", e)
