__author__ = 'ewoudwempe'

import numpy as np
from lenstronomy.Util.numba_util import jit


@jit()
def min_approx(x1, x2, x3, y1, y2, y3):
    """
    Get the x-value of the minimum of the parabola through the points (x1,y1), ...
    :param x1: x-coordinate point 1
    :param x2: x-coordinate point 2
    :param x3: x-coordinate point 3
    :param y1: y-coordinate point 1
    :param y2: y-coordinate point 2
    :param y3: y-coordinate point 3
    :return: x-location of the minimum
    """
    #
    div = (2.*(x3*(y1 - y2) + x1*(y2 - y3) + x2*(-y1 + y3)))
    return (x3**2*(y1 - y2) + x1**2*(y2 - y3) + x2**2*(-y1 + y3))/div


@jit()
def rotmat(th):
    """
    Calculates the rotation matrix
    :param th: angle
    :return: rotation matrix
    """
    return np.array([[np.cos(th), np.sin(th)], [-np.sin(th), np.cos(th)]])


@jit()
def cdot(a, b):
    """
    Calculates some complex dot-product that simplifies the math
    :param a: complex number
    :param b: complex number
    :return: dot-product
    """
    return a.real*b.real + a.imag*b.imag


@jit()
def ps(x, p):
    """
    A regularized power-law that gets rid of singularities, abs(x)**p*sign(x)
    :param x: x
    :param p: p
    :return:
    """
    return np.abs(x)**p*np.sign(x)


@jit()
def cart_to_pol(x, y):
    """
    Convert from cartesian to polar
    :param x: x-coordinate
    :param y: y-coordinate
    :return: tuple of (r, theta)
    """
    return np.sqrt(x**2+y**2), np.arctan2(y, x) % (2*np.pi)


@jit()
def pol_to_cart(r, th):
    """
    Convert from polar to cartesian
    :param r: r-coordinate
    :param th: theta-coordinate
    :return: tuple of (x,y)
    """
    return r*np.cos(th), r*np.sin(th)


@jit()
def pol_to_ell(r, theta, q):
    """Converts from polar to elliptical coordinates"""
    phi = np.arctan2(np.sin(theta), np.cos(theta)*q)
    rell = r*np.sqrt(q**2*np.cos(theta)**2+np.sin(theta)**2)
    return rell, phi


@jit()
def ell_to_pol(rell, theta, q):
    """Converts from elliptical to polar coordinates"""
    phi = np.arctan2(np.sin(theta)*q, np.cos(theta))
    r = rell*np.sqrt(1/q**2*np.cos(theta)**2+np.sin(theta)**2)
    return r, phi


def geomlinspace(a, b, N):
    """Constructs a geomspace from a to b, with a linspace prepended to it from 0 to a, with the same spacing as the
    geomspace would have at a"""
    delta = a*((b/a)**(1/(N-1))-1)
    return np.concatenate((np.linspace(0, a, int(a/delta), endpoint=False), np.geomspace(a, b, N)))


@jit()
def solvequadeq(a, b, c):
    """
    Solves a quadratic equation. Care is taken for the numerics, see also https://en.wikipedia.org/wiki/Loss_of_significance
    :param a: a
    :param b: b
    :param c: c
    :return: tuple of two solutions
    """
    sD = (b**2-4*a*c)**0.5
    x1 = (-b-np.sign(b)*sD)/(2*a)
    x2 = 2*c/(-b-np.sign(b)*sD)
    return np.where(b != 0, np.where(a != 0, x1, -c/b), -(-c/a)**0.5), \
                    np.where(b != 0, np.where(a != 0, x2, -c/b+1e-8), +(-c/a)**0.5)


def brentq_nojit(f, xa, xb, xtol=2e-14, rtol=16*np.finfo(float).eps, maxiter=100, args=()):
    """
    A numba-compatible implementation of brentq (largely copied from scipy.optimize.brentq).
    Unfortunately, the scipy verison is not compatible with numba, hence this reimplementation :(
    :param f: function to optimize
    :param xa: left bound
    :param xb: right bound
    :param xtol: x-coord root tolerance
    :param rtol: x-coord relative tolerance
    :param maxiter: maximum num of iterations
    :param args: additional arguments to pass to function in the form f(x, args)
    :return: 
    """
    xpre = xa
    xcur = xb
    xblk = 0.
    fblk = 0.
    spre = 0.
    scur = 0.
    fpre = f(xpre, args)
    fcur = f(xcur, args)
    funcalls = 2
    if fpre*fcur > 0:
        raise ValueError('Signs are not different')
    if fpre == 0:
        return xpre
    if fcur == 0:
        return xcur
    iterations = 0
    for i in range(maxiter):
        iterations += 1
        if fpre*fcur < 0:
            xblk = xpre
            fblk = fpre
            # spres = scur = xcur - xpre
        if abs(fblk) < abs(fcur):
            xpre = xcur
            xcur = xblk
            xblk = xpre
            
            fpre = fcur
            fcur = fblk
            fblk = fpre
            
        delta = (xtol + rtol*abs(xcur))/2
        sbis = (xblk - xcur)/2
        if fcur == 0 or abs(sbis) < delta:
            return xcur

        if abs(spre) > delta and abs(fcur) < abs(fpre):
            if xpre == xblk:
                stry = -fcur*(xcur - xpre)/(fcur - fpre)
            else:
                dpre = (fpre - fcur)/(xpre - xcur)
                dblk = (fblk - fcur)/(xblk - xcur)
                stry = -fcur*(fblk*dblk - fpre*dpre)\
                    / (dblk*dpre*(fblk - fpre))
            
            if 2*abs(stry) < min(abs(spre), 3*abs(sbis) - delta):
                spre = scur
                scur = stry
            else:
                spre = sbis
                scur = sbis
        else:
            spre = sbis
            scur = sbis

        xpre = xcur
        fpre = fcur
        if abs(scur) > delta:
            xcur += scur
        else:
            xcur += delta if sbis > 0 else -delta

        fcur = f(xcur, args)
        funcalls += 1
    
    return xcur


brentq_inline = jit(inline='always')(brentq_nojit)
