%matplotlib inline
from __future__ import division, print_function
import math
import numpy as np
import numpy.testing
import multiprocessing
import threading
import numba
from time import time
from fatiando.mesher import TesseroidMesh, Tesseroid
from fatiando import constants, utils, gridder
from fatiando.gravmag.tesseroid import gz as gz_cython
from fatiando.vis import mpl

def make_data():    
    model = TesseroidMesh((0, 5, 0, 10, 1000, -5000), (1, 5, 5))
    model.addprop('density', 1000*np.ones(model.size))
    x, y, z = gridder.regular((0, 5, 0, 10), (50, 50), z=50000)
    return x, y, z, model

def test_gz(func, **kwargs):
    print("Testing {}:".format(func.__name__))    
    x, y, z, model = make_data()
    # Test the func against the Cython version in Fatiando
    f1 = func(x, y, z, model, **kwargs)
    f2 = gz_cython(x, y, z, model)
    np.testing.assert_allclose(f1, f2, err_msg="Results don't match current implementation")
    print(" - Results match current implementation")
    # Time the execution
    times = []
    for i in xrange(5):
        t1 = time()
        func(x, y, z, model, **kwargs)
        t2 = time()
        times.append(t2 - t1)
    print(" - Execution time: {:f} s".format(np.median(times)))

nodes = np.array([-0.577350269189625731058868041146,
                  0.577350269189625731058868041146])
stack_size = 100

def _check_input(lon, lat, height, ratio, out):
    lon = np.radians(lon)
    radius = height + constants.MEAN_EARTH_RADIUS
    coslat = np.cos(np.radians(lat))
    sinlat = np.sin(np.radians(lat))
    if out is None:
        result = np.zeros_like(lon)
    else:
        result = out        
    return lon, coslat, sinlat, radius, ratio, result

def _get_engine(engine):
    if engine == 'numpy':
        func = _numpy_engine
    elif engine == 'numba':
        func = _numba_engine
    else:
        raise ValueError("engine trouble")
    return func

def _get_strides(size, njobs):    
    n = size//njobs
    strides = [(i*n, (i + 1)*n) for i in xrange(njobs - 1)]
    strides.append((strides[-1][-1], size))
    return strides

def _split_arrays(arrays, args, njobs):
    strides = _get_strides(arrays[0].size, njobs)
    return [[x[low:high] for x in arrays] + args
            for low, high in strides]

def _dispatcher(args):
    lon, coslat, sinlat, radius, result, model, engine, ratio = args  
    func = _get_engine(engine)
    for c in model:
        density = c.props['density']
        func(lon, coslat, sinlat, radius, c, density, ratio, result)
    return result

def gz(lon, lat, height, model, ratio=1.6, engine=None, njobs=None, out=None):
    lon, coslat, sinlat, radius, ratio, result = _check_input(lon, lat, height, ratio, out)
    if njobs is None or njobs == 1:
        _dispatcher([lon, coslat, sinlat, radius, result, model, engine, ratio])
    else:
        chunks = _split_arrays([lon, coslat, sinlat, radius, result], [model, engine, ratio], njobs)
        pool = multiprocessing.Pool(njobs)
        result = np.hstack(pool.map(_dispatcher, chunks))
        pool.close()
    result *= constants.G*constants.SI2MGAL
    return result

def _numpy_engine(lon, coslat, sinlat, radius, cell, density, ratio, result):
    for l in xrange(lon.size):
        stack = [cell]
        while stack:
            t = stack.pop()
            distance, Llon, Llat, Lr = distance_size_numpy(lon[l], coslat[l], sinlat[l], radius[l], t)
            nlon = 1 if distance/Llon > ratio else 2
            nlat = 1 if distance/Llat > ratio else 2
            nr = 1 if distance/Lr > ratio else 2
            new_cells = nlon*nlat*nr
            if new_cells > 1:
                if new_cells + len(stack) > stack_size:
                    print(distance, Llon, Llat, Lr)
                    raise OverflowError('Tesseroid stack overflowed')
                stack.extend(t.split(nlon, nlat, nr))
            else:
                lonc, sinlatc, coslatc, rc, scale = scale_nodes_numpy(t.get_bounds(), nodes)
                result[l] += scale*density*kernelz_numpy(lon[l], coslat[l], sinlat[l], radius[l], 
                                                         lonc, sinlatc, coslatc, rc)

def scale_nodes_numpy(bounds, nodes):
    w, e, s, n, top, bottom = bounds
    d2r = np.pi/180
    dlon = d2r*(e - w)
    dlat = d2r*(n - s)
    dr = top - bottom
    # Scale the GLQ nodes to the integration limits
    lonc = 0.5*dlon*nodes + d2r*0.5*(e + w)
    latc = 0.5*dlat*nodes + d2r*0.5*(n + s)
    sinlatc = np.sin(latc)
    coslatc = np.cos(latc)
    rc = (0.5*dr*nodes +
          0.5*(top + bottom + 2*constants.MEAN_EARTH_RADIUS))
    scale = dlon*dlat*dr*0.125
    return lonc, sinlatc, coslatc, rc, scale
    
def distance_size_numpy(lon, coslat, sinlat, radius, cell):
    w, e, s, n, top, bottom = cell.get_bounds()
    d2r = np.pi/180
    rt = 0.5*(top + bottom) + constants.MEAN_EARTH_RADIUS
    lont = d2r*0.5*(w + e)
    latt = d2r*0.5*(s + n)
    sinlatt = np.sin(latt)
    coslatt = np.cos(latt)
    cospsi = sinlat*sinlatt + coslat*coslatt*np.cos(lon - lont)
    distance = np.sqrt(radius**2 + rt**2 - 2*radius*rt*cospsi)
    # Calculate the dimensions of the tesseroid in meters
    rtop = top + constants.MEAN_EARTH_RADIUS
    Llon = rtop*np.arccos(sinlatt**2 + (coslatt**2)*np.cos(d2r*(e - w)))
    Llat = rtop*np.arccos(np.sin(d2r*n)*np.sin(d2r*s) + np.cos(d2r*n)*np.cos(d2r*s))
    Lr = top - bottom
    return distance, Llon, Llat, Lr

def kernelz_numpy(lon, coslat, sinlat, radius, lonc, sinlatc, coslatc, rc):
    r_sqr = radius**2
    result = 0
    for i in range(2):
        coslon = np.cos(lon - lonc[i])
        for j in range(2):
            cospsi = sinlat*sinlatc[j] + coslat*coslatc[j]*coslon
            for k in range(2):
                l_sqr = r_sqr + rc[k]**2 - 2*radius*rc[k]*cospsi
                kappa = (rc[k]**2)*coslatc[j]
                result += kappa*(rc[k]*cospsi - radius)/(l_sqr**1.5)
    result *= -1
    return result

@numba.jit(looplift=True, nogil=True)
def _numba_engine(lon, coslat, sinlat, radius, cell, density, ratio, result):
    bounds = np.array(cell.get_bounds())
    stack = np.empty((stack_size, 6))
    lonc = np.empty_like(nodes)
    sinlatc = np.empty_like(nodes)
    coslatc = np.empty_like(nodes)
    rc = np.empty_like(nodes)
    for l in range(result.size):
        for i in range(6):
            stack[0, i] = bounds[i]
        stktop = 0
        while stktop >= 0:
            w, e, s, n, top, bottom = stack[stktop, :]
            stktop -= 1
            distance, Llon, Llat, Lr = distance_size_numba(lon[l], coslat[l], sinlat[l], radius[l], 
                                                           w, e, s, n, top, bottom)
            nlon, nlat, nr, new_cells = divisions_per_dimension(distance, Llon, Llat, Lr, ratio)
            if new_cells > 1:
                if new_cells + (stktop + 1) > stack_size:
                    raise OverflowError
                stktop = split_numba(w, e, s, n, top, bottom, nlon, nlat, nr,
                                     stack, stktop)
            else:
                scale = scale_nodes_numba(w, e, s, n, top, bottom, nodes, 
                                          lonc, sinlatc, coslatc, rc)
                result[l] += density*scale*kernelz_numba(lon[l], coslat[l], sinlat[l], radius[l], 
                                                lonc, sinlatc, coslatc, rc)
    
@numba.jit(nopython=True, nogil=True)
def scale_nodes_numba(w, e, s, n, top, bottom, nodes, lonc, sinlatc, coslatc, rc):
    d2r = np.pi/180
    dlon = d2r*(e - w)
    dlat = d2r*(n - s)
    dr = top - bottom
    # Scale the GLQ nodes to the integration limits
    for i in range(len(nodes)):
        lonc[i] = 0.5*dlon*nodes[i] + d2r*0.5*(e + w)
        latc = 0.5*dlat*nodes[i] + d2r*0.5*(n + s)
        sinlatc[i] = np.sin(latc)
        coslatc[i] = np.cos(latc)
        rc[i] = (0.5*dr*nodes[i] +
              0.5*(top + bottom + 2*constants.MEAN_EARTH_RADIUS))
    scale = dlon*dlat*dr*0.125
    return scale
    
@numba.jit(nopython=True, nogil=True)
def distance_size_numba(lon, coslat, sinlat, radius, w, e, s, n, top, bottom):
    d2r = np.pi/180
    rt = 0.5*(top + bottom) + constants.MEAN_EARTH_RADIUS
    lont = d2r*0.5*(w + e)
    latt = d2r*0.5*(s + n)
    sinlatt = np.sin(latt)
    coslatt = np.cos(latt)
    cospsi = sinlat*sinlatt + coslat*coslatt*np.cos(lon - lont)
    distance = np.sqrt(radius**2 + rt**2 - 2*radius*rt*cospsi)
    # Calculate the dimensions of the tesseroid in meters
    rtop = top + constants.MEAN_EARTH_RADIUS
    Llon = rtop*np.arccos(sinlatt**2 + (coslatt**2)*np.cos(d2r*(e - w)))
    Llat = rtop*np.arccos(np.sin(d2r*n)*np.sin(d2r*s) + np.cos(d2r*n)*np.cos(d2r*s))
    Lr = top - bottom
    return distance, Llon, Llat, Lr

@numba.jit(nopython=True, nogil=True)
def kernelz_numba(lon, coslat, sinlat, radius, lonc, sinlatc, coslatc, rc):
    r_sqr = radius**2
    result = 0
    for i in range(2):
        coslon = np.cos(lon - lonc[i])
        for j in range(2):
            cospsi = sinlat*sinlatc[j] + coslat*coslatc[j]*coslon
            for k in range(2):
                l_sqr = r_sqr + rc[k]**2 - 2*radius*rc[k]*cospsi
                kappa = (rc[k]**2)*coslatc[j]
                result += kappa*(rc[k]*cospsi - radius)/(l_sqr**1.5)
    result *= -1
    return result

@numba.jit(nopython=True, nogil=True)
def split_numba(w, e, s, n, top, bottom, nlon, nlat, nr, stack, stktop):
    dlon = (e - w)/nlon
    dlat = (n - s)/nlat
    dr = (top - bottom)/nr
    for i in xrange(nlon):
        for j in xrange(nlat):
            for k in xrange(nr):
                stktop += 1
                stack[stktop, 0] = w + i*dlon
                stack[stktop, 1] = w + (i + 1)*dlon
                stack[stktop, 2] = s + j*dlat
                stack[stktop, 3] = s + (j + 1)*dlat
                stack[stktop, 4] = bottom + (k + 1)*dr
                stack[stktop, 5] = bottom + k*dr
    return stktop

@numba.jit(nopython=True, nogil=True)
def divisions_per_dimension(distance, Llon, Llat, Lr, ratio):
    nlon = 1 if distance/Llon > ratio else 2
    nlat = 1 if distance/Llat > ratio else 2
    nr = 1 if distance/Lr > ratio else 2
    return nlon, nlat, nr, nlon*nlat*nr

test_gz(gz_cython)

test_gz(gz, engine='numba')

test_gz(gz, engine='numpy')