import numpy as np
import math
import matplotlib.pyplot as plt
from IPython.display import clear_output

class PSPI:

    def __init__(self,nx,nf,dx,dt,fmin,fmax,dz):

        # determine original time length
        nt=2*(nf-1)
        
        # define some basic arrays and variables
        df = 1.0 /(nt*dt)
        ifmin=int(fmin/df)
        ifmax=int(fmax/df)
        nfeval = ifmax - ifmin + 1
        # define omega axis for range fmin-fmax
        dom = 2 * math.pi *df
        self.om = dom * np.arange(ifmin, ifmax+1)
        # define kx-axis
        nkx = nx
        dkx = 2*math.pi/(nkx*dx)
        self.kx = dkx * np.concatenate( (np.arange(0,nkx//2 + 1), np.arange(-(nkx//2-1),0)), axis=0)
        # initial mask
        self.mask=np.zeros((nkx,nfeval))+1.0
        # pass parameters to other functions
        self.noper_prev=0
        self.vprev=np.zeros(nx)
        self.nkx=nkx
        self.dz=dz
        self.ifmin=ifmin
        self.ifmax=ifmax

    def propagate(self,fdata,v,vstepref,fb):
        # some parameters
        dz = self.dz
        ifmin = self.ifmin
        ifmax = self.ifmax
        nfeval = ifmax - ifmin + 1
        nkx = self.nkx
        nx = fdata.shape[0]
        # wave propagation via PSPI
        vmin=np.min(v)
        vmax=np.max(v)
        noper=int(1+(vmax-vmin)/vstepref)
        if noper > 1:
            vstep=(vmax-vmin)/(noper-1)
        else:
            vstep=vstepref
        voper=vmin+np.arange(noper)*vstep
        # creat table for phase shift operators
        if noper != self.noper_prev:
            #print('found noper=%d<>noper_prev=%d' %(noper,self.noper_prev))
            self.operkz=np.zeros((noper,nkx,nfeval),dtype='complex128')
        # define difference with previous velocities
        err=np.sum((self.vprev-v)**2)
        # generate new operators only if velocity changed
        if err > 1:
            #print('new propagator with vmin=%g vmax=%g' %(vmin,vmax))
            for ioper in range(noper):
                c=voper[ioper]
                for ifreq in range(nfeval):
                    k=self.om[ifreq]/c
                    kz=np.real(np.sqrt(k*k-self.kx*self.kx+0j))
                    self.mask[:,ifreq]=self.mask[:,ifreq]*np.where(kz < 0.0001, 0.0, 1.0)
                    self.operkz[ioper,:,ifreq]=self.mask[:,ifreq]*np.exp(-1j*kz*dz)
        #else:
        #    print('keep same propagator with vmin=%g vmax=%g' %(vmin,vmax))
            
        # FFT from x to kx
        fkdata=np.fft.fft(fdata[:,ifmin:ifmax+1],axis=0)
        # generate table of wavefields in x and kx domain if needed
        if noper != self.noper_prev:
            self.fkdata_table=np.zeros((noper,nkx,nfeval),dtype='complex128')
            self.fdata_table=np.zeros((noper,nx,nfeval),dtype='complex128')
        # apply phase shift operators and inverse FFT to x-domain
        if fb == 'forward':
            for ioper in range(noper):
                self.fkdata_table[ioper,:,:]=fkdata*self.operkz[ioper,:,:]
                self.fdata_table[ioper,:,:]=np.fft.ifft(self.fkdata_table[ioper,:,:],axis=0)
        else:
            for ioper in range(noper):
                self.fkdata_table[ioper,:,:]=fkdata*np.conj(self.operkz[ioper,:,:])
                self.fdata_table[ioper,:,:]=np.fft.ifft(self.fkdata_table[ioper,:,:],axis=0)
        # interpolation per velocity value
        for ix in range(nx):
            vloc=v[ix]
            index0=int((vloc-vmin)/vstep)
            index1=min(index0+1,noper-1)
            v0=vmin+index0*vstep
            scl1=(vloc-v0)/vstep
            scl0=1.0-scl1
            fdata[ix,ifmin:ifmax+1]=scl0*self.fdata_table[index0,ix,:]+scl1*self.fdata_table[index1,ix,:]
        # update noper_prev
        self.noper_prev=noper
        self.vprev=v
        # return the extrapolated data
        return fdata
    
    def display_wavefields(self,fsource,fdata,image,ntdisp,z,rmin,rmax):
        source_extrap=np.fft.irfft(fsource)
        data_extrap=np.fft.irfft(fdata)
        vmax=np.max(data_extrap)
        if ntdisp < 700:
            vsize=8
        else:
            vsize=10
        clear_output(wait=True)
        plt.figure(figsize=(5,vsize))
        ax1 = plt.subplot2grid(shape=(3, 2), loc=(0, 0), rowspan=2)
        ax2 = plt.subplot2grid(shape=(3, 2), loc=(0, 1), rowspan=2)
        ax3 = plt.subplot2grid(shape=(3, 2), loc=(2, 0), colspan=2)

        ax1.imshow(source_extrap[:,:ntdisp].T,vmin=-vmax,vmax=vmax)
        ax1.set_title('Source field at z=%g' %(z),fontsize=8)
        ax2.imshow(data_extrap[:,:ntdisp].T,vmin=-vmax,vmax=vmax)
        ax2.set_title('Inverse extrap data at z=%g' %(z),fontsize=8)
        ax3.imshow(image, vmin=rmin, vmax=rmax)
        ax3.set_title('Image built until z=%g' %(z))
        plt.show()
        
    def mute_source(self,fsource,ntwin):
        eps=0.9
        nx=fsource.shape[0]
        nf=fsource.shape[1]
        nt=(nf-1)*2
        tmp=np.fft.irfft(fsource)
        taper1=np.array([0.80,0.6,0.4,0.2])
        taper2=np.array([0.20,0.4,0.6,0.8])
        for ix in range(nx):
            datmax=np.max(tmp[ix,:])
            for it in range(nt):
                if np.abs(tmp[ix,it])>eps*datmax:
                    tmp[ix,0:it-ntwin-4]=0.0
                    if it-ntwin-4 >= 0:
                        tmp[ix,it-ntwin-4:it-ntwin]=taper1*tmp[ix,it-ntwin-4:it-ntwin]
                    if it+ntwin < nt-4:
                        tmp[ix,it+ntwin:it+ntwin+4]=taper2*tmp[ix,it+ntwin:it+ntwin+4]
                    tmp[ix,it+ntwin+4:]=0.0
                    break
        fsource=np.fft.rfft(tmp)
        return fsource
