'''
===================
rpm.py
===================
Rock physics models
@author: Alessandro Amato del Monte
'''
import numpy as np


def critpor(K0, G0, phi, phic=0.4):
    """
    Critical porosity, Nur et al. (1991, 1995) written by Alessandro Amato
    del Monte (2015) from Mavko et al., Rock Physics Handbook, p.353.

    Args:
        K0, G0: float or array_like
            Mineral bulk & shear modulus in GPa
        phi: float or array_like
            porosity
        phic: float or array_like
            critical porosity (default 0.4)
    
    Returns:
        Tuple: K_DRY, G_DRY: dry rock bulk & shear modulus in GPa
    """
    K_DRY  = K0 * (1-phi/phic)
    G_DRY  = G0 * (1-phi/phic)
    return K_DRY, G_DRY


def hertzmindlin(K0, G0, phi, phic=0.4, Cn=8.6, P=10, f=1):
    '''
    Hertz-Mindlin model
    written by Alessandro Amato del Monte (2015)
    from Mavko et al., Rock Physics Handbook, p.246
    
    Args:
        K0, G0: mineral bulk & shear modulus in GPa
        phi: porosity
        phic: critical porosity (default 0.4)
        Cn: coordination nnumber (default 8.6)
        P: confining pressure in MPa (default 10)
        f: shear modulus correction factor, f=1 for dry pack with perfect adhesion between particles and f=0 for dry frictionless pack.
    
    Returns:
        Tuple: K_DRY, G_DRY: dry rock bulk & shear modulus in GPa
    '''
    P /= 1e3 # converts pressure in same units as solid moduli (GPa)
    PR0=(3*K0-2*G0)/(6*K0+2*G0) # poisson's ratio of mineral mixture
    K_HM = (P*(Cn**2*(1-phic)**2*G0**2) / (18*np.pi**2*(1-PR0)**2))**(1/3)
    G_HM = ((2+3*f-PR0*(1+3*f))/(5*(2-PR0))) * ((P*(3*Cn**2*(1-phic)**2*G0**2)/(2*np.pi**2*(1-PR0)**2)))**(1/3)
    return K_HM, G_HM

def softsand(K0, G0, phi, phic=0.4, Cn=8.6, P=10, f=1):
    '''
    Soft-sand (uncemented) model
    written by Alessandro Amato del Monte (2015)
    from Mavko et al., Rock Physics Handbook, p.258
    
    Args:
        K0, G0: mineral bulk & shear modulus in GPa
        phi: porosity
        phic: critical porosity (default 0.4)
        Cn: coordination nnumber (default 8.6)
        P: confining pressure in MPa (default 10)
        f: shear modulus correction factor, f=1 for dry pack with perfect adhesion between particles and f=0 for dry frictionless pack
    
    Returns:
        Tuple: K_DRY, G_DRY: dry rock bulk & shear modulus in GPa
    '''
    K_HM, G_HM = hertzmindlin(K0, G0, phi, phic, Cn, P, f)
    K_DRY =-4/3*G_HM + (((phi/phic)/(K_HM+4/3*G_HM)) + ((1-phi/phic)/(K0+4/3*G_HM)))**-1
    tmp = G_HM/6*((9*K_HM+8*G_HM) / (K_HM+2*G_HM))
    G_DRY = -tmp + ((phi/phic)/(G_HM+tmp) + ((1-phi/phic)/(G0+tmp)))**-1
    return K_DRY, G_DRY

def stiffsand(K0, G0, phi, phic=0.4, Cn=8.6, P=10, f=1):
    '''
    Stiff-sand model
    written by Alessandro Amato del Monte (2015)
    from Mavko et al., Rock Physics Handbook, p.260
    
    Args:
        K0, G0: mineral bulk & shear modulus in GPa
        phi: porosity
        phic: critical porosity (default 0.4)
        Cn: coordination nnumber (default 8.6)
        P: confining pressure in MPa (default 10)
        f: shear modulus correction factor, f=1 for dry pack with perfect adhesion between particles and f=0 for dry frictionless pack
    
    Returns
        Tuple: K_DRY, G_DRY: dry rock bulk & shear modulus in GPa
    '''
    K_HM, G_HM = hertzmindlin(K0, G0, phi, phic, Cn, P, f)
    K_DRY = -4/3*G0 + (((phi/phic)/(K_HM+4/3*G0)) + ((1-phi/phic)/(K0+4/3*G0)))**-1
    tmp = G0/6*((9*K0+8*G0) / (K0+2*G0))
    G_DRY = -tmp + ((phi/phic)/(G_HM+tmp) + ((1-phi/phic)/(G0+tmp)))**-1
    return K_DRY, G_DRY

def contactcement(K0, G0, phi, phic=0.4, Cn=8.6, Kc=37, Gc=45, scheme=2):
    '''
    Contact cement (cemented sand) model, Dvorkin-Nur (1996)
    written by Alessandro Amato del Monte (2015)
    from Mavko et al., Rock Physics Handbook, p.255
    
    Args:
        K0, G0: mineral bulk & shear modulus in GPa
        phi: porosity
        phic: critical porosity (default 0.4)
        Cn: coordination nnumber (default 8.6)
        Kc, Gc: cement bulk & shear modulus in GPa (default 37, 45 i.e. quartz)
        scheme: 1=cement deposited at grain contacts, 2=in uniform layer around grains (default 2)
    Returns:
        Tuple: K_DRY, G_DRY: dry rock bulk & shear modulus in GPa
    '''
    PR0=(3*K0-2*G0)/(6*K0+2*G0)
    PRc = (3*Kc-2*Gc)/(6*Kc+2*Gc)
    if scheme == 1: # scheme 1: cement deposited at grain contacts
        alpha = ((phic-phi)/(3*Cn*(1-phic))) ** (1/4)
    else: # scheme 2: cement evenly deposited on grain surface
        alpha = ((2*(phic-phi))/(3*(1-phic)))**(1/2)
    LambdaN = (2*Gc*(1-PR0)*(1-PRc)) / (np.pi*G0*(1-2*PRc))
    N1 = -0.024153*LambdaN**-1.3646
    N2 = 0.20405*LambdaN**-0.89008
    N3 = 0.00024649*LambdaN**-1.9864
    Sn = N1*alpha**2 + N2*alpha + N3
    LambdaT = Gc/(np.pi*G0)
    T1 = -10**-2*(2.26*PR0**2+2.07*PR0+2.3)*LambdaT**(0.079*PR0**2+0.1754*PR0-1.342)
    T2 = (0.0573*PR0**2+0.0937*PR0+0.202)*LambdaT**(0.0274*PR0**2+0.0529*PR0-0.8765)
    T3 = 10**-4*(9.654*PR0**2+4.945*PR0+3.1)*LambdaT**(0.01867*PR0**2+0.4011*PR0-1.8186)
    St = T1*alpha**2 + T2*alpha + T3
    K_DRY = 1/6*Cn*(1-phic)*(Kc+(4/3)*Gc)*Sn
    G_DRY = 3/5*K_DRY+3/20*Cn*(1-phic)*Gc*St
    return K_DRY, G_DRY
