using SymPy
using BenchmarkTools
using Plots
VERSION

# https://github.com/mitmath/18S096/blob/iap2017/pset3/pset3-solutions.ipynb

# SOLUTION code
# n coefficients of the Taylor series of E₁(z) + log(z), in type T:
function E₁_taylor_coefficients(::Type{T}, n::Integer) where T<:Number
    n < 0 && throw(ArgumentError("$n ≥ 0 is required"))
    n == 0 && return T[]
    n == 1 && return T[-MathConstants.eulergamma]
    # iteratively compute the terms in the series, starting with k=1
    term::T = 1
    terms = T[-MathConstants.eulergamma, term]
    for k=2:n
        term = -term * (k-1) / (k * k)
        push!(terms, term)
    end
    return terms
end

# SOLUTION code
macro E₁_taylor64(z, n::Integer)
    c = E₁_taylor_coefficients(Float64, n)
    taylor = Expr(:macrocall, Symbol("@evalpoly"),
        Base.LineNumberNode(@__LINE__, Symbol(@__FILE__)),
        :t, c...)
    quote
        let t = $(esc(z))
            $taylor - log(t)
        end
    end
end

# compute E₁ via n terms of the continued-fraction expansion, implemented
# in the simplest way:
function E₁_cf(z::Number, n::Integer)
    # starting with z seems to give many fewer terms for intermediate |z| ~ 3
    cf::typeof(inv(z)) = z
    for i = n:-1:1
        cf = z + (1+i)/cf
        cf = 1 + i/cf
    end
    return exp(-z) / (z + inv(cf))
end

# SOLUTION code
# for numeric-literal coefficients: simplify to a ratio of two polynomials:
import Polynomials
# return (p,q): the polynomials p(x) / q(x) corresponding to E₁_cf(x, a...),
# but without the exp(-x) term
function E₁_cfpoly(n::Integer, ::Type{T}=BigInt) where T<:Real
    q = Polynomials.PolyCompat.Poly(T[1])
    p = x = Polynomials.PolyCompat.Poly(T[0,1])
    for i = n:-1:1
        p, q = x*p+(1+i)*q, p # from cf = x + (1+i)/cf = x + (1+i)*q/p
        p, q = p + i*q, p     # from cf = 1 + i/cf = 1 + i*q/p
    end
    # do final 1/(x + inv(cf)) = 1/(x + q/p) = p/(x*p + q)
    return p, x*p + q
end
macro E₁_cf64(x, n::Integer)
    p,q = E₁_cfpoly(n, BigInt)
    evalpoly = Symbol("@evalpoly")
    num_expr = Expr(:macrocall, evalpoly, 
        Base.LineNumberNode(@__LINE__, Symbol(@__FILE__)),
        :t, Float64.(Polynomials.coeffs(p))...)
    den_expr = Expr(:macrocall, evalpoly,
        Base.LineNumberNode(@__LINE__, Symbol(@__FILE__)),
        :t, Float64.(Polynomials.coeffs(q))...)
    quote
        let t = $(esc(x))
            exp(-t) * $num_expr / $den_expr
        end
    end
end

# SOLUTION:
function E₁(z::Union{Float64,Complex{Float64}})
    x² = real(z)^2
    y² = imag(z)^2
    if x² + 0.233*y² ≥ 7.84 # use cf expansion, ≤ 30 terms
        if (x² ≥ 546121) & (real(z) > 0) # underflow
            return zero(z)
        elseif x² + 0.401*y² ≥ 58.0 # ≤ 15 terms
            if x² + 0.649*y² ≥ 540.0 # ≤ 8 terms
                x² + y² ≥ 4e4 && return @E₁_cf64 z 4
                return @E₁_cf64 z 8
            end
            return @E₁_cf64 z 15
        end
        return @E₁_cf64 z 30
    else # use Taylor expansion, ≤ 37 terms
        r² = x² + y²
        return r² ≤ 0.36 ? (r² ≤ 2.8e-3 ? (r² ≤ 2e-7 ? @E₁_taylor64(z,4) :
                                                       @E₁_taylor64(z,8)) :
                                         @E₁_taylor64(z,15)) :
                          @E₁_taylor64(z,37)
    end
end
E₁(z::Union{T,Complex{T},Rational{T},Complex{Rational{T}}}) where T<:Integer = E₁(float(z))

@vars z
E₁_cf(z, 4)

E₁_cf(z, 4).simplify()

@E₁_cf64(z, 4)

p, q = E₁_cfpoly(30, BigInt)
display(p)
display(q)

@vars z
print("p(z) = ")
show(@evalpoly z Float64.(Polynomials.coeffs(p))...)
print("\n\nq(z) = ")
show(@evalpoly z Float64.(Polynomials.coeffs(q))...)

r = Polynomials.PolyCompat.Poly(E₁_taylor_coefficients(Float64, 37))

@vars z
print("r(z) = ")
show(@evalpoly z Float64.(Polynomials.coeffs(r))...)

# https://github.com/uncorrelated/ExpIntCF/blob/master/e1_cf.jl

function E₁_cf_uc(z::Number, n::Integer)
    cf::typeof(z) = z
    for i = n:-1:1
        cf = z + (1+i)/cf
        cf = 1 + i/cf
    end
    return exp(-z) / (z + inv(cf))
end

function E₁_cf_uc(z::Number, reltol=1e-12)
    for n = 1:1000
        s = E₁_cf_uc(z, n)
        d = E₁_cf_uc(z, 2n)
        if abs(s - d) <= reltol*abs(d)
            return d
        end
    end
    error("iteration limit exceeded!")
end

# 再掲
# https://nbviewer.jupyter.org/github/stevengj/18S096/blob/iap2017/pset3/pset3-solutions.ipynb
# より

# compute E₁ via n terms of the continued-fraction expansion, implemented
# in the simplest way:
function E₁_cf(z::Number, n::Integer)
    # starting with z seems to give many fewer terms for intermediate |z| ~ 3
    cf::typeof(inv(z)) = z
    for i = n:-1:1
        cf = z + (1+i)/cf
        cf = 1 + i/cf
    end
    return exp(-z) / (z + inv(cf))
end

eps()

1e-16 < eps()

E₁_cf_uc(3, 4) # by uncorrelated

E₁_cf(3, 4) # solution of MIT homework

# 次の函数は
# https://github.com/uncorrelated/ExpIntCF/blob/master/e1_cf.jl
# より. 函数名に `_uc` を付け加えてある
function MakeMatrix_uc(s::Number, e::Number, length::Integer)
    x = range(s, e, length=length)
    return [E₁_cf_uc(x+y*im) for y in x, x in x]
end

M_uc = MakeMatrix_uc(0.1, 5, 100)
print("uncorrelated's code:")
@btime MakeMatrix_uc(0.1, 5, 100);

function MakeMatrix_mit(s::Number, e::Number, length::Integer)
    x = range(s, e, length=length)
    return [E₁(x+y*im) for y in x, x in x]
end

M_mit = MakeMatrix_mit(0.1, 5, 100)
print("MIT homework solution:")
@btime MakeMatrix_mit(0.1, 5, 100);

max_relerr = maximum(@. abs(M_uc/M_mit - 1))
@show max_relerr
x = y = range(0.1, 5, length=100)
heatmap(x, y, log10.(abs.(M_uc./M_mit .- 1)); size=(500, 450))

N_uc = MakeMatrix_uc(3, 20, 100)
print("uncorrelated's code:")
@btime MakeMatrix_uc(3, 20, 100);

N_mit = MakeMatrix_mit(3, 20, 100)
print("MIT homework solution:")
@btime MakeMatrix_mit(3, 20, 100);

max_relerr = maximum(@. abs(N_uc/N_mit - 1))
@show max_relerr
x = y = range(0.1, 5, length=100)
heatmap(x, y, log10.(abs.(N_uc./N_mit .- 1)); size=(500, 450))