import numpy as np
from numba import njit, cuda, float32
from numba.cuda.libdevice import sqrtf, sinf, cosf, atan2f, acosf, tanf

@njit(cache=True)
def cart2projection(x, y, z):
    r = np.sqrt(x**2 + y**2 + z**2)
    theta_x_rad = np.arctan2(x, z)
    theta_y_rad = np.arctan2(y, z)
    return r, theta_x_rad, theta_y_rad


@njit(cache=True)
def projection2cart(r, theta_x_rad, theta_y_rad):
    x_ = np.tan(theta_x_rad)
    y_ = np.tan(theta_y_rad)
    c = r / np.sqrt(1 + x_**2 + y_**2)

    x = x_ * c
    y = y_ * c
    z = c

    return x, y, z


@njit(cache=True)
def cart2spherical(x, y, z):
    r = np.sqrt(x**2 + y**2 + z**2)
    theta = np.arccos(z / r)
    phi = np.arctan2(-x, y)

    return r, theta, phi


@njit(cache=True)
def spherical2cart(r, theta_rad, phi_rad):
    x = -r * np.sin(theta_rad) * np.sin(phi_rad)
    y = r * np.sin(theta_rad) * np.cos(phi_rad)
    z = r * np.cos(theta_rad)

    return x, y, z

@njit(cache=True)
def projection2spherical(r, theta_x_rad, theta_y_rad):
    x, y, z = projection2cart(r, theta_x_rad, theta_y_rad)
    return cart2spherical(x, y, z)

@njit(cache=True)
def spherical2projection(r, theta_rad, phi_rad):
    x, y, z = spherical2cart(r, theta_rad, phi_rad)
    return cart2projection(x, y, z)

# @njit(cache=True)
# def R(zenith_rad, azimuth_rad):
#     R_z = np.array([[np.cos(azimuth_rad), -np.sin(azimuth_rad), 0.],
#                     [np.sin(azimuth_rad), np.cos(azimuth_rad), 0.],
#                     [0., 0., 1.]], dtype=np.float64)
#
#     R_x = np.array([[1., 0., 0.],
#                     [0., np.cos(zenith_rad), np.sin(zenith_rad)],
#                     [0., -np.sin(zenith_rad), np.cos(zenith_rad)]], dtype=np.float64)
#
#     return R_z @ R_x


@njit(cache=True)
def det2earth(x, y, z, zenith_rad, azimuth_rad):
    x_ = np.cos(azimuth_rad) * x - np.sin(azimuth_rad) * np.cos(zenith_rad) * y - np.sin(azimuth_rad) * np.sin(zenith_rad) * z
    y_ = np.sin(azimuth_rad) * x + np.cos(azimuth_rad) * np.cos(zenith_rad) * y + np.cos(azimuth_rad) * np.sin(zenith_rad) * z
    z_ = -np.sin(zenith_rad) * y + np.cos(zenith_rad) * z

    return x_, y_, z_


@njit(cache=True)
def earth2det(x, y, z, zenith_rad, azimuth_rad):
    x_ = np.cos(azimuth_rad) * x + np.sin(azimuth_rad) * y
    y_ = -np.sin(azimuth_rad) * np.cos(zenith_rad) * x + np.cos(azimuth_rad) * np.cos(zenith_rad) * y - np.sin(zenith_rad) * z
    z_ = -np.sin(azimuth_rad) * np.sin(zenith_rad) * x + np.cos(azimuth_rad) * np.sin(zenith_rad) * y + np.cos(zenith_rad) * z

    return x_, y_, z_
@njit(cache=True)
def det2zenith(theta_x_mrad, theta_y_mrad, zenith_rad, azimuth_rad):
    x, y, z = projection2cart(1, theta_x_mrad / 1000, theta_y_mrad / 1000)
    xe, ye, ze = det2earth(x, y, z, zenith_rad, azimuth_rad)
    _, theta_rad, _ = cart2spherical(xe, ye, ze)
    return theta_rad

def mrad2zenith(angle_deg, theta_rad, phi_rad):
    mrad = int(np.radians(angle_deg) * 1000)
    nx = np.arange(-mrad, mrad + 1)
    ny = np.arange(-mrad, mrad + 1)

    X, Y = np.meshgrid(nx, ny)

    zenith_angle_ = np.zeros(X.shape, dtype=np.float64)
    for i in range(X.shape[0]):
        for j in range(X.shape[1]):
            zenith_angle_[i, j] = det2zenith(X[i, j], Y[i, j], theta_rad, phi_rad)

    return zenith_angle_


# Small epsilon to avoid div-by-zero
EPS = 1e-8


# ------------------------------------------------------------
# cart2projection: (x,y,z) -> (r, theta_x, theta_y)
#   out[0] = r
#   out[1] = theta_x_rad
#   out[2] = theta_y_rad
# ------------------------------------------------------------
@cuda.jit("void(float32, float32, float32, float32[:])",
          device=True, inline=True, fastmath=True)
def cart2projection_device(x, y, z, out):
    r = sqrtf(x * x + y * y + z * z)
    theta_x = atan2f(x, z)  # arctan2(x, z)
    theta_y = atan2f(y, z)  # arctan2(y, z)

    out[0] = r
    out[1] = theta_x
    out[2] = theta_y


# ------------------------------------------------------------
# projection2cart: (r, theta_x, theta_y) -> (x,y,z)
#   out[0] = x
#   out[1] = y
#   out[2] = z
# ------------------------------------------------------------
@cuda.jit("void(float32, float32, float32, float32[:])",
          device=True, inline=True, fastmath=True)
def projection2cart_device(r, theta_x_rad, theta_y_rad, out):
    tx = tanf(theta_x_rad)
    ty = tanf(theta_y_rad)

    denom = sqrtf(1.0 + tx * tx + ty * ty)
    if denom < EPS:
        c = float32(0.0)
    else:
        c = r / denom

    out[0] = tx * c  # x
    out[1] = ty * c  # y
    out[2] = c       # z


# ------------------------------------------------------------
# cart2spherical: (x,y,z) -> (r, theta, phi)
#   theta: polar angle from +z
#   phi:   azimuth, same convention as your CPU version
#   out[0] = r
#   out[1] = theta
#   out[2] = phi
# ------------------------------------------------------------
@cuda.jit("void(float32, float32, float32, float32[:])",
          device=True, inline=True, fastmath=True)
def cart2spherical_device(x, y, z, out):
    r = sqrtf(x * x + y * y + z * z)
    if r < EPS:
        theta = float32(0.0)
    else:
        theta = acosf(z / r)

    # phi = arctan2(-x, y) (your convention)
    phi = atan2f(-x, y)

    out[0] = r
    out[1] = theta
    out[2] = phi


# ------------------------------------------------------------
# spherical2cart: (r, theta, phi) -> (x,y,z)
#   out[0] = x
#   out[1] = y
#   out[2] = z
# ------------------------------------------------------------
@cuda.jit("void(float32, float32, float32, float32[:])",
          device=True, inline=True, fastmath=True)
def spherical2cart_device(r, theta_rad, phi_rad, out):
    st = sinf(theta_rad)
    ct = cosf(theta_rad)
    sp = sinf(phi_rad)
    cp = cosf(phi_rad)

    x = -r * st * sp
    y =  r * st * cp
    z =  r * ct

    out[0] = x
    out[1] = y
    out[2] = z


# ------------------------------------------------------------
# projection2spherical: (r, theta_x, theta_y) -> (r_sph, theta_sph, phi_sph)
#   composition: projection2cart -> cart2spherical
# ------------------------------------------------------------
@cuda.jit("void(float32, float32, float32, float32[:])",
          device=True, inline=True, fastmath=True)
def projection2spherical_device(r, theta_x_rad, theta_y_rad, out):
    tmp_cart = cuda.local.array(3, dtype=float32)
    projection2cart_device(r, theta_x_rad, theta_y_rad, tmp_cart)
    cart2spherical_device(tmp_cart[0], tmp_cart[1], tmp_cart[2], out)
    # out[0] = r_sph, out[1] = theta_sph, out[2] = phi_sph


# ------------------------------------------------------------
# spherical2projection: (r, theta, phi) -> (r_proj, theta_x, theta_y)
#   composition: spherical2cart -> cart2projection
# ------------------------------------------------------------
@cuda.jit("void(float32, float32, float32, float32[:])",
          device=True, inline=True, fastmath=True)
def spherical2projection_device(r, theta_rad, phi_rad, out):
    tmp_cart = cuda.local.array(3, dtype=float32)
    spherical2cart_device(r, theta_rad, phi_rad, tmp_cart)
    cart2projection_device(tmp_cart[0], tmp_cart[1], tmp_cart[2], out)
    # out[0] = r_proj, out[1] = theta_x, out[2] = theta_y


# ------------------------------------------------------------
# det2earth: detector coords -> earth coords
#   (same formulas as your CPU version)
#   out[0] = x_e, out[1] = y_e, out[2] = z_e
# ------------------------------------------------------------
@cuda.jit("void(float32, float32, float32, float32, float32, float32[:])",
          device=True, inline=True, fastmath=True)
def det2earth_device(x, y, z, zenith_rad, azimuth_rad, out):
    ca = cosf(azimuth_rad)
    sa = sinf(azimuth_rad)
    cz = cosf(zenith_rad)
    sz = sinf(zenith_rad)

    x_e = ca * x - sa * cz * y - sa * sz * z
    y_e = sa * x + ca * cz * y + ca * sz * z
    z_e = -sz * y + cz * z

    out[0] = x_e
    out[1] = y_e
    out[2] = z_e


# ------------------------------------------------------------
# earth2det: earth coords -> detector coords
#   out[0] = x_d, out[1] = y_d, out[2] = z_d
# ------------------------------------------------------------
@cuda.jit("void(float32, float32, float32, float32, float32, float32[:])",
          device=True, inline=True, fastmath=True)
def earth2det_device(x, y, z, zenith_rad, azimuth_rad, out):
    ca = cosf(azimuth_rad)
    sa = sinf(azimuth_rad)
    cz = cosf(zenith_rad)
    sz = sinf(zenith_rad)

    x_d = ca * x + sa * y
    y_d = -sa * cz * x + ca * cz * y - sz * z
    z_d = -sa * sz * x + ca * sz * y + cz * z

    out[0] = x_d
    out[1] = y_d
    out[2] = z_d


# ------------------------------------------------------------
# det2zenith: (theta_x_mrad, theta_y_mrad, zenith_rad, azimuth_rad)
#   -> zenith angle in earth frame
#
# This is the GPU analogue of your det2zenith() that internally:
#   1) projection2cart(r=1, theta_x, theta_y)
#   2) det2earth(...)
#   3) cart2spherical(...) and returns theta
# ------------------------------------------------------------
@cuda.jit("float32(float32, float32, float32, float32)",
          device=True, inline=True, fastmath=True)
def det2zenith_device(theta_x_mrad, theta_y_mrad, zenith_rad, azimuth_rad):
    # Convert mrad -> rad (float32)
    thx = theta_x_mrad * 0.001
    thy = theta_y_mrad * 0.001

    # r = 1
    tmp_cart = cuda.local.array(3, dtype=float32)
    tmp_earth = cuda.local.array(3, dtype=float32)
    tmp_sph = cuda.local.array(3, dtype=float32)

    projection2cart_device(float32(1.0), thx, thy, tmp_cart)
    det2earth_device(tmp_cart[0], tmp_cart[1], tmp_cart[2],
                     zenith_rad, azimuth_rad, tmp_earth)
    cart2spherical_device(tmp_earth[0], tmp_earth[1], tmp_earth[2], tmp_sph)

    # theta (zenith) is index 1
    return tmp_sph[1]


if __name__ == "__main__":
    r = 1
    theta_x = np.radians(10)
    theta_y = np.radians(20)

    x, y, z = projection2cart(r, theta_x, theta_y)

    theta = np.radians(0)
    phi = np.radians(0)

    x_, y_, z_ = det2earth(x, y, z, theta, phi)

    zenith, azimuth = cart2spherical(x_, y_, z_)[1:]
    print(np.degrees(zenith), np.degrees(azimuth))
    print(x_)
    print(y_)
    print(z_)

