using Base.MathConstants
using Base64
using Printf
using Statistics
const e = ℯ
endof(a) = lastindex(a)
linspace(start, stop, length) = range(start, stop, length=length)

using Plots
#gr(); ENV["PLOTS_TEST"] = "true"
#pyplot(fmt=:svg)
pyplot()
#clibrary(:colorcet)
#clibrary(:misc)
default(fmt=:png)

function pngplot(P...; kwargs...)
    sleep(0.1)
    pngfile = tempname() * ".png"
    savefig(plot(P...; kwargs...), pngfile)
    showimg("image/png", pngfile)
end
pngplot(; kwargs...) = pngplot(plot!(; kwargs...))

showimg(mime, fn) = open(fn) do f
    base64 = base64encode(f)
    display("text/html", """<img src="data:$mime;base64,$base64">""")
end

using SymPy
#sympy.init_printing(order="lex") # default
#sympy.init_printing(order="rev-lex")

using SpecialFunctions
lgamma(x::Real) = logabsgamma(x)[1]

using QuadGK
using PyPlot: PyPlot, plt

# Override the Base.show definition of SymPy.jl:
# https://github.com/JuliaPy/SymPy.jl/blob/29c5bfd1d10ac53014fa7fef468bc8deccadc2fc/src/types.jl#L87-L105

@eval SymPy function Base.show(io::IO, ::MIME"text/latex", x::SymbolicObject)
    print(io, as_markdown("\\displaystyle " * sympy.latex(x, mode="plain", fold_short_frac=false)))
end
@eval SymPy function Base.show(io::IO, ::MIME"text/latex", x::AbstractArray{Sym})
    function toeqnarray(x::Vector{Sym})
        a = join(["\\displaystyle " * sympy.latex(x[i]) for i in 1:length(x)], "\\\\")
        """\\left[ \\begin{array}{r}$a\\end{array} \\right]"""
    end
    function toeqnarray(x::AbstractArray{Sym,2})
        sz = size(x)
        a = join([join("\\displaystyle " .* map(sympy.latex, x[i,:]), "&") for i in 1:sz[1]], "\\\\")
        "\\left[ \\begin{array}{" * repeat("r",sz[2]) * "}" * a * "\\end{array}\\right]"
    end
    print(io, as_markdown(toeqnarray(x)))
end

N = 10^8
SimpleSum = sum(n->1/n^2, 1:N)
TrueValue = π^2/6
@show SimpleSum
@show TrueValue
@show SimpleSum - TrueValue;

N = 50
sum(n->zeta(n)-1, 2:N)

l, n = symbols("l n", integer=true, positive=true) # l = k-1
s = sympy.Sum(1/(l+1)^n, (n,2,oo)).doit().simplify()
S = sympy.Sum(s, (l,1,oo)).doit()

N = 50
sum(n->(-1)^n*(zeta(n)-1), 2:N)

l, n = symbols("l n", integer=true, positive=true) # l = k-1
s = sympy.Sum((-1)^n/(l+1)^n, (n,2,oo)).doit().simplify()
S = sympy.Sum(s, (l,1,oo)).doit()

@time s = [sum(1/n for n in 10^(k-1):10^k-1 if !in('9', string(n))) for k in 1:8]

s[2:end] ./ s[1:end-1]

S = cumsum(s)

8 + 23//100 + (1 + 64//100)*(9//10)/(1 - 9//10)

n = symbols("n", integer=true)
S = sympy.Sum((-1)^(n-1)/n, (n, 1, oo)).doit()

N = 1000
SimpleSum = sum(n->(-1)^(n-1)/n, 1:N)
TrueValue = log(2)
@show SimpleSum
@show TrueValue
@show SimpleSum - TrueValue;

binom(n,k) = exp(lgamma(n+1) - lgamma(k+1) - lgamma(n-k+1))

struct EulerWeight{T} W::T end
function EulerWeight(; Lmax=2^6)
    W = zeros(Lmax,Lmax)
    for L in 1:Lmax
        for k in 0:L-1
            W[L,k+1] = sum(n->2.0^(-n-1)*binom(n,k), k:L-1)
        end
    end
    EulerWeight(W)
end
(eulerweight::EulerWeight)(L,k) = eulerweight.W[L,k+1]
eulerweight = EulerWeight()

L = 2^6
EulerTran = sum(k->eulerweight(L,k)*(-1)^k/(k+1), 0:L-1)
TrueValue = log(2)
@show EulerTran
@show TrueValue
@show EulerTran - TrueValue

f(s) = sum(k -> eulerweight(L,k)*(-1)^k/(1+k)^s, 0:L-1)/(1-2^(1-s))
x = 0.0:0.005:1.0
y = 13.75:0.005:14.75
z = x' .+ im.*y
@time w = f.(z)
# plot(size=(600, 500), title="first non-trivial zero of zeta(s)")
# heatmap!(x, y, abs.(w), color=:rainbow, xlabel="Re s", ylabel="Im s")
# vline!([0.5], color="gray", ls=:dot, label="")
cmap = ["prism", "CMRmap", "gist_rainbow_r"]
plt.figure(figsize=(6,5))
plt.title("first non-trivial zero of zeta(s)")
plt.pcolormesh(x, y, abs.(w), cmap=cmap[3])
plt.xlabel("Re s")
plt.ylabel("Im s")
plt.vlines([0.5], extrema(y)..., lw=0.2, ls="--")

x = 0:0.01:1
y = 12:0.005:45
s = x' .+ y*im
@time w = zeta.(s)

plt.figure(figsize=(4,10))

plt.subplot(121)
plt.pcolormesh(x, y, log.(abs.(w)), cmap="prism")
plt.colorbar()
plt.grid(ls=":")
plt.title("log(abs(ζ(s)))")

plt.subplot(122)
plt.pcolormesh(x, y, angle.(w), cmap="rainbow")
plt.colorbar()
plt.grid(ls=":")
plt.title("angle(ζ(s))")

plt.tight_layout()

n = symbols("n", integer=true)
S = sympy.Sum((-1)^(n-1)/(2n-1), (n, 1, oo)).doit()

N = 10^8
SimpleSum = sum(n->(-1)^(n-1)/(2n-1), 1:N)
TrueValue = π/4
@show SimpleSum
@show TrueValue
@show SimpleSum - TrueValue;

L = 2^6
EulerTran = sum(k->eulerweight(L,k)*(-1)^k/(2k+1), 0:L-1)
TrueValue = π/4
@show EulerTran
@show TrueValue
@show EulerTran - TrueValue

x = symbols("x", real=true, positive=true)
integrate(1/(1+x^3), (x,0,1))

x = symbols("x", real=true, positive=true)
integrate(x/(1+x^3), (x,0,1))

x = symbols("x", real=true, positive=true)
integrate(x^2/(1+x^3), (x,0,1))

x = symbols("x", real=true, positive=true)
integrate(1/(1+x^4), (x,0,1))

x = symbols("x", real=true, positive=true)
integrate(x/(1+x^4), (x,0,1))

x = symbols("x", real=true, positive=true)
integrate(x^2/(1+x^4), (x,0,1))

x = symbols("x", real=true, positive=true)
integrate(x^3/(1+x^4), (x,0,1))

x = symbols("x", real=true, positive=true)
s = symbols("s", real=true, positive=true)
J = simplify(integrate(1/(1+x^(1/s)), (x,0,1)))

sympy.lerchphi(Sym(:z), Sym(:s), Sym(:a))

x = symbols("x", real=true, positive=true)
s = symbols("s", real=true, positive=true)
J = simplify(integrate(1/(1+x^(1/s)), (x,0,oo)))

θ = 2π/10
a(n) = exp(im*n*θ)/n^1.5
sum_a(n) = iszero(n) ? Complex(0.0) : sum(k->a(k), 1:n)
N = 70
n = 0:N
z = sum_a.(n)
plot(size=(400, 400), xlims=(0,1.0), ylim=(0,1.3))
plot!(real.(z), imag.(z), legend=false)#, line=:arrow)

θ = 2π/10
a(n) = exp(im*n*θ)/n^1.5
sum_a(n) = iszero(n) ? Complex(0.0) : sum(k->a(k), 1:n)
@time anim = @animate for N in 0:70
    n = 0:N
    z = sum_a.(n)
    plot(size=(400, 400), xlims=(0,1.0), ylim=(0,1.3))
    plot!(real.(z), imag.(z), legend=false, line=:arrow)
end

gifname = "sum_a.gif"
@time gif(anim, gifname; fps=10)
sleep(0.1)
showimg("image/gif", gifname)

f(z) = -log(1-z)
g(z; N=1000) = sum(n->z^n/n, 1:N)

θ = 0.1:0.01:π
v = @.(f(exp(im*θ)))
w = @.(g(exp(im*θ)))

P1 = plot(xlims=(0,π), xlabel="t")
plot!(θ, real.(w), label="Re(sum exp(int)/n)",  lw=2)
plot!(θ, real.(v), label="Re(-log(1-exp(it)))", lw=2, ls=:dash)

P2 = plot(xlims=(0,π), xlabel="t")
plot!(θ, imag.(w), label="Im(sum exp(int)/n)",  lw=2)
plot!(θ, imag.(v), label="Im(-log(1-exp(it)))", lw=2, ls=:dash)

plot(P1, P2, size=(700,300))

θ = 2π/10
a(n) = exp(im*n*θ)/n
sum_a(n) = iszero(n) ? Complex(0.0) : sum(k->a(k), 1:n)
N = 400
n = 0:N
z = sum_a.(n)
plot(size=(400, 400), xlims=(0,1.0), ylim=(0,1.6))
plot!(real.(z), imag.(z), legend=false)#, line=:arrow)

θ = 2π/10
a(n) = exp(im*n*θ)/n
sum_a(n) = iszero(n) ? Complex(0.0) : sum(k->a(k), 1:n)
@time anim = @animate for N in 0:4:400
    n = 0:N
    z = sum_a.(n)
    plot(size=(400, 400), xlims=(0,1.0), ylim=(0,1.6))
    plot!(real.(z), imag.(z), legend=false, line=:arrow)
end

gifname = "sum_b.gif"
@time gif(anim, gifname; fps=10)
sleep(0.1)
showimg("image/gif", gifname)

Q(N) = sum(n->1/n^2, 1:N)
R(N) = sum(n->1/(2.0^(n-1)*n^2), 1:N) + log(2)^2
@show Q(10), Q(10)-π^2/6, R(10), R(10)-π^2/6
@show Q(30), Q(30)-π^2/6, R(30), R(30)-π^2/6
[(N, Q(N), Q(N)-π^2/6, R(N), R(N)-π^2/6) for N in 10:10:50]

partition_number(n) = sympy.functions.combinatorial.numbers.nT(n)
[(n, partition_number(n)) for n in 1:10]

q = symbols("q")
N = 10
series(prod(1/(1-q^n) for n in 1:N), n=N+1)

q, n = symbols("q n")
N = 5
M = N*(3N-1)÷2
@time series(sympy.summation((-1)^n*q^(n*(3n-1)÷2), (n, -M, M)), q, n=M+1) |>display
@time series(sympy.product(1 - q^n, (n, 1, M)).expand(), q, n=M+1) |> display