# Implementation of TOMS748
# https://na.math.kit.edu/alefeld/download/1995_Algorithm_748_Enclosing_Zeros_of_Continuous_Functions.pdf

# from math import sin, fabs, cos
from math import *

# checks that none of the values are the same
# returns true if at least two values are identical
def checkTwoValuesIdentical(fa, fb, fc, fd):
    same = False
    same |= fa == fb
    same |= fa == fc
    same |= fa == fd

    same |= fb == fc
    same |= fb == fd

    same |= fc == fd
    return same

# finite differences 1st derivative given a, b, fa, fb
def fbracket1(a, b, fa, fb):
    return (fb-fa)/(b-a)

# finite differences 2nd derivative given a, b, d, fa, fb, fd
def fbracket2(a, b, d, fa, fb, fd):
    return (fbracket1(b, d, fb, fd) - fbracket1(a, b, fa, fb))/(d-a)

# standard bracketing routine
# returns ahat, bhat, d
# d is a point outside the new interval
def bracket(a, b, c, fa, fb, fc):
    if fa*fc < 0:
        return (a, c, b, fa, fc, fb)
    else:
        return (c, b, a, fc, fb, fa)

# standard bracketing routine
# returns ahat, bhat, d
# d is a point outside the new interval
def bracketAndCheckConvergence(a, b, c, fa, fb, fc, tol):
    ok = False
    if fabs(fc) < tol:
        ok = True
    if fa*fc < 0:
        return (a, c, b, fa, fc, fb, ok)
    else:
        return (c, b, a, fc, fb, fa, ok)

# finds an approximate solution to the quadratic P(x) = fa + f[a,b]*(x-a) + f[a,b,d]*(x-a)(x-b)
# with f[a,b] = fbracket1(a,b) and f[a,b,d] = fbracket2(a,b,d)
def NewtonQuadratic(a, b, d, fa, fb, fd, k):
    A = fbracket2(a,b,d,fa,fb,fd)
    B = fbracket1(a,b,fa,fb)
    if A == 0:
        return a - fa/B
    if A*fa > 0:
        r = a
    else:
        r = b

    for i in range(k):
        r = r - B*r/(B + A*(2*r - a - b))
    return r

# Inverse cubic interpolation evaluated at 0 (modified Aitken-Neville interpolation)
def ipzero(a, b, c, d, fa, fb, fc, fd):
    Q11 = (c-d)*fc/(fd-fc)
    Q21 = (b-c)*fb/(fc-fb)
    Q31 = (a-b)*fa/(fb-fa)
    D21 = (b-c)*fc/(fc-fb)
    D31 = (a-b)*fb/(fb-fa)
    Q22 = (D21-Q11)*fb/(fd-fb)
    Q32 = (D31-Q21)*fa/(fc-fa)
    D32 = (D31-Q21)*fc/(fc-fa)
    Q33 = (D32-Q22)*fa/(fd-fa)
    return a + Q31 + Q32 + Q33

def TOMS748_solve1(f, a, b, tol, Nmax = 1000):
    mu = 0.5
    fa, fb = f(a), f(b)
    c = a - fa/fbracket1(a, b, fa, fb)                              # 4.1.1 secant method
    fc = f(c)
    a, b, d, fa, fb, fd, ok = bracketAndCheckConvergence(a, b, c, fa, fb, fc, tol)              # 4.1.2
    if(ok): return (c, fc)
    e = d
    fe = fd
    # ---
    for n in range(2,Nmax):                                           # 4.1.3
        if n == 2 or checkTwoValuesIdentical(fa, fb, fd, fe):
            c = NewtonQuadratic(a, b, d, fa, fb, fd, 2)
        else:
            c = ipzero(a, b, d, e, fa, fb, fd, fe)
            if (c-a)*(c-b) >= 0:
                c = NewtonQuadratic(a, b, d, fa, fb, fd, 2)
        # ---
        fc = f(c)
        a, b, dbar, fa, fb, fdbar, ok = bracketAndCheckConvergence(a, b, c, fa, fb, fc, tol)    # 4.1.4
        if(ok): return (c, fc)
        # ---
        if fabs(fa) < fabs(fb):                                     # 4.1.5
            u = a
            fu = fa
        else:
            u = b
            fu = fb
        # ---
        c = u - 2*fu/fbracket1(a, b, fa, fb)                        # 4.1.6
        # ---
        if fabs(c - u) > 0.5*(b - a):                               # 4.1.7
            c = 0.5*(b + a)
        # ---
        fc = f(c)
        a, b, dhat, fa, fb, fdhat, ok = bracketAndCheckConvergence(a, b, c, fa, fb, fc, tol)    # 4.1.8
        if(ok): return (c, fc)
        # ---
        if b - a < mu*(b - a):                                      # 4.1.9
            d = dhat
            e = dbar
            fd = fdhat
            fe = fdbar
        else:
            e = dhat
            fe = fdhat
        c = 0.5*(a+b)
        fc = f(c)
        a, b, d, fa, fb, fd, ok = bracketAndCheckConvergence(a, b, c, fa, fb, fc, tol)
        if(ok): return (c, fc)
    return (c, fc)# no solution found, return last estimate

def TOMS748_solve2(f, a, b, tol, Nmax = 1000):
    mu = 0.5
    fa, fb = f(a), f(b)
    c = a - fa/fbracket1(a, b, fa, fb)                              # 4.2.1 secant method
    fc = f(c)
    a, b, d, fa, fb, fd, ok = bracketAndCheckConvergence(a, b, c, fa, fb, fc, tol)              # 4.2.2
    if(ok): return (c, fc)
    e = d
    fe = fd
    # ---
    for n in range(2,Nmax):                                           # 4.2.3
        if n == 2 or checkTwoValuesIdentical(fa, fb, fd, fe):
            c = NewtonQuadratic(a, b, d, fa, fb, fd, 2)
        else:
            c = ipzero(a, b, d, e, fa, fb, fd, fe)
            if (c-a)*(c-b) >= 0:
                c = NewtonQuadratic(a, b, d, fa, fb, fd, 2)
        # ---
        e = d
        fe = fd
        fc = f(c)
        a, b, d, fa, fb, fd, ok = bracketAndCheckConvergence(a, b, c, fa, fb, fc, tol)          # 4.2.4
        if(ok): return (c, fc)
        # ---
        if checkTwoValuesIdentical(fa, fb, fd, fe):                      # 4.2.5
            c = NewtonQuadratic(a, b, d, fa, fb, fd, 3)
        else:
            c = ipzero(a, b, d, e, fa, fb, fd, fe)
            if (c-a)*(c-b) >= 0:
                c = NewtonQuadratic(a, b, d, fa, fb, fd, 3)
        # ---
        fc = f(c)
        a, b, dbar, fa, fb, fdbar, ok = bracketAndCheckConvergence(a, b, c, fa, fb, fc, tol)    # 4.2.6
        if(ok): return (c, fc)
        # ---
        if fabs(fa) < fabs(fb):                                     # 4.2.7
            u = a
            fu = fa
        else:
            u = b
            fu = fb
        # ---
        c = u - 2*fu/fbracket1(a, b, fa, fb)                        # 4.2.8
        # ---
        if fabs(c - u) > 0.5*(b - a):                               # 4.2.9
            c = 0.5*(b + a)
        # ---
        fc = f(c)
        a, b, dhat, fa, fb, fdhat, ok = bracketAndCheckConvergence(a, b, c, fa, fb, fc, tol)    # 4.2.10
        if(ok): return (c, fc)
        # ---
        if b - a < mu*(b - a):                                      # 4.2.11
            d = dhat
            e = dbar
            fd = fdhat
            fe = fdbar
        else:
            e = dhat
            fe = fdhat
        c = 0.5*(a+b)
        fc = f(c)
        a, b, d, fa, fb, fd, ok = bracketAndCheckConvergence(a, b, c, fa, fb, fc, tol)
        if(ok): return (c, fc)
    return (c, fc)# no solution found, return last estimate

# TEST
f = lambda E : E - 0.5*sin(E) - 0.3
a, b, c, d = 0, 1, 0.3, 1.2
fa, fb, fc, fd = f(a), f(b), f(c), f(d)
tol = 1e-12

print(fbracket1(a, b, fa, fb))
print(fbracket2(a, b, d, fa, fb, fd))
print(bracket(a, b, c, fa, fb, fc))
print(NewtonQuadratic(a, b, d, fa, fb, fd, 2))
print(ipzero(a, b, c, d, fa, fb, fc, fd))

# TTTEST
print(checkTwoValuesIdentical(1,2,3,4))
print(checkTwoValuesIdentical(1,2,3,1))
print(checkTwoValuesIdentical(1,1,3,4))
print(checkTwoValuesIdentical(1,2,2,4))
print(checkTwoValuesIdentical(1,2,3,3))
print(checkTwoValuesIdentical(1,2,1,4))
print(checkTwoValuesIdentical(1,2,3,2))
# TTTEST

# solution using bracket only :
print('bracket only')
for i in range(50):
    c = (a+b)/2
    fc = f(c)
    a, b, d, fa, fb, fd = bracket(a, b, c, fa, fb, fc)
    if fabs(b-a) < tol:
        break
    print(i)

print((a,b), (a+b)/2, b-a)

print('Newton with derivative')
df = lambda E : 1 - 0.5*cos(E)
x = 0.5
for i in range(50):
    delta = f(x)/df(x)
    x = x - delta
    if fabs(delta) < tol:
        break
    print(i)
print(x)

print('TOMS748_solve1 only')
def fct1(E):
    print("x = %20.15f" % E)
    return E - 0.5*sin(E) - 0.3

c, fc = TOMS748_solve1(fct1, 0, 1, tol)
print("Root found :")
print((c, fc))

print('TOMS748_solve2 only')
c, fc = TOMS748_solve2(fct1, 0, 1, tol)
print("Root found :")
print((c, fc))

# ---

print('TOMS748_solve1 only')
def fct2(x):
    print("x = %20.15f" % x)
    n = 20
    return x**2 - (1-x)**n

c, fc = TOMS748_solve1(fct2, 0, 1, tol)
print("Root found :")
print((c, fc))

print('TOMS748_solve2 only')
c, fc = TOMS748_solve2(fct2, 1e-100, 1, tol)# if a is 0 the algorithm does not work (division by zero, because newton step a = 0 finds d = 0). It seems to be just this one case...
print("Root found :")
print((c, fc))

print('Newton with derivative')
dfct2 = lambda x : 2*x + 20*(-x + 1)**19
x = 0.5
for i in range(50):
    delta = fct2(x)/dfct2(x)
    x = x - delta
    if fabs(delta) < tol:
        break
    print(i)
print(x)

# ---

print('--------------------------')
print('fct3')
print('TOMS748_solve1 only')
def fct3(x):
    print("x = %20.15f" % x)
    return x*exp(-x**2)
# a, b = -.5, 1.5
a, b = -1/sqrt(2), 2/sqrt(2)

c, fc = TOMS748_solve1(fct3, a, b, tol)
print("Root found :")
print((c, fc))

print('TOMS748_solve2 only')
c, fc = TOMS748_solve2(fct3, a, b, tol)# if a is 0 the algorithm does not work (division by zero, because newton step a = 0 finds d = 0). It seems to be just this one case...
print("Root found :")
print((c, fc))

print('Newton with derivative')
dfct3 = lambda x : (-2*x**2 + 1)*exp(-x**2)
x = (a+b)/2
for i in range(50):
    delta = fct3(x)/dfct3(x)
    x = x - delta
    if fabs(delta) < tol:
        break
    print(i)
print(x)