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

# # AVO projections
# 
# This notebook shows all the steps to implement _AVO projections_ in a routine quantitative interpretation workflow. Please refer to "Seismic Amplitude: An Interpreter's Handbook" (Simm & Bacon, Cambridge University Press, 2014) for a detailed discussion on AVO projections and the theory behind.
# 
# The data I use in this notebook comes (as usual) from the [Quantitative Seismic Interpretation book website](http://pangea.stanford.edu/researchgroups/srb/resources/books/quantitative-seismic-interpretation).
# 
# _Please note that this is work in progress; my intention is to better explain some of the key theoretical and practical aspects as I go along. However, everything seems to be working ok so I thought that other people may benefit from this._

# ## import libraries

# In[1]:


import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from bruges.reflection import akirichards, shuey
import segyio
import xarray as xr

# from scipy.optimize import curve_fit
from scipy.interpolate import griddata

get_ipython().run_line_magic('matplotlib', 'inline')
# comment out the following if you're not on a Mac with HiDPI display
get_ipython().run_line_magic('config', "InlineBackend.figure_format = 'retina'")

from ipywidgets import interact, interactive, fixed


# ## coordinate rotation, an example

# In[2]:


mean = [0, 0]
cov = [[8, -5], #  var_x, cov_xy
       [-5, 4]] #  cov_xy, var_y
x, y = np.random.multivariate_normal(mean, cov, 500).T

fit = np.polyfit(x,y,1)
XX=np.linspace(x.min(),x.max())
YY = np.polyval(fit,XX)

# add anomaly point
x=np.append(x,-6)
y=np.append(y,-9)

# compute rotation angle
alfa=np.deg2rad(-90)+fit[0]
print('angle of rotation: {:.1f}'.format(np.rad2deg(alfa)))

plt.figure()
plt.plot(x, y, '.k')
plt.plot(XX,YY, '-r')
plt.axis('square')
plt.grid()
plt.xlim(-15,15)
plt.ylim(-15,15)


# I will now rotate the coordinates by the angle given by the best fit line:

# In[3]:


x_new = x*np.cos(alfa)+y*np.sin(alfa)
y_new = -x*np.sin(alfa)+y*np.cos(alfa)

plt.figure()
plt.plot(x_new, y_new, '.k')
plt.axis('square')
plt.grid()
plt.xlim(-15,15)
plt.ylim(-15,15)


# We can make it interactive now:

# In[4]:


@interact(alfa=(-180,180,5))
def coordrot(alfa):
    alfa=np.deg2rad(alfa)
    x_new = x*np.cos(alfa)+y*np.sin(alfa)
    y_new = -x*np.sin(alfa)+y*np.cos(alfa)
    plt.figure()
    plt.plot(x_new, y_new, '.k')
    plt.axis('square')
    plt.grid()
    plt.xlim(-15,15)
    plt.ylim(-15,15)


# ## define a few functions
# 
# I will define here 4 functions to make the rest of the workflow smoother:
# 
# 1. `vpvs`: to analyze Vp-Vs relationships and return an empirical relation to be applied on well with no Vs log available;
# 2. `wellplot`: to display well log suite nicely formatted etc;
# 3. `avoplot`: to calculate and display single interface AVO and Intercept-Gradient crossplots;
# 4. `avoproj`: to calculate and display "extended" AVO responses and identify optimal angles for lithology and fluid projections.

# In[5]:


def vpvs(ww,vsh_cutoff,z1,z2):
    ss=(ww.VSH<=vsh_cutoff) & (ww.DEPTH>=z1) & (ww.DEPTH<=z2)
    sh=(ww.VSH>vsh_cutoff)  & (ww.DEPTH>=z1) & (ww.DEPTH<=z2)

    flagfrm=True if 'VP_FRMB' in ww.columns else False
    if flagfrm:
        vp_ss=ww.VP_FRMB[ss]
        vs_ss=ww.VS_FRMB[ss]
        vp_sh=ww.VP_FRMB[sh]
        vs_sh=ww.VS_FRMB[sh]
    else:
        vp_ss=ww.VP[ss]
        vs_ss=ww.VS[ss]
        vp_sh=ww.VP[sh]
        vs_sh=ww.VS[sh]

    x=np.linspace(1500,6000,100)
    par_ss = np.polyfit(vp_ss, vs_ss, 1)
    fit_y_ss = par_ss[0]*x+par_ss[1]
    par_sh = np.polyfit(vp_sh, vs_sh, 1)
    fit_y_sh = par_sh[0]*x+par_sh[1]

    lst = (-0.0551*(x/1e3)**2 + 1.016*(x/1e3) - 1.0305) * 1e3
    dol = 0.583*x - 78
    sst = 0.804*x - 856
    sh  = 0.769*x - 867

    plt.figure()
    plt.plot(vp_sh,vs_sh,'.g', alpha=0.5, label='sh', ms=5)
    plt.plot(vp_ss,vs_ss,'.k', alpha=0.5, label='sst', ms=5)
    plt.plot(x,fit_y_ss,'-k', lw=2, label='Vs (sst)=%.3f%+.3f*Vp' % (par_ss[1], par_ss[0]))
    plt.plot(x,fit_y_sh,'-g', lw=2, label='Vs (sh)=%.3f%+.3f*Vp' % (par_sh[1], par_sh[0]))
    plt.plot(x,lst,':r',label='G-C lst')
    plt.plot(x,sst,'-r',label='G-C sst')
    plt.plot(x,sh,'--r',label='G-C sh')
    plt.grid(), plt.xlabel('Vp'), plt.ylabel('Vs')
    plt.xlim(1000,6000), plt.ylim(500,3000)
    plt.legend(fontsize='small',loc='lower right')

    return par_ss, par_sh


# In[6]:


def wellplot(WW,ztop=None,zbot=None,name=None,tops=None):
    '''
    wellplot (c) aadm 2015
    Simple logview plot.

    INPUT
    WW: Pandas dataframe with VP, [VS], RHO, IP, [VPVS or PR], [SWE], PHIE, VSH
    ztop,zbot: depth range to plot (defaults to min,max depth)
    name: well name (or anything else) to print
    tops: dictionary containing stratigraphic tops, defined as
        a dictionary (dict):
        tops={'Trias': 1200,'Permian': 2310}
        or a Pandas series

    HISTORY
    2015-10-01 added option to read tops (either as dict or Pandas series)
    2015-09-29 also accepts logset with no Vs
    2015-08-04 updated with depth range input
    2015-06-30 first version
    '''
    flagtops=False
    flagvs=True if 'VS' in WW.columns else False
    if tops is None:
        flagtops=False
    else:
        flagtops=True
        if not isinstance(tops, pd.Series): # check if tops is a Pandas series
            tops = pd.Series(tops)
#         tmp_tops=tops.to_dict()
        
#     elif tops!=None: # then checks if tops exists (as a dict)
#         tmp_tops=tops
#         flagtops=True
    if ztop==None:
        ztop=WW.DEPTH.min()
    if zbot==None:
        zbot=WW.DEPTH.max()
    ff = (WW.DEPTH >= ztop) & (WW.DEPTH <= zbot)
    if flagvs:
        velmin=WW.VS[ff].min()
    else:
        velmin=WW.VP[ff].min()
#     velmax=WW[(WW.DEPTH>=ztop) & (WW.DEPTH<=zbot)].ix[:,['VP']].max().values
    velmax = WW.VP[ff].max()
    dmin = WW.RHO[ff].min()
    dmax = WW.RHO[ff].max()
    ipmin = WW.IP[ff].min()
    ipmax = WW.IP[ff].max()
    rmin = WW.VPVS[ff].min()
    rmax = WW.VPVS[ff].max()

#     dmin=WW[(WW.DEPTH>=ztop) & (WW.DEPTH<=zbot)].ix[:,['RHO']].min().values
#     dmax=WW[(WW.DEPTH>=ztop) & (WW.DEPTH<=zbot)].ix[:,['RHO']].max().values
#     ipmin=WW[(WW.DEPTH>=ztop) & (WW.DEPTH<=zbot)].ix[:,['IP']].min().values
#     ipmax=WW[(WW.DEPTH>=ztop) & (WW.DEPTH<=zbot)].ix[:,['IP']].max().values
#     rmin=WW[(WW.DEPTH>=ztop) & (WW.DEPTH<=zbot)].ix[:,['VPVS']].min().values
#     rmax=WW[(WW.DEPTH>=ztop) & (WW.DEPTH<=zbot)].ix[:,['VPVS']].max().values

    f, ax = plt.subplots(nrows=1,ncols=5,sharey=True,figsize=(8,6),facecolor='w')
    if 'SWE' in WW.columns:
        ax[0].plot(WW.SWE, WW.DEPTH, color='blue', alpha=0.8, lw=1, label='Swe')
    ax[0].plot(WW.VSH, WW.DEPTH, color='chocolate', alpha=0.8, lw=1, label='Vsh')
    ax[0].plot(WW.PHIE, WW.DEPTH, color='black', alpha=0.8, lw=1, label='phie')
    ax[1].plot(WW.VP, WW.DEPTH, 'k', lw=1, label='Vp')
    if flagvs:
        ax[1].plot(WW.VS, WW.DEPTH, color='lightsteelblue', lw=1, label='Vs')
    ax[2].plot(WW.RHO, WW.DEPTH, 'k', lw=1)
    ax[3].plot(WW.IP, WW.DEPTH, 'k', lw=1)
    if flagvs:
        if 'PR' in WW.columns:
            ax[4].plot(WW.PR, WW.DEPTH, 'k', lw=1)
            ax[4].set_xlabel('PR'), ax[4].set_xlim(0.0,0.5)
        else:
            ax[4].plot(WW.VPVS, WW.DEPTH, 'k', lw=1)
            ax[4].set_xlabel('Vp/Vs'), ax[4].set_xlim(rmin-rmin*.01,rmax+rmax*.01)
    else:
        ax[4].text(2.5, (zbot-ztop)/2+ztop, 'no Vs', fontsize=14,horizontalalignment='center')
        ax[4].set_xlim(0,5)
        ax[4].set_xticklabels([])
    ax[0].set_xlabel('Vcl/Sw/phi'),     ax[0].set_xlim(-0.1,1.1),                          ax[0].set_ylim(ztop,zbot)
    ax[1].set_xlabel('Velocity (m/s)'), ax[1].set_xlim(velmin-velmin*.1,velmax+velmax*.1), ax[1].set_ylim(ztop,zbot)
    ax[2].set_xlabel('Density (g/cc)'), ax[2].set_xlim(dmin-dmin*.01,dmax+dmax*.01),       ax[2].set_ylim(ztop,zbot)
    ax[3].set_xlabel('Ip (m/s*g/cc)'),  ax[3].set_xlim(ipmin-ipmin*.1,ipmax+ipmax*.1),     ax[3].set_ylim(ztop,zbot)
    ax[0].set_ylim(zbot,ztop)
    ax[0].legend(fontsize=8, loc='lower right')
    ax[1].legend(fontsize=8, loc='lower right')
    for aa in ax.flatten():
        aa.grid()
        aa.xaxis.tick_top()
        plt.setp(aa.xaxis.get_majorticklabels(), rotation=90, fontsize=8)
        if flagtops:
            for topn, topz in tops.iteritems():
                if (~np.isnan(topz)) & (topz>=ztop) & (topz<=zbot):
                    ax[0].text(x=min(ax[0].xaxis.get_data_interval())*1.0,
                               y=float(topz),s=topn,alpha=0.75,color='k',fontsize=10,ha='left',va='center',
                               bbox=dict(facecolor='white', alpha=1.0, lw = 0.25),
                               weight='light')
                    aa.axhline( y = float(topz), color='gold', lw=2,alpha=0.5,xmin=0.05,xmax=0.95)
    if name!=None:
        plt.suptitle(name, fontsize=14)


# In[7]:


def avoplot(WW, top, above, below, res='down', name=None):
    '''
    avoplot (c) aadm 2015
    AVO and Intercept-Gradient crossplots.

    INPUT
    WW: Pandas dataframe with GR, VP, VS, RHO
    top: interface for AVO calculations
    above,below: meters above/below interface to calculate average elastic properties
    res: 'up' or 'down', specifies where reservoir (and associated fluid changes) is located
    name: well name (or anything else) to print

    HISTORY
    2015-11-10 added possibility to specify whether reservoir is the upper or lower layer
    2015-10-02 first version
    '''
    flagfrm=True if 'VP_FRMB' in WW.columns else False
    mx = 0.4
    ang=np.linspace(0,40,41)
    z=WW.DEPTH
    up=(WW.DEPTH>=top-above) & (WW.DEPTH<=top)
    lo=(WW.DEPTH>=top)       & (WW.DEPTH<=top+below)
    plotmax=WW[(WW.DEPTH>=top-above*2) & (WW.DEPTH<=top+below*2)][['GR','IP']].max().values
    plotmin=WW[(WW.DEPTH>=top-above*2) & (WW.DEPTH<=top+below*2)][['GR','IP']].min().values
    el0=np.array([
              [np.nanmean(WW.VP[up]), np.nanmean(WW.VS[up]), np.nanmean(WW.RHO[up])],
              [np.nanmean(WW.VP[lo]), np.nanmean(WW.VS[lo]), np.nanmean(WW.RHO[lo])] ])
    print('... [upper layer] Vp:%8.2f, Vs:%8.2f, rho:%8.4f, Ip: %8.2f, Vp/Vs:%5.2f' % (el0[0,0],el0[0,1],el0[0,2],el0[0,0]*el0[0,2],el0[0,0]/el0[0,1]))
    print('... [lower layer] Vp:%8.2f, Vs:%8.2f, rho:%8.4f, Ip: %8.2f, Vp/Vs:%5.2f' % (el0[1,0],el0[1,1],el0[1,2],el0[1,0]*el0[1,2],el0[1,0]/el0[1,1]))
    rc     = akirichards(el0[0,0],el0[0,1],el0[0,2],el0[1,0],el0[1,1],el0[1,2],ang)
    tmp    = shuey(el0[0,0],el0[0,1],el0[0,2],el0[1,0],el0[1,1],el0[1,2],30,terms=True)
    I      = tmp[0]
    G      = tmp[1] / np.sin(np.radians(30))**2

    if flagfrm:
        if res=='up':
            el1=np.array([
                      [np.nanmean(WW.VP_FRMB[up]), np.nanmean(WW.VS_FRMB[up]), np.nanmean(WW.RHO_FRMB[up])],
                      [np.nanmean(WW.VP_FRMO[up]), np.nanmean(WW.VS_FRMO[up]), np.nanmean(WW.RHO_FRMO[up])],
                      [np.nanmean(WW.VP_FRMG[up]), np.nanmean(WW.VS_FRMG[up]), np.nanmean(WW.RHO_FRMG[up])] ])

            rc_bri = akirichards(el1[0,0],el1[0,1],el1[0,2],el0[1,0],el0[1,1],el0[1,2],ang)
            tmp    =      shuey2(el1[0,0],el1[0,1],el1[0,2],el0[1,0],el0[1,1],el0[1,2],30,terms=True)
            I_bri  = tmp[0]
            G_bri  = tmp[1] / np.sin(np.deg2rad(30))**2

            rc_oil = akirichards(el1[1,0],el1[1,1],el1[1,2],el0[1,0],el0[1,1],el0[1,2],ang)
            tmp    =      shuey2(el1[1,0],el1[1,1],el1[1,2],el0[1,0],el0[1,1],el0[1,2],30,terms=True)
            I_oil  = tmp[0]
            G_oil  = tmp[1] / np.sin(np.deg2rad(30))**2

            rc_gas = akirichards(el1[2,0],el1[2,1],el1[2,2],el0[1,0],el0[1,1],el0[1,2],ang)
            tmp    =      shuey2(el1[2,0],el1[2,1],el1[2,2],el0[1,0],el0[1,1],el0[1,2],30,terms=True)
            I_gas  = tmp[0]
            G_gas  = tmp[1] / np.sin(np.deg2rad(30))**2
            bb='upper layer'
        else:
            el1=np.array([
                        [np.nanmean(WW.VP_FRMB[lo]), np.nanmean(WW.VS_FRMB[lo]), np.nanmean(WW.RHO_FRMB[lo])],
                        [np.nanmean(WW.VP_FRMO[lo]), np.nanmean(WW.VS_FRMO[lo]), np.nanmean(WW.RHO_FRMO[lo])],
                        [np.nanmean(WW.VP_FRMG[lo]), np.nanmean(WW.VS_FRMG[lo]), np.nanmean(WW.RHO_FRMG[lo])] ])
            rc_bri = akirichards(el0[0,0],el0[0,1],el0[0,2],el1[0,0],el1[0,1],el1[0,2],ang)
            tmp    =      shuey(el0[0,0],el0[0,1],el0[0,2],el1[0,0],el1[0,1],el1[0,2],30,terms=True)
            I_bri  = tmp[0]
            G_bri  = tmp[1] / np.sin(np.deg2rad(30))**2

            rc_oil = akirichards(el0[0,0],el0[0,1],el0[0,2],el1[1,0],el1[1,1],el1[1,2],ang)
            tmp    =      shuey(el0[0,0],el0[0,1],el0[0,2],el1[1,0],el1[1,1],el1[1,2],30,terms=True)
            I_oil  = tmp[0]
            G_oil  = tmp[1] / np.sin(np.deg2rad(30))**2

            rc_gas = akirichards(el0[0,0],el0[0,1],el0[0,2],el1[2,0],el1[2,1],el1[2,2],ang)
            tmp    =      shuey(el0[0,0],el0[0,1],el0[0,2],el1[2,0],el1[2,1],el1[2,2],30,terms=True)
            I_gas  = tmp[0]
            G_gas  = tmp[1] / np.sin(np.deg2rad(30))**2
            bb='lower layer'
        print('... [%s] Vp:%8.2f, Vs:%8.2f, rho:%8.2f, Ip: %8.2f, Vp/Vs:%5.2f --> BRINE' % (bb,el1[0,0],el1[0,1],el1[0,2],el1[0,0]*el1[0,2],el1[0,0]/el1[0,1]))
        print('... [%s] Vp:%8.2f, Vs:%8.2f, rho:%8.2f, Ip: %8.2f, Vp/Vs:%5.2f --> OIL'   % (bb,el1[1,0],el1[1,1],el1[1,2],el1[1,0]*el1[1,2],el1[1,0]/el1[1,1]))
        print('... [%s] Vp:%8.2f, Vs:%8.2f, rho:%8.2f, Ip: %8.2f, Vp/Vs:%5.2f --> GAS'   % (bb,el1[2,0],el1[2,1],el1[2,2],el1[2,0]*el1[2,2],el1[2,0]/el1[2,1]))

    f=plt.subplots(figsize=(10, 5),facecolor='w')
    ax0 = plt.subplot2grid((1,6), (0,0), colspan=1)
    ax1 = plt.subplot2grid((1,6), (0,1), colspan=1)
    ax2 = plt.subplot2grid((1,6), (0,2), colspan=2)
    ax3 = plt.subplot2grid((1,6), (0,4), colspan=2)
    ax0.plot(WW.GR,WW.DEPTH,'-k')
    ax0.set_xlabel('GR'), ax0.set_xlim(0,plotmax[0]+plotmax[0]*.2)
    ax0.set_ylim(top-(above*2),top+(below*2))
    ax0.plot((6, 6), (top-above, top+below), 'c', lw=6, alpha=0.5)
    if flagfrm:
        ax1.plot(WW.VP_FRMB*WW.RHO_FRMB, WW.DEPTH, 'b')
        ax1.plot(WW.VP_FRMO*WW.RHO_FRMO, WW.DEPTH, 'g')
        ax1.plot(WW.VP_FRMG*WW.RHO_FRMG, WW.DEPTH, 'r')
    ax1.plot(WW.IP,WW.DEPTH,'-k')
    ax1.set_yticklabels([])
    ax1.set_xlabel('Ip (m/s*g/cc)'), ax1.set_xlim(plotmin[1]-plotmin[1]*.2,plotmax[1]+plotmax[1]*.2)
    ax1.set_ylim(top-(above*2),top+(below*2))
    for aa in [ax0, ax1]:
        aa.invert_yaxis()
        aa.grid()
        aa.xaxis.tick_top()
        plt.setp(aa.xaxis.get_majorticklabels(), rotation=90, fontsize=8)
        aa.axhline(y=top,color='gold',lw=2,alpha=0.5)
    ax2.plot(ang, rc,'-k', lw=4, label=str(top)+' m')
    # ax2.plot(ang, rc[1], '--k',  lw=2, label='Top +anis')
    if flagfrm:
        ax2.plot(ang, rc_bri, '-b', lw=2, label='brine')
        ax2.plot(ang, rc_oil, '-g', lw=2, label='oil')
        ax2.plot(ang, rc_gas, '-r', lw=2, label='gas')
    ax2.axhline(0,color='k',lw=2)
    ax2.set_xlabel('angle of incidence')
    ax2.legend(loc='lower right', fontsize='small')
    if flagfrm:
        ax3.plot(I_bri, G_bri, 'bo', ms=10, mec='b')
        ax3.plot(I_oil, G_oil, 'go', ms=10, mec='g')
        ax3.plot(I_gas, G_gas, 'ro', ms=10, mec='r')
    ax3.plot(I,G,'ko',ms=10,mfc='none',mew=2)
    ax3.axhline(0, color='k', lw=2), ax3.axvline(0, color='k', lw=2)
    ax3.set_xlim(-mx, mx)
    ax3.xaxis.set_label_position('top'), ax3.xaxis.tick_top()
    ax3.yaxis.set_label_position('right'), ax3.yaxis.tick_right()
    ax3.set_xlabel('intercept'), ax3.set_ylabel('gradient')
    for bb in [ax2, ax3]:
        bb.grid()
        bb.tick_params(axis='both', which='major', labelsize=8)
        bb.set_ylim(-mx, mx)
    if name!=None:
        plt.suptitle(name, fontsize=16, y=0.05, x=0.95, ha='right')


# In[8]:


def avoproj(WW, top, above, below, name=None, fluid='gas'):
    '''
    avoproj (c) aadm 2015
    AVO projections.

    INPUT
    WW: Pandas dataframe with GR, VP, VS, RHO
    top: interface for AVO calculations
    above,below: meters above/below interface to calculate average elastic properties
    name: well name (or anything else) to print
    fluid: select either gas or oil to determine best angle to maximise fluid separation

    HISTORY
    2015-11-06 removed I-G plot
    2015-10-21 first version
    '''
    from bruges.reflection import shuey
    ang=np.linspace(0,30)
    angl=np.sin(np.deg2rad(ang))**2
    chi=np.deg2rad(np.arange(-180,185,5))
    if fluid == 'oil':
        print('==> AVO projections to maximise brine-oil separation')
    elif fluid == 'gas':
        print('==> AVO projections to maximise brine-gas separation')
    else:
        print('==> please choose either gas or oil!')
        return
    z=WW.DEPTH
    up=(WW.DEPTH>=top-above) & (WW.DEPTH<=top)
    lo=(WW.DEPTH>=top)       & (WW.DEPTH<=top+below)
    el=np.array([ [np.nanmean(WW.VP[up]), np.nanmean(WW.VS[up]), np.nanmean(WW.RHO[up])],
              [np.nanmean(WW.VP[lo]), np.nanmean(WW.VS[lo]), np.nanmean(WW.RHO[lo])],
              [np.nanmean(WW.VP_FRMB[lo]), np.nanmean(WW.VS_FRMB[lo]), np.nanmean(WW.RHO_FRMB[lo])],
              [np.nanmean(WW.VP_FRMO[lo]), np.nanmean(WW.VS_FRMO[lo]), np.nanmean(WW.RHO_FRMO[lo])],
              [np.nanmean(WW.VP_FRMG[lo]), np.nanmean(WW.VS_FRMG[lo]), np.nanmean(WW.RHO_FRMG[lo])] ])
    rc_bri = shuey(el[0,0],el[0,1],el[0,2],el[2,0],el[2,1],el[2,2],ang,return_gradient=True)
    rc_oil = shuey(el[0,0],el[0,1],el[0,2],el[3,0],el[3,1],el[3,2],ang,return_gradient=True)
    rc_gas = shuey(el[0,0],el[0,1],el[0,2],el[4,0],el[4,1],el[4,2],ang,return_gradient=True)
    I_bri, G_bri = rc_bri[0], rc_bri[1]
    I_oil, G_oil = rc_oil[0], rc_oil[1]
    I_gas, G_gas = rc_gas[0], rc_gas[1]
    rc_bri_m  = I_bri*np.cos(chi)+G_bri*np.sin(chi)
    rc_oil_m  = I_oil*np.cos(chi)+G_oil*np.sin(chi)
    rc_gas_m  = I_gas*np.cos(chi)+G_gas*np.sin(chi)

    # calculates angle where distance brine-oil/gas is maximised/minimised
    if fluid == 'oil':
        idl=np.argmin(np.abs((rc_bri_m)-(rc_oil_m)))
        idf=np.argmax(np.abs((rc_bri_m)-(rc_oil_m)))
    else:
        idl=np.argmin(np.abs((rc_bri_m)-(rc_gas_m)))
        idf=np.argmax(np.abs((rc_bri_m)-(rc_gas_m)))
    # do the plot
    plt.figure(figsize=(6,6),facecolor='w')
    plt.axhline(0,color='k',lw=2), plt.axvline(0,color='k',lw=2)
    plt.plot(np.rad2deg(chi), rc_bri_m,'-b',lw=2)
    plt.plot(np.rad2deg(chi), rc_oil_m,'-g',lw=2)
    plt.plot(np.rad2deg(chi), rc_gas_m,'-r',lw=2)
    plt.axvline(np.rad2deg(chi[idl]), color='m',label=r'$\mathrm{litho}\; \chi=%+.0f^\circ$' % np.rad2deg(chi[idl]))
    plt.axvline(np.rad2deg(chi[idf]), color='c',label=r'$\mathrm{fluid}\; \chi=%+.0f^\circ$' % np.rad2deg(chi[idf]))
    plt.xlabel(r'$\chi$')
    plt.legend()
    plt.grid()
    plt.tick_params(axis='both', which='major', labelsize=8)
    plt.xticks(np.arange(-180,190,40))
    plt.xlim(-180, 180), plt.ylim(-.5, .5)
    if name!=None:
        plt.suptitle(name, fontsize=16, y=0.05, x=0.95, ha='right')


# ## step 1: load well data and do AVO at selected level

# In[9]:


w2=pd.read_csv('qsiwell2_frm.csv')


# In[10]:


par_ss,par_sh=vpvs(w2,0.1,2100,2300)

print("Vp-Vs relations:")
print("sand: Vs = {:0.3f}*Vp {:+0.3f}".format(par_ss[0],par_ss[1]))
print("shale: Vs = {:0.3f}*Vp {:+0.3f}".format(par_sh[0],par_sh[1]))


# In the following code I load data from well 3 (where no shear wave sonic log has been acquired), calculate shale volume and porosity(*) logs, then apply the empirical Vp-Vs relation from the previous step to calculate a Vs log and its derivates: Vp/Vs and shear impedance Is.
# 
# ***
# (*) 
# <i>The porosity log is not really used in this workflow except for display purposes (it is  required from my `wellplot` function). A larger discussion on well log data analysis in Python can be found [in this other notebook](http://nbviewer.ipython.org/github/aadm/geophysical_notes/blob/master/seismic_petrophysics.ipynb).
# </i>

# In[11]:


rho_qz=2.65;  k_qz=37;  mu_qz=44
rho_sh=2.81;  k_sh=15;  mu_sh=5
rho_w=1.09;   k_w=2.8
rho_o=0.78;   k_o=0.94 # # oil gravity: 32 API, GOR:    64

# quick petrophysical log calculation for Well 3
# and Vs calculation though empirical relation
w3raw=pd.read_table('well_3.txt', sep='\s+', header=None, skiprows=1, names=['DEPTH','VP','RHO','GR'])
w3raw.VP=w3raw.VP*1000

w3raw['VSH']=(w3raw.GR-w3raw.GR.min())/(w3raw.GR.max()-w3raw.GR.min())
w3raw['RHOm']=w3raw.VSH*rho_sh + (1-w3raw.VSH)*rho_qz
w3raw['RHOf']=np.ones(w3raw.DEPTH.size)*rho_w
w3raw['PHIE']= (w3raw.RHOm-w3raw.RHO) / ( w3raw.RHOm- w3raw.RHOf)
# compute Vs ...
vs_sst=par_ss[0]*w3raw.VP+par_ss[1]
vs=par_sh[0]*w3raw.VP+par_sh[1]
vs[w3raw.VSH<0.3]=vs_sst[w3raw.VSH<=0.3]

w3raw['VS']=vs
w3raw['VPVS']=w3raw.VP/w3raw.VS
w3raw['IP']=w3raw.VP*w3raw.RHO
w3raw['IS']=w3raw.VS*w3raw.RHO
w3=w3raw.copy()


# Now I will take a look at the well data (using the `wellplot` function previously defined):

# In[12]:


tops_w2={'Heimdal': 2153,'OWC': 2183}
wellplot(w2,ztop=2100,zbot=2300,name='Well 2',tops=tops_w2)


# In[13]:


tops_w3={'Heimdal': 2180}
wellplot(w3,ztop=2100,zbot=2300,name='Well 3',tops=tops_w3)


# The AVO signature at Top Heimdal on both wells can be calculated with `avoplot`, which is called with these parameters:
# 
#     avoplot(well_DataFrame, top, above, below, [name])
#     
# The function will calculate the amplitude variation with angles of incidence at the interface found at `top` meters of depth, using as elastic properties for the half-space above this interface the average Vp, Vs and density calculated between `top` and `top-above` meters, and for the half-space below the average over `top` and `top`+`below`.
# 
# The functions outputs a plot composed of gamma-ray (`GR`) and acoustic impedance (`Ip`) logs, and the AVO signatures in the angle of incidence versus reflectivity and in the Intercept versus Gradient domains. If fluid replaced logs are available, the function will also draw brine, oil and gas Ip logs as well as their AVO signatures.
# 
# Averages of the elastic properties for the two half-spaces are also output to the console.

# In[14]:


print('Well 2:')
avoplot(w2,2156,25,15,name='Well 2')

print('Well 3:')
avoplot(w3,2180,30,30,name='Well 3')


# The result is that well 2 and 3 have different AVO responses at Top Heimdal. The best thing to do would be fluid replacement on well 3 and understand if the AVO response we have here is affected by the fluid (I don't even remember now if well 3 is supposed to be dry or oil-bearing; remember this was just an exercise for me so relax!). If for example the Heimdal sand is brine-bearing then, if we do fluid substitution to oil we may have a drastic reduction in Ip and Vp/Vs ratio and a similar AVO response to well 2 in similar fluid conditions (i.e., oil). 
# 
# Anyway, let's move on and calculate AVO projection plots, to determine the projection angles to optimise fluid separation. I will use the `avoproj` function which is called with these parameters:
# 
#     avoproj(well_DataFrame, top, above, below, [name], [fluid='gas'])
# 
# By selecting `fluid=oil` we instruct the function to look for the angle where the distance between brine and oil is maximised (the function defaults to gas). This will be the _fluid projection angle_. The _lithology projection angle_ is the angle where the distance between the fluid curves is minimized. 

# In[15]:


avoproj(w2,2156,25,15,name='Well 2',fluid='oil')


# ## step 2: load seismic and interpretation

# In[17]:


file_3d = '3d_nearstack.sgy'

with segyio.open(file_3d, 'r', iline=41, xline=21) as segyfile:
    ss = segyio.tools.cube(segyfile)
    ntraces = segyfile.tracecount
    sr = segyio.tools.dt(segyfile)/1e3
    nsamples = segyfile.samples.size
    twt = segyfile.samples + 1500
    size_mb= ss.nbytes/1024**2
    inl = segyfile.ilines/1000
    crl = segyfile.xlines


# The following lines will load the Top Heimdal interpreted horizon, grid and display it:

# In[18]:


# load data
hrz=np.recfromtxt('Top_Heimdal_subset.txt', names=['il','xl','z'])

# define well locations
well2_il=1376
well2_xl=1776
well3_il=1468
well3_xl=1847

# gridding
xi = np.linspace(inl.min(), inl.max(),250)
yi = np.linspace(crl.min(), crl.max(),250)
X, Y = np.meshgrid(xi, yi)
Z = griddata((hrz['il'], hrz['xl']), hrz['z'], (X,Y), method='cubic')

# smoothing for better-looking contours
import scipy.ndimage as ndimage
Zf = ndimage.gaussian_filter(Z, sigma=2, order=0)

# make the plot
plt.figure(figsize=(8,8))
im0=plt.pcolormesh(X, Y, Z, cmap='gist_earth_r')
CS=plt.contour(X, Y, Zf, 8, colors='w')
plt.clabel(CS, inline=1, fontsize=10, fmt='%.0f')
plt.xlabel('IL'), plt.ylabel('XL')
plt.plot(well2_il, well2_xl, 'ok', ms=10, mew=2, mec='w')
plt.plot(well3_il, well3_xl, 'ok', ms=10, mew=2, mec='w')
plt.text(well2_il, well2_xl+9, 'Well 2', ha='center', fontsize=10)
plt.text(well3_il, well3_xl+9, 'Well 3', ha='center', fontsize=10)
plt.title('Top Heimdal')
plt.colorbar(im0, shrink=0.5)
plt.axes().set_aspect(0.4)


# ## step 3: AVO attributes
# 
# The original dataset also has AVO Intercept and Gradient at Top Heimdal; I will now load and grid them.
# 
# In a [separate notebook](http://nbviewer.ipython.org/github/aadm/geophysical_notes/blob/master/avo_attributes.ipynb) I also show a solution to calculate AVO attributes from Near and Far seismic data.

# In[19]:


ig=np.recfromtxt('RG_3dinv.txt', names=['i','g','z','xl','il'], skip_header=16)

II=griddata((ig['il'],ig['xl']),ig['i'],(X,Y),method='cubic')
GG=griddata((ig['il'],ig['xl']),ig['g'],(X,Y),method='cubic')


# In[20]:


f, ax = plt.subplots(nrows=1,ncols=2,figsize=(14,8))
map0=ax[0].pcolormesh(X, Y, II, cmap='cubehelix')
map1=ax[1].pcolormesh(X, Y, GG, cmap='cubehelix')
ax[0].set_title('Top Heimdal AVO Intercept')
ax[1].set_title('Top Heimdal AVO Gradient')
plt.colorbar(map0, ax=ax[0], shrink=0.5)
plt.colorbar(map1, ax=ax[1], shrink=0.5)
for aa in ax:
    CS=aa.contour(X, Y, Zf, 5, colors='w')
    plt.clabel(CS, inline=1, fontsize=10, fmt='%.0f')
    aa.plot(well2_il, well2_xl, 'ok', ms=10, mew=2, mec='w')
    aa.plot(well3_il, well3_xl, 'ok', ms=10, mew=2, mec='w')
    aa.set_xlabel('IL'), aa.set_ylabel('XL')
    aa.set_aspect(0.4)


# ## calculate AVO projection maps

# The final maps are calculated with this equation:
# 
# $$ R = I * \cos \chi + G * \sin \chi $$
# 
# According the analysis previously done on well data, we should have maximum separation between brine and oil at +40 degrees (_AVO fluid projection angle_) and minimum separation at -50 degrees (_AVO lithology projection angle_).
# 
# However, it's always best to get a feel for the various combination and "let the data speak", so we can build a little cycle to loop over various projection angles:

# In[21]:


ang=np.linspace(-180,180,9)
f, ax = plt.subplots(3,3, figsize=(14,12))
ax = ax.ravel()
for i,val in enumerate(ang):
    chi=np.deg2rad(val)
    AA = II*np.cos(chi)+GG*np.sin(chi)
    vmin = np.percentile(AA, 1)
    vmax = np.percentile(AA, 99)
    rr=ax[i].pcolormesh(X,Y,AA,vmin=vmin,vmax=vmax,cmap='cubehelix')
    ax[i].set_title('$\chi={0}$'.format(val))
    cb=plt.colorbar(rr, ax=ax[i], shrink=0.75)
    cb.ax.tick_params(labelsize=8)
for aa in ax:
    aa.set_xticklabels([])
    aa.set_yticklabels([])


# In a real world scenario I would run the previous analysis on the entire angle range at 10 degrees interval and then choose the best one, trusting to see something radically different around the angles suggested by the previous analysis done on well data.
# 
# For this dataset, well 2 gives $\chi = 50^\circ$ for the lithological projection and  $\chi = -40^\circ$ for the fluid projection. The following lines plot these two AVO projections.

# In[22]:


chi_litho=np.deg2rad(-50)
chi_fluid=np.deg2rad(40)

AVOL = II*np.cos(chi_litho)+GG*np.sin(chi_litho)
AVOF = II*np.cos(chi_fluid)+GG*np.sin(chi_fluid)

f, ax = plt.subplots(1,2,figsize=(14,8))
map0 = ax[0].pcolormesh(X, Y, AVOL, cmap='cubehelix')
map1 = ax[1].pcolormesh(X, Y, AVOF, cmap='cubehelix')
ax[0].set_title('AVO Lithology Projection')
ax[1].set_title('AVO Fluid Projection')
for aa in ax.flatten():
    CS=aa.contour(X,Y,Zf,5,colors='w')
    plt.clabel(CS, inline=1, fontsize=10, fmt='%.0f')
    aa.plot(well2_il, well2_xl, 'ok', ms=10, mew=2, mec='w')
    aa.plot(well3_il, well3_xl, 'ok', ms=10, mew=2, mec='w')
    aa.set_xlabel('IL'), aa.set_ylabel('XL')  
    aa.set_aspect(0.4)
    plt.colorbar(map0, ax=aa, shrink=0.5)


# I don't know if it is worth repeating, but please bear in mind that in this notebook I am not trying to make actual sense of this dataset (I think Avseth et al. have already done everything possible on this data).
# 
# If I were to present the results however I would not stop here and what I would do first of all is to compare these maps to the original inputs (Intercept and Gradient), as well as to simple Near and Far amplitude extractions. Then put together some explanation for the visible patterns. But that's another story.