# ------------------------------------------------------------------------------------------
# # Character based RNN language model trained on 'The Complete Works of William
# Shakespeare'
# Based on http://karpathy.github.io/2015/05/21/rnn-effectiveness
# ------------------------------------------------------------------------------------------

RNNTYPE = :lstm
BATCHSIZE = 64
SEQLENGTH = 64
INPUTSIZE = 512
VOCABSIZE = 128
HIDDENSIZE = 512
NUMLAYERS = 2
DROPOUT = 0.0
LR=0.001
BETA_1=0.9
BETA_2=0.999
EPS=1e-08
EPOCHS = 30;

base = joinpath(Base.JULIA_HOME, Base.DATAROOTDIR, "julia", "base")
text = ""
for (root,dirs,files) in walkdir(base)
    for f in files
        f[end-2:end] == ".jl" || continue
        text *= readstring(joinpath(root,f))
    end
    # println((root,length(files),all(f->contains(f,".jl"),files)))
end
length(text)

charcnt = Dict{Char,Int}()
for c in text; charcnt[c]=1+get(charcnt,c,0); end
chars = sort(collect(keys(charcnt)), by=(x->charcnt[x]), rev=true)
charid = Dict{Char,Int}()
for i=1:length(chars); charid[chars[i]]=i; end
hcat(chars, map(c->charcnt[c],chars))

data = map(c->charid[c], collect(text))
data[data .> VOCABSIZE] = VOCABSIZE
ntst = 1<<19
tst = data[1:ntst]
trn = data[1+ntst:end]
length.((data,trn,tst))

# Load 'The Complete Works of William Shakespeare'
using Knet
#include(Knet.dir("data","gutenberg.jl"))
#trn,tst,chars = shakespeare()
#map(summary,(trn,tst,chars))

# Print a sample
r = rand(1:(length(trn)-1000))
println(string(chars[trn[r:r+1000]]...))

# Minibatch data
function mb(a)
    N = div(length(a),BATCHSIZE)
    x = reshape(a[1:N*BATCHSIZE],N,BATCHSIZE)' # reshape full data to (B,N) with contiguous rows
    minibatch(x[:,1:N-1], x[:,2:N], SEQLENGTH) # split into (B,T) blocks
end
dtrn,dtst = mb(trn),mb(tst)
map(length, (dtrn,dtst))

# Define model
function initmodel()
    w(d...)=KnetArray(xavier(Float32,d...))
    b(d...)=KnetArray(zeros(Float32,d...))
    r,wr = rnninit(INPUTSIZE,HIDDENSIZE,rnnType=RNNTYPE,numLayers=NUMLAYERS,dropout=DROPOUT)
    wx = w(INPUTSIZE,VOCABSIZE)
    wy = w(VOCABSIZE,HIDDENSIZE)
    by = b(VOCABSIZE,1)
    return r,wr,wx,wy,by
end;

# Given the current character, predict the next character
function predict(ws,xs,hx,cx;pdrop=0)
    r,wr,wx,wy,by = ws
    x = wx[:,xs]                                    # xs=(B,T) x=(X,B,T)
    x = dropout(x,pdrop)
    y,hy,cy = rnnforw(r,wr,x,hx,cx,hy=true,cy=true) # y=(H,B,T) hy=cy=(H,B,L)
    y = dropout(y,pdrop)
    y2 = reshape(y,size(y,1),size(y,2)*size(y,3))   # y2=(H,B*T)
    return wy*y2.+by, hy, cy
end

# Define loss and its gradient
function loss(w,x,y,h;o...)
    py,hy,cy = predict(w,x,h...;o...)
    h[1],h[2] = getval(hy),getval(cy)
    return nll(py,y)
end

lossgradient = gradloss(loss);

function train(model,data,optim)
    hiddens = Any[nothing,nothing]
    Σ,N=0,0
    for (x,y) in data
        grads,loss1 = lossgradient(model,x,y,hiddens;pdrop=DROPOUT)
        update!(model, grads, optim)
        Σ,N=Σ+loss1,N+1
    end
    return Σ/N
end;

function test(model,data)
    hiddens = Any[nothing,nothing]
    Σ,N=0,0
    for (x,y) in data
        Σ,N = Σ+loss(model,x,y,hiddens), N+1
    end
    return Σ/N
end;

# Sample from trained model
function generate(model,n)
    function sample(y)
        p,r=Array(exp.(y-logsumexp(y))),rand()
        for j=1:length(p); (r -= p[j]) < 0 && return j; end
    end
    h,c = nothing,nothing
    chars = model[end]
    x = findfirst(chars,'\n')
    for i=1:n
        y,h,c = predict(model,[x],h,c)
        x = sample(y)
        print(chars[x])
    end
    println()
end;

#=
# Train model or load from file if exists
using JLD
model=optim=nothing; knetgc()
if true # !isfile("juliacharlm.jld")
    model = initmodel()
    optim = optimizers(model, Adam; lr=LR, beta1=BETA_1, beta2=BETA_2, eps=EPS);    info("Training...")
    @time for epoch in 1:EPOCHS
        @time trnloss = train(model,dtrn,optim) # ~18 seconds
        @time tstloss = test(model,dtst)        # ~0.5 seconds
        println((:epoch, epoch, :trnppl, exp(trnloss), :tstppl, exp(tstloss)))
    end
    save("juliacharlm.jld","model",model)
else
    model = load("juliacharlm.jld","model")
end
summary(model)
=#

# generate(model,1000)

function hyperloss(v)
    global HIDDENSIZE, INPUTSIZE, NUMLAYERS, DROPOUT, EPOCHS, model
    HIDDENSIZE = floor(Int, 128 * 2^v[1])
    INPUTSIZE = floor(Int, 128 * 2^v[2])
    NUMLAYERS = max(1,floor(Int,v[3]))
    DROPOUT = isa(v[4],Number) ? sigm(v[4]) : 0
    @show (HIDDENSIZE,INPUTSIZE,NUMLAYERS,DROPOUT)
    knetgc()
    model = initmodel()
    optim = optimizers(model, Adam; lr=LR, beta1=BETA_1, beta2=BETA_2, eps=EPS)
    for epoch in 1:10
        trnloss = train(model,dtrn,optim) # ~18 seconds
        tstloss = test(model,dtst)        # ~0.5 seconds
        println((:epoch, epoch, :trnppl, exp(trnloss), :tstppl, exp(tstloss)))
    end
    trnloss = test(model,dtrn)
    trnppl = exp(trnloss)
    @show (HIDDENSIZE,INPUTSIZE,NUMLAYERS,DROPOUT,trnppl)
    return trnppl
end

model=nothing; knetgc(); hyperloss([3,0,3,0])

#=

(HIDDENSIZE, INPUTSIZE, NUMLAYERS, DROPOUT) = (1024, 128, 2, 0.5)
(:epoch, 1, :trnppl, 9.589466f0, :tstppl, 4.9871907f0)
(:epoch, 2, :trnppl, 3.741831f0, :tstppl, 3.9569008f0)
(:epoch, 3, :trnppl, 3.111019f0, :tstppl, 3.6438942f0)
(:epoch, 4, :trnppl, 2.842651f0, :tstppl, 3.490925f0)
(:epoch, 5, :trnppl, 2.6792212f0, :tstppl, 3.406119f0)
(:epoch, 6, :trnppl, 2.567202f0, :tstppl, 3.337223f0)
(:epoch, 7, :trnppl, 2.484569f0, :tstppl, 3.2894032f0)
(:epoch, 8, :trnppl, 2.4184995f0, :tstppl, 3.2509122f0)
(:epoch, 9, :trnppl, 2.3658345f0, :tstppl, 3.2483802f0)
(:epoch, 10, :trnppl, 2.3232405f0, :tstppl, 3.2095058f0)
(HIDDENSIZE, INPUTSIZE, NUMLAYERS, DROPOUT, trnppl) = (1024, 128, 2, 0.5, 2.0069642f0)

(HIDDENSIZE, INPUTSIZE, NUMLAYERS, DROPOUT) = (4096, 128, 1, 0.5)
(:epoch, 1, :trnppl, 7.630349f0, :tstppl, 8.994627f0)
(:epoch, 2, :trnppl, 3.1094067f0, :tstppl, 3.7359138f0)
(:epoch, 3, :trnppl, 2.5211895f0, :tstppl, 3.406949f0)
(:epoch, 4, :trnppl, 2.24173f0, :tstppl, 3.3002508f0)
(:epoch, 5, :trnppl, 2.0695934f0, :tstppl, 3.2824345f0)
(:epoch, 6, :trnppl, 1.9510676f0, :tstppl, 3.2519991f0)
(:epoch, 7, :trnppl, 1.8637797f0, :tstppl, 3.2833624f0)
(:epoch, 8, :trnppl, 1.7990172f0, :tstppl, 3.2755644f0)
(:epoch, 9, :trnppl, 1.7516453f0, :tstppl, 3.3127692f0)
(:epoch, 10, :trnppl, 1.7112837f0, :tstppl, 3.3849468f0)
(HIDDENSIZE, INPUTSIZE, NUMLAYERS, DROPOUT, trnppl) = (4096, 128, 1, 0.5, 1.5795085f0)
1.5795085f0

(HIDDENSIZE, INPUTSIZE, NUMLAYERS, DROPOUT) = (2048, 128, 1, 0.5)
(:epoch, 1, :trnppl, 6.897045f0, :tstppl, 4.5084825f0)
(:epoch, 2, :trnppl, 3.1653059f0, :tstppl, 3.7902246f0)
(:epoch, 3, :trnppl, 2.6861045f0, :tstppl, 3.5171616f0)
(:epoch, 4, :trnppl, 2.4583292f0, :tstppl, 3.4100456f0)
(:epoch, 5, :trnppl, 2.3147066f0, :tstppl, 3.3488684f0)
(:epoch, 6, :trnppl, 2.2171206f0, :tstppl, 3.3439498f0)
(:epoch, 7, :trnppl, 2.1446044f0, :tstppl, 3.3066962f0)
(:epoch, 8, :trnppl, 2.0894327f0, :tstppl, 3.3102229f0)
(:epoch, 9, :trnppl, 2.0460289f0, :tstppl, 3.287103f0)
(:epoch, 10, :trnppl, 2.0100765f0, :tstppl, 3.3185039f0)
(HIDDENSIZE, INPUTSIZE, NUMLAYERS, DROPOUT, trnppl) = (2048, 128, 1, 0.5, 1.7895192f0)
1.7895192f0

(HIDDENSIZE, INPUTSIZE, NUMLAYERS, DROPOUT) = (2896, 128, 1, 0.5)
(:epoch, 1, :trnppl, 7.1333003f0, :tstppl, 4.4698453f0)
(:epoch, 2, :trnppl, 3.04668f0, :tstppl, 3.719867f0)
(:epoch, 3, :trnppl, 2.5365477f0, :tstppl, 3.439777f0)
(:epoch, 4, :trnppl, 2.295574f0, :tstppl, 3.3415797f0)
(:epoch, 5, :trnppl, 2.1462688f0, :tstppl, 3.27515f0)
(:epoch, 6, :trnppl, 2.044871f0, :tstppl, 3.2823677f0)
(:epoch, 7, :trnppl, 1.9718276f0, :tstppl, 3.2806444f0)
(:epoch, 8, :trnppl, 1.9180799f0, :tstppl, 3.2798069f0)
(:epoch, 9, :trnppl, 1.8782355f0, :tstppl, 3.2703993f0)
(:epoch, 10, :trnppl, 1.8446114f0, :tstppl, 3.3073356f0)
(HIDDENSIZE, INPUTSIZE, NUMLAYERS, DROPOUT, trnppl) =
1.6504527f0
(2896, 128, 1, 0.5, 1.6504527f0)

(HIDDENSIZE, INPUTSIZE, NUMLAYERS, DROPOUT) = (2048, 128, 1, 0)
(:epoch, 1, :trnppl, 17.54582f0, :tstppl, 9.120878f0)
(:epoch, 2, :trnppl, 5.3387966f0, :tstppl, 5.298701f0)
(:epoch, 3, :trnppl, 3.4404564f0, :tstppl, 4.2200007f0)
(:epoch, 4, :trnppl, 2.7447715f0, :tstppl, 3.8503358f0)
(:epoch, 5, :trnppl, 2.3742914f0, :tstppl, 3.695858f0)
(:epoch, 6, :trnppl, 2.1303215f0, :tstppl, 3.6608467f0)
(:epoch, 7, :trnppl, 1.9553958f0, :tstppl, 3.687474f0)
(:epoch, 8, :trnppl, 1.8226186f0, :tstppl, 3.7828565f0)
(:epoch, 9, :trnppl, 1.7212292f0, :tstppl, 3.8975382f0)
(:epoch, 10, :trnppl, 1.6417012f0, :tstppl, 4.0848227f0)
(HIDDENSIZE, INPUTSIZE, NUMLAYERS, DROPOUT, trnppl) =
1.7484467f0
(2048, 128, 1, 0, 1.7484467f0)

(HIDDENSIZE, INPUTSIZE, NUMLAYERS, DROPOUT) = (1024, 128, 2, 0.5)
(:epoch, 1, :trnppl, 22.390598f0, :tstppl, 13.189085f0)
(:epoch, 2, :trnppl, 7.4823937f0, :tstppl, 5.720241f0)
(:epoch, 3, :trnppl, 4.6992884f0, :tstppl, 4.448046f0)
(:epoch, 4, :trnppl, 3.7799704f0, :tstppl, 3.972568f0)
(:epoch, 5, :trnppl, 3.3514497f0, :tstppl, 3.7654393f0)
(:epoch, 6, :trnppl, 3.0903006f0, :tstppl, 3.5971427f0)
(:epoch, 7, :trnppl, 2.9079294f0, :tstppl, 3.484381f0)
(:epoch, 8, :trnppl, 2.7710462f0, :tstppl, 3.443712f0)
(:epoch, 9, :trnppl, 2.663027f0, :tstppl, 3.3361208f0)
(:epoch, 10, :trnppl, 2.5764978f0, :tstppl, 3.3100765f0)
(HIDDENSIZE, INPUTSIZE, NUMLAYERS, DROPOUT, trnppl) =
2.175947f0
(1024, 128, 2, 0.5, 2.175947f0)

(HIDDENSIZE, INPUTSIZE, NUMLAYERS, DROPOUT) = (1024, 128, 2, 0)
(:epoch, 1, :trnppl, 22.009565f0, :tstppl, 11.428911f0)
(:epoch, 2, :trnppl, 6.030965f0, :tstppl, 5.366185f0)
(:epoch, 3, :trnppl, 3.614693f0, :tstppl, 4.288177f0)
(:epoch, 4, :trnppl, 2.906743f0, :tstppl, 3.948376f0)
(:epoch, 5, :trnppl, 2.572164f0, :tstppl, 3.8014898f0)
(:epoch, 6, :trnppl, 2.3514614f0, :tstppl, 3.7232788f0)
(:epoch, 7, :trnppl, 2.1846726f0, :tstppl, 3.7083063f0)
(:epoch, 8, :trnppl, 2.056641f0, :tstppl, 3.7179863f0)
(:epoch, 9, :trnppl, 1.9550104f0, :tstppl, 3.768648f0)
(:epoch, 10, :trnppl, 1.8722341f0, :tstppl, 3.8625576f0)
(HIDDENSIZE, INPUTSIZE, NUMLAYERS, DROPOUT, trnppl) =
1.9898754f0

=#

#foo = download("http://ai.ku.edu.tr/models/nlp-demos/juliacharlm.jld")
using JLD
model = nothing; knetgc()
model = load(foo,"model")
summary.(model)

cpumodel = (model[1],map(Array,model[2:end])...)

generate(cpumodel,1000)

dump(cpumodel[1])

cpumodel[1].rnnDesc = cpumodel[1].dropoutDesc = nothing

dump(cpumodel[1])

save("juliacharlm.jld","model",cpumodel)

pwd()

foo

shake = load("shakespeare.jld","model")

generate(shake,1000)

cpumodel = (cpumodel..., chars)

include(Knet.dir("data","gutenberg.jl"))
_,_,chars2 = shakespeare()

shake = (shake..., chars2)

generate(shake,1000)

summary.(shake)

cpushake = (shake[1],map(Array,shake[2:end])...)

dump(cpushake[1])

cpushake[1].rnnDesc = cpushake[1].dropoutDesc = nothing

save("shakespeare.jld","model",cpushake)

foo1 = download("http://people.csail.mit.edu/deniz/models/nlp-demos/shakespeare.jld")

m1 = load(foo1,"model")

generate(m1,100)

generate(m2,100)

cpumodel

save("juliacharlm.jld","model",cpumodel)

foo2 = download("http://people.csail.mit.edu/deniz/models/nlp-demos/juliacharlm.jld")
m2 = load(foo2,"model")

generate(m2,100)

dump(m2[1])


import Base.Threads

export push_v!

function push_v!(p, kernel:: CloudInCell, m :: TwoDGrid, ex, ey, bz, dt)

    nbpart = size(p.array)[2]

    dx = m.dx
    dy = m.dy

    Threads.@threads for ipart = 1:nbpart

        v1 = p.array[3, ipart]
        v2 = p.array[4, ipart]

        xp = p.array[1, ipart] / dx
        yp = p.array[2, ipart] / dy

        i = floor(Int, xp) + 1
        j = floor(Int, yp) + 1

        dxp = xp - i + 1
        dyp = yp - j + 1

        a1 = (1 - dxp) * (1 - dyp)
        a2 = dxp * (1 - dyp)
        a3 = dxp * dyp
        a4 = (1 - dxp) * dyp

        e1 = a1 * ex[i, j] + a2 * ex[i+1, j] + a3 * ex[i+1, j+1] + a4 * ex[i, j+1]
        e2 = a1 * ey[i, j] + a2 * ey[i+1, j] + a3 * ey[i+1, j+1] + a4 * ey[i, j+1]
        b3 = a1 * bz[i, j] + a2 * bz[i+1, j] + a3 * bz[i+1, j+1] + a4 * bz[i, j+1]

        v1 += 0.5dt * e1
        v2 += 0.5dt * e2

        tantheta = 0.5dt * b3
        sintheta = 2 * tantheta / (1 + tantheta * tantheta)

        v1 += v2 * tantheta
        v2 += -v1 * sintheta
        v1 += v2 * tantheta

        p.array[3, ipart] = v1 + 0.5dt * e1
        p.array[4, ipart] = v2 + 0.5dt * e2

    end

end

export push_x!

function push_x!(p, mesh::TwoDGrid, dt::Float64)

    nbpart = size(p.array)[2]

    dimx, dimy = mesh.dimx, mesh.dimy

    Threads.@threads for i = 1:nbpart
        p1 = p.array[1, i] + dt * p.array[3, i]
        p2 = p.array[2, i] + dt * p.array[4, i]
        p1 > dimx && (p1 -= dimx)
        p2 > dimy && (p2 -= dimy)
        p1 < 0.0 && (p1 += dimx)
        p2 < 0.0 && (p2 += dimy)
        p.array[1, i] = p1
        p.array[2, i] = p2
    end

end


## Turing Complete Interpreter
### A Very Tiny Turing Complete Programming Language defined with denotational semantics

# semantica dalle dispense degano
using Metatheory

import Base.ImmutableDict
Mem = Dict{Symbol,Union{Bool,Int}}

read_mem = @theory v σ begin
	(v::Symbol, σ::Mem) => if v == :skip σ else σ[v] end
end

@testset "Reading Memory" begin
	ex = :((x), $(Mem(:x => 2)))
	@test true == areequal(read_mem, ex, 2)
end

arithm_rules = @theory a b σ begin
	(a + b, σ::Mem) 	--> (a, σ) + (b, σ)
	(a * b, σ::Mem) 	--> (a, σ) * (b, σ)
	(a - b, σ::Mem) 	--> (a, σ) - (b, σ)
	(a::Int, σ::Mem) 	--> a
	(a::Int + b::Int) 	=> a + b
	(a::Int * b::Int) 	=> a * b
	(a::Int - b::Int) 	=> a - b
end

@testset "Arithmetic" begin
	@test areequal(read_mem ∪ arithm_rules, :((2 + 3), $(Mem())), 5)
end

# don't need to access memory
bool_rules = @theory a b σ begin
	(a < b, σ::Mem) --> (a, σ) < (b, σ)
	(a ∨ b, σ::Mem) --> (a, σ) ∨ (b, σ)
	(a ∧ b, σ::Mem) --> (a, σ) ∧ (b, σ)
	(¬(a), σ::Mem) 	--> ¬((a, σ))

	(a::Bool, σ::Mem) 		=> a
	(¬a::Bool) 				=> !a
	(a::Bool ∨ b::Bool) 	=> (a || b)
	(a::Bool ∧ b::Bool) 	=> (a && b)
	(a::Int < b::Int) 		=> (a < b)
end

t = read_mem ∪ arithm_rules ∪ bool_rules

@testset "Booleans" begin
	@test areequal(t, :((false ∨ false), $(Mem())), false)

	exx = :((false ∨ false) ∨ ¬(false ∨ false), $(Mem(:x => 2)))
	g = EGraph(exx)
	saturate!(g, t)
	ex = extract!(g, astsize)
	@test ex == true
	params = SaturationParams(timeout=12)
	@test areequal(t, exx, true; params=params)

	@test areequal(t, :((2 < 3) ∧ (3 < 4), $(Mem(:x => 2))), true)
	@test areequal(t, :((2 < x) ∨ ¬(3 < 4), $(Mem(:x => 2))), false)
	@test  areequal(t, :((2 < x) ∨ ¬(3 < 4), $(Mem(:x => 4))), true)
end

if_rules = @theory guard t f σ begin
	(if guard; t end) --> (if guard; t else :skip end)
	(if guard; t else f end, σ::Mem) --> (if (guard, σ); t else f end, σ)
	(if true; t else f end, σ::Mem) --> (t, σ)
	(if false; t else f end, σ::Mem) --> (f, σ)
end

if_language = read_mem ∪ arithm_rules ∪ bool_rules ∪ if_rules

@testset "If Semantics" begin
	@test areequal(if_language, 2, :(if true x else 0 end, 	$(Mem(:x => 2))))
	@test areequal(if_language, 0, :(if false x else 0 end, 	$(Mem(:x => 2))))
	@test areequal(if_language, 2, :(if ¬(false) x else 0 end, $(Mem(:x => 2))))
	params = SaturationParams(timeout=10)
	@test areequal(if_language, 0, :(if ¬(2 < x) x else 0 end, $(Mem(:x => 3))); params=params)
end

while_rules = @theory a b σ begin
	(:skip, σ::Mem) 	--> σ
	((a; b), σ::Mem) 	--> ((a, σ); b)
	(a::Int; b) 		--> b
	(a::Bool; b) 		--> b
	(σ::Mem; b) 		--> (b, σ)
	(while a b end, σ::Mem) -->
		(if a; (b; while a b end) else :skip end, σ)
end

write_mem = @theory sym val σ begin
	(sym::Symbol = val, σ::Mem) --> (sym = (val, σ), σ)
	(sym::Symbol = val::Int, σ::Mem) => merge(σ, Dict(sym => val))
end

while_language = if_language ∪ write_mem ∪ while_rules;

@testset "While Semantics" begin
	exx = :((x = 3), $(Mem(:x => 2)))
	g = EGraph(exx)
	saturate!(g, while_language)
	ex = extract!(g, astsize)

	@test areequal(while_language, Mem(:x => 3), exx)

	exx = :((x = 4; x = x + 1), $(Mem(:x => 3)))
	g = EGraph(exx)
	saturate!(g, while_language)
	ex = extract!(g, astsize)

	params = SaturationParams(timeout=10)
	@test areequal(while_language, Mem(:x => 5), exx; params=params)

	params = SaturationParams(timeout=14)
	exx = :((if x < 10 x = x + 1 else skip end), $(Mem(:x => 3)))
	@test areequal(while_language, Mem(:x => 4), exx; params=params)

	exx = :(
		(while x < 10
			x = x + 1
		end; x), $(Mem(:x => 3)))
	g = EGraph(exx)
	params = SaturationParams(timeout=100, scheduler=Schedulers.SimpleScheduler)
	saturate!(g, while_language, params)
	@test 10 == extract!(g, astsize)
end


a = 0; b = 1
nₓ = 100
m(x) = x
σ = 1
s = 2
V(x) = 0.
left_bc = DirichletBC(0); right_bc = DirichletBC(0)
f(x) = 1.

problem = RandomODEProblem(a, b, nₓ, m, σ, s, left_bc, right_bc, f, potential = V);
assemble_system!(problem)
solve!(problem)

true_solution(x) = -x/2*(x - 1)

error = abs.(true_solution.(problem.x) - problem.u)
norm(error, Inf) < 1e-14


#= test/cur.jl
=#

  @testset "cur" begin
    m = 128
    n =  64
    M = matrixlib(:fourier, rand(m), rand(n))
    Mh = M[1:n,1:n]; Mh += Mh';
    Ms = M[1:n,1:n]; Ms += transpose(Ms);
    opts = LRAOptions(maxdet_tol=0., sketch_randn_niter=1)

    @time for (t, s) in ((:none,                 :none ),
                         (:RandomGaussian,       :randn),
                         (:RandomSubset,         :sub  ),
                         (:SRFT,                 :srft ),
                         (:SparseRandomGaussian, :sprn ))
      opts.sketch = s
      for T in (Float32, Float64, ComplexF32, ComplexF64)
        let A
            println("  $t/$T")

            rtol = 5*eps(real(T))
            approx_rtol = 1000*rtol
            opts.rtol = rtol

            for (N, ρ, q) in ((M, m, n), (Mh, n, n), (Ms, n, n))
              A = convert(Array{T}, T <: Real ? real(N) : N)

              U  = curfact(A, opts)
              F  = CUR(A, U)
              @test norm(A - Matrix(F)) < approx_rtol*norm(A)

              xp = rand(T, ρ)
              xq = rand(T, q)
              y = A  *xq; @test norm(y - F  *xq) < approx_rtol*norm(y)
              y = A' *xp; @test norm(y - F' *xp) < approx_rtol*norm(y)
              y = transpose(A)*xp; @test norm(y - transpose(F)*xp) < approx_rtol*norm(y)

              Bp = rand(T, ρ, ρ)
              Bq = rand(T, q, q)
              C = A   *Bq  ; @test norm(C - F   *Bq  ) < approx_rtol*norm(C)
              C = A   *Bq' ; @test norm(C - F   *Bq' ) < approx_rtol*norm(C)
              C = A   *transpose(Bq); @test norm(C - F   *transpose(Bq)) < approx_rtol*norm(C)
              C = A'  *Bp  ; @test norm(C - F'  *Bp  ) < approx_rtol*norm(C)
              C = A'  *Bp' ; @test norm(C - F'  *Bp' ) < approx_rtol*norm(C)
              C = transpose(A) *Bp  ; @test norm(C - transpose(F) *Bp  ) < approx_rtol*norm(C)
              C = transpose(A) *transpose(Bp); @test norm(C - transpose(F) *transpose(Bp)) < approx_rtol*norm(C)
              C = Bp  *A   ; @test norm(C - Bp  *F   ) < approx_rtol*norm(C)
              C = Bq  *A'  ; @test norm(C - Bq  *F'  ) < approx_rtol*norm(C)
              C = Bq  *transpose(A) ; @test norm(C - Bq  *transpose(F) ) < approx_rtol*norm(C)
              C = Bp' *A   ; @test norm(C - Bp' *F   ) < approx_rtol*norm(C)
              C = Bq' *A'  ; @test norm(C - Bq' *F'  ) < approx_rtol*norm(C)
              C = transpose(Bp)*A   ; @test norm(C - transpose(Bp)*F   ) < approx_rtol*norm(C)
              C = transpose(Bq)*transpose(A) ; @test norm(C - transpose(Bq)*transpose(F) ) < approx_rtol*norm(C)
          end
        end
      end
    end
  end


# by Lazaro Alonso
using GeoMakie, Proj4, CairoMakie
CairoMakie.activate!()
let
    lons = -180:180
    lats = -90:90
    field = [exp(cosd(l)) + 3(y/90) for l in lons, y in lats]

    trans = Proj4.Transformation("+proj=longlat +datum=WGS84", "+proj=wintri",
        always_xy=true)
    # see https://proj.org/operations/projections/index.html for more options
    ptrans = Makie.PointTrans{2}(trans)

    fig = Figure(resolution = (1250,700), fontsize = 22)
    ax = Axis(fig[1,1], aspect = DataAspect(),
        title = "Projection: Winkel Tripel, Field")
    # all input data coordinates are projected using this function
    ax.scene.transformation.transform_func[] = ptrans
    # add some limits, still it needs to be manual

    points = [Point2f0(lon, lat) for lon in lons, lat in lats]
    rectLimits = FRect2D(Makie.apply_transform(ptrans, points))
    limits!(ax, rectLimits)
    # now the plot
    hm = surface!(ax, lons, lats, field, shading = false)
    Colorbar(fig[1,2], hm, label = "variable, color code", height = Relative(0.65))
    lines!(ax, GeoMakie.coastlines(), color = :black, overdraw = true)

    hidedecorations!(ax, ticks = false, label = false, grid=false)
    hidespines!(ax)
    fig
    save(joinpath(@__DIR__, "output", "geoCoastlinesField.png"), fig) # HIDE
end
using Pkg # HIDE
Pkg.status(["GeoMakie", "Proj4", "CairoMakie"]) # HIDE


using Revise
using AdFem

mesh = Mesh3(10,10,10,0.1)
nu = constant(rand(get_ngauss(mesh)))
u = constant(rand(mesh.ndof))

f = compute_fem_laplace_term1(u, nu, mesh)

sess = Session(); init(sess)
run(sess, f)


function solve_007()
    # init sieve vector
    numbers = Set(2:150000)
    primes = Set()

    # run sieve
    for candidate = 2:150000
        for num in numbers
            if num%candidate == 0 && num ≠ candidate delete!(numbers, num) end
        end
        if candidate % 100 == 0 println(candidate) end
    end

    return numbers
end


#*******************************************************/
#* Copyright(c) 2018 by Artelys                        */
#* This source code is subject to the terms of the     */
#* MIT Expat License (see LICENSE.md)                  */
#*******************************************************/

#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
#  Knitro example defining a single callback function for function
#  and gradient(i.e. EVALFCGA) requests, and a second callback for
#  Hessian(i.e. EVALH / EVALHV) requests.  This model is test
#  problem HS40 from the Hock & Schittkowski collection.
#
#  max   x0*x1*x2*x3        (obj)
#  s.t.  x0^3 + x1^2 = 1    (c0)
#        x0^2*x3 - x2 = 0   (c1)
#        x3^2 - x1 = 0      (c2)
#
#  See exampleNLP2.c to see how this is solved in the standard
#  way with separate callbacks for functions and gradients.
#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

using KNITRO
using Test

function example_fcga(; verbose=true)
    #*------------------------------------------------------------------*/
    #*     FUNCTION callbackEvalFCGA                                    */
    #*------------------------------------------------------------------*/
    # The signature of this function matches KNITRO.KN_eval_callback in knitro.h.
    # To compute the functions and gradient together, set "obj", "c",
    # "objGrad" and "jac" in the KNITRO.KN_eval_result structure when there
    # is a function+gradient evaluation request(i.e. EVALFCGA).
    #
    # NOTE: It is generally more efficient and recommended to have
    # separate callback routines for functions and gradients since a
    # gradient evaluation is not always needed for every function
    # evaluation.  However, in some cases, it may be more convenient
    # to compute them together if most of the work for computing the
    # gradients is already done for the function evaluation.
    function callbackEvalFCGA(kc, cb, evalRequest, evalResult, userParams)
        x = evalRequest.x

        # Evaluate nonlinear term in objective
        evalResult.obj[1] = x[1] * x[2] * x[3] * x[4]

        # Evaluate nonlinear terms in constraints
        evalResult.c[1] = x[1] * x[1] * x[1]
        evalResult.c[2] = x[1] * x[1] * x[4]

        # Evaluate nonlinear term in objective gradient
        evalResult.objGrad[1] = x[2] * x[3] * x[4]
        evalResult.objGrad[2] = x[1] * x[3] * x[4]
        evalResult.objGrad[3] = x[1] * x[2] * x[4]
        evalResult.objGrad[4] = x[1] * x[2] * x[3]

        # Evaluate nonlinear terms in constraint gradients(Jacobian)
        evalResult.jac[1] = 3.0 * x[1] * x[1] # derivative of x0^3 term  wrt x0
        evalResult.jac[2] = 2.0 * x[1] * x[4] # derivative of x0^2*x3 term  wrt x0
        evalResult.jac[3] = x[1] * x[1]     # derivative of x0^2*x3 terms wrt x3

        return 0
    end

    #*------------------------------------------------------------------*/
    #*     FUNCTION callbackEvalH                                       */
    #*------------------------------------------------------------------*/
    # The signature of this function matches KNITRO.KN_eval_callback in knitro.h.
    # Only "hess" and "hessVec" are set in the KNITRO.KN_eval_result structure.
    function callbackEvalH(kc, cb, evalRequest, evalResult, userParams)
        x = evalRequest.x
        lambda_ = evalRequest.lambda
        # Scale objective component of Hessian by sigma
        sigma = evalRequest.sigma

        # Evaluate nonlinear terms in the Hessian of the Lagrangian.
        # Note: If sigma=0, some computations can be avoided.
        if sigma > 0.0  # Evaluate the full Hessian of the Lagrangian
            evalResult.hess[1] = lambda_[1] * 6.0 * x[1] + lambda_[2] * 2.0 * x[4]
            evalResult.hess[2] = sigma * x[3] * x[4]
            evalResult.hess[3] = sigma * x[2] * x[4]
            evalResult.hess[4] = sigma * x[2] * x[3] + lambda_[2] * 2.0 * x[1]
            evalResult.hess[5] = sigma * x[1] * x[4]
            evalResult.hess[6] = sigma * x[1] * x[3]
            evalResult.hess[7] = sigma * x[1] * x[2]
        else            # sigma=0, do not include objective component
            evalResult.hess[1] = lambda_[1] * 6.0 * x[1] + lambda_[2] * 2.0 * x[4]
            evalResult.hess[2] = 0.0
            evalResult.hess[3] = 0.0
            evalResult.hess[4] = lambda_[2] * 2.0 * x[1]
            evalResult.hess[5] = 0.0
            evalResult.hess[6] = 0.0
            evalResult.hess[7] = 0.0
        end

        return 0
    end

    #*------------------------------------------------------------------*
    #*     main                                                         *
    #*------------------------------------------------------------------*

    # Create a new Knitro solver instance.
    kc = KNITRO.KN_new()

    # Initialize Knitro with the problem definition.

    # Add the variables and specify initial values for them.
    # Note: any unset lower bounds are assumed to be
    # unbounded below and any unset upper bounds are
    # assumed to be unbounded above.
    vars = KNITRO.KN_add_vars(kc, 4)
    for x in vars
        KNITRO.KN_set_var_primal_init_value(kc, x, 0.8)
    end

    # Add the constraints and set the rhs and coefficients
    KNITRO.KN_add_cons(kc, 3)
    KNITRO.KN_set_con_eqbnds_all(kc, [1.0, 0.0, 0.0])

    # Coefficients for 2 linear terms
    lconIndexCons = Int32[1, 2]
    lconIndexVars = Int32[2, 1]
    lconCoefs = [-1.0, -1.0]
    KNITRO.KN_add_con_linear_struct(kc, lconIndexCons, lconIndexVars, lconCoefs)

    # Coefficients for 2 quadratic terms

    qconIndexCons = Int32[0, 2]
    qconIndexVars1 = Int32[1, 3]
    qconIndexVars2 = Int32[1, 3]
    qconCoefs = [1.0, 1.0]

    KNITRO.KN_add_con_quadratic_struct(
        kc,
        qconIndexCons,
        qconIndexVars1,
        qconIndexVars2,
        qconCoefs,
    )

    # Add callback to evaluate nonlinear(non-quadratic) terms in the model:
    #    x0*x1*x2*x3  in the objective
    #    x0^3         in first constraint c0
    #    x0^2*x3      in second constraint c1
    cIndices = Int32[0, 1]
    cb = KNITRO.KN_add_eval_callback(kc, true, cIndices, callbackEvalFCGA)

    # Set obj. gradient and nonlinear jac provided through callbacks.
    # Mark objective gradient as dense, and provide non-zero sparsity
    # structure for constraint Jacobian terms. Set the gradient
    # evaluations to use the same callback routine as used for
    # function evaluations.
    cbjacIndexCons = Int32[0, 1, 1]
    cbjacIndexVars = Int32[0, 0, 3]
    KNITRO.KN_set_cb_grad(
        kc,
        cb,
        nothing,
        jacIndexCons=cbjacIndexCons,
        jacIndexVars=cbjacIndexVars,
    )

    # Set nonlinear Hessian provided through callbacks. Since the
    # Hessian is symmetric, only the upper triangle is provided.
    # The upper triangular Hessian for nonlinear callback structure is:
    #
    #  lambda0*6*x0 + lambda1*2*x3     x2*x3    x1*x3    x1*x2 + lambda1*2*x0
    #               0                    0      x0*x3         x0*x2
    #                                             0           x0*x1
    #                                                          0
    # (7 nonzero elements)
    cbhessIndexVars1 = Int32[0, 0, 0, 0, 1, 1, 2]
    cbhessIndexVars2 = Int32[0, 1, 2, 3, 2, 3, 3]
    KNITRO.KN_set_cb_hess(
        kc,
        cb,
        length(cbhessIndexVars1),
        callbackEvalH,
        hessIndexVars1=cbhessIndexVars1,
        hessIndexVars2=cbhessIndexVars2,
    )

    # Set minimize or maximize(if not set, assumed minimize)
    KNITRO.KN_set_obj_goal(kc, KNITRO.KN_OBJGOAL_MAXIMIZE)

    # Set option to print output after every iteration.
    kn_outlev = verbose ? KNITRO.KN_OUTLEV_ITER : KNITRO.KN_OUTLEV_NONE
    KNITRO.KN_set_param(kc, KNITRO.KN_PARAM_OUTLEV, kn_outlev)

    # Set option to tell Knitro that the gradients are being provided
    # with the functions in one callback.
    KNITRO.KN_set_param(kc, KNITRO.KN_PARAM_EVAL_FCGA, KNITRO.KN_EVAL_FCGA_YES)

    # Solve the problem.
    #
    # Return status codes are defined in "knitro.h" and described
    # in the Knitro manual.
    nStatus = KNITRO.KN_solve(kc)
    nStatus, objSol, x, lambda_ = KNITRO.KN_get_solution(kc)

    # An example of obtaining solution information.
    if verbose
        println("Knitro converged with final status = ", nStatus)
        println("  optimal objective value  = ", objSol)
        println("  optimal primal values x  = ", x)
        println("  feasibility violation    = ", KNITRO.KN_get_abs_feas_error(kc))
        println("  KKT optimality violation = ", KNITRO.KN_get_abs_opt_error(kc))
    end

    # Delete the Knitro solver instance.
    KNITRO.KN_free(kc)

    @testset "Example HS40 FCGA" begin
        @test nStatus == 0
        @test objSol ≈ 0.25
        @test x ≈ [0.793701, 0.707107, 0.529732, 0.840896] atol = 1e-5
    end
end

example_fcga(; verbose=isdefined(Main, :KN_VERBOSE) ? KN_VERBOSE : true)


# generate a tridiagonal matrix
function tridiagm(a, b, c, n)
    dd, du = ones(n), ones(n - 1)
    b*diagm(0 => dd) + a*diagm(-1 => du) + c*diagm(1 => du)
end

# support function
SUITE["ρ"] = BenchmarkGroup()

# support vector
SUITE["σ"] = BenchmarkGroup()

N = Float64
for set_type in (Ball1, BallInf, Hyperrectangle)
    for n in (2, 10, 100, 1000)
        rng = MersenneTwister(n)
        d = rand(rng, n)
        X = rand(set_type, dim=n, rng=rng)
        SUITE["ρ"][string(set_type), "dense", n] = @benchmarkable ρ($d, $X)
        SUITE["σ"][string(set_type), "dense", n] = @benchmarkable σ($d, $X)
        A = tridiagm(1, -2, 0.5, n)
        Y = A * X
        SUITE["ρ"]["LinearMap of $(string(set_type))", "dense", n] = @benchmarkable ρ($d, $Y)
        SUITE["σ"]["LinearMap of $(string(set_type))", "dense", n] = @benchmarkable σ($d, $Y)
    end
end


function absorption_line(x::T; mid=zero(T), width=one(T), depth=one(T)) where T<:AF
    return one(T) - depth * exp(-((x-mid)/width)^2.0/2.0)
end

function gaussian_line(x::T; mid=zero(T), width=one(T), depth=one(T)) where T<:AF
    return absorption_line(x, mid=mid, width=width, depth=depth)
end

function lorentzian_line(x::T; mid=zero(T), width=one(T), depth=one(T)) where T<:AF
    y = (x - mid)/(width/2)
    return one(T) - (depth / (one(T) + y^2))
end

# calculate doppler factor
function doppler_factor(vel::T) where T<:AF
    num = one(T) + vel/c_ms
    den = one(T) - vel/c_ms
    return sqrt(num/den)
end

# Width of line due to thermal broadening
function width_thermal(; λ::T1=1.0, M::T1=1.0, T::T1=5778.0, v_turb::T1=0.0) where T1<:AF
    return sqrt((2.0*kB/mH) * (T/M) + v_turb^2) * (λ/c)
end

# Quadratic limb darkening law.
# Takes μ = cos(heliocentric angle) and LD parameters, u1 and u2.
function quad_limb_darkening(μ::T, u1::T, u2::T) where T<:AF
    return !iszero(μ) * (one(T) - u1*(one(T)-μ) - u2*(one(T)-μ)^2)
end

function quad_limb_darkening(x::T, y::T, u1::T, u2::T) where T<:AF
    return quad_limb_darkening(calc_mu(x,y), u1, u2)
end

"""
    rotation_period(sin_lat; A=14.713, B=2.396, C=1.787)

Calculate stellar rotation period in days at given sine of latitude.

# Arguments
- `sin_lat::Float64`: Sine of stellar latitude.
- `A::Float64=14.713`: Coefficient from Snodgrass & Ulrich (1990)
- `B::Float64=2.396`: Coefficient from Snodgrass & Ulrich (1990)
- `C::Float64=1.787`: Coefficient from Snodgrass & Ulrich (1990)
"""
function rotation_period(sin_lat::T; A::T=14.713, B::T=-2.396, C::T=-1.787) where T<:AF
    @assert -1.0 <= sin_lat <= 1.0
    return 360.0/(A + B * sin_lat^2 + C * sin_lat^4)
end

"""
    patch_velocity_los(x, y; rstar, pole)

Compute line of sight velocity of a patch of stellar surface given by x,y (assumed in [-1,1]).
Return value is in (Rsol/day)/speed of light (i.e., dimensionless like z = v/c)

# Arguments
- `rstar::Float64=1.0`: in R_sol (affects return velocity, but not x,y)
- `pole = (0,1,0)`: unit vector for stellar rotation axis (default is equator-on)
"""
function patch_velocity_los(x::T,y::T; rstar=one(T), pole=(zero(T), one(T), zero(T))) where T<:AF
    polex, poley, polez = pole
    v0 = 0.000168710673 # in (Rsol/day)/speed_of_light
    z = sqrt(rstar - calc_r2(x,y))
    sin_lat = (x*polex) + (y*poley) + (z*polez)
    vmax = v0*rstar/rotation_period(sin_lat)
    return vmax * (polex*y - poley*x)
end

function patch_velocity_los(t::Tuple{T,T}; rstar=one(T), pole=(zero(T), one(T), zero(T))) where T<:AF
    return patch_velocity_los(t[1], t[2], rstar=rstar, pole=pole)
end


using Dispersal, Test

@testset "test mating population" begin
    init =  [10.0 10.0 10.0;
             10.0 10.0 10.0;
             10.0 10.0 10.0]

    output = ArrayOutput(init; tspan=1:3)
    rule = Ruleset(MatingPopulation(Moore(1), 1.0, 0.0))
    sim!(output, rule)
    @test output[1] == init
    @test output[2] == init
    @test output[3] == init

    output = ArrayOutput(init; tspan=1:3)
    rule = Ruleset(MatingPopulation(Moore(1), 0.0, 1.0))
    sim!(output, rule)
    @test output[1] == init
    @test output[2] == [30/8 50/8 30/8;
                  50/8 80/8 50/8;
                  30/8 50/8 30/8]
    @test output[3] ==  [180/64  240/64 180/64;
                   240/64  320/64 240/64;
                   180/64  240/64 180/64]

    output = ArrayOutput(init; tspan=1:3)
    rule = Ruleset(MatingDispersal(Moore(1), 0.5, 0.5))
    sim!(output, rule)
    @test output[1] == init
    @test output[2] == 0.5.*init + 0.5.* [30/8 50/8 30/8;
                                          50/8 80/8 50/8;
                                          30/8 50/8 30/8]
    @test output[3] == 0.5.*output[2] + 0.5.*[ 26.25/8  40/8 26.25/8;
                                               40/8  60.0/8 40/8;
                                               26.25/8  40/8 26.25/8]

end

using Dispersal
@testset "test mating dispersal population" begin

init =  [0.0 0.0 0.0;
         0.0 1.0 1.0;
         0.0 0.0 0.0]

output = ArrayOutput(init; tspan=1:3)
rule = Ruleset(MatingDispersal(-1,-1))
sim!(output, rule)
output[2]
output[3]
end


function holtslag_psis(xi)
    #println("xi")
    ##println(xi)
    #println("-------------")
if (xi>=0) & (xi<=1.)
    y=1+5*xi;
elseif xi>=1
    y=5+xi
else xi<0
    y=(1-15*xi)^(-1/2)
end
return y
end


# # Tutorial 01 Rijke Tube

# ## Introduction
#
# A Rijke tube is the most simple thermo-acoustic configuration. It comprises a
# tube with an unsteady heat source somewhere inside the tube. This example will
# walk you through the basic steps of setting up a Rijke tube in a
# Helmholtz-solver stability analysis.

# ### The Helmholtz equation
# The Helmholtz equation is the Fourier transform of the acoustic wave equation.
# Given that there is a heat source, it reads:
#
# ∇⋅(c² ∇ p̂) + ω² p̂ = -iω(γ-1)q̂
#
# Here c is speed-of-sound-field, p̂ the (complex) pressure fluctuation amplitude
# ω the (complex) frequency of the problem, i the imaginary unit, γ the ratio of
# specifc heats, and q̂ the amplitude of the heat release fluctuation.
#
# (Note that the Fourier transform here follows a f'(t) --> f̂(ω)exp(+iωt) convention.)
#
# Together with some boundary conditions the Helmholtz equation models
# thermo-acoustic resonance in a cavity. The minimum required inputs to specify
# a problem are
#
# 1. the shape of the domain
# 2. the speed-of-sound field
# 3. the boundary conditions
#
# In case of an active flame (q̂≠0). We will also need
#
# 4. some additional gas porperties and
# 5. a flame response model linking the heat release fluctuations q̂ to the pressure fluctuations p̂
#
# Once you completed this tutorial you will know the basic steps of how to
# conduct a thermo-acoustic stability analysis

## #jl

# ## Header
# First you will need to load the Helmholtz solver. The following line brings all
# necessary tools into scope:
using WavesAndEigenvalues.Helmholtz

## #jl
# ## Mesh
# Now we can load the mesh file. It is the essential piece of information
# defining the domain shape . This tutorial's mesh has been specified in mm
# which is why we scale it by `0.001` to load the coordinates as m:
mesh=Mesh("Rijke_mm.msh",scale=0.001)

# You can have a look at some basic mesh data by printing it
print(mesh)
# In an interactive session, is suffices to just type the mesh's variable name
# without the enclosing `print` command.
#
#
# This info tells us that the mesh features `1006` points as vertices
# `1562` (addressable) triangles on its surface and `3380`tetrahedra forming
# the tube. Note, that there are currently no lines stored. This is because
# line information is only needed when using second-order finite elements. To
# save memory this information is only assembled and stored when explicitly
# requested. We will come back to this aspect in a later tutorial.
#
# You can also see that the mesh features several named domains such as
# `"Interior"`,`"Flame"`, `"Inlet"` and `"Outlet"`. These domains are used to
# specify certain conditions for our model.
#
# A model descriptor is always given as a dictionairy featureing domain names
# as keys and tuples as values. The first tuple entry is a Julia symbol
# specifying the operator to be build on the associated domain, while the second
# entry is again a tuple holding information that is specific to the chosen
# operator.

## #src
# ## Model set-up
# For instance, we certainly want to build the wave operator on the entire mesh.
# So first, we initialize an empty dictionairy
dscrp=Dict()

# And then specify that the wave operator should be build everywhere. For the
# given mesh the domain-specifier to address the entire resononant cavity is
# `"Interior"` and in general the symbol to create the wave operator is
# `:interior`. It requires no options so the options tuple remains empty.
dscrp["Interior"]=(:interior, ())

## src
# ## Boundary Conditions
# The tube should have a pressure node at its outlet, in order to model an
# open-end boundary condition. Boundary conditions are specified in terms of
# admittance values. A pressure note corresponds to an infinite admittance.
# To represent this on a computer we just give it a crazily high value like
# `1E15`. We will also need to specify a variable name for our admmittance value.
# This will allow to quickly change the value after discretization of the
# model by addressing it by this very name. This feature is one of the core
# concepts of the solver. As will be demonstrated later.
#
# The complete description of our boundary condition reads
dscrp["Outlet"]=(:admittance, (:Y,1E15))

# You may wonder why we do not specify conditions for the other boundaries of
# the model. This is because the discretization is based on Bubnov-Galerkin
# finite elements. All unspecified, boundaries will therefore *naturally* be
# discretized as solid walls.

## #src
# ## Flame
# The Rijke tube's main feauture is a domain of heat release. In the current
# mesh there is a designated domain `"Flame"` addressing a thin volume at the
# center of the tube. We can use this key to specify a flame with simple
# n-tau-dynamics. Therefore, we first need to define some basic gas properties.

γ=1.4 #ratio of specific heats
ρ=1.225 #density at the reference location upstream to the flame in kg/m^3
Tu=300.0    #K unburnt gas temperature
Tb=1200.0    #K burnt gas temperature
P0=101325.0 # ambient pressure in Pa
A=pi*0.025^2 # cross sectional area of the tube
Q02U0=P0*(Tb/Tu-1)*A*γ/(γ-1) #the ratio of mean heat release to mean velocity Q02U0

# We will also need to provide the position where the reference velocity has
# been taken and the direction of that velocity
x_ref=[0.0; 0.0; -0.00101]
n_ref=[0.0; 0.0; 1.00] #must be a unit vector
# And of course, we need some values for n and tau
n=0.0 #interaction index
τ=0.001 #time delay
# With these values the specification of the flame reads
dscrp["Flame"]=(:flame,(γ,ρ,Q02U0,x_ref,n_ref,:n,:τ,n,τ))
# Note that the order of the values in the options tuple is important. Also note
# that we assign the symbols `:n`and `:τ` for later analysis.

## #jl
# ## Speed of Sound
# The description of the Rijke tube model is nearly complete. We just need to
# specify the speed of sound field.  For this example, the field is fairly
# simple and can be specified analytically, using the `generate_field` function
# and a custom three-parameter function.
R=287.05 # J/(kg*K) specific gas constant (air)
function speedofsound(x,y,z)
    if z<0.
        return sqrt(γ*R*Tu)#m/s
    else
        return sqrt(γ*R*Tb)#m/s
    end
end
c=generate_field(mesh,speedofsound)
# Note that in more elaborate models, you may read the field `c` from a file
# containing simulation or experimental data, rather than specifying it
# analytically.

## #src
# ## Model Discretization
# Now we have all ingredients together and we can discretize the model.
L=discretize(mesh,dscrp,c)
# The return value `L` here is a family of linear operators. You can display
# an algebraic representation of the family plus a list of associated parameters
# by just printing it
print(L)
# (In an interactive session the enclosing print function is not necessary.)
#
# You might notice that the list contains all parameters that have been
# specified during the model description (`n`,`τ`,`Y`). However, there are also
# two parameters that were added by default: `ω` and `λ`. `ω` is the complex
# frequency of the model and `λ` an auxiliary value that is important for the
# eigenfrequency computation.
## #src
# ## Solving the Model
# ### Global Solver
# We can solve the model for some eigenvalues using one of the eigenvalue
# solvers. The package provides you with two types of eigenvalue solvers. A global
# contour-integration-based solver. That finds you all eigenvalues inside of a
# specified contour Γ in the complex plane.
Γ=[150.0+5.0im, 150.0-5.0im, 1000.0-5.0im, 1000.0+5.0im].*2*pi #corner points for the contour (in this case a rectangle)
Ω, P = beyn(L,Γ,l=10,N=64, output=true)

# The found eigenvalues are stored in the array `Ω`. The corresponding
# eigenvectors are the columns of `P`.
# The huge advantage of the global eigenvalue solver is that it finds you
# multiple eigenvalues. Nonetheless, its accuracy may be low and sometimes it
# provides you with outright wrong solutions.
# However, for the current case the method works as we can varify that in our
# search window there are two eigenmodes oscilating at 272 and 695 Hz
# respectively:
for ω in Ω
    println("Frequency: $(ω/2/pi)")
end
## #src
# ### Local Solver
# To get high accuracy eigenvalues , there is also local iterative eigenvalue
# solver. Based on an initial guess, it only finds you one eigenvalue at a time
# but with machine-precise accuracy.
sol,nn,flag=householder(L,250*2*pi,output=true);

# The return values of this solver are a solution object `sol`, the number of
# iterations `nn` performed by the solver and an error flag `flag` providing
# information on the quality of the solution. If `flag>0` the solution has
# converged. If `flag<0`something went wrong. And if `flag==0`... well, probably
# the solution is fine but you should check it.
#
# In the current case the flag is -1 indicating that the maximum number of
# iterations has been reached. This is because we haven't specified a stopping
# criterion and the iteration just ran for the maximum number of iterations.
# Nevertheless, the 272 Hz mode is correct. Indeed, as you can see from the
# printed output it already converged to machine-precision after 5 iterations.

## #src
# ## Changing a specified parameter.
#
# Remember that we have specified the parameters `Y`,`n`, and `τ`?. We can change
# these easily without rediscretizing our model. For instance currently `n==0.0`
# so the solution is purely accoustic. The effect of the flame just enters the
# problem via the speed of sound field (this is known as passive flame). By
# setting the interaction index to `1.0` we activate the flame.
L.params[:n]=1
# Now we can just solve the model again to get the active solutions
sol_actv,nn,flag=householder(L,245*2*pi-82im*2*pi,output=true,order=3);
# The eigenfrequency is now complex:
sol_actv.params[:ω]/2/pi
# with a growth rate `-imag(sol_actv.params[:ω]/2/pi)≈  59.22`

# Instead of recomputing the eigenvalue by one of the solvers. We can also
# approximate the eigenvalue as a function of one of the model parameters
# utilizing high-order perturbation theory. For instance this gives you
# the 30th order diagonal Padé estimate expanded from the passive solution.

perturb_fast!(sol,L,:n,16) #compute the coefficients
freq(n)=sol(:n,n,8,8)/2/pi #create some function for convenience
freq(1) #evaluate the function at any value for `n` you are interested in.

# The method is slow when only computing one eigenvalue. But its computational
# costs are almost constant when computing multiple eigenvalues. Therefore,
# for larger parametric studies this is clearly the method of choice.

## #src
# ## VTK Output for Paraview
#
# To visualizte your solutions you can store them as `"*.vtu"` files to open
# them with paraview. Just, create a dictionairy, where your modes are stored as
# fields.
data=Dict()
data["speed_of_sound"]=c
data["abs"]=abs.(sol_actv.v)/maximum(abs.(sol_actv.v)) #normalize so that max=1
data["phase"]=angle.(sol_actv.v)

vtk_write("tutorial_01", mesh, data) # Write the solution to paraview

# The `vtk_write` function automatically sorts your data by its interpolation
# scheme. Data that is constant on a tetrahedron will be written to a file
# `*_const.vtu`, data that is linearly interpolated on a tetrahedron will go
# to `*_lin.vtu`, and data that is quadratically interpolated to `*_quad.vtu`.
# The current example uses constant speed of sound on a tetrahedron and linear
# finite elements for the discretization of p. Therefore, two files are
# generated, namely `tutorial_01_const.vtu` containing the speed-of-sound field
# and `tutorial_01_lin.vtu` containing the mode shape. Open them with paraview
# and have a look!.

## src
# ## Summary
#
# You learnt how to set-up a simple Rijke-tube study. This  already introduced
# all basic steps that are needed to work with the Helmholtz solver. However,
# there are a lot of details you can fine tune and even features that weren't
# mentioned in this tutorial. The next tutorials will introduce these aspects in
# more detail.


"""
# Usage

    solve(prob::FDDEMatrixProblem, h, MatrixForm())

### Reference

https://github.com/mandresik/system-of-linear-fractional-differential-delayed-equations
"""
struct MatrixForm <: FractionalDiffEqAlgorithm end

function solve(prob::FDDEMatrixProblem, h, ::MatrixForm)
    @unpack α, τ, A, B, f, x0, tspan = prob
    t0 = tspan[1]; T = tspan[2]
    limit = 100
    var_num = length(A[:, 1])
    m = ceil(Int, α)

    t_tau = t0-τ
    temp = collect(t0-h:-h:t_tau)
    t = [temp[end:-1:1]; collect(t0:h:T)]

    index = length(temp)
    N = length(t)-index-1

    x = zeros(length(t), var_num)
    b = zeros(1, N)
    F = zeros(var_num, N)
    x0TP = x0(t0)

    @fastmath @inbounds @simd for i=1:index+1
        x[i, :] = x0(t[i])*[1; zeros(m-1, 1)]
    end

    function function_values(M, limit)
        cols::Int64 = size(M, 2)
        risk_index = zeros(1, var_num)
        col = 0
        M0 = copy(M)

        @fastmath @inbounds @simd for r=1:var_num
            @fastmath @inbounds @simd for s=1:cols
                if M0[r, s]==Inf || M0[r, s]==-Inf
                    risk_index = hcat(risk_index[:, 1:col], [r; s])
                    col += 1
                end
            end
        end
        Mt = zeros(var_num, cols*length(t))

        for j=1:length(t)
            Mt[:, (cols*j-cols+1):(cols*j)] = M
        end

        if risk_index !== zeros(var_num, 1)
            risk_start_index = judgeeqsum(t, -1)
            inv_num = judgesum(t, 1)-judgeeqsum(t, -1)
            risks_num = size(risk_index, 2)
            @fastmath @inbounds @simd for j=1:inv_num
                @fastmath @inbounds @simd for k=1:risks_num
                    id = risk_start_index*cols + risk_index[2, k] + cols*(j-1)
                    if Mt[risk_index[1, k], id] > limit
                        Mt[risk_index[1, k], id] = limit
                    elseif Mt[risk_index[1, k], id] < -limit
                        M[risk_index[1, k], id] = -limit
                    end
                end
            end
        end
        return Mt
    end

    At = function_values(A, limit)
    Bt = function_values(B, limit)
    ft = function_values(f, limit)

    @fastmath @inbounds @simd for n=1:N
        col_num = size(A, 2)
        ind = col_num*(n+index)
        b[n] = (n+1)^α-n^α
        F[:, n] = At[:, ind-col_num+1:ind]*x[n+index, :] + Bt[:, ind-col_num+1:ind]*x[n, :]+ft[:, n+index]
        TP = zeros(var_num, 1)
        @fastmath @inbounds @simd for k=0:m-1
            TP += ((t[n+index+1]-t0)^k).*x0TP[:, k+1]./factorial(k)
        end
        sum = zeros(var_num, 1)
        @fastmath @inbounds @simd for j=1:n
            sum += b[n-j+1].*F[:, j]
        end
        x[n+index+1, :] = TP + h^α.*sum
    end
    return x
end

function judgesum(t, thres)
    result = 0
    for i in t
        i < thres ? result+=i : continue
    end
    return result
end

function judgeeqsum(t, thres)
    result = 0
    for i in t
        i <= thres ? result+=i : continue
    end
    return result
end


#   Author: V. Vitale
#   Feb 2022
#   1 site and 2 sites DMRG

include("mps.jl")
include("mpo.jl")
include("contractions.jl")
using TensorOperations
using KrylovKit
using LinearAlgebra

function one_site_dmrg!(psi::MPS,
                        W::MPO,
                        sweeps::Int64)

    right_normalize!(psi)

    d = dims(psi)[2][2]
    chi = dims(psi)[2][1]
    L = construct_L(psi, W)
    R = construct_R(psi, W)

    #global psi,L,R
    Energy=0
    for sweep in 1:Int(sweeps/2)
        for i in 1:psi.N-1
            Energy,psi.data[i],psi.data[i+1] = optimize_one_site(psi.data[i], psi.data[i+1], W.data[i], L[i], R[i], "right", chi)
            L[i+1] = contract_from_left(L[i], psi.data[i], W.data[i])
        end

        for i in psi.N:-1:2
            Energy,psi.data[i-1],psi.data[i] = optimize_one_site( psi.data[i-1],psi.data[i], W.data[i], L[i], R[i], "left", chi)
            R[i-1] = contract_from_right(R[i], psi.data[i], W.data[i])
        end
    end
    println("Done! Energy= ",real(Energy),"; Variance: ",real(psi*(W*(W*psi))-(psi*(W*psi))^2))
end

function optimize_one_site(A::Array, B::Array, W::Array, E::Array, F::Array, dir::String, chi::Int )
    function H_lin(v)
        ##     +--A--+
        ##     |  |  |
        ##     L--W--R
        ##     |  |  |
        #println("E ",size(E))
        #println("MPS ",size(v))
        #println("F ",size(F))
        #println("W ",size(W))
        @tensor temp_1[ :] := E[ 1, -2, -3] *  v[ 1, -1, -4]
        @tensor temp_2[:] :=  temp_1[-1,-2,-3,1] * F[1,-4,-5]
        @tensor temp[:] := temp_2[1, 2, -1, 3,-3] * W[ 2 3 ; 1 -2 ]
        return temp
    end

    if (dir == "right" )
        sA=size(A)

        en,V,info = eigsolve( H_lin , A ,  1  , :SR ; ishermitian = true  )
        en = en[1]
        V = V[1]
        sV = size(V)

        V=reshape(V,(sV[1]*sV[2],sV[3]))#s i j -> i*s, j
        A,S,V = svd(V,full=false)
        V=V'
        A = reshape( A , ( sA[1], sA[2], :) )
        #"ij,jl,slk->sik"
        S=diagm(S)
        @tensor B[:] := S[-1,1 ] * V[ 1,2 ] * B[2,-2,-3]
    else
        sB=size(B)
        en,V,info = eigsolve( H_lin , B ,  1  , :SR ; ishermitian = true  )
        en = en[1]
        V = V[1]

        sV=size(V)
        V=reshape(V,(sV[1],sV[2]*sV[3]))  #s i j -> i s j -> i, s*j

        U,S,B = svd(V,full=false)
        B=B'
        B = reshape(B,(:,sB[2],sB[3]))
        #"sij,jk,kl->sil"
        S=diagm(S)
        @tensor A[:] := A[-1,-2,2] * U[ 2,3 ] * S[ 3,-3 ]
    end

    return en, A, B

end

function two_sites_dmrg!(psi::MPS,
                        W::MPO,
                        sweeps::Int)

    right_normalize!(psi)

    d = dims(psi)[2][2]
    L = construct_L(psi, W)
    R = construct_R(psi, W)

    #global psi,L,R
    Energy=0
    for sweep in 1:Int(sweeps)
        println("Sweep: ",sweep)
        for i in 1:psi.N-1
            Energy,psi.data[i],psi.data[i+1] = two_sites_swipe_right( psi.data[i],
                                                                        psi.data[i+1],
                                                                        W.data[i],
                                                                        W.data[i+1],
                                                                        L[i],
                                                                        R[i+1])
            L[i+1] = contract_from_left(L[i], psi.data[i], W.data[i])
        end

        for i in psi.N:-1:2
            Energy,psi.data[i-1],psi.data[i] = two_sites_swipe_left(  psi.data[i-1],
                                                                        psi.data[i],
                                                                        W.data[i-1],
                                                                        W.data[i],
                                                                        L[i-1],
                                                                        R[i])
            R[i-1] = contract_from_right(R[i], psi.data[i], W.data[i])
        end

    end
    println("Done! Energy= ",real(Energy),"; Variance: ",real(psi*(W*(W*psi))-(psi*(W*psi))^2))
end

function two_sites_swipe_right(AL::Array, AR::Array, WL::Array, WR::Array, E::Array, F::Array)
    tol=1e-15

    sAL = size(AL)
    sAR = size(AR)
    @tensor A[:] := AL[-1,-2,1]*AR[1,-3,-4]
    A = reshape(A,(sAL[1],sAL[2]*sAR[2],sAR[3]))
    sWL = size(WL)
    sWR = size(WR)

    @tensor M[:] := WL[-1,1,-3,-5]*WR[1,-2,-4,-6]
    M = reshape(M,(sWL[1],sWR[2],sWL[3]*sWR[3],sWL[4]*sWR[4]))
    function H_lin(v)
        @tensor temp_1[ :] := E[ 1, -2, -3] *  v[ 1, -1, -4]
        @tensor temp_2[:] :=  temp_1[-1,-2,-3,1] * F[1,-4,-5]
        @tensor temp[:] := temp_2[1, 2, -1, 3,-3] * M[ 2 3 ; 1 -2 ]
        return temp
    end

    en,V,info = eigsolve( H_lin , A ,  1  , :SR ; ishermitian = true  )
    en = en[1]
    V = V[1]

    V=reshape(V,(sAL[1]*sAL[2],sAR[2]*sAR[3]))
    U,S,V = svd(V,full=false)
    V=V'
    S=S/norm(S)
    indices = findall(1 .-cumsum(S.^2) .< tol)
    if length(indices)>0
        chi = indices[1]+1
    else
        chi = size(S)[1]
    end

    if size(S)[1] > chi
        U = U[:,1:chi]
        S = S[1:chi]
        V =  V[1:chi,:]
    end

    S /= norm(S)

    AL = reshape( U , ( sAL[1], sAL[2], :) )
    #"ij,jl,slk->sik"
    S=diagm(S)
    @tensor AR[:] := S[-1,1 ] * V[ 1,-2 ]
    AR=reshape(AR,(:,sAR[2],sAR[3]))
    return en, AL, AR
end

function two_sites_swipe_left(AL::Array, AR::Array, WL::Array, WR::Array, E::Array, F::Array)
    tol=1e-15

    sAL = size(AL)
    sAR = size(AR)

    @tensor A[:] :=AL[-1,-2,1]*AR[1,-3,-4]
    A = reshape(A,(sAL[1],sAL[2]*sAR[2],sAR[3]))

    sWL = size(WL)
    sWR = size(WR)

    @tensor M[:] := WL[-1,1,-3,-5]*WR[1,-2,-4,-6]
    M = reshape(M,(sWL[1],sWR[2],sWL[3]*sWR[3],sWL[4]*sWR[4]))

    function H_lin(v)
            @tensor temp_1[ :] := E[ 1, -2, -3] *  v[ 1, -1, -4]
            @tensor temp_2[:] :=  temp_1[-1,-2,-3,1] * F[1,-4,-5]
            @tensor temp[:] := temp_2[1, 2, -1, 3,-3] * M[ 2 3 ; 1 -2 ]
        return temp
    end

    en,V,info = eigsolve( H_lin , A ,  1  , :SR ; ishermitian = true  )
    en = en[1]
    V = V[1]
    sV=size(V)
    V=reshape(V,(sAL[1]*sAL[2],sAR[2]*sAR[3]))

    U,S,V = svd(V,full=false)
    V=V'

    S /=norm(S)
    indices = findall(1 .-cumsum(S.^2) .< tol)
    if length(indices)>0
        chi = indices[1]+1
    else
        chi = size(S)[1]
    end
    if size(S)[1] > chi
        U = U[:,1:chi]
        S = S[1:chi]
        V =  V[1:chi,:]
    end

    S /= norm(S)

    AR = reshape(V,(:,sAR[2],sAR[3]))

    S=diagm(S)
    @tensor AL[:] :=  U[ -1,1 ] * S[ 1,-2 ]
    AL = reshape( AL , ( sAL[1], sAL[2], :) )
    return en, AL, AR
end


using NumCME
using Test
using BenchmarkTools
using SparseArrays

𝕊 = [[1] [-1]]
λ = 10.0
γ = 5.0
a1 = propensity() do x, p
    p[1]
end
a2 = propensity() do x, p
    p[2]
end
sensmodel = CmeModelWithSensitivity(CmeModel(𝕊, [a1, a2], [λ, γ]))
test_space = StateSpaceSparse(𝕊, [[i] for i in 1:100])
sensmat = ForwardSensFspMatrixSparse{Float64}(sensmodel, test_space)
v = ones(3*get_rowcount(sensmat.fspmatrix))
out = similar(v)
matvec!(out, 0.0, sensmat, v)
@test sum(out) ≈ 0.0


using ATA
using JuMP
# add Cbc by running import Pkg; Pkg.add("Cbc")
using Cbc
# cd("where your input files are")

# 1. Start ATA and add file with custom settings (Needed)
ata_model = start_ata(
    settings_file = "settings_ata minimax.jl",
    bank_file = "data/bank.csv",
    bank_delim = ";",
)

# 2. Add categorical constraints (Optional)
add_constraints!(ata_model; constraints_file = "constraints.csv", constraints_delim = ";")

# 3. Add objective function (Optional)
add_obj_fun!(ata_model)
print_infos(ata_model)

# Assembly settings

# Set the solver, "siman" for simulated annealing, "jump" for MILP solver.
solver = "jump"

# MILP (Not suggested for large scale ATA)
# Select the solver, Cbc as open-source is a good option.
optimizer_constructor = "Cbc"
# #Optimizer attributes
optimizer_attributes = [("seconds", 500), ("logLevel", 1)]

# 4. assemble
assemble!(
    ata_model;
    solver = solver,
    optimizer_attributes = optimizer_attributes,
    optimizer_constructor = optimizer_constructor,
)

# All the settings and outputs from optimization are in ata_model object.
# See the struct in ATA.jl to understand how to retrieve all the information.
# A summary of the resulting tests is available in results_folder/results.txt
# If siman is chosen, the optimality and feasibility of the best neighbourhood
# is reported in "results/results_ata.jl"

print_results(ata_model; results_folder = "results")

#]add https://github.com/giadasp/ATAPlot.jl
using ATAPlot
plot_results(ata_model;  plots_folder = "plots")


abstract type AbstractLocalTurbulenceIntensityModel end

"""
    LocalTIModelNoLocalTI()

Don't calculate local turbulence intensity. Ambient TI will be used instead for all points

"""
struct LocalTIModelNoLocalTI{} <: AbstractLocalTurbulenceIntensityModel

end

"""
    LocalTIModelMaxTI(astar, bstar, k1, k2)

Calculate local turbulence intensity using the model presented in Niayifar and
Porte Agel (2015, 2016)

# Arguments
- `astar::Float`: wake spreading parameter from Bastankhah and Porte-Agel Gaussian wake model
- `bstar::Float`: wake spreading parameter from Bastankhah and Porte-Agel Gaussian wake model
- `k1::Float`: slope of k vs TI curve
- `k2::Float`: vertical offset of k vs TI curve
"""
struct LocalTIModelMaxTI{TF} <: AbstractLocalTurbulenceIntensityModel
    astar::TF
    bstar::TF
    k1::TF
    k2::TF
end
LocalTIModelMaxTI(x, y) = LocalTIModelMaxTI(x, y, 0.3837, 0.003678)
LocalTIModelMaxTI() = LocalTIModelMaxTI(2.32, 0.154, 0.3837, 0.003678)

"""
    calculate_local_ti(turbine_x, turbine_y, ambient_ti, rotor_diameter, hub_height, turbine_yaw, turbine_local_ti, sorted_turbine_index,
    turbine_inflow_velcities, turbine_ct, ti_model::LocalTIModelNoLocalTI; turbine_id=1, tol=1E-6)

Returns ambient turbulence intesity value whenever local turbulence intensity is requestesd

# Arguments
- `turbine_x::Array{Float,nTurbines}`: turbine wind direction locations in the wind direction
    reference frame
- `turbine_y::Array{Float,nTurbines}`: turbine cross wind locations in the wind direction
    reference frame
- `ambient_ti::Float`: ambient turbulence intensity
- `rotor_diameter::Array{Float,nTurbines}`: rotor diameters of all turbines
- `hub_height::Array{Float,nTurbines}`: hub heights of all turbines relative to the ground
- `turbine_yaw::Array{Float,nTurbines}`: yaw of all turbines for the current wind state in radians
- `turbine_local_ti::Array{Float,nTurbines}`: local turbulence intensity of all turbines for the current wind state`
- `sorted_turbine_index::Array{Float,nTurbines}`: turbine north-south locations in the
    global reference frame
- `turbine_inflow_velcities::Array{Float,nTurbines}`: effective inflow wind speed at each turbine for given state
- `turbine_ct::Array{Float,nTurbines}`: thrust coefficient of each turbine for the given state
- `ti_model::LocalTIModelNoLocalTI`: contains a struct defining the desired turbulence intensity model, no local TI in this case
- `turbine_id::Int`: index of wind turbine of interest. Provide 1 as default.
- `tol::Float`: How far upstream a turbine should be before being included in TI calculations
"""
function calculate_local_ti(turbine_x, turbine_y, ambient_ti, rotor_diameter, hub_height, turbine_yaw, turbine_local_ti, sorted_turbine_index,
                    turbine_inflow_velcities, turbine_ct, ti_model::LocalTIModelNoLocalTI; turbine_id=1, tol=1E-6)
    return ambient_ti
end

"""
    _k_star_func(ti_ust,k1, k2)

Calculate local wake spreading rate based on turbulence intensity using the model presented
in Niayifar and Porte Agel (2015, 2016)

# Arguments
- `ti_ust::Float`: upstream turbine local turbulence intensity
- `k1::Float`: slope of k vs TI curve
- `k2::Float`: vertical offset of k vs TI curve
"""
# compute wake spread parameter based on local turbulence intensity
function _k_star_func(ti_ust,k1,k2)
    # calculate wake spread parameter from Niayifar and Porte Agel (2015, 2016)

    k_star_ust = k1*ti_ust + k2

    return k_star_ust

end
_k_star_func(x) = _k_star_func(x, 0.3837, 0.003678)

"""
    _niayifar_added_ti_function(x, d_dst, d_ust, h_ust, h_dst, ct_ust, kstar_ust, delta_y,
        ti_amb, ti_ust, ti_dst, ti_area_ratio_in; s=700.0)

Main code for calculating the local turbulence intensity at a turbine using the method of
    Niayifar and Porte Agel (2015, 2016).

# Arguments
- `x::Float`: downstream distance from turbine to point of interest
- `d_dst::Float`: downstream turbine rotor diameter
- `d_ust::Float`: upstream turbine rotor diameter
- `h_ust::Float`: upstream turbine hub height
- `h_dst::Float`: downstream turbine hub height
- `ct_ust::Float`: upstream turbine thrust coefficient
- `kstar_ust::Float`: upstream turbine wake expansion rate
- `delta_y::Float`: cross wind separation from turbine to point of interest
- `ti_amb::Float`: ambient turbulence intensity
- `ti_ust::Float`: upstream turbine local turbulence intensity
- `ti_dst::Float`: downstream turbine local turbulence intensity
- `ti_area_ratio_in::Float`: current value of TI-area ratio for use in calculatin local TI
- `s::Float`: smooth max smootheness parameter
"""
function _niayifar_added_ti_function(x, d_dst, d_ust, h_ust, h_dst, ct_ust, kstar_ust, delta_y, ti_amb, ti_ust, ti_dst, ti_area_ratio_in; s=700.0)
    # Niayifar and Porte Agel 2015, 2016 using smooth max on area TI ratio

    # calculate axial induction based on the Ct value
    axial_induction_ust = _ct_to_axial_ind_func(ct_ust)

    # calculate BPA spread parameters Bastankhah and Porte Agel 2014
    beta = 0.5*((1.0 + sqrt(1.0 - ct_ust))/sqrt(1.0 - ct_ust))
    epsilon = 0.2*sqrt(beta)

    # calculate wake spread for TI calcs
    sigma = kstar_ust*x + d_ust*epsilon
    d_w = 4.0*sigma

    # calculate wake overlap ratio
    wake_overlap = overlap_area_func(0.0, h_dst, d_dst, delta_y, h_ust, d_w)

    # initialize the wake/rotor area overlap ratio
    ti_area_ratio = 0.0

    # only include turbines with area overlap in the softmax
    if wake_overlap > 0.0

        # Calculate the turbulence added to the inflow of the downstream turbine by the
        # wake of the upstream turbine
        ti_added = 0.73*(axial_induction_ust^0.8325)*(ti_ust^0.0325)*((x/d_ust)^(-0.32))

        rotor_area_dst = 0.25*pi*d_dst^2
        ti_area_ratio_tmp = ti_added*(wake_overlap/rotor_area_dst)

        # Run through the smooth max to get an approximation of the true max TI area ratio
        ti_area_ratio = smooth_max(ti_area_ratio_in, ti_area_ratio_tmp, s=s)

        # Calculate the total turbulence intensity at the downstream turbine based on
        # the result of the smooth max function
        ti_dst = sqrt(ti_amb^2 + ti_area_ratio^2)

    end

    return ti_dst, ti_area_ratio

end

"""
    calculate_local_ti(turbine_x, turbine_y, ambient_ti, rotor_diameter, hub_height, turbine_yaw, turbine_local_ti, sorted_turbine_index,
    turbine_inflow_velcities, turbine_ct, ti_model::LocalTIModelMaxTI; turbine_id=1, tol=1E-6)

Returns local turbulence intensity calculated using Niayifar and Porte Agel 2015, 2016 using smooth max on area TI ratio

# Arguments
- `turbine_x::Array{Float,nTurbines}`: turbine wind direction locations in the wind direction
    reference frame
- `turbine_y::Array{Float,nTurbines}`: turbine cross wind locations in the wind direction
    reference frame
- `ambient_ti::Float`: ambient turbulence intensity
- `rotor_diameter::Array{Float,nTurbines}`: rotor diameters of all turbines
- `hub_height::Array{Float,nTurbines}`: hub heights of all turbines relative to the ground
- `turbine_yaw::Array{Float,nTurbines}`: yaw of all turbines for the current wind state in radians
- `turbine_local_ti::Array{Float,nTurbines}`: local turbulence intensity of all turbines for the current wind state`
- `sorted_turbine_index::Array{Float,nTurbines}`: turbine north-south locations in the
    global reference frame
- `turbine_inflow_velcities::Array{Float,nTurbines}`: effective inflow wind speed at each turbine for given state
- `turbine_ct::Array{Float,nTurbines}`: thrust coefficient of each turbine for the given state
- `ti_model::LocalTIModelMaxTI`: contains a struct defining the desired turbulence intensity model, no local TI in this case
- `turbine_id::Int`: index of wind turbine of interest. Provide 1 as default.
- `tol::Float`: How far upstream a turbine should be before being included in TI calculations
"""
function calculate_local_ti(turbine_x, turbine_y, ambient_ti, rotor_diameter, hub_height, turbine_yaw, turbine_local_ti, sorted_turbine_index,
                    turbine_inflow_velcities, turbine_ct, ti_model::LocalTIModelMaxTI; turbine_id=1, tol=1E-6)

    # calculate local turbulence intensity at turbI

    # initialize the ri_dst and ti_area_ratio to 0.0 for current turbine
    ti_area_ratio = 0.0
    ti_dst = copy(ambient_ti)

    # extract downstream turbine information
    d_dst = rotor_diameter[turbine_id]
    h_dst = hub_height[turbine_id]

    # extract the number of turbines
    nturbines = length(rotor_diameter)

    # loop over upstream turbines
    for u=1:nturbines

        # get index of upstream turbine
        turb = sorted_turbine_index[u]

        # skip turbine's influence on itself
        if turb == turbine_id; continue; end

        # calculate downstream distance between wind turbines
        x = turbine_x[turbine_id] - turbine_x[turb]

        if x > tol

            # extract state and design info for current upstream turbine
            d_ust = rotor_diameter[turb]
            h_ust = hub_height[turb]
            yaw_ust = turbine_yaw[turb]
            ti_ust = turbine_local_ti[turb]
            ct_ust = turbine_ct[turb]

            # determine the far-wake onset location
            astar = ti_model.astar
            bstar = ti_model.bstar
            x0 = _gauss_yaw_potential_core(d_ust, yaw_ust, ct_ust, astar, ti_ust, bstar)

            # calculate the distance from the onset of far-wake
            deltax0 = x - x0

            # calculate wake spread rate for current upstream turbine
            kstar_ust = _k_star_func(ti_ust,ti_model.k1,ti_model.k2)

            # calculate horizontal and vertical spread standard deviations
            sigmay = sigmaz = _gauss_yaw_spread(x, x0, kstar_ust, d_ust, yaw_ust)

            # determine the initial wake angle at the onset of far wake
            theta0 = _bpa_theta_0(yaw_ust, ct_ust)

            # horizontal cross-wind wake displacement from hub
            wake_offset = _bpa_deflection(d_ust, ct_ust, yaw_ust, kstar_ust, kstar_ust, sigmay, sigmaz, theta0, x0)

            # cross wind distance from point location to upstream turbine wake center
            delta_y = turbine_y[turbine_id]  - (turbine_y[turb] + wake_offset)

            # save ti_area_ratio and ti_dst to new memory locations to avoid
            # aliasing during differentiation
            ti_area_ratio_tmp = deepcopy(ti_area_ratio)

            # update local turbulence intensity
            ti_dst, ti_area_ratio = _niayifar_added_ti_function(x, d_dst, d_ust, h_ust, h_dst, ct_ust, kstar_ust, delta_y, ambient_ti, ti_ust, ti_dst, ti_area_ratio_tmp)

        end

    end

    return ti_dst

end

"""
    GaussianTI(loc,turbine_x, turbine_y, rotor_diameter, hub_height, turbine_ct,
        sorted_turbine_index, ambient_ti; div_sigma=2.5, div_ti=1.2)

Calculate local turbulence intensity based on "On wake modeling, wind-farm gradients and AEP
    predictions at the Anholt wind farm" by Pena Diaz, Alfredo; Hansen, Kurt Schaldemose;
    Ott, Søren; van der Laan, Paul ??

# Arguments
- `loc::Array{Float,3}`: [x,y,z] location of point of interest in wind direction ref. frame
- `turbine_x::Array{Float,nTurbines}`: turbine wind direction locations in the wind direction
    reference frame
- `turbine_y::Array{Float,nTurbines}`: turbine cross wind locations in the wind direction
    reference frame
- `rotor_diameter::Array{Float,nTurbines}`: rotor diameters of all turbines
- `hub_height::Array{Float,nTurbines}`: hub heights of all turbines relative to the ground
- `turbine_ct::Array{Float,nTurbines}`: thrust coefficient of each turbine for the given state
- `sorted_turbine_index::Array{Float,nTurbines}`: turbine north-south locations in the
    global reference frame
- `ambient_ti::Float`: ambient turbulence intensity
- `div_sigma::Float`: ?
- `div_ti::Float`: ?
"""
function GaussianTI(loc,turbine_x, turbine_y, rotor_diameter, hub_height, turbine_ct, sorted_turbine_index, ambient_ti; div_sigma=2.5, div_ti=1.2)

    added_ti = 0.0
    e = 1.0*ambient_ti^0.1
    nturbines = length(turbine_x)

    for u=1:nturbines

        # get index of upstream turbine
        turb = sorted_turbine_index[u]

        # calculate downstream distance between wind turbines
        dx = loc[1] - turbine_x[turb]

        if dx > 1e-6
            dy = loc[2] - turbine_y[turb]
            dz = loc[3] - hub_height[turb]
            r = sqrt(dy^2 + dz^2)
            ct = turbine_ct[turb]

            kstar = 0.11*ct^1.07*ambient_ti^0.2
            epsilon = 0.23*ct^-0.25*ambient_ti^0.17
            d = 2.3*ct^-1.2
            f = 0.7*ct^-3.2*ambient_ti^-0.45

            dist = 0.5
            if r/rotor_diameter[turb] <= dist
                k1 = cos(pi/2.0*(r/rotor_diameter[turb]-dist))^2
                k2 = cos(pi/2.0*(r/rotor_diameter[turb]+dist))^2
            else
                k1 = 1.0
                k2 = 0.0
            end

            sigma = kstar*dx + epsilon*rotor_diameter[turb]
            if dz >= 0.0
                delta = 0.0
            else
                delta = ambient_ti*sin(pi*dz/hub_height[turb])^2
            end

            #2.5 for low TI 2.0 for high TI
            sigma = sigma/div_sigma

            new_ex = -2.0 #orig -2.0
            p1 = 1.0/(d + e*dx/rotor_diameter[turb] + f*(1.0+dx/rotor_diameter[turb])^new_ex)
            p2 = k1*exp(-(r-rotor_diameter[turb]/2.0)^2/(2.0*sigma^2)) + k2*exp(-(r+rotor_diameter[turb]/2.0)^2/(2.0*sigma^2))
            dI = p1*p2 - delta
            #1.2 for low TI 2.0 for high TI
            added_ti += dI/div_ti
        end
    end
    return ambient_ti + added_ti
end

# function GaussianTI_stanley(loc,windfarm,windfarmstate,ambient_ti)
#
#     added_ti = 0.0
#     e = 1.0*ambient_ti^0.1
#     nturbines = length(windfarm.turbine_x)
#
#     for u=1:nturbines
#
#         # get index of upstream turbine
#         turb = windfarmstate.sorted_turbine_index[u]
#
#         # calculate downstream distance between wind turbines
#         dx = loc[1] - windfarmstate.turbine_x[turb]
#
#         if dx > 1e-6
#             turbine_type = windfarm.turbine_definition_ids[turb]
#             rotor_diameter = windfarm.turbine_definitions[turbine_type].rotor_diameter[1]
#             hub_height = windfarm.turbine_definitions[turbine_type].hub_height[1]
#             dy = loc[2] - windfarmstate.turbine_y[turb]
#             dz = loc[3] - hub_height
#             r = sqrt(dy^2 + dz^2)
#
#             ct = windfarmstate.turbine_ct[turb]
#             kstar = 0.11*ct^1.07*ambient_ti^0.2
#             epsilon = 0.23*ct^-0.25*ambient_ti^0.17
#             d = 2.3*ct^-1.2
#             f = 0.7*ct^-3.2*ambient_ti^-0.45
#
#             # dist = 0.0
#             # if r/rotor_diameter <= dist
#             #     k1 = cos(pi/2.0*(r/rotor_diameter-dist))^2
#             #     k2 = cos(pi/2.0*(r/rotor_diameter+dist))^2
#             # else
#             #     k1 = 1.0
#             #     k2 = 0.0
#             # end
#             k1 = 1.0
#             k2 = 0.0
#
#             sigma = kstar*dx + epsilon*rotor_diameter
#             if dz >= 0.0
#                 delta = 0.0
#             else
#                 delta = ambient_ti*sin(pi*dz/hub_height)^2
#             end
#
#
#             sigma = sigma*0.5
#             p1 = 1.0/(d + e*dx/rotor_diameter + f*(1.0+dx/rotor_diameter)^-2.0)
#             # p2 = k1*exp(-(r-rotor_diameter/2.0)^2/(2.0*sigma^2)) + k2*exp(-(r+rotor_diameter/2.0)^2/(2.0*sigma^2))
#             p2 = exp(-(r/2.0)^2/(2.0*sigma^2))
#             dI = p1*p2 - delta
#             # if r < rotor_diameter*4.0/5.0
#             added_ti += dI
#             # end
#         end
#     end
#     return ambient_ti + added_ti
# end


using Flux.Tracker

# This replicates the housing data example from the Knet.jl readme. Although we
# could have reused more of Flux (see the mnist example), the library's
# abstractions are very lightweight and don't force you into any particular
# strategy.

cd(@__DIR__)

isfile("housing.data") ||
  download("https://raw.githubusercontent.com/MikeInnes/notebooks/master/housing.data",
           "housing.data")

rawdata = readdlm("housing.data")'

# The last feature is our target -- the price of the house.

x = rawdata[1:13,:]
y = rawdata[14:14,:]

# Normalise the data
x = (x .- mean(x,2)) ./ std(x,2)

# The model

W = param(randn(1,13)/10)
b = param([0.])

# using CuArrays
# W, b, x, y = cu.((W, b, x, y))

predict(x) = W*x .+ b
meansquarederror(ŷ, y) = sum((ŷ .- y).^2)/size(y, 2)
loss(x, y) = meansquarederror(predict(x), y)

function update!(ps, η = .1)
  for w in ps
    w.data .-= w.grad .* η
    w.grad .= 0
  end
end

for i = 1:10
  back!(loss(x, y))
  update!((W, b))
  @show loss(x, y)
end

predict(x[:,1]) / y[1]


function quantnormal(
    obs::Vector{Float64}, mu::Float64, sigma::Float64, q::Float64, sample_size::Int,
    config::Config
)::NTuple{4, Vector{Float64}}

    weights, mus, sigmas = adaptive_parzen_normal(obs, mu, sigma, config)
    post = GMM.GMM1(weights, mus, sigmas, q, sample_size)

    return post, weights, mus, sigmas
end


using DReal
using Base.Test
reset_ctx!()

a = Var(Bool,"a")
b = Var(Bool,"b")
add!(a → b) # \rightarrow - equivalent to implies(a,b)
a,b =  model(a,b)
@show a,b
@test !a | b


### QuantumOperations.jl
# Convenience functionality for handling common operations
###

"""
    bra(a::QKet)

Construct a dual-vector (one-form) from a ket vector `a`.

Such a one-form has type: (b::QKet) -> dot(a,b)
"""
bra(a::QKet) = (b::QKet) -> dot(a,b)

"""
    comm(a, b) = a ⊖ b

Commutator of operators `a` and `b`.

### Returns:
  - Anti-Hermitian operator: a * b' - b * a'
"""
comm(a, b) = a * b' - b * a'
⊖ = comm

"""
    scomm(a)

Superoperator for commutator with operator `a`.
Assumes Hermitian superket.

### Returns:
  - Superoperator: scomm(a) * superket(b) == superket(a * b - b * a')
"""
scomm(a) = superopl(a) - superopr(a')

"""
    acomm(a, b) = a ⊕ b

Anti-commutator of operators `a` and `b`.

### Returns:
  - Hermitian operator: a * b' + b * a'
"""
acomm(a, b) = a * b' + b * a'
⊕ = acomm

"""
    sacomm(a)

Superoperator for anticommutator with operator `a`.
Assumes Hermitian superket.

### Returns:
  - Superoperator: scomm(a) * superket(b) == superket(a * b + b * a')
"""
sacomm(a) = superopl(a) + superopr(a')

"""
    sand(a, b)

Sandwich `b` operator with `a`.

### Returns:
  - Operator: a * b * a'
"""
sand(a, b) = a * b * a'

"""
    ssand(a)

Superoperator for sandwich with operator `a`.

### Returns:
  - Superoperator: ssand(a) * superket(b) == superket(a * b * a')
"""
ssand(a) = superopl(a) * superopr(a')

"""
    diss(a)

Dissipation function for `a` action.

### Returns:
  - Function: ρ -> sand(a, ρ) - acomm(at*a, ρ)/2
"""
diss(a) = ρ -> sand(a, ρ) - acomm(a'*a, ρ)/2

"""
    sdiss(a)

Dissipation superoperator for `a` action.

### Returns:
  - ssand(a) - sacomm(at*a)/2
"""
sdiss(a) = ssand(a) - sacomm(a'*a)/2

"""
    inn(a)

Innovation function for `a` action.

### Returns:
  - Function: ρ -> acomm(a, ρ) - tr(acomm(a, ρ))*ρ
"""
inn(a)  = ρ -> let ac = acomm(a, ρ);  ac - tr(ac)*ρ end

"""
    expect(ρ, op)

Expectation value of `op` in state `ρ`.

### Returns:
  - tr(ρ * op) / tr(ρ)
"""
expect(ρ, op) = tr(ρ * op) / tr(ρ)

"""
    expectvec(ψ, op)

Expectation value of `op` in state vector `ψ`.

### Returns:
  - dot(ψ, op * ψ) / dot(ψ, ψ)
"""
expectvec(ψ, op) = dot(ψ, op * ψ) / dot(ψ, ψ)

"""
    weakvalue(ρi, ρf, op)

Generalized weak value of `op` with initial state `ρi`
and final state `ρf`.

### Returns:
  - tr(ρf * op * ρi) / tr(ρf * ρi)
"""
weakvalue(ρi, ρf, op) = tr(ρf * op * ρi) / tr(ρf * ρi)

"""
    weakvaluevec(ψi, ψf, op)

Weak value of `op` with initial state `ψi` and final state `ψf`.

### Returns:
  - dot(ψf, op * ψi) / dot(ψf, ψi)
"""
weakvaluevec(ψi, ψf, op) = dot(ψf, op * ψi) / dot(ψf, ψi)


"""
    CDFT(Yc, Xc, Xp)

"""
struct CDFT <: ContinuousUnivariateDistribution
    Yc::ContinuousUnivariateDistribution  # Distribution des données observées pour la période de calibration
    Xc::ContinuousUnivariateDistribution  # Distribution des données du modèle pour la période de calibration
    Xp::ContinuousUnivariateDistribution  # Distribution des données du modèle pour la période projetée
end

"""
    T_transform(m::CDFT, x::Real)

"""
function T_transform(m::CDFT, x::Real)

    F_Xp = cdf(m.Xp, x)  # Évaluation en x de la CDF des données du modèle pour la période projeté

    Fm1_Xc = quantile(m.Xc, F_Xp)  # Quantile correspondant de la distribution des données du modèle pour la période de calibration

    F_Yp = cdf(m.Yc,  Fm1_Xc)  # Évaluation de la CDF des données observées pour la période de calibration au point correspondant

    return F_Yp  # Valeur de la CDF évaluée au point x
end

#### Evaluation

cdf(m::CDFT, x::Real) = T_transform(m, x)

function pdf(m::CDFT, x::Real)
    # très approximatif... à revoir
    # Inf en zéro (?)

    δ = 0.00001
    x₁ = x - δ
    x₂ = x + δ

    return (cdf(m, x₂) - cdf(m, x₁)) / (x₂ - x₁)
end

function quantile(m::CDFT,
                  p::Real;
                  interval::Array{Float64,1}=[0.0, 500.0])

    @assert zero(p)<p<one(p) "the quantile level should be between 0 and 1."

    res = optimize(θ -> (p - cdf(m, θ))^2, interval[1], interval[2])

    return Optim.minimizer(res)[1]
end

#### Sampling

function rand(rng::AbstractRNG, m::CDFT)
    # Generate a Float64 random number uniformly in (0,1].
    u = 1 - rand(rng)

    return quantile(m, u)
end

sampler(m::CDFT) = m    # à vérifier


"""
Background Newman-Penrose scalars
"""
module BackgroundNP

struct NP_0
   mu_0  ::Array{ComplexF64,2}
   tau_0 ::Array{ComplexF64,2}
   pi_0  ::Array{ComplexF64,2}
   rho_0 ::Array{ComplexF64,2}
   eps_0 ::Array{ComplexF64,2}
   psi2_0::Array{ComplexF64,2}

   function NP_0(;
         Rvals::Vector{Float64},
         Yvals::Vector{Float64},
         Cvals::Vector{Float64},
         Svals::Vector{Float64},
         bhm::Float64,
         bhs::Float64,
         cl::Float64)

      nx = length(Rvals)
      ny = length(Yvals)

      mu_0   = zeros(ComplexF64,nx,ny)
      tau_0  = zeros(ComplexF64,nx,ny)
      pi_0   = zeros(ComplexF64,nx,ny)
      rho_0  = zeros(ComplexF64,nx,ny)
      eps_0  = zeros(ComplexF64,nx,ny)
      psi2_0 = zeros(ComplexF64,nx,ny)

      for j=1:ny
         sy = Svals[j]
         cy = Cvals[j]

         for i=1:nx
            R = Rvals[i]

            mu_0[i,j]   = 1.0 / (-(cl^2) + im*bhs*R*cy)

            tau_0[i,j]  = (im*bhs*sy/sqrt(2.0)) / ((cl^2 - im*bhs*R*cy)^2)

            pi_0[i,j]   = - (im*bhs*sy/sqrt(2.0)) / (cl^4 + (bhs*cy*R)^2)

            rho_0[i,j]  = (
               -
               0.5*(
                    cl^4 - 2.0*(cl^2)*bhm*R + (bhs*R)^2
               )/(
                  ((cl^2 - im*bhs*R*cy)^2)*(cl^2 + im*bhs*R*cy)
               )
              )

            eps_0[i,j]  = (
               0.5*(
                  (cl^2)*bhm - (bhs^2)*R - im*bhs*(cl^2-bhm*R)*cy
               )/(
                  ((cl^2 - im*bhs*R*cy)^2)*(cl^2 + im*bhs*R*cy)
               )
              )

            psi2_0[i,j] = - bhm / ((cl^2 - im*bhs*R*cy)^3)
         end
      end
      return new(mu_0,tau_0,pi_0,rho_0,eps_0,psi2_0)
   end
end

end


using BEAST
using CompScienceMeshes
using LinearAlgebra
using SparseArrays

import Plotly
import Plots
Plots.plotly()

function Base.:+(a::S, b::S) where {S<:BEAST.Space}
    @assert geometry(a) == geometry(b)
    S(geometry(a), [a.fns; b.fns], [a.pos; b.pos])
end

l = w = 1.0 #Length and width of capacitor plates
d = 0.1     #seperation of plates
h = 1/36    #size of meshes

Γ₀ = meshrectangle(l,w,h)
Γ₁ = CompScienceMeshes.translate(Γ₀, point(0.0,0.0,d))

γ₀ = meshsegment(l,h,3)
γ₁ = CompScienceMeshes.translate(γ₀, point(0.0,0.0,d))

Γ = weld(Γ₁, Γ₀)

κ = 2pi / 100
V₀ = 1.0
f = ScalarTrace(p -> V₀)

edges_all = skeleton(Γ, 1)
edges_int = submesh(!in(boundary(Γ)), edges_all)

on_γ₀ = overlap_gpredicate(γ₀)
edges_pt0 = submesh(edges_all) do edge
    ch = chart(edges_all, edge)
    return on_γ₀(ch)
end

on_γ₁ = overlap_gpredicate(γ₁)
edges_pt1 = submesh(edges_all) do edge
    ch = chart(edges_all, edge)
    return on_γ₁(ch)
end

RT_int = raviartthomas(Γ, edges_int)
RT_pt0 = raviartthomas(Γ, edges_pt0)
RT_pt1 = raviartthomas(Γ, edges_pt1)

verts_pt0_int = submesh(!in(boundary(edges_pt0)), skeleton(edges_pt0,0))
verts_pt1_int = submesh(!in(boundary(edges_pt1)), skeleton(edges_pt1,0))

C0 = connectivity(verts_pt0_int, edges_pt0)
C1 = connectivity(verts_pt1_int, edges_pt1)

X = RT_int
Y0 = RT_pt0 * C0
Y1 = RT_pt1 * C1

# This is ugly and will fail if the two ports are not equal...
D = sparse(reshape([fill(+1.0, length(edges_pt0)); fill(-1.0, length(edges_pt1))],length(edges_pt0)+length(edges_pt1),1))
RT_pt = RT_pt0 + RT_pt1
Z = RT_pt * D

@hilbertspace x y0 y1 z
@hilbertspace ξ η0 η1 ζ

T = Maxwell3D.singlelayer(wavenumber=κ)
trc = X->ntrace(X,γ₁)

efie = @discretise( @varform(
    T[ξ,x]  + T[ξ,y0]  + T[ξ,y1]  + T[ξ,z]  +
    T[η0,x] + T[η0,y0] + T[η0,y1] + T[η0,z] +
    T[η1,x] + T[η1,y0] + T[η1,y1] + T[η1,z] +
    T[ζ,x]  + T[ζ,y0]  + T[ζ,y1]  + T[ζ,z] == f[trc(ζ)]),
    ξ∈X, η0∈Y0, η1∈Y1, ζ∈Z,
    x∈X, y0∈Y0, y1∈Y1, z∈Z)
u = solve(efie)

# You can access the current coefficients pertaining to subspaces
# using the Hilbert space 'placeholder'
@show length(u[x])
@show length(u[y0])
@show length(u[y1])
@show length(u[z])

S = X + Y0 + Y1 + Z
fcr, geo = facecurrents(u, S)
Plotly.plot(patch(geo, norm.(fcr))) |> display

# Compute the Scalar Potential across the ports. Note how a voltage
# gap of 1V comes out as a 'natural' condition on the solution.
zs = range(-0.5, stop=0.5, length=100)
pts = [point(0.5,0.5,z) for z ∈ zs]
Φ = potential(MWSingleLayerPotential3D(κ), pts, u, S, type=ComplexF64)
Plots.plot(zs, real(Φ), xlabel="height",ylabel="scalar potential",label=false) |> display

# Plot current along γ₀ and γ₁. Note: in this symmetric situation, we
# expect these to be opposite on corresponding points of the two ports.
# In general this is not true; in fact, the two ports could have different shape and size
traceS0 = ntrace(S, γ₀)
traceS1 = ntrace(S, γ₁)
fcr0, geo0 = facecurrents(u, traceS0)
fcr1, geo1 = facecurrents(u, traceS1)
plot()
plot!(imag.(fcr0))
plot!(imag.(fcr1)) |> display


################################################################################
# Convenience functions for kinematics
################################################################################

#..oooOOO0Q0Q0OOOooo...oooOOO0Q0Q0OOOooo...oooOOO0Q0Q0OOOooo...oooOOO0Q0Q0OOOooo

"""
	β(γ::Number)

Returns β given γ.
"""
β(γ::Number) = √(γ^2 - 1) / γ

export β
#..oooOOO0Q0Q0OOOooo...oooOOO0Q0Q0OOOooo...oooOOO0Q0Q0OOOooo...oooOOO0Q0Q0OOOooo
"""
    function momentum(E::Quantity,id::ParticleID)

Returns momentum given total energy and rest mass.
"""
function momentum(E::Quantity, id::ParticleID)
    @assert typeof(E) <: Unitful.Energy "E requires units of energy"
    m = mass(id)
    return uconvert(u"GeV/c",√(E^2 - m^2 * Unitful.c^4) / Unitful.c)
end #momentum

"""
	function momentum(E::Quantity,m::Quantity)

Returns momentum given total energy and rest mass.
"""
function momentum(E::Quantity, m::Quantity)
    @assert typeof(E) <: Unitful.Energy "E requires units of energy"
    @assert typeof(m) <: Unitful.Mass "m requires units of mass"
    return uconvert(u"GeV/c",√(E^2 - m^2 * Unitful.c^4) / Unitful.c)
end #momentum

"""
    function momentumsi(E::Float64,id::ParticleID)

Returns momentum given total energy and a ParticleID. SI units used.
"""
function momentum_si(E::Float64, id::ParticleID)
    m = usistrip(mass(id))
    momentum_si(E,m)
end #momentum

"""
	function momentum(E::Number,m::Number)

Returns momentum given total energy and rest mass in SI units.
"""
function momentum_si(E::Number, m::Number)
    return √(E^2 - m^2 * Unitful.c0.val^4) / Unitful.c0.val
end #momentum

export momentum, momentum_si
#..oooOOO0Q0Q0OOOooo...oooOOO0Q0Q0OOOooo...oooOOO0Q0Q0OOOooo...oooOOO0Q0Q0OOOooo

"""
	γFromβ(β::Number)

Returns γ given β
"""
function γFromβ(β::Number)
    @assert β <= 1.0 "β ≦ 1 required"
    return 1 / sqrt(1 - β^2)
end #γFromβ

#..oooOOO0Q0Q0OOOooo...oooOOO0Q0Q0OOOooo...oooOOO0Q0Q0OOOooo...oooOOO0Q0Q0OOOooo

"""
    γFromP(p::Array{T}, m::Quantity) where {T<:Quantity}

Returns gamma given momentum and rest mass.
"""
function γFromP(p::Array{T}, m::Quantity) where {T<:Quantity}
    pnorm = norm(p)
    γFromP(pnorm,m)
end #γFromP

"""
    γFromP(p::Array{T}, m::Number) where {T<:Number}

Returns gamma given momentum and rest mass in SI units.
"""
function γFromP(p::Array{T}, m::Number) where {T<:Number}
    pnorm = norm(p)
    γFromP(pnorm,m)
end #γFromP

"""
γFromP(p::Quantity,m::Quantity)

Returns gamma given momentum and rest mass.
"""
function γFromP(p::Quantity, m::Quantity)
    @assert typeof(p) <: Unitful.Momentum "p requires units of momentum"
    @assert typeof(m) <: Unitful.Mass "m requires units of mass"
    c = Unitful.c
    βγ = upreferred(p / (m * c))
    return upreferred(sqrt(1 + βγ^2))
end #γFromP

"""
	γFromP(p::Number,m::Number)

Returns gamma given momentum and rest mass in SI units.
"""
function γFromP(p::Number, m::Number)
    c = Unitful.c0.val
    return sqrt(1 + (p / (m * c))^2)
end #γFromP

export γFromP
#..oooOOO0Q0Q0OOOooo...oooOOO0Q0Q0OOOooo...oooOOO0Q0Q0OOOooo...oooOOO0Q0Q0OOOooo
"""
	βFromP(p::Quantity,m::Quantity)

Returns β given particle momentum and rest mass.
"""
function βFromP(p::Quantity, m::Quantity)
    @assert typeof(p) <: Unitful.Momentum "p requires units of momentum"
    @assert typeof(m) <: Unitful.Mass "m requires units of mass"
    c = Unitful.c
    η = upreferred(p / (m * c))
    return sqrt(η^2 / (1 + η^2))
end

"""
	βFromP(p::Number,m::Number)

Returns β given particle momentum and rest mass in SI units
"""
function βFromP(p::Number, m::Number)
    c = Unitful.c0.val
    η = p / (m * c)
    return sqrt(η^2 / (1 + η^2))
end

export βFromP
#..oooOOO0Q0Q0OOOooo...oooOOO0Q0Q0OOOooo...oooOOO0Q0Q0OOOooo...oooOOO0Q0Q0OOOooo
"""
    vFromP(p::Array{T}, m::U) where {T <: Quantity, U <: Quantity}

Returns velocity given particle momentum and rest mass in SI units
"""
function vFromP(p::Array{T}, m::U) where {T <: Quantity, U <: Quantity}
    β = βFromP(norm(p),m)
    c = Unitful.c
    phat = p / norm(p)
    return upreferred.(c * β * phat)
end

"""
    vFromP(p::Array{T}, m::T) where {T <: Number}

Returns velocity given particle momentum and rest mass in SI units
"""
function vFromP(p::Array{T}, m::T) where {T <: Number}
    β = βFromP(norm(p),m)
    c = Unitful.c0.val
    phat = p / norm(p)
    return c * β * phat
end

"""
    vFromP(p::Quantity, m::Quantity)

Returns velocity given particle momentum and rest mass in SI units
"""
function vFromP(p::Quantity, m::Quantity)
    β = βFromP(p,m)
    return upreferred.(Unitful.c * β )
end

"""
    vFromP(p::Array{T}, m::T) where {T <: Number}

Returns velocity given particle momentum and rest mass in SI units
"""
function vFromP(p::Number, m::Number)
    β = βFromP(p,m)
    c = Unitful.c0.val
    return c * β
end

export vFromP
#..oooOOO0Q0Q0OOOooo...oooOOO0Q0Q0OOOooo...oooOOO0Q0Q0OOOooo...oooOOO0Q0Q0OOOooo
"""
	gyroradius(p::Quantity,Z::Quantity,B::Quantity)

Returns gyroradius given dimensionally correct momentum, charge, and B field strength
"""
function gyroradius(p::Quantity, Q::Quantity, B::Quantity)
    @assert typeof(p) <: Unitful.Momentum "p must have units of momentum"
    @assert typeof(Q) <: Unitful.Charge "Z must have units of charge"
    @assert typeof(B*1u"m^2") <: Unitful.MagneticFlux "B must have units of magnetic flux density"

    return upreferred(p/(Q*B))
end
# TODO: add functionality to allow for vector p and B

"""
	gyroradius(p::Number, Z::Number, B::Number)

Returns gyroradius (m) given momentum, atomic (charge) number, and B field strength

# Arguments
- `p::Number`: momentum of particle in GeV/c
- `Z::Number`: charge of particle in units of q_e
- `B::Number`: magnetic flux density in gauss
"""
function gyroradius(p::Number, Z::Number, B::Number)
    33356.40951981521 * p / (Z * B)
end

export gyroradius


using DelimitedFiles: readdlm
using LinearAlgebra: dot, I

using KalmanFilters.UnscentedKalmanFilter
using KalmanFilters.SigmaPoints

function normalize_angle(α::Float64)::Float64
    m = 2π
    α = α > 0 ? α % m : m + (α % m)
    if α > π
        α -= m
    end
    α
end

function state_residual(
    y::AbstractArray{Float64, 1}, x::AbstractArray{Float64, 1}
)::AbstractArray{Float64, 1}
    r = y - x
    r[3] = normalize_angle(r[3])
    r
end

function state_residual(
    y::AbstractArray{Float64, 2}, x::AbstractArray{Float64, 1}
)::AbstractArray{Float64, 2}
    R = y .- x'
    R[:, 3] .= normalize_angle.(@view(R[:, 3]))
    R
end

function state_add(
    y::AbstractArray{Float64, 1}, x::AbstractArray{Float64, 1}
)::AbstractArray{Float64, 1}
    r = y + x
    r[3] = normalize_angle(r[3])
    r
end

function state_mean(
    Σ::Array{Float64, 2}, m_w::Array{Float64, 1}
)::Array{Float64, 1}
    [
        dot(@view(Σ[:, 1]), m_w),
        dot(@view(Σ[:, 2]), m_w),
        atan(dot(sin.(@view(Σ[:, 3])), m_w), dot(cos.(@view(Σ[:, 3])), m_w)),
    ]
end

hx(;x::AbstractArray{Float64, 1})::AbstractArray{Float64, 1} = x

function fx(
    ;x::Array{Float64, 1}, δt::Float64,
    μ::Array{Float64, 1}, wheelbase::Float64
)::Array{Float64, 1}
    ν, α = μ
    if abs(ν) < 0.3
        return x
    end

    θ = x[3]
    s = ν * δt
    if abs(α) < 0.017
        return x + [s * cos(θ), s * sin(θ), 0.0]
    end

    β = (s / wheelbase) * tan(α)
    r = s / β
    x + [-r * sin(θ) + r * sin(θ + β), r * cos(θ) - r * cos(θ + β), β]
end

load_from_txt(path::String) = readdlm(path, ' ', Float64, '\n')

function main()
    base_path = raw"C:\Users\tonys\projects\carla-dataset\high-quality\recording-2020-05-01-13-52-24-005689"
    positions_time = load_from_txt(joinpath(base_path, "gnss", "timestamp"))
    positions = load_from_txt(joinpath(base_path, "gnss", "enu"))[:, 1:2]

    imu_time = load_from_txt(joinpath(base_path, "imu", "timestamp"))
    compass = load_from_txt(joinpath(base_path, "imu", "compass"))
    speed = load_from_txt(joinpath(base_path, "imu", "speed"))
    steer = load_from_txt(joinpath(base_path, "imu", "steer"))
    # Convert data
    measurement_noise = 1.5
    compass *= π / 180
    steer *= -70.0 * π / 180
    wheelbase = 2.9591
    positions += randn(size(positions)...) * measurement_noise
    println("Loaded data")

    σ_parameters = MerweScaled(n=3, α=1e-3, β=2.0, κ=0.0, residual_x=state_residual)
    ukf = UKFState(
        dim_x=3, dim_z=3, σ_parameters=σ_parameters,
        add_x=state_add,
        mean_x=state_mean, mean_z=state_mean,
        residual_x=state_residual, residual_z=state_residual
    )
    ukf.x .= [positions[1, 1], positions[1, 2], compass[1]]
    ukf.P .*= 0.1

    δt = 1.0 / 50
    μ = zeros(Float64, 2)
    z = zeros(Float64, 3)
    R = Matrix{Float64}(I, 3, 3) * measurement_noise
    Q = Float64[[4e-12 4e-10 2e-8]; [4e-10 4e-8 2e-6]; [2e-8 2e-6 1e-4]]

    track = zeros(Float64, size(imu_time, 1), 2)
    pos_i::Int64 = 1
    can_update::Bool = false

    for (i, t) = enumerate(imu_time)
        μ[1] = speed[i]
        μ[2] = steer[i]
        predict!(ukf=ukf, δt=δt, Q=Q, fx=fx, μ=μ, wheelbase=wheelbase)

        can_update = pos_i < size(positions_time, 1) && (
            abs(t - positions_time[pos_i]) < δt
            || t > positions_time[pos_i]
        )
        if can_update
            z[1:2] = positions[pos_i, :]
            z[3] = compass[i]
            update!(ukf=ukf, z=z, R=R, hx=hx)
            pos_i += 1
        end
        track[i, :] = ukf.x[1:2]
    end

    println(ukf.x)
end

main()


using SampleChainsDynamicHMC
using Soss

g = @model N begin
    n ~ Poisson(10) |> iid(N)
    p ~ Uniform()
    y ~ For(n) do nj Binomial(nj, p) end
end;

obs = predict(g, N=5, p=0.2)

pairs(obs)

m = predictive(g, :n)

post = m(n=[6,12,7,8,12]) | (y=[1,3,1,2,3],)

s = sample(DynamicHMCChain, post)

predict(m(n=obs.n), s)

drawsamples!(s, 1000)

a1 = ash(getchains(s)[1].p)
a2 = ash(getchains(s)[2].p)
a3 = ash(getchains(s)[3].p)
a4 = ash(getchains(s)[4].p)

plot(a1.rng, a1.density)
plot!(a2.rng, a2.density)
plot!(a3.rng, a3.density)
plot!(a4.rng, a4.density)

drawsamples!(s, 10000)

p = 0.001:0.001:0.999
plot(p, exp.([ℓ((p=pj,)) for pj in p]))

using TransformVariables

ℓ(x) = logdensity(Beta(4,3), x.p)
t = xform(post)
chain = initialize!(DynamicHMCChain, ℓ, t)
drawsamples!(chain, 10000)
plot(ash(chain.p))


"""
    CoxIngersollRoss()

[Cox-Ingersoll-Ross model](https://en.wikipedia.org/wiki/Cox%E2%80%93Ingersoll%E2%80%93Ross_model).
"""
struct CoxIngersollRoss <: Model end

"""
    evaluate(IRD, CoxIngersollRoss(), n)

Evaluate interest rate derivative `IRD` using `CoxIngersollRoss` model.

# Arguments
- `IRD::InterestRateDerivative`: interest rate derivative
- `n`: number of paths to simulate
"""
function evaluate(IRD::InterestRateDerivative, m::CoxIngersollRoss, n::Int64 = 12)
    Δt = IRD.t / n
    rates = [IRD.r]
    for i = 0:n
        dr = IRD.k * (IRD.θ - rates[end]) * Δt + IRD.σ * √rates[end] * randn()
        append!(rates, rates[end] + dr)
    end
    return rates
end


include("common.jl")

Δt = 0.01
mmesh = Mesh(30, 30, 1/30)
bdnode = bcnode(mmesh)
bdnode = [bdnode; bdnode .+ mmesh.ndof]
nu = constant(0.0001ones(get_ngauss(mmesh)))
nodes = fem_nodes(mmesh)
u = @. sin(2π * nodes[:,1])
v = @. sin(2π * nodes[:,2])
u0 = [u;v]
u0[bdnode] .= 0.0
us = solve_burgers(u0, 10, nu)

sess = Session(); init(sess)
U = run(sess, us)

matwrite("fenics/fwd1.mat", Dict("U"=>U))
close("all")
figure(figsize=(10,3))
subplot(121)
title("u displacement")
visualize_scalar_on_fem_points(U[end,1:mmesh.nnode], mmesh)
subplot(122)
title("v displacement")
visualize_scalar_on_fem_points(U[end,mmesh.ndof + 1:mmesh.ndof + mmesh.nnode], mmesh)
savefig("fenics/fwd1.png")

close("all")
visualize_vector_on_fem_points(U[end,1:mmesh.nnode], U[end,mmesh.ndof + 1:mmesh.ndof + mmesh.nnode], mmesh)
savefig("fenics/fwd1_quiver.png")