#!/usr/bin/env python
# coding: utf-8

# In[1]:


get_ipython().run_line_magic('pylab', 'inline')


# In[2]:


from pkprocess import *


# In[3]:


sd=read_su('data/marmousi.su') # Seismic Data


# In[4]:


get_ns(sd),get_ntr(sd),get_dt(sd)


# In[5]:


print_range(sd)


# # Header information

# In[6]:


plot(ntr_per_shot(sd),'o')
xlabel('Shot gather numbers')
ylabel('Number of traces/shot gather')
grid()


# In[7]:


plot(get_key(sd,'sx'))
xlabel('Trace number')
ylabel('Sources x-axis locations (m)')


# In[10]:


print(get_key(sd,'sx'))
print(get_key(sd,'gx'))


# In[11]:


plot(get_key(sd,'sdepth'))
xlabel('Trace number')
ylabel('Sources depth (m)')


# In[12]:


plot(get_key(sd,'gelev'))
xlabel('Trace number')
ylabel('Receivers depth (m)')


# In[14]:


stacking_chart(sd)
savefig('stacking_chart.png')


# # Displaying seismic data

# In[15]:


get_key_unique(sd,'fldr')


# In[16]:


shot5000=window(sd,'fldr',5000)
plot_wiggle(shot5000,perc=99)


# In[17]:


plot_image(shot5000,perc=99)


# In[18]:


plot_image(shot5000,perc=99,cmap='bwr')


# In[19]:


plot_image(sd,figsize=[15,10],perc=99)


# In[20]:


get_key_unique(sd,'offset')


# ## Nearest-offset traces

# In[21]:


near=window(sd,'offset',-200)


# In[22]:


plot_wiggle(near,figsize=[10,10],perc=99)


# ## Gain control

# In[23]:


shot5000_tpow=gain(shot5000,tpow=1)
shot5000_epow=gain(shot5000,epow=1)
plot_comp((shot5000_tpow,shot5000_epow),plot='wiggle',perc=99)


# In[24]:


sd_gained=gain(sd,tpow=1)
sd_gained.print_log()


# In[25]:


seis_env_dB(shot5000,shot5000_tpow)


# ## FX, FK Spectrum

# In[26]:


specfx(shot5000_tpow)


# In[27]:


specfk(shot5000_tpow)


# ## Autocorrelograms

# In[28]:


sd3=window(sd_gained,'fldr',[5000, 5025, 5050])
max_lag=0.2
auto=auto_correlation_map(sd3,max_lag)
plot_wiggle(auto,figsize=[10,10],perc=99)
ylabel('Time lag (s)')


# In[29]:


plot_wiggle(sd3,figsize=[15,10],perc=99)


# ## Spiking deconvolution

# In[30]:


mu=0.1
sd_decon=spiking_decon(sd_gained,max_lag,mu)


# In[32]:


#%timeit -r 1 -n 1 spiking_decon(sd_gained,max_lag,mu)


# In[33]:


sd3=window(sd_decon,'fldr',[5000, 5025, 5050])
plot_wiggle(sd3,figsize=[15,10],perc=99)


# ## Automatic Gain Control

# In[34]:


sd_agc=gain(sd_decon,agc=True,agc_gate=0.5,norm='rms')


# In[35]:


plot_wiggle(window(sd_agc,'fldr',[5000,5025,5050]),figsize=[15,10],perc=99)


# In[36]:


sd_agc.print_log()


# ## CDP Sorting

# In[37]:


sd_sort=trace_sort(sd_agc,('+cdp','-offset'))


# In[38]:


cmps=get_key(sd_sort,'cdp')
cmpu=get_key_unique(sd_sort,'cdp')
fold_num = [sum(icmp==cmps) for icmp in cmpu]

plot(cmpu,fold_num)
xlabel('CMP numbers')
ylabel('Fold')


# In[39]:


sd_sort.print_log()


# ## Output for the velocity analysis

# In[40]:


sd_sort.write("marm_gain_decon_agc_sort.sd")


# ## Velocity analysis
# 
# 1. Open and edit `Par_velan.py` file.
# 2. Run `Run_velan.py`. (`$ python Run_velan.py`)
#     - Left click: Add a point
#     - Right click: Remove a point

# In[41]:


from IPython.display import YouTubeVideo
display(YouTubeVideo("zWt5lbWEx3g",width=775,height=542))


# ## After the velocity analysis

# In[42]:


sd_nmo=read("marm_gain_decon_agc_sort_nmo.sd")


# In[43]:


sd_nmo.print_log(nmo=True)


# In[44]:


plot_image(sd_nmo,figsize=[15,10],perc=99)


# ## Stacking

# In[45]:


sd_stk=stack(sd_nmo)


# In[46]:


plot_image(sd_stk,figsize=[15,10],perc=99,key='cdp')


# In[47]:


print_range(sd_stk)


# In[48]:


get_key(sd_stk,'shortpad') # shortpad keyword contains the fold number


# In[49]:


plot_image(window(sd_stk,'shortpad',48),figsize=[15,10],perc=99,key='cdp')


# ## Stolt migration

# In[50]:


v=3000.
dx=12.5
sd_mig=stolt_mig(sd_stk,v,dx)


# In[51]:


plot_image(sd_mig,figsize=[15,10],perc=99,key='cdp')


# In[52]:


plot_image(window(sd_mig,'shortpad',48),figsize=[15,10],perc=99,key='cdp')


# In[53]:


plot_comp((sd_stk,sd_mig),figsize=[12,10],perc=99,key='cdp')


# In[54]:


sd_mig.print_log()


# In[51]:


#sd_mig.write('mig.sd')
#sd_mig.write_su('mig.su')