using MPIReco, JLD, SensorSelection, PyPlot

# estimate measurement noise from the training data
function varMeas(x::Matrix{T},rs::Int64) where T
  roi = vec(x[1:rs,1:rs].-mean(x[1:rs,1:rs]))
  append!(roi,x[end-rs+1:end,1:rs].-mean(x[end-rs+1:end,1:rs]))
  append!(roi,x[1:rs,end-rs+1:end].-mean(x[1:rs,end-rs+1:end]))
  append!(roi,x[end-rs+1:end,end-rs+1:end].-mean(x[end-rs+1:end,end-rs+1:end]))

  return sum(abs.(vec(roi)).^2)/(length(roi)-1)
end

####################
# load system matrix
####################
# Dataset store handling
# uncomment the next two lines if running this script separately
# datadir = artifact"MDFStore"
# store = MDFDatasetStore(datadir)

# system matrix
datadirSF1 = joinpath(calibdir(store), "1.mdf")
fSF1 = MPIFile(datadirSF1)
nx,ny,nz = calibSize(fSF1)
S1 = reshape( getSystemMatrix(fSF1,bgCorrection=true), nx, ny, :)
snr1 = MPIReco.getSNRAllFrequencies(fSF1)

# select a slice and apply SNR thresholding
snr_idx1 = findall(x->x>3.0, vec(snr1))[2:end]
S1 = S1[:,:,snr_idx1]

# normalize the frequency components
for k=1:size(S1,3)
  S1[:,:,k] ./= norm(S1[:,:,k])
end

# chose frequency components for training
Random.seed!(1234)
p = shuffle(collect(4:2:size(S1,3)))
St = S1[:,:,p[1:30]]  # training data

########################################################
# setup up measurement operators and covariance matrices
########################################################

# mean amplitudes of measurements
st_amp = zeros(nx,ny,size(St,3))
for k=1:size(St,3)
  st_amp[:,:,k] .= sqrt.( 1/size(St,3)*dropdims(sum(abs.(St[:,:,:]).^2,dims=3),dims=3) )
end
st_amp = reshape(st_amp,nx*ny,:)

# measurement matrix
dctOp = DCTOp(ComplexF64,(nx,ny))
H = Matrix(adjoint(dctOp))
Ht = collect(transpose(real.(H))) # build transpose of measurement matrix

# DCT representation of training data
St_dct = zeros(ComplexF64,nx*ny,size(St,3));
for k=1:size(St,3)
    St_dct[:,k] .= dctOp*vec(St[:,:,k])
end

# obtain sparse measurement matrix
thresh = 1.e-1  # portion of the total number of sparse coefficients to use for optimization
numCoeff = floor(Int64,thresh*nx*ny)
Ht_sub = zeros(numCoeff,size(Ht,2),size(St,3))
for k=1:size(St,3)
  Sk_dct = vec(abs.(St_dct[:,k]))
  idx_sig = maxIndices(Sk_dct,numCoeff)
  Ht_sub[:,:,k] .= Ht[idx_sig,:]
end

# estimate measurement noise for each frequency component
measNoise = [varMeas(St[:,:,k],3) for k=1:size(St,3)]
r = zeros(nx*ny,size(St,3))
for k=1:size(St,3)
  r[:,k] .= fill(measNoise[k],nx*ny)
end

######################
# Forward optimization
######################
λ=3.0                             # proportionality factor for the prior covariance matrices
numSamples = floor(Int64, nx*ny)  # number of samples to add to the sampling pattern
w = zeros(Bool,nx*ny)             # initial sampling pattern (empty)
m = SensorSelection.Experiment(Ht_sub, r, Ht_sub, λ.*st_amp.^2,w)
@time idx_opt = optSamplingFW!(m,numSamples,4,180)

#############################
# plot some sampling patterns
#############################
# precomputed Poisson disk patterns
msk2_pd = zeros(nx,ny); msk2_pd[load(datadir*"/samplingPatterns/pd_r2.jld","idx")] .= 1
msk6_pd = zeros(nx,ny); msk6_pd[load(datadir*"/samplingPatterns/pd_r6.jld","idx")] .= 1
msk10_pd = zeros(nx,ny); msk10_pd[load(datadir*"/samplingPatterns/pd_r10.jld","idx")] .= 1
# optimized patterns
msk2_opt = zeros(nx,ny); msk2_opt[idx_opt[1:div(nx*ny,2)]] .= 1
msk6_opt = zeros(nx,ny); msk6_opt[idx_opt[1:div(nx*ny,6)]] .= 1
msk10_opt = zeros(nx,ny); msk10_opt[idx_opt[1:div(nx*ny,10)]] .= 1

figure("sampling Patterns", figsize=(8,6))
subplot(2,3,1); imshow(msk2_pd); title("r=2"); ylabel("PD"); xticks([]); yticks([])
subplot(2,3,2); imshow(msk6_pd); title("r=6"); xticks([]); yticks([])
subplot(2,3,3); imshow(msk10_pd); title("r=10"); xticks([]); yticks([])
subplot(2,3,4); imshow(msk2_opt); ylabel("Opt"); xticks([]); yticks([])
subplot(2,3,5); imshow(msk6_opt); xticks([]); yticks([])
subplot(2,3,6); imshow(msk10_opt); xticks([]); yticks([])
tight_layout()


using AcFDTD
using Base.Test
using Base.Profile
srand(123)

geo = CuboidRoom(10,15,16,[10.;20.;30.;40.;50.;60.], FDTDEnv(0.1,IWB()))
Nt = 200
s_t = [randn(Nt,1) randn(Nt,1)]
S = randn(size(geo)...,Nt)
xs = [(5, 5, 5),(5,2,4)]
xr = [(5, 8, 7),(2,2,2)]
f = FDTD(geo)

p0ic = randn(10,15,16)
p1ic = randn(10,15,16)

println("test calls: \n")
println("----------      no IC")
p    = fdtd(xr, xs, Nt,geo, s_t)
p    = fdtd(xr, xs, Nt,  f, s_t)
p2 = copy(p)
fdtd!(p2, xr, xs, Nt,  f, s_t)
@time p    = fdtd(xr, xs, Nt,  f, s_t)
println("                       in-place call")
@time        fdtd!(p2,xr, xs, Nt,  f, s_t)

@test norm(p-p2) < 1e-8

println("----------         IC")
pic  = fdtd(xr, xs, Nt,geo, s_t, p0ic, p1ic)
ppic = fdtd(xr, xs, Nt,  f, s_t, p0ic, p1ic)
pic2 = copy(pic)
fdtd!(pic2, xr, xs, Nt,  f, s_t, p0ic, p1ic)
@time ppic = fdtd(xr, xs, Nt,  f, s_t, p0ic, p1ic)
println("                       in-place call")
@time        fdtd!(pic2, xr, xs, Nt,  f, s_t, p0ic, p1ic)

@test norm(pic-pic2) < 1e-8
@test norm(pic-ppic) < 1e-8

println("---------- full no IC")
P    = fdtd(    xs, Nt,geo, s_t)
P    = fdtd(    xs, Nt,  f, s_t)
P2 = copy(P)
fdtd!(P2,    xs, Nt,  f, s_t)
@time P    = fdtd(    xs, Nt,  f, s_t)
println("                       in-place call")
@time        fdtd!(P2, xs, Nt,  f, s_t)

@test vecnorm(P-P2) < 1e-8

println("---------- full    IC")
PIC  = fdtd(    xs, Nt,geo, s_t, p0ic, p1ic)
PIC  = fdtd(    xs, Nt,  f, s_t, p0ic, p1ic)
PIC2 = copy(PIC)
fdtd!(PIC2,    xs, Nt,  f, s_t, p0ic, p1ic)
@time PIC  = fdtd(    xs, Nt,  f, s_t, p0ic, p1ic)
println("                       in-place call")
@time      fdtd!(PIC2,xs, Nt,  f, s_t, p0ic, p1ic)

@test vecnorm(PIC-PIC2) < 1e-8

println("---------- source everywhere, no IC")
ps    = fdtd(    xr, Nt,geo, S)
ps    = fdtd(    xr, Nt,  f, S)
ps2 = copy(ps)
fdtd!(ps2, xr, Nt,  f, S)
@time ps    = fdtd(      xr, Nt,  f, S)
println("                       in-place call")
@time        fdtd!(ps2,  xr, Nt,  f, S)

println("---------- source everywhere, IC")
psic    = fdtd(    xr, Nt,geo, S, p0ic, p1ic)
psic    = fdtd(    xr, Nt,  f, S, p0ic, p1ic)
psic2 = copy(ps)
fdtd!(psic2, xr, Nt,  f, S, p0ic, p1ic)
@time psic    = fdtd(      xr, Nt,  f, S, p0ic, p1ic)
println("                       in-place call")
@time        fdtd!(psic2,  xr, Nt,  f, S, p0ic, p1ic)

@test vecnorm(psic-psic2) < 1e-8

println("---------- full no IC, source everywhere")
Ps    = fdtd(        Nt,geo, S)
Ps    = fdtd(        Nt,  f, S)
Ps2 = copy(Ps)
fdtd!(Ps2,    Nt,  f, S)
@time Ps    = fdtd(        Nt,  f, S)
println("                       in-place call")
@time        fdtd!(Ps2,    Nt,  f, S)

@test vecnorm(Ps-Ps2) < 1e-8

println("---------- full IC, source everywhere")
Psic    = fdtd(        Nt,geo, S, p0ic, p1ic)
Psic    = fdtd(        Nt,  f, S, p0ic, p1ic)
Psic2 = copy(Ps)
fdtd!(Psic2,    Nt,  f, S)
@time Psic    = fdtd(        Nt,  f, S, p0ic, p1ic)
println("                       in-place call")
@time        fdtd!(Psic2,    Nt,  f, S, p0ic, p1ic)

@test vecnorm(Psic-Psic2) < 1e-8

println("\ntest equivalence with matrix formulation. \n")

B = get_B(geo,Nt)
E = get_Exp(geo,Nt,xs)
F = get_Sel(geo,Nt,xr)
s = E*s_t[:]
s =[p0ic[:];p1ic[:];s[2*geo.Nx*geo.Ny*geo.Nz+1:end]]
S =[p0ic[:];p1ic[:];S[:]]
pm = F*(B\s)
println("sparse matrix      call")
@time pm = F*(B\s)
println("sparse matrix full call")
Pm = B\s
@time Pm = B\s

Pms = B\S

@test norm(PIC[:]-Pm[2*geo.Nx*geo.Ny*geo.Nz+1:end]) <1e-8
@test norm(pic[:]-pm[:]) <1e-8

@test norm(Psic[:]-Pms[2*geo.Nx*geo.Ny*geo.Nz+1:end]) <1e-8

println("\ntest resonances on cuboid room are correct \n")
Nt = 3000
X = 0.4
geo = CuboidRoom(10, 11, 12, 1e8*ones(6), FDTDEnv(X,IISO()))

using DSP
s_t = zeros(Nt)   #source signal
s_t[3] = 1
f2 = geo.env.Fs/2*0.14#cut-off frequency of source
filt!(s_t,digitalfilter(Bandpass(10,f2;fs = geo.env.Fs),Butterworth(15)),s_t)

xr = (2, 2, 2)
xs = (geo.Nx-1, geo.Ny-1, geo.Nz-1)
p_t2 = fdtd(xr, xs, Nt, geo, s_t)

f = linspace(0,geo.env.Fs/2,div(Nt,2)+1)
fres  = zeros(5,5,5)
for l = 0:4,m=0:4,i =0:4

	fres[l+1,m+1,i+1] = 343/(2*π)*sqrt(  (l*π/(geo.Lx))^2+(m*π/(geo.Ly))^2+(i*π/(geo.Lz))^2 )
end
fres = sort(fres[:])[2:7]

P = 10*log10(abs(rfft(p_t2)))

##using PyPlot
##figure()
##plot(f,P)
##plot(fres[:],5.* ones(length(fres[:])), "r*" )
##xlim([0,400])

@test norm(f[P.>5.5][1:6]-fres)<0.3 #test on resonance frequencies

println("\ntest acoustic delay is correct \n")
X = 0.4
Nt = 500
geo = CuboidRoom(20, 21, 24, 4*ones(6), FDTDEnv(X,IWB()))
s_t = zeros(Nt)   #source signal
s_t[3] = 1
f2 = geo.env.Fs/2*0.1#cut-off frequency of source
filt!(s_t,digitalfilter(Bandpass(10,f2;fs = geo.env.Fs),Butterworth(2)),s_t)

xs = (div(geo.Nx,2), div(geo.Ny,2)  , div(geo.Nz,2))
xr = (div(geo.Nx,2), div(geo.Ny,2)+4, div(geo.Nz,2))
d  = norm(X*[ xs[1]-xr[1],xs[2]-xr[2],xs[3]-xr[3]])

p_t2 = fdtd(xr, xs, Nt, geo, s_t)
t = 1/geo.env.Fs:1/geo.env.Fs:Nt/geo.env.Fs

##using PyPlot
##plot(t,p_t2)
##plot(d/343.+2*X/343, 0., "r*")#+2 due to initial conditions

#test delay
@test norm((indmax(xcorr(p_t2[:],s_t)[end-length(p_t2):end])-1)/geo.env.Fs-d/343.)<1e-3


################################################################################
# MIT License
#
# Copyright (c) 2021 Hajime Nakagami<nakagami@gmail.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
################################################################################
using SHA

const SRP_KEY_SIZE = 128
const SRP_SALT_SIZE = 32

function pad(n::BigInt)::Vector{UInt8}
    v::Vector{UInt8} = []
    for x in 1:SRP_KEY_SIZE
        push!(v, UInt8(n & 255))
        n >>= 8
        if n == 0
            break
        end
    end

    reverse!(v)
    v
end

function bigint_to_bytes(v::BigInt)::Vector{UInt8}
    pad(v)
end

function bytes_to_bigint(v::Vector{UInt8})::BigInt
    r::BigInt = 0
    for ui in v
        r = r * 256 + ui
    end
    r
end

function bigint_to_sha1(n::BigInt)::Vector{UInt8}
    v::Vector{UInt8} = []
    while n > 0
        push!(v, UInt8(n & 255))
        n >>= 8
    end
    reverse!(v)
    sha1(v)
end

function get_prime()::Tuple{BigInt, BigInt, BigInt}
    N = big"0xE67D2E994B2F900C3F41F08F5BB2627ED0D49EE1FE767A52EFCD565CD6E768812C3E1E9CE8F0A8BEA6CB13CD29DDEBF7A96D4A93B55D488DF099A15C89DCB0640738EB2CBDD9A8F7BAB561AB1B0DC1C6CDABF303264A08D1BCA932D1F1EE428B619D970F342ABA9A65793B8B2F041AE5364350C16F735F56ECBCA87BD57B29E7"
    g = big"2"
    k = big"1277432915985975349439481660349303019122249719989"
    N, g, k
end

function get_scramble(keyA::BigInt, keyB::BigInt)::BigInt
    # keyA:A client public ephemeral values
    # keyB:B server public ephemeral values
    bytes_to_bigint(sha1(vcat(pad(keyA), pad(keyB))))
end

function get_string_hash(s::AbstractString)::BigInt
    bytes_to_bigint(sha1(s))
end

function get_user_hash(salt::Vector{UInt8}, user::AbstractString, password::AbstractString)::BigInt
    hash1 = sha1(string(user, ":", password))
    hash2 = sha1(vcat(salt, hash1))
    bytes_to_bigint(hash2)
end

function get_client_seed(keya::BigInt)::Tuple{BigInt, BigInt}
    prime, g, _ = get_prime()
    keyA = powermod(g, keya, prime)
    (keyA, keya)
end

function get_client_seed()::Tuple{BigInt, BigInt}
    get_client_seed(BigInt(rand(1:340282366920938463463374607431768211456)))
end

function get_salt()::Vector{UInt8}
    buf::Vector{UInt8} = []
    for _ in 1:SRP_SALT_SIZE
        push!(buf, rand(0:255))
    end
    buf
end

function get_verifier(user::AbstractString, password::AbstractString, salt::Vector{UInt8})::BigInt
    prime, g, _ = get_prime()
    x = get_user_hash(salt, user, password)
    powermod(g, x, prime)
end

function get_server_seed(v::BigInt, keyb::BigInt)::Tuple{BigInt, BigInt}
    prime, g, k = get_prime()
    gb = powermod(g, keyb, prime)
    kv = (k * v) % prime
    keyB = (kv + gb) % prime
    keyB, keyb
end

function get_server_seed(v::BigInt)::Tuple{BigInt, BigInt}
    get_server_seed(v, BigInt(rand(1:340282366920938463463374607431768211456)))
end

function get_client_session(user::AbstractString, password::AbstractString, salt::Vector{UInt8}, keyA::BigInt, keyB::BigInt, keya::BigInt)::Vector{UInt8}
    prime, g, k = get_prime()
    u = get_scramble(keyA, keyB)
    x = get_user_hash(salt, user, password)
    gx = powermod(g, x, prime)
    kgx = (k * gx) % prime
    diff = (keyB - kgx) % prime
    ux = (u * x) % prime
    aux = (keya + ux) % prime
    session_secret = powermod(diff, aux, prime)
    bigint_to_sha1(session_secret)
end

function get_server_session(user::AbstractString, password::AbstractString, salt::Vector{UInt8}, keyA::BigInt, keyB::BigInt, keyb::BigInt)::Vector{UInt8}
    prime, _, _ = get_prime()
    u = get_scramble(keyA, keyB)
    v = get_verifier(user, password, salt)
    vu = powermod(v, u, prime)
    avu = (keyA * vu) % prime
    session_secret = powermod(avu, keyb, prime)
    bigint_to_sha1(session_secret)
end

function get_client_proof(user::AbstractString, password::AbstractString, salt::Vector{UInt8}, keyA::BigInt, keyB::BigInt, keya::BigInt, plugin_name::AbstractString)::Tuple{Vector{UInt8}, Vector{UInt8}}
    prime, g, _ = get_prime()
    keyK = get_client_session(user, password, salt, keyA, keyB, keya)

    n1 = bytes_to_bigint(bigint_to_sha1(prime))
    n2 = bytes_to_bigint(bigint_to_sha1(g))
    n3 = powermod(n1, n2, prime)
    n4 = get_string_hash(user)
    buf = vcat(
        bigint_to_bytes(n3),
        bigint_to_bytes(n4),
        salt,
        bigint_to_bytes(keyA),
        bigint_to_bytes(keyB),
        keyK
    )

    if plugin_name == "Srp"
        keyM = sha1(buf)
    elseif plugin_name == "Srp256"
        keyM = sha256(buf)
    else
        throw(DomainError(plugin_name, "unknown plugin"))
    end

    keyM, keyK
end


@testset "Lifting Line" begin

    section_models = fill(Steady(), 2)
    submodel = Submodel(LiftingLine(section_models))
    submodel_jacobian_tests(submodel)
    submodel_io_tests(submodel)

    section_models = fill(QuasiSteady(), 2)
    submodel = Submodel(LiftingLine(section_models))
    submodel_jacobian_tests(submodel)
    submodel_io_tests(submodel)

    section_models = fill(Wagner(), 2)
    submodel = Submodel(LiftingLine(section_models))
    submodel_jacobian_tests(submodel)
    submodel_io_tests(submodel)

    section_models = fill(Peters{4}(), 2)
    submodel = Submodel(LiftingLine(section_models))
    submodel_jacobian_tests(submodel)
    submodel_io_tests(submodel)

end


using Thebes, Luxor, ColorSchemes

function makemobius()
    x(u, v) = (1 + (v/2 * cos(u/2))) * cos(u)
    y(u, v) = (1 + (v/2 * cos(u/2))) * sin(u)
    z(u, v) = v/2 * sin(u/2)
    # v goes across
    # u goes round
    w = 1
    st = 2π/50
    Δ = .05
    result = []
    for u in 0:st:2π# -st
        for v in -w:Δ:w
            p1 = Point3D(x(u,      v + Δ),   y(u,      v + Δ),     z(u,       v + Δ))
            p2 = Point3D(x(u + st, v + Δ),   y(u + st, v + Δ),     z(u  + st, v + Δ))
            p3 = Point3D(x(u + st, v),       y(u + st, v),         z(u + st,  v))
            p4 = Point3D(x(u,      v),       y(u,      v),         z(u,       v))
            push!(result, [p1, p2, p3, p4])
        end
    end
    return result
end

function frame(scene, framenumber, mobius)
    background("black")
    # flyaround
    theta = rescale(framenumber, 1, scene.framerange.stop, 0, 2pi)
    perspective(400)
    eyepoint(Point3D(250(cos(theta)), 250(sin(theta)), 200))
    setline(.5)
    for (n, p) in enumerate(mobius)
        sethue(get(ColorSchemes.cyclic_wrwbw_40_90_c42_n256_s25, n/length(mobius)))
        pin(100p, gfunction  = (p3l, p2l) -> begin
            poly(p2l, :stroke, close=true)
        end)
    end
end

function  main()
    width, height = (800, 400)
    mobiusmovie = Movie(width, height, "mobiusmovie")

    mb = makemobius()
    gif = animate(mobiusmovie, [
        Scene(mobiusmovie, (s, f) -> frame(s, f, mb), 1:100)
        ],
        creategif=true,
        framerate=20,
        pathname="/tmp/mobiusmovie.gif")
    return gif
end

main()


using Plots
using LLLplus
using Printf
using Random
using LinearAlgebra
using DoubleFloats
using Quadmath

include("lrtest.jl")

# getIntType is used to indicate what type we want the unimodular integer
# matrix to have. We need to add methods of getIntType to handle the
# DoubleFloats and Quadmath, since they are external packages that LLLplus
# doesn't know about. So yes, we do a bit of type piracy.
import LLLplus.getIntType
getIntType(Td::Type{Tr}) where {Tr<:Float128} = Int128
getIntType(Td::Type{Tr}) where {Tr<:Double64} = Int128

lrtest(40,2 .^[7],[100],[Int32,Int64,Int128,Float32,Float64,Float128,Double64,BigInt,BigFloat],"rand")
savefig("perfVsDataType.png")

lrtest(40,2 .^[0:8;],[1],[Float64],"randn")
savefig("perfVsNfloat64.png")


module wilsonmod

include("general.jl")
using Unitful
using NLsolve

∑(x)=sum(x)
λ(T::Quantity,a,b,c,i::Int,j::Int) = a[i,j]u"cal/mol"+ b[i,j]u"cal/(mol*K)"*T + c[i,j]u"cal/(mol*K^2)"*T^2
λ(T::Float64,a,b,c,i::Int,j::Int) = a[i,j]+ b[i,j]*T + c[i,j]*T^2
Λ(T::Quantity,a,b,c,v,i::Int,j::Int) = (v[j]/v[i])*exp(-λ(T,a,b,c,i,j)/(1.98721u"cal/(mol*K)"*T))
Λ(T::Float64,a,b,c,v,i::Int,j::Int) = (v[j]/v[i])*exp(-λ(T,a,b,c,i,j)/(1.98721*T))
γ(T,x,a,b,c,v,i,n)=exp(-log(∑(x[j]*Λ(T,a,b,c,v,i,j) for j=1:n)) + 1 -∑((x[k]*Λ(T,a,b,c,v,k,i))/(∑(x[j]*Λ(T,a,b,c,v,k,j) for j=1:n)) for k=1:n))
P(T,x,CA,a,b,c,v,i,n) = x[i]*γ(T,x,a,b,c,v,i,n)*general.P_Antoine(CA,T,i)
Pₜₒₜ(T,x,CA,a,b,c,v,n) = ∑(P(T,x,CA,a,b,c,v,i,n) for i=1:n)
y_calc(T,x,CA,a,b,c,v,i,n) = P(T,x,CA,a,b,c,v,i,n)/Pₜₒₜ(T,x,CA,a,b,c,v,n)

"""
    pxy(n,T,CA,a,b,c,v,xx;uni=u"Torr")

**Los campos de entrada son:**
- `n :: Int` Es el numero de componentes
- `T :: Float` Es la temperatura del sistema
- `CA :: n×3 Array{Float,2}` Constantes de Antoine para los `n` componentes
- `a :: n×n Array{Float,2}` Parametros de interacción a en ``\\frac{cal}{mol}``
- `b :: n×n Array{Float,2}` Parametros de interacción b en ``\\frac{cal}{mol K}``
- `c :: n×n Array{Float,2}` Parametros de interacción c en ``\\frac{cal}{mol K^2}``
- `v :: n-element Array{Float,1}` Volumen molar
- `xx :: StepRange ó Vector` Son los puntos en el liquido donde se buscara el equilibrio con el vapor
- `uni` Son las unidades en las que se desea el resultado, por defecto esta en Torr

**Salida:**

Regresa las fracciones de liquido `x`,
las fracciones del vapor `y` y la presión `P` en Torr (`m` es la longitud de `xx`).

- `x :: n×m Array{Float64,2}` liquido
- `y :: n×m Array{Float64,2}` vapor
- `P :: m-element Array{Float64,1}` presión
"""
function pxy(n::Int,T,CA::Array,a::Array,b::Array,c::Array,v::Vector,xx;uni=u"Torr")
    T = T|> u"K"
    xx = general.nc(xx...)
    m = size(xx,2)
    y = Array{Float64}(undef,n,m)
    P = Vector{Float64}(undef,m)*uni
    for i=1:m
        x = xx[:,i]
        for j=1:n
            y[j,i] = y_calc(T,x,CA,a,b,c,v,j,n)
        end
        P[i] = Pₜₒₜ(T,x,CA,a,b,c,v,n) |> uni
    end

    return xx,y,P

end

"""
    txy(n,Pobj,CA,a,b,c,v,xx;uni=u"K")

**Los campos de entrada son:**
- `n :: Int` Es el numero de componentes
- `Pobj :: Float` Es la presión del sistema
- `CA :: n×3 Array{Float,2}` Constantes de Antoine para los `n` componentes
- `a :: n×n Array{Float,2}` Parametros de interacción a en ``\\frac{cal}{mol}``
- `b :: n×n Array{Float,2}` Parametros de interacción b en ``\\frac{cal}{mol K}``
- `c :: n×n Array{Float,2}` Parametros de interacción c en ``\\frac{cal}{mol K^2}``
- `v :: n-element Array{Float,1}` Volumen molar
- `xx :: StepRange ó Vector` Son los puntos en el liquido donde se buscara el equilibrio con el vapor
- `uni` Son las unidades en las que se desea el resultado, por defecto esta en grados kelvin

**Salida:**

Regresa las fracciones de liquido `x`,
las fracciones del vapor `y` y la temperatura `T` en grados kelvin (`m` es la longitud de `xx`).

- `x :: n×m Array{Float64,2}` liquido
- `y :: n×m Array{Float64,2}` vapor
- `T :: m-element Array{Float64,1}` temperatura
"""
function txy(n::Int,Pobj,CA::Array,a::Array,b::Array,c::Array,v::Vector,xx;uni=u"K")
    Pobj = Pobj |> u"Torr"
    xx = general.nc(xx...)
    #=
    if n==3
        xx = general.nc_ord(xx...)
    else
        xx = general.nc(xx...)
    end=#
    m = size(xx,2)
    y = Array{Float64}(undef,n,m)
    T = Vector{Float64}(undef,m)*u"K"
    Tᵢ = 300.0
    #temperatura inicial
    for i=1:n
        if xx[i,1]==1
            T[1] = general.T_Antoine(CA,Pobj,i)
            Tᵢ  = ustrip(u"K",T[1])
        end
    end
    for i = 1:m
        x = xx[:,i]
        P2ₜₒₜ(T) = Pₜₒₜ(T,x,CA,a,b,c,v,n) - ustrip(u"Torr",Pobj)
        #println(i)
        Pₛₒₗ = nlsolve(n_ary(P2ₜₒₜ),[Tᵢ])
        Tᵢ = Pₛₒₗ.zero[1]
        T[i] = Tᵢ*u"K"
        for j = 1:n
            y[j,i] = y_calc(T[i],x,CA,a,b,c,v,j,n)
        end
    end

    return xx,y,T .|> uni

end
end  # module


mutable struct GravTreeGather{N,T}
	acc::SVector{N,T}
	pot::T
    mass::T
    #nodecenters::Vector{Vector{T}} #for debugging/visualization
    #nodelengths::Vector{Vector{T}} #for debugging/visualization
end

GravTreeGather{N,T}() where {N,T} = GravTreeGather{N,T}(zero(SVector{N,T}),0,0)

#when changing the argument of treewalk, remember to also change the recursive treewalk call inside the function body!
function gravity_treewalk!(ga::GravTreeGather{N,T}, p::SVector{N,T}, node::Node{N,T,D},
	openingangle::T, softening::T, boxsizes::SVector{N,T}) where {N,T,D}
    if isLeaf(node)
        #println("in a leaf node")
        if node.p != nothing
            #println("nonempty leaf")
            ga.mass += node.p.mass
            dx = nearest.(node.p.pos - p, boxsizes)
			r2 = sum(dx.^2)
			#@show r2, node.p.pos, p
			if r2 > 0.0 #exclude self contribution (r2=0)
				r_inv = sqrt(1.0 / (r2 + softening^2))
				r3_inv = r_inv^3
				ga.acc += node.p.mass * r3_inv * dx
				ga.pot += node.p.mass * r_inv
	            #push!(ga.nodecenters, node.p.pos)
	            #push!(ga.nodelengths, zeros(N))
			end
            #@show "it's a particle..." node.center, node.length
            #r2 = get_distance2(node.part, p, boxsizes, true)
        end
    else
        #println("This is a node... check the opening angle")
		#r2 = get_distance2(pos_c, p, boxsizes, true)
		#don't use get_distance2 as we need both r2 and dx (for vec2pix)
		dx = nearest.(node.n.pos - p, boxsizes)
		r2 = sum(dx.^2)
		if r2 > (node.length[1] / openingangle)^2
			#println("skip node ", i)
			ga.mass += node.n.mass
			r_inv = sqrt(1.0 / (r2 + softening^2))
			r3_inv = r_inv^3
			ga.acc += node.n.mass * r3_inv * dx
			ga.pot += node.n.mass * r_inv
			#push!(ga.nodecenters, node.center)
			#push!(ga.nodelengths, node.length)
			#@show "use this node!" node.center, node.length
		else
	        @inbounds for i in 1:2^N
	            #println("open this node")
	            gravity_treewalk!(ga, p, node.child[i], openingangle, softening, boxsizes)
	        end
		end
    end
end

function gravity_bruteforce!(ga::GravTreeGather{N,T}, p::SVector{N,T}, X::Vector{SVector{N,T}},
							 mass::Vector{T}, boxsizes::SVector{N,T}) where {N,T}
	for i in eachindex(X)
		ga.mass += mass[i]
		dx = nearest.(X[i] - p, boxsizes)
		r2 = sum(dx.^2)
		if r2 > 0.0 #exclude self contribution (r2=0)
			r_inv = sqrt(1.0 / r2)
			r3_inv = r_inv^3
			ga.acc += mass[i] * r3_inv * dx
			ga.pot += mass[i] * r_inv
		end
	end
end


module WNNMDenoise

using LinearAlgebra
using ImageCore
using ImageCore: NumberLike, GenericGrayImage, GenericImage
using ImageDistances
using Statistics
using BlockMatching
using Test
# using Dates # for benchmark

using ImageQualityIndexes # TODO: remove this
using ProgressMeter # TODO: remove this

include("utilities.jl")

export WNNM

"""
    WNNM(noise_level; kwargs...)

Weighted nuclear norm minimization denoising algorithm.

!!! note
    To keep consistent with the original implementation[2], here we stick
    to the 0-255 value range world. If the image is 0-1 range, then the default
    parameters may not work at all.

# Arguments

- `noise_level`: the (estimated) gaussian noise level, used in the first WNNM iteration.

# Keywords

The default values of these keywords come from the original reference
implementation [2]. Except that `patch_size` in our implementation has
to be odd number.

- `K`: the number of WNNM iterations
- `δ`: step value for each WNNM iteration

keywords that control the block matching step, it can be scalar or a vector of length `K`. If it is
a vector then it specifies the values for each WNNM iteration

- `num_patches`: the number of matched patches.
- `patch_size`: patch size.
- `patch_stride`: the stride of each sliding window.
- `window_size`: non-local search size in block matching.

keywords that control the internal WNNM solver, which is a small threshold-svd loop
on the extracted block matching results.

- `λ`: weight constant used to estimate the remained noise level of each patch.
- `C`: weight constant in WNNM solver used by the original implementation.

# Examples

```julia
using TestImages, WNNMDenoise, ImageTransformations

# 1. load a clean image
#    0-255 range is required
clean_img = Float64.(imresize(testimage("cameraman") .* 255, ratio=0.5));

# 2. add guassian noise
σₙ = 40;
noisy_img = clean_img .+ σₙ .* randn(Float32, size(clean_img))

# 3. denoise using the WNNM denoiser
f = WNNM(40)
denoised_img = f(copy(noisy_img), noisy_img)
```

# References

[1] Gu, Shuhang, et al. "Weighted nuclear norm minimization with application to image denoising." _Proceedings of the IEEE conference on computer vision and pattern recognition_. 2014.
[2] The MATLAB reference implementation: http://www4.comp.polyu.edu.hk/~cslzhang/code/WNNM_code.zip

"""
struct WNNM
    noise_level::Float64
    K::Int
    δ::Float64
    num_patches::Vector{Int}
    patch_size::Vector{Int}
    patch_stride::Vector{Int}
    λ::Float64
    C::Float64
    window_size::Int
end

function WNNM(noise_level;
              δ=0.1,
              C=2sqrt(2),
              window_size=60,
              patch_size=nothing,
              num_patches=nothing,
              K=nothing,
              λ=nothing,
              patch_stride=nothing)
    if noise_level <= 20
        isnothing(patch_size) && (patch_size = 5)
        isnothing(num_patches) && (num_patches = 70)
        isnothing(K) && (K = 8)
        isnothing(λ) && (λ = 0.56)
    elseif noise_level <= 40
        isnothing(patch_size) && (patch_size = 7)
        isnothing(num_patches) && (num_patches = 90)
        isnothing(K) && (K = 12)
        isnothing(λ) && (λ = 0.56)
    elseif noise_level <= 60
        isnothing(patch_size) && (patch_size = 7)
        isnothing(num_patches) && (num_patches = 120)
        isnothing(K) && (K = 14)
        isnothing(λ) && (λ = 0.58)
    else
        isnothing(patch_size) && (patch_size = 9)
        isnothing(num_patches) && (num_patches = 140)
        isnothing(K) && (K = 14)
        isnothing(λ) && (λ = 0.58)
    end

    @assert isodd(patch_size) "patch size is expected to be odd number, instead it is $patch_size"
    patch_size isa Number && (patch_size = fill(patch_size, K))
    isnothing(patch_stride) && (patch_stride = @. max(1, patch_size ÷ 2 - 1))
    patch_stride isa Number && (patch_stride = fill(patch_stride, K))

    num_patches = fill(num_patches - 10, K)
    drop_freq = 2
    for k in 2:K
        # drop by 10 for every 2 iteration
        num_patches[k] = (k - 1) % drop_freq == 0 ? num_patches[k - 1] - 10 : num_patches[k - 1]
    end

    WNNM(noise_level, K, δ, num_patches, patch_size, patch_stride, λ, C, window_size)
end

## Implementation

function (f::WNNM)(imgₑₛₜ, imgₙ; clean_img=nothing)
    # FIXME: this if branch is not type stable
    # if imgₑₛₜ === imgₙ
    #     imgₑₛₜ = copy(imgₙ)
    # else
    #     copyto!(imgₑₛₜ, imgₙ)
    # end
    copyto!(imgₑₛₜ, imgₙ)

    T = eltype(eltype(imgₑₛₜ))
    # outfile = open(joinpath("benchmark", "julia_$(Int(f.noise_level)).csv"), "w") # for benchmark
    # println(outfile, "iter,psnr,runtime") # for benchmark
    for iter in 1:f.K
        @. imgₑₛₜ = imgₑₛₜ + f.δ * (imgₙ - imgₑₛₜ) # This iteration can be done more sophisticatedly

        # The noise level for the first iteration is known (whether it is estimated outside or a
        # white noise). The noise is removed in each iteration, so we have to estimate a noise level
        # at a patch level; the denoising performance on each patch can be different, which means a
        # global noise level can be misleading.
        σₚ = iter == 1 ? f.noise_level : zero(f.noise_level)

        # calculating svd using blas threads is not an optimal parallel strategy
        # start_time = now() # for benchmark
        imgₑₛₜ .= with_blas_threads(1) do
            _estimate_img(imgₑₛₜ, imgₙ,
                f.patch_size[iter],
                f.patch_stride[iter],
                f.num_patches[iter],
                f.window_size;
                noise_level=f.noise_level,
                λ=T(f.λ),
                C=T(f.C),
                σₚ=σₚ,
            )
        end
        # duration = now() - start_time # for benchmark

        # TODO: remove this logging part when it is ready
        if !isnothing(clean_img)
            cur_psnr = assess_psnr(clean_img, imgₑₛₜ, 255)
            @info "Result" iter psnr = cur_psnr num_patches = f.num_patches[iter]
            # println(outfile, "$(iter), $(cur_psnr), $(duration.value/1000.)") # for benchmark
            display(Gray.(imgₑₛₜ ./ 255))
            sleep(0.1)
        end
    end
    # close(outfile) # for benchmark
    return imgₑₛₜ
end

function _estimate_img(imgₑₛₜ::AbstractMatrix, imgₙ,
        patch_size::Int,
        patch_stride::Int,
        num_patches::Int,
        window_size::Int;
        kwargs...)
    patch_size = (patch_size, patch_size)

    # We set stride in both dimension instead only in column dimension.
    # This gives less computation but the similar output speaking of PSNR.
    Δ = CartesianIndex((patch_stride, 1)) # original version in MATLAB
    Δ = CartesianIndex((patch_stride, patch_stride))

    r = CartesianIndex(patch_size .÷ 2)
    R = CartesianIndices(imgₑₛₜ)
    R = first(R) +r:Δ:last(R) -r

    imgₑₛₜ⁺ = zeros(eltype(imgₑₛₜ), axes(imgₑₛₜ))
    W = zeros(Int, axes(imgₑₛₜ))

    progress = Progress(length(R[1:patch_stride:end]))
    out_buffer = [Matrix{eltype(imgₑₛₜ)}(undef, prod(patch_size), num_patches) for i in 1:Threads.nthreads()]
    patch_group_buffer = [Matrix{eltype(imgₑₛₜ)}(undef, prod(patch_size), num_patches) for i in 1:Threads.nthreads()]
    patch_q_indices_buffer = [Matrix{CartesianIndex{2}}(undef, prod(patch_size), num_patches) for i in 1:Threads.nthreads()]
    m_buffer = [Vector{eltype(imgₑₛₜ)}(undef, prod(patch_size)) for i in 1:Threads.nthreads()]
    dist_buffer = [Vector{Float64}(undef, (window_size+1)^2) for i in 1:Threads.nthreads()]

    Threads.@threads for p in R
        tid = Threads.threadid()
        out = out_buffer[tid]
        patch_q_indices = patch_q_indices_buffer[tid]
        patch_group = patch_group_buffer[tid]
        m = m_buffer[tid]
        dist = dist_buffer[tid]

        fill!(out, zero(eltype(out)))
        patch_q_indices = _estimate_patch!(
            patch_q_indices, out, patch_group, m, dist,
            imgₑₛₜ, imgₙ, p,
            patch_size,
            num_patches,
            window_size,
            ;kwargs...)

        # Technically, there will be data racing here if multiple threads are involved, but as we've
        # observed, this doesn't affect the overall performance.
        view(W, patch_q_indices) .+= 1
        view(imgₑₛₜ⁺, patch_q_indices) .+= out
        next!(progress)
    end

    return imgₑₛₜ⁺ ./ max.(W, 1)
end

function _estimate_patch!(patch_q_indices, out, patch_group, m, dist,
                          imgₑₛₜ, imgₙ, p,
                          patch_size::Tuple,
                          num_patches::Int,
                          window_size::Int;
                          noise_level,
                          λ,
                          C,
                          σₚ=nothing)
    rₚ = CartesianIndex(patch_size .÷ 2)
    p_indices = p - rₚ:p + rₚ

    alg = FullSearch{Cityblock, 2}(
        Cityblock(),
        rₚ,
        CartesianIndex(window_size÷2, window_size÷2),
        CartesianIndex(1, 1)
    )
    q_inds = multi_match(alg, imgₑₛₜ, imgₑₛₜ, p, dist; num_patches=num_patches)
    # the memory allocation will be a hospot if we directly generate `patch_q_indices` using `vcat`,
    # thus we pre-allocate it and then copyto the buffer.
    @inbounds for i = 1:length(q_inds)
        q = q_inds[i]
        R = q - rₚ:q + rₚ
        copyto!(view(patch_q_indices, :, i), R)
        copyto!(view(patch_group, :, i), view(imgₑₛₜ, R))
    end
    # patch_group .= @view imgₑₛₜ[patch_q_indices]
    mean2!(m, patch_group)

    if σₚ == 0
        # Try: use the mean estimated σₚ of each patch
        σₚ = _estimate_noise_level(view(imgₑₛₜ, p_indices), view(imgₙ, p_indices), noise_level; λ=λ)
    end
    @. out = patch_group - m
    WNNM_optimizer!(out, out, eltype(out)(σₚ); C=C)
    out .+= m

    return patch_q_indices
end

@doc raw"""
    WNNM_optimizer(Y, σₚ; C, rank, fixed_point_num_iters=3)

Optimizes the weighted nuclear norm minimization problem with a fixed point estimation

```math
    \min_X \lVert Y - X \rVert^2_{F} + \lVert X \rVert_{w, *}
```

The weight `w` is specially chosen so that it satisfies the condition of Corollary 1 in [1].

# References

[1] Gu, S., Zhang, L., Zuo, W., & Feng, X. (2014). Weighted nuclear norm minimization with application to image denoising. In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ (pp. 2862-2869).

"""
function WNNM_optimizer!(out, Y, σₚ; C, fixed_point_num_iters=3)
    # Apply Corollary 1 in [1] for image denoise purpose
    # Note: this solver is reserved to the denoising method and is not supposed to be used in other
    #       applications; it simply isn't designed so.

    n = size(Y, 2)
    F = svd!(Y)

    # For image denoising problems, it is natural to shrink large singular value less, i.e., to set
    # smaller weight to large singular value. For this reason, it uses `w = (C * sqrt(n))/(ΣX + eps())`
    # as the weights; inversely propotional to ΣX. With singular values ΣX sorted ascendingly, the
    # condition for Corollary 1 holds, and thus we could directly get the desired solution in a single
    # step.

    # Here we iterate more than once because we don't know what ΣX is; we have to iterate it a while
    # from ΣY to get an relatively good estimation of it.
    # TODO: could we set default σₚ as 0?
    # TODO: is this the best initialization we can get?
    σₚ² = σₚ*σₚ
    nσₚ² = n*σₚ²
    Csnσₚ² = C * sqrt(n) * σₚ²

    ΣX = @. sqrt(max(F.S*F.S - nσₚ², zero(σₚ)))
    for _ in 1:fixed_point_num_iters
        # the iterative algorithm proposed in section 2.2.2 in [1]
        # Step 1 in the iterative algorithm becomes trivial and a no-op

        # Step 2 degenerates to a soft thresholding; both P and Q are identity matrix.
        # all in one line to avoid unnecessary allocation for temporarily variable w
        @. ΣX = soft_threshold(F.S, Csnσₚ² / (ΣX + eps()))
    end

    rmul!(F.U, Diagonal(ΣX))
    mul!(out, F.U, F.Vt)
end

function _estimate_noise_level(patchₑₛₜ, patchₙ, σₙ; λ=0.56)
    # Estimate the noise level of given patch during the WNNM iteration
    # we still need to know the input noisy level σₙ to give an estimation
    λ * sqrt(abs(σₙ^2 - mse(patchₑₛₜ, patchₙ)))
end

end


using Pkg
Pkg.activate(".")
using Soss

using SymPy
import PyCall

stats = PyCall.pyimport_conda("sympy.stats", "sympy")
import_from(stats)

# macro symport(s)
#     str = "sympy.$s"
#     quote
#         $s = PyCall.pyimport_conda($str, "sympy")
#         import_from($s)
#     end
# end

import MacroTools: postwalk, @capture
using MLStyle

using MacroTools: @capture, postwalk, @q

function symlogpdf(expr, ctx)
    @assert @capture(expr, v_ ~ dist_)

    if @capture(dist, d_(dargs__))
        newargs = map(dargs) do arg
            postwalk(arg) do x
                get(ctx, x, x)
            end
        end

        s = ctx[d](ctx[v],newargs)
        s = :(SymPy.density($s).pdf($(ctx[v])) |> log)
        s = @q expand($s,log=true,deep=false,force=true)
        eval(s)

    else
        expr
    end

end

function symlogpdf(m::Model)
    m = canonical(m)
    vs = variables(m)
    ctx = Dict{Symbol, Any}(v => SymPy.symbols(v) for v in vs )

    stochs = []
    for arg in m.body.args
        if @capture(arg, v_=x_)
            merge!(ctx, Dict(v => x ))
        else push!(stochs, arg)
        end
    end

    merge!(ctx, Dict(
        :Normal => (v,args) -> :(SymPy.Normal($v,$(args...)))
      , :HalfNormal => (v,args) -> :(SymPy.Normal($v,$(args...)))
      , :TDist => (v,args) -> :(SymPy.StudentT($v,$(args...)))
      , :Cauchy => (v,args) -> :(SymPy.Cauchy($v,$(args...)))
      , :HalfCauchy => (v,args) -> :(SymPy.Cauchy($v,0,$(args...)))
    #   , :For => (f,xs) => :()
      , :iid => (v, args) -> begin

                :()
            end
    ))

    sum(symlogpdf(expr,ctx) for expr in stochs)
end

a = @model begin
    x ~ Normal(0,1)
    y ~ Normal(0,1)
end

m = @model x,y begin
    α ~ Cauchy(0,5)
    β ~ Normal(0,1)
    σ ~ HalfCauchy(3)
    x0 ~ Normal(0,1)
    ε ~ HalfNormal(0,1)
    x ~ Normal(x0, ε)
    y0 ~ Normal(α + β * x , σ)
    y ~ Normal(y0, ε)
end

symlogpdf(m)


using Random
using StatsBase
using Distributions
using StatsPlots

include("helper_functions.jl")

begin
    λ1 = 1
    λ2 = 1.5
    α = 1
    β = 1
    μ(α) = 1 / (1.5 * (α))
    ν(α) = 1 / (0.7 * (α))
    parameters = (λ1, λ2, α, μ, ν)
    he_parameters = (λ1, λ2, α, μ, ν, β)
    # opponent_parameters = (λ2, λ1, α, μ, ν)
    # opponent_he_parameters = (λ2, λ1, α, μ, ν, β)
    # qs = 0:0.001:1
end

Hyperexponential(parameters, weights) = MixtureModel(Exponential, parameters, weights)
RTDist(p, q, λ1, λ2, α, μ, ν) = Hyperexponential([(1 / (ν(α) - λ1 * p - λ2 * q)), (1 / (μ(α) - λ1 * (1 - p)))], [p, 1 - p])
RT = RTDist(0.57, 0.57, 1, 1, 1, μ, ν)
RT2 = RTDist(0.57, 0.57, 0.5, 0.5, 2, μ, ν)
density(log.(rand(RT, 1000000)), lw = 4, label = "Theoretical Response Distribution")
density!(log.(rand(RT2, 1000000)), lw = 4, label = "Theoretical Response Distribution")

JSs = 0.1:0.01:1
AR_JS_ratio = 1
ARs = AR_JS_ratio ./ JSs

begin
RTs = RTDist.(0.57, 0.57, ARs, ARs, JSs, μ, ν)
plot_mean = plot(JSs,
    mean.(RTs),
    title = "Mean Response Time by Job Size",
    xlabel = "Job Size",
    ylabel = "Mean Response Time",
    # label = "Theoretical Mean",
    size = [400, 300],
    legend = :none)

plot_variance = plot(JSs,
    var.(RTs),
    title = "Variance by Job Size",
    xlabel = "Job Size",
    ylabel = "Variance of Response Time",
    # label = "Theoretical Variance",
    size = [400, 300],
    legend = :none)

mean_var_plot = plot(plot_mean, plot_variance, size = [800, 300], margin = 12pt)
end

savefig("outputs/theoretical_mean_var.pdf")


module ECDSA

using BitConverter
using Secp256k1: Point, KeyPair, Signature, N, G
export KeyPair

KeyPair{:ECDSA}(𝑑::Integer) = 𝑑 ∉ 1:N-1 ? throw(NotInField()) : KeyPair{:ECDSA}(𝑑, 𝑑 * G)

"""
    ECDSA.sign(kp::KeyPair{:ECDSA}, 𝑧::Integer) -> Signature{:ECDSA}

Returns a Signature{:ECDSA} for a given `KeyPair{:ECDSA}` and data `𝑧` and in
which 𝑠 = (𝑧 + 𝑟𝑑) / 𝑘, 𝑘 being a random integer.
"""
function sign(kp::KeyPair{:ECDSA}, 𝑧::Integer)
    𝑘 = rand(big.(0:N))
    𝑟 = (𝑘 * G).𝑥.𝑛
    𝑘⁻¹ = powermod(𝑘, N - 2, N)
    𝑠 = mod((𝑧 + 𝑟 * kp.𝑑) * 𝑘⁻¹, N)
    if 𝑠 > N / 2
        𝑠 = N - 𝑠
    end
    return Signature{:ECDSA}(𝑟, 𝑠)
end

"""
    verify(𝑄::Point, 𝑧::Integer, sig::Signature{:ECDSA}) -> Bool

Returns true if Signature{:ECDSA} is valid for 𝑧 given 𝑄, false if not
"""
function verify(𝑄::Point, 𝑧::Integer, sig::Signature{:ECDSA})
    𝑠⁻¹ = powermod(sig.𝑠, N - 2, N)
    𝑢 = mod(𝑧 * 𝑠⁻¹, N)
    𝑣 = mod(sig.𝑟 * 𝑠⁻¹, N)
    𝑅 = 𝑢 * G + 𝑣 * 𝑄
    return 𝑅.𝑥.𝑛 == sig.𝑟
end

end  # module ECDSA


function calcRowJl!(
    M :: Matrix, A :: Matrix, F :: Matrix, W :: Matrix,
    Out :: Matrix,
    m :: Real, _a :: Real, f :: Real, w :: Real)

    d1, d2 = size(M)

    for j in 1:d2
      for i in 1:d1
        Out[i,j] = pow(M[i, j], m) * pow(A[i, j], _a) * pow(F[i, j], f) * pow(W[i, j], w)
      end
    end
end

pow(x, y) = x ^ y


using BoardGamesStrategies
using Test

include("TicTacToe.jl")

b = playgame(TicTacToeBoard(), MCTSStrategy{TicTacToe}(1000, √2), MCTSStrategy{TicTacToe}(1000, √2));

b = TicTacToeBoard()
b = play(b, 1)
b = play(b, 2)

function heuristic(b::TicTacToeBoard, player::Int64)
    b.winner * -((player - 2) * 2 + 1)
end

@testset "BoardGamesStrategies.jl" begin
    #a bit random but player 1 should win
    @test 0 <= winner(playgame(b, RandomStrategy{TicTacToe}(),
                                  RandomStrategy{TicTacToe}())) <= 2

    @test winner(playgame(b, MCTSStrategy{TicTacToe}(10000, √2),
                             MCTSStrategy{TicTacToe}(10000, √2))) == 1

    @test winner(playgame(b, Minimax{TicTacToe}(10,heuristic),
                             Minimax{TicTacToe}(10,heuristic))) == 1

    @test winner(playgame(b, MinimaxAlphaBeta{TicTacToe}(10, heuristic),
                             MinimaxAlphaBeta{TicTacToe}(10, heuristic))) == 1

end


using LinearElasticity

# Define problem
nels = (60,20,4)
sizes = (1.0,1.0,1.0)
E = 1.0;
ν = 0.3;
force = -1.0;
problem = PointLoadCantilever(nels, sizes, E, ν, force)

# Build element stiffness matrices and force vectors
einfo = ElementFEAInfo(problem);

# Assemble global stiffness matrix and force vector
ginfo = assemble(problem, einfo);

# Solve for node displacements
u = ginfo.K \ ginfo.f

mesh = problem.ch.dh.grid
node_dofs = problem.metadata.node_dofs

node_displacements = reshape(u[node_dofs], 3, JuAFEM.getnnodes(mesh))

using Makie
using GeometryTypes

# Overloads to get the fem nodes working!

function AbstractPlotting.to_vertices(cells::AbstractVector{<: JuAFEM.Node{N, T}}) where {N, T}
    convert_arguments(nothing, cells)[1]
end
function AbstractPlotting.to_gl_indices(cells::AbstractVector{<: JuAFEM.Cell})
    tris = GLTriangle[]
    for cell in cells
        to_triangle(tris, cell)
    end
    tris
end

function to_triangle(tris, cell::JuAFEM.Cell{3, 8, 6})
    nodes = cell.nodes
    push!(tris, GLTriangle(nodes[1], nodes[2], nodes[5]))
    push!(tris, GLTriangle(nodes[5], nodes[2], nodes[6]))

    push!(tris, GLTriangle(nodes[6], nodes[2], nodes[3]))
    push!(tris, GLTriangle(nodes[3], nodes[6], nodes[7]))

    push!(tris, GLTriangle(nodes[7], nodes[8], nodes[3]))
    push!(tris, GLTriangle(nodes[3], nodes[8], nodes[4]))

    push!(tris, GLTriangle(nodes[4], nodes[8], nodes[5]))
    push!(tris, GLTriangle(nodes[5], nodes[4], nodes[1]))

    push!(tris, GLTriangle(nodes[1], nodes[2], nodes[3]))
    push!(tris, GLTriangle(nodes[3], nodes[1], nodes[4]))
end

function AbstractPlotting.convert_arguments(P, x::AbstractVector{<: JuAFEM.Node{N, T}}) where {N, T}
    convert_arguments(P, reinterpret(Point{N, T}, x))
end
#TODO make this work without creating a Node
cnode = Node(zeros(Float32, length(mesh.nodes)))
scene = Makie.mesh(mesh.nodes, mesh.cells, color = cnode, colorrange = (0.0, 33.0), shading = false)
mplot = scene[end]
displacevec = reinterpret(Vec{3, Float64}, node_displacements, (size(node_displacements, 2),))
displace = norm.(displacevec)

new_nodes = broadcast(1:length(mesh.nodes), mesh.nodes) do i, node
    JuAFEM.Node(ntuple(Val(3)) do j
        node.x[j] + node_displacements[j, i]
    end)
end

mesh!(mesh.nodes, mesh.cells, color = (:gray, 0.4))

scatter!(Point3f0.(getfield.(new_nodes, :x)), markersize = 0.1)
# TODO make mplot[1] = new_nodes work

mplot.input_args[1][] = new_nodes
# TODO make mplot[:color] = displace work
push!(cnode, displace)

points = reinterpret(Point{3, Float64}, mesh.nodes)
arrows!(points, displacevec, linecolor = (:black, 0.3))
scene


#   Unit "mdm.jl" of the PosDefManifoldML Package for Julia language

#   MIT License
#   Copyright (c) 2019,
#   Marco Congedo, CNRS, Grenoble, France:
#   https://sites.google.com/site/marcocongedo/home
#   Saloni Jain, Indian Institute of Technology, Kharagpur, India

# ? CONTENTS :
#   This unit implements the Riemannian minimum distance to mean
#   machine learning classifier using package PosDefManifold.

"""
Abstract type for MDM (Minimum Distance to Mean)
machine learning models
"""
abstract type MDMmodel<:PDmodel end

"""
```
mutable struct MDM <: MDMmodel
    metric  :: Metric = Fisher;
    featDim :: Int
    means   :: ℍVector
    imeans  :: ℍVector
end
```

MDM machine learning models are incapsulated in this
mutable structure. MDM models have three fields:
`.metric`, `.featDim` and `.means`.

The field `metric`, of type
[Metric](https://marco-congedo.github.io/PosDefManifold.jl/dev/MainModule/#Metric::Enumerated-type-1),
is to be specified by the user.
It is the metric that will be adopted to compute the class means
and the distances to the mean.

The field `featDim` is the dimension of the manifold in which
the model acts. This is given by ``n(n+1)/2``, where ``n``
is the dimension of the PD matrices.
This field is not to be specified by the user, instead,
it is computed when the MDM model is fit using the [`fit`](@ref)
function and is accessible only thereafter.

The field `means` is an
[ℍVector](https://marco-congedo.github.io/PosDefManifold.jl/dev/MainModule/#%E2%84%8DVector-type-1)
holding the class means, i.e., one mean for each class.
This field is not to be specified by the user, instead,
the means are computed when the MDM model is fitted using the
[`fit`](@ref) function and are accessible only thereafter.

The field `imeans` is an ℍVector holding the inverse of the
matrices in `means`. This also is not to be specified by the user,
is computed when the model is fitted and is accessible only thereafter.
It is used to optimize the computation of distances if the
model is fitted useing the Fisher metric (default).

**Examples**:
```
using PosDefManifoldML, PosDefManifold

# create an empty model
m = MDM(Fisher)

# since the Fisher metric is the default metric,
# this is equivalent to
m = MDM()
```

Note that in general you need to invoke these constructors
only when an MDM model is needed as an argument to a function,
otherwise you can more simply create and fit an MDM model using
the [`fit`](@ref) function.

"""
mutable struct MDM <: MDMmodel
    metric :: Metric
    featDim
    means
    imeans
    function MDM(metric :: Metric = Fisher;
              featDim = nothing,
              means   = nothing,
              imeans  = nothing)
        new(metric, featDim, means)
    end
end

"""
```
function fit(model :: MDMmodel,
              𝐏Tr   :: ℍVector,
              yTr   :: IntVector;
       w        :: Vector = [],
       ✓w       :: Bool  = true,
       meanInit :: Union{ℍVector, Nothing} = nothing,
       tol      :: Real  = 1e-5,
       verbose  :: Bool  = true,
       ⏩       :: Bool  = true)
```

Fit an [`MDM`](@ref) machine learning model,
with training data `𝐏Tr`, of type
[ℍVector](https://marco-congedo.github.io/PosDefManifold.jl/dev/MainModule/#%E2%84%8DVector-type-1),
and corresponding labels `yTr`, of type [IntVector](@ref).
Return the fitted model.

Labels must be provided using the natural numbers, i.e.,
`1` for the first class, `2` for the second class, etc.

Fitting an MDM model involves only computing a mean of all the
matrices in each class. Those class means are computed according
to the metric specified by the [`MDM`](@ref) constructor.

Optional keyword argument `w` is a vector of non-negative weights
associated with the matrices in `𝐏Tr`.
This weights are used to compute the mean for each class.
See method (3) of the [mean](https://marco-congedo.github.io/PosDefManifold.jl/dev/riemannianGeometry/#Statistics.mean)
function for the meaning of the arguments
`w`, `✓w` and `⏩`, to which they are passed.
Keep in mind that here the weights should sum up to 1
separatedly for each class, which is what is ensured by this
function if `✓w` is true.

Optional keyword argument `tol` is the tolerance required for those algorithms
that compute the mean iteratively (they are those adopting the Fisher, logde0
or Wasserstein metric). It defaults to 1e-5. For details on this argument see
the functions that are called for computing the means:
- Fisher metric: [gmean](https://marco-congedo.github.io/PosDefManifold.jl/dev/riemannianGeometry/#PosDefManifold.geometricMean)
- logdet0 metric: [ld0mean](https://marco-congedo.github.io/PosDefManifold.jl/dev/riemannianGeometry/#PosDefManifold.logdet0Mean)
- Wasserstein metric: [Wasmean](https://marco-congedo.github.io/PosDefManifold.jl/dev/riemannianGeometry/#PosDefManifold.wasMean).
For those algorithm an initialization can be provided with optional keyword
argument `meanInit`. If provided, this must be a vector of `Hermitian`
matrices of the [ℍVector](https://marco-congedo.github.io/PosDefManifold.jl/dev/MainModule/#%F0%9D%95%84Vector-type-1)
type and must contain as many initializations as classes, in the
natural order corresponding to the class labels (see above).

If `verbose` is true (default), information is printed in the REPL.
This option is included to allow repeated calls to this function
without crowding the REPL.

**See**: [notation & nomenclature](@ref), [the ℍVector type](@ref).

**See also**: [`predict`](@ref), [`cvAcc`](@ref).

**Examples**
```
using PosDefManifoldML, PosDefManifold

# generate some data
PTr, PTe, yTr, yTe=gen2ClassData(10, 30, 40, 60, 80, 0.25)

# create and fit a model:
m=fit(MDM(Fisher), PTr, yTr)
```

"""
function fit(model :: MDMmodel,
              𝐏Tr   :: ℍVector,
              yTr   :: IntVector;
       w        :: Vector = [],
       ✓w       :: Bool   = true,
       meanInit :: Union{ℍVector, Nothing} = nothing,
       tol      :: Real   = 1e-5,
       verbose  :: Bool   = true,
       ⏩       :: Bool   = true)

    ⌚=now()

    ℳ=deepcopy(model) # output model

    k=length(𝐏Tr) # number of matrices
    !_check_fit(ℳ, k, length(yTr), length(w), 0, "MDM") && return

    verbose && println(greyFont, "Computing class means...")
    z = length(unique(yTr)) # number of classes
    𝐏 = [ℍ[] for i = 1:z]
    W = [Float64[] for i = 1:z]
    for j = 1:k push!(𝐏[yTr[j]], 𝐏Tr[j]) end
    if !isempty(w) for j = 1:k push!(W[yTr[j]], w[j]) end end

    meanInit==nothing ? ℳ.means = ℍVector([getMean(ℳ.metric, 𝐏[i];
                                              w=W[i], ✓w=✓w, tol=tol, ⏩=⏩) for i=1:z]) :
                        ℳ.means = ℍVector([getMean(ℳ.metric, 𝐏[i];
                                             w=W[i], ✓w=✓w, meanInit=meanInit[i], tol=tol, ⏩=⏩) for i=1:z])

    # store the inverse of the means for optimizing distance computations
    # if the metric is Fisher and the matrices are small
    if ℳ.metric==Fisher && size(𝐏Tr[1], 1)<=100
        if ⏩
            ℳ.imeans=ℍVector(undef, length(ℳ.means))
            @threads for i=1:length(ℳ.means) @inbounds ℳ.imeans[i]=inv(ℳ.means[i]) end
        else
            ℳ.imeans=ℍVector([inv(G) for G ∈ ℳ.means])
        end
    end

    ℳ.featDim =_triNum(𝐏Tr[1])

    verbose && println(defaultFont, "Done in ", now()-⌚,".")
    return ℳ
end

"""
```
function predict(model  :: MDMmodel,
                 𝐏Te    :: ℍVector,
                 what   :: Symbol = :labels;
        verbose :: Bool = true,
        ⏩     :: Bool = true)
```
Given an [`MDM`](@ref) `model` trained (fitted) on ``z`` classes
and a testing set of ``k`` positive definite matrices `𝐏Te` of type
[ℍVector](https://marco-congedo.github.io/PosDefManifold.jl/dev/MainModule/#%E2%84%8DVector-type-1),

if `what` is `:labels` or `:l` (default), return
the predicted **class labels** for each matrix in `𝐏Te`,
as an [IntVector](@ref).
For MDM models, the predicted class 'label' of an unlabeled matrix is the
serial number of the class whose mean is the closest to the matrix
(minimum distance to mean).
The labels are '1' for class 1, '2' for class 2, etc;

if `what` is `:probabilities` or `:p`, return the predicted **probabilities**
for each matrix in `𝐏Te` to belong to a all classes, as a ``k``-vector
of ``z`` vectors holding reals in ``[0, 1]`` (probabilities).
The 'probabilities' are obtained passing to a
[softmax function](https://en.wikipedia.org/wiki/Softmax_function)
minus the squared distances of each unlabeled matrix to all class means;

if `what` is `:f` or `:functions`, return the **output function** of the model.
The ratio of the squared distance to all classes to
their geometric mean gives the 'functions'.

If `verbose` is true (default), information is printed in the REPL.
This option is included to allow repeated calls to this function
without crowding the REPL.

It f `⏩` is true (default), the computation of distances is multi-threaded.

**See**: [notation & nomenclature](@ref), [the ℍVector type](@ref).

**See also**: [`fit`](@ref), [`cvAcc`](@ref), [`predictErr`](@ref).

**Examples**
```
using PosDefManifoldML, PosDefManifold

# generate some data
PTr, PTe, yTr, yTe=gen2ClassData(10, 30, 40, 60, 80)

# craete and fit an MDM model
m=fit(MDM(Fisher), PTr, yTr)

# predict labels
yPred=predict(m, PTe, :l)

# prediction error
predErr=predictErr(yTe, yPred)

# predict probabilities
predict(m, PTe, :p)

# output functions
predict(m, PTe, :f)

```
"""
function predict(model  :: MDMmodel,
                 𝐏Te    :: ℍVector,
                 what   :: Symbol = :labels;
            verbose :: Bool = true,
            ⏩     :: Bool = true)

    if !_whatIsValid(what, "predict (MDM model)") return end
    ⌚=now()

    verbose && println(greyFont, "Computing distances...")
    D = getDistances(model.metric, model.means, 𝐏Te; imeans=model.imeans, ⏩=⏩)
    (z, k)=size(D)

    verbose && println("Predicting...")
    if     what == :functions || what == :f
           gmeans=[PosDefManifold.mean(Fisher, D[:, j]) for j = 1:k]
           func(j::Int)=[D[i, j]/gmeans[j] for i=1:z]
           🃏 = [func(j) for j = 1:k]
    elseif what == :labels || what == :l
           🃏 = [findmin(D[:, j])[2] for j = 1:k]
    elseif what == :probabilities || what == :p
           🃏 = [softmax(-D[:, j]) for j = 1:k]
    end

    verbose && println(defaultFont, "Done in ", now()-⌚,".")
    verbose && println(titleFont, "\nPredicted ",_what2Str(what),":", defaultFont)
    return 🃏
end

"""
```
function getMean(metric :: Metric,
                 𝐏      :: ℍVector;
              w        :: Vector = [],
              ✓w       :: Bool    = true,
              meanInit :: Union{ℍ, Nothing} = nothing,
              tol      :: Real   = 0.,
              ⏩      :: Bool    = true)
```

Typically, you will not need this function as it is called by the
[`fit`](@ref) function.

Given a `metric` of type
[Metric](https://marco-congedo.github.io/PosDefManifold.jl/dev/MainModule/#Metric::Enumerated-type-1),
an [ℍVector](https://marco-congedo.github.io/PosDefManifold.jl/dev/MainModule/#%E2%84%8DVector-type-1)
of Hermitian matrices `𝐏` and an optional
non-negative real weights vector `w`,
return the (weighted) mean of the matrices in `𝐏`.
This is used to fit MDM models.

This function calls the appropriate mean functions of package
[PostDefManifold](https://marco-congedo.github.io/PosDefManifold.jl/dev/),
depending on the chosen `metric`,
and check that, if the mean is found by an iterative algorithm,
then the iterative algorithm converges.

See method (3) of the [mean](https://marco-congedo.github.io/PosDefManifold.jl/dev/riemannianGeometry/#Statistics.mean)
function for the meaning of the optional keyword arguments
`w`, `✓w`, `meanInit`, `tol` and `⏩`, to which they are passed.

The returned mean is flagged by Julia as an Hermitian matrix
(see [LinearAlgebra](https://docs.julialang.org/en/v1/stdlib/LinearAlgebra/)).

"""
function getMean(metric :: Metric,
                 𝐏      :: ℍVector;
              w        :: Vector = [],
              ✓w       :: Bool    = true,
              meanInit :: Union{ℍ, Nothing} = nothing,
              tol      :: Real   = 0.,
              ⏩      :: Bool    = true)

  tol==0. ? tolerance = √eps(real(eltype(𝐏[1]))) : tolerance = tol

  if      metric == Fisher
              G, iter, convergence = gMean(𝐏; w=w, ✓w=✓w, init=meanInit, tol=tolerance, ⏩=⏩)
  elseif  metric == logdet0
              G, iter, convergence = ld0Mean(𝐏; w=w, ✓w=✓w, init=meanInit, tol=tolerance, ⏩=⏩)
  elseif  metric == Wasserstein
              G, iter, convergence = wasMean(𝐏; w=w, ✓w=✓w, init=meanInit, tol=tolerance, ⏩=⏩)
  else        G = mean(metric, 𝐏, w=w, ✓w=✓w, ⏩=⏩)
  end

  if metric ∈ (Fisher, logdet0, Wasserstein) && convergence > tolerance
      tolerance == 0. ? toltype="default" : toltype="chosen"
      @error 📌*", getMean function: the iterative algorithm for computing
      the means did not converge using the "*toltype*" tolerance.
      Check your data and try an higher tolerance (with the `tol`=... argument)."
  else
      return G
  end
end

"""
```
function getDistances(metric :: Metric,
                      means  :: ℍVector,
                      𝐏      :: ℍVector;
                imeans  :: Union{ℍVector, Nothing} = false,
                scale   :: Bool = false,
                ⏩      :: Bool = true)
```
Typically, you will not need this function as it is called by the
[`predict`](@ref) function.

Given an [ℍVector](https://marco-congedo.github.io/PosDefManifold.jl/dev/MainModule/#%E2%84%8DVector-type-1)
`𝐏` holding ``k`` Hermitian matrices and
an ℍVector `means` holding ``z`` matrix means,
return the *square of the distance* of each matrix in `𝐏` to the means
in `means`.

The squared distance is computed according to the chosen `metric`, of type
[Metric](https://marco-congedo.github.io/PosDefManifold.jl/dev/MainModule/#Metric::Enumerated-type-1).
See [metrics](https://marco-congedo.github.io/PosDefManifold.jl/dev/introToRiemannianGeometry/#metrics-1)
for details on the supported distance functions.

The computation of distances is optimized for the Fisher metric
if an ℍVector holding the inverse of the means in `means` is passed as
optional keyword argument `imeans`.

If `scale` is true,
the distances are divided by the size of the matrices in `𝐏`.
This is used to compare disctances computed on manifolds with
different dimensions.

If `⏩` is true, the distances are computed using multi-threading,
unless the number of threads Julia is instructed to use is <2 or <3k.

The result is a ``z``x``k`` matrix of squared distances.

"""
function getDistances(metric :: Metric,
             means  :: ℍVector,
             𝐏      :: ℍVector;
          imeans :: Union{ℍVector, Nothing} = false,
          scale  :: Bool = false,
          ⏩    :: Bool = true)

    z, k = length(means), length(𝐏)
    if ⏩
        D = Matrix{eltype(𝐏[1])}(undef, z, k)
        threads, ranges = _GetThreadsAndLinRanges(length(𝐏), "getDistances")

        dist(i::Int, r::Int) =
           if metric==Fisher && imeans≠nothing
               for j in ranges[r] D[i, j]=sum(log.(eigvals(imeans[i]*𝐏[j]; permute=false, scale=false)).^2) end
           else
               for j in ranges[r] D[i, j]=distance²(metric, 𝐏[j], means[i]) end
           end

        for i=1:z @threads for r=1:length(ranges) dist(i, r) end end
    else
        D=[distance²(metric, 𝐏[j], means[i]) for i=1:z, j=1:k]
    end
    # optimize in PosDefManifold, don't need to compute all distances for some metrics
    return scale ? D./size(𝐏[1], 1) : D
end

# ++++++++++++++++++++  Show override  +++++++++++++++++++ # (REPL output)
function Base.show(io::IO, ::MIME{Symbol("text/plain")}, M::MDM)
    if M.means==nothing
        println(io, greyFont, "\n↯ MDM machine learning model")
        println(io, "⭒  ⭒    ⭒       ⭒          ⭒")
        println(io, ".metric : ", string(M.metric), defaultFont)
        println(io, "Unfitted model")
    else
        println(io, titleFont, "\n↯ MDM machine learning model")
        println(io, separatorFont, "⭒  ⭒    ⭒       ⭒          ⭒", defaultFont)
        nc=length(M.means)
        n=size(M.means[1], 1)
        println(io, "type    : PD Manifold model")
        println(io, "features: $(n)x$(n) Hermitian matrices")
        println(io, "classes : $(nc)")
        println(io, separatorFont, "Fields  : ")
        # # #
        println(io, greyFont, " MDM Parametrization", defaultFont)
        println(io, separatorFont," .metric  ", defaultFont, string(M.metric))
        println(io, separatorFont," .means   ", defaultFont, "vector of $(nc) Hermitian matrices")
        println(io, separatorFont," .imeans  ", defaultFont, "vector of $(nc) Hermitian matrices")
        println(io, separatorFont," .featDim ", defaultFont, "$(M.featDim) ($(n)*($(n)+1)/2)")
    end
end


using DynamicHMCModels
#=
using LinearAlgebra
using StaticArrays
using TransformVariables
using LogDensityProblems
using DynamicHMC
using Parameters
using Random
using MCMCChains
=#

function LogLikelihoodMin(ω,dyDep)

 dt = 0.1
 cD = [0 -1im ; 1im 0]
 H = (ω/2.) * cD
 rho=[0.05 0 ; 0 0.95]
 M = I - ((cD'*cD)/2) * dt + (cD * dyDep) - 1im * H * dt
 newRho = M * rho * M'
 lklhood = real(tr(newRho))- (ω* dt/2)^2
 return log(lklhood)

end

struct Experiment
      dyDep::Float64
end

function (problem::Experiment)((ω,)::NamedTuple{(:ω,)})
      @unpack dyDep  = problem        # extract the data
      LogLikelihoodMin(ω,dyDep)
end

dyDepObs=0.690691
p1 = Experiment(dyDepObs)

println(p1((ω=.4,)))

trans_single = as((ω=as(Real, 2, 4),))
P1 = TransformedLogDensity(trans_single, p1)
∇P1 = ADgradient(:ForwardDiff, P1)

# Sample 4 chains

a3d = Array{Float64, 3}(undef, 1000, 1, 4);
for j in 1:4
  global results = mcmc_with_warmup(Random.GLOBAL_RNG, ∇P1, 1000;
    reporter = NoProgressReport())
  global posterior = P1.transformation.(results.chain)

  for i in 1:1000
    a3d[i, 1, j] = values(posterior[i].ω)
  end
end

# Create MCMCChains object

parameter_names = ["ω"]
sections =   Dict(
  :parameters => parameter_names,
)
chns = create_mcmcchains(a3d, parameter_names, sections, start=1)
show(chns)

println()
DynamicHMC.Diagnostics.EBFMI(results.tree_statistics) |> display

println()
DynamicHMC.Diagnostics.summarize_tree_statistics(results.tree_statistics)