"""
In this module the availability of numba is first tested.
If not available, a dummy '@jit' decorator is applied

The @jit decoration can be suppressed by disabling numba in Roll's settings menu.
This will set 'DISABLE_JIT' to True in the module numba.config.
As a result, the decorated functions are no longer precompiled.

See: https://numba.readthedocs.io/en/stable/user/jit.html for preferred way of using @jit
See: https://stackoverflow.com/questions/57774497/how-do-i-make-a-dummy-do-nothing-jit-decorator
"""
import numpy as np
from qgis.PyQt.QtCore import QLineF, QRectF  # needed for pointsInRect

LOG_RESPONSE_FLOOR = 1.0e-12

try:
    from numba import jit, prange
except ImportError:
    prange = range

    def jit(*args, **kwargs):
        if args and callable(args[0]) and not kwargs:
            return args[0]

        def decorator(func):
            return func

        return decorator

# @jit((nb.types.Array(nb.float32, 3, 'C'), nb.boolean), nopython=True)      # numba needs array specs to work properly


@jit(nopython=True)
def numbaFilterSlice2D(slice2D: np.ndarray, unique=False):
    # size = slice3D.shape[0]
    # fold = slice3D.shape[1]
    # para = slice3D.shape[2]

    # slice2D = slice3D.reshape(size * fold, para)                              # convert to 2D

    if unique is True:                                                          # we'd like to use unique offsets
        havUnique = slice2D[:, 15]                                              # get all available cmp values belonging to this row
        useUnique = True if havUnique.min() == -1 else False                    # are there any -1 records ?
    else:
        useUnique = False                                                       # unique not required or not available

    if useUnique:
        included = (slice2D[:, 2] > 0) & (slice2D[:, 15] == -1)                 # fold > 0 AND unique == -1
    else:
        included = slice2D[:, 2] > 0                                            # fold > 0

    slice2D = slice2D[included, :]                                              # filter the 2D slice
    return slice2D                                                              # return array, potentially with shape[0,16] when empty


@jit(nopython=True)
def numbaSlice3D(slice3D: np.ndarray, unique=False):
    if unique:                                                                  # we'd like to use unique offsets; but does it make sense ?
        havUnique = slice3D[:, :, 15]                                           # get all available cmp values belonging to this slice
        useUnique = True if havUnique.min() == -1 else False                    # are there any -1 records ?
    else:
        useUnique = False                                                       # unique not required or not available

    if useUnique:
        included = (slice3D[:, :, 2] > 0) & (slice3D[:, :, 15] == -1)           # fold > 0 AND unique == -1
    else:
        included = slice3D[:, :, 2] > 0                                         # fold > 0

    return (slice3D, included)                                                  # tuple of information (array, mask)


# it would be nice to cache the outcome of the following routine.
# See: https://github.com/numba/numba/issues/4062
# See: https://stackoverflow.com/questions/74491198/how-to-set-signature-if-use-np-array-as-an-input-in-numba
# See: https://stackoverflow.com/questions/67213411/initizalizing-a-return-list-of-tuples-in-numba
# See: https://stackoverflow.com/questions/30363253/multiple-output-and-numba-signatures
# See: https://stackoverflow.com/questions/55765255/how-do-i-specify-a-tuple-in-a-numba-vectorize-signature

# @jit(nb.types.Tuple((nb.float32[:], nb.float32[:], nb.boolean))(nb.float32[:, :, :, :], nb.boolean), nopython=True)                    # numba needs array specs to work
def numbaSliceStats(slice4D: np.ndarray, unique=False):

    # fmt: off
    fold    = slice4D[:, :, :,  2]                                              # noqa: E221, E241, # we are left with 1 dimension
    offsets = slice4D[:, :, :, 13]                                              # noqa: E221, E241, # we are left with 1 dimension
    azimuth = slice4D[:, :, :, 14]                                              # noqa: E221, E241, # we are left with 1 dimension
    include = slice4D[:, :, :, 15]                                              # noqa: E221, E241, # we are left with 1 dimension
    # fmt: on

    if unique is True:                                                          # we'd like to use unique offsets
        useUnique = True if include.min() == -1 else False                      # are there any -1 records ?
    else:
        useUnique = False                                                       # unique not required or not available

    if useUnique:
        included = (fold > 0) & (include == -1)                                 # fold > 0 AND unique == -1
    else:
        included = fold > 0                                                     # fold value > 0

    # the following optimization was proposed by Gemeni; it is supposedly faster than the np.count_nonzero
    # if np.count_nonzero(included) == 0:                                       # nothing to show here
    # if not np.any(included):                                                  # optimized check for empty slice

    if not np.any(included):                                                    # optimized check for empty slice
        return (None, None, True)

    offsets = offsets[included]                                                 # filter the 1D slice
    azimuth = azimuth[included]                                                 # filter the 1D slice
    return (offsets, azimuth, False)


# @jit(nopython=True)
def numbaNdft1D(kMax: float, dK: float, slice3D: np.ndarray, inclu3D: np.ndarray):
    kR = np.arange(0, kMax, dK)                                                 # numpy array with k-values [0 ... kMax]
    nK = kR.shape[0]                                                            # number of points along y-axis (size of kR array)
    nP = slice3D.shape[0]                                                       # number of points along x-axis

    radialStk = np.zeros(shape=(nP, nK), dtype=np.float32)                      # start with empty array of the right size and type

    for p in range(nP):                                                         # iterate over all points in the current direction
        offRadial = slice3D[p, :, 13]                                           # get all available offsets belonging to this row
        incRadial = inclu3D[p, :]                                               # not all points will be valid, in case of unique offsets
        n = np.count_nonzero(incRadial)                                         # normalize by actual  nr of available traces
        a = 1 / n if n > 0 else 0                                               # response will be zero for n = zero
        response = np.dot(incRadial, np.exp(2j * np.pi * kR * offRadial[:, np.newaxis])) * a
        absResponse = np.maximum(np.abs(response), LOG_RESPONSE_FLOOR)
        logResponse = np.log(absResponse) * 20.0
        radialStk[p, :] = logResponse

    return radialStk


# there are three nested for loops in the following function; *** SLOW ***. Improve performance by:
# 1) use numba to get optimized precompiled code
#    see: https://www.infoworld.com/article/3622013/speed-up-your-python-with-numba.html
# 2) use numpy arrays, similar as done in the 1D transform
# 3) spawn a different process.
#    get result in an async manner http://pymotw.com/2/multiprocessing/communication.html
# 4) save results in a cache to prevent recalculation
#    see: https://docs.python.org/dev/library/functools.html#functools.lru_cache

# there are three nested for loops in the following function; *** SLOW ***. Improve performance by:
# 1) use numba to get optimized precompiled code
#    see: https://www.infoworld.com/article/3622013/speed-up-your-python-with-numba.html
# 2) use numpy arrays, similar as done in the 1D transform
# 3) spawn a different process.
#    get result in an async manner http://pymotw.com/2/multiprocessing/communication.html
# 4) save results in a cache to prevent recalculation
#    see: https://docs.python.org/dev/library/functools.html#functools.lru_cache
@jit(nopython=True, parallel=True)
def numbaNdft2D(kMin: float, kMax: float, dK: float, offsetX: np.ndarray, offsetY: np.ndarray):
    kX = np.arange(kMin, kMax, dK)
    kY = np.arange(kMin, kMax, dK)
    nX = kX.shape[0]
    nY = kY.shape[0]
    nP = offsetX.shape[0]

    xyCellStk = np.zeros(shape=(nX, nY), dtype=np.float32)  # start with empty array of the right size and type

    for x in prange(nX):                                                        # execute parallel loops
        for y in range(nY):
            # Vectorized inner sum over all points p
            response = np.sum(np.exp(2j * np.pi * (kX[x] * offsetX + kY[y] * offsetY)))
            absResponse = np.maximum(np.abs(response) / nP, LOG_RESPONSE_FLOOR)
            # logResponse = np.round(np.log(absResponse) * 20.0, 2)           # np.round is only handy to review array contents
            logResponse = np.log(absResponse) * 20.0
            xyCellStk[x, y] = logResponse

    return xyCellStk

# the following function is a non-parallel version of the above function; it is used for testing purposes to check if the parallelization is working properly


@jit(nopython=True)
def numbaNdft2DBeforeGemini(kMin: float, kMax: float, dK: float, offsetX: np.ndarray, offsetY: np.ndarray):
    kX = np.arange(kMin, kMax, dK)
    kY = np.arange(kMin, kMax, dK)
    nX = kX.shape[0]
    nY = kY.shape[0]
    nP = offsetX.shape[0]

    xyCellStk = np.zeros(shape=(nX, nY), dtype=np.float32)  # start with empty array of the right size and type

    for x in range(nX):                                                         # execute 3 loops; x, y, fold
        for y in range(nY):
            response = 0.0
            for p in range(nP):
                e = np.exp(2j * np.pi * (kX[x] * offsetX[p] + kY[y] * offsetY[p]))
                response += e
            absResponse = np.maximum(np.abs(response) / nP, LOG_RESPONSE_FLOOR)
            # logResponse = np.round(np.log(absResponse) * 20.0, 2)           # np.round is only handy to review array contents
            logResponse = np.log(absResponse) * 20.0
            xyCellStk[x, y] = logResponse

    return xyCellStk


@jit(nopython=True)
def numbaOffInline(slice2D: np.ndarray, ox: float, component: int = 0):
    xInline = slice2D[:, 9]                                                     # get all available cmp values belonging to this row
    if component == 1:
        oInline = slice2D[:, 6] - slice2D[:, 3]                                 # inline offset component
    elif component == 2:
        oInline = slice2D[:, 7] - slice2D[:, 4]                                 # x-line offset component
    elif component == 3:
        oInline = slice2D[:, 12]                                                # TWT [ms]
    else:
        oInline = slice2D[:, 13]                                                # absolute offset magnitude

    x = np.empty((2 * xInline.size), dtype=xInline.dtype)
    x[0::2] = xInline - ox
    x[1::2] = xInline + ox

    y = np.empty((2 * oInline.size), dtype=oInline.dtype)
    y[0::2] = oInline
    y[1::2] = oInline

    return (x, y)


@jit(nopython=True)
def numbaOffXline(slice2D: np.ndarray, oy: float, component: int = 0):
    yInline = slice2D[:, 10]                                                    # get all available cmp values belonging to this row
    if component == 1:
        oInline = slice2D[:, 6] - slice2D[:, 3]                                 # inline offset component
    elif component == 2:
        oInline = slice2D[:, 7] - slice2D[:, 4]                                 # x-line offset component
    elif component == 3:
        oInline = slice2D[:, 12]                                                # TWT [ms]
    else:
        oInline = slice2D[:, 13]                                                # absolute offset magnitude

    x = np.empty((2 * yInline.size), dtype=yInline.dtype)
    x[0::2] = yInline - oy
    x[1::2] = yInline + oy

    y = np.empty((2 * oInline.size), dtype=oInline.dtype)
    y[0::2] = oInline
    y[1::2] = oInline

    return (x, y)


@jit(nopython=True)
def numbaAziInline(slice2D: np.ndarray, ox: float):
    xInline = slice2D[:, 9]                                                   # get all available cmp values belonging to this row
    oInline = slice2D[:, 14]                                                  # get all available azimuths belonging to this row

    x = np.empty((2 * xInline.size), dtype=xInline.dtype)
    x[0::2] = xInline - ox
    x[1::2] = xInline + ox

    y = np.empty((2 * oInline.size), dtype=oInline.dtype)
    y[0::2] = oInline
    y[1::2] = oInline

    return (x, y)


@jit(nopython=True)
def numbaAziXline(slice2D: np.ndarray, oy: float):
    yInline = slice2D[:, 10]                                                  # get all available cmp values belonging to this row
    oInline = slice2D[:, 14]                                                  # get all available azimuths belonging to this row

    x = np.empty((2 * yInline.size), dtype=yInline.dtype)
    x[0::2] = yInline - oy
    x[1::2] = yInline + oy

    y = np.empty((2 * oInline.size), dtype=oInline.dtype)
    y[0::2] = oInline
    y[1::2] = oInline

    return (x, y)


@jit(nopython=True)
def numbaOffsetBin(slice2D: np.ndarray, unique=False):
    if unique is True:                                                          # we'd like to use unique offsets
        havUnique = slice2D[:, 15]                                              # get all available cmp values belonging to this row
        useUnique = True if havUnique.min() == -1 else False                    # are there any -1 records ?
    else:
        useUnique = False                                                       # unique not required or not available

    if useUnique:
        included = (slice2D[:, 2] > 0) & (slice2D[:, 15] == -1)                 # fold > 0 AND unique == -1
    else:
        included = slice2D[:, 2] > 0                                            # fold > 0

    noData = np.count_nonzero(included) == 0
    if noData:
        empty = np.empty(0, dtype=slice2D.dtype)
        return (empty, empty, noData)                                           # nothing to show; return

    slice2D = slice2D[included, :]                                              # filter the 2D slice

    offsetX = slice2D[:, 6] - slice2D[:, 3]                                     # x-component of available offsets
    offsetY = slice2D[:, 7] - slice2D[:, 4]                                     # y-component of available offsets

    return (offsetX, offsetY, noData)


@jit(nopython=True)
def numbaSpiderBin(slice2D: np.ndarray):                                        # slicing should already have reduced fold to account for unique fold, if need be

    foldX2 = slice2D.shape[0] * 2

    spiderSrcX = np.zeros(shape=foldX2, dtype=np.float32)                       # needed to display data points
    spiderSrcY = np.zeros(shape=foldX2, dtype=np.float32)                       # needed to display data points
    spiderSrcZ = np.zeros(shape=foldX2, dtype=np.float32)                       # needed to display data points (3D mode)
    spiderRecX = np.zeros(shape=foldX2, dtype=np.float32)                       # needed to display data points
    spiderRecY = np.zeros(shape=foldX2, dtype=np.float32)                       # needed to display data points
    spiderRecZ = np.zeros(shape=foldX2, dtype=np.float32)                       # needed to display data points (3D mode)

    spiderSrcX[0::2] = slice2D[:, 3]                                            # src-x
    spiderSrcX[1::2] = slice2D[:, 9]                                            # cmp-x

    spiderSrcY[0::2] = slice2D[:, 4]                                            # src-y
    spiderSrcY[1::2] = slice2D[:, 10]                                           # cmp-y

    spiderSrcZ[0::2] = slice2D[:, 5]                                            # src-z
    spiderSrcZ[1::2] = slice2D[:, 11]                                           # cmp-z

    spiderRecX[0::2] = slice2D[:, 6]                                            # rec-x
    spiderRecX[1::2] = slice2D[:, 9]                                            # cmp-x

    spiderRecY[0::2] = slice2D[:, 7]                                            # rec-y
    spiderRecY[1::2] = slice2D[:, 10]                                           # cmp-y

    spiderRecZ[0::2] = slice2D[:, 8]                                            # rec-z
    spiderRecZ[1::2] = slice2D[:, 11]                                           # cmp-z

    return (spiderSrcX, spiderSrcY, spiderSrcZ, spiderRecX, spiderRecY, spiderRecZ)


def pointsInRect(pointArray: np.ndarray, rect: QRectF):
    lt = rect.left()
    rt = rect.right()
    tp = rect.top()
    bm = rect.bottom()

    return numbaPointsInRect(pointArray, lt, rt, tp, bm)


@jit(nopython=True)
def numbaPointsInRect(pointArray: np.ndarray, lt: float, rt: float, tp: float, bm: float):
    included = (pointArray[:, 0] >= lt) & (pointArray[:, 0] <= rt) & (pointArray[:, 1] >= tp) & (pointArray[:, 1] <= bm)
    return included


# See: https://stackoverflow.com/questions/49907604/setting-structured-array-field-in-numba
# See: https://stackoverflow.com/questions/52409479/python-numba-accessing-structured-numpy-array-elements-as-fast-as-possible
# See: https://stackoverflow.com/questions/58786392/how-to-create-a-list-of-structured-scalars-in-numba
# See: https://stackoverflow.com/questions/60118008/accessing-structured-data-types-in-numba-vs-numpy
# See: https://stackoverflow.com/questions/74438714/numba-signature-for-structured-arrays
# See: https://stackoverflow.com/questions/53175601/what-the-best-way-to-get-structured-array-dataframe-like-structures-in-numba
# See: https://stackoverflow.com/questions/72747804/how-to-return-a-1d-structured-array-mixed-types-from-a-numba-jit-compiled-func
# See: http://numba.pydata.org/numba-doc/0.13/arrays.html#array-creation-loop-jitting
# See: https://stackoverflow.com/questions/73892609/how-exactly-to-work-with-string-arrays-in-numba
# See: https://numba.pydata.org/numba-doc/dev/reference/pysupported.html#typed-dict
# See: https://numba.pydata.org/numba-doc/0.15.1/tutorial_firststeps.html

# numType1 = nb.from_dtype(pntType1)  # create a Numba type corresponding to the Numpy dtype
# unfortunately, this isn't recognized in time for the following jit decoration
# @jit('void(numType1[:], i8, f8, f8, i8, f8, f8, f4[3])', nopython=True)
# likewise, the following description isn't recognized at numba compile time
# numType1 = 'Record(  Line[type=float32;offset=0],   \
#                     Point[type=float32;offset=4],   \
#                     Index[type=int32;offset=8],     \
#                     Code[type=[unichr x 2];offset=12],  \
#                     Depth[type=float32;offset=20],  \
#                     East[type=float32;offset=24],   \
#                     North[type=float32;offset=28],  \
#                     Elev[type=float32;offset=32],   \
#                     Uniq[type=int32;offset=36],     \
#                     InXps[type=int32;offset=40],    \
#                     LocX[type=float32;offset=44],   \
#                     LocY[type=float32;offset=48];52;False)'


@jit(nopython=True)
def numbaSetPointRecord(array: np.ndarray, index: int, line: float, point: float, block: int, east: float, north: float, pnt: np.ndarray) -> None:
    # pntType1 = np.dtype(
    #     [
    #         ('Line', 'f4'),  # F10.2
    #         ('Point', 'f4'),  # F10.2
    #         ('Index', 'i4'),  # I1
    #         ('Code', 'U2'),  # A2
    #         ('Depth', 'f4'),  # I4
    #         ('East', 'f4'),  # F9.1
    #         ('North', 'f4'),  # F10.1
    #         ('Elev', 'f4'),  # F6.1
    #         ('Uniq', 'i4'),  # check if record is unique
    #         ('InXps', 'i4'),  # check if record is orphan
    #         ('InUse', 'i4'),  # check if record is in use
    #         ('LocX', 'f4'),  # F9.1
    #         ('LocY', 'f4'),  # F10.1
    #     ]
    # )
    array[index]['Line'] = float(int(line))
    array[index]['Point'] = float(int(point))
    array[index]['Index'] = block % 10 + 1                                      # the single digit point index is used to indicate block nr
    array[index]['East'] = east
    array[index]['North'] = north
    array[index]['LocX'] = pnt[0]                                               # x-component of 3D-location
    array[index]['LocY'] = pnt[1]                                               # y-component of 3D-location
    # Elevation = height relative to a datum (e.g. mean sea level). As we
    # have no datum information available, leave Elev at 0.0 per SPS convention.
    array[index]['Elev'] = 0.0
    # Depth = burial depth below surface (positive number for sub-surface points).
    # By convention z < 0 below surface (e.g. well TVDss), so depth = -min(z, 0).
    z = pnt[2]
    array[index]['Depth'] = -z if z < 0.0 else 0.0
    array[index]['Uniq'] = 1                                                    # later, we want to use Uniq == 1 to remove empty records at the end
    array[index]['InUse'] = 1                                                   # later, we want to use InUse == 1 to check if the point is active
    array[index]['InXps'] = 1                                                   # later, we want to use InXps == 1 to check for any xps orphns


@jit(nopython=True)
def numbaSetRelationRecord(array: np.ndarray, index: int, srcLin: float, srcPnt: float, srcInd: int, shtRec: int, recLin: float, recMin: float, recMax: float):
    array[index]['SrcLin'] = float(int(srcLin))
    array[index]['SrcPnt'] = float(int(srcPnt))
    array[index]['SrcInd'] = srcInd
    array[index]['RecNum'] = shtRec
    array[index]['RecLin'] = int(recLin)
    array[index]['RecMin'] = float(int(recMin))
    array[index]['RecMax'] = float(int(recMax))
    array[index]['RecInd'] = srcInd
    array[index]['Uniq'] = 1                                                    # needed for compacting array later (remove empty records)


@jit(nopython=True)
def numbaBinBatchParallel(
    srcBatch,
    relIndices,
    recLocs,
    recInUse,
    recPointI,
    relRecMinI,
    relRecMaxI,
    relRecStartI,
    relRecEndI,
    binOutput,
    minOffset,
    maxOffset,
    anaOutput,
    binMat,
    st2Mat,
    fullAnalysis,
    maxFold
):
    numShots = len(srcBatch)
    nx_max, ny_max = binOutput.shape

    for i in range(numShots):
        src = srcBatch[i]
        minR = relIndices[i, 0]
        maxR = relIndices[i, 1]

        if maxR <= minR:
            continue

        # Optimization: Only iterate over relations for this shot
        for relIdx in range(minR, maxR):
            rMin = relRecMinI[relIdx]
            rMax = relRecMaxI[relIdx]
            rStart = relRecStartI[relIdx]
            rEnd = relRecEndI[relIdx]
            # srcLin = srcBatch[i, 3] if srcBatch.shape[1] > 3 else 0 # fallback if not passed
            # srcPnt = srcBatch[i, 4] if srcBatch.shape[1] > 4 else 0

            # ALGORITHMIC FIX: Only iterate over receivers known to be on this line
            for rIdx in range(rStart, rEnd):
                if recPointI[rIdx] >= rMin and recPointI[rIdx] <= rMax:

                    if recInUse[rIdx] == 0:
                        continue

                    rec = recLocs[rIdx]
                    cmpX = (src[0] + rec[0]) * 0.5
                    cmpY = (src[1] + rec[1]) * 0.5

                    # Apply Bin Transform
                    nx = int(binMat[0, 0] * cmpX + binMat[0, 1] * cmpY + binMat[0, 2])
                    ny = int(binMat[1, 0] * cmpX + binMat[1, 1] * cmpY + binMat[1, 2])

                    if 0 <= nx < nx_max and 0 <= ny < ny_max:
                        dist = ((src[0] - rec[0]) ** 2 + (src[1] - rec[1]) ** 2) ** 0.5
                        if fullAnalysis:
                            fold = binOutput[nx, ny]
                            if fold < maxFold:
                                # Line & Stake calculation
                                stkX = int(st2Mat[0, 0] * cmpX + st2Mat[0, 1] * cmpY + st2Mat[0, 2])
                                stkY = int(st2Mat[1, 0] * cmpX + st2Mat[1, 1] * cmpY + st2Mat[1, 2])

                                anaOutput[nx, ny, fold, 0] = stkX
                                anaOutput[nx, ny, fold, 1] = stkY
                                anaOutput[nx, ny, fold, 2] = fold + 1
                                anaOutput[nx, ny, fold, 3] = src[0]
                                anaOutput[nx, ny, fold, 4] = src[1]
                                anaOutput[nx, ny, fold, 5] = src[2]
                                anaOutput[nx, ny, fold, 6] = rec[0]
                                anaOutput[nx, ny, fold, 7] = rec[1]
                                anaOutput[nx, ny, fold, 8] = rec[2]
                                anaOutput[nx, ny, fold, 9] = cmpX
                                anaOutput[nx, ny, fold, 10] = cmpY
                                anaOutput[nx, ny, fold, 11] = (src[2] + rec[2]) * 0.5
                                anaOutput[nx, ny, fold, 13] = dist

                        binOutput[nx, ny] += 1
                        if dist < minOffset[nx, ny]: minOffset[nx, ny] = dist  # noqa: E701
                        if dist > maxOffset[nx, ny]: maxOffset[nx, ny] = dist  # noqa: E701


@jit(nopython=True)
def numbaFixRelationRecord(array: np.ndarray, index: int, recStkX: float):
    recMin = min(array[index]['RecMin'], int(recStkX))
    recMax = max(array[index]['RecMax'], int(recStkX))

    array[index]['RecMin'] = recMin
    array[index]['RecMax'] = recMax


def clipLineF(line: QLineF, border: QRectF) -> QLineF:

    if border.isNull():
        return QLineF(line)                                                     # don't clip against an empty rect !
    if line.isNull():
        return QLineF()                                                         # null in ?  null out !

    x1 = line.x1()                                                              # copy the line elements to individual float values
    y1 = line.y1()
    x2 = line.x2()
    y2 = line.y2()

    yMin = border.top()                                                        # copy the rect elements to individual float values
    xMin = border.left()
    yMax = border.bottom()
    xMax = border.right()

    x1, y1, x2, y2 = numbaClipLineF(x1, y1, x2, y2, xMin, xMax, yMin, yMax)
    return QLineF(x1, y1, x2, y2)                                               # return the clipped line


@jit(nopython=True)
def numbaClipLineF(x1: float, y1: float, x2: float, y2: float, xMin: float, xMax: float, yMin: float, yMax: float) -> tuple:

    # Python routine to implement Cohen Sutherland algorithm for line clipping.
    # See: https://www.geeksforgeeks.org/line-clipping-set-1-cohen-sutherland-algorithm/
    # See: https://en.wikipedia.org/wiki/Cohen%E2%80%93Sutherland_algorithm
    # See: https://www.geeksforgeeks.org/line-clipping-set-2-cyrus-beck-algorithm/?ref=rp

    # Define region codes
    INSIDE = 0                                                                  # 0000
    LEFT = 1                                                                    # 0001
    RIGHT = 2                                                                   # 0010
    BOTTOM = 4                                                                  # 0100
    TOP = 8  	                                                                # 1000

    # Inner function to compute the region code for a point(x, y) relative to the border of the rectangle
    def computeCode(x, y):
        code = INSIDE
        if x < xMin:                                                           # to the left of rectangle
            code |= LEFT
        elif x > xMax:                                                         # to the right of rectangle
            code |= RIGHT
        if y < yMin:                                                           # below the rectangle
            code |= BOTTOM
        elif y > yMax:                                                         # above the rectangle
            code |= TOP
        return code

    # Compute region codes for P1, P2
    code1 = computeCode(x1, y1)
    code2 = computeCode(x2, y2)
    accept = False

    while True:                                                                 # Keep doing this, till we can escape

        if code1 == 0 and code2 == 0:                                           # both endpoints lie within rectangle
            accept = True
            break

        if (code1 & code2) != 0:                                                # both endpoints are outside rectangle
            break

        # If we get here, the line needs clipping
        # At least one of the points is outside of rect, select it

        x = 1.0
        y = 1.0

        if code1 != 0:
            codeOut = code1
        else:
            codeOut = code2

        # Find intersection point using formulas
        #   y = y1 + slope * (x - x1),
        #   x = x1 + (1 / slope) * (y - y1)
        if codeOut & TOP:  		                                            # point is above the clip rectangle
            x = x1 + (x2 - x1) * (yMax - y1) / (y2 - y1)
            y = yMax

        elif codeOut & BOTTOM:  		                                        # point is below the clip rectangle
            x = x1 + (x2 - x1) * (yMin - y1) / (y2 - y1)
            y = yMin

        elif codeOut & RIGHT:  		                                        # point is to the right of the clip rectangle
            y = y1 + (y2 - y1) * (xMax - x1) / (x2 - x1)
            x = xMax

        elif codeOut & LEFT:  		                                            # point is to the left of the clip rectangle
            y = y1 + (y2 - y1) * (xMin - x1) / (x2 - x1)
            x = xMin

        # Now an intersection point (x, y) has been found,
        # we replace the point outside clipping rectangle
        # by the intersection point
        if codeOut == code1:
            x1 = x
            y1 = y
            code1 = computeCode(x1, y1)
        else:
            x2 = x
            y2 = y
            code2 = computeCode(x2, y2)

    if accept:
        return (x1, y1, x2, y2)                                                 # return the clipped line

    return (0.0, 0.0, 0.0, 0.0)                                                 # return a null line


@jit(nopython=True, parallel=True)
def pointInPolygon(xy, poly):
    # See: https://stackoverflow.com/questions/52471590/use-multithreading-in-numba
    D = np.empty(len(xy), dtype=bool)
    n = len(poly)
    for i in range(1, len(D) - 1):
        inside = False
        p2x = 0.0
        p2y = 0.0
        xints = 0.0
        p1x, p1y = poly[0]
        x = xy[i][0]
        y = xy[i][1]
        for i in range(n + 1):
            p2x, p2y = poly[i % n]
            if y > min(p1y, p2y):
                if y <= max(p1y, p2y):
                    if x <= max(p1x, p2x):
                        if p1y != p2y:
                            xints = (y - p1y) * (p2x - p1x) / (p2y - p1y) + p1x
                        if p1x == p2x or x <= xints:
                            inside = not inside
            p1x, p1y = p2x, p2y
        D[i] = inside
    return D


# ----------------------------------------------------------------------------
# Hot-path bin-update kernel for full-analysis binning runs.
#
# This is the numba equivalent of the per-trace Python ``for k in range(N):``
# loop in ``RollSurvey._applyBinUpdatesVectorized`` (writeAnalysis branch).
# It performs the same per-trace work in compiled code:
#   * read fold = binOutput[x, y]
#   * if fold < maxFold, write 15 anaOutput columns (cols 0..14) at slot fold
#   * binOutput[x, y] += 1
#   * update minOffset[x, y] / maxOffset[x, y] with hypArray[k]
#
# Semantics are byte-for-byte identical to the Python loop; only the loop
# body runs in compiled code. parallel=True is intentionally NOT used here
# because the per-(x, y) fold counter creates a write hazard across threads.
#
# cache=True persists the compiled artifact to __pycache__/ so the
# 20-40 s first-call JIT compile is paid only once per (numba, numpy,
# python) version triple, not on every QGIS restart.
# ----------------------------------------------------------------------------
@jit(nopython=True, cache=False)
def numbaApplyBinUpdatesAnalysis(
    nx,                # int64[N]
    ny,                # int64[N]
    stkX,              # int32[N]
    stkY,              # int32[N]
    cmpPoints,         # float32[N, 3]
    recPoints,         # float32[N, 3]
    hypArray,          # float32[N]
    aziArray,          # float32[N]
    totalTime,         # float32[N] (zeros if caller did not compute travel time)
    src,               # float32[3]
    binOutput,         # uint32[NX, NY]
    minOffset,         # float32[NX, NY]
    maxOffset,         # float32[NX, NY]
    anaOutput,         # float32[NX, NY, MAXFOLD, 16]
    maxFold,           # int
):
    n = nx.shape[0]
    for k in range(n):
        x = nx[k]
        y = ny[k]
        fold = binOutput[x, y]
        if fold < maxFold:
            anaOutput[x, y, fold, 0] = stkX[k]
            anaOutput[x, y, fold, 1] = stkY[k]
            anaOutput[x, y, fold, 2] = fold + 1
            anaOutput[x, y, fold, 3] = src[0]
            anaOutput[x, y, fold, 4] = src[1]
            anaOutput[x, y, fold, 5] = src[2]
            anaOutput[x, y, fold, 6] = recPoints[k, 0]
            anaOutput[x, y, fold, 7] = recPoints[k, 1]
            anaOutput[x, y, fold, 8] = recPoints[k, 2]
            anaOutput[x, y, fold, 9] = cmpPoints[k, 0]
            anaOutput[x, y, fold, 10] = cmpPoints[k, 1]
            anaOutput[x, y, fold, 11] = cmpPoints[k, 2]
            anaOutput[x, y, fold, 12] = totalTime[k]
            anaOutput[x, y, fold, 13] = hypArray[k]
            anaOutput[x, y, fold, 14] = aziArray[k]
        binOutput[x, y] = fold + 1
        h = hypArray[k]
        if h < minOffset[x, y]:
            minOffset[x, y] = h
        if h > maxOffset[x, y]:
            maxOffset[x, y] = h


# ----------------------------------------------------------------------------
# Batched variant of numbaApplyBinUpdatesAnalysis used by binTemplate10.
#
# Identical body to numbaApplyBinUpdatesAnalysis except `srcArr` is a per-trace
# (M, 3) float32 array instead of a single shared float32[3]. Each trace k
# stores its own (srcArr[k, 0..2]) into anaOutput[..., 3..5], so a whole
# (Ns * Nr) template batch can be dispatched in one kernel call without
# losing the src-per-trace info that the existing per-source kernel takes
# from the shared `src[3]` argument.
#
# Same write hazard caveat: parallel=False because per-(x, y) fold counter
# is not safe to race across threads. cache=True keeps the JIT artifact.
# ----------------------------------------------------------------------------
@jit(nopython=True, cache=False)
def numbaApplyBinUpdatesAnalysisBatch(
    nx,                # int64[N]
    ny,                # int64[N]
    stkX,              # int32[N]
    stkY,              # int32[N]
    cmpPoints,         # float32[N, 3]
    recPoints,         # float32[N, 3]
    hypArray,          # float32[N]
    aziArray,          # float32[N]
    totalTime,         # float32[N]
    srcArr,            # float32[N, 3]  -- per-trace src (the only difference)
    binOutput,         # uint32[NX, NY]
    minOffset,         # float32[NX, NY]
    maxOffset,         # float32[NX, NY]
    anaOutput,         # float32[NX, NY, MAXFOLD, 16]
    maxFold,           # int
):
    n = nx.shape[0]
    for k in range(n):
        x = nx[k]
        y = ny[k]
        fold = binOutput[x, y]
        if fold < maxFold:
            anaOutput[x, y, fold, 0] = stkX[k]
            anaOutput[x, y, fold, 1] = stkY[k]
            anaOutput[x, y, fold, 2] = fold + 1
            anaOutput[x, y, fold, 3] = srcArr[k, 0]
            anaOutput[x, y, fold, 4] = srcArr[k, 1]
            anaOutput[x, y, fold, 5] = srcArr[k, 2]
            anaOutput[x, y, fold, 6] = recPoints[k, 0]
            anaOutput[x, y, fold, 7] = recPoints[k, 1]
            anaOutput[x, y, fold, 8] = recPoints[k, 2]
            anaOutput[x, y, fold, 9] = cmpPoints[k, 0]
            anaOutput[x, y, fold, 10] = cmpPoints[k, 1]
            anaOutput[x, y, fold, 11] = cmpPoints[k, 2]
            anaOutput[x, y, fold, 12] = totalTime[k]
            anaOutput[x, y, fold, 13] = hypArray[k]
            anaOutput[x, y, fold, 14] = aziArray[k]
        binOutput[x, y] = fold + 1
        h = hypArray[k]
        if h < minOffset[x, y]:
            minOffset[x, y] = h
        if h > maxOffset[x, y]:
            maxOffset[x, y] = h
