# distutils: include_dirs = LIB_DIR
# distutils: libraries = STD_LIBS
# distutils: sources = FIXED_INTERP
# distutils: language = c++
# distutils: extra_compile_args = CPP14_FLAG
# distutils: extra_link_args = CPP14_FLAG
"""
Simple integrators for the radiative transfer equation



"""


import numpy as np

cimport cython
cimport numpy as np
from libc.math cimport atan2, cos, fabs, floor, sin, sqrt

from yt.utilities.lib.fp_utils cimport fmin

from .fixed_interpolator cimport *


@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
cdef int walk_volume(VolumeContainer *vc,
                     np.float64_t v_pos[3],
                     np.float64_t v_dir[3],
                     sampler_function *sample,
                     void *data,
                     np.float64_t *return_t = NULL,
                     np.float64_t max_t = 1.0) noexcept nogil:
    cdef int cur_ind[3]
    cdef int step[3]
    cdef int x, y, i, hit, direction
    cdef np.float64_t intersect_t = 1.1
    cdef np.float64_t iv_dir[3]
    cdef np.float64_t tmax[3]
    cdef np.float64_t tdelta[3]
    cdef np.float64_t exit_t = -1.0, enter_t = -1.0
    cdef np.float64_t tl, temp_x, temp_y = -1
    if max_t > 1.0: max_t = 1.0
    direction = -1
    if vc.left_edge[0] <= v_pos[0] and v_pos[0] < vc.right_edge[0] and \
       vc.left_edge[1] <= v_pos[1] and v_pos[1] < vc.right_edge[1] and \
       vc.left_edge[2] <= v_pos[2] and v_pos[2] < vc.right_edge[2]:
        intersect_t = 0.0
        direction = 3
    for i in range(3):
        if (v_dir[i] < 0):
            step[i] = -1
        elif (v_dir[i] == 0.0):
            step[i] = 0
            continue
        else:
            step[i] = 1
        iv_dir[i] = 1.0/v_dir[i]
        if direction == 3: continue
        x = (i+1) % 3
        y = (i+2) % 3
        if step[i] > 0:
            tl = (vc.left_edge[i] - v_pos[i])*iv_dir[i]
        else:
            tl = (vc.right_edge[i] - v_pos[i])*iv_dir[i]
        temp_x = (v_pos[x] + tl*v_dir[x])
        temp_y = (v_pos[y] + tl*v_dir[y])
        if fabs(temp_x - vc.left_edge[x]) < 1e-10*vc.dds[x]:
            temp_x = vc.left_edge[x]
        elif fabs(temp_x - vc.right_edge[x]) < 1e-10*vc.dds[x]:
            temp_x = vc.right_edge[x]
        if fabs(temp_y - vc.left_edge[y]) < 1e-10*vc.dds[y]:
            temp_y = vc.left_edge[y]
        elif fabs(temp_y - vc.right_edge[y]) < 1e-10*vc.dds[y]:
            temp_y = vc.right_edge[y]
        if vc.left_edge[x] <= temp_x and temp_x <= vc.right_edge[x] and \
           vc.left_edge[y] <= temp_y and temp_y <= vc.right_edge[y] and \
           0.0 <= tl and tl < intersect_t:
            direction = i
            intersect_t = tl
    if enter_t >= 0.0: intersect_t = enter_t
    if not ((0.0 <= intersect_t) and (intersect_t < max_t)): return 0
    for i in range(3):
        # Two things have to be set inside this loop.
        # cur_ind[i], the current index of the grid cell the ray is in
        # tmax[i], the 't' until it crosses out of the grid cell
        tdelta[i] = step[i] * iv_dir[i] * vc.dds[i]
        if i == direction and step[i] > 0:
            # Intersection with the left face in this direction
            cur_ind[i] = 0
        elif i == direction and step[i] < 0:
            # Intersection with the right face in this direction
            cur_ind[i] = vc.dims[i] - 1
        else:
            # We are somewhere in the middle of the face
            temp_x = intersect_t * v_dir[i] + v_pos[i] # current position
            temp_y = ((temp_x - vc.left_edge[i])*vc.idds[i])
            # There are some really tough cases where we just within a couple
            # least significant places of the edge, and this helps prevent
            # killing the calculation through a segfault in those cases.
            if -1 < temp_y < 0 and step[i] > 0:
                temp_y = 0.0
            elif vc.dims[i] - 1 < temp_y < vc.dims[i] and step[i] < 0:
                temp_y = vc.dims[i] - 1
            cur_ind[i] =  <int> (floor(temp_y))
        if step[i] > 0:
            temp_y = (cur_ind[i] + 1) * vc.dds[i] + vc.left_edge[i]
        elif step[i] < 0:
            temp_y = cur_ind[i] * vc.dds[i] + vc.left_edge[i]
        tmax[i] = (temp_y - v_pos[i]) * iv_dir[i]
        if step[i] == 0:
            tmax[i] = 1e60
    # We have to jumpstart our calculation
    for i in range(3):
        if cur_ind[i] == vc.dims[i] and step[i] >= 0:
            return 0
        if cur_ind[i] == -1 and step[i] <= -1:
            return 0
    enter_t = intersect_t
    hit = 0
    while 1:
        hit += 1
        if tmax[0] < tmax[1]:
            if tmax[0] < tmax[2]:
                i = 0
            else:
                i = 2
        else:
            if tmax[1] < tmax[2]:
                i = 1
            else:
                i = 2
        exit_t = fmin(tmax[i], max_t)
        sample(vc, v_pos, v_dir, enter_t, exit_t, cur_ind, data)
        cur_ind[i] += step[i]
        enter_t = tmax[i]
        tmax[i] += tdelta[i]
        if cur_ind[i] < 0 or cur_ind[i] >= vc.dims[i] or enter_t >= max_t:
            break
    if return_t != NULL: return_t[0] = exit_t
    return hit

def hp_pix2vec_nest(long nside, long ipix):
    raise NotImplementedError
    # cdef double v[3]
    # healpix_interface.pix2vec_nest(nside, ipix, v)
    # cdef np.ndarray[np.float64_t, ndim=1] tr = np.empty((3,), dtype='float64')
    # tr[0] = v[0]
    # tr[1] = v[1]
    # tr[2] = v[2]
    # return tr

def arr_pix2vec_nest(long nside,
                     np.ndarray[np.int64_t, ndim=1] aipix):
    raise NotImplementedError
    # cdef int n = aipix.shape[0]
    # cdef int i
    # cdef double v[3]
    # cdef long ipix
    # cdef np.ndarray[np.float64_t, ndim=2] tr = np.zeros((n, 3), dtype='float64')
    # for i in range(n):
    #     ipix = aipix[i]
    #     healpix_interface.pix2vec_nest(nside, ipix, v)
    #     tr[i,0] = v[0]
    #     tr[i,1] = v[1]
    #     tr[i,2] = v[2]
    # return tr

def hp_vec2pix_nest(long nside, double x, double y, double z):
    raise NotImplementedError
    # cdef double v[3]
    # v[0] = x
    # v[1] = y
    # v[2] = z
    # cdef long ipix
    # healpix_interface.vec2pix_nest(nside, v, &ipix)
    # return ipix

def arr_vec2pix_nest(long nside,
                     np.ndarray[np.float64_t, ndim=1] x,
                     np.ndarray[np.float64_t, ndim=1] y,
                     np.ndarray[np.float64_t, ndim=1] z):
    raise NotImplementedError
    # cdef int n = x.shape[0]
    # cdef int i
    # cdef double v[3]
    # cdef long ipix
    # cdef np.ndarray[np.int64_t, ndim=1] tr = np.zeros(n, dtype='int64')
    # for i in range(n):
    #     v[0] = x[i]
    #     v[1] = y[i]
    #     v[2] = z[i]
    #     healpix_interface.vec2pix_nest(nside, v, &ipix)
    #     tr[i] = ipix
    # return tr

def hp_pix2ang_nest(long nside, long ipnest):
    raise NotImplementedError
    # cdef double theta, phi
    # healpix_interface.pix2ang_nest(nside, ipnest, &theta, &phi)
    # return (theta, phi)

def arr_pix2ang_nest(long nside, np.ndarray[np.int64_t, ndim=1] aipnest):
    raise NotImplementedError
    # cdef int n = aipnest.shape[0]
    # cdef int i
    # cdef long ipnest
    # cdef np.ndarray[np.float64_t, ndim=2] tr = np.zeros((n, 2), dtype='float64')
    # cdef double theta, phi
    # for i in range(n):
    #     ipnest = aipnest[i]
    #     healpix_interface.pix2ang_nest(nside, ipnest, &theta, &phi)
    #     tr[i,0] = theta
    #     tr[i,1] = phi
    # return tr

def hp_ang2pix_nest(long nside, double theta, double phi):
    raise NotImplementedError
    # cdef long ipix
    # healpix_interface.ang2pix_nest(nside, theta, phi, &ipix)
    # return ipix

def arr_ang2pix_nest(long nside,
                     np.ndarray[np.float64_t, ndim=1] atheta,
                     np.ndarray[np.float64_t, ndim=1] aphi):
    raise NotImplementedError
    # cdef int n = atheta.shape[0]
    # cdef int i
    # cdef long ipnest
    # cdef np.ndarray[np.int64_t, ndim=1] tr = np.zeros(n, dtype='int64')
    # cdef double theta, phi
    # for i in range(n):
    #     theta = atheta[i]
    #     phi = aphi[i]
    #     healpix_interface.ang2pix_nest(nside, theta, phi, &ipnest)
    #     tr[i] = ipnest
    # return tr

@cython.boundscheck(False)
@cython.cdivision(False)
@cython.wraparound(False)
def pixelize_healpix(long nside,
                     np.ndarray[np.float64_t, ndim=1] values,
                     long ntheta, long nphi,
                     np.ndarray[np.float64_t, ndim=2] irotation):
    raise NotImplementedError
    # # We will first to pix2vec, rotate, then calculate the angle
    # cdef int i, j, thetai, phii
    # cdef long ipix
    # cdef double v0[3], v1[3]
    # cdef double pi = 3.1415926
    # cdef np.float64_t pi2 = pi/2.0
    # cdef np.float64_t phi, theta
    # cdef np.ndarray[np.float64_t, ndim=2] results
    # cdef np.ndarray[np.int32_t, ndim=2] count
    # results = np.zeros((ntheta, nphi), dtype="float64")
    # count = np.zeros((ntheta, nphi), dtype="int32")

    # cdef np.float64_t phi0 = 0
    # cdef np.float64_t dphi = 2.0 * pi/(nphi-1)

    # cdef np.float64_t theta0 = 0
    # cdef np.float64_t dtheta = pi/(ntheta-1)
    # # We assume these are the rotated theta and phi
    # for thetai in range(ntheta):
    #     theta = theta0 + dtheta * thetai
    #     for phii in range(nphi):
    #         phi = phi0 + dphi * phii
    #         # We have our rotated vector
    #         v1[0] = cos(phi) * sin(theta)
    #         v1[1] = sin(phi) * sin(theta)
    #         v1[2] = cos(theta)
    #         # Now we rotate back
    #         for i in range(3):
    #             v0[i] = 0
    #             for j in range(3):
    #                 v0[i] += v1[j] * irotation[j,i]
    #         # Get the pixel this vector is inside
    #         healpix_interface.vec2pix_nest(nside, v0, &ipix)
    #         results[thetai, phii] = values[ipix]
    #         count[i, j] += 1
    # return results, count

def healpix_aitoff_proj(np.ndarray[np.float64_t, ndim=1] pix_image,
                        long nside,
                        np.ndarray[np.float64_t, ndim=2] image,
                        np.ndarray[np.float64_t, ndim=2] irotation):
    raise NotImplementedError
    # cdef double pi = np.pi
    # cdef int i, j, k, l
    # cdef np.float64_t x, y, z, zb
    # cdef np.float64_t dx, dy, inside
    # cdef double v0[3], v1[3]
    # dx = 2.0 / (image.shape[1] - 1)
    # dy = 2.0 / (image.shape[0] - 1)
    # cdef np.float64_t s2 = sqrt(2.0)
    # cdef long ipix
    # for i in range(image.shape[1]):
    #     x = (-1.0 + i*dx)*s2*2.0
    #     for j in range(image.shape[0]):
    #         y = (-1.0 + j * dy)*s2
    #         zb = (x*x/8.0 + y*y/2.0 - 1.0)
    #         if zb > 0: continue
    #         z = (1.0 - (x/4.0)**2.0 - (y/2.0)**2.0)
    #         z = sqrt(z)
    #         # Longitude
    #         phi = (2.0*atan(z*x/(2.0 * (2.0*z*z-1.0))) + pi)
    #         # Latitude
    #         # We shift it into co-latitude
    #         theta = (asin(z*y) + pi/2.0)
    #         # Now to account for rotation we translate into vectors
    #         v1[0] = cos(phi) * sin(theta)
    #         v1[1] = sin(phi) * sin(theta)
    #         v1[2] = cos(theta)
    #         for k in range(3):
    #             v0[k] = 0
    #             for l in range(3):
    #                 v0[k] += v1[l] * irotation[l,k]
    #         healpix_interface.vec2pix_nest(nside, v0, &ipix)
    #         #print("Rotated", v0[0], v0[1], v0[2], v1[0], v1[1], v1[2], ipix, pix_image[ipix])
    #         image[j, i] = pix_image[ipix]

def arr_fisheye_vectors(int resolution, np.float64_t fov, int nimx=1, int
        nimy=1, int nimi=0, int nimj=0, np.float64_t off_theta=0.0, np.float64_t
        off_phi=0.0):
    # We now follow figures 4-7 of:
    # http://paulbourke.net/miscellaneous/domefisheye/fisheye/
    # ...but all in Cython.
    cdef np.ndarray[np.float64_t, ndim=3] vp
    cdef int i, j
    cdef np.float64_t r, phi, theta, px, py
    cdef np.float64_t fov_rad = fov * np.pi / 180.0
    cdef int nx = resolution//nimx
    cdef int ny = resolution//nimy
    vp = np.zeros((nx,ny, 3), dtype="float64")
    for i in range(nx):
        px = (2.0 * (nimi*nx + i)) / resolution - 1.0
        for j in range(ny):
            py = (2.0 * (nimj*ny + j)) / resolution - 1.0
            r = sqrt(px*px + py*py)
            if r > 1.01:
                vp[i,j,0] = vp[i,j,1] = vp[i,j,2] = 0.0
                continue
            phi = atan2(py, px)
            theta = r * fov_rad / 2.0
            theta += off_theta
            phi += off_phi
            vp[i,j,0] = sin(theta) * cos(phi)
            vp[i,j,1] = sin(theta) * sin(phi)
            vp[i,j,2] = cos(theta)
    return vp
