
using HyperCubicRoots
# Creates HgF as an alias for HypergeometricFunctions._₂F₁
# import HyperCubicRoots: HyperCubicRoots, HgF, solve_real_cubic_roots, solve_all_cubic_roots

import Polynomials
import PolynomialRoots

function test_identity(f1,f2,test_array;tol=10*eps())
    A = f1.(test_array)
    B = f2.(test_array)
    D = abs.(A .- B)
    M = maximum(D)

    return (M < tol), M, argmax(D)
end

aa = 0.5:1e-2:100;

# # Testing on a known case: 
L(z) = log(1+z)/z
G(z) = HgF(1,1,2,-Complex(z))

test_identity(L,G,aa)

## Cubic-related identity:
P(z) = ∛(1 + z) + ∛(1-z);
H(z) = 2HgF(-1/6, 1/3, 1/2, Complex(z)^2)

bb = 0:1e-3:0.999

test_identity(H,P,bb, tol=100*eps())

# Note: we need to ensure that |z| < 1 in the argument of HgF
P(1.01) ≈ H(1.01)  #  isapprox(P(1.01), H(1.01))   false
# (This is done via the Kummer transformation in the cubic roots derivation.)

##
cc = [-72, -36, 0, 1]

p1 = -12
q1 = -36
Del1 = q1^2 + p1^3
s1 = Del1/p1^3

r_num = PolynomialRoots.roots(cc)

r_hyp = [
     2q1/p1*HgF(1/3, 2/3, 1/2, s1),
   -(2q1/3p1)*HgF(2/3, 4/3, 3/2, 1/2*(1 + sqrt(s1))),
   -(2q1/3p1)*HgF(2/3, 4/3, 3/2, 1/2*(1 - sqrt(s1))),
]

r_trig = [
    2*sqrt(12)*cos(  π/18),
    2*sqrt(12)*cos(13π/18),
    2*sqrt(12)*cos(25π/18),
]

# # Leading coefficient is zero:
aa = [-417.47, 843.75, 521.67, 0.0]
try
    rr_hyp = HyperCubicRoots.solve_real_cubic_roots(aa)
catch e
    @warn e
end
    

real_to_tolerance(vec,ϵ) = real.(filter(z -> abs(imag(z)) < ϵ, vec))

# # Three real roots example:
# a = [-23.88, -8.04, 81.02, 13.17]

# # One real roots example:
# a = [81.80, -59.66, -76.47, -64.32]

# # General comparision:
# Compare the roots obtained from PolynomialRoots to this method.

# Notes: the below test compares the numerical roots from PolynomialRoots
# with those generated by the hypergeometric method.
# Experimental observations show that the number of real roots and their
# approximate value are in agreement, 
# when using a tolerance (test_leading_tol) of 1e-5 or larger.
# However, the evaluation of the roots in the polynomial function indicate that the 
# PolynomialRoots solutions are much more accurate:
# p = Polynomials.Polynomial([-590.92, 213.72, -419.35, 274.27])
# e.g. hypergeometric: p(1.7742160152725672) =  0.00263
#     PolynomialRoots: p(1.7742140121531602) = -3.96224e-13
# The trade-off is that the hypergeometric evaluation is about a factor of 2 faster.
# It may be worth "polishing" the root is greater accuracy is required.

test_leading_tol = 1e-5
# Notes:
#  1e-5 is experimentally accurate for Float64 calculations
#  1e-15 works for BigFloat calculations, albeit much slower calculations.

test_lead_zero_tol = 1e-8
# test_lead_zero_tol = big"1e-8"
test_real_tol = 1e-8
N_iter = 100_000
global count = 0
for i in 1:N_iter
    global count
    a = rand(-1000:0.01:1000,4)  # randomly generate the coefficients of the cubic.
    # a = rand(big"-1000":0.01:big"1000",4)  # randomly generate the coefficients of the cubic: BigFloat
    if abs(a[4]) < test_leading_tol
        @info("The leading cubic coefficient is approximately zero. Skipping iteration $(i).")
        println("a =", a)
        continue
    end
    # fa = Polynomials.Polynomial(a)
    R_hyp = sort(HyperCubicRoots.solve_real_cubic_roots(a,leading_tol=test_lead_zero_tol))

    R_num = PolynomialRoots.roots(a)
    R_num_real = sort(real_to_tolerance(R_num,test_real_tol))

    try
        @assert all(isapprox.(sort(R_hyp), sort(R_num_real), atol=test_leading_tol))
    catch e
        @error e
        println(e)
        println("Coefficients: ")
        println(a)
        println("HyperCubicRoots : R_hyp = ", R_hyp)
        println("PolynomialRoots : R_num = ", R_num_real)
        break
    end 

    ## Display examples of computing 3 real roots:
    if mod(i, N_iter ÷ 10)  == 0 || length(R_hyp) == 3
        length(R_hyp) == 3 && count <= 10 ? (count = count + 1) : continue
        println()
        println("Coefficients: ")
        println(a)
        println("HyperCubicRoots:")
        @show(R_hyp)
        println("PolynomialRoots")
        @show R_num
    end

end

function polyeval(aa, root)
    p = Polynomials.Polynomial(aa)
    return p.(root)
end

# # We see that the accuray of the root degenerates as the leading coefficient
# # becomes smaller. Hence we set the accuracy of the method to about 1e-7.
# TYPE = Float16 # errors due to overflow in calculation.
# TYPE = Float32
TYPE = Float64 
# TYPE = BigFloat
for i in 1.0:1:8.0
    x = convert(Vector{TYPE},[-1.,2.1,-3.2,10^-i])
    p = Polynomials.Polynomial(x)
    r = HyperCubicRoots.solve_real_cubic_roots(x, leading_tol=TYPE(1e-16))
    println( (i, 10^-i, r[1]))
    println(p(r[1]))
end

# # The relative magnitudes of the polynomial coefficients plays a role here.
# # The various methods from 
for i in 1.0:1:10.0
    x = [-1.,34.1,-34.2,10^-i]
    p = Polynomials.Polynomial(x)
    r = HyperCubicRoots.solve_real_cubic_roots(x, leading_tol=HyperCubicRoots.CUBIC_ATOL)        
    # r = PolynomialRoots.roots(x)
    # r = PolynomialRoots.solve_cubic_eq(Vector{Complex{Float64}}(x))
    # println(r)
    if !isempty(r)
        println( (i, 10^-i, r[1], p(r[1])))
    end
end

## Timing comparison: 
# # Cubics:
trials = 1_000_000
@time for _ in 1:trials
    a = rand(0.01:0.01:1000.01,4)
    R_hyp = HyperCubicRoots.solve_real_cubic_roots(a)
end

@time for _ in 1:trials
    a = rand(-1000:0.01:1000.01,4)
    R_num = PolynomialRoots.roots(a)
end

# # Quartics:
trials = 1_000_000
@time for _ in 1:trials
    a = rand(0.01:0.01:1000.01,5)
    R_hyp = HyperCubicRoots.solve_all_quartic_roots(a; warn_scaling=false)
end

@time for _ in 1:trials
    a = rand(-1000:0.01:1000.01,5)
    R_num = PolynomialRoots.roots(a)
end

## Comparing our solution functions to PolynomialRoots:
comp_tol = 1e-3  # comparison difference tolerance
t = Vector{Complex{Float64}}()  # vector to hold
r = Vector{Complex{Float64}}()
n = 4  # test n = 2 (quadratic), 3 (cubic), 4 (quartic)
rdvec = []

for i in 1:100_000
    rdvec = rand(-1000:0.01:1000.01,n+1)

    if n == 2
        t = solve_quadratic_roots(rdvec)
    elseif n == 3
        t = solve_all_cubic_roots(rdvec)
    elseif n == 4
        t = solve_all_quartic_roots(rdvec)
    end

    r = PolynomialRoots.roots(rdvec)

    # print("HyperCubicRoots: ")
    # display(t)
    # print("PolynomialRoots: ")
    # display(r)
    # println("")

    t_r = sort(real.(t))
    t_i = sort(imag.(t))
    P_r = sort(real.(r))
    P_i = sort(imag.(r))
    d_r = isapprox.(t_r,P_r; atol=comp_tol)
    d_i = isapprox.(t_i,P_i; atol=comp_tol)

    # if length(t) == 3 && sum(imag.(r)) > 1e-20
    if !all(d_r)
        @warn "Real component mismatch."
        display(t)
        display(r)
        display(t_r .- P_r)
        break
    elseif !all(d_i)
        @warn "Imaginary component mismatch."
        display(t)
        display(r)
        display(t_i .- P_i)
        break
    end
    
end

# # Evaluation comparison:
polyeval(rdvec,solve_all_quartic_roots(rdvec))
polyeval(rdvec,PolynomialRoots.roots(rdvec))

## Fail case:
# # Constant term is zero:
rdvec = [
    0.0
  388.54
 -442.51
  445.24
]

# # Largish error compared to PolynomialRoots:
rdvec = [
    -729.31
    346.23
   -544.86
    721.04
     -0.01
]

## Fail cases for symmetric roots bug:
# # They seem to just be the case of two real roots.
# # We needed to fix the symmetries:

polycoeff = [
    -93.12
    -738.32
    -477.59
    237.14
    487.27
]
polycoeff = [
  936.22
  165.74
  749.85
  229.61
 -553.5
 ]

t = solve_all_quartic_roots(polycoeff)
r = PolynomialRoots.roots(polycoeff)

t_r = sort(real.(t))
t_i = sort(imag.(t))
P_r = sort(real.(r))
P_i = sort(imag.(r))
d_r = isapprox.(t_r,P_r; atol=comp_tol)
d_i = isapprox.(t_i,P_i; atol=comp_tol)

t_r .- P_r
t_i .- P_i

a = polycoeff[1:4]./polycoeff[5]
-a[4]/4

warn_scaling = true; leading_tol=1e-8; coeff_tol=1.0/leading_tol;

rd = []
h =[]
println("Start\n")
for _ in 1:10
    rd = rand(0.01:0.01:100.01,5)
    h = solve_all_quartic_roots(rd; leading_tol=1e-8, coeff_tol=1e10)
    if !isapprox(imag(h), zero(imag(h))) && real(h[1]) > 0
        display(h)
    end
end
