using Base
using MatlabCompat
using MatlabCompat.Io
using MatlabCompat.ImageTools
import MatlabCompat.MathTools
import MatlabCompat.MathTools: max
using Base.Test
using Images
using FixedPointNumbers
using Tk

const TEST_FOLDER = dirname(@__FILE__)
const JANUS_PATH = joinpath(TEST_FOLDER, "janus.m");
const TESTMAT_PATH = joinpath(TEST_FOLDER, "test.mat");

# basic test
println("foo test")
x = function ()
    return 1 == 1
end
@test x()

################################
#     ImageTools Module Tests
################################

# make a test image
array = [0.1 0.001 0 0; 0.1 0.401 0 0; 0.1 0.01 0 0; 0.1 0.901 0 0.6];
img = grayim(array);
properties(img)["IMcs"] = "Gray";

# test the fuctions
println("setting true tests results");
trueThreshold = 0.11100752480522075;
trueBWSum= 3;

# println("testing graythresh()");
# @test graythresh(img) == trueThreshold;

println("testing im2bw()");
@test sum(reinterpret(Float32, float32(im2bw(img, trueThreshold)))[:]) == trueBWSum;

println("testing imread()");
img2 = MatlabCompat.imread("http://matlabcompat.github.io/img/example.tif");
expected = nothing
if VERSION < v"0.4-"
    expected = "Image{Gray{UfixedBase{Uint8,8}},2,Array{Gray{UfixedBase{Uint8,8}},2}}";
else
    expected = "Images.Image{ColorTypes.Gray{FixedPointNumbers.UFixed{UInt8,8}},2,Array{ColorTypes.Gray{FixedPointNumbers.UFixed{UInt8,8}},2}}"
	end
@test "$(typeof(img2))" == expected;

################################
#     Support Module Tests
################################

println("testing im2mat()");
mat = im2mat(img)
@test "$(typeof(mat))" == "Array{Float32,2}";

println("testing mat2im()");
expected = nothing
if (VERSION < v"0.4-")
    expected = "Image{Float64,2,Array{Float64,2}}"
else
    expected = "Images.Image{Float64,2,Array{Float64,2}}"
end
@test "$(typeof(mat2im(img)))" == expected;

println("testing rosetta()");
@test "$(typeof(rosetta(JANUS_PATH)))" == "Array{UTF8String,1}";

################################
#     Io Module Tests
################################

println("testing load()");
println("$(typeof(MatlabCompat.load(TESTMAT_PATH)))")
@test "$(typeof(MatlabCompat.load(TESTMAT_PATH)))" == "Dict{ASCIIString,Any}";
# "$(typeof(A))" == "Dict{Any,Any}"

################################
#     StringTools Module Tests
################################
println("testing disp()");
@test disp("test") == true;

println("testing num2str()");
@test num2str(3) == "3";

println("testing strcat()");
@test strcat("1", "2", "3") == "123";

################################
#     MathTools Module Tests
################################
matrix = [1 2 3];

println("testing numel()");
@test numel(matrix[:]) == 3;

println("testing max()");

@test max(matrix) == maximum(matrix);


mutable struct DE <: AbstractParameters
    N::Int
    F::Float64
    CR::Float64
    CR_min::Float64
    CR_max::Float64
    F_min::Float64
    F_max::Float64
    strategy::Symbol
end

"""
    DE(;
        N  = 0,
        F  = 1.0,
        CR = 0.9,
        strategy = :rand1,
        information = Information(),
        options = Options()
    )

Parameters for Differential Evolution (DE) algorithm: step-size `F`,`CR` controlls the binomial
crossover, `N` is the population size. The parameter `trategy` is related to the variation
operator (`:rand1`, `:rand2`, `:best1`, `:best2`, `:randToBest1`).

# Example

```jldoctest
julia> f(x) = sum(x.^2)
f (generic function with 1 method)

julia> optimize(f, [-1 -1 -1; 1 1 1.0], DE())
+=========== RESULT ==========+
  iteration: 1000
    minimum: 0
  minimizer: [0.0, 0.0, 0.0]
    f calls: 30000
 total time: 0.0437 s
+============================+

julia> optimize(f, [-1 -1 -1; 1 1 1.0], DE(N=50, F=1.5, CR=0.8))
+=========== RESULT ==========+
  iteration: 600
    minimum: 8.68798e-25
  minimizer: [3.2777877981303293e-13, 3.7650459509488005e-13, -7.871487597385812e-13]
    f calls: 30000
 total time: 0.0319 s
+============================+
```

"""
function DE(;
    N::Int = 0,
    F = 1.0,
    CR = 0.9,
    CR_min = CR,
    CR_max = CR,
    F_min = F,
    F_max = F,
    strategy::Symbol = :rand1,
    information = Information(),
    options = Options(),
)

    parameters =
        DE(N, promote(F, CR, CR_min, CR_max, F_min, F_max)..., strategy)

    Algorithm(
        parameters,
        information = information,
        options = options,
    )

end

function update_state!(
        status,
        parameters::DE,
        problem::AbstractProblem,
        information::Information,
        options::Options,
        args...;
        kargs...
)
    population = status.population
    current_pop = copy(population)

    F = parameters.F
    CR = parameters.CR

    # stepsize
    if parameters.F_min < parameters.F_max
        F = parameters.F_min + (F_max - parameters.F_min) * rand()
    end

    if parameters.CR_min < parameters.CR_max
        CR =
            parameters.CR_min + (parameters.CR_max - parameters.CR_min) * rand()
    end

    N = parameters.N
    strategy = parameters.strategy

    xBest = get_position(status.best_sol)
    best_ind = 1

    for i in 1:N
        x = get_position(current_pop[i])

        u = DE_mutation(current_pop, F, strategy, best_ind)
        v = DE_crossover(x, u, CR)

        # instance child
        v = evo_boundary_repairer!(v, xBest, problem.bounds)
        h = create_solution(v, problem,ε=options.h_tol)

        # select survivals
        if is_better(h, current_pop[i])
            population[i] = h

            if is_better(h, status.best_sol)
                status.best_sol = h
                best_ind = i
            end
        end

        stop_criteria!(status, parameters, problem, information, options)
        if status.stop
            break
        end
    end

end

function initialize!(
        status,
        parameters::DE,
        problem::AbstractProblem,
        information::Information,
        options::Options,
        args...;
        kargs...
)
    D = size(problem.bounds, 2)

    if parameters.N <= 5
        parameters.N = 10 * D
    end

    if parameters.CR < 0 || parameters.CR > 1
        parameters.CR = 0.5
        options.debug &&
            @warn("CR should be from interval [0,1]; set to default value 0.5")
    end

    if options.f_calls_limit == 0
        options.f_calls_limit = 10000D
        options.debug &&
            @warn( "f_calls_limit increased to $(options.f_calls_limit)")
    end

    if options.iterations == 0
        options.iterations = div(options.f_calls_limit, parameters.N) + 1
    end

    return gen_initial_state(problem,parameters,information,options,status)
end

function final_stage!(
        status,
        parameters::DE,
        problem::AbstractProblem,
        information::Information,
        options::Options,
        args...;
        kargs...
    )
    status.final_time = time()
end


@testset "GR4J.jl" begin
    @testset "Constructors" begin
		model = GR4J()
		model2 = GR4J(model)
		@test model == model2
	end

    @testset "Check param or state values when creating new instance" begin
        @test_throws DomainError GR4J(-350.0, 0.0, 40.0, 0.5)
        @test_throws DomainError GR4J(350.0, 0.0, -40.0, 0.5)
        @test_throws DomainError GR4J(350.0, 0.0, 40.0, -0.5)
        @test_throws DomainError GR4J(350.0, 0.0, 40.0, 0.5, 351.0, 30.0)
        @test_throws DomainError GR4J(350.0, 0.0, 40.0, 0.5, 0.0, 41.0)
    end

    @testset "Check values when a parameter or state value is set" begin
        model = GR4J()
        @test_throws DomainError model.X1 = -350.0
        @test_throws DomainError model.X3 = -0.1
        @test_throws DomainError model.X4 = -0.1
        @test_throws DomainError model.Sp = 500.0
        @test_throws DomainError model.Sr = 50.0
    end

    @testset "Stash and recover a model for hot start" begin
        model = GR4J()
        d = stash(model)
        model2 = GR4J(d)
        @test model == model2
    end

    @testset "Print a model" begin
    	model = GR4J()
    	println(model)
	end

    @testset "Test one step running" begin
        model = GR4J()
        Q1, AET1 = simulate(model, 10.0, 5.0)
        Q2, AET2 = simulate(model, 0.0, 5.0)

        model = GR4J()
        Qs, AETs = simulate(model, [10.0, 0.0], [5.0, 5.0])

        @test Qs[2] == Q2
        @test AETs[2] == AET2
    end

    @testset "Check againt GR4J MS Excel modelling results" begin
        data = dataset("gr4j_sample")

        X1 = 320.0
        X2 = 2.5
        X3 = 69.5
        X4 = 1.5

        Sp0 = 0.60*X1
        Sr0 = 0.70*X3

        model = GR4J(X1, X2,X3, X4, Sp0, Sr0)

        Q, AET = simulate(model, data.P, data.PET)
        dQ = Q - data.Qsim
        @test maximum(abs.(dQ)) < 1E-3
  	end

end


function johnsontrotter(low, high)
    function permutelevel(vec)
        if length(vec) < 2
            return [vec]
        end
        sequences = []
        endint = vec[end]
        smallersequences = permutelevel(vec[1:end-1])
        leftward = true
        for seq in smallersequences
            for pos in (leftward ? (length(seq)+1:-1:1): (1:length(seq)+1))
                push!(sequences, insert!(copy(seq), pos, endint))
            end
            leftward = !leftward
        end
        sequences
    end
    permutelevel(collect(low:high))
end

for (i, sequence) in enumerate(johnsontrotter(1,4))
    println("""$sequence, $(i & 1 == 1 ? "+1" : "-1")""")
end


using Test

@testset "mini-FL Tests" begin
   @testset "comp" begin
		@test Plasm.comp([acos,atan,tan,sin])(pi/4) * 4 - pi < 0.1e-14
		@test Plasm.comp([acos,atan,tan,sin])(1)==(acos∘atan∘tan∘sin)(1)
		@test Plasm.comp([atan,tan])(pi/4)*4 > 3.14159265358979
   end

   @testset "cons Tests" begin
		@test ((Plasm.cons([acos,atan,tan,sin])(pi/4) .<=
		[0.667458,0.665774,1.0,0.707107]))==[true,true,true,true]
		@test (Plasm.cons([sin,cos])(pi/2).>=[1.0,0.0])==[true,true]
		@test Plasm.cons([+,-])([2,3])==[[2,3],[-2,-3]]
		@test typeof(Plasm.cons([+,-])([2,3]))==Array{Array{Int64,1},1}
   end

   @testset "k Tests" begin
		@test Plasm.k(10)(100)==10
		@test Plasm.k(10)(sin)==10
		@test Plasm.k(sin)(cos)==sin
		@test Plasm.k(sin)(1000)==sin
		@test Plasm.k([1,2,3])(100)==[1,2,3]
		@test Plasm.k([1,2,3])(sin)==[1,2,3]
   end

   @testset "aa Tests" begin
		@test (Plasm.aa(sin)([0.,pi/2])==[0.0,1.0])
		@test typeof(Plasm.aa(sin)([0.,pi/6,pi/2]))==Array{Float64,1}
		@test Plasm.aa(sum)([[1,2],[3,4],[5,6],[7,8]])==[3,7,11,15]
   end
end


# proof-of-principle implementation of non-equilibrium muca for standard SIR dynamics
using MonteCarloX
using Random
using Distributions
using StatsBase
using LinearAlgebra
using StaticArrays

using Printf
using DelimitedFiles
using ProgressMeter

function generate_data(path_out::String;
                       do_muca::Bool=true,
                       do_naive::Bool=false,
                       mu=1/8, t=14, epsilon = 0.0,
                       num_iterations::Int = 10,
                       num_updates_iteration::Int = Int(1e6),
                       num_updates_production::Int = Int(1e7)
                      )
    # define system
    I0 = 1
    S0 = Int(1e5) - I0
    R0 = 0
    lambda = mu
    system = SIR(lambda, mu, epsilon, S0, I0, R0)

    # define muca bounds
    I_max = 300
    I_bound = 100
    range_I = 0:I_max
    hist = Histogram(0:I_max+1)
    logW = Histogram(0:I_max+1, Float64)

    # i-o stuff
    list_hist = zeros(length(hist.weights), num_iterations)
    list_logW = zeros(length(logW.weights), num_iterations)

    # total time of evaluation on my machine: 0:02:01
    if do_muca
        #display(plot())
        # iteration with histograms & weights written to file
        for iteration in 1:num_iterations
            println("iteration ", iteration)
            #store weights
            list_logW[:,iteration] .= logW.weights

            muca_iteration!(hist, logW, system, t, num_updates_iteration)

            # store histogram and weights
            list_hist[:,iteration] .= hist.weights

            #display(plot!(log.(hist.weights), xlims=(0,120)))
            # update weights
            if iteration != num_iterations
                muca_update!(logW, hist, I_bound)
            end
        end
        # write histogram and weights
        filename = @sprintf("%s/pop_muca_sir_critical_iterations_hist.dat",path_out)
        open(filename,"w") do fout
            write(fout, "#I\t H^1(I)\t ... \n")
            writedlm(fout, [collect(range_I) list_hist])
        end
        filename = @sprintf("%s/pop_muca_sir_critical_iterations_logW.dat",path_out)
        open(filename,"w") do fout
            write(fout, "#I\t ln(W^1(I))\t ...\n")
            writedlm(fout, [collect(range_I) list_logW])
        end

        # production run with final histogram and final reweighting
        println("production run  ")
        muca_iteration!(hist, logW, system, t, num_updates_production)
        log_dist = log.(hist.weights) .- logW.weights
        norm = sum(exp.(log_dist))
        log_dist = log_dist .- log(norm)
        filename = @sprintf("%s/pop_muca_sir_critical_production.dat",path_out)
        open(filename,"w") do fout
            write(fout, "#I\t P(I)\t ln(W(I))\t H(I)\n")
            writedlm(fout, [collect(range_I) exp.(log_dist) logW.weights hist.weights])
        end
    end

    # analytic solution for comparisson
    dist = distribution_analytic_critical(range_I, mu, t)
    filename = @sprintf("%s/pop_muca_sir_critical_analytical.dat",path_out)
    open(filename,"w") do fout
        write(fout, "#I\t P(I)\n")
        writedlm(fout, [collect(range_I) dist])
    end

    # naive estimation with same "computing time" (skip for now because muca is not optimized)
    # Time of evaluation on my machine: 0:03:35
    if do_naive
        num_trajectories = Int(1e7)
        simulation_naive!(hist, system, t, num_trajectories)
        filename = @sprintf("%s/pop_muca_sir_critical_naive.dat",path_out)
        open(filename,"w") do fout
            write(fout, "#I\t P(I)\n")
            writedlm(fout, [collect(range_I) hist.weights./num_trajectories])
        end
    end
end

"""
Analytical solution frrom Raúl Toral (online lecture course)
https://ifisc.uib-csic.es/raul/CURSOS/SP/Master_equations.pdf
Eq. (5.31)-(5.32)
"""
function distribution_analytic_critical(range_I, mu, t)
    b_t = mu*t

    dist = zeros(length(range_I))
    for (i,I) in enumerate(range_I)
        if I==0
            dist[i] = b_t/(1+b_t)
        else
            dist[i] = b_t^(I - 1) / ((1+b_t)^(I + 1))
        end
    end
    return dist
end

function simulation_naive!(hist_I, system, t, num_trajectories;
                            dT::Float64 = 1.0,
                            seed = 1000)
    hist_I.weights .= 0
    @showprogress 1 for i in 1:num_trajectories
        rng = MersenneTwister(seed+i);
        reset!(system)
        # evolution
        rates = current_rates(system)
        pass_update!(rates, event) = update!(rates, event, system)
        dT_sim = advance!(KineticMonteCarlo(), rng, rates, pass_update!, Float64(t))
        if system.measure_I in hist_I.edges[1]
            hist_I[system.measure_I] += 1
        end
    end
end

function muca_iteration!(hist_I, logW, system, t, num_updates;
                        dT::Float64 = 1.0,
                        epsilon::Float64 = 0.0,
                        updates_therm::Int = Int(1e5),
                        seed_mc = 1000, seed_dynamics = 2000)
    rng_mc = MersenneTwister(seed_mc);
    rng_dynamics = MutableRandomNumbers(MersenneTwister(seed_dynamics), mode=:dynamic)

    current_I = system.I
    hist_I.weights .= 0
    @showprogress 1 for i in 1:(updates_therm + num_updates)
        # update (todo: make incremental update of rng)
        rand_index = rand(rng_mc, 1:length(rng_dynamics))
        old_rng = rng_dynamics[rand_index]
        rng_dynamics[rand_index] = rand(rng_mc)

        # reset dynamics and rerun (has a LOT OF POTENTIAL for optimization)
        MonteCarloX.reset(rng_dynamics)
        reset!(system)
        rates = current_rates(system)
        pass_update!(rates, event) = update!(rates, event, system)
        dT_sim = advance!(KineticMonteCarlo(), rng_dynamics, rates, pass_update!, Float64(t))
        new_I = system.measure_I

        # acceptance
        accept = false
        if new_I in logW.edges[1]
            if rand(rng_mc) < exp(logW[new_I]-logW[current_I])
                #accept
                current_I = new_I
                accept = true
            end
        end
        if !accept
            rng_dynamics[rand_index] = old_rng
        end
        if i > updates_therm
            hist_I[current_I] += 1
        end
    end
end

"""
Naive muca update:
W^{n_1}(I) = W^{n}(I)/H(I) for all I in range [0,I_bound)
if H(I) > 10 [arbitrary choice]
"""
function muca_update!(logW, hist_I, I_bound)
    for I in logW.edges[1]
        if I < I_bound && hist_I[I] > 1
            logW[I] = logW[I] - log(hist_I[I])
        elseif I >= I_bound
            logW[I] = logW[I-1]
        end
    end
end

mutable struct SIR
    lambda::Float64
    mu::Float64
    epsilon::Float64
    N::Int
    S::Int
    I::Int
    R::Int
    S0::Int
    I0::Int
    R0::Int
    measure_S::Int
    measure_I::Int
    measure_R::Int

    function SIR(lambda, mu, epsilon, S0::Int, I0::Int, R0::Int)
        N = S0 + I0 + R0
        new(lambda, mu, epsilon, N, S0, I0, R0, S0, I0, R0, S0, I0, R0)
    end
end

function reset!(system::SIR)
    system.S = system.measure_S = system.S0
    system.I = system.measure_I = system.I0
    system.R = system.measure_R = system.R0
end

function update!(rates::AbstractVector, event::Int, system::SIR)
    system.measure_S = system.S
    system.measure_I = system.I
    system.measure_R = system.R
    if event == 1 # recovery
        system.I -= 1
        system.R += 1
    elseif event == 2 # infection
        system.S -= 1
        system.I += 1
    elseif event == 0 # absorbing state of zero infected

    else
        throw(UndefVarError(:event))
    end

    rates .= current_rates(system)
end

"""
evaluate current rates of SIR system
"""
function current_rates(system::SIR)
    rate_recovery = system.mu * system.I
    rate_infection = system.lambda* system.I* system.S/system.N + system.epsilon
    return MVector{2,Float64}(rate_recovery, rate_infection)
end


using StitchDPI, LinearAlgebra, Test

#initial_guess_tfm
ysz_chip = 1000
ysz_roi = 500
tfm0 = StitchDPI.initial_guess_tfm(ysz_chip, ysz_roi)
tfm_tol = 1e-6

#fake split
ft, rt, fb, rb = fake_split(8, 4; frac_overlap = 0.0, xsz = 4)
@test all(ft[:,3:6].==rt)
@test all(fb[:,3:6].==rb)
@test all(ft[:,7:8].==rb[:,1:2])
@test all(fb[:,1:2].==rt[:,3:4])
#one pixel overlap
ft, rt, fb, rb = fake_split(8, 4; frac_overlap = 0.25, xsz = 4)
@test all(ft[:,3:6].==rt)
@test all(fb[:,3:6].==rb)
@test all(rt[:,4].==rb[:,1])

#stitching
ft, rt, fb, rb = fake_split(80, 40; frac_overlap = 0.0, xsz = 40)
tfm = StitchDPI.initial_guess_tfm(size(ft,2), size(rt,2))
#imgst = stitch(rt, StitchDPI.flipy(rb), tfm, size(ft,2); flip_y_bottom = true)
imgst = stitch(rt, StitchDPI.flipy(rb), tfm; flip_y_bottom = true)
@test all(imgst[:,1:40].==rt)
@test all(imgst[:,41:80].==rb)

#stitch transform generation
tfmfound, mm = stitch_tfm(StitchDPI.flipy(fb), rt, StitchDPI.flipy(rb); maxevals=500)
@test sum(abs.(tfmfound.linear - LinearAlgebra.I(2))) < tfm_tol
@test sum(abs.(tfmfound.translation - tfm.translation)) < tfm_tol
#simg = stitch(rt, StitchDPI.flipy(rb), tfmfound, size(ft,2); flip_y_bottom = true)
#simg = stitch(rt, StitchDPI.flipy(rb), tfmfound; flip_y_bottom = true)
#@test all(simg.==imgst)

#with ten pixels of overlap
ft, rt, fb, rb = fake_split(80, 40; frac_overlap = 0.25, xsz = 40)
tfmfound, mm = stitch_tfm(StitchDPI.flipy(fb), rt, StitchDPI.flipy(rb); maxevals=500)
@test sum(abs.(tfmfound.linear - LinearAlgebra.I(2))) < tfm_tol
@test sum(abs.(tfmfound.translation - (tfm.translation.+[0.0; 10.0]))) < tfm_tol
#simg = stitch(rt, StitchDPI.flipy(rb), tfmfound, size(ft,2); flip_y_bottom = true)
#simg = stitch(rt, StitchDPI.flipy(rb), tfmfound; flip_y_bottom = true)
#@test all(simg[:,31:40].==(rt[:,31:40].*2))


export LogNum

struct LogNum{T}
    l::T # log of the absolute value
    s::Int # sign

    function LogNum{T}(l::T, s::Int) where T
        s = sign(s)
        if isnan(l)
            s = 1
        elseif isinf(l) && l < 0
            s = 0
        elseif s == 0
            l = -inf(T)
        end
        new(l, s)
    end
end

LogNum(l::T, s::Int) where T = LogNum{T}(l, s)
LogNum() = LogNum(-Inf, 0)
LogNum(x::LogNum) = x

function LogNum(x::T) where T
    if isnan(x)
        LogNum(T(NaN), 1)
    else
        LogNum(log(abs(x)), sign(x) ≥ 0 ? 1 : -1)
    end
end

function Base.convert(::Type{T}, x::LogNum{T})::T where T
    x.s * exp(x.l)
end

function Base.convert(::Type{LogNum{T}}, x::T)::LogNum{T} where T
    LogNum(x)
end

Base.promote_rule(::Type{T}, ::Type{LogNum{T}}) where T = LogNum{T}

function Base.zero(::Type{LogNum{T}})::LogNum{T} where T
    LogNum(zero(T))
end

function Base.one(::Type{LogNum{T}})::LogNum{T} where T
    LogNum(one(T))
end

Base.zero(x::LogNum{T}) where T = zero(LogNum{T})
Base.one(x::LogNum{T}) where T = one(LogNum{T})

Base.iszero(x::LogNum) = iszero(x.s)
Base.isfinite(x::LogNum) = x.l < Inf
Base.isinf(x::LogNum) = x.l == Inf
Base.isnan(x::LogNum) = isnan(x.l)

#Base.:(==)(x::LogNum, y::LogNum) = x.s == y.s && x.l == y.l
Base.:(==)(x::LogNum, y::Real) = x == LogNum(y)
Base.:(==)(x::Real, y::LogNum) = LogNum(x) == y

Base.isless(x::LogNum, y::LogNum) = !isnan(x) && !isnan(y) && (x.s < y.s || x.s == y.s == 1 && x.l < y.l || x.s == y.s == -1 && x.l > y.l)
Base.isless(x::Real, y::LogNum) = LogNum(x) < y
Base.isless(x::LogNum, y::Real) = x < LogNum(y)

Base.isapprox(x::LogNum, y::LogNum; atol::Real = 0, rtol = sqrt(eps(Float64))) = x == y || (isfinite(x) && isfinite(y) && abs(x-y) <= atol + rtol*max(abs(x), abs(y)))
Base.isapprox(x::LogNum, y::Real; atol::Real = 0, rtol = sqrt(eps(Float64))) = isapprox(x, LogNum(y); atol=atol, rtol=rtol)
Base.isapprox(x::Real, y::LogNum; atol::Real = 0, rtol = sqrt(eps(Float64))) = isapprox(LogNum(x), y; atol=atol, rtol=rtol)

Base.abs(x::LogNum) = LogNum(x.l, 1)
Base.sign(x::LogNum) = x.s

function Base.exp(x::LogNum{T})::LogNum{T} where T
    LogNum(convert(T, x), 1)
end

function Base.log(x::LogNum{T})::LogNum{T} where T
    if isnan(x)
        LogNum(T(NaN))
    elseif x < 0
        throw(DomainError())
    else
        LogNum(log(abs(x.l)), (x - 1).s)
    end
end

function Base.:(+)(x::LogNum{T}, y::LogNum{T})::LogNum{T} where {T}
    if isnan(x) || isnan(y)
        LogNum(T(NaN))
    elseif x.s == y.s
        LogNum(logadd(x.l, y.l), x.s)
    else
        LogNum(logsub(x.l, y.l), x.l > y.l ? x.s : y.s)
    end
end

Base.:(+)(x::LogNum, y::Real) = x + LogNum(y)
Base.:(+)(x::Real, y::LogNum) = LogNum(x) + y
Base.:(-)(x::LogNum) = LogNum(x.l, -x.s)
Base.:(-)(x::LogNum, y::LogNum) = x + (-y)
Base.:(-)(x::LogNum, y::Real) = x - LogNum(y)
Base.:(-)(x::Real, y::LogNum) = LogNum(x) - y

function Base.:(/)(x::LogNum{T}, y::LogNum{T})::LogNum{T} where {T}
    if isnan(x) || isnan(y) || iszero(x) && iszero(y)
        LogNum(T(NaN))
    elseif iszero(y)
        LogNum(T(Inf), x.s)
    else
        LogNum(x.l - y.l, x.s * y.s)
    end
end

Base.:(/)(x::LogNum, y::Real) = x / LogNum(y)
Base.:(/)(x::Real, y::LogNum) = LogNum(x) / y

function Base.:(*)(x::LogNum{T}, y::LogNum{T})::LogNum{T} where T
    LogNum(x.l + y.l, x.s * y.s)
end

Base.:(*)(x::LogNum, y::Real) = x * LogNum(y)
Base.:(*)(x::Real, y::LogNum) = LogNum(x) * y

function Base.:(^)(x::LogNum{T}, y::LogNum{T})::LogNum{T} where T
    if iszero(x) && iszero(y)
        LogNum(one(x))
    else
        exp(y * log(x))
    end
end

Base.:(^)(x::LogNum, y::AbstractFloat) = x ^ LogNum(y)
Base.:(^)(x::Real, y::LogNum) = LogNum(x) ^ y


using Distributions, Random
using LinearAlgebra
using DataFrames
using Plots
using RCall
using CSV
include("funcs.jl")

data = DataFrame(
        m20 = [3, 11, 10, 7, 4, 3, 2, 1, 0, 0, 0, 1, 1, 101],
        f20 = [0, 2, 4, 8, 9, 3, 0, 0, 0, 0, 0, 0, 0, 126],
        m32 = [7, 10, 11, 16, 3, 2, 1, 0, 0, 0, 0, 0, 0, 19],
        f32 = [1, 5, 11, 10, 5, 1, 0, 1, 0, 0, 0, 0, 0, 47],
        m50 = [5, 8, 11, 15, 4, 2, 1, 1, 0, 0, 0, 0, 0, 7],
        f50 = [0, 4, 6, 6, 3, 1, 1, 4, 0, 0, 0, 1, 1, 17],
        m80 = [4, 10, 8, 14, 8, 2, 1, 0, 0, 1, 0, 0, 0, 2],
        f80 = [2, 7, 15, 9, 3, 4, 1, 1, 0, 1, 0, 0, 0, 4]
)

β_p = Exponential(0.25)
α_p = Uniform(0,1)
γ_p = Gamma(1.5,1)
p_p = Gamma(1.5,1)

α_i = rand(α_p)
β_i = 0.00000001
γ_i = rand(γ_p)
p_i = rand(p_p)

# test of simple model
test = beetle_sim(144, 0.6, 5, 2, 2, 0.2)
println(test)
println(sum(test))

function test_prop()
    tot = zeros(14)
    for i = 1:10000
        tot += beetle_sim(144, 0.5, 5, 0.08, 0.31, 0.2)
    end
    return tot ./ 10000
end
# proportions from test of simple model
test_proportions = test_prop()

α_i = rand(α_p)
p_i = rand(p_p)
# test of extended model
ext_test = beetle_sim_ext(144, α_i, β_i, γ_i, p_i, 0.2)
println(ext_test)
println(sum(ext_test))

diggle = [3, 11, 10, 7, 4, 3, 2, 1, 0,0, 0, 1, 1, 101]

# abc simulation
sim = abc_sim(diggle, 0.2, 10^5, 144, α_p, β_p, γ_p, p_p)


filename() = "day03_sample.txt"

# Read file. Repeat horizontally certain number of times.
lines(repeat) = [r^repeat for r in readlines(filename())]

# Part1
# Walk down the grid according to "right 3, 1 down" rule.
# Count how many trees (`#` sign) that I encounter.
# 32 times because 323 * 3 / 31 = 31.26
part1() = count([val[(row-1)*3+1] == '#' for (row,val) in enumerate(lines(32))])

# Part2
# Customized rule with `right` and `down` values.
function count_trees(lines, right, down)
    col = 1
    cnt = 0
    for row in 1:down:length(lines)
        row > 1 || continue
        col += right
        if lines[row][col] == '#'
            cnt += 1
        end
    end
    return cnt
end

function part2()
    grid = lines(73)
    counts = [count_trees(grid, x[1], x[2]) for x in ((1,1),(3,1),(5,1),(7,1),(1,2))]
    return prod(counts)
end


f = "eleven.txt"

function flash(mat,i,j)
    if i > 1
        if j > 1
            mat[i-1][j-1] += 1  # nw
        end
        if j < 10
            mat[i-1][j+1] += 1  # sw
        end
        mat[i-1][j] += 1        # w
    end
    if i < 10
        if j > 1
            mat[i+1][j-1] += 1  # ne
        end
        if j < 10
            mat[i+1][j+1] += 1  # se
        end
        mat[i+1][j] += 1        # e
    end
    if j > 1
        mat[i][j-1] += 1        # n
    end
    if j < 10
        mat[i][j+1] += 1        # s
    end
    mat[i][j] = 0
    mat
end

function done(mat)
    ! any(map(line -> any(x->x >= 10, line), mat))
end

mat = map(x -> map(y -> parse(Int64, y), split(x, "")), split(strip(read(f, String)), "\n"))

cntA = 0
# Part A:
# for s in 1:100
for s in 1:500
    global mat = map(x -> map(y -> y+1, x), mat)
    flashed = []
    while !done(mat)
        for (i) in eachindex(mat)
            for (j) in eachindex(mat[i])
                if !((i,j) in flashed) && mat[i][j] >= 10
                    global mat = flash(mat,i,j)
                    global cntA += 1
                    push!(flashed, (i,j))
                end
            end
        end
    end
    for (x,y) in flashed
        global mat[x][y] = 0
    end
    # Part B:
    if all(map(line -> all(x->x == 0, line), mat))
        println("Part B: ",s)
        break
    end
end
# print("Part A: ",cntA)


######################################################################
# alpha_vector.jl
#
# implements policy that is a set of alpha vectors
######################################################################

struct AlphaVectorPolicy{P<:Union{POMDP,IPOMDP,RPOMDP,RIPOMDP}, A} <: Policy
    pomdp::P
    alphas::Vector{Vector{Float64}}
    action_map::Vector{A}
end
function AlphaVectorPolicy(pomdp::Union{POMDP,IPOMDP,RPOMDP,RIPOMDP}, alphas)
    AlphaVectorPolicy(pomdp, alphas, ordered_actions(pomdp))
end
# assumes alphas is |S| x |A|
function AlphaVectorPolicy(p::Union{POMDP,IPOMDP,RPOMDP,RIPOMDP}, alphas::Matrix{Float64}, action_map)
    # turn alphas into vector of vectors
    num_actions = size(alphas, 2)
    alpha_vecs = Vector{Float64}[]
    for i = 1:num_actions
        push!(alpha_vecs, vec(alphas[:,i]))
    end

    AlphaVectorPolicy(p, alpha_vecs, action_map)
end

updater(p::AlphaVectorPolicy) = DiscreteUpdater(p.pomdp, p.alphas)

value(p::AlphaVectorPolicy, b::DiscreteBelief) = value(p, b.b)
function value(p::AlphaVectorPolicy, b::Vector{Float64})
    maximum(dot(b,a) for a in p.alphas)
end

function action(p::AlphaVectorPolicy, b::DiscreteBelief)
    num_vectors = length(p.alphas)
    best_idx = 1
    max_value = -Inf
    for i = 1:num_vectors
        temp_value = dot(b.b, p.alphas[i])
        if temp_value > max_value
            max_value = temp_value
            best_idx = i
        end
    end
    return p.action_map[best_idx]
end

function Base.push!(p::AlphaVectorPolicy, alpha::Vector{Float64}, a)
    push!(p.alphas, alpha)
    push!(p.action_map, a)
end


using Latexify
using LaTeXStrings

test_array = ["x/y * d" :x ; :( (t_sub_sub - x)^(2*p) ) 3//4 ]
@test latexinline.(test_array) == [
    L"$\frac{x}{y} \cdot d$"                          L"$x$"         ;
    L"$\left( t_{sub\_sub} - x \right)^{2 \cdot p}$"  L"$\frac{3}{4}$"
    ]

# @test_throws ErrorException latexify("x/y"; bad_kwarg="should error")


# sample_target_mismatch-gaussian_vs_t-kernel_vs_graph
#
# This script compares the values of the univariate graph Stein
# discrepancy and kernel discrepancy obtained when samples are drawn i.i.d.
# from their targets.

using SteinDistributions:
    SteinGaussian,
    SteinScaleLocationStudentT,
    gradlogdensity
using SteinDiscrepancy:
    SteinGaussianKernel,
    ksd,
    gsd

include("experiment_utils.jl")

# random number generator seed
# we allow this to be set on the commandline
@parseintcli seed "s" "seed" 7
# set the target distribution
@parsestringcli distname "t" "target" "gaussian"
# set the discrepancy
@parsestringcli discrepancytype "d" "discrepancy" "gaussiankernel"
# Specify dimension of sampled vectors
d = 1
# Specify a maximum number of samples
maxn = round(Int, 5e5)
# gaussian target
gaussian_target = SteinGaussian(d)
# the graphsolver
graphsolver = "gurobi"
#graphsolver = "clp"
# setup kernel (if necessary)
kernel = nothing
if discrepancytype == "gaussiankernel"
    kernel = SteinGaussianKernel()
end
# setup target
if distname == "studentt"
    df = 10.0
    target = SteinScaleLocationStudentT(df, 0.0, sqrt((df-2.0)/df))
elseif distname == "gaussian"
    target = SteinGaussian(d)
end
# define gradlogp
function gradlogp(x::Array{Float64,1})
    gradlogdensity(target, x)
end
# Draw samples from target distribution
X = @setseed rand(target,maxn)
# Sample sizes at which optimization problem will be solved
ns = vcat(100:100:min(1000,maxn),
          2000:1000:min(10000,maxn),
          20000:10000:min(100000,maxn),
          200000:100000:maxn)

# Solve optimization problem at each sample size
for i in ns
    @printf("[Beginning n=%d]:\n", i)
    # Create program with data subset
    Xi = X[1:i,:]
    discrepancy = solvetime = nothing
    if kernel != nothing
        @printf("[Beginning kernel computations n=%d]\n", i)
        result = ksd(points=Xi,
                     gradlogdensity=gradlogp,
                     kernel=kernel)
        discrepancy = sqrt(result.discrepancy2)
        solvetime = result.solvetime
    elseif discrepancytype == "graph"
        @printf("[Beginning graph solver n=%d]\n", i)
        result = gsd(points=Xi,
                     gradlogdensity=gradlogp,
                     solver=graphsolver)
        discrepancy = sum(result.objectivevalue)
        solvetime = sum(result.solvetime)
    end
    # save data
    instance_data = Dict{Any, Any}(
        "target" => distname,
        "n" => i,
        "d" => d,
        "seed" => seed,
        "discrepancytype" => discrepancytype,
        "discrepancy" => discrepancy,
        "solvetime" => solvetime,
    )
    save_json(
        instance_data;
        dir="sample_target_mismatch-gaussian_vs_t-kernel_vs_graph",
        target=distname,
        discrepancytype=discrepancytype,
        n=i,
        d=d,
        seed=seed
    )
end

@printf("COMPLETE!")


#------------------------------------------------------------------------------#
#                                   Package                                    #
#------------------------------------------------------------------------------#

"""
# 1. Description

    module HermitePolyTensor

Symbolic Hermite Polynomial Tensors — or N-dimensional Hermite polynomials [1] —
in one to three Euclidean space dimensions.

# 2. Usage

```julia-repl
julia> import HermitePolyTensor

```

# 3. References

[1]: H. Grad, “Note on N-dimensional  Hermite  polynomials,”  Communications  on
Pure and Applied Mathematics, vol. 2, no. 4, pp. 325–330, 1949.
"""
module HermitePolyTensor

#------------------------------------------------------------------------------#
#                                   Imports                                    #
#------------------------------------------------------------------------------#

using Reexport

#------------------------------------------------------------------------------#
#                                   Includes                                   #
#------------------------------------------------------------------------------#

include("version.jl")

include("HPT.jl")
@reexport using .HPT

end # module


###
# Optimization problems for VoPP Upper Bounds
###

####Separation oracles
#hbar(t) = M * abs(t-1)/2 - D
#notice no dependence on mu.
#maximum occurs at an end point or 1
function sep_MAD(a, b, lam1, lam2, M, D)
	lhs = lam1 * a + lam2/2 * abs(a-1) * M - lam2 * D
	rhs = lam1 * b + lam2/2 * abs(b-1) * M - lam2 * D
	mid = a <= 1. <= b ? lam1 - lam2 * D : -Inf
	val = max(lhs, rhs, mid)

	if lhs >= val
		return val, a
	elseif rhs >= val
		return val, b
	elseif mid >= val
		return val, 1.
	end
	throw("Unreachable")
	return -1., -1.
end

#hbar(t) = M^2 (t - 1)^2 - C^2
#Again, no dependence on mu
#maximum occurs at an end point or at unique minimizer
function sep_CV(a, b, lam1, lam2, M, C)
	lhs = lam1 * a + lam2 * M^2 * (a - 1)^2 - lam2 * C^2
	rhs = lam1 * b + lam2 * M^2 * (b - 1)^2 - lam2 * C^2
	vstar = 1 - lam1 / 2 / lam2 / M^2
	mid = a <= vstar <= b ? lam1 * vstar + lam2 * M^2 * (vstar - 1)^2 - lam2 * C^2 : -Inf
	val = max(lhs, rhs, mid)

	if lhs >= val
		return val, a
	elseif rhs >= val
		return val, b
	elseif mid >= val
		return val, vstar
	end
	throw("Unreachable")
	return -1., -1.
end

#h(t) = -log(t/mu) + log (B/mu)   equivalent to Exp[log[V]]= log[B]
#hbar(t)  = -log(M(t -1 ) + 1) + log (B/mu)
function sep_GM(a, b, lam1, lam2, M, mu, B)
	lhs = lam1 * a - lam2 * log(M * (a - 1) + 1 ) + lam2 * log(B/mu)
	rhs = lam1 * b - lam2 * log(M * (b - 1) + 1 ) + lam2 * log(B/mu)

	vstar = lam2 <= 0 ? 1 - 1/M + lam2 / lam1 : a - 1e-6
	mid = a <= vstar <= b ? lam1 * vstar - lam2 * log(M * (vstar - 1) + 1) + lam2 * log(B/mu) : -Inf
	val = max(lhs, rhs, mid)

	if lhs >= val
		return val, a
	elseif rhs >= val
		return val, b
	elseif mid >= val
		return val, vstar
	end
	throw("Unreachable")
	return -1., -1.
end

#### Functions based on a generic separation argument
#h 1d function that represents moment constraint (includes moment):  E[h(V)] = 0
#sep_fun(a, b, lam1, lam2, M, mu) maximizes of lam1 * v  + lam2 h(mu * (Mv + 1-M) ).  returns optval, vstar
function _vopp_moment_dual(S, M, mu, h, sep_fun; N, numCuts=100, TOL =1e-6, print_trace=false)
	#translate everything to the Vc space
	hbar(t) = h( mu * (M * t + (1 - M)) )
	Sc = vopp.comp_Sc(S, M)

	@assert isfinite(Sc) "Currenty only supports finite Sc"

	p_grid = range(0, stop=Sc, length=N+1) #indexes from 1, doc indexes from 0
	m = Model(solver=GurobiSolver(OutputFlag = false))
	@variable(m, theta)
	@variable(m, lam[1:2])
	@variable(m, Qs[1:N+1] >= 0)  #indexes form 1, doc indexes from 0

	@constraint(m, sum(Qs) == 1)
	@constraint(m, theta + lam[1] * Sc + lam[2] * hbar(Sc) <= sum([p_grid[i] * Qs[i] for i = 1:N+1]))
	@objective(m, Max, theta + lam[1])

	for iter = 1:numCuts
		iter == numCuts && throw("Max Iteration Count Reached")
		status = solve(m)  #check status?
		feasible = true
		#check if any cut is violated
		for k = 2:N+1
			val, vstar = sep_fun(p_grid[k-1], p_grid[k], getvalue(lam[1]), getvalue(lam[2]), M, mu)
			lhs = val + getvalue(theta)
			rhs = sum([p_grid[j] * getvalue(Qs[j]) for j = 1:k-1])

			if lhs >= rhs + TOL
				@constraint(m, theta + vstar * lam[1] + hbar(vstar) * lam[2] <= sum( [p_grid[j] * Qs[j] for j = 1:k-1]) )
				feasible = false
				print_trace && println("iter $(iter): \t $lhs \t $rhs")
				continue
			end
		end
		if feasible
			@assert status == :Optimal "Cutting Procedure failed: $(status)"
			return 1/(getobjectivevalue(m) - TOL)
		end
	end
	##unreachable
	return -1.
end

######################
####
# Specialized function for MAD
# Not using separation.
# Algorithm relies on fact that minimizers are always at endpt or 1.
####
function vopp_ub_MAD_opt(S, M, D; N = 500)
	#Solve the two cases for lambda2
	val_neg = _vopp_ub_MAD(S, M, D, false, N)
	val_pos = _vopp_ub_MAD(S, M, D, true, N)

	#take minimum bc we already took reciprocals
	min(val_neg, val_pos), val_neg, val_pos
end

function _vopp_ub_MAD(S, M, D, is_lam2_pos, N)
	Sc = comp_Sc(S, M)
	@assert isfinite(Sc) "Currenty only supports finite Sc"

	p_grid = range(0, stop=Sc, length=N+1) #indexes from 1, doc indexes from 0
	m = Model(solver=GurobiSolver(OutputFlag=false))
	@variable(m, theta)
	@variable(m, lam[1:2])
	@variable(m, Qs[1:N+1] >= 0)  #indexes form 1, doc indexes from 0

	if is_lam2_pos
		@constraint(m, lam[2] >= 0)
	else
		@constraint(m, lam[2] <= 0)
	end

	hbar(t) = M * abs(t-1)/2 - D

	@constraint(m, sum(Qs) == 1)
	@constraint(m, theta + lam[1] * Sc + lam[2] * hbar(Sc) <= sum([p_grid[i] * Qs[i] for i = 1:N+1]))
	@objective(m, Max, theta + lam[1])

	for k = 2:N+1
		@constraint(m, theta + p_grid[k-1] * lam[1] + hbar(p_grid[k-1]) * lam[2] <= sum( [p_grid[j] * Qs[j] for j = 1:k-1]) )
		@constraint(m, theta + p_grid[k]   * lam[1] + hbar(p_grid[k])   * lam[2] <= sum( [p_grid[j] * Qs[j] for j = 1:k-1]) )

		if !is_lam2_pos && p_grid[k-1] <= 1 <= p_grid[k]
			@constraint(m, theta + lam[1] + hbar(1.) * lam[2] <= sum( [p_grid[j] * Qs[j] for j = 1:k-1]) )
		end
	end
	status = solve(m)
	@assert status == :Optimal "Problem Failed $status"
	return 1/(getobjectivevalue(m))
end

### Specialized function for CV
# useSep determines use separator or else Reformulation
# Notice that the minimizer is not always at 1 because of linear offset
#Weirdly slow for some reason.
function vopp_ub_CV(S, M, mu, C; method=:ConstraintGeneration, N=500, numCuts=100, TOL=1e-6, print_trace=false)
	if method == :Slemma
		return _vopp_CV_SLemma(S, M, mu, C, N)
	elseif method == :ConstraintGeneration
		h(v) = (v - mu)^2 / mu^2 - C^2
		sep_fun(a, b, lam1, lam2, M, mu) = sep_CV(a, b, lam1, lam2, M, C)
		return _vopp_moment_dual(S, M, mu, h, sep_fun, N=N, numCuts=numCuts, TOL=TOL, print_trace=print_trace)
	elseif method == :Reformulation
		return min( _vopp_ub_CV(S, M, mu, C, true, N), _vopp_ub_CV(S, M, mu, C, false, N) )
	else
		throw("Method must be one of :Slemma , :ConstraintGeneration , :Reformulation")
	end
end

function _vopp_ub_CV(S, M, mu, C, is_lam2_pos, N=500)
	#translate everything to the Vc space
	hbar(t) = M^2 * (t - 1)^2 - C^2  ##This functional form implicitly coded into dual constraints
	Sc = vopp.comp_Sc(S, M)
	@assert isfinite(Sc) "Currenty only supports finite Sc"

	p_grid = range(0, stop=Sc, length=N+1) #indexes from 1, doc indexes from 0
	m = Model(solver=GurobiSolver(OutputFlag = false))

	@variable(m, theta)
	@variable(m, lam[1:2])

	if is_lam2_pos
		@constraint(m, lam[2] >= 0)  #means solutions only occur at endpoints
	else
		@constraint(m, lam[2] <= 0)
	end

	@variable(m, Qs[1:N+1] >= 0)  #indexes form 1, doc indexes from 0
	@constraint(m, sum(Qs) == 1)

	@constraint(m, theta + lam[1] * Sc + lam[2] * hbar(Sc) <= sum([p_grid[i] * Qs[i] for i = 1:N+1]))

	#Add the auxiliary variables for each constraint in lam_2_neg
	if !is_lam2_pos
		@variable(m, alphas[2:N+1])
		@variable(m, betas[2:N+1])
		@variable(m, ts[2:N+1] >= 0)
	end

	for k = 2:N+1
		#always need the end points
		@constraint(m, theta + p_grid[k-1] * lam[1] + hbar(p_grid[k-1]) * lam[2] <= sum( [p_grid[j] * Qs[j] for j = 1:k-1]) )
		@constraint(m, theta + p_grid[k] * lam[1] + hbar(p_grid[k]) * lam[2] <= sum( [p_grid[j] * Qs[j] for j = 1:k-1]) )

		if !is_lam2_pos #concave case
			#use auxiliary variables to specify dual constraints
			@constraint(m, alphas[k] - lam[1] == betas[k])
			@constraint(m, betas[k]^2 <= -4 * M^2 * lam[2] * ts[k])
			@constraint(m, ts[k] - alphas[k] + lam[1] + alphas[k] * p_grid[k-1] <=
								C^2 * lam[2] - theta + sum([p_grid[i] * Qs[i] for i = 1:k-1]))

			@constraint(m, ts[k] - alphas[k] + lam[1] + alphas[k] * p_grid[k] <=
								C^2 * lam[2] - theta + sum([p_grid[i] * Qs[i] for i = 1:k-1]))
		end
	end
	@objective(m, Max, theta + lam[1])

	##Solve and extract solution
	status = solve(m)
	@assert status == :Optimal "Reformulation failed: $(status)"

	return 1/getobjectivevalue(m)
end

function _vopp_CV_SLemma(S, M, mu, C, N)
	#translate everything to the Vc space
	hbar(t) = M^2 * (t - 1)^2 - C^2  ##This functional form implicitly coded into dual constraints
	Sc = vopp.comp_Sc(S, M)
	@assert isfinite(Sc) "Currenty only supports finite Sc"

	ps = range(0, stop=Sc, length=N+1) #indexes from 1, doc indexes from 0
	m = Model(solver=GurobiSolver(OutputFlag = false))

	@variable(m, theta)
	@variable(m, lam[1:2])
	@variable(m, Qs[1:N+1] >= 0)  #doc indexes from 0 to N
	@constraint(m, sum(Qs) == 1)

	@constraint(m, theta + lam[1] * Sc + lam[2] * hbar(Sc) <= sum([ps[i] * Qs[i] for i = 1:N+1]))

	for k = 2:N+1
		#add auxiliary variables and specify modified S-Lemma
		x = @variable(m)
		y = @variable(m)
		z = @variable(m)

		@constraint(m, x >= 0)
		@constraint(m, y >= 0)
		@constraint(m, y >= -lam[2] * M^2)

		@constraint(m, ps[k-1] * ps[k] * (y + lam[2]*M^2) - x
						>=
						theta - sum(ps[j] * Qs[j] for j = 1:k) +
						lam[2] * M^2 - lam[2] * C^2
					)

		@constraint(m, 4 * y * x >= z^2)
		@constraint(m, z == 2 * lam[2] * M^2 - lam[1] -
							(ps[k-1] + ps[k]) * (y + lam[2] * M^2)
					)
	end
	@objective(m, Max, theta + lam[1])

	##Solve and extract solution
	status = solve(m)
	@assert status == :Optimal "Reformulation failed: $(status)"

	return 1/getobjectivevalue(m)
end

### Specialized function for BM
# useSep determines use separator or else Reformulation
function vopp_ub_GM(S, M, mu, B; useSep=false, N=500, numCuts=100, TOL=1e-6, print_trace=false)
	if useSep
		h(v) = -log(v/mu) + log(B/mu)
		sep_fun(a, b, lam1, lam2, M, mu) = sep_GM(a, b, lam1, lam2, M, mu, B)
		return _vopp_moment_dual(S, M, mu, h, sep_fun, N=N, numCuts=numCuts, TOL=TOL, print_trace=print_trace)
	end
	return min( _vopp_ub_GM(S, M, mu, B, true, N), _vopp_ub_GM(S, M, mu, B, false, N) )
end

### Specialized function for GM using convex conjugates
#uses a combination of IPOPT and prayer.
#h(t) = -log(t/mu) + log(B/mu)
function _vopp_ub_GM(S, M, mu, B, is_lam2_pos, N)
	#translate everything to the Vc space
	hbar(t) = -log(M*t + 1- M) + log(B/mu)  ##This functional form implicitly coded into dual constraints

	Sc = vopp.comp_Sc(S, M)
	@assert isfinite(Sc) "Currenty only supports finite Sc"

	p_grid = range(0, stop=Sc, length=N+1) #indexes from 1, doc indexes from 0
	m = Model(solver=IpoptSolver())

	@variable(m, theta)
	@variable(m, lam[1:2])

	if is_lam2_pos
		@constraint(m, lam[2] >= 0)  #means solutions only occur at endpoints
	else
		@constraint(m, lam[2] <= 0)
	end

	@variable(m, Qs[1:N+1] >= 0)  #indexes form 1, doc indexes from 0
	@constraint(m, sum(Qs) == 1)

	@constraint(m, theta + lam[1] * Sc + lam[2] * hbar(Sc) <= sum([p_grid[i] * Qs[i] for i = 1:N+1]))

	#Add the auxiliary variables for each constraint in lam_2_neg
	if !is_lam2_pos
		@variable(m, rhos[2:N+1])
		@variable(m, ts[2:N+1])
	end

	for k = 2:N+1
		#always need the end points
		@constraint(m, theta + p_grid[k-1] * lam[1] + hbar(p_grid[k-1]) * lam[2] <= sum( [p_grid[j] * Qs[j] for j = 1:k-1]) )
		@constraint(m, theta + p_grid[k]   * lam[1] + hbar(p_grid[k])   * lam[2] <= sum( [p_grid[j] * Qs[j] for j = 1:k-1]) )

		if !is_lam2_pos #concave case
			#use auxiliary variables to specify dual constraints
			@constraint(m, rhos[k] >= lam[1])
			@NLconstraint(m, -lam[2] * log(-M * lam[2] / (rhos[k] - lam[1])) <= ts[k])
			@constraint(m, (1-M)/M * (rhos[k] - lam[1]) + lam[2] + lam[2]*log(B/mu) + ts[k] +
						rhos[k] * p_grid[k-1] <= sum([p_grid[i] * Qs[i] for i = 1:k-1]) - theta)
			@constraint(m, (1-M)/M * (rhos[k] - lam[1]) + lam[2] + lam[2]*log(B/mu) + ts[k] +
						rhos[k] * p_grid[k] <= sum([p_grid[i] * Qs[i] for i = 1:k-1]) - theta)

			# @constraint(m, -lam[2] * log( -rhos[k] / M /lam[2] ) <= ts[k])
			# @constraint(m, lam[1]* p_grid[k-1] - rhos[k] * p_grid[k-1] - lams[2] * log(B/mu) - lams[2] +
			# 				(1 - 1/M) * rhos[k] + ts[k] <= sum([p_grid[i] * Qs[i] for i = 1:k-1]) - theta )
			# @constraint(m, lam[1]* p_grid[k] - rhos[k] * p_grid[k] - lams[2] * log(B/mu) - lams[2] +
			# 				(1 - 1/M) * rhos[k] + ts[k] <= sum([p_grid[i] * Qs[i] for i = 1:k-1]) - theta )
		end
	end
	@objective(m, Max, theta + lam[1])

	##Solve and extract solution
	status = solve(m)
	@assert status == :Optimal "Reformulation failed: $(status)"

	return 1/getobjectivevalue(m)
end

#h(t) = I( V >= phat mu) - q
#hbar(t) = I( t >= v0) - q where v0 = (phat -1)/M + 1
#sep_fun(a, b, lam1, lam2, M, mu) maximizes of lam1 * v  + lam2 hbar(mu * (Mv + 1-M) ).  returns optval, vstar
function sep_IC(a, b, lam1, lam2, M, mu, phat, q; tol=1e-10)
	v0 = (phat - 1)/M + 1
	lhs(v) =  v >= v0 ? lam1 * v + lam2 * (1 - q) : lam1 * v - lam2 * q

	#check end points
	tstar = a
	val = lhs(tstar)

	if lhs(b) > val
		tstar = b
		val = lhs(tstar)
	end

	if a <= v0 < b
		if lam2 >= 0
			#check v0
			if lhs(v0) > val
				tstar = v0
				val = lhs(tstar)
			end
		else #lam2 < 0
			#check v0-tol
			if lhs(v0 -tol) > val
				tstar = v0-tol
				val = lhs(tstar)
			end
		end
	end
	val, tstar
end

### Specialized function for Incumbent Price
function vopp_ub_IC(S, M, mu, phat, q; N=500,numCuts=100, TOL=1e-6, print_trace=false)
	v0 = (phat - 1)/M + 1
	h(v) = v >= v0 ?  1 - q : -q
	sep_fun(a, b, lam1, lam2, M, mu) = sep_IC(a, b, lam1, lam2, M, mu, phat, q)
	return _vopp_moment_dual(S, M, mu, h, sep_fun, N=N, numCuts=numCuts, TOL=TOL, print_trace=print_trace)
end


#=
This file contains code for separating a timeseries into its seasonal and residual
fluctuations. It interfaces SignalSeparation.jl and also uses some analytic knowledge
of the form of the incoming solar radiation
=#
using SignalDecomposition

export decompose, Fourier, MeanAlbedo, ProductInversion, fielddecompose, Sinusoidal
export nrmse, rmse, TimeAnomaly

# my own decomposition method that dispatches on dimensional array for input method

function fielddecompose(A::AbDimArray, method::Decomposition)
    t = dims(A, Time)
    fielddecompose(t, A, method)
end
function fielddecompose(A::AbDimArray, method::Fourier)
    t = 1:size(A, Time)
    fielddecompose(t, A, method)
end

function fielddecompose(t::AbstractVector, A::AbDimArray, method::Decomposition)
    seasonal = DimensionalArray(copy(Array(A)), dims(A), A.name*"x")
    residual = DimensionalArray(copy(Array(A)), dims(A), A.name*"r")
    for i in spatialidxs(A)
        y = Array(A[i...])
        sea, res = SignalDecomposition.decompose(t, y, method)
        seasonal[i...] .= sea
        residual[i...] .= res
    end
    return seasonal, residual
end

SignalDecomposition.decompose(t::Vector{<:TimeType}, s, m::Fourier) =
decompose(1:length(t), s, m)

SignalDecomposition.decompose(s::AbDimArray, m::Sinusoidal) =
decompose(dims(s, Time).val, Array(s), m)

function SignalDecomposition.decompose(t::AbstractVector{<:TimeType}, s, m::Sinusoidal)
    truetime = realtime_days(t)
    decompose(truetime, s, m)
end

#########################################################################
# Mean Albedo
#########################################################################
"""
    MeanAlbedo(S::Vector) <: Decomposition
    MeanAlbedo(t::Vector, φ::Real)
Decomposition that splits a timeseries `R` into its mean albedo component
``\\bar{\\alpha}S`` with ``S`` the Solar
insolation (either given directly as `S` or calculated analytically with given time
`t` and latitude `φ`) and ``\\bar{\\alpha} = `` `timeaverage(R)/timeaverage(S)`,
and the residuals of that.

It is supposed to approximately purely internal dynamical fluctuations.
"""
struct MeanAlbedo{T<:Real} <: Decomposition
    solar::Vector{T}
end
MeanAlbedo(t::Vector, φ::Real) = MeanAlbedo(insolation.(t, φ))
MeanAlbedo(x) = MeanAlbedo(Array(x))

# Extention for convenient usage in decompose below
ClimateBase.timeagg(f, t::Vector, a::AbDimArray{<:Any, 1}) = timeagg(f, t, a.data)

function SignalDecomposition.decompose(t, x, method::MeanAlbedo)
    @assert length(x) == length(method.solar)
    if eltype(t) <: TimeType
        meanα = timeagg(mean, t, x)/timeagg(mean, t, method.solar) # average albedo
    else
        meanα = mean(x)/mean(method.solar)
    end
    seasonal = meanα .* method.solar
    return seasonal, x .- seasonal
end

function meanalbedo_decompose(A::AbDimArray, S)
    solar = DimensionalArray(copy(Array(A)), dims(A), "MeanAlbedo αS")
    residual = DimensionalArray(copy(Array(A)), dims(A), "MeanAlbedo r")
    for i in spatialidxs(A)
        y = Array(A[i...])
        s = Array(S[i...])
        method = MeanAlbedo(s)
        sea, res = SignalDecomposition.decompose(y, method)
        solar[i...] .= sea
        residual[i...] .= res
    end
    return solar, residual
end


using VegaLite
using DataFrames

## histograms by group

df= DataFrame(group=rand(0:1, 200))
df[:x] = df[:group]*2 + randn(size(df,1))

data_values(df) +
    mark_bar() +
    encoding_x_quant(:x; bin=Dict(:maxbins=>15)) +
    encoding_y_quant(:*, aggregate="count") +
    encoding_row_nominal(:group) +
    encoding_color_nominal(:group)


module GinzburgLandau

# package code goes here

dct_nodes(N, L) = (0:N-1)*L/N + L/(2N)

AdamsMoulton = [1          0          0         0         0          0;
                1/2        1/2        0         0         0          0;
                5/12       8/12      -1/12      0         0          0;
                9/24       19/24     -5/24      1/24      0          0;
                251/720    646/720   -264/720   106/720  -19/720     0;
                475/1440   1427/1440 -798/1440  482/1440 -173/1440   27/1440]

AdamsBashforth=[1          0          0          0        0          0;
                3/2        -1/2       0          0        0          0;
                23/12      -16/12     5/12       0        0          0;
                55/24      -59/24     37/24     -9/24     0          0;
                1901/720   -2774/720  2616/720  -1274/720  251/720   0;
                4277/1440  -7923/1440 9982/1440 -7298/1440 2877/1440 -475/1440]

type GLECoeffs
    σ::Float64
    γ::Float64
    c₁::Float64
    c₃::Float64
    k::Vector{Complex{Float64}}
end

function GLECoeffs(σ, γc, c₁, c₃)
    k = zeros(Complex{Float64}, 3)
    γ = γc/c₁
    k[1] = σ
    k[2] = γ*(1 + im*c₁)
    k[3] = -σ*(1 - im*c₃)
    GLECoeffs(σ, γ, c₁, c₃, k)
end

import Base.getindex

getindex(coeffs::GLECoeffs, i) = coeffs.k[i]

type GLESolver1
    L::Float64
    N::Int
    Δt::Float64
    δ::Vector{Complex{Float64}}
    g::Vector{Complex{Float64}}
end

function GLESolver1(L, N, Δt, k::GLECoeffs)

    δ = zeros(Complex{Float64}, N)

    for i = 1:N
        w = π/L*(i-1)
        δ[i] = 1.0 - Δt*(k[1] - k[2]*w^2)
    end

    g = zeros(Complex{Float64}, N)
    GLESolver1(L, N, Δt, δ, g)
end

nmodes(slv::GLESolver1) = slv.N
nodes(slv::GLESolver1) = dct_nodes(slv.N, slv.L)
tstep(slv::GLESolver1, i) = (i-1)*slv.Δt

function gle_step!(slv::GLESolver1, k::GLECoeffs, a₀::AbstractVector{Complex{Float64}},
                   a₁::AbstractVector{Complex{Float64}})
    # Calcularo f
    g = slv.g
    dt = slv.Δt
    N = slv.N
    δ = slv.δ

    copy!(a₁, a₀)
    idct!(a₁)

    for i = 1:N
        g[i] = k[3] * abs(a₁[i])^2 * a₁[i]
    end

    dct!(g)
    for i = 1:N
        a₁[i] = ( a₀[i] + dt*g[i] ) / δ[i]
    end

    idct!(a₁)
    for i = 1:N
        g[i] = k[3] * abs(a₁[i])^2 * a₁[i]
    end
    dct!(g)

    for i = 1:N
        a₁[i] = ( a₀[i] + dt*g[i] ) / δ[i]
    end

    return nothing

end

type GLEAdamsBashforth
    L::Float64
    N::Int
    s::Int
    Δt::Float64
    f::Vector{Vector{Complex{Float64}}
    a₀::Vector{Complex{Float64}}
    idx::Vector{Integer}
    step::Int
    function GLEAdamsBashforth(L, N, s, Δt)
        f = Vector{Vector{Complex{Float64}}}(s)
        for i = 1:s
            f[i] = zeros(Complex{Float64}, N)
        end
        a₀ = zeros(Complex{Float64}, N)
        new(L, N, s, Δt, f, a₀, zeros(Int, s), 0)
    end

end
nmodes(slv::GLEAdamsBashforth) = slv.N
nodes(slv::GLEAdamsBashforth) = dct_nodes(slv.N, slv.L)
tstep(slv::GLEAdamsBashforth, i) = (i-1)*slv.Δt
nstep(slv::GLEAdamsBashforth) = slv.s

function init_gle!(slv::GLEAdamsBashforth, a::AbstractVector)
    s = nstep(slv)
    slv.step = 0
    fill!(slv.idx, 0)
    for i = 1:s
        fill!(slv.f[i], 0)
    end
    fill!(slv.a₀, a)

end

function simple_step!(slv::GLEAdamsBashforth, a::AbstractVector{Complex{Float64}})

end # module


function gradfunc{T <: Number}(
    f::Function,
    x::T;
    method::Symbol = :finite,
    direction::Symbol = :central,
) # -> Function
    if method == :finite
        function g(x::T)
            return finitegrad(f, x, direction)
        end
    elseif method == :ad
        function g(x::T)
            return autograd(f, x)
        end
    else
        throw(ArgumentError("method must be :finite or :ad"))
    end
    return g
end

function gradfunc!{T <: Number}(
    f::Function,
    x::Vector{T};
    method::Symbol = :finite,
    direction::Symbol = :central,
) # -> Function
    if method == :finite
        function g!(gr::Vector{T}, x::Vector{T})
            return finitegrad!(gr, f, x, direction)
        end
    elseif method == :ad
        x_d = Array(Dual{T}, length(x))
        function g!(gr::Vector{T}, x::Vector{T})
            return autograd!(gr, x_d, f, x)
        end
    else
        throw(ArgumentError("method must be :finite or :ad"))
    end
    return g!
end

function gradfunc{T <: Number}(
    f::Function,
    x::Vector{T};
    method::Symbol = :finite,
    direction::Symbol = :central,
) # -> Function
    if method == :finite
        function g(x::Vector{T})
            gr = similar(x)
            finitegrad!(gr, f, x, direction)
            return gr
        end
    elseif method == :ad
        x_d = Array(Dual{T}, length(x))
        function g(x::Vector{T})
            gr = similar(x)
            autograd!(gr, x_d, f, x)
            return gr
        end
    else
        throw(ArgumentError("method must be :finite or :ad"))
    end
    return g
end


using EngEconomics, Roots, Plots

# Given


# ---
# title: 1310. XOR Queries of a Subarray
# id: problem1310
# author: Indigo
# date: 2021-05-14
# difficulty: Medium
# categories: Bit Manipulation
# link: <https://leetcode.com/problems/xor-queries-of-a-subarray/description/>
# hidden: true
# ---
#
# Given the array `arr` of positive integers and the array `queries` where
# `queries[i] = [Li, Ri]`, for each query `i` compute the **XOR** of elements
# from `Li` to `Ri` (that is, `arr[Li] **xor** arr[Li+1] **xor** ... **xor**
# arr[Ri]` ). Return an array containing the result for the given `queries`.
#
#
#
# **Example 1:**
#
#
#
#     Input: arr = [1,3,4,8], queries = [[0,1],[1,2],[0,3],[3,3]]
#     Output: [2,7,14,8]
#     Explanation:
#     The binary representation of the elements in the array are:
#     1 = 0001
#     3 = 0011
#     4 = 0100
#     8 = 1000
#     The XOR values for queries are:
#     [0,1] = 1 xor 3 = 2
#     [1,2] = 3 xor 4 = 7
#     [0,3] = 1 xor 3 xor 4 xor 8 = 14
#     [3,3] = 8
#
#
# **Example 2:**
#
#
#
#     Input: arr = [4,8,2,10], queries = [[2,3],[1,3],[0,0],[0,3]]
#     Output: [8,0,4,4]
#
#
#
#
# **Constraints:**
#
#   * `1 <= arr.length <= 3 * 10^4`
#   * `1 <= arr[i] <= 10^9`
#   * `1 <= queries.length <= 3 * 10^4`
#   * `queries[i].length == 2`
#   * `0 <= queries[i][0] <= queries[i][1] < arr.length`
#
#
## @lc code=start
using LeetCode

function xor_queries(arr::Vector{Int}, queries::Vector{Vector{Int}})
    pushfirst!(arr, 0)
    off_arr = OffsetArray(arr, -1)
    for i in 1:(length(arr) - 1)
        @inbounds off_arr[i] = off_arr[i] ⊻ off_arr[i - 1]
    end
    res = fill(0, length(queries))
    for (idx, q) in enumerate(queries)
        @inbounds res[idx] = off_arr[q[2] + 1] ⊻ off_arr[q[1]]
    end
    return res
end
## @lc code=end


module BarotropicQGQL

export
  Problem,
  set_zeta!,
  updatevars!,

  energy,
  enstrophy,
  dissipation,
  work,
  drag

using
  FFTW,
  CUDA,
  Reexport,
  DocStringExtensions

@reexport using FourierFlows

using LinearAlgebra: mul!, ldiv!
using FourierFlows: parsevalsum, parsevalsum2
import FFTW: rfft

abstract type BarotropicQGQLVars <: AbstractVars end

nothingfunction(args...) = nothing

"""
    Problem(dev::Device=CPU();
                      nx = 256,
                      ny = nx,
                      Lx = 2π,
                      Ly = Lx,
                       β = 0.0,
      deformation_radius = Inf,
                     eta = nothing,
                       ν = 0.0,
                      nν = 1,
                       μ = 0.0,
                      dt = 0.01,
                 stepper = "RK4",
                   calcF = nothingfunction,
              stochastic = false,
        aliased_fraction = 1/3,
                       T = Float64)

Construct a quasi-linear barotropic quasi-geostrophic `problem` on device `dev`.

Keyword arguments
=================
    - `dev`: (required) `CPU()` or `GPU()`; computer architecture used to time-step `problem`.
    - `nx`: Number of grid points in ``x``-domain.
    - `ny`: Number of grid points in ``y``-domain.
    - `Lx`: Extent of the ``x``-domain.
    - `Ly`: Extent of the ``y``-domain.
    - `β`: Planetary vorticity ``y``-gradient.
    - `eta`: Topographic potential vorticity.
    - `ν`: Small-scale (hyper)-viscosity coefficient.
    - `nν`: (Hyper)-viscosity order, `nν```≥ 1``".
    - `μ`: Linear drag coefficient.
    - `dt`: Time-step.
    - `stepper`: The extent of the ``y``-domain.
    - `calcF`: Function that calculates the Fourier transform of the forcing, ``F̂``.
    - `stochastic`: `true` or `false`; boolean denoting whether `calcF` is temporally stochastic.
    - `aliased_fraction`: the fraction of high-wavenubers that are zero-ed out by `dealias!()`.
    - `T`: `Float32` or `Float64`; floating point type used for `problem` data.
"""
function Problem(dev::Device=CPU();
  # Numerical parameters
                nx = 256,
                ny = nx,
                Lx = 2π,
                Ly = Lx,
  # Physical parameters
                 β = 0.0,
               eta = nothing,
  # Drag and/or hyper-/hypo-viscosity
                 ν = 0.0,
                nν = 1,
                 μ = 0.0,
  # Timestepper and equation options
                dt = 0.01,
           stepper = "RK4",
             calcF = nothingfunction,
        stochastic = false,
  # Float type and dealiasing
  aliased_fraction = 1/3,
                 T = Float64)

  # the grid
  grid = TwoDGrid(dev, nx, Lx, ny, Ly; aliased_fraction=aliased_fraction, T=T)
  x, y = gridpoints(grid)

  # topographic PV
  eta === nothing && ( eta = zeros(dev, T, (nx, ny)) )

  params = !(typeof(eta)<:ArrayType(dev)) ?
           Params(grid, β, eta, μ, ν, nν, calcF) :
           Params(β, eta, rfft(eta), μ, ν, nν, calcF)

  vars = calcF == nothingfunction ? DecayingVars(dev, grid) : stochastic ? StochasticForcedVars(dev, grid) : ForcedVars(dev, grid)

  equation = BarotropicQGQL.Equation(params, grid)

  FourierFlows.Problem(equation, stepper, dt, grid, vars, params, dev)
end

# ----------
# Parameters
# ----------

"""
    Params{T, Aphys, Atrans}(β, eta, etah, μ, ν, nν, calcF!)

A struct containing the parameters for a barotropic QL QG problem. Included are:

$(TYPEDFIELDS)
"""
struct Params{T, Aphys, Atrans} <: AbstractParams
    "planetary vorticity y-gradient"
       β :: T
    "topographic potential vorticity"
     eta :: Aphys
    "Fourier transform of topographic potential vorticity"
    etah :: Atrans
    "linear drag coefficient"
       μ :: T
    "small-scale (hyper)-viscosity coefficient"
       ν :: T
    "(hyper)-viscosity order, `nν```≥ 1``"
      nν :: Int
    "function that calculates the Fourier transform of the forcing, ``F̂``"
  calcF! :: Function
end

"""
    Params(grid, β, eta::Function, μ, ν, nν, calcF)

Return the `params` for barotropic QL QG problem on `grid` and with topographic PV prescribed
as function, `eta(x, y)`.
"""
function Params(grid::AbstractGrid{T, A}, β, eta::Function, μ, ν, nν, calcF) where {T, A}
  eta_on_grid = FourierFlows.on_grid(eta, grid)
  etah_on_grid = rfft(eta_on_grid)

  return Params(β, A(eta_on_grid), A(etah_on_grid), μ, ν, nν, calcF)
end

# ---------
# Equations
# ---------

"""
    Equation(params, grid)

Return the equation for two-dimensional barotropic QG QL problem with parameters `params` and
on `grid`. Linear operator ``L`` includes bottom drag ``μ``, (hyper)-viscosity of order ``n_ν``
with coefficient ``ν`` and the ``β`` term:
```math
L = - μ - ν |𝐤|^{2 n_ν} + i β k_x / |𝐤|² .
```
Nonlinear term is computed via `calcN!` function.
"""
function Equation(params::Params, grid::AbstractGrid)
  L = @. - params.μ - params.ν * grid.Krsq^params.nν + im * params.β * grid.kr * grid.invKrsq
  CUDA.@allowscalar L[1, 1] = 0

  return FourierFlows.Equation(L, calcN!, grid)
end

# ----
# Vars
# ----

"""
    Vars{Aphys, Atrans, F, P}(u, v, U, uzeta, vzeta, zeta, Zeta, psi, Psi, N, NZ, uh, vh, Uh, zetah, Zetah, psih, Psih, Fh, prevsol)

The variables for barotropic QL QG:

$(FIELDS)
"""
mutable struct Vars{Aphys, Atrans, F, P} <: BarotropicQGQLVars
    "x-component of small-scale velocity"
        u :: Aphys
    "y-component of small-scale velocity"
        v :: Aphys
    "x-component of large-scale velocity"
        U :: Aphys
    "small-scale u′ζ′"
    uzeta :: Aphys
    "small-scale v′ζ′"
    vzeta :: Aphys
    "small-scale relative vorticity"
     zeta :: Aphys
    "large-scale relative vorticity"
     Zeta :: Aphys
    "small-scale relative vorticity"
      psi :: Aphys
    "large-scale relative vorticity"
      Psi :: Aphys
    "small-scale nonlinear term"
       Nz :: Atrans
    "large-scale nonlinear term"
       NZ :: Atrans
    "Fourier transform of x-component of small-scale velocity"
       uh :: Atrans
    "Fourier transform of y-component of small-scale velocity"
       vh :: Atrans
    "Fourier transform of x-component of large-scale velocity"
       Uh :: Atrans
    "Fourier transform of small-scale relative vorticity"
    zetah :: Atrans
    "Fourier transform of large-scale relative vorticity"
    Zetah :: Atrans
    "Fourier transform of small-scale relative vorticity"
     psih :: Atrans
    "Fourier transform of large-scale relative vorticity"
     Psih :: Atrans
    "Fourier transform of forcing"
       Fh :: F
    "`sol` at previous time-step"
  prevsol :: P
end

const DecayingVars = Vars{<:AbstractArray, <:AbstractArray, Nothing, Nothing}
const ForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, Nothing}
const StochasticForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray}

"""
    DecayingVars(dev, grid)

Return the vars for unforced two-dimensional quasi-linear barotropic QG problem on device `dev`
and with `grid`.
"""
function DecayingVars(dev::Dev, grid::AbstractGrid) where Dev
  T = eltype(grid)
  @devzeros Dev T (grid.nx, grid.ny) u v U uzeta vzeta zeta Zeta psi Psi
  @devzeros Dev Complex{T} (grid.nkr, grid.nl) N NZ uh vh Uh zetah Zetah psih Psih

  return Vars(u, v, U, uzeta, vzeta, zeta, Zeta, psi, Psi, N, NZ, uh, vh, Uh, zetah, Zetah, psih, Psih, nothing, nothing)
end

"""
    ForcedVars(dev, grid)

Return the `vars` for forced two-dimensional quasi-linear barotropic QG problem on device
`dev` and with `grid`.
"""
function ForcedVars(dev::Dev, grid::AbstractGrid) where Dev
  T = eltype(grid)
  @devzeros Dev T (grid.nx, grid.ny) u v U uzeta vzeta zeta Zeta psi Psi
  @devzeros Dev Complex{T} (grid.nkr, grid.nl) N NZ uh vh Uh zetah Zetah psih Psih Fh

  return Vars(u, v, U, uzeta, vzeta, zeta, Zeta, psi, Psi, N, NZ, uh, vh, Uh, zetah, Zetah, psih, Psih, Fh, nothing)
end

"""
    StochasticForcedVars(dev, grid)

Return the `vars` for stochastically forced two-dimensional quasi-linear barotropic QG problem
on device `dev` and with `grid`.
"""
function StochasticForcedVars(dev::Dev, grid::AbstractGrid) where Dev
  T = eltype(grid)
  @devzeros Dev T (grid.nx, grid.ny) u v U uzeta vzeta zeta Zeta psi Psi
  @devzeros Dev Complex{T} (grid.nkr, grid.nl) N NZ uh vh Uh zetah Zetah psih Psih Fh prevsol

  return Vars(u, v, U, uzeta, vzeta, zeta, Zeta, psi, Psi, N, NZ, uh, vh, Uh, zetah, Zetah, psih, Psih, Fh, prevsol)
end

# -------
# Solvers
# -------

"""
    calcN_advection!(N, sol, t, clock, vars, params, grid)

Calculate the Fourier transform of the advection term for quasi-linear barotropic QG dynamics.
"""
function calcN_advection!(N, sol, t, clock, vars, params, grid)
  @. vars.zetah = sol
  CUDA.@allowscalar @. vars.zetah[1, :] = 0
  @. vars.Zetah = sol
  CUDA.@allowscalar @. vars.Zetah[2:end, :] = 0

  @. vars.uh =  im * grid.l  * grid.invKrsq * vars.zetah
  @. vars.vh = -im * grid.kr * grid.invKrsq * vars.zetah
  @. vars.Uh =  im * grid.l  * grid.invKrsq * vars.Zetah

  ldiv!(vars.zeta, grid.rfftplan, vars.zetah)
  ldiv!(vars.u, grid.rfftplan, vars.uh)
  ldiv!(vars.v, grid.rfftplan, vars.vh)

  ldiv!(vars.Zeta, grid.rfftplan, vars.Zetah)
  ldiv!(vars.U, grid.rfftplan, vars.Uh)

  @. vars.uzeta = vars.u * vars.zeta           # u*ζ
  @. vars.vzeta = vars.v * vars.zeta           # v*ζ

  mul!(vars.uh, grid.rfftplan, vars.uzeta)     # \hat{u*q}
  @. vars.NZ = -im * grid.kr * vars.uh         # -∂[u*q]/∂x
  mul!(vars.vh, grid.rfftplan, vars.vzeta)     # \hat{v*q}
  @. vars.NZ += - im * grid.l * vars.vh        # -∂[v*q]/∂y
  CUDA.@allowscalar @. vars.NZ[2:end, :] = 0

  @. vars.U = vars.U * vars.zeta                # U*ζ
  @. vars.u = vars.u * vars.Zeta                # u*Ζ
  @. vars.v = vars.v * vars.Zeta                # v*Ζ

  mul!(vars.uh, grid.rfftplan, vars.U + vars.u) # \hat{U*ζ + u*Ζ}
  @. vars.Nz = -im * grid.kr*vars.uh            # -∂[U*ζ + u*Ζ]/∂x
  mul!(vars.vh, grid.rfftplan, vars.v)          # \hat{v*Z}
  @. vars.Nz += - im * grid.l*vars.vh           # -∂[v*Z]/∂y
  CUDA.@allowscalar @. vars.Nz[1, :] = 0

  @. N = vars.NZ + vars.Nz

  return nothing
end

"""
    calcN!(N, sol, t, clock, vars, params, grid)

Calculate the nonlinear term, that is the advection term and the forcing,

```math
N = - \\widehat{𝖩(ψ, ζ + η)}^{\\mathrm{QL}} + F̂ .
```
"""
function calcN!(N, sol, t, clock, vars, params, grid)
  dealias!(sol, grid)

  calcN_advection!(N, sol, t, clock, vars, params, grid)

  addforcing!(N, sol, t, clock, vars, params, grid)

  return nothing
end

"""
    addforcing!(N, sol, t, clock, vars, params, grid)

When the problem includes forcing, calculate the forcing term ``F̂`` and add it to the
nonlinear term ``N``.
"""
addforcing!(N, sol, t, cl, vars::Vars, params, grid) = nothing

function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid)
  params.calcF!(vars.Fh, sol, t, clock, vars, params, grid)
  @. N += vars.Fh

  return nothing
end

function addforcing!(N, sol, t, clock, vars::StochasticForcedVars, params, grid)
  if t == clock.t # not a substep
    @. vars.prevsol = sol # sol at previous time-step is needed to compute budgets for stochastic forcing
    params.calcF!(vars.Fh, sol, t, clock, vars, params, grid)
  end
  @. N += vars.Fh

  return nothing
end

# ----------------
# Helper functions
# ----------------

"""
    updatevars!(sol, vars, params, grid)
    updatevars!(prob)

Update the `vars` of a problem `prob` that has `grid` and `params` with the solution in `sol`.
"""
function updatevars!(sol, vars, params, grid)
  dealias!(sol, grid)

  CUDA.@allowscalar sol[1, 1] = 0
  @. vars.zetah = sol
  CUDA.@allowscalar @. vars.zetah[1, :] = 0
  @. vars.Zetah = sol
  CUDA.@allowscalar @. vars.Zetah[2:end, :] = 0

  @. vars.Psih = -vars.Zetah * grid.invKrsq
  @. vars.psih = -vars.zetah * grid.invKrsq
  @. vars.uh = -im * grid.l  * vars.psih
  @. vars.vh =  im * grid.kr * vars.psih
  @. vars.Uh =  im * grid.l  * vars.Zetah * grid.invKrsq

  ldiv!(vars.zeta, grid.rfftplan, deepcopy(vars.zetah))
  ldiv!(vars.Zeta, grid.rfftplan, deepcopy(vars.Zetah))
  ldiv!(vars.psi, grid.rfftplan, vars.psih)
  ldiv!(vars.Psi, grid.rfftplan, vars.Psih)
  ldiv!(vars.u, grid.rfftplan, deepcopy(vars.uh))
  ldiv!(vars.v, grid.rfftplan, deepcopy(vars.vh))
  ldiv!(vars.U, grid.rfftplan, deepcopy(vars.Uh))

  return nothing
end

updatevars!(prob) = updatevars!(prob.sol, prob.vars, prob.params, prob.grid)

"""
    set_zeta!(prob, zeta)
    set_zeta!(sol, vars, grid, zeta)

Set the solution `sol` as the transform of `zeta` and update variables `vars` on the `grid`.
"""
function set_zeta!(sol, vars, params, grid, zeta)
  mul!(vars.zetah, grid.rfftplan, zeta)
  CUDA.@allowscalar vars.zetah[1, 1] = 0.0
  @. sol = vars.zetah

  updatevars!(sol, vars, params, grid)
  return nothing
end

set_zeta!(prob, zeta) = set_zeta!(prob.sol, prob.vars, prob.params, prob.grid, zeta)

"""
    energy(sol, grid)
    energy(prob)

Return the domain-averaged kinetic energy of `sol`.
"""
@inline energy(sol, grid::AbstractGrid) =
  0.5 * (parsevalsum2(grid.kr .* grid.invKrsq .* sol, grid)
        + parsevalsum2(grid.l .* grid.invKrsq .* sol, grid)) / (grid.Lx * grid.Ly)

energy(prob) = energy(prob.sol, prob.grid)

"""
    enstrophy(sol, grid, vars)
    enstrophy(prob)

Return the domain-averaged enstrophy of `sol`.
"""
function enstrophy(sol, grid::AbstractGrid, vars::AbstractVars)
  @. vars.uh = sol
  CUDA.@allowscalar vars.uh[1, 1] = 0

  return 0.5 * parsevalsum2(vars.uh, grid) / (grid.Lx * grid.Ly)
end
enstrophy(prob) = enstrophy(prob.sol, prob.grid, prob.vars)

"""
    dissipation(prob)
    dissipation(sol, vars, params, grid)

Return the domain-averaged energy dissipation rate. `nν` must be >= 1.
"""
@inline function dissipation(sol, vars, params, grid)
  @. vars.uh = grid.Krsq^(params.nν - 1) * abs2(sol)
  CUDA.@allowscalar vars.uh[1, 1] = 0

  return params.ν / (grid.Lx * grid.Ly) * parsevalsum(vars.uh, grid)
end

@inline dissipation(prob) = dissipation(prob.sol, prob.vars, prob.params, prob.grid)

"""
    work(prob)
    work(sol, vars, params, grid)

Return the domain-averaged rate of work of energy by the forcing, `params.Fh`.
"""
@inline function work(sol, vars::ForcedVars, grid)
  @. vars.uh = grid.invKrsq * sol * conj(vars.Fh)

  return 1 / (grid.Lx * grid.Ly) * parsevalsum(vars.uh, grid)
end

@inline function work(sol, vars::StochasticForcedVars, grid)
  @. vars.uh = grid.invKrsq * (vars.prevsol + sol) / 2 * conj(vars.Fh)

  return 1 / (grid.Lx * grid.Ly) * parsevalsum(vars.uh, grid)
end

@inline work(prob) = work(prob.sol, prob.vars, prob.grid)

"""
    drag(prob)
    drag(sol, vars, params, grid)

Return the extraction of domain-averaged energy by drag `μ`.
"""
@inline function drag(prob)
  sol, vars, params, grid = prob.sol, prob.vars, prob.params, prob.grid
  @. vars.uh = grid.invKrsq * abs2(sol)
  CUDA.@allowscalar vars.uh[1, 1] = 0

  return params.μ / (grid.Lx * grid.Ly) * parsevalsum(vars.uh, grid)
end

end # module


"""
*Gamble*

`Gamble` constructs a gamble object with probability vector `p` and outcome vector `v`.

- `p`: probability vector
- `v`: outcome vector

Constructor
````julia
Gamble(;p=[.5,.5], v=[10.0,0.0])
````
"""
mutable struct Gamble{T1,T2}
    p::T1
    v::T2
end

function Gamble(;p=[.5,.5], v=[10.0,0.0])
    return Gamble(p, v)
end


# Lebesgue measure

export Lebesgue
struct Lebesgue{X} <: AbstractMeasure end

Lebesgue(X) = Lebesgue{X}()

basemeasure(μ::Lebesgue{X}) where {X} = μ

isprimitive(::Lebesgue) = true

sampletype(::Lebesgue{X}) where{X} = X


struct AE
	encoder
	decoder
end

Flux.@functor AE

function (ae::AE)(x)
	z = ae.encoder(x)
	x̂ = ae.decoder(z)
	return x̂
end

# ae_loss already in svdd

function conv_ae_constructor(
	;idim=(2,2,1),
	zdim::Int=1,
	activation="relu",
	kernelsizes=(1,1),
	channels=(1,1),
	scalings=(1,1),
	init_seed=nothing,
	batchnorm=false,
	kwargs...
)
	# if seed is given, set it
	(init_seed !== nothing) ? Random.seed!(init_seed) : nothing
	# encoder
	encoder = Chain(
		conv_encoder(idim, zdim, kernelsizes, channels, scalings; activation=activation, batchnorm=batchnorm)...
	)
	# decoder
	decoder = Chain(
		conv_decoder(idim, zdim, reverse(kernelsizes), reverse(channels), reverse(scalings),
			activation=activation, vec_output=false, vec_output_dim=nothing, batchnorm=batchnorm)...,
			x->sigmoid.(x) # our dataset images are scaled from 0-1
	)
	# reset seed
	(init_seed !== nothing) ? Random.seed!() : nothing

	model = AE(encoder, decoder)
end

function ae_constructor(
	;idim::Int=1,
	zdim::Int=1,
	hdim::Int=64,
	activation = "relu",
	nlayers::Int=3,
	init_seed=nothing,
	kwargs...
	)
	(init_seed !== nothing) ? Random.seed!(init_seed) : nothing

	encoder = build_mlp(idim, hdim, zdim, nlayers, activation=activation, lastlayer="linear")

	decoder = build_mlp(zdim, hdim, idim, nlayers, activation=activation, lastlayer="linear")

	(init_seed !== nothing) ? Random.seed!() : nothing

	return AE(encoder, decoder)
end

function anomaly_score(ae::AE, real_input; dims=(1,2,3), to_testmode::Bool=true)
	(to_testmode == true) ? Flux.testmode!(ae) : nothing
	score = mse_score(ae(real_input), real_input, dims=dims)
	(to_testmode == true) ? Flux.testmode!(ae, false) : nothing
	return score
end

function anomaly_score_gpu(ae::AE, real_input; batch_size=64, dims=(1,2,3), to_testmode::Bool=true)
	real_input = Flux.Data.DataLoader(real_input, batchsize=batch_size)
	(to_testmode == true) ? Flux.testmode!(ae) : nothing
	ae = ae |> gpu
	scores = Array{Float32}([])
	for x in real_input
		x = x |> gpu
		score = mse_score(ae(x), x, dims=dims) #vec(Flux.sum((ae(x) .- x).^2, dims=dims))
		scores = cat(scores, score |> cpu, dims=1)
	end
	(to_testmode == true) ? Flux.testmode!(ae, false) : nothing
	return scores
end

function StatsBase.fit!(ae::AE, data, params)
	# prepare batches & loaders
	train_loader, val_loader = prepare_dataloaders(data, params)
	history = MVHistory()
	# prepare for early stopping
	best_val_loss = Inf

	patience = params.patience
	val_batches = length(val_loader)

	if ndims(first(train_loader)) == 2
		dims = 1
	elseif ndims(first(train_loader)) == 4
		dims = (1,2,3)
	else
		error("unknown data type (no 2D or 4D tensor)")
	end

	#model to gpu
	ae = ae |> gpu
	# ADAMW(η = 0.001, β = (0.9, 0.999), decay = 0) = Optimiser(ADAM(η, β), WeightDecay(decay))
	opt_ae = haskey(params, :decay) ? ADAMW(params.lr, (0.9, 0.999), params.decay) : ADAM(params.lr)

	best_ae = deepcopy(ae)

	# train "AE"
	progress = Progress(params.iters)
	ps_ae = Flux.params(AE)

	for (iter, X) in enumerate(train_loader)
		X = getobs(X)|>gpu
		loss_ae, back = Flux.pullback(ps_ae) do
			ae_loss(ae(X), X, dims=dims)
		end
		grad_ae = back(1f0)
		Flux.Optimise.update!(opt_ae, ps_ae, grad_ae)

		push!(history, :loss_ae, iter, loss_ae)

		next!(progress; showvalues=[
			(:iters, "$(iter)/$(params.iters)"),
			(:ae_loss, loss_ae)
			])

		if mod(iter, params.check_every) == 0
			total_val_loss = 0
			Flux.testmode!(ae)
			for X_val in val_loader
				X_val = X_val |> gpu
				total_val_loss += ae_loss(ae(X), X)
			end
			Flux.testmode!(ae, false)
			push!(history, :validation_ae_loss, iter, total_val_loss/val_batches)
			if total_val_loss < best_val_loss
				best_val_loss = total_val_loss
				patience = params.patience
				best_ae = deepcopy(ae)
			else
				patience -= 1
				if patience == 0
					@info "Stopped training of Autoencoder after $(iter) iterations"
					global iters = iter - params.check_every*params.patience
					break
				end
			end
		end
		if iter == params.iters
			global iters = params.iters
		end
	end
	return history, best_ae, sum(map(p->length(p), Flux.params(ae))), iters
end


using StabTrussTopOpt
sto = StabTrussTopOpt
using Makie
using Test

# function example_relaxed_stab_opt()

file_dir = joinpath(pwd(), "test", "instances")
mesh_file = "tim.json"
load_file = "tim_load_support_case.json"

mesh_file = joinpath(file_dir, mesh_file)
load_file = joinpath(file_dir, load_file)

# parse ground mesh, supp, load from json
par = sto.parse_TO_problem(mesh_file, load_file)

# draw ground mesh
init_sc = Makie.Scene()
sto.draw_truss!(init_sc, par.X, par.T, par.Bound)
sto.draw_load!(init_sc, par.X, par.F)
init_sc

# define TO problem
# Steel-S235
E = 2.1e8 # kN/m^2 (210GPa)
τ = 1.0
σ_c = 3.5e5 # kN/m^2 (350MPa)
σ_t = 3.5e5
problem = sto.TO_define(par, E=E, τ=τ, σ_c=σ_c, σ_t=σ_t)

opt_a, opt_q = sto.TO_solve_relaxed_stab(problem)

# a_max = maximum(opt_a)
# for i=1:length(a)
#     if opt_a[i] < 0.001 * a_max
#         opt_a[i] = 0
#     end
# end

# draw solution
final_sc = Scene()
sto.draw_truss!(final_sc, par.X, par.T, par.Bound, area = opt_a, line_width=1e4, stress=opt_q)
# TODO: plot stress color
sto.draw_load!(final_sc, par.X, par.F, load_scale=1e-2)
display(final_sc)
# end

# example_relaxed_stab_opt()


__precompile__()

module CFET_OrderConditions

using Giac
using Combinatorics

export commutator, LCCC
export gen_exp_tBt_derivatives_coeffs
export gen_C_derivatives_at_0
export gen_C_derivatives_at_0_in_terms_of_A
export gen_exp_tBt_derivatives_at_0_in_terms_of_A
export expand_commutators
export gen_CFET_order_conditions

export number_commutators_of_grade_equal_to
export number_commutators_of_grade_equal_or_less_than

export integer_vector_weighted, lyndon_words, gen_hall_basis_elements_of_grade
export gen_hall_basis_elements_of_grades_leq
export extend_by_rightmost_subwords
export gen_CFET_order_conditions_without_redundancies

const Commutator = Array{Int64,1}
LCCC{T} = Dict{Array{Array{Int64,1},1},T} # Linear Combination of Commutator Chains

function gen_exp_tBt_derivatives_coeffs(N::Integer)
    A = [Dict{Array{Int64,1},Int64}([]=>1),    #Id
         Dict{Array{Int64,1},Int64}([0]=>1)]
    for n = 2:N
        a0 = A[n]
        # a1 = a0*C:
        a1 = [vcat(key, 0)=>val for (key,val) in a0]
        # a1 = a1 + (d/dt)a0:
        for (p,c) in a0
            for j = 1:length(p)
                q = copy(p)
                q[j] += 1
                d = get(a1, q, 0)
                a1[q] = d+c
            end
        end
        push!(A, a1)
    end
    A
end

function gen_C_derivatives_at_0(N::Integer)
    A = Array{Dict{Array{Int64,1},Int64},1}[[Dict(vcat(zeros(Int64,k),1)=>1)] for k=1:N-1]
    for k = 1:N-1
        for l = 1:N-k-1
            a0 = A[k][l]
            a1 = Dict{Array{Int64,1},Int64}()
            for (p,c) in a0
                for j=1:length(p)
                    q = copy(p)
                    q[j] += 1
                    if q[end] != q[end-1]
                        d = get(a1, q, 0)
                        a1[q] = d+c
                    end
                end
            end
            push!(A[k], a1)
        end
    end
    C = [Dict{Array{Int64,1},Int64}([n]=>(n+1)) for n=0:N]
    for n = 2:N
        a1 = C[n+1]
        for k=1:n-1
            f = binomial(n,k+1)
            a0 = A[k][n-k-1+1]
            for (p,c) in a0
                d = get(a1, p, 0//1)
                a1[p] = d+f*c
            end
        end
        C[n+1] = a1
    end
    C
end

function gen_C_derivatives_at_0_in_terms_of_A{T}(N::Integer, b::Array{T,1})
    C_derivatives_at_0=gen_C_derivatives_at_0(N)
    C1 = [Dict{Array{Int64,1},T}() for n=0:N]
    for n = 0:N
        for (p,coeff) in C_derivatives_at_0[n+1]
            C1[n+1][p] = coeff * prod([b[p[j]+1]  for j =1:length(p)])
        end
    end
    C1
end

Id(T::Type; f=one(T)) = LCCC{T}(Array[Int64[]]=>f)
is_c_number(x) = length(x)==1 && length(first(keys(x))[1])==0

function mult{T}(x::LCCC{T}, y::Dict{Array{Int64,1},T})
    r = LCCC{T}()
    for (p,c) in x
        for (q, d) in y
            pq = copy(p)
            push!(pq, q)
            r[pq] = c*d
        end
    end
    r
end

function mult{T}(x::LCCC{T}, y::LCCC{T})
    if is_c_number(x)
        c = first(values(x))
        return LCCC{T}([p=>c*d for (p,d) in y])
    elseif is_c_number(y)
        c = first(values(y))
        return LCCC{T}([p=>c*d for (p,d) in x])
    else
        r = LCCC{T}()
        for (p,c) in x
            for (q, d) in y
                pq = vcat(p,q)
                r[pq] = c*d
            end
        end
        return r
    end
end

function add{T}(x::LCCC{T}, y::LCCC{T};
                fx=one(T), fy=one(T))
    r = copy(x)
    for (p,c) in y
        d = get(r, p, 0)
        r[p] = fy*c + fx*d
    end
    r
end

function gen_exp_tBt_derivatives_at_0_in_terms_of_A{T}(N::Integer, b::Array{T,1})
    C_derivatives_at_0_in_terms_of_A=gen_C_derivatives_at_0_in_terms_of_A(N-1, b)
    exp_tBt_derivatives_coeffs=gen_exp_tBt_derivatives_coeffs(N)
    r = [LCCC{T}(Array[Int64[]]=>one(T))] # Id
    for n = 1:N
        y = LCCC{T}()
        for (p,d) in exp_tBt_derivatives_coeffs[n+1]
            x = LCCC{T}(
            [Array[p1]=>d*d1 for (p1,d1) in C_derivatives_at_0_in_terms_of_A[p[1]+1]])
            for i = 2:length(p)
                x = mult(x, C_derivatives_at_0_in_terms_of_A[p[i]+1])
            end
            y = add(y, x)
        end
        push!(r,y)
    end
    r
end

function expand_commutators{T}(a::Array{Int64,1}, f::T; coeff_type=T)
    if length(a)>=3
        b = LCCC{coeff_type}(Array[[a[1]]]=>one(T))
        c = expand_commutators(a[2:end], f, coeff_type=coeff_type)
        return add(mult(b,c), mult(c,b), fy=-one(T))
    elseif length(a)==2
        return LCCC{coeff_type}(Array[[a[1]], [a[2]]]=>f, Array[[a[2]], [a[1]]]=>-f)
    else
        return LCCC{coeff_type}(Array[a]=>f)
    end
end

function expand_commutators{T}(a::LCCC{T})
    b = LCCC{T}()
    for (p,c) in a
        q = expand_commutators(p[1], one(T), coeff_type=T)
        for i = 2:length(p)
            q = mult(q, expand_commutators(p[i], one(T), coeff_type=T))
        end
        b = add(b, q, fy=c)
    end
    b
end

multinomial_coeff(q::Int, k::Array{Int,1}) = div(factorial(q), prod([factorial(i) for i in k]))

function gen_CFET_order_conditions(N::Integer,J::Integer)
    b = [giac[giac(string("b",j,"_",n)) for n=0:N] for j=1:J]
    C = gen_exp_tBt_derivatives_at_0_in_terms_of_A(N, b[1])
    C = [expand_commutators(c) for c in C]
    C = [LCCC{giac}([key=>factor(val) for (key,val) in c]) for c in C]
    CC = Array[C]
    for j=2:J
        push!(CC, [LCCC{giac}([key=>subst(val,b[1],b[j]) for (key,val) in c]) for c in C])
    end
    Y = Dict{Array{Int64,1},giac}[]
    for q=0:N-1
        y = LCCC{giac}()
        for k0 in Combinatorics.WithReplacementCombinations(1:J,q+1)
            k = zeros(Int,J)
            for j in k0
                k[j]+=1
            end
            coeff = multinomial_coeff(q+1, k)
            x = Id(giac, f=giac(coeff))
            for j = J:-1:1
                x = mult(x, CC[j][k[j]+1])
            end
            y = add(y, x)
        end
        for k0 in Combinatorics.WithReplacementCombinations(1:J+1,q)
            k = zeros(Int,J+1)
            for j in k0
                k[j]+=1
            end
            coeff = multinomial_coeff(q, k)
            x = LCCC{giac}(Array[[k[J+1]]]=>coeff)
            for j = J:-1:1
                x = mult(x, CC[j][k[j]+1])
            end
            y = add(y, x, fy=-one(giac))

        end
        push!(Y, Dict{Array{Int64,1},giac}([[k[1] for k in key]=>factor(val) for (key,val) in y]))
    end
    Y
end

function number_commutators_of_grade_equal_to(n::Integer)
    if n==1
        return 1
    end
    f = factor(n)
    p = collect(keys(f))
    k = collect(values(f))
    nu = 0
    for l = 0:(2^length(p)-1)
        lb = bin(l,length(p))
        mu = (-1)^count_ones(l)
        pot = prod([p[j]^(k[j]-(lb[j]=='0'?0:1)) for j=1:length(p)])
        nu += mu*2^pot
    end
    div(nu,n)
end

number_commutators_of_grade_equal_or_less_than(n::Integer) =
    sum([num_of_eqs(k) for k=1:n])

# The following functions are stolen from MuPAD-Combinat,
# see http://mupad-combinat.sourceforge.net

# mupad/COMBINAT/words.mu/evaluationTable:
function evaluation_table(word::Vector)
    t = Dict{Any, Int}()
    for l in word
        if haskey(t, l)
            t[l] += 1
        else
            t[l] = 1
        end
    end
    t
end

# mupad/COMBINAT/words.mu/standard
function standard(word::Vector)
    t = evaluation_table(word)
    kk = sort(collect(keys(t)))
    offset = 0
    tmp = 0
    for k in kk
        tmp = t[k]
        t[k] = offset
        offset += tmp
    end
    r = zeros(Int, length(word))
    i = 1
    for l in word
        t[l] += 1
        r[i] = t[l]
        i += 1
    end
    r
end

# mupad/COMBINAT/integerVectorsWeighted.mu/
# (There it is a local function inside function list
function recfun(n::Int, l1::Vector{Int})
    result = []
    l = copy(l1)
    w = pop!(l)
    if length(l)==0
        d = div(n, w)
        if mod(n, w)==0
            return Vector{Int}[[d]]
        else
            return Vector{Int}[]
        end
    end
    for d = div(n, w):-1:0
        result = vcat(result, [vcat(x,d) for x in recfun(n-d*w, l)])
    end
    result
end

# mupad/COMBINAT/permutations.mu/mult2
function mult2(x::Vector{Int}, y::Vector{Int})
    nx = length(x)
    ny = length(y)
    # complete with fixed points if the permutations are not
    # of the same length
    x = vcat(x, nx+1:ny)
    y = vcat(y, ny+1:nx)
    result = x
    for i=1:length(x)
        result[y[i]] = x[i]
    end
    result
end

# mupad/COMBINAT/integerVectorsWeighted.mu/list
function integer_vector_weighted(n::Int, l::Vector{Int})
    if length(l)==0
        if n==0
            return Vector{Int}[[]]
        else
            return Vector{Int}[]
        end
    end

    # Heuristic: we sort the weights increasingly to limit
    # the backtracking, and reorder back the result afterward
    perm = standard(l)
    sort!(l)

    [mult2(x, perm) for x in recfun(n,l)]
end

# For the following 2 functions, see:
#   Sawada, Joe.  "A fast algorithm to generate necklaces with fixed content",
#   Source Theoretical Computer Science archive Volume 301 , Issue 1-3 (May 2003)
# and sage/src/sage/combinat/necklace.py

function simple_fixed_content1(res::Vector{Vector{Int}}, a::Vector{Int},
            content::Vector{Int}, t::Int, p::Int, k::Int)
    n = length(a)
    if t>n
        if n==p
            push!(res, a)
        end
    else
        for j=a[t-p]:k-1
            if content[j+1]!=0
                a[t] = j
                content[j+1] -= 1
                if j==a[t-p]
                    simple_fixed_content1(res, copy(a), content, t+1, p, k)
                else
                    simple_fixed_content1(res, copy(a), content, t+1, t, k)
                end
                content[j+1] += 1
            end
        end
    end
end

function simple_fixed_content(content::Vector{Int})
    a = zeros(Int, sum(content))
    content[1] -= 1
    k = length(content)
    res = Vector{Int}[]
    simple_fixed_content1(res, a, content, 2, 1, k)
    res
end

#The following function is stolen from sage/combinat/necklace.py

function lyndon_words(e::Vector{Int})
    k=1
    while e[k]==0
        k += 1
    end
    Vector{Int}[[i+k for i in z] for z in simple_fixed_content(e[k:end])]
end

#The following is stolen from http://mathoverflow.net/questions/97703/list-of-hall-basis

function gen_hall_basis_elements_of_grade(g::Int)
    L = integer_vector_weighted(g, collect(1:g))
    res = Vector{Int}[]
    for l in L
        for w in lyndon_words(l)
            push!(res, w-1)
        end
    end
    res
end

gen_hall_basis_elements_of_grades_leq(p) = vcat([gen_hall_basis_elements_of_grade(k) for k=1:p]...)

function gen_CFET_order_conditions_without_redundancies(N::Integer,J::Integer)
    C0 = gen_CFET_order_conditions(N, J)
    Dict{Array{Int, 1}, giac}[Dict{Array{Int, 1}, giac}(
        [b=>C0[k][b] for b in gen_hall_basis_elements_of_grade(k)]) for k=1:length(C0)]
end

coeff_rhs(w::Array{Int64,1}) = 1//prod([sum(w[j:end]+1) for j=1:length(w)])

function extend_by_rightmost_subwords(W::Array{Array{Int64,1},1})
    WW=Dict{Array{Int64,1},Int}(Int64[]=>1)
    for w in W
        for l=1:length(w)
            WW[w[l:end]] = 1
        end
    end
    return collect(keys(WW))
end

function gen_CFET_order_conditions(W::Array{Array{Int64,1},1}, J::Integer)
    W1 = extend_by_rightmost_subwords(W)

    qmax = maximum([maximum(w) for w in W])
    b = giac[giac(string("b",j, "_", q)) for q=0:qmax, j=1:J]
    bj = giac[giac(string("bj_", q)) for q=0:qmax]

    m=length(W1)
    M=zeros(giac, m, m)
    for k=1:m
        for l=1:m
            r = length(W1[k])-length(W1[l])
            if r>=0 && W1[k][r+1:end]==W1[l]
                w = W1[k][1:r]
                M[k,l]  = factor(prod(bj[w+1]))/factorial(length(w))
            end
        end
    end

    c = map(x->subst(x, bj, b[:,1]), M[:,findfirst(W1, Int64[])])
    for j=2:J
        Mj = map(x->subst(x, bj, b[:,j]), M)
        c = Mj*c
    end

    c = [c[findfirst(W1, w)] for w in W]

    eqs = [ c[j]-coeff_rhs(W[j]) for j=1:length(W) ]
    vars = vcat(b...)
    eqs, vars
end

end # module


"""
	unfold(ref, to_unf, samps=nothing; isomap=:default, optim=:default)

Unfold the input points `to_unf` based on the reference points informed in `ref`.
Returns a coordinate matrix with the unfolded points. Or a tuple of matrices if
`to_unf` is a list of points.

## Parameters:

* `to_unf`  - coordinate matrix with the points that will be unfolded
* `ref`     - coordinate matrix with the reference points for unfolding
* `isomap`  - additional parameters for isomap step; either :default or a
  NamedTuple with the keys shown below
* `optim`   - additional parameters for optimization step; either :default or a
  NamedTuple with the keys shown below

### `isomap` parameters:

*Default:* isomap = (search=:knn, neigh=16, anchors=1500, reftol=0.01)
* `search`  - search type to build neighbors graph for Isomap (:knn for
  k-nearest neighbor or :radius for radius search).
* `neigh`   - number of neighbors (for search=:knn) or radius distance
  (for search=:radius) to build neighbors graph for Isomap.
* `anchors` - number of anchors for landmark isomap
* `reftol`  - distance to which reference points are considered duplicates

### `optim` parameters:

*Default:* optim = (search=:knn, neigh=50, maxerr=5, nchunks=4)
* `search`  - search type to optimize locally (:knn for k-nearest neighbor or
  :radius for radius search).
* `neigh`   - number of neighbors (for search=:knn) or radius distance
  (for search=:radius) to local optimization.
* `maxerr`  - the maximum accepted distortion for unfolding optimization.
* `nchunks` - number of rounds of optimization.
"""
function unfold(to_unf::AbstractMatrix, ref::AbstractMatrix;
	            isomap=:default, optim=:default)

	# read isomap and optim parameters or assign the defaults below
	ipars = (search=:knn, neigh=16, anchors=1500, reftol=0.01)
	opars = (search=:knn, neigh=50, maxerr=5, nchunks=4)

	ipars = updatepars(ipars, isomap)
	opars = updatepars(opars, optim)

	# conversions if necessary
	!(ref[1] isa Float64) && (ref = Float64.(ref))
	!(to_unf[1] isa Float64) && (to_unf = Float64.(to_unf))

	# random seed
	Random.seed!(1234567890)

	# pre-process reference points
	ref = remove_duplicates(ref, tol=ipars.reftol)
	resol   = get_resolution(ref)

	# do landmark isomap at reference points
	unf_ref = landmark_isomap(ref, ipars.search, ipars.neigh, ipars.anchors)

	# get initial guess
	normals, good = getnormals(ref, ipars.search, ipars.neigh)
	unf = firstguess(to_unf, view(ref,:,good), view(unf_ref,:,good), view(normals,:,good))

	# unfold points in random chunks
	ndom = size(to_unf,2)
	shuffled_ids = shuffle(1:ndom)
	nchunks = ceil(Int, ndom/opars.nchunks)
	chunks  = collect(Iterators.partition(shuffled_ids, nchunks))

	for (i, ids) in enumerate(chunks)
		# only reference surface as conditioner
		known_org = ref[:,good]
		known_unf = unf_ref[:,good]
		ids_to_unf = ids

		# add unfolded from previous chunks
		if i > 1
			extra_ids = vcat(chunks[1:(i-1)]...)
			extra_org = view(to_unf,:,extra_ids)
			extra_unf = view(unf,:,extra_ids)
			good_, bad_ = error_ids(extra_org, extra_unf, opars.search,
			                        opars.neigh, opars.maxerr)

			known_org = hcat(known_org, view(extra_org, :, good_))
			known_unf = hcat(known_unf, view(extra_unf, :, good_))
			ids_to_unf = Int.(union(ids, view(extra_ids, bad_)))
		end

		# unfold points
		initguess    = view(unf, :, ids_to_unf)
		chunk_to_unf = view(to_unf, :, ids_to_unf)
		unf[:,ids_to_unf] .= opt(known_org, known_unf, chunk_to_unf,
		                             opars.search, opars.neigh, initguess)
	end

	unf
end

function unfold(to_unf::AbstractVector, ref::AbstractMatrix;
	            isomap=:default, optim=:default)

	merged_to_unf = reduce(hcat, to_unf)
	unf = unfold(merged_to_unf, ref, isomap=isomap, optim=optim)
	j_ = cumsum(size.(to_unf, 2))
	i_ = [i == 1 ? 1 : j_[i-1]+1 for i in 1:length(j_)]
	[unf[:,i:j] for (i,j) in zip(i_,j_)]
end

"""
	unfold(to_unf, ref, unf_ref; optim=:default)

Unfold the input points `to_unf` based on points already unfolded `unf_ref` and
their correspondent in the original space `ref`. Useful to unfold samples after
the domain is already unfolded, for example. Returns a coordinate matrix with
the unfolded points. Or a tuple of matrices if `to_unf` is a list of points.

## Parameters:

* `to_unf`  - coordinate matrix with the points that will be unfolded
* `ref`     - coordinate matrix with the reference points for unfolding
* `unf_ref` - coordinate matrix with the unfolded reference points
* `optim`   - additional parameters for optimization step; either :default or a
  NamedTuple with the keys shown below

### `optim` parameters:

*Default:* optim = (search=:knn, neigh=50, maxerr=5, nchunks=1)
* `search`  - search type to optimize locally (:knn for k-nearest neighbor or
  :radius for radius search).
* `neigh`   - number of neighbors (for search=:knn) or radius distance
  (for search=:radius) to local optimization.
* `maxerr`  - the maximum accepted distortion for unfolding optimization.
* `nchunks` - number of rounds of optimization.
"""
function unfold(to_unf::AbstractMatrix, ref::AbstractMatrix, unf_ref::AbstractMatrix;
	optim=:default)

	# read optim parameters or assign the defaults below
	opars = (search=:knn, neigh=50, maxerr=5, nchunks=2)
	opars = updatepars(opars, optim)

	# conversions if necessary
	!(to_unf[1] isa Float64) && (to_unf = Float64.(to_unf))

	# unfold points in random chunks
	ndom = size(to_unf,2)
	shuffled_ids = shuffle(1:ndom)
	nchunks = ceil(Int, ndom/opars.nchunks)
	chunks  = collect(Iterators.partition(shuffled_ids, nchunks))
	unf = firstguess(to_unf, ref, unf_ref)

	for (i, ids) in enumerate(chunks)
		# only reference surface as conditioner
		known_org = ref
		known_unf = unf_ref
		ids_to_unf = ids

		# add unfolded from previous chunks
		if i > 1
			extra_ids = vcat(chunks[1:(i-1)]...)
			extra_org = view(to_unf,:,extra_ids)
			extra_unf = view(unf,:,extra_ids)
			good_, bad_ = error_ids(extra_org, extra_unf, opars.search,
									opars.neigh, opars.maxerr)

			known_org = hcat(known_org, view(extra_org, :, good_))
			known_unf = hcat(known_unf, view(extra_unf, :, good_))
			ids_to_unf = Int.(union(ids, view(extra_ids, bad_)))
		end

		# unfold points
		initguess    = view(unf, :, ids_to_unf)
		chunk_to_unf = view(to_unf, :, ids_to_unf)
		unf[:,ids_to_unf] .= opt(known_org, known_unf, chunk_to_unf,
									 opars.search, opars.neigh, initguess)
	end
	unf
end

function unfold(to_unf::AbstractVector, ref::AbstractMatrix,
	            unf_ref::AbstractMatrix; optim=:default)

	merged_to_unf = reduce(hcat, to_unf)
	unf = unfold(merged_to_unf, ref, unf_ref, optim=optim)
	j_ = cumsum(size.(to_unf, 2))
	i_ = [i == 1 ? 1 : j_[i-1]+1 for i in 1:length(j_)]
	[unf[:,i:j] for (i,j) in zip(i_,j_)]
end

function updatepars(default, newpars)
	if newpars != :default
		for (k,v) in zip(keys(newpars), values(newpars))
			default = Setfield.setindex(default, v, k)
		end
	end
	default
end


# fourier series
# Pkg.add("Javis")

using Colors, Luxor
function ground(args...)
    background("black")
    sethue("white")
    Javis.translate(-args[1].width/2+50, 0)
    Javis.scale(700, 150)
end

function wave(xs, ys, opacity, color)
    setline(1.5)
    setopacity(opacity)
    sethue(color)
    points = Vector{Luxor.Point}(undef, length(xs))
    for (x, y, i) in zip(xs, ys, 1:length(xs))
        points[i] = Luxor.Point(x, y)
    end
    Javis.prettypoly(points, :stroke, ()->())
end

function term(x, i)
    return 4/π * sin.(2π*(2i-1)*x)/(2i-1)
end

function sum_term(x, k)
    k == 0 && return zeros(length(x))
    return sum(term(x, i) for i in 1:k)
end

nframes = 300
frames_per_wave = 40
video = Video(800, 600)
Background(1:nframes, ground)

x = 0.0:0.001:1.0
k = 6
colors = [RGB(0.0, 1.0, 0.4), RGB(0, 1.0, 1.0), RGB(1.0, 0.25, 0.25), RGB(1.0, 1.0, 0.0), RGB(1.0, 0.5, 1.0), RGB(0.75, 0.75, 1.0)]
waves = Object[]
for i = 1:k
    frames = frames_per_wave*(i-1)+1:nframes
    push!(waves, Object(frames, (args...; y=term(x,i)) -> wave(x, y, 0.5, colors[i])))
    act!(waves[end], Action(5:frames_per_wave,
                    change(:y, term(x,i) => sum_term(x, i))))
end

sum_wave = Object(1:nframes, (args...; y=zeros(length(x)))->wave(x, y, 1.0, "white"))
for i = 1:k
    act!(sum_wave, Action(frames_per_wave*(i-1)+1:frames_per_wave*i,
                         change(:y, sum_term(x, i-1) => sum_term(x, i))))
end

render(video; pathname="fourier_1D.gif")


using BayesianOptimization, GaussianProcesses
const BO = BayesianOptimization
using Test, Random
include("acquisitionfunctions.jl")
include("acquisition.jl")
include("branin.jl")
include("warmstart.jl")
include("utils.jl")