using FrankWolfe
using LinearAlgebra
using Test
using SparseArrays

@testset "Testing Blended Pairwise Conditional Gradients" begin
    f(x) = norm(x)^2
    function grad!(storage, x)
        @. storage = 2x
        return nothing
    end
    lmo_prob = FrankWolfe.ProbabilitySimplexOracle(4)
    x0 = FrankWolfe.compute_extreme_point(lmo_prob, zeros(10))
    res_bpcg = FrankWolfe.blended_pairwise_conditional_gradient(
        f,
        grad!,
        lmo_prob,
        x0,
        max_iteration=6000,
        line_search=FrankWolfe.Adaptive(),
        verbose=true,
        epsilon=3e-7,
    )
    res_afw = FrankWolfe.away_frank_wolfe(
        f,
        grad!,
        lmo_prob,
        x0,
        max_iteration=6000,
        line_search=FrankWolfe.Adaptive(),
        print_iter=100,
        verbose=true,
        epsilon=3e-7,
    )
    @test res_afw[3] ≈ res_bpcg[3]
    @test norm(res_afw[1] - res_bpcg[1]) ≈ 0 atol=1e-6
    res_bpcg2 = FrankWolfe.blended_pairwise_conditional_gradient(
        f,
        grad!,
        lmo_prob,
        x0,
        max_iteration=6000,
        line_search=FrankWolfe.Adaptive(),
        verbose=false,
        lazy=true
    )
    @test res_bpcg2[3] ≈ res_bpcg[3] atol=1e-5
    active_set_afw = res_afw[end]
    storage = copy(active_set_afw.x)
    grad!(storage, active_set_afw.x)
    epsilon = 1e-6
    @inferred FrankWolfe.afw_step(active_set_afw.x, storage, lmo_prob, active_set_afw, epsilon)
    @inferred FrankWolfe.lazy_afw_step(active_set_afw.x, storage, lmo_prob, active_set_afw, epsilon, epsilon)
end


module TestGraphGenerators
using Test

using Catlab.Graphs.BasicGraphs, Catlab.Graphs.GraphGenerators
using Catlab.CategoricalAlgebra: homomorphisms

# Path graphs
#------------

n = 5
g = path_graph(Graph, n)
@test (nv(g), ne(g)) == (n, n-1)
g = path_graph(SymmetricGraph, n)
@test (nv(g), ne(g)) == (n, 2(n-1))

# Cycle graphs
#-------------

g = cycle_graph(Graph, n)
@test (nv(g), ne(g)) == (n, n)
g = cycle_graph(SymmetricGraph, n)
@test (nv(g), ne(g)) == (n, 2n)

# Complete graphs
#----------------

for T in (Graph, SymmetricGraph)
  g = complete_graph(T, n)
  @test (nv(g), ne(g)) == (n, n*(n-1))
end
for T in (ReflexiveGraph, SymmetricReflexiveGraph)
  g = complete_graph(T, n)
  @test (nv(g), ne(g)) == (n, n*n)
end

# Star graphs
#------------

g = star_graph(Graph, n)
@test (nv(g), ne(g)) == (n, n-1)
@test length(unique(src(g,e) for e in edges(g))) == 1
g = star_graph(SymmetricGraph, n)
@test (nv(g), ne(g)) == (n, 2(n-1))

# Wheel graphs
#-------------

g = wheel_graph(Graph, n)
@test (nv(g), ne(g)) == (n, 2(n-1))
triangle = Graph(3)
add_edges!(triangle, [1,2,1], [2,3,3])
@test length(homomorphisms(triangle, g)) == n-1

g = wheel_graph(SymmetricGraph, n)
@test (nv(g), ne(g)) == (n, 4(n-1))
triangle = cycle_graph(SymmetricGraph, 3)
@test length(homomorphisms(triangle, g)) == 6(n-1) # == 3! * (n-1)

# Parallel arrows
#----------------

g = parallel_arrows(Graph, n)
@test (nv(g), ne(g)) == (2, n)
@test all(==(1), src(g))
@test all(==(2), tgt(g))

# Erdos-Renyi graphs
#-------------------

g = erdos_renyi(Graph, 100, 0.0)
@test (nv(g), ne(g)) == (100, 0)
g = erdos_renyi(Graph, 100, 0)
@test (nv(g), ne(g)) == (100, 0)
g = erdos_renyi(Graph, 100, 5)
@test (nv(g), ne(g)) == (100, 5)
g = erdos_renyi(Graph, 100, 1.0)
@test g == complete_graph(Graph, 100)

# Expected Degree Graph
#----------------------

g = expected_degree_graph(Graph, 5 .* ones(100))
@test nv(g) == 100

# Watts Strogatz Graph
#---------------------

g = watts_strogatz(Graph, 100, 4, 0.2)
@test (nv(g), ne(g)) == (100, 200)

end


using LinearAlgebra

function time_slice_spectrum()
    H_matrix = kinetic_matrix_imgtime * auxiliary_field_matrix() * kinetic_matrix_imgtime
end

function sub_partition_weight(L)
    matrix_product = Matrix{Complex{Float64}}(I, num_sites, num_sites)
    for i = 1 : L
        matrix_product *= time_slice_spectrum()
    end
    #return matrix_product
    #if flag_background_subtraction == 1
        #matrix_product *= exp(sum(β * (U/2) * background_dist.^2))
        #return matrix_product
    #end
    effective_spectrum, u = eigen(matrix_product) # spectrum element: e^(-βϵᵢ)
    u_inv = inv(u)
    U_array = complex(zeros(num_sites, num_sites, num_sites))
    for i = 1 : num_sites
        for j = 1 : num_sites
            for k = 1 : num_sites
                U_array[i, j, k] = u_inv[k, i] * u[j, k]
            end
        end
    end
    return effective_spectrum, U_array
end

function recursion_partition_function(effective_spectrum)
    Z = complex(zeros(num_bosons))
    zₖ = complex(zeros(num_bosons))
    for k = 1 : num_bosons
        zₖ[k] = sum(effective_spectrum .^ k)
    end
    for N = 1 : num_bosons
        for M = 1 : N
            if M == N
                Z[N] += 1 * zₖ[M]
            else
                Z[N] += Z[N-M] * zₖ[M]
            end
        end
        Z[N] = Z[N] / N
    end
    return Z
end

function recursion_OccNum(Z, effective_spectrum) # Calculating <nᵢ>
    n_i = complex(zeros(num_bosons, num_sites))
    n_i[1, :] = effective_spectrum
    for N = 2 : num_bosons
        n_i[N, :] = effective_spectrum .* (Z[N-1] .+ n_i[N-1, :])
    end
    return n_i
end

function recursion_double_OccNum_eq(Z, effective_spectrum, n_i) # Calculating <nᵢ^2>
    n²_i = complex(zeros(num_bosons, num_sites)) # <nᵢ^2> array
    n²_i[1, :] = effective_spectrum
    for N = 2 : num_bosons
        n²_i[N, :] = effective_spectrum .* (Z[N-1] .+ 2 * n_i[N-1, :] .+ n²_i[N-1, :])
    end
    return n²_i[num_bosons, :]
end

function recursion_double_OccNum_neq(Z, effective_spectrum, n_i) # Calculating <nᵢnⱼ>
    n²_i = complex(zeros(num_bosons, num_sites, num_sites)) # <nᵢnⱼ> array
    for i = 1 : num_sites
        for j = i + 1 : num_sites
            for N = 2 : num_bosons
                n²_i[N, i, j] = effective_spectrum[j] / 2 * (n_i[N-1, i] + n²_i[N-1, i, j]) +
                            + effective_spectrum[i] / 2 * (n_i[N-1, j] + n²_i[N-1, i, j])
            end
            n²_i[num_bosons, j, i] = n²_i[num_bosons, i, j]
        end
    end
    return n²_i[num_bosons, :, :]
end

#function measure_condensate_fraciton()
#end

function measure_TotEnergy(num_iterations)
    partition_func_array = Vector{Complex{Float64}}()
    Ek_array = Vector{Complex{Float64}}()
    Ep_array = Vector{Complex{Float64}}()
    E_tot_array = Vector{Complex{Float64}}()
    for N = 1 : num_iterations
        effective_spectrum, U_array = sub_partition_weight(L)
        Z = recursion_partition_function(effective_spectrum)
        n_i = recursion_OccNum(Z, effective_spectrum)
        n²_i_diag = recursion_double_OccNum_eq(Z, effective_spectrum, n_i)
        n²_i_offdiag = recursion_double_OccNum_neq(Z, effective_spectrum, n_i)
        E_onebody, E_twobody = 0, 0
        # one-body energy calculation
        for i = 1 : num_sites
            for j = 1 : num_sites
                if abs(i - j) == 1 || abs(i - j) == num_sites - 1
                    D_ij = sum(U_array[i, j, :] .* n_i[num_bosons, :])
                    E_onebody += -t * D_ij
                end
            end
        end
        # two-body energy calculation
        for i = 1 : num_sites
            U_matrix = complex(zeros(num_sites, num_sites))
            D_iiii = 0
            for p = 1 :num_sites
                for r = 1 : num_sites
                    if p != r
                        U_matrix[p, r] = U_array[i, i, p] * U_array[i, i, r]
                        D_iiii += U_matrix[p, r] * n_i[num_bosons, p]
                    end
                end
                D_iiii += U_array[i, i, p]^2 * n²_i_diag[p]
            end
            D_iiii += sum( 2 * U_matrix .* n²_i_offdiag)
            E_twobody += U/2 * D_iiii
        end
        push!(partition_func_array, Z[num_bosons])
        push!(Ek_array, E_onebody)
        push!(Ep_array, E_twobody)
        push!(E_tot_array, E_onebody + E_twobody)
    end
    Z_mean = mean(partition_func_array)
    Ek_mean = mean(Ek_array) / Z_mean
    Ek_error = std(Ek_array) / (sqrt(num_iterations) * real(Z_mean))
    Ep_mean = mean(Ep_array) / Z_mean
    Ep_error = std(Ep_array) / (sqrt(num_iterations) * real(Z_mean))
    E_tot_mean = mean(E_tot_array) / Z_mean
    E_tot_error_bar = std(E_tot_array) / (sqrt(num_iterations) * real(Z_mean))
    return Ek_mean, Ek_error, Ep_mean, Ep_error, E_tot_mean, E_tot_error_bar
end


include("../src/CleverAlgorithms.jl")

using CA

val = CA.Stochastic.randdom_search((x) -> x^2, (-5,5))
println(val)
val = CA.Stochastic.randdom_search((x) -> x^2, (-5,5); op = >)
println(val)


 #=

Implementations of cg and pcg that use BLAS when they can.

Started by Dan Spielman.
Contributors:

=#

"""
`cg(mat, b; tol, maxits, maxtime, verbose)` solves a symmetric linear system.
`tol` is set to 1e-6 by default,
`maxits` defaults to Inf
`maxtime` defaults to Inf.  It measures seconds.
`verbose` defaults to false
"""
function cg end

"""
`pcg(mat, b, pre; tol, maxits, maxtime, verbose)` solves a symmetric linear system using preconditioner `pre`.
`pre` can be a function or a matrix.  If a matrix, a function to solve it is created with cholFact.
`tol` is set to 1e-6 by default,
`maxits` defaults to Inf
`maxtime` defaults to Inf.  It measures seconds.
`verbose` defaults to false
"""
function pcg end

"""
`cgSolver(mat; tol, maxits, maxtime, verbose)` creates a solver for a PSD system `mat`.
The parameters are as described in cg.
"""
function cgSolver end

"""
`pcgSolver(mat, pre; tol, maxits, maxtime, verbose)` creates a solver for a PSD system using preconditioner `pre`.
The parameters are as described in pcg.
"""
function pcgSolver end

function cg(mat, b::Array{Float64,1}; tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false)
    cgBLAS(mat, b, tol=tol, maxits=maxits, maxtime=maxtime, verbose=verbose)
end

function cg(mat, b::Array{Float32,1}; tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false)
    cgBLAS(mat, b, tol=tol, maxits=maxits, maxtime=maxtime, verbose=verbose)
end

function cg(mat, b; tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=true)
    cgSlow(mat, b, tol=tol, maxits=maxits, maxtime=maxtime, verbose=verbose)
end

function cgSolver(mat; tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false)
    f(b; tol::Real=1e-6, maxits=maxits, maxtime=maxtime, verbose=verbose) = cg(mat, b, tol=tol, maxits=maxits, maxtime=maxtime, verbose=verbose)
end

function pcg(mat, b, pre::Union{AbstractArray,Matrix}; tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false)
    fact = cholfact(pre)
    F = x->(fact \ x)
    pcg(mat, b, F; tol=tol, maxits=maxits, maxtime=maxtime, verbose=verbose)
end

function pcg(mat, b::Array{BigFloat,1}, pre::Function; tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false)
    pcgBLAS(mat, b, pre, tol=tol, maxits=maxits, maxtime=maxtime, verbose=verbose)
end

function pcg(mat, b::Array{Float64,1}, pre::Function; tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false)
    pcgBLAS(mat, b, pre, tol=tol, maxits=maxits, maxtime=maxtime, verbose=verbose)
end

function pcg(mat, b::Array{Float32,1}, pre::Function; tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false)
    pcgBLAS(mat, b, pre, tol=tol, maxits=maxits, maxtime=maxtime, verbose=verbose)
end

function pcg(mat, b, pre::Function; tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false)
    pcgSlow(mat, b, pre, tol=tol, maxits=maxits, maxtime=maxtime, verbose=verbose)
end

function pcgSolver(mat, pre; tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false)
    f(b; tol::Real=1e-6, maxits=maxits, maxtime=maxtime, verbose=verbose) = pcg(mat, b, pre, tol=tol, maxits=maxits, maxtime=maxtime, verbose=verbose)
end

"""Create a solver that uses cg to solve Laplacian systems in mat. Specialized for the case when pre is a Laplacian matrix.  Fix the default parameters of the solver as given"""
function pcgLapSolver(mat, pre; tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false)
    fact = lapChol(pre)
    f(b; tol::Real=1e-6, maxits=maxits, maxtime=maxtime, verbose=verbose) = pcg(mat, b, fact, tol=tol, maxits=maxits, maxtime=maxtime, verbose=verbose)

end

# uses BLAS.  As fast as Matlab's pcg.
# set to use similar paramaters
function cgBLAS{Tval}(mat, b::Array{Tval,1};
        tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false)

    n = size(mat,2)

    x = zeros(Tval,n)

    tol = tol * norm(b)
    r = copy(b)
    p = copy(r)

    # the following would have higher accuracy
    #       rho = sum(r.^2)
    rho = norm(r)^2

    t1 = time()

    itcnt = 0
    while itcnt < maxits
        itcnt = itcnt+1

        q = mat*p

        al = rho/dot(p, q)

        BLAS.axpy!(al,p,x)  # x = x + al * p

        BLAS.axpy!(-al,q,r)  # r -= al*q

        if norm(r) < tol #Converged?
            break
        end

        oldrho = rho
        rho = norm(r)^2

        # the following would have higher accuracy
        #       rho = sum(r.^2)

        beta = rho/oldrho

        # p = r + beta*p
        BLAS.scal!(n,beta,p,1) # p *= beta
        BLAS.axpy!(1.0,r,p) # p += r

        if (time() - t1) > maxtime
            if verbose
                println("CG stopped at maxtime.")
            end
            break
        end
      end

    if verbose
        println("CG stopped after: ", itcnt, " iterations with relative error ", (norm(r)/norm(b)), ".")
    end

    return x
end

# For when you can't use BLAS.
# set to use similar paramaters similar to MATLAB
function cgSlow{Tval}(mat, b::Array{Tval,1};
        tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false)

    n = size(mat,2)

    x = zeros(Tval,n)

    tol = tol * norm(b)
    r = copy(b)
    p = copy(r)

    # the following would have higher accuracy
    #       rho = sum(r.^2)
    rho = norm(r)^2

    t1 = time()

    itcnt = 0
    while itcnt < maxits
        itcnt = itcnt+1

        q = mat*p

        al = rho/dot(p, q)

        # x = x + al * p
        for i in 1:n
            x[i] = x[i] + al*p[i]
        end

        # r = r - al*q
        for i in 1:n
            r[i] = r[i] - al*q[i]
        end

        if norm(r) < tol
            break
        end

        oldrho = rho
        rho = norm(r)^2

        # the following would have higher accuracy
        #       rho = sum(r.^2)

        beta = rho/oldrho

        # p = r + beta*p
        for i in 1:n
            p[i] = r[i] + beta*p[i]
        end

        if (time() - t1) > maxtime
            if verbose
                println("CG stopped at maxtime.")
            end
            break
        end
      end

    if verbose
        println("CG stopped after: ", itcnt, " iterations with relative error ", (norm(r)/norm(b)), ".")
    end

    return x
end

#====================================
   PCG Code
=====================================#

# uses BLAS.  As fast as Matlab's pcg.
# set to use similar parameters "
function pcgBLAS{Tval}(mat, b::Array{Tval,1}, pre;
        tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false)

    n = size(mat,2)

    x = zeros(Tval,n)

    tol = tol * norm(b)
    r = copy(b)
    z = pre(r)
    p = copy(z)

    rho = dot(r, z) # BLAS.dot does not seem faster

    t1 = time()

    itcnt = 0
    while itcnt < maxits
        itcnt = itcnt+1

        q = mat*p

        if dot(p,q) == 0
            break
        end

        al = rho/dot(p, q)

        BLAS.axpy!(al,p,x)  # x = x + al * p

        BLAS.axpy!(-al,q,r)  # r -= al*q

        if norm(r) < tol #Converged?
            break
        end

        # here is the top of the code in numerical templates

        z = pre(r)

        oldrho = rho
        rho = dot(z, r)

        # the following would have higher accuracy
        #       rho = sum(r.^2)

        beta = rho/oldrho

        p = z + beta*p
        # BLAS.scal!(n,beta,p,1) # p *= beta
        # BLAS.axpy!(1.0,z,p) # p += z

        if (time() - t1) > maxtime
            if verbose
                println("PCG stopped at maxtime.")
            end
            break
        end

      end

    if verbose
        println("PCG stopped after: ", itcnt, " iterations with relative error ", (norm(r)/norm(b)), ".")
    end

    return x
end

function pcgSlow{Tval}(mat, b::Array{Tval,1}, pre;
        tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false)

    n = size(mat,2)

    x = zeros(Tval,n)

    tol = tol * norm(b)
    r = copy(b)
    z = pre(r)
    p = copy(z)

    rho = dot(r, z)

    t1 = time()

    itcnt = 0
    while itcnt < maxits
        itcnt = itcnt+1

        q = mat*p

        al = rho/dot(p, q)

        # x = x + al * p
        for i in 1:n
            x[i] = x[i] + al*p[i]
        end

        # r = r - al*q
        for i in 1:n
            r[i] = r[i] - al*q[i]
        end

        if norm(r) < tol #Converged?
            break
        end

        # here is the top of the code in numerical templates

        z = pre(r)

        oldrho = rho
        rho = dot(z, r)

        # the following would have higher accuracy
        #       rho = sum(r.^2)

        beta = rho/oldrho

        # p = z + beta*p
        for i in 1:n
            p[i] = z[i] + beta*p[i]
        end

        if (time() - t1) > maxtime
            if verbose
                println("PCG stopped at maxtime.")
            end
            break
        end

    end

    if verbose
        println("PCG stopped after: ", itcnt, " iterations with relative error ", (norm(r)/norm(b)), ".")
    end

    return x
end


"""
Rational Krylov helper module. Contains functionality needed for methods
like NLEIGS and CORK.
"""
module RKHelper

include("inpolygon.jl")
include("discretizepolygon.jl")
include("rk_nep.jl")
include("rk_utils.jl")
include("linsolvercache.jl")

end


#GracePlot demonstration 5: Another multi-plot
#-------------------------------------------------------------------------------

using GracePlot

#==Input data
===============================================================================#
x=collect(-10:0.1:10)
y = []
for i in 1:4
	push!(y, (x.^i)./(10^i)) #Normalized powers of x
end
titles = ["Linear", "Quadratic", "Cubic", "Quartic"]

#=="Defaults"
===============================================================================#

pdefaults = defaults(font="Courier-BoldOblique",
	linewidth=2.5, linestyle=:solid,
	charsize=2, symbolsize=.75
)

#==Generate plot
===============================================================================#
plot = GracePlot.new()
set(plot, pdefaults)

#Add 4 subplots:
for gidx in 0:3
	g = add(plot, subtitle = titles[gidx+1])

	#Add 10 datasets:
	for i in 1:10
		ds = add(g, x, i.*y[gidx+1], glyph(shape=i, skip=20))
	end

	autofit(g)
end

#Make it a multi-graph plot:
#arrange(plot, (1, 4))
arrange(plot, (2, 3))

#Finalize:
redraw(plot)

#Last line


using FactCheck, Base.Test, Images

srand(2015)

facts("Distances") do
  context("Hausdorff") do
    A = eye(3); B = copy(A); C = copy(A)
    B[1,2] = 1; C[1,3] = 1
    @fact hausdorff_distance(A,A) == 0 --> true
    @fact hausdorff_distance(A,B) == hausdorff_distance(B,A) --> true
    @fact hausdorff_distance(A,B) < hausdorff_distance(A,C) --> true
    A = rand([0,1],10,10)
    B = rand([0,1],10,10)
    @fact hausdorff_distance(A,B) ≥ 0 --> true
  end
end


module FunctionWrapper

using JavaCall

const Abs = @jimport org.hipparchus.analysis.function.Abs
const Acos = @jimport org.hipparchus.analysis.function.Acos
const Acosh = @jimport org.hipparchus.analysis.function.Acosh
const Add = @jimport org.hipparchus.analysis.function.Add
const Asin = @jimport org.hipparchus.analysis.function.Asin
const Asinh = @jimport org.hipparchus.analysis.function.Asinh
const Atan = @jimport org.hipparchus.analysis.function.Atan
const Atan2 = @jimport org.hipparchus.analysis.function.Atan2
const Atanh = @jimport org.hipparchus.analysis.function.Atanh
const Cbrt = @jimport org.hipparchus.analysis.function.Cbrt
const Ceil = @jimport org.hipparchus.analysis.function.Ceil
const Class = @jimport java.lang.Class
const Constant = @jimport org.hipparchus.analysis.function.Constant
const Cos = @jimport org.hipparchus.analysis.function.Cos
const Cosh = @jimport org.hipparchus.analysis.function.Cosh
const Derivative = @jimport org.hipparchus.analysis.differentiation.Derivative
const DerivativeStructure = @jimport org.hipparchus.analysis.differentiation.DerivativeStructure
const Divide = @jimport org.hipparchus.analysis.function.Divide
const Exp = @jimport org.hipparchus.analysis.function.Exp
const Expm1 = @jimport org.hipparchus.analysis.function.Expm1
const Floor = @jimport org.hipparchus.analysis.function.Floor
const Gaussian = @jimport org.hipparchus.analysis.function.Gaussian
const HarmonicOscillator = @jimport org.hipparchus.analysis.function.HarmonicOscillator
const Identity = @jimport org.hipparchus.analysis.function.Identity
const Inverse = @jimport org.hipparchus.analysis.function.Inverse
const JString = @jimport java.lang.String
const Log = @jimport org.hipparchus.analysis.function.Log
const Log10 = @jimport org.hipparchus.analysis.function.Log10
const Log1p = @jimport org.hipparchus.analysis.function.Log1p
const Logistic = @jimport org.hipparchus.analysis.function.Logistic
const Logit = @jimport org.hipparchus.analysis.function.Logit
const Max = @jimport org.hipparchus.analysis.function.Max
const Min = @jimport org.hipparchus.analysis.function.Min
const Minus = @jimport org.hipparchus.analysis.function.Minus
const Multiply = @jimport org.hipparchus.analysis.function.Multiply
const Object = @jimport java.lang.Object
const Pow = @jimport org.hipparchus.analysis.function.Pow
const Power = @jimport org.hipparchus.analysis.function.Power
const Rint = @jimport org.hipparchus.analysis.function.Rint
const Sigmoid = @jimport org.hipparchus.analysis.function.Sigmoid
const Signum = @jimport org.hipparchus.analysis.function.Signum
const Sin = @jimport org.hipparchus.analysis.function.Sin
const Sinc = @jimport org.hipparchus.analysis.function.Sinc
const Sinh = @jimport org.hipparchus.analysis.function.Sinh
const Sqrt = @jimport org.hipparchus.analysis.function.Sqrt
const StepFunction = @jimport org.hipparchus.analysis.function.StepFunction
const Subtract = @jimport org.hipparchus.analysis.function.Subtract
const Tan = @jimport org.hipparchus.analysis.function.Tan
const Tanh = @jimport org.hipparchus.analysis.function.Tanh
const Ulp = @jimport org.hipparchus.analysis.function.Ulp

export Abs
export Acos
export Acosh
export Add
export Asin
export Asinh
export Atan
export Atan2
export Atanh
export Cbrt
export Ceil
export Constant
export Cos
export Cosh
export Divide
export Exp
export Expm1
export Floor
export Gaussian
export HarmonicOscillator
export Identity
export Inverse
export Log
export Log10
export Log1p
export Logistic
export Logit
export Max
export Min
export Minus
export Multiply
export Pow
export Power
export Rint
export Sigmoid
export Signum
export Sin
export Sinc
export Sinh
export Sqrt
export StepFunction
export Subtract
export Tan
export Tanh
export Ulp
export equals
export get_class
export hash_code
export notify
export notify_all
export to_string
export value
export wait

include("abs.jl")
include("acos.jl")
include("acosh.jl")
include("add.jl")
include("asin.jl")
include("asinh.jl")
include("atan.jl")
include("atan2.jl")
include("atanh.jl")
include("cbrt.jl")
include("ceil.jl")
include("constant.jl")
include("cos.jl")
include("cosh.jl")
include("divide.jl")
include("exp.jl")
include("expm1.jl")
include("floor.jl")
include("gaussian.jl")
include("harmonic_oscillator.jl")
include("identity.jl")
include("inverse.jl")
include("log.jl")
include("log10.jl")
include("log1p.jl")
include("logistic.jl")
include("logit.jl")
include("max.jl")
include("min.jl")
include("minus.jl")
include("multiply.jl")
include("pow.jl")
include("power.jl")
include("rint.jl")
include("sigmoid.jl")
include("signum.jl")
include("sin.jl")
include("sinc.jl")
include("sinh.jl")
include("sqrt.jl")
include("step_function.jl")
include("subtract.jl")
include("tan.jl")
include("tanh.jl")
include("ulp.jl")

end


@testset "Random walks" begin
    """is_nonbacktracking: a predicate that tests if a walk is nonbacktracking.
        That is that for no index i walk[i+2]==walk[i]
    """
    function is_nonbacktracking(walk)
        n = length(walk)
        if n < 3
           return true
        end
        for i in 1:(n - 2)
            if walk[i + 2] == walk[i]
                return false
            end
        end
        return true
    end

    """test_nbw: check that a walk is nonbacktracking and print it if is isn't.
        Used only for testing and debugging.
    """
    function test_nbw(g, start, len)
        w = @inferred(non_backtracking_randomwalk(g, start, len))
        if is_nonbacktracking(w)
            return true
        else
            print("walk was:\n  $w")
        end
        return false
    end
    gx = PathDiGraph(10)
    for g in testdigraphs(gx)
      @test @inferred(randomwalk(g, 1, 5)) == [1:5;]
      @test @inferred(randomwalk(g, 2, 100)) == [2:10;]
      @test_throws BoundsError randomwalk(g, 20, 20)
      @test @inferred(non_backtracking_randomwalk(g, 10, 20)) == [10]
      @test @inferred(non_backtracking_randomwalk(g, 1, 20)) == [1:10;]
    end

    gx = PathGraph(10)
    for g in testgraphs(gx)
      @test @inferred(saw(g, 1, 20)) == [1:10;]
      @test_throws BoundsError saw(g, 20, 20)
      @test @inferred(non_backtracking_randomwalk(g, 1, 20)) == [1:10;]
      @test_throws BoundsError non_backtracking_randomwalk(g, 20, 20)
    end

    gx = SimpleDiGraph(PathGraph(10))
    for g in testdigraphs(gx)
      @test @inferred(non_backtracking_randomwalk(g, 1, 20)) == [1:10;]
      @test_throws BoundsError non_backtracking_randomwalk(g, 20, 20)
    end

    gx = CycleGraph(10)
    for g in testgraphs(gx)
      visited = @inferred(non_backtracking_randomwalk(g, 1, 20))
      @test visited == [1:10; 1:10;] || visited == [1; 10:-1:1; 10:-1:2;]
    end

    gx = CycleDiGraph(10)
    for g in testdigraphs(gx)
      @test @inferred(non_backtracking_randomwalk(g, 1, 20)) == [1:10; 1:10;]
    end

    n = 10
    gx = CycleGraph(n)
    for k = 3:(n - 1)
        add_edge!(gx, 1, k)
    end

    for g in testgraphs(gx)
      for len = 1:(3 * n)
        @test test_nbw(g, 1, len)
        @test test_nbw(g, 2, len)
      end
    end
    #test to make sure it works with self loops.
    add_edge!(gx, 1, 1)
    for g in testgraphs(gx)
      for len = 1:(3 * n)
          @test test_nbw(g, 1, len)
          @test test_nbw(g, 2, len)
      end
    end
end


# filling

function fills(filler::T, data::AbstractVector{T}) where {T}
    nvals  = axes(data)[1].stop
    return fill(filler, nvals)
end

function fills(filler::T1, data::AbstractVector{T2}) where {T1, T2}
    elemtype = Union{T1,T2}
    nvals  = axes(data)[1].stop
    result = Array{elemtype,1}(undef, nvals)
    result[:] = filler
    return result
end

# number of values to be obtained

function nrolled(seqlength::T, windowspan::T) where {T<:Signed}
    (0 < windowspan <= seqlength) || throw(SpanError(seqlength,windowspan))

    return seqlength - windowspan + 1
end

# number of values to be imputed

function nfilled(windowspan::T) where {T<:Signed}
    windowspan < 1 && throw(SpanError(seqlength,windowspan))

    return windowspan - 1
end

# local exceptions

SpanError(seqlength, windowspan) =
    ErrorException("\n\tBad window span ($windowspan) for length $seqlength.\n" )

WeightsError(nweighting, windowspan) =
    ErrorException("\n\twindowspan ($windowspan) != length(weighting) ($nweighting))).\n" )


using MatrixProfile
# Download data from https://sites.google.com/site/snippetfinderinfo/home/MixedBag.zip
datapath = "/home/fredrikb/Downloads/MixedBag/"
cd(datapath)

## Some data
T = parse.(Float64, split(join(Char.(read("01911m_02019m_III_7680_200.txt"))), ','))
snips = snippets(T, 3, 200, m=50)
plot(snips)

## Power demand
T = parse.(Float64, readlines("Powerdemand_12_4500_200.txt"))
snips = snippets(T, 4, 24, m=24)
plot(snips)

profile = matrix_profile(T, 24)
mot = motifs(profile, 4, r=3, th=48)
plot(profile, mot, legend=false)
plot(mot)

## Walking styles
T = parse.(Float64, readlines("PAMAP_Subject4_NormalWalking_NordicWalking_17001_500.txt"))
T = T[1:2:end]
snips = snippets(T, 4, 200)
plot(snips)

profile = matrix_profile(T, 50)
mot = motifs(profile, 4, r=3)
plot(mot)
plot(profile, mot, legend=false)

## Robot dog
T = parse.(Float64, readlines("RoboticDogActivityY_64_4000_400.txt"))
snips = snippets(T, 4, 100)
plot(snips)

profile = matrix_profile(T, 50)
mot = motifs(profile, 4, r=2, th=200)
plot(profile, mot, legend=false)
plot(mot)

using DynamicAxisWarping

dist = DTW(3)
normalizer = ZNormalizer
profile2 = matrix_profile(T, 50, dist, normalizer=normalizer)

using Distances
struct ZWrapper{D} <: Distances.Metric
    dist::D
end
function Distances.evaluate(d::ZWrapper, x, y)
    evaluate(d.dist, znorm(x), znorm(y))
end
(d::ZWrapper)(x,y) = evaluate(d,x,y)

mot2 = motifs(profile2, 4, r=3, th=200, dist=ZWrapper(dist))
plot(profile2, mot2, legend=false)
plot(mot2)

snips = snippets(T, 4, 100, dist)
plot(snips)

@btime snippets($(T[1:2000]), 4, 100)
@profiler snippets((T[1:2000]), 4, 100)


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    MetricBall(radii, [angles]; convention=TaitBryanExtr)
    MetricBall(radius, metric=Euclidean())

A metric ball is a neighborhood that can be expressed in terms
of a metric and a set of `radii`. The two main examples are the
Euclidean ball an the Mahalanobis (ellipsoid) ball.

When multiple `radii` are provided, they can be rotated by `angles`
according to a given rotation `convention`. Alternatively, a metric
from the [Distances.jl](https://github.com/JuliaStats/Distances.jl)
package can be specified together with a single `radius`.

## Examples

N-dimensional Euclidean ball with radius `1.0`:

```julia
julia> euclidean = MetricBall(1.0)
```

Axis-aligned 3D ellispoid with radii `[3.0, 2.0, 1.0]`:

```julia
julia> mahalanobis = MetricBall([3.0, 2.0, 1.0])
```

See also [`mahalanobis`](@ref).
"""
struct MetricBall{R,M} <: Neighborhood
  radii::R
  metric::M
end

function MetricBall(radii::AbstractVector{T}, angles=nothing;
                    convention=TaitBryanExtr) where {T<:Real}
  ndims  = length(radii)
  angles = isnothing(angles) ? (ndims == 2 ? [zero(T)] : zeros(T, 3)) : angles
  metric = mahalanobis(radii, angles, convention)
  MetricBall{typeof(radii),typeof(metric)}(radii, metric)
end

function MetricBall(radius::T, metric=Euclidean()) where {T<:Real}
  MetricBall{typeof(radius),typeof(metric)}(radius, metric)
end

"""
    radii(ball)

Return the radii of the metric `ball`.
"""
radii(ball::MetricBall) = ball.radii

"""
    metric(ball)

Return the metric of the metric `ball`.
"""
metric(ball::MetricBall) = ball.metric

"""
    range(ball)

Return the range of the metric `ball`, i.e.
the smallest value `r` such that `||v|| ≤ r`
for any `v ∈ ball`.
"""
Base.range(ball::MetricBall{<:Real}) = ball.radii
Base.range(::MetricBall{R,<:Mahalanobis}) where {R} = one(eltype(R))

function Base.show(io::IO, ball::MetricBall)
  r = length(ball.radii) > 1 ? Tuple(ball.radii) : ball.radii
  m = nameof(typeof(ball.metric))
  print(io, "MetricBall($r, $m)")
end


using Logging
global_logger(ConsoleLogger(stderr, Logging.Debug)) # modify the debugging level
# this file is best used interactively
using Distributed
addprocs(4)
using ScatteringTransform
@everywhere using Revise
@everywhere using ScatteringTransform
using Plots
using Statistics
# simple example; a sum of sine waves that has a discontinuous jump
t = 0.001:6π/100:6π
f = 2 .* max.(0,-(t.-3*π))./(t.-3*π) .* (sin.(2π*t)) .+ 10 .* max.(0,t.-3*π)./(t.-3*π) .+ max.(0,t.-3*π)./(t.-3*π).* sin.(4π*t.+π)
plot(t, f)

# First create the transform configuration
layers = stParallel(2, length(f))

# two major versions:
###########################################################################################################
#####    the first outputs a datatype specifically for working with Scattering transforms, and has both all internal results and the output. It is not particulary appropriate for use on large datasets
###########################################################################################################

@time output = st(f, layers, absType(), thin=false) # run twice to get the time. The first is compile time

###########################################################################################################
####    the second outputs a single array where the last dimension corresponds to the transformed data. Useful for large datasets, and more optimized
###########################################################################################################
@time thinOp = st(f, layers, absType(), outputSubsample=(-1,3), thin=true); # run twice to get the time. The first is compile time

# computing the morlet wavelet coefficients subsampled for comparison
waves = wavelet(WT.Morlet(5.0), averagingType = WT.Father())
cwtf = cwt(f, waves)
firstLayer = zeros(eltype(cwtf), 25, size(cwtf, 2))
for i=1:size(firstLayer,2)
  firstLayer[:,i] = ScatteringTransform.bspline(ScatteringTransform.bspline(abs.(cwt(f, waves)[:,i]), 2), 2)
end

plot(abs.(firstLayer)[:,end,1])
# let's take a look at some of the coefficients
# first, for comparison, the cwt using the same Morlet wavelets is
heatmap(abs.(firstLayer)[:,2:end,1]', xaxis="Space",yaxis="Frequency", title="Morlet coefficients of f")
# where the discontinuity is clearly visible, as is the difference in frequency
# For comparison, the first layer (the difference in magnitude is due to normalization)
heatmap(output[2][:,:,1]', xlabel="space", ylabel="Frequency",title="First Layer Output (smoothed version of the previous)")
# the discontinuity is less visible here, but it remains in the next layer

# let's take a closer look at the 6th frequency, which is where the sinusoid above is represented
p = pathType([6]) # an object that we can use as an index into the output
plot(output[p], title="6th frequency coefficients") # biased towards the left, where the oscillation is present
# now it's children
p = pathType([6,:])
heatmap(output[p][:,:,1]', xlabel="space", ylabel="Frequency", title="6th frequency variation")
# now the highest frequency, which should reflect the discontinuity
p = pathType([15,:])
heatmap(output[p][:,:,1]', xlabel="space", ylabel="Frequency", title="15th frequency decomposition")
p = pathType([14,:])
heatmap(output[p][:,:,1]', xlabel="space", ylabel="Frequency", title="15th frequency decomposition")

# An example from a sonar dataset
using JLD2
@load "examples/resp1999.0_10.2.jld2"
resp = permutedims(resp, (1,3,2))[:,:,1:45]
# this example is the reflection off of a sharkfin whose interior liquid has speed 1999.0 at distance 10.2 meters
# it is now ordered as (range)×(cross-range)×(angle)

layers = stParallel(2, size(resp,1))
size(resp)
# since this is going to be pretty large, we should subsample heavily (see the docs via ?st for options)
@time tmpOutput = st(resp[:,:,1], layers, absType(), outputSubsample=(-1,5))
output = zeros(size(tmpOutput)..., size(resp,3))
for i=1:size(resp,3)
    @time output[:,:,i] = st(resp[:,:,i], layers, absType(), outputSubsample=(-1,5))
end
output
# output is now prepared. At this point, the preprocessing is done, and you could do regression, clustering, etc.


using
  PyPlot,
  JLD2,
  Printf,
  FourierFlows

using FFTW: ifft

import GeophysicalFlows.MultilayerQG
import GeophysicalFlows.MultilayerQG: energies, fluxes

# Numerical parameters and time-stepping parameters
nx = 128          # 2D resolution = nx^2
ny = nx

stepper = "FilteredRK4"   # timestepper
dt  = 5e-3     # timestep
nsteps = 40000 # total number of time-steps
nsubs  = 250   # number of time-steps for plotting
               # (nsteps must be multiple of nsubs)

# Physical parameters
 Lx  = 2π      # domain size
   μ = 5e-2    # bottom drag
beta = 5       # the y-gradient of planetary PV

nlayers = 2       # these choice of parameters give the
f0, g = 1, 1      # desired PV-streamfunction relations
  H = [0.2, 0.8]  # q1 = Δψ1 + 25*(ψ2-ψ1), and
rho = [4.0, 5.0]  # q2 = Δψ2 + 25/4*(ψ1-ψ2).
  U = zeros(nlayers)
  U[1] = 1.0
  U[2] = 0.0

gr = TwoDGrid(nx, Lx)
x, y = gridpoints(gr)
k0, l0 = gr.k[2], gr.l[2] # fundamental wavenumbers

# Initialize problem
prob = MultilayerQG.Problem(nlayers=nlayers, nx=nx, Lx=Lx, f0=f0, g=g, H=H, rho=rho, U=U, dt=dt, stepper=stepper, μ=μ, beta=beta)
sol, cl, pr, vs, gr = prob.sol, prob.clock, prob.params, prob.vars, prob.grid

# Files
filepath = "."
plotpath = "./plots_2layer"
plotname = "snapshots"
filename = joinpath(filepath, "2layer.jld2")

# File management
if isfile(filename); rm(filename); end
if !isdir(plotpath); mkdir(plotpath); end

# Initialize with zeros
MultilayerQG.set_q!(prob, 1e-2randn((nx, ny, nlayers)))

# Create Diagnostics
E = Diagnostic(energies, prob; nsteps=nsteps)
diags = [E]

# Create Output
get_sol(prob) = sol # extracts the Fourier-transformed solution
function get_u(prob)
  @. v.qh = sol
  streamfunctionfrompv!(v.psih, v.qh, p.invS, g)
  @. v.uh = -im*g.l *v.psih
  invtransform!(v.u, v.uh, p)
  v.u
end
out = Output(prob, filename, (:sol, get_sol), (:u, get_u))

function plot_output(prob, fig, axs; drawcolorbar=false)

  # Plot the PV field and the evolution of energy and enstrophy.

  sol, v, p, g = prob.sol, prob.vars, prob.params, prob.grid
  MultilayerQG.updatevars!(prob)

  for j in 1:nlayers
    sca(axs[j])
    pcolormesh(x, y, v.q[:, :, j])
    axis("square")
    xlim(-Lx/2, Lx/2)
    ylim(-Lx/2, Lx/2)
    title(L"$q_"*string(j)*L"$")
    if drawcolorbar==true
      colorbar()
    end

    sca(axs[j+2])
    cla()
    contourf(x, y, v.psi[:, :, j])
    contour(x, y, v.psi[:, :, j], colors="k")
    axis("square")
    xlim(-Lx/2, Lx/2)
    ylim(-Lx/2, Lx/2)
    title(L"$\psi_"*string(j)*L"$")
    if drawcolorbar==true
      colorbar()
    end
  end

  sca(axs[5])
  plot(μ*E.t[E.i], E.data[E.i][1][1], ".", color="b", label=L"$KE_1$")
  plot(μ*E.t[E.i], E.data[E.i][1][2], ".", color="r", label=L"$KE_2$")
  xlabel(L"\mu t")
  ylabel(L"KE")
  # legend()

  sca(axs[6])
  plot(μ*E.t[E.i], E.data[E.i][2][1], ".", color="k", label=L"$PE_{3/2}$")
  xlabel(L"\mu t")
  ylabel(L"PE")
  # legend()

  pause(0.001)
end

fig, axs = subplots(ncols=3, nrows=2, figsize=(15, 8))
plot_output(prob, fig, axs; drawcolorbar=true)

# Step forward
startwalltime = time()

while cl.step < nsteps
  stepforward!(prob, diags, nsubs)

  # Message
  log = @sprintf("step: %04d, t: %d, KE1: %.4f, KE2: %.4f, PE: %.4f, τ: %.2f min", cl.step, cl.t, E.data[E.i][1][1], E.data[E.i][1][2], E.data[E.i][2][1], (time()-startwalltime)/60)

  println(log)

  plot_output(prob, fig, axs; drawcolorbar=false)

end
println((time()-startwalltime))

plot_output(prob, fig, axs; drawcolorbar=false)

savename = @sprintf("%s_%09d.png", joinpath(plotpath, plotname), cl.step)
savefig(savename, dpi=240)


using FTCTests
const FTC = FaultTolerantControl
using LinearAlgebra
using Plots
using UnPack
using Debugger
using ReferenceFrameRotations

"""
# Refs
[1] M. Tahavori and A. Hasan, “Fault recoverability for nonlinear systems with application to fault tolerant control of UAVs,” Aerosp. Sci. Technol., vol. 107, p. 106282, 2020, doi: 10.1016/j.ast.2020.106282.

# Notes
- x = state vector, [eta, omega, v, xi]^T
    - eta = Euler angle vector, [phi, theta, psi]^T
    - omega = angular rate vector, [p, q, r]^T
    - v = velocity vector, [u, v, w]^T
    - xi = position vector, [x, y, z]^T
- u = control input vector, thrust vector, [omega_1, omega_2, omega_3, omega_4]^T
- y = output vector, [eta, omega, xi]^T
"""
function compute_minHSV_example(lambda, num::Int; dt=0.01, tf=1.0)
    @assert lambda >= 0.0 && lambda <= 1.0
    Λ = diagm(ones(4))
    Λ[num, num] = lambda
    # @show Λ
    # Quadcopter state space model [1]
    gc = 9.81
    m = 0.65
    l = 0.023
    k = 31.1e-5
    b = 7.5e-7
    J = Diagonal([7.5e-3, 7.4e-3, 1.3e-2])
    Jinv = inv(J)
    A = [zeros(3, 3) Matrix(I, 3, 3) zeros(3, 3) zeros(3, 3);
    zeros(3, 3) zeros(3, 3) zeros(3, 3) zeros(3, 3);
    zeros(3, 3) zeros(3, 3) zeros(3, 3) zeros(3, 3);
    zeros(3, 3) zeros(3, 3) Matrix(I, 3, 3) zeros(3, 3)]
    B = zeros(12, 4)
    for i = 1:12
        for j = 1:4
            if i == 5
                if j == 1
                    B[i, j] = -l*k/J[2, 2]
                elseif j == 3
                    B[i, j] = l*k/J[2, 2]
                end
            elseif i == 4
                if j == 2
                    B[i, j] = -l*k/J[1, 1]
                elseif j == 4
                    B[i, j] = l*k/J[1, 1]
                end
            elseif i == 6
                if j == 1 || j == 3
                    B[i, j] = -b/J[3, 3]
                elseif j == 2 || j == 4
                    B[i, j] = b/J[3, 3]
                end
            elseif i == 9
                if j == 1 || j == 2 || j == 3 || j == 4
                    B[i, j] = k/m
                end
            else
                B[i, j] = 0
            end
        end
    end
    C = [Matrix(I, 3, 3) zeros(3, 3) zeros(3, 3) zeros(3, 3);
    zeros(3, 3) Matrix(I, 3, 3) zeros(3, 3) zeros(3, 3);
    zeros(3, 3) zeros(3, 3) zeros(3, 3) Matrix(I, 3, 3)]

    T(x) = [1 sin(x[1])*tan(x[2]) cos(x[1])*tan(x[2]);
            0 cos(x[1]) -sin(x[1]);
            0 sec(x[2])*sin(x[1]) cos(x[1])*sec(x[2])]
    R(x) = [cos(x[2])*cos(x[3]) sin(x[1])*sin(x[2])*cos(x[3])-cos(x[1])*sin(x[3]) cos(x[1])*sin(x[2])*cos(x[3])+sin(x[1])*sin(x[3]);
            cos(x[2])*sin(x[3]) sin(x[1])*sin(x[2])*sin(x[3])-cos(x[1])*cos(x[3]) cos(x[1])*sin(x[2])*sin(x[3])-sin(x[1])*cos(x[3]);
            -sin(x[2]) sin(x[1])*cos(x[2]) cos(x[1])*cos(x[2])]
    F(x) = vec([
        T(x) * x[4:6]
        Jinv * cross(-x[4:6], J * x[4:6])
        -cross(x[4:6], x[7:9]) + gc * R(x)' * [0; 0; 1]
        R(x) * x[7:9]
    ])

    f(x, u, p, t) = A*x + F(x) + B*(Λ*u).^2
    g(x, u, p, t) = C*x

	# System dimension
	n, m = size(B)
	l = size(C)[1]

    # Initial settings
    x0 = zeros(n)
    u0 = zeros(m)
    pr = zeros(4, 1)

    Wc = FTC.empirical_gramian(f, g, m, n, l; opt=:c, dt=dt, tf=tf, pr=pr, xs=x0, us=u0)
    Wo = FTC.empirical_gramian(f, g, m, n, l; opt=:o, dt=dt, tf=tf, pr=pr, xs=x0, us=u0)
    # minHSV = FTC.min_HSV(Wc, Wo)
    eigvals_Wc = Wc |> LinearAlgebra.eigvals |> minimum |> sqrt
end

# function Dynamics!(multicopter::IslamQuadcopter)
#     @unpack B = multicopter
#     function dynamics!(dx, x, param, t; u, Λ)
#         ν = B * Λ * u
#         f, M = ν[1], ν[2:4]
#         @nested_log FSimZoo.__Dynamics!(multicopter)(dx, x, (), t; f=f, M=M)
#     end
# end

# function test_model()
#     multicopter = IslamQuadcopter()
#     X = State(multicopter)()
#     Λ = diagm(ones(4))
#     u = (multicopter.m * multicopter.g / multicopter.kf) / 4 * ones(4)
#     dX = State(multicopter)()
#     Dynamics!(multicopter)(dX, X, (), 0.0; u=u, Λ=Λ)
#     dX
# end

# function compute_minHSV_Islam(lambda, num::Int; dt=0.01, tf=1.0)
#     @assert lambda >= 0.0 && lambda <= 1.0
#     Λ = diagm(ones(4))
#     Λ[num, num] = lambda
#     @show Λ
#     multicopter = IslamQuadcopter()
#     function f(x, u, param, t)
#         p = x[1:3]
#         v = x[4:6]
#         quat = x[7:10]
#         ω = x[11:13]
#         R = quat_to_dcm(Quaternion(quat...))
#         X = State(multicopter)(p, v, R, ω)
#         dX = State(multicopter)()  # initialisation
#         Dynamics!(multicopter)(dX, X, param, t; u=u, Λ=Λ)
#         dX_quat = dcm_to_quat(DCM(dX.R))
#         dx = [dX.p..., dX.v..., dX_quat.q0, dX_quat.q1, dX_quat.q2, dX_quat.q3, dX.ω...]
#         return dx
#     end
#     g(x, u, p, t) = x[7:10]  # different from [1]
#     n, m, l = 13, 4, 4
#     X0 = State(multicopter)()
#     quat = dcm_to_quat(DCM(X0.R))
#     _quat = [quat.q0, quat.q1, quat.q2, quat.q3]
#     x0 = [X0.p..., X0.v..., _quat..., X0.ω...]
#     # u0 = (multicopter.m * multicopter.g / multicopter.kf) / 4 * ones(4)
#     u0 = zeros(4)
#     pr = zeros(4, 1)

#     Wc = FTC.empirical_gramian(f, g, m, n, l; opt=:c, dt=dt, tf=tf, pr=pr, xs=x0, us=u0, xm=1.0, um=10000.0)
#     Wo = FTC.empirical_gramian(f, g, m, n, l; opt=:o, dt=dt, tf=tf, pr=pr, xs=x0, us=u0, xm=1.0, um=10000.0)
#     _eigvals_Wc = Wc |> LinearAlgebra.eigvals
#     eigvals_Wc = []
#     for eigval in _eigvals_Wc
#         if abs(imag(eigval)) < 1e-6
#             push!(eigvals_Wc, real(eigval))
#         else
#             error("Too large imaginary part")
#         end
#     end
#     min_HSV = eigvals_Wc |> minimum |> sqrt # not min_HSV
#     # minHSV = FTC.min_HSV(Wc, Wo)
# end

function plotting(rotor_idx)
    lambda = 0:0.10:1 |> collect
    HSVs = []
    for i = 1:length(lambda)
        HSVs = push!(HSVs, compute_minHSV_example(lambda[i], rotor_idx))
        # HSVs = push!(HSVs, compute_minHSV_Islam(lambda[i], rotor_idx))
    end
    return plot(lambda,
                HSVs,
                xlabel = "Actuator effectiveness",
                ylabel = "Minimum HSV",
                label = nothing,
               )
end

function main()
    figs = []
    for rotor_idx in 1:4
        fig = plotting(rotor_idx)
        figs = push!(figs, fig)
    end
    plot(figs..., layout=(2, 2))
end


# various functions for analyzing the performance characteristics of RP trees

function local_cov_dim(eigs, eps)
    d = 1
    while !has_local_cov_dim(eigs, d, eps)
        d += 1
    end
    return d
end

has_local_cov_dim(eigs, d, eps) = sum(eigs[end-(d-1):end].^2) >= (1 - eps)*sum(eigs.^2)


function limit_features(doc_terms::DocumentTermMatrix,
                        high::Int,
                        low::Int)
    doc_freqs = vec(sum(doc_terms.dtm, dims=2))

    # build mask to restrict terms
    mask = trues(length(doc_freqs))
    if high < 1
        mask .&= (doc_freqs .<= high)
    end
    if low > 0
        mask .&= (doc_freqs .>= low)
    end

    new_terms = doc_terms.terms[mask]

    return (doc_terms.dtm[mask, :], new_terms)
end

## Helper functions to build Corpus ##
_convert_bag_of_words(X::Dict{<:NGram, <:Integer}) =
    Dict(join(k, " ") => v for (k, v) in X)

build_corpus(X::Vector{<:Dict{<:NGram, <:Integer}}) =
    build_corpus(_convert_bag_of_words.(X))
build_corpus(X::Vector{<:Dict{S, <:Integer}}) where {S <: AbstractString} =
    Corpus(NGramDocument.(X))
build_corpus(X) = Corpus(TokenDocument.(X))

# based on https://github.com/zgornel/StringAnalysis.jl/blob/master/src/dtm.jl
# and https://github.com/JuliaText/TextAnalysis.jl/blob/master/src/dtm.jl
build_dtm(docs::Corpus) = build_dtm(docs, sort(collect(keys(lexicon(docs)))))
function build_dtm(docs::Corpus, terms::Vector{T}) where {T}
    # we are flipping the orientation of this matrix
    # so we get the `columnindices` from the TextAnalysis API
    row_indices = TextAnalysis.columnindices(terms)

    m = length(terms) # terms are rows
    n = length(docs)  # docs are columns

    rows = Vector{Int}(undef, 0) # terms
    columns = Vector{Int}(undef, 0) # docs
    values = Vector{Int}(undef, 0)
    for i in eachindex(docs.documents)
        doc = docs.documents[i]
        ngs = ngrams(doc)
        for ngram in keys(ngs)
            j = get(row_indices, ngram, 0)
            v = ngs[ngram]
            if j != 0
                push!(columns, i)
                push!(rows, j)
                push!(values, v)
            end
        end
    end
    if length(rows) > 0
        doc_term_mat = sparse(rows, columns, values, m, n)
    else
        doc_term_mat = spzeros(Int, m, n)
    end
    DocumentTermMatrix(doc_term_mat, terms, row_indices)
end

## General method to calculate IDF vector ##
function compute_idf(smooth_idf::Bool, doc_term_mat::SparseMatrixCSC{T}) where {T <: Real}
    n = size(doc_term_mat, 2) # docs are columns

    smooth_idf = Int(smooth_idf)
    documents_containing_term = vec(sum(doc_term_mat .> 0, dims=2)) .+ smooth_idf
    idf = log.((n + smooth_idf) ./ documents_containing_term) .+ 1

    return idf
end


# This file can be used to create the validation file for deterministic scc values.
# It will create a validation file containing all possibilities of parameter values
# defined in the specs dictionary below produce the same results.

using MimiDICE2010
using DataFrames
using Query
using CSVFiles
using Test

specs = Dict([
    :year => [2015, 2055],
    :eta => [0, 1.5],
    :prtp => [0.015, 0.03],
    :last_year => [2295, 2595],
    :pulse_size => [1e3, 1e7, 1e10]
])

results = DataFrame(year = [], eta = [], prtp = [], last_year = [], pulse_size = [], SC = [])

for year in specs[:year]
    for eta in specs[:eta]
        for prtp in specs[:prtp]
            for last_year in specs[:last_year]
                for pulse_size in specs[:pulse_size]
                    sc = MimiDICE2010.compute_scc(year=Int(year), eta=eta, prtp=prtp, last_year=Int(last_year), pulse_size=pulse_size)
                    push!(results, (year, eta, prtp, last_year, pulse_size, sc))
                end
            end
        end
    end
end

path = joinpath(@__DIR__, "deterministic_sc_values_v1-0-1.csv")
save(path, results)


using ConvDiffMIPDECO
using jInv.Mesh
using jInv.ForwardShare
using jInv.InverseSolve
using jInv.LinearSolvers
using MUMPSjInv
using MAT

filename= joinpath(dirname(pathof(ConvDiffMIPDECO)),"..","examples","Sources3D.mat")
file = matread(filename)
wtrue = vec(file["W"]) ## Sources on finest mesh
domain = file["domain"]
m      = vec(file["m"])

M      = getRegularMesh(domain,m)

rec = rand(100,2) .*[domain[2] domain[4] ]

Ainv = getMUMPSsolver()

# build linear interpolation matrix from nodes to receiver locations
x1c,x2c,x3c = getCellCenteredAxes(M)
rec3D = [kron(ones(length(x3c)),rec) kron(x3c,ones(size(rec,1)))]

x1,x2,x3 = getNodalAxes(M)
P = interpmat(x1,x2,x3, rec3D);

# build param that holds forward problem
pFor = getConvDiffFEMParam(M,Ainv=Ainv,P=P);
dtrue, pFor = getData(vec(wtrue),pFor);
utrue = pFor.Fields

file["utrue"]  = reshape(utrue,tuple(M.n.+1...));
file["dtrue"]  = dtrue
file["rec"]    = rec
file["sig"]    = pFor.sig
file["v"]      = pFor.v

matwrite(filename,file)


### Struct variable
struct OptimalPowerFlowDC
    main::Array{Float64,2}
    flow::Array{Float64,2}
    generation::Array{Float64,2}
end

### DC optimal power flow, created by Ognjen Kundacina
function rundcopf(system, num, settings, info)
    printstyled("Algorithm: DC optimal power flow\n"; bold = true)

    ########## Pre-processing ##########
    busi, type, Pload, Gshunt, Tini, Pmax, Pmin, reactance, transTap, transShift, branchOn, longTermRating, angleMin, angleMax,
    costModel, numOfCoeffs, coeffs = read_opfdcsystem(system, num)

    main = fill(0.0, num.Nbus, 6)
    flow = fill(0.0, num.Nbranch, 5)
    gene = fill(0.0, num.Ngen, 2)
    Va, Pinj, Pbus, Pij, Pji, Pgen = write_dcsystem(system, num, main, flow, gene)

  algtime = @elapsed begin
    ########## Convert in integers ##########
    geni = convert(Array{Int64,1}, system.generator[:, 1])
    genOn = convert(Array{Int64,1}, system.generator[:, 8])
    fromi = convert(Array{Int64,1}, system.branch[:, 2])
    toi = convert(Array{Int64,1}, system.branch[:, 3])

    ########## Numbering and slack bus ##########
    numbering = false
    bus = collect(1:num.Nbus)
    slack = 0
    for i = 1:num.Nbus
        if bus[i] != busi[i]
            numbering = true
        end
        if type[i] == 3
            slack = bus[i]
        end
    end
    if numbering
        println("The new bus numbering is running.")
    end
    if slack == 0
        slack = 1
        println("The slack bus is not found. Slack bus is the first bus.")
    end
    gen = renumber(geni, busi, bus, numbering)
    from = renumber(fromi, busi, bus, numbering)
    to = renumber(toi, busi, bus, numbering)

    ########## Optimization model ##########
    model = Model(Ipopt.Optimizer)
    @variable(model, thetas[bus])
    @constraint(model, thetas[slack] == 0.0)
    @variable(model, Pmin[i] / system.basePower <= pGen[i = 1:num.Ngen] <= Pmax[i] / system.basePower)
    @inbounds for i = 1:num.Ngen
        if genOn[i] == 0
            @constraint(model, pGen[i] == 0.0)
        else
            set_start_value(pGen[i], (Pmax[i] + Pmin[i]) / (2 * system.basePower))
        end
    end

    ########## Ybus matrix ##########
    Ybus, admitance, Pshift = ybusdc(system, num, bus, from, to, branchOn, transTap, reactance, transShift)

    ########## Branch flow constraints ##########
    @inbounds for i = 1:num.Nbranch
        if longTermRating[i] ≉  0 && longTermRating[i] < 10^10
            upf = longTermRating[i] / system.basePower + admitance[i] * transShift[i] * pi / 180
            upt = longTermRating[i] / system.basePower - admitance[i] * transShift[i] * pi / 180
            @constraint(model, -upt <= admitance[i] * thetas[from[i]] - admitance[i] * thetas[to[i]] <= upf)
        end
    end

    ########## Branch angle constraints ##########
    @inbounds for i = 1:num.Nbranch
        if angleMin[i] > -360 && angleMax[i] < 360
            angMinRad = angleMin[i] * (pi / 180)
            angMaxRad = angleMax[i] * (pi / 180)
            @constraint(model, angMinRad <= thetas[from[i]] - thetas[to[i]] <= angMaxRad)
        end
    end

    ########## Power mismatch constraints ##########
    Amis = [Ybus sparse(gen, collect(1:num.Ngen), -1, num.Nbus, num.Ngen)]
    bmis = - (Pload + Gshunt) / system.basePower - Pshift;
    @constraint(model, Amis * [thetas; pGen] .== bmis)

    Qpg, cpg, kpg, nPWL, pwlIds = process_cost_table(system, num, costModel, coeffs, numOfCoeffs)
    if nPWL == 0
        @objective(model, Min, sum(0.5 * Qpg[i] * pGen[i]^2 + cpg[i] * pGen[i] + kpg[i] for i = 1:num.Ngen))
    elseif nPWL > 0
        @variable(model, pwlVars[i = 1:nPWL])
        pgbas = 1
        qgbas = zeros(Int64, 0)
        nq = 0
        ybas = 1 + num.Ngen + nq
        Apwl, bpwl = makeApwl(system, num, nPWL, pwlIds, pgbas, qgbas, ybas, numOfCoeffs)
        for i = 1:(length(bpwl))
            @constraint(model, dot(Apwl[i,:], [pGen; pwlVars]) <= bpwl[i])
        end
        @objective(model, Min, sum(0.5 * Qpg[i] * pGen[i]^2 + cpg[i] * pGen[i] + kpg[i] for i = 1:num.Ngen) + sum(pwlVars[i] for i = 1:nPWL))
    end

    optimize!(model)
    Ti = value.(thetas) .+ (pi / 180) * Tini[slack]

    ########## Post-processing ##########
    for i = 1:num.Nbranch
        if branchOn[i] == 1
            Pij[i] = admitance[i] * (Ti[from[i]] - Ti[to[i]] - (pi / 180) * transShift[i]) * system.basePower
            Pji[i] = -Pij[i]
        end
    end

    Pgen[:] = value.(pGen) .* system.basePower
    for (k, i) in enumerate(gen)
        if genOn[k] == 1
            Pbus[i] += Pgen[k]
        end
    end

    for i = 1:num.Nbus
        Va[i] = (180 / pi) * Ti[i]
        Pinj[i] = Pbus[i] - Pload[i]
        if i == slack
            I = 0.0
            for j in Ybus.colptr[i]:(Ybus.colptr[i + 1] - 1)
                row = Ybus.rowval[j]
                I += Ybus[slack, row] * Ti[row]
            end
            Pinj[slack] = (I + Pshift[slack]) * system.basePower + Gshunt[slack]
        end
    end

 end # algtime

    ########## Results ##########
    results = OptimalPowerFlowDC(main, flow, gene)
    header, group = results_flowdc(system, num, settings, results, slack, algtime)
    if !isempty(settings.save)
        savedata(results, system; info = info, group = group, header = header, path = settings.save)
    end

    return results
end

### Process data needed for the objective function
function process_cost_table(system, num, costModel, coeffs, numOfCoeffs)
    linTerm = fill(0.0, num.Ngen); quadTerm = fill(0.0, num.Ngen); constTerm = fill(0.0, num.Ngen)
    nPWL = 0; pwlIds = zeros(Int64, 0)

    @inbounds for i = 1:num.Ngen
        if costModel[i] == 1
            ########## PWL cost model ##########
            if numOfCoeffs[i] == 2
                ########## Single block pwl cost, convert to polynomial ##########
                x0 = system.generatorcost[i, 5]
                y0 = system.generatorcost[i, 6]
                x1 = system.generatorcost[i, 7]
                y1 = system.generatorcost[i, 8]
                m = (y1 - y0) / (x1 - x0)
                b = y0 - m * x0
                linTerm[i] = m
                constTerm[i] = b
            else
                nPWL = nPWL + 1
                append!(pwlIds, i)
            end
        elseif costModel[i] == 2
            ########## Polynomial ##########
            if numOfCoeffs[i] == 2
                linTerm[i] = coeffs[i, 1]
                constTerm[i] = coeffs[i, 2]
            elseif numOfCoeffs[i] == 3
                ########## Quadratic generator costs ##########
                quadTerm[i] = coeffs[i, 1]
                linTerm[i] = coeffs[i, 2]
                constTerm[i] = coeffs[i, 3]
            end
        end
    end
    Qpg = 2 * quadTerm * system.basePower^2
    cpg = linTerm * system.basePower
    kpg = constTerm

    return Qpg, cpg, kpg, nPWL, pwlIds
end

### Construct parameters for linear basin constraints
function makeApwl(system, num, nPWL, pwlIds, pgbas, qgbas, ybas, numOfCoeffs)
    if nPWL == 0
        Apwl = sparse([], [], [], 0, ybas + nPWL - 1, 0)
        bpwl = zeros(Int64, 0)

        return Apwl, bpwl
    end

    numCostPoints = sum(numOfCoeffs[i] for i in pwlIds)
    Apwl = zeros(Float64, numCostPoints - nPWL, ybas + nPWL - 1)
    bpwl = zeros(Int64, 0)

    k = 1
    @inbounds for i in pwlIds
        ns = numOfCoeffs[i]
        p = @view(system.generatorcost[i, 5:2:(5 + 2 * ns - 1)]) / system.basePower
        c = @view(system.generatorcost[i, 6:2:(5 + 2 * ns)])
        m = diff(c) ./ diff(p)
        if any(diff(p) == 0)
            println("makeApwl: Bad x axis data in the generator cost matrix, row: ", i)
        end
        b = m .* p[1:(ns - 1)] - c[1:(ns - 1)]
        bpwl = [bpwl; b]
        ########## Q and P costs ##########
        if i > num.Ngen
            sidx = qgbas + (i-ng) - 1
        else
            sidx = pgbas + i - 1
        end
        @inbounds for j = 1:(ns - 1)
            Apwl[j + k - 1, sidx] = m[j]
        end
        k = k + ns - 1
    end

    k = 1; j = 1
    @inbounds for i in pwlIds
        ns = numOfCoeffs[i]
        @inbounds for it = k:(k + ns - 2)
            Apwl[it, ybas + j - 1] = -1.0
        end
        k = k + ns - 1
        j = j + 1
    end

    return Apwl, bpwl
end

### Read data
function read_opfdcsystem(system, num)
    busi = @view(system.bus[:, 1])
    type = @view(system.bus[:, 2])
    Pload = @view(system.bus[:, 3])
    Gshunt = @view(system.bus[:, 5])
    Tini = @view(system.bus[:, 9])

    Pmax = @view(system.generator[:, 9])
    Pmin = @view(system.generator[:, 10])

    reactance = @view(system.branch[:, 5])
    transTap = @view(system.branch[:, 10])
    transShift = @view(system.branch[:, 11])
    branchOn = @view(system.branch[:, 12])
    longTermRating = @view(system.branch[:, 7])
    angleMin = @view(system.branch[:, 13])
    angleMax = @view(system.branch[:, 14])

    costModel = @view(system.generatorcost[:, 1])
    numOfCoeffs = convert(Array{Int64,1}, system.generatorcost[:, 4])
    coeffs = @view(system.generatorcost[:, 5:end])

    return busi, type, Pload, Gshunt, Tini,
            Pmax, Pmin,
            reactance, transTap, transShift, branchOn, longTermRating, angleMin, angleMax,
            costModel, numOfCoeffs, coeffs
end


@doc raw"
    einsum(::EinCode{ixs, iy}, xs, size_dict) where {ixs, iy}

return the tensor that results from contracting the tensors `xs` according
to their indices `ixs`, where all indices that do not appear in the output `iy` are
summed over.
The result is permuted according to `out`.

- `ixs` - tuple of tuples of index-labels of the input-tensors `xs`

- `iy` - tuple of index-labels of the output-tensor

- `xs` - tuple of tensors

- `size_dict` - `IndexSize`-object that maps index-labels to their sizes

# example

```jldoctest; setup = :(using OMEinsum)
julia> a, b = rand(2,2), rand(2,2);

julia> einsum(EinCode((('i','j'),('j','k')),('i','k')), (a, b)) ≈ a * b
true

julia> einsum(EinCode((('i','j'),('j','k')),('k','i')), (a, b)) ≈ permutedims(a * b, (2,1))
true
```
"
@generated function einsum(code::EinCode{ixs, iy}, xs, size_dict) where {ixs, iy}
    rule = match_rule(ixs, iy)
    :(einsum($rule, code, xs, size_dict))
end

function einsum(::Tr, ::EinCode, xs, size_dict)
    @debug "Tr" size.(xs)
    asarray(tr(xs[1]), xs[1])
end

using TensorOperations

function einsum(::PTrace, ::EinCode{ixs,iy}, xs::NTuple{<:Any, AbstractArray{<:BlasFloat}}, size_dict) where {ixs, iy}
    @debug "PTrace tensortrace" ixs => iy size.(xs)
    asarray(tensortrace(xs[1], ixs[1], iy), xs[1])
end

# note that dispatching to Hadamard if some `ixs` are permuted has inferior
# performance compared to the fallback
function einsum(::Hadamard, ::EinCode{ixs, iy}, xs, size_dict) where {ixs, iy}
    @debug "Hadamard" ixs => iy size.(xs)
    asarray(broadcast(*, xs...), xs[1])
end

function einsum(::PairWise, ::EinCode{ixs, iy}, xs::NTuple{NT,AbstractArray}, size_dict) where {ixs, iy, NT}
    @debug "PairWise optcontract" ixs => iy size.(xs)
    optcontract(ixs, xs, iy)
end

function einsum(sm::Sum, code::EinCode{ixs, iy}, xs, size_dict) where {ixs, iy}
    dims = (findall(i -> i ∉ iy, ixs[1])...,)
    (ix1,) = ixs
    ix1f = filter!(i -> i in iy, collect(ix1))
    perm = map(i -> findfirst(==(i), ix1f), iy)
    res = dropdims(sum(xs[1], dims=dims), dims=dims)
    if perm == iy
        @debug "Sum" ixs => iy size.(xs)
        res
    else
        @debug "Sum permutedims" ixs => iy size.(xs) perm
        tensorpermute(res, perm)
    end
end

@generated function einsum(sm::MatMul, code::EinCode{ixs, iy}, xs, size_dict) where {ixs, iy}
    ix1, ix2 = ixs
    l = ifelse(ix1[1] in ix2, ix1[1], ix1[2])
    if ix1[2] == l && ix2[1] == l
        if iy == (ix1[1], ix2[2])
            return :(@debug "MatMul1 (ij,jk -> ik)" ixs => iy size.(xs);
                xs[1] * xs[2]
            )
        else
            return :(@debug "MatMul2 (ij,jk -> ki) permutedims(ik)" ixs => iy size.(xs);
                permutedims(xs[1] * xs[2])
            )
        end
    elseif ix1[1] == l && ix2[1] == l
        if iy == (ix1[2], ix2[2])
            return :(@debug "MatMul3 (ji,jk -> ik) transpose(ji)" ixs => iy size.(xs);
                transpose(xs[1]) * xs[2]
            )
        else
            return :(@debug "MatMul4 (ji,jk -> ki) transpose(jk)" ixs => iy size.(xs);
                transpose(xs[2]) * xs[1]
            )
        end
    elseif ix1[2] == l && ix2[2] == l
        if iy == (ix1[1], ix2[1])
            return :(@debug "MatMul5 (ij,kj -> ik) transpose(kj)" ixs => iy size.(xs);
                xs[1] * transpose(xs[2])
            )
        else
            return :(@debug "MatMul6 (ij,kj -> ki) transpose(ij)" ixs => iy size.(xs);
                xs[2] * transpose(xs[1])
            )
        end
    else #ix1[1] == l && ix2[2] == l
        if iy == (ix1[2], ix2[1])
            return :(@debug "MatMul7 (ji,kj -> ik) permutedims(ki)" ixs => iy size.(xs);
                permutedims(xs[2] * xs[1])
            )
        else
            return :(@debug "MatMul8 (ji,kj -> ki)" ixs => iy size.(xs);
                xs[2] * xs[1]
            )
        end
    end
end

function einsum(::Permutedims, code::EinCode{ixs, iy}, xs, size_dict) where {ixs, iy}
    (ix,) = ixs
    (x,) = xs
    perm = map(i -> findfirst(==(i), ix), iy)
    @debug "Permutedims" ix => iy size(xs[1]) perm
    return tensorpermute(x, perm)
end

function einsum(::Identity, ::EinCode{ixs, iy}, xs, size_dict) where {ixs, iy}
    xs[1]
end

# the fallback
function einsum(::DefaultRule, code::EinCode{ixs, iy}, xs, size_dict) where {ixs, iy}
    @debug "DefaultRule loop_einsum" ixs => iy size.(xs)
    loop_einsum(code, xs, size_dict)
end

function einsum(::PTrace, code::EinCode{ixs, iy}, xs::NTuple{NT, Any}, size_dict) where {ixs, iy, NT}
    @debug "PTrace loop_einsum" ixs => iy size.(xs)
    loop_einsum(code, xs, size_dict)
end

function einsum(::BatchedContract, code::EinCode{ixs, iy}, xs::NTuple{NT, Any}, size_dict) where {ixs, iy, NT}
    @debug "BatchedContract loop_einsum" ixs => iy size.(xs)
    loop_einsum(code, xs, size_dict)
end

@generated function _preprocess_dupindices(::Val{ix}, x) where {ix}
    if length(tunique(ix)) != length(ix)
        iy = [l for l in ix if count(==(l), ix) == 1]
        :(($(Val((iy...,))), einsum($(EinCode((ix,), (iy...,))), (x,), get_size_dict(($ix,), (x,)))))
    else
        :(($(Val(ix)), x))
    end
end

@generated function einsum(::BatchedContract, ::EinCode{ixs,iy}, xs::NTuple{<:Any, AbstractArray{<:BlasFloat}}, size_dict) where {ixs, iy}
    quote
        ixs1, xs1 = _preprocess_dupindices($(Val(ixs[1])), xs[1])
        ixs2, xs2 = _preprocess_dupindices($(Val(ixs[2])), xs[2])
        @debug "BatchedContract" ixs => iy Tuple(ixs1) Tuple(ixs2) size.((xs1, xs2))
        batched_contract(ixs1, xs1, ixs2, xs2, $(Val(iy)))
    end
end


module MLCourse

import Pkg

project_relative_path(xs...) = normpath(joinpath(dirname(dirname(pathof(MLCourse))), xs...))

include("notebooks.jl")
include_dependency(joinpath("..", "Project.toml"))
const _VERSION = VersionNumber(Pkg.TOML.parsefile(project_relative_path("Project.toml"))["version"])

using Zygote, Plots, MLJ, MLJLinearModels, MLJGLMInterface, Markdown, DataFrames, Base64
export plot_residuals!, fitted_linear_func, grid, biplot

function plot_residuals!(x, y, f; kwargs...)
    for (xi, yi) in zip(x, y)
        plot!(fill(xi, 2), [yi, f(xi)]; color = :red, linewidth = 2,
              label = nothing, kwargs...)

    end
end
function plot_residuals!(x1, x2, y, f; kwargs...)
    for (x1i, x2i, yi) in zip(x1, x2, y)
        plot!(fill(x1i, 2), fill(x2i, 2), [yi, f(x1i, x2i)];
              color = :red, linewidth = 2, label = nothing, kwargs...)
    end
end

_slope(::MLJLinearModels.LinearRegressor, θ̂) = θ̂.coefs[1][2]
_slope(::MLJGLMInterface.LinearRegressor, θ̂) = θ̂.coef[1]

function fitted_linear_func(mach)
    θ̂ = fitted_params(mach)
    θ̂₀ = θ̂.intercept
    θ̂₁ = _slope(mach.model, θ̂)
    x -> θ̂₀ + θ̂₁ * x
end;

function gradient_descent(f, x₀, η, T; callback = x -> nothing)
    x = copy(x₀) # use copy to not overwrite the input
    for t in 1:T
        x .-= η * gradient(f, x)[1] # update parameters in direction of -∇f
        callback(x) # the callback will be used to save intermediate values
    end
    x
end

_grid(x, y) = (repeat(x, length(y)), repeat(y, inner = length(x)))
function grid(x, y; names = (:X, :Y), output_format = NamedTuple)
    t = NamedTuple{names}(_grid(x, y))
    output_format(t)
end

polysymbol(x, d) = Symbol(d == 0 ? "" : d == 1 ? "$x" : "$x^$d")
function poly(data, degree)
    res = DataFrame([data.x .^ k for k in 1:degree],
                    [polysymbol("x", k) for k in 1:degree])
    if hasproperty(data, :y)
        res.y = data.y
    end
    res
end
function poly2(data, degree)
    res = DataFrame([data.X1 .^ d1 .* data.X2 .^ d2
                     for d1 in 0:degree, d2 in 0:degree if 0 < d1 + d2 ≤ degree],
                    [Symbol(polysymbol("x₁", d1), polysymbol("x₂", d2))
                     for d1 in 0:degree, d2 in 0:degree if 0 < d1 + d2 ≤ degree])
    if hasproperty(data, :y)
        res.y = data.y
    end
    res
end

Base.@kwdef mutable struct PolynomialRegressor{T <: Deterministic} <: Deterministic
    degree::Int = 3
    regressor::T = LinearRegressor()
end
function MLJ.MLJBase.fit(model::PolynomialRegressor, verbosity, X, y)
    Xpoly = poly(X, model.degree)
    MLJ.MLJBase.fit(model.regressor, verbosity, Xpoly, y)
end
function MLJ.MLJBase.predict(model::PolynomialRegressor, fitresult, Xnew)
    Xpoly = poly(Xnew, model.degree)
    MLJ.MLJBase.predict(model.regressor, fitresult, Xpoly)
end

function biplot(m; pc = 1:2)
    scores = MLJ.transform(m, m.data[1])
    p = scatter(getproperty(scores, Symbol(:x, pc[1])),
                getproperty(scores, Symbol(:x, pc[2])),
                label = nothing, aspect_ratio = 1, size = (600, 600),
                xlabel = "PC$(pc[1])", ylabel = "PC$(pc[2])",
                framestyle = :axis, markeralpha = 0,
                txt = text.(1:nrows(scores), 8, :gray), markersize = 0)
    plot!(p[1], inset = (1, bbox(0, 0, 1, 1)),
          right_margin = 10Plots.mm, top_margin = 10Plots.mm)
    p2 = p[1].plt.subplots[end]
    plot!(p2, aspect_ratio = 1, mirror = true, legend = false)
    params = fitted_params(m)
    loadings = if hasproperty(params, :pca)
        params.pca.projection
    else
        params.projection
    end
    n = names(m.data[1])
    for i in 1:length(n)
        plot!(p2, [0, .9*loadings[i, pc[1]]], [0, .9*loadings[i, pc[2]]],
              c = :red, arrow = true)
        annotate!(p2, [(loadings[i, pc[1]], loadings[i, pc[2]],
                        (n[i], :red, :center, 8))])
    end
    scatter!(p2, 1.1*loadings[:, pc[1]], 1.1*loadings[:, pc[2]],
             markersize = 0, markeralpha = 0) # dummy for sensible xlims and xlims
    plot!(p2,
          background_color_inside = RGBA{Float64}(0, 0, 0, 0),
          tickfontcolor = RGB(1, 0, 0))
    p
end

function embed_figure(name)
    "![](data:img/png; base64,
         $(open(base64encode,
                project_relative_path("notebooks", "figures", name))))"
end

function start()
    sysimg = project_relative_path("precompile", "mlcourse.so")
    root = project_relative_path()
    exe = joinpath(Sys.BINDIR, "julia")
    script = :(using Pkg;
               Pkg.activate($root);
               using Pluto;
               Pluto.run(notebook = $(joinpath(root, "index.jl"))))
    if isfile(sysimg)
        run(`$exe -J$sysimg -e $script`)
    else
        run(`$exe -e $script`)
    end
end

function update()
    @info "Performing an automatic update while keeping local changes.
    If this fails, please run manually `git pull` in the directory
    `$(project_relative_path())`."
    current_dir = pwd()
    cd(project_relative_path())
    if !isempty(readlines(`git diff --stat`))
        run(`git add -u`)
        run(`git commit -m "automatic commit of local changes"`)
    end
    run(`git pull -s recursive -X patience -X ours -X ignore-all-space --no-edit`)
    cd(current_dir)
    Pkg.activate(project_relative_path())
    Pkg.instantiate()
end

function create_sysimage()
    exe = joinpath(Sys.BINDIR, "julia")
    run(`$exe $(project_relative_path("precompile", "precompile.jl"))`)
end

if isfile(project_relative_path("precompile", "mlcourse.so"))
    @warn "You may have to create a new system image with this update."
end
@info """\n
    Welcome to the Machine Learning Course v$(_VERSION)!

    Get started with:

    julia> MLCourse.start()

    Or create a new system image for faster loading of the notebooks.
    This might take several minutes (~30 minutes on my laptop).
    During this process you may be asked to download data sets (answer: yes :))
    and a browser window will open, which you can close again. Do not interrupt
    the process. If all goes well it shows at the end for several minutes:
    `[ Info: PackageCompiler: creating system image object file, this might take a while...`

    julia> MLCourse.create_sysimage() # Be patient!
\n"""

end # module


@testset "core" begin

  # Accessor functions
  img = fill(1.0, (4,4))
  img[1:2,1:2] .= 2.0
  img[1:2,3:4] .= 3.0
  seg = fast_scanning(img, 0.5)

  @test labels_map(seg) == seg.image_indexmap
  @test segment_labels(seg) == seg.segment_labels
  @test segment_mean(seg) == seg.segment_means
  @test all([segment_mean(seg, i) == seg.segment_means[i] for i in keys(seg.segment_means)])
  @test segment_pixel_count(seg) == seg.segment_pixel_count
  @test all([segment_pixel_count(seg, i) == seg.segment_pixel_count[i] for i in keys(seg.segment_pixel_count)])

  img = [1 1.1 0.9 10 10.2 9.8; 1 1 1.1 4 3.8 4]
  r = fuzzy_cmeans(img, 2, 2.0)

  @test segment_labels(r) == [1,2]
  @test segment_mean(r, 1) == r.centers[:,1]
  @test segment_mean(r) == Dict(1=>r.centers[:,1], 2=>r.centers[:,2])
  @test all([segment_pixel_count(r,i) == 6 for i in 1:2])
  @test segment_pixel_count(r) == Dict(1=>6, 2=>6)

  # RAG
  img = fill(1.0, (10,10))
  img[:, 5:10] .= 2.0
  seg = fast_scanning(img, 0.5)
  g, vm = region_adjacency_graph(seg, (i,j)->sum(abs2, seg.segment_means[i]-seg.segment_means[j]))

  expectedg = SimpleWeightedGraph(2)
  add_edge!(expectedg, 1, 2, 1.0)
  expectedvm = Dict(1=>1, 2=>2)

  @test g == expectedg
  @test vm == expectedvm

  img = fill(1.0, (10,10))
  img[3:8,3:8] .= 2.0
  img[5:7,5:7] .= 3.0
  seg = fast_scanning(img, 0.5)
  g, vm = region_adjacency_graph(seg, (i,j)->1)

  expectedg = SimpleWeightedGraph(3)
  add_edge!(expectedg,1,2)
  add_edge!(expectedg,2,3)
  expectedvm = Dict(1=>1, 2=>2, 3=>3)

  @test g == expectedg
  @test vm == expectedvm

  # remove_segment
  img = fill(1.0, (10,10))
  img[1:4,:] .= 2.0
  img[:,5:10] .= 4.0
  seg = fast_scanning(img, 0.5)
  new_seg = remove_segment(seg, 1, (i,j)->(-seg.segment_pixel_count[j]))

  expected = fill(3, (10,10))
  expected[5:10,1:4] .= 2
  expected_labels = [2,3]
  expected_means = Dict(2=>1.0, 3=>68/19)
  expected_count = Dict(2=>24, 3=>76)

  @test all(label->(label in expected_labels), new_seg.segment_labels)
  @test all(label->(label in new_seg.segment_labels), expected_labels)
  @test expected_count == new_seg.segment_pixel_count
  @test all(label->(expected_means[label] ≈ new_seg.segment_means[label]), new_seg.segment_labels)
  @test new_seg.image_indexmap == expected

  new_seg = remove_segment(seg, 1, (i,j)->sum(abs2, seg.segment_means[i]-seg.segment_means[j]))

  expected = fill(3, (10,10))
  expected[:,1:4] .= 2
  expected_means = Dict(2=>1.4, 3=>4.0)
  expected_count = Dict(2=>40, 3=>60)

  @test all(label->(label in expected_labels), new_seg.segment_labels)
  @test all(label->(label in new_seg.segment_labels), expected_labels)
  @test expected_count == new_seg.segment_pixel_count
  @test all(label->(expected_means[label] ≈ new_seg.segment_means[label]), new_seg.segment_labels)
  @test new_seg.image_indexmap == expected

  remove_segment!(seg, 1, (i,j)->sum(abs2, seg.segment_means[i]-seg.segment_means[j]))

  @test all(label->(label in expected_labels), seg.segment_labels)
  @test all(label->(label in seg.segment_labels), expected_labels)
  @test expected_count == seg.segment_pixel_count
  @test all(label->(expected_means[label] ≈ seg.segment_means[label]), seg.segment_labels)
  @test seg.image_indexmap == expected

  # prune_segments
  img = fill(1.0, (10,10))
  img[3,3] = 2.0
  img[5,5] = 2.0
  img[8:9,7:8] .= 3.0
  seg = fast_scanning(img, 0.5)
  new_seg = prune_segments(seg, l->(seg.segment_pixel_count[l] < 2), (i,j)->sum(abs2, seg.segment_means[i]-seg.segment_means[j]))

  expected = fill(1, (10,10))
  expected[8:9,7:8] .= 4
  expected_labels = [1,4]
  expected_means = Dict(4=>3.0, 1=>49/48)
  expected_count = Dict(4=>4, 1=>96)

  @test all(label->(label in expected_labels), new_seg.segment_labels)
  @test all(label->(label in new_seg.segment_labels), expected_labels)
  @test expected_count == new_seg.segment_pixel_count
  @test all(label->(expected_means[label] ≈ new_seg.segment_means[label]), new_seg.segment_labels)
  @test new_seg.image_indexmap == expected

  new_seg = prune_segments(seg, [2,3], (i,j)->sum(abs2, seg.segment_means[i]-seg.segment_means[j]))

  @test all(label->(label in expected_labels), new_seg.segment_labels)
  @test all(label->(label in new_seg.segment_labels), expected_labels)
  @test expected_count == new_seg.segment_pixel_count
  @test all(label->(expected_means[label] ≈ new_seg.segment_means[label]), new_seg.segment_labels)
  @test new_seg.image_indexmap == expected

  # region_adjacency_graph
  img = fill(1.0, (10,10))
  img[4:6, 4:6] .= 0

  weight_fn(I, J) = exp(-abs(I.second - J.second)) / ((I.first[1]-J.first[1])^2 + (I.first[2]-J.first[2])^2)
  G, vertex2cartesian = region_adjacency_graph(img, weight_fn, 2)

  @test nv(G) == 100
  @test G.weights[45,46] == 1
  @test G.weights[43,44] == exp(-1)

  #test that all edges weights equal weights computed using weight_fn
  @test all(map(edge-> edge.weight == weight_fn(vertex2cartesian[edge.src]=>img[vertex2cartesian[edge.src]], vertex2cartesian[edge.dst]=>img[vertex2cartesian[edge.dst]]), edges(G)))

  #test that all edges are between nodes less than r pixels away
  @test all(map(edge-> max(abs(vertex2cartesian[edge.src][1] - vertex2cartesian[edge.dst][1]), abs(vertex2cartesian[edge.src][2] - vertex2cartesian[edge.dst][2])) <= 2 , edges(G)))

  @test region_adjacency_graph(img, weight_fn, 2) == region_adjacency_graph(img, weight_fn, CartesianIndex(2,2))

  G, vertex2cartesian = region_adjacency_graph(img, weight_fn, CartesianIndex(1,2))
  @test G.weights[45,47] == 0
  @test G.weights[45,65] != 0

end


using Test
using LogicCircuits.Logical
using LogicCircuits.Utils

@testset "Trimmed SDD Test" begin

    num_vars = 7
    mgr = balanced_vtree(TrimMgrNode, num_vars)

    @test num_variables(mgr) == num_vars
    @test num_nodes(mgr) == 2*num_vars-1
    @test num_edges(mgr) == 2*num_vars-2
    @test mgr isa TrimSddMgr

    @test descends_from(mgr[1], mgr[end])
    @test descends_from(mgr[end].left, mgr[end])
    @test descends_from(mgr[end].right, mgr[end])
    @test descends_left_from(mgr[end].left, mgr[end])
    @test descends_left_from(mgr[end].left.left, mgr[end])
    @test descends_left_from(mgr[end].left.right, mgr[end])
    @test descends_right_from(mgr[end].right, mgr[end])
    @test descends_right_from(mgr[end].right.right, mgr[end])
    @test descends_right_from(mgr[end].right.left, mgr[end])

    @test !descends_from(mgr[end], mgr[1])
    @test !descends_left_from(mgr[end].right, mgr[end])
    @test !descends_left_from(mgr[end].right.left, mgr[end])
    @test !descends_left_from(mgr[end].right.right, mgr[end])
    @test !descends_left_from(mgr[end], mgr[end])
    @test !descends_left_from(mgr[end], mgr[end].left)
    @test !descends_left_from(mgr[end], mgr[end].right)
    @test !descends_right_from(mgr[end].left, mgr[end])
    @test !descends_right_from(mgr[end].left.right, mgr[end])
    @test !descends_right_from(mgr[end].left.left, mgr[end])
    @test !descends_right_from(mgr[end], mgr[end])
    @test !descends_right_from(mgr[end], mgr[end].left)
    @test !descends_right_from(mgr[end], mgr[end].right)

    x = Var(1)
    y = Var(2)

    x_c = compile(mgr, x)
    y_c = compile(mgr, y)

    @test x_c != y_c

    @test variable(x_c) == x
    @test literal(x_c) == var2lit(x)
    @test vtree(x_c) ∈ mgr
    @test positive(x_c)
    @test x_c == compile(mgr, x)

    @test variable(y_c) == y
    @test literal(y_c) == var2lit(y)
    @test vtree(y_c) ∈ mgr
    @test positive(y_c)
    @test y_c == compile(mgr, y)

    notx = -var2lit(x)

    notx_c = compile(mgr,notx)

    @test sat_prob(node2dag(x_c)) == 1//2
    @test sat_prob(node2dag(notx_c)) == 1//2
    @test model_count(node2dag(x_c),num_vars) == BigInt(2)^(num_vars-1)
    @test model_count(node2dag(notx_c),num_vars) == BigInt(2)^(num_vars-1)

    @test variable(notx_c) == x
    @test literal(notx_c) == notx
    @test vtree(notx_c) ∈ mgr
    @test negative(notx_c)
    @test notx_c == compile(mgr, notx)

    true_c = compile(true)

    @test is_true(true_c)
    @test constant(true_c) == true

    false_c = compile(false)

    @test is_false(false_c)
    @test constant(false_c) == false

    @test !true_c == false_c
    @test !false_c == true_c
    @test !x_c == notx_c
    @test !notx_c == x_c

    @test false_c & true_c == false_c
    @test false_c & notx_c == false_c
    @test false_c & x_c == false_c
    @test true_c & notx_c == notx_c
    @test x_c & notx_c == false_c
    @test true_c & true_c == true_c
    @test false_c & false_c == false_c
    @test x_c & x_c == x_c
    @test !x_c & !x_c == !x_c

    @test false_c | true_c == true_c
    @test false_c | notx_c == notx_c
    @test false_c | x_c == x_c
    @test true_c | notx_c == true_c
    @test x_c | notx_c == true_c
    @test true_c | true_c == true_c
    @test false_c | false_c == false_c
    @test x_c | x_c == x_c
    @test !x_c | !x_c == !x_c

    @test model_count(node2dag(true_c),num_vars) == BigInt(2)^(num_vars)
    @test model_count(node2dag(false_c),num_vars) == BigInt(0)

    v1 = compile(mgr, Var(1))
    v2 = compile(mgr, Var(2))
    v3 = compile(mgr, Var(3))
    v4 = compile(mgr, Var(4))
    v5 = compile(mgr, Var(5))
    v6 = compile(mgr, Var(6))
    v7 = compile(mgr, Var(7))

    p1 = XYPartition([Element(true_c,v3)])
    @test canonicalize(p1) === v3
    p2 = XYPartition([Element(v1,true_c), Element(!v1,false_c)])
    @test canonicalize(p2) === v1

    p3 = XYPartition([Element(v1,v3), Element(!v1,v7)])
    n1 = canonicalize(p3)
    p4 = XYPartition([Element(!v1,v7), Element(v1,v3)])
    n2 = canonicalize(p4)
    @test n1 === n2

    t1 = v1 & v3
    t2 = v3 & v1

    @test t1 === t2
    @test model_count(node2dag(t1),num_vars) == BigInt(2)^(num_vars-2)

    @test model_count(node2dag(v1 & v2 & v6),num_vars) == BigInt(2)^(num_vars-3)

    c1 = v1 | v3
    c2 = v3 | v1

    @test c1 === c2
    @test model_count(node2dag(c1),num_vars) == BigInt(2)^(num_vars-2) * 3

    f1 = (c1 & c2)

    @test f1 === (c2 & c1)
    @test f1 === (c1 & c2 & c2)
    @test f1 === (c1 & c2 & c2 & true_c)

    @test (v3 | v7) !== false_c
    @test (v3 | v7) & (!v3 | v7) !== false_c
    @test (v3 | v7) & (!v3 | v7) & (v3 | !v7) !== false_c
    @test (v3 | v7) & (!v3 | v7) & (v3 | !v7) & (!v3 | !v7) === false_c

    f2 = (c1 | c2)

    @test f2 === (c2 | c1)
    @test f2 === (c1 | c2 | c2)
    @test f2 === (c1 | c2 | c2 | false_c)

    @test (v3 & v7) !== true_c
    @test (v3 & v7) | (!v3 & v7) !== true_c
    @test (v3 & v7) | (!v3 & v7) | (v3 & !v7) !== true_c
    @test (v3 & v7) | (!v3 & v7) | (v3 & !v7) | (!v3 & !v7) === true_c

    @test (v1 & v7 & v3) === !(!v1 | !v7 | !v3)

    @test f2 & !f2 === false_c
    @test f2 | !f2 === true_c

    f3 = (v1 | !v3 | v5) & (!v2 | !v3 | v7) & (v2 | !v6 | !v5) & (v1 | v2 | v3) & (!v5 | v6 | v7) & (v2 | !v3 | !v6) & (!v1 | !v6 | !v7)

    @test f3 & !f3 === false_c
    @test f3 | !f3 === true_c
    @test f3 & !(v2 | !v6 | !v5) === false_c

    f4 = ((v1 | !v3 | v5) & (!v5 | v6 | v7| v1) & (!v5 | v6 | v7)  & (!v1 | !v6 | !v7)  & (!v2 | !v3 | v7)) & ((v2 | !v3 | !v6)  & (v2 | !v6 | !v5) & (v1 | v2 | v3) & (v1 | v2 | v3 | !v6) )

    @test f3 === f4

    fΔ = node2dag(f4, TrimNode)

    @test TrimSdd <: Sdd
    @test fΔ isa TrimSdd
    @test fΔ isa Sdd

end


#=
@author : Spencer Lyon
          Victoria Gregory

@date: 07/09/2014
=#

using Plots
pyplot()

const r = 0.05
const beta = 1.0 / (1.0 + r)
const T = 20  # Time horizon
const S = 5   # Impulse date
const sigma1 = 0.15
const sigma2 = 0.15

function time_path(permanent=false)
    w1 = zeros(T+1)
    w2 = zeros(T+1)
    b = zeros(T+1)
    c = zeros(T+1)

    if permanent === false
        w2[S+2] = 1.0
    else
        w1[S+2] = 1.0
    end

    for t=2:T
        b[t+1] = b[t] - sigma2 * w2[t]
        c[t+1] = c[t] + sigma1 * w1[t+1] + (1 - beta) * sigma2 * w2[t+1]
    end

    return b, c
end

function main()
    L = 0.175

    b1, c1 = time_path(false)
    b2, c2 = time_path(true)
    p = plot(0:T, [c1 c2 b1 b2], layout=(2, 1),
             color=[:green :green :blue :blue],
             label=["consumption" "consumption" "debt" "debt"])
    t = ["impulse-response, transitory income shock"
         "impulse-response, permanent income shock"]
    plot!(title=t', xlabel="Time", ylims=(-L, L), legend=[:topright :bottomright])
    vline!([S S], color=:black, layout=(2, 1), label="")

    return p
end


# This is just a bit of code to set up a temporary package environment since not everyone
# will have the same environment set up; this is largely unnecessary when working only on
# your own code.
# See https://bkamins.github.io/julialang/2020/06/28/automatic-project-environments.html
using Pkg: Pkg
# Change to a newly created temporary directory
cd(mktempdir()) do
    Pkg.activate(".")  # activate a new project in the temporary directory
    Pkg.UPDATED_REGISTRY_THIS_SESSION[] = true  # don't update the registry; on Windows it's slow...
    # Add the required packages
    Pkg.add(Pkg.PackageSpec(name="Plots"))
end

# Load required packages
using Plots

function integrate_sde(sde, x0, t1; dt=0.01)
    # Euler-Maruyama for a 1D SDE
    x = x0
    for t in range(0, t1, step=dt)
        (f, g) = sde(x)
        x = x + f*dt + g*sqrt(dt)*randn()
    end
    return x
end

function integrate_many(sde, x0, t1; n = 1, kwargs...)
    # Integrate many sample paths using Euler-Maruyama
    x = zeros(n)
    Threads.@threads for i in eachindex(x)  # use threads to speed up the computation
        x[i] = integrate_sde(sde, x0, t1; kwargs...)
    end
    return x
end

ornstein_uhlenbeck(x, θ, σ) = (-θ*x, σ) # Ornstein-Uhlenbeck equation (this is short hand notation for a function)
cubic(x, θ, σ) = (θ*x - x^3, σ) # SDE with two potential wells

println("Using $(Threads.nthreads()) threads")

x = integrate_many(x->ornstein_uhlenbeck(x, 1.0, 0.1), 0.0, 10.0, n = 1_000_000) # integrate a million sample paths
plt1 = histogram(x)
title!("Ornstein-Uhlenbeck: θ=1.0, σ=0.1")

x = integrate_many(x->cubic(x, 1.0, 0.1), 0.0, 10.0, n = 1_000_000) # integrate a million sample paths
plt2 = histogram(x)
title!("Cubic: θ=1.0, σ=0.1")

x = integrate_many(x->cubic(x, 1.0, 1.5), 0.0, 10.0, n = 1_000_000) # integrate a million sample paths
plt3 = histogram(x)
title!("Cubic: θ=1.0, σ=1.5")

# Use `gui(plt1)` etc to show the different figures


using Random
using LinearAlgebra
using GalacticOptim,Optim,NLopt#,BlackBoxOptim
using DifferentialEquations#,DiffEqFlux
#using ForwardDiff
using Zygote
using Plots
#using ReverseDiff
include("../src/GLOQ.jl")
pyplot()
# BlackBoxOptim somehow downgrade some packages and as a result breaks the auto-differentiation with Zygote
# we should avoid it.
# NLopt seems to be really fast. Their LBFGS with box constrains is faster than Optim's FMinbox(LBFGS)

N_states = 4;
freqs = [4.10817777; 3.88303940; 3.61937188]
omegas = 2*pi.*freqs
gamma1   = [2.59451982e-05; 4.54296537e-05; 0.0]
gamma2   = [2.58005604e-05; 9.00000000e-05; 0.0]
omr_ramsey = 2.0*pi*(3.8830355 - 2.0e-3)
omr_echo = 2.0*pi*3.8830355
omr_T1 = 2.0*pi*3.8830355

width = 17.00
TC = 2.5*width
Ω = 0.5*pi/(TC*sqrt(2.0))
θ = 0.0

fromState = 1
rho_u0 = [0.0;0.0;0.0;0.0]
rho_v0 = [0.0;0.0;0.0;0.0]
rho_u0[fromState+1] = 1.0

T_Ramsey = 1.0*GLOQ.GLOQ_MICRO_SEC#10.0 * 1000
N_Ramsey = 101
t_ramsey = zeros(Float64,N_Ramsey)
dt_ramsey = T_Ramsey/(N_Ramsey-1)
for i = 1:N_Ramsey
	dark_time = (i-1)*dt_ramsey
	t_ramsey[i] = dark_time
end

# Ramsey Forward Solve
@time rho_ramsey_u_exp,rho_ramsey_v_exp = GLOQ.RamseyForwardSolve(rho_u0,rho_v0,
			     omegas,omr_ramsey,
				 gamma1,gamma2,
				 fromState,
				 TC,t_ramsey,N_states;
				 initial_type="states")
@time rho_ramsey_u,rho_ramsey_v = GLOQ.RamseyForwardSolve(rho_u0,rho_v0,
			     omegas,omr_ramsey,
				 gamma1,gamma2,
				 fromState,
				 TC,t_ramsey,N_states;
				 method = Trapezoid(),
				 abstol = 1e-10,
				 restol = 1e-10
				 )
population_Ramsey = GLOQ.get_population(rho_ramsey_u)
population_Ramsey_exp = GLOQ.get_population(rho_ramsey_u_exp)

println("Ramsey max error:",maximum(abs.(population_Ramsey-population_Ramsey_exp)) )

fig=plot(t_ramsey./GLOQ.GLOQ_MICRO_SEC,population_Ramsey)
plot!(fig,t_ramsey./GLOQ.GLOQ_MICRO_SEC,population_Ramsey_exp)
display(fig)

fig2=plot(t_ramsey./GLOQ.GLOQ_MICRO_SEC,population_Ramsey-population_Ramsey_exp)
display(fig2)

# Echo Forward Solve
T_Echo = 10.0*GLOQ.GLOQ_MICRO_SEC#10.0 * 1000
N_Echo = 1001
t_echo = zeros(Float64,N_Echo)
dt_echo = T_Echo/(N_Echo-1)
for i = 1:N_Echo
	dark_time = (i-1)*dt_echo
	t_echo[i] = dark_time
end
@time rho_echo_u_exp,rho_echo_v_exp = GLOQ.EchoForwardSolve(rho_u0,rho_v0,
			     omegas,omr_echo,
				 gamma1,gamma2,
				 fromState,
				 TC,t_echo,N_states;
				 initial_type = "states")
@time rho_echo_u,rho_echo_v = GLOQ.EchoForwardSolve(rho_u0,rho_v0,
			     omegas,omr_echo,
				 gamma1,gamma2,
				 fromState,
				 TC,t_echo,N_states;
				 initial_type = "states",
				 method = Trapezoid(),
				 abstol = 1e-10,
				 restol = 1e-10)
population_Echo = GLOQ.get_population(rho_echo_u)
population_Echo_exp = GLOQ.get_population(rho_echo_u_exp)

println( maximum(abs.(population_Echo-population_Echo_exp)) )

fig_echo=plot(t_echo./GLOQ.GLOQ_MICRO_SEC,population_Echo)
plot!(fig_echo,t_echo./GLOQ.GLOQ_MICRO_SEC,population_Echo_exp,line=(:dash))
display(fig_echo)

fig_echo2=plot(t_echo./GLOQ.GLOQ_MICRO_SEC,population_Echo-population_Echo_exp)
display(fig_echo2)

# T1 Forward Solve
T_T1 = 10.0*1000
N_T1 = 1001
t_T1 = zeros(Float64,N_T1)
dt_T1 = T_T1/(N_T1-1)
for i = 1:N_T1
	dark_time = (i-1)*dt_T1
	t_T1[i] = dark_time
end
@time rho_t1_u_exp,rho_t1_v_exp = GLOQ.T1ForwardSolve(rho_u0,rho_v0,
			     omegas,omr_T1,
				 gamma1,gamma2,
				 fromState,
				 TC,t_T1,N_states;
				 initial_type = "states")
population_T1_exp = GLOQ.get_population(rho_t1_u_exp)
@time rho_t1_u,rho_t1_v = GLOQ.T1ForwardSolve(rho_u0,rho_v0,
			     omegas,omr_T1,
				 gamma1,gamma2,
				 fromState,
				 TC,t_T1,N_states;
				 initial_type = "states",
				 method = Trapezoid(),
				 abstol = 1e-10,
				 restol = 1e-10)
population_T1 = GLOQ.get_population(rho_t1_u)

println( maximum(abs.(population_T1-population_T1_exp)) )

fig_t1=plot(t_T1./GLOQ.GLOQ_MICRO_SEC,population_T1)
plot!(fig_t1,t_T1./GLOQ.GLOQ_MICRO_SEC,population_T1_exp,line=(:dash))
display(fig_t1)

fig_t1_error=plot(t_T1./GLOQ.GLOQ_MICRO_SEC,population_T1-population_T1_exp)
display(fig_t1_error)


using Test
using BenchmarkTools

include(joinpath(dirname(dirname(@__FILE__)), "examples", "cartesian.jl"))

using .Cartesian

const nx = 10
const h = 2/(nx-1)
println("building grid")
const cgrid = CartesianGrid(nx,nx,h)

println("tests...")
@test nb_nodes(cgrid) == nx*nx
const crds = coordinates(cgrid)
@test crds[1,1:2] == [-1.0, -1.0]
@test crds[end,1:2] == [1.0, 1.0]

println("nb neighbors...")
@btime nb_neighbors(cgrid, 4, (g,I) -> I[1] == 1)

# if isinteractive()
#   using Plots
#   scatter(crds[:,1], crds[:,2])
# end