"""Fused stitch blend kernel (Numba JIT when available)."""

from __future__ import annotations

import numpy as np

try:
    from numba import njit

    _NUMBA_AVAILABLE = True
except ImportError:
    _NUMBA_AVAILABLE = False

    def njit(*args, **kwargs):
        def decorator(fn):
            return fn

        return decorator


if _NUMBA_AVAILABLE:

    @njit(cache=True)
    def blend_kernel(blend_patch, wsum_patch, cov_patch, src_patch, w_patch):
        """Fused incremental weighted-average composite (numba-accelerated)."""
        sh = src_patch.shape[0]
        sw = src_patch.shape[1]
        for y in range(sh):
            for x in range(sw):
                w = w_patch[y, x]
                s0 = np.float32(src_patch[y, x, 0])
                s1 = np.float32(src_patch[y, x, 1])
                s2 = np.float32(src_patch[y, x, 2])
                if cov_patch[y, x] == 0:
                    blend_patch[y, x, 0] = s0
                    blend_patch[y, x, 1] = s1
                    blend_patch[y, x, 2] = s2
                    wsum_patch[y, x] = w
                else:
                    ws = wsum_patch[y, x] + w
                    if ws < 1e-6:
                        ws = 1e-6
                    a = w / ws
                    blend_patch[y, x, 0] += (s0 - blend_patch[y, x, 0]) * a
                    blend_patch[y, x, 1] += (s1 - blend_patch[y, x, 1]) * a
                    blend_patch[y, x, 2] += (s2 - blend_patch[y, x, 2]) * a
                    wsum_patch[y, x] = ws
                cov_patch[y, x] = 1

else:
    blend_kernel = None