using SymbolicUtils: Symbolic, Sym, FnType, Term, Add, Mul, Pow, symtype, operation, arguments
using SymbolicUtils
using IfElse: ifelse
using Test

@testset "@syms" begin
    let
        @syms a b::Float64 f(::Real) g(p, h(q::Real))::Int

        @test a isa Sym{Number}
        @test a.name === :a

        @test b isa Sym{Float64}
        @test b.name === :b

        @test f isa Sym{FnType{Tuple{Real}, Number}}
        @test f.name === :f

        @test g isa Sym{FnType{Tuple{Number, FnType{Tuple{Real}, Number}}, Int}}
        @test g.name === :g

        @test f(b) isa Term
        @test symtype(f(b)) === Number
        @test_throws ErrorException f(a)

        @test g(b, f) isa Term
        @test_throws ErrorException g(b, a)

        @test symtype(g(b, f)) === Int

        # issue #91
        @syms h(a,b,c)
        @test isequal(h(1,2,3), h(1,2,3))
    end
end

@testset "hashing" begin
    @syms a b f(x, y)
    @test hash(a) == hash(a)
    @test hash(a) != hash(b)
    @test hash(a+1) == hash(a+1)
    @test hash(sin(a+1)) == hash(sin(a+1))
    @test hash(f(1,a)) == hash(f(1, a))

    c = a
    g = f
    @syms a f(x, y)
    @test hash(a) == hash(c)
    @test hash(g(a, b)) == hash(f(a,b))
    @test hash(f(a, b)) == hash(f(c,b))
    @test hash(sin(a+1)) == hash(sin(c+1))

    ex = sin(a+1)
    h = hash(ex, UInt(0))
    @test ex.hash[] == h
    ex1 = sin(a+1)
    hash(asin(ex1), UInt(0))
    @test ex1.hash[] == h
end

struct Ctx1 end
struct Ctx2 end

@testset "metadata" begin
    @syms a b c
    for a = [a, sin(a), a+b, a*b, a^3]

        a′ = setmetadata(a, Ctx1, "meta_1")

        @test hasmetadata(a′, Ctx1)
        @test !hasmetadata(a′, Ctx2)

        a′ = setmetadata(a′, Ctx2, "meta_2")

        @test hasmetadata(a′, Ctx1)
        @test hasmetadata(a′, Ctx2)

        @test getmetadata(a′, Ctx1) == "meta_1"
        @test getmetadata(a′, Ctx2) == "meta_2"
    end

    # In substitute #283
    #
    @syms f(t) t
    f = setmetadata(f(t), Ctx1, "yes")
    hasmetadata(f, Ctx1) # true
    newf = substitute(f, Dict(a=>b)) # unrelated substitution
    @test hasmetadata(newf, Ctx1)
    @test getmetadata(newf, Ctx1) == "yes"

    @test isequal(substitute(1+sqrt(a), Dict(a => 2), fold=false),
                  1 + term(sqrt, 2, type=Number))
    @test substitute(1+sqrt(a), Dict(a => 2), fold=true) isa Float64
end

@testset "Base methods" begin
    @syms w::Complex z::Complex a::Real b::Real x

    @test isequal(w + z, Add(Number, 0, Dict(w=>1, z=>1)))
    @test isequal(z + a, Add(Number, 0, Dict(z=>1, a=>1)))
    @test isequal(a + b, Add(Real, 0, Dict(a=>1, b=>1)))
    @test isequal(a + x, Add(Number, 0, Dict(a=>1, x=>1)))
    @test isequal(a + z, Add(Number, 0, Dict(a=>1, z=>1)))

    foo(w, z, a, b) = 1.0
    SymbolicUtils.promote_symtype(::typeof(foo), args...) = Real
    @test SymbolicUtils._promote_symtype(foo, (w, z, a, b,)) === Real

    # promote_symtype of identity
    @test isequal(Term(identity, [w]), Term{Complex}(identity, [w]))
    @test isequal(+(w), w)
    @test isequal(+(a), a)

    @test isequal(rem2pi(a, RoundNearest), Term{Real}(rem2pi, [a, RoundNearest]))

    # bool
    for f in [(==), (!=), (<=), (>=), (<), (>)]
        @test isequal(f(a, 0), Term{Bool}(f, [a, 0]))
        @test isequal(f(0, a), Term{Bool}(f, [0, a]))
        @test isequal(f(a, a), Term{Bool}(f, [a, a]))
    end

    @test symtype(ifelse(true, 4, 5)) == Int
    @test symtype(ifelse(a < 0, b, w)) == Union{Real, Complex}
    @test SymbolicUtils.promote_symtype(ifelse, Bool, Int, Bool) == Union{Int, Bool}
    @test_throws MethodError w < 0
    @test isequal(w == 0, Term{Bool}(==, [w, 0]))

    @eqtest x // 5 == (1 // 5) * x
    @eqtest (1//2 * x) / 5 == (1 // 10) * x
    @eqtest x // Int16(5) == Rational{Int16}(1, 5) * x
    @eqtest 5 // x == 5 / x
    @eqtest x // a == x / a
end

@testset "array-like operations" begin
    abstract type SquareDummy end
    Base.:*(a::Symbolic{SquareDummy}, b) = b^2
    @syms s t a::SquareDummy A[1:2, 1:2]

    @test isequal(ndims(A), 2)
    @test_broken isequal(a.*[1 (s+t); t pi], [1 (s+t)^2; t^2 pi^2])
    @test isequal(s.*[1 (s+t); t pi], [s s*(s+t); s*t s*pi])
end

@testset "err test" begin
    @syms t()
    @test_throws ErrorException t(2)
end

@testset "substitute" begin
    @syms a b
    @test substitute(a, Dict(a=>1)) == 1
    @test isequal(substitute(sin(a+b), Dict(a=>1)), sin(b+1))
    @test substitute(a+b, Dict(a=>1, b=>3)) == 4
    @test substitute(exp(a), Dict(a=>2)) ≈ exp(2)
end

@testset "occursin" begin
    @syms a b c
    @test occursin(a, a + b)
    @test !occursin(sin(a), a + b + c)
    @test occursin(sin(a),  a * b + c + sin(a^2 * sin(a)))
end

@testset "printing" begin
    @syms a b c
    @test repr(a+b) == "a + b"
    @test repr(-a) == "-a"
    @test repr(-a + 3) == "3 - a"
    @test repr(-(a + b)) == "-a - b"
    @test repr((2a)^(-2a)) == "(2a)^(-2a)"
    @test repr(1/2a) == "1 / (2a)"
    @test repr(2/(2*a)) == "1 / a"
    @test repr(Term(*, [1, 1])) == "1"
    @test repr(Term(*, [2, 1])) == "2*1"
    @test repr((a + b) - (b + c)) == "a - c"
    @test repr(a + -1*(b + c)) == "a - b - c"
    @test repr(a + -1*b) == "a - b"
    @test repr(-1^a) == "-(1^a)"
    @test repr((-1)^a) == "(-1)^a"
end

@testset "similarterm" begin
    @syms a b c
    @test isequal(SymbolicUtils.similarterm((b + c), +, [a,  (b+c)]).dict, Dict(a=>1,b=>1,c=>1))
    @test isequal(SymbolicUtils.similarterm(b^2, ^, [b^2,  1//2]), b)
end

toterm(t) = Term{symtype(t)}(operation(t), arguments(t))

@testset "diffs" begin
    @syms a b c
    @test isequal(toterm(-1c), Term{Number}(*, [-1, c]))
    @test isequal(toterm(-1(a+b)), Term{Number}(+, [-1a, -b]))
    @test isequal(toterm((a + b) - (b + c)), Term{Number}(+, [a, -1c]))
end

@testset "hash" begin
    @syms a b
    @test hash(a + b, UInt(0)) === hash(a + b) === hash(a + b, UInt(0)) # test caching
    @test hash(a + b, UInt(2)) !== hash(a + b)
end

@testset "methoderror" begin
    @syms a::Any b::Any

    @test_throws MethodError a * b
    @test_throws MethodError a + b
end

@testset "canonical form" begin
    @syms a b c
    for x in [a, a*b, a+b, a-b, a^2, sin(a)]
        @test isequal(x * 1, x)
        @test x * 0 === 0
        @test isequal(x + 0, x)
        @test isequal(x + x, 2x)
        @test isequal(x + 2x, 3x)
        @test x - x === 0
        @test isequal(-x, -1x)
        @test isequal(x^1, x)
        @test isequal((x^-1)*inv(x^-1), 1)
    end
end

@testset "isequal" begin
    @syms a b c
    @test isequal(a + b, a + b + 0.01 - 0.01)
end

@testset "subtyping" begin
    T = FnType{Tuple{T,S,Int} where {T,S}, Real}
    s = Sym{T}(:t)
    @syms a b c::Int
    @test isequal(arguments(s(a, b, c)), [a, b, c])
end

@testset "div" begin
    @syms x::SafeReal y::Real
    @test (2x/2y).num isa Sym
    @test (2x/3y).num.coeff == 2
    @test (2x/3y).den.coeff == 3
    @test (2x/-3x).num.coeff == -2
    @test (2x/-3x).den.coeff == 3
    @test (2.5x/3x).num.coeff == 2.5
    @test (2.5x/3x).den.coeff == 3
    @test (x/3x).den.coeff == 3

    @syms x y
    @test (2x/2y).num isa Sym
    @test (2x/3y).num.coeff == 2
    @test (2x/3y).den.coeff == 3
    @test (2x/-3x) == -2//3
    @test (2.5x/3x).num == 2.5
    @test (2.5x/3x).den == 3
    @test (x/3x) == 1//3
end

@testset "LiteralReal" begin
    @syms x::LiteralReal y::LiteralReal z::LiteralReal
    @test repr(x+x) == "x + x"
    @test repr(x*x) == "x*x"
    @test repr(x*x + x*x) == "x*x + x*x"
    for ex in [sin(x), x+x, x*x, x\x, x/x]
        @test typeof(sin(x)) <: Term{LiteralReal}
    end
    @test repr(sin(x) + sin(x)) == "sin(x) + sin(x)"
end


using Base.Filesystem:ispath
using Printf:@sprintf
using Plots
function download_dataset(filename)
    url = "http://elib.zib.de/pub/mp-testdata/tsp/tsplib/tsp/"*filename
    #println(url)
    if !isfile(filename)
        download(url,filename)
    else
        message = @sprintf("The file %s is exists!",filename)
        @info message
    end
end
function get_data(filename)
    data_list = []
    open(filename,"r") do stream
        flag = false
        for line in eachline(stream)
            if occursin("EOF",line)
                break
            elseif occursin("NODE_COORD_SECTION",line)
                flag = true
            elseif flag
                tp_list = Int64[]
                output = split(line,['\t',' '])
                vala = parse(Float32,output[1])
                valb = parse(Float32,output[2])
                valc = parse(Float32,output[3])
                append!(tp_list,vala)
                append!(tp_list,valb)
                append!(tp_list,valc)
                append!(data_list,tp_list)
            end
        end
        return data_list
    end
    arr_len = Int64(length(data_list)/3)
    data_list = reshape(data_list,3,arr_len)
    return data_list'
end
function main()
    filename = "eil51.tsp"
    download_dataset(filename)
    citys_mat = get_data(filename)
    # 初始化参数
    t0, tf, alpha, markov_len = 100, 1, 0.99, 10000
    # 获取城市的个数
    citys_num = size(citys_mat)[1]
    # 获取邻接矩阵
    citys_x = reshape(citys_mat[:,2],citys_num,1).*ones(1,citys_num)
    citys_y = reshape(citys_mat[:,3],citys_num,1).*ones(1,citys_num)
    citys_index = citys_mat[:,1]
    citys_distance = sqrt.((citys_x-citys_x').^2+(citys_y-citys_y').^2)
    E_current = Inf32 # 当前解对应的距离
    E_best = Inf32 # 表示的是最优解
    sol_new = [i for i in range(1,stop=citys_num,step=1)] # 初始化解
    sol_current =copy(sol_new)
    sol_best = copy(sol_new)
    len_list = Float32[]
    t = t0
    epoches = 0
    while t>=tf
        for step=1:markov_len
            # 产生随机扰动
            if rand() < 0.5
                # 两交换
                ind1,ind2 = 0,0
                while (ind1 == ind2)
                    ind1 = Int(ceil(rand()*citys_num))
                    ind2 = Int(ceil(rand()*citys_num))
                    sol_new[ind1],sol_new[ind2] = sol_new[ind2],sol_new[ind1]
                end
            else
                ind1,ind2,ind3 = 0,0,0
                while (ind1 == ind2) || (ind1==ind3)|| (ind2==ind3)
                    ind1 = Int(ceil(rand()*citys_num))
                    ind2 = Int(ceil(rand()*citys_num))
                    ind3 = Int(ceil(rand()*citys_num))
                end
                tp_list = [ind1,ind2,ind3]
                tp_list = sort(tp_list)
                ind1,ind2,ind3 = tp_list[1],tp_list[2],tp_list[3]
                tpa = copy(sol_new[ind1:ind2-1])
                tpb = copy(sol_new[ind2:ind3-1])
                sol_new[ind1:ind1+ind3-ind2-1],sol_new[ind1+ind3-ind2:ind3-1] = tpb,tpa
            end
            # 计算目标函数值
            E_new = 0
            for k in range(1,stop=citys_num-1,step=1)
                E_new+= citys_distance[sol_new[k],sol_new[k+1]]
            end
            E_new += citys_distance[sol_new[1],sol_new[end]]
            if E_new <E_current
                E_current = E_new
                sol_current = copy(sol_new)
                if E_new < E_best
                    E_best = E_new
                    sol_best = copy(sol_new)
                end
            else
                if rand()<exp(-(E_new-E_current)/t)
                    E_current = E_new
                    sol_current = copy(sol_new)
                else
                    sol_new = copy(sol_current)
                end
            end
        end
        append!(len_list,E_best)
        epoches += 1
        message = @sprintf("epoches: %d,temperature:%.2f,length:%0.4f",epoches,t,E_best)
        if E_best<0.1
            break
        end
        t = t*alpha
        @info message
    end
    #draw(len_list,sol_best,citys_mat)
end
#=
function draw(len_list,points_list,citys_mat)
    #plot(len_list,linewidth=2,title="The best length")
    #savefig("length.png")
    #closeall()
    len_arr = length(points_list)
    for k in range(1,stop=len_arr-1,step=1)
        start_id = points_list[k]
        end_id = points_list[k+1]
        tp_x = [citys_mat[start_id,2],citys_mat[end_id,2]]
        tp_y = [citys_mat[start_id,3],citys_mat[end_id,3]]
        plot!(tp_x,tp_y,linewidth=2)
    end
    start_id = points_list[end]
    end_id = points_list[1]
    tp_x = [citys_mat[start_id,2],citys_mat[end_id,2]]
    tp_y = [citys_mat[start_id,3],citys_mat[end_id,3]]
    title!("The path of the route")
    savefig("points.png")
    closeall()
end
=#
main()


#-------------------------------------------------------------------------------
#transformed coordinates are handy

#upward calculations
P2ω(P) = -sqrt(P)
ω2P(ω) = ω*ω
dωfac(P) = 2*sqrt(P)
P2ω(Pₛ, Pₜ) = P2ω(Pₛ), P2ω(Pₜ)

#downward calculations
P2ι(P) = sqrt(P)
ι2P(ι) = ι*ι
dιfac(P) = 2*sqrt(P)
P2ι(Pₜ, Pₛ) = P2ι(Pₜ), P2ι(Pₛ)

#-------------------------------------------------------------------------------

export pressuregrid
function pressuregrid(Pₜ, Pₛ, n)
    @assert Pₛ > Pₜ
    @assert n >= 3
    exp.(chebygrid(log(Pₜ), log(Pₛ), n))
end

export trapz
function trapz(x::AbstractVector, y::AbstractVector)
    @assert length(x) == length(y) "vectors must be equal length"
    s = zero(eltype(y))
    @inbounds for i ∈ 1:length(x) - 1
        s += (x[i+1] - x[i])*(y[i] + y[i+1])/2
    end
    return s
end

export meshgrid
function meshgrid(x::AbstractVector, y::AbstractVector)
    X = x' .* ones(length(y))
    Y = y .* ones(length(x))'
    return X, Y
end

export logrange
function logrange(a, b, N::Int=101, γ::Real=1)
    ((10 .^ LinRange(0, γ, N)) .- 1)*(b - a)/(10^γ - 1) .+ a
end

#evaluates the 0th, 1st, and 2nd derivative of a parabola through three points, anywhere
function quaddiff(x, x₁, x₂, x₃, y₁, y₂, y₃)
    z₁ = x₁^2 - x₂^2
    z₂ = x₂^2 - x₃^2
    w = z₁/z₂
    b = (y₁ - y₂ - (y₂ - y₃)*w)/(x₁ - x₂ - (x₂ - x₃)*w)
    a = (y₂ - y₃ - b*(x₂ - x₃))/z₂
    c = y₁ - (a*x₁^2 + b*x₁)
    f = a*x^2 + b*x + c
    f′ = 2*a*x + b
    f′′ = 2*a
    return f, f′, f′′
end

function quaddiff(x, xₚ, yₚ)
    @assert length(xₚ) == length(yₚ) == 3
    @inbounds quaddiff(x, xₚ[1], xₚ[2], xₚ[3], yₚ[1], yₚ[2], yₚ[3])[2]
end

export deriv
function deriv(x, y, T)
    d = similar(x)
    d[1] = (y[2] - y[1])/(x[2] - x[1])
    for i in 2:length(d)-1
        #d₁ = (y[i] - y[i-1])/(x[i] - x[i-1])
        #d₂ = (y[i+1] - y[i])/(x[i+1] - x[i])
        #if i == 25
        #    d[i] = (d₁ + d₂)/2
        #elseif i > 25
        #    d[i] = d₁
        #else
        #    d[i] = d₂
        #end
        #if d₁*d₂ > 0
        #    d[i] = (d₁ + d₂)/2
        #else
        #    if T[i-1] < T[i+1]
        #        d[i] = d₁
        #    else
        #        d[i] = d₂
        #    end
        #end
        d[i] = quaddiff(x[i], view(x,i-1:i+1), view(y,i-1:i+1))
    end
    d[end] = (y[end] - y[end-1])/(x[end] - x[end-1])
    return d
end

#------------------------------------------------------------------------------
# a couple of root finding methods

terminate(a, b, tol)::Bool = abs(a - b) < (tol + tol*abs(b)) ? true : false

function terminate(x₁, x₂, y₁, y₂, tol)::Bool
    terminate(x₁, x₂, tol) && terminate(y₁, y₂, tol) && return true
    return false
end

export regulafalsi
function regulafalsi(F, x₁, x₂, p=nothing; tol=1e-6)
    @assert x₁ != x₂ "starting points must not be identical"
    y₁ = F(x₁, p)
    y₁ == 0 && return x₁
    y₂ = F(x₂, p)
    y₂ == 0 && return x₂
    @assert sign(y₁) != sign(y₂) "regula falsi non-bracketing"
    yₘ = floatmax(y₂)
    yₚ = zero(y₂)
    n = 0
    while !terminate(x₁, x₂, yₚ, yₘ, tol) || (n < 2)
        #store previous evaluation
        yₚ = yₘ
        #approximate zero
        xₘ = x₁ - y₁*(x₂ - x₁)/(y₂ - y₁)
        yₘ = F(xₘ, p)
        #reduce the interval
        if y₁*yₘ > 0
            x₁ = xₘ
        else
            x₂ = xₘ
        end
        #count
        n += 1
    end
    return (x₁ + x₂)/2
end

export secant
function secant(F, x₁, x₂, p=nothing; tol=1e-6)
    @assert x₁ != x₂ "starting points must not be identical"
    y₁ = F(x₁, p)
    y₁ == 0 && return x₁
    y₂ = F(x₂, p)
    y₂ == 0 && return x₂
    y₃ = floatmax(y₂)
    x₃ = zero(x₂)
    n = 0
    while !terminate(x₁, x₂, y₁, y₂, tol) || (n < 2)
        #approximate zero
        x₃ = x₁ - y₁*(x₂ - x₁)/(y₂ - y₁)
        y₃ = F(x₃, p)
        #swap values
        x₁, x₂ = x₂, x₃
        y₁, y₂ = y₂, y₃
        #count
        n += 1
    end
    return x₃
end


# Using IterationControl.jl to control a particle tracker

using Pkg
Pkg.activate(@__DIR__)
Pkg.instantiate()

using Plots
pyplot(size = (600, 300*(sqrt(5) -1)))

using IterationControl

# Consider an object that tracks a particle in one dimension, moving
# in discrete time at a speed proportional to the distance away from
# some moving target. We are viewing this particle as a toy example of
# online learning and call the constant of proportionality `η` a
# *learning rate*.

# The particle is initially at rest at the origin, which is also where
# the target begins.

# Calling `train!` on the particle moves it along for the specified
# number of time steps, while calling `ingest!` updates the target
# position ("data ingestion").

mutable struct Particle
    position::Float64
    target::Float64
    velocity::Float64
    η::Float64 # learning rate
    Particle(η) = new(0.0, 0.0, 0.0, η)
end

Particle(; η=0.1) = Particle(η)

loss(model::Particle) =  abs(model.target - model.position)

# A train method for one time step:

function train!(model::Particle)
    model.velocity = model.η*(model.target - model.position)
    model.position = model.position + model.velocity
end

# And for multiple time steps:

function train!(model, n)
    foreach(1:n) do _
        train!(model)
    end
    return nothing
end

# A method for updating the target:

function ingest!(model::Particle, target)
    model.target = target
    return nothing
end

# Lifting the `train!`, `loss` and `ingest!` methods:

IterationControl.train!(model::Particle, n) = train!(model, n)
IterationControl.loss(model::Particle) = loss(model)
IterationControl.ingest!(model::Particle, datum) = ingest!(model, datum)

# Lets `train` our particle by alternatively setting the target
# position to -1 or +1 every 5 time steps, for a total of 40 time
# steps (four -1/+1 cycles). We will then then allow our particle to "cruise"
# with the target fixed thereafter (the provided `data` iterator
# having been exhausted). This cruising is allowed to continue until
# the particle is within 0.01 of the target.

# We use a learning rate of 0.1.

model = Particle(0.1)
data = repeat([-1, 1], outer=4)

#-

target_positions = Float64[]
particle_positions = Float64[]

function callback!(model)
    push!(target_positions, model.target)
    push!(particle_positions, model.position)
end

IterationControl.train!(model,
                        Data(data),
                        Step(5),
                        Callback(callback!),
                        Threshold(0.01),
                        TimeLimit(0.005));

scatter(target_positions, label="target", ms=6)
scatter!(particle_positions, label="particle", ms=6)

#-

loss(model)

# Switching the target just once more to the -1 position and
# continuing training until the particle and target are withing 0.01:

data = [-1,]

IterationControl.train!(model,
                        Data(data),
                        Step(5),
                        Callback(callback!),
                        Threshold(0.01),
                        TimeLimit(0.005));

scatter(target_positions, label="target", ms=6)
scatter!(particle_positions, label="particle", ms=6)

using Literate #src
Literate.markdown(@__FILE__, @__DIR__, execute=false) #src
Literate.notebook(@__FILE__, @__DIR__, execute=true) #src


# source code representation cannot be recovered, but code_lowered yes
# see for example Symbolics.jl or Cassette.jl

# unary minus, plus

# some unary functions:
# sqrt, cbrt, exp, log, log2, log10

# some trig functions:
# sin    cos    tan    cot    sec    csc
# sinh   cosh   tanh   coth   sech   csch
# asin   acos   atan   acot   asec   acsc
# asinh  acosh  atanh  acoth  asech  acsch
# sinc   cosc

# Implement:
# unary minus, plus
# sqrt, cbrt, exp, log, log2, log10
# sin    cos    tan

# simplifying unary minus

module symbolics

function default_variable_value()
    return 0.0
end

abstract type AbstractVariable end
abstract type AbstractSymbolicExpression end
abstract type AbstractConstant end

include("is_operators.jl")
include("constant.jl")
include("variable.jl")
include("symbolic_expression.jl")
include("basic_operations.jl")
include("to_string.jl")
include("derivative.jl")

function is_equal(x, y)
    return false
end

function is_equal(x::Constant, y::Constant)
    return x.value == y.value
end

function is_equal(x::Variable, y::Variable)
    return x.name == y.name
end

function is_equal(x::SymbolicExpression, y::SymbolicExpression)
    # how to deal with permutations on expressions?
    # e.g.
    # a + b ?= b + a, should be, but
    # (a + b) + c ?= a + (b + c); this should be also true, but much harder to operate on
    # a + b + c + d -> possibly could be represented as +(a, b, c, d)
    # but then it is required to rethink the binary operations
    return false
end

# this is not what I want
# I dont want to compare variables by pointer -> already provided by Julia
# but I'd like to compare whether they are equaivalent
#
# function Base.:(==)(x::Constant, y::Constant)
#     # by value
#     return x.value == y.value
# end

# function Base.:(==)(x::Variable, y::Variable)
#     # by pointer
#     return x == y
# end

# function Base.:(==)(x::SymbolicExpression, y::SymbolicExpression)
#     # by pointer
#     return x == y
# end

function __simplify_plus(ex)
    # op +
    # ex + 0 -> ex
    # 0 + ex -> ex
    # c1 + c2 -> c3(c1.value + c2.value)
    # x + x -> 2 * x
    # ex + ex -> 2 * ex

    is_zero = is_zero_constant.(ex.args)
    if all(is_zero)
        return Constant(0.0)
    end

    is_not_zero = .!is_zero
    if any(is_not_zero)
        arg_vec = ex.args[is_not_zero]
        if length(arg_vec) == 1
            simpli_expr = arg_vec[1]
        else
            b_consts = is_constant.(arg_vec)
            idx      = findall(b_consts)
            if !isempty(idx)
                c        = reduce(+, arg_vec[idx])
                arg_vec  = [arg_vec[.!b_consts]; c]
            end
            simpli_expr = SymbolicExpression(:+, arg_vec...)
        end
        return simpli_expr
    end

    return ex
end

function __simplify_minus(ex)
    # op -
    # ex - 0 -> ex
    # 0 - ex -> -ex
    # c1 - c2 -> c3(c1.value - c2.value)
    # x - x -> 0
    # ex - ex -> 0

    L = length(ex.args)
    if L == 1
        if is_constant(ex.args[1])
            return Constant( -ex.args[1].value )
        else
            return ex
        end
    end

    if is_zero_constant(ex.args[1])
        return -ex.args[2]
    end

    if is_zero_constant(ex.args[2])
        return ex.args[1]
    end

    if is_equal(ex.args[1], ex.args[2])
        return Constant(0.0)
    end

    lhs = ex.args[1]
    rhs = ex.args[2]
    if is_constant(lhs) && is_constant(rhs)
        return lhs - rhs
    end

    return ex
end

function __simplify_product(ex)
    # op *
    # 0 * ex -> 0
    # ex * 0 -> 0
    # 1 * ex -> ex
    # ex * 1 -> ex

    is_zero = is_zero_constant.(ex.args)
    if any(is_zero)
        return Constant(0.0)
    end

    is_one     = is_one_constant.(ex.args)
    is_not_one = .!is_one
    if any(is_not_one)
        arg_vec = ex.args[is_not_one]
        if length(arg_vec) == 1
            simpli_expr = arg_vec[1]
        else

            b_consts = is_constant.(arg_vec)
            idx      = findall(b_consts)
            if !isempty(idx)
                c        = reduce(*, arg_vec[idx])
                arg_vec  = [arg_vec[.!b_consts]; c]
            end
            simpli_expr = SymbolicExpression(:*, arg_vec...)
        end
        return simpli_expr
    end

    return ex
end

function __simplify_divide(ex)
    # op /
    # ex / ex -> 1
    # 0 / rhs -> 0
    # rhs / 0 -> inf
    #
    # if (subexpr1 * subexpr2) / (subexpr1 * subexpr3)
    # -> subexpr2 / subexpr3

    if is_zero_constant(ex.args[1])
        return Constant(0.0)
    end

    if is_zero_constant(ex.args[2])
        return Constant(Inf)
    end

    if is_equal(ex.args[1], ex.args[2])
        return Constant(1.0)
    end

    lhs = ex.args[1]
    rhs = ex.args[2]
    if is_constant(lhs) && is_constant(rhs)
        return lhs / rhs
    end

    return ex
end

function __simplify_power(ex)
    # op ^
    # x^i / x^j -> x^(i-j)
    # x^0 -> 1
    # x^1 -> x
    # 1^x

    if is_zero_constant(ex.args[2])
        return simpli_expr = Constant(1.0)
    end

    if is_one_constant(ex.args[1])
        return simpli_expr = Constant(1.0)
    end

    if is_one_constant(ex.args[2])
        return simpli_expr = ex.args[1]
    end

    lhs = ex.args[1]
    rhs = ex.args[2]
    if is_constant(lhs) && is_constant(rhs)
        return lhs^rhs
    end

    return ex
end

function simplify(ex::Constant)
    return ex
end

function simplify(ex::Variable)
    return ex
end

function simplify(ex::SymbolicExpression)

    # rewriting rules
    # zero element
    # one element
    #
    # harder ones
    # commutative laws
    # associative laws
    #
    # a + b -> b + a
    # (a + b) + c -> a + (b + c)

    return __rec_simplify(ex)
end

function __rec_simplify(ex::Constant)
    return ex
end

function __rec_simplify(ex::Variable)
    return ex
end

function __rec_simplify(ex::SymbolicExpression)

    L           = length(ex.args)
    sym         = ex.head
    simpli_args = Vector{Any}(undef, L)

    for ii = 1:L
        simpli_args[ii] = __rec_simplify(ex.args[ii])
    end

    simpli_expr = SymbolicExpression(sym, simpli_args...)

    if is_plus_operator(sym)
        simpli_expr = __simplify_plus(simpli_expr)
    elseif is_minus_operator(sym)
        simpli_expr = __simplify_minus(simpli_expr)
    elseif is_product_operator(sym)
        simpli_expr = __simplify_product(simpli_expr)
    elseif is_divide_operator(sym)
        simpli_expr = __simplify_divide(simpli_expr)
    elseif is_power_operator(sym)
        simpli_expr = __simplify_power(simpli_expr)
    # else
    #     simpli_expr = ex
    end

    return simpli_expr
end

function simplify(ex_vec::Vector)
    L = length(ex_vec)
    simpli_vec = Vector{Any}(undef, L)
    for ii = 1:L
        simpli_vec[ii] = simplify(ex_vec[ii])
    end
    return simpli_vec
end

function simplify(ex_mat::Matrix)
    L = length(ex_mat)
    (n_row, n_col) = size(ex_mat)
    simpli_mat = Matrix{Any}(undef, n_row, n_col)
    for ii = 1:L
        simpli_mat[ii] = simplify(ex_mat[ii])
    end
    return simpli_mat
end

end # module

module stest
using ..symbolics

# symbolics.@variable x 10
# symbolics.@variable y 0.2
# symbolics.@variable z 3
# e1 = x + y
# e2 = x / y
# e3 = z ^ (x + y)
# e4 = -x
# e5 = -x - y
# e6 = sin(x + y - z * x)
# e3_val = symbolics.evaluate(e3)
# str1 = symbolics.to_string( stest.e1 )
# str2 = symbolics.to_string( stest.e2 )
# str3 = symbolics.to_string( stest.e3 )
# str4 = symbolics.to_string( stest.e4 )
# str5 = symbolics.to_string( stest.e5 )
# str6 = symbolics.to_string( stest.e6 )
# e7 = y * 10
# str7 = symbolics.to_string( stest.e7 )
# e8 = 10 * y
# str8 = symbolics.to_string( stest.e8 )
# e9 = 10 + (x - y)
# str9 = symbolics.to_string( stest.e9 )
# # !!!
# # this doesnt work correctly now
# # e10 = y * 10 * (x * 2 - 3*z)
# e10 = x * y * x * (x * 2 - 3*z)
# str10 = symbolics.to_string( stest.e10 )

function paren_test()

    symbolics.@variable x 10
    symbolics.@variable y 0.2
    symbolics.@variable z 3

    e1 = x + y
    e2 = x / y

    e3 = z ^ (x + y)

    e4 = -x
    e5 = -x - y
    e6 = sin(x + y - z * x)

    # e3_val = symbolics.evaluate(e3)

    # str1 = symbolics.to_string( stest.e1 )
    # str2 = symbolics.to_string( stest.e2 )
    # str3 = symbolics.to_string( stest.e3 )
    # str4 = symbolics.to_string( stest.e4 )
    # str5 = symbolics.to_string( stest.e5 )
    # str6 = symbolics.to_string( stest.e6 )

    # e7 = y * 10
    # str7 = symbolics.to_string( stest.e7 )
    # e8 = 10 * y
    # str8 = symbolics.to_string( stest.e8 )

    # e9 = 10 + (x - y)
    # str9 = symbolics.to_string( stest.e9 )

    # !!!
    # this doesnt work correctly now
    # e7 = y * 10 * (x * 2 - 3*z)
    e7 = x * y * x * (x * 2 - 3*z)
    str7 = symbolics.to_string( stest.e7 )

    # str = symbolics.to_string( stest.e6 )

    # x_val = symbolics.evaluate(x)
    # y_val = symbolics.evaluate(x)

    # e1_val = symbolics.evaluate(e1)

end

function test_derivative_simpli()
    symbolics.@variable x 10
    symbolics.@variable y 0.2
    symbolics.@variable z 3
    symbolics.@variable w 5

    u = -x

    ex1 = x + y
    ex2 = x * y + z
    ex3 = x + y + z
    ex4 = x * y * z
    ex5 = sin(x)

    # str = symbolics.to_string(ex2)

    der_ex3       = symbolics.derivative(stest.ex3, stest.x)
    der_str3      = symbolics.to_string(der_ex3)
    @time s_der_ex3     = symbolics.simplify(der_ex3)
    str_s_der_ex3 = symbolics.to_string(s_der_ex3)

    der_ex4       = symbolics.derivative(stest.ex4, stest.x)
    der_str4      = symbolics.to_string(der_ex4)
    @time s_der_ex4     = symbolics.simplify(der_ex4)
    str_s_der_ex4 = symbolics.to_string(s_der_ex4)
end

function harmonic_osc(der_q, time, q)
    der_q[1] = q[2]
    der_q[2] = -5.0 * q[1] - q[2]

    return der_q
end

# symbolics.@variable x 10
# symbolics.@variable v 0.2

# q = Vector{symbolics.Variable}(undef, 2)
# der_q = Vector{Any}(undef, 2)

# q[1] = x
# q[2] = v

# der_q = harmonic_osc(der_q, 0.0, q)

# jac = symbolics.jacobian(der_q, q)
# sjac = symbolics.simplify(jac)
# sjac = symbolics.simplify(sjac)

# symbolics.@variable a 30.0

# expr1 = a + 1.0

end


# In this file, doctests which differ in the printed Float32 values won't fail
```@meta
DocTestFilters = r"[0-9\.]+f0"
```

"""
    mae(ŷ, y; agg = mean)

Return the loss corresponding to mean absolute error:

    agg(abs.(ŷ .- y))

# Example
```jldoctest
julia> y_model = [1.1, 1.9, 3.1];

julia> Flux.mae(y_model, 1:3)
0.10000000000000009
```
"""
function mae(ŷ, y; agg = mean)
  _check_sizes(ŷ, y)
  agg(abs.(ŷ .- y))
end

"""
    mse(ŷ, y; agg = mean)

Return the loss corresponding to mean square error:

    agg((ŷ .- y) .^ 2)

See also: [`mae`](@ref), [`msle`](@ref), [`crossentropy`](@ref).

# Example
```jldoctest
julia> y_model = [1.1, 1.9, 3.1];

julia> y_true = 1:3;

julia> Flux.mse(y_model, y_true)
0.010000000000000018
```
"""
function mse(ŷ, y; agg = mean)
  _check_sizes(ŷ, y)
  agg(abs2.(ŷ .- y))
end

"""
    msle(ŷ, y; agg = mean, ϵ = eps(ŷ))

The loss corresponding to mean squared logarithmic errors, calculated as

    agg((log.(ŷ .+ ϵ) .- log.(y .+ ϵ)) .^ 2)

The `ϵ` term provides numerical stability.
Penalizes an under-estimation more than an over-estimatation.

# Example
```jldoctest
julia> Flux.msle(Float32[1.1, 2.2, 3.3], 1:3)
0.009084041f0

julia> Flux.msle(Float32[0.9, 1.8, 2.7], 1:3)
0.011100831f0
```
"""
function msle(ŷ, y; agg = mean, ϵ = epseltype(ŷ))
  _check_sizes(ŷ, y)
  agg((log.((ŷ .+ ϵ) ./ (y .+ ϵ))) .^2 )
end

"""
    huber_loss(ŷ, y; δ = 1, agg = mean)

Return the mean of the [Huber loss](https://en.wikipedia.org/wiki/Huber_loss)
given the prediction `ŷ` and true values `y`.

                 | 0.5 * |ŷ - y|^2,            for |ŷ - y| <= δ
    Huber loss = |
                 |  δ * (|ŷ - y| - 0.5 * δ), otherwise
"""
function huber_loss(ŷ, y; agg = mean, δ = ofeltype(ŷ, 1))
   _check_sizes(ŷ, y)
   abs_error = abs.(ŷ .- y)
   #TODO: remove dropgrad when Zygote can handle this function with CuArrays
   temp = Zygote.dropgrad(abs_error .<  δ)
   x = ofeltype(ŷ, 0.5)
   agg(((abs_error .^ 2) .* temp) .* x .+ δ * (abs_error .- x * δ) .* (1 .- temp))
end

"""
    label_smoothing(y::Union{Number, AbstractArray}, α; dims::Int=1)

Returns smoothed labels, meaning the confidence on label values are relaxed.

When `y` is given as one-hot vector or batch of one-hot, its calculated as

    y .* (1 - α) .+ α / size(y, dims)

when `y` is given as a number or batch of numbers for binary classification,
its calculated as

    y .* (1 - α) .+ α / 2

in which case the labels are squeezed towards `0.5`.

α is a number in interval (0, 1) called the smoothing factor. Higher the
value of α larger the smoothing of `y`.

`dims` denotes the one-hot dimension, unless `dims=0` which denotes the application
of label smoothing to binary distributions encoded in a single number.

# Example
```jldoctest
julia> y = Flux.onehotbatch([1, 1, 1, 0, 1, 0], 0:1)
2×6 OneHotMatrix(::Vector{UInt32}) with eltype Bool:
 ⋅  ⋅  ⋅  1  ⋅  1
 1  1  1  ⋅  1  ⋅

julia> y_smoothed = Flux.label_smoothing(y, 0.2f0)
2×6 Matrix{Float32}:
 0.1  0.1  0.1  0.9  0.1  0.9
 0.9  0.9  0.9  0.1  0.9  0.1

julia> y_sim = softmax(y .* log(2f0))
2×6 Matrix{Float32}:
 0.333333  0.333333  0.333333  0.666667  0.333333  0.666667
 0.666667  0.666667  0.666667  0.333333  0.666667  0.333333

julia> y_dis = vcat(y_sim[2,:]', y_sim[1,:]')
2×6 Matrix{Float32}:
 0.666667  0.666667  0.666667  0.333333  0.666667  0.333333
 0.333333  0.333333  0.333333  0.666667  0.333333  0.666667

julia> Flux.crossentropy(y_sim, y) < Flux.crossentropy(y_sim, y_smoothed)
true

julia> Flux.crossentropy(y_dis, y) > Flux.crossentropy(y_dis, y_smoothed)
true
```
"""
function label_smoothing(y::Union{AbstractArray,Number}, α::Number; dims::Int = 1)
    if !(0 < α < 1)
        throw(ArgumentError("α must be between 0 and 1"))
    end
    if dims == 0
        y_smoothed = y .* (1 - α) .+ α*1//2
    elseif dims == 1
        y_smoothed = y .* (1 - α) .+ α* 1 // size(y, 1)
    else
        throw(ArgumentError("`dims` should be either 0 or 1"))
    end
    return y_smoothed
end

"""
    crossentropy(ŷ, y; dims = 1, ϵ = eps(ŷ), agg = mean)

Return the cross entropy between the given probability distributions;
calculated as

    agg(-sum(y .* log.(ŷ .+ ϵ); dims))

Cross entropy is typically used as a loss in multi-class classification,
in which case the labels `y` are given in a one-hot format.
`dims` specifies the dimension (or the dimensions) containing the class probabilities.
The prediction `ŷ` is supposed to sum to one across `dims`,
as would be the case with the output of a [`softmax`](@ref) operation.

For numerical stability, it is recommended to use [`logitcrossentropy`](@ref)
rather than `softmax` followed by `crossentropy` .

Use [`label_smoothing`](@ref) to smooth the true labels as preprocessing before
computing the loss.

See also: [`logitcrossentropy`](@ref), [`binarycrossentropy`](@ref), [`logitbinarycrossentropy`](@ref).

# Example
```jldoctest
julia> y_label = Flux.onehotbatch([0, 1, 2, 1, 0], 0:2)
3×5 OneHotMatrix(::Vector{UInt32}) with eltype Bool:
 1  ⋅  ⋅  ⋅  1
 ⋅  1  ⋅  1  ⋅
 ⋅  ⋅  1  ⋅  ⋅

julia> y_model = softmax(reshape(-7:7, 3, 5) .* 1f0)
3×5 Matrix{Float32}:
 0.0900306  0.0900306  0.0900306  0.0900306  0.0900306
 0.244728   0.244728   0.244728   0.244728   0.244728
 0.665241   0.665241   0.665241   0.665241   0.665241

julia> sum(y_model; dims=1)
1×5 Matrix{Float32}:
 1.0  1.0  1.0  1.0  1.0

julia> Flux.crossentropy(y_model, y_label)
1.6076053f0

julia> 5 * ans ≈ Flux.crossentropy(y_model, y_label; agg=sum)
true

julia> y_smooth = Flux.label_smoothing(y_label, 0.15f0)
3×5 Matrix{Float32}:
 0.9   0.05  0.05  0.05  0.9
 0.05  0.9   0.05  0.9   0.05
 0.05  0.05  0.9   0.05  0.05

julia> Flux.crossentropy(y_model, y_smooth)
1.5776052f0
```
"""
function crossentropy(ŷ, y; dims = 1, agg = mean, ϵ = epseltype(ŷ))
  _check_sizes(ŷ, y)
  agg(.-sum(xlogy.(y, ŷ .+ ϵ); dims = dims))
end

"""
    logitcrossentropy(ŷ, y; dims = 1, agg = mean)

Return the cross entropy calculated by

    agg(-sum(y .* logsoftmax(ŷ; dims); dims))

This is mathematically equivalent to `crossentropy(softmax(ŷ), y)`,
but is more numerically stable than using functions [`crossentropy`](@ref)
and [`softmax`](@ref) separately.

See also: [`binarycrossentropy`](@ref), [`logitbinarycrossentropy`](@ref), [`label_smoothing`](@ref).

# Example
```jldoctest
julia> y_label = Flux.onehotbatch(collect("abcabaa"), 'a':'c')
3×7 OneHotMatrix(::Vector{UInt32}) with eltype Bool:
 1  ⋅  ⋅  1  ⋅  1  1
 ⋅  1  ⋅  ⋅  1  ⋅  ⋅
 ⋅  ⋅  1  ⋅  ⋅  ⋅  ⋅

julia> y_model = reshape(vcat(-9:0, 0:9, 7.5f0), 3, 7)
3×7 Matrix{Float32}:
 -9.0  -6.0  -3.0  0.0  2.0  5.0  8.0
 -8.0  -5.0  -2.0  0.0  3.0  6.0  9.0
 -7.0  -4.0  -1.0  1.0  4.0  7.0  7.5

julia> Flux.logitcrossentropy(y_model, y_label)
1.5791205f0

julia> Flux.crossentropy(softmax(y_model), y_label)
1.5791197f0
```
"""
function logitcrossentropy(ŷ, y; dims = 1, agg = mean)
  _check_sizes(ŷ, y)
  agg(.-sum(y .* logsoftmax(ŷ; dims = dims); dims = dims))
end

"""
    binarycrossentropy(ŷ, y; agg = mean, ϵ = eps(ŷ))

Return the binary cross-entropy loss, computed as

    agg(@.(-y * log(ŷ + ϵ) - (1 - y) * log(1 - ŷ + ϵ)))

Where typically, the prediction `ŷ` is given by the output of a [`sigmoid`](@ref) activation.
The `ϵ` term is included to avoid infinity. Using [`logitbinarycrossentropy`](@ref) is recomended
over `binarycrossentropy` for numerical stability.

Use [`label_smoothing`](@ref) to smooth the `y` value as preprocessing before
computing the loss.

See also: [`crossentropy`](@ref), [`logitcrossentropy`](@ref).

# Examples
```jldoctest
julia> y_bin = Bool[1,0,1]
3-element Vector{Bool}:
 1
 0
 1

julia> y_prob = softmax(reshape(vcat(1:3, 3:5), 2, 3) .* 1f0)
2×3 Matrix{Float32}:
 0.268941  0.5  0.268941
 0.731059  0.5  0.731059

julia> Flux.binarycrossentropy(y_prob[2,:], y_bin)
0.43989f0

julia> all(p -> 0 < p < 1, y_prob[2,:])  # else DomainError
true

julia> y_hot = Flux.onehotbatch(y_bin, 0:1)
2×3 OneHotMatrix(::Vector{UInt32}) with eltype Bool:
 ⋅  1  ⋅
 1  ⋅  1

julia> Flux.crossentropy(y_prob, y_hot)
0.43989f0
```
"""
function binarycrossentropy(ŷ, y; agg = mean, ϵ = epseltype(ŷ))
  _check_sizes(ŷ, y)
  agg(@.(-xlogy(y, ŷ + ϵ) - xlogy(1 - y, 1 - ŷ + ϵ)))
end

"""
    logitbinarycrossentropy(ŷ, y; agg = mean)

Mathematically equivalent to
[`binarycrossentropy(σ(ŷ), y)`](@ref) but is more numerically stable.

See also: [`crossentropy`](@ref), [`logitcrossentropy`](@ref).

# Examples
```jldoctest
julia> y_bin = Bool[1,0,1];

julia> y_model = Float32[2, -1, pi]
3-element Vector{Float32}:
  2.0
 -1.0
  3.1415927

julia> Flux.logitbinarycrossentropy(y_model, y_bin)
0.160832f0

julia> Flux.binarycrossentropy(sigmoid.(y_model), y_bin)
0.16083185f0
```
"""
function logitbinarycrossentropy(ŷ, y; agg = mean)
  _check_sizes(ŷ, y)
  agg(@.((1 - y) * ŷ - logσ(ŷ)))
end

"""
    kldivergence(ŷ, y; agg = mean, ϵ = eps(ŷ))

Return the
[Kullback-Leibler divergence](https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence)
between the given probability distributions.

The KL divergence is a measure of how much one probability distribution is different
from the other. It is always non-negative, and zero only when both the distributions are equal.

# Example
```jldoctest
julia> p1 = [1 0; 0 1]
2×2 Matrix{Int64}:
 1  0
 0  1

julia> p2 = fill(0.5, 2, 2)
2×2 Matrix{Float64}:
 0.5  0.5
 0.5  0.5

julia> Flux.kldivergence(p2, p1) ≈ log(2)
true

julia> Flux.kldivergence(p2, p1; agg = sum) ≈ 2log(2)
true

julia> Flux.kldivergence(p2, p2; ϵ = 0)  # about -2e-16 with the regulator
0.0

julia> Flux.kldivergence(p1, p2; ϵ = 0)  # about 17.3 with the regulator
Inf
```
"""
function kldivergence(ŷ, y; dims = 1, agg = mean, ϵ = epseltype(ŷ))
  _check_sizes(ŷ, y)
  entropy = agg(sum(xlogx.(y), dims = dims))
  cross_entropy = crossentropy(ŷ, y; dims = dims, agg = agg, ϵ = ϵ)
  return entropy + cross_entropy
end

"""
    poisson_loss(ŷ, y)

# Return how much the predicted distribution `ŷ` diverges from the expected Poisson
# distribution `y`; calculated as `sum(ŷ .- y .* log.(ŷ)) / size(y, 2)`.

[More information.](https://peltarion.com/knowledge-center/documentation/modeling-view/build-an-ai-model/loss-functions/poisson).
"""
function poisson_loss(ŷ, y; agg = mean)
  _check_sizes(ŷ, y)
  agg(ŷ .- xlogy.(y, ŷ))
end

"""
    hinge_loss(ŷ, y; agg = mean)

Return the [hinge_loss loss](https://en.wikipedia.org/wiki/Hinge_loss) given the
prediction `ŷ` and true labels `y` (containing 1 or -1); calculated as
`sum(max.(0, 1 .- ŷ .* y)) / size(y, 2)`.

See also: [`squared_hinge_loss`](@ref)
"""
function hinge_loss(ŷ, y; agg = mean)
  _check_sizes(ŷ, y)
  agg(max.(0, 1 .- ŷ .* y))
end

"""
    squared_hinge_loss(ŷ, y)

Return the squared hinge_loss loss given the prediction `ŷ` and true labels `y`
(containing 1 or -1); calculated as `sum((max.(0, 1 .- ŷ .* y)).^2) / size(y, 2)`.

See also: [`hinge_loss`](@ref)
"""
function squared_hinge_loss(ŷ, y; agg = mean)
  _check_sizes(ŷ, y)
  agg((max.(0, 1 .- ŷ .* y)) .^ 2)
end

"""
    dice_coeff_loss(ŷ, y; smooth = 1)

Return a loss based on the dice coefficient.
Used in the [V-Net](https://arxiv.org/abs/1606.04797) image segmentation
architecture.
Similar to the F1_score. Calculated as:

    1 - 2*sum(|ŷ .* y| + smooth) / (sum(ŷ.^2) + sum(y.^2) + smooth)
"""
function dice_coeff_loss(ŷ, y; smooth = ofeltype(ŷ, 1.0))
  _check_sizes(ŷ, y)
  1 - (2 * sum(y .* ŷ) + smooth) / (sum(y .^ 2) + sum(ŷ .^ 2) + smooth) #TODO agg
end

"""
    tversky_loss(ŷ, y; β = 0.7)

Return the [Tversky loss](https://arxiv.org/abs/1706.05721).
Used with imbalanced data to give more weight to false negatives.
Larger β weigh recall more than precision (by placing more emphasis on false negatives)
Calculated as:
    1 - sum(|y .* ŷ| + 1) / (sum(y .* ŷ + β*(1 .- y) .* ŷ + (1 - β)*y .* (1 .- ŷ)) + 1)
"""
function tversky_loss(ŷ, y; β = ofeltype(ŷ, 0.7))
    _check_sizes(ŷ, y)
    #TODO add agg
    num = sum(y .* ŷ) + 1
    den = sum(y .* ŷ + β * (1 .- y) .* ŷ + (1 - β) * y .* (1 .- ŷ)) + 1
    1 - num / den
end

"""
    binary_focal_loss(ŷ, y; agg=mean, γ=2, ϵ=eps(ŷ))

Return the [binary_focal_loss](https://arxiv.org/pdf/1708.02002.pdf)
The input, 'ŷ', is expected to be normalized (i.e. [`softmax`](@ref) output).

For `γ == 0`, the loss is mathematically equivalent to [`Losses.binarycrossentropy`](@ref).

# Example
```jldoctest
julia> y = [0  1  0
            1  0  1]
2×3 Matrix{Int64}:
 0  1  0
 1  0  1

julia> ŷ = [0.268941  0.5  0.268941
            0.731059  0.5  0.731059]
2×3 Matrix{Float64}:
 0.268941  0.5  0.268941
 0.731059  0.5  0.731059

julia> Flux.binary_focal_loss(ŷ, y) ≈ 0.0728675615927385
true
```

See also: [`Losses.focal_loss`](@ref) for multi-class setting

"""
function binary_focal_loss(ŷ, y; agg=mean, γ=2, ϵ=epseltype(ŷ))
    _check_sizes(ŷ, y)
    ŷ = ŷ .+ ϵ
    p_t = y .* ŷ  + (1 .- y) .* (1 .- ŷ)
    ce = -log.(p_t)
    weight = (1 .- p_t) .^ γ
    loss = weight .* ce
    agg(loss)
end

"""
    focal_loss(ŷ, y; dims=1, agg=mean, γ=2, ϵ=eps(ŷ))

Return the [focal_loss](https://arxiv.org/pdf/1708.02002.pdf)
which can be used in classification tasks with highly imbalanced classes.
It down-weights well-classified examples and focuses on hard examples.
The input, 'ŷ', is expected to be normalized (i.e. [`softmax`](@ref) output).

The modulating factor, `γ`, controls the down-weighting strength.
For `γ == 0`, the loss is mathematically equivalent to [`Losses.crossentropy`](@ref).

# Example
```jldoctest
julia> y = [1  0  0  0  1
            0  1  0  1  0
            0  0  1  0  0]
3×5 Matrix{Int64}:
 1  0  0  0  1
 0  1  0  1  0
 0  0  1  0  0

julia> ŷ = softmax(reshape(-7:7, 3, 5) .* 1f0)
3×5 Matrix{Float32}:
 0.0900306  0.0900306  0.0900306  0.0900306  0.0900306
 0.244728   0.244728   0.244728   0.244728   0.244728
 0.665241   0.665241   0.665241   0.665241   0.665241

julia> Flux.focal_loss(ŷ, y) ≈ 1.1277571935622628
true
```

See also: [`Losses.binary_focal_loss`](@ref) for binary (not one-hot) labels

"""
function focal_loss(ŷ, y; dims=1, agg=mean, γ=2, ϵ=epseltype(ŷ))
    _check_sizes(ŷ, y)
    ŷ = ŷ .+ ϵ
    agg(sum(@. -y * (1 - ŷ)^γ * log(ŷ); dims=dims))
end

"""
    siamese_contrastive_loss(ŷ, y; margin = 1, agg = mean)

Return the [contrastive loss](http://yann.lecun.com/exdb/publis/pdf/hadsell-chopra-lecun-06.pdf)
which can be useful for training Siamese Networks. It is given by

    agg(@. (1 - y) * ŷ^2 + y * max(0, margin - ŷ)^2)

Specify `margin` to set the baseline for distance at which pairs are dissimilar.

"""
function siamese_contrastive_loss(ŷ, y; agg = mean, margin::Real = 1)
    _check_sizes(ŷ, y)
    margin < 0 && throw(DomainError(margin, "Margin must be non-negative"))
    return agg(@. (1 - y) * ŷ^2 + y * max(0, margin - ŷ)^2)
end

```@meta
DocTestFilters = nothing
```


using DataStructure101
const DS = DataStructure101

export dfs

function dfs!(g::DS.AbstractGraph, v::Int64, traversed::Vector{Int64}=Int[])
    push!(traversed, v)

    neighbors = DS.neighbor(g, v)
    for n in neighbors
        !(n in traversed) && dfs!(g, n, traversed)
    end

    return traversed
end

dfs(g::DS.AbstractGraph, v::Int64) = dfs!(g, v)


# ---
# title: 381. Insert Delete GetRandom O(1) - Duplicates allowed
# id: problem381
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Array, Hash Table, Design
# link: <https://leetcode.com/problems/insert-delete-getrandom-o1-duplicates-allowed/description/>
# hidden: true
# ---
#
# Design a data structure that supports all following operations in _average_
# **O(1)** time.
#
# **Note: Duplicate elements are allowed.**
#
#   1. `insert(val)`: Inserts an item val to the collection.
#   2. `remove(val)`: Removes an item val from the collection if present.
#   3. `getRandom`: Returns a random element from current collection of elements. The probability of each element being returned is **linearly related** to the number of same value the collection contains.
#
# **Example:**
#
#
#
#     // Init an empty collection.
#     RandomizedCollection collection = new RandomizedCollection();
#
#     // Inserts 1 to the collection. Returns true as the collection did not contain 1.
#     collection.insert(1);
#
#     // Inserts another 1 to the collection. Returns false as the collection contained 1. Collection now contains [1,1].
#     collection.insert(1);
#
#     // Inserts 2 to the collection, returns true. Collection now contains [1,1,2].
#     collection.insert(2);
#
#     // getRandom should return 1 with the probability 2/3, and returns 2 with the probability 1/3.
#     collection.getRandom();
#
#     // Removes 1 from the collection, returns true. Collection now contains [1,2].
#     collection.remove(1);
#
#     // getRandom should return 1 and 2 both equally likely.
#     collection.getRandom();
#
#
#
## @lc code=start
using LeetCode

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


### A Pluto.jl notebook ###
# v0.12.20

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 126b1cca-2476-11eb-087b-474174c4c1d0
begin
begin
begin
	import Pkg
	Pkg.add("Plots")
end

using Plots
end
	plotly()
end

# ╔═╡ e0937cd4-239d-11eb-141d-1be13a185dca
md"# _Ideal Gas Laws!_
All gases are affected by several variables: Volume, temperature, amount, and pressure. Together, they all affect how a gas behaves. But, before we look at those variables, lets learn the basic rules of gases."

# ╔═╡ 17827eb0-239f-11eb-3266-2d71e0aedd9c
md"## Kinetic Theory of Gases
1. **Gases are mostly empty space** ~ the total number of molecules in any space is much smaller compared to the volume of that space
2. **Gas molecules are in constant, chaotic motion** ~ the molecules constantly collide with eachother and the container walls
3. **Collisions are elastic** ~ there is no energy loss during collisions nor are there attractive forces that keep the molecules together
4. **Gas pressure is caused by collisions of molecules against the walls of a container** ~ Pressure increases as the energy and frequency of the collisions increase.
"

# ╔═╡ e51f5ba4-239f-11eb-0d73-cbb2f0b727b7
md" Now that we understand how gases behave generally, lets look at specific variables that impact gases"

# ╔═╡ 07010ba2-23a0-11eb-2656-b12946848a01
md"## Relationship: Volume and Temperature
Before we look into the relationship between the two, it is important to understand several distinctions about each variable.

**Volume of a Gas equals Volume of a Container** : Gases will expand uniformly to fill any containers that it was placed in which means the volume of the container equals the volume of the gas."

# ╔═╡ 3b34ea3a-23a1-11eb-31b5-3171d41bc984
md"**Temperature in Kelvin**: Whenever discussing the behavior of gas, it is very important to use kelvin temperature. To relate this back to celsius, lets look at a conversion equation."

# ╔═╡ 1b88d926-23a3-11eb-0d2e-7f8a2fa4eb1b
md"""Move the slider to change the degrees in Celsius and to see how that affects the temperature in kelvin: """

# ╔═╡ 9f4d4afa-23a4-11eb-082b-cf05aa408069
 @bind C html"<input type=range>"

# ╔═╡ a34f586e-23a4-11eb-3ac3-13713ac16268
md"$C degrees Celsius equates to"

# ╔═╡ 65a88f26-23a4-11eb-02af-8998dbbbcf7f
TemperatureK= C + 273.15

# ╔═╡ c3f552c0-23a5-11eb-2174-33ebd9de1cff
begin

	md" The equation to calculate kelvin from celsius is T(k) = C + 273.15"
end

# ╔═╡ cc464fea-23a5-11eb-3e9e-87f62b2a4faf
md" Now that we know the basics of Volume and temperature, lets examine their relationship!"

# ╔═╡ e450fb58-23a5-11eb-2caa-639830646520
md".............................................................................................................................................................................................................."

# ╔═╡ 466d76b8-2475-11eb-32e4-19a694079c2c
md"**VOLUME is DIRECTLY PROPORTIONAL to absolute TEMPERATURE**"

# ╔═╡ a50225d4-2475-11eb-3c89-9b50dc329236
md" Volume does not depend on the type of gas you have because all the gases behave the same. So, for all gases, as the temperature increases, so does the volume."

# ╔═╡ 4a55c8d8-2476-11eb-07d3-9b01d5e70a1e
Temperature = [0,50,100,150,200,250,300]

# ╔═╡ 008ce65e-2490-11eb-3b8a-f91d3cb8e452
Volume = [0,1,2,3,4,5,6]

# ╔═╡ bc1d2f24-2476-11eb-0c81-17c4da48b190

# ╔═╡ 68c8e54c-2490-11eb-3d2b-dd73b9af11cf
plot(Temperature, Volume,legend=false, xlabel="Temperature(K)", ylabel="Volume(L)", title="Correlation between Temperature and Volume")

# ╔═╡ be952f68-2476-11eb-049b-1beb8f592f70
md"Since the relationship is directly proportional, we can use this principle and apply it into the equation `V1/T1 = V2/T2`. This equation is apart of a larger equation which we will cover later on in this section"

# ╔═╡ bf35d8e4-2476-11eb-078b-3b6b1f5fa04e
md"Now, lets move on to a different relationship: the relationship between amount and volume."

# ╔═╡ e5a02fd2-2475-11eb-3f1e-152a7a3a43ce
md"## Relationship: Volume and Amount
We understand a lot about volume, but before we can look into the relationship between amount and volume, we need to understand amount and the general assumptions we make about the relationship between volume and amount."

# ╔═╡ dfb0faac-24e3-11eb-350d-d7b5fed50b49
md"First, amount usually expressed in **moles**, the international unit for measuring the amount of a substance. This can be converted into grams(which depends on the substance you are looking for) or moles(which has the conversion factor of 6.022*10^23). Second, when we examine the relationship between volume and amount, it is assumed that the volume of the container is NOT fixed. This is important because if the volume is fixed, other variables are changed and you are no longer examining the relationship between amount and volume. Now lets talk about their relationship!"

# ╔═╡ dc0e97cc-24e5-11eb-215e-5dd5e41f7946
md"**VOLUME is DIRECTLY PROPORTIONAL to AMOUNT**"

# ╔═╡ f5e00744-24e5-11eb-3760-a75ea6b16b9c
md" The relationship between volume and amount is similar to volume and temperature. Assuming that the volume is not fixed, the more moles of a substance you have, the greater the volume. This is because as you increase the number of particles, the volume has to expand to keep everything constant."

# ╔═╡ 89333174-24e6-11eb-36d3-71f27b566e4f
begin
	plotly()
	NumberofMoles = [0.0,0.05,0.1,0.15,0.2,0.25,0.3]
	Volume2 = [0.0,1.0,2.0,3.0,4.0,5.0,6.0]
end

# ╔═╡ ee6eac6c-24e6-11eb-2df9-1fb31e1f2c96
plot(NumberofMoles, Volume2,legend=false, xlabel="Number of Moles", ylabel="Volume(L)", title="Correlation between Amount and Volume")

# ╔═╡ d504768c-27b9-11eb-2a2c-e5ccce269690
md"Since the relationship is directly proportional, we can use this principle and apply it into the equation `V1/N1 = V2/N2`. This equation is also apart of a larger equation which we will cover later on in this section"

# ╔═╡ e82d8afa-27b9-11eb-0ad1-417d616337dd
md"Now lets discuss the next component: Pressure and Temperature"

# ╔═╡ 075ec916-27ba-11eb-3ea2-a516b703aa90
md"## Relationship: Pressure and Temperature
Before we look at the relationship between temperature and pressure, lets understand some basic facts about pressure"

# ╔═╡ 2f624a46-27ba-11eb-205a-3b78744f49af
md" Pressure is defined as **force per unit area**. In the context of gases, as mentioned in the kinetic theory of gases, pressure in gases is caused by collisions into particles or the walls of a container. There are several ways to measure pressure, and each way depends on the conditions in which the pressure is measured. A **manometer** measures pressure in a closed system. This means that it will measure the pressure of a sealed container. A **barometer** measures pressure in a open system, like the atmosphere."

# ╔═╡ 82beda50-27bb-11eb-11af-c75e38f6fa2a
md"Pressure forms based on the number of collisions in particles. The greater the number of collisions, the higher the pressure. Since more collisions are a direct result of increase in temperature, we can say that the higher the temperature, the higher the pressure.(This is assuming that volume remains constant)"

# ╔═╡ 59e62f4e-2f62-11eb-097e-f32238a77ec4
Temperature2 = [0,50,100,150,200,250,300]

# ╔═╡ 9762122c-2f62-11eb-03f7-ff2dc854fe40
Pressure = [0,1,2,3,4,5,6]

# ╔═╡ a00adb0c-2f62-11eb-2dec-23a060be15cc
plot(Temperature2, Pressure,legend=false, xlabel="Temperature(Kelvin)", ylabel="Pressure(atm)", title="Correlation between Pressure and Temperature")

# ╔═╡ d10de92c-2f62-11eb-15c6-9b30bca82516
md"Pressure is directly proportional to Temperature so when the temperature increases, pressure also increases. The equation for this is `P1/T1 = P2/T2`. Now lets look at the final relationship before putting everything together"

# ╔═╡ bd83ed76-2f66-11eb-08df-338b5985c137
md" ## Relationship: Volume and Pressure
This relationship is different from the rest because the relationship between volume and temperature is inversly proportional to pressure. As the volume increases, the pressure decreases and as the volume decreases, the pressure increases. This happens because as the volume increase, the molecules have a lot more room to move around so less collisions happen which means less pressure. On the flipside, when volume decreases, there is less space for the molecules to move around so they collide more often and those collisions result in an increase in pressure."

# ╔═╡ 8d3afe88-2f67-11eb-308e-ad7440ee187c
Volume3 = [11,6,3,2,1.5,1.25,1]

# ╔═╡ 28b7c798-2f67-11eb-1ee7-ff083e19acaf
plot(Pressure, Volume3 ,legend=false, xlabel="Pressure(atm)", ylabel="Volume(L)", title="Correlation between Volume and Pressure")

# ╔═╡ d3129b2a-2f67-11eb-15b7-c9bac218038c
md"Now we understand each component, Volume, pressure, temperature, and moles, and their relationships with eachother, we can put all the factors together to form one equation."

# ╔═╡ f8bbb5ba-2f67-11eb-32ef-bbc3f87fa010
md" ## Ideal Gas Law Formula
The Ideal Gas formula is
$\ce{PV=nRT}$ , where P stands for pressure, V stands for volume, n stands for moles, and T stands for temperature. The R is a special value called the gas constant, which stays the same no matter the calculation. R = 0.0821. Understanding this, you can now solve for each individual variable. "

# ╔═╡ Cell order:
# ╟─e0937cd4-239d-11eb-141d-1be13a185dca
# ╟─17827eb0-239f-11eb-3266-2d71e0aedd9c
# ╟─e51f5ba4-239f-11eb-0d73-cbb2f0b727b7
# ╟─07010ba2-23a0-11eb-2656-b12946848a01
# ╟─3b34ea3a-23a1-11eb-31b5-3171d41bc984
# ╠═1b88d926-23a3-11eb-0d2e-7f8a2fa4eb1b
# ╠═9f4d4afa-23a4-11eb-082b-cf05aa408069
# ╠═a34f586e-23a4-11eb-3ac3-13713ac16268
# ╠═65a88f26-23a4-11eb-02af-8998dbbbcf7f
# ╟─c3f552c0-23a5-11eb-2174-33ebd9de1cff
# ╟─cc464fea-23a5-11eb-3e9e-87f62b2a4faf
# ╟─e450fb58-23a5-11eb-2caa-639830646520
# ╟─466d76b8-2475-11eb-32e4-19a694079c2c
# ╟─a50225d4-2475-11eb-3c89-9b50dc329236
# ╟─126b1cca-2476-11eb-087b-474174c4c1d0
# ╠═4a55c8d8-2476-11eb-07d3-9b01d5e70a1e
# ╠═008ce65e-2490-11eb-3b8a-f91d3cb8e452
# ╟─bc1d2f24-2476-11eb-0c81-17c4da48b190
# ╟─68c8e54c-2490-11eb-3d2b-dd73b9af11cf
# ╟─be952f68-2476-11eb-049b-1beb8f592f70
# ╟─bf35d8e4-2476-11eb-078b-3b6b1f5fa04e
# ╟─e5a02fd2-2475-11eb-3f1e-152a7a3a43ce
# ╟─dfb0faac-24e3-11eb-350d-d7b5fed50b49
# ╟─dc0e97cc-24e5-11eb-215e-5dd5e41f7946
# ╟─f5e00744-24e5-11eb-3760-a75ea6b16b9c
# ╟─89333174-24e6-11eb-36d3-71f27b566e4f
# ╟─ee6eac6c-24e6-11eb-2df9-1fb31e1f2c96
# ╟─d504768c-27b9-11eb-2a2c-e5ccce269690
# ╟─e82d8afa-27b9-11eb-0ad1-417d616337dd
# ╟─075ec916-27ba-11eb-3ea2-a516b703aa90
# ╟─2f624a46-27ba-11eb-205a-3b78744f49af
# ╟─82beda50-27bb-11eb-11af-c75e38f6fa2a
# ╟─59e62f4e-2f62-11eb-097e-f32238a77ec4
# ╟─9762122c-2f62-11eb-03f7-ff2dc854fe40
# ╟─a00adb0c-2f62-11eb-2dec-23a060be15cc
# ╟─d10de92c-2f62-11eb-15c6-9b30bca82516
# ╟─bd83ed76-2f66-11eb-08df-338b5985c137
# ╟─8d3afe88-2f67-11eb-308e-ad7440ee187c
# ╟─28b7c798-2f67-11eb-1ee7-ff083e19acaf
# ╟─d3129b2a-2f67-11eb-15b7-c9bac218038c
# ╟─f8bbb5ba-2f67-11eb-32ef-bbc3f87fa010


# Create line data with NaN as polyline separator
x = [randn(10000); NaN; randn(10000) .+ 5 ]
y = [randn(10000); NaN; randn(10000) .+ 5]
# Draw shade
shade(x, y, xform="loglog")


import DelimitedFiles
import Statistics
import NPZ
import RestrictedBoltzmannMachines as RBMs

using Test: @test, @testset
using Statistics: mean, std

#=
Comparisons against https://github.com/jertubiana/PGM.

I computed the energy, free_energy, and so on, using the PGM
code. So here we compare if using the same RBM I obtain the same
values of energies, etc. with my code.

The only difference is in the definition of some θ's (for Gaussian, ReLU, and dReLU layers).
I wanted to be consistent with the definition of the external fields, which appear with
a minus sign in the energy for discrete units. Therefore in the following tests pay
attention to these minus signs.
=#

@testset "Binary / Binary RBM" begin
    g = vec(DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_bin_g_v.txt"))
    θ = vec(DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_bin_g_h.txt"))
    w = DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_bin_w.txt")'

    v = DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_bin_v.txt")'
    h = DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_bin_h.txt")'
    E = vec(DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_bin_E.txt"))
    F = vec(DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_bin_F.txt"))

    rbm = RBMs.RBM(RBMs.Binary(g), RBMs.Binary(θ), w)
    @test RBMs.energy(rbm, v, h) ≈ E rtol=1e-5
    @test RBMs.free_energy(rbm, v) ≈ F rtol=1e-5
end

@testset "Binary / Gaussian RBM" begin
    #= Note the minus sign in front of θ =#
    θ = -vec(DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_gauss_theta_h.txt"))
    γ =  vec(DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_gauss_gamma_h.txt"))
    g = vec(DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_gauss_g_v.txt"))
    w = DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_gauss_w.txt")'

    v = DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_gauss_v.txt")'
    h = DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_gauss_h.txt")'
    Ev = vec(DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_gauss_Ev.txt"))
    Eh = vec(DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_gauss_Eh.txt"))
    E = vec(DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_gauss_E.txt"))
    F = vec(DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_gauss_F.txt"))

    rbm = RBMs.RBM(RBMs.Binary(g), RBMs.Gaussian(θ, γ), w)
    @test RBMs.energy(rbm.visible, v) ≈ Ev rtol=1e-5
    @test RBMs.energy(rbm.hidden,  h) ≈ Eh rtol=1e-5
    @test RBMs.energy(rbm, v, h) ≈ E rtol=1e-5
    @test RBMs.free_energy(rbm, v) ≈ F rtol=1e-5
end

@testset "Binary / dReLU RBM" begin
    #= Note the minus sign in front of θp. For Gaussian, ReLU, and dReLU layers
    (as well as for the reparameterized pReLU), my θ's are minus those of Jerome.
    I did this because I wanted to be consistent with the external field of Binary,
    Spin and Potts layer, for which the field appears with this sign. =#
    θp = -vec(DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_dReLU_thetap_h.txt"))
    θn =  vec(DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_dReLU_thetam_h.txt"))
    γp = vec(DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_dReLU_gammap_h.txt"))
    γn = vec(DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_dReLU_gammam_h.txt"))
    g = vec(DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_dReLU_g_v.txt"))
    w = DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_dReLU_w.txt")'

    v = DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_dReLU_v.txt")'
    h = DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_dReLU_h.txt")'
    Ev = vec(DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_dReLU_Ev.txt"))
    Eh = vec(DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_dReLU_Eh.txt"))
    E = vec(DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_dReLU_E.txt"))
    F = vec(DelimitedFiles.readdlm("compare_to_pgm/RBM_bin_dReLU_F.txt"))

    rbm = RBMs.RBM(RBMs.Binary(g), RBMs.dReLU(θp, θn, γp, γn), w)
    @test RBMs.energy(rbm.visible, v) ≈ Ev rtol=1e-5
    @test RBMs.energy(rbm.hidden,  h) ≈ Eh rtol=1e-5
    @test RBMs.energy(rbm, v, h) ≈ E rtol=1e-5
    @test RBMs.free_energy(rbm, v) ≈ F rtol=1e-3
end

@testset "Binary pseudolikelihood" begin
    w = Matrix(DelimitedFiles.readdlm("compare_to_pgm/PL/RBM_Bernoulli_weights.txt")')
    visible_g = vec(DelimitedFiles.readdlm("compare_to_pgm/PL/RBM_Bernoulli_visible_fields.txt"))
    hidden_g = vec(DelimitedFiles.readdlm("compare_to_pgm/PL/RBM_Bernoulli_hidden_fields.txt"))
    v = BitMatrix(DelimitedFiles.readdlm("compare_to_pgm/PL/RBM_Bernoulli_data.txt", Bool)')
    pl = vec(DelimitedFiles.readdlm("compare_to_pgm/PL/RBM_Bernoulli_PL.txt"))

    rbm = RBMs.RBM(RBMs.Binary(visible_g), RBMs.Binary(hidden_g), w)
    pl_rbm = RBMs.log_pseudolikelihood(rbm, v)

    @test mean(pl_rbm) ≈ mean(pl) rtol=0.05
    @test std(pl_rbm)  ≈ std(pl)  rtol=0.1
end

@testset "Potts pseudolikelihood" begin
    # https://numpy.org/doc/stable/reference/generated/numpy.savez_compressed.html
    d = NPZ.npzread("compare_to_pgm/PL/RBM_Potts_PL.npz")
    rbm = RBMs.RBM(
        RBMs.Potts(Matrix(d["g"]')),
        RBMs.dReLU(-d["theta_plus"], d["theta_minus"], d["gamma_plus"], d["gamma_minus"]),
        permutedims(d["weights"], (3,2,1))
    )
    data = RBMs.onehot_encode(d["data"]')
    pl = RBMs.log_pseudolikelihood(rbm, data)

    # TODO: #7 see if we can decrease rtol
    @test mean(pl) ≈ mean(d["PL"]) rtol=0.1
    @test std(pl) ≈ std(d["PL"]) rtol=0.1
end


# algorithm

abstract type QuadraticOT end

# solver

struct QuadraticOTSolver{A<:QuadraticOT,M,N,CT,E<:Real,T<:Real,R<:Real,C1,C2}
    source::M
    target::N
    C::CT
    eps::E
    alg::A
    atol::T
    rtol::R
    maxiter::Int
    check_convergence::Int
    cache::C1
    convergence_cache::C2
end

struct QuadraticOTConvergenceCache{N<:Real}
    norm_source::N
    norm_target::N
end

function build_convergence_cache(::Type{T}, μ::AbstractVector, ν::AbstractVector) where {T}
    norm_μ = norm(μ, Inf)
    norm_ν = norm(ν, Inf)
    return QuadraticOTConvergenceCache(norm_μ, norm_ν)
end

function build_solver(
    μ::AbstractVector,
    ν::AbstractVector,
    C::AbstractMatrix,
    ε::Real,
    alg::QuadraticOT;
    atol=nothing,
    rtol=nothing,
    check_convergence=1,
    maxiter::Int=100,
)
    # check that source and target marginals have the correct size and are balanced
    checksize(μ, ν, C)
    checkbalanced(μ, ν)
    # do not use checksize2 since for quadratic OT (at least for now) we do not support batch computations

    # compute type
    T = float(Base.promote_eltype(μ, ν, one(eltype(C)) / ε))

    # build caches
    cache = build_cache(T, alg, μ, ν, C, ε)
    convergence_cache = build_convergence_cache(T, μ, ν)

    # set tolerances
    _atol = atol === nothing ? 0 : atol
    _rtol = rtol === nothing ? (_atol > zero(_atol) ? zero(T) : sqrt(eps(T))) : rtol

    # create solver
    solver = QuadraticOTSolver(
        μ, ν, C, ε, alg, _atol, _rtol, maxiter, check_convergence, cache, convergence_cache
    )

    return solver
end

"""
    quadreg(μ, ν, C, ε, alg::QuadraticOT; kwargs...)

Computes the optimal transport plan of histograms `μ` and `ν` with cost matrix `C` and quadratic regularization parameter `ε`.

The optimal transport plan `γ` is of the same size as `C` and solves

```math
\\inf_{\\gamma \\in \\Pi(\\mu, \\nu)} \\langle \\gamma, C \\rangle
+ \\varepsilon \\Omega(\\gamma),
```
where ``\\Omega(\\gamma) = \\frac{1}{2} \\sum_{i,j} \\gamma_{i,j}^2`` is the quadratic
regularization term.

Every `check_convergence` steps it is assessed if the algorithm is converged by checking if
the iterate of the transport plan `γ` satisfies
```julia
    norm_diff < max(atol, rtol * max(norm(μ, Inf), norm(ν, Inf)))
```
where
```math
    \\text{normdiff} = \\max\\{ \\| \\gamma \\mathbf{1} - \\mu \\|_\\infty , \\|  \\gamma^\\top \\mathbf{1} - \\nu \\|_\\infty  \\} .
```
After `maxiter` iterations, the computation is stopped.

Note that unlike in the case of Sinkhorn's algorithm for the entropic regularisation, batch computation of optimal transport is not supported for the quadratic regularisation.

See also: [`sinkhorn`](@ref), [`QuadraticOTNewton`](@ref)
"""
function quadreg(μ, ν, C, ε, alg::QuadraticOT=QuadraticOTNewton(); kwargs...)
    # build solver
    solver = build_solver(μ, ν, C, ε, alg; kwargs...)

    # perform Sinkhorn algorithm
    solve!(solver)

    # compute optimal transport plan
    γ = plan(solver)

    return γ
end


# Define abstract interface
"""
$(TYPEDEF)

Abstract supertype for magnetization fields.
"""
abstract type Magnetizations{T} end

"""
$(SIGNATURES)

Return the magnetization field M [T/μ₀] (μ₀ being the vacuum permeability)
as `SVector{3,T}` for a given magnetic field H [T/μ₀], location r [m] and
time t [s].
"""
function magnetizationField(magnetization::Magnetizations{T},H::SVector{3,T},r::SVector{3,T},t::T) where {T}
    error("magnetizationField is not implemented for magnetization::$(typeof(magnetization)).")
end

"""
$(SIGNATURES)

Return if (`Bool`) magnetization field is zero at a given location r [m] at all
times.
"""
function isZero(magnetization::Magnetizations{T},r::SVector{3,T}) where {T}
    error("isZero is not implemented for magnetization::$(typeof(magnetization)).")
end

"""
$(SIGNATURES)

Return `QuadratureNodes` required to integrate over the magnetization M.
"""
function QuadratureNodes(magnetization::Magnetizations{T},density) where {T}
    error("QuadratureNodes is not implemented for magnetization::$(typeof(magnetization)).")
end

# Define specific subtypes
## separable magnetization
struct SeparableMagnetization{T,U<:TracerDistributions{T},V<:MagneticMoments{T}} <: Magnetizations{T}
    tracerDistribution::U
    magneticMoment::V
end

@inline isZero(magnetization::SeparableMagnetization,r::SVector{3,T}) where {T} = isZero(magnetization.tracerDistribution,r)

function magnetizationField(magnetization::SeparableMagnetization{T},H::SVector{3,T},r::SVector{3,T},t::T) where {T}
    return tracerConcentration(magnetization.tracerDistribution,r,t)*meanMagneticMoment(magnetization.magneticMoment,H,r,t)
end

function QuadratureNodes(magnetization::SeparableMagnetization{T},density=10/1e-3) where {T}
    return QuadratureNodes(magnetization.tracerDistribution,density)
end


export HermitianOperator, SymmetricOperator, AdjointOperator, TransposeOperator

struct HermitianOperator{T<:Number,B<:Operator} <: Operator{T}
    op::B
    uplo::Char
end

HermitianOperator(B::Operator{T}, uplo::Symbol=:U) where {T<:Number} = HermitianOperator{T,typeof(B)}(B, char_uplo(uplo))
convert(::Type{Operator{T}},A::HermitianOperator) where {T}=HermitianOperator(convert(Operator{T},A.op), A.uplo)

domainspace(P::HermitianOperator)=domainspace(P.op)
rangespace(P::HermitianOperator)=rangespace(P.op)
domain(P::HermitianOperator)=domain(P.op)
bandwidths(P::HermitianOperator) = (b = bandwidth(P.op, P.uplo == 'L' ? 1 : P.uplo == 'U' ? 2 : 0); (b, b))

function getindex(P::HermitianOperator,k::Integer,j::Integer)
    if P.uplo == 'L'
        if j > k
            conj(P.op[j,k])
        else
            P.op[k,j]
        end
    elseif P.uplo == 'U'
        if j < k
            conj(P.op[j,k])
        else
            P.op[k,j]
        end
    end
end

Hermitian(A::Operator, uplo::Symbol)=HermitianOperator(A, uplo)
Hermitian(A::Operator)=HermitianOperator(A)

struct SymmetricOperator{T<:Number,B<:Operator} <: Operator{T}
    op::B
    uplo::Char
end

SymmetricOperator(B::Operator{T}, uplo::Symbol=:U) where {T<:Number} = SymmetricOperator{T,typeof(B)}(B, char_uplo(uplo))
convert(::Type{Operator{T}},A::SymmetricOperator) where {T}=SymmetricOperator(convert(Operator{T},A.op), A.uplo)

domainspace(P::SymmetricOperator)=domainspace(P.op)
rangespace(P::SymmetricOperator)=rangespace(P.op)
domain(P::SymmetricOperator)=domain(P.op)
bandwidths(P::SymmetricOperator) = (b = bandwidth(P.op, P.uplo == 'L' ? 1 : P.uplo == 'U' ? 2 : 0); (b, b))

function getindex(P::SymmetricOperator,k::Integer,j::Integer)
    if P.uplo == 'L'
        if j > k
            P.op[j,k]
        else
            P.op[k,j]
        end
    elseif P.uplo == 'U'
        if j < k
            P.op[j,k]
        else
            P.op[k,j]
        end
    end
end

Symmetric(A::Operator, uplo::Symbol)=SymmetricOperator(A, uplo)
Symmetric(A::Operator)=SymmetricOperator(A)

struct AdjointOperator{T<:Number,B<:Operator} <: Operator{T}
    op::B
end

AdjointOperator(B::Operator{T}) where {T<:Number}=AdjointOperator{T,typeof(B)}(B)
convert(::Type{Operator{T}},A::AdjointOperator) where {T}=AdjointOperator(convert(Operator{T},A.op))

domainspace(P::AdjointOperator)=rangespace(P.op)
rangespace(P::AdjointOperator)=domainspace(P.op)
domain(P::AdjointOperator)=domain(P.op)
bandwidths(P::AdjointOperator) = reverse(bandwidths(P.op))

getindex(P::AdjointOperator,k::Integer,j::Integer) = conj(P.op[j,k])

function BandedMatrix(S::SubOperator{T,TO}) where {T,TO<:AdjointOperator}
    kr,jr=parentindices(S)
    adjoint(BandedMatrix(view(parent(S).op,jr,kr)))
end

adjoint(A::Operator)=AdjointOperator(A)

struct TransposeOperator{T<:Number,B<:Operator} <: Operator{T}
    op::B
end

TransposeOperator(B::Operator{T}) where {T<:Number}=TransposeOperator{T,typeof(B)}(B)
convert(::Type{Operator{T}},A::TransposeOperator) where {T}=TransposeOperator(convert(Operator{T},A.op))

domainspace(P::TransposeOperator)=rangespace(P.op)
rangespace(P::TransposeOperator)=domainspace(P.op)
domain(P::TransposeOperator)=domain(P.op)
bandwidths(P::TransposeOperator) = reverse(bandwidths(P.op))

getindex(P::TransposeOperator,k::Integer,j::Integer) = P.op[j,k]

function BandedMatrix(S::SubOperator{T,TO}) where {T,TO<:TransposeOperator}
    kr,jr=parentindices(S)
    transpose(BandedMatrix(view(parent(S).op,jr,kr)))
end

transpose(A::Operator)=TransposeOperator(A)


# ============================================================================
# Thermostat model
# See https://fenix.tecnico.ulisboa.pt/downloadFile/3779579688470/lygeros.pdf,
# Section 1.3.4.
# ============================================================================

function thermostat()
    c_a = 0.1

    # automaton structure
    automaton = LightAutomaton(2)

    # mode on
    A = hcat(c_a)
    inv = HalfSpace([1.0], 22.0)  # x ≤ 22
    # @vars x v
    # @set x ≤ 22
    m_on = ConstrainedLinearContinuousSystem(A, inv)

    # mode off
    A = hcat(-c_a)
    inv = HalfSpace([-1.0], -18.0)  # x ≥ 18
    m_off = ConstrainedLinearContinuousSystem(A, inv)

    # modes
    modes = [m_on, m_off]

    # transition from on to off
    add_transition!(automaton, 1, 2, 1)
    guard = HalfSpace([-1.0], -21.0)  # x ≥ 21
    t_on2off = ConstrainedIdentityMap(2, guard)

    # transition from off to on
    add_transition!(automaton, 2, 1, 2)
    guard = HalfSpace([1.0], 19.0)  # x ≤ 19
    t_off2on = ConstrainedIdentityMap(2, guard)

    # transition annotations
    resetmaps = [t_on2off, t_off2on]

    # switching
    switchings = [AutonomousSwitching()]

    ℋ = HybridSystem(automaton, modes, resetmaps, switchings)

    # initial condition in mode on
    X0 = Singleton([18.0])
    initial_condition = [(1, X0)]

    return InitialValueProblem(ℋ, initial_condition)
end


# NOTE
# Did a prior analysis for a version of repFAM where Z is
# a function of H, v. The results were underwhelming, and we couldn't
# achieve repulsiveness. Need to give up this idea.

import Pkg; Pkg.activate("../../")
include("HMC.jl")
using Distributions
using StatsBase

function grad_logprobH(H, j, k, phi)
  J, K = size(H)
  out = -H[j, k]
  for m in 1:K
    if m != k
      dmk = mean((H[:, k] - H[:, m]) .^ 2)
      tmp = zero(H[1])
      tmp += mean(H[:, k] - H[:, m]) / (exp(dmk / phi) - 1)
      out += tmp * 2 / phi
    end
  end
  return out
end

function grad_logprobH(H, phi)
  J, K = size(H)
  return [grad_logprobH(H, j, k, phi) for j in 1:J, k in 1:K]
end

"""
log1m(x) = log1p(-x) = log(1 - x)
"""
log1m(x) = log1p(-x)

function logprobH(H, phi)
  J, K = size(H)
  normal_kernel = -sum((H .^ 2)) / 2
  penalty = zero(H[1])

  for m in 1:(K - 1)
    for n in (m + 1):K
      d = mean((H[:, m] - H[:, n]) .^ 2)
      penalty += log1m(exp(-d / phi))
    end
  end

  return normal_kernel + penalty
end

struct State H end

function drawH(J, K, phi)
  if phi == 0
    randn(J, K)
  else
    # TODO
    0
  end
end

function drawH(n, J, K, phi)
  if phi == 0
    [randn(J, K) for _ in 1:n]
  else
    # TODO

  end
end

"""
NOTE:
Let H be a JxK continuous matrix, with H_{jk} \\in \\mathbb{R}.
Let Z be a JxK binary matrix.
Let v be a K-dimensional continuous vector, bounded within the unit interval.

Model:
Z_{jk} = 1(v_{k} > \\Phi(H_{jk}))
v_k | \\alpha ~ Beta(\\alpha/K, 1)
p(H) \\propto \\prod_{j,k} exp(-H_{jk}^2/2) *
              \\prod{k\\ne k'} 1 - \\exp(-d_{k,k'} / \\phi)
"""
function drawZ(J, K, alpha, phi; stickbreak=false)
  if stickbreak
    v = cumprod(rand(Beta(alpha, 1), 1, K), dims=2)
  else
    v = rand(Beta(alpha / K, 1), 1, K)
  end
  H = drawH(J, K, phi)
  return v .> cdf.(Normal(0, 1), H)
end

function drawH(n, J, K, phi; num_leapfrog_steps=2^6, eps=.01, burn=0)
  lp(H) = logprobH(H, phi)
  grad_lp(H) = grad_logprobH(H, phi)

  H = randn(J, K)
  for _ in 1:burn
    H, _ = HMC.hmc_update(H, lp, grad_lp, num_leapfrog_steps, eps)
  end

  Hs = [zeros(J, K) for _ in 1:n]
  lpHs = zeros(n)

  for i in 1:n
    H, lpH = HMC.hmc_update(H, lp, grad_lp, num_leapfrog_steps, eps)
    Hs[i] .= copy(H)
    lpHs[i] = lpH
  end

  return Hs, lpHs
end

function drawZ(N, J, K, alpha, phi;
               stickbreak=false, num_leapfrog_steps=2^6, eps=.01, burn=0)
  if stickbreak
    vs = [cumprod(rand(Beta(alpha, 1), 1, K), dims=2) for _ in 1:N]
  else
    vs = [rand(Beta(alpha / K, 1), 1, K) for _ in 1:N]
  end

  Hs, lpH = drawH(N, J, K, phi,
                  num_leapfrog_steps=num_leapfrog_steps, eps=eps, burn=burn)

  return [vs[i] .> cdf.(Normal(0, 1), Hs[i]) for i in 1:N]
end

# Simulate
N, alpha, phi, J, K, burn = (10000, 0.5, 100.0, 2, 2, 100)
@time Zs = drawZ(N, J, K, alpha, phi, stickbreak=true, burn=burn)
Zcounts = countmap(Zs)
sort(collect(Zcounts), by=x -> -x.second)


# Pkg.build("QuantEcon")
import Dolo
path = Dolo.pkg_path
using AxisArrays

###############################################################################
## Markov Chain
###############################################################################
filename = joinpath(path,"examples","models","rbc_mc.yaml")
# model = Dolo.Model(joinpath(Dolo.pkg_path, "examples", "models", "rbc_mc.yaml"), print_code=true)
model = Dolo.yaml_import(filename)

@time dr_ITI  = Dolo.improved_time_iteration(model; verbose=true, tol = 1e-06, smaxit=100)

dprocess = Dolo.discretize( model.exogenous )
# init_dr = Dolo.ConstantDecisionRule(model.calibration[:controls])
# @time dr_ITI_2  = Bruteforce_module.improved_time_iteration(model, dprocess,init_dr)

@time dr_TI  = Dolo.time_iteration(model;tol_η=1e-08, maxit=1000)

################################################################################
# function profile_ITI(m)
#     dr_ITI  = Bruteforce_module.improved_time_iteration(m)
#     return dr_ITI
# end
#
#
# using ProfileView
#
# profile_ITI(model)
# Profile.clear()
# @profile profile_ITI(model)
# ProfileView.view()
#

###############################################################################
#Plotting the DR

# ITIT method
df_ITI = Dolo.tabulate(model, dr_ITI.dr, :k)

import PyPlot
plt = PyPlot;
fig = plt.figure("Decision Rule, ITI-method",figsize=(8.5,5))

plt.subplot(1,2,1)
plt.plot(df_ITI[:k], df_ITI[:n])
plt.ylabel("Hours");
plt.xlabel("state = k");
plt.title("Decision Rule");

plt.subplots_adjust(wspace=1)
# plt.tight_layout()

plt.subplot(1,2,2)
plt.plot(df_ITI[:k], df_ITI[:i])
plt.ylabel("Investment");
plt.xlabel("state = k");
plt.title("Decision Rule");
###############################################
# TI method
df = Dolo.tabulate(model, dr_TI.dr, :k)

fig = plt.figure("Decision Rule, IT-method",figsize=(8.5,5))

plt.subplot(1,2,1)
plt.plot(df[:k], df[:n])
plt.ylabel("Hours");
plt.xlabel("state = k");
plt.title("Decision Rule");

plt.subplots_adjust(wspace=1)
# plt.tight_layout()

plt.subplot(1,2,2)
plt.plot(df[:k], df[:i])
plt.ylabel("Investment");
plt.xlabel("state = k");
plt.title("Decision Rule");

################################################################################
# RBC with an AR process
################################################################################

filename = joinpath(path,"examples","models","rbc.yaml")
model_ar1 = Dolo.yaml_import(filename)
@time dr_ITI_ar1  = Dolo.improved_time_iteration(model_ar1; verbose=true, tol = 1e-06, smaxit=50)
@time dr_TI_ar1  = Dolo.time_iteration(model_ar1; tol_η=1e-08, maxit=1000)

dprocess = Dolo.discretize( model_ar1.exogenous , [3], [2])
@time dr2_ITI_ar1  = Dolo.improved_time_iteration(model_ar1, dprocess; verbose=true, tol = 1e-06, smaxit=50)
@time dr2_TI_ar1  = Dolo.time_iteration(model_ar1, dprocess; tol_η=1e-08, maxit=1000)

################################################################################
df_ITI = Dolo.tabulate(model_ar1, dr2_ITI_ar1.dr, :k)

import PyPlot
plt = PyPlot;
fig = plt.figure("Decision Rule, ITI-method",figsize=(8.5,5))

plt.subplot(1,2,1)
plt.plot(df_ITI[:k], df_ITI[:n])
plt.ylabel("Hours");
plt.xlabel("state = k");
plt.title("Decision Rule");

plt.subplots_adjust(wspace=1)
# plt.tight_layout()

plt.subplot(1,2,2)
plt.plot(df_ITI[:k], df_ITI[:i])
plt.ylabel("Investment");
plt.xlabel("state = k");
plt.title("Decision Rule");

df = Dolo.tabulate(model_ar1, dr2_TI_ar1.dr, :k)

fig = plt.figure("Decision Rule, IT-method",figsize=(8.5,5))

plt.subplot(1,2,1)
plt.plot(df[:k], df[:n])
plt.ylabel("Hours");
plt.xlabel("state = k");
plt.title("Decision Rule");

plt.subplots_adjust(wspace=1)
# plt.tight_layout()

plt.subplot(1,2,2)
plt.plot(df[:k], df[:i])
plt.ylabel("Investment");
plt.xlabel("state = k");
plt.title("Decision Rule");

################################################################################
# RBC with an IID process
################################################################################

filename = joinpath(path,"examples","models","rbc_iid.yaml")
model_iid = Dolo.yaml_import(filename)
@time dr_ITI_iid  = Bruteforce_module.improved_time_iteration(model_iid; verbose=true, tol = 1e-06, smaxit=50)
@time dr_TI_iid  = Dolo.time_iteration(model_iid; tol_η=1e-08, maxit=1000)

df_ITI = Dolo.tabulate(model_iid, dr_ITI_iid.dr, :k)

import PyPlot
plt = PyPlot;
fig = plt.figure("Decision Rule, ITI-method",figsize=(8.5,5))

plt.subplot(1,2,1)
plt.plot(df_ITI[:k], df_ITI[:n])
plt.ylabel("Hours");
plt.xlabel("state = k");
plt.title("Decision Rule");

plt.subplots_adjust(wspace=1)
# plt.tight_layout()

plt.subplot(1,2,2)
plt.plot(df_ITI[:k], df_ITI[:i])
plt.ylabel("Investment");
plt.xlabel("state = k");
plt.title("Decision Rule");

df = Dolo.tabulate(model_iid, dr_TI_iid.dr, :k)

fig = plt.figure("Decision Rule, IT-method",figsize=(8.5,5))

plt.subplot(1,2,1)
plt.plot(df[:k], df[:n])
plt.ylabel("Hours");
plt.xlabel("state = k");
plt.title("Decision Rule");

plt.subplots_adjust(wspace=1)
# plt.tight_layout()

plt.subplot(1,2,2)
plt.plot(df[:k], df[:i])
plt.ylabel("Investment");
plt.xlabel("state = k");
plt.title("Decision Rule");

################################################################################
# RBC with an AR1 with 2 shocks
################################################################################

# filename = joinpath(path,"examples","models","rbc_dtcc_ar1_2shocks.yaml")
# model_ar1_2 = Dolo.yaml_import(filename)
# @time dr_ITI_ar1_2 = Bruteforce_module.improved_time_iteration(model_ar1_2; verbose=true, tol = 1e-06, smaxit=50)
# @time dr_TI_ar_2  = Dolo.time_iteration(model_ar1_2; tol_η=1e-08, maxit=1000)
#

################################################################################
# RBC with a sudden stop. Complementarities
################################################################################

filename = joinpath(path,"examples","models","sudden_stop.yaml")
model_Sudden = Dolo.yaml_import(filename)
@time dr_ITI_s  = Dolo.improved_time_iteration(model_Sudden; complementarities = true, verbose=true, tol = 1e-06, smaxit=50)
@time dr_TI_s  = Dolo.time_iteration(model_Sudden; tol_η=1e-08, maxit=1000)

df_ITI = Dolo.tabulate(model_Sudden, dr_ITI_s.dr, :l)

import PyPlot
plt = PyPlot;
fig = plt.figure("Decision Rule, ITI-method",figsize=(8.5,5))

plt.subplot(1,2,1)
plt.plot(df_ITI[:l], df_ITI[:b])
plt.ylabel("b(t-1)");
plt.xlabel("state = b");
plt.title("Decision Rule");

plt.subplots_adjust(wspace=1)
# plt.tight_layout()

plt.subplot(1,2,2)
plt.plot(df_ITI[:l], df_ITI[:lam])
plt.ylabel("b/c");
plt.xlabel("state = b");
plt.title("Decision Rule");

#
df = Dolo.tabulate(model_Sudden, dr_TI_s.dr, :l)

fig = plt.figure("Decision Rule, IT-method",figsize=(8.5,5))

plt.subplot(1,2,1)
plt.plot(df[:l], df[:b])
plt.ylabel("Hours");
plt.xlabel("state = k");
plt.title("Decision Rule");

plt.subplots_adjust(wspace=1)
# plt.tight_layout()

plt.subplot(1,2,2)
plt.plot(df[:l], df[:lam])
plt.ylabel("Investment");
plt.xlabel("state = k");
plt.title("Decision Rule");


include("setupFWItests.jl")

using Test
using LinearAlgebra

println("\tTest getData serial and parallel")
Ds,pFor  = getData(vec(m),pFor)
Dp,pForp = getData(vec(m),pForp,M2Mp,true)
for k=1:length(Dp)
	@test norm(vec(Ds[:,:,k]-fetch(Dp[k])))/norm(vec(Ds[:,:,k])) < 1e-10
end

# check sensitivities
function dFun(m,pFor,v=[])
	D,pFor = getData(m,pFor)
	if isempty(v)
		return vec(D)
	else
		dD = getSensMatVec(v,m,pFor)
		return vec(D),vec(dD)
	end
end

println("\tderivative tests")
m0 = vec(m) + .001*randn(prod(n))
f(x,v=[]) = dFun(x,pFor,v)
chkDer1, = checkDerivative(f,m0,out=false)
@test chkDer1 == true

println("\tadjoint tests")
v1 = randn(length(Ds))
v2 = randn(length(m))

t1 = dot(v1,getSensMatVec(v2,vec(m),pFor))
t2 = dot(v2,getSensTMatVec(v1,vec(m),pFor))
@test norm(t1-t2)/norm(t1)<1e-8


@testset "Peak width" begin
    fs = 100
    T = 1/fs
    t = 0:T:6pi+T
    sint = sin.(t)

    @testset "minima/maxima" begin
        sinpks = argmaxima(sint)
        _, widths, _, _ = peakwidths(sinpks, sint, sint[sinpks]; strict=false, relheight=1)
        @test widths ≈ fill(pi*100, length(sinpks)) atol=.01

        sinpks = argminima(sint)
        _, widths, _, _ = peakwidths(sinpks, sint, abs.(sint[sinpks]); strict=false, relheight=1)
        @test widths ≈ fill(pi*100, length(sinpks)) atol=.01
    end

    _, widths, _, _ = peakwidths([2], [0.,1.,0.], [1.])
    @test widths == [1.]
    _, widths, _, _ = peakwidths([2], [0.,1.,0.], [NaN])
    @test widths[1] === NaN
    _, widths, _, _ = peakwidths([2], [0.,1.,0.], [missing])
    @test widths[1] === missing
    _, widths, _, _ = peakwidths([2], [0.,1.,NaN], [1.]; strict=true)
    @test widths[1] === NaN
    _, widths, _, _ = peakwidths([2], [0.,1.,0.,-1.], [1.]; strict=false)
    _, widthsnan, _, _ = peakwidths([2], [0.,1.,NaN,-1.], [1.]; strict=false)
    @test widths == widthsnan

    @testset "strict" begin
        sinpks = argmaxima(sint)
        _, widths, _, _ = peakwidths(sinpks, sint, ones(length(sinpks)); strict=true, relheight=1)
        @test first(widths) === NaN
        _, widths, _, _ = peakwidths(sinpks, sint, ones(length(sinpks)); strict=false, relheight=1)
        @test first(widths) !== NaN
    end

    @testset "Min/max width" begin
        sinpks = argmaxima(sint)
        _, proms = peakproms(sinpks, sint)

        @test length(first(peakwidths(sinpks, sint, proms; minwidth=pi*75))) == 2
        @test length(first(peakwidths(sinpks, sint, proms; maxwidth=pi*75))) == 1

        @test_throws ArgumentError peakwidths(1:3, ones(3), ones(3); maxwidth=0.1, minwidth=1)
    end

end


include("../examples/rijke.jl")
include("../src/clvs.jl")
using OrdinaryDiffEq
using PyPlot
	s = [7.0, 0.2]
	#n = 100000
	n = 50000
	d = N
	nRunup = 400000
	u = Rijke(rand(d),s,nRunup)
	du_trj = zeros(d,d,n)
	duT_trj = zeros(d,d,n)
	f_trj = zeros(d,n)
	for i = 1:n
		f!(view(f_trj,:,i), u, s, 1.)
		du_trj[:,:,i] = dRijke(u, s, 1.e-8)
		duT_trj[:,:,n+1-i] = du_trj[:,:,i]'
		u .= Rijke(u, s, 1)
	end
	d_u = 20
	les, tan_clvs = clvs(du_trj, d_u)
	les1, adj_clvs = clvs(duT_trj, d_u)

	tan_angles = zeros(d_u, d_u, n)
	adj_angles = zeros(d_u, d_u, n)

	for i = 1:n
		tan_angles[:,:,i] = tan_clvs[:,:,i]'*tan_clvs[:,:,i]
		adj_angles[:,:,n+1-i] = adj_clvs[:,:,i]'*
								adj_clvs[:,:,i]
	end
	n_spinup = 5000
	tan_adj_angles = zeros(d_u, d_u, n-1)
	for i = 2:n
		tan_adj_angles[:,:,i-1] = tan_clvs[:,:,i]'*
		adj_clvs[:,:,n+2-i]
	end
	tan_adj_angles = tan_adj_angles[:,:,n_spinup:n-1-n_spinup]
	mean_tan_adj_angles = 180*acos.(sum(tan_adj_angles,
										dims=3)/n)/pi
	mean_tan_adj_angles = mean_tan_adj_angles[:,:,1]

	mean_tan_angles = sum(tan_angles,dims=3)/n
	mean_tan_angles = mean_tan_angles[1:6,1:6,1]
	mean_tan_angles = 180*acos.(mean_tan_angles)/pi
	mean_adj_angles = sum(adj_angles,dims=3)/n
	mean_adj_angles = mean_adj_angles[1:6,1:6,1]
	mean_adj_angles = 180*acos.(mean_adj_angles)/pi

	fig, ax = subplots(1,1)
	ax.xaxis.set_tick_params(labelsize=25)
	ax.yaxis.set_tick_params(labelsize=25)
	angle_plot = ax.matshow(mean_tan_angles)
	cbar = fig.colorbar(angle_plot)
	cbar.ax.tick_params(labelsize=25)
	fig2, ax2 = subplots(1,1)
	ax2.xaxis.set_tick_params(labelsize=25)
	ax2.yaxis.set_tick_params(labelsize=25)
	angle_plot2 = ax2.matshow(mean_adj_angles)
	cbar = fig2.colorbar(angle_plot2)
	cbar.ax.tick_params(labelsize=25)

	fig1, ax1 = subplots(1,1)
	ax1.xaxis.set_tick_params(labelsize=25)
	ax1.yaxis.set_tick_params(labelsize=25)
	ax1.set_title(L"Probability distribution of
				  angles between $V^1$ and $V^2$",
				  fontsize=25)
	ax1.hist(acos.(tan_angles[1,2,:])*180/pi,
			 density=true,bins=100)
	# minimum angle between 1 and 2 is 15.6 degrees
	ax1.set_xlabel("Angle in degrees",fontsize=25)
	ax1.grid(true)
	fig3, ax3 = subplots(1,1)
	ax3.xaxis.set_tick_params(labelsize=25)
	ax3.yaxis.set_tick_params(labelsize=25)
	ax3.set_title(L"Probability distribution of
				  angles between $W^1$ and $W^2$",
				  fontsize=25)
	ax3.hist(acos.(adj_angles[1,2,:])*180/pi,
			 density=true,bins=100)
	# minimum angle between 1 and 2 is 4 degrees
	ax3.set_xlabel("Angle in degrees",fontsize=25)
	ax3.grid(true)

	fig4, ax4 = subplots(1,1)
	ax4.xaxis.set_tick_params(labelsize=25)
	ax4.yaxis.set_tick_params(labelsize=25)
	angle_plot4 = ax4.matshow(mean_tan_adj_angles)
	cbar = fig4.colorbar(angle_plot4)
	cbar.ax.tick_params(labelsize=25)


using Distributed
using Plots
addprocs(8)

@everywhere using DifferentialEquations
@everywhere using DiffEqParamEstim
@everywhere using BlackBoxOptim

# Define function to fix parameters

function fixvalues(vs, fixspec, dims)
    newvs = zeros(Float64, dims)
    vsidx = fixidx = 1
    for i in 1:dims
        if (fixidx <= length(fixspec) && i == fixspec[fixidx][1])
            newvs[i] = fixspec[fixidx][2]
            fixidx += 1
        else
            newvs[i] = vs[vsidx]
            vsidx += 1
        end
    end
    newvs
end

# Define DAE System: Lorenz Attractor + Extra Variable to illustarte DAE syntax
function g(du,u,p,t)
  σ,ρ,β = p
  x,y,z,w = u
  du[1] = σ*(y-x)
  du[2] = x*(ρ-z) - y
  du[3] = x*y - β*z
  du[4] = w - (x + y + z)
end

# Define Limits
lower=Float64.([0, 0, 0])
upper=Float64.([20, 30, 10])

# Define Initial Conditions as well as time span
u0 = [1.0;0.0;0.0;1.0]
t = 0.0:0.05:1
tspan = (0.0,1.0)

# Generate Mass Matrox MM to transfrom ODE into DAE
MM=zeros(4,4)
MM[1,1]=MM[2,2]=MM[3,3]=1

# Generate Equation system
ODESystem = ODEFunction(g;mass_matrix=MM)

# Generate Function that creates ODE problem as function of parameters (optmization argument)
model_ode(p_) = ODEProblem(ODESystem, u0, tspan,p_)

# Generate Function that solves system as function of parameters, saves every s interval and save only vars we want.
solve_model(mp_,s, i) = DifferentialEquations.solve(model_ode(mp_), Rodas5(),saveat=s, save_idxs=i)

# Generate Clean mock data / will add nisy mock data later
mock_data = Array(solve_model([10.0,28.0,8/3],0.05,[1,2,3,4]))#,4

# Create Loss function
function L(t,dat,w=nothing)
    return L2Loss(t,dat;data_weight=w)
end

# Create objective function
loss_objective(mp_, dat)=build_loss_objective(model_ode(mp_), Rodas5(), L(t,dat))

# Simplify fitness/objective function to depend only on parameter we want to vary
function origfitness(args)
  return loss_objective(args, mock_data)(args)
end

dims = 3

# In specific opt problem A there is one var (at position 1) with fixed value of 10.

problemfixed = [(1, 10.0)]

# Create objective function with fix value =10 at position 1
fitnessFix(vs) = origfitness(fixvalues(vs, problemfixed, dims))

# Generate options for optimization for problem with free parameters
opt_free = bbsetup(origfitness; Method=:separable_nes, SearchRange = (collect(zip(lower,upper))),
               NumDimensions = 3, MaxFuncEvals = 100000, workers=workers(), TargetFitness=100.0 )

# Generate options for optimization for problem with fix parameters (note dimensions are 2 instead of 3)
opt_fix = bbsetup(fitnessFix; Method=:separable_nes, SearchRange = (collect(zip(lower[2:3],upper[2:3]))),
               NumDimensions = 2, MaxFuncEvals = 100000, workers=workers())

# Call bboptim as usual and optimizing for fitnessFix and origfitness.
bbfix=bboptimize(opt_fix);
bbfree=bboptimize(opt_free);
bsfix=fixvalues(best_candidate(bbfix), problemfixed,dims) # then is our best solution to the general DAE problem with a fixed param.
bsfree=best_candidate(bbfree) # then is our best solution to the general DAE problem with full freedom.

# Note that within error bsfix and bs free are equal

(bsfree, origfitness(bsfree)), (bsfix, origfitness(bsfix))

# Plot results
pyplot()
# Shot Times to see Fitness to Data
Sol1=solve(ODEProblem(ODESystem, u0, (0,1.0), bsfix),Rodas5())
plot(Sol1, layout=(2,2), xlabel="Time", ylabel=[:x :y :z :w], label=["" "" "" "w=x+y+z"])
scatter!(t,[mock_data[i,:] for i in 1:4],layout=(2,2), label="")

# Long times to see attractor behavior
Sol2=solve(ODEProblem(ODESystem, u0, (0,100.0), bsfix),Rodas5())
#plot(Sol,layout=(2,2), labels=[:w :x :y "z=w+x+y"])
P1=plot(Sol2.t, [Sol2[i,:] for i in 1:4], layout=(2,2), xlabel="Time", ylabel=[:x :y :z :w], labels=["" "" "" "w=x+y+z"])
P2=plot(Sol2,vars=(1,2,3),xlabel="x", ylabel="y", zlabel="z", labels="")
l = @layout([a;b{0.5h}])
PF=plot(P1,P2, layout=l, size=(900,600))


module Models

function System(func_generic, params)
   func = (t, u)->func_generic(t, u, params)
   return func
end

function lorenz63(t, u, p)
   du = similar(u)
   du[1] = p["σ"]*(u[2]-u[1])
   du[2] = u[1]*(p["ρ"]-u[3]) - u[2]
   du[3] = u[1]*u[2] - p["β"]*u[3]

   return copy(du)
end

lorenz63_true = System(lorenz63, Dict("σ" => 10, "β" => 8/3, "ρ" => 28))

lorenz63_err = System(lorenz63, Dict("σ" => 10, "β" => 8/3 - 0.1, "ρ" => 28))
lorenz63_err2 = System(lorenz63, Dict("σ" => 10 - 0.3, "β" => 8/3, "ρ" => 28.1))
lorenz63_err3 = System(lorenz63, Dict("σ" => 10.00001, "β" => 8/3, "ρ" => 28))
lorenz63_err4 = System(lorenz63, Dict("σ" => 10, "β" => (1 + 0.1)*(8/3), "ρ" => 28))

function lorenz96(t, u, p)
   N = p["N"]

   du = similar(u)

   for i=1:N
      if i <= 10
         F = p["F1"]
      elseif i <= 20
         F = p["F2"]
      elseif i <= 30
         F = p["F3"]
      else
         F = p["F4"]
      end
      du[i] = (u[mod(i+1, 1:N)] - u[mod(i-2, 1:N)])*u[mod(i-1, 1:N)] - u[i] + F
   end

   return copy(du)
end

lorenz96_true = System(lorenz96, Dict("F1" => 8, "F2" => 10, "F3" => 12, "F4" => 14, "N" => 40))

lorenz96_err = System(lorenz96, Dict("F1" => 8, "F2" => 8, "F3" => 8, "F4" => 8, "N" => 40))
lorenz96_err2 = System(lorenz96, Dict("F1" => 12, "F2" => 12, "F3" => 12, "F4" => 12, "N" => 40))
lorenz96_err3 = System(lorenz96, Dict("F1" => 14, "F2" => 14, "F3" => 14, "F4" => 14, "N" => 40))
lorenz96_err4 = System(lorenz96, Dict("F1" => 10, "F2" => 10, "F3" => 10, "F4" => 10, "N" => 40))

lorenz96_half_true = System(lorenz96, Dict("F1" => 8, "F2" => 10, "F3" => 12, "F4" => 14, "N" => 20))

function lorenz96_twoscale(t, u, p)
   N = p["N"]
   n = p["n"]

   dx = zeros(N)
   dy = zeros(n, N)

   u = reshape(u, n + 1, N)
   x = u[1, :]
   y = u[2:end, :]

   for i=1:N
      if i <= 10
         F = p["F1"]
      elseif i <= 20
         F = p["F2"]
      elseif i <= 30
         F = p["F3"]
      else
         F = p["F4"]
      end
      dx[i] = (x[mod(i+1, 1:N)] - x[mod(i-2, 1:N)])*x[mod(i-1, 1:N)] - x[i] + F - p["h"]*p["c"]/p["b"]*sum(y[:, i])

      for j=1:n
         dy[j, i] = p["c"]*p["b"]*y[mod(j+1, 1:n), i]*(y[mod(j-1, 1:n), i] - y[mod(j+2, 1:n), i]) - p["c"]*y[j, i] + p["h"]*p["c"]/p["b"]*x[i]
      end
   end

   du = vec([dx dy']')

   return du
end

lorenz96_twoscale_true = System(lorenz96_twoscale, Dict("F1" => 8, "F2" => 10, "F3" => 12, "F4" => 14, "b" => 10, "c" => 10, "h" => 1.0, "N" => 20, "n" => 10))
lorenz96_twoscale_err = System(lorenz96_twoscale, Dict("F1" => 8, "F2" => 8, "F3" => 8, "F4" => 8, "b" => 10, "c" => 10, "h" => 1.0, "N" => 20, "n" => 10))

end


"""
    Module EntropyStableEuler

Includes general math tools
"""

module EntropyStableEuler
using StaticArrays

const γ = 1.4

export logmean, γ
include("./logmean.jl")

# # submodules
# export Fluxes1D,Fluxes2D,Fluxes3D

# export entropy_scaling,scale_entropy_input,scale_entropy_output
# changes definition of entropy variables by a constant scaling
#const entropy_scaling = 1/γ # implies vE = -1 /(ι * γ) and γι = non-dim temperature if R* = cp
const entropy_scaling = 1
scale_entropy_output(V...) = (x->@. x*entropy_scaling).(V)
scale_entropy_input(V...) = (x->@. x/entropy_scaling).(V)

#####
##### one-dimensional fluxes
#####

module Fluxes1D
import ..γ
import ..entropy_scaling
import ..EntropyStableEuler: logmean
import ..EntropyStableEuler: scale_entropy_output, scale_entropy_input

export wavespeed_1D

export primitive_to_conservative, conservative_to_primitive_beta
export u_vfun, v_ufun
export euler_fluxes
export Sfun,pfun,betafun
include("./euler_fluxes_1D.jl")
include("./entropy_variables.jl")
end

#####
##### two-dimensional fluxes
#####
module Fluxes2D
import ..γ
import ..entropy_scaling
import ..EntropyStableEuler: logmean, entropy_scaling
import ..EntropyStableEuler: scale_entropy_output, scale_entropy_input

export primitive_to_conservative,conservative_to_primitive_beta
export u_vfun, v_ufun
export Sfun,pfun,betafun
export euler_fluxes

include("./entropy_variables.jl")
include("./euler_fluxes_2D.jl")
end

#####
##### three-dimensional fluxes
#####
module Fluxes3D
import ..γ
import ..entropy_scaling
import ..EntropyStableEuler: logmean, entropy_scaling
import ..EntropyStableEuler: scale_entropy_output, scale_entropy_input

export primitive_to_conservative,conservative_to_primitive_beta
export u_vfun, v_ufun
export Sfun,pfun,betafun
export euler_fluxes

include("./entropy_variables.jl")
include("./euler_fluxes_3D.jl")
end

# export u_vfun, v_ufun, betafun, pfun, rhoe_ufun
# export dVdU_explicit, dUdV_explicit
# export wavespeed # c
# export Sfun, sfun # math/physical entropies
# export u_vfun1D, v_ufun1D, betafun1D # specialization to 1D
# export primitive_to_conservative
# include("./euler_variables.jl")

# export euler_fluxes
# export euler_flux_x, euler_flux_y # separate x,y fluxes for faster implicit assembly using ForwardDiff
# include("./euler_fluxes.jl")

# export vortex
# include("./analytic_solutions.jl")

end


using Revise
using AdFem

mesh = Mesh(8,8,1/8)
ρ = ones(get_ngauss(mesh))
B = compute_fem_advection_matrix1(2*ρ, 3*ρ, mesh)

ρ = constant(ones(get_ngauss(mesh)))
B2 = compute_fem_advection_matrix1(2*ρ, 3*ρ, mesh)

sess = Session(); init(sess)
maximum(abs.(run(sess, B2) - B))


using MFCC, MFCC.Mel, MFCC.Filterbanks, Plots

windows, windowstarts = mffilters(10, 22050, 22050)
filterbank = Filterbank(22050, windows, windowstarts)
melscale = mel.(range(1, length=22050, 22050))

# Plot the curve of the hertz scale by the mel scale
melhz = plot(melscale, xlabel="Frequency(Hz)", ylabel="Frequency(Mel)")
savefig(melhz, "docs/assets/mel_hz_mapping.svg")

# Plot the filterbank in the mel scale
melfb = plot(melscale, getfilters(filterbank), xlabel="Frequency(Mel)")
savefig(melfb, "docs/assets/mel_scale_filterbank.svg")

# Plot the filterbank in the hertz scale
hzfb = plot(getfilters(filterbank), xlabel="Frequency(Hz)")
savefig(hzfb, "docs/assets/hz_scale_filterbank.svg")


# This file was generated, do not modify it. # hide
using DelimitedFiles

function cnt(z,c)
    k::UInt8 = 0
    while (k < 255)
        z = z^2 + c
        if abs2(z)>4.
            return k
        end
        k += 1
    end
    return k
end

function mand(M, N)
    grid = Array{UInt8}(undef, M, N)
    init_z = complex(0.0, 0.0)
    xs = range(-2,2,length=N)
    ys = range(-2,2,length=M)
    grid = cnt.(init_z,complex.(xs',ys))
    return grid
end

function main()
    s = time()
    M = N = 5000
    grid = mand(M, N)
    t = time()
    @show(t - s)
    writedlm("/tmp/resultJ2.txt",grid,",")
end

main()

using BenchmarkTools
@btime mand(5000, 5000);