"""Tkinter UI: remove the background gradient from the image currently
loaded in Siril.

Run from Siril's own menu: Scripts > Python Scripts. Siril launches this as
a subprocess with its bundled Python and hands it a live connection (the
sirilpy module auto-detects it); no files, no manual pipes. The corrected
pixels are pushed straight back into Siril's in-memory image over shared
memory, so the open window updates in place.

Three background models are available:

  membrane  MEMBRANE, the robust cascadic-multigrid membrane/plate fit
            (default). Minimises alpha*|grad u|^2 + beta*|lap u|^2 plus an
            asymmetric redescending data term over a coarse grid, solved
            coarse-to-fine so the cost is linear in grid cells. Follows
            complex/asymmetric light pollution without the O(n^3) solve of
            an RBF and without a global shape prior, and the reweighting
            loop keeps extended nebulosity out of the model.
  rbf       thin-plate spline through the samples (classic DBE-like).
  poly      global polynomial, the most conservative option.

The correction can be subtracted or divided out. Light pollution adds, so it
is subtracted; vignetting is a transmission loss that scales whatever came
through the optics, so it -- and the signal riding on it -- must be divided.
Subtracting a vignette leaves the target dimmed in the corners by exactly the
falloff you thought you removed.

Self-contained: no other project script is imported or shelled out to.
"""

from __future__ import annotations

import importlib
import subprocess
import sys
import threading
import time
import tkinter as tk
import traceback
import warnings
from concurrent.futures import ThreadPoolExecutor
from tkinter import scrolledtext, ttk


def _ensure_packages(packages: dict[str, str]) -> None:
    """Import-or-pip-install each {import_name: pip_name} into this interpreter.

    Siril runs scripts with its own bundled venv; numpy/scipy are normally
    preinstalled there, but this makes the script self-sufficient if not.
    """
    missing = []
    for module_name, pip_name in packages.items():
        try:
            importlib.import_module(module_name)
        except ImportError:
            missing.append(pip_name)
    if not missing:
        return

    print(f"installing missing packages into {sys.executable}: {', '.join(missing)}")
    result = subprocess.run(
        [sys.executable, "-m", "pip", "install", "--quiet", *missing],
        capture_output=True, text=True,
    )
    if result.returncode != 0:
        print(result.stdout)
        print(result.stderr, file=sys.stderr)
        raise RuntimeError(
            f"failed to install {', '.join(missing)} into {sys.executable}; "
            "install them manually with that interpreter's pip"
        )
    importlib.invalidate_caches()
    for module_name in packages:
        importlib.import_module(module_name)


_ensure_packages({"numpy": "numpy", "scipy": "scipy"})

import numpy as np  # noqa: E402
from scipy import ndimage  # noqa: E402
from scipy.interpolate import RBFInterpolator  # noqa: E402

try:
    import sirilpy as sp  # noqa: E402
    from sirilpy import tksiril  # noqa: E402
except ImportError:
    print("sirilpy not found -- this script must be run from Siril's own "
          "Scripts > Python Scripts menu, which provides it automatically.",
          file=sys.stderr)
    raise


# --------------------------------------------------------------------------- #
# Bayer handling
# --------------------------------------------------------------------------- #


def split_bayer(data: np.ndarray) -> list[np.ndarray]:
    """Split a CFA mosaic into its 4 sub-channels (2x2 super-pixel positions)."""
    h, w = data.shape
    h -= h % 2
    w -= w % 2
    view = data[:h, :w]
    # views, not copies: the samplers only read, and the models are built fresh
    return [view[dy::2, dx::2] for dy in (0, 1) for dx in (0, 1)]


def merge_bayer(planes: list[np.ndarray], base: np.ndarray) -> np.ndarray:
    """Write the 4 corrected sub-channels back over a copy of the mosaic."""
    out = base.astype(np.float32, copy=True)
    h, w = base.shape
    h -= h % 2
    w -= w % 2
    for plane, (dy, dx) in zip(planes, [(0, 0), (0, 1), (1, 0), (1, 1)]):
        out[dy:h:2, dx:w:2] = plane
    return out


# --------------------------------------------------------------------------- #
# Background sampling (fully vectorised over tiles)
# --------------------------------------------------------------------------- #


def _tile_blocks(plane: np.ndarray, ny: int, nx: int, max_px: int
                 ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Reshape `plane` into (ny*nx, P) tile blocks, sub-sampled to <=max_px each.

    Returns (blocks, inside_count, y_centres, x_centres). Padding introduced to
    make the tiling regular is NaN and is excluded from `inside_count`.
    """
    h, w = plane.shape
    bh = -(-h // ny)
    bw = -(-w // nx)
    pad_y, pad_x = ny * bh - h, nx * bw - w
    if pad_y or pad_x:
        plane = np.pad(plane, ((0, pad_y), (0, pad_x)),
                       mode="constant", constant_values=np.nan)

    # stride inside each tile so no tile contributes more than ~max_px pixels
    step = max(1, int(np.ceil(np.sqrt(bh * bw / float(max_px)))))
    rows = np.arange(0, bh, step)
    cols = np.arange(0, bw, step)

    blocks = plane.reshape(ny, bh, nx, bw)[:, ::step, :, ::step]
    blocks = blocks.transpose(0, 2, 1, 3).reshape(ny * nx, rows.size * cols.size)

    ty = np.arange(ny)[:, None] * bh
    tx = np.arange(nx)[:, None] * bw
    rows_in = (ty + rows[None, :] < h).sum(axis=1)          # (ny,)
    cols_in = (tx + cols[None, :] < w).sum(axis=1)          # (nx,)
    inside = (rows_in[:, None] * cols_in[None, :]).ravel().astype(np.float64)

    ycen = (ty[:, 0] + np.minimum(ty[:, 0] + bh, h) - 1) / 2.0
    xcen = (tx[:, 0] + np.minimum(tx[:, 0] + bw, w) - 1) / 2.0
    ycen = np.repeat(ycen, nx)
    xcen = np.tile(xcen, ny)
    return blocks, inside, ycen, xcen


def _clipped_tile_stats(blocks: np.ndarray, iters: int = 4
                        ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Asymmetric sigma-clipped median/sigma/count for every tile at once.

    Stars and nebulosity are the bright tail, so the upper cut is tighter
    (1.5 sigma) than the lower one (3 sigma) -- same policy as before, but
    computed for all tiles in a handful of vectorised passes instead of one
    astropy call per tile.
    """
    v = blocks.astype(np.float32, copy=True)
    med, count = _row_median(v)
    sigma = np.zeros_like(med)
    for _ in range(iters):
        mad, _ = _row_median(np.abs(v - med[:, None]))
        sigma = 1.4826 * mad
        bad = ~np.isfinite(sigma) | (sigma <= 0)
        if bad.any():
            sigma[bad] = _nanstd(v[bad], axis=1)
        sigma = np.where(np.isfinite(sigma), sigma, 0.0)
        lo = med - 3.0 * sigma
        hi = med + 1.5 * sigma
        keep = (v >= lo[:, None]) & (v <= hi[:, None])
        v = np.where(keep, v, np.nan)
        new_med, count = _row_median(v)
        if np.allclose(new_med, med, rtol=0, atol=1e-7, equal_nan=True):
            med = new_med
            break
        med = new_med

    return med.astype(np.float64), sigma.astype(np.float64), count.astype(np.float64)


def _row_median(v: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Per-row median ignoring non-finite entries, plus the per-row valid count.

    np.nanmedian is avoided deliberately: it switches strategy at
    shape[axis] == 600, taking a per-row Python loop below that and a
    memory-hungry vectorised path above it. Tile blocks sit right on that
    boundary, so the cost jumped around with the tile setting and the
    per-plane threads stopped scaling. Sorting a +inf-filled copy once is
    2-5x faster than either path and behaves the same at every tile size.
    """
    valid = np.isfinite(v)
    n = np.count_nonzero(valid, axis=1)
    x = np.where(valid, v, np.inf)
    x.sort(axis=1)                       # missing values land at the end
    lo = np.maximum(n - 1, 0) // 2
    hi = np.maximum(n // 2, 0)
    med = 0.5 * (np.take_along_axis(x, lo[:, None], 1)[:, 0] +
                 np.take_along_axis(x, hi[:, None], 1)[:, 0])
    return np.where(n > 0, med, np.nan), n


def _nanstd(a: np.ndarray, axis: int) -> np.ndarray:
    with warnings.catch_warnings(), np.errstate(invalid="ignore", all="ignore"):
        warnings.simplefilter("ignore", RuntimeWarning)
        return np.nanstd(a, axis=axis)


def _local_trend(field: np.ndarray, valid: np.ndarray, sigma: float,
                 iters: int = 4) -> np.ndarray:
    """Robust, LOCAL trend of a tile-median field by reweighted smoothing.

    Rejecting tiles against a global plane is what makes a sampler eat real
    structure: any genuine curvature in the light pollution -- a bright corner,
    an off-frame moon, an airglow band -- reads as a positive outlier and gets
    thrown away, and the model then reproduces a gradient it was never shown.
    Smoothing the tile field itself with normalised (missing-data aware)
    convolution gives a trend that follows the sky, so only structure smaller
    than the smoothing scale is ever called contamination.
    """
    base = np.where(valid, field, 0.0)
    w = valid.astype(np.float64)
    trend = field
    for _ in range(iters):
        num = ndimage.gaussian_filter(base * w, sigma, mode="nearest")
        den = ndimage.gaussian_filter(w, sigma, mode="nearest")
        trend = num / np.maximum(den, 1e-12)
        resid = np.where(valid, field - trend, 0.0)
        live = resid[valid]
        scale = 1.4826 * float(np.median(np.abs(live - np.median(live))))
        if not np.isfinite(scale) or scale <= 0:
            break
        k = np.where(resid > 0, 1.0 * scale, 3.0 * scale)
        w = valid / (1.0 + (resid / k) ** 2)
    return trend


def sample_background(plane: np.ndarray, n_tiles: int, tolerance: float,
                      min_frac: float = 0.15, max_px: int = 1600
                      ) -> tuple[np.ndarray, np.ndarray, int]:
    """Robust sky level per tile, with object-contaminated tiles rejected.

    Returns (points[N, 2] as (y, x) in pixels, values[N], tiles_examined).
    """
    h, w = plane.shape
    long_axis = max(h, w)
    # keep tiles at least ~8 px a side, else there is nothing to clip against
    ny = max(2, min(int(round(n_tiles * h / long_axis)), h // 8))
    nx = max(2, min(int(round(n_tiles * w / long_axis)), w // 8))

    blocks, inside, ycen, xcen = _tile_blocks(plane, ny, nx, max_px)
    med, std, count = _clipped_tile_stats(blocks)

    usable = (count >= 8) & (count >= 0.4 * np.maximum(inside, 1)) & np.isfinite(med)
    if not usable.any():
        raise ValueError("no usable background samples")
    total = int(usable.sum())

    # Reject tiles sitting on extended objects: their level sits far above the
    # LOCAL trend of the tile field. The smoothing scale sets the size of the
    # thing being rejected -- roughly an eighth of the frame -- so anything
    # larger is treated as sky here and left to the model's own reweighting.
    grid = med.reshape(ny, nx)
    valid = usable.reshape(ny, nx)
    sigma = max(1.5, 0.12 * max(ny, nx))
    trend = _local_trend(np.where(valid, grid, 0.0), valid, sigma)
    resid = (grid - trend).ravel()

    live = resid[usable]
    scale = 1.4826 * float(np.median(np.abs(live - np.median(live))))
    noise = float(np.median(std[usable][std[usable] > 0])) if np.any(std[usable] > 0) else 0.0
    scale = max(scale, noise, 1e-6)
    keep = usable & (resid < tolerance * scale) & (resid > -4.0 * scale)
    if keep.sum() < max(10, int(min_frac * total)):
        keep = usable                       # rejection ran away; trust the tiles

    return np.column_stack([ycen[keep], xcen[keep]]), med[keep], total


# --------------------------------------------------------------------------- #
# Model: polynomial
# --------------------------------------------------------------------------- #


def _design(points: np.ndarray, degree: int, scale: np.ndarray) -> np.ndarray:
    y = points[:, 0] / scale[0]
    x = points[:, 1] / scale[1]
    cols = [(x ** i) * (y ** j)
            for total in range(degree + 1)
            for i in range(total + 1)
            for j in [total - i]]
    return np.column_stack(cols)


def _poly_fit(fit_pts: np.ndarray, fit_vals: np.ndarray, degree: int
              ) -> tuple[np.ndarray, np.ndarray]:
    scale = np.array([max(fit_pts[:, 0].max(), 1.0), max(fit_pts[:, 1].max(), 1.0)])
    A = _design(fit_pts, degree, scale)
    coef, *_ = np.linalg.lstsq(A, fit_vals, rcond=None)
    return coef, scale


def model_poly(points: np.ndarray, values: np.ndarray,
               shape: tuple[int, int], degree: int) -> np.ndarray:
    """Evaluate the fit separably: no (H*W, n_coef) design matrix is ever built."""
    h, w = shape
    coef, scale = _poly_fit(points, values, degree)

    ypow = [np.ones(h, dtype=np.float64)]
    xpow = [np.ones(w, dtype=np.float64)]
    ybase = np.arange(h, dtype=np.float64) / scale[0]
    xbase = np.arange(w, dtype=np.float64) / scale[1]
    for _ in range(degree):
        ypow.append(ypow[-1] * ybase)
        xpow.append(xpow[-1] * xbase)

    out = np.zeros((h, w), dtype=np.float32)
    k = 0
    for total in range(degree + 1):
        for i in range(total + 1):
            j = total - i
            c = coef[k]
            k += 1
            if c != 0.0:
                out += (c * ypow[j][:, None] * xpow[i][None, :]).astype(np.float32)
    return out


# --------------------------------------------------------------------------- #
# Model: thin-plate spline (RBF)
# --------------------------------------------------------------------------- #


def model_rbf(points: np.ndarray, values: np.ndarray, shape: tuple[int, int],
              smooth: float, coarse: int = 192) -> np.ndarray:
    """Thin-plate-spline interpolation, evaluated coarse then zoomed (DBE-like)."""
    h, w = shape
    norm = np.array([h, w], dtype=np.float64)
    p = points / norm
    offset = float(np.median(values))
    v = values - offset
    lam = smooth * max(len(v), 1)

    # The dense solve is O(n^3) but each evaluation point is then a single
    # matvec; the local variant re-solves per query, so it only pays off once
    # the dense factorisation dominates. Measured crossover is ~2500 samples.
    neighbors = None if p.shape[0] <= 2500 else 64
    interp = RBFInterpolator(p, v, kernel="thin_plate_spline",
                             smoothing=lam, neighbors=neighbors)

    ch, cw = _coarse_shape(h, w, coarse)
    gy = np.linspace(0, h - 1, ch) / norm[0]
    gx = np.linspace(0, w - 1, cw) / norm[1]
    gyy, gxx = np.meshgrid(gy, gx, indexing="ij")
    coarse_model = interp(np.column_stack([gyy.ravel(), gxx.ravel()])).reshape(ch, cw)
    coarse_model += offset
    return _upsample(coarse_model.astype(np.float32), (h, w))


# --------------------------------------------------------------------------- #
# Model: MEMBRANE -- cascadic-multigrid membrane fit
# --------------------------------------------------------------------------- #
#
# The samples are binned onto a coarse grid, giving a data field d with a
# per-cell confidence weight w (0 where no sample landed). The background is
# the field u minimising
#
#     E(u) = sum ( alpha*|grad u|^2 + beta*|lap u|^2 )
#            + lambda * sum w * rho(u - d)
#
# i.e. the stiffest surface that still honours the sky samples. The first
# term is the membrane (surface tension), the second the thin plate
# (bending); their ratio is the "rigidity" control. Pure membrane has a kink
# at every data point, pure plate is a thin-plate spline; a blend gives a
# TPS-smooth surface that still cannot ring or overshoot into empty regions.
#
# rho is not a square but a redescending (Cauchy) loss, applied through
# iteratively reweighted least squares and made asymmetric: a cell sitting
# ABOVE the surface is probably nebulosity or an unmasked halo and gets
# demoted hard, a cell below it is probably just noise and is kept. That
# single loop is what stops the fit from eating extended signal, which is
# the thing a plain least-squares background fit always gets wrong.
#
# Its normal equations are solved by a cascadic multigrid sweep: relax on the
# coarsest grid, prolong the result as the initial guess for the next finer
# grid, relax again, and so on. Low-frequency error -- the part plain
# relaxation is hopeless at, and the part a light-pollution gradient is made
# of -- is resolved on grids where it is cheap, so the total cost is linear
# in cells and the result is independent of the sample count.
#
# Versus the other two models: no dense n x n solve (RBF) and no global
# functional form to over/under-fit (poly). Empty regions are filled by the
# harmonic (membrane) term rather than extrapolated, so masked-out object
# areas stay smooth instead of blowing up.


def _coarse_shape(h: int, w: int, coarse: int) -> tuple[int, int]:
    long_axis = max(h, w)
    ch = max(8, min(h, int(round(coarse * h / long_axis))))
    cw = max(8, min(w, int(round(coarse * w / long_axis))))
    return ch, cw


def _nb_avg(u: np.ndarray) -> np.ndarray:
    """4-neighbour average with replicated (zero-flux) borders."""
    p = np.pad(u, 1, mode="edge")
    out = p[:-2, 1:-1] + p[2:, 1:-1]
    out += p[1:-1, :-2]
    out += p[1:-1, 2:]
    out *= 0.25
    return out


def _laplacian(u: np.ndarray) -> np.ndarray:
    """-h^2 * Laplacian: 4u - sum(neighbours), zero-flux borders. SPD, diag 4."""
    return 4.0 * (u - _nb_avg(u))


def _apply_operator(u: np.ndarray, w: np.ndarray, alpha: float, beta: float,
                    lam: float) -> np.ndarray:
    """A u for A = alpha*L + beta*L^2 + lam*diag(w)."""
    out = alpha * _laplacian(u)
    if beta > 0.0:
        out += beta * _laplacian(_laplacian(u))
    out += (lam * w) * u
    return out


def _relax(u: np.ndarray, d: np.ndarray, w: np.ndarray, lam: float, sweeps: int,
           alpha: float = 1.0, beta: float = 0.0) -> np.ndarray:
    """Damped-Jacobi sweeps on (alpha*L + beta*L^2 + lam*w) u = lam*w*d.

    L has diagonal 4 and spectral radius 8, L^2 has 20 and 64, so the damping
    that keeps Jacobi a smoother (omega < 2/rho(D^-1 A)) is set from the mix
    rather than hard-coded: 0.8 for a pure membrane, 0.5 for a pure plate.
    """
    lw = lam * w
    diag = alpha * 4.0 + beta * 20.0 + lw
    rhs = lw * d
    omega = 1.6 * (alpha * 4.0 + beta * 20.0) / (alpha * 8.0 + beta * 64.0)
    for _ in range(sweeps):
        resid = rhs - _apply_operator(u, w, alpha, beta, lam)
        u += omega * resid / diag
    return u


def _restrict(d: np.ndarray, w: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """2x2 box restriction that conserves total weight (weights add, data averages)."""
    h, w_ = d.shape
    ph, pw = h % 2, w_ % 2
    if ph or pw:
        d = np.pad(d, ((0, ph), (0, pw)), mode="edge")
        w = np.pad(w, ((0, ph), (0, pw)), mode="constant")
    wd = (w * d).reshape(d.shape[0] // 2, 2, d.shape[1] // 2, 2).sum(axis=(1, 3))
    ws = w.reshape(w.shape[0] // 2, 2, w.shape[1] // 2, 2).sum(axis=(1, 3))
    dc = np.zeros_like(ws)
    nz = ws > 0
    dc[nz] = wd[nz] / ws[nz]
    return dc, ws


def _upsample(a: np.ndarray, shape: tuple[int, int], order: int = 3) -> np.ndarray:
    h, w = shape
    if a.shape == (h, w):
        return a
    return ndimage.zoom(a, (h / a.shape[0], w / a.shape[1]),
                        order=order, mode="nearest", prefilter=order > 1)


def _solve_cascadic(data: np.ndarray, weight: np.ndarray, lam: float,
                    alpha: float, beta: float, sweeps: int) -> np.ndarray:
    """Coarse-to-fine solve of (alpha*L + beta*L^2 + lam*w) u = lam*w*d.

    Levels are related by 2x2 agglomeration, so cell spacing doubles each
    step down. Discretised on the cell grid, the membrane energy is spacing
    invariant and the data term carries its own factor (restriction SUMS the
    weights, i.e. four fine cells' worth of evidence), but the plate energy
    picks up 1/h^2 -- so beta is halved-squared per level and the coarsest
    grids relax as an almost pure membrane, which is exactly where the plate
    term would have been an ill-conditioned nuisance anyway.
    """
    pyramid = [(data, weight)]
    while min(pyramid[-1][0].shape) > 4:
        pyramid.append(_restrict(*pyramid[-1]))

    top = len(pyramid) - 1
    d0, w0 = pyramid[top]
    start = float(np.average(d0, weights=w0)) if w0.sum() > 0 else 0.0
    u = np.full(d0.shape, start)
    u = _relax(u, d0, w0, lam, sweeps * 4, alpha, beta / 4.0 ** top)
    for level in range(top - 1, -1, -1):
        d_lvl, w_lvl = pyramid[level]
        u = _upsample(u, d_lvl.shape, order=1)
        u = _relax(u, d_lvl, w_lvl, lam, sweeps, alpha, beta / 4.0 ** level)
    return u


def _robust_weights(resid: np.ndarray, weight: np.ndarray, k_hi: float = 1.2,
                    k_lo: float = 3.5, floor: float = 0.12) -> np.ndarray:
    """Asymmetric redescending (Cauchy) weights on the fit residual.

    A cell above the surface is unmodelled signal -- nebulosity, a galaxy
    halo, a star's outer glow -- and must lose its vote quickly. A cell below
    it is almost always noise or a dust mote, so it keeps most of its weight.

    No cell is ever silenced completely (`floor`): a genuinely bright patch of
    sky -- an airglow band, the glow of an off-frame moon -- also sits above
    the surface, and a hard rejection there would leave the model blind in
    exactly the region it most needs to bend.
    """
    live = weight > 0
    if live.sum() < 8:
        return weight
    r = resid[live]
    sigma = 1.4826 * float(np.median(np.abs(r - np.median(r))))
    if not np.isfinite(sigma) or sigma <= 0:
        sigma = float(np.std(r)) or 1e-9
    k = np.where(resid > 0, k_hi * sigma, k_lo * sigma)
    return np.maximum(weight / (1.0 + (resid / k) ** 2), floor * weight)


def model_membrane(points: np.ndarray, values: np.ndarray,
                   shape: tuple[int, int], smooth: float, rigidity: float = 0.25,
                   robust_iters: int = 1, coarse: int = 256, sweeps: int = 32
                   ) -> np.ndarray:
    """MEMBRANE: robust cascadic-multigrid membrane/plate fit.

    Linear in grid cells, no dense solve, and the outer IRLS loop keeps
    extended objects that survived tile rejection from dragging the surface up.
    """
    h, w = shape
    ch, cw = _coarse_shape(h, w, coarse)

    # Solve on a grid padded past the frame. The relaxation's zero-flux border
    # forces the surface flat where it sits, and a real gradient is not flat at
    # the frame edge -- that shows up as a bowed-down rim on every corrected
    # image. Pushing the wall a few cells outside the data, into cells with no
    # weight at all, lets the fit keep its slope right up to the edge.
    pad = max(2, int(round(0.06 * max(ch, cw))))
    gh, gw = ch + 2 * pad, cw + 2 * pad

    # bin the samples onto the working grid
    iy = np.clip((points[:, 0] * ch / h).astype(np.intp), 0, ch - 1) + pad
    ix = np.clip((points[:, 1] * cw / w).astype(np.intp), 0, cw - 1) + pad
    flat = iy * gw + ix
    ch, cw = gh, gw
    cells = ch * cw
    weight = np.bincount(flat, minlength=cells).astype(np.float64).reshape(ch, cw)
    offset = float(np.median(values))
    acc = np.bincount(flat, weights=values - offset, minlength=cells).reshape(ch, cw)
    data = np.zeros_like(acc)
    nz = weight > 0
    data[nz] = acc[nz] / weight[nz]

    # normalise so lambda means the same thing whatever the tile grid is
    weight /= max(weight[nz].mean(), 1e-12)
    lam = 1.0 / (1e-3 + 100.0 * max(smooth, 0.0))
    rigidity = float(np.clip(rigidity, 0.0, 1.0))
    alpha, beta = 1.0, 4.0 * rigidity

    u = _solve_cascadic(data, weight, lam, alpha, beta, sweeps)
    for _ in range(max(0, robust_iters)):
        rw = _robust_weights(data - u, weight)
        if not np.isfinite(rw).all() or rw.sum() <= 0:
            break
        u_new = _solve_cascadic(data, rw, lam, alpha, beta, sweeps)
        converged = np.allclose(u_new, u, rtol=0,
                                atol=1e-4 * max(float(np.ptp(u_new)), 1e-9))
        u = u_new
        if converged:
            break

    # the cell grid is the model's real resolution; a light blur removes the
    # last of the per-cell staircase before the cubic zoom to full size
    u = ndimage.gaussian_filter(u, 0.8, mode="nearest")
    u = u[pad:ch - pad, pad:cw - pad]
    u += offset
    return _upsample(u.astype(np.float32), (h, w))


MODELS = ("membrane", "rbf", "poly")
CORRECTIONS = ("subtract", "divide")


# --------------------------------------------------------------------------- #
# Driver
# --------------------------------------------------------------------------- #


def collect_samples(plane: np.ndarray, params: dict
                    ) -> tuple[np.ndarray, np.ndarray, int]:
    """Robust sky samples for `plane`, honouring the ignore_below/above cuts."""
    sampling = plane
    lo, hi = params.get("ignore_below"), params.get("ignore_above")
    if lo is not None or hi is not None:
        sampling = plane.astype(np.float32, copy=True)
        if lo is not None:
            sampling[plane <= lo] = np.nan
        if hi is not None:
            sampling[plane >= hi] = np.nan
    return sample_background(sampling, params["tiles"], params["tolerance"])


def build_model(plane: np.ndarray, params: dict, samples=None
                ) -> tuple[np.ndarray, int, int]:
    """Sample `plane`, fit the chosen model, return (model, n_total, n_kept)."""
    pts, vals, total = samples if samples is not None else collect_samples(plane, params)
    name = params["model"]
    if name == "poly":
        model = model_poly(pts, vals, plane.shape, params["degree"])
    elif name == "rbf":
        model = model_rbf(pts, vals, plane.shape, params["smooth"])
    else:
        model = model_membrane(pts, vals, plane.shape, params["smooth"],
                               rigidity=params.get("rigidity", 0.25),
                               robust_iters=params.get("robust_iters", 1))
    return model, total, int(vals.size)


def _correct_plane(plane: np.ndarray, params: dict) -> tuple[np.ndarray, str]:
    model, total, kept = build_model(plane, params)
    level = float(np.median(model))
    mode = params.get("correction", "subtract")

    if mode == "divide" and level > 1e-6:
        # Treat the model as a flat frame: every pixel is scaled by how much
        # illumination it lost, so the target brightens back with the sky
        # instead of keeping the dimming a subtraction would leave behind.
        gain = level / np.maximum(model, level * 1e-3)
        out = plane * gain
        pedestal = level
        if not params["keep_level"]:
            out -= level
            pedestal = 0.0
    else:
        if mode == "divide":
            mode = "subtract (model level too low to divide by)"
        pedestal = level if params["keep_level"] else 0.0
        out = plane - model
        if pedestal:
            out += pedestal

    amp = float(model.max() - model.min())
    note = (f"{mode}, samples {kept}/{total} kept, amplitude {amp:.4g}, "
            f"pedestal {pedestal:.4g}")
    return out.astype(np.float32, copy=False), note


def remove_gradient_array(data: np.ndarray, bayer: str | None, params: dict,
                          log=None) -> np.ndarray:
    """Apply the background model to every independent plane of `data`.

    `data` is a 2-D CFA mosaic, 2-D mono, or 3-D (y, x, c). Returns an array
    of the same shape, dtype float32. Planes are fitted concurrently -- the
    heavy steps are numpy/scipy calls that release the GIL.
    """
    if bayer and data.ndim == 2:
        planes, layout = split_bayer(data), "cfa"
    elif data.ndim == 3:
        planes, layout = [data[..., c] for c in range(data.shape[2])], "rgb"
    else:
        planes, layout = [data], "mono"

    if len(planes) == 1:
        results = [_correct_plane(planes[0], params)]
    else:
        with ThreadPoolExecutor(max_workers=len(planes)) as pool:
            results = list(pool.map(lambda p: _correct_plane(p, params), planes))

    out_planes = []
    for i, (corrected, note) in enumerate(results):
        out_planes.append(corrected)
        if log is not None:
            log(f"  plane {i}: {note}")

    if layout == "cfa":
        return merge_bayer(out_planes, data)
    if layout == "rgb":
        return np.stack(out_planes, axis=-1)
    return out_planes[0]


def auto_clip_ceiling(data: np.ndarray) -> float | None:
    """Detect a hard saturation ceiling so star cores are not sampled as sky."""
    peak = float(np.nanmax(data))
    if not np.isfinite(peak):
        return None
    frac = float(np.count_nonzero(data >= peak)) / data.size
    return peak if 0 < frac < 0.01 else None


# --------------------------------------------------------------------------- #
# Tkinter UI
# --------------------------------------------------------------------------- #


class App(tk.Tk):
    def __init__(self, siril: sp.SirilInterface):
        super().__init__()
        self.siril = siril
        self.title("Gradient Removal")
        self.geometry("500x780")
        self.resizable(False, True)

        self._apply_dark_theme()

        self.model = tk.StringVar(value="membrane")
        self.degree = tk.IntVar(value=2)
        self.tiles = tk.IntVar(value=32)
        self.tolerance = tk.DoubleVar(value=1.5)
        self.smooth = tk.DoubleVar(value=0.001)
        self.rigidity = tk.DoubleVar(value=0.25)
        self.protect = tk.IntVar(value=1)
        self.correction = tk.StringVar(value="subtract")
        self.keep_level = tk.BooleanVar(value=True)
        self.clip_negative = tk.BooleanVar(value=False)
        self.auto_clip = tk.BooleanVar(value=True)
        self.save_as = tk.StringVar(value="")
        self.status = tk.StringVar(value="ready")

        self._build()
        self._sync_model_widgets()
        tksiril.elevate(self)

    # -- theme -------------------------------------------------------------- #

    def _apply_dark_theme(self) -> None:
        bg, bg2, fg, muted = "#1e1e1e", "#2b2b2b", "#e6e6e6", "#9a9a9a"
        accent, accent_active, border = "#3b82f6", "#2563eb", "#3f3f3f"

        self.configure(bg=bg)
        self._dark = {"bg": bg, "bg2": bg2, "fg": fg, "muted": muted}

        style = ttk.Style(self)
        style.theme_use("clam")  # 'aqua'/'default' ignore most custom colors
        style.configure(".", background=bg, foreground=fg, fieldbackground=bg2,
                        bordercolor=border, lightcolor=bg, darkcolor=bg)
        for widget in ("TFrame", "TLabel", "TCheckbutton"):
            style.configure(widget, background=bg, foreground=fg)
        style.configure("TLabelframe", background=bg, foreground=fg, bordercolor=border)
        style.configure("TLabelframe.Label", background=bg, foreground=fg)
        style.configure("TButton", background=accent, foreground="white", bordercolor=accent)
        style.map("TButton", background=[("active", accent_active), ("disabled", "#444")])
        style.configure("Muted.TLabel", foreground=muted)
        style.configure("TProgressbar", background=accent, troughcolor=bg2,
                        bordercolor=border, lightcolor=accent, darkcolor=accent)
        for widget in ("TCombobox", "TSpinbox", "TEntry"):
            style.configure(widget, fieldbackground=bg2, background=bg2, foreground=fg,
                            arrowcolor=fg, bordercolor=border, insertcolor=fg)
        style.map("TCombobox", fieldbackground=[("readonly", bg2)], foreground=[("readonly", fg)])
        self.option_add("*TCombobox*Listbox.background", bg2)
        self.option_add("*TCombobox*Listbox.foreground", fg)

    # -- layout ------------------------------------------------------------- #

    def _build(self) -> None:
        pad = {"padx": 12, "pady": 6}

        model_frame = ttk.LabelFrame(self, text="Background model")
        model_frame.pack(fill="x", **pad)
        ttk.Label(model_frame, text="Model").grid(row=0, column=0, sticky="w", padx=8, pady=4)
        box = ttk.Combobox(model_frame, textvariable=self.model, state="readonly",
                           values=list(MODELS), width=12)
        box.grid(row=0, column=1, padx=8, pady=4, sticky="w")
        box.bind("<<ComboboxSelected>>", lambda _e: self._sync_model_widgets())

        self.model_hint = ttk.Label(model_frame, text="", style="Muted.TLabel",
                                    wraplength=440, justify="left")
        self.model_hint.grid(row=1, column=0, columnspan=2, sticky="w", padx=8, pady=(0, 4))

        self.degree_label = ttk.Label(model_frame, text="Polynomial degree")
        self.degree_label.grid(row=2, column=0, sticky="w", padx=8, pady=4)
        self.degree_spin = ttk.Spinbox(model_frame, textvariable=self.degree,
                                       from_=1, to=6, width=10)
        self.degree_spin.grid(row=2, column=1, padx=8, pady=4, sticky="w")

        self.smooth_label = ttk.Label(model_frame, text="Smoothing / stiffness")
        self.smooth_label.grid(row=3, column=0, sticky="w", padx=8, pady=4)
        self.smooth_spin = ttk.Spinbox(model_frame, textvariable=self.smooth, from_=0.0,
                                       to=1.0, increment=0.001, width=10)
        self.smooth_spin.grid(row=3, column=1, padx=8, pady=4, sticky="w")

        self.rigidity_label = ttk.Label(model_frame, text="Rigidity (membrane -> plate)")
        self.rigidity_label.grid(row=4, column=0, sticky="w", padx=8, pady=4)
        self.rigidity_spin = ttk.Spinbox(model_frame, textvariable=self.rigidity,
                                         from_=0.0, to=1.0, increment=0.05, width=10)
        self.rigidity_spin.grid(row=4, column=1, padx=8, pady=4, sticky="w")

        self.protect_label = ttk.Label(model_frame, text="Object protection (passes)")
        self.protect_label.grid(row=5, column=0, sticky="w", padx=8, pady=4)
        self.protect_spin = ttk.Spinbox(model_frame, textvariable=self.protect,
                                        from_=0, to=5, width=10)
        self.protect_spin.grid(row=5, column=1, padx=8, pady=4, sticky="w")

        corr_frame = ttk.LabelFrame(self, text="Correction")
        corr_frame.pack(fill="x", **pad)
        ttk.Label(corr_frame, text="Apply as").grid(row=0, column=0, sticky="w",
                                                    padx=8, pady=4)
        ttk.Combobox(corr_frame, textvariable=self.correction, state="readonly",
                     values=list(CORRECTIONS), width=12).grid(
            row=0, column=1, padx=8, pady=4, sticky="w")
        ttk.Label(corr_frame,
                  text="Subtract for light pollution, which adds to the frame. "
                       "Divide for vignetting and any other transmission "
                       "falloff: it scales the target too, so subtracting it "
                       "leaves the corners dimmed by exactly the amount you "
                       "meant to remove.",
                  style="Muted.TLabel", wraplength=440, justify="left").grid(
            row=1, column=0, columnspan=2, sticky="w", padx=8, pady=(0, 6))

        sample_frame = ttk.LabelFrame(self, text="Sampling")
        sample_frame.pack(fill="x", **pad)
        ttk.Label(sample_frame, text="Tiles (long axis)").grid(
            row=0, column=0, sticky="w", padx=8, pady=4)
        ttk.Spinbox(sample_frame, textvariable=self.tiles, from_=4, to=256, width=10).grid(
            row=0, column=1, padx=8, pady=4, sticky="w")
        ttk.Label(sample_frame, text="Rejection tolerance (sigma)").grid(
            row=1, column=0, sticky="w", padx=8, pady=4)
        ttk.Spinbox(sample_frame, textvariable=self.tolerance, from_=0.1, to=10,
                    increment=0.1, width=10).grid(row=1, column=1, padx=8, pady=4, sticky="w")
        ttk.Checkbutton(sample_frame, text="Ignore clipped highlights and black borders",
                        variable=self.auto_clip).grid(row=2, column=0, columnspan=2,
                                                      sticky="w", padx=8, pady=(6, 2))
        ttk.Checkbutton(sample_frame, text="Keep sky pedestal (restore background median)",
                        variable=self.keep_level).grid(row=3, column=0, columnspan=2,
                                                       sticky="w", padx=8, pady=2)
        ttk.Checkbutton(sample_frame, text="Clamp negative pixels to zero",
                        variable=self.clip_negative).grid(row=4, column=0, columnspan=2,
                                                          sticky="w", padx=8, pady=(2, 6))

        out_frame = ttk.LabelFrame(self, text="Output")
        out_frame.pack(fill="x", **pad)
        ttk.Label(out_frame, text="Also save corrected image as (optional)").pack(
            anchor="w", padx=8, pady=(4, 0))
        ttk.Entry(out_frame, textvariable=self.save_as).pack(fill="x", padx=8, pady=(0, 4))
        ttk.Label(out_frame, text="The image open in Siril is always updated in memory "
                                  "either way.", style="Muted.TLabel",
                  wraplength=440).pack(anchor="w", padx=8, pady=(0, 6))

        bar = ttk.Frame(self)
        bar.pack(fill="x", padx=12, pady=(4, 6))
        self.run_button = ttk.Button(bar, text="Remove Gradient", command=self.on_run)
        self.run_button.pack(side="left")
        ttk.Label(bar, textvariable=self.status, style="Muted.TLabel").pack(
            side="left", padx=10)
        self.progress = ttk.Progressbar(bar, mode="indeterminate", length=120)
        self.progress.pack(side="right")

        self.log_box = scrolledtext.ScrolledText(
            self, height=12, font=("Menlo", 10), state="disabled",
            bg=self._dark["bg2"], fg=self._dark["fg"], insertbackground=self._dark["fg"],
            selectbackground="#3b82f6", relief="flat", borderwidth=0,
        )
        self.log_box.pack(fill="both", expand=True, padx=12, pady=(0, 12))

    HINTS = {
        "membrane": "MEMBRANE -- stiffest surface that still honours the sky "
                    "samples, solved coarse-to-fine and re-weighted against its "
                    "own residual so extended objects lose their vote. Rigidity 0 "
                    "is a pure membrane, 1 a thin plate (smoothest). Object "
                    "protection is how many times the fit is reweighted against "
                    "its own residual: raise it if the model eats your target, "
                    "lower it if a genuinely bright patch of sky survives.",
        "rbf": "Thin-plate spline through the samples. Flexible, slower on "
               "dense tile grids.",
        "poly": "Global polynomial. The safest, most conservative model.",
    }

    def _sync_model_widgets(self) -> None:
        name = self.model.get()
        self.model_hint.configure(text=self.HINTS.get(name, ""))
        poly = "normal" if name == "poly" else "disabled"
        smooth = "disabled" if name == "poly" else "normal"
        membrane = "normal" if name == "membrane" else "disabled"
        self.degree_spin.configure(state=poly)
        self.degree_label.configure(style="TLabel" if name == "poly" else "Muted.TLabel")
        self.smooth_spin.configure(state=smooth)
        self.smooth_label.configure(style="Muted.TLabel" if name == "poly" else "TLabel")
        self.rigidity_spin.configure(state=membrane)
        self.protect_spin.configure(state=membrane)
        for label in (self.rigidity_label, self.protect_label):
            label.configure(style="TLabel" if name == "membrane" else "Muted.TLabel")

    # -- logging ------------------------------------------------------------ #

    def _log(self, text: str) -> None:
        self.log_box.configure(state="normal")
        self.log_box.insert("end", text + "\n")
        self.log_box.see("end")
        self.log_box.configure(state="disabled")

    def log(self, text: str) -> None:
        """Thread-safe log entry point for the worker."""
        self.after(0, self._log, text)

    def _set_status(self, text: str) -> None:
        self.after(0, self.status.set, text)

    # -- run ---------------------------------------------------------------- #

    def on_run(self) -> None:
        self.run_button.configure(state="disabled")
        self.progress.start(12)
        self.status.set("working...")
        self.log_box.configure(state="normal")
        self.log_box.delete("1.0", "end")
        self.log_box.configure(state="disabled")

        params = {
            "model": self.model.get(),
            "degree": self.degree.get(),
            "tiles": self.tiles.get(),
            "tolerance": self.tolerance.get(),
            "smooth": self.smooth.get(),
            "rigidity": self.rigidity.get(),
            "robust_iters": self.protect.get(),
            "correction": self.correction.get(),
            "keep_level": self.keep_level.get(),
            "clip_negative": self.clip_negative.get(),
            "auto_clip": self.auto_clip.get(),
            "save_as": self.save_as.get().strip(),
        }
        threading.Thread(target=self._run_worker, args=(params,), daemon=True).start()

    def _run_worker(self, params: dict) -> None:
        try:
            if not self.siril.is_image_loaded():
                raise RuntimeError("no image loaded in Siril")

            t0 = time.perf_counter()

            # set_image_pixeldata() requires the processing thread claimed first --
            # everything from reading the image to writing it back must happen
            # inside this lock (Siril's own image_lock() context manager).
            with self.siril.image_lock():
                self.siril.undo_save_state("remove_gradient")

                ffit = self.siril.get_image(with_pixels=False)
                bayer = (ffit.keywords.bayer_pattern or "").strip().upper() or None

                data = self.siril.get_image_pixeldata()  # (C,H,W)|(H,W), float32|uint16
                orig_dtype = data.dtype
                channels_first = data.ndim == 3
                if channels_first:
                    data = np.moveaxis(data, 0, -1)  # -> (H,W,C) for processing
                if data.ndim != 2:
                    bayer = None
                work = data.astype(np.float32, copy=False)

                if params["auto_clip"]:
                    ceiling = auto_clip_ceiling(work)
                    if ceiling is not None:
                        params["ignore_above"] = ceiling
                        self.log(f"ignoring pixels >= {ceiling:.4g} (clipped highlights)")
                    if float(np.count_nonzero(work <= 0)) / work.size > 0.001:
                        params["ignore_below"] = 0.0
                        self.log("ignoring pixels <= 0 (stacking borders)")

                self.log(f"input: shape {work.shape}, dtype {orig_dtype}, "
                         f"bayer={bayer}, median={float(np.median(work)):.4g}")
                self._set_status(f"fitting ({params['model']})...")

                t1 = time.perf_counter()
                result = remove_gradient_array(work, bayer, params, log=self.log)
                fit_ms = (time.perf_counter() - t1) * 1e3
                self.log(f"model fitted in {fit_ms:.0f} ms; "
                         f"corrected median={float(np.median(result)):.4g}")

                if params["clip_negative"]:
                    np.clip(result, 0, None, out=result)

                if channels_first:
                    result = np.moveaxis(result, -1, 0)  # back to (C,H,W)
                if orig_dtype == np.uint16:
                    result = np.clip(np.rint(result), 0, 65535).astype(np.uint16)
                else:
                    result = np.ascontiguousarray(result, dtype=np.float32)

                self.siril.set_image_pixeldata(result)

            self.siril.log(
                f"remove_gradient: background model {params['correction']}ed "
                f"({params['model']}, tiles={params['tiles']}, "
                f"tolerance={params['tolerance']})")

            if params["save_as"]:
                self.log(f"saving as {params['save_as']}...")
                self.siril.cmd("save", params["save_as"])

            total_ms = (time.perf_counter() - t0) * 1e3
            self.log(f"\ndone in {total_ms:.0f} ms -- image updated in Siril.")
            self._set_status("done")

        except Exception:  # noqa: BLE001
            self.log(f"\nERROR:\n{traceback.format_exc()}")
            self._set_status("failed")
        finally:
            self.after(0, self._finish)

    def _finish(self) -> None:
        self.progress.stop()
        self.run_button.configure(state="normal")


def main() -> None:
    siril = sp.SirilInterface()
    siril.connect()
    try:
        App(siril).mainloop()
    finally:
        siril.disconnect()


if __name__ == "__main__":
    main()