import tensorflow as tf

from . import math as tk
from .utils import get_deriv_func, extend, find_zero


def sign_bessel_k(v, x, m=0, n=0, name=None):
    with tf.name_scope(name or 'sign_bessel_k_tk2'):
        dtype = tk.common_dtype([v, x])
        x = tf.convert_to_tensor(x, dtype)
        v = tf.convert_to_tensor(v, dtype)
        m = tf.convert_to_tensor(m, tf.int32)
        n = tf.convert_to_tensor(n, tf.int32)
        return _sign_bessel_k_naive(v, x, m, n)


def log_bessel_k(v, x, m=0, n=0, name=None):
    with tf.name_scope(name or 'log_bessel_k_tk2'):
        dtype = tk.common_dtype([v, x])
        x = tf.convert_to_tensor(x, dtype)
        v = tf.convert_to_tensor(v, dtype)
        m = tf.convert_to_tensor(m, tf.int32)
        n = tf.convert_to_tensor(n, tf.int32)
        return _log_bessel_k_custom_gradient(v, x, m, n)


def slog_bessel_k(v, x, m=0, n=0, name=None):
    with tf.name_scope(name or 'slog_bessel_k_tk2'):
        dtype = tk.common_dtype([v, x])
        x = tf.convert_to_tensor(x, dtype)
        v = tf.convert_to_tensor(v, dtype)
        m = tf.convert_to_tensor(m, tf.int32)
        n = tf.convert_to_tensor(n, tf.int32)
        sign = _sign_bessel_k_naive(v, x, m, n)
        logk = _log_bessel_k_custom_gradient(v, x, m, n)
        return sign, logk


def bessel_ke(v, x, m=0, n=0, name=None):
    with tf.name_scope(name or 'bessel_ke_tk2'):
        dtype = tk.common_dtype([v, x])
        x = tf.convert_to_tensor(x, dtype)
        v = tf.convert_to_tensor(v, dtype)
        m = tf.convert_to_tensor(m, tf.int32)
        n = tf.convert_to_tensor(n, tf.int32)
        sign = _sign_bessel_k_naive(v, x, m, n)
        logk = _log_bessel_k_custom_gradient(v, x, m, n)
        return sign * tk.exp(logk + x)


def bessel_k_ratio(v, x, d=1, m=0, n=0, name=None):
    with tf.name_scope(name or 'bessel_k_ratio_tk2'):
        dtype = tk.common_dtype([v, x])
        x = tf.convert_to_tensor(x, dtype)
        v = tf.convert_to_tensor(v, dtype)
        d = tf.convert_to_tensor(d, dtype)
        m = tf.convert_to_tensor(m, tf.int32)
        n = tf.convert_to_tensor(n, tf.int32)
        signd = _sign_bessel_k_naive(v + d, x, m, n)
        logkd = _log_bessel_k_custom_gradient(v + d, x, m, n)
        sign = _sign_bessel_k_naive(v, x, m, n)
        logk = _log_bessel_k_custom_gradient(v, x, m, n)
        return signd * sign * tf.exp(logkd - logk)


@tf.custom_gradient
def _log_bessel_k_custom_gradient(v, x, m, n):
    def _grad(u):
        sign = _sign_bessel_k_naive(v, x, m, n)
        sign_dv = _sign_bessel_k_naive(v, x, m + 1, n)
        logk_dv = _log_bessel_k_custom_gradient(v, x, m + 1, n)
        sign_dx = _sign_bessel_k_naive(v, x, m, n + 1)
        logk_dx = _log_bessel_k_custom_gradient(v, x, m, n + 1)
        return (
            u * sign_dv * sign * tk.exp(logk_dv - logk),
            u * sign_dx * sign * tk.exp(logk_dx - logk),
            None, None,
        )
    logk = _log_bessel_k_naive(v, x, m, n)
    return logk, _grad


def _sign_bessel_k_naive(v, x, m, n):
    dtype = tk.common_dtype([v, x])
    convv = tf.equal(m % 2, 1) & (v < 0)
    convx = tf.equal(n % 2, 1)
    return tf.where(convv ^ convx, tf.cast(-1, dtype), tf.cast(1, dtype))


def _log_bessel_k_naive(
    v, x, m, n, mask=None, dt0=1., tol=1., max_iter=10, bins=128):

    def func(t):
        out = tf.where(
            tf.equal(m % 2, 0), tk.log_cosh(v * t), tk.log_sinh(v * t),
        )
        out -= x * tk.cosh(t)
        out += tf.where(m > 0, mf * tk.log(t), tf.cast(0, dtype))
        out += tf.where(n > 0, nf * tk.log_cosh(t), tf.cast(0, dtype))
        return out

    def func_mth(t):
        return func(t) - th

    dtype = tk.common_dtype([v, x])
    shape = tf.shape(v * x)

    v = tk.abs(v)
    mf = tf.cast(m, dtype)
    nf = tf.cast(n, dtype)

    if mask is None:
        mask = tf.ones(shape, tf.bool)
    condzero = tf.equal(m % 2, 1) & tf.equal(v, 0)
    condinf = tf.equal(x, 0)
    condnan = x < 0
    mask &= ~condzero & ~condinf & ~condnan

    zero_peak = tf.equal(m, 0) & (tf.square(v) + mf < x)
    zero_inf = m > 0

    eps = tk.epsilon(dtype)
    scale = tf.cast(dt0, dtype)
    deriv = get_deriv_func(func)

    dt = tf.where(~zero_peak & mask, scale, tf.cast(0, dtype))
    ts, te = extend(deriv, tf.zeros(shape, dtype), dt)
    tp = find_zero(deriv, te, ts, tol, max_iter)
    fp = func(tp)

    th = fp + tk.log(eps)
    tpm = tk.minimum(tp - bins * eps, tf.cast(0, dtype))
    t0 = find_zero(func_mth, tf.zeros(shape, dtype), tpm, tol, max_iter)

    tpp = tk.maximum(tp + bins * eps, tp * (1 + bins * eps))
    zero_exists = (func_mth(tpp) > 0) & mask
    dt = tf.where(zero_exists, dt0 * scale, tf.cast(0, dtype))
    ts, te = extend(func_mth, tpp, dt)
    t1 = find_zero(func_mth, te, ts, tol, max_iter)

    t = tf.linspace(t0, t1, 2 * bins + 1, axis=0)[1:-1:2]
    h = (t1 - t0) / bins
    sum_eft = h * tk.sum(tk.exp(func(t) - fp), axis=0)
    out = fp + tk.log(sum_eft)

    out = tf.where(condzero, tf.cast(-tk.inf, dtype), out)
    out = tf.where(condinf, tf.cast(tk.inf, dtype), out)
    out = tf.where(condnan, tf.cast(tk.nan, dtype), out)
    return out
