import numpy as np
import os
import time
from tqdm import tqdm
import matplotlib.pyplot as plt
from IPython.display import clear_output
import boundaries

class FDTDEL:

    def __init__(self, cp_grid, cs_grid, rho_grid, dx, dt, src_data, src_x_index, 
                 src_z_index, src_type, src_fmax, boundary_method_order, 
                 boundary_n_gridpoints, boundary_absorbing_function):

        courant=0.707
        gridpoints_per_min_wavelength=5
        
        if boundary_method_order <= 1:
            self.boundary_method_order = boundary_method_order
        else:
            print('For elastic FD modeling we can only use boundary_metohd_order=1')
            self.boundary_method_order = 1

        self.boundary_n_gridpoints = boundary_n_gridpoints
        self.boundary_absorbing_function = boundary_absorbing_function

        # Calculate grid spacing from src_fmax
        # Starting with 'lambda_min' as min wavelength
        # make sure that we do not include grid points where c=0
        cpmax=cp_grid.max()
        csmax=cs_grid.max()
        cp_grid_cond=np.where(cp_grid<1.0,0.0*cp_grid+cpmax,cp_grid)
        cpmin=cp_grid_cond.min()
        cs_grid_cond=np.where(cs_grid<1.0,0.0*cs_grid+csmax,cs_grid)
        csmin=cs_grid_cond.min()
        lambda_min = csmin / src_fmax

        # Grid and time spacing in distance / time units respectively
        self.dx = dx
        dx_opt = lambda_min / gridpoints_per_min_wavelength
        self.dt = dt
        dt_opt = courant / cp_grid.max() * dx
        
        # Grid dimensions without boundary zone from the input model
        nx_orig  = cp_grid.shape[1]
        nz_orig  = cp_grid.shape[0]
        
        # Same with boundary zone
        self.nx = nx_orig + 2 * self.boundary_n_gridpoints
        self.nz = nz_orig + 2 * self.boundary_n_gridpoints
        
        # display some info on the modeling
        print('Starting FD module with dx=%f dt=%f' %(dx,dt))
        if cpmin != cp_grid.min():
            print('Found free surface boundary in the velocity model')
        print('Velocity model nx=%d z=%d csmin=%f csmax=%f cpmin=%f cpmax=%f'\
              %(nx_orig,nz_orig,csmin,csmax,cpmin,cpmax))
        print('Given csmin, maximum dx=%f' %(dx_opt))
        print('Given cpmax, maximum dt=%f' %(dt_opt))

        # extend model with absorbing boundaries, copy last values
        cp_grid_ext=np.pad(cp_grid,((boundary_n_gridpoints,boundary_n_gridpoints),\
                             (boundary_n_gridpoints,boundary_n_gridpoints)),mode='edge')
        cs_grid_ext=np.pad(cs_grid,((boundary_n_gridpoints,boundary_n_gridpoints),\
                             (boundary_n_gridpoints,boundary_n_gridpoints)),mode='edge')
    
        rho_grid_ext=np.pad(rho_grid,((boundary_n_gridpoints,boundary_n_gridpoints),\
                            (boundary_n_gridpoints,boundary_n_gridpoints)),mode='edge')
    
        # pass velocity and density model with extended boundaries
        self.cp = cp_grid_ext
        self.cs = cs_grid_ext
        self.rho = rho_grid_ext
        
        # we also need Lame parameters: mu=rho*cs^2, lambda+2mu=rho*cp^2 and lambda
        self.mu = self.rho*self.cs*self.cs
        self.la2mu = self.rho*self.cp*self.cp
        self.la = self.la2mu - 2*self.mu
        
        # set absorbing boundary layer
        self.sigma = boundaries.get_sigma((self.nz, self.nx), boundary_n_gridpoints, boundary_absorbing_function)

        # define the initial wavefields
        self.tzz = np.zeros((self.nz, self.nx))
        self.txx = np.zeros((self.nz, self.nx))
        self.txz = np.zeros((self.nz, self.nx))
        self.vz = np.zeros((self.nz, self.nx))
        self.vx = np.zeros((self.nz, self.nx))
        self.time_index = 0
        
        # defined the source in the original grid (including extra boundaries)
        self.src_x_index = src_x_index + boundary_n_gridpoints
        self.src_z_index = src_z_index + boundary_n_gridpoints
        self.src_data = src_data
        self.src_type = src_type

    def calc_wave_equation(self):
        nx=self.nx
        nz=self.nz
        # all variables txx, tzz,txz,vx,vz are updated
        # first from tau's to v's over 0.5dt
        fact=self.dt/(self.dx*self.rho)
        # step 1: dvx/dt = (1/rho)[dtxx/dx+dtxz/dz]
        self.vx[1:nz,0:-1] = self.vx[1:nz,0:-1]+\
            fact[1:nz,0:-1]*(self.txx[1:nz,1:nx]-self.txx[1:nz,0:-1]+\
                             self.txz[1:nz,0:-1]-self.txz[0:-1,0:-1])
        # step 2: dvz/dt = (1/rho)[dtxz/dx+dtzz/dz]
        self.vz[0:-1,1:nx] = self.vz[0:-1,1:nx]+\
            fact[0:-1,1:nx]*(self.txz[0:-1,1:nx]-self.txz[0:-1,0:-1]+\
                             self.tzz[1:nz,1:nx]-self.tzz[0:-1,1:nx])
        # next, from v's to tau's over 0.5dt
        # step 3: dtxx/dt = (lambda+2mu)*dvx/dx+lambda*dvz/dz
        fact1=(self.dt/self.dx)*self.la2mu
        fact2=(self.dt/self.dx)*self.la
        self.txx[1:nz,1:nx] = self.txx[1:nz,1:nx]+\
            fact1[1:nz,1:nx]*(self.vx[1:nz,1:nx]-self.vx[1:nz,0:-1])+\
            fact2[1:nz,1:nx]*(self.vz[1:nz,1:nx]-self.vz[0:-1,1:nx])
        # step 4: dtzz/dt = (lambda+2mu)*dvz/dz+lambda*dvx/dx
        self.tzz[1:nz,1:nx] = self.tzz[1:nz,1:nx]+\
            fact1[1:nz,1:nx]*(self.vz[1:nz,1:nx]-self.vz[0:-1,1:nx])+\
            fact2[1:nz,1:nx]*(self.vx[1:nz,1:nx]-self.vx[1:nz,0:-1])
        # step 5: dtxz/dt = mu*(dvx/dz+dvz/dx)
        fact=self.dt/(self.dx)*self.mu
        self.txz[0:-1,0:-1] = self.txz[0:-1,0:-1]+\
            fact[0:-1,0:-1]*(self.vx[1:nz,0:-1]-self.vx[0:-1,0:-1]+\
                             self.vz[0:-1,1:nx]-self.vz[0:-1,0:-1])
        # this function has no return variable, it just updated all wavefields

    def step(self):
        # calculate wave equation with finite differences
        self.calc_wave_equation()
        # add source value at source location (if sample index is not too large)
        if isinstance(self.src_x_index, np.ndarray):
            # array of sources
            if self.time_index < self.src_data.shape[1]:
                for k in range(len(self.src_x_index)):
                    ix=self.src_x_index[k]
                    if self.src_type == 'monopole':
                        self.txx[self.src_z_index, ix] += self.src_data[k,self.time_index]
                        self.tzz[self.src_z_index, ix] += self.src_data[k,self.time_index]
                    else:
                        self.txx[self.src_z_index+1, ix] += self.src_data[k,self.time_index]
                        self.tzz[self.src_z_index+1, ix] += self.src_data[k,self.time_index]
                        self.txx[self.src_z_index-1, ix] -= self.src_data[k,self.time_index]
                        self.tzz[self.src_z_index-1, ix] -= self.src_data[k,self.time_index]
        else:
            # single source
            if self.time_index < len(self.src_data):
                if self.src_type == 'monopole':
                    self.txx[self.src_z_index, self.src_x_index] += self.src_data[self.time_index]
                    self.tzz[self.src_z_index, self.src_x_index] += self.src_data[self.time_index]
                else:
                    self.txx[self.src_z_index+1, self.src_x_index] += self.src_data[self.time_index]
                    self.tzz[self.src_z_index+1, self.src_x_index] += self.src_data[self.time_index]
                    self.txx[self.src_z_index-1, self.src_x_index] -= self.src_data[self.time_index]
                    self.tzz[self.src_z_index-1, self.src_x_index] -= self.src_data[self.time_index]

        # NEED TO BE UPDATED
        # calculate absorbing boundary conditions
        # Clayton & Engquist 1977 https://authors.library.caltech.edu/34673/
        tzz_boundary = boundaries.calculate_boundaries(self.tzz, None, self.dt, self.dx, self.cp,
                                                     self.boundary_n_gridpoints,
                                                     method='clayton-engquist',
                                                     order=self.boundary_method_order)

        # apply hybrid absorbing boundary conditions
        # Liu & Sen 2011 https://www.researchgate.net/publication/228978155_A_hybrid_scheme_for_absorbing_edge_reflections_in_numerical_modeling_of_wave_propagation
        self.tzz = self.sigma*self.tzz + (1-self.sigma)*tzz_boundary

        # also for txx
        txx_boundary = boundaries.calculate_boundaries(self.txx, None, self.dt, self.dx, self.cp,
                                                     self.boundary_n_gridpoints,
                                                     method='clayton-engquist',
                                                     order=self.boundary_method_order)

        self.txx = self.sigma*self.txx + (1-self.sigma)*txx_boundary

        txz_boundary = boundaries.calculate_boundaries(self.txz, None, self.dt, self.dx, self.cp,
                                                     self.boundary_n_gridpoints,
                                                     method='clayton-engquist',
                                                     order=self.boundary_method_order)

        self.txz = self.sigma*self.txz + (1-self.sigma)*txz_boundary

        vz_boundary = boundaries.calculate_boundaries(self.vz, None, self.dt, self.dx, self.cp,
                                                     self.boundary_n_gridpoints,
                                                     method='clayton-engquist',
                                                     order=self.boundary_method_order)

        self.vz = self.sigma*self.vz + (1-self.sigma)*vz_boundary

        vx_boundary = boundaries.calculate_boundaries(self.vx, None, self.dt, self.dx, self.cp,
                                                     self.boundary_n_gridpoints,
                                                     method='clayton-engquist',
                                                     order=self.boundary_method_order)

        self.vx = self.sigma*self.vx + (1-self.sigma)*vx_boundary

        # update
        #self.pold, self.p = np.copy(self.p), np.copy(pnew)
        self.time_index += 1
        
        
    def run_simulation(self, n_steps, int_QC):

        time_Start = time.time()
        bg = self.boundary_n_gridpoints
        synth_seis = np.zeros((n_steps,self.nz, self.nx))

        print("Running simulation in Numpy\n")
        progress_bar = tqdm(range(n_steps))
        
        plt.figure()
        for index in progress_bar:
            progress_bar.set_postfix({'t=': f'{self.time_index*self.dt:03f} s',
                                      't_index': self.time_index})

            # make plot of snapshot if desired
            if int_QC > 0 and index > 0:
                if index % int_QC == 0:
                    clear_output(wait=True)
                    plt.imshow(self.tzz[bg:-bg,bg:-bg],cmap='gray')
                    plt.title('Snapshot at index=%d t=%g' %(index,index*self.dt))
                    plt.show()


            self.step()
            synth_seis[index,:,:]=self.tzz

        synth_seis=synth_seis[:,bg:-bg,bg:-bg]
        print("\nDuration: {:.2f} sec\n".format(time.time() - time_Start))
       
        return synth_seis
    