setprecision(100)
T = BigFloat
PI = T(π)
πi = PI * 1im
πi2 = 2πi

function γ(z, ϵ)
    z = z - 1
    half = T(1) / T(2)
    a = floor(max(real(z), 2)) + half
    while true
        ϵ̂ = √(a) / (2 * PI)^(a + half) / (real(z) + a)
        ϵ̂ < ϵ && break
        a += 1
    end

    coef = 1
    k = 1
    fac = T(1)
    while k < a
        c_k = (a - k)^(k - half) * exp(a - k) * (-1)^(k - 1) / √(2 * PI) / fac
        fac *= k
        coef += c_k / (z + k)
        k += 1
    end
    coef * (z + a)^(z + half) * √(2 * PI) / exp(z + a)
end

function brentq(f, xa, xb, xtol, rtol, iter)
    xpre, xcur, xblk, fpre, fcur, fblk, spre, scur = xa, xb, 0, f(xa), f(xb), 0, 0, 0

    fpre * fcur > 0 && return 0
    fpre == 0 && return xpre
    fcur == 0 && return xcur

    for _ in 1:iter
        if fpre * fcur < 0
            xblk = xpre
            fblk = fpre
            spre = xcur - xpre
            scur = spre
        end

        if abs(fblk) < abs(fcur)
            xpre = xcur
            xcur = xblk
            xblk = xpre

            fpre = fcur
            fcur = fblk
            fblk = fpre
        end

        Δ = (xtol + rtol * abs(xcur)) / 2
        sbis = (xblk - xcur) / 2
        (fcur == 0 || abs(sbis) < Δ) && return xcur

        if abs(spre) < Δ && abs(fcur) < abs(fpre)
            if xpre == xblk
                stry = -fcur * (xcur - xpre) / (fcur - fpre)
            else
                dpre = (fpre - fcur) / (xpre - xcur);
                dblk = (fblk - fcur) / (xblk - xcur);
                stry = -fcur * (fblk * dblk - fpre * dpre) / (dblk * dpre * (fblk - fpre))
            end

            if 2 * abs(stry) < min(abs(spre), 3 * abs(sbis) - Δ)
                spre = scur
                scur = stry
            else
                spre, scur = sbis, sbis
            end
        else
            spre, scur = sbis, sbis
        end

        xpre = xcur
        fpre = fcur

        if abs(scur) > Δ
            xcur += scur
        elseif sbis > 0
            xcur += Δ
        else
            xcur -= Δ
        end

        fcur = f(xcur)
    end

    xcur
end

function I0(s, ϵ)
#     println("[I0] s = $s")
    g(z, s) = z^(-s) * exp(z * z * πi)
    H(w) = 1 / (1 - exp(πi2 * w))

    function f(z, s)
        a = exp(z * πi)
        g(z, s) / (a - 1 / a)
    end

    σ = real(s)
    z_s = √(s / πi2)
    n = max(floor(real(z_s) - imag(z_s)), 1)
#     println("z_s = $z_s")

    Δ = max(0, log(n^(-σ) / ϵ) / (2 * PI))
#     println("σ = $σ, n = $n, n^-σ = $(n^-σ)")
    direction = (1 + 1im) / √ T(2)

    criterion(α) = abs(f(n + 0.5 + direction * α * Δ, s)) - ϵ
    α_1 = brentq(criterion, 0,  2, ϵ, ϵ, 20)
    α_2 = brentq(criterion, 0, -2, ϵ, ϵ, 20)

    z_1 = n + 1/2 + direction * α_1 * Δ
    z_2 = n + 1/2 + direction * α_2 * Δ
#     println("α_1 = $α_1, α_2 = $α_1, $(Δ), $σ, $(n^(-σ))")

    z_l = 0
    z_r = 0
    h = 0

    m = 1
    found = false
    for i in 1:1000
        h = (z_2 - z_1) / m
        inv_2h = 1 / (2 * h)
        z_l = sqrt(inv_2h^2 + s / πi2) + inv_2h
        z_r = sqrt(inv_2h^2 + s / πi2) - inv_2h
        err_l = g(z_l, s) * exp((z_1 - z_l) / h * πi2)
        err_r = g(z_r, s) * exp((z_r - z_1) / h * πi2)
        if abs(err_l) ≤ ϵ && abs(err_r) ≤ ϵ
            found = true
            break
        end
        m = ceil(m * 1.1)
#         println("err_l = $err_l, err_r = $err_r, g(z_l, s) = $(g(z_l, s)), z_l = $z_l, z_r = $z_r")
    end
    @assert found

    n_l = ceil(real(z_l) - imag(z_l))
    n_r = floor(real(z_r) - imag(z_r))

    fs2(i) = i^(-s) * H((i - z_1) / h)
    fs3(i) = i^(-s) * H((z_1 - i) / h)
#     println("[ζ] n = $n, m = $m, n_l = $n_l, n_r = $n_r")

    s0 = sum(@. (1:n)^(-s))
    s1 = sum(@. f((0:m) * h + z_1, s))
    s2 = sum(@. fs2(n_l:n))
    s3 = sum(@. fs3((n+1):n_r))

    s0 + s1 * h - s2 + s3
end

χ(s, ϵ) = (2 * PI)^s / cos(PI * s / 2) / γ(s, ϵ) / 2
# ζ(s, ϵ) = I0(s, ϵ / 2) + χ(s, ϵ) * conj(I0(1 - conj(s), ϵ / 2))
function ζ(s, ϵ)
    ret = I0(s, ϵ / 2) + χ(s, ϵ) * conj(I0(1 - conj(s), ϵ / 2))
    println("[ζ] s = $s, ϵ = $ϵ, ret = $ret, $(χ(s, ϵ))")
    ret
end

function erfc(z, ϵ)
    if abs(z) > 2
        h = π / √(log(6 / ϵ))
        K = ceil(√(log(1 / ϵ)) / h)
        ret = 1 / z^2
        PI = T(π)
        for k in 1:K
            ret += 2 * exp(-h^2 * k^2) / (z^2 + h^2 * k^2)
        end
        return ret * exp(-z^2) * h * z / PI + 2 / (1 - exp(2 * PI * z / h))
    else
        if z ≥ 0
            s = 1
        else
            s = -1
            z = -z
        end

        ϵ0 = ϵ / 2
        K = ceil(log2(1 / ϵ) / 2)
        ϵ1 = ϵ0 / 2K
        ϵ1 = 2^floor(log2(ϵ1))

        t = z
        k = T(0)
        S = T(0)
        while abs(t) / (2k + 1) > ϵ0
            S = S + t / (2k + 1)
            t = -z^2 * t / (k + 1) # error in paper
            k += 1
        end
        S += t / (2k + 1)
        1 - s * S * 2 / √π
    end
end

using SpecialFunctions

function find_T(x, λ, σ, ϵ)
#     println("[find T] x = $x, λ = $λ, σ = $σ, ϵ = $ϵ")
    # only to determine bounds, so it's fine to use Float64 here.

    # σ is real and 1 ≤ σ ≤ 2

    E_Σ(T, x, λ, σ) = exp(λ^2 * σ^2 / 2) / 2π * x^σ * log(zeta(σ)) * expint(λ^2 * T^2 / 2)

    # now we solve min T s.t. E_Σ(T, x, λ, σ) < 0.75ϵ
    # T = √(2u) / λ where E_1(u) < 0.75 ϵ / exp(λ^2 * σ^2 / 2) / (2 * π) * log(gamma(σ))
    limit = 0.75ϵ / (exp(λ^2 * σ^2 / 2) / (2 * π) * log(zeta(σ))) / x^σ
#     println("[find T] limit = $limit, expint(0) = $(expint(0))")
    f(u) = expint(u) - limit

    u = brentq(f, λ^2 / 2, -log(limit), ϵ, ϵ, 100)
#     println("[debug] est u = $(-log(limit)), u = $u")
    T = √(2u) / λ
    println("[param] T = $T, residual = $(E_Σ(T, x, λ, σ)) < ϵ = $ϵ")

    T
end

function quad_π_star(x, σ, λ, ϵ)
    ϵ = ϵ / 4
    T = find_T(x, λ, σ, ϵ)

    h_L = 2 * π * (σ - 1) / (log(x / ϵ) + λ^2 / 2 + 1 / x)
    h_R = 2 * π / (log(x / 2) + σ * λ^2 + λ * √(2 * σ * log(x / 2) + σ^2 * λ^2 + 2 * log(3.4 / ϵ)))
    h = min(h_L, h_R)
    println("[param] h = $h, h_L = $h_L, h_R = $h_R")

    ϵ_0 = ϵ / x^σ / log(x)

    Ψ(s, x, λ, ϵ) = exp((λ * s)^2 / 2) * x^s * log(ζ(s, ϵ)) / s

    results = @. real(Ψ(σ + 1im * h * (1:ceil(T/h)+10), x, λ, ϵ_0))
#     println(results)
    s = sum(results)
    h / π * (Ψ(σ, x, λ, ϵ_0) / 2 + s)
end

function get_primes(l, r)
    a = ones(Bool, r - l + 1)
    for i in 2:Int64(floor(√r))
        li = max((l - 1) ÷ i + 1, 2)
        ri = r ÷ i
        for j in li : ri
            a[i * j - l + 1] = 0
        end
    end

    ret = []
    for i in max(l, 2):r
        if a[i - l + 1]
            push!(ret, i)
        end
    end
    ret
end

function Δ(x, λ, ϵ)
    Φ(ρ, ϵ) = erfc(ρ / √2, ϵ) / 2
    ϕ(u, x, λ, ϵ) = Φ(log(u / x) / λ, ϵ)
    E_plus(u, x, λ, ϵ) = x * exp(λ^2 / 2) * Φ(log(u / x) / λ - λ, ϵ / x) - u * Φ(log(u / x) / λ, ϵ / 8 / u)
    E_minus(u, x, λ, ϵ) = u * Φ(-log(u / x) / λ, ϵ / 8 / u) - x * exp(λ^2 / 2) * Φ(λ - log(u / x) / λ, ϵ / x)

    ϵ = ϵ / 2
    x1 = floor(x)
    if E_minus(x1, x, λ, ϵ) > ϵ
        g(u) = E_minus(u, x, λ, ϵ) - 0.75 * ϵ
        u = brentq(g, 2, x, 1e-8, 1e-8, 100)
        x1 = floor(u)
    end

    x2 = ceil(x)
    if E_plus(x2, x, λ, ϵ) > ϵ
        f(u) = E_plus(u, x, λ, ϵ) - 0.75 * ϵ
        u = brentq(f, x * 2, x, 1e-8, 1e-8, 100)
        x2 = ceil(u)
    end
    println("[Δ] x1 = $x1, x2 = $x2")
    println("$(E_minus(900, x, λ, ϵ))")
    println("$(E_plus(x2, x, λ, ϵ))")

    N = (x2 - x1 + 1) + ceil(√x2)
    ϵ_1 = ϵ * 2^ceil(-log2(N))

    ret = 0
    for p in get_primes(Int64(x1), Int64(x2))
        ret -= ϕ(T(p), x, λ, ϵ_1)
        if p ≤ x
            ret += 1
        end
    end
    for p in get_primes(2, Int64(floor(√x2)))
        m = 2
        while p^m < x2
            if p^m < x1
                ret -= T(1 / m)
            else
                ret -= ϕ(p^m, x, λ, ϵ_1) / m
            end
            m += 1
        end
    end

    ret
end

function analytic_π(x)
    μ = 1 + 0
    σ_LR = 2
    σ_0 = (μ + 1/2) / (μ - 1/2)
#     σ = min(σ_LR, σ_0)
    σ = 1.5
    println("[param] σ = $σ")

    λ = 1 / (√x * log(x)^0.25)
    println("[param] λ = $λ")

    ϵ = 1/4
    pi_star = quad_π_star(x, σ, λ, ϵ / 2)
    Delta = Δ(x, λ, ϵ / 2)
    println("pi_star = $pi_star, Delta = $Delta")
    round(Delta + pi_star)
end

print(analytic_π(10^4))


using DIVAnd, PyPlot

mask,(pm,pn),(xi,yi) = DIVAnd_squaredom(2,-1:0.01:1)

x        = [-0.5,0.5]; # x-positions
y        = [ 0. ,0. ]; # y-positions
f        = [ 1. ,1. ]; # values
epsilon2 = [ 0.5,1. ]; # expected error variances (normalized)

len = 0.2 # correlation length-scale
fi,s = DIVAndrun(mask,(pm,pn),(xi,yi),(x,y),f,len,epsilon2);

# Plotting
clf();
pcolor(xi,yi,fi; cmap="jet");
axis("equal")
colorbar();


#
using DiffMesh
using ModelingToolkit, NeuralPDE
using DiffEqFlux, LinearAlgebra, Plots
using ParameterSchedulers:Scheduler,Sin,SinDecay2,SinExp
using GalacticOptim
#----------------------------------------------------------------------------#
nprt = 2
sdf = (x,y) -> @. (x^2 + y^2) - 1

nprt = 6
pou  = nnPOU2D(nprt,sdf;depth=4,act=relu)

nx = ny = 4
bsp = TensorProdCosines(nprt,nx,ny)
net = MeshNet2D(pou,bsp)

p,re = Flux.destructure(net)
#--------------------------------------#
function utrue(x,y)

    r2 = @. x*x + y*y

    u .= @. 1 - r2
    u .= @. sin(pi*x)
    u .= @. x*y^3 + sin(r2)  + exp(x^2*y)
    return u
end

it = 0 # iter
#----------------------------------------------------------------------------#
function makeplot(phi)

    fpp = phi(xpp,ypp)
    fp  = reshape(fpp,np,np)
    ep  = up .- fp

    p1 = heatmap(Ωp.*fp)
    p2 = heatmap(Ωp.*up)
    p3 = heatmap(Ωp.*ep)

    p1 = plot!(p1,title="Pred")
    p2 = plot!(p2,title="True")
    p3 = plot!(p3,title="Err ")

    plt = plot(p1,p2,p3)

    display(plt)
    return
end
#--------------------------------------#
function callback(p,l)
    phi = re(p)

    leastSq!(phi,xe,ye,ue)

    fpp = phi(xpp,ypp)
    epp = (upp .- fpp) .* Ωpp
    er  = norm(epp,Inf)

    global it += 1
    println("it: $it, loss: $l, pointwise er: $er")
    (it % 10) == 0 && makeplot(phi)
    (l < 1e-6) && return true
    return false
end
#----------------------------------------------------------------------------#
# NeuralPDE setup

@parameters x y
@variables u(..)
Dx  = Differential(x)
Dy  = Differential(y)
Dxx = Dx ∘ Dx
Dyy = Dy ∘ Dy

eq = [Dxx(u(x,y)) + Dyy(u(x,y))  ~ 1f0]
bc = [u(-1f0) ~ 1f0
     ,u( 1f0) ~ 1f0]

domain = [x ∈ ModelingToolkit.Interval(-1f0,1f0),
          y ∈ ModelingToolkit.Interval(-1f0,1f0)]

strategy = GridTraining(0.1)
discretization = PhysicsInformedNN(net,strategy)
pdesys = PDESystem(eq,bc,domain,[x,y],[u],name=:PoissonDisk)
prob = discretize(pdesys,discretization)
#--------------------------------------#
res = GalacticOptim.solve(prob,ADAM(1e-2),cb=callback,maxiters=200)
#----------------------------------------------------------------------------#
return


using Planck
using Test
using Unitful

@testset "Planck.jl" begin
    @testset "Consistent unit conversions" begin
        @test ustrip(u"W/m^3", blackbody(545u"nm", 6000u"K")) ≈ blackbody(545e-9, 6000)
        @test ustrip(u"W/m^2/Hz", blackbody(500u"THz", 6000u"K")) ≈ 3.441e-8 rtol=1e-3
    end
    @testset "struct interface" begin
        bb = Blackbody(6000)
        @test bb(545e-9) ≈ 3.079e13 rtol=1e-3
        @test bb(Float32, 545e-9) ≈ 3.079f13 rtol=1f-3

        qbb = Blackbody(6000u"K")
        @test qbb(545u"nm") ≈ 3.079e13u"W/m^3" rtol=1e-3
        @test qbb(500u"THz") ≈ 3.441e-8u"W/m^2/Hz" rtol=1e-3
    end
end


import QuadratureRules: GaussLegendreQuadrature, LobattoLegendreQuadrature

"""
Lagrange basis on the interval [0..1].
"""
struct Lagrange{T, BT, XT <: AbstractVector{T}} <: Basis{T}
    b::BT
    x::XT

    denom::XT
    diffs::Matrix{T}
    vdminv::Matrix{T}

    function Lagrange{T}(x::XT) where {T, XT <: SVector}
        n = length(x)
        denom = zeros(n)
        diffs = zeros(n,n)

        for i in eachindex(x)
            local p = 1
            for j in eachindex(x)
                diffs[i,j] = x[i] - x[j]
                if i ≠ j
                    p *= diffs[i,j]
                end
            end
            denom[i] = 1/p
        end

        sdenom = SVector{n}(denom)

        b = collect(y -> sdenom[j] * mapreduce(i -> i ≠ j ? (y - x[i]) : one(T), *, eachindex(x)) for j in eachindex(sdenom))

        new{T, typeof(b), typeof(x)}(b, x, sdenom, diffs, vandermonde_matrix_inverse(x))
    end

    Lagrange{T}(x::Vector) where {T} = Lagrange{T}(SVector{length(x),T}(x))
    Lagrange{T}(x::AbstractVector) where {T} = Lagrange{T}(collect(x))
end

Lagrange(x::AbstractVector{T}) where {T} = Lagrange{T}(x)

LagrangeGauß(n) = Lagrange(GaussLegendreQuadrature(n).nodes)
LagrangeLobatto(n) = Lagrange(LobattoLegendreQuadrature(n).nodes)

(L::Lagrange)(x::Number, j::Integer) = L.b[j](x)

basis(L::Lagrange) = L.b
nodes(L::Lagrange) = L.x
nbasis(L::Lagrange) = length(basis(L))
nnodes(L::Lagrange) = length(nodes(L))
order(L::Lagrange)  = nnodes(L)
degree(L::Lagrange) = nnodes(L) - 1

Base.eltype(::Lagrange{T}) where {T} = T
Base.eachindex(L::Lagrange) = eachindex(L.b)
Base.axes(L::Lagrange) = (Inclusion(0..1), eachindex(L))
ContinuumArrays.grid(L::Lagrange) = nodes(L)

Base.hash(L::Lagrange, h::UInt) = hash(L.x, h)
Base.:(==)(L1::Lagrange, L2::Lagrange) = (L1.x == L2.x)
Base.isequal(L1::Lagrange{T1}, L2::Lagrange{T2}) where {T1,T2} = (T1 == T2 && L1 == L2)
Base.isapprox(L1::Lagrange, L2::Lagrange; kwargs...) = isapprox(L1.x, L2.x; kwargs...)

Base.getindex(L::Lagrange, x::Number, j::Integer) = L(x,j)
Base.getindex(L::Lagrange, x::Number,  ::Colon) = [b(x) for b in L.b]
Base.getindex(L::Lagrange, X::AbstractVector, j::Integer) = L.(X,j)
Base.getindex(L::Lagrange, X::AbstractVector,  ::Colon) = [b(x) for x in X, b in L.b]

## Derivative

@simplify *(D::Derivative, L::Lagrange) = Mul(D,L)

const LagrangeDerivative = QMul2{<:Derivative,<:Lagrange}

function _eval(D::LagrangeDerivative, x::DT, j::Int) where {DT}
    local L = D.B
    local T = promote_type(eltype(L), DT)
    local d::T = 0

    for l in eachindex(L)
        if l ≠ j
            z = 1 / L.diffs[j,l]
            for i in eachindex(L)
                i ≠ j && i ≠ l ? z *= (x - L.x[i]) / L.diffs[j,i] : nothing
            end
            d += z
        end
    end
    return d
end

Base.getindex(D::LagrangeDerivative, x::Number, j::Integer) = _eval(D, x, j)
Base.getindex(D::LagrangeDerivative, x::Number,  ::Colon) = [_eval(D, x, j) for j in eachindex(D.B)]
Base.getindex(D::LagrangeDerivative, X::AbstractVector, j::Integer) = [_eval(D, x, j) for x in X]
Base.getindex(D::LagrangeDerivative, X::AbstractVector,  ::Colon) = [_eval(D, x, j) for x in X, j in eachindex(D.B)]

Base.adjoint(L::Lagrange) = Derivative(axes(L,1)) * L


using StaticArrays, NLSolvers, LinearAlgebra
OPT_PROBS = Dict()

#### Himmelblau

OPT_PROBS["himmelblau"] = Dict()
OPT_PROBS["himmelblau"]["array"] = Dict()

function himmelblau!(x)
    fx = (x[1]^2 + x[2] - 11)^2 + (x[1] + x[2]^2 - 7)^2
    return fx
end
function himmelblau_batched_f!(X)
    F = map(himmelblau!, X)
    return F
end
function himmelblau_batched_f!(F, X)
    map!(himmelblau!, F, X)
    return F
end
function himmelblau_g!(∇f, x)
    ∇f[1] = 4.0 * x[1]^3 + 4.0 * x[1] * x[2] -
    44.0 * x[1] + 2.0 * x[1] + 2.0 * x[2]^2 - 14.0
    ∇f[2] = 2.0 * x[1]^2 + 2.0 * x[2] - 22.0 +
    4.0 * x[1] * x[2] + 4.0 * x[2]^3 - 28.0 * x[2]
    ∇f
end

function himmelblau_fg!(∇f, x)
    ∇f = himmelblau_g!(∇f, x)
    fx = himmelblau!(x)
    return fx, ∇f
end
function himmelblau_fgh!(∇f, ∇²f, x)
    ∇²f = himmelblau_h!(∇²f, x)
    fx, ∇f = himmelblau_fg!(∇f, x)
    return fx, ∇f, ∇²f
end
function himmelblau_h!(∇²f, x)
    ∇²f[1, 1] = 12.0 * x[1]^2 + 4.0 * x[2] - 44.0 + 2.0
    ∇²f[1, 2] = 4.0 * x[1] + 4.0 * x[2]
    ∇²f[2, 1] = ∇²f[1, 2]
    ∇²f[2, 2] = 2.0 + 4.0 * x[1] + 12.0 * x[2]^2 - 28.0
    return ∇²f
end
function himmelblau_hv!(hv, x, v)
    hv[1] = (12.0 * x[1]^2 + 4.0 * x[2] - 44.0 + 2.0)*v[1] + (4.0 * x[1] + 4.0 * x[2])*v[2]
    hv[2] =  (4.0 * x[1] + 4.0 * x[2])*v[1] + (2.0 + 4.0 * x[1] + 12.0 * x[2]^2 - 28.0)*v[2]
    return hv
end

OPT_PROBS["himmelblau"]["array"]["x0"] = [3.0, 1.0]
OPT_PROBS["himmelblau"]["array"]["mutating"] = ScalarObjective(himmelblau!, himmelblau_g!, himmelblau_fg!, himmelblau_fgh!, himmelblau_h!, himmelblau_hv!, himmelblau_batched_f!, nothing)

OPT_PROBS["himmelblau"]["staticarray"] = Dict()

function himmelblau_f(x)
    fx = (x[1]^2 + x[2] - 11)^2 + (x[1] + x[2]^2 - 7)^2
    return fx
end

function himmelblau_fg(∇f, x)
    ∇f = himmelblau_g(∇f, x)
    fx = himmelblau_f(x)
    return fx, ∇f
end

function himmelblau_g(∇f, x)
    ∇f1 = 4.0 * x[1]^3 + 4.0 * x[1] * x[2] -
    44.0 * x[1] + 2.0 * x[1] + 2.0 * x[2]^2 - 14.0
    ∇f2 = 2.0 * x[1]^2 + 2.0 * x[2] - 22.0 +
    4.0 * x[1] * x[2] + 4.0 * x[2]^3 - 28.0 * x[2]
    ∇f = @SVector([∇f1, ∇f2])
    return ∇f
end

function himmelblau_fgh(∇f, ∇²f, x)
    ∇²f11 = 12.0 * x[1]^2 + 4.0 * x[2] - 42.0
    ∇²f12 = 4.0 * x[1] + 4.0 * x[2]
    ∇²f21 = 4.0 * x[1] + 4.0 * x[2]
    ∇²f22 = 12.0 * x[2]^2 + 4.0 * x[1] - 26.0
    ∇²f = @SMatrix([∇²f11 ∇²f12; ∇²f21 ∇²f22])
    fx, ∇f = himmelblau_fg(∇f, x)

    return fx, ∇f, ∇²f
end

sl = @SVector([3.0,1.0])
state0newton = (@SVector([2.0,2.0]), I+sl*sl')
OPT_PROBS["himmelblau"]["staticarray"]["x0"] = sl
OPT_PROBS["himmelblau"]["staticarray"]["state0"] = state0newton
OPT_PROBS["himmelblau"]["staticarray"]["static"] = ScalarObjective(himmelblau_f, himmelblau_g, himmelblau_fg, himmelblau_fgh, nothing, nothing, nothing, nothing)

### Exponential

exponential!(x) = exp((2.0 - x[1])^2) + exp((3.0 - x[2])^2)
function exponential_g!(g, x)
    g[1] = -2.0 * (2.0 - x[1]) * exp((2.0 - x[1])^2)
    g[2] = -2.0 * (3.0 - x[2]) * exp((3.0 - x[2])^2)
    return g
end
function exponential_h!(H, x)
    H[1, 1] = 2.0 * exp((2.0 - x[1])^2) * (2.0 * x[1]^2 - 8.0 * x[1] + 9)
    H[1, 2] = 0.0
    H[2, 1] = 0.0
    H[2, 2] = 2.0 * exp((3.0 - x[2])^2) * (2.0 * x[2]^2 - 12.0 * x[2] + 19)
    return H
end
function exponential_hv!(Hv, x, v)
    Hv[1, 1] = (2.0 * exp((2.0 - x[1])^2) * (2.0 * x[1]^2 - 8.0 * x[1] + 9))*v[1]
    Hv[2, 2] = (2.0 * exp((3.0 - x[2])^2) * (2.0 * x[2]^2 - 12.0 * x[2] + 19))*v[2]
    return Hv
end
function exponential_fg!(g, x)
    fx = exponential!(x)
    g = exponential_g!(g, x)
    return fx, g
end
function exponential_fgh!(g, H, x)
    fx, g = exponential_fg!(g, x)
    H = exponential_h!(H, x)
    return fx, g, H
end

OPT_PROBS["exponential"] = Dict()
OPT_PROBS["exponential"]["array"] = Dict()
# Byttet om på x og H
OPT_PROBS["exponential"]["array"]["x0"] = [0.0, 0.0]
OPT_PROBS["exponential"]["array"]["mutating"] = ScalarObjective(exponential!, exponential_g!, exponential_fg!, exponential_fgh!, exponential_h!, exponential_hv!, nothing, nothing)

OPT_PROBS["laplacian"] = Dict()
OPT_PROBS["laplacian"]["array"] = Dict()
# Byttet om på x og H
OPT_PROBS["laplacian"]["array"]["x0(n)"] = n->zeros(n)

plap(U; n=length(U)) = (n-1) * sum((0.1 .+ diff(U).^2).^2) - sum(U) / (n-1)
plap1(U; n=length(U), dU = diff(U), dW = 4 .* (0.1 .+ dU.^2) .* dU) =
(n - 1) .* ([0.0; dW] .- [dW; 0.0]) .- ones(n) / (n-1)
precond(x::Vector) = precond(length(x))
precond(n::Number) = spdiagm(-1 => -ones(n-1), 0 => 2*ones(n), 1 => -ones(n-1)) * (n+1)
_f(x) = plap([0;x;0])
function _fg(∇f, x)
    fx = _f(x)
    copyto!(∇f, (plap1([0;x;0]))[2:end-1])
    fx, ∇f
end

OPT_PROBS["laplacian"]["array"]["mutating"] = OptimizationProblem(ScalarObjective(_f, nothing, _fg, nothing, nothing, nothing, nothing, nothing))

function theta(x)
    if x[1] > 0
        return atan(x[2] / x[1]) / (2.0 * pi)
    else
        return (pi + atan(x[2] / x[1])) / (2.0 * pi)
    end
end

function fletcher_powell_f(x)
    theta_x = theta(x)
    fx = 100.0 * ((x[3] - 10.0 * theta_x)^2 + (sqrt(x[1]^2 + x[2]^2) - 1.0)^2) + x[3]^2
    return fx
end
function fletcher_powell_g(∇f, x)
    T = eltype(x)
    theta_x = theta(x)

    if x[1]^2 + x[2]^2 == 0
        dtdx1 = T(0)
        dtdx2 = T(0)
    else
        dtdx1 = - x[2] / ( T(2) * pi * ( x[1]^2 + x[2]^2 ) )
        dtdx2 =   x[1] / ( T(2) * pi * ( x[1]^2 + x[2]^2 ) )
    end
    ∇f1 = -2000.0*(x[3]-10.0*theta_x)*dtdx1 +
        200.0*(sqrt(x[1]^2+x[2]^2)-1)*x[1]/sqrt( x[1]^2+x[2]^2 )
    ∇f2 = -2000.0*(x[3]-10.0*theta_x)*dtdx2 +
        200.0*(sqrt(x[1]^2+x[2]^2)-1)*x[2]/sqrt( x[1]^2+x[2]^2 )
    ∇f3 =  200.0*(x[3]-10.0*theta_x) + 2.0*x[3];
    ∇f = @SVector[∇f1, ∇f2, ∇f3]

    return ∇f
end
function fletcher_powell_fg(∇f, x)
    T = eltype(x)
    theta_x = theta(x)

    if x[1]^2 + x[2]^2 == 0
        dtdx1 = T(0)
        dtdx2 = T(0)
    else
        dtdx1 = - x[2] / ( T(2) * pi * ( x[1]^2 + x[2]^2 ) )
        dtdx2 =   x[1] / ( T(2) * pi * ( x[1]^2 + x[2]^2 ) )
    end
    ∇f1 = -2000.0*(x[3]-10.0*theta_x)*dtdx1 +
        200.0*(sqrt(x[1]^2+x[2]^2)-1)*x[1]/sqrt( x[1]^2+x[2]^2 )
    ∇f2 = -2000.0*(x[3]-10.0*theta_x)*dtdx2 +
        200.0*(sqrt(x[1]^2+x[2]^2)-1)*x[2]/sqrt( x[1]^2+x[2]^2 )
    ∇f3 =  200.0*(x[3]-10.0*theta_x) + 2.0*x[3];
    ∇f = @SVector[∇f1, ∇f2, ∇f3]

    fx = 100.0 * ((x[3] - 10.0 * theta_x)^2 + (sqrt(x[1]^2 + x[2]^2) - 1.0)^2) + x[3]^2

    return fx, ∇f
end
const sv3 = @SVector[0.1,0.0,0.0]

const fp_static = ScalarObjective(fletcher_powell_f, fletcher_powell_g, fletcher_powell_fg, nothing, nothing, nothing, nothing, nothing)

const state0 = (@SVector[-0.5, 0.0, 0.0], I+sv3*sv3')
OPT_PROBS["fletcher_powell"] = Dict()
OPT_PROBS["fletcher_powell"]["staticarray"] = Dict()
# Byttet om på x og H
OPT_PROBS["fletcher_powell"]["staticarray"]["x0"] = state0
OPT_PROBS["fletcher_powell"]["staticarray"]["static"] = OptimizationProblem(fp_static; inplace=false)


const libsigma = Pkg.dir("SDE", "deps", "libsigma")

#void fe_mu(double *mu, double *y, int N, int L, int K)
function fe_mu_c(y::Vector{Float64}, L, K)

	n = 2^(L-1)
	mu = zeros(Float64,2n-1 + K)
	N = length(y)
	product = ccall( (:fe_mu_at, libsigma),
                  Void,
                  (Ptr{Float64}, Ptr{Float64}, Int32, Int32),
                  mu, y, N, L)
	return mu;
end

#void fe_Sigma_at(double *S, double *y, int N, double dt, int L)
function fe_Sigma_c(y::Vector{Float64}, dt::Float64, L)
	n = 2^(L-1)
	N = length(y)
	S = zeros(Float64,(2n-1)^2)
	product = ccall( (:fe_Sigma_at,  libsigma),
                  Void,
                  (Ptr{Float64}, Ptr{Float64}, Int32, Float64, Int32),
                  S, y, N, dt, L)
	return reshape(S, 2n-1 , 2n-1)

end
function fe_SigmaB1_c(y::Vector{Float64}, dt::Float64, L)
	n = 2^(L-1)
	N = length(y)
	S = zeros(Float64,(2n)^2)
	product = ccall( (:fe_SigmaB1_at,  libsigma),
                  Void,
                  (Ptr{Float64}, Ptr{Float64}, Int32, Float64, Int32),
                  S, y, N, dt, L)
	return reshape(S, 2n , 2n)

end

fe_mu = fe_mu_c
fe_SigmaB1 = fe_SigmaB1_c
fe_Sigma = fe_Sigma_c


# function to build the fast application of the dual Lippmann-Schwinger equation
# we suppose u(x) = \int_D G(x,y) \sigma(y) dy
# and we solve ( \Delta + omega^2(1+q)) u = - omega^2q u^{inc}

type FastMDual
    # type to encapsulate the fast application of M = I + omega^2G*spadiagm(nu)
    GFFT :: Array{Complex128,2}
    nu :: Array{Float64,1}
    # number of points in the extended domain
    ne :: Int64
    me :: Int64
    # number of points in the original domain
    n  :: Int64
    m  :: Int64
    # frequency
    omega :: Float64
    quadRule :: String
    function FastMDual(GFFT,nu,ne,me,n,m,k; quadRule::String = "trapezoidal")
      return new(GFFT,nu,ne,me,n, m, k, quadRule)
    end
end

# we need to overload these functions for gmres
import Base.*
import Base.A_mul_B!
import Base.size

function size(M::FastMDual, dim::Int64)
  # function to returns the size of the underliying matrix (M.m*M.n)^2
  return M.m*M.n
end

function *(M::FastMDual, b::Array{Complex128,1})
    # function to overload the applyication of
    # M using a Toeplitz reduction via a FFT
    # dummy function to call fastconvolution
    return fastconvolution(M,b)
end

function A_mul_B!(Y,
                  M::FastMDual,
                  V)
    # in place matrix matrix multiplication
    @assert(size(Y) == size(V))
    # print(size(V))
    for ii = 1:size(V,2)
        Y[:,ii] = M*V[:,ii]
    end
end

@inline function fastconvolution(M::FastMDual, b::Array{Complex128,1})
    # function to overload the applyication of
    # M using a Toeplitz reduction via a FFT

    # computing omega^2 nu G*(b)
    B = M.omega^2*(M.nu.*FFTconvolution(M,b))

    # returning -b + omega^2(nu G*b)
    return (-b + B[:])
end

@inline function FFTconvolution(M::FastMDual, b::Array{Complex128,1})
    # function to overload the applyication of
    # convolution of b times G

    if M.quadRule == "trapezoidal"

      #obtaining the middle index
      indMiddle = round(Integer, M.n)

      # Allocate the space for the extended B
      BExt = zeros(Complex128,M.ne, M.me);
      # Apply spadiagm(nu) and ented by zeros
      BExt[1:M.n,1:M.m]= reshape(M.nu.*b,M.n,M.m) ;

      # Fourier Transform
      BFft = fft(BExt)
      # Component-wise multiplication
      BFft = M.GFFT.*BFft
      # Inverse Fourier Transform
      BExt = ifft(BFft)

      # multiplication by omega^2
      B = (BExt[M.n:M.n+M.n-1, M.m:M.m+M.m-1]);

    elseif M.quadRule == "Greengard_Vico"
      # for this we use the Greengard Vico method in the
      # frequency domain

      # Allocate the space for the extended B
      BExt = zeros(Complex128,M.ne, M.me);
      # Apply spadiagm(nu) and ented by zeros
      BExt[1:M.n,1:M.m]= reshape(M.nu.*b,M.n,M.m) ;

      # Fourier Transform
      BFft = fftshift(fft(BExt))
     # Component-wise multiplication
      BFft = M.GFFT.*BFft
      # Inverse Fourier Transform
      BExt = ifft(ifftshift(BFft))

      # multiplication by omega^2
      B = BExt[1:M.n, 1:M.m];
    end
    return  B[:]
end

function buildFastConvolutionDual(x::Array{Float64,1},y::Array{Float64,1},
                              h::Float64,k,nu::Function; quadRule::String = "trapezoidal")

  if quadRule == "trapezoidal"

    (ppw,D) = referenceValsTrapRule();
    D0      = D[round(Int,k*h)];
    (n,m) = length(x), length(y)
    Ge    = buildGConv(x,y,h,n,m,D0,k);
    GFFT  = fft(Ge);
    X = repmat(x, 1, m)[:]
    Y = repmat(y', n,1)[:]

    return FastMDual(GFFT,nu(X,Y),2*n-1,2*m-1,n, m, k);

  elseif quadRule == "Greengard_Vico"

      Lp = 4*(abs(x[end] - x[1]) + h)
      L  =   (abs(x[end] - x[1]) + h)*1.5
      (n,m) = length(x), length(y)
      X = repmat(x, 1, m)[:]
      Y = repmat(y', n,1)[:]

      # this is depending if n is odd or not
      if mod(n,2) == 0
        kx = (-(2*n):1:(2*n-1));
        ky = (-(2*m):1:(2*m-1));

        KX = (2*pi/Lp)*repmat(kx, 1, 4*m);
        KY = (2*pi/Lp)*repmat(ky', 4*n,1);

        S = sqrt.(KX.^2 + KY.^2);

        GFFT = Gtruncated2D(L, k, S)
        return FastMDual(GFFT, nu(X,Y), 4*n, 4*m,
                     n, m, k , quadRule="Greengard_Vico");
      else
        # kx = (-2*(n-1):1:2*(n-1) )/4;
        # ky = (-2*(m-1):1:2*(m-1) )/4;

        # KX = (2*pi/Lp)*repmat(kx, 1, 4*m-3);
        # KY = (2*pi/Lp)*repmat(ky', 4*n-3,1);

        # S = sqrt.(KX.^2 + KY.^2);

        # GFFT = Gtruncated2D(L, k, S)

        # return FastM(GFFT,nu(X,Y),4*n-3,4*m-3,
        #              n,m, k,quadRule = "Greengard_Vico");

        kx = (-2*n:1:2*n-1);
        ky = (-2*m:1:2*m-1);

        KX = (2*pi/Lp)*repmat( kx, 1,4*m);
        KY = (2*pi/Lp)*repmat(ky',4*n,  1);

        S = sqrt.(KX.^2 + KY.^2);

        GFFT = Gtruncated2D(L, k, S)

        return FastMDual(GFFT,nu(X,Y),4*n,4*m,
                     n,m, k,quadRule = "Greengard_Vico");

    end
  end
end

##########################################################################
########  Function fo build the preconditioner for the  ##################
########        Dual Lippmann-Schwinger equation    ######################
##########################################################################

function buildGSparseAandG(k::Float64,X::Array{Float64,1},Y::Array{Float64,1},
                          D0::Complex128, n::Int64 ,m::Int64, nu::Function ; method::String = "normal")

    nuArray = nu(X,Y)
    Ind = reshape(collect(1:n*m),n,m);

    if method=="normal"
      (Indices, Values) = entriesSparseA(k,X,Y,D0, n ,m);
    elseif method == "randomized"
      (Indices, Values) = entriesSparseARand(k,X,Y,D0, n ,m);
    end

    Entries = entriesSparseG(k,X,Y,D0, n ,m);

    (rowA, colA, valA) = createIndices(Ind[2:end-1,2:end-1][:],
                                       Indices[1][:], Values[1][:]);

    (rowG, colG, valG) = createIndices(Ind[2:end-1,2:end-1][:],
                                       Indices[1][:],Values[1][:], Entries[1],nuArray);

    (Row, Col, Val) = createIndices(Ind[1,2:end-1][:],
                                    Indices[2][:], Values[2][:]);

    rowA = vcat(rowA,Row);
    colA = vcat(colA,Col);
    valA = vcat(valA,Val);

    (Row, Col, Val) = createIndices(Ind[end,2:end-1][:],
                                    Indices[3][:], Values[3][:]);

    rowA = vcat(rowA,Row)
    colA = vcat(colA,Col)
    valA = vcat(valA,Val)

    (Row, Col, Val) = createIndices(Ind[2:end-1,1][:],
                                    Indices[4][:], Values[4][:]);

    rowA = vcat(rowA,Row)
    colA = vcat(colA,Col)
    valA = vcat(valA,Val)

    (Row, Col, Val) = createIndices(Ind[2:end-1,end][:],
                                    Indices[5][:], Values[5][:]);

    rowA = vcat(rowA,Row)
    colA = vcat(colA,Col)
    valA = vcat(valA,Val)

    (Row, Col, Val) = createIndices(Ind[1,1],
                                    Indices[6][:], Values[6][:]);

    rowA = vcat(rowA,Row)
    colA = vcat(colA,Col)
    valA = vcat(valA,Val)

    (Row, Col, Val) = createIndices(Ind[end,1],
                                    Indices[7][:], Values[7][:]);

    rowA = vcat(rowA,Row)
    colA = vcat(colA,Col)
    valA = vcat(valA,Val)

    (Row, Col, Val) = createIndices(Ind[1,end],
                                    Indices[8][:], Values[8][:]);

    rowA = vcat(rowA,Row)
    colA = vcat(colA,Col)
    valA = vcat(valA,Val)

    (Row, Col, Val) = createIndices(Ind[end,end],
                                    Indices[9][:], Values[9][:]);

    rowA = vcat(rowA,Row)
    colA = vcat(colA,Col)
    valA = vcat(valA,Val)

    A = sparse(rowA,colA,valA);

    (Row, Col, Val) = createIndices(Ind[1,2:end-1][:],
                                    Indices[2][:],Values[2][:], Entries[2],nuArray);

    rowG = vcat(rowG,Row);
    colG = vcat(colG,Col);
    valG = vcat(valG,Val);

    (Row, Col, Val) = createIndices(Ind[end,2:end-1][:],
                                    Indices[3][:], Values[3][:], Entries[3],nuArray);
    rowG = vcat(rowG,Row)
    colG = vcat(colG,Col)
    valG = vcat(valG,Val)

    (Row, Col, Val) = createIndices(Ind[2:end-1,1][:],
                                    Indices[4][:], Values[4][:], Entries[4],nuArray);

    rowG = vcat(rowG,Row)
    colG = vcat(colG,Col)
    valG = vcat(valG,Val)

    (Row, Col, Val) = createIndices(Ind[2:end-1,end][:],
                                    Indices[5][:], Values[5][:], Entries[5],nuArray);
    rowG = vcat(rowG,Row)
    colG = vcat(colG,Col)
    valG = vcat(valG,Val)

    (Row, Col, Val) = createIndices(Ind[1,1],
                                    Indices[6][:], Values[6][:], Entries[6],nuArray);
    rowG = vcat(rowG,Row)
    colG = vcat(colG,Col)
    valG = vcat(valG,Val)

    (Row, Col, Val) = createIndices(Ind[end,1],
                                    Indices[7][:], Values[8][:], Entries[8],nuArray);

    rowG = vcat(rowG,Row)
    colG = vcat(colG,Col)
    valG = vcat(valG,Val)

    (Row, Col, Val) = createIndices(Ind[1,end],
                                    Indices[8][:], Values[8][:], Entries[7],nuArray);

    rowG = vcat(rowG,Row)
    colG = vcat(colG,Col)
    valG = vcat(valG,Val)

    (Row, Col, Val) = createIndices(Ind[end,end],
                                    Indices[9][:], Values[9][:], Entries[9],nuArray);

    rowG = vcat(rowG,Row)
    colG = vcat(colG,Col)
    valG = vcat(valG,Val)

    ADGs = sparse(rowG,colG,valG);

    return (A, ADGs)

end

function createIndices(row::Array{Int64,1}, col::Array{Int64,1},
                       valA::Array{Complex128,1},
                       valG::Array{Complex128,2}, nu::Array{Float64,1})
  # function to create the indices for a sparse matrix
  @assert length(col) == length(valA)
  nn = length(col);
  mm = length(row);

  ll = length(valA)

  Row = kron(row, ones(Int64, nn));
  Col = kron(ones(Int64,mm), col) + Row;
  Val = nu[Col].*(kron(ones(Int64,mm), valA))

  # multiplying by the blocks of G.
  for ii = 1:mm
    Val[(ii-1)*ll+1:ii*ll] = (Val[(ii-1)*ll+1:ii*ll].')*valG
  end

  return (Row,Col,Val)
end

function createIndices(row::Int64, col::Array{Int64,1},
                       valA::Array{Complex128,1},
                       valG::Array{Complex128,2}, nu::Array{Float64,1})

  @assert length(col) == length(valA)
  nn = length(col);
  mm = 1;
  ll = length(valA)

  Row = kron(row, ones(Int64, nn));
  Col = kron(ones(Int64,mm), col) + Row;
  Val = nu[Col].*(kron(ones(Int64,mm), valA))

  print(size(Val))

  # multiplying by the blocks of G.
  for ii = 1:mm
    Val[(ii-1)*ll+1:ii*ll] = (Val[(ii-1)*ll+1:ii*ll].')*valG
  end

  return (Row,Col,Val)
end

# function sparsifiedMatrix()

#     As = buildSparseA(k,X,Y,D0, n ,m);


using Test
using PolynomialBasis
using ImplicitDomainQuadrature
# using Revise
using CutCellDG
include("../useful_routines.jl")

polyorder = 1
numqp = 2
elasticitybasis = LagrangeTensorProductBasis(2, polyorder)
levelsetbasis = HermiteTensorProductBasis(2)
levelset = InterpolatingPolynomial(1, levelsetbasis)

x0 = [0.0, 0.0]
meshwidths = [3.0, 1.0]
nelements = [3, 1]
dim,numinterp = size(interpolation_points(levelsetbasis))
points = interpolation_points(elasticitybasis)
dgmesh = CutCellDG.DGMesh(x0, meshwidths, nelements, points)
cgmesh = CutCellDG.CGMesh(x0, meshwidths, nelements, numinterp)

normal = [1.0, 0.0]
xI = [0.5, 0.0]
tol = 1e-3
perturbation = 1e-3

levelset = CutCellDG.LevelSet(
    x -> plane_distance_function(x, normal, xI),
    cgmesh,
    levelsetbasis,
)

cutmesh = CutCellDG.CutMesh(dgmesh, levelset)

cellquads = CutCellDG.CellQuadratures(cutmesh, levelset, numqp)
@test length(cellquads.quads) == 3
testcelltoquad = [
    2 1 1
    3 0 0
]
@test allequal(cellquads.celltoquad, testcelltoquad)


using StaticArrays
using ClimateMachine.VariableTemplates
using ClimateMachine.BalanceLaws:
    BalanceLaw,
    Prognostic,
    Auxiliary,
    Gradient,
    GradientFlux,
    GradientLaplacian,
    Hyperdiffusive
import ClimateMachine.BalanceLaws:
    vars_state,
    flux_first_order!,
    flux_second_order!,
    source!,
    compute_gradient_argument!,
    compute_gradient_flux!,
    nodal_init_state_auxiliary!,
    init_state_prognostic!,
    boundary_conditions,
    boundary_state!,
    wavespeed,
    transform_post_gradient_laplacian!

using ClimateMachine.Mesh.Geometry: LocalGeometry
using ClimateMachine.DGMethods.NumericalFluxes:
    NumericalFluxFirstOrder, NumericalFluxSecondOrder

abstract type HyperDiffusionProblem end
struct HyperDiffusion{dim, P} <: BalanceLaw
    problem::P
    function HyperDiffusion{dim}(
        problem::P,
    ) where {dim, P <: HyperDiffusionProblem}
        new{dim, P}(problem)
    end
end

vars_state(::HyperDiffusion, ::Auxiliary, FT) = @vars(D::SMatrix{3, 3, FT, 9})
#
# Density is only state
vars_state(::HyperDiffusion, ::Prognostic, FT) = @vars(ρ::FT)

# Take the gradient of density
vars_state(::HyperDiffusion, ::Gradient, FT) = @vars(ρ::FT)
# Take the gradient of laplacian of density
vars_state(::HyperDiffusion, ::GradientLaplacian, FT) = @vars(ρ::FT)

vars_state(::HyperDiffusion, ::GradientFlux, FT) = @vars()
# The hyperdiffusion DG auxiliary variable: D ∇ Δρ
vars_state(::HyperDiffusion, ::Hyperdiffusive, FT) = @vars(σ::SVector{3, FT})

function flux_first_order!(m::HyperDiffusion, _...) end

"""
    flux_second_order!(m::HyperDiffusion, flux::Grad, state::Vars,
                     auxDG::Vars, auxHDG::Vars, aux::Vars, t::Real)

Computes diffusive flux `F` in:

```
∂ρ
-- = - ∇ • (σ) = - ∇ • F
∂t
```
Where

 - `σ` is hyperdiffusion DG auxiliary variable (`σ = D ∇ Δρ` with D being the hyperdiffusion tensor)
"""
function flux_second_order!(
    m::HyperDiffusion,
    flux::Grad,
    state::Vars,
    auxDG::Vars,
    auxHDG::Vars,
    aux::Vars,
    t::Real,
)
    σ = auxHDG.σ
    flux.ρ += σ
end

"""
    compute_gradient_argument!(m::HyperDiffusion, transform::Vars, state::Vars,
                   aux::Vars, t::Real)

Set the variable to take the gradient of (`ρ` in this case)
"""
function compute_gradient_argument!(
    m::HyperDiffusion,
    transform::Vars,
    state::Vars,
    aux::Vars,
    t::Real,
)
    transform.ρ = state.ρ
end

compute_gradient_flux!(m::HyperDiffusion, _...) = nothing
function transform_post_gradient_laplacian!(
    m::HyperDiffusion,
    auxHDG::Vars,
    gradvars::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
)
    ∇Δρ = gradvars.ρ
    D = aux.D
    auxHDG.σ = D * ∇Δρ
end

"""
    source!(m::HyperDiffusion, _...)

There is no source in the hyperdiffusion model
"""
source!(m::HyperDiffusion, _...) = nothing

function init_state_prognostic!(
    m::HyperDiffusion,
    state::Vars,
    aux::Vars,
    localgeo,
    t::Real,
)
    initial_condition!(m.problem, state, aux, localgeo, t)
end

boundary_conditions(::HyperDiffusion) = ()
boundary_state!(nf, ::HyperDiffusion, _...) = nothing


using DifferentialEquations, ModelingToolkit, DiffEqParamEstim
using Turing, Distributions, StatsPlots
using CSV, DataFrames

include("mcmcODE.jl")

#=
1. Set up the ODE
2. Import and process the data
3. Transform the ODE output to match the data
4. Generate a loss function
5. Run the MCMC algorithm + Postprocessing
=#

###############################################################
                    # 1. Set up the ODE
###############################################################

#Constants in the model
cellVol = 3e-12 #Cell Volume (liters)
Na = 6.02e23 #Avagadro's number

#Defining non zero species in the model
m2c(molecule) = @. 1e9*molecule/(cellVol*Na) #Converts molecules to nM

#Initial Conditions
intConditions = (
cGAS = m2c(1e3),
DNA = m2c(1e3),
Sting = m2c(1e3),
cCAMP = 0.0,
IRF3 = m2c(1e4),
IFNbm = 0.0,
IFNb = 0.0,
STAT = 0.0,
SOCSm = 0.0,
IRF7m = 0.0,
TREX1m = 0.0,
IRF7 = 0.0,
TREX1 = 0.0)

#Algebraic Equations (mass balances)
cGAStot, Stingtot, IRF3tot = (
intConditions[:cGAS],
intConditions[:Sting],
intConditions[:IRF3])

## Define the ODE model
@parameters t k1f k1r k3f k3r k4f kcat5 Km5 k5r kcat6 Km6 kcat7 Km7 kcat8 Km8 k8f k9f k10f1 k10f2 k11f k12f k13f k6f kcat2 Km2 τ4 τ6 τ7 τ8 τ9 τ10 τ11 τ12 τ13
@variables cGAS(t) DNA(t) Sting(t) cGAMP(t) IRF3(t) IFNβm(t) IFNβ(t) STAT(t) SOCSm(t) IRF7m(t) TREX1m(t) IRF7(t) TREX1(t)
@derivatives D'~t

#Write the equations for the states
eqs= [
  D(cGAS) ~ -k1f*cGAS*DNA + k1r*(cGAStot - cGAS)
  D(DNA) ~ -k1f*cGAS*DNA + k1r*(cGAStot - cGAS) - kcat2*TREX1*DNA / (Km2 + DNA)
  D(Sting) ~ -k3f*cGAMP*Sting + k3r*(Stingtot - Sting)
  D(cGAMP) ~ k4f*(cGAStot - cGAS) - k3f*cGAMP*Sting + k3f*(Stingtot - Sting) - τ4*cGAMP
  D(IRF3) ~ -kcat5*IRF3*(Stingtot - Sting) / (Km5 +IRF3) + k5r*(IRF3tot - IRF3)
  D(IFNβm) ~ kcat6*(IRF3tot - IRF3) / (Km6 + (IRF3tot - IRF3)) + k6f*IRF7 - τ6*IFNβm
  D(IFNβ) ~ kcat7*IFNβm / (Km7 + IFNβm) - τ7*IFNβ
  D(STAT) ~ kcat8*IFNβ / (Km8 + IFNβ) * 1.0/(1.0+k8f*SOCSm) - τ8*STAT
  D(SOCSm) ~ k9f*STAT - τ9*SOCSm
  D(IRF7m) ~ k10f1*STAT + k10f2*IRF7 - τ10*IRF7m
  D(TREX1m) ~ k11f*STAT - τ11*TREX1m
  D(IRF7) ~ k12f*IRF7m - τ12*IRF7
  D(TREX1) ~ k13f*TREX1m - τ13*TREX1]

#Generate a function that evaluates the ODE
de = ODESystem(eqs)
f = ODEFunction(de,[cGAS, DNA, Sting, cGAMP, IRF3, IFNβm, IFNβ, STAT, SOCSm, IRF7m, TREX1m, IRF7, TREX1],[ k1f, k1r, k3f, k3r, k4f, kcat5, Km5, k5r, kcat6, Km6, kcat7, Km7, kcat8, Km8, k8f, k9f, k10f1, k10f2, k11f, k12f, k13f, k6f, kcat2, Km2, τ4, τ6, τ7, τ8, τ9, τ10, τ11, τ12, τ13])

#Contruct the ODE Problem
#Time to run the simulation
tspan = (0.0, 48.0)
#Initial guesses for the parameters
pNames = (:k1f, :k1r, :k3f, :k3r, :k4f, :kcat5, :Km5, :k5r, :kcat6, :Km6, :kcat7, :Km7, :kcat8, :Km8, :k8f, :k9f, :k10f1, :k10f2, :k11f, :k12f, :k13f, :k6f, :kcat2, :Km2, :τ4, :τ6, :τ7, :τ8, :τ9, :τ10, :τ11, :τ12, :τ13)

p = [2.6899, 4.8505, 0.0356, 7.487, 517.4056, 22328.3852, 11226.3682,0.9341,
         206.9446, 10305.461, 47639702.95,3.8474, 13.006, 78.2048, 0.0209,
         0.0059, 0.001, 0.0112, 0.001, 99.9466, 15.1436,0.0276, 237539.3249,
         61688.259, 0.96, 1.347, 12242.8736,1.2399, 1.5101, 0.347, 0.165, 6.9295,
         0.0178]

#Get the values for the initial condtions
u0 = collect(intConditions)

#Contruct the problem
prob = ODEProblem(f,u0,tspan,p)

#Choose an ODE algorithm and solve for the solution
alg = Rodas5()
sol = solve(prob,alg)

###############################################################
                    # 2. Import the data
###############################################################

#Read CSV in
data = CSV.read("NormData.csv",missingstring = "-")

#Two seperate error functions are used, one comparing normalized simulations and data, and the other comparing desired peaks (max/min)

#Define desired peaks
desiredScale = m2c([0.8*1e3,eps(),0.8*1e3, 1e5, 3e3,100.0,1e5,1e4,
                    30.0,30.0,30.0,1e4,1e4])

#Transform dataframe into Array of Arrays
timePoints = [data[data.Experiment .==i,2] for i ∈ unique(data.Experiment)]
experiments = [data[data.Experiment .==i,3:end] for i ∈ unique(data.Experiment)]

#Convert into matrices and vectors
timePoints = convert.(Vector,timePoints)
experiments = convert.(Matrix,experiments)

#Scale data by the desired value
dataIndex = [6,7,10,9,11,5,3]

scaledExperiments = [ desiredScale[dataIndex]'.*exp for exp in experiments]

###############################################################
        # 3. Transform the ODE output to match the data
###############################################################

#Depending on the data you might want to convert your ODE solution to match the data (convert to log fold change, measuring different species in mass balance, etc)

#Example expIdx
#expIdx = @. !ismissing(experiments[1][1,:])

function Observables(sol,intConditions,dataIndex,t)
  obversedSpecies = sol(t)[dataIndex,:]

  if isa(t,Number)
    #Have data for activated protiens, ode returns inactive
    obversedSpecies[end] = intConditions[:Sting] - obversedSpecies[end]
    obversedSpecies[end-1] = intConditions[:IRF3] - obversedSpecies[end-1]
  else
    obversedSpecies[end,:] = intConditions[:Sting] .- obversedSpecies[end,:]
    obversedSpecies[end-1,:] = intConditions[:IRF3] .- obversedSpecies[end-1,:]
  end

  #What species are being measured?
  return obversedSpecies
end

function heuristic(sol,desiredScale)
  peaksModel = maximum(sol[:,:],dims=2)
  peaksModel[[1,3,5]] = minimum(sol[[1,3,5],:],dims=2)
  h = @. (1.0-peaksModel/desiredScale)^2
  h[2] = 0.0
  return sum(h)
end
#What parameters are we changing?
#Need to pass the parameter values and the indicies of the parameters you want to change
parVary = (:k1f, :k1r, :k3f, :k3r, :k4f, :kcat5, :Km5, :k5r, :kcat6, :Km6, :k11f, :k13f, :kcat2, :Km2, :τ7, :kcat7, :Km7, :kcat8, :Km8, :k8f, :k9f, :k10f1, :k10f2, :k12f, :k6f)

#Number of parameters in the model to vary
parNum = length(parVary)

parVaryIdx = findall(par -> par ∈ parVary, pNames)

parChange = NamedTuple{parVary}(parVaryIdx)

###############################################################
                    # 4. Run the MCMC algorithm
###############################################################

#Provide any prior knowledge for the parameters
priors = fill(Uniform(0.0,1.0),parNum)

#include("mcmcODE.jl")
result = mcmcODE(prob,alg,timePoints,experiments,priors,parVary,parVaryIdx;num_samples=10_000)

###############################################################
                # 5. Postprocessing & Plotting
###############################################################

#Summary of the MCMC results
#describe(result)

#Save the parameter chains and information about the MCMC run
chainParameters = DataFrame(result,:parameters)
CSV.write("Parameters.csv",chainParameters)

chainInternals = DataFrame(result,:internals)
CSV.write("Internals.csv",chainInternals)

#Make a plot of the chains
chainsPlot = plot(result)
savefig(chainsPlot,"ChainsODE.pdf")

#Plot the data against the simulation

dataPlots = Vector(undef,length(dataIndex))
#loop through the data
for (i,protein) in enumerate(names(data)[3:end])
  dataPlots[i] = @df data plot(:Time,data[!,protein].*desiredScale[dataIndex[i]],
                 group=:Experiment,title=protein,markershape=:auto)
end

plot(dataPlots...,legend=false)

#Get the best parameter through the lowest log Posterior
bestPar = copy(p)
bestParIdx = argmax(chainInternals.lp)

#Why are the parameters out of order, why do you hate my Turing.jl?
for par in names(chainParameters)
  if par ∈ pNames
  bestPar[findfirst(x-> x==par,pNames)] = chainParameters[bestParIdx,par]
  end
end

newProb = remake(prob, p=bestPar)
newSol = solve(newProb,alg)

t = range(tspan[1],tspan[2],length=100)
obs = Observables(sol,intConditions,dataIndex,t)

plot!(t,obs',layout=length(dataIndex),legend=false,framestyle=:box)


import SGP4
using Base.Test

line1 = "1 00005U 58002B   00179.78495062  .00000023  00000-0  28098-4 0  4753"
line2 = "2 00005  34.2682 348.7242 1859667 331.7664  19.3264 10.82419157413667"

wgs72 = SGP4.GravityModel("wgs72")
satellite = SGP4.twoline2rv(line1, line2, wgs72)

@test satellite[:satnum] == 5

t = Dates.DateTime(2000, 6, 29, 12, 50, 19)
(position, velocity) = SGP4.propagate( satellite, 2000, 6, 29, 12, 50, 19)
(position2, velocity2) = SGP4.propagate( satellite, t )

@test satellite[:error] == 0

@test_approx_eq_eps position[1] position2[1] eps()
@test_approx_eq_eps position[2] position2[2] eps()
@test_approx_eq_eps position[3] position2[3] eps()

@test_approx_eq_eps velocity[1] velocity2[1] eps()
@test_approx_eq_eps velocity[2] velocity2[2] eps()
@test_approx_eq_eps velocity[3] velocity2[3] eps()

@test_approx_eq_eps position[1] 5576.056952 1e-6
@test_approx_eq_eps position[2] -3999.371134 1e-6
@test_approx_eq_eps position[3] -1521.957159 1e-6

@test_approx_eq_eps velocity[1] 4.772627 1e-6
@test_approx_eq_eps velocity[2] 5.119817 1e-6
@test_approx_eq_eps velocity[3] 4.275553 1e-6

sats = [satellite; satellite; satellite]

(positions, velocities) = SGP4.propagate( sats, t )

@test_approx_eq_eps positions[1,1] positions[1,3] eps()
@test_approx_eq_eps positions[2,1] positions[2,3] eps()
@test_approx_eq_eps positions[3,1] positions[3,3] eps()

@test_approx_eq_eps velocities[1,1] velocities[1,3] eps()
@test_approx_eq_eps velocities[2,1] velocities[2,3] eps()
@test_approx_eq_eps velocities[3,1] velocities[3,3] eps()


function hyperbolic_cross_gradient(weights::Array{T,1},node::Array{R,1},multi_index::Array{S,2}) where {T<:AbstractFloat,R<:Number,S<:Integer}

  d = length(node)
  gradient = Array{R,2}(undef,1,d)

  for i = 1:d
    gradient[i] = hyperbolic_cross_derivative(weights,node,multi_index,i)
  end

  return gradient

end

function hyperbolic_cross_gradient(weights::Array{T,1},node::Array{R,1},multi_index::Array{S,2},domain::Union{Array{T,1},Array{T,2}}) where {T<:AbstractFloat,R<:Number,S<:Integer}

  d = length(node)
  gradient = Array{R,2}(undef,1,d)

  for i = 1:d
    gradient[i] = hyperbolic_cross_derivative(weights,node,multi_index,domain,i)
  end

  return gradient

end


function testidea2(posterior)

    rg = MersenneTwister(1)

    pdf_without_scale = MvNormalPartition(posterior, [2;3;4;5;6])

    pdf_b = MvNormalPartition(pdf_without_scale, [3;4;5])

    pdf_a_given_b = MvNormalConditional(pdf_without_scale, [1;2], [3;4;5])

    # draw samples for b
    B = [rand(pdf_b) for i=1:200]

    figure(1)
    cla()
    xrange = collect(LinRange(-2,2,100))

    for i=1:5

        a, b = rand(rg, pdf_a), rand(rg, pdf_b)

        for x in xrange
               local p = b + a*x
               plot3D(p[1], p[2], p[3], "r.", alpha=0.3)
        end
    end

    function test(x)

        l = zeros(length(B))

        for (i,b) in enumerate(B)

            local v = x - b

            l[i] = logpdf(pdf_a_norm, v / norm(v))

        end

        return logsumexp(l) - log(length(B))

    end

end


using Hecke
using Base.Test

#include("NfMaxOrd-test.jl")
#include("NfMaxOrd.jl")
include("NfOrd.jl")

#test_pseudohnf()

set_verbose_level(:UnitGroup, 1)
set_verbose_level(:ClassGroup, 0)

test_NfOrd()
test_NfOrd_Elem()
test_NfOrd_Idl()
test_NfOrd_FracIdl()

Qx, x = PolynomialRing(QQ, "x")

K, a = NumberField(x^2 - 2)
print("Testing $(x^2 - 2)\n")
O = maximal_order(K)
#
# x^5 + 514944*x^2 + 123904 test prime decomposition with this (2 is index divisor and only one prime ideal over 2)
#
print("Tentative class and unit group computation ... \n")
c, U, b = Hecke._class_unit_group(O);
print("Saturating the tentative unit group ... \n")
Hecke._refine_with_saturation(c, U)

K, a = NumberField(x^3 - 2)
print("Testing $(x^3 - 2)\n")
O = maximal_order(K)
#
# x^5 + 514944*x^2 + 123904 test prime decomposition with this (2 is index divisor and only one prime ideal over 2)
#
print("Tentative class and unit group computation ... \n")
c, U, b = Hecke._class_unit_group(O);
print("Saturating the tentative unit group ... \n")
Hecke._refine_with_saturation(c, U)

K, a = NumberField(x^5 - 11^2 * 7)
print("Testing $(x^5 - 11^2 * 7)\n")
O = maximal_order(K)
#
# x^5 + 514944*x^2 + 123904 test prime decomposition with this (2 is index divisor and only one prime ideal over 2)
@test discriminant(O) == 109853253125
#
print("Tentative class and unit group computation ... \n")
c, U, b = Hecke._class_unit_group(O);
print("Saturating the tentative unit group ... \n")
Hecke._refine_with_saturation(c, U)
##
f = x^18 + 18*x^16 + 135*x^14 + 192*x^12 - 2961*x^10 - 17334*x^8+ 20361*x^6 +  315108*x^4 + 514944*x^2 + 123904
#
print("Testing $f\n")
println(signature(f))
K, a = NumberField(f)
O = maximal_order(K)
#
print("Tentative class and unit group computation ... \n")
c, U, b = Hecke._class_unit_group(O);
print("Saturating the tentative unit group ... \n")
Hecke._refine_with_saturation(c, U)
#
@test c.h == 36
@test Hecke._validate_class_unit_group(c, U) == 1
##
#f = x^5 + 514944*x^2 + 123904
##
#print("Testing $f\n")
#println(signature(f))
#K, a = NumberField(f)
#O = maximal_order(K)
##
#print("Tentative class and unit group computation ... \n")
#c, U, b = Hecke._class_unit_group(O);
#print("Saturating the tentative unit group ... \n")
#Hecke._refine_with_saturation(c, U)
#
#@test c.h == 2


using NetworkEpidemics
using LightGraphs
using Statistics
using Plots
using ColorSchemes
using Random

Random.seed!(2021)

CorrectedMetapopulation(χ::Real, mp::Metapopulation{SI}) = Metapopulation(mp.h, mp.D, SI(χ*mp.dynamics.β))

N = 1000
M = 10

h = complete_graph(M)
g = watts_strogatz(N, 100, 0.0)

k = degree(g)

χ = mean(k)/N # doesn't change so only compute it once

β = 0.1
D = 0.1

mp = Metapopulation(h, D, SI(β))
mpc = CorrectedMetapopulation(χ, mp)

mpx = Metaplex(g, h, D, SI(β))

x0_i = fill(1, N)
x0_i[1:10] .= 2
x0_μ = rand(1:M, N)
x0_mp = zeros(Int, M, 2)
for i in 1:N
    x0_mp[x0_μ[i], x0_i[i]] += 1
end

tmax = 20.0
nmax = 10000
nsims = 200
nbins = 200

ts_mp, output_mp = gillespie(mp, x0_mp, tmax=tmax, nmax=nmax)
ts_mpx, output_mpx = gillespie(mpx, [x0_i,x0_μ], tmax=tmax, nmax=nmax)

plot(ts_mp, output_mp[2], label="")
plot(ts_mpx, output_mpx[2], label="")

ts = LinRange(0.0, tmax, nbins)
ps = LinRange(0.0, 0.4, 5)

ts_mf_mp, u_mf_mp = meanfield(mp, x0_mp, tmax=tmax, saveat=ts)
ts_mf_mpc, u_mf_mpc = meanfield(mpc, x0_mp, tmax=tmax, saveat=ts)

u_mf_mpx = zeros(length(ps), length(ts))
for i in 1:length(ps), j in 1:100
    g = watts_strogatz(N, 100, ps[i])
    mpx = Metaplex(g, h, D, SI(β))
    t, u = meanfield(mpx, [x0_i,x0_μ], tmax=tmax, saveat=ts)
    u_mf_mpx[i,:] += sum(u[2], dims=2)
end
u_mf_mpx /= 100

colors = ColorSchemes.GnBu_9;
orange = ColorSchemes.Oranges_3[3];

plot(ts, sum(u_mf_mpc[2], dims=2)/N,
    #label="Meanfield (No correction)",
    label = "Mean degree correction",
    xlabel="Time",
    ylabel="Fraction of infected individuals",
    fontfamily="Deja Vu Sans",
    legend=:bottomright,
    color = orange
    )
    plot!(ts, u_mf_mpx[1,:]/N,
    label = "p=0.0",
    color = colors[5]
    )
    plot!(ts, u_mf_mpx[2,:]/N,
    label = "p=0.1",
    color = colors[6]
    )
    plot!(ts, u_mf_mpx[3,:]/N,
    label = "p=0.2",
    color = colors[7]
    )
    plot!(ts, u_mf_mpx[4,:]/N,
    label = "p=0.3",
    color = colors[8]
    )
    plot!(ts, u_mf_mpx[5,:]/N,
    label = "p=0.4",
    color = colors[9]
    )
    plot!(ts, sum(u_mf_mpc[2], dims=2)/N,
    label="",
    color = orange
    )

plot(ts, (sum(u_mf_mpc[2], dims=2) - u_mf_mpx[1,:])/N,
    xlabel="Time",
    ylabel="Residual",
    label = "p=0.0",
    color = colors[5]
    )
    plot!(ts, (sum(u_mf_mpc[2], dims=2) - u_mf_mpx[2,:])/N,
    label = "p=0.1",
    color = colors[6]
    )
    plot!(ts, (sum(u_mf_mpc[2], dims=2) - u_mf_mpx[3,:])/N,
    label = "p=0.2",
    color = colors[7]
    )
    plot!(ts, (sum(u_mf_mpc[2], dims=2) - u_mf_mpx[4,:])/N,
    label = "p=0.3",
    color = colors[8]
    )
    plot!(ts, (sum(u_mf_mpc[2], dims=2) - u_mf_mpx[5,:])/N,
    label = "p=0.4",
    color = colors[9]
    )

cols = map(x->[x.r, x.g, x.b], colors.colors)
or = [orange.r, orange.g, orange.b]

using DelimitedFiles

writedlm("umpc.dat", sum(u_mf_mpc[2], dims=2)/N)
writedlm("umpx.dat", u_mf_mpx/N)
writedlm("ts.dat", ts)
writedlm("colors.dat", [cols..., or])

using MATLAB;

mat"""
figure
xlabel('Time')
ylabel('Fraction of Infected Individuals')
hold on
plot($ts, $(sum(u_mf_mpc[2], dims=2)/N), 'Color', $(or))
plot($ts, $(sum(u_mf_mpx[1], dims=2)/N), 'Color', $(cols[5]))
plot($ts, $(sum(u_mf_mpx[2], dims=2)/N), 'Color', $(cols[6]))
plot($ts, $(sum(u_mf_mpx[3], dims=2)/N), 'Color', $(cols[7]))
plot($ts, $(sum(u_mf_mpx[4], dims=2)/N), 'Color', $(cols[8]))
plot($ts, $(sum(u_mf_mpx[5], dims=2)/N), 'Color', $(cols[9]))
hold off
legend('Mean degree correction', 'p=0.0', 'p=0.1', 'p=0.2', 'p=0.3', 'p=0.4')
""";


using LinearAlgebra
using Combinatorics, DataStructures

function row_echelon_linsolve(A::Matrix{T}, b::Vector{T}) where {T}
    j = 0
    x = map(1:size(A, 1)) do i
        while j < size(A, 2)
            j += 1
            if isone(A[i, j]) && all(k -> k == i || iszero(A[k, j]), 1:size(A, 1))
                return b[j]
            end
        end
        error("Not in row_echelon_form, cannot find for `$i`th entry.")
    end
    @assert transpose(A) * x == b
    return x
end

function ordered_block_diag(As, d)
    U = block_diag(As, d)
    U === nothing && return nothing
    iU = U'
    @assert iU ≈ inv(U)
    Bs = [iU * A * U for A in As]
    @assert all(Bs) do B
        isblockdim(B, d)
    end
    refs = [B[1:d, 1:d] for B in Bs]
    offset = d
    for offset in d:d:(size(U, 1) - d)
        I = offset .+ (1:d)
        Cs = [B[I, I] for B in Bs]
        σ_ok = nothing
        for σ in permutations(1:d)
            if all(zip(refs, Cs)) do refC
                ref, C = refC
                isapprox(ref[σ, σ], C, rtol=1e-8)
            end
                σ_ok = σ
                break
            end
        end
        if σ_ok === nothing
            error("No permutation can make $refs and $Cs match")
        end
        U[:, I] = U[:, I[σ_ok]]
        offset += d
    end
    ordered_block_check(U, As, d)
    return U
end

function ordered_block_check(U, As, d)
    iU = U'
    @assert iU ≈ inv(U)
    @assert all(As) do A
        is_ordered_blockdim(iU * A * U, d)
    end
end

function block_diag(As, d)
    for A in As
        #T = eigen(A).vectors
        T = schur(A).vectors
        iT = T'
        @assert iT ≈ inv(T)
        n = LinearAlgebra.checksquare(A)
        union_find = IntDisjointSets(n)
        Bs = [iT * A * T for A in As]
        for B in Bs
            merge_sparsity!(union_find, B)
        end
        blocks = Dict{Int,Vector{Int}}()
        for i in 1:n
            r = find_root!(union_find, i)
            blocks[r] = push!(get(blocks, r, Int[]), i)
        end
        if length(blocks) > 1
            U = similar(T)
            offset = 0
            for v in values(blocks)
                @assert iszero(length(v) % d)
                if length(v) == d
                    V = T[:, v]
                else
                    Cs = [B[v, v] for B in Bs]
                    V = block_diag(Cs, d)
                    V === nothing && break
                    V *= transpose(T[:, v])
                end
                U[:, offset .+ eachindex(v)] = V
                offset += length(v)
            end
            if offset == n
                return U
            end
        end
    end
end

function merge_sparsity!(union_find::IntDisjointSets, A, tol=1e-8)
    for I in CartesianIndices(A)
        i, j = I.I
        if abs(A[I]) > tol
            union!(union_find, i, j)
        end
    end
end

function is_ordered_blockdim(A, d, tol=1e-8)
    B = A[1:d, 1:d]
    return isblockdim(A, d, tol) &&
        all(d:d:(size(A, 1) - d)) do offset
            I = offset .+ (1:d)
            isapprox(B, A[I, I], rtol=tol)
        end
end
function isblockdim(A, d, tol=1e-8)
    for I in CartesianIndices(A)
        i, j = I.I
        if abs(i - j) >= d && abs(A[I]) > tol
            return false
        end
    end
    return true
end


using LinearAlgebra
using FermiCG
using Printf
using Test
using Arpack

@testset "fci" begin
    atoms = []
    push!(atoms,Atom(1,"H",[0,0,0]))
    push!(atoms,Atom(2,"H",[0,0,1]))
    #push!(atoms,Atom(3,"H",[0,0,2]))
    #push!(atoms,Atom(4,"H",[0,0,3]))
    #basis = "6-31g"
    basis = "sto-3g"
    mol     = Molecule(0,1,atoms,basis)

    n_elec_a = 1
    n_elec_b = 1

    mf = FermiCG.pyscf_do_scf(mol)
    nbas = size(mf.mo_coeff)[1]
    ints = FermiCG.pyscf_build_ints(mol,mf.mo_coeff, zeros(nbas,nbas));
    e_fci, d1_fci, d2_fci = FermiCG.pyscf_fci(ints,1,1)
    @printf(" FCI Energy: %12.8f\n", e_fci)

    norb = size(ints.h1)[1]
    problem = FermiCG.StringCI.FCIProblem(norb, n_elec_a, n_elec_b)

    display(problem)

    @time Hmat = FermiCG.StringCI.build_H_matrix(ints, problem)
    Ssq = FermiCG.StringCI.build_S2_matrix(problem)
    println("Ssq")
    println(Ssq)
    F = eigen(Hmat+0.01*Ssq)
    e = F.values
    v = F.vectors
    e = diag(v'*Hmat*v)
    ss = diag(v'*Ssq*v)

    e = real(e)
    for (i,ei) in enumerate(e)
        @printf(" Energy: %12.8f   <S2>: %12.8f\n",ei+ints.h0,ss[i])
    end
    @test isapprox(e[1], e_fci , atol=1e-10)
end


abstract type SignedDistanceField <: Function end

struct SDFKDTree <: SignedDistanceField
    tree::KDTree
    extrusion::Float64
end

function (sdf::SignedDistanceField)(point::AbstractPoint)
    @show "No sdf distance implemented for this type:", typeof(sdf), dump(sdf)
end

@inline function (sdf::SDFKDTree)(point::AbstractPoint)
    distance = knn(sdf.tree, point, 1)[2][1] - sdf.extrusion
end

function curve_to_kdtree_sdf(γ::FunctionCurve, extrusion::Float64, arc_length_sampling_freq::Float64)
    t_to_s, s_to_t, γs, sspan = arc_length_parameterization(γ.γ, γ.tspan)
    mins, maxs = sspan[1:2]

    γsamples = (mins:arc_length_sampling_freq:maxs)
    γsampled_points = map(γs, γsamples)

    γkdtree = KDTree(γsampled_points)
    SDFKDTree(γkdtree, extrusion)
end

struct UnionSDF{SDF1 <: SignedDistanceField, SDF2 <: SignedDistanceField} <: SignedDistanceField
    sdf1::SDF1
    sdf2::SDF2
end

@inline function (sdf::UnionSDF)(point::AbstractPoint)
    distance = min(sdf.sdf1(point), sdf.sdf2(point))
end

struct IntersectSDF{SDF1 <: SignedDistanceField, SDF2 <: SignedDistanceField} <: SignedDistanceField
    sdf1::SDF1
    sdf2::SDF2
end

@inline function (sdf::IntersectSDF)(point::AbstractPoint)
    distance = max(sdf.sdf1(point), sdf.sdf2(point))
end

struct FunctionSDF <: SignedDistanceField
    sdf::Function
end

@inline function (sdf::FunctionSDF)(point::AbstractPoint)
    distance = sdf.sdf(point)
end

function curves_to_union_sdf(γs::Vector{Function}, tspans::Vector{Tuple{Float64, Float64}},
        extrusions::Vector{Float64}, arc_length_parameterization::Float64; boundary_sdf=nothing)
    γcurves = [FunctionCurve(γ, tspan) for (γ, tspan) in zip(γs, tspans)]
    γsdfs = [curve_to_kdtree_sdf(γ, extrusion, arc_length_parameterization) for (γ, extrusion) in zip(γcurves, extrusions)]
    unionsdf = reduce((sdf1, sdf2) -> UnionSDF(sdf1, sdf2), γsdfs)
    if !(boundary_sdf isa Nothing) && boundary_sdf isa SignedDistanceField
        sdf = IntersectSDF(unionsdf, boundary_sdf)
    else
        sdf = unionsdf
    end
end


using PyPlot, Revise, HiQGA.transD_GP, Random, SpecialFunctions
## model
zfixed   = [-1e5,   0,    ]
rho      = [1e20,   1,  ]
nmax = 200
##  geometry
rTx = 15
rRx = 0.001
zRx = -0.001
zTx = -0.002
doconvramp = false
modelprimary = true
nkᵣeval = 200
ntimesperdecade = 15
nfreqsperdecade = 15
times = 10 .^LinRange(-6,1,40)
## build operator
F = transD_GP.AEM_VMD_HMD.HFieldDHT(
                      ntimesperdecade = ntimesperdecade,
                      nfreqsperdecade = nfreqsperdecade,
                      times  = times,
                      zTx    = zTx,
                      rRx    = rRx,
                      rTx    = rTx,
                      zRx    = zRx,
                      doconvramp = doconvramp,
                      modelprimary = modelprimary,
                      nkᵣeval = nkᵣeval)
## get TD field
transD_GP.AEM_VMD_HMD.getfieldTD!(F, zfixed, rho)
## Analytic loop
μ = transD_GP.AEM_VMD_HMD.μ
sig0 = 1/rho[2]
theta = sqrt.(μ*sig0./4F.interptimes)
a     = rTx
dbzdt = -1 ./(sig0*a^3)*(3*erf.(theta*a) - 2/sqrt(pi)*theta.*a.*(3 .+ 2*theta.^2*a^2).*exp.(-theta.^2*a^2))
dbzdt = dbzdt/(pi*a^2)
##
figure()
s1 = subplot(211)
loglog(F.interptimes,μ*abs.(F.HTD_z_interp), label="loop W&H")
loglog(F.interptimes,abs.(dbzdt), "--", label="loop W&H transient half space ")
xlim(extrema(times))
ylabel("dBzdt V/Am^4")
grid(which="both")
legend()
s2 = subplot(212, sharex=s1)
# loglog(F.interptimes, 100*abs.(-AEM_VMD_HMD.μ*F.HTDinterp-dbzdt)./abs.(dbzdt))
# signs are supposed to be the same at loop centre ...
loglog(F.interptimes, 100*abs.(abs.(μ*F.HTD_z_interp)-abs.(dbzdt))./abs.(dbzdt))
grid(which="both")
ylabel("% difference")
xlabel("time s")


#-----------------------------------------------------------------------
#   MINKOWSKI METRIC
#-----------------------------------------------------------------------
η   = x->ημν()

#-----------------------------------------------------------------------
#   INITIAL DATA GENERATION
#-----------------------------------------------------------------------

function idg( ne::Int , Xc::RealVec=rand(4)/100 , spacelike::Bool=true )
    V = rand(4,ne)

    if spacelike
        V[1,:] = -ones(ne)

        for i=1:ne
            V[2:4,i] = V[2:4,i]/norm(V[2:4,i])
        end
        for i=1:ne
            @test ηdot(V[:,i],V[:,i]) ≈ 0 atol=1e-15
        end

        Y   = zeros(4,ne)
        Zi   = zeros(8,ne)

        for i=1:ne
            Y[:,i] = V[:,i] + Xc
        end
        for i=1:ne
        Zi[:,i] = vcat( Y[:,i] , Xc - Y[:,i] )
        end
    else
        V[4,:] = ones(ne)

        for i=1:ne
            V[1,i] = √( 1 + V[2,i]^2 + V[3,i]^2 )
        end
        for i=1:ne
            @test ηdot(V[:,i],V[:,i]) ≈ 0 atol=1e-15
        end

        Y   = zeros(4,ne)
        Zi   = zeros(8,ne)

        for i=1:ne
            Y[:,i] = V[:,i] + Xc
        end
        for i=1:ne
        Zi[:,i] = vcat( Y[:,i] , Xc - Y[:,i] )
        end
    end

    return (Y,Xc,Zi,V)
end


using StructuresKit

shape_name = "W14X90"
#define discretization along member
zcoords = 0:1:20
n=(4, 4, 4, 4, 4)

shape_info = StructuresKit.CrossSection.AISC(shape_name)

#A Ix Iy J Cw xc yc xs ys
section_properties = [(shape_info.A, shape_info.Ix, shape_info.Iy, shape_info.J, shape_info.Cw, shape_info.d/2, 0.0, shape_info.d/2, 0.0)]

#E  ν
material_properties = [(29000,0.30)]

#kx ky kϕ hx hy
springs = [0.0, 0.0, 0.0, 0.0, 0.0]

#member information
#L dL SectionProperties MaterialProperties Springs
member_definitions = [(13.0*12, (13*12)/12,1,1,1)]

num_nodes = Int(member_definitions[1][1]/member_definitions[1][2]) + 1

#P qx qy ax ay
loads = [(1000.0 * ones(num_nodes)),(0.0 * ones(num_nodes)), (0.0 * ones(num_nodes)),(0.0 * ones(num_nodes)),(0.0 * ones(num_nodes))]

#end boundary conditions
#type=1 u''=v''=ϕ''=0 (simply supported), type=2 u'=v'=ϕ'=0  (fixed), type=3 u''=v''=ϕ''=u'''=v'''=ϕ'''=0 (free end, e.g., a cantilever)
end_boundary_conditions = [1 1]

#supports
#location where u=v=ϕ=0
supports = [0.0 13.0*12]

#imperfections
L = member_definitions[1][1]
dL = member_definitions[1][2]
z = 0: dL :L
Δo = L / 1000 * sin.(π*z/L)
uo = Δo
vo = zeros(Float64, num_nodes)
ϕo = zeros(Float64, num_nodes)

imperfections = [uo, vo, ϕo]

u, v, ϕ, properties = BeamColumn.solve(member_definitions, section_properties, material_properties, loads, springs, end_boundary_conditions, supports)

using Plots
plot(properties.z, u)
plot!(properties.z, Δo)


using Gen
import Random

##############################################################
# variational inference with multi-sample gradient estimator #
##############################################################

Random.seed!(1)

@gen function model()
    slope = @trace(normal(-1, exp(0.5)), :slope)
    intercept = @trace(normal(1, exp(2.0)), :intercept)
end

@gen function approx()
    @param slope_mu::Float64
    @param slope_log_std::Float64
    @param intercept_mu::Float64
    @param intercept_log_std::Float64
    @trace(normal(slope_mu, exp(slope_log_std)), :slope)
    @trace(normal(intercept_mu, exp(intercept_log_std)), :intercept)
end
init_param!(approx, :slope_mu, 0.)
init_param!(approx, :slope_log_std, 0.)
init_param!(approx, :intercept_mu, 0.)
init_param!(approx, :intercept_log_std, 0.)

observations = choicemap()
update = ParamUpdate(GradientDescent(1., 1000), approx)
black_box_vi!(model, (), observations, approx, (), update, 20;
    iters=10000, samples_per_iter=100, verbose=true, geometric=false)
slope_mu = get_param(approx, :slope_mu)
slope_log_std = get_param(approx, :slope_log_std)
intercept_mu = get_param(approx, :intercept_mu)
intercept_log_std = get_param(approx, :intercept_log_std)
println((slope_mu, slope_log_std, intercept_mu, intercept_log_std))
@assert isapprox(slope_mu, -1., atol=0.01)
@assert isapprox(slope_log_std, 0.5, atol=0.01)
@assert isapprox(intercept_mu, 1., atol=0.01)
@assert isapprox(intercept_log_std, 2.0, atol=0.01)

#######################################################
# variational inference with score function estimator #
#######################################################

Random.seed!(1)

@gen function model()
    slope = @trace(normal(-1, exp(0.5)), :slope)
    intercept = @trace(normal(1, exp(2.0)), :intercept)
end

@gen function approx()
    @param slope_mu::Float64
    @param slope_log_std::Float64
    @param intercept_mu::Float64
    @param intercept_log_std::Float64
    @trace(normal(slope_mu, exp(slope_log_std)), :slope)
    @trace(normal(intercept_mu, exp(intercept_log_std)), :intercept)
end
init_param!(approx, :slope_mu, 0.)
init_param!(approx, :slope_log_std, 0.)
init_param!(approx, :intercept_mu, 0.)
init_param!(approx, :intercept_log_std, 0.)

observations = choicemap()
update = ParamUpdate(GradientDescent(1., 500), approx)
black_box_vi!(model, (), observations, approx, (), update;
    iters=500, samples_per_iter=100, verbose=false)
slope_mu = get_param(approx, :slope_mu)
slope_log_std = get_param(approx, :slope_log_std)
intercept_mu = get_param(approx, :intercept_mu)
intercept_log_std = get_param(approx, :intercept_log_std)
println((slope_mu, slope_log_std, intercept_mu, intercept_log_std))
@assert isapprox(slope_mu, -1., atol=0.01)
@assert isapprox(slope_log_std, 0.5, atol=0.01)
@assert isapprox(intercept_mu, 1., atol=0.01)
@assert isapprox(intercept_log_std, 2.0, atol=0.01)


"""
    quadratic objective
"""
struct QuadraticObjective <: Objective
    Q
    q
    R
    r
    c
end

function quadratic_tracking_objective(Q, R, x, u)
    T = length(Q)
    q = [-2.0 * Q[t] * x[t] for t = 1:T]
    r = [-2.0 * R[t] * u[t] for t = 1:T-1]
    c = [(t == T  ? x[t]' * Q[t] * x[t]
         : x[t]' * Q[t] * x[t] + u[t]' * R[t] * u[t]) for t = 1:T]

    return QuadraticObjective(Q, q, R, r, c)
end

function objective(Z, obj::QuadraticObjective, idx, T)
    J = 0.0

    for t = 1:T
        x = view(Z, idx.x[t])
        Q = obj.Q[t]
        q = obj.q[t]
        c = obj.c[t]

        J += x' * Q * x + q' * x
        J += c

        t == T && continue

        u = view(Z, idx.u[t])
        R = obj.R[t]
        r = obj.r[t]

        J += u' * R * u + r' * u
    end

    return J
end

function objective_gradient!(∇J, Z, obj::QuadraticObjective, idx, T)

    for t = 1:T
        x = view(Z, idx.x[t])
        Q = obj.Q[t]
        q = obj.q[t]
        c = obj.c[t]

        ∇J[idx.x[t]] += 2.0 * Q * x + q

        t == T && continue

        u = view(Z, idx.u[t])
        R = obj.R[t]
        r = obj.r[t]

        ∇J[idx.u[t]] += 2.0 * R * u + r
    end

    return nothing
end

"""
    quadratic time objective
"""
struct QuadraticTimeObjective{T} <: Objective
    obj::QuadraticObjective
    c::T
end

function quadratic_time_tracking_objective(Q, R, x, u, c)
    return QuadraticTimeObjective(quadratic_tracking_objective(Q, R, x, u),c)
end

function objective(Z, obj::QuadraticTimeObjective, idx, T)
    J = 0.0

    for t = 1:T-1
        x = view(Z, idx.x[t])
        Q = obj.obj.Q[t]
        q = obj.obj.q[t]

        u = view(Z, idx.u[t])
        R = obj.obj.R[t]
        r = obj.obj.r[t]

        h = u[end]

        c = obj.obj.c[t]

        J += (x' * Q * x + q' * x) * h
        J += (u' * R * u + r' * u) * h
        J += (c + obj.c) * h
    end

    x = view(Z, idx.x[T])
    Q = obj.obj.Q[T]
    q = obj.obj.q[T]

    J += (x' * Q * x + q' * x)

    return J
end

function objective_gradient!(∇J, Z, obj::QuadraticTimeObjective, idx, T)

    J = 0.0

    for t = 1:T-1
        x = view(Z, idx.x[t])
        Q = obj.obj.Q[t]
        q = obj.obj.q[t]

        u = view(Z, idx.u[t])
        R = obj.obj.R[t]
        r = obj.obj.r[t]

        h = u[end]

        c = obj.obj.c[t]

        ∇J[idx.x[t]] += (2.0 * Q * x + q) * h
        ∇J[idx.u[t]] += (2.0 * R * u + r) * h
        ∇J[idx.u[t][end]] += x' * Q * x + q' * x + u' * R * u + r' * u + c + obj.c
    end

    x = view(Z, idx.x[T])
    Q = obj.obj.Q[T]
    q = obj.obj.q[T]

    ∇J[idx.x[T]] += (2.0 * Q * x + q)

    return nothing
end


# =======================================
# Functionality to project the flowpipe
# =======================================
import LazySets.Approximations: project

# add a "time" variable by taking the cartesian product of the flowpipe ℱ with each time lapse
function add_time(ℱ::Vector{ReachSet{Zonotope{Float64}}})
    ℱ_with_time = Vector{ReachSet{Zonotope{Float64}}}(undef, length(ℱ))
    @inbounds for i in eachindex(ℱ)
        t0, t1 = ℱ[i].t_start, ℱ[i].t_end
        radius = (t1 - t0)/2.0
        Xi = ℱ[i].X × Zonotope([t0 + radius], hcat(radius))
        Xi = convert(Zonotope, Xi)
        ℱ_with_time[i] = ReachSet{Zonotope{Float64}}(Xi, t0, t1)
    end
    return ℱ_with_time
end

function project(sol::ReachSolution{Zonotope{Float64}})
    N = length(sol.Xk)  # number of reach sets
    n = dim(first(sol.Xk).X) # state space dimension
    options = copy(sol.options)
    πℱ = Vector{ReachSet{Zonotope{Float64}}}(undef, N) # preallocated projected reachsets
    πvars = sol.options[:plot_vars] # variables for plotting
    @assert length(πvars) == 2

    if 0 ∈ πvars
        # add the time variable to the flowpipe (assuming it's not already
        # part of the model)
        ℱ = add_time(sol.Xk)
        n += 1
        options[:n] += 1 # TODO : remove when option is removed
        πvars = copy(πvars)
        πvars[first(indexin(0, πvars))] = n # time index is added in the end
    else
        ℱ = sol.Xk
    end

    M = sparse([1, 2], πvars, [1.0, 1.0], 2, n)
    for i in eachindex(ℱ)
        t0, t1 = ℱ[i].t_start, ℱ[i].t_end
        πℱ_i = linear_map(M, ℱ[i].X)
        πℱ_i = Zonotope(πℱ_i.center, πℱ_i.generators)
        πℱ_i = reduce_order(πℱ_i, options[:max_order])
        πℱ[i] = ReachSet{Zonotope{Float64}}(πℱ_i, t0, t1)
    end
    return ReachSolution(πℱ, options)
end


######################################################################
# TODO: Take methods from docs/jmd/extremumEstimation.jmd and put here
######################################################################

######################################################################

# EL inference

"""
    gel_pλ(model::GMMModel, h::Function=log)

Return a function that given parameters θ solves

`` max_p sum h(p) s.t. sum(p) = 1, sum(p*gi(θ)) = 0 ``

The returned function gives (p(θ),λ(θ))

The returned function is not thread-safe.
"""
function gel_pλ(model::GMMModel, h::Function=log)
  n = number_observations(model)
  d = number_parameters(model)
  k = number_moments(model)
  gi = get_gi(model)
  g = gi(zeros(d))
  p = Convex.Variable(n)
  el = sum(h.(p))
  problem = Convex.maximize(el, sum(p)==1, p>=0, g'*p==0)
  function(θ)
    g = gi(θ)
    problem.constraints[3] =  g'*p==0
    Convex.solve!(problem, ECOSSolver(verbose=false), #warmstart=false,
                  verbose=false)
    (p=p.value ,λ=problem.constraints[3].dual)
  end
end

"""
    gel_tests(θ, model::GMMModel, h::Function=log)

Computes GEL test statistics for H₀ : θ = θ₀. Returns a tuple
containing statistics and p-values.
"""
function gel_tests(θ, model::GMMModel, h::Function=log)
  d = number_parameters(model)
  k = number_moments(model)
  p, λ = gel_pλ(model, h)(θ)
  # Compared to Andrews & Guggenberger, this λ = -n*λ(AG)
  n = length(p)
  λ /= -n
  lr = -2*sum(h.(p) .- h(1/n))

  gi = get_gi(model)
  iΔ = inv(cov(gi(θ)))
  D = ForwardDiff.jacobian(x->gi(x)'*p, θ)
  S = (n*λ'*D*inv(D'*iΔ*D)*D'*λ)[1]
  return(GELR=lr, S=S, plr=cdf(Chisq(k),lr), ps=cdf(Chisq(d),S))
end


# ENV["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4,5,6,7"
# using Revise
using ADSeismic
using ADCME
using MAT
using PyPlot
using Random
using DelimitedFiles
using Optim
using Dates
matplotlib.use("Agg")
close("all")

gpu = has_gpu() ? true : false

data_dir = "data/acoustic"
method = "NNFWI/marmousi"

figure_dir = string("figure/",method)
result_dir = string("result/",method)
model_dir = string("nn_model/",method)
loss_file = joinpath(result_dir, "loss_$(Dates.now()).txt")

check_path(dir) = !ispath(dir) ? mkpath(dir) : nothing
check_path(figure_dir)
check_path(result_dir)
check_path(model_dir)

################### Inversion using Automatic Differentiation #####################
reset_default_graph()
model_name = "models/marmousi2-model-smooth.mat"
# model_name = "models/marmousi2-model-smooth-large.mat"
# model_name = "models/BP-model-smooth.mat"

## load model setting
params = load_params(model_name, vp_ref=1e3)
src = load_acoustic_source(model_name)
rcv = load_acoustic_receiver(model_name)
vp0 = matread(model_name)["vp"]

## original fwi
# vp = Variable(vp0)

## fwi with fc
# nx, ny = params.NX, params.NY
# dx, dy = params.DELTAX, params.DELTAY
# z = zeros((nx+2)*(ny+2),2)
# for i = 1:nx+2
#  for j = 1:ny+2
#     z[(i-1)*(ny+2)+j, :] = [(i-1)*dx (j-1)*dy]
#   end
# end
# init_guess = fc_init([2, 20, 20, 20, 1])
# θ = Variable(init_guess)
# c = fc(z, [20, 20, 20, 1], θ) * 100.0
# c = reshape(c, (nx+2, ny+2))
# vp = vp0 + c

## fwi with cnn
Random.seed!(123)
# tf.set_random_seed(123)
z = constant(rand(Float32, 1, 8))
# x = Generator(z, ratio=size(vp0)[1]/size(vp0)[2], base=8)
num_layer = 4
Δvp = Generator(z, w0=ceil(Int, size(vp0)[1]/(2^num_layer)), h0=ceil(Int, size(vp0)[2]/(2^num_layer)), num_layer=num_layer, vmin=-1.5, vmax=1.5)
Δvp = tf.cast(Δvp, tf.float64)*1000.0
# Δvp = Variable(zeros(size(vp0)))
vp = add_initial_model(Δvp, constant(vp0))
if exists(matopen(model_name), "mask")
  mask = matread(model_name)["mask"]
  vp = mask .* vp + (1.0.-mask) .* vp0
end

## assemble acoustic propagator model
model = x->AcousticPropagatorSolver(params, x, vp)
vars = get_collection()

## load data
std_noise = 0
Random.seed!(1234);
Rs = Array{Array{Float64,2}}(undef, length(src))
for i = 1:length(src)
  Rs[i] = readdlm(joinpath(data_dir, "marmousi-r$i.txt"))
    # Rs[i] = readdlm(joinpath(data_dir, "BP-r$i.txt"))
  Rs[i] .+= randn(size(Rs[i])) .* std(Rs[i]) .* std_noise
end

## calculate loss and gradient
if gpu
  losses, grads = compute_loss_and_grads_GPU(model, src, rcv, Rs, vars)
  grad = sum(grads)
  loss = sum(losses)
else
  [SimulatedObservation!(model(src[i]), rcv[i]) for i = 1:length(src)]
  loss = sum([sum((rcv[i].rcvv-Rs[i])^2) for i = 1:length(rcv)])
  grad = gradients(loss, vars)
end

## optimization using Adam
# global_step = tf.Variable(0, trainable=false)
max_iter = 10000
# lr_decayed = tf.train.cosine_decay(0.001, global_step, max_iter)
lr_decayed = placeholder(0.0)
# opt = AdamOptimizer(learning_rate=lr_decayed).minimize(loss, global_step=global_step, colocate_gradients_with_ops=true)
opt = AdamOptimizer(learning_rate=lr_decayed).minimize(loss, colocate_gradients_with_ops=true)

## run inversion
sess = Session(); init(sess)
loss0 = run(sess, loss)
@info "Initial loss: ", loss0

fp = open(loss_file, "w")
write(fp, "0,$loss0\n")
iter0 = 0
function callback(_, iter, loss)
  iter += iter0
  if iter%10==0
    x = run(sess, vp)'
    clf()
    pcolormesh([0:params.NX+1;]*params.DELTAX/1e3,[0:params.NY+1;]*params.DELTAY/1e3,  x, shading="auto", cmap="jet")
    axis("scaled")
    colorbar(shrink=0.4)
    xlabel("x (km)")
    ylabel("z (km)")
    gca().invert_yaxis()
    title("Iteration = $iter")
    savefig(joinpath(figure_dir, "inv_$(lpad(iter,5,"0")).png"), bbox_inches="tight")
    writedlm(joinpath(result_dir, "inv_$(lpad(iter,5,"0")).txt"), x)
  end
  if iter%100 == 0
    ADCME.save(sess, joinpath(model_dir, "NNFWI_$(lpad(iter,5,"0")).mat"))
  end
  write(fp, "$iter,$loss\n")
  flush(fp)
end

## optimization using Adam
time = 0
for iter = 0:max_iter
  # global time += @elapsed  _, ls, lr = run(sess, [opt, loss, lr_decayed])
  # if iter < 500
  #   lr = iter/500 * 1e-3
  # else
  #   # lr = (max_iter - iter)/(max_iter - 500) * 1e-3
  #   lr = 1e-3
  # end
  lr = 1e-3
  t = @elapsed  _, ls = run(sess, [opt, loss], feed_dict=Dict(lr_decayed=>lr))
  callback(nothing, iter, ls)
  println("   $iter\t$ls\t$lr\t$t")
end

## optimization using BFGS
iter0 = max_iter+1
Optimize!(sess, loss, 1000, vars=vars, grads=grad, callback=callback)

close(fp)


"""
    cauchyendgame(H::AbstractHomotopy, J_H, ∂H∂t, x, R, algorithm, kwargs..., pathtrackingkwargs...)

Execute a cauchy endgame for `H` starting at radius `R`. This assumes that ``H(x,t)=0``.
It is based on "A Parallel Endgame " by Bates, Hauenstein and Sommese. [^1]

## Optional arguments
* `prediction_tolerance=1e-4`:
* `geometric_series_factor=0.5`:
* `endgame_tolerance=1e-8`:
* `samples_per_loop=8`:
* `max_winding_number=8`:
* `loopclosed_tolerance=1e-3`
*  optional arguments to [`pathtracking`](@ref)

[^1]: Bates, Daniel J., Jonathan D. Hauenstein, and Andrew J. Sommese. "A Parallel Endgame." Contemp. Math 556 (2011): 25-35.
"""
function cauchyendgame(H::AbstractHomotopy{T}, J_H, ∂H∂t, x, t::Float64, algorithm::APCA{true};
    prediction_tolerance=1e-8,
    geometric_series_factor=0.5,
    endgame_tolerance=1e-8,
    samples_per_loop::Int=8,
    max_winding_number::Int=8,
    loopclosed_tolerance=1e-3,
    tolerance_infinity=1e-8,
    pathtrackingkwargs...
) where {T}
    λ = geometric_series_factor
    R = t
    steps = [R]
    xs = [x]
    lastprediction = Nullable(x, false)
    endgame_started = false
    k = 1
    while R > eps(Float64)
        res = trackpath(H, J_H, ∂H∂t, last(xs), algorithm, R, λ*R; pathtrackingkwargs...)
        if nosuccess(res)
            return CauchyEndgameResult(get(lastprediction, last(xs)), :IllConditionedZone, k, R, λ, xs, steps, Nullable{ConvergentCluster{T}}())
        end
        R = λ*R
        push!(steps, R)
        push!(xs, res.result)

        # we need at least 4 points to start with our first heuristic
        if length(xs) < 4
            continue
        end

        if atinfinity(get(lastprediction, last(xs)), tolerance_infinity)
            return CauchyEndgameResult(get(lastprediction, last(xs)), :AtInfinity, k, R, λ, xs, steps, Nullable{ConvergentCluster{T}}())
        end

        if endgame_started || firstheuristic(R, λ, xs[end-3], xs[end-2], xs[end-1], xs[end])
            retcode, samples =
                loop(H, J_H, ∂H∂t, last(xs), R, algorithm, samples_per_loop, max_winding_number,
                    loopclosed_tolerance, endgame_tolerance, apply_heuristic=!endgame_started)
            if retcode == :Success
                endgame_started = true
                prediction = predict_with_cif(samples)
                if !isnull(lastprediction)
                    # check whether to predictions are close enough, so we can say that
                    # the endgame terminates
                    Δprediction = projectivenorm(prediction, get(lastprediction))
                    if Δprediction < prediction_tolerance
                        cluster = ConvergentCluster(samples[1:samples_per_loop:end], prediction, R)
                        return CauchyEndgameResult(prediction, :Success, k, R, λ, xs, steps, Nullable(cluster))
                    end
                end
                lastprediction = Nullable(prediction)
            end
        end

        k += 1
    end

    CauchyEndgameResult(get(lastprediction, last(xs)), :MachineEpsilon, k, R, λ, xs, steps, Nullable{ConvergentCluster{T}}())
end

atinfinity(x, tolerance_infinity) = norm(normalize(x)[1]) < tolerance_infinity

"""
    loop(H, J_H, ∂H∂t, x, radius, algorithm, samples_per_loop, max_winding_number, loopclosed_tolerance, apply_heuristic=true, pathtrackingkwargs...)

Tracks the implicit defined path z(t) around the `n`-gon with vertices
``r⋅exp(i2πk/n)`` where `n=samples_per_loop`.
"""
function loop(
    H::AbstractHomotopy{T},
    J_H,
    ∂H∂t,
    x,
    radius::Real,
    algorithm::APCA,
    samples_per_loop::Int,
    max_winding_number::Int,
    loopclosed_tolerance::Float64,
    endgame_tolerance::Float64;
    apply_heuristic=true,
    pathtrackingkwargs...
) where {T}
    unitroots = UnitRootsIterator(radius, samples_per_loop)
    samples = [x]
    start = first(unitroots)
    for (k, finish) in enumerate(Iterators.drop(unitroots, 1))
        pathresult = trackpath(H, J_H, ∂H∂t, samples[end], algorithm, start, finish; pathtrackingkwargs...)
        if pathresult.returncode != :Success
            return (:BadBranch, samples)
        end

        # apply the second heuristic to tell us whether it makes sense to start the endgame
        if apply_heuristic && k == samples_per_loop &&
            !secondheuristic(radius, samples, β=endgame_tolerance)
            return (:HeuristicFailed, samples)
        end

        # check whether the loop is closed. This can only happen if we finished one circle.
        if k % samples_per_loop == 0
            loopnorm = projectivenorm(first(samples), pathresult.result)
            if loopnorm < loopclosed_tolerance
                break
            end
        end

        if k > max_winding_number * samples_per_loop
            return (:WindingNumberTooHigh, samples)
        end

        push!(samples, pathresult.result)
        start = finish
    end
    (:Success, samples)
end

"""
    firstheuristic(t, x_R, x_λR, x_λ2R, x_λ3R; L = 0.75)

This is the first heuristic when to start the endgame.
It checks whether two consectuive appromixations of the first nonvanishing
monomial of the Puiseux series expansion of x(t) at 0 agree.
For more details see "A Parallel Endgame " by Bates, Hauenstein and Sommese [^1],
page 8 and 9.

[^1]: Bates, Daniel J., Jonathan D. Hauenstein, and Andrew J. Sommese. "A Parallel Endgame." Contemp. Math 556 (2011): 25-35.
"""
function firstheuristic(t, λ, x_R, x_λR, x_λ2R, x_λ3R; L = 0.75)
    # failsafe mechanism
    if t < 1e-8
        return true
    end
    v = rand(Complex128, length(x_R))
    g(x_R, x_λR, x_λ2R) = log(abs2((v⋅x_λR - v⋅x_λ2R) / (v⋅x_R - v⋅x_λR))) / log(λ)

    g_R = g(x_R, x_λR, x_λ2R)
    g_λR = g(x_λR, x_λ2R, x_λ3R)

    if g_R <= 0
        return false
    end

    L < g_R / g_λR < 1 / L
end

"""
    secondheuristic(t, samplepoints, β; K=0.5)

This is the second heuristic when to start the endgame.
It enforces that the values around the lopp do not differ radically.
This functions assumes that `samplepoints` has been collected on ``x(Re^{iθ})`` for ``θ∈[0,2π]``.
See "A Parallel Endgame " by Bates, Hauenstein and Sommese [^1], page 9 for more details.

[^1]: Bates, Daniel J., Jonathan D. Hauenstein, and Andrew J. Sommese. "A Parallel Endgame." Contemp. Math 556 (2011): 25-35.
"""
function secondheuristic(radius, samplepoints; β=1e-8, K=0.5)
    # failsafe
    if radius < 1e-14
        true
    end

    norms = norm.(samplepoints)
    m = minimum(norms)
    M = maximum(norms)
    (M - m < β) || (m / M > K)
end

"""
    projectivenorm(a, b)

Calculate |a-b|, but first bring `a` and `b` on the same patch.
Brings fist `a` and `b` on the same patch by finding diving `a` through it's maximum value
(w.r.t. to the absolute value) with index `i` and diving `b` through `b[i]`.
Then computes the norm of the differences.
"""
function projectivenorm(a::Vector{<:Number}, b::Vector{<:Number})
    _, i = findmax(abs2.(a))
    norm(a ./ a[i] .- b ./ b[i])
end

"""
    predict_with_cif(samplepoints)

Predict the value of a function based on the samplepoints using the cauchy integral formula.
This function assumes that the samplepoints are sampled uniformly from a circle around 0,
since then the integral (approximated with the trapezoidal rule) is simply the mean of
all samplepoints.
"""
predict_with_cif(samplepoints) = mean(samplepoints)

"""
    UnitRootsIterator(r, n)

Construct an infinite iterator which returns the `n`-th scalded unit roots, i.e.
the values ``r⋅exp(i2πk/n)`` for ``k=0,1,...``.
"""
struct UnitRootsIterator
    radius::Float64
    order::Float64
end
UnitRootsIterator(r::Real, order::Real) = UnitRootsIterator(float(r), float(order))
Base.start(::UnitRootsIterator) = 0
Base.next(loop::UnitRootsIterator, k::Int) = (loop.radius *  exp(im * 2π * k / loop.order), k + 1)
Base.done(::UnitRootsIterator, ::Int) = false
Base.iteratorsize(::UnitRootsIterator) = Base.IsInfinite()
Base.iteratoreltype(::UnitRootsIterator) = Base.HasEltype()
Base.eltype(::UnitRootsIterator) = Complex128


using OrdinaryDiffEq, OpenQuantumTools
using Plots, StatsBase

H = DenseHamiltonian([(s)->1.0], [-σz], unit=:ħ)
u0 = PauliVec[1][1]

coupling = ConstantCouplings(["Z"], unit=:ħ)
# number of fluctuators
num = 10
# The values of b created here are in angular frequency units:
bvec = 0.01 * ones(num)
# γᵢ
γvec = log_uniform(0.01, 1, num)
# create the fluctuator coupling interaction
fluctuator_ensemble = EnsembleFluctuator(bvec, γvec);
interaction_fluctuator = Interaction(coupling, fluctuator_ensemble)
# create the Ohmic coupling interaction
ohmic_bath = Ohmic(1e-4, 4, 16)
interaction_ohmic = Interaction(coupling, ohmic_bath)
# merge these two bath objects into `InteractionSet`
interactions = InteractionSet(interaction_fluctuator, interaction_ohmic)
annealing = Annealing(H, u0, interactions=interactions)

tf = 100
prob = build_ensembles(annealing, tf, :ame, ω_hint = range(-2, 2, length=100))
t_list = range(0,tf,length=200)
sol = solve(prob, Tsit5(), EnsembleSerial(), trajectories=1000, reltol=1e-6, saveat=t_list)

dataset = []
st(s, so) = normalize(so(s, continuity=:left))
for s in t_list
    push!(dataset, [real(st(s, so)' * σx * st(s, so)) for so in sol])
end

pop_mean = []
pop_rmse = []
for data in dataset
    p_mean = sum(data)/1000
    p_rmse = sqrt(sum((x)->(x-p_mean)^2, data))/1000
    push!(pop_mean, p_mean)
    push!(pop_rmse, 2*p_rmse)
end

a_list = range(0,tf,length=400)
annealing_ame = Annealing(H, u0, coupling=coupling, bath=ohmic_bath)
sol_ame = solve_ame(annealing_ame, tf, alg=Tsit5(), ω_hint = range(-2, 2, length=100),
    reltol=1e-6, saveat=a_list)
ame_x = [real(tr(u*σx)) for u in sol_ame.u]

plot(a_list, ame_x, label="Ohmic", linewidth=2)
plot!(t_list, pop_mean, ribbon=pop_rmse, label="Hyrbrid", linewidth=2)
xlabel!("s (ns)")
ylabel!("<X>")

using HOQSTTutorials
HOQSTTutorials.tutorial_footer(WEAVE_ARGS[:folder],WEAVE_ARGS[:file])


@testset "'Design' ............................. " begin

  srand(1234)

  function simonsDesign(r1, n1, r, n)
    nvec = [[n1 for x1 in 0:r1]; [n for x1 in (r1 + 1):n1]]
    cvec = [[Inf for x1 in 0:r1]; [r for x1 in (r1 + 1):n1]]
    return Design(nvec, cvec)
  end

  # Simon's designs for beta = .2, alpha = .05, p1 = p0 +0.2
  # cf. Simon, R "Optimal Two-Stage Designs for Phase II Clinical Trials",
  # Controlled Clinical Trials 10:1-10 (1989). (p. 4)
  p0 = collect(linspace(.1, .7, 7))
  sd = [
    simonsDesign( 1, 10,  5, 29), # p0 = 0.1
    simonsDesign( 3, 13, 12, 43), # p0 = 0.2
    simonsDesign( 5, 15, 18, 46), # p0 = 0.3
    simonsDesign( 7, 16, 23, 46), # p0 = 0.4
    simonsDesign( 8, 15, 26, 43), # p0 = 0.5
    simonsDesign( 7, 11, 30, 43), # p0 = 0.6
    simonsDesign( 4,  6, 22, 27)  # p0 = 0.7
  ]

  @test typeof(sd[1]) == Design{typeof(1), typeof(Inf), NoParameters}

  @test DataFrames.DataFrame(sd[1]) == DataFrames.DataFrame(
        x1 = 0:interimsamplesize(sd[1]),
        n  = samplesize(sd[1]),
        c  = criticalvalue(sd[1])
    )

  @test interimsamplesize(sd[1]) == 10

  @test parameters(sd[1]) == NoParameters()

  @test samplesize(sd[1]) == [[10 for x1 in 0:1]; [29 for x1 in (1 + 1):10]]
  @test samplesize(sd[1], 1) == 10
  @test samplesize(sd[1], 2) == 29

  @test criticalvalue(sd[1]) == [[Inf for x1 in 0:1]; [5 for x1 in (1 + 1):10]]
  @test criticalvalue(sd[1], 1) == Inf
  @test criticalvalue(sd[1], 2) == 5

  supp_ = support(sd[1])
  for p in linspace(0, 1, 11)

    @test pdf(sd[1], 1, 2, p) == 0.0
    pdf_ = pdf.(sd[1], support(sd[1])[:,1], support(sd[1])[:,2], p)
    @test all(pdf_ .>= 0.0)
    @test isapprox(sum(pdf_), 1.0, atol = 0.00001)

  end

  @test power(sd[1], p0[1]) <= .05
  @test power(sd[1], p0[1] + .2) >= .8

  prior = p -> 1
  @test expectedpower(sd[1], prior, mcrv = p0[1] + .1) < power(sd[1], .8)
  @test expectedpower(sd[1], prior, mcrv = p0[1] + .1) > power(sd[1], .2)
  @test expectedpower(sd[1], 1, prior, mcrv = p0[1] + .1) < expectedpower(sd[1], 2, prior, mcrv = p0[1] + .1)

  @test isapprox(
    stoppingforfutility(sd[1], .2), pdf(sd[1], 0, 0, .2) + pdf(sd[1], 1, 0, .2), atol = .00001
  )

  @test !test(sd[1], 0, 0)
  @test test(sd[1], 5, 11)

  sim_ = simulate(sd[1], p0[1], 1000)
  @test mean(sim_[:rejectedH0]) < .05
  @test all(sim_[:n] .<= 29)
  @test all(sim_[:n] .>= 0)

  @test isapprox(quadgk(jeffreysprior(sd[1]), 0, 1)[1], 1.0, atol = 0.0001)

  save("test.jls", sd[1])
  @test true
  rm("test.jls")
  writecsv("test.csv", sd[1]; label = "test")
  @test true
  rm("test.csv")

end


function build_astar_path(came_from::Dict, t::Integer, v::Integer, t0::Integer)
    path = Path()
    (τ, u) = (t, v)
    pushfirst!(path, (τ, u))
    while τ > t0
        (τ, u) = came_from[(τ, u)]
        pushfirst!(path, (τ, u))
    end
    return path
end

function temporal_astar(
    g::AbstractGraph{V},
    s::Integer,
    d::Integer,
    t0::Integer;
    edge_indices::Dict,
    edge_weights::Vector{W},
    heuristic=v -> 0.0,
    reservation::Reservation=Reservation(),
    conflict_price=Inf,
) where {V,W}
    T = Int
    safe_conflict_price = conflict_price * (conflict_price < Inf)

    # Initialize storage
    came_from = Dict{Tuple{T,V},Tuple{T,V}}()
    distance = Dict{Tuple{T,V},W}()
    conflicts = Dict{Tuple{T,V},Int}()
    queue = PriorityQueue{Tuple{T,V},Float64}()

    # Add first node to storage
    distance_s = zero(W)
    conflicts_s = Int((t0, s) in reservation)
    if conflicts_s == 0 || conflict_price < Inf
        priority_s = safe_conflict_price * conflicts_s + heuristic(s)
        distance[(t0, s)] = distance_s
        conflicts[(t0, s)] = conflicts_s
        enqueue!(queue, (t0, s), priority_s)
    end

    # Explore
    nodes_explored = 0
    while !isempty(queue)
        (t, v) = dequeue!(queue)
        nodes_explored += 1
        (v == d) && return build_astar_path(came_from, t, v, t0)

        for w in outneighbors(g, v)
            heur_w = heuristic(w)

            e_vw = edge_indices[(v, w)]
            weight_vw = edge_weights[e_vw]
            conflict_vw = Int((t + 1, w) in reservation)

            distance_w = distance[(t, v)] + weight_vw
            conflicts_w = conflicts[(t, v)] + conflict_vw

            if conflicts_w == 0 || conflict_price < Inf
                old_distance_w = get(distance, (t + 1, w), Inf)
                old_conflicts_w = get(conflicts, (t + 1, w), typemax(Int) ÷ 2)

                old_cost_w = (safe_conflict_price * old_conflicts_w + old_distance_w)
                cost_w = (safe_conflict_price * conflicts_w + distance_w)

                if cost_w < old_cost_w
                    came_from[(t + 1, w)] = (t, v)
                    distance[(t + 1, w)] = distance_w
                    conflicts[(t + 1, w)] = conflicts_w
                    queue[(t + 1, w)] = cost_w + heur_w
                end
            end
        end
    end
    return Path()
end


# module for calculating std of sensor observation
# observe a landmark by a sensor
# iterate for 1000 times

module SensorTestStatic
  using DataFrames, Statistics

  include(joinpath(split(@__FILE__, "src")[1], "src/model/world/world.jl"))
  include(joinpath(split(@__FILE__, "src")[1], "src/model/robot/differential_wheeled_robot/differential_wheeled_robot.jl"))

  function main()
    m = Map()
    add_object(m, Object(1.0, 0.0))

    distances = []
    directions = []

    for i in 1:1000
      s = Sensor(m, dist_noise_rate=0.1,
                 dir_noise=pi/90, dist_bias_rate_stddev=0.1,
                 dir_bias_stddev=pi/90)
      data(s, [0.0, 0.0, 0.0])
      if length(s.last_data) > 0
        push!(distances, s.last_data[1][1][1])
        push!(directions, s.last_data[1][1][2])
      end
    end

    df = DataFrame(distance=distances, direction=directions)
    println("Observation data: ", df)
    println("standard deviation of distance: ", std(df.distance, corrected=false))
    println("standard deviation of direction: ", std(df.direction, corrected=false))
  end
end