''' A package with Python functions for forward and inverse calculations
of the EM-34 response. Written by Vincent Post, April 2007.'''

from numpy import array, sqrt
from scipy import optimize, hstack, r_, arange

def EMresponse(spacing, z, sigma, print_output = True, data = array([])):
    '''Function that returns the theoretical response of the EM-34
    assuming a horizontally-layered subsurface. Formulas from McNeill (1980).
    Input:
        spacing: array with loop spacings
        z: array with depths to bottoms of layers
        sigma: array with conductivities of the layers
        print_output (optional): boolean to output results to screen
        data (optional): array with the measured data to be printed
    Output:
        Array with 6 numbers. The first 3 numbers represent the theoretical
        conductivity recorded by the EM-34 for loop spacings of 10, 20 and 40 m
        for a vertical dipole (i.e., horizontal loops), the last 3 numbers
        represent the theoretical conductivity recorded by the EM-34 for loop
        spacings of 10, 20 and 40 m for a vertical dipole (i.e., horizontal
        loops).'''

    hloops = vloops = array([])

    for s in spacing:
        # For horizontal loops
        z_s = z / s
        z_contr = 1 / sqrt(1 + 4 * (z_s ** 2))
        z_contr = hstack((1.0, z_contr, 0.0))
        z_contr = z_contr[:-1] - z_contr[1:]
        z_contr = z_contr * sigma
        hloops = r_[hloops, sum(z_contr)]
        #print 'hor. loops (vert. dipole)', s, 'm: ', sum(z_contr)

        # For vertical loops
        z_s = z / s
        z_contr = sqrt(1 + 4 * (z_s ** 2)) - 2 * z_s
        z_contr = hstack((1.0, z_contr, 0.0))
        z_contr = z_contr[:-1] - z_contr[1:]
        z_contr = z_contr * sigma
        vloops = r_[vloops, sum(z_contr)]
        #print 'vert. loops (hor. dipole)', s, 'm: ', sum(z_contr)

    if print_output:
        print_response(spacing, hloops, vloops)
        if (len(data)):
            print_data(spacing, data)
        print_model(z, sigma)
        
    return r_[hloops, vloops]

def collect_unknowns(z, sigma):
    '''Function that looks for negative numbers in the z and sigma arrays. These are used as flags that indicate
    they have to be optimized.
    Input:
        z: array with depths to bottoms of layers
        sigma: array with conductivities of the layers
    Output:
        Array with the unknown depths followed by the unknown conductivities. These are passed on in the
        optimize.leastsq function'''
    
    return_value = array([])
    for a in z:
        if (a < 0):
            return_value = r_[return_value, -a]
        
    for a in sigma:
        if (a < 0):
            return_value = r_[return_value, -a]

    return return_value

def EMresponse_wrapper(p, spacing, z, sigma):
    '''This function is a wrapper that first updates the z and sigma arrays (or rather, local copies thereof) with the
    numbers that are being optimized. It then calls EMresponse().
    Input:
        p: Array with the unknown depths followed by the unknown conductivities.
        spacing: array with loop spacings
        z: array with depths to bottoms of layers
        sigma: array with conductivities of the layers
    Output:
        Same as EMresponse function. See above'''

    z_tmp, sigma_tmp = update_model(p, z, sigma)
    return EMresponse(spacing, z_tmp, sigma_tmp, False)
    
def update_model(p, z, sigma):
    '''This function looks for negative numbers in the z and sigma arrays and inserts the numbers
    that are being optimized (from the p array) at the appropriate positions.
    Input:
        p: Array with the unknown depths followed by the unknown conductivities.
        z: array with depths to bottoms of layers
        sigma: array with conductivities of the layers
    Output:
        Local copies of the updated z and sigma arrays'''

    i = 0
    # Note that we use local copies of z and sigma
    z_tmp = z.copy()
    for j in arange(len(z_tmp)):
        if (z_tmp[j] < 0):
            z_tmp[j] = p[i]
            i = i + 1

    sigma_tmp = sigma.copy()
    for j in arange(len(sigma_tmp)):
        if (sigma_tmp[j] < 0):
            sigma_tmp[j] = p[i]
            i = i + 1
            
    return z_tmp, sigma_tmp

def fit_model(spacing, z, sigma, data):
    '''Function that optimizes the layered conductivity model by minimizing the difference between the calculated
    theoretical EM-34 response and the measured conductivities.
    Input:
        spacing: array with loop spacings
        z: array with depths to bottoms of layers
        sigma: array with conductivities of the layers
        data: Array with 6 numbers. The first 3 numbers represent the EM-34 measured conductivity for loop spacings
        of 10, 20 and 40 m for a vertical dipole (i.e., horizontal loops), the last 3 numbers represent the EM-34
        measured conductivity for loop spacings of 10, 20 and 40 m for a vertical dipole (i.e., horizontal loops).
    Output:
        One array with the depths to bottom of the layers (optimized z) and one array with the conductivities (optimized
        sigma) of the layers.'''

    # First determine which dephts and conductivities need to be optimized and store in p0
    p0 = collect_unknowns(z, sigma)
    # Define the target function to which the field data will be fitted
    fitfunc = lambda p, w, x, y: EMresponse_wrapper(p, w, x, y)
    # Define the objective function
    errfunc = lambda p, w, x, y, z: fitfunc(p, w, x, y) - z
    # Perform the actual optimization using the optimize object
    p1, success = optimize.leastsq(errfunc, p0, args = (spacing, z, sigma, data))

    return update_model(p1, z, sigma)
        
def print_model(z, sigma):
    '''Procedure thats prints the layered conductivity model as columns.
    Input:
        z: array with depths to bottoms of layers
        sigma: array with conductivities of the layers'''

    print ''
    print 'Layered conductivity model'
    print '=========================='
    s = "%s\t%s\t%s"  % ('Conduc', 'Resis', 'Depth')
    s = "%s\n%s\t%s"  % (s, 'tivity', 'tivity')
    print s
    print '---------------------'    
    i = 0
    for v in sigma:
        if (i < len(z)):
            s = "%5.1f\t%7.1f\t%5.1f"  % (v, 1000.0 / v, z[i])
        else:
            s = "%5.1f\t%7.1f"  % (v, 1000.0 / v)
        i = i + 1
        print s
        
    print ''

def print_response(spacing, hloops, vloops):
    '''Procedure thats prints the calculated theoretical EM-34 response on the screen.
    Input:
        spacing: Array with loop spacing
        hloops: Array with calculated EM-34 response for horizontal loops (vertical dipole)
        vloops: Array with calculated EM-34 response for vertical loops (horizontal dipole)'''

    print ''
    print 'Calculated theoretical EM-34 response'
    print '====================================='
    s = "\t\t%s"  % ('  loop spacing (m) ')
    print s
    s = "\t\t%d\t%d\t%d"  % (spacing[0], spacing[1], spacing[2])
    print s
    s = "\t\t%s"  % ('---------------------')
    print s

    s = "%s\t%5.1f\t%5.1f\t%5.1f"  % ('hor. loops', hloops[0], hloops[1], hloops[2])
    s = "%s\n%s\t%5.1f\t%5.1f\t%5.1f"  % (s, 'vert. loops', vloops[0], vloops[1], vloops[2])
    print s
    print ''

def print_data(spacing, data):
    '''Procedure thats prints the calculated theoretical EM-34 response on the screen.
    Input:
        spacing: Array with loop spacing
        hloops: Array with calculated EM-34 response for horizontal loops (vertical dipole)
        vloops: Array with calculated EM-34 response for vertical loops (horizontal dipole)'''

    print ''
    print 'Measured EM-34 response'
    print '====================================='
    s = "\t\t%s"  % ('  loop spacing (m) ')
    print s
    s = "\t\t%d\t%d\t%d"  % (spacing[0], spacing[1], spacing[2])
    print s
    s = "\t\t%s"  % ('---------------------')
    print s

    s = "%s\t%5.1f\t%5.1f\t%5.1f"  % ('hor. loops', data[0], data[1], data[2])
    s = "%s\n%s\t%5.1f\t%5.1f\t%5.1f"  % (s, 'vert. loops', data[3], data[4], data[5])
    print s
    print ''