# -*- coding: utf-8 -*-
# 2009-11-27 Moritz

import numpy as np
import matplotlib
matplotlib.rc('figure.subplot', hspace=0.8) # set default
import matplotlib.pyplot as plt
plt.close('all')


def ricker(f, dt):
    """
    Generate Ricker Wavelet with central frequency f

    written by M.D.Sacchi, last modified December 10,  1998
    Copyright (C) 1998 Seismic Processing and Imaging Group
                       Department of Physics
                       The University of Alberta

    @param f: central frequency
    @param dt: sampling interval in [s]
    @return: Ricker wavelet as numpy.ndarray dtype float64
    """
    nw=6./f/dt
    nw=2*np.floor(nw/2)+1
    nc=np.floor(nw/2)
    i=np.arange(1,nw,dtype='float64')
    alpha=(nc-i+1)*f*dt*np.pi
    beta=alpha**2
    return (1.0-beta*2.0)*np.exp(-beta)



#################################################################################
#                          Start of main Program                                #
#################################################################################


# Generate ricker wavelet
wavelet = ricker(2, 0.05)

# Generate greensfunction
N = 200
greenfct = np.zeros(N,dtype='float64')
greenfct[np.array([48,50,54,60,65,76,110,140])] = 1
greenfct[np.array([49,52,57,61,63,70,90,120,170])] = -1

# Calculate length
Nwav = len(wavelet)
Ngreen = len(greenfct)
Ntot = Ngreen + Nwav


#
# Time domain
#
seis_time = np.zeros(Ntot)
for t in xrange(Ntot):
    _sum = 0.0
    for ts in xrange(Ngreen):
        if t-ts < 0 or t-ts >= Nwav:
            continue
        _sum += greenfct[ts]*wavelet[t-ts]
    seis_time[t] = _sum


#
# Frequency domain
#
fgreen = np.fft.rfft(greenfct, n=Ntot)
fwav = np.fft.rfft(wavelet,n=Ntot)
seis_freq = np.fft.irfft(fgreen*fwav)


# Buildtin function
#seis_build = np.convolve(greenfct, wavelet)

#
# Plot all the results
#
plt.clf()
plt.subplot(311)
plt.plot(wavelet)
plt.title("Ricker Wavelet")
plt.subplot(312)
plt.stem(np.arange(N),greenfct,markerfmt='.',basefmt='b-')
plt.title("Greens Function")
plt.ylim(-1.5,1.5)
plt.subplot(313)
plt.plot(seis_time)
plt.title("Convolution of Ricker Wavelet with Greens Function")
plt.plot(seis_freq, 'k-', label='Time Domain',lw=1)
plt.plot(seis_freq, 'r--', label='Frequency Domain',lw=2)
plt.legend()
plt.show()