
import numpy as np

import scipy.signal as signal


def sbandpass(s, h, freqlo, freqhi):
    nr = h.nr
    dt = h.dt
    fs = 1/dt

    for ir in range(0, nr):
        s[:,ir] = bandpass(s[:,ir],freqlo,freqhi,fs)
    return s


def sconvolve(s, h, w, inplace=True):
    nt = h.nt
    nr = h.nr

    if inplace:
        for ir in range(nr):
            s[:,ir] = np.convolve(s[:,ir], w, 'same')
        return s
    else:
        s2 = np.zeros((nt,nr))
        for ir in range(nr):
            s2[:,ir] = np.convolve(s[:,ir], w, 'same')
        return s2


def slowpass(s, h, freq):
    raise NotImplementedError


def shighpass(s, h, freq):
    raise NotImplementedError


def smute(s, h, vel, toff, xoff=0, constant_spacing=False):
    nt = h.nt
    dt = h.dt
    nr = h.nr

    # construct tapered window
    length = 400
    win = np.sin(np.linspace(0, 1, 2*length))
    win = win[0:length]

    for ir in range(0, nr):

        # calculate slope
        if vel!=0:
            slope = 1./vel
        else:
            slope = 0

        # calculate offsets
        if constant_spacing:
            itoff = toff/dt
            ixoff = (ir-xoff)/dt
        else:
            itoff = toff/dt
            ixoff = (h.rx[ir]-h.sx[0]-xoff)/dt

        itmin = int(np.ceil(slope*abs(ixoff)+itoff)) - length/2
        itmax = itmin + length

        # apply window
        if 1 < itmin and itmax < nt:
            s[0:itmin,ir] = 0.
            s[itmin:itmax,ir] = win*s[itmin:itmax,ir]
        elif itmin < 1 <= itmax:
            s[1:itmax,ir] = win[length-itmax+1:length]*s[1:itmax,ir]
        elif itmin < nt < itmax:
            s[0:itmin,ir] = 0.
            s[itmin:nt,ir] = win[1:nt-itmin+1]*s[itmin:nt,ir]
        elif itmin > nt:
            s[:,ir] = 0.

    return s


def swindow(s, h, tmin, tmax, alpha=0.05, units='samples'):
    nt = h.nt
    dt = h.dt
    nr = h.nr
    t0 = h.t0

    if units == 'time':
        itmin = int((tmin-t0)/dt)
        itmax = int((tmax-t0)/dt)

    elif units == 'samples':
        itmin = int(tmin)
        itmax = int(tmax)

    else:
        raise ValueError

    win = tukeywin(nt, itmin, itmax, alpha)
    for ir in range(0,nr):
        s[:,ir] = win*s[:,ir]

    return s


# --

def correlate(u, v):
    w = np.convolve(u, np.flipud(v))
    return


def bandpass(w, freqlo, freqhi, fs, npass=2):
    wn = [2*freqlo/fs, 2*freqhi/fs]
    b, a = signal.butter(npass, wn, btype='band')
    w = signal.lfilter(b, a, w)
    return w


def highpass(w, freq, df, npass=2):
    raise Exception('Not yet implemented.')


def lowpass(w, freq, df, npass=2):
    raise Exception('Not yet implemented.')


def window(nt, type='sine', **kwargs):
    if type == 'sine':
        return np.sin(np.linspace(0, 1, 2*length))
    elif type == 'tukey':
        return tukeywin


def tukeywin(nt, imin, imax, alpha=0.05):
    t = np.linspace(0,1,imax-imin)
    w = np.zeros(imax-imin)
    p = alpha/2.
    lo = np.floor(p*(imax-imin-1))+1
    hi = imax-imin-lo
    w[:lo] = (1+np.cos(np.pi/p*(t[:lo]-p)))/2
    w[lo:hi] = np.ones((hi-lo))
    w[hi:] = (1+np.cos(np.pi/p*(t[hi:]-p)))/2
    win = np.zeros(nt)
    win[imin:imax] = w
    return win


@@ -45,16 +45,100 @@
 signal%0A%0A
+from seisflows.tools.config import SeisflowsParameters%0APAR = SeisflowsParameters()%0A%0A
 %0Adef sba
@@ -926,33 +926,37 @@
 (np.linspace(0, 
-1
+np.pi
 , 2*length))%0A   
@@ -1292,37 +1292,211 @@
    i
-xoff = (h.rx%5Bir%5D-h.sx%5B0%5D-xoff
+f PAR.SOLVER == 'specfem2d':%0A                ixoff = (h.rx%5Bir%5D-h.sx%5B0%5D-xoff)/dt  %0A            elif PAR.SOLVER == 'specfem3d':%0A                ixoff = np.sqrt((h.rx%5Bir%5D-h.sx%5B0%5D)**2 + (h.ry%5Bir%5D-h.sy%5B0%5D)**2
 )/dt
@@ -1785,17 +1785,17 @@
       s%5B
-1
+0
 :itmax,i
@@ -1815,18 +1815,16 @@
 th-itmax
-+1
 :length%5D
@@ -1826,17 +1826,17 @@
 ngth%5D*s%5B
-1
+0
 :itmax,i
@@ -1829,32 +1829,33 @@
 h%5D*s%5B0:itmax,ir%5D
+ 
 %0A        elif it
@@ -1936,17 +1936,17 @@
 %5D = win%5B
-1
+0
 :nt-itmi
@@ -1946,18 +1946,16 @@
 nt-itmin
-+1
 %5D*s%5Bitmi
@@ -1958,24 +1958,25 @@
 itmin:nt,ir%5D
+ 
 %0A        eli
@@ -2015,16 +2015,92 @@
 %5D = 0.%0A%0A
+        # inner mute %0A        if ixoff*dt %3C xoff:%0A            s%5B:,ir%5D = 0.%0A%0A
     retu

add muting for 3D inversions