export linearInter

function linearInter(T::Array{Float64},domain::Array{Float64},xin::Array{Float64};doDerivative::Bool=false)

	dim = round(Int,length(domain)/2)
	n   = (dim==1) ? length(T) : size(T)
	h   = sdiag(1./collect(n)) * vec(domain[2:2:end]-domain[1:2:end])
	np  = round(Int,length(xin)/dim)
	x   = reshape(copy(xin),np,dim)

	# map x from [h/2,domain-h/2] -> [1,m],
	for i=1:dim
		x[:,i] = (x[:,i]-domain[2*i-1])/h[i] + 0.5;
	end
	Tc = zeros(np)                                # initialize output
	Valid(j) = (0.<x[:,j]) .& (x[:,j].<(n[j]+1))      # determine indices of valid points

	if dim==1
		valid = find(Valid(1))
	elseif dim==2
		valid = find( Valid(1) .& Valid(2) );
	elseif dim==3
		valid = find( Valid(1) .& Valid(2) .& Valid(3) );
	end

	if isempty(valid)
		if !doDerivative
			return Tc
		else
			return Tc, spzeros(np,dim*np)
		end
	end

	pad = 1; TP = zeros(tuple(collect(n)+2*pad...));                 # pad data to reduce cases
	P     = round.(Int,floor.(x)); x = x-P;                        # split x into integer/remainder
	p(j)  = P[valid,j]
	xi(j) = x[valid,j]

	# increments for linearized ordering
	i1 = 1; i2 = size(T,1)+2*pad; i3 = (size(T,1)+2*pad)*(size(T,2)+2*pad);

	if dim==1
		TP[pad+(1:n)] = vec(T)
		p = pad + p(1);
	    Tc[valid] = TP[p].* (1-xi(1)) + TP[p+1].*xi(1);   # compute weighted sum

	    if doDerivative
			dT = zeros(np)
			dT[valid] = TP[p+1]-TP[p]
			dT = sdiag(dT[:,1]/h[1])
		end;
	elseif dim==2
	    TP[pad+(1:n[1]),pad+(1:n[2])] = T
	    TP = vec(TP)
	    p  = (pad + p(1)) + i2*(pad + p(2) - 1)

	    # compute Tc as weighted sum
	    Tc[valid] = (TP[p]   .* (1-xi(1))  + TP[p+i1]    .*xi(1)) .* (1-xi(2))  + (TP[p+i2] .* (1-xi(1)) + TP[p+i1+i2] .*xi(1)) .* xi(2)

		if doDerivative
			dT = zeros(np,dim)
			dT[valid,1] = (TP[p+i1]-TP[p]).*(1-xi(2)) + (TP[p+i1+i2]-TP[p+i2]).*xi(2)
			dT[valid,2] = (TP[p+i2]-TP[p]).*(1-xi(1)) + (TP[p+i1+i2]-TP[p+i1]).*xi(1)
			dT = spdiagm((dT[:,1]/h[1], dT[:,2]/h[2]),(0,np), np, 2*np)
			# dT = [sdiag(dT[:,1]/h[1])  sdiag(dT[:,2]/h[2])]
		end
	  elseif dim==3
	    TP[pad+(1:n[1]),pad+(1:n[2]),pad+(1:n[3])] = T
	    p  = (pad + p(1)) + i2*(pad + p(2) - 1) + i3*(pad + p(3) -1);
	    # compute Tc as weighted sum
	    Tc[valid] = ((TP[p].*(1-xi(1))+TP[p+i1].*xi(1)).*(1-xi(2)) +
		  (TP[p+i2].*(1-xi(1))+TP[p+i1+i2].*xi(1)).*(xi(2))).*(1-xi(3)) +
	      +((TP[p+i3].*(1-xi(1))+TP[p+i1+i3].*xi(1)).*(1-xi(2)) +
	      +(TP[p+i2+i3].*(1-xi(1))+TP[p+i1+i2+i3].*xi(1)).*(xi(2))).*(xi(3))

	    if doDerivative
			dT = zeros(np,dim)
			dT[valid,1] = ((TP[p+i1]-TP[p]).*(1-xi(2))+(TP[p+i1+i2]-TP[p+i2]).*xi(2)).*(1-xi(3)) +
	                      ((TP[p+i1+i3]-TP[p+i3]).*(1-xi(2))+(TP[p+i1+i2+i3]-TP[p+i2+i3]).*xi(2)).*(xi(3))
		    dT[valid,2] = ((TP[p+i2]-TP[p]).*(1-xi(1))+(TP[p+i1+i2]-TP[p+i1]).*xi(1)).*(1-xi(3)) +
	                      ((TP[p+i2+i3]-TP[p+i3]).*(1-xi(1))+(TP[p+i1+i2+i3]-TP[p+i1+i3]).*xi(1)).*(xi(3))
		    dT[valid,3] = ((TP[p+i3].*(1-xi(1))+TP[p+i1+i3].*xi(1)).*(1-xi(2)) +
	                      (TP[p+i2+i3].*(1-xi(1))+TP[p+i1+i2+i3].*xi(1)).*(xi(2))) -
	                      ((TP[p].*(1-xi(1))+TP[p+i1].*xi(1)).*(1-xi(2)) +
	                      (TP[p+i2].*(1-xi(1))+TP[p+i1+i2]*xi(1)).*(xi(2)));
			dT = [sdiag(dT[:,1]/h[1])  sdiag(dT[:,2]/h[2])  sdiag(dT[:,3]/h[3])]
		end
	end
	if doDerivative
		return Tc,dT
	else
		return Tc
	end
end


module AudioDeclipping

using BenchmarkTools
using StructuredOptimization, AbstractOperators
using DSP, WAV
using PyPlot

function set_up()

	Nl = 2^10                           # frames length
	Fs = 16e3                           # sampling frequency
	x0, Fs0 = wavread("demo.wav")       # load input signal
	x0 = x0[:,1]
	x0 = resample(x0,Int(Fs)//Int(Fs0))
	normalize!(x0,Inf)
	Nt = length(x0)                     # signal length
	x0 .*= 0.8

	C = 0.02                                       # clip threshold
	xc = copy(x0)
	xc[abs.(xc).>=C] = C.*sign.(xc[abs.(xc).>=C])  # clip original signal

	#time win options
	K = 2 # K = 2 hanning
	win = sqrt.(tukey(Nl+1,2/K)[1:Nl])
	overlap = div(Nl,K)

	xd = zeros(x0)                    # allocate output
	x, y = Variable(Nl), Variable(Nl) # optimization variables
	cf = ls( idct(x) - y )            # cost function
	yw = zeros(Nl)

	return Fs, x, x0, xd, xc, y, yw, Nl, Nt, C, win, overlap, cf
end

function run_demo()
	slv = PANOC(tol = 1e-5, verbose = 0)
	setup = set_up()
	@time solve_problem!(slv, setup..., true)

	Fs, x, x0, xd, xc, y, yw, Nl, Nt, C, overlap, cf = setup
	Irp = find(xc.>=C)  #positive clipping indices
	Irn = find(xc.<=-C) #negative clipping indices

	SDRy  = 20*log10(norm(x0[[Irp;Irn]])/norm(x0[[Irp;Irn]]-xc[[Irp;Irn]]))
	SDRyd = 20*log10(norm(x0[[Irp;Irn]])/norm(x0[[Irp;Irn]]-xd[[Irp;Irn]]))

	println("SDRy: $(SDRy) SDRyd: $(SDRyd) diff: $(SDRyd-SDRy)")
	return setup
end

function solve_problem!(slv, Fs, x, x0, xd, xc, y, yw, Nl, Nt, C, win, overlap, cf, verb)
	its = Int64[]
	fit_tol = 1e-5
	z = 0                             # frame initial intex
	counter = 0
    winC = C.*win
	#weighted Overlap-Add
	while z+Nl < Nt
		counter += 1
		push!(its,0)

		fill!(x.x,0.) # initialize variables
		fill!(y.x,0.)

		Irp = sort(find(     xc[z+1:z+Nl]  .>=  C) ) #positive clipping indices
		Irn = sort(find(     xc[z+1:z+Nl]  .<= -C) ) #negative clipping indices
		Im  = sort(find(abs.(xc[z+1:z+Nl]) .<   C)) #non clipped indices
		yw .= xc[z+1:z+Nl].*win

		Mp = GetIndex((Nl,),Irp)
		Mn = GetIndex((Nl,),Irn)
		M  = GetIndex((Nl,),Im )

		if isempty(Irp) && isempty(Irn)
			xd[z+1:z+Nl] .+= yw.*win
			if verb @printf("%7d / %7d |           no_clip          | \n", z, Nt) end
		else
			N = 0 # number of active components in DCT
			for N = 30:30:30*div(Nl,30)
				cstr = (norm(x,0) <= N,
					norm(M*y-M*yw) <= sqrt(fit_tol),
					Mp*y >=  Mp*winC,
                    Mn*y <= -Mn*winC)
                it, = @minimize cf st cstr with slv
				its[counter] += it
                if norm(idct(~x)- ~y) <= sqrt(fit_tol) break end
			end
			@views xd[z+1:z+Nl] .+= (~y).*win
			if verb @printf("%7d / %7d | N: %7d /  %7d | \n", z, Nt, N, Nl) end
		end

		z += Nl-overlap                      # update index
	end
	return mean(its)
end

function benchmark(;verb = 0, samples = 1, seconds = 100)

	suite = BenchmarkGroup()

	tol = 1e-4
	solvers = [
               "PANOC",
               "ZeroFPR",
               "PG"
          ]
	slv_opt = ["(verbose = $verb, tol = $tol)",
		   "(verbose = $verb, tol = $tol)",
		   "(verbose = $verb, tol = $tol)"]

	its = Dict([(sol,0.) for sol in solvers])
	for i in eachindex(solvers)

		setup = set_up()
		solver = eval(parse(solvers[i]*slv_opt[i]))

		suite[solvers[i]] =
		@benchmarkable(it = solve_problem!(solver, setup..., false),
			       setup = (
					it = 0;
					setup = deepcopy($setup);
					solver = deepcopy($solver) ),
			       teardown = (
					  $its[$solvers[$i]] = it;
					  ),
			       evals = 1, samples = samples, seconds = seconds)
	end

	results = run(suite, verbose = (verb != 0))
	println("AudioDecliping (mean of per-frame) its")
	println(its)
	return results
end

function show_results(Fs, x, x0, xd, xc, y, yw, Nl, Nt, C, win, overlap, cf)
	idxs = 2*Nl+1:4*Nl
	t = linspace(0,length(idxs)/Fs,length(idxs))
	figure()
	plot(t,x0[idxs])
	plot(t,xd[idxs])
	plot(t,xc[idxs])
end

function save_wav(Fs, x, x0, xd, xc, y, yw, Nl, Nt, C, win, overlap, cf)
    wavwrite(x0 ,"unclipped.wav"; Fs = Fs, nbits = 16, compression=WAVE_FORMAT_PCM)
    wavwrite(xc ,"clipped.wav";   Fs = Fs, nbits = 16, compression=WAVE_FORMAT_PCM)
    wavwrite(0.8*normalize(xd,Inf) ,"declipped.wav"; Fs = Fs, nbits = 16, compression=WAVE_FORMAT_PCM)
end

end


####
#   normalform (reduction to Zumbroich basis)

"""
    normalform(α::Cyclotomic)
Reduces `α` to the Zumbroich basis.
"""
normalform(α::Cyclotomic) = normalform!(deepcopy(α))

"""
    normalform!(α::Cyclotomic[, tmp=dense(α)]; basis_forbidden=...)
Reduces `α` to the Zumbroich basis in-place. Note that unless `α` has dense
storage the operation will need at least one dense `tmp::Cyclotomic`.
"""
function normalform!(
    α::Cyclotomic{T},
    tmp = dense(α);
    basis_forbidden = zumbroich_viacomplement(conductor(α)),
) where {T}
    # @debug "normalizing:" conductor(α) coeffs(α)
    isnormalized(α, first(basis_forbidden)) && return α

    copyto!(coeffs(tmp), coeffs(α))
    normalform!(tmp, basis_forbidden = basis_forbidden)
    copyto!(coeffs(α), coeffs(tmp))

    return α
end

function normalform!(
    α::Cyclotomic{T,<:DenseVector};
    basis_forbidden = zumbroich_viacomplement(conductor(α)),
) where {T}

    basis, forbidden = basis_forbidden
    isnormalized(α, basis) && return α
    for fb in forbidden
        for exp in exponents(α)
            exp in basis && continue
            _replace_exponent!(α, exp, fb)
        end
    end
    return α
end

function _replace_exponent!(α::Cyclotomic, exp::Integer, fb)
    val = α[exp]
    p, q, forbidden_exps = fb

    # @debug p forbidden_exps
    if exp % q ∈ forbidden_exps
        # @debug "removing $(α[exp])*ζ^$exp from" α
        m = conductor(α) ÷ p
        # @debug " " collect(α[exp .+ 1:m:m*(p-1)])

        # TODO: this doesn't work
        # α[exp .+ 1:m:m*(p-1)] .-= val

        for i = 1:p-1
            α[exp+i*m] -= val
            # exp + i*m is not guaranteed to be allowed, but it's larger than exp and hence will be dealt later
        end
        α[exp] = 0
        # @debug "after removal:" α
    end
    return α
end

"""
    embed(α::Cyclotomic, m::Integer)
Embed `α` into the `m`-th cyclotomic field. `m` can be either larger, or smaller
than the conductor of `α`, however either `conductor(α)` must divide `m`, or
the other way around.
"""
function embed(α::Cyclotomic, m::Integer)
    conductor(α) == m && return deepcopy(α)
    if conductor(α) > m
        @assert conductor(α) % m == 0 "Embeding of ℚ(ζ$(subscriptify(m))) ↪ ℚ(ζ$(subscriptify(conductor(α)))) is not possible."
        return reduced_embedding(α, m)
    else
        k, _r = divrem(m, conductor(α))
        @assert _r == 0 "Embeding of ℚ(ζ$(subscriptify(conductor(α)))) ↪ ℚ(ζ$(subscriptify(m))) is not possible."

        res = zero!(similar(α, valtype(α), m))

        @inbounds for e in exponents(α)
            res[e*k] = α[e]
        end
        return res
    end
end

function _tmp_for_reduced_embedding(α::Cyclotomic{T}) where {T}

    all(iszero, exponents(α)) && return Cyclotomic{T,Vector{T}}(1, [α[0]])

    k = gcd(conductor(α), exponents(α)...)
    @debug "gcd(exponents(α)) = $k" collect(exponents(α))

    tmp = if k > 1
        d = div(conductor(α), k)
        @debug "Reducing the embedding ring to $T(ζ$(subscriptify(d)))"
        tmp = Cyclotomic{T,Vector{T}}(d, zeros(T, d))
    else
        n = conductor(α)
        tmp = Cyclotomic{T,Vector{T}}(n, zeros(T, n))
    end
    # now tmp has dense storage

    if k > 1
        # the trivial reduction E(n)^e → E(n÷k)^(e÷k)
        @debug "Performing trivial reduction from ℚ(ζ$(subscriptify(conductor(α)))) → ℚ(ζ$(subscriptify(conductor(tmp))))"
        @inbounds for (e, c) in exps_coeffs(α)
            tmp[div(e, k)] = c
        end
    else
        @debug "No trivial reduction is possible"
        copyto!(tmp.coeffs, coeffs(α))
    end
    return tmp
end

"""
    reduced_embedding(α::Cyclotomic{T,V}[, m::Integer=1])
Return the reduced embedding of `α` into `m`-th cyclotomic field.

When `m=1`, the embedding into possibly smallest degree is returned.
"""
function reduced_embedding(α::Cyclotomic{T,V}, m::Integer = 1) where {T,V}

    tmp = _tmp_for_reduced_embedding(normalform!(α))

    if conductor(tmp) == 1
        res = Cyclotomic{T,V}(conductor(tmp), coeffs(tmp))
        return res
    end

    basis, forbidden = Cyclotomics.zumbroich_viacomplement(conductor(tmp))

    normalform!(tmp, basis_forbidden = (basis, forbidden))

    phi_nc = length(basis)
    nz = count(!iszero, coeffs(tmp))

    if all(p == q for (p, q, _) in forbidden) && # conductor(tmp) is square-free
       phi_nc == nz && # tmp is supported on φ(n) elements
       _all_equal(tmp, exponents(tmp)) # all nz coeffs of tmp are equal

        val = if iseven(length(forbidden))
            tmp[first(exponents(tmp))]
        else
            -tmp[first(exponents(tmp))]
        end
        @debug "Cyclotomic is real:" α = val
        res = Cyclotomic{T,V}(m, similar(coeffs(α), m))
        zero!(res)
        res[0] = val
        return res
    end

    for (p, q, fb_res) in forbidden
        p == 2 && continue
        n = conductor(tmp)
        nz = count(!iszero, coeffs(tmp))

        p == q || continue
        @debug "$(p)² does NOT divide conductor $(conductor(tmp))" p q

        m % p != 0 || continue
        @debug "prime $p does not divide $m so embedding is possible"

        # there are n/p residue classes, each contains p-1 elts, and tmp is constant on each
        nnz, _r = divrem(nz, p - 1)
        @debug "the number of non-zero coeffs $nz ≡ $_r mod $(p-1)"
        _r == 0 || continue

        # Check that coeffs(tmp) are constant on each congruence classes.
        # To illustrate by example write:
        #
        # 1 = - ζ₅ - ζ₅² - ζ₅³ - ζ₅⁴
        #
        # and consider
        #
        # c*ζ₉^e = c*ζ₉^e * 1 = c*ζ₉^e * (- ζ₅ - ζ₅² - ζ₅³ - ζ₅⁴) =
        # - c*ζ₄₅^(5e+1*9) - c*ζ₄₅^(5e+2*9) - c*ζ₄₅^(5e+3*9) - c*ζ₄₅^(5e+4*9)
        #
        # If tmp ∈ ℚ(ζ_n) can be reduced to ℚ(ζ_{n÷p})
        # (n = 45, p = 5 in the example), we need to check that all coefficients
        #
        # tmp[p*e + i*n÷p] are equal (for 1 ≤ i ≤ p-1)
        #
        # where exponents 0 ≤ e ≤ n-1 represent congruence classes modulo n÷p

        n_p = n ÷ p

        equal_on_classes = true
        for pe = 0:p:(n-1)
            equal_on_classes || break
            if !_all_equal(tmp, pe .+ (n_p .* (1:p-1)))
                equal_on_classes = false
                @debug "skipping reduction over prime $p as cyclotomic is not constant over class $(pe + (n_p)) (mod $(n_p)) " (
                    pe + n_p,
                    tmp[pe+n_p],
                ) (pe .+ (n_p .* (2:p-1)), tmp[pe.+(n_p.*(2:p-1))])
            end
        end
        equal_on_classes || continue # to the next prime

        @debug "cyclotomic is constant over residue classes (mod $(n_p)), reducing"

        for pe = 0:p:(n-1)
            # replace sum_i(tmp[p*e + i*n÷p]) (1 ≤ i ≤ p-1)
            # inverse to normalform! over p
            @debug "replacing $pe-th power by $(pe + n_p)-th:" tmp[pe] -tmp[pe+n_p]
            tmp[pe] = -tmp[pe+n_p] # tmp[pe + 0*n_p] = -tmp[pe + 1*n_p]
            tmp[pe.+(n_p.*(1:p-1))] .= zero(T)
        end

        for i = 1:n_p
            tmp[i] = tmp[i*p]
            tmp[i*p] = zero(T)
        end

        tmp = Cyclotomic{T,Vector{T}}(n_p, resize!(tmp.coeffs, n_p))
    end

    res = Cyclotomic{T,V}(conductor(tmp), coeffs(tmp))

    return res
end


@testset "ece.jl" begin
    @testset "Trivial tests" begin
        ece = ECE(UniformBinning(10))

        # categorical distributions
        for predictions in ([[0, 1], [1, 0]], ColVecs([0 1; 1 0]), RowVecs([0 1; 1 0]))
            @test iszero(@inferred(ece(predictions, [2, 1])))
        end
        for predictions in (
            [[0, 1], [0.5, 0.5], [0.5, 0.5], [1, 0]],
            ColVecs([0 0.5 0.5 1; 1 0.5 0.5 0]),
            RowVecs([0 1; 0.5 0.5; 0.5 0.5; 1 0]),
        )
            @test iszero(@inferred(ece(predictions, [2, 2, 1, 1])))
        end

        # probabilities
        for predictions in ([0, 1], [0.0, 1.0])
            @test iszero(@inferred(ece(predictions, [false, true])))
        end
        @test iszero(@inferred(ece([0, 0.5, 0.5, 1], [false, false, true, true])))
    end

    @testset "Uniform binning: Basic properties" begin
        ece = ECE(UniformBinning(10))
        estimates = Vector{Float64}(undef, 1_000)

        # categorical distributions
        for nclasses in (2, 10, 100)
            dist = Dirichlet(nclasses, 1.0)

            predictions = [Vector{Float64}(undef, nclasses) for _ in 1:20]
            targets = Vector{Int}(undef, 20)

            for i in 1:length(estimates)
                rand!.(Ref(dist), predictions)
                targets .= rand.(Categorical.(predictions))

                estimates[i] = ece(predictions, targets)
            end

            @test all(x -> zero(x) < x < one(x), estimates)
        end

        # probabilities
        predictions = Vector{Float64}(undef, 20)
        targets = Vector{Bool}(undef, 20)
        for i in 1:length(estimates)
            rand!(predictions)
            map!(targets, predictions) do p
                rand() < p
            end
            estimates[i] = ece(predictions, targets)
        end
        @test all(x -> zero(x) < x < one(x), estimates)
    end

    @testset "Median variance binning: Basic properties" begin
        ece = ECE(MedianVarianceBinning(10))
        estimates = Vector{Float64}(undef, 1_000)

        for nclasses in (2, 10, 100)
            dist = Dirichlet(nclasses, 1.0)

            predictions = [Vector{Float64}(undef, nclasses) for _ in 1:20]
            targets = Vector{Int}(undef, 20)

            for i in 1:length(estimates)
                rand!.(Ref(dist), predictions)
                targets .= rand.(Categorical.(predictions))

                estimates[i] = ece(predictions, targets)
            end

            @test all(x -> zero(x) < x < one(x), estimates)
        end
    end
end


#=
Copyright Alec Hammond 2022

Simple fdtd in 2d excercise.

Assumes TMᶻ polarization. Relevant physics:
    −σₐHx − μ∂Hx/∂t = ∂Ez/∂y
     σₐHy + μ∂Hy/∂t = ∂Ez/∂x
     σEz + ε∂Ez/∂t = ∂Hy/∂x - ∂Hx/∂y
where
    σₐ is the magnetic conductivity
    σ is the electric conductivity
    ε is the permittivity
    μ is the permeability

Define Ez in lower left corner: Ez(m,n)
Define Hy just to the right:    Hy(m+1/2,n)
Define Hx just above:           Hx(m,n+1/2)

Nodes at top of grid lack Hx
Nodes at right of grid lack Hy

References
-------------------
* https://eecs.wsu.edu/~schneidj/ufdtd/chap8.pdf

=#

const USE_GPU = false
using ImplicitGlobalGrid
using ParallelStencil
using ParallelStencil.FiniteDifferences2D
using Plots, Printf, Statistics

type = Float64
dims = 2
@static if USE_GPU
    @init_parallel_stencil(CUDA, type, dims);
else
    @init_parallel_stencil(Threads, type, dims);
end

@parallel function step_H!(Hx::Data.Array,Hy::Data.Array,Ez::Data.Array,dt::Data.Number,dx::Data.Number,dy::Data.Number)
    # assume σₐ=0
    # assume μ=1
    @all(Hx) = @all(Hx) - dt/dy*@d_ya(Ez)
    @all(Hy) = @all(Hy) + dt/dx*@d_xa(Ez)
    return
end

@parallel function step_E!(Hx::Data.Array,Hy::Data.Array,Ez::Data.Array,ε::Data.Array,dt::Data.Number, dx::Data.Number, dy::Data.Number)
    # assume σ=0
    @inn(Ez) = @inn(Ez) + (dt/dx./@inn(ε)).*(@d_xi(Hy)) - (dt/dy./@inn(ε)).*(@d_yi(Hx))
    return
end

@parallel_indices (ix,iy) function step_E_sources!(Ez::Data.Array,n::Int64,dt::Data.Number, dx::Data.Number, dy::Data.Number)
    # point source in center of cell
    if (ix==size(Ez,1)/2 && iy==size(Ez,2)/2)
        t = n*dt
        Ez[ix,iy] = cos(20*t)*exp((t-5)^2)
    end
    return
end

function fdtd2D_step!(Hx::Data.Array,Hy::Data.Array,Ez::Data.Array,ε::Data.Array,n::Int64,dt::Data.Number, dx::Data.Number, dy::Data.Number)
    @parallel step_H!(Hx,Hy,Ez,dt,dx,dy)
    @parallel step_E!(Hx,Hy,Ez,ε,dt,dx,dy)
    @parallel step_E_sources!(Ez,n,dt,dx,dy)
    return
end

function fdtd2D()

    # Physics
    lx, ly = 1.0, 1.0;

    # Numerics
    C      = 0.5;        # Courant factor
    nx, ny = 256, 256;
    nt     = 250;

    # Derived numerics
    dx, dy    = lx/(nx-1), ly/(ny-1) # cell sizes
    dt        = min(dx,dy)*C

    # Array initializations
    Hx       = @zeros(nx, ny-1);
    Hy       = @zeros(nx-1, ny);
    Ez       = @zeros(nx, ny);
    ε        = @ones(nx, ny);
    @printf("Starting...\n")
    # Time loop
    for it = 1:nt
        if (it==11)  global wtime0 = Base.time()  end
        fdtd2D_step!(Hx,Hy,Ez,ε,it,dt,dx,dy)
    end
    @printf("Stopping...\n")

    # Performance
    wtime    = Base.time() - wtime0
    A_eff    = (2*2+2)/1e9*nx*ny*sizeof(Data.Number)  # Effective main memory access per iteration [GB] (Lower bound of required memory access: H and dHdτ have to be read and written (dHdτ for damping): 4 whole-array memaccess; B has to be read: 1 whole-array memaccess)
    wtime_it = wtime/(nt-10)                          # Execution time per iteration [s]
    T_eff    = A_eff/wtime_it                         # Effective memory throughput [GB/s]

    @printf("Total steps=%d, time=%1.3e sec (@ T_eff = %1.2f GB/s) \n", nt, wtime, round(T_eff, sigdigits=2))
    gr()
    hm = heatmap(Ez, color = :bluesreds, aspect_ratio=:equal, show = true)
    png("Test")
    return
end

fdtd2D()


#
# Create initial point and minimize
#

using Random
using Printf
Random.seed!(1234567)

function initial(input :: MDInput)

  # To clear out the code
  n = input.n

  # Structures to save data
  atoms = Atoms(n)
  traj = MDTraj(n,input.nsave)

  # Local names for code simplicity
  x = atoms.x # this does not not copy x, so that is fine
  if input.uf == "uf_QUAD!"
    uf! = uf_QUAD!
  else
    uf! = input.uf
  end

  # Creating random initial coordinates
  i = 0
  while i < n
    i = i + 1
    x[i,1] = -input.side/2.0 + input.side*rand()
    x[i,2] = -input.side/2.0 + input.side*rand()
  end

  xtrial = similar(x)
  forces = Forces(input)
  f = forces.f

  # Minimizing the energy of the initial point

  println(" ----------------------------------------------------------------")
  println(" ENERGY MINIMIZATION")
  println(" ----------------------------------------------------------------")
  ulast = uf!(n,x,forces,input)
  println(" Energy before minimization: ", ulast)

  dx = 1.0
  fnorm = 1.0
  istep = 0
  while fnorm > input.tol
    istep = istep + 1

    # Compute gradient

    u = uf!(n,x,forces,input)
    fnorm = 0.
    for i in 1:n
      fnorm = fnorm + f[i,1]^2 + f[i,2]^2
    end
    fnorm = fnorm / n

    # Compute trial point ( xtrial = x - (dU/dx)*dx )

    for i in 1:n
      xtrial[i,1] = x[i,1] + f[i,1]*dx
      xtrial[i,2] = x[i,2] + f[i,2]*dx
      xtrial[i,1] = image(xtrial[i,1],input)
      xtrial[i,2] = image(xtrial[i,2],input)
    end
    ustep = uf!(n,xtrial,forces,input)

    # If energy decreased, accept, if not, reject and decrease dx

    if ustep < ulast
      for i in 1:n
        x[i,1] = xtrial[i,1]
        x[i,2] = xtrial[i,2]
      end
      ulast = ustep
      dx = dx * 2.
    else
      dx = dx / 2.
    end

    if istep%input.iprint == 0
      println(@sprintf(" MINIM: U = %12.5e G_norm = %12.5e", ustep, fnorm))
    end

  end

  println(" Energy after minimization: ", ulast)

  # Contamining some individuals
  ncont = max(n*input.Si,1)
  nc = 0
  while nc < ncont
    i = trunc(Int64,n*rand())+1
    if atoms.status[i] == 0
      atoms.status[i] = 1
      nc = nc + 1
    end
  end

  return atoms, traj

end


function estimateBM(xs::Vector{Float64}, b::Int64)
  lxs::Int64 = length(xs)
  a = floor(Int64, lxs/b)
  @assert a > 1
  n::Int64 = a*b
  start::Int64 = length(xs) - n
  overallMean::Float64 = 0.0
  for i = 1:n
    overallMean += xs[start+i]
  end
  overallMean /= n
  acc::Float64 = 0.0
  for i = 1:a
    batchAcc::Float64 = 0.0
    batchStart::Int64 = start + (i-1)*b
    for j = 1:b
      batchAcc += xs[batchStart + j]
    end
    tmp::Float64 = batchAcc/b
    tmp -= overallMean
    acc += tmp * tmp
  end
  return b/(a-1)*acc
end

## Basic batch means estimation of the asymptotic variance
function estimateBM(xs::Vector{Float64})
  return estimateBM(xs, floor(Int64, sqrt(length(xs))))
end


for N in [Float64, Rational{Int}, Float32]
    # ==============================
    # Testing box approximation
    # ==============================
    # Box approximation of a 2D square
    b = BallInf(N[1, 1], N(0.1))
    h = box_approximation(b)
    hexp = Hyperrectangle(N[1, 1], N[0.1, 0.1])
    @test h.center ≈ hexp.center
    @test h.radius ≈ hexp.radius

    # Box approximation of a 2D unit ball in the 1-norm
    b = Ball1(N[1, -2], N(0.2))
    h = box_approximation(b)
    hexp = Hyperrectangle(N[1, -2], N[0.2, 0.2])
    @test h.center ≈ hexp.center
    @test h.radius ≈ hexp.radius

    # Box approximation of a 3D unit ball in the 1-norm
    b = Ball1(N[1, 2, 0], N(1))
    h = box_approximation(b)
    hexp = Hyperrectangle(N[1, 2, 0], N[1, 1, 1])
    @test h.center ≈ hexp.center
    @test h.radius ≈ hexp.radius

    # Box approximation of a Hyperrectangle
    hexp = Hyperrectangle(N[0, 0], N[1, 1])
    h = box_approximation(hexp)
    @test h.center ≈ hexp.center
    @test h.radius ≈ hexp.radius

    # empty set
    E = EmptySet{N}()
    @test box_approximation(E) == E

    # rectification
    @test box_approximation(Rectification(EmptySet{N}())) isa EmptySet{N}
    r = Rectification(Ball1(N[0, 0], N(1)))
    @test box_approximation(r) == Hyperrectangle(low=N[0, 0], high=N[1, 1])

    # ===================================================================
    # Testing box_approximation_symmetric (= symmetric interval hull)
    # ===================================================================
    # Box approximation of a 2D square
    b = BallInf(N[1, 1], to_N(N, 0.1))
    h = box_approximation_symmetric(b)
    hexp = Hyperrectangle(N[0, 0], to_N(N, [1.1, 1.1]))
    @test h.center ≈ hexp.center
    @test h.radius ≈ hexp.radius

    # Box approximation of a 2D unit ball in the 1-norm
    b = Ball1(N[1, -2], to_N(N, 0.2))
    h = box_approximation_symmetric(b)
    hexp = Hyperrectangle(N[0, 0], to_N(N, [1.2, 2.2]))
    @test h.center ≈ hexp.center
    @test h.radius ≈ hexp.radius

    # Box approximation of a 3D unit ball in the 1-norm
    b = Ball1(N[1, 2, 0], to_N(N, 0.1))
    h = box_approximation_symmetric(b)
    hexp = Hyperrectangle(N[0, 0, 0], to_N(N, [1.1, 2.1, 0.1]))
    @test h.center ≈ hexp.center
    @test h.radius ≈ hexp.radius

    # Box approximation of a 4D hyperrectangle
    b = Hyperrectangle(to_N(N, [-1.5, -2.5, 2.4, -0.4]),
                       to_N(N, [0.1, 0.2, 0.3, 0.4]))
    h = box_approximation_symmetric(b)
    hexp = Hyperrectangle(zeros(N, 4), to_N(N, [1.6, 2.7, 2.7, 0.8]))
    @test h.center ≈ hexp.center
    @test h.radius ≈ hexp.radius

    # Testing alias symmetric_interval_hull
    h = symmetric_interval_hull(b)

    # empty set
    E = EmptySet{N}()
    @test box_approximation_symmetric(E) == E
end


module Determinism
export randBool, resetSeed

"""
seed used for the deterministic stochastic arithmetic
should be changed from one sample to another
"""
DSAseed = rand(Int64)

"""
changes the seed in order to run the computation and get a new result
"""
function resetSeed()
  global DSAseed = rand(Int64)
end

"""
Function that returns a boolean pseudo-randomly
garantees *determinism*: the same inputs will always produce the same output
and *diffusion*: all bits in the output have a 50% probability of being flipped if any bit of an input change
"""
function randBool(args...)
  global DSAseed # put global variable in scope
  pseudoRandomNumber = hash((DSAseed, args...)) # hash used as a random number
  return isodd(pseudoRandomNumber) # turn hash into a bool
end

end # module

# NOTES:
# diffusion is an important property in cryptographic cyphers (see paper "random as easy as 1, 2, 3" for concept and fast PRNG that have this property)
# means that a change of 1 bit in the input guarantees that all bits in the output have a 50% proba of changing
# splittable RNG paper (hash functions used as RNG): http://www.thesalmons.org/john/random123/papers/random123sc11.pdf


"""
Snape's Riddle - Author JK Rowling
Danger lies before you, while safety lies behind,
Two of us will help you, whichever you would find,
One among us seven will let you move ahead,
Another will transport the drinker back instead,
Two among our number hold only nettle wine,
Three of us are killers, waiting hidden in line.
Choose, unless you wish to stay here for evermore,
To help you in your choice, we give you these clues four:
First, however slyly the poison tries to hide
You will always find some on nettle wine’s left side;
Second, different are those who stand at either end,
But if you would move onwards, neither is your friend;
Third, as you see clearly, all are different size,
Neither dwarf nor giant holds death in their insides;
Fourth, the second left and the second on the right
Are twins once you taste them, though different at first sight.
"""
using AbstractLogic

# Codes:
# 1 is netted wine - NW
# 2 move ahead     - MA
# 3 move back      - MB
# 4 poison         - PO

# Seven Potions a,b,c,d,e,f,g each have 1 of 4 possible values
# a is left of b which is left of c etc.
hpset = logicalparse("a, b, c, d, e, f, g  ∈  NW, MA, MB, PO")

# Two among our number hold only nettle wine, #code 1
hpset = logicalparse("{{i}} == 'NW' {{2}}", logicset=hpset); showfeasible(hpset)
# Equivalent to logicalparse(["a|b|c|d|e|f|g == 1 {2}"], logicset=hpset)

# One among us seven will let you move ahead, #code 2
hpset = logicalparse("{{i}} == 'MA' {{1}}", logicset=hpset); showfeasible(hpset)
#hpset = logicalparse(["a|b|c|d|e|f|g == 2 {1}"], logicset=hpset)

#Another will transport the drinker back instead, #code 3
hpset = logicalparse("{{i}} == 'MB' {{1}}", logicset=hpset); showfeasible(hpset)

# Three of us are killers #code 4
hpset = logicalparse("{{i}} == 'PO' {{3}}", logicset=hpset); showfeasible(hpset)
# This last specification does not reduce the feasible set

# First, however slyly the poison tries to hide
# You will always find some on nettle wine’s left side;
# So if b is poison then a is nettle wine
# The i-1! indicates that outofbounds matches are considered "false"
hpset = logicalparse("{{i}} == 'NW' ==> {{i-1!}} == 'PO'", logicset=hpset); showfeasible(hpset)
# 60 possible outcomes

# Second, different are those who stand at either end,
hpset = logicalparse("a != g", logicset=hpset)

# But if you would move onwards, neither is your friend;
hpset = logicalparse("a,g != 'MA'", logicset=hpset); showfeasible(hpset)
# 30 possible outcomes

# Third, as you see clearly, all are different size,
# Neither dwarf nor giant holds death in their insides;
# c dwarf, f are giant
hpset = logicalparse("c,f != 'PO'", logicset=hpset); showfeasible(hpset)

# Fourth, the second left and the second on the right
# Are twins once you taste them, though different at first sight.
hpset = logicalparse("b == f", logicset=hpset)

search("='MA'", hpset) # c is a match, for the third bottle
search("='MB'", hpset) # g is a match, for the last bottle


import Base.LinAlg.BlasFloat

export eigvals_lanczos

function lanczos!{T}(K::KrylovSubspace{T})
    m = K.n
    αs = Array(T, m)
    βs = Array(T, m-1)
    for j=1:m-1
        w = nextvec(K)
        if j>1 w -= βs[j-1]*K.v[1] end
        w, y = orthogonalize(w, K, 1)
        αs[j] = y[1]
        βs[j] = convert(T, norm(w))
        append!(K, w/βs[j])
    end
    αs[m]= dot(nextvec(K), lastvec(K))
    SymTridiagonal(αs, βs)
end

function eigvals_lanczos(A, neigs::Int=size(A,1); tol::Real = size(A,1)^3*eps(real(eltype(A))), maxiter::Integer = size(A,1))
    history = ConvergenceHistory()
    history[:tol] = tol
    reserve!(history,:resnorm, maxiter)
    e1 = eiglancz_method!(history, A, neigs; maxiter=maxiter, tol=tol)
    e1, history
end

function eiglancz_method!(log::ConvergenceHistory, A, neigs::Int=size(A,1); tol::Real = size(A,1)^3*eps(real(eltype(A))), maxiter::Integer = size(A,1))
    K = KrylovSubspace(A, size(A, 1), 2) #In Lanczos, only remember the last two vectors
    initrand!(K)
    e1 = eigvals(lanczos!(K), 1:neigs)
    for iter=1:maxiter
        nextiter!(log,mvps=1)
        e0, e1 = e1, eigvals(lanczos!(K), 1:neigs)
        resnorm = norm(e1-e0)
        push!(log, :resnorm, resnorm)
        resnorm < tol && (setconv(log, resnorm>=0); break)
    end
    shrink!(log)
    e1
end


using ArnoldiMethod
#computes normalized cut cost for partition `cut`
function _normalized_cut_cost(cut, W::AbstractMatrix, D)
    cut_cost = 0
    for j in axes(W, 2)
        for i in axes(W, 1)
            if cut[i] != cut[j]
                cut_cost += W[i, j]
            end
        end
    end
    cut_cost /= 2
    return cut_cost / sum(D * cut) + cut_cost / sum(D * (.~cut))
end

function _normalized_cut_cost(cut, W::SparseMatrixCSC, D)
    cut_cost = 0
    rows = rowvals(W)
    vals = nonzeros(W)
    n = size(W, 2)
    for i = 1:n
        for j in nzrange(W, i)
            row = rows[j]
            if cut[i] != cut[row]
                cut_cost += vals[j] / 2
            end
        end
    end
    return cut_cost / sum(D * cut) + cut_cost / sum(D * (.~cut))
end

function _partition_weightmx(cut, W::AbstractMatrix)
    nv = length(cut)
    nv2 = sum(cut)
    nv1 = nv - nv2
    newvid = Vector{Int}(undef, nv)
    vmap1 = Vector{Int}(undef, nv1)
    vmap2 = Vector{Int}(undef, nv2)
    j1 = 1
    j2 = 1
    for i in eachindex(cut)
        if cut[i] == false
            newvid[i] = j1
            vmap1[j1] = i
            j1 += 1
        else
            newvid[i] = j2
            vmap2[j2] = i
            j2 += 1
        end
    end

    W1 = similar(W, (nv1, nv1))
    W2 = similar(W, (nv2, nv2))

    for j in axes(W, 2)
        for i in axes(W, 1)
            if cut[i] == cut[j] == false
                W1[newvid[i], newvid[j]] = W[i, j]
            elseif cut[i] == cut[j] == true
                W2[newvid[i], newvid[j]] = W[i, j]
            end
        end
    end

    return (W1, W2, vmap1, vmap2)
end

function _partition_weightmx(cut, W::SparseMatrixCSC)
    nv = length(cut)
    nv2 = sum(cut)
    nv1 = nv - nv2
    newvid = Vector{Int}(undef, nv)
    vmap1 = Vector{Int}(undef, nv1)
    vmap2 = Vector{Int}(undef, nv2)
    j1 = 1
    j2 = 1
    for i in eachindex(cut)
        if cut[i] == false
            newvid[i] = j1
            vmap1[j1] = i
            j1 += 1
        else
            newvid[i] = j2
            vmap2[j2] = i
            j2 += 1
        end
    end

    rows = rowvals(W)
    vals = nonzeros(W)
    I1 = Vector{Int}(); I2 = Vector{Int}()
    J1 = Vector{Int}(); J2 = Vector{Int}()
    V1 = Vector{Float64}(); V2 = Vector{Float64}()
    for i = 1:nv
        for j in nzrange(W, i)
            row = rows[j]
            if cut[i] == cut[row] == false
                push!(I1, newvid[i])
                push!(J1, newvid[row])
                push!(V1, vals[j])
            elseif cut[i] == cut[row] == true
                push!(I2, newvid[i])
                push!(J2, newvid[row])
                push!(V2, vals[j])
            end
        end
    end
    W1 = sparse(I1, J1, V1)
    W2 = sparse(I2, J2, V2)
    return (W1, W2, vmap1, vmap2)
end

function _recursive_normalized_cut(W, thres=thres, num_cuts=num_cuts)
    m, n = size(W)
    D = Diagonal(vec(sum(W, dims=2)))

    m == 1 && return [1]

    #get eigenvector corresponding to second smallest eigenvalue
    # v = eigs(D-W, D, nev=2, which=SR())[2][:,2]
    # At least some versions of ARPACK have a bug, this is a workaround
    invDroot = sqrt.(inv(D)) # equal to Cholesky factorization for diagonal D
    if n > 12
        λ, Q = eigs(invDroot' * (D - W) * invDroot, nev=12, which=SR())
        ret = real(Q[:,2])
    else
        ret = eigen(Matrix(invDroot' * (D - W) * invDroot)).vectors[:,2]
    end
    v = invDroot * ret

    #perform n-cuts with different partitions of v and find best one
    min_cost = Inf
    best_thres = -1
    for t in range(minimum(v), stop=maximum(v), length=num_cuts)
        cut = v .> t
        cost = _normalized_cut_cost(cut, W, D)
        if cost < min_cost
            min_cost = cost
            best_thres = t
        end
    end

    if min_cost < thres
        #split graph, compute normalized_cut for each subgraph recursively and merge indices.
        cut = v .> best_thres
        W1, W2, vmap1, vmap2 = _partition_weightmx(cut, W)
        labels1 = _recursive_normalized_cut(W1, thres, num_cuts)
        labels2 = _recursive_normalized_cut(W2, thres, num_cuts)

        labels = Vector{Int}(undef, m)
        offset = maximum(labels1)

        for i in eachindex(labels1)
            labels[vmap1[i]] = labels1[i]
        end
        for i in eachindex(labels2)
            labels[vmap2[i]] = labels2[i] + offset
        end

        return labels
    else
        return ones(Int, m)
    end
end

"""
    normalized_cut(g, thres, distmx=weights(g), [num_cuts=10]);

Perform [recursive two-way normalized graph-cut](https://en.wikipedia.org/wiki/Segmentation-based_object_categorization#Normalized_cuts)
on a graph, partitioning the vertices into disjoint sets.
Return a vector that contains the set index for each vertex.

It is important to identify a good threshold for your application. A bisection search over the range (0,1) will help determine a good value of thres.

### Keyword Arguments
- `thres`: Subgraphs aren't split if best normalized cut is above this threshold
- `num_cuts`: Number of cuts performed to determine optimal cut

### References
"Normalized Cuts and Image Segmentation" - Jianbo Shi and Jitendra Malik
"""
function normalized_cut(g::AbstractGraph,
    thres::Real,
    W::AbstractMatrix{T}=adjacency_matrix(g),
    num_cuts::Int=10) where T <: Real

    return _recursive_normalized_cut(W, thres, num_cuts)
end


function test_combination_combinations()
    @test collect(combinations(Int[], 0, Val{:unique})) == Array{Int,1}[[]]
    @test length(combinations(Int[], 0, Val{:unique})) == 1
    @test eltype(combinations(Int[], 0, Val{:unique})) == Array{Int,1}

    @test collect(combinations(Int[], 1, Val{:unique})) == Array{Int,1}[]
    @test length(combinations(Int[], 1, Val{:unique})) == 0
    @test eltype(combinations(Int[], 1, Val{:unique})) == Array{Int,1}

    @test collect(combinations(["a", "b", "c", "d"], 2, Val{:unique})) == Array[
        ["a", "b"],
        ["a", "c"],
        ["a", "d"],
        ["b", "c"],
        ["b", "d"],
        ["c", "d"],
    ]
    @test length(combinations(["a", "b", "c", "d"], 2, Val{:unique})) == 6
    @test eltype(combinations(["a", "b", "c", "d"], 2, Val{:unique})) == Array{String,1}

    @test collect(combinations(["a", "b"], 3, Val{:unique})) == Array{String,1}[]
    @test length(combinations(["a", "b"], 3, Val{:unique})) == 0
end

function test_combination_combinations_with_repetition()
    @test collect(combinations(Int[], 0, Val{:repeated})) == Array{Int,1}[[]]
    @test length(combinations(Int[], 0, Val{:repeated})) == 1
    @test eltype(combinations(Int[], 0, Val{:repeated})) == Array{Int,1}

    @test collect(combinations(Int[], 1, Val{:repeated})) == Array{Int,1}[]
    @test length(combinations(Int[], 1, Val{:repeated})) == 0
    @test eltype(combinations(Int[], 1, Val{:repeated})) == Array{Int,1}

    @test collect(combinations(["a", "b", "c", "d"], 2, Val{:repeated})) == Array[
        ["a", "a"],
        ["a", "b"],
        ["a", "c"],
        ["a", "d"],
        ["b", "b"],
        ["b", "c"],
        ["b", "d"],
        ["c", "c"],
        ["c", "d"],
        ["d", "d"],
    ]
    @test length(combinations(["a", "b", "c", "d"], 2, Val{:repeated})) == 10
    @test eltype(combinations(["a", "b", "c", "d"], 2, Val{:repeated})) == Array{String,1}

    @test collect(combinations(["a", "b"], 3, Val{:repeated})) == Array[
        ["a", "a", "a"],
        ["a", "a", "b"],
        ["a", "b", "b"],
        ["b", "b", "b"],
    ]
    @test length(combinations(["a", "b"], 3, Val{:repeated})) == 4
end

function test_combination_combinations_unknown_mode()
    @test_throws MethodError combinations(["a", "b", "c", "d"], 2, Val{:unknown})
end

function test_combination_all()
    print(rpad("Math.Combinatorics.Combination...", 50, ' '))

    test_combination_combinations()
    test_combination_combinations_with_repetition()
    test_combination_combinations_unknown_mode()

    println("PASS")
end


@testset "runge_kutta" begin
    function f(t, y, θ)
        return y
    end

    x0 = 1.0
    out = rk4(f, 1.0, 100, x0)
    @test abs(run(sess, out)[end]-exp(1))<1e-5
    out = ode45(f, 1.0, 100, x0)
    @test abs(run(sess, out)[end]-exp(1))<1e-5

    function f(t, y, θ)
        return y
    end

    x0 = [1.0;1.0]
    out = rk4(f, 1.0, 100, x0)
    @test norm(run(sess, out)[end,:] .- exp(1))<1e-5
    out = ode45(f, 1.0, 100, x0)
    @test norm(run(sess, out)[end,:] .- exp(1))<1e-5

    x0 = [1.0 1.0
          1.0 1.0]
    out = rk4(f, 1.0, 100, x0)
    @test norm(run(sess, out)[end,:,:] .- exp(1))<1e-5
    out = ode45(f, 1.0, 100, x0)
    @test norm(run(sess, out)[end,:,:] .- exp(1))<1e-5

end

@testset "alpha scheme" begin
    d0 = zeros(2)
    v0 = [2.;3.]
    a0 = zeros(2)
    Δt = 1e-3
    n = 9001
    Δt = Δt * ones(n)
    M = sparse([1.0 0.0;0.0 1.0])
    C = sparse(zeros(2,2))
    K = sparse([4.0 0.0;0.0 9.0])
    F = zeros(n, 2)

    d, v, a = αscheme(M, C, K, F, d0, v0, a0, Δt)
    d_, v_, a_ = run(sess, [d, v, a])

    tspan = 0:1e-3:(n-1)*1e-3
    @test norm(d_[1:end-1,1] - sin.(2tspan))<1e-3
    @test norm(d_[1:end-1,2] - sin.(3tspan))<1e-3
end

@testset "tr_bdf2" begin
    D0 = sparse([1.0 0.0;1.0 1.0])
    D1 = sparse([2.0 0.0; 0.0 1.0])
    ts = collect(LinRange(0, 1, 100))
    Δt = ts[2]-ts[1]
    F = zeros(199, 2)
    for i = 1:100
        t = ts[i]
        F[2*i-1, :] = [
           t^2 + 5t + 2
           t^2 + 2t + 1
        ]
        if i<100
            t = t + Δt/2
            F[2*i, :] = [
                t^2 + 5t + 2
                t^2 + 2t + 1
            ]
        end
    end
    td = TR_BDF2(D0, D1, Δt)
    @test td.symbolic == false
    u = td(zeros(2), F)

    @test abs(u[end, 1]-2.0)<1e-5
    @test abs(u[end, 2]-1.0)<1e-5

    td = constant(td)
    ua = td(zeros(2), F)
    @test run(sess, ua)≈u
end

@testset "ExplicitNewmark" begin

end


using MAT
using PyPlot
using JOLI
using GenSPGL
using SeisJOLI
using Arpack
using LinearAlgebra

###load the weighted algorithm function
include("../script/Weighted_LR.jl")

###load the true data
datafile = MAT.matopen("../data/Full.mat")
data_full = read(datafile,"Full_Data")
datafull = copy(data_full[:,1:355,1:355])
nff,ns,nr = size(datafull)

###loaded the jittered index for subsampling
sub_index = MAT.matopen("../data/ind.mat")
ind = read(sub_index,"ind")
close(sub_index)

###define operators subsampling and joMH(source-receiver => midpoint-offset)
RM = joKron(joRestriction(ns,dropdims(trunc.(Int,ind),dims=1);DDT = Complex{Float64}), joDirac(nr;DDT = Complex{Float64}))
MH = joSRtoCMO(nr,ns)

###define rank information if rank=r_sub(convertional weighted method)
rank = 85;
r_sub = 25;

###define the final result to save
Final_result = zeros(Complex{Float64},nr,ns,nff)
Final_SNR = zeros(Float64,nff,1)

###weighted LR to interpolate the data
for sliceNum = 31:304  #(round([7Hz,75Hz]/df+1=[30,304]))

    if sliceNum == 31
	# load unweighted result as prior information
	pre_int = MAT.matopen(string("../data/Idx_", string(sliceNum-1),".mat"))
	pre_L = read(pre_int,"L1")
	pre_R = read(pre_int,"R1")
	close(pre_int)

	# using limited-subspace weighted method
    	SNR, Result, L1, R1 = Weighted_LR(pre_L, pre_R, datafull[sliceNum,:,:], RM, MH, rank, r_sub)
    	println("Num=",sliceNum)
    else
	# load prior information(neighbor low frequency results)
	pre_int = MAT.matopen(string("../julia_result/Idx_", string(sliceNum-1),".mat"))
	pre_L = read(pre_int,"L1")
	pre_R = read(pre_int,"R1")
	close(pre_int)

	# using limited-subspace weighted method
	SNR, Result, L1, R1 = Weighted_LR(pre_L, pre_R, datafull[sliceNum,:,:], RM, MH, rank, r_sub)
    	println("Num=",sliceNum)
    end
    Final_SNR[sliceNum,1] = SNR;
    Final_result[:,:,sliceNum] = Result;

    # save each frequency result
    file = matopen(string("../julia_result/Idx_",string(sliceNum),".mat"), "w")
    write(file,string("Final_result",string(sliceNum)),Result)
    write(file,"SNR",SNR)
    write(file,"L1",L1)
    write(file,"R1",R1)
    close(file);
end

### save all limited-subspace weighted result
file = matopen("../julia_result/recursive_data.mat", "w")
write(file,"Final_result",Final_result)
write(file,"Final_SNR",Final_SNR)
close(file)


# Implements Candid Covariance-Free Incremental PCA algorithm of Weng et al:
# Weng, Juyang, Yilu Zhang, and Wey-Shiuan Hwang. "Candid covariance-free
# incremental principal component analysis." IEEE Transactions on Pattern
# Analysis and Machine Intelligence 25.8 (2003): 1034-1040.
#
# A comparison of online PCA algorithms in 2018 showed it to be a good
# trade-off in that it is fast and has low reconstruction error:
# H. Cardot and D. Degras, "Online Principal Component Analysis in
# High Dimension: Which Algorithm to Choose?", Int. Statistical Review, 2018-

"""
    CCIPCA(outdim::Int, indim; l::Int)

Online PCA with the CCIPCA (Candid Covariance-free Incremental PCA) algorithm,
where indim is the length of incoming vectors, outdim is the number of
dimension to project to, and l is the level of amnesia. Give values of l in
the range 2-4 if you want old vectors to be gradually less important, i.e.
latest vectors added get more weight.

If no indim is specified it will be set later, on first call to fit. After
that it is fixed and cannot change.

The CCIPCA is a very fast, simple, and online approximation of PCA. It can be used
for Dimensionality Reduction to project high-dimensional vectors into a
low-dimensional (typically 2D or 3D) space. This algorithm has shown very
good properties in comparative studies; it is both fast and give a good
approximation to (batch) PCA.

# Example

    o = CCIPCA(2, 10)                # Project 10-dimensional vectors into 2D
    u1 = rand(10)
    fit!(o, u1)                      # Fit to u1
    u2 = rand(10)
    fit!(o, u2)                      # Fit to u2
    u3 = rand(10)
    OnlineStats.transform(o, u3)     # Project u3 into PCA space fitted to u1 and u2 but don't change the projection
    u4 = rand(10)
    OnlineStats.fittransform!(o, u4) # Fit u4 and then project u4 into the space
    sort!(o)                         # Sort from high to low eigenvalues
    o[1]                             # Get primary (1st) eigenvector
    OnlineStats.variation(o)         # Get the variation (explained) "by" each eigenvector
"""
mutable struct CCIPCA <: OnlineStat{AbstractVector{<:Real}}
    lambda::Vector{Float64} # Eigenvalues, one per outdim
    U::Matrix{Float64}      # Eigenvectors, one row per indim, one column per outdim
    center::Vector{Float64} # Center/mean, one per indim
    xi::Vector{Float64}     # temp var used in calc, saved here for speed
    v::Vector{Float64}      # temp var used in calc, saved here for speed
    l::Int
    n::Int
    indimdet::Bool          # True iff indim has been determined (if not it is set on first fit)
end
function CCIPCA(outdim::Int, indim::Union{Nothing,Int} = nothing; l::Int=0)
    if isnothing(indim)
        indim = 1 # Dummy value, since we will update on first _fit
        indimdet = false
    elseif indim > 0
        @assert outdim < indim "indim ($indim) must be larger than outdim ($outdim)"
        indimdet = true
    else
        error("indim cannot take a value below 1, was set to $indim")
    end
    @assert l >= 0
    CCIPCA(setup_internal_vars(outdim, indim)..., l, 0, indimdet)
end
function setup_internal_vars(outdim::Int, indim::Int)
    lambda = zeros(Float64, outdim)
    U      = zeros(Float64, indim, outdim)
    center = zeros(Float64, indim)
    xi     = zeros(Float64, indim)
    v      = zeros(Float64, indim)
    return (lambda, U, center, xi, v)
end
Base.length(o::CCIPCA) = outdim(o) # Number of eigen-vectors
Base.getindex(o::CCIPCA, i) = o.U[:, i]
Base.lastindex(o::CCIPCA) = outdim(o)
Base.size(o::CCIPCA) = (indim(o), outdim(o))
@inline outdim(o::CCIPCA) = length(o.lambda)
@inline indim(o::CCIPCA) = length(o.center)
function _fit!(o::CCIPCA, x::Vector{Float64})
    if o.indimdet == false
        # Now indim is known so we set up accordingly
        idim = length(x)
        o.indimdet = true
        odim = size(o.U, 2)
        @assert odim < idim "indim ($idim) must be larger than outdim ($odim)"
        o.lambda, o.U, o.center, o.xi, o.v = setup_internal_vars(odim, idim)
    else
        @assert length(x) == OnlineStats.indim(o)
    end
    n = o.n + 1
    # update center with new observation:
    #o.center = (o.n * o.center .+ x)/n
    o.center += (x .- o.center) ./ n
    # center the new observation, unless this is the first observation:
    o.xi = (o.n > 0) ? (x .- o.center) : deepcopy(x)
    # Now recalc eigen-values and -vectors given the new observation:
    f = (1.0+o.l)/n
    @inbounds for i in 1:outdim(o)
        if i == n
            o.lambda[i] = norm(o.xi)
            o.U[:, i] = o.xi / o.lambda[i]
            break
        end
        o.v = (1-f) * o.lambda[i] * o.U[:, i] + f * dot(o.U[:, i], o.xi) * o.xi
        o.lambda[i] = norm(o.v)
        o.U[:, i] = o.v/o.lambda[i]
        o.xi = o.xi .- (dot(o.U[:, i], o.xi) * o.U[:, i])
    end
    o.n = n
end
"""
    transform(o::CCIPCA, u::AbstractArray{Float64})

Transform (i.e. project) the vector `u` into the PCA space
represented by `o`.
"""
function transform(o::CCIPCA, u::AbstractArray{Float64})
    @assert indim(o) == length(u)
    (u .- o.center)' * o.U
end
# We can also call the CCIPCA object itself to transform:
(o::CCIPCA)(u::AbstractArray{Float64}) = transform(o, u)
"""
    reconstruct(o::CCIPCA, uproj::AbstractArray{Float64})

Reconstruct the (projected) vector `uproj` back to the original
space from which `o` has been fitted.
"""
function reconstruct(o::CCIPCA, uproj::AbstractArray{Float64})
    @assert outdim(o) == length(uproj)
    o.center .+ (o.U * uproj')
end
"""
    fittransform!(o::CCIPCA, u::Vector{Float64})

First `fit!` and then `transform` the vector `u` into the PCA
space represented by `o`.
"""
function fittransform!(o::CCIPCA, u::Vector{Float64})
    _fit!(o, u)
    transform(o, u)
end
"""
    eigenvalue(o::CCIPCA, i::Int)

Get the `i`th eigenvalue of `o`. Also called principal variance
for PCA.
"""
function eigenvalue(o::CCIPCA, i::Int)
    @assert 1 <= i <= outdim(o) "This CCIPCA has $(outdim(o)) eigenvalues, cannot return eigenvalue $(i)!"
    o.lambda[i]
end
"""
    eigenvalue(o::CCIPCA, i::Int)

Get the `i`th eigenvalue of `o`. Also called principal variance
for PCA.
"""
principalvar(o::CCIPCA, i::Int) = eigenvalue(o, i)
"""
    eigenvector(o::CCIPCA, i::Int)

Get the `i`th eigenvector of `o`.
"""
function eigenvector(o::CCIPCA, i::Int)
    @assert 1 <= i <= outdim(o) "This CCIPCA has $(outdim(o)) eigenvectors, cannot return eigenvector $(i)!"
    o.U[:, i]
end
eigenvalues(o::CCIPCA) = Float64[eigenvalue(o, i) for i in 1:outdim(o)]
principalvars(o::CCIPCA) = eigenvalues(o)
"""
    relativevariances(o::CCIPCA)

Get the relative variance (explained) in the direction of each
eigenvector. Returns a vector of zeros if no vectors have yet been fitted.
Note that this does not inclue the residual variance that is not captured
in the eigenvectors.
"""
function relativevariances(o::CCIPCA)
    if o.n == 0
        return zeros(Float64, outdim(o))
    else
        return o.lambda ./ sum(o.lambda)
    end
end
"""
    sort!(o::CCIPCA)

Sort eigenvalues and their eigenvectors of `o` so highest ones come first.
Useful before visualising since it ensures most variation is on the
first (X) axis.
"""
function Base.sort!(o::CCIPCA)
    perm = sortperm(o.lambda, rev=true)
    o.lambda = o.lambda[perm]
    o.U = o.U[:, perm]
end


""" Simulate heteroscedastic data with known errors. """
function sim_data(F_X::Distribution, F_Y::Distribution,
                  σ_X_distr::Distribution, σ_Y_distr::Distribution,
                  n_X::Int, n_Y::Int)
    σ_X = rand(σ_X_distr, n_X)
    σ_Y = rand(σ_Y_distr, n_Y)
    X = rand(F_X, n_X)
    Y = rand(F_Y, n_Y)
    ϵX = σ_X .* randn(n_X)
    ϵY = σ_Y .* randn(n_Y)
    Xtilde = X .+ ϵX
    Ytilde = Y .+ ϵY
    return Dict(
        # inputs
        :F_X => F_X,
        :F_Y => F_Y,
        :σ_X_distr => σ_X_distr,
        :σ_Y_distr => σ_Y_distr,
        :n_X => n_X,
        :n_Y => n_Y,
        # outputs
        :σ_X => σ_X,
        :σ_Y => σ_Y,
        :X => X,
        :Y => Y,
        :ϵX => ϵX,
        :ϵY => ϵY,
        :Xtilde => Xtilde,
        :Ytilde => Ytilde,
    )
end

function sim_homo(F_X::Distribution, F_Y::Distribution,
                             σ::Real, n_X::Int, n_Y::Int)
    X = rand(F_X, n_X)
    Y = rand(F_Y, n_Y)
    ϵX = σ .* randn(n_X)
    ϵY = σ .* randn(n_Y)
    Xtilde = X .+ ϵX
    Ytilde = Y .+ ϵY

    σ_X = ones(n_X) .* σ
    σ_Y = ones(n_Y) .* σ

    return Dict(
        # inputs
        :F_X => F_X,
        :F_Y => F_Y,
        :σ => σ,
        :n_X => n_X,
        :n_Y => n_Y,
        # outputs
        :σ_X => σ_X,
        :σ_Y => σ_Y,
        :X => X,
        :Y => Y,
        :ϵX => ϵX,
        :ϵY => ϵY,
        :Xtilde => Xtilde,
        :Ytilde => Ytilde,
    )
end


using ScikitLearn
using ScikitLearn.GridSearch
using PyPlot
using GaussianProcesses: GPE, MeanZero, SE
using Random
import Random
Random.seed!(42)
# Training d:w
n = 10
x = 2π * rand(n, 1)
y = sin.(x[:, 1]) + 0.05*randn(n)

# Select mean and covariance function
mZero = MeanZero()                   # Zero mean function
kern = SE(0.0,0.0)                   # Squared exponential kernel with parameters
                                     # log(ℓ) = 0.0, log(σ) = 0.0
gp = fit!(GPE(m=mZero,k=kern, logNoise=-1.0), x,y);

gp_cv = fit!(GridSearchCV(GPE(m=mZero,k=SE(0.0,0.0)), Dict(:logNoise=>collect(-10:0.3:10), :k_lσ=>collect(0:0.1:5))), x, y);
best_gp = gp_cv.best_estimator_;
@show get_params(best_gp)[:logNoise] get_params(best_gp)[:k_lσ]
nothing

xx = -5:0.1:10
plot(xx, predict(gp, reshape(collect(xx), length(xx), 1)), label="hand-specified")
plot(xx, predict(best_gp, reshape(collect(xx), length(xx), 1)), label="gridsearch-optimized")
plot(x, y, "bo")
legend();

get_params(gp)


using jInv.Utils
using Base.Test
using OED
using jInv.ForwardShare
using jInv.Mesh

n = 100
nrhs = 2
# A = abs(sprandn(N,n,.2))
A = sdiag(1./linspace(1,n,n).^4)
A = A[1:end-2,:]
N = size(A,1)
x = rand(n,nrhs); x[4:7,:] = 1;
D = A*x  + 0* randn(N,nrhs)/1e3
OEDparamA = OEDTypeAParam(A,D,zeros(0,0))

# setup OED type B param
r = 20;
l = 8;
Mimg  = getRegularMesh([-1 1 -1 1.],[10,10])
Mgrid = getRegularMesh([-1.2 2.4 -1.6 1.4],[20,20])
OEDparamB = getOEDTypeBParam(r,l,x,Mimg,Mgrid,noiseLevel=0.)

OEDs = (OEDparamA, OEDparamB)

for k=1:2
	display(OEDs[k])

	L = 1*speye(n)
	Ce = ones(1,n)
	ce = Ce*x
	pFor = EqconstrainedOEDParam(OEDs[k], L,Ce,ce, zeros(0,0), [])

	println("--- solve LS problem ---")
	display(pFor)
	np  =  getNumberOfDesignParameters(OEDs[k])
	p0      = rand(np)
	xr, = getData(p0,pFor)
	@test norm(Ce*xr-ce) < 1e-12
	println("rel.error: $(norm(xr-x)/norm(x))")

	println("--- test derivative of LS solution w.r.t. design ---")
	pass, = checkDerivative(rand(np),pFor)
	@test pass

	println("--- adjoint test ---")
	nw,nv = getSensMatSize(pFor)
	v = randn(nv)
	w = randn(nw)
	t1 = dot(w,getSensMatVec(v,p0,pFor))
	t2 = dot(v,getSensTMatVec(w,p0,pFor))
	@test abs(t1-t2)/abs(t1) < 1e-10
end


using Images, ShapeFromShading
using Test

@testset "ShapeFromShading" begin
    include("deterministic.jl")
    include("deterministic_custom.jl")
    include("discreteshape.jl")
    include("discreteshapebound.jl")
    include("multiresolutionhybrid.jl")
    include("pentland.jl")
    include("photometric.jl")
    include("shah.jl")
    # include("SimulatedAnnealing.jl")
end

@testset "Anisotropicdiffusion" begin
    include("anisotropicdiffusion.jl")
end
@testset "Durou" begin
    include("durou.jl")
end
@testset "Frankot" begin
    include("frankot.jl")
end
@testset "Horn" begin
    include("horn.jl")
end
@testset "Mumfordshah" begin
    include("mumfordshah.jl")
end
@testset "Nonconvex1" begin
    include("nonconvex1.jl")
end
@testset "Path" begin
    include("path.jl")
end
@testset "Splitpath" begin
    include("splitpath.jl")
end
@testset "Totalvariation" begin
    include("totalvariation.jl")
end
@testset "Quadratic" begin
    include("quadratic.jl")
end

@testset "SyntheticSurface" begin
    include("cake.jl")
    include("cake2.jl")
    include("dem.jl")
    include("prism.jl")
    include("ripple.jl")
    include("ripple2.jl")
    include("supergaussian.jl")
    include("synthsphere.jl")
    include("synthvase.jl")
    include("synthgaussian.jl")
    include("tent.jl")
end

@testset "ImageGeneration" begin
    include("estimatealbedo.jl")
    include("generate_surface.jl")
end

@testset "OtherFunctions" begin
    include("mask.jl")
end

@testset "Benchmarks" begin
    include("benchmark_shape.jl")
    include("benchmark_integration.jl")
    include("benchmark_normals.jl")
    include("benchmark_image.jl")
end


# Tests for Profit Additive DEA Models
@testset "ProfitAdditiveDEAModel" begin

    # Test using Book data
    X = [2; 4; 8; 12; 6; 14; 14; 9.412];
    Y = [1; 5; 8; 9; 3; 7; 9; 2.353];
    W = [1; 1; 1; 1; 1; 1; 1; 1];
    P = [2; 2; 2; 2; 2; 2; 2; 2];

    # With :Ones
    deaprofitaddones = deaprofitadd(X, Y, W, P, :Ones)

    @test typeof(deaprofitaddones) == ProfitAdditiveDEAModel

    @test nobs(deaprofitaddones) == 8
    @test ninputs(deaprofitaddones) == 1
    @test noutputs(deaprofitaddones) == 1
    @test ismonetary(deaprofitaddones) == false

    @test efficiency(deaprofitaddones) == efficiency(deaprofitaddones, :Economic)
    @test efficiency(deaprofitaddones, :Economic)   ≈ [8.0; 2.0; 0.0; 2.0; 8.0; 8.0; 4.0; 12.706]
    @test efficiency(deaprofitaddones, :Technical)  ≈ [0.0; 0.0; 0.0; 0.0; 4.0; 22/3; 2.0; 8059/1000]
    @test efficiency(deaprofitaddones, :Allocative) ≈ [8.0; 2.0; 0.0; 2.0; 4.0; 2/3; 2.0; 4647/1000]
    @test normfactor(deaprofitaddones) ≈ [1.0; 1.0; 1.0; 1.0; 1.0; 1.0; 1.0; 1.0]

    # Default is :Ones
    @test efficiency(deaprofitadd(X, Y, W, P)) == efficiency(deaprofitaddones)

    # With :MIP
    deaprofitaddmip = deaprofitadd(X, Y, W, P, :MIP)

    @test efficiency(deaprofitaddmip, :Economic)   ≈ [4.0; 0.5; 0.0; 1/6; 4/3; 4/7; 2/7; 2.699958] atol = 1e-5
    @test efficiency(deaprofitaddmip, :Technical)  ≈ [0.0; 0.0; 0.0; 0.0; 7/6; 4/7; 1/7; 2.549936] atol = 1e-5
    @test efficiency(deaprofitaddmip, :Allocative) ≈ [4.0; 0.5; 0.0; 1/6; 1/6; 0.0; 1/7; 0.150021] atol = 1e-5
    @test normfactor(deaprofitaddmip) ≈ [2.0; 4.0; 8.0; 12.0; 6.0; 14.0; 14.0; 4.706]

    # With :MIP and monetary = true
    deaprofitaddmipmon = deaprofitadd(X, Y, W, P, :MIP, monetary = true)

    @test ismonetary(deaprofitaddmipmon) == true

    @test efficiency(deaprofitaddmipmon) ≈ efficiency(deaprofitaddmip) .* normfactor(deaprofitaddmip)
    @test efficiency(deaprofitaddmipmon, :Technical) ≈ efficiency(deaprofitaddmip, :Technical) .* normfactor(deaprofitaddmip)
    @test efficiency(deaprofitaddmipmon, :Allocative) ≈ efficiency(deaprofitaddmip, :Allocative) .* normfactor(deaprofitaddmip)

    # With custom weights
    deaprofitaddcustom = deaprofitadd(X, Y, W, P, rhoX = 1 ./ X, rhoY = 1 ./ Y)

    @test efficiency(deaprofitaddcustom) ≈ efficiency(deaprofitaddmip)

    # Print
    show(IOBuffer(), deaprofitaddones)

    # Test errors
    @test_throws DimensionMismatch deaprofitadd([1; 2 ; 3], [4; 5], [1; 1; 1], [1; 1], :Ones) #  Different number of observations
    @test_throws DimensionMismatch deaprofitadd([1; 2; 3], [4; 5; 6], [1; 2; 3; 4], [1; 1; 1], :Ones) # Different number of observation in input prices
    @test_throws DimensionMismatch deaprofitadd([1; 2; 3], [4; 5; 6], [1; 1; 1], [1; 2; 3; 4], :Ones) # Different number of observation in output prices
    @test_throws DimensionMismatch deaprofitadd([1 1; 2 2; 3 3 ], [4; 5; 6], [1 1 1; 2 2 2; 3 3 3], [1; 1; 1], :Ones) # Different number of input prices and inputs
    @test_throws DimensionMismatch deaprofitadd([1; 2; 3 ], [4; 5; 6], [1; 1; 1], [1 1 1; 2 2 2; 3 3 3], :Ones) # Different number of output prices and outputs
    @test_throws DimensionMismatch deaprofitadd([1; 2; 3] , [4; 5; 6], [1; 1; 1], [1; 1; 1], rhoX = [1 1; 2 2; 3 3], rhoY = [1; 1; 1]) # Different size of inputs weights
    @test_throws DimensionMismatch deaprofitadd([1; 2; 3] , [4; 5; 6], [1; 1; 1], [1; 1; 1], rhoX = [1; 1; 1], rhoY = [1 1; 2 2; 3 3]) # Different size of output weights
    @test_throws ArgumentError deaprofitadd([1; 2; 3], [1; 2; 3], [1; 1; 1], [1; 1; 1], :Error) # Invalid inuts weights
    @test_throws ArgumentError deaprofitadd([1; 2; 3], [1; 2; 3], [1; 1; 1], [1; 1; 1], :Ones, rhoX= [1; 1; 1], rhoY = [1; 1; 1]) # Weights not allowed if model != :Custom

end


co9(x) = x == 9  ? 0 :
         1<=x<=8 ? x :
         co9(sum(digits(x)))


### A Pluto.jl notebook ###
# v0.16.4

using Markdown
using InteractiveUtils

# ╔═╡ 9edb35ce-4ea7-11eb-353f-53459c337880
begin
	import Pkg
	Pkg.activate(Base.current_project())
	using ReinforcementLearning
	using Statistics
	using Flux
	using Plots
end

# ╔═╡ 6ce0a200-4ea7-11eb-123c-73e3a260351c
md"""
# Chapter 6.7 Maximization Bias and Double Learning

In example 6.7, authors introduced a MDP problem to compare the different performance of **Q-Learning** and **Double-Q-Learning**. This environment is kind of special compared to the environments we have seen before. In the first step, only `LEFT` and `RIGHT` are allowed. In the second step, if the `LEFT` is chosen previously, then we have 10 valid actions. We call this kind of environment is of `FULL_ACTION_SET`.
"""

# ╔═╡ 172236ac-4ea8-11eb-1393-5bc3683bf6b0
begin
	"""
	states:
	1:A
	2:B
	3:terminal
	actions:
	1: left
	2: right
	"""
	Base.@kwdef mutable struct MaximizationBiasEnv <: AbstractEnv
		position::Int = 1
		reward::Float64 = 0.0
	end
	function (env::MaximizationBiasEnv)(a::Int)
		if env.position == 1
			if a == 1
				env.position = 2
				env.reward = 0.0
			else
				env.position = 3
				env.reward = 0.0
			end
		elseif env.position == 2
			env.position = 3
			env.reward = randn() - 0.1
		end
		nothing
	end
end

# ╔═╡ 19876a70-4ea8-11eb-2b70-9d48875f7ae7
RLBase.state_space(env::MaximizationBiasEnv) = Base.OneTo(3)

# ╔═╡ 247ada66-4ea8-11eb-3c0a-af131b73e9f3
RLBase.action_space(env::MaximizationBiasEnv) = Base.OneTo(10)

# ╔═╡ 26b9f6f2-4ea8-11eb-38f3-3f1822c4acd2
RLBase.ActionStyle(env::MaximizationBiasEnv) = FULL_ACTION_SET

# ╔═╡ 296e0ab6-4ea8-11eb-33c5-71be241708b1
const LEFT = 1

# ╔═╡ 2c06d25a-4ea8-11eb-26e3-01d81d714859
const RIGHT = 2

# ╔═╡ 2e11adfc-4ea8-11eb-02bc-d550994774f8
function RLBase.legal_action_space(env::MaximizationBiasEnv)
    if env.position == 1
        (LEFT, RIGHT)
    else
        Base.OneTo(10)
    end
end

# ╔═╡ 3148b560-4ea8-11eb-0fc1-f77efb282ba4
function RLBase.legal_action_space_mask(env::MaximizationBiasEnv)
    m = fill(false, 10)
    if env.position == 1
        m[LEFT] = true
        m[RIGHT] = true
    else
        m .= true
    end
    m
end

# ╔═╡ 3a674788-4ea8-11eb-1931-b39142fffb47
function RLBase.reset!(env::MaximizationBiasEnv)
    env.position = 1
    env.reward = 0.0
    nothing
end

# ╔═╡ 3ec32b6c-4ea8-11eb-1575-319c528b4f02
RLBase.reward(env::MaximizationBiasEnv) = env.reward

# ╔═╡ 4127acde-4ea8-11eb-0e47-7d82bb5cb310
RLBase.is_terminated(env::MaximizationBiasEnv) = env.position == 3

# ╔═╡ 4309b632-4ea8-11eb-1d46-edd93efd6925
RLBase.state(env::MaximizationBiasEnv) = env.position

# ╔═╡ 45650d82-4ea8-11eb-39c2-71fc325d6c13
md"""
Now the environment is well defined.
"""

# ╔═╡ 56f8ba1c-4ea8-11eb-2ce7-8357675ee6ef
world = MaximizationBiasEnv()

# ╔═╡ 657eca7c-4ea8-11eb-3c42-b1a3085f6129
NS, NA = length(state_space(world)), length(action_space(world))

# ╔═╡ 8d15e872-4ea8-11eb-2fc2-7f3c4a90a85a
md"""
To calculate the percentage of chosing `LEFT` action in the first step, we'll create customized hook here:
"""

# ╔═╡ 6f798d14-4ea8-11eb-2825-8f6f17ff4f05
begin
	Base.@kwdef mutable struct CountOfLeft <: AbstractHook
		counts::Vector{Bool} = []
	end
	function (f::CountOfLeft)(::PreActStage, agent, env, action)
		if state(env) == 1
			push!(f.counts, action == LEFT)
		end
	end
end

# ╔═╡ d026cbb8-4ea8-11eb-29a2-dd75efe64466
md"""
Next we create two agent factories, Q-Learning and Double-Q-Learning.
"""

# ╔═╡ ea42c3c6-4ea8-11eb-1d82-0566785eeef3
create_double_Q_agent() = Agent(
    policy=QBasedPolicy(
        learner=DoubleLearner(
            L1=TDLearner(
                approximator=TabularQApproximator(
                    n_state=NS,
                    n_action=NA,
                    opt=Descent(0.1),
                ),
                method=:SARS,
                γ=1.,
                n=0,
            ),
            L2=TDLearner(
                approximator=TabularQApproximator(;
                    n_state=NS,
                    n_action=NA,
                    opt=Descent(0.1),
                ),
                method=:SARS,
                γ=1.,
                n=0,
            ),
            ),
        explorer=EpsilonGreedyExplorer(0.1;is_break_tie=true)
    ),
    trajectory=VectorSARTTrajectory()
)

# ╔═╡ eccb09aa-4ea8-11eb-08a0-2f6344c38669
create_Q_agent() = Agent(
    policy=QBasedPolicy(
        learner=TDLearner(
            approximator=TabularQApproximator(;n_state=NS, n_action=NA, opt=Descent(0.1)),
            method=:SARS,
            γ=1.0,
            n=0
            ),
        explorer=EpsilonGreedyExplorer(0.1;is_break_tie=true)
        ),
    trajectory=VectorSARTTrajectory()
)

# ╔═╡ 0e4b6994-4ea9-11eb-111f-7b37d6a716e5
begin
	DQ_stats = []
	for _ in 1:1000
		hook = CountOfLeft()
		run(create_double_Q_agent(), world, StopAfterEpisode(300),hook)
		push!(DQ_stats, hook.counts)
	end
	plot(mean(DQ_stats)*100, legend=:topright, label="double q", xlabel="Episodes", ylabel="% left actions from A")

	Q_stats = []
	for _ in 1:1000
		hook = CountOfLeft()
		run(create_Q_agent(), world, StopAfterEpisode(300),hook)
		push!(Q_stats, hook.counts)
	end
	plot!(mean(Q_stats)*100, legend=:topright, label="q")
	hline!([5], linestyle=:dash, label="optimal")
end

# ╔═╡ Cell order:
# ╟─6ce0a200-4ea7-11eb-123c-73e3a260351c
# ╠═9edb35ce-4ea7-11eb-353f-53459c337880
# ╠═172236ac-4ea8-11eb-1393-5bc3683bf6b0
# ╠═19876a70-4ea8-11eb-2b70-9d48875f7ae7
# ╠═247ada66-4ea8-11eb-3c0a-af131b73e9f3
# ╠═26b9f6f2-4ea8-11eb-38f3-3f1822c4acd2
# ╠═296e0ab6-4ea8-11eb-33c5-71be241708b1
# ╠═2c06d25a-4ea8-11eb-26e3-01d81d714859
# ╠═2e11adfc-4ea8-11eb-02bc-d550994774f8
# ╠═3148b560-4ea8-11eb-0fc1-f77efb282ba4
# ╠═3a674788-4ea8-11eb-1931-b39142fffb47
# ╠═3ec32b6c-4ea8-11eb-1575-319c528b4f02
# ╠═4127acde-4ea8-11eb-0e47-7d82bb5cb310
# ╠═4309b632-4ea8-11eb-1d46-edd93efd6925
# ╟─45650d82-4ea8-11eb-39c2-71fc325d6c13
# ╠═56f8ba1c-4ea8-11eb-2ce7-8357675ee6ef
# ╠═657eca7c-4ea8-11eb-3c42-b1a3085f6129
# ╟─8d15e872-4ea8-11eb-2fc2-7f3c4a90a85a
# ╠═6f798d14-4ea8-11eb-2825-8f6f17ff4f05
# ╟─d026cbb8-4ea8-11eb-29a2-dd75efe64466
# ╠═ea42c3c6-4ea8-11eb-1d82-0566785eeef3
# ╠═eccb09aa-4ea8-11eb-08a0-2f6344c38669
# ╠═0e4b6994-4ea9-11eb-111f-7b37d6a716e5