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 FDTDAC:

    def __init__(self, c_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
        
        self.boundary_method_order = boundary_method_order
        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
        cmax=c_grid.max()
        c_grid_cond=np.where(c_grid<1.0,0.0*c_grid+cmax,c_grid)
        cmin=c_grid_cond.min()
        lambda_min = cmin / 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 / c_grid.max() * dx
        
        # Grid dimensions without boundary zone from the input model
        nx_orig  = c_grid.shape[1]
        nz_orig  = c_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 cmin != c_grid.min():
            print('Found free surface boundary in the velocity model')
        print('Velocity model nx=%d z=%d cmin=%f cmax=%f' %(nx_orig,nz_orig,cmin,cmax))
        print('Given cmin, maximum dx=%f' %(dx_opt))
        print('Given cmax, maximum dt=%f' %(dt_opt))

        # extend model with absorbing boundaries, copy last values
        c_grid_ext=np.pad(c_grid,((boundary_n_gridpoints,boundary_n_gridpoints),\
                             (boundary_n_gridpoints,boundary_n_gridpoints)),mode='edge')
        
        if rho_grid is None:
            print('No density model provided')
        else:
            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.c = c_grid_ext
        self.rho = None if rho_grid is None else rho_grid_ext

        # set absorbing boundary layer
        self.sigma = boundaries.get_sigma((self.nz, self.nx), boundary_n_gridpoints, boundary_absorbing_function)

        # define the initial wavefields
        self.p = np.zeros((self.nz, self.nx))
        self.pold = 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 return_solution(self):
        return self.p

    def calc_wave_equation(self):
        p_we = np.zeros((self.nz, self.nx))
        # O(t^2, x^2) discretization scheme
        if self.rho is None:

            pxx = (self.p[1:-1, 2:] - 2 * self.p[1:-1, 1:-1] + self.p[1:-1, 0:-2]) / (self.dx ** 2)
            pzz = (self.p[0:-2, 1:-1] - 2 * self.p[1:-1, 1:-1] + self.p[2:, 1:-1]) / (self.dx ** 2)

            p_we[1:-1, 1:-1] = 2 * self.p[1:-1, 1:-1] - self.pold[1:-1, 1:-1] + \
                               self.dt ** 2 * self.c[1:-1, 1:-1] ** 2 * (pxx + pzz)
        else:
            x_pos = (slice(1, self.nz - 1), slice(2, self.nx))
            x_neg = (slice(1, self.nz - 1), slice(0, self.nx - 2))
            z_pos = (slice(2, self.nz), slice(1, self.nx - 1))
            z_neg = (slice(0, self.nz - 2), slice(1, self.nx - 1))
            p_we[1:-1, 1:-1] = (2 * self.p[1:-1, 1:-1] - self.pold[1:-1, 1:-1]
                  + self.rho[1:-1, 1:-1] * self.c[1:-1, 1:-1] ** 2 * self.dt ** 2 / (2 * self.dx ** 2) *
                  (
                        + (1 / self.rho[x_pos] + 1 / self.rho[1:-1, 1:-1]) * self.p[x_pos]
                        - (1 / self.rho[x_pos] + 2 / self.rho[1:-1, 1:-1] + 1 / self.rho[x_neg]) * self.p[1:-1, 1:-1]
                        + (1 / self.rho[1:-1, 1:-1] + 1 / self.rho[x_neg]) * self.p[x_neg]
                        + (1 / self.rho[z_pos] + 1 / self.rho[1:-1, 1:-1]) * self.p[z_pos]
                        - (1 / self.rho[z_pos] + 2 / self.rho[1:-1, 1:-1] + 1 / self.rho[z_neg]) * self.p[1:-1, 1:-1]
                        + (1 / self.rho[1:-1, 1:-1] + 1 / self.rho[z_neg]) * self.p[z_neg]
                    ))
        return p_we

    def step(self):
        # calculate wave equation with finite differences
        p_we = 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':
                        p_we[self.src_z_index, ix] += self.src_data[k,self.time_index]
                    else:
                        p_we[self.src_z_index+1, ix] += self.src_data[k,self.time_index]
                        p_we[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':
                    p_we[self.src_z_index, self.src_x_index] += self.src_data[self.time_index]
                else:
                    p_we[self.src_z_index+1, self.src_x_index] += self.src_data[self.time_index]
                    p_we[self.src_z_index-1, self.src_x_index] -= self.src_data[self.time_index]

        # calculate absorbing boundary conditions
        # Clayton & Engquist 1977 https://authors.library.caltech.edu/34673/
        p_boundary = boundaries.calculate_boundaries(self.p, self.pold, self.dt, self.dx, self.c,
                                                     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
        pnew = self.sigma*p_we + (1-self.sigma)*p_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.p[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.return_solution()

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