"""
utils.py

Mathematical and signal processing utilities for acoustic tracking.

(c) Walter M.X. Zimmer (2025)
Original code produced for the CIAN Tool Kit, associated with the manuscript: 
Berkenbaum, L., Glotin, H., Sarano, F., Zimmer, W.M.X., Sarano, V., Adam, O., Giraudet, P. (2026). "Audiovisual diarization of overlapping click trains in sperm whale (Physeter macrocephalus) vocal sparring using a three-hydrophone array," J. Acoust. Soc. Am.

Includes shared constants for the OPALE array configuration.

"""

import numpy as np
import soundfile as sf
from numba import jit

# =================================================================
# SHARED CONSTANTS (OPALE ARRAY & ENVIRONMENT)
# =================================================================
C_SOUND = 1541.0                           # Local sound speed in m/s
HSEL = np.array([[0, 1], [0, 2], [1, 2]])  # Sensor pairs for cross-correlation

# =================================================================
# FUNCTIONS
# =================================================================
def wav_read(filename, t1, dt):
    with sf.SoundFile(filename) as sound:
        fs = sound.samplerate
        nch = sound.channels
        nd = sound.frames
        sound.seek(int(fs*t1))
        data = sound.read(int(fs*dt))
        if data.ndim == 1: 
            data = data.reshape(-1,1)
    return data, fs, nd

def fft_filt(b, x, zi=None, nh=0):
    if x.ndim == 1: x = x.reshape(-1,1)
    if b.ndim == 1: b = b.reshape(-1,1)

    L_I, N_I = b.shape
    L_sig, N_sig = x.shape
    N_chan = max(N_sig, N_I)

    L_F = 2<<(L_I-1).bit_length()
    L_S = L_F - L_I + 1
    if L_sig < L_F:
        offsets = range(1)
    else:
        offsets = range(0, L_sig, L_S)

    if np.iscomplexobj(b) or np.iscomplexobj(x):
        fft_func = np.fft.fft
        ifft_func = np.fft.ifft
        res = np.zeros((L_sig+L_F, N_chan), dtype=np.complex128)
    else:
        fft_func = np.fft.rfft
        ifft_func = np.fft.irfft
        res = np.zeros((L_sig+L_F, N_chan))

    FDir = fft_func(b, n=L_F, axis=0)

    for n in offsets:
        u1 = fft_func(x[n:n+L_S,:], n=L_F, axis=0)
        u2 = u1 * FDir
        res[n:n+L_F,:] += ifft_func(u2, axis=0)

    if zi is not None:
        res[:zi.shape[0],:] = res[:zi.shape[0],:] + zi
        zi = res[L_sig:,:]
        return res[nh:nh+L_sig,:], zi 
    else:
        return res[nh:nh+L_sig,:]

@jit(nopython=True, cache=True)
def bit_length(n): 
    bits = 0
    if n == 0: return 0
    while n: 
        n >>= 1
        bits += 1
    return bits

def stack(uu):
    return uu + np.ones((uu.shape[0],1)) * range(uu.shape[1])

def fftCorr(zz, hsel, nfft=None):
    win = 1
    nz = np.shape(zz)[0]
    nc = np.shape(hsel)[0]
    vv = np.zeros((nz, nc))
    if nfft == None: nfft = nz
    for jj in range(nc):
        u0 = zz[:, hsel[jj,0]]
        u1 = zz[:, hsel[jj,1]]
        v0 = np.fft.rfft(u0*win, n=nfft, axis=0)
        v1 = np.fft.rfft(u1*win, n=nfft, axis=0)
        wo = v1 * np.conjugate(v0)
        wo = wo / np.sqrt(np.abs(wo))
        ww = np.fft.irfft(wo, axis=0)
        vv[:, jj] = np.fft.fftshift(ww)[:nz]
    return vv

@jit(nopython=True, cache=True)
def xCorr(uu, hsel, mc):
    nc = np.shape(hsel)[0]
    vv = np.zeros((1+mc+mc, nc))

    for jj in range(nc):
        u0 = uu[:, hsel[jj,0]].copy()
        u1 = uu[:, hsel[jj,1]].copy()
        for ii in range(mc):
            vv[ii, jj] = np.dot(u1[:-(mc-ii)], u0[(mc-ii):])
        vv[mc, jj] = np.dot(u1, u0)
        for ii in range(1, mc+1):
            vv[mc+ii, jj] = np.dot(u1[ii:], u0[:-ii])
            
    return vv

def quadInt(xx, idl):
    nx, nc = xx.shape
    ni = len(idl)
    if ni > nc: return None

    idl = np.minimum(nx-2, np.maximum(1, idl))
    uo = xx[idl, range(ni)]
    um = xx[idl-1, range(ni)]
    up = xx[idl+1, range(ni)]

    b = (up + um - 2*uo) / 2
    xo = (um - up) / (4*b)
    yo = uo - b*xo**2 

    jdl = idl.astype('float') + xo
    adl = yo
    return jdl, adl

def baseband_filter(xx, fs, fk, Q):
    N = np.ceil(fs/fk * Q).astype(int)
    W = np.zeros(N, dtype=complex)
    W = np.hanning(N) / N * np.exp(-2 * np.pi * 1j * Q * np.arange(N)/N)
    yy = fft_filt(W, xx, zi=None, nh=0)
    return yy

@jit(nopython=True, cache=True)
def leak(zz, zo, aa):
    bb = np.exp(np.log(aa)/zo)
    ny = zz.shape[0]
    for ii in range(1, ny):
        ax = zz[ii-1,:] * bb
        zz[ii,:] = np.maximum(ax, zz[ii,:])  
    return(zz)  

def extract(uu, ii, i0=-1000, n0=2000):
    i1 = ii + i0
    i2 = i1 + n0
    return uu[i1:i2,:]

def roty(X, a):
    Y = X.copy()
    Y[:,0] = X[:,0] * np.cos(a) - X[:,2] * np.sin(a)
    Y[:,2] = X[:,0] * np.sin(a) + X[:,2] * np.cos(a)
    return Y