#===============================================================
# This code is based on Matlab demo of JMI by Xander Staal
# These are the subroutines for the JMI code
# 
# Author : Aayush Garg 
# Date   : January 2020
# Update : Eric Verschuur
# Date   : May 2021
#          Included comments for clarification 
#          Included option for primaries-only, the "_prim" versions
# 
# Copyright 2020 - Delft University of Technology
# This code belongs to the Delphi Consortium
# It cannot be distributed to third parties without
# written approval from Delphi via d.j.verschuur@tudelft.nl
# It can be used and modified by members of the Delphi Consortium
#===============================================================

import numpy as np
#import scipy.io as mat_io
import math

#--------------------------------------------------------------------
# create a Ricker wavelet
#--------------------------------------------------------------------
def ricker(f0, nt, dt):
    nf = nt//2 + 1
    df = 1/(nt*dt)

    wave=np.zeros((1,nf))

    for ifr in range(nf):
        f = ifr * df
        wave[0,ifr] = 2 * f**2 / (np.sqrt(math.pi) * f0**3) * np.exp(-1 * f**2 / (f0**2))

    wave = nt * wave / (2 * np.sum(wave))

    return wave

#--------------------------------------------------------------------
# make source
#--------------------------------------------------------------------
def make_S(wave,nx,nt,dt):
    nf = nt//2 + 1
    dom = 2*math.pi/(nt*dt)
    S = np.zeros((nf,nx,nx), dtype=complex)
    for ifr in range(1,nf):
        om = ifr * dom
        S[ifr,:,:] = wave[0,ifr] * np.eye(nx) / np.sqrt(1j * om)

    return S

#--------------------------------------------------------------------
# apply taper to the propagator matrix
#--------------------------------------------------------------------
def make_prop_taper(angle,nxtap,nx,dx,dz):

    nxflat = np.int32(dz*np.tan(angle*math.pi/180)/dx)
    ntap = np.int32((nx-2*nxflat)/2)

    nx2 = nx + 2*nxtap

    tap = 0.5 + 0.5 * np.cos(-1 *math.pi * np.arange(1,ntap+1) / ntap)
    wtaper = np.concatenate((np.ones(nxflat), tap), axis=0)

    Wtaper = np.diag(np.ones(nx2))


    # print(Wtaper.shape)
    for idiag in range(wtaper.size):
        Wtaper = Wtaper + wtaper[idiag] * np.diag(np.ones(nx2-idiag-1), idiag+1)
        Wtaper = Wtaper + wtaper[idiag] * np.diag(np.ones(nx2-idiag-1), -idiag-1)


    tap2 = 0.5 + 0.5 * np.cos(-1 *math.pi * np.arange(1,nxtap+1) / nxtap)
    wtaper2 = np.concatenate((tap2[::-1], np.ones(nx), tap2), axis=0)

    Wtaper = np.diag(wtaper2) @  Wtaper @ np.diag(wtaper2)

    return Wtaper

#--------------------------------------------------------------------
# initialize the propagator matrices
#--------------------------------------------------------------------
def init_prop(nx,dx,nz,ifmin,ifmax):
    '''
    Function to initialise global parameters for the wave propagation
    '''

    global Wkxshift
    global Wvelglobal
    global Wxglobal
    global Gvelglobal
    global Gxglobal


    nfeval = ifmax - ifmin + 1

    # define global parameters for propagation operators

    nkx = (2**math.ceil(math.log(nx,2)))
    dkx = 2*math.pi/(nkx*dx)
    kx = dkx * np.concatenate( (np.arange(0,nkx//2 + 1), np.arange(-(nkx//2-1),0)), axis=0)
    kx3D = kx[None, :, None] + np.zeros((nfeval, nkx, nx))
    #print(kx3D.shape)

    x0 = np.arange(0,nx) * dx
    x03D = x0[None, None, :] + np.zeros((nfeval, nkx, nx))
    Wkxshift = np.exp(-1j * x03D * kx3D)

    Wvelglobal = np.zeros((nz,nx))
    Wxglobal = np.zeros((nz, nfeval, nx, nx), dtype=complex)
    Gvelglobal = np.zeros((nz,nx))
    Gxglobal = np.zeros((nz, nfeval, nx, nx), dtype=complex)

#--------------------------------------------------------------------
# calculate the downgoing wavefields Pplus and Qplus
#--------------------------------------------------------------------
def make_Pplus(S,Pmin,imag,vel,Wtaper,nx,dx,nz,dz,nt,dt,ifmin,ifmax):

    nfeval = ifmax - ifmin + 1
    nsrc = S.shape[2]

    dom = 2 * math.pi / (nt*dt)
    om = dom * np.arange(ifmin, ifmax+1)

    nx2 = (2**math.ceil(math.log(nx,2)))
    dkx = 2*math.pi/(nx2*dx)
    kx = dkx * np.concatenate( (np.arange(0,nx2//2 + 1), np.arange(-(nx2//2-1),0)), axis=0)

    # preallocate wavefields
    Pplus = np.zeros(Pmin.shape, dtype=complex)
    Qplus = np.zeros(Pmin.shape, dtype=complex)
    pextr = np.zeros((nfeval, nx, nsrc), dtype=complex)

    for iz in range(nz):

        ptmp = np.squeeze(Pmin[iz,:,:,:])
        r = imag[iz,:]
        R = r[None,:,None] + np.zeros((nfeval, nx, nsrc))

        if iz == 0:
            pextr = (1 + R)*pextr - R*ptmp + S
        else:
            pextr = (1 + R)*pextr - R*ptmp

        Qplus[iz,:,:,:] = pextr

        v = vel[iz,:]
        pextr = WFprop(pextr,v,Wtaper,om,kx,nx,nx2,dx,dz,ifmin,ifmax);
        Pplus[iz+1,:,:,:] = pextr


    return Pplus, Qplus

#--------------------------------------------------------------------
# calculate the downgoing wavefields without internal multiples/transm.
# therefore, as input Pmin and imag are not required
#--------------------------------------------------------------------
def make_Pplus_prim(S,vel,Wtaper,nx,dx,nz,dz,nt,dt,ifmin,ifmax):

    nfeval = ifmax - ifmin + 1
    nsrc = S.shape[2]

    dom = 2 * math.pi / (nt*dt)
    om = dom * np.arange(ifmin, ifmax+1)

    nx2 = (2**math.ceil(math.log(nx,2)))
    dkx = 2*math.pi/(nx2*dx)
    kx = dkx * np.concatenate( (np.arange(0,nx2//2 + 1), np.arange(-(nx2//2-1),0)), axis=0)

    # preallocate wavefields
    Pplus = np.zeros((nz+1,nfeval,nx,nsrc), dtype=complex)
    Qplus = np.zeros((nz+1,nfeval,nx,nsrc), dtype=complex)
    
    # start with Pplus=S at surface
    pextr = S
    
    for iz in range(nz):

        Qplus[iz,:,:,:] = pextr

        v = vel[iz,:]
        pextr = WFprop(pextr,v,Wtaper,om,kx,nx,nx2,dx,dz,ifmin,ifmax);
        Pplus[iz+1,:,:,:] = pextr

    return Pplus, Qplus

#--------------------------------------------------------------------
# apply forward wavefield propagation of Pin over one depth level
#--------------------------------------------------------------------
def WFprop(Pin,vel,Wtaper,om,kx,nx,nx2,dx,dz,ifmin,ifmax):
    global Wkxshift
    global Wvelglobal
    global Wxglobal

    s = 1 / vel
    nfeval = ifmax - ifmin + 1

    nz = Wvelglobal.shape[0]
    vv = vel[None,:] + np.zeros((nz,vel.size))
    temp = np.sum((Wvelglobal-vv)**2, axis=1)
    dvmin = np.min(temp)
    izmin = np.argmin(temp)
    dvmax = np.max(temp)
    izmax = np.argmax(temp)



    if dvmin > 1.0:

        om23D = om[:,None,None]**2 + np.zeros((nfeval, nx2, s.size))
        kx3D = kx[None, :, None] + np.zeros((nfeval, kx.shape[0], nx))
        s23D = s[None, None, :]**2 + np.zeros((nfeval, nx2, nx))

        kz3D = np.conj(np.sqrt((om23D * s23D - kx3D**2).astype(complex)))
        Wkx = np.exp(-1j * kz3D * dz) * Wkxshift

        Wxtmp = np.fft.ifft(Wkx,axis=1)
        Wx = Wxtmp[:,:nx,:]

        Wxglobal[izmax,:,:,:]=Wx
        Wvelglobal[izmax,:]=vel

        #print('make in WFprop new global operator for izmax=%d' % izmax)

    else:
        Wx = Wxglobal[izmin,:,:,:]
        #print('use WFprop global operator izmin=%d' % izmin)

    tap3D = Wtaper[None,:,:] + np.zeros((nfeval, Wtaper.shape[0], Wtaper.shape[1]))
    Wx = tap3D * Wx

    #Pout = np.zeros(Pin.shape, dtype=complex)
    Pout = Wx @ Pin

    return Pout

#--------------------------------------------------------------------
# 2D inverse FFT from wavenumber to space
#--------------------------------------------------------------------
def ifft_2(in_arr,nf,ifmin,ifmax):
    nt = (nf-1)*2;
    out_arr = np.zeros((nt,in_arr.shape[1]), dtype=complex)
    out_arr[ifmin:ifmax+1, :] = in_arr
    out_arr[nf:nt, :] = np.conj(out_arr[nf-2:0:-1, :])
    out_arr = np.fft.ifft(out_arr, axis=0).real

    return out_arr


#--------------------------------------------------------------------
# calculate the upgoing wavefields Pmin and Qmin
#--------------------------------------------------------------------
def make_Pmin(Pplus,imag,vel,Wtaper,nx,dx,nz,dz,nt,dt,ifmin,ifmax):


    nfeval = ifmax - ifmin + 1
    nsrc = Pplus.shape[3]

    dom = 2 * math.pi / (nt*dt)
    om = dom * np.arange(ifmin, ifmax+1)

    nx2 = (2**math.ceil(math.log(nx,2)))
    dkx = 2*math.pi/(nx2*dx)
    kx = dkx * np.concatenate( (np.arange(0,nx2//2 + 1), np.arange(-(nx2//2-1),0)), axis=0)

    # preallocate wavefields
    Pmin = np.zeros(Pplus.shape, dtype=complex)
    Qmin = np.zeros(Pplus.shape, dtype=complex)
    pextr = np.zeros((nfeval, nx, nsrc), dtype=complex)

    for iz in range(nz-1,-1,-1):

        ptmp = np.squeeze(Pplus[iz+1,:,:,:])

        r = imag[iz+1,:]
        R = r[None,:,None] + np.zeros((nfeval, nx, nsrc))

        pextr = (1 - R)*pextr + R*ptmp

        Qmin[iz+1,:,:,:] = pextr

        v = vel[iz,:]
        pextr = WFprop(pextr,v,Wtaper,om,kx,nx,nx2,dx,dz,ifmin,ifmax);
        Pmin[iz,:,:,:] = pextr

    return Pmin, Qmin

#--------------------------------------------------------------------
# calculate the upgoing wavefields without internal multiples/transm.
#--------------------------------------------------------------------
def make_Pmin_prim(Pplus,imag,vel,Wtaper,nx,dx,nz,dz,nt,dt,ifmin,ifmax):


    nfeval = ifmax - ifmin + 1
    nsrc = Pplus.shape[3]

    dom = 2 * math.pi / (nt*dt)
    om = dom * np.arange(ifmin, ifmax+1)

    nx2 = (2**math.ceil(math.log(nx,2)))
    dkx = 2*math.pi/(nx2*dx)
    kx = dkx * np.concatenate( (np.arange(0,nx2//2 + 1), np.arange(-(nx2//2-1),0)), axis=0)

    # preallocate wavefields
    Pmin = np.zeros(Pplus.shape, dtype=complex)
    Qmin = np.zeros(Pplus.shape, dtype=complex)
    pextr = np.zeros((nfeval, nx, nsrc), dtype=complex)

    for iz in range(nz-1,-1,-1):

        ptmp = np.squeeze(Pplus[iz+1,:,:,:])

        r = imag[iz+1,:]
        R = r[None,:,None] + np.zeros((nfeval, nx, nsrc))

        pextr = pextr + R*ptmp

        Qmin[iz+1,:,:,:] = pextr

        v = vel[iz,:]
        pextr = WFprop(pextr,v,Wtaper,om,kx,nx,nx2,dx,dz,ifmin,ifmax);
        Pmin[iz,:,:,:] = pextr

    return Pmin, Qmin


#--------------------------------------------------------------------
# Calculate the reflectivity image gradient
#--------------------------------------------------------------------
def make_gradientA(Res,Pplus,Pmin,imag,vel,Wtaper,nx,dx,nz,dz,nt,dt,ifmin,ifmax):


    nfeval = ifmax - ifmin + 1
    dom = 2 * math.pi / (nt*dt)
    om = dom * np.arange(ifmin, ifmax+1)

    nx2 = (2**math.ceil(math.log(nx,2)))
    dkx = 2*math.pi/(nx2*dx)
    kx = dkx * np.concatenate( (np.arange(0,nx2//2 + 1), np.arange(-(nx2//2-1),0)), axis=0)

    nsrc, nrcv = Pplus.shape[3], Pplus.shape[2]

    dimag = np.zeros((nz+1,nx))
    sens = np.zeros((nz+1,nx))
    wextr = np.zeros((nfeval,nx,nx), dtype=complex)
    pextr = np.zeros((nfeval,nrcv,nsrc), dtype=complex)
    # needed?
    # Wmin = np.zeros((nz+1,nfeval,nx,nx), dtype=complex)

    for iz in range(nz):
        pplus = np.squeeze(Pplus[iz+1, :, :, :])

        v = vel[iz,:]
        r = imag[iz,:]
        R = r[None,:,None] + np.zeros((nfeval,nx,nx))

        if iz == 0:
            wextr = (1 - R) * wextr + np.eye(nx).reshape(1,nx,nx)
            #print(R.shape)
            #print(wextr.shape)
        else:
            wextr = (1 - R) * wextr

        wextr = WBprop(wextr,v,Wtaper,om,kx,nx,nx2,dx,dz,ifmin,ifmax)
        #Wmin[iz+1,:,:,:] = wextr

        pextr = wextr @ Res

        dimag[iz+1,:] = 4 * np.sum(np.real(pextr * np.conj(pplus)), axis=(0,2))
        sens[iz+1,:] = 4 * np.sum(np.abs(pplus)**2, axis=(0,2)) * np.sum(np.abs(wextr)**2, axis=(0,2))

    return dimag, sens

#--------------------------------------------------------------------
# Calculate the reflectivity image gradient primaries-only option
# therefore, parameters Pmin and imag are not required
#--------------------------------------------------------------------
def make_gradientA_prim(Res,Pplus,vel,Wtaper,nx,dx,nz,dz,nt,dt,ifmin,ifmax):


    nfeval = ifmax - ifmin + 1
    dom = 2 * math.pi / (nt*dt)
    om = dom * np.arange(ifmin, ifmax+1)

    nx2 = (2**math.ceil(math.log(nx,2)))
    dkx = 2*math.pi/(nx2*dx)
    kx = dkx * np.concatenate( (np.arange(0,nx2//2 + 1), np.arange(-(nx2//2-1),0)), axis=0)

    nsrc, nrcv = Pplus.shape[3], Pplus.shape[2]

    dimag = np.zeros((nz+1,nx))
    sens = np.zeros((nz+1,nx))
    wextr = np.zeros((nfeval,nx,nx), dtype=complex)
    pextr = np.zeros((nfeval,nrcv,nsrc), dtype=complex)

    for iz in range(nz):
        pplus = np.squeeze(Pplus[iz+1, :, :, :])

        v = vel[iz,:]

        if iz == 0:
            wextr = wextr + np.eye(nx).reshape(1,nx,nx)

        wextr = WBprop(wextr,v,Wtaper,om,kx,nx,nx2,dx,dz,ifmin,ifmax)

        pextr = wextr @ Res

        dimag[iz+1,:] = 4 * np.sum(np.real(pextr * np.conj(pplus)), axis=(0,2))
        sens[iz+1,:] = 4 * np.sum(np.abs(pplus)**2, axis=(0,2)) * np.sum(np.abs(wextr)**2, axis=(0,2))

    return dimag, sens

#--------------------------------------------------------------------
# apply backward wavefield propagation to wavefield Pin over one dz
#--------------------------------------------------------------------
def WBprop(Pin,vel,Wtaper,om,kx,nx,nx2,dx,dz,ifmin,ifmax):
    global Wkxshift
    global Wvelglobal
    global Wxglobal

    s = 1 / vel
    nfeval = ifmax - ifmin + 1

    nz = Wvelglobal.shape[0]
    vv = vel[None,:] + np.zeros((nz,vel.size))
    temp = np.sum((Wvelglobal-vv)**2, axis=1)
    dvmin = np.min(temp)
    izmin = np.argmin(temp)
    dvmax = np.max(temp)
    izmax = np.argmax(temp)



    if dvmin > 1.0:

        om23D = om[:,None,None]**2 + np.zeros((nfeval, nx2, s.size))
        kx3D = kx[None, :, None] + np.zeros((nfeval, kx.shape[0], nx))
        s23D = s[None, None, :]**2 + np.zeros((nfeval, nx2, nx))

        kz3D = np.conj(np.sqrt((om23D * s23D - kx3D**2).astype(complex)))
        Wkx = np.exp(-1j * kz3D * dz) * Wkxshift

        Wxtmp = np.fft.ifft(Wkx,axis=1)
        Wx = Wxtmp[:,:nx,:]

        Wxglobal[izmax,:,:,:]=Wx
        Wvelglobal[izmax,:]=vel

        #print('make in WFprop new global operator for izmax=%d' % izmax)

    else:
        Wx = Wxglobal[izmin,:,:,:]
        #print('use WFprop global operator izmin=%d' % izmin)

    tap3D = Wtaper[None,:,:] + np.zeros((nfeval, Wtaper.shape[0], Wtaper.shape[1]))
    Wx = tap3D * Wx

    Pout = np.conj(Wx).transpose((0,2,1)) @ Pin

    return Pout


#--------------------------------------------------------------------
# apply band-limitation to the velocity gradient (smoothing)
#--------------------------------------------------------------------
def bandlimit_vel_gradient(grad,fmax,vel,dx,dz):

    nz, nx = grad.shape
    dkx = 2 * np.pi / (nx*dx)
    dkz = 2 * np.pi / (nz*dz)
    nkz = int(np.ceil(nz/2))
    nkx = int(np.ceil(nx/2))

    kmax = 2 * np.pi * fmax / np.min(vel)

    # define the vertical filter
    Filterz = np.zeros(nz)
    for ikz in range(nkz):
        kz = ikz * dkz
        Filterz[ikz] = np.exp(-1 * kz**2 / (kmax**2))

    Filterz[nz-nkz+1:] = np.conj(Filterz[nkz-1:0:-1])
    filterz = np.real(np.fft.fft(Filterz))/nz
    filterz = np.roll(filterz, -nkz)

    # define the vertical filter
    Filterx = np.zeros(nx)
    for ikx in range(nkx):
        kx = ikx * dkx
        Filterx[ikx] = np.exp(-1 * kx**2 / (kmax**2))

    Filterx[nx-nkx+1:] = np.conj(Filterx[nkx-1:0:-1])
    filterx = np.real(np.fft.fft(Filterx))/nx
    filterx = np.roll(filterx, -nkx)

    # extend the model
    gradtmp = np.zeros((nz+2*nkz, nx+2*nkx))
    gradtmp[nkz:nkz+nz, nkx:nkx+nx] = grad

    # apply filtering
    grad_bltmp1 = np.zeros((2*nz+2*nkz-1, nx+2*nkx))
    for ix in range(nx+2*nkx):
        grad_bltmp1[:,ix] = np.convolve(gradtmp[:,ix], filterz[nz-1::-1])

    grad_bltmp2 = np.zeros((2*nz+2*nkz-1, 2*nx+2*nkx-1))
    for iz in range(2*nz+2*nkz-1):
        grad_bltmp2[iz,:] = np.convolve(grad_bltmp1[iz,:], filterx[nx-1::-1])


    grad_bl = grad_bltmp2[2*nkz-1:2*nkz+nz-1, 2*nkx-1:2*nkx+nx-1]

    return grad_bl

#--------------------------------------------------------------------
# calculate velocity gradient
#--------------------------------------------------------------------
def make_gradientB(Res,Qplus,Qmin,imag,vel,Wtaper,nx,dx,nz,dz,nt,dt,ifmin,ifmax):

    nfeval = ifmax - ifmin + 1
    dom = 2 * math.pi / (nt*dt)
    om = dom * np.arange(ifmin, ifmax+1)

    nx2 = (2**math.ceil(math.log(nx,2)))
    dkx = 2*math.pi/(nx2*dx)
    kx = dkx * np.concatenate( (np.arange(0,nx2//2 + 1), np.arange(-(nx2//2-1),0)), axis=0)

    nsrc, nrcv = Qmin.shape[3], Qmin.shape[2]
    stab = 0.01

    dchi = np.zeros((nz,nx))
    sens = np.zeros((nz,nx))
    wextr = np.zeros((nfeval,nx,nx), dtype=complex)
    dpextr = np.zeros((nfeval,nrcv, nsrc), dtype=complex)
    Wmin = np.zeros((nz+1,nfeval,nx,nx), dtype=complex)

    for iz in range(nz):
        qmin = np.squeeze(Qmin[iz+1, :, :, :])

        v = vel[iz,:]
        r = imag[iz,:]
        R = r[None,:,None] + np.zeros((nfeval, nx, nx))

        if iz == 0:
            wextr = (1 - R) * wextr + np.eye(nx).reshape(1,nx,nx)
        else:
            wextr = (1 - R) * wextr

        dwextr = GBprop(wextr,v,Wtaper,om,kx,nx,nx2,dx,dz,ifmin,ifmax,stab)

        dpextr = dwextr @ Res

        dchi[iz,:] = 4 * np.sum(np.real(dpextr * np.conj(qmin)), axis=(0,2))
        sens[iz,:] = 4 * np.sum(np.abs(qmin)**2, axis=(0,2)) * np.sum(np.abs(dwextr)**2, axis=(0,2))

        wextr = WBprop(wextr,v,Wtaper,om,kx,nx,nx2,dx,dz,ifmin,ifmax)
        Wmin[iz+1,:,:,:] = wextr

    #print(np.max(dchi))
    wextr = 0 * wextr

    for iz in range(nz-1,-1,-1):

        qplus = np.squeeze(Qplus[iz,:,:,:])
        wmin = np.squeeze(Wmin[iz+1,:,:,:])

        v = vel[iz, :]
        r = imag[iz+1, :]
        R = r[None,:,None] + np.zeros((nfeval,nx, nx))

        wextr = (1 + R) * wextr + R * wmin
        dwextr=GBprop(wextr,v,Wtaper,om,kx,nx,nx2,dx,dz,ifmin,ifmax,stab)

        dpextr = dwextr @ Res

        dchi[iz,:] = dchi[iz,:] + 4 * np.sum(np.real(dpextr * np.conj(qplus)), axis=(0,2))
        sens[iz,:] = sens[iz,:] + 4 * np.sum(np.abs(qplus)**2, axis=(0,2)) * np.sum(np.abs(dwextr)**2, axis=(0,2))

        wextr = WBprop(wextr,v,Wtaper,om,kx,nx,nx2,dx,dz,ifmin,ifmax)

    return dchi, sens

#--------------------------------------------------------------------
# calculate velocity gradient for primaries-only option
#--------------------------------------------------------------------
def make_gradientB_prim(Res,Qplus,Qmin,imag,vel,Wtaper,nx,dx,nz,dz,nt,dt,ifmin,ifmax):

    nfeval = ifmax - ifmin + 1
    dom = 2 * math.pi / (nt*dt)
    om = dom * np.arange(ifmin, ifmax+1)

    nx2 = (2**math.ceil(math.log(nx,2)))
    dkx = 2*math.pi/(nx2*dx)
    kx = dkx * np.concatenate( (np.arange(0,nx2//2 + 1), np.arange(-(nx2//2-1),0)), axis=0)

    nsrc, nrcv = Qmin.shape[3], Qmin.shape[2]
    stab = 0.01

    dchi = np.zeros((nz,nx))
    sens = np.zeros((nz,nx))
    wextr = np.zeros((nfeval,nx,nx), dtype=complex)
    dpextr = np.zeros((nfeval,nrcv, nsrc), dtype=complex)
    Wmin = np.zeros((nz+1,nfeval,nx,nx), dtype=complex)

    for iz in range(nz):
        qmin = np.squeeze(Qmin[iz+1, :, :, :])

        v = vel[iz,:]

        if iz == 0:
            wextr = wextr + np.eye(nx).reshape(1,nx,nx)

        dwextr = GBprop(wextr,v,Wtaper,om,kx,nx,nx2,dx,dz,ifmin,ifmax,stab)

        dpextr = dwextr @ Res

        dchi[iz,:] = 4 * np.sum(np.real(dpextr * np.conj(qmin)), axis=(0,2))
        sens[iz,:] = 4 * np.sum(np.abs(qmin)**2, axis=(0,2)) * np.sum(np.abs(dwextr)**2, axis=(0,2))

        wextr = WBprop(wextr,v,Wtaper,om,kx,nx,nx2,dx,dz,ifmin,ifmax)
        Wmin[iz+1,:,:,:] = wextr

    #print(np.max(dchi))
    wextr = 0 * wextr

    for iz in range(nz-1,-1,-1):

        qplus = np.squeeze(Qplus[iz,:,:,:])
        wmin = np.squeeze(Wmin[iz+1,:,:,:])

        v = vel[iz, :]
        r = imag[iz+1, :]
        R = r[None,:,None] + np.zeros((nfeval,nx, nx))

        wextr = wextr + R * wmin
        dwextr=GBprop(wextr,v,Wtaper,om,kx,nx,nx2,dx,dz,ifmin,ifmax,stab)

        dpextr = dwextr @ Res

        dchi[iz,:] = dchi[iz,:] + 4 * np.sum(np.real(dpextr * np.conj(qplus)), axis=(0,2))
        sens[iz,:] = sens[iz,:] + 4 * np.sum(np.abs(qplus)**2, axis=(0,2)) * np.sum(np.abs(dwextr)**2, axis=(0,2))

        wextr = WBprop(wextr,v,Wtaper,om,kx,nx,nx2,dx,dz,ifmin,ifmax)

    return dchi, sens

#--------------------------------------------------------------------
# apply the backward propagation for the velocity gradient
#--------------------------------------------------------------------
def GBprop(Pin,vel,Wtaper,om,kx,nx,nx2,dx,dz,ifmin,ifmax,stab):

    global Wkxshift
    global Gvelglobal
    global Gxglobal

    s = 1 / vel
    nfeval = ifmax - ifmin + 1

    nz = Gvelglobal.shape[0]
    vv = vel[None, :] + np.zeros((nz,vel.size))

    temp = np.sum((Gvelglobal-vv)**2, axis=1)
    dvmin = np.min(temp)
    izmin = np.argmin(temp)
    dvmax = np.max(temp)
    izmax = np.argmax(temp)

    if dvmin > 1.0:

        om3D = om[:,None,None] + np.zeros((nfeval, nx2, s.size))
        kx3D = kx[None, :, None] + np.zeros((nfeval, kx.shape[0], nx))
        s3D = np.transpose(s[None,None,:], (1, 0, 2)) + np.zeros((nfeval, nx2, nx))

        # calculate propagation operator
        kz3D = np.conj(np.sqrt((om3D**2 * s3D**2 - kx3D**2).astype(complex)))
        eps = stab * om3D**2 * s3D**2

        Gkx = -1j * dz * om3D**2 * s3D * np.conj(kz3D) / (np.abs(kz3D)**2+eps) * np.exp(-1j * kz3D * dz) * Wkxshift
        Gxtmp = np.fft.ifft(Gkx,axis=1)
        Gx = Gxtmp[:,:nx,:]

        # store operator in the table
        Gxglobal[izmax,:,:,:] = Gx
        Gvelglobal[izmax,:] = vel

    else:
        Gx = Gxglobal[izmin,:,:,:]

    tap3D = Wtaper[None,:,:] + np.zeros((nfeval, Wtaper.shape[0], Wtaper.shape[1]))
    Gx = tap3D * Gx

    Pout = np.conj(Gx).transpose((0,2,1)) @ Pin

    return Pout

#--------------------------------------------------------------------
# apply the forward propagation for the velocity gradient
#--------------------------------------------------------------------
def GFprop(Pin,vel,Wtaper,om,kx,nx,nx2,dx,dz,ifmin,ifmax,stab):

    global Wkxshift
    global Gvelglobal
    global Gxglobal

    s = 1 / vel
    nfeval = ifmax - ifmin + 1

    nz = Gvelglobal.shape[0]
    vv = vel[None, :] + np.zeros((nz,vel.size))

    temp = np.sum((Gvelglobal-vv)**2, axis=1)
    dvmin = np.min(temp)
    izmin = np.argmin(temp)
    dvmax = np.max(temp)
    izmax = np.argmax(temp)

    if dvmin > 1.0:

        om3D = om[:,None,None] + np.zeros((nfeval, nx2, s.size))
        kx3D = kx[None, :, None] + np.zeros((nfeval, kx.shape[0], nx))
        s3D = np.transpose(s[None,None,:], (1, 0, 2)) + np.zeros((nfeval, nx2, nx))

        # calculate propagation operator
        kz3D = np.conj(np.sqrt((om3D**2 * s3D**2 - kx3D**2).astype(complex)))
        eps = stab * om3D**2 * s3D**2

        Gkx = -1j * dz * om3D**2 * s3D * np.conj(kz3D) / (np.abs(kz3D)**2+eps) * np.exp(-1j * kz3D * dz) * Wkxshift
        Gxtmp = np.fft.ifft(Gkx,axis=1)
        Gx = Gxtmp[:,:nx,:]

        # store operator in the table
        Gxglobal[izmax,:,:,:] = Gx
        Gvelglobal[izmax,:] = vel

    else:
        Gx = Gxglobal[izmin,:,:,:]

    tap3D = Wtaper[None,:,:] + np.zeros((nfeval, Wtaper.shape[0], Wtaper.shape[1]))
    Gx = tap3D * Gx

    Pout = Gx @ Pin

    return Pout

#--------------------------------------------------------------------
# calculate the wavefield perturbation due to the image update
#--------------------------------------------------------------------
def make_dPmin0A(Pplus,Pmin,dimag,imag,vel,Wtaper,nx,dx,nz,dz,nt,dt,ifmin,ifmax):

    nfeval = ifmax - ifmin + 1
    nsrc = Pplus.shape[3]

    dom = 2 * math.pi / (nt*dt)
    om = dom * np.arange(ifmin, ifmax+1)

    nx2 = (2**math.ceil(math.log(nx,2)))
    dkx = 2*math.pi/(nx2*dx)
    kx = dkx * np.concatenate( (np.arange(0,nx2//2 + 1), np.arange(-(nx2//2-1),0)), axis=0)

    dpextr = np.zeros((nfeval, nx, nsrc), dtype=complex)
    dPplus = np.zeros(Pmin.shape, dtype=complex)

    for iz in range(nz-1,-1,-1):
        pplus = np.squeeze(Pplus[iz+1,:,:,:])
        dpplus = np.squeeze(dPplus[iz+1,:,:,:])

        # scattering due to contrasts
        dr = dimag[iz+1,:]
        dR = dr[None,:,None] + np.zeros((nfeval, nx, nsrc))

        r = imag[iz+1,:]
        R = r[None,:,None] + np.zeros((nfeval, nx, nsrc))

        dpextr = (1-R) * dpextr + R * dpplus + dR * pplus

        v = vel[iz,:]
        dpextr = WFprop(dpextr,v,Wtaper,om,kx,nx,nx2,dx,dz,ifmin,ifmax)

    dPmin0 = dpextr

    return dPmin0

#--------------------------------------------------------------------
# calculate the wavefield perturbation due to the image update
# primaries-only option, negelcting internal multiples+transmission
#--------------------------------------------------------------------
def make_dPmin0A_prim(Pplus,dimag,vel,Wtaper,nx,dx,nz,dz,nt,dt,ifmin,ifmax):

    nfeval = ifmax - ifmin + 1
    nsrc = Pplus.shape[3]

    dom = 2 * math.pi / (nt*dt)
    om = dom * np.arange(ifmin, ifmax+1)

    nx2 = (2**math.ceil(math.log(nx,2)))
    dkx = 2*math.pi/(nx2*dx)
    kx = dkx * np.concatenate( (np.arange(0,nx2//2 + 1), np.arange(-(nx2//2-1),0)), axis=0)

    dpextr = np.zeros((nfeval, nx, nsrc), dtype=complex)

    for iz in range(nz-1,-1,-1):
        pplus = np.squeeze(Pplus[iz+1,:,:,:])

        # scattering due to contrasts
        dr = dimag[iz+1,:]
        dR = dr[None,:,None] + np.zeros((nfeval, nx, nsrc))

        dqextr = dpextr + dR * pplus

        v = vel[iz,:]
        dpextr = WFprop(dqextr,v,Wtaper,om,kx,nx,nx2,dx,dz,ifmin,ifmax)

    dPmin0 = dpextr

    return dPmin0

#--------------------------------------------------------------------
# calculate the wavefield perturbation due to the velocity update
#--------------------------------------------------------------------
def make_dPmin0B(Qplus,Qmin,imag,vel,chi_v,Wtaper,nx,dx,nz,dz,nt,dt,ifmin,ifmax):
    nfeval = ifmax - ifmin + 1
    nsrc, nrcv = Qmin.shape[3], Qmin.shape[2]

    dom = 2 * math.pi / (nt*dt)
    om = dom * np.arange(ifmin, ifmax+1)

    nx2 = (2**math.ceil(math.log(nx,2)))
    dkx = 2*math.pi/(nx2*dx)
    kx = dkx * np.concatenate( (np.arange(0,nx2//2 + 1), np.arange(-(nx2//2-1),0)), axis=0)

    stab = 0.01

    pextr = np.zeros((nfeval, nx, nsrc), dtype=complex)
    dPplus = np.zeros(Qmin.shape, dtype=complex)

    for iz in range(nz):
        qplus = np.squeeze(Qplus[iz+1,:,:,:])

        v = vel[iz,:]
        r = imag[iz,:]
        R = r[None,:,None] + np.zeros((nfeval, nx, nsrc))

        # scattering in the background model
        qextr = (1 + R) * pextr

        # propagate
        pextr = WFprop(qextr,v,Wtaper,om,kx,nx,nx2,dx,dz,ifmin,ifmax)

        # scattering due to the contrasts
        cv = chi_v[iz,:]
        CV = cv[None,:,None] + np.zeros((nfeval, nx, nsrc))
        dq = CV * qplus
        dp = GFprop(dq,v,Wtaper,om,kx,nx,nx2,dx,dz,ifmin,ifmax,stab)

        pextr = pextr + dp
        dPplus[iz+1, :, :, :] = pextr

    pextr = 0 * pextr

    for iz in range(nz-1,-1,-1):
        qmin = np.squeeze(Qmin[iz+1,:,:])
        dpplus = np.squeeze(dPplus[iz+1,:,:,:])

        v = vel[iz,:]
        r = imag[iz+1,:]
        R = r[None,:,None] + np.zeros((nfeval, nx, nsrc))

        # Scattering in the background model
        qextr = (1-R) * pextr + R * dpplus

        # propagate
        pextr = WFprop(qextr,v,Wtaper,om,kx,nx,nx2,dx,dz,ifmin,ifmax)

        # scattering due to contrasts
        cv = chi_v[iz,:]
        CV = cv[None,:,None] + np.zeros((nfeval, nx, nsrc))
        dq = CV * qmin
        dp = GFprop(dq,v,Wtaper,om,kx,nx,nx2,dx,dz,ifmin,ifmax,stab)

        pextr = pextr + dp

    return pextr

#--------------------------------------------------------------------
# calculate the wavefield perturbation due to the velocity update
# primaries_only option, without internal multiples and transmission
#--------------------------------------------------------------------
def make_dPmin0B_prim(Qplus,Qmin,imag,vel,chi_v,Wtaper,nx,dx,nz,dz,nt,dt,ifmin,ifmax):
    nfeval = ifmax - ifmin + 1
    nsrc, nrcv = Qmin.shape[3], Qmin.shape[2]

    dom = 2 * math.pi / (nt*dt)
    om = dom * np.arange(ifmin, ifmax+1)

    nx2 = (2**math.ceil(math.log(nx,2)))
    dkx = 2*math.pi/(nx2*dx)
    kx = dkx * np.concatenate( (np.arange(0,nx2//2 + 1), np.arange(-(nx2//2-1),0)), axis=0)

    stab = 0.01

    pextr = np.zeros((nfeval, nx, nsrc), dtype=complex)
    dPplus = np.zeros(Qmin.shape, dtype=complex)

    for iz in range(nz):
        qplus = np.squeeze(Qplus[iz+1,:,:,:])

        v = vel[iz,:]

        # propagate
        pextr = WFprop(pextr,v,Wtaper,om,kx,nx,nx2,dx,dz,ifmin,ifmax)

        # scattering due to the contrasts
        cv = chi_v[iz,:]
        CV = cv[None,:,None] + np.zeros((nfeval, nx, nsrc))
        dq = CV * qplus
        dp = GFprop(dq,v,Wtaper,om,kx,nx,nx2,dx,dz,ifmin,ifmax,stab)

        pextr = pextr + dp
        dPplus[iz+1, :, :, :] = pextr

    pextr = 0 * pextr

    for iz in range(nz-1,-1,-1):
        qmin = np.squeeze(Qmin[iz+1,:,:])
        dpplus = np.squeeze(dPplus[iz+1,:,:,:])

        v = vel[iz,:]
        r = imag[iz+1,:]
        R = r[None,:,None] + np.zeros((nfeval, nx, nsrc))

        # Scattering in the background model
        pextr = pextr + R * dpplus

        # propagate
        pextr = WFprop(pextr,v,Wtaper,om,kx,nx,nx2,dx,dz,ifmin,ifmax)

        # scattering due to contrasts
        cv = chi_v[iz,:]
        CV = cv[None,:,None] + np.zeros((nfeval, nx, nsrc))
        dq = CV * qmin
        dp = GFprop(dq,v,Wtaper,om,kx,nx,nx2,dx,dz,ifmin,ifmax,stab)

        pextr = pextr + dp

    return pextr
