################################
# Solving the Replicator Equator equation,
# appearing in Corrolary 2.2 of
#  https://www.biorxiv.org/content/10.1101/623330v1
################################

cd(@__DIR__)
using HighDimPDE
using Random
using Test
using Flux
using Revise
using PyPlot
using UnPack
plotting = true

tspan = (0f0,15f-2)
dt = 5f-2 # time step
μ(X,p,t) = 0f0 # advection coefficients
σ(X,p,t) = 1f-1 # diffusion coefficients
d = 5
ss0 = 5f-2#std g0
U = 25f-2
u_domain = (fill(-U, d), fill(U, d))

##############################
####### Neural Network #######
##############################
batch_size = 10000
train_steps = 2000
K = 1

hls = d + 50 #hidden layer size

nn_batch = Flux.Chain(
        # BatchNorm(d,affine = true, dim = 1),
        Dense(d, hls, tanh),
        # BatchNorm(hls,affine = true, dim = 1),
        Dense(hls,hls,tanh),
        # BatchNorm(hls, affine = true, dim = 1),
        # Dense(hls,hls,tanh),
        Dense(hls, 1, x->x^2)) # Neural network used by the scheme, with batch normalisation

opt = ADAM(1e-2)#optimiser
alg = DeepSplitting(nn_batch, K=K, opt = opt, mc_sample = NormalSampling(1f0) )

##########################
###### PDE Problem #######
##########################
g(x) = Float32((2*π)^(-d/2)) * ss0^(- Float32(d) * 5f-1) * exp.(-5f-1 *sum(x .^2f0 / ss0, dims = 1)) # initial condition
m(x) = - 5f-1 * sum(x.^2, dims=1)
vol = prod(u_domain[2] - u_domain[1])
f(y, z, v_y, v_z, ∇v_y, ∇v_z, p, t) =  v_y .* (m(y) .- vol * v_z .* m(z) ) # nonlocal nonlinear part of the

# defining the problem
prob = PIDEProblem(g, f, μ, σ, tspan,
                    u_domain = u_domain
                    )
# solving
@time xgrid,ts,sol = solve(prob,
                alg,
                dt,
                verbose = true,
                # abstol = 1f0,
                maxiters = train_steps,
                batch_size = batch_size,
                use_cuda = true
                )

###############################
######### Plotting ############
###############################
if plotting
        clf()
        fig, ax = plt.subplots(1,2, sharey = true)

        xgrid1 = collect((-U:5f-3:U))
        xgrid = [reshape(vcat(x, fill(0f0,d-1)),:,1) for x in xgrid1]

        # Analytic sol
        function _SS(x, t, p)
                d = length(x)
                MM = σ(x, p, t) * ones(d)
                SSt = MM .* ((MM .* sinh.(MM *t) .+ ss0 .* cosh.( MM * t)) ./ (MM .* cosh.(MM * t ) .+ ss0 .* sinh.(MM * t)))
                return SSt
        end

        function uanal(x, t, p)
                d = length(x)
                return (2*π)^(-d/2) * prod(_SS(x, t, p) .^(-1/2)) * exp(-0.5 *sum(x .^2 ./ _SS(x, t, p)) )
        end

        for t in collect(tspan[1]: dt : tspan[2])
                ys = uanal.(xgrid, t, Ref(Dict()))
                ax[2].plot(xgrid1, reduce(hcat,ys)[:], label = "t = $t")
        end
        ax[2].set_title("Analytical solution")
        gcf()

        #Deepsplitting sol
        for i in 1:length(sol)
                ax[1].scatter(xgrid1, reduce(vcat,sol[i].(xgrid)), s = .2, label="t = $(dt * (i-1))")
        end
        gcf()

        ax[1].set_title("DeepSplitting")

        for _a in ax[1:1]
                _a.legend()
        end
        gcf()
        savefig("hamel_$(d)d.pdf")

        #####
        # other DimensionMismatch
        #####
        if false
                dx = 0.05
                x = u_domain[1,1]:dx:u_domain[1,2]
                plt.contourf(x,x,g.(repeat(x,2)))
        end
end
ax[1].set_ylabel(L"u(t,(x_1,0,\dots,0))")
ax[1].set_xlabel(L"x_1")
ax[2].set_xlabel(L"x_1")

fig.tight_layout()
gcf()

fig.savefig("hamel_5d.pdf", dpi=800)


function largez(k::Integer)
#
#  test polynomial suggested by Z Zeng
#
    p = [
         1.0
        -0.7
        -0.19
         0.177
        -0.7364
        -0.43780
        -0.952494
        -0.2998258
        -0.00322203
        -0.328903811
        -0.4959527435
        -0.9616679762
         0.4410459281
         0.1090273141
         0.6868094008
         0.0391923826
         0.0302248540
         0.6603775863
        -0.1425784968
        -0.3437618593
         0.4357949015
        ];
    z = [
         0.5 + im
         0.5 - im
        -1 + 0.2im
        -1 - 0.2im
        -0.1 + im
        -0.1 - im
         0.8 + 0.6im
         0.8 - 0.6im
        -0.7 + 0.7im
        -0.7 - 0.7im
         1.4
        -0.4 + 0.9im
        -0.4 - 0.9im
         0.9
        -0.8 + 0.3im
        -0.8 - 0.3im
         0.3 + 0.8im
         0.3 - 0.8im
         0.6 + 0.4im
         0.6 - 0.4im
        ]

    mult = 1
    for i = 1:k
        p = conv(p, p)
        mult *= 2
    end

    z = [z mult*ones(20)]
    p, PolyZeros(z)
end

largez01() = largez(0)
largez02() = largez(1)
largez03() = largez(2)
largez04() = largez(3)
largez05() = largez(4)
largez06() = largez(5)


imu = -6 ± 0.4

u = complex(5 ± 0.3, imu)
v = complex(imu, 8 ± 0.9)

@testset "Addition and subtraction" begin
    @test @inferred(2u - u - 3v + 2v + pi) ≈ @inferred(u - v + pi)
    @test @inferred(u + v) == @inferred(v + u)
    for a in (u, v, z)
        @test @inferred(a + 0) == @inferred(0 + a) == a
        @test @inferred(a - 0) == @inferred(-(0 - a)) == a
        @test @inferred(a + float(im)) == @inferred(float(im) + a)
    end
end

@testset "Multiplication and power" begin
    @test @inferred(v*v*u + u*u*v) ≈ @inferred(u^2*v + v^2*u)
    @test @inferred(u * v) == @inferred(v * u)
    for a in (u, v, z)
        @test @inferred(a * 1) == @inferred(1 * a) == a
        @test @inferred(a ^ 1) == a
        @test @inferred(a ^ 0) == one(a)
        @test @inferred(a * im) == @inferred(im * a)
    end
end

@testset "Division" begin
    @test @inferred(v/v + u/u) ≈ complex(2 ± 0) atol = 5e-18
    for a in (u, v, z)
        @test @inferred(a / 1) == a
        @test @inferred(a / im) == @inferred(-im * a)
    end
end

@testset "Inverse" begin
    @test @inferred(inv(v) + inv(u)) ≈ @inferred((v + u)/(v*u))
end

@testset "Trig functions" begin
    for a in (u, v)
        @test @inferred(cos(acos(a)))   ≈ a
        @test @inferred(cosh(acosh(a))) ≈ a
        @test @inferred(sin(asin(a)))   ≈ a
        @test @inferred(sinh(asinh(a))) ≈ a
        @test @inferred(tan(a))         ≈ @inferred(sin(a)/cos(a))
        @test @inferred(tan(atan(a)))   ≈ a
        @test @inferred(tanh(a))        ≈ @inferred(sinh(a)/cosh(a))
        @test @inferred(tanh(atanh(a))) ≈ a
        @test @inferred(sinpi(a))       ≈ sin(pi * a)
        @test @inferred(cospi(a))       ≈ cos(pi * a)
        @test @inferred(sinc(a))        ≈ (sin(pi * a) / (pi * a))
        @test @inferred(cosc(a))        ≈ ((cos(pi * a) - sin(pi * a) / (pi * a)) / a)
    end
end

@testset "Exponential and logarithm" begin
    @test @inferred(u ^ v) ≈ @inferred(exp(v * log(u)))
    for a in (u, v, z)
        @test @inferred(exp(a))          ≈ ℯ^a
        @test @inferred(expm1(a))        ≈ exp(a) - 1
        @test @inferred(exp10(a))        ≈ 10^a
        @test @inferred(exp2(a))         ≈ 2^a
        @test @inferred(exp(log(a)))     ≈ a
        @test @inferred(log(ℯ, a))       ≈ log(a)
        @test @inferred(log(10, a))      ≈ log10(a)
        @test @inferred(log(2, a))       ≈ log2(a)
        @test @inferred(log1p(a))        ≈ log(1 + a)
        @test @inferred(exp2(log2(a)))   ≈ a
        @test @inferred(exp10(log10(a))) ≈ a
        @test @inferred(expm1(log1p(a))) ≈ a
    end
end

@testset "Square root" begin
    for a in (u, v, z)
        @test @inferred(sqrt(a))         ≈ @inferred(a^0.5)
        @test @inferred(sqrt(a)*sqrt(a)) ≈ a
        @test @inferred(sqrt(big(a)))    ≈ sqrt(a)
    end
end

@testset "Absolute value" begin
    for a in (u, v, z); @test @inferred(abs(a)) ≈ @inferred(sqrt(abs2(a))); end
end

@testset "Representation" begin
    # Make sure the printed representation of a Measurement object is correctly parsed as
    # the same object (well, the tag will be different, but that's not important here).
    for a in (u, v, z); @test eval(Meta.parse(repr(a))) == a; end
end

@testset "Other" begin
    ⊕(a::Complex{<:Measurement}, b::Integer) =
        Measurements.result(Measurements.value(a) + b, (1, 0, 0, 1), a)
    ⊗(a::Complex{<:Measurement}, b::Integer) =
        Measurements.result(Measurements.value(a) * b, (b, 0, 0, b), a)
    for a in (u, v, z)
        @test a ⊕ 5 == a + 5
        @test a ⊗ 5 == a * 5
    end
end

# TODO: Define a method for lgamma with complex argument.
# # Factorial and gamma
# for a in (u, v); test_approx_eq(gamma(a), factorial(a - 1)); end
# for a in (u, v); test_approx_eq(gamma(a + 1), factorial(a)); end
# for a in (u, v); test_approx_eq(lgamma(a), log(gamma(a))); end


# Power cone constrained problems
using Test, LinearAlgebra, SparseArrays
using COSMO

# This test is precision agnostic and will be run with T precision
if !@isdefined(UnitTestFloats)
    UnitTestFloats = [Float64] #if not run in full test setup, just do it for one float type
end

@testset "Power cone problems" begin

for T in UnitTestFloats

    @testset "Power cone problems (T = $(T))" begin
        if precision(T) >= precision(Float64)
            @testset "Feasible" begin

                # solve the following power cone problem
                # max  x1^0.6 y^0.4 + x2^0.1
                # s.t. x1, y, x2 >= 0
                #      x1 + 2y  + 3x2 == 3
                # which is equivalent to
                # max z1 + z2
                # s.t. (x1, y, z1) in K_pow(0.6)
                #      (x2, 1, z2) in K_pow(0.1)
                #      x1 + 2y + 3x2 == 3

                # x = (x1, y, z1, x2, y2, z2)
                n = 6
                P = spzeros(T, n, n)
                q = zeros(T, 6)
                q[3] = q[6] = -one(T)

                # (x1, y, z1) in K_pow(0.6)
                A1 = spdiagm(0 => ones(T, 3))
                b1 = zeros(T, 3)
                cs1 = COSMO.Constraint(A1, b1, COSMO.PowerCone(T(0.6)), 6, 1:3)

                # (x2, 2, z2) in K_pow(0.1)
                A2 = spdiagm(0 => ones(T, 3))
                b2 = zeros(T, 3)
                cs2 = COSMO.Constraint(A2, b2, COSMO.PowerCone(T(0.1)),6, 4:6)

                # x1 + 2y + 3x2 == 3
                cs3 = COSMO.Constraint(T[1.0 2.0 0 3.0 0 0], T[-3.], COSMO.ZeroSet)

                # y2 == 1
                cs4 = COSMO.Constraint(T[0 0 0 0 1.0 0], T[-1.], COSMO.ZeroSet)

                model = COSMO.Model{T}()
                assemble!(model, P, q, [cs1; cs2; cs3; cs4], settings = COSMO.Settings{T}(max_iter = 5000))

                res = COSMO.optimize!(model)
                @test res.status == :Solved
                @test isapprox(res.obj_val, -1.8458, atol = 1e-3)

                # using JuMP
                # m = Model(with_optimizer(COSMO.Optimizer))
                # @variable(m, x1)
                # @variable(m, y)
                # @variable(m, z1)
                # @variable(m, x2)
                # @variable(m, y2 == 1.)
                # @variable(m, z2)

                # @objective(m, Max, z1 + z2)
                # @constraint(m, x1 + 2 * y + 3 * x2 == 3)
                # @constraint(m, [x1; y; z1] in MOI.PowerCone(0.6))
                # @constraint(m, [x2; y2; z2] in MOI.PowerCone(0.1))
                # status = JuMP.optimize!(m)

            end
    end
        # Run infeasibility tests only if precision(T) is high enough
        if precision(T) >= precision(Float64)
            @testset "Primal infeasible" begin
                # max  z
                # s.t. x^0.8 * y^0.2 >= z
                #       x = y == 1
                #       z = 2

                P = spzeros(T, 3, 3)
                q = T[0; 0; -1.0]

                cs1 = COSMO.Constraint(spdiagm(0 => ones(T, 3)), zeros(T, 3), COSMO.PowerCone{T}(0.8))
                cs2 = COSMO.Constraint(spdiagm(0 => ones(T, 3)), T[-1.; -1.; -2.], COSMO.ZeroSet)

                model = COSMO.Model{T}()
                assemble!(model, P, q, [cs1; cs2])

                res = COSMO.optimize!(model)

                @test res.status == :Primal_infeasible
            end

            @testset "Dual infeasible" begin
                # min z
                # s.t. x^0.8 * y^0.2 >= z

                P = spzeros(T, 3, 3)
                q = T[0; 0; 1.0]

                cs1 = COSMO.Constraint(spdiagm(0 => ones(T, 3)), zeros(T, 3), COSMO.PowerCone(T(0.8)))

                model = COSMO.Model{T}()
                assemble!(model, P, q, cs1)

                res = COSMO.optimize!(model)
                @test res.status == :Dual_infeasible
            end
        end

    end

    @testset "Dual power cone problems" begin
        @testset "Feasible" begin
            # max  z
            # s.t. (x/0.8)^0.8 * (y/0.2)^0.2 >= z
            #       x == 0.8,  y == 0.2
            P = spzeros(T, 3, 3)
            q = T[0.; 0.; -1]

            # (x, y, z) in K^*_pow(0.8)
            A1 = spdiagm( 0 => ones(T, 3))
            b1 = zeros(T, 3)
            cs1 = COSMO.Constraint(A1, b1, COSMO.DualPowerCone(T(0.8)))

            # x == 0.8, y == 0.2
            cs2 = COSMO.Constraint(T[1. 0 0; 0 1 0], T[-0.8; -0.2], COSMO.ZeroSet)

            model = COSMO.Model{T}()
            assemble!(model, P, q, [cs1; cs2])

            res = COSMO.optimize!(model)
            @test res.status == :Solved
            @test isapprox(res.obj_val, -one(T), atol=1e-3)

        end
    end
end
end
nothing


# ---
# title: 829. Consecutive Numbers Sum
# id: problem829
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Math
# link: <https://leetcode.com/problems/consecutive-numbers-sum/description/>
# hidden: true
# ---
#
# Given a positive integer `N`, how many ways can we write it as a sum of
# consecutive positive integers?
#
# **Example 1:**
#
#
#
#     Input: 5
#     Output: 2
#     Explanation: 5 = 5 = 2 + 3
#
# **Example 2:**
#
#
#
#     Input: 9
#     Output: 3
#     Explanation: 9 = 9 = 4 + 5 = 2 + 3 + 4
#
# **Example 3:**
#
#
#
#     Input: 15
#     Output: 4
#     Explanation: 15 = 15 = 8 + 7 = 4 + 5 + 6 = 1 + 2 + 3 + 4 + 5
#
# **Note:**  `1 <= N <= 10 ^ 9`.
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


using PyPlot
using Glob
using JLD2

defaultcolors = plt.rcParams["axes.prop_cycle"].by_key()["color"]
removespine(side) = gca().spines[side].set_visible(false)
removespines(sides...) = [removespine(side) for side in sides]

include("ConvergenceTests/ConvergenceTests.jl")

using .ConvergenceTests
using .ConvergenceTests: compute_errors, extract_sizes

filenames = glob("taylor_green*.jld2", joinpath(@__DIR__, "data"))

errors = compute_errors(
            (x, y, z, t) -> ConvergenceTests.DoublyPeriodicTaylorGreen.u(x, y, t),
            filenames...)

sizes = extract_sizes(filenames...)

Nx = map(sz -> sz[1], sizes)
L₁ = map(err -> err.L₁, errors)
L∞ = map(err -> err.L∞, errors)

close("all")
fig, ax = subplots()

ax.loglog(Nx, L₁, basex=2, linestyle="None", marker="o", label="error, \$L_1\$-norm")
ax.loglog(Nx, L∞, basex=2, linestyle="None", marker="^", label="error, \$L_\\infty\$-norm")
ax.loglog(Nx, L₁[end] * (Nx[end] ./ Nx).^2, "k-", basex=2, linewidth=1, alpha=0.6,  label=L"\sim N_x^{-2}")

legend()

title("Convergence for freely-decaying Taylor-Green vortex")
removespines("top", "right")
ylabel("Norms of the absolute error, \$ | u_\\mathrm{simulation} - u_\\mathrm{analytical} | \$")
xlabel(L"N_x")

filepath = joinpath(@__DIR__, "figs", "taylor_green_convergence.png")
savefig(filepath, dpi=480)

test_rate_of_convergence(L₁, Nx, expected=-2.0, atol=0.001, name="Taylor-Green L₁")
test_rate_of_convergence(L∞, Nx, expected=-2.0, atol=0.050, name="Taylor-Green L∞")


# heatmaps

# Plots
import Pkg
Pkg.add("Plots")
Pkg.add("PyPlot")
using Plots

using Random
Random.seed!(2018)

a = rand(5,5)
xlabel = string.(collect('A':'E'))
ylabel = string.(collect('a':'e'))
Plots.heatmap(a, xticks=(1:5, xlabel), yticks=(1:5, ylabel),
    fill_z=a, aspect_ratio=:equal)

fontsize = 15
nrow, ncol = size(a)
ann = [(i,j, Plots.text(round(a[i,j], digits=2), fontsize, :white, :center))
            for i in 1:nrow for j in 1:ncol]
annotate!(ann, linecolor=:white)

# makie
using Random

let
    Random.seed!(123)
    fig = Figure(resolution = (600, 400))
    ax = Axis(fig[1, 1]; xlabel = "x", ylabel = "y")
    hmap = heatmap!(2rand(20, 20) .- 1; colormap = :RdPu_9)
    Colorbar(fig[1, 2], hmap; label = "values", width = 15, ticksize = 15, tickalign = 1)
    colsize!(fig.layout, 1, Aspect(1, 1.0))
    colgap!(fig.layout, 7)
    display(fig)
end;

# gadfly
using Gadfly

parameters = ["ξ₁", "η₁",  "ξ₂", "η₂"]
palettef = Scale.lab_gradient("darkgreen", "orange", "blue")

Gadfly.spy(rand(4,4), Scale.y_discrete(labels = i->parameters[i]), Scale.x_discrete,
Guide.ylabel("Parameters"), Guide.xlabel("Mode"), Scale.color_continuous(colormap=palettef, minvalue=0, maxvalue=1.5))


function processinterval(   𝑖::Int,
                            𝓟::Vector{T},
                            sample_position_template::Vector{T},
                            𝑒::Function,
                            L::Int,
                            verbose_flag) where T
    @assert 1 < 𝑖 <= length(𝓟)

    a = 𝓟[𝑖-1]
    b = 𝓟[𝑖]

    # build approximation of f over current interval.
    X = samplepoints( sample_position_template, a, b )

    𝑒_itp_X = xx->Barycentric2nditp(xx, X, 𝑒.(X))

    # fit Chebyshevinterpolator over this approximation.
    d, ν = getChebyshevinterpolator(𝑒_itp_X, L, a, b)

    𝑒_extrema_positions_X = scalefromChebyinterval.(findChebyitpextremawithchk(d), a, b)

    # d and X can be used for plotting or debugging purposes.
    return 𝑒_extrema_positions_X, d, X
end

# sample_position_template is between [-1,1]. For each interval, rescale it to the interval at hand.
function getextremafromroots(  𝑒::Function,
                                L::Int,
                                𝓟_in::Vector{T},
                                sample_position_template::Vector{T},
                                tol_𝓧_spacing, # was defaulted to convert(T,1e-6/length(𝓟_in)
                                tol_derivative_zero,
                                verbose_flag = true) where T
    𝓟 = Vector{T}(undef,length(𝓟_in)+2)
    𝓟[1] = -one(T)
    𝓟[2:end-1] = 𝓟_in
    𝓟[end] = one(T)

    makenotclose!(𝓟, convert(T,1e-14/length(𝓟_in))) # just to be sure positions are spread out enough.

    # prepare outputs.
    out = Vector{T}(undef,0)
    d_set = Vector{Vector{T}}(undef,length(𝓟)-1)
    X_set = Vector{Vector{T}}(undef,length(𝓟)-1)

    # work on intervals.
    for i = 2:length(𝓟)

        # build approximation of f over current interval.
        X = samplepoints( sample_position_template,𝓟[i-1],𝓟[i] )

        𝑒_itp_X = xx->Barycentric2nditp(xx, X, 𝑒.(X))

        # fit Chebyshevinterpolator over this approximation.
        d, ν = getChebyshevinterpolator(𝑒_itp_X, L, 𝓟[i-1], 𝓟[i])

        if verbose_flag
            Printf.@printf("working on %d-th interval\n",i-1)
        end
        𝑒_extrema_positions_X = scalefromChebyinterval.(findChebyitpextremawithchk(d), 𝓟[i-1], 𝓟[i])

        push!(out, 𝑒_extrema_positions_X...)

        # for plotting or debugging.
        d_set[i-1] = d
        X_set[i-1] = X
    end

    # remove extrema search results that are too close to zero.
    #   These are probably incorrect results, since the local extremas
    #   should be greater or equal to the level error, which is non-zero
    #   unless the target function is in the space of Chebyshev approximations.
    while sign(𝑒(out[1])) == zero(T) && !isempty(out)
        popfirst!(out)
    end

    return out, d_set, X_set
end

function getextremafromrootsparallel(𝑒::Function,
                                L::Int,
                                𝓧::Vector{T},
                                sample_position_template::Vector{T},
                                tol_𝓧_spacing, # was defaulted to convert(T,1e-6/length(𝓟_in)
                                tol_derivative_zero,
                                verbose_flag = true) where T

    ## construct the boundary point sets that will dictate the piece-wise proxy intervals.
    𝓑 = copy(𝓧)
    if !isnumericallyclose(-one(T), 𝓧[1])
        insert!(𝓑,1,-one(T))
    end
    if !isnumericallyclose(one(T), 𝓧[end])
        push!(𝓑,one(T))
    end

    # work on intervals.
    workerfunc = xx->processinterval(xx, 𝓑, sample_position_template, 𝑒, L, verbose_flag)
    sol = pmap(workerfunc, collect(2:length(𝓑)) )

    # package up outputs.
    A = sol
    k = 1
    len = sum( length(A[i][k]) for i = 1:length(A))
    out = Vector{T}(undef,len)

    st = 1
    for i = 1:length(A)
        fin = st + length(A[i][k]) -1
        out[st:fin] = A[i][k]

        st = fin + 1
    end

    d_set = collect( A[i][2] for i = 1:length(A) )
    X_set = collect( A[i][3] for i = 1:length(A) )

    return out, d_set, X_set
end


using LinearAlgebra: diag, eigen, eigvals, eigvecs, Symmetric, Diagonal, tr, norm, cholesky, LowerTriangular, UpperTriangular

@doc doc"""
    CholeskySpace{N} <: Manifold

the manifold of lower triangular matrices with positive diagonal and
a metric based on the cholesky decomposition. The formulae for this manifold
are for example summarized in Table 1 of
> Lin, Zenhua: "Riemannian Geometry of Symmetric Positive Definite Matrices via
> Cholesky Decomposition", arXiv: [1908.09326](https://arxiv.org/abs/1908.09326).
"""
struct CholeskySpace{N} <: Manifold end
CholeskySpace(n::Int) = CholeskySpace{n}()

# two small helper for strictly lower and upper triangulars
strictlyLowerTriangular(x) = LowerTriangular(x) - Diagonal(diag(x))
strictlyUpperTriangular(x) = UpperTriangular(x) - Diagonal(diag(x))

@generated representation_size(::CholeskySpace{N}) where N = (N,N)

@generated manifold_dimension(::CholeskySpace{N}) where N = div(N*(N+1), 2)

@doc doc"""
    inner(M,x,v,w)

computes the inner product on the [`CholeskySpace`](@ref) `M` at the
lower triangular matric with positive diagonal `x` and the two tangent vectors
`v`,`w`, i.e they are both lower triangular matrices with arbitrary diagonal.
The formula reads

````math
    g_{x}(v,w) = \sum_{i>j} v_{ij}w_{ij} + \sum_{j=1}^m v_{ii}w_{ii}x_{ii}^{-2}
````
"""
inner(::CholeskySpace{N},x,v,w) where N = sum(strictlyLowerTriangular(v).*strictlyLowerTriangular(w)) + sum(diag(v).*diag(w)./( diag(x).^2 ))

@doc doc"""
    distance(M,x,y)

computes the Riemannian distance on the [`CholeskySpace`](@ref) `M` between two
matrices `x`, `y` that are lower triangular with positive diagonal. The formula
reads

````math
d_{\mathcal M}(x,y) = \sqrt{
\sum_{i>j} (x_{ij}-y_{ij})^2 +
\sum_{j=1}^m (\log x_{jj} - \log y_{jj})^2
}
````
"""
distance(::CholeskySpace{N},x,y) where N = sqrt(
  sum( (strictlyLowerTriangular(x) - strictlyLowerTriangular(y)).^2 ) + sum( (log.(diag(x)) - log.(diag(y))).^2 )
)

@doc doc"""
    exp!(M,y,x,v)

compute the exponential map on the [`CholeskySpace`](@ref) `M` eminating from
the lower triangular matrix with positive diagonal `x` towards the lower triangular
matrix `v` and return the result in `y`. The formula reads

````math
\exp_x v = \lfloor x \rfloor + \lfloor v \rfloor
+\operatorname{diag}(x)\operatorname{diag}(x)\exp\bigl( \operatorname{diag}(v)\operatorname{diag}(x)^{-1}\bigr)
````
where $\lfloor x\rfloor$ denotes the strictly lower triangular matrix of $x$ and
$\operatorname{diag}(x)$ the diagonal matrix of $x$
"""
function exp!(::CholeskySpace{N},y,x,v) where N
    y .= strictlyLowerTriangular(x) + strictlyLowerTriangular(v) + Diagonal(x)*Diagonal(exp.(diag(v)./diag(x)))
    return y
end
@doc doc"""
    log!(M,v,x,y)

compute the exponential map on the [`CholeskySpace`](@ref) `M` eminating from
the lower triangular matrix with positive diagonal `x` towards the lower triangular
matrx `v` and return the result in `y`. The formula reads

````math
\log_x v = \lfloor x \rfloor - \lfloor y \rfloor
+\operatorname{diag}(x)\log\bigl(\operatorname{diag}(y)\operatorname{diag}(x)^{-1}\bigr)
````
where $\lfloor x\rfloor$ denotes the strictly lower triangular matrix of $x$ and
$\operatorname{diag}(x)$ the diagonal matrix of $x$
"""
function log!(::CholeskySpace{N},v,x,y) where N
    v .= strictlyLowerTriangular(y) - strictlyLowerTriangular(x) + Diagonal(diag(x))*Diagonal(log.(diag(y)./diag(x)))
    return v
end

@doc doc"""
    zero_tangent_vector!(M,v,x)

returns the zero tangent vector on the [`CholeskySpace`](@ref) `M` at `x` in
the variable `v`.
"""
function zero_tangent_vector!(M::CholeskySpace{N},v,x) where N
    fill!(v,0)
    return v
end

@doc doc"""
    vector_transport!(M,vto,x,v,y)

parallely transport the tangent vector `v` at `x` along the geodesic to `y`
on respect to the [`CholeskySpace`](@ref) manifold `M`. The formula reads

````math
    \mathcal P_{x\to y}(v) = \lfloor v \rfloor  + \operatorname{diag}(y)\operatorname{diag}(x)^{-1}\operatorname{diag}(v),
````
where $\lfloor\cdot\rfloor$ denotes the strictly lower triangular matrix,
and $\operatorname{diag}$ extracts the diagonal matrix.
"""
function vector_transport_to!(::CholeskySpace{N}, vto, x, v, y, ::ParallelTransport) where N
    vto .= strictlyLowerTriangular(x) + Diagonal(diag(y))*Diagonal(1 ./ diag(x))*Diagonal(v)
    return vto
end

@doc doc"""
    check_manifold_point(M,x;kwargs...)

check whether the matrix `x` lies on the [`CholeskySpace`](@ref) `M`, i.e.
it's size fits the manifold, it is a lower triangular matrix and has positive
entries on the diagonal.
The tolerance for the tests can be set using the `kwargs...`.
"""

function check_manifold_point(M::CholeskySpace{N}, x; kwargs...) where N
    if size(x) != representation_size(M)
        return DomainError(size(x),"The point $(x) does not lie on $(M), since its size is not $(representation_size(M)).")
    end
    if !isapprox( norm(strictlyUpperTriangular(x)), 0.; kwargs...)
        return DomainError(norm(UpperTriangular(x) - Diagonal(x)), "The point $(x) does not lie on $(M), since it strictly upper triangular nonzero entries")
    end
    if any( diag(x) .<= 0)
        return DomainError(min(diag(x)...), "The point $(x) does not lie on $(M), since it hast nonpositive entries on the diagonal")
    end
    return nothing
end
"""
    check_tangent_vector(M,x,v; kwargs... )

checks whether `v` is a tangent vector to `x` on the [`CholeskySpace`](@ref) `M`, i.e.
atfer [`check_manifold_point`](@ref)`(M,x)`, `v` has to be of same dimension as `x`
and a symmetric matrix.
The tolerance for the tests can be set using the `kwargs...`.
"""
function check_tangent_vector(M::CholeskySpace{N}, x,v; kwargs...) where N
    mpe = check_manifold_point(M, x)
    if mpe !== nothing
        return mpe
    end
    if size(v) != representation_size(M)
        return DomainError(size(v),
            "The vector $(v) is not a tangent to a point on $(M) since its size does not match $(representation_size(M)).")
    end
    if !isapprox(norm(v-transpose(v)), 0.; kwargs...)
        return DomainError(size(v),
            "The vector $(v) is not a tangent to a point on $(M) (represented as an element of the Lie algebra) since its not symmetric.")
    end
    return nothing
end


using Parameters
using PyPlot

include("lv_model.jl")

function iso1(u2, p)
    @unpack α12, K1 = p
    return K1 - α12 * u2 / K1
end

function iso2(u2, p)
    @unpack α21, K2 = p
    return (K2 - u2) / α21
end

let
    par = LVPar(K1 = 1.5)
    u2range = 0.0:0.1:3.0
    lv_plot = figure()
    plot(u2range, [iso1(u2, par) for u2 in u2range])
    plot(u2range, [iso2(u2, par) for u2 in u2range])
    xlim(0,3)
    ylim(0,3)
    return lv_plot
end


using FastGMM
using Distributions
using StatsBase

d = 2
n = 1000000

# start very simple
m = MvNormal([10, -10], [5, 0.2])

X = Float32.(rand(m, n))
ω = Float32.(repeat([1], n))

gmm = fitGMM(X, ω, 1)

K = 4

# build a model to sample from
m = MixtureModel(MvNormal, [
    ([0, 0], [1, 1]),
    ([5, 5], [0.2, 0.4]),
    ([10, 0], [0.8, 0.1]),
    ([1, 24], [0.2, 2])],
    [0.4, 0.2, 0.1, 0.3])

X = Float32.(rand(m, n))
ω = Float32.(repeat([1], n))

gmm = fitGMM(X, ω, K)


"""
# module Force

# Examples

```jldoctest
julia>
```
"""
module Force

using Distances: euclidean

using VCSMD.Geometry

export lennard_jones_scalar,
    lennard_jones_vector

function lennard_jones_scalar(ε::Float64, σ::Float64)::Function
    function (r::Float64)
        x = σ / r
        24 * ε / σ * (2 * x^13 - x^7)
    end
end

function lennard_jones_vector(ε::Float64, σ::Float64)::Function
    f = lennard_jones_scalar(ε, σ)

    function (ri::Point{N}, rj::Point{N}) where {N}
        r = euclidean(rj, ri)
        (rj - ri) / r * f(r)
    end
end

end


@testset "normal.jl" begin
    @testset "SKCE: basic example" begin
        skce = UnbiasedSKCE(
            ExponentialKernel(; metric=Wasserstein()) ⊗ SqExponentialKernel()
        )

        # only two predictions, i.e., one term in the estimator
        normal1 = Normal(0, 1)
        normal2 = Normal(1, 2)
        @test @inferred(skce([normal1, normal1], [0, 0])) ≈ 1 - sqrt(2) + 1 / sqrt(3)
        @test @inferred(skce([normal1, normal2], [1, 0])) ≈
            exp(-sqrt(2)) *
              (exp(-1 / 2) - 1 / sqrt(2) - 1 / sqrt(5) + exp(-1 / 12) / sqrt(6))
        @test @inferred(skce([normal1, normal2], [0, 1])) ≈
            exp(-sqrt(2)) * (
            exp(-1 / 2) - exp(-1 / 4) / sqrt(2) - exp(-1 / 10) / sqrt(5) +
            exp(-1 / 12) / sqrt(6)
        )
    end

    @testset "SKCE: basic example (transformed)" begin
        skce = UnbiasedSKCE(
            ExponentialKernel(; metric=Wasserstein()) ⊗
            (SqExponentialKernel() ∘ ScaleTransform(0.5)),
        )

        # only two predictions, i.e., one term in the estimator
        normal1 = Normal(0, 1)
        normal2 = Normal(1, 2)
        @test @inferred(skce([normal1, normal1], [0, 0])) ≈
            1 - 2 / sqrt(1.25) + 1 / sqrt(1.5)
        @test @inferred(skce([normal1, normal2], [1, 0])) ≈
            exp(-sqrt(2)) *
              (exp(-1 / 8) - 1 / sqrt(1.25) - 1 / sqrt(2) + exp(-1 / 18) / sqrt(2.25))
        @test @inferred(skce([normal1, normal2], [0, 1])) ≈
            exp(-sqrt(2)) * (
            exp(-1 / 8) - exp(-1 / 10) / sqrt(1.25) - exp(-1 / 16) / sqrt(2) +
            exp(-1 / 18) / sqrt(2.25)
        )
    end

    @testset "SKCE: basic properties" begin
        skce = UnbiasedSKCE(
            ExponentialKernel(; metric=Wasserstein()) ⊗ SqExponentialKernel()
        )

        estimates = map(1:10_000) do _
            predictions = map(Normal, randn(20), rand(20))
            targets = map(rand, predictions)
            return skce(predictions, targets)
        end

        @test any(x -> x > zero(x), estimates)
        @test any(x -> x < zero(x), estimates)
        @test mean(estimates) ≈ 0 atol = 1e-4
    end

    @testset "UCME: basic example" begin
        # one test location
        ucme = UCME(
            ExponentialKernel(; metric=Wasserstein()) ⊗ SqExponentialKernel(),
            [Normal(0.5, 0.5)],
            [1],
        )

        # two predictions
        normal1 = Normal(0, 1)
        normal2 = Normal(1, 2)
        @test @inferred(ucme([normal1, normal2], [0, 0.5])) ≈
            (
            exp(-1 / sqrt(2)) * (exp(-1 / 2) - exp(-1 / 4) / sqrt(2)) +
            exp(-sqrt(5 / 2)) * (exp(-1 / 8) - 1 / sqrt(5))
        )^2 / 4

        # two test locations
        ucme = UCME(
            ExponentialKernel(; metric=Wasserstein()) ⊗ SqExponentialKernel(),
            [Normal(0.5, 0.5), Normal(-1, 1.5)],
            [1, -0.5],
        )
        @test @inferred(ucme([normal1, normal2], [0, 0.5])) ≈
            (
            (
                exp(-1 / sqrt(2)) * (exp(-1 / 2) - exp(-1 / 4) / sqrt(2)) +
                exp(-sqrt(5 / 2)) * (exp(-1 / 8) - 1 / sqrt(5))
            )^2 +
            (
                exp(-sqrt(5) / 2) * (exp(-1 / 8) - exp(-1 / 16) / sqrt(2)) +
                exp(-sqrt(17) / 2) * (exp(-1 / 2) - exp(-9 / 40) / sqrt(5))
            )^2
        ) / 8
    end

    @testset "kernels with input transformations" begin
        nsamples = 100
        ntestsamples = 5

        # create predictions and targets
        predictions = map(Normal, randn(nsamples), rand(nsamples))
        targets = randn(nsamples)

        # create random test locations
        testpredictions = map(Normal, randn(ntestsamples), rand(ntestsamples))
        testtargets = randn(ntestsamples)

        for γ in (1.0, rand())
            kernel1 =
                ExponentialKernel(; metric=Wasserstein()) ⊗
                (SqExponentialKernel() ∘ ScaleTransform(γ))
            kernel2 =
                ExponentialKernel(; metric=Wasserstein()) ⊗
                (SqExponentialKernel() ∘ ARDTransform([γ]))

            # check evaluation of the first two observations
            p1 = predictions[1]
            p2 = predictions[2]
            t1 = targets[1]
            t2 = targets[2]
            for f in (
                CalibrationErrors.unsafe_skce_eval_targets,
                CalibrationErrors.unsafe_ucme_eval_targets,
            )
                out1 = f(kernel1.kernels[2], p1, t1, p2, t2)
                out2 = f(kernel2.kernels[2], p1, t1, p2, t2)
                @test out2 ≈ out1
                if isone(γ)
                    @test f(SqExponentialKernel(), p1, t1, p2, t2) ≈ out1
                end
            end

            # check estimates
            for estimator in (UnbiasedSKCE, x -> UCME(x, testpredictions, testtargets))
                estimate1 = estimator(kernel1)(predictions, targets)
                estimate2 = estimator(kernel2)(predictions, targets)
                @test estimate2 ≈ estimate1
                if isone(γ)
                    @test estimator(
                        ExponentialKernel(; metric=Wasserstein()) ⊗ SqExponentialKernel()
                    )(
                        predictions, targets
                    ) ≈ estimate1
                end
            end
        end
    end
end


struct LeastSquaresMonteCarlo{M<:MonteCarloModel}
    m::M
    degree::Int
end

function value(m::LeastSquaresMonteCarlo, c::AnytimeBefore{C}) where C
    uc = c.c # underlying contract
    uc1 = uc.c1
    uc2 = uc.c2
    N = date2index(m.m, maturitydate(c))
    R = discount(m.m.core.yieldcurve, m.m.dates[2]) # TODO: could be better...

    V = [valueat(ms, uc, N)*R for ms in scenarios(m.m)]
    for n = N-1:-1:2
        I = V .> 0 # in the money options
        A = [x^d for d = 0:m.degree, x in m.m.paths[:,n]] # design matrix
        β = A[:,I]' \ V[I] # least squares regression
        cV = A'*β          # estimated continuation values
        for (p,ms) in enumerate(scenarios(m.m))
            ev = valueat(ms,uc,n)  # exercise value
            if I[p] && cV[p] < ev
                # in-the-money and exercise is greater than continuation
                V[p] = ev*R
            else
                V[p] *= R
            end
        end
    end
    return max(mean(V), mean(valueat(ms, uc, N) for ms in scenarios(m.m)))
end

function value(m::GeomBMModel, c::Contract, ::Type{LeastSquaresMonteCarlo}, npaths::Integer, degree::Integer)
    mcm = montecarlo(m, startdate(m):Day(1):maturitydate(c), npaths)
    value(LeastSquaresMonteCarlo(mcm, degree), c)
end


#Reineke SDI calculation
function sdi(s::Stand)
    return s.trees_per_area * ((s.qmd / 10.0) ^ 1.605)
end

function qmd(s::Stand)
    sum_d2 = 0.0
    @simd for i in eachindex(s.treelist)
    @inbounds sum_d2 += (s.treelist[i].diameter ^ 2 * s.treelist[i].expansion)
    end
    qmd = ismissing(s.trees_per_area) ? sqrt(sum_d2 / s.trees_per_area) : sqrt(sum_d2 / trees_per_area(s))
    return qmd
end

function qmd!(s::Stand)
    s.qmd = qmd(s)
end

function basal_area(s::Stand)
    ba = 0.0
    @simd for i in eachindex(s.treelist)
    @inbounds ba += basal_area(s.treelist[i]) * s.treelist[i].expansion
    end
    return ba
end

function basal_area!(s::Stand)
    s.basal_area = basal_area(s)
end

function basal_area!(s::Stand,m::Metric)
    s.basal_area = basal_area(s, metric)
end

function basal_area(t::Tree)
    return ((t.diameter ^ 2) * K)
end

function basal_area(t::Tree,m::Metric)
    return ((t.diameter ^ 2) * KMETRIC)
end

function trees_per_area(s::Stand)
    tpa = 0.0
    @simd for i in eachindex(s.treelist)
    @inbounds tpa += s.treelist[i].expansion
    end
    return tpa
end

function trees_per_area!(s::Stand)
    s.trees_per_area = trees_per_area(s)
end


# basic test of multiplying big and small together should give small

using ApproxManifoldProducts
using Test
using Manifolds
using TensorCast

##

@testset "multiply big and small" begin

##

M = SpecialEuclidean(2)
N = 100
u0 = ProductRepr([0;0.0],[1 0; 0 1.0])
ϵ = identity_element(M, u0)

X1 = [exp(M, ϵ, hat(M, ϵ, randn(3))) for i in 1:N]
X2 = [exp(M, ϵ, hat(M, ϵ, 0.1.*randn(3))) for i in 1:N]

# test get_coordinates
testval = vee(M, ϵ, log(M, ϵ, X1[1]))
@test length(testval) === 3
@test all(abs.(testval) .< 10.0)

##

p = manikde!(M, X1)
q = manikde!(M, X2)

# check new MKD have right type info cached
@test (p._u0 |> typeof) == typeof(u0)
@test (p._u0 |> typeof) == typeof(u0)

pq = manifoldProduct([p;q], M)

# check new product also has right point type info cached
@test (pq._u0 |> typeof) == typeof(u0)

##

X12_ = getPoints(pq)

X12 = AMP._pointsToMatrixCoords(pq.manifold, X12_)

@test 0.7*N < sum(abs.(X12[1,:]) .< 0.3)
@test 0.7*N < sum(abs.(X12[2,:]) .< 0.3)
@test 0.7*N < sum(abs.(X12[3,:]) .< 0.3)

##

end


using Turing
using LinearAlgebra

using BenchmarkHelper

include("lr_helper.jl")

if !haskey(BenchmarkSuite, "nuts")
    BenchmarkSuite["nuts"] = BenchmarkGroup(["nuts"])
end

X, Y = readlrdata()

@model function lr_nuts(x, y, σ)

    N,D = size(x)

    α ~ Normal(0, σ)
    β ~ MvNormal(zeros(D), σ^2 * I)

    for n = 1:N
        y[n] ~ BinomialLogit(1, dot(x[n,:], β) + α)
    end
end

# Sampling parameter settings
n_samples = 1_000
n_adapts = 1_000

# Sampling
BenchmarkSuite["nuts"]["lr"] = @benchmarkable sample(lr_nuts(X, Y, 100),
                                                     NUTS(0.65), n_samples)


#! /usr/bin/julia

# Rosetta Code, Cholesky decomposition

a = [25 15 5; 15 18 0; -5 0 11]
b = [18 22 54 22; 22 70 86 62; 54 86 174 134; 42 62 134 106]

println(a, "\n => \n", chol(a, :L))
println(b, "\n => \n", chol(b, :L))


include("jset.jl")
include("pgmfile.jl")

h = 400; w = 800; m = Array(Int64, h, w);
c0 = -0.8+0.16im;
pgm_name = "julia.pgm";

t0 = time();
for y=1:h, x=1:w
    c = complex((x-w/2)/(w/2), (y-h/2)/(w/2))
    m[y,x] = juliaset(c, c0, 256)
end
t1 = time();

create_pgmfile(m, pgm_name);
print("Written $pgm_name\nFinished in $(t1-t0) seconds.\n");


using LinearAlgebra, Zygote, Plots

const n = 100
S = circshift(Matrix{Float64}(I, n, n),(1,0))

function loss(W)
    norm(W*S-S*W)/norm(W)
end

function step!(W;lr=0.003)
    # computing current loss and backprop
    current_loss, back_loss = pullback(w -> loss(w),W)
    # computing gradient
    grads = back_loss(1)[1]
    # updating W
    W .-= lr .*grads
end

W = randn(n,n)
W ./= norm(W)

heatmap(W,clims=(-0.03,0.03),legend=:none,axis=nothing)

for i=1:1000
    step!(W)
end

heatmap(W,clims=(-0.03,0.03),legend=:none,axis=nothing)


using JuMP, SCS, LinearAlgebra, Test

"""
    example_robust_uncertainty()
Computes the Value at Risk for a data-driven uncertainty set; see "Data-Driven
Robust Optimization" (Bertsimas 2013), section 6.1 for details. Closed-form
expressions for the optimal value are available.
"""
function example_robust_uncertainty()
    R = 1
    d = 3
    𝛿 = 0.05
    ɛ = 0.05
    N = ceil((2 + 2 * log(2 / 𝛿))^2) + 1

    c = randn(d)

    μhat = rand(d)
    M = rand(d, d)
    Σhat = 1 / (d - 1) * (M - ones(d) * μhat')' * (M - ones(d) * μhat')

    Γ1(𝛿, N) = R / sqrt(N) * (2 + sqrt(2 * log(1 / 𝛿)))
    Γ2(𝛿, N) = 2 * R^2 / sqrt(N) * (2 + sqrt(2 * log(2 / 𝛿)))

    model = Model(with_optimizer(SCS.Optimizer, verbose = 10))

    @variable(model, Σ[1:d, 1:d], PSD)
    @variable(model, u[1:d])
    @variable(model, μ[1:d])
    @constraint(model, [Γ1(𝛿 / 2, N); μ - μhat] in SecondOrderCone())
    @constraint(model, [Γ2(𝛿 / 2, N); vec(Σ - Σhat)] in SecondOrderCone())
    @SDconstraint(model, [((1 - ɛ) / ɛ) (u - μ)'; (u - μ) Σ] >= 0)
    @objective(model, Max, dot(c, u))

    JuMP.optimize!(model)

    exact = dot(μhat, c) + Γ1(𝛿 / 2, N) * norm(c) + sqrt((1 - ɛ) / ɛ) *
        sqrt(dot(c, (Σhat + Γ2(𝛿 / 2, N) * Matrix(1.0I, d, d)) * c))
    @test JuMP.objective_value(model) ≈ exact atol = 1e-3
end

example_robust_uncertainty()


###############################################
#FEniCS tutorial demo program: Diffusion of a Gaussian hill.
#  u'= Laplace(u) + f  in a square domain
#  u = u_D             on the boundary
#  u = u_0             at t = 0
#  u_D = f = 0
#The initial condition u_0 is chosen as a Gaussian hill.
###############################################

module ft04

using FenicsPy

T = 2.0            # final time
num_steps = 50     # number of time steps
Δt = T / num_steps # time step size

# Create mesh and define function space
nx = ny = 30
mesh = RectangleMesh(Point(-2, -2), Point(2, 2), nx, ny)
V = FunctionSpace(mesh, "P", 1)

# Define boundary condition
bc = DirichletBC(V, Constant(0), "on_boundary")

# Define initial value
u_0 = Expression("exp(-a*pow(x[0], 2) - a*pow(x[1], 2))",
                 degree=2, a=5)
u_n = interpolate(u_0, V)

# Define variational problem
u = TrialFunction(V)
v = TestFunction(V)
f = Constant(0)

F = u*v*dx + Δt*dot(grad(u), grad(v))*dx - (u_n + Δt*f)*v*dx
a, L = lhs(F), rhs(F)

# Create VTK file for saving solution
vtkfile = File("heat_gaussian/solution.pvd")

# Time-stepping
u = FeFunction(V)
t = 0
for n = 1:num_steps

    # Update current time
    global t += Δt

    # Compute solution
    solve(a == L, u, bc)

    # Save to file and plot solution
    vtkfile << (u, t)
    #plot(u)

    # Update previous solution
    u_n.assign(u)
end

end # module ft04


#=
-------------------------------------------------------------------------------------------------------------
This file contains the Student template for solving Project 5: the 2D Poisson Equation using the CG method with GGP storage
introduced in F.X. Giraldo's Introduction to Element-based Galerkin Methods using
Tensor-Product Bases: Analysis, Algorithms, and Applications.

The approached follows Algorithm 12.18 in the book.

Written by F.X. Giraldo on July 10, 2021.
Department of Applied Mathematics
Naval Postgraduate School
Monterey, CA 93943

The interpolation points used are the following:
ipoints=1: Lobatto
ipoints=2: Legendre
ipoints=3: Chebyshev
ipoints=4: Equi-spaced

The integration points used are the following:
qpoints=1: Lobatto
qpoints=2: Legendre
-------------------------------------------------------------------------------------------------------------
=#

using Plots
using LinearAlgebra

include("QuadraturePoints.jl")
include("element_matrices.jl")
include("global_matrices.jl")
include("create_grid.jl")
include("compute_metrics.jl")
include("global_matrices.jl")
include("exact_solution.jl")
include("compute_norm.jl")

#Some Constants
DFloat = Float64
machine_zero=eps(DFloat)

#Main Function
function main()

    #-----------------------------Only Change these Input parameters---------------------------------#
    Nex=8
    Ney=8
    N=4
    Q=N+1
    ipoints=1
    qpoints=1
    c=DFloat(1) #Constant in Exact solution
    plot_grid=true
    plot_solution=true
    warp_grid=false
    #-----------------------------Only Change these Input parameters---------------------------------#

    space_method="CG" #only capable of doing CG
    Np=N+1
    Nq=Q+1
    Npts=Np^2
    Nqs=Nq^2
    Ne=Nex*Ney
    Nx=Nex*N+1
    Ny=Ney*N+1
    Npoin=Nx*Ny
    Nboun=2*Nx + 2*(Ny-2)
    @show(N,Q,Ne,Npoin,Nboun,c)

    #Select Interpolation Points
    if ipoints == 1
        (ξ,ω) = QuadraturePoints.lobatto_gauss(Np,DFloat) #Lobatto
    elseif ipoints == 2
        (ξ,ω) = QuadraturePoints.legendre_gauss(Np,DFloat) #Legendre
    elseif ipoints == 3
        (ξ,ω) = QuadraturePoints.chebyshev_gauss(Np,DFloat) #Chebyshev
    elseif ipoints == 4
        (ξ,ω) = QuadraturePoints.equispaced_points(Np,DFloat) #Equi-spaced
    end

    #Select Quadrature Points
    if qpoints == 1
        (ξq,ωq) = QuadraturePoints.lobatto_gauss(Nq,DFloat) #Lobatto
    elseif qpoints == 2
        (ξq,ωq) = QuadraturePoints.legendre_gauss(Nq,DFloat) #Legendre
    end

    #Construct Lagrange Polynomials
    (ψ,dψ) = QuadraturePoints.lagrange_basis(Np,Nq,ξ,ξq,DFloat)

    #Construct Grid and (i,e)->I map
    (coord,intma,iboun) = create_grid(Np,Npoin,Ne,Nboun,Nex,Ney,ξ,plot_grid,warp_grid,DFloat)

    #Construct Metric Terms
    (ξ_x,ξ_y,η_x,η_y,jac) = compute_metrics(coord,intma,ψ,dψ,Ne,Np,Nq,DFloat)

    #-----------------------------Students Add their Functions inside ti_LSRK--------------------------#
    #-----------------------------Students Add their Functions inside ti_LSRK--------------------------#
    #Construct Element Matrices
    (Me,Le) = element_matrices(ψ,dψ,ξ_x,ξ_y,η_x,η_y,jac,ωq,Ne,Np,Nq,DFloat)

    #Construct Global Matrices
    (M,L) = global_matrices(Me,Le,intma,Ne,Np,Npoin,DFloat)
    #-----------------------------Students Add their Functions inside ti_LSRK--------------------------#
    #-----------------------------Students Add their Functions inside ti_LSRK--------------------------#

    #Compute Exact Solution
    (qe,fe) = exact_solution(coord,Npoin,c,DFloat)

    #Apply Dirichlet BCs: Impose exact solution
    R=M*fe
    for i=1:Nboun
        I=iboun[i]
        L[I,:] = zeros(DFloat,Npoin)
        L[I,I] = DFloat(-1)
        R[I]=qe[I]
    end#i

    #Solve Matrix problem
    q=-L\R

    #Compute L2 Norm
    norm2 = compute_norm(q,qe,Npoin,DFloat)
    println("norm.L2 = ",norm2)

    #Plot Solution
    if (plot_solution)
        x=reshape(coord[1,:],(Nx,Ny))
        y=reshape(coord[2,:],(Nx,Ny))
        qq=reshape(q,(Nx,Ny))
        xx=x[:,1]
        yy=y[1,:]

        p1 = contourf(xx,yy,qq)
        title!(p1,"Numerical Solution")
        xlabel!(p1,"x")
        ylabel!(p1,"y")

        qq=reshape(qe,(Nx,Ny))
        p2 = contourf(xx,yy,qq)
        title!(p2,"Exact Solution")
        xlabel!(p2,"x")
        ylabel!(p2,"y")

        figure_1 = plot(p1,p2,layout = (1,2))
        display(figure_1)
    end

    println("Done") #output

end

#----------------------------------#
# Run the main function
#----------------------------------#
main()


function half_adder(A, B)
    S = xor(A, B)
    Cout = A & B
    Cout, S
end

function full_adder(A, B, Cin)
    T = xor(A, B)
    S = xor(Cin, T)
    Cout = (A & B) | (Cin & T)
    Cout, S
end

function nibble_adder(A4, A3, A2, A1, B4, B3, B2, B1)
    carry, S1 = half_adder(A1, B1)
    carry, S2 = full_adder(A2, B2, carry)
    carry, S3 = full_adder(A3, B3, carry)
    carry, S4 = full_adder(A4, B4, carry)

    S4, S3, S2, S1
end

function binary_multiply(input, counter)
    accumulator = 0
    input += 0b0000_0000_0000_0000
    while counter > 0
        if counter & 1 == 1
            accumulator += input
        end
        counter >>= 1
        input <<= 1
    end
    accumulator
end

binary_multiply(0b0000_1100,0b1110_1010)


# Poisson distribution

export Poisson
import Base
using SpecialFunctions: logfactorial

@parameterized Poisson(λ) ≪ CountingMeasure(ℤ[0:∞])

Base.eltype(::Type{P}) where {P<:Poisson} = Int

function logdensity(d::Poisson{(:λ,)}, y)
    λ = d.λ
    return y * log(λ) - λ - logfactorial(y)
end

function logdensity(d::Poisson{(:logλ,)}, y)
    return y * d.logλ - exp(d.logλ) - logfactorial(y)
end

asparams(::Type{<:Poisson}, ::Val{:λ}) = asℝ₊
asparams(::Type{<:Poisson}, ::Val{:logλ}) = asℝ

sampletype(::Poisson) = Int

Base.rand(rng::AbstractRNG, T::Type, d::Poisson{(:λ,)}) = rand(rng, Dists.Poisson(d.λ))
Base.rand(rng::AbstractRNG, T::Type, d::Poisson{(:logλ,)}) = rand(rng, Dists.Poisson(exp(d.logλ)))

≪(::Poisson, ::IntegerRange{lo,hi}) where {lo, hi} = lo ≤ 0 && isinf(hi)


reduce(+, map(x->x^2, 1:100)) #338350
reduce(+, 1:100)^2 #25502500


function op_K( pw_grid::PWGrid, psi::Array{ComplexF64,2} )
    out = zeros(ComplexF64,size(psi))
    Ncol = size(psi,2)
    Ω = pw_grid.Ω
    G2 = pw_grid.G2
    Npoints = pw_grid.Npoints
    for is = 1:Ncol
        for ip = 1:Npoints
            out[ip,is] = psi[ip,is]*G2[ip]
        end
    end
    return 0.5*out # two minus signs
end

function op_K( pw_grid::PWGrid, psi::Array{ComplexF64,1} )
    G2 = pw_grid.G2
    Npoints = pw_grid.Npoints
    out = zeros( size(psi) )
    for ip = 1:Npoints
        out[ip] = psi[ip]*G2[ip]
    end
    return 0.5*out
end


using NonlinearEigenproblemsTest
using NonlinearEigenproblems
using Test
using LinearAlgebra

@bench @testset "broyden" begin
    dep=nep_gallery("dep1");
    S,V=broyden(dep)
    # Broyden returns a Schur factorization so check with MM
    @test opnorm(compute_MM(dep,S,V))<sqrt(eps())
    # test addconj
    S,V=broyden(dep,addconj=true,pmax=5)
    # Test by computing the eigenpairs
    λv,X = eigen(S)
    V_eigvecs=V*X;
    verify_lambdas(4, dep, λv, V_eigvecs, sqrt(eps()))
end


using Geodesy
using Base.Test

# Construction

x, y = (rand(2) - .5) * 10_000

@test ENU(x, y) == ENU(x, y, 0.0)

@test LLA(x, y) == LLA(x, y, 0.0)

ECEF(x, y, 0.0)

# get* methods

ll = LL(x, y)
lla = LLA(x, y, rand())
@test LLA(getX(lla), getY(lla), getZ(lla)) == lla
@test getY(ll) == y
@test getX(ll) == x

enu = ENU(x, y, rand())
@test ENU(getX(enu), getY(enu), getZ(enu)) == enu


"Square the sum of the numbers up to the given number"
function square_of_sum(n::Int)

end

"Sum the squares of the numbers up to the given number"
function sum_of_squares(n::Int)

end

"Subtract sum of squares from square of sums"
function difference(n::Int)

end


Base.parent(A::StarAlgebra{<:Group}) = A.object

function StarAlgebras.AlgebraElement(
    χ::AbstractClassFunction,
    RG::StarAlgebra{<:Group},
)
    G = parent(RG)
    @assert G === parent(χ)
    b = basis(RG)

    dim = degree(χ)
    ord = order(Int, G)
    c = dim // ord

    T = Base._return_type(*, Tuple{typeof(c), eltype(χ)})

    x = AlgebraElement(zeros(T, length(b)), RG)

    for (v, cc) in zip(values(χ), conjugacy_classes(χ))
        for g in cc
            x[inv(g)] = c * v
        end
    end
    return x
end

function algebra_elt_from_support(support, RG::StarAlgebra)
    b = basis(RG)
    I = [b[s] for s in support]
    V = fill(1, length(support))
    return AlgebraElement(sparsevec(I, V, length(b)), RG)
end
struct CyclicSubgroups{Gr, GEl}
    group::Gr
    seen::Dict{Int,Set{GEl}}
    min_order::Int
    max_order::Int

    function CyclicSubgroups(G::Group; min_order=1, max_order=order(Int, G))
        seen = Dict{Int, Set{eltype(G)}}()
        return new{typeof(G), eltype(G)}(G, seen, min_order, max_order)
    end
end

Base.eltype(citr::CyclicSubgroups) = valtype(citr.seen)
Base.IteratorSize(::CyclicSubgroups) = Base.SizeUnknown()

function Base.iterate(citr::CyclicSubgroups)
    g, state = iterate(citr.group) # g is identity here
    @assert isone(g)
    if citr.min_order ≤ 1 ≤ citr.max_order
        citr.seen[1] = Set([g])
        return Set([g]), state
    end
    return iterate(citr, state)
end

function Base.iterate(citr::CyclicSubgroups, state)
    k = iterate(citr.group, state)
    k === nothing && return nothing
    g, state = k
    ord = order(Int, g)
    if citr.min_order ≤ ord ≤ citr.max_order
        if !haskey(citr.seen, ord)
            citr.seen[ord] = Set([g])
            return Set(g^i for i in 1:ord), state
        else
            if any(g^i in citr.seen[ord] for i in 1:ord-1)
                return iterate(citr, state)
            else
                push!(citr.seen[ord], g)
                return Set(g^i for i in 1:ord), state
            end
        end
    end
    return iterate(citr, state)
end

function small_idempotents(
    RG::StarAlgebra{<:Group},
    subgroups = CyclicSubgroups(parent(RG), min_order = 2),
)
    return (algebra_elt_from_support(H, RG) // length(H) for H in subgroups)
end

function (χ::AbstractClassFunction)(α::AlgebraElement{<:StarAlgebra{<:Group}})
    @assert parent(χ) === parent(parent(α))
    iszero(α) && return zero(eltype(StarAlgebras.coeffs(α)))*zero(eltype(χ))
    return sum(α(g) * χ(g) for g in supp(α))
end

function minimal_rank_projection(χ::Character, RG::StarAlgebra{<:Group};
    idems = small_idempotents(RG), # idems are AlgebraElements over Rational{Int}
    iters=3
)
    degree(χ) == 1 && return one(RG, Rational{Int}), true

    for µ in idems
        isone(χ(µ)) && µ^2 == µ && return µ, true
    end

    for n in 2:iters
        for (i, elts) in enumerate(Iterators.product(fill(idems, n)...))
            µ = *(elts...)
            isone(χ(µ)) && µ^2 == µ && return µ, true
        end
    end
    @debug "Could not find minimal projection for $χ"
    return one(RG, Rational{Int}), false
end

function minimal_projection_system(
    chars::AbstractVector{<:AbstractClassFunction},
    RG::StarAlgebra{<:Group}
)
    res = fetch.([Threads.@spawn minimal_rank_projection(χ, RG) for χ in chars])

    r1p, simple = first.(res), last.(res) # rp1 are sparse storage

    mps = [isone(µ) ? AlgebraElement(χ, RG) : µ*AlgebraElement(χ, RG)
        for (µ,χ) in zip(r1p, chars)] # dense storage

    return mps, simple
end


#tracywidomlis.jl
#Algorithm 2.7 of Random Eigenvalues by Alan Edelman

#Experiment:    Sample random permutations
#Plot:          Histogram of lengths of longest increasing subsequences
#Theory:        The Tracy-Widom law

## Parameters
t = 10000       # number of trials
n = 6^6         # length of permutations
dx = 1/6        # bin size

## Experiment
require("patiencesort.jl")
function tracywidomlis(t, n, dx)
    ## single processor: 1x speed
    #v = [patiencesort(randperm(n)) for i = 1:t]

    ## simple parallelism: Nx speed * 130%
    v = pmap((i)->patiencesort(randperm(n)), 1:t)

    ## maximum parallelism: Nx speed
    #grouped = floor(t/(nprocs()-2))
    #println(grouped)
    #v = vcat(pmap((i)->[patiencesort(randperm(n)) for j = 1:grouped], 1:t/grouped)...)

    w = (v-2sqrt(n))/n^(1/6)
    return hist(w, -5:dx:2)
end
@time grid, count = tracywidomlis(t, n, dx)

## Plot
using Winston
p = FramedPlot()
h = Histogram(count/(t*dx), step(grid))
h.x0 = first(grid)
include("../1/tracywidom.jl")
add(p, h)
if isinteractive()
    Winston.display(p)
else
    file(p, "tracywidomlis.png")
end


struct OrnsteinUhlenbeck{T1,T2,T3}
  Θ::T1
  μ::T2
  σ::T3
end
# http://www.math.ku.dk/~susanne/StatDiff/Overheads1b.pdf
function (X::OrnsteinUhlenbeck)(dW,W,dt,u,p,t,rng) #dist
  if typeof(dW) <: AbstractArray
    rand_val = wiener_randn(rng,dW)
  else
    rand_val = wiener_randn(rng,typeof(dW))
  end
  drift = X.μ .+ (W[end] .- X.μ) .* exp.(-X.Θ*dt)
  diffusion = X.σ .* sqrt.((1 .- exp.(-2X.Θ*dt))./(2X.Θ))
  drift .+ rand_val .* diffusion .- W[end]
end

#=
http://www.tandfonline.com/doi/pdf/10.1080/14697688.2014.941913?needAccess=true
http://www.tandfonline.com/doi/full/10.1080/14697688.2014.941913?src=recsys
=#

#=
(qX + r) = Θ(μ-x) = Θμ - Θx
q = -Θ
r = Θμ
https://arxiv.org/pdf/1011.0067.pdf page 18
=#
function ou_bridge(dW,ou,W,W0,Wh,q,hu,p,t,rng) end
function ou_bridge!(rand_vec,ou,W,W0,Wh,q,h,u,p,t,rng) end

function OrnsteinUhlenbeckProcess(Θ,μ,σ,t0,W0,Z0=nothing;kwargs...)
  ou = OrnsteinUhlenbeck(Θ,μ,σ)
  NoiseProcess{false}(t0,W0,Z0,ou,nothing;kwargs...)
end

struct OrnsteinUhlenbeck!{T1,T2,T3}
  Θ::T1
  μ::T2
  σ::T3
end

function (X::OrnsteinUhlenbeck!)(rand_vec,W,dt,u,p,t,rng) #dist!
  wiener_randn!(rng,rand_vec)
  @.. rand_vec = X.μ+(W[end]-X.μ)*exp(-X.Θ*dt) + rand_vec*X.σ*sqrt((1-exp.(-2*X.Θ.*dt))/(2*X.Θ)) - W[end]
end
function OrnsteinUhlenbeckProcess!(Θ,μ,σ,t0,W0,Z0=nothing;kwargs...)
  ou = OrnsteinUhlenbeck!(Θ,μ,σ)
  NoiseProcess{true}(t0,W0,Z0,ou,nothing;kwargs...)
end


using ArrayFire

maxIterations = 500
gridSize = 2048
xlim = [-0.748766713922161, -0.748766707771757]
ylim = [ 0.123640844894862,  0.123640851045266]

x = range( xlim[1], xlim[2], length = gridSize )
y = range( ylim[1], ylim[2], length = gridSize )

xGrid = [i for i in x, j in y]
yGrid = [j for i in x, j in y]

z0 = xGrid + im*yGrid

function mandelbrotCPU(z0, maxIterations)
    z = copy(z0)
    count = ones( size(z) )

    for n in 1:maxIterations
        z .= z.*z .+ z0
        count .+= abs.( z ).<=2
    end
    count = log.( count )
end

function mandelbrotGPU(z0, maxIterations)
    z = z0
    count = ones(AFArray{Float32}, size(z) )

    for n in 1:maxIterations
        z = z .* z .+ z0
        count = count + (abs(z)<= 2)
    end
    sync(log( count ))
end

# warmup
count = mandelbrotCPU(z0, 1)

cpu_time = @elapsed count = mandelbrotCPU(z0, maxIterations)

count .-= minimum(count)
count ./= maximum(count)
img = AFArray(Array{Float32}(count))

ArrayFire.image(img)

count = mandelbrotGPU(AFArray(z0), 1)
count = mandelbrotGPU(AFArray(z0), maxIterations)
gpu_time = @elapsed count = mandelbrotGPU(AFArray(z0), maxIterations)

ArrayFire.figure(2)
count -= min_all(count)[1]
count /= max_all(count)[1]
img = AFArray{Float32}(count)

ArrayFire.image(img)

@show cpu_time, gpu_time, cpu_time/gpu_time


using Gen
using GenDirectionalStats
import Rotations
import Random
using Plots

###########################################
# plots of distributions on 3D directions #
###########################################

function show_uniform_3d_direction()
    p = plot(aspect_ratio=:equal, title="uniform_3d_dir()")
    for i in 1:1000
        x = uniform_3d_direction()
        scatter!(p, [x.v[1]], [x.v[2]], [x.v[3]], w=1, label=nothing, color="red",
            markersize=5, markerstrokealpha=0.0, alpha=0.1)
    end
    xlims!(-1.0, 1.0)
    ylims!(-1.0, 1.0)
    zlims!(-1.0, 1.0)
    return p
end

function show_vmf_3d_direction(mu, k)
    p = plot([0.0, mu.v[1]], [0.0, mu.v[2]], [0.0, mu.v[3]], w=3, label=nothing, color="blue", aspect_ratio=:equal, title="vmf_3d_dir($(mu.v), $k)")
    for i in 1:1000
        x = vmf_3d_direction(mu, k)
        scatter!(p, [x.v[1]], [x.v[2]], [x.v[3]], w=1, label=nothing, color="red",
            markersize=5, markerstrokealpha=0.0, alpha=0.1)
    end
    xlims!(-1.0, 1.0)
    ylims!(-1.0, 1.0)
    zlims!(-1.0, 1.0)
    return p
end

function show_3d_direction_dists()
    println("generating plots for 3D directions...")
    mu = UnitVector3(1.0, 0.0, 0.0)
    p1 = show_vmf_3d_direction(mu, 0.1)
    p2 = show_vmf_3d_direction(mu, 1.0)
    p3 = show_vmf_3d_direction(mu, 10.0)
    p4 = show_vmf_3d_direction(mu, 100.0)
    mu = UnitVector3(0.0, -1.0, 0.0)
    p5 = show_vmf_3d_direction(mu, 0.1)
    p6 = show_vmf_3d_direction(mu, 1.0)
    p7 = show_vmf_3d_direction(mu, 10.0)
    p8 = show_vmf_3d_direction(mu, 100.0)
    plot(
        show_uniform_3d_direction(), p1, p2, p3, p4,
        show_uniform_3d_direction(), p5, p6, p7, p8, layout=(2, 5), size=(2000, 800))
    savefig("vmf_3d_direction.png")
end

show_3d_direction_dists()

##########################################
# plots of distributions on 3D rotations #
##########################################

function scatter_rot3!(p, rot::Rot3)
    x_axis = rot[:,1]
    y_axis = rot[:,2]
    z_axis = rot[:,3]
    scatter!(p, [x_axis[1]], [x_axis[2]], [x_axis[3]], w=1, label=nothing, color="red",
        markersize=5, markerstrokealpha=0.0, alpha=0.1)
    scatter!(p, [y_axis[1]], [y_axis[2]], [y_axis[3]], w=1, label=nothing, color="green",
        markersize=5, markerstrokealpha=0.0, alpha=0.1)
    scatter!(p, [z_axis[1]], [z_axis[2]], [z_axis[3]], w=1, label=nothing, color="blue",
        markersize=5, markerstrokealpha=0.0, alpha=0.1)
end

function plot_rot3!(p, rot::Rot3; alpha=1.0, w=5)
    x_axis = rot[:,1]
    y_axis = rot[:,2]
    z_axis = rot[:,3]
    plot!(p, [0.0, x_axis[1]], [0.0, x_axis[2]], [0.0, x_axis[3]], label=nothing, color="red", w=w, alpha=alpha)
    plot!(p, [0.0, y_axis[1]], [0.0, y_axis[2]], [0.0, y_axis[3]], label=nothing, color="green", w=w, alpha=alpha)
    plot!(p, [0.0, z_axis[1]], [0.0, z_axis[2]], [0.0, z_axis[3]], label=nothing, color="blue", w=w, alpha=alpha)
end

function show_uniform_3d_rotation()
    p = plot(aspect_ratio=:equal, title="uniform_rot3()")
    for i in 1:250
        rot = uniform_rot3()
        scatter_rot3!(p, rot)
        if i == 1
            plot_rot3!(p, rot; alpha=0.5, w=2)
        end
    end
    xlims!(-1.0, 1.0)
    ylims!(-1.0, 1.0)
    zlims!(-1.0, 1.0)
    return p
end

function show_vmf_3d_rotation(mu, k)
    p = plot(aspect_ratio=:equal, title="vmf_rot3(I, $k)")
    plot_rot3!(p, mu)
    for i in 1:250
        rot = vmf_rot3(mu, k)
        scatter_rot3!(p, rot)
        if i == 1
            plot_rot3!(p, rot; alpha=0.5, w=2)
        end
    end
    xlims!(-1.0, 1.0)
    ylims!(-1.0, 1.0)
    zlims!(-1.0, 1.0)
    return p
end

function show_uniform_vmf_3d_rotation(mu, k, prob_outlier)
    p = plot(aspect_ratio=:equal, title="uniform_vmf_rot3(I, $k, $prob_outlier)")
    plot_rot3!(p, mu)
    for i in 1:250
        rot = uniform_vmf_rot3(mu, k, prob_outlier)
        scatter_rot3!(p, rot)
        if i == 1
            plot_rot3!(p, rot; alpha=0.5, w=2)
        end
    end
    xlims!(-1.0, 1.0)
    ylims!(-1.0, 1.0)
    zlims!(-1.0, 1.0)
    return p
end

function show_3d_rotation_dists()
    println("generating plots for 3D rotations...")
    mu = Rot3(1.0, 0.0, 0.0, 0.0)
    p1 = show_vmf_3d_rotation(mu, 0.1)
    p2 = show_vmf_3d_rotation(mu, 1.0)
    p3 = show_vmf_3d_rotation(mu, 10.0)
    p4 = show_vmf_3d_rotation(mu, 100.0)
    p5 = show_vmf_3d_rotation(mu, 1000.0)
    p6 = show_uniform_vmf_3d_rotation(mu, 1000.0, 0.5)
    plot(
        p1, p2, p3, p4, p5,
        layout=(1, 5), size=(2000, 400))
    savefig("vmf_3d_rotation_1.png")
    plot(
        show_uniform_3d_rotation(), p6,
        layout=(1, 2), size=(1000, 400))
    savefig("vmf_3d_rotation_2.png")
end

show_3d_rotation_dists()

#############################################
# plots of distributions on plane rotations #
#############################################

function scatter_rot2!(p, rot::Rot2)
    rotmat = Rotations.RotMatrix{2}(rot)
    x_axis = rotmat[:,1]
    y_axis = rotmat[:,2]
    scatter!(p, [x_axis[1]], [x_axis[2]], w=1, label=nothing, color="red",
        markersize=5, markerstrokealpha=0.0, alpha=0.1)
    scatter!(p, [y_axis[1]], [y_axis[2]], w=1, label=nothing, color="green",
        markersize=5, markerstrokealpha=0.0, alpha=0.1)
end

function plot_rot2!(p, rot::Rot2; alpha=1.0, w=5)
    x_axis = rot[:,1]
    y_axis = rot[:,2]
    plot!(p, [0.0, x_axis[1]], [0.0, x_axis[2]], label=nothing, color="red", w=w, alpha=alpha)
    plot!(p, [0.0, y_axis[1]], [0.0, y_axis[2]], label=nothing, color="green", w=w, alpha=alpha)
end

function show_uniform_2d_rotation()
    p = plot(aspect_ratio=:equal, title="uniform_rot2()")
    for i in 1:250
        rot = uniform_rot2()
        scatter_rot2!(p, rot)
        if i == 1
            plot_rot2!(p, rot; alpha=0.5, w=2)
        end
    end
    xlims!(-1.1, 1.1)
    ylims!(-1.1, 1.1)
    return p
end

function show_vmf_2d_rotation(mu, k)
    p = plot(aspect_ratio=:equal, title="von_mises_rot2(I, $k)")
    plot_rot2!(p, mu)
    for i in 1:250
        rot = von_mises_rot2(mu, k)
        scatter_rot2!(p, rot)
        if i == 1
            plot_rot2!(p, rot; alpha=0.5, w=2)
        end
    end
    xlims!(-1.1, 1.1)
    ylims!(-1.1, 1.1)
    return p
end

function show_2d_rotation_dists()
    println("generating plots for 2D rotations...")
    mu = Rotations.RotMatrix{2}(0.0)
    @time p1 = show_vmf_2d_rotation(mu, 0.1)
    @time p2 = show_vmf_2d_rotation(mu, 1.0)
    @time p3 = show_vmf_2d_rotation(mu, 10.0)
    @time p4 = show_vmf_2d_rotation(mu, 100.0)
    @time p5 = show_vmf_2d_rotation(mu, 1000.0)
    @time plot(
        show_uniform_2d_rotation(), p1, p2, p3, p4, p5,
        layout=(2, 3), size=(900, 600))
    @time savefig("vmf_2d_rotation.png")
end

show_2d_rotation_dists()


using JuMP
function test_jump_sdp(solver)
    @testset "Semidefinite Programming" begin
        @testset "SDP1" begin
            # Problem SDP1 - sdo1 from MOSEK docs
            # From Mosek.jl/test/mathprogtestextra.jl, under license:
            #   Copyright (c) 2013 Ulf Worsoe, Mosek ApS
            #   Permission is hereby granted, free of charge, to any person obtaining a copy of this
            #   software and associated documentation files (the "Software"), to deal in the Software
            #   without restriction, including without limitation the rights to use, copy, modify, merge,
            #   publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons
            #   to whom the Software is furnished to do so, subject to the following conditions:
            #   The above copyright notice and this permission notice shall be included in all copies or
            #   substantial portions of the Software.
            #
            #     | 2 1 0 |
            # min | 1 2 1 | . X + x1
            #     | 0 1 2 |
            #
            #
            # s.t. | 1 0 0 |
            #      | 0 1 0 | . X + x1 = 1
            #      | 0 0 1 |
            #
            #      | 1 1 1 |
            #      | 1 1 1 | . X + x2 + x3 = 1/2
            #      | 1 1 1 |
            #
            #      (x1,x2,x3) in C^3_q
            #      X in C_sdp

            m = Model(solver=solver)

            @variable(m, x[1:3])
            @constraint(m, x in MOI.SecondOrderCone(3))
            @variable(m, X[1:3, 1:3], PSD)

            C = [2 1 0
                 1 2 1
                 0 1 2]
            A1 = [1 0 0
                  0 1 0
                  0 0 1]
            A2 = [1 1 1
                  1 1 1
                  1 1 1]

            @objective(m, Min, vecdot(C, X) + x[1])

            @constraint(m, vecdot(A1, X) + x[1] == 1)
            @constraint(m, vecdot(A2, X) + x[2] + x[3] == 1/2)

            JuMP.solve(m)

            @test JuMP.isattached(m)
            @test JuMP.hasvariableresult(m)
            @test JuMP.terminationstatus(m) == MOI.Success
            @test JuMP.primalstatus(m) == MOI.FeasiblePoint

            @test JuMP.objectivevalue(m) ≈ 0.705710509 atol=1e-6

            xv = JuMP.resultvalue.(x)
            Xv = JuMP.resultvalue.(X)
            @test vecdot(C, Xv) + xv[1] ≈ 0.705710509 atol=1e-6

            @test eigmin(Xv) > -1e-6
        end

        @testset "Nonsensical SDPs" begin
            m = Model()
            @test_throws ErrorException @variable(m, unequal[1:5,1:6], PSD)
            # Some of these errors happen at compile time, so we can't use @test_throws
            @test macroexpand(:(@variable(m, notone[1:5,2:6], PSD))).head == :error
            @test macroexpand(:(@variable(m, oneD[1:5], PSD))).head == :error
            @test macroexpand(:(@variable(m, threeD[1:5,1:5,1:5], PSD))).head == :error
            @test macroexpand(:(@variable(m, psd[2] <= rand(2,2), PSD))).head == :error
            @test macroexpand(:(@variable(m, -ones(3,4) <= foo[1:4,1:4] <= ones(4,4), PSD))).head == :error
            @test macroexpand(:(@variable(m, -ones(3,4) <= foo[1:4,1:4] <= ones(4,4), Symmetric))).head == :error
            @test macroexpand(:(@variable(m, -ones(4,4) <= foo[1:4,1:4] <= ones(4,5), Symmetric))).head == :error
            @test macroexpand(:(@variable(m, -rand(5,5) <= nonsymmetric[1:5,1:5] <= rand(5,5), Symmetric))).head == :error
        end

        # min o                    max y + X11
        # Q11 - 1   = Q22        [y-X12-X21  0     [0 0
        #                             0     -y] <=  0 0]
        # [1   Q11
        #  Q11 o  ] >= 0          -X[2,2] = 1
        # Q >= 0                        y free
        # o free                        X <= 0
        #   @testset "Just another SDP" begin
        #       model = Model(solver=CSDPSolver(printlevel=0))
        #       @variable(model, Q[1:2, 1:2], PSD)
        #       c1 = @constraint(model, Q[1,1] - 1 == Q[2,2])
        #       @variable(model, objective)
        #       T = [1 Q[1,1]; Q[1,1] objective]
        #       @test_throws ErrorException SDConstraint(T, 1)
        #       c2 = JuMP.addconstraint(model, SDConstraint(T, 0))
        #       @objective(model, Min, objective)
        #
        #       JuMP.solve(m)
        #
        #       @test JuMP.terminationstatus(m) == MOI.Success
        #       @test JuMP.primalstatus(m) == MOI.FeasiblePoint
        #
        #       @test JuMP.resultvalue(Q) ≈ [1 0; 0 0] atol=1e-3
        #       @test JuMP.objectivevalue(model) ≈ 1 atol=1e-4
        #       @test JuMP.resultvalue(objective) ≈ 1 atol=1e-4
        #       @test getdual(objective) ≈, 0, atol=1e-5
        #       @test getdual(Q) ≈ [0 0; 0 2] atol=1e-3
        #       @test getdual(c1) ≈ 2, atol=1e-4 # y
        #       @test getdual(c2) ≈ [-1 1; 1 -1] atol=1e-3 # X
        #   end

        # The four following tests are from Example 2.11, Example 2.13 and Example 2.27 of:
        # Blekherman, G., Parrilo, P. A., & Thomas, R. R. (Eds.).
        # Semidefinite optimization and convex algebraic geometry SIAM 2013

        # Example 2.11
        @testset "SDP variable and optimal objective not rational" begin
            #       solver = fixscs(solver, 7000000)
            m = Model(solver=solver)
            @variable(m, X[1:2,1:2], PSD)
            c = @constraint(m, X[1,1]+X[2,2] == 1)
            @objective(m, Min, 2*X[1,1]+2*X[1,2])
            #       @test all(isnan.(getdual(X)))

            JuMP.solve(m)

            @test JuMP.terminationstatus(m) == MOI.Success
            @test JuMP.primalstatus(m) == MOI.FeasiblePoint

            @test JuMP.objectivevalue(m) ≈ 1-sqrt(2) atol=1e-5
            @test JuMP.resultvalue.(X) ≈ [(2-sqrt(2))/4 -1/(2*sqrt(2)); -1/(2*sqrt(2)) (2+sqrt(2))/4] atol=1e-4
            #       @test getdual(X) ≈ [1+sqrt(2) 1; 1 sqrt(2)-1] atol=1e-4
            #       @test getdual(c) ≈ 1-sqrt(2) atol=1e-5
        end

        #   # Example 2.13
        #   @testset "SDP constraint and optimal objective not rational with $solver" for solver in sdp_solvers
        #       solver = fixscs(solver, 7000000)
        #       m = Model(solver=solver)
        #       @variable(m, y)
        #       c = @SDconstraint(m, [2-y 1; 1 -y] >= 0)
        #       @objective(m, Max, y)
        #       @test all(isnan, getdual(c))
        #       status = solve(m)
        #
        #       @test status == :Optimal
        #       @test isapprox(getobjectivevalue(m), 1-sqrt(2), atol=1e-5)
        #       @test isapprox(getvalue(y), 1-sqrt(2), atol=1e-5)
        #
        #       X = getdual(c)
        #       @test isapprox(getdual(c), [(2-sqrt(2))/4 -1/(2*sqrt(2)); -1/(2*sqrt(2)) (2+sqrt(2))/4], atol=1e-4)
        #       @test isapprox(getdual(y), 0, atol=1e-5)
        #   end
        #
        #   # Example 2.27
        #   # min X[1,1]   max y
        #   # 2X[1,2] = 1  [0 y     [1 0
        #   # X ⪰ 0         y 0] ⪯   0 0]
        #   # The dual optimal solution is y=0 and there is a primal solution
        #   # [ eps  1/2
        #   #   1/2  1/eps]
        #   # for any eps > 0 however there is no primal solution with objective value 0.
        #   @testset "SDP with dual solution not attained with $solver" for solver in sdp_solvers
        #       solver = fixscs(solver, 7000000)
        #       m = Model(solver=solver)
        #       @variable(m, y)
        #       c = @SDconstraint(m, [0 y; y 0] <= [1 0; 0 0])
        #       @objective(m, Max, y)
        #       @test all(isnan, getdual(c))
        #       status = solve(m)
        #
        #       if contains(string(typeof(solver)),"MosekSolver")
        #           # Mosek returns Stall on this instance
        #           # Hack until we fix statuses in MPB
        #           JuMP.fillConicDuals(m)
        #       else
        #           @test status == :Optimal
        #       end
        #       @test isapprox(getobjectivevalue(m), 0, atol=1e-5)
        #       @test isapprox(getvalue(y), 0, atol=1e-5)
        #
        #       X = getdual(c)
        #       @test isapprox(X[1,1], 0, atol=1e-5)
        #       @test isapprox(X[1,2], 1/2, atol=1e-5)
        #       @test isapprox(X[2,1], 1/2, atol=1e-5)
        #       @test isapprox(getdual(y), 0, atol=1e-5)
        #   end

        @testset "SDP with primal solution not attained" begin
            #       solver = fixscs(solver, 7000000)
            m = Model(solver=solver)
            @variable(m, X[1:2,1:2], PSD)
            c = @constraint(m, 2*X[1,2] == 1)
            @objective(m, Min, X[1,1])
            #       @test all(isnan, getdual(X))
            status = solve(m)

            @test JuMP.terminationstatus(m) in [ MOI.Success, MOI.SlowProgress ]
            @test JuMP.primalstatus(m) == MOI.FeasiblePoint

            @test JuMP.objectivevalue(m) ≈ 0 atol=1e-5
            Xval = JuMP.resultvalue.(X)
            #Mosek.writedata(m.solverinstance.task,"jump_sdo_1_stalls.task")
            #showall(m.solverinstance.task)
            #showall(m.solverinstance.task[Sol(Mosek.MSK_SOL_ITR)])

            @test Xval[1,1] ≈ 0 atol=1e-5
            @test Xval[1,2] ≈ 1/2 atol=1e-5
            @test Xval[2,1] ≈ 1/2 atol=1e-5

            #       @test isapprox(getdual(X), [1 0; 0 0], atol=1e-4)
            #       @test isapprox(getdual(c), 0, atol=1e-5)
        end

        @testset "No constraint" begin
            m = Model(solver=solver)
            @variable(m, X[1:3,1:3], PSD)
            @objective(m, Min, trace(X))

            JuMP.solve(m)

            status = solve(m)

            @test JuMP.terminationstatus(m) == MOI.Success
            @test JuMP.primalstatus(m) == MOI.FeasiblePoint

            @test abs(JuMP.objectivevalue(m)) < 1e-5
            @test norm(JuMP.resultvalue.(X)) < 1e-5
            #@test isapprox(getdual(X), eye(3), atol=1e-5)
        end

    end
end


@testset "Linear algebra" begin

    @testset "SVector as a (mathematical) vector space" begin
        c = 2
        v1 = @SVector [2,4,6,8]
        v2 = @SVector [4,3,2,1]

        @test v1 + c === @SVector [4,6,8,10]
        @test v1 - c === @SVector [0,2,4,6]
        @test v1 * c === @SVector [4,8,12,16]
        @test v1 / c === @SVector [1.0,2.0,3.0,4.0]

        @test v1 + v2 === @SVector [6, 7, 8, 9]
        @test v1 - v2 === @SVector [-2, 1, 4, 7]
    end

    @testset "eye()" begin
        @test eye(SMatrix{2,2,Int}) === @SMatrix [1 0; 0 1]
        @test eye(SMatrix{2,2}) === @SMatrix [1.0 0.0; 0.0 1.0]
        @test eye(SMatrix{2}) === @SMatrix [1.0 0.0; 0.0 1.0]

        @test eye(MMatrix{2,2,Int})::MMatrix == @MMatrix [1 0; 0 1]
        @test eye(MMatrix{2,2})::MMatrix == @MMatrix [1.0 0.0; 0.0 1.0]
        @test eye(MMatrix{2})::MMatrix == @MMatrix [1.0 0.0; 0.0 1.0]
    end

    @testset "cross()" begin
        @test cross(SVector(1,2,3), SVector(4,5,6)) === SVector(-3, 6, -3)
    end

    @testset "transpose() and conj()" begin
        @test conj(SVector(1+im, 2+im)) === SVector(1-im, 2-im)

        @test @SVector([1, 2, 3]).' === @SMatrix([1 2 3])
        @test @SMatrix([1 2; 0 3]).' === @SMatrix([1 0; 2 3])
        @test @SMatrix([1 2 3; 4 5 6]).' === @SMatrix([1 4; 2 5; 3 6])

        @test @SVector([1, 2, 3])' === @SMatrix([1 2 3])
        @test @SMatrix([1 2; 0 3])' === @SMatrix([1 0; 2 3])
        @test @SMatrix([1 2 3; 4 5 6])' === @SMatrix([1 4; 2 5; 3 6])
    end

    @testset "vcat() and hcat()" begin
        @test vcat(SVector(1,2,3), SVector(4,5,6)) === SVector(1,2,3,4,5,6)
        @test hcat(SVector(1,2,3), SVector(4,5,6)) === @SMatrix [1 4; 2 5; 3 6]

        @test vcat(@SMatrix([1;2;3]), SVector(4,5,6)) === @SMatrix([1;2;3;4;5;6])
        @test vcat(SVector(1,2,3), @SMatrix([4;5;6])) === @SMatrix([1;2;3;4;5;6])
        @test hcat(@SMatrix([1;2;3]), SVector(4,5,6)) === @SMatrix [1 4; 2 5; 3 6]
        @test hcat(SVector(1,2,3), @SMatrix([4;5;6])) === @SMatrix [1 4; 2 5; 3 6]

        @test vcat(@SMatrix([1;2;3]), @SMatrix([4;5;6])) === @SMatrix([1;2;3;4;5;6])
        @test hcat(@SMatrix([1;2;3]), @SMatrix([4;5;6])) === @SMatrix [1 4; 2 5; 3 6]
    end
end


# an example on mixed integer programming
#
#   maximize x + 2 y + 5 z
#
#   s.t.  x + y + z <= 10
#         x + 2 y + z <= 15
#
#         x is continuous: 0 <= x <= 5
#         y is integer: 0 <= y <= 10
#         z is binary
#

using CPLEX
using JuMP

m = Model(solver=CplexSolver())
@defVar(m, 5 >= x >= 0)
@defVar(m, 10 >= y >= 0, Int)
@defVar(m, z, Bin)
@setObjective(m, Max, x + 2*y + 5*z)
@addConstraint(m, x + y + z <= 10)
@addConstraint(m, x + 2*y + z <= 15)
status = solve(m)
println("walltime = $(MathProgBase.getsolvetime(getInternalModel(m)))")


using Revise
using AdFem
using PyCall
using LinearAlgebra
using PyPlot
using SparseArrays
using MAT
np = pyimport("numpy")

λ = 2.0
μ = 0.5
η = 1/12

β = 1/4; γ = 1/2
a = b = 0.1
m = 10
n = 5
h = 0.01
NT = 20
Δt = 1/NT
bdedge = []
for j = 1:n
  push!(bdedge, [(j-1)*(m+1)+m+1 j*(m+1)+m+1])
end
bdedge = vcat(bdedge...)

bdnode = Int64[]
for j = 1:n+1
  push!(bdnode, (j-1)*(m+1)+1)
end

G = [1/Δt+μ/η -μ/3η 0.0
  -μ/3η 1/Δt+μ/η-μ/3η 0.0
  0.0 0.0 1/Δt+μ/η]
S = [2μ/Δt+λ/Δt λ/Δt 0.0
    λ/Δt 2μ/Δt+λ/Δt 0.0
    0.0 0.0 μ/Δt]
invG = inv(G)
H = invG*S

M = compute_fem_mass_matrix1(m, n, h)
Zero = spzeros((m+1)*(n+1), (m+1)*(n+1))
M = [M Zero;Zero M]

K = compute_fem_stiffness_matrix(H, m, n, h)
C = a*M + b*K # damping matrix

L = M + γ*Δt*C + β*Δt^2*K
L, Lbd = fem_impose_Dirichlet_boundary_condition(L, bdnode, m, n, h)

Forces = zeros(NT, 2(m+1)*(n+1))
for i = 1:NT
  T = eval_f_on_boundary_edge((x,y)->0.1, bdedge, m, n, h)
  T = [T zeros(length(T))]
  rhs = compute_fem_traction_term(T, bdedge, m, n, h)

  if i*Δt>0.5
    rhs = zero(rhs)
  end
  Forces[i, :] = rhs
end

a = zeros(2(m+1)*(n+1))
v = zeros(2(m+1)*(n+1))
d = zeros(2(m+1)*(n+1))
U = zeros(2(m+1)*(n+1),NT+1)
Sigma = zeros(NT+1, 4m*n, 3)
Varepsilon = zeros(NT+1, 4m*n, 3)
for i = 1:NT
  @info i
    global a, v, d
    F = compute_strain_energy_term(Sigma[i,:,:]*invG/Δt, m, n, h) - K * U[:,i]
    # @show norm(compute_strain_energy_term(Sigma[i,:,:]*invG/Δt, m, n, h)), norm(K * U[:,i])
    rhs = Forces[i,:] - Δt^2 * F

    td = d + Δt*v + Δt^2/2*(1-2β)*a
    tv = v + (1-γ)*Δt*a
    rhs = rhs - C*tv - K*td
    rhs[[bdnode; bdnode.+(m+1)*(n+1)]] .= 0.0

    a = L\rhs
    d = td + β*Δt^2*a
    v = tv + γ*Δt*a
    U[:,i+1] = d

    Varepsilon[i+1,:,:] = eval_strain_on_gauss_pts(U[:,i+1], m, n, h)
    Sigma[i+1,:,:] = Sigma[i,:,:]*invG/Δt +  (Varepsilon[i+1,:,:]-Varepsilon[i,:,:])*(invG*S)
end

matwrite("U.mat", Dict("U"=>U'|>Array))

# visualize_displacement(U, m, n, h; name = "_eta$η", xlim_=[-0.01,0.5], ylim_=[-0.05,0.22])
# visualize_displacement(U, m, n, h;  name = "_viscoelasticity")
# visualize_stress(H, U, m, n, h;  name = "_viscoelasticity")

close("all")
figure(figsize=(15,5))
subplot(1,3,1)
idx = div(n,2)*(m+1) + m+1
plot((0:NT)*Δt, U[idx,:])
xlabel("time")
ylabel("\$u_x\$")

subplot(1,3,2)
idx = 4*(div(n,2)*m + m)
plot((0:NT)*Δt, Sigma[:,idx,1])
xlabel("time")
ylabel("\$\\sigma_{xx}\$")

subplot(1,3,3)
idx = 4*(div(n,2)*m + m)
plot((0:NT)*Δt, Varepsilon[:,idx,1])
xlabel("time")
ylabel("\$\\varepsilon_{xx}\$")
savefig("visco_eta$η.png")


using ForwardDiff

export  FermiOccupation,
        ∂ϵFermiOccupation,
        ∂μFermiOccupation,
        ∂TFermiOccupation

function FermiOccupation(ϵ,μ,T)
    kb = 8.6*10^(-5) #eV / K
    1/(exp( (ϵ-μ)/(kb*T)) +1 )
end

∂ϵFermiOccupation(ϵ,μ,T) = ForwardDiff.derivative( x -> FermiOccupation(x,μ,T))
∂μFermiOccupation(ϵ,μ,T) = ForwardDiff.derivative( x -> FermiOccupation(ϵ,x,T))
∂TFermiOccupation(ϵ,μ,T) = ForwardDiff.derivative( x -> FermiOccupation(ϵ,μ,x))

function Dx(kx,ky,Ham!,Matrix1,Matrix2,diff)
    dϵ = 1<<15
    ϵ = 1/dϵ
    Ham!(kx+ϵ,ky,Matrix1) ; Ham!(kx-ϵ,ky,Matrix2)
    @. diff = (Matrix1-Matrix2)*dϵ/2
end

function Dy(kx,ky,Ham!,Matrix1,Matrix2,diff)
    dϵ = 1<<15
    ϵ = 1/dϵ
    Ham!(kx,ky+ϵ,Matrix1) ; Ham!(kx,ky-ϵ,Matrix2)
    @. diff = (Matrix1-Matrix2)*dϵ/2
end


const unary_function_list = [:+,:-,:cos,:sin,:tan,:acos,:asin,:atan,:cosh,:sinh,:tanh,:acosh,:asinh,:atanh,:sqrt,:exp,:log,:expm1,:erf,:erfc]
const binary_function_list = [:+,:-,:*,:/,:^,:hypot,:atan2]
const unary_diff = Dict(
    (+) => x -> 1,
    (-) => x -> -1,
    (cos) => x -> -sin(x),
    (sin) => x -> cos(x),
    (tan) => x -> tan(x)^2 + 1,
    (acos) => x -> -1/sqrt(-x^2 + 1),
    (asin) => x -> 1/sqrt(-x^2 + 1),
    (atan) => x -> 1/(x^2 + 1),
    (cosh) => x -> sinh(x),
    (sinh) => x -> cosh(x),
    (tanh) => x -> -tanh(x)^2 + 1,
    (acosh) => x -> 1/sqrt(x^2 - 1),
    (asinh) => x -> 1/sqrt(x^2 + 1),
    (atanh) => x -> 1/(-x^2 + 1),
    (sqrt) => x -> 1/(2*sqrt(x)),
    (exp) => x -> exp(x),
    (log) => x -> 1/x,
    (expm1) => x -> exp(x),
    (erf) => x -> 2*exp(-x^2)/sqrt(pi),
    (erfc) => x -> -2*exp(-x^2)/sqrt(pi))

const binary1_diff = Dict(
    (+) => ( x, y ) -> 1,
    (-) => ( x, y ) -> 1,
    (*) => ( x, y ) -> y,
    (/) => ( x, y ) -> 1/y,
    (^) => ( x, y ) -> x^y*y/x,
    (hypot) => ( x, y ) -> x/sqrt(x^2 + y^2),
    (atan2) => ( x, y ) -> 1/(y*(x^2/y^2 + 1)))

const binary2_diff = Dict(
    (+) => ( x, y ) -> 1,
    (-) => ( x, y ) -> -1,
    (*) => ( x, y ) -> x,
    (/) => ( x, y ) -> -x/y^2,
    (^) => ( x, y ) -> x^y*log(x),
    (hypot) => ( x, y ) -> y/sqrt(x^2 + y^2),
    (atan2) => ( x, y ) -> -x/(y^2*(x^2/y^2 + 1)))


α = 0.3
δ = 0.1
β = 0.9
k_l = 2
k_u = 0.001
n = 10
kgrid = collect(range(k_l, stop = k_u, length =n))

v = zeros(n)
tolerance = 0.001
imax = 1000
iter = 1
vnew = zeros(n)
while maximum(abs.(v-vnew) > tolerance) && iter < imax
v=vnew
c = zeros(n,n)
for i in 1:n
    for j in 1:n
    c[i,j] = kgrid[i]^α + (1-δ)*kgrid[i] - kgrid[j]
        if c[i,j]<0
            c[i,j]=0
        end
    end
end
vnew = maximum(log.(c) .+ β*v' , dim = 2)
iter += 1
end


@doc raw"""
Penning trap with magnetic bottle in (x,y,z) coordinates.
Based on Yanyan Shi, Yajuan Sun, Yulei Wang, Jian Liu, Study of adaptive symplectic methods for
      simulating charged particle dynamics, Journal of Computational Dynamics 6, 429-448, 2019.

The covariant components of the vector potential are given by
```math
A (x,y,z) = B_0 / 2 \, ( -y , x, 0)^T - B_1 \, (xz, yz, (x^2 + y^2)/2 - z^2)^T,
```
resulting in the magnetic field with covariant components
```math
B (x,y,z) = B_0 \, ( 0, 0, 1)^T - B_1 \, ( yz^2 - y^3 / 6, x^3 / 6, xyz )^T ,
```
and the electrostatic potential given by
```math
\varphi (x,y,z) = E_0 \, ( x^2 / 2 + y^2 / 2 - z^2) ,
```
resulting in the electric field with covariant components
```math
E (x,y,z) = E_0 \, ( x, y, - 2 z)^T .
```

Parameters:
* `B₀`: B-field strength
* `Bₚ`: B-field perturbation strength
* `E₀`: E-field strength
"""
module PenningTrapBottle

    using RecipesBase

    import ..ElectromagneticFields
    import ..ElectromagneticFields: CartesianEquilibrium, code
    import ..AnalyticCartesianField: X, Y, Z

    export  PenningTrapBottleEquilibrium

    const DEFAULT_B₀ = 100.0
    const DEFAULT_Bₚ = 200.0
    const DEFAULT_E₀ = 10.0

    struct PenningTrapBottleEquilibrium{T <: Number} <: CartesianEquilibrium
        name::String
        B₀::T
        Bₚ::T
        E₀::T

        function PenningTrapBottleEquilibrium{T}(B₀::T, Bₚ::T, E₀::T) where T <: Number
            new("PenningTrapBottleEquilibrium", B₀, Bₚ, E₀)
        end
    end

    PenningTrapBottleEquilibrium(B₀::T=DEFAULT_B₀, Bₚ::T=DEFAULT_Bₚ, E₀::T=DEFAULT_E₀) where T <: Number = PenningTrapBottleEquilibrium{T}(B₀, Bₚ, E₀)

    function init(B₀=DEFAULT_B₀, Bₚ=DEFAULT_Bₚ, E₀=DEFAULT_E₀)
        PenningTrapBottleEquilibrium(B₀, E₀, Bₚ)
    end

    macro code(B₀=DEFAULT_B₀, Bₚ=DEFAULT_Bₚ, E₀=DEFAULT_E₀)
        code(init(B₀, Bₚ, E₀); escape=true)
    end

    function Base.show(io::IO, equ::PenningTrapBottleEquilibrium)
        print(io, "Penning trap with magnetic bottle in (x,y,z) coordinates with\n")
        print(io, "  B₀ = ", equ.B₀, "\n")
        print(io, "  Bₚ = ", equ.Bₚ, "\n")
        print(io, "  E₀ = ", equ.E₀)
    end

    ElectromagneticFields.A₁(x::AbstractVector, equ::PenningTrapBottleEquilibrium) = - equ.B₀ / 2 * Y(x,equ) - equ.Bₚ * (Y(x,equ) * Z(x,equ)^2 - Y(x,equ)^3 / 6)
    ElectromagneticFields.A₂(x::AbstractVector, equ::PenningTrapBottleEquilibrium) = + equ.B₀ / 2 * X(x,equ) - equ.Bₚ * X(x,equ)^3 / 6
    ElectromagneticFields.A₃(x::AbstractVector, equ::PenningTrapBottleEquilibrium) = - equ.Bₚ * X(x,equ) * Y(x,equ) * Z(x,equ)
    ElectromagneticFields.φ(x::AbstractVector, equ::PenningTrapBottleEquilibrium) = - equ.E₀ * (X(x,equ)^2 / 2 + Y(x,equ)^2 / 2 - Z(x,equ)^2)

    ElectromagneticFields.get_functions(::PenningTrapBottleEquilibrium) = (X=X, Y=Y, Z=Z)

end


# Package StatsBase, Distributions, JuMP, Gurobi are needed!
function BernsteinEstimate_MIQO(Y,m,a,b,k,e,T)
     n = size(Y,1)
     M = Model(with_optimizer(Gurobi.Optimizer, TimeLimit = T, MIPGap = e))
     #M = Model(solver = GurobiSolver(TimeLimit = T, MIPGap = e))

     @variable(M, w[1:m]>=0)
     @variable(M, z[1:m-1], Bin)
     @variable(M, σ_plus[1:m-2], Bin)
     @variable(M, σ_minus[1:m-2], Bin)
     @variable(M, α[1:m-2,1:2], Bin)

    # Uncomment this line for doing a warmstart
    #setvalue(w, BernsteinEstimate_MD(Y,m,a,b,0)')

    @constraint(M, sum(w)==1);
    # A big constant
    K = 1
    # k-modality constraint
    if k>=1
        for j=1:m-1
            @constraint(M, w[j]<=w[j+1]+K*z[j]);
            @constraint(M, w[j]>=w[j+1]-K*(1-z[j]));
        end
        for j=1:m-2
            @constraint(M, z[j+1]-z[j]==σ_plus[j]-σ_minus[j]);
            @constraint(M, σ_plus[j]<=K*α[j,1]);
            @constraint(M, σ_minus[j]<=K*α[j,2]);
            @constraint(M, α[j,1]+α[j,2]==1);
        end
        @constraint(M, sum(σ_plus)<=k);
        @constraint(M, sum(σ_minus)<=k);
    end

    ϵ = 3/(8*n)
    Ecdf = ecdf(vec(Y))

    B = zeros(n,m)
    f = zeros(n,1)
    for i = 1:n
        f[i] = Ecdf(Y[i])
    end

    diag_L = (1 ./((Ecdf(Y).+ϵ).*(1+ϵ.-Ecdf(Y)))).^0.5
    L = diagm(0 => diag_L)

    for i=1:n
        for j=1:m
            B[i,j] = betacdf(j, m-j+1, (Y[i]-a)/(b-a))
        end
    end

    #svd approximation - for numerical stability
    Q1 =(L*B)';
    Q = Q1*Q1';

    SVD_Q1 = svd(Q1)
    ids = SVD_Q1.S.>10e-3;
    vals = SVD_Q1.S[ids];
    uu = SVD_Q1.U[:,ids];

    Qhat = uu*diagm(0 => vals.^2)*uu'
    Qhat = (Qhat + Qhat')/2 + 0.001*Matrix{Float64}(I,m,m);

    V = (L*f)'*(L*B)

    c1 = sum(diag(V))
    c2 = sum(diag(Qhat))

    @objective(M, Min, -2*sum(V[i]*w[i] for i=1:m) + sum(Qhat[i,j]*w[i]*w[j] for i=1:m for j=1:m) )

    optimize!(M)

    return JuMP.value.(w)

end


function BinomPutDelta(K, T, S, sig, r, N)

  PutArray, AssetArray = (GenAmericanBinomialPut(K, T, S, sig, r, N))

  delta = (PutArray[1,2] - PutArray[2,2])/(AssetArray[1,2]-AssetArray[2,2])
  delta

end


using LinearAlgebra, ForwardDiff
using Convex, SCS, ECOS

n = 3
m = 3
k = n + m + 1

c = zeros(n)
A = Diagonal(ones(n))
b = zeros(n)

"Convex.jl"
x = Variable(n)
prob = minimize(c' * x)
prob.constraints += norm(b - A * x) <= 0.0
@time solve!(prob, ECOS.Optimizer)

@show prob.status
@show x.value
@show prob.constraints[1].dual
prob.optval

Q = Array([zeros(n, n) A' c;
     -A zeros(m, m) b;
     -c' -b'      0.0])
Q_vec = vec(Q)

function F(z, Q_vec)
    ũ = z[1:k]
    u = z[k .+ (1:k)]
    v = z[2 * k .+ (1:k)]

    [(I + reshape(Q_vec, k, k)) * ũ - (u + v);
     u - P_soc(ũ - v);
     ũ - u]
end

z = rand(3k)
F(z, Q_vec)
Fz(x) = F(x, Q_vec)
FQ(x) = F(z, x)
Jz = ForwardDiff.jacobian(Fz, z)
JQ = ForwardDiff.jacobian(FQ, Q_vec)

norm((Jz' * Jz) \ (Jz' * JQ), Inf)


"""
This function takes in some reactance (at frequency fₒ) and determines the values for a series LC pair that resonates at fᵣ and has the same reactance,X at fₒ

    For future reference, see Everitt, Communication Engineering, P. 267
"""
function findEquivLC(X,fₒ,fᵣ)
    ωₒ = 2*π*fₒ
    ωᵣ = 2*π*fᵣ
    L = ωₒ/(ωₒ^2 - ωᵣ^2)*X
    C = (ωₒ^2 - ωᵣ^2) / (ωₒ*ωᵣ^2*X)
    return L,C
end


#=
import Pkg; Pkg.activate(".")
=#
using Distributions
import Random
using RCall
import LinearAlgebra

include("MCMC.jl")

prettymat(X) = begin show(stdout, "text/plain", X); println(); end

### MAIN ###
Random.seed!(1)

# easy = true  # bivariate normal
easy = false  # 250-variate normal

if easy
  K = 3
  S = MCMC.eye(K)
  S[1,2] = S[2,1] = -0.8
  S[1,3] = S[3,1] = 0.5
else
  K = 50
  S = rand(InverseWishart(K, MCMC.eye(K)))
end

mvn = MvNormal(randn(K) * 10, S)
logprob(x::Vector{Float64}) = logpdf(mvn, x)

init = randn(K)

lp(s::Vector{Float64}) = 0.0
out, propcov = MCMC.mcmc(init, logprob, lp, thin=10, batchsize=100,
                         temper_power=3,
                         propcov_init=MCMC.eye(length(init)) * 1e-3,
                         nburn=10000, nmcmc=1000, max_temper=1000.0,
                         verbose=1)

acc_rate = size(unique(out, dims=1), 1) / size(out, 1)

ll = [logprob(out[i, :]) for i in 1:size(out, 1)]
R"plot($ll, type='l')"

G = 8
println("cov(out):")
prettymat(round.(cov(out[:, 1:G]), digits=3))
println("cov(mvn):")
prettymat(round.(cov(mvn)[1:G, 1:G], digits=3))

[mean(out, dims=1)' mean(mvn)]

R"rcommon::plotPosts($(out)[, 1:5])";
println("Acceptance rate: $(acc_rate)")


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENCE in the project root.
# ------------------------------------------------------------------

@userplot ParetoPlot

@recipe function f(p::ParetoPlot)
  # get user input
  data = p.args[1]

  x = sort(data, rev=true)
  n = length(x)
  logp = log.([i/(n+1) for i in 1:n])
  logx = log.(x)

  seriestype --> :scatter
  xguide --> "-log(i/(n+1))"
  yguide --> "log(x*)"

  -logp, logx
end


"""
`Base.length(tab::ODERKTableau)`

Defines the length of a Runge-Kutta method to be the number of stages.
"""
Base.length(tab::ODERKTableau) = tab.stages

"""
`stability_region(z,tab::ODERKTableau)`

Calculates the stability function from the tableau at `z`. Stable if <1.

```math
r(z) = \\frac{\\det(I-zA+zeb^T)}{\\det(I-zA)}
```
"""
stability_region(z,tab::ODERKTableau) = det(eye(tab.stages)- z*tab.A + z*ones(tab.stages)*tab.α')/det(eye(tab.stages)-z*tab.A)


# Some basic functions to interact with the data.

function gencov(d::Integer)
  A = rand(d, d)
  eye(d) + (A + A')/2
end

function plotclouds(cloudA, cloudB)
  scatter(cloudA[1,:], cloudA[2,:], label="-1")
  scatter!(cloudB[1,:], cloudB[2,:], label="1")
end

function errorrate(classA::MvNormal,
                   classB::MvNormal,
                   w::AbstractVector, n::Integer)
  testA = rand(classA, n)
  testB = rand(classB, n)
  X = [[testA testB]' ones(2*n)]
  Y = [-1.*ones(n); ones(n)]
  sum(sign.(X*w).≠Y)/n
end

function drawborder(f)
  xs = linspace(xlims()..., 100)
  ys = linspace(ylims()..., 100)
  contour!(xs, ys, f, levels=[0], colorbar=false, ls=:dash)
end

function plotdualitygap(m::Integer,
                        μ::Real,
                        numstepsarray::AbstractVector{T}) where T<:Integer
  dualgaps = [m/μ^i for i=0:length(numstepsarray)]
  numstepssumsarray = reduce(
    (arr, n) -> [arr; arr[end]+n] ,[0], numstepsarray)
  dualgaps2 = repeat(dualgaps, inner=2)[1:end-1]
  numstepssumsarray2 = repeat(numstepssumsarray, inner=2)[2:end]
  plot(numstepssumsarray2, dualgaps2, yscale=:log10,
       xlabel="Newton iterations", ylabel="duality gap")
end


function create_graph(start_node, end_node, link_length)
	@assert length(start_node)==length(end_node)

	no_node = max(maximum(start_node), maximum(end_node))
	no_arc = length(start_node)

	graph = Graph(no_node)
	distmx = Inf*ones(no_node, no_node)
	for i=1:no_arc
		add_edge!(graph, start_node[i], end_node[i])
		distmx[start_node[i], end_node[i]] = link_length[i]
	end
	return graph, distmx
end

function get_shortest_path(start_node::Array, end_node::Array, link_length::Array, origin::Int, destination::Int)
	@assert length(start_node)==length(end_node)
	@assert length(start_node)==length(link_length)

	graph, distmx = create_graph(start_node, end_node, link_length)

	state = dijkstra_shortest_paths(graph, origin, distmx)

	path = get_path(state, origin, destination)
	x = get_vector(path, start_node, end_node)

	len = state.dists[destination]

	return path, x, len
end

function get_path(state, origin, destination)
	_current = destination
	_rpath = _current
	_parent = -1

	while _parent != origin
		_parent = state.parents[_current]
		_rpath = [_rpath; _parent]
		_current = _parent
	end

	_path = zeros(Int, length(_rpath))
	for i=length(_rpath):-1:1
		_path[length(_rpath)-i+1] = _rpath[i]
	end
	return _path
end

function get_vector(path, start_node, end_node)
	x = zeros(Int, length(start_node))

	for i=1:length(path)-1
		st = path[i]
		en = path[i+1]

		for j=1:length(start_node)
			if start_node[j]==st && end_node[j]==en
				x[j] = 1
				break
			end
		end

	end
	return x
end


#=
#rigi according to 2016 paper
y=679520.05
x=212273.44
#using BenchmarkTools
@btime latLong($x,$y)

@code_warntype latLong(x,y)
lat,long=latLong(x,y)
latA,longA=approxLatLong(x,y)
lat_seconds=((lat-47)*60-3)*60
long_seconds=((long-8)*60-29)*60
latA_seconds=((latA-47)*60-3)*60
longA_seconds=((longA-8)*60-29)*60
-3*60

latLong(x,y)
a,b=latLong(x,y)
a,b=latLong(200_000,600_000)

a,b=latLong(-100_000,100_000)
compound(a),compound(b)

a,b=latLong(191775.030,602030.680)
compound(a),compound(b)
#we have CH1903+ ellipsoid coordinates
X,Y,Z=geoCentricCartesianCoordinates(a,b,897.361)
#X,Y,Z=geoCentricCartesianCoordinates(a,b,500)

#convert to ETRS89
Xnew=X+674.374
Ynew=Y+15.056
Znew=Z+405.346

aNew,bNew,hNew=geoCentricCartesianToEllipsoidCoordinates(Xnew,Ynew,Znew)
compound(aNew),compound(bNew)

@code_warntype geoCentricCartesianToEllipsoidCoordinates(Xnew,Ynew,Znew)
@btime geoCentricCartesianToEllipsoidCoordinates($Xnew,$Ynew,$Znew)

xx=191775.030
yy=602030.680
@btime LV03toETRS89($xx,$yy)
@code_warntype LV03toETRS89(xx,yy)

=#

#=
#Bern, Sidlerstrasse 5
#approxCHCoord(lat,long)
#http://geodesy.geo.admin.ch/reframe/lv03towgs84?easting=600000&northing=200000&altitude=550.0&format=json

lat,long = 46.95108288705888,7.438632502714563
chCoord(lat,long)

y,x=600000,200000
lat_res,long_res=approxLatLong(x,y)
lat-lat_res
long-long_res

approxLatLong(200_000,600_000)

=#

#=
    @btime approxCHCoord($lat,$long)
    @code_warntype approxCHCoord(lat,long)

    x,y=approxCHCoord(lat,long)
    approxLatLong(x,y)

    @btime approxLatLong($x,$y)
    @code_warntype approxLatLong(x,y)
=#

#=
#Rigi
#47°03'28.95659233"N 8°29'11.11127154"E
lat=47+(3*60+28.95659233)/3600
long=8+(29*60+11.11127154)/3600
#y,x=679602.034,212422.148
x,y=chCoord(lat,long)
lat_res,long_res=approxLatLong(x,y)
lat-lat_res
long-long_res
x-chCoord(lat_res,long_res)[1]
y-chCoord(lat_res,long_res)[2]
=#
#=
    @code_warntype chCoord(lat,long)
    @btime chCoord($lat,$long) #198.879ns (199.165,199.165)
=#