# This file is a part of Julia. License is MIT: https://julialang.org/license

using Base.Test
import Base.LinAlg: BlasFloat, BlasComplex, SingularException, A_rdiv_B!, A_rdiv_Bt!,
    A_rdiv_Bc!

n=12 #Size of matrix problem to test
srand(1)

@testset for relty in (Float32, Float64, BigFloat), elty in (relty, Complex{relty})
    d=convert(Vector{elty}, randn(n))
    v=convert(Vector{elty}, randn(n))
    U=convert(Matrix{elty}, randn(n,n))
    if elty <: Complex
        d+=im*convert(Vector{elty}, randn(n))
        v+=im*convert(Vector{elty}, randn(n))
        U+=im*convert(Matrix{elty}, randn(n,n))
    end
    D = Diagonal(d)
    DM = diagm(d)

    @testset "constructor" begin
        for x in (d, GenericArray(d))
            @test Diagonal(x)::Diagonal{elty,typeof(x)} == DM
            @test Diagonal(x).diag === x
            @test Diagonal{elty}(x)::Diagonal{elty,typeof(x)} == DM
            @test Diagonal{elty}(x).diag === x
        end
    end

    @testset "Basic properties" begin
        @test eye(Diagonal{elty},n) == Diagonal(ones(elty,n))
        @test_throws ArgumentError size(D,0)
        @test typeof(convert(Diagonal{Complex64},D)) <: Diagonal{Complex64}
        @test typeof(convert(AbstractMatrix{Complex64},D)) <: Diagonal{Complex64}

        @test Array(real(D)) == real(DM)
        @test Array(abs.(D)) == abs.(DM)
        @test Array(imag(D)) == imag(DM)

        @test parent(D) == d
        @test diag(D) == d
        @test D[1,1] == d[1]
        @test D[1,2] == 0

        @test issymmetric(D)
        @test istriu(D)
        @test istril(D)
        if elty <: Real
            @test ishermitian(D)
        end
    end

    @testset "Simple unary functions" begin
        for op in (-,)
            @test op(D)==op(DM)
        end

        for func in (det, trace)
            @test func(D) ≈ func(DM) atol=n^2*eps(relty)*(1+(elty<:Complex))
        end
        if relty <: BlasFloat
            for func in (expm,)
                @test func(D) ≈ func(DM) atol=n^3*eps(relty)
            end
            @test logm(Diagonal(abs.(D.diag))) ≈ logm(abs.(DM)) atol=n^3*eps(relty)
        end
        if elty <: BlasComplex
            for func in (logdet, sqrtm)
                @test func(D) ≈ func(DM) atol=n^2*eps(relty)*2
            end
        end
    end

    @testset "Linear solve" begin
        let vv = v, UU = U
            @testset for atype in ("Array", "SubArray")
                if atype == "Array"
                    v = vv
                    U = UU
                else
                    v = view(vv, 1:n)
                    U = view(UU, 1:n, 1:2)
                end

                @test D*v ≈ DM*v atol=n*eps(relty)*(1+(elty<:Complex))
                @test D*U ≈ DM*U atol=n^2*eps(relty)*(1+(elty<:Complex))

                @test U.'*D ≈ U.'*Array(D)
                @test U'*D ≈ U'*Array(D)

                if relty != BigFloat
                    atol_two = 2n^2 * eps(relty) * (1 + (elty <: Complex))
                    atol_three = 2n^3 * eps(relty) * (1 + (elty <: Complex))
                    @test D\v ≈ DM\v atol=atol_two
                    @test D\U ≈ DM\U atol=atol_three
                    @test A_ldiv_B!(D, copy(v)) ≈ DM\v atol=atol_two
                    @test At_ldiv_B!(D, copy(v)) ≈ DM\v atol=atol_two
                    @test Ac_ldiv_B!(conj(D), copy(v)) ≈ DM\v atol=atol_two
                    @test A_ldiv_B!(D, copy(U)) ≈ DM\U atol=atol_three
                    @test At_ldiv_B!(D, copy(U)) ≈ DM\U atol=atol_three
                    @test Ac_ldiv_B!(conj(D), copy(U)) ≈ DM\U atol=atol_three
                    Uc = adjoint(U)
                    target = scale!(Uc, inv.(D.diag))
                    @test A_rdiv_B!(Uc, D) ≈ target atol=atol_three
                    @test_throws DimensionMismatch A_rdiv_B!(eye(elty, n-1), D)
                    @test_throws SingularException A_rdiv_B!(Uc, zeros(D))
                    @test A_rdiv_Bt!(Uc, D) ≈ target atol=atol_three
                    @test A_rdiv_Bc!(Uc, conj(D)) ≈ target atol=atol_three
                    @test A_ldiv_B!(D, eye(D)) ≈ D\eye(D) atol=atol_three
                    @test_throws DimensionMismatch A_ldiv_B!(D, ones(elty, n + 1))
                    @test_throws SingularException A_ldiv_B!(Diagonal(zeros(relty, n)), copy(v))
                    b = rand(elty, n, n)
                    b = sparse(b)
                    @test A_ldiv_B!(D, copy(b)) ≈ Array(D)\Array(b)
                    @test_throws SingularException A_ldiv_B!(Diagonal(zeros(elty, n)), copy(b))
                    b = view(rand(elty, n), collect(1:n))
                    b2 = copy(b)
                    c = A_ldiv_B!(D, b)
                    d = Array(D)\b2
                    for i in 1:n
                        @test c[i] ≈ d[i]
                    end
                    @test_throws SingularException A_ldiv_B!(Diagonal(zeros(elty, n)), b)
                    b = rand(elty, n+1, n+1)
                    b = sparse(b)
                    @test_throws DimensionMismatch A_ldiv_B!(D, copy(b))
                    b = view(rand(elty, n+1), collect(1:n+1))
                    @test_throws DimensionMismatch A_ldiv_B!(D, b)
                end
            end
        end
    end
    d = convert(Vector{elty}, randn(n))
    D2 = Diagonal(d)
    DM2= diagm(d)
    @testset "Binary operations" begin
        for op in (+, -, *)
            @test Array(op(D, D2)) ≈ op(DM, DM2)
        end
        @testset "with plain numbers" begin
            a = rand()
            @test Array(a*D) ≈ a*DM
            @test Array(D*a) ≈ DM*a
            @test Array(D/a) ≈ DM/a
            if relty <: BlasFloat
                b = rand(elty,n,n)
                b = sparse(b)
                @test A_mul_B!(copy(D), copy(b)) ≈ Array(D)*Array(b)
                @test At_mul_B!(copy(D), copy(b)) ≈ Array(D).'*Array(b)
                @test Ac_mul_B!(copy(D), copy(b)) ≈ Array(D)'*Array(b)
            end
        end

        #a few missing mults
        bd = Bidiagonal(D2)
        @test D*D2.' ≈ Array(D)*Array(D2).'
        @test D2*D.' ≈ Array(D2)*Array(D).'
        @test D2*D' ≈ Array(D2)*Array(D)'

        #division of two Diagonals
        @test D/D2 ≈ Diagonal(D.diag./D2.diag)
        @test D\D2 ≈ Diagonal(D2.diag./D.diag)

        # Performance specialisations for A*_mul_B!
        vv = similar(v)
        @test (r = full(D) * v   ; A_mul_B!(vv, D, v)  ≈ r ≈ vv)
        @test (r = full(D)' * v  ; Ac_mul_B!(vv, D, v) ≈ r ≈ vv)
        @test (r = full(D).' * v ; At_mul_B!(vv, D, v) ≈ r ≈ vv)

        UU = similar(U)
        @test (r = full(D) * U   ; A_mul_B!(UU, D, U) ≈ r ≈ UU)
        @test (r = full(D)' * U  ; Ac_mul_B!(UU, D, U) ≈ r ≈ UU)
        @test (r = full(D).' * U ; At_mul_B!(UU, D, U) ≈ r ≈ UU)

        # make sure that A_mul_B{c,t}! works with B as a Diagonal
        VV = Array(D)
        DD = copy(D)
        r  = VV * full(D)
        @test Array(A_mul_B!(VV, DD)) ≈ r ≈ Array(D)*Array(D)
        DD = copy(D)
        r  = VV * (Array(D).')
        @test Array(A_mul_Bt!(VV, DD)) ≈ r
        DD = copy(D)
        r  = VV * (Array(D)')
        @test Array(A_mul_Bc!(VV, DD)) ≈ r
    end
    @testset "triu/tril" begin
        @test istriu(D)
        @test istril(D)
        @test triu(D,1)  == zeros(D)
        @test triu(D,0)  == D
        @test triu(D,-1) == D
        @test tril(D,1)  == D
        @test tril(D,-1) == zeros(D)
        @test tril(D,0)  == D
        @test_throws ArgumentError tril(D,n+1)
        @test_throws ArgumentError triu(D,n+1)
    end

    # factorize
    @test factorize(D) == D

    @testset "Eigensystem" begin
        eigD = eigfact(D)
        @test Diagonal(eigD[:values]) ≈ D
        @test eigD[:vectors] == eye(D)
    end

    @testset "ldiv" begin
        v = rand(n + 1)
        @test_throws DimensionMismatch D\v
        v = rand(n)
        @test D\v ≈ DM\v
        V = rand(n + 1, n)
        @test_throws DimensionMismatch D\V
        V = rand(n, n)
        @test D\V ≈ DM\V
    end

    @testset "conj and transpose" begin
        @test transpose(D) == D
        if elty <: BlasComplex
            @test Array(conj(D)) ≈ conj(DM)
            @test adjoint(D) == conj(D)
        end
        # Translates to Ac/t_mul_B, which is specialized after issue 21286
        @test(D' * v == conj(D) * v)
        @test(D.' * v == D * v)
    end

    #logdet
    if relty <: Real
        ld=convert(Vector{relty},rand(n))
        @test logdet(Diagonal(ld)) ≈ logdet(diagm(ld))
    end

    @testset "similar" begin
        @test isa(similar(D), Diagonal{elty})
        @test isa(similar(D, Int), Diagonal{Int})
        @test isa(similar(D, (3,2)), Matrix{elty})
        @test isa(similar(D, Int, (3,2)), Matrix{Int})
    end

    # Issue number 10036
    # make sure issymmetric/ishermitian work for
    # non-real diagonal matrices
    @testset "issymmetric/hermitian for complex Diagonal" begin
        @test issymmetric(D2)
        @test ishermitian(D2)
        if elty <: Complex
            dc = d + im*convert(Vector{elty}, ones(n))
            D3 = Diagonal(dc)
            @test issymmetric(D3)
            @test !ishermitian(D3)
        end
    end

    @testset "svd (#11120/#11247)" begin
        U, s, V = svd(D)
        @test (U*Diagonal(s))*V' ≈ D
        @test svdvals(D) == s
        @test svdfact(D)[:V] == V
    end
end

@testset "svdvals and eigvals (#11120/#11247)" begin
    D = Diagonal(Matrix{Float64}[randn(3,3), randn(2,2)])
    @test sort([svdvals(D)...;], rev = true) ≈ svdvals([D.diag[1] zeros(3,2); zeros(2,3) D.diag[2]])
    @test [eigvals(D)...;] ≈ eigvals([D.diag[1] zeros(3,2); zeros(2,3) D.diag[2]])
end

@testset "isposdef" begin
    @test isposdef(Diagonal(1.0 + rand(n)))
    @test !isposdef(Diagonal(-1.0 * rand(n)))
end

@testset "getindex" begin
    d = randn(n)
    D = Diagonal(d)
    # getindex bounds checking
    @test_throws BoundsError D[0, 0]
    @test_throws BoundsError D[-1, -2]
    @test_throws BoundsError D[n, n + 1]
    @test_throws BoundsError D[n + 1, n]
    @test_throws BoundsError D[n + 1, n + 1]
    # getindex on and off the diagonal
    for i in 1:n, j in 1:n
        @test D[i, j] == (i == j ? d[i] : 0)
    end
end

@testset "setindex!" begin
    d = randn(n)
    D = Diagonal(d)
    # setindex! bounds checking
    @test_throws BoundsError D[0, 0] = 0
    @test_throws BoundsError D[-1 , -2] = 0
    @test_throws BoundsError D[n, n + 1] = 0
    @test_throws BoundsError D[n + 1, n] = 0
    @test_throws BoundsError D[n + 1, n + 1] = 0
    for i in 1:n, j in 1:n
        if i == j
            # setindex on! the diagonal
            @test ((D[i, j] = i) == i; D[i, j] == i)
        else
            # setindex! off the diagonal
            @test ((D[i, j] = 0) == 0; iszero(D[i, j]))
            @test_throws ArgumentError D[i, j] = 1
        end
    end
end

@testset "inverse" begin
    for d in (randn(n), [1, 2, 3], [1im, 2im, 3im])
        D = Diagonal(d)
        @test inv(D) ≈ inv(Array(D))
    end
    @test_throws SingularException inv(Diagonal(zeros(n)))
    @test_throws SingularException inv(Diagonal([0, 1, 2]))
    @test_throws SingularException inv(Diagonal([0im, 1im, 2im]))
end

# allow construct from range
@test all(Diagonal(linspace(1,3,3)) .== Diagonal([1.0,2.0,3.0]))

# Issue 12803
for t in (Float32, Float64, Int, Complex{Float64}, Rational{Int})
    @test Diagonal(Matrix{t}[ones(t, 2, 2), ones(t, 3, 3)])[2,1] == zeros(t, 3, 2)
end

# Issue 15401
@test eye(5) \ Diagonal(ones(5)) == eye(5)

@testset "Triangular and Diagonal" begin
    for T in (LowerTriangular(randn(5,5)), LinAlg.UnitLowerTriangular(randn(5,5)))
        D = Diagonal(randn(5))
        @test T*D   == Array(T)*Array(D)
        @test T'D   == Array(T)'*Array(D)
        @test T.'D  == Array(T).'*Array(D)
        @test D*T'  == Array(D)*Array(T)'
        @test D*T.' == Array(D)*Array(T).'
        @test D*T   == Array(D)*Array(T)
    end
end

let D1 = Diagonal(rand(5)), D2 = Diagonal(rand(5))
    @test_throws MethodError A_mul_B!(D1,D2)
    @test_throws MethodError At_mul_B!(D1,D2)
    @test_throws MethodError Ac_mul_B!(D1,D2)
end

@testset "multiplication of QR Q-factor and Diagonal (#16615 spot test)" begin
    D = Diagonal(randn(5))
    Q = qrfact(randn(5, 5))[:Q]
    @test D * Q' == Array(D) * Q'
    Q = qrfact(randn(5, 5), Val(true))[:Q]
    @test_throws MethodError A_mul_B!(Q, D)
end

@testset "block diagonal matrices" begin
    D = Diagonal([[1 2; 3 4], [1 2; 3 4]])
    Dherm = Diagonal([[1 1+im; 1-im 1], [1 1+im; 1-im 1]])
    Dsym = Diagonal([[1 1+im; 1+im 1], [1 1+im; 1+im 1]])
    @test D' == Diagonal([[1 3; 2 4], [1 3; 2 4]])
    @test D.' == Diagonal([[1 3; 2 4], [1 3; 2 4]])
    @test Dherm' == Dherm
    @test Dherm.' == Diagonal([[1 1-im; 1+im 1], [1 1-im; 1+im 1]])
    @test Dsym' == Diagonal([[1 1-im; 1-im 1], [1 1-im; 1-im 1]])
    @test Dsym.' == Dsym

    v = [[1, 2], [3, 4]]
    @test Dherm' * v == Dherm * v
    @test D.' * v == [[7, 10], [15, 22]]

    @test issymmetric(D) == false
    @test issymmetric(Dherm) == false
    @test issymmetric(Dsym) == true

    @test ishermitian(D) == false
    @test ishermitian(Dherm) == true
    @test ishermitian(Dsym) == false

    @test expm(D) == Diagonal([expm([1 2; 3 4]), expm([1 2; 3 4])])
    @test logm(D) == Diagonal([logm([1 2; 3 4]), logm([1 2; 3 4])])
    @test sqrtm(D) == Diagonal([sqrtm([1 2; 3 4]), sqrtm([1 2; 3 4])])
end

@testset "multiplication with Symmetric/Hermitian" begin
    for T in (Float64, Complex128)
        D = Diagonal(randn(T, n))
        A = randn(T, n, n); A = A'A
        S = Symmetric(A)
        H = Hermitian(A)
        for f in (*, Ac_mul_B, A_mul_Bc, Ac_mul_Bc, At_mul_B, A_mul_Bt, At_mul_Bt)
            @test f(D, S) ≈ f(Matrix(D), Matrix(S))
            @test f(D, H) ≈ f(Matrix(D), Matrix(H))
            @test f(S, D) ≈ f(Matrix(S), Matrix(D))
            @test f(S, H) ≈ f(Matrix(S), Matrix(H))
        end
    end
end

@testset "multiplication of transposes of Diagonal (#22428)" begin
    for T in (Float64, Complex{Float64})
        D = Diagonal(randn(T, 5, 5))
        B = Diagonal(randn(T, 5, 5))
        DD = Diagonal([randn(T, 2, 2), rand(T, 2, 2)])
        BB = Diagonal([randn(T, 2, 2), rand(T, 2, 2)])
        fullDD = copy!(Matrix{Matrix{T}}(2, 2), DD)
        fullBB = copy!(Matrix{Matrix{T}}(2, 2), BB)
        for f in (*, Ac_mul_B, A_mul_Bc, Ac_mul_Bc, At_mul_B, A_mul_Bt, At_mul_Bt)
            @test f(D, B)::typeof(D) == f(Matrix(D), Matrix(B))
            @test f(DD, BB)::typeof(DD) == f(fullDD, fullBB)
        end
    end
end


module Prim

using ..Omega
using ..Omega: Ω, RandVar, URandVar, MaybeRV, ID, lift, uid, elemtype, isconstant
import ..Omega: params, name, ppapl, apl, reify
import Statistics: mean, var, quantile
import ..Causal: ReplaceRandVar
using ..Util
using Spec
import Distributions
const Djl = Distributions
import Base: minimum, maximum
import Random
using DocStringExtensions: SIGNATURES

export  bernoulli,
        betarv,
        β,
        categorical,
        # dirichlet,
        exponential,
        gammarv,
        Γ,
        invgamma,
        kumaraswamy,
        logistic,
        # mvnormal,
        normal,
        poisson,
        rademacher,
        uniform,

        mean

"Primitive random variable of known distribution"
abstract type PrimRandVar <: RandVar end

"Name of a distribution"
function name end

name(t::T) where {T <: PrimRandVar} = T.name.name

"Parameters of `rv`"
@generated function params(rv::PrimRandVar)
  fields = [Expr(:., :rv, QuoteNode(f)) for f in fieldnames(rv) if f !== :id]
  Expr(:tuple, fields...)
end

ppapl(rv::PrimRandVar, ωπ) = rvtransform(rv)(ωπ, reify(ωπ, params(rv))...)

@generated function anysize(args::Union{<:AbstractArray, Real}...)
  isarr = (arg -> arg <: AbstractArray).([args...])
  firstarr = findfirst(isarr)
  if isempty(isarr)
    :(())
  else
    :(size(args[$firstarr]))
  end
end
@spec same(size.(filter(a -> a isa AbstractArray, args)))

include("univariate.jl")      # Univariate Distributions
include("multivariate.jl")    # Multivariate Distributions
include("statistics.jl")      # Distributional properties: mean, variance, etc
include("randx.jl")      # Distributional properties: mean, variance, etc
export  succprob,
        failprob,
        maximum,
        minimum,
        islowerbounded,
        isupperbounded,
        isbounded,
        std,
        median,
        mode,
        modes,

        skewness,
        kurtosis,
        isplatykurtic,
        ismesokurtic,

        isleptokurtic,
        entropy,
        mean,
        samplemean,
        samplemeanᵣ,
        sampleprob,
        sampleprobᵣ,
        prob,
        lprob

# Lifted distributional functions
export  lsuccprob,
        lfailprob,
        lmaximum,
        lminimum,
        lislowerbounded,
        lisupperbounded,
        lisbounded,
        lstd,
        lmedian,
        lmode,
        lmodes,

        lskewness,
        lkurtosis,
        lisplatykurtic,
        lismesokurtic,

        lisleptokurtic,
        lentropy,
        lmean
include("djl.jl")             # Distributions.jl interop

end


using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation
using Plasm
import Base.show

function show(filename)
	V, EV = Plasm.svg2lar(filename)
	#Plasm.view(Plasm.numbering(.075)((V,[[[k] for k=1:size(V,2)], EV])))
	V, EV = Plasm.svg2lar(filename, flag=false)
	#Plasm.view(Plasm.numbering(2)((V,[[[k] for k=1:size(V,2)], EV])))
	return V, EV
end

#show("/Users/paoluzzi/Documents/dev/Plasm.jl/test/svg/new.svg")
#show("/Users/paoluzzi/Documents/dev/Plasm.jl/test/svg/curved.svg")
#show("/Users/paoluzzi/Documents/dev/Plasm.jl/test/svg/twopaths.svg")
#show("/Users/paoluzzi/Documents/dev/Plasm.jl/test/svg/paths.svg")
#show("/Users/paoluzzi/Documents/dev/Plasm.jl/test/svg/boundarytest2.svg")
#show("/Users/paoluzzi/Documents/dev/Plasm.jl/test/svg/tile.svg")
#show("/Users/paoluzzi/Documents/dev/Plasm.jl/test/svg/interior.svg")
#show("/Users/paoluzzi/Documents/dev/Plasm.jl/test/svg/holes.svg")

V,EV = show("/Users/paoluzzi/Documents/dev/Plasm.jl/test/svg/holes.svg")

V = Plasm.normalize(V,flag=true)
Plasm.view(Plasm.numbering(.125)((V,[[[k] for k=1:size(V,2)], EV])))

# subdivision of input edges
W = convert(Lar.Points, V')
cop_EV = Lar.coboundary_0(EV::Lar.Cells)
cop_EW = convert(Lar.ChainOp, cop_EV)
V, copEV = Lar.planar_arrangement_1(W::Lar.Points, cop_EW::Lar.ChainOp)

# compute containment graph of components
bicon_comps = Lar.Arrangement.biconnected_components(copEV)

# visualization of component graphs
EW = Lar.cop2lar(copEV)
W = convert(Lar.Points, V')
hpcs = [ Plasm.lar2hpc(W,[EW[e] for e in comp]) for comp in bicon_comps ]
Plasm.view([ Plasm.color(Plasm.colorkey[(k%12)==0 ? 12 : k%12])(hpcs[k]) for k=1:(length(hpcs)) ])
Plasm.view(Plasm.numbering(.125)((W,[[[k] for k=1:size(W,2)], EW])))

# computation of 2-cells and 2-boundary
V, copEV, copFE = Lar.planar_arrangement_2(V, copEV, bicon_comps)

# visualization of numbered arrangement
Plasm.view( Plasm.numbering1(0.125)((V, copEV, copFE)) )

# final solid visualization
triangulated_faces = Lar.triangulate2D(V, [copEV, copFE])
V = convert(Lar.Points, V')
FVs = convert(Array{Lar.Cells}, triangulated_faces)
Plasm.viewcolor(V::Lar.Points, FVs::Array{Lar.Cells})

EVs = Lar.FV2EVs(copEV, copFE) # polygonal face boundaries
EVs = convert(Array{Array{Array{Int64,1},1},1}, EVs)
Plasm.viewcolor(V::Lar.Points, EVs::Array{Lar.Cells})

model = V,EVs
Plasm.view(Plasm.lar_exploded(model)(1.2,1.2,1.2))


"""
    (fx,ft)=fft3d(z)

Compute spatial (`fx`) and temporal (`ft`) Fourier spectra of field `z`.
"""
        function fft3d(z)

@static        if VERSION < v"0.7.0-DEV.2005"
           function sum(a;dims=1)
              return Base.sum(a,dims);
           end
        end

#       function [fx,fy,ft]=fft3d(ss);
        nss=size(z); ns=nss[1];
        if(length(nss)>=3)
            nt=nss[3]
	else
	    nt=1
        end

        ns2=div(ns,2);
        nt2=div(nt,2);
        if((nt>1)&&(mod(nt,2)!=0))
            z=z[:,:,1:(nt-1)];
            nt=nt-1;
        end

        zf=abs.(fft(z)/(ns*ns*nt)).^2;
        zf0=zf;
        zf[ns2+1,:,:]=zf0[ns2+1,:,:]/2;
        zf[:,ns2+1,:]=zf0[:,ns2+1,:]/2;
        if(nt>1)
          zf[:,:,nt2+1]=zf0[:,:,nt2+1]/2;
        end

@compat        fs=reshape(sum(zf,dims=3),ns,ns);
@compat        ft=reshape(sum(sum(zf,dims=1),dims=2),nt,1);

        fs=fftshift(fs) ;
        fs=fs/nt;

        nn=zeros(ns,1);
        fx=zeros(ns,1);
        for j=-ns2:ns2-1
           for i=-ns2:ns2-1
              k2=sqrt(i*i+j*j);
              ik2=floor(Integer,k2+1.5);
              if(ik2>1)
                fx[ik2]=fx[ik2]+fs[i+ns2+1,j+ns2+1];
                nn[ik2]=nn[ik2]+1;
              end
           end
        end
    if(nt>1)
        ft=ft[2:nt2+1];
    end
   fx=fx[2:ns2+1]./nn[2:ns2+1];
   return (fx,ft,fs)
   end


module TestWiringDiagramLayouts

using Test

using Catlab.Theories, Catlab.WiringDiagrams
using Catlab: Graphics
using Catlab.Graphics.WiringDiagramLayouts: LeftToRight, position

layout_diagram(f; kw...) = Graphics.layout_diagram(f; outer_ports_layout=:fixed, kw...)
box_layouts(d::WiringDiagram) = map(box -> box.value, boxes(d))
box_values(d::WiringDiagram) = map(layout -> layout.value, box_layouts(d))

# Generators
A, B = Ob(FreeSymmetricMonoidalCategory, :A, :B)
f, g = Hom(:f, A, B), Hom(:g, B, A)

d = layout_diagram(f, orientation=LeftToRight)
@test box_values(d) == [f]
fv = first(box_ids(d))
((f_in_x, f_in_y),) = map(position, input_ports(box(d,fv)))
((f_out_x, f_out_y),) = map(position, output_ports(box(d,fv)))
@test f_in_x < f_out_x
@test f_in_y == f_out_y

# Composition
d = layout_diagram(compose(f,g), orientation=LeftToRight)
@test box_values(d) == [f,g]
(fx, fy), (gx, gy) = map(position, box_layouts(d))
@test fx < gx
@test fy == gy

# Products
d = layout_diagram(otimes(f,g), orientation=LeftToRight)
@test box_values(d) == [f,g]
(fx, fy), (gx, gy) = map(position, box_layouts(d))
@test fx == gx
@test fy < gy

# Identities and braidings
@test box_values(layout_diagram(id(A))) == []
@test box_values(layout_diagram(braid(A,B))) == []

# Diagonals and codiagonals
A, B, C = Ob(FreeBiproductCategory, :A, :B, :C)

for expr in (mcopy(A), delete(A), mmerge(A), create(A))
  @test map(l -> l.shape, box_layouts(layout_diagram(expr))) == [:junction]
end

for expr in (mcopy(A⊗B), delete(A⊗B), mmerge(A⊗B), create(A⊗B))
  d = layout_diagram(expr)
  @test map(l -> l.shape, box_layouts(d)) == repeat([:junction], 2)
end

for expr in (mcopy(A⊗B⊗C), delete(A⊗B⊗C), mmerge(A⊗B⊗C), create(A⊗B⊗C))
  d = layout_diagram(expr)
  @test map(l -> l.shape, box_layouts(d)) == repeat([:junction], 3)
end

end


using Aqua
using DSP: conv
using FFTW
using LinearAlgebra
using QSM
using Random
using SparseArrays
using Statistics: mean
using Test

@testset "Unwrap" begin
    include("unwrap/unwrap.jl")
end

@testset "Utils" begin
    include("utils/utils.jl")
end

#@testset "QSM.jl" begin
#    Aqua.test_all(QSM)
#end


push!(LOAD_PATH, joinpath(@__DIR__, "..", ".."))

import ClimaCore.Geometry, LinearAlgebra, UnPack
import ClimaCore:
    Fields,
    Domains,
    Topologies,
    Meshes,
    DataLayouts,
    Operators,
    Geometry,
    Spaces

using OrdinaryDiffEq: OrdinaryDiffEq, ODEProblem, solve, SSPRK33

using Logging: global_logger
using TerminalLoggers: TerminalLogger
global_logger(TerminalLogger())

const FT = Float64

# https://github.com/CliMA/CLIMAParameters.jl/blob/master/src/Planet/planet_parameters.jl#L5
const MSLP = 1e5 # mean sea level pressure
const grav = 9.8 # gravitational constant
const R_d = 287.058 # R dry (gas constant / mol mass dry air)
const γ = 1.4 # heat capacity ratio
const C_p = R_d * γ / (γ - 1) # heat capacity at constant pressure
const C_v = R_d / (γ - 1) # heat capacit at constant volume
const R_m = R_d # moist R, assumed to be dry

const f = 5e-5
const ν = 0.01
const L = 2e2
const nelems = 30
const Cd = ν / (L / nelems)
const ug = 1.0
const vg = 0.0
const d = sqrt(2 * ν / f)
domain = Domains.IntervalDomain(
    Geometry.ZPoint{FT}(0.0),
    Geometry.ZPoint{FT}(L);
    boundary_tags = (:bottom, :top),
)
#mesh = Meshes.IntervalMesh(domain, Meshes.ExponentialStretching(7.5e3); nelems = 30)
mesh = Meshes.IntervalMesh(domain; nelems = nelems)

cspace = Spaces.CenterFiniteDifferenceSpace(mesh)
fspace = Spaces.FaceFiniteDifferenceSpace(cspace)

# https://github.com/CliMA/Thermodynamics.jl/blob/main/src/TemperatureProfiles.jl#L115-L155
# https://clima.github.io/Thermodynamics.jl/dev/TemperatureProfiles/#DecayingTemperatureProfile
function adiabatic_temperature_profile(z; T_surf = 300.0, T_min_ref = 230.0)
    u = FT(ug)
    v = FT(vg)
    return (u = u, v = v)
end

zc = Fields.coordinate_field(cspace)
Yc = adiabatic_temperature_profile.(zc.z)
w = Geometry.Cartesian3Vector.(zeros(Float64, fspace))

Y_init = copy(Yc)
w_init = copy(w)

# Y = (Yc, w)

function tendency!(dY, Y, _, t)

    (Yc, w) = Y.x
    (dYc, dw) = dY.x

    UnPack.@unpack u, v = Yc
    du = dYc.u
    dv = dYc.v

    # S 4.4.1: potential temperature density
    # Mass conservation

    u_1 = parent(u)[1]
    v_1 = parent(v)[1]
    u_wind = sqrt(u_1^2 + v_1^2)
    A = Operators.AdvectionC2C(
        bottom = Operators.SetValue(0.0),
        top = Operators.SetValue(0.0),
    )

    # u-momentum
    bcs_bottom =
        Operators.SetValue(Geometry.Cartesian3Vector(Cd * u_wind * u_1))  # Eq. 4.16
    bcs_top = Operators.SetValue(FT(ug))  # Eq. 4.18
    gradc2f = Operators.GradientC2F(top = bcs_top)
    divf2c = Operators.DivergenceF2C(bottom = bcs_bottom)
    @. du = divf2c(ν * gradc2f(u)) + f * (v - vg) - A(w, u)   # Eq. 4.8

    # v-momentum
    bcs_bottom =
        Operators.SetValue(Geometry.Cartesian3Vector(Cd * u_wind * v_1))  # Eq. 4.17
    bcs_top = Operators.SetValue(FT(vg))  # Eq. 4.19
    gradc2f = Operators.GradientC2F(top = bcs_top)
    divf2c = Operators.DivergenceF2C(bottom = bcs_bottom)
    @. dv = divf2c(ν * gradc2f(v)) - f * (u - ug) - A(w, v)   # Eq. 4.9

    return dY
end

using LinearAlgebra
using RecursiveArrayTools

Y = ArrayPartition(Yc, w)
dY = tendency!(similar(Y), Y, nothing, 0.0)

Δt = 2.0
ndays = 0
# Solve the ODE operator
prob = ODEProblem(tendency!, Y, (0.0, 60 * 60 * 50))
sol = solve(
    prob,
    SSPRK33(),
    dt = Δt,
    saveat = 600, # save 10 min
    progress = true,
    progress_message = (dt, u, p, t) -> t,
);

ENV["GKSwstype"] = "nul"
import Plots
Plots.GRBackend()

dirname = "ekman"
path = joinpath(@__DIR__, "output", dirname)
mkpath(path)

z_centers = parent(Fields.coordinate_field(cspace))
z_faces = parent(Fields.coordinate_field(fspace))

function ekman_plot(u; title = "", size = (1024, 600))
    u_ref =
        ug .-
        exp.(-z_centers / d) .*
        (ug * cos.(z_centers / d) + vg * sin.(z_centers / d))
    sub_plt1 = Plots.plot(
        u_ref,
        z_centers,
        marker = :circle,
        xlabel = "u",
        label = "Ref",
    )
    sub_plt1 =
        Plots.plot!(sub_plt1, parent(u.x[1].u), z_centers, label = "Comp")

    v_ref =
        vg .+
        exp.(-z_centers / d) .*
        (ug * sin.(z_centers / d) - vg * cos.(z_centers / d))
    sub_plt2 = Plots.plot(
        v_ref,
        z_centers,
        marker = :circle,
        xlabel = "v",
        label = "Ref",
    )
    sub_plt2 =
        Plots.plot!(sub_plt2, parent(u.x[1].v), z_centers, label = "Comp")

    return Plots.plot(
        sub_plt1,
        sub_plt2,
        title = title,
        layout = (1, 2),
        size = size,
    )
end

anim = Plots.@animate for (i, u) in enumerate(sol.u)
    ekman_plot(u, title = "Hour $(i)")
end
Plots.mp4(anim, joinpath(path, "ekman.mp4"), fps = 10)

Plots.png(ekman_plot(sol[end]), joinpath(path, "ekman_end.png"))

function linkfig(figpath, alt = "")
    # buildkite-agent upload figpath
    # link figure in logs if we are running on CI
    if get(ENV, "BUILDKITE", "") == "true"
        artifact_url = "artifact://$figpath"
        print("\033]1338;url='$(artifact_url)';alt='$(alt)'\a\n")
    end
end

linkfig("output/$(dirname)/ekman_end.png", "ekman End")


# SimpleTiger Problem
# POMDP, Robust POMDP, and Robust ρPOMDP formulations

struct SimpleTigerRPOMDP <: RPOMDP{Symbol, Symbol, Symbol}
    discount::Float64
    usize::Float64
end

const discount_simpletiger = 0.95
const states_simpletiger = [:tigerleft, :tigerright]
const actions_simpletiger = [:listen, :openleft, :openright]
const observations_simpletiger = [:tigerleft, :tigerright]

SimpleTigerRPOMDP(err) = SimpleTigerRPOMDP(discount_simpletiger, err)
SimpleTigerRPOMDP() = SimpleTigerRPOMDP(discount_simpletiger, 0.025)

states(::SimpleTigerRPOMDP) = states_simpletiger
actions(::SimpleTigerRPOMDP) = actions_simpletiger
observations(::SimpleTigerRPOMDP) = observations_simpletiger

n_states(::SimpleTigerRPOMDP) = 2
n_actions(::SimpleTigerRPOMDP) = 3
n_observations(::SimpleTigerRPOMDP) = 2

state_index(::SimpleTigerRPOMDP, s::Symbol) = s == :tigerleft ? 1 : 2
action_index(::SimpleTigerRPOMDP, a::Symbol) = a == :listen ? 1 : a == :openleft ? 2 : 3
observation_index(::SimpleTigerRPOMDP, z::Symbol) = z == :tigerleft ? 1 : 2
obs_index(prob::SimpleTigerRPOMDP, z::Symbol) = observation_index(prob, z)

initial_state_distribution(::SimpleTigerRPOMDP) = SparseCat([:tigerleft, :tigerright], [0.5, 0.5])
initial_belief(::SimpleTigerRPOMDP) = [0.5, 0.5]

# Nominal transition function distributions
const pϵ_simpletiger = 1e-6
const tdist_simpletiger = [SparseCat(states_simpletiger, [1.0, 0.0]) SparseCat(states_simpletiger, [0.5, 0.5]) SparseCat(states_simpletiger, [0.5, 0.5]);
                    SparseCat(states_simpletiger, [0.0, 1.0]) SparseCat(states_simpletiger, [0.5, 0.5]) SparseCat(states_simpletiger, [0.5, 0.5])]

# Nominal transition function array tarray_simpletiger[s,a,s'] = Pr(s' | s ,a)
const tarray_simpletiger = cat(3, [1.0 0.5 0.5; 0.0 0.5 0.5], [0.0 0.5 0.5; 1.0 0.5 0.5]) # P(tigerleft | s, a), P(tigerright | s, a)

# Robust transitions
# P(z,t|s,a) = P(z|s,a,t)P(t|s,a)
# PztLower = PzLower * PtLower
# PtLower = Plower(t|s,a) = t_lower[s,a,sp]
function transition(prob::SimpleTigerRPOMDP, s::Symbol, a::Symbol)
    si, ai = state_index(prob, s), action_index(prob, a)
    plower = max.(tarray_simpletiger[si,ai,:] - pϵ_simpletiger, 0.0 + pϵ_simpletiger)
    pupper = min.(tarray_simpletiger[si,ai,:] + pϵ_simpletiger, 1.0 - pϵ_simpletiger)
    plower, pupper
end
function transition(prob::SimpleTigerRPOMDP, s::Symbol, a::Symbol, sp::Symbol)
    si, ai, spi = state_index(prob, s), action_index(prob, a), state_index(prob, sp)
    plower = max.(tarray_simpletiger[si, ai, spi] - pϵ_simpletiger, 0.0 + pϵ_simpletiger)
    pupper = min.(tarray_simpletiger[si, ai, spi] + pϵ_simpletiger, 1.0 - pϵ_simpletiger)
    plower, pupper
end

# Nominal observation function distributions
const odist_simpletiger = [SparseCat(observations_simpletiger, [0.85 0.15]) SparseCat(observations_simpletiger, [0.15 0.85]);
                    SparseCat(observations_simpletiger, [0.5 0.5]) SparseCat(observations_simpletiger, [0.5 0.5]);
                    SparseCat(observations_simpletiger, [0.5 0.5]) SparseCat(observations_simpletiger, [0.5 0.5])]

# Nominal observation function array oarray_simpletiger[a,sp,z] = Pr(z|a,sp)
const oarray_simpletiger = cat(3, [0.85 0.15; 0.5 0.5; 0.5 0.5], [0.15 0.85; 0.5 0.5; 0.5 0.5]) # Pr(tigerleft|a,sp), Pr(tigerright|a,sp)

# Robust observations
# P(z,t|s,a) = P(z|s,a,t)P(t|s,a)
# PztLower = PzLower * PtLower
# PzLower = Plower(z|s,a,t) = o_lower[a,sp,z]
function observation(prob::SimpleTigerRPOMDP, a::Symbol, sp::Symbol)
    ai, spi = action_index(prob, a), state_index(prob, sp)
    plower = max.(oarray_simpletiger[ai,spi,:] - prob.usize, 0.0 + pϵ_simpletiger)
    pupper = min.(oarray_simpletiger[ai,spi,:] + prob.usize, 1.0 - pϵ_simpletiger)
    plower, pupper
end
function observation(prob::SimpleTigerRPOMDP, a::Symbol, sp::Symbol, z::Symbol)
    ai, spi, zi = action_index(prob, a), state_index(prob, sp), observation_index(prob, z)
    plower = max.(oarray_simpletiger[ai, spi, zi] - prob.usize, 0.0 + pϵ_simpletiger)
    pupper = min.(oarray_simpletiger[ai, spi, zi] + prob.usize, 1.0 - pϵ_simpletiger)
    plower, pupper
end

function dynamics(prob::SimpleTigerRPOMDP)
    ns = n_states(prob)
    nz = n_observations(prob)
    na = n_actions(prob)
    tarr = states(prob)
    sarr = states(prob)
    aarr = actions(prob)
    zarr = observations(prob)
    p = Array{Float64}(ns,nz,ns,na)
    plower = Array{Float64}(ns,nz,ns,na)
    pupper = Array{Float64}(ns,nz,ns,na)
    for tind in 1:ns, zind in 1:nz, sind in 1:ns, aind in 1:na
        ptl, ptu = transition(prob, sarr[sind], aarr[aind], tarr[tind])
        pzl, pzu = observation(prob, aarr[aind], tarr[tind], zarr[zind])
        plower[tind, zind, sind, aind] = pzl * ptl
        pupper[tind, zind, sind, aind] = pzu * ptu
    end
    plower, pupper
end

function reward(prob::SimpleTigerRPOMDP, s::Symbol, a::Symbol)
    r = 0.0
    if a == :listen
        r = -1.0
    elseif a == :openleft
        r = (s == :tigerleft) ? -100.0 : 10.0
    elseif a == :openright
        r = (s == :tigerright) ? -100.0 : 10.0
    end
    return r
end
reward(prob::SimpleTigerRPOMDP, b::Vector{Float64}, s::Symbol, a::Symbol, sp::Symbol) = reward(prob,s,a)

discount(p::SimpleTigerRPOMDP) = p.discount

# Simulation Functions
function initial_state(prob::SimpleTigerRPOMDP, rng::AbstractRNG)
    d = initial_state_distribution(prob)
    return rand(rng, d)
end

function generate_sor(prob::SimpleTigerRPOMDP, b, s, a, rng::AbstractRNG)
    tdist = SparseCat(states(prob), psample(transition(prob, s, a)...))
    sp = rand(rng, tdist)
    odist = SparseCat(observations(prob), psample(observation(prob, a, sp)...))
    o = rand(rng, odist)
    r = reward(prob, b, s, a, sp)
    sp, o, r
end


using TriangleIntersect
using Base.Test

a = Point(0,0,0)
b = Point(1,0,0)
c = Point(0,1,0)
t = Triangle(a,b,c)
r = Ray(Point(0.1,0.1,1), Point(0.1,0.1,-1))

@test intersect(r, t) == Intersection(Point(0.2,0.2,0.0), 1.02, true)

# yeah, more tests are needed :)


export dice_image

"""
    dice_image(sheet::Image, rows, cols)
Dice (or split) and image made up of many smaller images into multiple separate images
based on the `rows` and `cols` we have provided as arguments. This is useful because
game assets often come as sprite sheets containing multiple images of the same character
of different tiles.
"""
function dice_image(sheet::Image, rows::Integer, cols::Integer)
    sheet_width, sheet_height = size(sheet)
    w, h = sheet_width/cols, sheet_height/rows

    frames = Image[]
    for row in 1:rows
        y0 = row*h
        for col in 1:cols
            x0 = col*w
            push!(frames, Image[x0:x0+w, y0:y0+h])
        end
    end

    frames
end


struct CompatibleMLE{TI<:Integer,TR<:Real,TD<:AbstractDesign} <: Estimator
    design::TD
    sample_space::Array{TI,2}
    estimates::Array{TR,1}
end

string(estimator::CompatibleMLE) = "CompatibleMLE"

function (estimator::CompatibleMLE)(x1::TI, x2::TI, design::TD) where {TI<:Integer,TD<:AbstractDesign}
    for i in 1:size(estimator.sample_space, 1)
        if [x1, x2] == estimator.sample_space[i,:]
            return estimator.estimates[i]
        end
    end
    error("(x1,x2) not found in sample space, valid observation?")
end

function CompatibleMLE(design::TD; verbosity = 0,  lambda = .5, epsilon = 1e-4, smoothmax_scale = 10., max_iter = 10^4) where {TD<:AbstractDesign}

    λ = lambda
    ϵ = epsilon
    @assert (0 <= verbosity <= 5) "Ipopt takes levels between 0 and 5"
    XX = sample_space(design)
    # precompute test decisions and get indices for rejection/non-rejection
    # regions into XX
    decisions       = reject.(XX[:, 1], XX[:, 2], design)
    inds_reject     = findall(decisions .== 1)
    inds_not_reject = findall(decisions .== 0)
    # precompute mles
    mle  = MaximumLikelihoodEstimator()
    mles = mle.(XX[:, 1], XX[:, 2], design)
    nn   = size(XX, 1)
    # define versions for automatic differentiation, ll = loglikelihood
    dbinom(x, n, p)       = gamma(n + 1)/gamma(x + 1)/gamma(n - x + 1)*(1 - p)^(n - x)*p^x
    likelihood(p, x1, x2) = dbinom(x1, n1(design), p) * dbinom(x2, n2(design, x1), p)
    llresidual(p1, p2, i) = likelihood(p1, XX[i,1], XX[i,2]) - likelihood(p2, XX[i,1], XX[i,2])
    smoothmax(x, scale)   = sum( x .* exp.(scale .* x) ./ sum(exp.(scale .* x)) )
    # objective
    function f(phat...)
        [λ*(phat[i] - mles[i])^2 + (1 - λ)*llresidual(phat[i], mles[i], i)^2 for i in 1:length(phat)] |>
        residuals -> smoothmax(residuals, smoothmax_scale)
    end
    # set up JuMP model
    m = Model(
        optimizer_with_attributes(Ipopt.Optimizer,
                              "max_iter" => max_iter,
                       "constr_viol_tol" => ϵ/100,
            "acceptable_constr_viol_tol" => ϵ/100,
                           "print_level" => verbosity
        )
    )
    # define all variables upfront, start with estimates
    @variable(m, 0 <= phat[i = 1:nn] <= 1, start = mles[i])
    # auxiliary variables for compatibility constraint
    @variable(m, a)
    @variable(m, b)
    # define compatibility constraints
    @constraint(m, a_constraint[i = inds_reject], phat[i] >= a)
    @constraint(m, b_constraint[i = inds_not_reject], phat[i] <= b)
    @constraint(m, a >= b + ϵ)
    # define stage-wise consistency constraints
    for i in 1:nn, j in 1:nn
        x1, x2   = XX[i, :]
        x1_, x2_ = XX[j, :]
        # if both are early stopping, and one has more responses we want that
        # it has a larger estimate
        if (n2(design, x1) == 0) & (n2(design, x1_) == 0) & (x1 > x1_)
                mles[i] <= mles[j] ? println([x1, x2]) : nothing
                @constraint(m, phat[i] >= phat[j] + ϵ)
        end
        # we have same x1 (=> same overall sample size),
        # estimates should be orderd by x2
        if (x1 == x1_) & (x2 > x2_)
                mles[i] <= mles[j] ? println([x1, x2]) : nothing
                @constraint(m, phat[i] >= phat[j] + ϵ)
        end
    end
    # define objective
    register(m, :f, nn, f, autodiff = true)
    @NLobjective(m, Min, f(phat...) )
    optimize!(m)
    estimates = value.(phat)
    # build object and return
    return CompatibleMLE{eltype(XX),eltype(estimates),typeof(design)}(
        design, XX, estimates
    )
end


using FractionalDiffEq, Plots
plotlyjs()
(X, Y) = meshgrid(collect(-6:0.2:6), collect(-6:0.2:6))

Z = X.+Y.*im

L = mittleff(0.8, 0.9, Z)
L1=real.(L)
id = findall(x -> x>3, L1)
L1[id].=3
ix = findall(x-> x<-3, L1)
L1[ix].=-3
plot(X, Y, L1, st=:surface)


module QuantumTimeSteppers

using ApproxFun, FastTransforms, Plots, LinearAlgebra, FastGaussQuadrature
import ApproxFun: points, plan_transform, plan_itransform, TransformPlan, ITransformPlan, domain, canonicaldomain, spacescompatible, evaluate
import ApproxFun: Multiplication, ConcreteMultiplication, Derivative, ConcreteDerivative, space, rangespace, bandwidths
import Base: *, first, last, getindex
export HermiteFSE
export Lie_Trotter_evolve, Strang_evolve, Zassenhaus2_evolve

include("splitting_methods.jl")
include("HermiteFSE.jl")

end


module SkewDist

using Distributions, StatsFuns, MultivariateStats, PDMats, Optim, Roots
import Base: length, mean, show, rand, var, cov, *
import Distributions: sampler, _rand!, pdf, _logpdf, _logpdf!, dof, cdf, quantile, minimum, maximum, mgf

export MvSkewNormal, MvSkewTDist, SkewNormal, SkewTDist, fit_MvSkewNormal, fit_MvSkewTDist, marginals, mgf
export pdf, dof, cdf, quantile

include("utils.jl")
include("mvskewnormal.jl")
include("mvskewtdist.jl")
include("skewnormal.jl")
include("skewtdist.jl")
include("operators.jl")
include("fit_mvskewtdist.jl")

end # module


export modifiedMatsuno

"""
    modifiedMatsuno(param;kwargs)

Define an object of type `AbstractModel` in view of solving the initial-value problem for
the modified Matsuno model

# Argument
`param` is of type `NamedTuple` and must contain
- dimensionless parameters `ϵ` (nonlinearity) and `μ` (dispersion);
- numerical parameters to construct the mesh of collocation points as `mesh = Mesh(param)`

## Optional keyword arguments
- `ν`: shallow/deep water multiplication factor. By default, `ν=1` if `μ≦1` and `ν=1/√μ` otherwise. Set the infinite-layer case if `ν=0` (or `μ=Inf`).
- `ktol`: tolerance of the low-pass Krasny filter (default is `0`, i.e. no filtering);
- `dealias`: dealiasing with Orlicz rule `1-dealias/(dealias+2)` (default is `0`, i.e. no dealiasing);
- `verbose`: prints information if `true` (default is `true`).

# Return values
Generate necessary ingredients for solving an initial-value problem via `solve!` and in particular
1. a function `modifiedMatsuno.f!` to be called in the time-integration solver;
2. a function `modifiedMatsuno.mapto` which from `(η,v)` of type `InitialData` provides the raw data matrix on which computations are to be executed;
3. a function `modifiedMatsuno.mapfro` which from such data matrix returns the Tuple of real vectors `(η,v)`, where
    - `η` is the surface deformation;
    - `v` is the derivative of the trace of the velocity potential.

"""
mutable struct modifiedMatsuno <: AbstractModel

    label   :: String
	f!		:: Function
	mapto	:: Function
	mapfro	:: Function
	param	:: NamedTuple
	kwargs	:: NamedTuple

    function modifiedMatsuno(param::NamedTuple;
							ν=nothing,ktol=0,dealias=0, verbose=true)

		label = "modified Matsuno"

		μ 	= param.μ
		ϵ 	= param.ϵ
		if ν == nothing
			if μ > 1
				ν = 1/sqrt(μ)
			else
				ν = 1
			end
		end

		mesh = Mesh(param)

		param = ( ϵ = ϵ, μ = μ, xmin = mesh.xmin, xmax = mesh.xmax, N = mesh.N )
		kwargs = (ν=ν,dealias=dealias,ktol=ktol,verbose=verbose)

		x = mesh.x
		k = copy(mesh.k)
		K = mesh.kmax * (1-dealias/(2+dealias))
		Π⅔ 	= abs.(mesh.k) .<= K # Dealiasing low-pass filter
		if dealias == 0
			if verbose @info "no dealiasing" end
			Π⅔ 	= ones(size(mesh.k))
		elseif verbose
			@info string("dealiasing : spectral scheme for power ", dealias + 1," nonlinearity ")
		end
		if μ == Inf || ν==0
			∂ₓF₀ 	= 1im * sign.(mesh.k)
			G₀ 	= abs.(mesh.k)
			μ = 1
			ν = 1
		else
        	∂ₓF₀ 	= 1im* sign.(mesh.k) .* tanh.(sqrt(μ)*abs.(mesh.k))
			G₀ 	= sqrt(μ)*abs.(mesh.k).*tanh.(sqrt(μ)*abs.(mesh.k))
		end
		∂ₓ	=  1im * sqrt(μ)* mesh.k            # Differentiation
		z = zeros(Complex{Float64}, mesh.N)
		η = copy(z) ; v = copy(z) ; F = copy(z) ;
		fftη = copy(z) ; fftv = copy(z) ; Q = copy(z) ; R = copy(z) ;

		# Discrete Fourier transform with, possibly, dealiasing and Krasny filter.
		function mapto(data::InitialData)
			U = [Π⅔ .* fft(data.η(x)) Π⅔ .*fft(data.v(x))]
			U[abs.(U).< ktol ].=0
			return U
		end

		# Return `(η,v)`, where
		# - `η` is the surface deformation;
		# - `v` is the derivative of the trace of the velocity potential.
		# Inverse Fourier transform and takes the real part.
		function mapfro(U)
			real(ifft(U[:,1])),real(ifft(U[:,2]))
		end

		# Evolution equations are ∂t U = f!(U)
		function f!(U)

			fftη .= U[:,1]
		    fftv .= U[:,2]
			η  .= ifft(U[:,1])
			v  .= ifft(U[:,2])
			Q .= -∂ₓF₀.*fftv
			Q += -ϵ *∂ₓ.*fft(η.*ifft(fftv)) .+ ϵ*∂ₓF₀.*fft(η.*ifft(G₀.*fftv))
			F .= exp.(-ϵ*ifft(G₀.*fftη))
			R .= -fft(F.*ifft(∂ₓ.*fftη))*ν
			R += ϵ/2*∂ₓ.*fft(-v.^2)

			U[:,1] .= Π⅔.*Q/sqrt(μ)/ν
		   	U[:,2] .= Π⅔.*R/sqrt(μ)/ν
			U[abs.(U).< ktol ].=0

		end

        new(label, f!, mapto, mapfro, param, kwargs )
    end
end


## Wetted-area method
#======================================================#

# Raymer form factor for wing, pylon, nacelle, fuselage? (Eq. 12.30)
form_factor_wing(x_c, t_c, Λ_m, M) = (1 + 0.6t_c / x_c + 100t_c^4) * (1.34M^0.18 * cos(Λ_m)^0.28)

function form_factor(wing :: HalfWing, M)
    # (x/c)_max, (t/c)_max, (t/c)_max sweep angles
    xcs, tcs, sweeps = max_thickness_to_chord_ratio_sweeps(wing, 60)
    Kf  = form_factor_wing.(xcs, tcs, sweeps, M)
end

form_factor(wing :: Wing, M) = (form_factor(wing.left, M) + form_factor(wing.right, M)) / 2

# Schlichting averaged skin-friction coefficients
cf_lam(Re_c, k_lam = 1.) = 1.328 / √(Re_c * k_lam)
cf_turb(Re_c, M) = 0.455 / log10(Re_c)^2.58 / (1 + 0.144M^2)^0.65
cf_schlichting(Re, Re_xtr, M) = max(cf_lam(Re), cf_turb(Re, M) - (Re_xtr / 320 - 39) / Re)

function skin_friction_coefficients(mean_chords, x_tr, V, ρ, M, μ)
    # Skin-friction coefficients
    Re_c    = reynolds_number.(ρ, V, mean_chords, μ)
    Re_xtr  = reynolds_number.(ρ, V, mean_chords .* x_tr, μ)
    cfs     = cf_schlichting.(Re_c, Re_xtr, M)
end

function wetted_area_drag(mean_chords, S_wets, K_fs, x_tr, V, ρ, M, μ)
    cfs = skin_friction_coefficients(mean_chords, x_tr, V, ρ, M, μ)

    # Profile drag
    Dp_by_q = sum(@. cfs * S_wets * K_fs)
end

function wetted_area_drag(wing :: HalfWing, x_tr, V, ρ, a_ref = 330., μ = 1.5e-5)
    # Chord processing
    mean_chords = (forward_sum ∘ chords)(wing) / 2

    # Wetted areas
    S_wets  = @. mean_chords * wing.spans / cos(wing.dihedrals)

    # Form factors
    M       = V / a_ref
    K_fs    = form_factor(wing, M)

    wetted_area_drag(mean_chords, S_wets, K_fs, x_tr, V, ρ, M, μ)
end

function wetted_area_drag(wing :: WingMesh, x_tr, V, ρ, a_ref = 330., μ = 1.5e-5)
    # Chord processing
    mean_chords = (forward_sum ∘ chords)(wing.surf) / 2

    # Wetted areas
    surf_pans = camber_panels(wing)
    S_wets    = sum(panel_area, surf_pans, dims = 1)

    # Form factors
    M       = V / a_ref
    K_fs    = form_factor(wing.surf, M)

    wetted_area_drag(mean_chords, S_wets, K_fs, x_tr, V, ρ, M, μ)
end

# Sato's local-friction and local-dissipation based on power balance method from Mark Drela, Flight Vehicle Aerodynamics, eq. 4.115.
function local_dissipation_drag(wing :: Wing, wetted_areas, ρ_es, u_es, x_tr, V, ρ, M, μ)
    # Chord processing
    mean_chords = (forward_sum ∘ chords)(wing) / 2

    # Compute weighted wetted areas based on inviscid edge velocity distribution.
    weighted_S_wets = sum(@. ρ_es * norm(u_es)^3 * wetted_areas; dims = 1) ./ (ρ * V^3)

    wetted_area_drag(mean_chords, weighted_S_wets, 1., x_tr, V, ρ, M, μ)
end

function local_dissipation_drag(wing :: WingMesh, ρ_es, u_es, x_tr, V, ρ, M, μ)
    # Chord processing
    mean_chords = (forward_sum ∘ chords)(wing.surf) / 2

    # Compute weighted wetted areas based on inviscid edge velocity distribution.
    S_wets = panel_area.(camber_panels(wing))
    weighted_S_wets = sum(@. ρ_es * u_es^3 * S_wets; dims = 1) ./ (ρ * V^3)

    wetted_area_drag(mean_chords, weighted_S_wets, 1., x_tr, V, ρ, M, μ)
end

abstract type AbstractProfileDrag end

struct FormFactor <: AbstractProfileDrag end
struct Dissipation <: AbstractProfileDrag end

profile_drag_coefficient(wing :: HalfWing, x_tr, V, rho_ref, a_ref, area_ref, μ) = wetted_area_drag(wing, x_tr, V, rho_ref, a_ref, μ) / area_ref
profile_drag_coefficient(wing :: WingMesh, x_tr, V, rho_ref, a_ref, area_ref, μ) = wetted_area_drag(wing, x_tr, V, rho_ref, a_ref, μ) / area_ref
profile_drag_coefficient(wing :: Wing, x_tr, V, rho_ref, a_ref, area_ref, μ) = profile_drag_coefficient(wing.left, x_tr, V, rho_ref, a_ref, area_ref, μ) + profile_drag_coefficient(wing.right, x_tr, V, rho_ref, a_ref, area_ref, μ)
profile_drag_coefficient(wing :: AbstractWing, x_tr, refs :: References) = profile_drag_coefficient(wing, x_tr, refs.speed, refs.density, refs.sound_speed, refs.area, refs.viscosity)

function wave_drag(M, Λ, t_by_c, Cl, κ_A)
    M_drag_divergence = κ_A / cos(Λ) - t_by_c / cos(Λ)^2 - Cl / (10 * cos(Λ)^3) # Drag divergence Mach number
    M_crit   = M_drag_divergnece - (0.1 / 80)^(1/3) # Critical Mach number
    CD_wave = ifelse(M > M_crit, 20 * (Mach - M_crit)^4, 0.)
end


module mvib1
#using DelimitedFiles
using FinEtools
using FinEtoolsDeforLinear
using FinEtoolsDeforLinear.AlgoDeforLinearModule
using FinEtoolsVibInFluids.LaplBEM
using LinearAlgebra
using Arpack
using Test

E = 206000*phun("MPa");# Young's modulus
nu = 0.3;# Poisson ratio
rho = 7850*phun("KG*M^-3");# mass density
Length= 10e3*phun("MM"); Width= 10e3*phun("MM"); Height= 0.238e3*phun("MM");
rhow = 1000*phun("KG*M^-3");
# E = 206000.0;# Young's modulus
# nu = 0.3;# Poisson ratio
# rho = 7.850e-9;# mass density
# Length= 10.0e3; Width= 10.0e3; Height= 0.238e3;
# rhow = 1.000e-9;
nHeight = 4; nLength  = 4*3; nWidth = nLength;
tolerance = Height/nHeight/100;
OmegaShift = (2*pi*0.0) ^ 2; # to resolve rigid body modes
neigvs = 5;

function plate_free_vibration_solver()
	# Free-vibration solution for a clamped plate
	# Reference: Fu, Y., and Price, W. G., 1987, “Interactions Between a Partially
	# or Totally Immersed Vibrating Cantilever Plate and the Surrounding Fluid,”
	# J. Sound Vib., 118(3), pp. 495–513.
	# Dry Natural frequencies [radians per second] (nLength=10,nWidth=10)
	# 12.45153      29.44122      75.04792      94.26498      106.7881

	Mshift =0;

    MR = DeforModelRed3D
    fens,fes  = T4block(Length, Width, Height, nLength, nWidth, nHeight)

    geom = NodalField(fens.xyz)
    u = NodalField(zeros(size(fens.xyz,1),3)) # displacement field
    nl = selectnode(fens, box=[0.0 0.0 -Inf Inf -Inf Inf], inflate=tolerance)
    setebc!(u, nl, true, collect(1:3))
    applyebc!(u)
    numberdofs!(u)

    material = MatDeforElastIso(MR, rho, E, nu, 0.0)

    femm = FEMMDeforLinearESNICET4(MR, IntegDomain(fes, NodalSimplexRule(3)), material)
    associategeometry!(femm,  geom)
    K = stiffness(femm, geom, u)
    M = mass(femm, geom, u)

    d, v, nconv = eigs(Symmetric(K+OmegaShift*M), Symmetric(M); nev=neigvs, which=:SM,  explicittransform=:none)
    d = d .- OmegaShift;
    d = real.(d)
    v = real.(v)
    #for i in 1:length(d)
    #    @show v[:, i]' * K * v[:, i]
    #end
    #for i in 1:length(d)
    #    @show v[:, i]' * M * v[:, i]
    #end
    #open("d.txt", "w") do io
    #        writedlm(io, d)
    #    end
    #    open("v.txt", "w") do io
    #        writedlm(io, v)
    #    end

    return Dict("fens"=>fens, "femm"=>femm, "geom"=>geom, "u"=>u, "eigenvectors"=>v, "eigenvalues"=>d)
end # plate_free_vibration_solver

function plate_completely_wet()
	# Reference 7.35, 20.2, 50.45, 70.41, 78.85 [angular frequency, radians per second]

	model = plate_free_vibration_solver()

	fens = model["fens"]
	femm = model["femm"]
	d = model["eigenvalues"]
	v = model["eigenvectors"]
	u = model["u"]

	neigvs = size(v, 2)

	fes = femm.integdomain.fes

	bfes = meshboundary(fes)
	tolerance = Height/1000;
	rl = vcat(selectelem(fens, bfes, box = [-Inf Inf -Inf Inf Height Height], inflate = tolerance),
		selectelem(fens, bfes, box = [-Inf Inf -Inf Inf 0.0 0.0], inflate = tolerance));
	# These are the wet surface panels
	wbfes = subset(bfes, rl);
	nnperel = nodesperelem(wbfes)

	pnormals = LaplBEM.allnormals(LaplBEM.allelxs(fens.xyz, connasarray(wbfes)))
	vn = fill(0.0, count(wbfes), neigvs)
	uk = fill(0.0, 3)
	for mode in 1:neigvs
		for i in 1:count(wbfes)
			vn[i, mode] = 0.0
			for k in wbfes.conn[i]
				for m in 1:3
					dof = u.dofnums[k, m]
					uk[m] = (dof != 0 ? v[dof, mode] : 0.0)
				end
				vn[i, mode] += dot(pnormals[i], uk)
			end
			vn[i, mode] /= nnperel
		end
	end

	# scalars = []
	# for mode in 1:neigvs
	# 	push!(scalars, ("vn$mode", deepcopy(vn[:, mode])))
	# end
	# File =  "plate_wet-vn.vtk"
	# vtkexportmesh(File, fens, wbfes; scalars=scalars)
	# @async run(`"paraview.exe" $File`)

	A = fill(0.0, count(wbfes), count(wbfes))
	A = LaplBEM.doublelayer!(A, fens.xyz, connasarray(wbfes), TriRule(3))
	B = fill(0.0, count(wbfes), count(wbfes))
	B = LaplBEM.singlelayer!(B, fens.xyz, connasarray(wbfes), TriRule(3), TriRule(3))
    #open("A.txt", "w") do io
    #    writedlm(io, A)
    #end
    #open("B.txt", "w") do io
    #    writedlm(io, B)
    #end
    #open("vn.txt", "w") do io
    #    writedlm(io, vn)
    #end
	pp = A \ (B * vn)

	areas = fill(0.0, count(wbfes))
	n = fill(0.0, 3)
	for i in 1:count(wbfes)
		ix = collect(wbfes.conn[i])
		LaplBEM.trinml!(n, @view fens.xyz[ix, :])
		areas[i] = norm(n) / 2
	end

	rK = diagm(vec(d))
	rM = I
	raM = fill(0.0, neigvs, neigvs);
	for i in 1:neigvs
	    for j in 1:neigvs
	        raM[i,j] = rhow*dot(vn[:,i].*areas, pp[:,j]);
	    end
	end
	raM = (raM +raM')/2;# it should be symmetric, but make sure

    #@show rK, rM, raM
	decomp = eigen(rK, rM+raM)
	wv = v*decomp.vectors;
	return sqrt.(decomp.values)
end # plate_completely_wet

function test()
	angular_frequencies = plate_completely_wet()
	@test norm(angular_frequencies - [8.64673912672756, 22.463163136168948, 54.24478379306806, 72.16460600046824, 81.89085075349693]) <= 1.0e-5
	true
end
end
using .mvib1
mvib1.test()


include("ImageCreator.jl")
include("dipole.jl")

const c = 299792458.
const mu0 = 4*pi*1e-7
const eps0 = 1/(mu0*c^2)

#cavity
L=5
l=4
h=3
losses= 0.998
order=10

const freq=.5e9

##dipole
#dipole position
X=1.2
Y=1
Z=1.3
tilt=pi/2-acos(sqrt(2/3));
azimut=pi/4
phase=0
#dipole moment
#total time averaged radiated power P= 1 W dipole moment => |p|=sqrt(12πcP/µOω⁴)
Pow=1
amplitude=sqrt(12*pi*c*Pow/(mu0*(2*pi*freq)^4))

#Images
POS=IC(L,l,h,X,Y,Z,tilt,azimut,phase,amplitude,order)

#using HDF5,JLD
#@load "POS.jld" POS

numberofimages=1+2*order+(2*order*(order+1)*(2*order+1))/3
POS=POS[1:numberofimages,:]

#observation points
dx=0.02
dy=0.02
dz=0.02
const  nx=int(L/dx)
const  xmax=L
const  ny=int(h/dy)
const  ymax=l
const  x=linspace(0,L,nx)
const  z=1
const  y=linspace(0,l,ny)

println("Computing the radiation...")
Z=zeros(nx,ny)
Et=zeros(nx,ny)
Bt=zeros(nx,ny)
for i=1:nx
  perc=round(i/nx*1000)/10
  println("$perc %")
  for j=1:ny
    r=[x[i],y[j],z] #position of the receiver
    E=zeros(Complex128,3)
    B=zeros(Complex128,3)
    for m=1:numberofimages
      p=vec(POS[m,1:3]) #image dipole moment
      R=vec(POS[m,4:6]) #image dipole position
      ord=POS[m,7] #order of the dipole
      Ed,Bd=Hertz_dipole (r, p*losses^ord, R, phase, freq) #fields computation
      E+=Ed
      B+=Bd
    end
    Z[i,j]=sqrt(sum(abs(E).^2))/sqrt(sum(abs(B).^2))*mu0 #real(E).^2#0.5*numpy.cross(E.T,conjugate(B.T)) #impedance
    Et[i,j]=sqrt(sum(abs(E).^2)) #total E-field
    Bt[i,j]=sqrt(sum(abs(B).^2)) #total B-field
  end
end

#Some figures
using PyPlot

figure(1) #total E-field
title("\$\\vert E\\vert\$ (V/m)")
pcolor(x,y,Et',cmap="jet")
#clim(0,20)
colorbar()
axis("scaled")
xlim(0,L)
ylim(0,l)
grid()
xlabel("\$x\$/m")
ylabel("\$y\$/m")
savefig("E_$order.png",bbox="tight")
clf()

figure(2) #total H-field
title("\$\\vert H\\vert\$ (A/m)")
pcolor(x,y,Bt'/mu0,cmap="jet")
#clim(0,.2)
colorbar()
axis("scaled")
xlim(0,L)
ylim(0,l)
grid()
xlabel("x/m")
ylabel("\$y\$/m")
savefig("H_$order.png",bbox="tight")
clf()

figure(3) #wave impedance
title("\$Z\$ (\$ \\Omega\$)")
pcolor(x,y,Z',cmap="jet")
clim(0,1000)
colorbar(ticks=[0,100,200,300,120*pi,500,60*pi^2,700,800,900,1000])
axis("scaled")
xlim(0,L)
ylim(0,l)
grid()
xlabel("x/m")
ylabel("\$y\$/m")
savefig("Z_$order.png",bbox="tight")
clf()


using DynamicalSystems, TimeseriesSurrogates

D = 7
lo = Systems.lorenz96(D, range(0; length = D, step = 0.1); F = 8.0)
X = regularize(trajectory(lo, 1000, dt = 0.1, Ttr = 100.0))

e = 10.0 .^ range(-4, 1, length = 22)
CX = correlationsum(X, e; w = 5)

figure()
le = log10.(e)
plot(le, log10.(CX))

i = findfirst(z -> z > 0, CX)

@show linear_region(le[i:end], log10.(CX)[i:end])

sg = surrogenerator(X, ShuffleDimensions())
for i in 1:10
    Z = sg()
    CZ = correlationsum(Z, e)
    plot(le, log10.(CZ), alpha = 0.5, lw = 0.5, color = "k", ls = ":")
    @show linear_region(le, log10.(CZ))[2]
end
xlabel("log(e)")
ylabel("log(C)")


# Julia code to calculate spectrum of tridiagonal 'DoS' TB matrices with Sturm sequences
# See section 5, p38: http://arxiv.org/pdf/math-ph/0510043v2.pdf

module Sturm

println("Loading Jarv's Sturm Library...")

export sturm,randexp,randH

# Calculates number of eigenvalues less than 'sigma' in tridiagonal matrix
# described by: diagm(E.^0.5,-1)+diagm(D)+diagm(E.^0.5,1)
# Nb: Off digonal elements E are supplied elementally squared for computational speed
function sturm(D,E_squared,sigma)
    t=0.0
    countnegatives=0

    t=D[1]-sigma #first term of sequence calculated differently, to avoid needing E[0], t[0]
    if t<0.0
        countnegatives=countnegatives+1
    end

    for i=2:length(D)
        t=D[i]-sigma-E_squared[i-1]/t   # Sturm sequence, overwriting temporary values...
        if t<0.0                     # if t<0, we've found another eigenvalue
            countnegatives=countnegatives+1
        end
    end

    return countnegatives
end

function randexp(N) # random exponential with the ln(X) 0<X<1 method
    return (log(rand(N)))
end

# Generate a random tridiagonal TightBinding Hamiltonian, in a form suitable for the Sturm sequence
# Given:
#   SiteEnergy - scalar eV; reference for site energy of monomer
#   disorder - scalar eV ; amount of Gaussian / normal energetic disorder, for trace of Hamiltonian
#   modelJ(theta) - function, takes degrees, returns eV ; model for the transfer integral (e.g. E=(J0*cos(thetas*pi/180)).^2 )
#   B - scalar (units?); Thermodynamic (B)eta, used to populate Probability Density Function
#   Z - scalar (units?); Partition function, weighting for absolute Boltzmann populations
#   U - function(theta angle); Free energy function, used to generate Bolztmann populations
#   N - integar ; size of diagonal of Hamiltonian
function randH(SiteEnergy, disorder, modelJ, B,Z,U,N)
# Random Trace / diagonal elements
    D=SiteEnergy + disorder*randn(N)
# Random Off-diag elements
#E=0.1+0.05*randn(N-1)

#Generate thetas... by rejection sampling. Might want to rewrite this to use ApproxFun sampling methods? Seem generally superior.
#    thetas=randexp(N-1)
    thetas=Float64[]
    for i=1:N-1 #number of thetas we need for off-diagonal element
        theta=0.0
        while true # this is a do-while loop, Julia styleee
            theta=360.0 * rand()     #random theta angle [DEGREES]; rand is on [0,1]
            p=exp(-U(theta)*B)/Z  #probability by stat mech
            p>1.0 && @printf("Eeek! Rejection sampling p>1.0; your distribution is biased.\n")
            p>rand() && break     #rejection sampling of distribution
        end
        push!(thetas,theta)       #tack theta onto end of list of thetas
    end

#Transfer integral from
#    J0=0.800 #Max Transfer integral
#    E=(J0*cos(thetas*pi/180)).^2
    E=modelJ(thetas)
# Squared (element wise) for the Sturm algorithm...
    E= E.^2
    return (D,E)
end

end


using PyPlot, DelimitedFiles, Random, Revise,
      HiQGA.transD_GP
##
Random.seed!(23)
rMin = 500 #m
rMax = 5200  #m
nRx  = 32
freqs   = [0.25, 0.75, 1.75, 3.25] #Hz
rRx    = collect(LinRange(rMin,rMax,nRx))  # Ranges to receivers   (m)
zRx    = 1000-eps(1.)
zTx    = 975.
RxAzim = 0.
TxDip  = 0.
# Note that the receiver depth needs to be in same model layer as transmitter.
zfixed = [-1e6,    0      ]
ρfixed = [1e13,    0.3    ]
zstart = zRx + eps(1.)
extendfrac, dz = 1.008, 12.
zall, znall, zboundaries = transD_GP.setupz(zstart, extendfrac, dz=dz, n=126)
z, ρ, nfixed = transD_GP.makezρ(zboundaries; zfixed=zfixed, ρfixed=ρfixed)
##
F = transD_GP.CSEM1DEr.RadialErLagged(zTx    = [zTx],
                      rRx    = rRx,
                      freqs  = freqs,
                      zRx    = [zRx],
                      RxAzim = RxAzim,
                      TxDip  = TxDip)
##
ρ[(z.>=zstart) .& (z.<1400)] .= 10^0.0
ρ[(z.>=1400) .&(z.<1700)] .= 10^0.3
ρ[(z.>=1700) .&(z.<1800)] .= 10^0.5
ρ[(z.>=1800) .&(z.<1900)] .= 10^0.2
ρ[(z.>=1900) .&(z.<2000)] .= 10^1.4 # 1.4 or 0.5
ρ[(z.>=2000) .&(z.<2400)] .= 10^0.0
ρ[(z.>=2400) .&(z.<2500)] .= 10^0.3
ρ[(z.>=2500) .&(z.<3500)] .= 10^0.7
ρ[(z.>=3500)]              .= 10^1.0
##
transD_GP.CSEM1DInversion.plotmodelfield!(F, z, ρ)
d, σ = transD_GP.addnoise(F, z, ρ, noisefrac=0.05,
               dz=dz, extendfrac=extendfrac, nfixed=nfixed)
##
csem = transD_GP.CSEMRadialEr(F, d, σ, z=z, ρ=ρ, nfixed=nfixed)


"""
    LIF

A leaky-integrate-fire neuron described by the following differential equation

``\\frac{\\mathrm{d}v}{\\mathrm{d}t} = \\frac{R}{\\tau} I - \\lambda``

# Fields
- `voltage::VT`: membrane potential
- `current::VT`: injected (unprocessed) current
- `lastt::IT`: the last time this neuron processed a spike
- `τm::VT`: membrane time constant
- `vreset::VT`: reset voltage potential
- `R::VT`: resistive constant (typically = 1)
"""
mutable struct LIF{VT<:Real, IT<:Integer} <: AbstractCell
    # required fields
    voltage::VT
    current::VT

    # model specific fields
    lastt::IT
    τm::VT
    vreset::VT
    R::VT
end

Base.show(io::IO, ::MIME"text/plain", neuron::LIF) =
    print(io, """LIF:
                     voltage: $(neuron.voltage)
                     current: $(neuron.current)
                     τm:      $(neuron.τm)
                     vreset:  $(neuron.vreset)
                     R:       $(neuron.R)""")
Base.show(io::IO, neuron::LIF) =
    print(io, "LIF(voltage = $(neuron.voltage), current = $(neuron.current))")

"""
    LIF(τm, vreset, vth, R = 1.0)

Create a LIF neuron with zero initial voltage and empty current queue.
"""
LIF(τm::Real, vreset::Real, R::Real = 1.0) = LIF{Float32, Int}(vreset, 0, 0, τm, vreset, R)

isactive(neuron::LIF, t::Integer; dt::Real = 1.0) = (neuron.current > 0)
getvoltage(neuron::LIF) = neuron.voltage

"""
    excite!(neuron::LIF, current)
    excite!(neurons::AbstractArray{<:LIF}, current)

Excite a `neuron` with external `current`.
"""
excite!(neuron::LIF, current) = (neuron.current += current)
excite!(neurons::T, current) where T<:AbstractArray{<:LIF} = (neurons.current .+= current)

"""
    spike!(neuron::LIF, t::Integer; dt::Real = 1.0)
    spike!(neurons::AbstractArray{<:LIF}, spikes; dt::Real = 1.0)

Record a output spike from the threshold function with the `neuron` body.
Sets `neuron.lastt`.
"""
spike!(neuron::LIF, t::Integer; dt::Real = 1.0) = (t > 0) && (neuron.voltage = neuron.vreset)
function spike!(neurons::T, spikes; dt::Real = 1.0) where T<:AbstractArray{<:LIF}
    map!((x, y, s) -> (s > 0) ? y : x, neurons.voltage, neurons.voltage, neurons.vreset, spikes)
end

"""
    evaluate!(neuron::LIF, t::Integer; dt::Real = 1.0)
    (neuron::LIF)(t::Integer; dt::Real = 1.0)
    evaluate!(neurons::AbstractArray{<:LIF}, t::Integer; dt::Real = 1.0)

Evaluate the neuron model at time `t`. Return the resulting membrane potential.
"""
function evaluate!(neuron::LIF, t::Integer; dt::Real = 1.0)
    neuron.voltage = SpikingNNFunctions.Neuron.lif((t - neuron.lastt) * dt, neuron.current, neuron.voltage; R = neuron.R, tau = neuron.τm)
    neuron.lastt = t
    neuron.current = 0

    return neuron.voltage
end
(neuron::LIF)(t::Integer; dt::Real = 1.0) = evaluate!(neuron, t; dt = dt)
function evaluate!(neurons::T, t::Integer; dt::Real = 1.0) where T<:AbstractArray{<:LIF}
    SpikingNNFunctions.Neuron.lif!((t .- neurons.lastt) .* dt, neurons.current, neurons.voltage; R = neurons.R, tau = neurons.τm)
    neurons.lastt .= t
    neurons.current .= 0

    return neurons.voltage
end

"""
    reset!(neuron::LIF)
    reset!(neurons::AbstractArray{<:LIF})

Reset `neuron` by setting the membrane potential to `neuron.vreset`.
"""
function reset!(neuron::LIF)
    neuron.lastt = 0
    neuron.voltage = neuron.vreset
    neuron.current = 0
end
function reset!(neurons::T) where T<:AbstractArray{<:LIF}
    neurons.lastt .= 0
    neurons.voltage .= neurons.vreset
    neurons.current .= 0
end


using Agents
using AgentsPlots
using Random

mutable struct SugarSeeker <: AbstractAgent
	id::Int
	pos::Tuple{Int, Int}
	vision::Int
	metabolic_rate::Int
	age::Int
	max_age::Int
	wealth::Int
end

function distances(pos, sugar_peaks, max_sugar)
	all_dists = Array{Int, 1}(undef, length(sugar_peaks))
	for (ind, peak) in enumerate(sugar_peaks)
		d = round(Int, sqrt(sum((pos .- peak).^2)))
		all_dists[ind] = d
	end
	return minimum(all_dists)
end

function sugar_caps(dims, sugar_peaks, max_sugar, dia=4)
	sugar_capacities = zeros(Int, dims)
	for i in 1:dims[1], j in 1:dims[2]
		sugar_capacities[i, j] = distances((i,j), sugar_peaks, max_sugar)
	end
	for i in 1:dims[1]
		for j in 1:dims[2]
			sugar_capacities[i, j] = max(0, max_sugar - (sugar_capacities[i, j] ÷ dia))
		end
	end
	return sugar_capacities
end

"Start a sugarscape simulation"
function sugarscape(;dims=(50,50), sugar_peaks=((10, 40), (40, 10)), growth_rate=1, N=250, w0_dist=(5, 25), metabolic_rate_dist=(1,4), vision_dist=(1,6), 	max_age_dist=(60, 100), max_sugar = 4)
	sugar_capacities = sugar_caps(dims, sugar_peaks, max_sugar, 6)
	sugar_values = deepcopy(sugar_capacities)
	space = GridSpace(dims, periodic=true, moore=true)
	properties = Dict(
		:growth_rate => growth_rate,
		:N => N,
		:w0_dist => w0_dist,
		:metabolic_rate_dist => metabolic_rate_dist,
		:vision_dist => vision_dist,
		:max_age_dist => max_age_dist,
		:sugar_values => sugar_values,
		:sugar_capacities => sugar_capacities
		)
	model = AgentBasedModel(SugarSeeker, space, scheduler=random_activation,
		properties=properties)
	for ag in 1:N
		add_agent_single!(model,
			rand(vision_dist[1]:vision_dist[2]),
			rand(metabolic_rate_dist[1]:metabolic_rate_dist[2]),
			0,
			rand(max_age_dist[1]:max_age_dist[2]),
			rand(w0_dist[1]:w0_dist[2]),
		)
	end
	return model, sugarscape_agent_step!, sugarscape_env!
end

function sugarscape_env!(model)
	## At each cell, sugar grows back at a rate of $\alpha$ units per time-step up to the cell's capacity c.
	togrow = findall(x -> model.sugar_values[x] < model.sugar_capacities[x], 1:prod(model.space.dimensions))
	model.sugar_values[togrow] .+= model.growth_rate
end

function movement!(agent, model)
	posvertex = coord2vertex(agent.pos, model)
	newsite = posvertex
	## find all unoccupied cells within vision
	neighbors = node_neighbors(coord2vertex(agent.pos, model), model, agent.vision)
	empty_nodes = [i for i in neighbors if isempty(i, model)]
	if length(empty_nodes) > 0
		## identify the one(s) with greatest amount of sugar
		maxsugar = maximum(model.sugar_values[empty_nodes])
		if maxsugar > 0
		sugary_sites_inds = findall(x -> x == maxsugar, model.sugar_values[empty_nodes])
		sugary_sites = empty_nodes[sugary_sites_inds]
		## select the nearest one (randomly if more than one)
		for dia in 1:agent.vision
			nn = node_neighbors(posvertex, model, dia)
			suitable = intersect(nn, sugary_sites)
			if length(suitable) > 0
				newsite = rand(suitable)
				break
			end
		end
		## move there and collect all the sugar in it
		newsite != posvertex && move_agent!(agent, newsite, model)
		end
	end
	## update wealth (collected - consumed)
	agent.wealth += (model.sugar_values[newsite] - agent.metabolic_rate)
	model.sugar_values[newsite] = 0
	## age
	agent.age += 1
end

function replacement!(agent, model)
	## If the agent's sugar wealth become zero or less, it dies
	if agent.wealth <= 0 || agent.age >= agent.max_age
		kill_agent!(agent, model)
		## Whenever an agent dies, a young one is added to a random pos.
		## New agent has random attributes
		add_agent_single!(model,
			rand(model.vision_dist[1]:model.vision_dist[2]),
			rand(model.metabolic_rate_dist[1]:model.metabolic_rate_dist[2]),
			0,
			rand(model.max_age_dist[1]:model.max_age_dist[2]),
			rand(model.w0_dist[1]:model.w0_dist[2]),
		)
	end
end

function sugarscape_agent_step!(agent, model)
	movement!(agent, model)
	replacement!(agent, model)
end


################################################################################
#  Copyright 2021, Tom Van Acker                                               #
################################################################################
# StochasticPowerModels.jl                                                     #
# An extention package of PowerModels.jl for Stochastic (Optimal) Power Flow   #
# See http://github.com/timmyfaraday/StochasticPowerModels.jl                  #
################################################################################

"expected cost of active power generation"
function objective_min_expected_generation_cost(pm::AbstractPowerModel; kwargs...)
    gen_cost = Dict()

    T2 = pm.data["T2"]

    for (g, gen) in _PM.ref(pm, :gen, nw=1)
        pg = Dict(nw => _PM.var(pm, nw, :pg, g) for nw in _PM.nw_ids(pm))

        if length(gen["cost"]) == 1
            gen_cost[g] = gen["cost"][1]
        elseif length(gen["cost"]) == 2
            gen_cost[g] = gen["cost"][1]*pg[1] +
                          gen["cost"][2]
        elseif length(gen["cost"]) == 3
            gen_cost[g] = gen["cost"][1]*sum(T2.get([n-1,n-1]) * pg[n]^2 for n in _PM.nw_ids(pm)) +
                          gen["cost"][2]*pg[1] +
                          gen["cost"][3]
        else
            gen_cost[g] = 0.0
        end
    end

    return JuMP.@objective(pm.model, Min,
            sum(gen_cost[g] for g in _PM.ids(pm, :gen, nw=1))
    )
end


import MittagLefflerFunctions.MLF.mlf3
import MittagLefflerFunctions.MLF.radius
import MittagLefflerFunctions.Γ

α = 0.5
β = 2.0
z = 8.7 * exp(Complex(0.0, 3π/4))
p = 10
@assert (   (  abs(angle(z)) > π*α )
         && ( abs(abs(angle(z)) - (π*α)) > 10.0^(-p) ) )

r0 = radius(α, β, z, p)
val = mlf3(α, β, z, p, r0)
ref = ( erfcx(-z) - 1.0 - z / Γ(3/2) ) / z^2

@test abs(val-ref) ≤ abs(ref)


using Test, SpecialFunctions, Statistics

function test_digitize(sep; C=1000, K=1000)
    cal0 = randn(C)
    cal1 = randn(C) .+ sep

    # flip 1000 coins, and add noise
    rbits  = round.(Int, rand(K))
    rvolts = sep*rbits + randn(K)

    # estimate of prob of being correct
    dm, pe0m, pe1m = digitize(rvolts, cal0, cal1, mode=:max)
    de, pe0e, pe1e = digitize(rvolts, cal0, cal1, mode=:equal)
    # compute theoretical prob of being correct
    pe = 1/2 - 1/2*erf(sep/(2*sqrt(2)))

    return pe,
           sum(abs.(dm-rbits))/length(rbits),
           sum(abs.(de-rbits))/length(rbits),
           pe0m,
           pe1m,
           pe0e,
           pe1e

end

# run a bunch of tests, store each as a column of a matrix, take the mean
# the rows are:
#   - predicted error rate (asymptotic)
#   - observed error rate for :max mode
#   - observed error rate for :equal mode
#   - estimated error rate based on cals for :max mode (0 preps)
#   - estimated error rate based on cals for :max mode (1 preps)
#   - estimated error rate based on cals for :equal mode (0 preps)
#   - estimated error rate based on cals for :equal mode (1 preps)
test1 = mean(reduce(hcat,[[test_digitize(0.75)...] for _ in 1:1000]))
test2 = mean(reduce(hcat,[[test_digitize(0.50)...] for _ in 1:1000]))
# make sure these are all close to the prediction
@test all(map(x->isapprox(x,0.35,atol=5e-2),test1))
@test all(map(x->isapprox(x,0.40,atol=5e-2),test2))


using SurrogateCube
# Define some Parameters
Ncomp = 3 #We want 3 independent components
Nvar  = 10 # we want 10 measured variables
noise = WhiteNoise(0.2) #Each component has a noise with sigma=0.5
means = [TrendBaseline(0.0,0.0,2.0),ConstantBaseline(0.0),ConstantBaseline(0.0)] # All Baselines are flat
dists = [TrendEvent(CubeEvent(0.2)),EmptyEvent(),EmptyEvent()] # Only one of the components is disturbed with a size of 20% in each dimension
varNoise = WhiteNoise(0.1) # each variable generated has some independent white noise
kb=0.0
kn=0.0
ks=-5.0  # How strongly do we want to perturb
#Set dimensions
Nlon = 10
Nlat = 10
Ntime =100
ds = genDataCube(means,noise,dists,varNoise,Ncomp,Nvar,Ntime,Nlat,Nlon,kb,kn,ks);
#Save the datacube to a file
save_cube(ds)
#Plot something
using PyPlot
pygui(true)
f=figure()
p=subplot()
for i=1:Nvar
plot(ds.variables[:,5,5,i][:])
end
p[:set_title]("Variables")
p[:set_xlabel]("Time")

f=figure()
p=subplot()
for i=1:Ncomp
plot(ds.components[:,5,5,i][:])
end
p[:set_title]("Components")
p[:set_xlabel]("Time")

f=figure()
p=subplot()
for i=1:Ncomp
plot(ds.eventmap[:,5,5,i][:])
end
p[:set_title]("Events")
p[:set_xlabel]("Time")
p[:set_ylim]([-0.1,2.1])


using Flux
using Flux: @epochs
using Statistics
using Random
using Parameters: @with_kw
using IterTools: ncycle

@with_kw mutable struct Args
    lr::Float64 = 0.5
    batch_size::Int = 256
    repeat::Int = 20
end

images = Flux.Data.FashionMNIST.images()
labels = Flux.Data.FashionMNIST.labels()

function get_fashion_label(labels)
	text_label = ["t-shirt", "trouser", "pullover", "dress", "coat", "sandal", "shirt", "sneaker", "bag", "ankle boot"]
	return [text_label[i+1] for i in labels]
end

#get train_data

data_x = rand(784, 6000)
data_y = []

for i in 1:6000
	data_x[:,i] = Float64.(reshape(images[i],(784,1)))
	push!(data_y, get_fashion_label(labels)[i])
end

label = sort(unique(data_y))
data_onehot_labels = Flux.onehotbatch(data_y, label)

train_x = data_x[:, [1:3:6000 ; 2:3:6000]]
train_y = data_onehot_labels[:, [1:3:6000 ; 2:3:6000]]

test_x = data_x[:, 3:3:6000]
test_y = data_onehot_labels[:, 3:3:6000]

train_data = Flux.Data.DataLoader((train_x, train_y), batchsize=Args().batch_size, shuffle=true)

#model
model = Chain(
    Dense(784, 256, relu),
    Dense(256, 10)
    )

#define loss function: cross entropy function
loss(x, y) = Flux.logitcrossentropy(model(x), y)

# params
ps = Flux.params(model)

#SDG
opt = Descent(Args().lr)

@time Flux.train!(loss, ps, ncycle(train_data, Args().repeat),opt)

accuracy(x, y, model) = Flux.mean(Flux.onecold(model(x)) .== Flux.onecold(y))

print("train loss: ", accuracy(train_x, train_y, model), ", test loss:", accuracy(test_x, test_y, model))


export CGPFunctions

module CGPFunctions

global arity = Dict()

SorX = Union{Symbol, Expr}

function scaled(x::Float64)
    if isnan(x)
        return 0.0
    end
    min(max(x, -1.0), 1.0)
end

function fgen(name::Symbol, ar::Int, s1::SorX; safe::Bool=false)
    if safe
        @eval function $name(x::Float64, y::Float64)::Float64
            try
                return $s1
            catch
                return x
            end
        end
    else
        @eval function $name(x::Float64, y::Float64)::Float64
            $s1
        end
    end
    arity[String(name)] = ar
end

# Mathematical
fgen(:f_add, 2, :((x + y) / 2.0))
fgen(:f_subtract, 2, :(abs(x - y) / 2.0))
fgen(:f_mult, 2, :(x * y))
fgen(:f_div, 2, :(scaled(x / y)))
fgen(:f_abs, 1, :(abs(x)))
fgen(:f_sqrt, 1, :(sqrt(abs(x))))
fgen(:f_pow, 2, :(abs(x) ^ abs(y)))
fgen(:f_exp, 1, :((2 * (exp(x+1)-1.0))/(exp(2.0)-1.0) -1))
fgen(:f_sin, 1, :(sin(x)))
fgen(:f_cos, 1, :(cos(x)))
fgen(:f_tanh, 1, :(tanh(x)))
fgen(:f_sqrt_xy, 2, :(sqrt(x^2 + y^2) / sqrt(2.0)))
fgen(:f_lt, 2, :(Float64(x < y)))
fgen(:f_gt, 2, :(Float64(x > y)))

# Logical
fgen(:f_and, 2, :(Float64((&)(Int(round(x)), Int(round(y))))))
fgen(:f_or, 2, :(Float64((|)(Int(round(x)), Int(round(y))))))
fgen(:f_xor, 2, :(Float64(xor(Int(abs(round(x))), Int(abs(round(y)))))))
fgen(:f_not, 1, :(1 - abs(round(x))))

end


using Luna
import Test: @test, @testset
import HCubature: hquadrature
import SpecialFunctions: besselj
import GSL: sf_bessel_zero_Jnu

@testset "On-axis intensity statistics" begin
    # Manually normalise the field distribution of HE11 to find scaling factor between
    # power and intensity
    a = 13e-6
    unm = sf_bessel_zero_Jnu(0, 1)
    E(r) = besselj(0, unm*r/a)
    norm, err = hquadrature(0, a) do r
        2π*r * abs2(E(r))
    end

    energy = 1e-6
    gas = :Ar
    pres = 5
    τ = 30e-15
    λ0 = 800e-9
    grid = Grid.RealGrid(15e-2, 800e-9, (160e-9, 3000e-9), 1e-12)
    m = Capillary.MarcatiliMode(a, gas, pres, loss=false)
    aeff(z) = Modes.Aeff(m, z=z)
    energyfun, energyfunω = Fields.energyfuncs(grid)

    dens0 = PhysData.density(gas, pres)
    densityfun(z) = dens0
    responses = (Nonlinear.Kerr_field(PhysData.γ3_gas(gas)),)
    linop, βfun!, frame_vel, αfun = LinearOps.make_const_linop(grid, m, λ0)
    inputs = Fields.GaussField(λ0=λ0, τfwhm=τ, energy=1e-6)
    Eω, transform, FT = Luna.setup(
        grid, densityfun, responses, inputs, βfun!, aeff)

    statsfun = Stats.default(grid, Eω, m, linop, transform; gas=gas, onaxis=true)
    output = Output.MemoryOutput(0, grid.zmax, 201, statsfun)
    Luna.run(Eω, grid, linop, transform, FT, output, status_period=5)

    @test all(output["stats"]["peakintensity"] .≈ output["stats"]["peakpower"]/norm)
end


# Bernoulli is good sanity check since many quantities have explicit representations
# and so we can double check if things work as expected.

using Random
using Empirikos
using Hypatia
using Setfield

Random.seed!(1)

n = 200

Zs = BinomialSample.(Int.(rand(Bernoulli(0.7), 200)), 1)
Zs_summary = summarize(Zs)

@test Zs_summary(BinomialSample(1,1)) == sum( ==(1), response.(Zs))
@test Zs_summary[BinomialSample(1,1)] == Zs_summary(BinomialSample(1,1))

@test length(Zs_summary) == 2
@test nobs(Zs_summary) == n

α = 0.11
floc_dkw = DvoretzkyKieferWolfowitz(;α=α)
floc_chisq = Empirikos.ChiSquaredFLocalization(;α=α)

pmf_at_1 = MarginalDensity(BinomialSample(1,1))

floc_dkw_interval = FLocalizationInterval(; flocalization = floc_dkw,
                                            solver = Hypatia.Optimizer,
                                            convexclass = DiscretePriorClass(0:0.001:1) )

floc_chisq_interval = FLocalizationInterval(; flocalization = floc_chisq,
                                            solver = Hypatia.Optimizer,
                                            convexclass = DiscretePriorClass(0:0.001:1) )

ci_at_1_dkw = confint(floc_dkw_interval, pmf_at_1, Zs_summary)

dkw_error = sqrt(log(2/α)/(2n))
hatf_1 =  Zs_summary(BinomialSample(1,1))/n

@test hatf_1 + dkw_error ≈ ci_at_1_dkw.upper
@test hatf_1 - dkw_error ≈ ci_at_1_dkw.lower

ci_at_1_chisq = confint(floc_chisq_interval, pmf_at_1, Zs_summary)

τsq = quantile(Chisq(1), 1-α)
chisq_f_center = (hatf_1 + τsq/(2n))/(1 + τsq/(n))
chisq_error = sqrt(τsq/n)/(1 + τsq/n)*sqrt( hatf_1*(1-hatf_1) +  τsq/(4n))

@test chisq_f_center +  chisq_error ≈ ci_at_1_chisq.upper
@test !(hatf_1 +  chisq_error ≈ ci_at_1_chisq.upper)
@test chisq_f_center -  chisq_error ≈ ci_at_1_chisq.lower

postmean_at_1 = PosteriorMean(BinomialSample(1,1))

ci_postmean_at_1_dkw = confint(floc_dkw_interval, postmean_at_1, Zs_summary)

@test ci_postmean_at_1_dkw.upper ≈ 1.0 atol = 1e-6
@test ci_postmean_at_1_dkw.lower ≈ hatf_1 - dkw_error atol = 1e-6

ci_postmean_at_1_chisq = confint(floc_chisq_interval, postmean_at_1, Zs_summary)

@test ci_postmean_at_1_chisq.upper ≈ 1.0 atol = 1e-6
@test ci_postmean_at_1_chisq.lower ≈ chisq_f_center -  chisq_error atol = 1e-6
@test !(ci_postmean_at_1_chisq.lower ≈ ci_postmean_at_1_dkw.lower)

Random.seed!(1)

n = 200

Zs = BinomialSample.(Int.(rand(Bernoulli(0.7), 200)), 1)
Zs_summary = summarize(Zs)

comp_Zs = Empirikos.compound(Zs)
@test comp_Zs[1].probs == [1.0]
@test comp_Zs[1].vec == [BinomialSample(1)]
@test response(comp_Zs[1]) == response(Zs[1])

het_Zs = Empirikos.heteroskedastic(Zs)
@test het_Zs.probs == [1.0]
@test het_Zs.vec == [BinomialSample(1)]

amari_with_F = AMARI(;flocalization = (@set floc_dkw.α = 0.01),
                      solver = Hypatia.Optimizer,
                      convexclass = DiscretePriorClass(0:0.1:1),
                      modulus_model = Empirikos.ModulusModelWithF,
                      discretizer = Empirikos.integer_discretizer(0:1))

amari_without_F = AMARI(;flocalization = (@set floc_dkw.α = 0.01),
                      solver = Hypatia.Optimizer,
                      convexclass = DiscretePriorClass(0:0.1:1),
                      modulus_model = Empirikos.ModulusModelWithoutF,
                      discretizer = Empirikos.integer_discretizer(0:1))

for amari_ in (amari_with_F, amari_without_F)
    amari_fit_priormean = fit(amari_, Empirikos.PriorMean(), Zs)
    ci_priormean = confint(amari_fit_priormean, Empirikos.PriorMean(), Zs; level=1-α)
    @test confint(amari_, Empirikos.PriorMean(), Zs; level=1-α) == ci_priormean

    @test amari_fit_priormean.Q(BinomialSample(1,1)) ≈ 1.0 atol = 1e-6
    @test amari_fit_priormean.Q(BinomialSample(0,1)) ≈ 0.0 atol = 1e-6

    @test ci_priormean.estimate ≈ mean(response.(Zs)) atol = 1e-6
    @test ci_priormean.maxbias ≈ 0 atol = 1e-7
    @test amari_fit_priormean.max_bias ≈ 0 atol = 1e-7
    @test ci_priormean.se ≈ std(response.(Zs))/sqrt(nobs(Zs)) atol = 1e-7

    @test ci_priormean.lower ≈ ci_priormean.estimate - quantile(Normal(), 1-α/2)*ci_priormean.se
    @test ci_priormean.upper ≈ ci_priormean.estimate + quantile(Normal(), 1-α/2)*ci_priormean.se

    @test amari_fit_priormean.target === Empirikos.PriorMean()
    @test amari_fit_priormean.unit_var_proxy ≈ var(response.(Zs)) atol=0.01
    @test amari_fit_priormean.unit_var_proxy ≈ mean(response.(Zs))*(1- mean(response.(Zs))) atol=1e-5

    amari_fit_prior_second_mean = fit(amari_, Empirikos.PriorSecondMoment(), Zs)
    ci_second_mean = confint(amari_fit_prior_second_mean, Empirikos.PriorSecondMoment(), Zs; level=1-α)

    @test_throws ErrorException confint(amari_fit_prior_second_mean, Empirikos.PriorMean(), Zs; level=1-α)
    tmp_fit = fit((@set floc_dkw.α = 0.01), Zs)
    dkw_lb = mean(response, Zs) - tmp_fit.band

    @test amari_fit_prior_second_mean.max_bias ≈ dkw_lb*(1-dkw_lb)/2 rtol = 0.005
    @test amari_fit_prior_second_mean.Q(BinomialSample(1,1)) ≈ 1.0 - dkw_lb*(1-dkw_lb)/2  atol = 1e-3
    @test amari_fit_prior_second_mean.Q(BinomialSample(0,1)) ≈ - dkw_lb*(1-dkw_lb)/2  atol = 1e-3
    @test ci_second_mean.estimate ≈ mean(response.(Zs)) - dkw_lb*(1-dkw_lb)/2 atol = 1e-3
    @test ci_second_mean.lower < ci_second_mean.estimate - quantile(Normal(), 1-α/2)*ci_second_mean.se
    @test ci_second_mean.lower > ci_second_mean.estimate - quantile(Normal(), 1-α/2)*ci_second_mean.se - ci_second_mean.maxbias
    @test ci_second_mean.se ≈ std(response.(Zs))/sqrt(nobs(Zs)) atol = 1e-7
    @test amari_fit_prior_second_mean.unit_var_proxy ≈ mean(response.(Zs))*(1- mean(response.(Zs))) atol=1e-5

    both_cis = confint.(amari_, [Empirikos.PriorMean(), Empirikos.PriorSecondMoment()], Zs)
    @test getfield.(both_cis, :α) ≈ [0.05; 0.05]

    both_cis_kw = confint.(amari_, [Empirikos.PriorMean(), Empirikos.PriorSecondMoment()], Zs; level=0.95)
    @test both_cis_kw == both_cis

    both_cis_at_α = confint.(amari_, [Empirikos.PriorMean(), Empirikos.PriorSecondMoment()], Zs; level=1-α)
    @test both_cis_at_α == [ci_priormean; ci_second_mean]

    # some work towards AMARI with posterior mean
    c = 0.6
    postmean_lin = Empirikos.PosteriorTargetNullHypothesis(PosteriorMean(BinomialSample(1,1)), c)

    amari_fit_postmean_lin = fit(amari_, postmean_lin, Zs)

    @test amari_fit_postmean_lin.max_bias ≈ dkw_lb*(1-dkw_lb)/2 rtol = 0.005
    @test amari_fit_postmean_lin.Q(BinomialSample(1,1)) ≈ (1-c) - dkw_lb*(1-dkw_lb)/2  atol = 1e-3
    @test amari_fit_postmean_lin.Q(BinomialSample(0,1)) ≈ - dkw_lb*(1-dkw_lb)/2  atol = 1e-3

    ci_postmean_lin = confint(amari_fit_postmean_lin, postmean_lin, Zs; level=1-α)

    @test ci_postmean_lin.estimate ≈ (1-c)*mean(response.(Zs)) - dkw_lb*(1-dkw_lb)/2 atol = 1e-3
    @test ci_postmean_lin.lower < ci_postmean_lin.estimate - quantile(Normal(), 1-α/2)*ci_postmean_lin.se
    @test ci_postmean_lin.lower > ci_postmean_lin.estimate - quantile(Normal(), 1-α/2)*ci_postmean_lin.se - ci_postmean_lin.maxbias
    @test amari_fit_postmean_lin.unit_var_proxy ≈ (1-c)^2*mean(response.(Zs))*(1- mean(response.(Zs))) atol=1e-5

    function tmp_f(c)
        barf = mean(response, Zs)
        _est = (1-c)*barf - dkw_lb*(1-dkw_lb)/2
        _se = (1-c)*sqrt(barf*(1-barf)/ nobs(Zs))
        _maxbias =  dkw_lb*(1-dkw_lb)/2
        _pm = Empirikos.gaussian_ci(_se; maxbias=_maxbias, α=α)
        _est - _pm
    end

    cs = 0:0.0001:1
    idx = findfirst( tmp_f.(cs) .<= 0)
    c_left = cs[idx]

    ci_postmean = confint(amari_, PosteriorMean(BinomialSample(1,1)), Zs; level=1-α)
    @test ci_postmean.upper ≈ 1.0
    @test ci_postmean.lower ≈ c_left atol = 0.0015

    Zs_flip = BinomialSample.(1 .- response.(Zs), 1)
    ci_postmean_0_flip = confint(amari_, PosteriorMean(BinomialSample(0,1)), Zs_flip; level=1-α)
    @test ci_postmean_0_flip.lower ≈ 1 - ci_postmean.upper atol = 1e-5
    @test ci_postmean_0_flip.upper ≈ 1 - ci_postmean.lower atol = 1e-5

    ci_postmean_0 = confint(amari_, PosteriorMean(BinomialSample(0,1)), Zs; level=1-α)
    @test ci_postmean_0.lower ≈ 0 atol=1e-8

    both_cis_postmean = confint.(amari_, identity(PosteriorMean.(BinomialSample.([0,1],1))), Zs)
    @test getfield.(both_cis_postmean, :α) ≈ [0.05; 0.05]

    both_cis_postmean_kw = confint.(amari_, identity(PosteriorMean.(BinomialSample.([0,1],1))), Zs; level=0.95)
    @test both_cis_postmean_kw == both_cis_postmean

    both_cis_postmean_at_α = confint.(amari_, identity(PosteriorMean.(BinomialSample.([0,1],1))), Zs; level=1-α)
    @test both_cis_postmean_at_α == [ci_postmean_0; ci_postmean]

end


export SpikedWigner

struct SpikedWigner <: ContinuousMatrixDistribution
    beta::Integer
    n::Integer
    spikes::Array{Float64, 1}
    scaled::Bool
end

SpikedWigner(beta, n, spikes; scaled=false) = SpikedWigner(beta, n, spikes, scaled)
SpikedWigner(beta, n; scaled=false) = SpikedWigner(beta, n, [], scaled)

# SAMPLERS

function scaling(d::SpikedWigner)
    d.scaled ? 1/sqrt(d.n) : 1
end

function randreduced(d::SpikedWigner)
    if length(d.spikes) <= 1
        return randtridiagonal(d)
    else
        throw(">1 spike not yet implemented for Wigner matrices")
    end
end

function randtridiagonal(d::SpikedWigner)

    # diagonal and superdiagonal of B
    dv = rand(Normal(0, sqrt(2)), d.n)/sqrt(d.beta)
    ev = [rand(Chi(d.beta*(d.n-k))) for k in 1:(d.n-1)]/sqrt(d.beta)

    if length(d.spikes) == 1
        dv[1] += d.spikes[1] * sqrt(d.n)
    end

    SymTridiagonal(dv, ev) * scaling(d)
end


using BeastUtils.RTrees

newick = "(((A:0.5,B:0.2):0.1,C:0.3):1.0,(D:0.3,E:0.4):0.7);"
tree = RTrees.parse_newick(newick)

taxa1 = ["A", "B", "C", "D", "E"]

Ψ1 = RTrees.tree_variance(tree, taxa1)

Ψ1_true = [
    1.6 1.1 1.0 0.0 0.0;
    1.1 1.3 1.0 0.0 0.0;
    1.0 1.0 1.3 0.0 0.0;
    0.0 0.0 0.0 1.0 0.7;
    0.0 0.0 0.0 0.7 1.1
    ]

tol = 1e-14

diff1 = maximum(abs.(Ψ1 - Ψ1_true))

@test diff1 < tol

perm = [5, 2, 3, 1, 4]
taxa2 = taxa1[perm]
Ψ2_true = Ψ1_true[perm, perm]
Ψ2 = RTrees.tree_variance(tree, taxa2)

diff2 = maximum(abs.(Ψ2 - Ψ2_true))

@test diff2 < tol


"""
    phase_diff(field, dir)

Compute point-to-point difference of phase on a grid along a direction

Returns an array of size `(size(field)[1], size(field)[2], size(field)[3])`
containing differences in direction `dir`.
"""
function phase_diff(field, dir)
    out = similar(field, Float64)

    out .= angle.(circshift(field, dir) ./  field)

    out
end

"""
    phase_diff(field::PencilArray, dir)

Compute point-to-point difference of phase on a grid along a direction

Returns an array of size `(size(field)[1], size(field)[2], size(field)[3])`
containing differences in direction `dir`.
"""
function phase_diff(field::PencilArray, dir)
    out = similar(field, Float64)

    # TODO diff at boundaries
    out .= 0
    out[1:end-dir[1], 1:end-dir[2], 1:end-dir[3]] .= angle.(circshift(field[1:end-dir[1], 1:end-dir[2], 1:end-dir[3]], dir) ./ field[1+dir[1]:end, 1+dir[2]:end, 1+dir[3]:end])

    out
end

"""
    normed_max_phase_grad(grids)

Compute maximum phase gradient of a grid

Normalised to ignore large gradients in regions with low density.  These tend
to be anomalous.
"""
function max_normed_phase_diff(psi, rho, density_threshold)
    tmp = similar(psi, Float64)
    n = ndims(tmp)

    max_grads = zeros(n)

    shift = zeros(Int, n)
    shift[1] = 1

    for i in 1:n
        dir = circshift(shift, i)
        tmp .= abs.(phase_diff(psi, dir))
        tmp[rho / maximum(rho) .< density_threshold] .= NaN
        max_grads[i] = maximum(tmp)
    end

    maximum(max_grads)

end

"""
    bad_phase_diff(grids, config)

Check if the phase of the ψ field is in a trustable regime.

Returns
"""
function good_phase_diff(grids, config)
    density_threshold = config.density_threshold
    if max_normed_phase_diff(grids.ψx, grids.ρx, density_threshold) > config.phase_grad_limit
        false
    else
        true
    end
end


using NLsolve
using Dolang
using DataFrames

function newton(f, init; tol=1e-8, maxit=500, backsteps=5, verbose=false)

    line = 0.5.^(0:(backsteps-1))
    x0 = init[:]
    x = copy(x0)
    N = length(x0)
    it = 0
    err = 1.0
    while (it<maxit) & (err>tol)
        it += 1
        ress, jacc = f(reshape(x0,size(init)...))
        res = reshape(ress,N)
        jac = reshape(jacc,N,N)
        err = maximum(abs(ress))
        ul = 0.1
        if err>tol
            dx = -jac\res
            nerr = 2.0
            for l in line
                ul = l
                try
                    x = x0+l*dx
                    nr, junk = f(reshape(x,size(init)...))
                    nerr = maximum(abs(nr))
                catch
                    println("skip")
                end
                if nerr<err
                    break
                end
            end
            if nerr>err
                println("Backtracking didn't work")
            end
        end
        x0 = x
        if verbose
            println(it, ": ", err, ":", ul)
        end
    end
    return reshape(x0,size(init)...)
end

function deterministic(model, exogenous_series; verbose=false, T=100)

    N=T

    symbols = model.symbols
    endogenous = symbols[:endogenous]
    exogenous = symbols[:exogenous]
    parameters = symbols[:parameters]

    it = Dolang.IncidenceTable(model.equations)

    fut_syms = it.by_date[1]
    cur_syms = it.by_date[0]
    past_syms = it.by_date[-1]

    fut_inds = [i for (i,v) in enumerate(endogenous) if v in it.by_date[1]]
    cur_inds = [i for (i,v) in enumerate(endogenous) if v in it.by_date[0]]
    past_inds = [i for (i,v) in enumerate(endogenous) if v in it.by_date[-1]]
    shocks_inds = [i for (i,v) in enumerate(exogenous) if v in it.by_date[0]]
    parms_inds = [i for (i,v) in enumerate(parameters) if v in it.by_date[0]]

    P = length(symbols[:endogenous])
    indices = fut_inds
    indices = cat(1,indices, P+cur_inds)

    indices = cat(1,indices, 2*P+past_inds)
    findices = cat(1,indices, 3*P+shocks_inds)

    ne = length(model.symbols[:exogenous])

    T = size(exogenous_series,1)
    exogenous = zeros(N, ne)
    exogenous[1:T,:] = exogenous_series
    for t=T+1:N
        exogenous[t,:] = exogenous[T,:]
    end

    f_s = model.functions.f_static
    f_d = model.functions.f_dynamic

    y0 = [model.calibration[e] for e in model.symbols[:endogenous]]
    e0 = [model.calibration[e] for e in model.symbols[:exogenous]]
    p = [model.calibration[e] for e in model.symbols[:parameters]]

    e_start = exogenous[1,:]
    e_final = exogenous[end,:]

    function fobj_start(x, out)
        out[:] = f_s(cat(1, x, e_start) , p)
    end

    function fobj_final(x, out)
        out[:] = f_s(cat(1, x, e_final) , p)
    end

    fdiff_start = NLsolve.DifferentiableMultivariateFunction(fobj_start)
    sol_start = nlsolve(fdiff_start,y0)
    y_start = sol_start.zero
    fdiff_final = NLsolve.DifferentiableMultivariateFunction(fobj_final)
    sol_final = nlsolve(fdiff_final,y0)
    y_final = sol_final.zero

    y_init_0 = [(1-l)*y_start + l*y_final for l in linspace(0,1,N)]
    y_init = cat(1, [e' for e in y_init_0]...) :: Array{Float64, 2}

    n_y = length(model.symbols[:endogenous])
    n_e = length(model.symbols[:exogenous])

    function residual(y)

        a = y[1:end-2,:]   # y{t-1}
        b = y[2:end-1,:]   # y_t
        c = y[3:end,:]     # y_{t+1}
        d = exogenous[1:end-2,:]

        x = cat(2,c,b,a,d)
        xx = x[:,findices]

        N,P = size(y)

        res = zeros(N,P)
        jac0 = zeros(N,P,3*P)
        for n = 2:N-1
            res[n,:] = f_d(Der{0},xx[n-1,:],p)
            jac0[n,:,indices] = f_d(Der{1},xx[n-1,:],p)[:,1:length(indices)]
        end
        jac = zeros(N,P,N,P)
        for n in 2:N-1
            jac[n,:,n,:] += jac0[n,:,(P+1):2*P]  # y_t
            jac[n,:,n+1,:] += jac0[n,:,1:P] # y_{t+1}
            jac[n,:,n-1,:] +=  jac0[n,:,2*P+1:3*P]  # y_{t-1}
        end
        jac[1,:,1,:] = eye(P)
        jac[end,:,end,:] = eye(P)
        res[1,:] = y[1,:]-y_start
        res[end,:] = y[end,:]-y_final
        return res, jac
    end
    a,b = residual(y_init)
    if verbose
        println("Initial: ", maximum(abs(a)))
    end
    sol = newton(residual, y_init, verbose=verbose)

    tvec = collect(0:(size(sol,1)-1))
    sol = cat(2,tvec,sol,exogenous)

    colnames = cat(1, [:t], model.symbols[:endogenous], model.symbols[:exogenous])
    columns = Dict(colnames[i] => sol[:,i] for i=1:size(sol,2))
    return DataFrame(columns)
end


# requires LinearAlgebra, Images, ImageView

include("GrahamScan.jl")

using ImageView, Images
using SparseArrays
using Random
using LinearAlgebra
using Statistics

function FindRoughEdge(A)
    # Find the rough coordinates of the non-horizontal edges of shapes in A
    B = Float64.(A);
    abs.(B[:, 2:end] - B[:, 1:end - 1]) |>
        rough_outline -> findall(px -> abs.(px - 1.0) < 0.8, rough_outline)
end

function FitLinear(x, y, deg)
    # Use linear regression to approximate f(x) = y with a polynomial of degree deg
    # returns coefficients of polynomial
    vs = hcat(map(p -> x.^p, (deg:-1:0))...);
    Q, R = qr(vs);
    R_pad = [R; zeros(size(Q, 1) - size(R, 1), size(R, 2))];
    params = R_pad \ (Q' * y);
    params
end

function FitLinearModel(x, y, deg)
    # Use linear regression to approximate f(x) = y with a polynomial of degree deg
    # returns a ready to use function
    params = FitLinear(x, y, deg);
    t -> (t.^(deg:-1:0))' * params
end

function CleanUpEdges(sensitivity, contour)
    # Try to remove unnecessary points from contour
    # contour is a matrix where the rows are x,y coordinate pairs
    reduced_contour = [];
    current_xs = Array{eltype(contour),1}();
    current_ys = Array{eltype(contour),1}();
    for point in eachrow(contour) # the first element is always equal to GrahamScan.EdgePoint(contour)
        x, y = point[1], point[2]
        if isempty(current_xs)
            push!(current_xs, x);
            push!(current_ys, y);
        else
            next_xs = [x; current_xs];
            next_ys = [y; current_ys];
            # approximate all points in current segment by a line
            approx = FitLinearModel(next_xs, next_ys, 1);
            # calculate deviations from approximation and actual values
            # and find the maximum deviation
            md = vcat(approx.(next_xs)...) - next_ys |>
                ds -> abs.(ds) |>
                maximum;
            if md <= sensitivity # if the maximum deviation doesn't exceed `sensitivity` pixels
                # try the next line
                current_xs = next_xs;
                current_ys = next_ys;
                continue
            else
                # otherwise save ends of the last line to the reduced contour
                push!(reduced_contour, [current_xs[1], current_ys[1]]);
                push!(reduced_contour, [current_xs[end], current_ys[end]]);
                # and begin a new segment starting at the current point
                current_xs = [x];
                current_ys = [y];
            end
        end
    end
    reduced_contour |> unique |>
        x -> hcat(x...)'
end

function Center(C)
    # find "center" of a point cloud
    mean.(eachcol(C))'
end

function Triangulate(center, points)
    start = points[1, :];
    current = start;
    Channel() do channel
        for point in eachrow(points[2:end, :])
            put!(channel, (current', center, point'));
            current = point;
        end
        put!(channel, (current', center, start'));
    end
end

function CalculateArea(C)
    # calculate area of star domain with center Center(C)
    triangleArea((a, b, c)) = abs(det([a 1; b 1; c 1])) / 2;
    map(triangleArea, Triangulate(Center(C), C)) |> sum
end

function CartesianToCoords(idxs)
    # convert CartesianIndex array to (no of pixels) x (2) Matrix
    hcat(map(idx -> [idx[1], idx[2]], idxs)...)'
end

function CoordsToImage(A)
    S = sparse(A[:,1], A[:,2], ones(size(A)[1]));
    I, J, V = findnz(S);
    I = I .- (minimum(I) - 1);
    J = J .- (minimum(J) - 1);
    sparse(I, J, V) |> Matrix
end

function Show(A)
    imshow(CoordsToImage(A))
    A
end

if isempty(ARGS)
    norm_path = "norm.bmp";
    img_path = "in.bmp";
else
    norm_path = ARGS[1];
    img_path = ARGS[2];
end

# calculate area of a square we know to have an area of 1mm²
# (this could of course be replaced with a arbitrary shape of known area)
area_square = load(norm_path) |>
    img -> Gray.(img) |>
    FindRoughEdge |>
    CartesianToCoords |>
    GrahamScan.FindConvexHull |> Show |>
    CalculateArea

pixelsToMillimeters(x) = x / sqrt(area_square)

area(path) = load(path) |>
    img -> Gray.(img) |>
    FindRoughEdge |>
    CartesianToCoords |>
    GrahamScan.FindConvexHull |> Show |>
    hull -> CleanUpEdges(10, hull) |> Show |>
    pixelsToMillimeters |>
    CalculateArea |>
    area ->
        println("The area is $(area) mm². ~$(Int64(round(sqrt(area_square)))) pixels are 1 mm.")

area(img_path)

println("Press any button to continue...")
readline()