@testset "Graph coloring" begin

function normal_49_states(;reverse_constraint = false)
    m = Model(with_optimizer(CS.Optimizer, keep_logs=true))
    num_colors = 20

    @variable(m, 1 <= states[1:50] <= num_colors, Int)
    @variable(m, 1 <= max_color <= num_colors, Int)

    washington,montana,maine,north_dakota,south_dakota,wyoming,wisconsin,idaho,vermont,minnesota,oregon,new_hampshire,
    iowa,massachusetts,nebraska,new_york,pennsylvania,connecticut,rhode_island,new_jersey,indiana,nevada,utah,california,ohio,
    illinois,washington_dc,delaware,west_virginia,maryland,colorado,kentucky,kansas,virginia,missouri,arizona,oklahoma,north_carolina,
    tennessee,texas,new_mexico,alabama,mississippi,georgia,south_carolina,arkansas,louisiana,florida,michigan = states

    @constraint(m, washington != oregon)
    @constraint(m, washington != idaho)
    @constraint(m, oregon != idaho)
    @constraint(m, oregon != nevada)
    @constraint(m, oregon != california)
    @constraint(m, california != nevada)
    @constraint(m, california != arizona)
    @constraint(m, nevada != idaho)
    @constraint(m, nevada != utah)
    @constraint(m, nevada != arizona)
    @constraint(m, idaho != montana)
    @constraint(m, idaho != wyoming)
    @constraint(m, idaho != utah)
    @constraint(m, utah != wyoming)
    @constraint(m, utah != colorado)
    @constraint(m, utah != new_mexico)
    @constraint(m, utah != arizona)
    @constraint(m, arizona != colorado)
    @constraint(m, arizona != new_mexico)
    @constraint(m, montana != north_dakota)
    @constraint(m, montana != south_dakota)
    @constraint(m, montana != wyoming)
    @constraint(m, wyoming != south_dakota)
    @constraint(m, wyoming != nebraska)
    @constraint(m, wyoming != colorado)
    @constraint(m, colorado != nebraska)
    @constraint(m, colorado != kansas)
    @constraint(m, colorado != oklahoma)
    @constraint(m, colorado != new_mexico)
    @constraint(m, new_mexico != oklahoma)
    @constraint(m, new_mexico != texas)
    @constraint(m, north_dakota != minnesota)
    @constraint(m, north_dakota != south_dakota)
    @constraint(m, south_dakota != minnesota)
    @constraint(m, south_dakota != iowa)
    @constraint(m, south_dakota != nebraska)
    @constraint(m, nebraska != iowa)
    @constraint(m, nebraska != missouri)
    @constraint(m, nebraska != kansas)
    @constraint(m, kansas != missouri)
    @constraint(m, kansas != oklahoma)
    @constraint(m, oklahoma != arkansas)
    @constraint(m, oklahoma != texas)
    @constraint(m, texas != arkansas)
    @constraint(m, texas != louisiana)
    @constraint(m, minnesota != wisconsin)
    @constraint(m, minnesota != iowa)
    @constraint(m, iowa != wisconsin)
    @constraint(m, iowa != illinois)
    @constraint(m, iowa != missouri)
    @constraint(m, missouri != illinois)
    @constraint(m, missouri != kentucky)
    @constraint(m, missouri != tennessee)
    @constraint(m, missouri != arkansas)
    @constraint(m, arkansas != tennessee)
    @constraint(m, arkansas != mississippi)
    @constraint(m, arkansas != louisiana)
    @constraint(m, louisiana != mississippi)
    @constraint(m, wisconsin != illinois)
    @constraint(m, wisconsin != michigan)
    @constraint(m, illinois != indiana)
    @constraint(m, illinois != kentucky)
    @constraint(m, kentucky != indiana)
    @constraint(m, kentucky != ohio)
    @constraint(m, kentucky != west_virginia)
    @constraint(m, kentucky != virginia)
    @constraint(m, kentucky != tennessee)
    @constraint(m, tennessee != virginia)
    @constraint(m, tennessee != north_carolina)
    @constraint(m, tennessee != georgia)
    @constraint(m, tennessee != alabama)
    @constraint(m, tennessee != mississippi)
    @constraint(m, mississippi != alabama)
    @constraint(m, michigan != indiana)
    @constraint(m, michigan != ohio)
    @constraint(m, indiana != ohio)
    @constraint(m, alabama != georgia)
    @constraint(m, alabama != florida)
    @constraint(m, ohio != pennsylvania)
    @constraint(m, ohio != west_virginia)
    @constraint(m, maine != new_hampshire)
    @constraint(m, new_hampshire != vermont)
    @constraint(m, new_hampshire != massachusetts)
    @constraint(m, vermont != massachusetts)
    @constraint(m, vermont != new_york)
    @constraint(m, massachusetts != new_york)
    @constraint(m, massachusetts != rhode_island)
    @constraint(m, massachusetts != connecticut)
    @constraint(m, connecticut != rhode_island)
    @constraint(m, new_york != connecticut)
    @constraint(m, new_york != pennsylvania)
    @constraint(m, new_york != new_jersey)
    @constraint(m, pennsylvania != new_jersey)
    @constraint(m, pennsylvania != delaware)
    @constraint(m, pennsylvania != maryland)
    @constraint(m, pennsylvania != west_virginia)
    @constraint(m, new_jersey != delaware)
    @constraint(m, maryland != washington_dc)
    @constraint(m, maryland != west_virginia)
    @constraint(m, maryland != virginia)
    @constraint(m, washington_dc != virginia)
    @constraint(m, west_virginia != virginia)
    @constraint(m, virginia != north_carolina)
    @constraint(m, north_carolina != south_carolina)
    @constraint(m, south_carolina != georgia)
    @constraint(m, georgia != florida)

    if reverse_constraint
        @constraint(m, max_color .>= states)
    else
        @constraint(m, states .<= max_color)
    end

    @objective(m, Min, max_color)

    optimize!(m)

    status = JuMP.termination_status(m)

    com = JuMP.backend(m).optimizer.model.inner

    CS.save_logs(com, "graph_color_optimize.json")
    rm("graph_color_optimize.json")

    @test status == MOI.OPTIMAL
    @test com.best_sol == 4 == JuMP.objective_value(m) == JuMP.objective_bound(m)
    @test maximum([JuMP.value(var) for var in states]) == 4 == JuMP.value(max_color)

    return com
end

@testset "49 US states + DC" begin
    com1 = normal_49_states(;reverse_constraint=true)
    com2 = normal_49_states()
    info_1 = com1.info
    info_2 = com2.info
    @test info_1.pre_backtrack_calls == info_2.pre_backtrack_calls
    @test info_1.backtrack_fixes == info_2.backtrack_fixes
    @test info_1.in_backtrack_calls == info_2.in_backtrack_calls
    @test info_1.backtrack_reverses == info_2.backtrack_reverses
    logs_1 = CS.get_logs(com1)
    logs_2 = CS.get_logs(com2)
    @test CS.same_logs(logs_1[:tree], logs_2[:tree])
end

@testset "49 US states + DC without sorting" begin
    m = Model(with_optimizer(CS.Optimizer, keep_logs=true, backtrack_sorting=false))
    num_colors = 8
    @variable(m, 1 <= max_color <= num_colors, Int)

    @variable(m, 1 <= states[1:50] <= num_colors, Int)
    washington,montana,maine,north_dakota,south_dakota,wyoming,wisconsin,idaho,vermont,minnesota,oregon,new_hampshire,
    iowa,massachusetts,nebraska,new_york,pennsylvania,connecticut,rhode_island,new_jersey,indiana,nevada,utah,california,ohio,
    illinois,washington_dc,delaware,west_virginia,maryland,colorado,kentucky,kansas,virginia,missouri,arizona,oklahoma,north_carolina,
    tennessee,texas,new_mexico,alabama,mississippi,georgia,south_carolina,arkansas,louisiana,florida,michigan = states

    @constraint(m, washington != oregon)
    @constraint(m, washington != idaho)
    @constraint(m, oregon != idaho)
    @constraint(m, oregon != nevada)
    @constraint(m, oregon != california)
    @constraint(m, california != nevada)
    @constraint(m, california != arizona)
    @constraint(m, nevada != idaho)
    @constraint(m, nevada != utah)
    @constraint(m, nevada != arizona)
    @constraint(m, idaho != montana)
    @constraint(m, idaho != wyoming)
    @constraint(m, idaho != utah)
    @constraint(m, utah != wyoming)
    @constraint(m, utah != colorado)
    @constraint(m, utah != new_mexico)
    @constraint(m, utah != arizona)
    @constraint(m, arizona != colorado)
    @constraint(m, arizona != new_mexico)
    @constraint(m, montana != north_dakota)
    @constraint(m, montana != south_dakota)
    @constraint(m, montana != wyoming)
    @constraint(m, wyoming != south_dakota)
    @constraint(m, wyoming != nebraska)
    @constraint(m, wyoming != colorado)
    @constraint(m, colorado != nebraska)
    @constraint(m, colorado != kansas)
    @constraint(m, colorado != oklahoma)
    @constraint(m, colorado != new_mexico)
    @constraint(m, new_mexico != oklahoma)
    @constraint(m, new_mexico != texas)
    @constraint(m, north_dakota != minnesota)
    @constraint(m, north_dakota != south_dakota)
    @constraint(m, south_dakota != minnesota)
    @constraint(m, south_dakota != iowa)
    @constraint(m, south_dakota != nebraska)
    @constraint(m, nebraska != iowa)
    @constraint(m, nebraska != missouri)
    @constraint(m, nebraska != kansas)
    @constraint(m, kansas != missouri)
    @constraint(m, kansas != oklahoma)
    @constraint(m, oklahoma != arkansas)
    @constraint(m, oklahoma != texas)
    @constraint(m, texas != arkansas)
    @constraint(m, texas != louisiana)
    @constraint(m, minnesota != wisconsin)
    @constraint(m, minnesota != iowa)
    @constraint(m, iowa != wisconsin)
    @constraint(m, iowa != illinois)
    @constraint(m, iowa != missouri)
    @constraint(m, missouri != illinois)
    @constraint(m, missouri != kentucky)
    @constraint(m, missouri != tennessee)
    @constraint(m, missouri != arkansas)
    @constraint(m, arkansas != tennessee)
    @constraint(m, arkansas != mississippi)
    @constraint(m, arkansas != louisiana)
    @constraint(m, louisiana != mississippi)
    @constraint(m, wisconsin != illinois)
    @constraint(m, wisconsin != michigan)
    @constraint(m, illinois != indiana)
    @constraint(m, illinois != kentucky)
    @constraint(m, kentucky != indiana)
    @constraint(m, kentucky != ohio)
    @constraint(m, kentucky != west_virginia)
    @constraint(m, kentucky != virginia)
    @constraint(m, kentucky != tennessee)
    @constraint(m, tennessee != virginia)
    @constraint(m, tennessee != north_carolina)
    @constraint(m, tennessee != georgia)
    @constraint(m, tennessee != alabama)
    @constraint(m, tennessee != mississippi)
    @constraint(m, mississippi != alabama)
    @constraint(m, michigan != indiana)
    @constraint(m, michigan != ohio)
    @constraint(m, indiana != ohio)
    @constraint(m, alabama != georgia)
    @constraint(m, alabama != florida)
    @constraint(m, ohio != pennsylvania)
    @constraint(m, ohio != west_virginia)
    @constraint(m, maine != new_hampshire)
    @constraint(m, new_hampshire != vermont)
    @constraint(m, new_hampshire != massachusetts)
    @constraint(m, vermont != massachusetts)
    @constraint(m, vermont != new_york)
    @constraint(m, massachusetts != new_york)
    @constraint(m, massachusetts != rhode_island)
    @constraint(m, massachusetts != connecticut)
    @constraint(m, connecticut != rhode_island)
    @constraint(m, new_york != connecticut)
    @constraint(m, new_york != pennsylvania)
    @constraint(m, new_york != new_jersey)
    @constraint(m, pennsylvania != new_jersey)
    @constraint(m, pennsylvania != delaware)
    @constraint(m, pennsylvania != maryland)
    @constraint(m, pennsylvania != west_virginia)
    @constraint(m, new_jersey != delaware)
    @constraint(m, maryland != washington_dc)
    @constraint(m, maryland != west_virginia)
    @constraint(m, maryland != virginia)
    @constraint(m, washington_dc != virginia)
    @constraint(m, west_virginia != virginia)
    @constraint(m, virginia != north_carolina)
    @constraint(m, north_carolina != south_carolina)
    @constraint(m, south_carolina != georgia)
    @constraint(m, georgia != florida)

    @constraint(m, max_color .>= states)

    @objective(m, Min, max_color)

    optimize!(m)

    status = JuMP.termination_status(m)

    com = JuMP.backend(m).optimizer.model.inner

    CS.save_logs(com, "graph_color_optimize.json")
    rm("graph_color_optimize.json")

    @test status == MOI.OPTIMAL
    @test com.best_sol == 4
    @test all([length(CS.values(m, var)) == 1 for var in states])
    @test maximum([JuMP.value(var) for var in states]) == JuMP.value(max_color) == 4

end

@testset "49 US states + DC only 3 colors" begin
    m = Model(with_optimizer(CS.Optimizer))
    num_colors = 3

    @variable(m, 1 <= states[1:50] <= num_colors, Int)
    washington,montana,maine,north_dakota,south_dakota,wyoming,wisconsin,idaho,vermont,minnesota,oregon,new_hampshire,
    iowa,massachusetts,nebraska,new_york,pennsylvania,connecticut,rhode_island,new_jersey,indiana,nevada,utah,california,ohio,
    illinois,washington_dc,delaware,west_virginia,maryland,colorado,kentucky,kansas,virginia,missouri,arizona,oklahoma,north_carolina,
    tennessee,texas,new_mexico,alabama,mississippi,georgia,south_carolina,arkansas,louisiana,florida,michigan = states

    @constraint(m, washington != oregon)
    @constraint(m, washington != idaho)
    @constraint(m, oregon != idaho)
    @constraint(m, oregon != nevada)
    @constraint(m, oregon != california)
    @constraint(m, california != nevada)
    @constraint(m, california != arizona)
    @constraint(m, nevada != idaho)
    @constraint(m, nevada != utah)
    @constraint(m, nevada != arizona)
    @constraint(m, idaho != montana)
    @constraint(m, idaho != wyoming)
    @constraint(m, idaho != utah)
    @constraint(m, utah != wyoming)
    @constraint(m, utah != colorado)
    @constraint(m, utah != new_mexico)
    @constraint(m, utah != arizona)
    @constraint(m, arizona != colorado)
    @constraint(m, arizona != new_mexico)
    @constraint(m, montana != north_dakota)
    @constraint(m, montana != south_dakota)
    @constraint(m, montana != wyoming)
    @constraint(m, wyoming != south_dakota)
    @constraint(m, wyoming != nebraska)
    @constraint(m, wyoming != colorado)
    @constraint(m, colorado != nebraska)
    @constraint(m, colorado != kansas)
    @constraint(m, colorado != oklahoma)
    @constraint(m, colorado != new_mexico)
    @constraint(m, new_mexico != oklahoma)
    @constraint(m, new_mexico != texas)
    @constraint(m, north_dakota != minnesota)
    @constraint(m, north_dakota != south_dakota)
    @constraint(m, south_dakota != minnesota)
    @constraint(m, south_dakota != iowa)
    @constraint(m, south_dakota != nebraska)
    @constraint(m, nebraska != iowa)
    @constraint(m, nebraska != missouri)
    @constraint(m, nebraska != kansas)
    @constraint(m, kansas != missouri)
    @constraint(m, kansas != oklahoma)
    @constraint(m, oklahoma != arkansas)
    @constraint(m, oklahoma != texas)
    @constraint(m, texas != arkansas)
    @constraint(m, texas != louisiana)
    @constraint(m, minnesota != wisconsin)
    @constraint(m, minnesota != iowa)
    @constraint(m, iowa != wisconsin)
    @constraint(m, iowa != illinois)
    @constraint(m, iowa != missouri)
    @constraint(m, missouri != illinois)
    @constraint(m, missouri != kentucky)
    @constraint(m, missouri != tennessee)
    @constraint(m, missouri != arkansas)
    @constraint(m, arkansas != tennessee)
    @constraint(m, arkansas != mississippi)
    @constraint(m, arkansas != louisiana)
    @constraint(m, louisiana != mississippi)
    @constraint(m, wisconsin != illinois)
    @constraint(m, wisconsin != michigan)
    @constraint(m, illinois != indiana)
    @constraint(m, illinois != kentucky)
    @constraint(m, kentucky != indiana)
    @constraint(m, kentucky != ohio)
    @constraint(m, kentucky != west_virginia)
    @constraint(m, kentucky != virginia)
    @constraint(m, kentucky != tennessee)
    @constraint(m, tennessee != virginia)
    @constraint(m, tennessee != north_carolina)
    @constraint(m, tennessee != georgia)
    @constraint(m, tennessee != alabama)
    @constraint(m, tennessee != mississippi)
    @constraint(m, mississippi != alabama)
    @constraint(m, michigan != indiana)
    @constraint(m, michigan != ohio)
    @constraint(m, indiana != ohio)
    @constraint(m, alabama != georgia)
    @constraint(m, alabama != florida)
    @constraint(m, ohio != pennsylvania)
    @constraint(m, ohio != west_virginia)
    @constraint(m, maine != new_hampshire)
    @constraint(m, new_hampshire != vermont)
    @constraint(m, new_hampshire != massachusetts)
    @constraint(m, vermont != massachusetts)
    @constraint(m, vermont != new_york)
    @constraint(m, massachusetts != new_york)
    @constraint(m, massachusetts != rhode_island)
    @constraint(m, massachusetts != connecticut)
    @constraint(m, connecticut != rhode_island)
    @constraint(m, new_york != connecticut)
    @constraint(m, new_york != pennsylvania)
    @constraint(m, new_york != new_jersey)
    @constraint(m, pennsylvania != new_jersey)
    @constraint(m, pennsylvania != delaware)
    @constraint(m, pennsylvania != maryland)
    @constraint(m, pennsylvania != west_virginia)
    @constraint(m, new_jersey != delaware)
    @constraint(m, maryland != washington_dc)
    @constraint(m, maryland != west_virginia)
    @constraint(m, maryland != virginia)
    @constraint(m, washington_dc != virginia)
    @constraint(m, west_virginia != virginia)
    @constraint(m, virginia != north_carolina)
    @constraint(m, north_carolina != south_carolina)
    @constraint(m, south_carolina != georgia)
    @constraint(m, georgia != florida)

    optimize!(m)

    @test JuMP.termination_status(m) == MOI.INFEASIBLE
end

@testset "Maximization objective" begin
    m = Model(with_optimizer(CS.Optimizer))
    num_colors = 20

    @variable(m, 1 <= states[1:50] <= num_colors, Int)
    @variable(m, 1 <= max_color <= num_colors, Int)

    washington,montana,maine,north_dakota,south_dakota,wyoming,wisconsin,idaho,vermont,minnesota,oregon,new_hampshire,
    iowa,massachusetts,nebraska,new_york,pennsylvania,connecticut,rhode_island,new_jersey,indiana,nevada,utah,california,ohio,
    illinois,washington_dc,delaware,west_virginia,maryland,colorado,kentucky,kansas,virginia,missouri,arizona,oklahoma,north_carolina,
    tennessee,texas,new_mexico,alabama,mississippi,georgia,south_carolina,arkansas,louisiana,florida,michigan = states

    @constraint(m, washington != oregon)
    @constraint(m, washington != idaho)
    @constraint(m, oregon != idaho)
    @constraint(m, oregon != nevada)
    @constraint(m, oregon != california)
    @constraint(m, california != nevada)
    @constraint(m, california != arizona)
    @constraint(m, nevada != idaho)
    @constraint(m, nevada != utah)
    @constraint(m, nevada != arizona)
    @constraint(m, idaho != montana)
    @constraint(m, idaho != wyoming)
    @constraint(m, idaho != utah)
    @constraint(m, utah != wyoming)
    @constraint(m, utah != colorado)
    @constraint(m, utah != new_mexico)
    @constraint(m, utah != arizona)
    @constraint(m, arizona != colorado)
    @constraint(m, arizona != new_mexico)
    @constraint(m, montana != north_dakota)
    @constraint(m, montana != south_dakota)
    @constraint(m, montana != wyoming)
    @constraint(m, wyoming != south_dakota)
    @constraint(m, wyoming != nebraska)
    @constraint(m, wyoming != colorado)
    @constraint(m, colorado != nebraska)
    @constraint(m, colorado != kansas)
    @constraint(m, colorado != oklahoma)
    @constraint(m, colorado != new_mexico)
    @constraint(m, new_mexico != oklahoma)
    @constraint(m, new_mexico != texas)
    @constraint(m, north_dakota != minnesota)
    @constraint(m, north_dakota != south_dakota)
    @constraint(m, south_dakota != minnesota)
    @constraint(m, south_dakota != iowa)
    @constraint(m, south_dakota != nebraska)
    @constraint(m, nebraska != iowa)
    @constraint(m, nebraska != missouri)
    @constraint(m, nebraska != kansas)
    @constraint(m, kansas != missouri)
    @constraint(m, kansas != oklahoma)
    @constraint(m, oklahoma != arkansas)
    @constraint(m, oklahoma != texas)
    @constraint(m, texas != arkansas)
    @constraint(m, texas != louisiana)
    @constraint(m, minnesota != wisconsin)
    @constraint(m, minnesota != iowa)
    @constraint(m, iowa != wisconsin)
    @constraint(m, iowa != illinois)
    @constraint(m, iowa != missouri)
    @constraint(m, missouri != illinois)
    @constraint(m, missouri != kentucky)
    @constraint(m, missouri != tennessee)
    @constraint(m, missouri != arkansas)
    @constraint(m, arkansas != tennessee)
    @constraint(m, arkansas != mississippi)
    @constraint(m, arkansas != louisiana)
    @constraint(m, louisiana != mississippi)
    @constraint(m, wisconsin != illinois)
    @constraint(m, wisconsin != michigan)
    @constraint(m, illinois != indiana)
    @constraint(m, illinois != kentucky)
    @constraint(m, kentucky != indiana)
    @constraint(m, kentucky != ohio)
    @constraint(m, kentucky != west_virginia)
    @constraint(m, kentucky != virginia)
    @constraint(m, kentucky != tennessee)
    @constraint(m, tennessee != virginia)
    @constraint(m, tennessee != north_carolina)
    @constraint(m, tennessee != georgia)
    @constraint(m, tennessee != alabama)
    @constraint(m, tennessee != mississippi)
    @constraint(m, mississippi != alabama)
    @constraint(m, michigan != indiana)
    @constraint(m, michigan != ohio)
    @constraint(m, indiana != ohio)
    @constraint(m, alabama != georgia)
    @constraint(m, alabama != florida)
    @constraint(m, ohio != pennsylvania)
    @constraint(m, ohio != west_virginia)
    @constraint(m, maine != new_hampshire)
    @constraint(m, new_hampshire != vermont)
    @constraint(m, new_hampshire != massachusetts)
    @constraint(m, vermont != massachusetts)
    @constraint(m, vermont != new_york)
    @constraint(m, massachusetts != new_york)
    @constraint(m, massachusetts != rhode_island)
    @constraint(m, massachusetts != connecticut)
    @constraint(m, connecticut != rhode_island)
    @constraint(m, new_york != connecticut)
    @constraint(m, new_york != pennsylvania)
    @constraint(m, new_york != new_jersey)
    @constraint(m, pennsylvania != new_jersey)
    @constraint(m, pennsylvania != delaware)
    @constraint(m, pennsylvania != maryland)
    @constraint(m, pennsylvania != west_virginia)
    @constraint(m, new_jersey != delaware)
    @constraint(m, maryland != washington_dc)
    @constraint(m, maryland != west_virginia)
    @constraint(m, maryland != virginia)
    @constraint(m, washington_dc != virginia)
    @constraint(m, west_virginia != virginia)
    @constraint(m, virginia != north_carolina)
    @constraint(m, north_carolina != south_carolina)
    @constraint(m, south_carolina != georgia)
    @constraint(m, georgia != florida)

    @constraint(m, max_color .<= states)

    @objective(m, Max, max_color)

    optimize!(m)

    status = JuMP.termination_status(m)

    com = JuMP.backend(m).optimizer.model.inner

    @test status == MOI.OPTIMAL
    @test com.best_sol == 17
    @test minimum([JuMP.value(var) for var in states]) == 17 == JuMP.value(max_color)
end

end


using SequentialMonteCarlo
using RNGPool
using SMCExamples.MVLinearGaussian
import SMCExamples.Particles.MVFloat64Particle
using StaticArrays
using Plots
Plots.gr()

setRNGs(0)

model, theta, ys, ko = MVLinearGaussian.defaultMVLGModel(2, 10)

nsamples = 2^14

smcio = SMCIO{model.particle, model.pScratch}(16, model.maxn, 1, true, 2.0)

v = Vector{MVFloat64Particle{2}}(undef, 10)
for p = 1:10
  v[p] = MVFloat64Particle{2}()
  v[p].x .= zeros(MVector{2, Float64})
end

function smoothingMeans(model, smcio, m, v)
  tmp = Vector{MVector{2,Float64}}(undef, 10)
  for p = 1:10
    tmp[p] = zeros(MVector{2, Float64})
  end
  for i = 1:m
    csmc!(model, smcio, v, v)
    for p = 1:10
      tmp[p] .+= v[p].x
    end
  end
  return tmp ./ m
end

@time out = smoothingMeans(model, smcio, nsamples, v)
plot((x -> x[1]).(out))
plot!((x -> x[1]).(ko.smoothingMeans), color=:red)

plot((x -> x[2]).(out))
plot!((x -> x[2]).(ko.smoothingMeans), color=:red)


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

"""
    DensityRatioWeighting(tdata, [vars]; [options])

Density ratio weights based on empirical distribution of
variables in target data `tdata`. Default to all variables.

## Optional parameters

* `estimator` - Density ratio estimator (default to `LSIF()`)
* `optlib`    - Optimization library (default to `default_optlib(estimator)`)

### Notes

Estimators from `DensityRatioEstimation.jl` are supported.
"""
struct DensityRatioWeighting <: WeightingMethod
  tdata
  vars
  dre
  optlib
end

function DensityRatioWeighting(tdata, vars=nothing; estimator=LSIF(),
                               optlib=default_optlib(estimator))
  validvars = collect(name.(variables(tdata)))
  wvars = isnothing(vars) ? validvars : vars
  @assert wvars ⊆ validvars "invalid variables ($wvars) for spatial data"
  DensityRatioWeighting(tdata, wvars, estimator, optlib)
end

function weight(sdata, method::DensityRatioWeighting)
  # retrieve method parameters
  tdata  = method.tdata
  vars   = method.vars
  dre    = method.dre
  optlib = method.optlib

  @assert vars ⊆ name.(variables(sdata)) "invalid variables ($vars) for spatial data"

  # numerator and denominator samples
  Ωnu = view(tdata, vars)
  Ωde = view(sdata, vars)
  xnu = collect(Tables.rows(values(Ωnu)))
  xde = collect(Tables.rows(values(Ωde)))

  # perform denstiy ratio estimation
  ratios = densratio(xnu, xde, dre, optlib=optlib)

  SpatialWeights(domain(sdata), ratios)
end


function featurize(df)
  X = zeros(size(df, 1), 0)
  function onehot(column::DataArray)
    categories = unique(column)
    p = length(categories)
    n = size(column, 1)
    result = zeros(n, p)
    for (i, x) in enumerate(column)
      result[i,:] = eye(p)[:,findfirst(categories, x)]
    end
    return result
  end

  for feature in names(df)
    column = df[feature]
    if eltype(column) <: Number
      X = [X (column - mean(column))/std(column)]
    else
      X = [X onehot(column)]
    end
  end
  return X
end

function class(column)
  result = zeros(length(column))
  categories = unique(column)
  @assert length(categories) == 2
  for (i, x) in enumerate(column)
    if findfirst(categories, x) == 1
      result[i] = -1
    else
      @assert findfirst(categories, x) == 2
      result[i] = 1
    end
  end
  return result
end


using JustJoshing
using Test
using ForwardDiff: derivative

const EPS = 0.0001

@testset "Call option consistency" begin
    # I: see code

    # II: Price of call option monotonically decreases with strike price
    @test all(derivative.(x->C(1,0.5;K=x),rand(100_000)*10) .<= EPS)

    # III: Call option should be between S and S - Ke(-r(T-t)) for all inputs
    #
    # (if call option is more expensive than stock, it is clearly better
    # to just buy the stock; if it is cheaper than the difference between
    # the current value of the stock and the discounted strike price, we
    # can arbitrage by selling the stock now, buying the option, putting the
    # strike price in a risk free bond, and then exercising the option
    # (so we end up with the stock back and some profit))
    #
    # NB: Joshi refers to time-to-maturity, i.e. t_joshi = T-t :)
    @test [
        let (S,t,K,r,σ) = rand(5), T = t + rand()
            S + EPS >= C(S,t;T=T,K=K,r=r,σ=σ) >= (S - K*exp(-r*(T-t))) - EPS
        end
    for _ in 1:100] |> all

    # IV: Price of call option should monotonically increase with volatility
    @test all(derivative.(x->C(1,0.5;σ=x),rand(1000)) .>= -EPS)

    # V: if d=0 (always true in our model) then price should increase as we get closer to expiry
    @test all(derivative.(T->C(1,0;T=T),rand(1000)) .>= -EPS)

    # VI: Call option should be a convex function of strike - i.e. gradient is an increasing function of strike
    #   TODO: work out if there's a more ergonomic way of calculating higher order derivatives
    @test all(derivative.(x->derivative.(K->C(1,0;K=K),x),rand(1000)) .>= -EPS)

    # VII: The price of a call-spread should approximate the price of a digital-call option
    # TODO - NOT IMPLEMENTED - 1) find out what call-spread is and price it; 2) find out to calculate price of digital call (it was really simple IIRC); 3) check they are roughly the same

    # VIII: The price of a digital-call option plus a digital-put option is equal to the price of a zero-coupon bond
    # TODO - Not implemented: see VII
end

@testset "Monte-Carlo validation" begin
    # Validation with Monte-Carlo
    # S-K: forward contract
    let S_t=100, t=0.01, K=100, r=0.02, σ=0.05, T=1, trials=100_000_000
        exact = S_t - K*exp(-r*(T-t))

        mc = mc_pricer(S->S-K;S_0=S_t,t=t,T=T,K=K,r=r,σ=σ,trials=trials) |> first

        @test mc ≈ exact rtol=1e-3
    end

    # max(S-K,0): call option
    let S_t=100, t=0.01, K=100, r=0.02, σ=0.05, T=1, trials=100_000_000
        bs = C(S_t,t;T=T,K=K,r=r,σ=σ)

        mc = mc_pricer(S->max(S-K,0);S_0=S_t,t=t,T=T,K=S_t,r=r,σ=σ,trials=trials) |> first

        @test mc ≈ bs rtol=1e-3
    end

    # Int(S>E): binary option
    let S_t=100, t=0.01, E=100, r=0.02, σ=0.05, T=1, trials=100_000_000
        exact = binary(S_t,t;E=E,r=r,σ=σ,T=T)

        mc = mc_pricer(S->Int(S>E);S_0=S_t,t=t,T=T,K=E,r=r,σ=σ,trials=trials) |> first

        @test mc ≈ exact rtol=1e-3
    end
end


using StochasticSemiDiscretizationMethod

function createSLDOProblem(A,B,ζ,α,β,σ)
    AMx =  ProportionalMX(@SMatrix [0. 1.;-A -2ζ]);
    τ1=2π
    BMx1 = DelayMX(τ1,@SMatrix [0. 0.; B 0.]);
    cVec = Additive(2)
    noiseID = 1
    αMx1 = stCoeffMX(noiseID,ProportionalMX(@SMatrix [0. 0.; α 0.]))
    βMx11 = stCoeffMX(noiseID,DelayMX(τ1,@SMatrix [0. 0.; β 0.]))
    σVec = stAdditive(1,Additive(@SVector [0., σ]))
    LDDEProblem(AMx,[BMx1],[αMx1],[βMx11],cVec,[σVec])
end

SLDOP_lddep=createSLDOProblem(1.,0.1,0.1,0.1,0.1,0.5); # LDDE problem for Hayes equation
method=SemiDiscretization(5,(2π+100eps())/10) # 5th order semi discretization with Δt=2π/10
τmax=2π # the largest τ of the system
# Second Moment mapping
mapping=DiscreteMapping_M2(SLDOP_lddep,method,τmax,n_steps=10,calculate_additive=true); #The discrete mapping of the system

@show spectralRadiusOfMapping(mapping); # spectral radius ρ of the mapping matrix (ρ>1 unstable, ρ<1 stable)
statM2=VecToCovMx(fixPointOfMapping(mapping), length(mapping.M1_Vs[1])); # stationary second moment matrix of the hayes equation (equilibrium position)
@show statM2[1,1] |> sqrt;

# Stability chart for the Delay Oscillator
using MDBM
using Plots
gr();
using LaTeXStrings

method=SemiDiscretization(5,2π/30);
τmax=2π+100eps()
idxs=[1,2,3:2:StochasticSemiDiscretizationMethod.rOfDelay(τmax,method)*2...]

# ζ=0.05, α=0.3*A, β=0.3*B
foo(A,B) = log(spectralRadiusOfMapping(DiscreteMapping_M2(createSLDOProblem(A,B,0.05,0.3*A,0.3*B,0.),method,τmax,idxs,
    n_steps=30),nev=8)); # No additive term calculated

axis=[Axis(-1.0:0.6:5.0,:A),
    Axis(LinRange(-1.5,1.5,12),:B)]

iteration=4;
stab_border_points=getinterpolatedsolution(solve!(MDBM_Problem(foo,axis),iteration));

scatter(stab_border_points...,xlim=(-1.,5.),ylim=(-1.5,1.4),
    label="",title="2nd moment stability border of the Delay Oscillator",xlabel=L"A",ylabel=L"$B$",
    guidefontsize=14,tickfont = font(10),markersize=2,markerstrokewidth=0)


@testset "sNES" begin

function assign_weights_wrapper(candi_ixs::Vector{Int})
    candidates = BlackBoxOptim.Candidate{Float64}[BlackBoxOptim.Candidate{Float64}([0.0], i, NaN) for i in candi_ixs]
    u = BlackBoxOptim.fitness_shaping_utilities_linear(length(candi_ixs))
    BlackBoxOptim.assign_weights!(similar(u), candidates, u)
end

@testset "assign_weights!()" begin
    @testset "when indices are already ordered" begin
        u = assign_weights_wrapper([1, 2])
        @test length(u) == 2
        @test isapprox(sum(u), 0.0; atol=1E-10)
        @test u[1] == 0.5
        @test u[2] == -0.5

        u = assign_weights_wrapper([1, 2, 3])
        @test length(u) == 3
        @test isapprox(sum(u), 0.0; atol=1E-10)
        @test u[1] == 1.0 - 1/3
        @test u[2] == -1/3
        @test u[3] == -1/3

        u = assign_weights_wrapper([1, 2, 3, 4])
        @test length(u) == 4
        @test isapprox(sum(u), 0.0; atol=1E-10)
        @test isapprox(u[1], 2/3-1/4)
        @test isapprox(u[2], 1/3-1/4)
        @test u[3] == -1/4
        @test u[4] == -1/4

        u = assign_weights_wrapper([1, 2, 3, 4, 5])
        @test length(u) == 5
        @test isapprox(sum(u), 0.0; atol=1E-10)
        @test isapprox(u[1], 2/3-1/5)
        @test isapprox(u[2], 1/3-1/5)
        @test u[3] == -1/5
        @test u[4] == -1/5
        @test u[5] == -1/5

        u = assign_weights_wrapper([1, 2, 3, 4, 5, 6])
        @test length(u) == 6
        @test isapprox(sum(u), 0.0; atol=1E-10)
        @test isapprox(u[1], (3/3)/(3/3+2/3+1/3)-1/6)
        @test isapprox(u[2], (2/3)/(3/3+2/3+1/3)-1/6)
        @test isapprox(u[3], (1/3)/(3/3+2/3+1/3)-1/6)
        @test u[4] == -1/6
        @test u[5] == -1/6
        @test u[6] == -1/6
    end

    @testset "when indices are not ordered" begin
        u = assign_weights_wrapper([4, 1, 2, 3])
        @test isapprox(sum(u), 0.0; atol=1E-10)
        @test isapprox(u[4], 2/3-1/4)
        @test isapprox(u[1], 1/3-1/4)
        @test u[2] == -1/4
        @test u[3] == -1/4
    end
end

end


using BAT
using Test
using Statistics
using StatsBase
using Distributions
using HypothesisTests

@testset "Gaussian Shell Distribution" begin
    #Tests different ways of instantiate GaussianShell Distribution
    @test @inferred(BAT.GaussianShell(r = 3., w = 0.1, n = 3)) isa BAT.GaussianShell
    @test @inferred(BAT.GaussianShell()) isa BAT.GaussianShell

    gs = BAT.GaussianShell(r = 3., w = 0.1, n = 3)

    @test @inferred(Base.length(gs)) == 3
    @test @inferred(Base.eltype(gs)) == Float64

    @test @inferred(StatsBase.params(gs)) == (3.0, 0.1, [0.0, 0.0, 0.0])

    @test isapprox(@inferred(Statistics.cov(gs)), [3.02718 0.00166016 0.0113455;
    0.00166016 3.01293 0.00453035; 0.0113455   0.00453035  3.01046], atol = 1e-1)

    @test isapprox(@inferred(Distributions._logpdf(gs, [0., 0., 0.])), -453.34571, atol = 1e-5)

    #In the limit when r is close to zero, the distribution should be a Gaussian
    gs = BAT.GaussianShell(r = 5e-10, w = 1., n = 1)
    ks_test = HypothesisTests.ExactOneSampleKSTest(rand(gs, 10^6)[:], Normal(0, 1.))
    @test pvalue(ks_test) > 0.05

end


@testset "6.6 Hyperbolic cosecant" begin
    include("6.6.1 (c+d x)^m (a+b csch)^n.jl")
    include("6.6.2 (e x)^m (a+b csch(c+d x^n))^p.jl")
    include("6.6.3 Hyperbolic cosecant functions.jl")
    include("6.6.7 (d hyper)^m (a+b (c csch)^n)^p.jl")
end

test/6 Hyperbolic functions/6.6 Hyperbolic cosecant/.jl


rationalize(0.9054054)
rationalize(0.518518)
rationalize(0.75)


"""This collection of methods and definitions allow for easy location determination and comparison between differnt points on any sphere defined by their latitude and longitude."""
module EarthSphere
abstract Coordinate
#A coordinate describes a point in any coordinate system.

abstract LatLong_Coord <: Coordinate
#This is a coordinate in a standard latitude-longitude coordinate system

type Degree
	#Abstraction for dealing with sexagesimal units of subtending the circle
	degree::Float64
	minute::Float64
	second::Float64
end
export Degree

type Sphere_LatLong_Coord <: LatLong_Coord
	#Deals with a spherical latlong coordinate, similar to that used on good old planet earth
	latitude::Float64
	longitude::Float64
	#Example bearing: SW, NE
	bearing::String
	radius::Float64 # 6371000 Meters for earth
end
function Earth_LatLong_Coord(latitude::Degree,longitude::Degree,bearing::String)
	latnumber::Float64 = degree_to_decimal_degree(latitude)
	longnumber::Float64 = degree_to_decimal_degree(longitude)
	Sphere_LatLong_Coord(latnumber,longnumber,bearing,6371)
end
export Earth_LatLong_Coord

Earth_LatLong_Coord(latitude::Float64,longitude::Float64,bearing::Float64) = Sphere_LatLong_Coord(latitude,longitude,bearing,6371)
function Earth_LatLong_Coord(latitude::Float64,longitude::Float64)
	if sign(latitude) == 1
		inclination = "N"
	elseif sign(latitude) == -1
		inclination = "S"
	else
		inclination = ""
	end
	if sign(longitude) == 1
		azimuth = "E"
	elseif sign(longitude) == -1
		azimuth = "W"
	else
		azimuth = ""
	end
	my_sphere = Sphere_LatLong_Coord(abs(latitude),abs(longitude),inclination*azimuth,6371)
	return my_sphere
end

type Spherical_Coord <: Coordinate
	#Mathematical formulation for a 3D spherical coordinate system
	radius::Float64
	theta::Float64
	phi::Float64
end
export Spherical_Coord

function sign_lat_long(latitudeangle::Float64,longitudeangle::Float64,bearing::String)
	"""Returns a signed version of latitude and longitude based on the bearing given,
	uses a coordinate system where east is positive, west is negative, north is positive
	south is negative."""
	"""This will break if you feed it a bearing like "ES" or "WN", but such a case would be
	trivial to build up, build against, and is also not idiomatic"""
	"""UPDATE: This will no longer break if you do that thanks to a cheap and poorly executed hack. Please feel free to make this clean"""
	if length(bearing) == 1
		if bearing[1] == 'S'
			latitudeangle = (-1)*latitudeangle
		elseif bearing[1] == 'W'
			longitudeangle = (-1)*longitudeangle
		else
			return (latitudeangle,longitudeangle)
		end
	end
	if bearing[1] == 'E' | bearing[1] == 'W'
		bearing[1] = temp
		bearing[1] = bearing[2]
		bearing[2] = bearing[1]
	end

	if bearing[1] == 'S'
		latitudeangle = (-1)*latitudeangle
	end
	if bearing[2] == 'W'
		longitudeangle = (-1)*longitudeangle
	end
	return (latitudeangle,longitudeangle)
end

function sign_lat_long(latlong::LatLong_Coord)
	"Signs a latlong coord using the general sign_lat_long"
	(newlat::Float64,newlong::Float64) = sign_lat_long(latlong.latitude,latlong.longitude,latlong.bearing)
end

function degree_to_decimal_degree(degree::Degree)
	"Takes a degree in sextagesimal and returns a decimal degree"
	return degree.degree + degree.minute/60 + degree.second/(60^2)
end

function convert(Spherical_Coord,latlong::LatLong_Coord)
	(signedlat,signedlong) = sign_lat_long(latlong.latitude,latlong.longitude,latlong.bearing)
	phi::Float64 = to_radians(signedlat)
	theta::Float64 =  to_radians(signedlong)
	return Spherical_Coord(latlong.radius,theta,phi)
end
export convert
function tuple(s::Sphere_LatLong_Coord)
	#Takes a Spherical Coordinate and returns the signed latitude and longitude
	first = s.longitude
	second = s.latitude
	if ismatch(r"S",s.bearing)
		first = (-1)*first
	end
	if ismatch(r"W",s.bearing)
		second = (-1)*second
	end
	return (first,second)
end

function to_radians(degree::Real)
	"Takes a degree, returns radians"
	_sign = sign(degree)
	reduced_degree::Float64 = mod(abs(degree),360)
	radians::Float64 = reduced_degree*(pi/180)
	return radians*_sign
end

function get_distance(c1::Spherical_Coord,c2::Spherical_Coord)
	"Gets the distance between two different points on a sphere."
	if c1.radius != c2.radius
		error("These two points are not located on the same sphere")
	end
	#We're using the Vincenty formula, because it appears to be very computationally and algorithmically reasonable.
	#=central_angle = spherical_law_of_cosines(c1,c2)=#
	central_angle::Float64 = vincenty(c1,c2)
	#=central_angle = haversine(c1,c2)=#
	#=central_angle = alternative_haversine(c1,c2)=#
	distance::Float64 = c1.radius*central_angle

	return distance
end
function get_distance(c1::Sphere_LatLong_Coord,c2::Sphere_LatLong_Coord)
	c1s = convert(Spherical_Coord,c1)
	c2s = convert(Spherical_Coord,c2)
	return get_distance(c1s,c2s)::Float64
end

export get_distance

function vincenty(c1::Spherical_Coord,c2::Spherical_Coord)
	#An implementation of the Vincenty formula: http://en.wikipedia.org/wiki/Great-circle_distance
	longdiff::Float64 = c1.theta-c2.theta
	#We could take absolute value of long (the difference of longitudes) but we only ever take the cosine of it, so we can just leave it (cosine is even, sign doesn't matter)
	vin_numerator::Float64 = ( (cos(c2.phi)*sin(longdiff))^2 + ( cos(c1.phi)*sin(c2.phi) - sin(c1.phi)*cos(c2.phi)*cos(longdiff))^2 )^(1/2)
	vin_denominator::Float64 = sin(c1.phi)*sin(c2.phi)+cos(c1.phi)*cos(c2.phi)*cos(longdiff)
	central_angle::Float64 = atan2(vin_numerator,vin_denominator)
	return central_angle
end
vincenty(c1,c2) = vincenty(convert(Spherical_Coord,c1),convert(Spherical_Coord,c2))
spherical_law_of_cosines(c1,c2) = spherical_law_of_cosines(convert(Spherical_Coord,c1),convert(Spherical_Coord,c2))
function spherical_law_of_cosines(c1::Spherical_Coord,c2::Spherical_Coord)
	#Implements Spherical Law of cosines, neither vincenty nor this are getting me close to google maps
	#UPDATE!: It's okay, I just accidentally had latitude for longitude and longitude for latitude, not really a big deal except it broke everything
	longdiff::Float64 = abs(c1.theta-c2.theta)
	central_angle::Float64 = acos(sin(c1.phi)*sin(c2.phi)+cos(c1.phi)*cos(c2.phi)*cos(longdiff))
	return central_angle
end
function haversine(c1::Spherical_Coord,c2::Spherical_Coord)
	longdiff::Float64 = c1.theta-c2.theta
	latdiff::Float64 = c1.phi-c2.phi
	a::Float64 = (sin(latdiff/2))^2 + cos(c1.phi)*cos(c2.phi)*(sin(longdiff/2))^2
	central_angle::Float64 = 2*atan2(sqrt(a),sqrt(1-a))
	return central_angle
end
function alternative_haversine(c1::Spherical_Coord,c2::Spherical_Coord)
	#This is another formulation of the haversine formula from wikipedia

	latdiff::Float64 = abs(c1.phi-c2.phi)
	longdiff::Float64 = abs(c1.theta-c2.theta)
	argument::Float64 = sqrt(sin((latdiff)/2)*sin((latdiff)/2) - cos(c1.phi)*cos(c2.phi)*sin((longdiff)/2)*sin((longdiff)/2))
	central_angle::Float64 = 2*asin(argument)
	return central_angle
end

alternative_haversine(c1,c2) =alternative_haversine(convert(Spherical_Coord,c1),convert(Spherical_Coord,c2))
haversine(c1,c2) = haversine(convert(Spherical_Coord,c1),convert(Spherical_Coord,c2))

end #End module


# sample from the Muller potential using BBK
using Printf
using Random
using LinearAlgebra
using ForwardDiff
using BasicMD
using Statistics

include("potentials.jl")

β = 0.1;
γ = 0.1;
M = [1.5, 0.75];
q₀ = [-0.75, 1.5];
p₀ = [0.0, 0.0];
Δt = 1e-2;
n_iters = 10^6;

V = x->Muller(x);
cfg = ForwardDiff.GradientConfig(V, q₀);
gradV! = (gradV, x)-> ForwardDiff.gradient!(gradV, V, x, cfg);

# second moment and energy obsrevables
f₁ = x-> x[1][1]^2;
f₂ = x-> V(x[1]);
observables = (f₁, f₂);

sampler = BBK(gradV!, β, γ, M, Δt)

Random.seed!(100);
Xvals = sample_trajectory([q₀, p₀], sampler, options=MDOptions(n_iters=n_iters,n_save_iters=n_save_iters));
f₁_estimate = mean(f₁.(Xvals));
f₂_estimate = mean(f₂.(Xvals));
@printf("f₁ estimate with %d samples: %g\n",length(Xvals), f₁_estimate);
@printf("f₂ estimate with %d samples: %g\n", length(Xvals), f₂_estimate);
Random.seed!(100);
observable_samples = sample_observables([q₀, p₀], sampler, observables, options=MDOptions(n_iters=n_iters,n_save_iters=n_save_iters));
@printf("f₁ estimate with %d samples: %g\n",length(observable_samples), mean(observable_samples[1,:]));
@printf("f₂ estimate with %d samples: %g\n", length(observable_samples), mean(observable_samples[2,:]));


using BenchmarkTools, ModelingToolkit

include(joinpath(pwd(), "src/models/car.jl"))

x = zeros(model.n)
u = zeros(model.m)
w = zeros(model.d)
h = 0.0
t = 0
q = zeros(model.nq)

@benchmark M_func($model, $q)
@benchmark C_func($model, $q, $q)
@benchmark fd($model, $x, $x, $u, $w, $h, $t)

n = 15
A = Diagonal(@SVector ones(n))

L(q) = transpose(q) * A * q

f(q) = ForwardDiff.gradient(L, q)
g(q) = 2.0 * A * q
x = rand(n)
f(x) == g(x)
@benchmark f($x)
@benchmark g($x)

# generate fast functions
@variables x_sym[1:n]

F = g(x_sym);
fast = eval(ModelingToolkit.build_function(F, x_sym,
            parallel=ModelingToolkit.MultithreadedForm())[1])
fast_sparsity = ModelingToolkit.sparsejacobian(F, x_sym)
∇fast! = eval(ModelingToolkit.build_function(fast_sparsity, x_sym,
            parallel=ModelingToolkit.MultithreadedForm())[2])
∇fast = similar(fast_sparsity, Float64)

y = ones(n)
@benchmark fast!($y, $x)
@benchmark fast($x)

l = L(x_sym);
fast = eval(ModelingToolkit.build_function([l], x_sym,
            parallel=ModelingToolkit.MultithreadedForm())[1])
fast_sparsity = ModelingToolkit.sparsejacobian([l], x_sym)
dfast = eval(ModelingToolkit.build_function(fast_sparsity, x_sym,
            parallel=ModelingToolkit.MultithreadedForm())[1])

l2 = dfast(x_sym)

fast2 = eval(ModelingToolkit.build_function(simplify.(l2), x_sym,
            parallel=ModelingToolkit.MultithreadedForm())[1])
fast2_sparsity = ModelingToolkit.sparsejacobian(simplify.(l2), x_sym)
∇fast2! = eval(ModelingToolkit.build_function(fast2_sparsity, x_sym,
            parallel=ModelingToolkit.MultithreadedForm())[2])
∇fast2 = similar(fast2_sparsity, Float64)

# pendulum

L(q, q̇) = 0.5 * 1.0 * q̇^2.0 - 1.0 * 9.81 * 1.0 * cos(q)


using CSV, StateSpaceModels, Plots, Statistics, Dates

# Load the AirPassengers dataset
AP = CSV.read("AirPassengers.csv")

# Take the log of the series
logAP = log.(Vector{Float64}(AP[:Passengers]))

# Plot the data
p1 = plot(AP[:Date], logAP, label = "Log-airline passengers", legend = :topleft, color = :black)

# Specify the state-space model
model = structural(logAP, 12)

# Estimate the state-space model
ss = statespace(model)

# Analyze its decomposition in trend and seasonal
p2 = plot(AP[:Date], [ss.smoother.alpha[:, 1] ss.smoother.alpha[:, 3]], layout = (2, 1),
            label = ["Trend component" "Seasonal component"], legend = :topleft)

# Forecast 24 months ahead
N = 24
pred, dist = forecast(ss, N)

# Define forecasting dates
firstdate = AP[:Date][end] + Month(1)
newdates = collect(firstdate:Month(1):firstdate + Month(N - 1))

p3 = plot!(p1, newdates, pred, label = "Forecast")


using Plots
using DifferentialEquations

labs = ["Prey", "Predator"];
function f(du,u,p,t)
    du[1] = (p[1]*u[1]) - (p[2]*u[1]*u[2])
    du[2] = (p[2]*u[1]*u[2]) - (p[3]*u[2])
end
u0 = [10.0;4.0]
p = [0.5;0.05;0.05]
tspan = (0.0, 500.0)
prob = ODEProblem(f,u0,tspan,p)

# You need to solve the problem.  i.e. you need to run the simulation:
res = solve(prob)

# Then you need to get the simulated values in a nice format for plotting
t = res.t
prey = [values[1] for values in res.u]
pred = [values[2] for values in res.u]

# Plot the dynamics as usual
p1 = plot(t,pred)
plot!(t,prey)

# Plot the phase-space diagram
p2 = plot(pred,prey)

# Combine the two plots side by side
plot(p1,p2,layout=(1,2),legend=false,linewidth=3 )


using StructuralCausalModels, Test

ProjDir = @__DIR__

a= [
  0    0    0    0;
  1    0    0    1;
  0    1    0    0;
  0    0    0    0
]

vars = [Symbol("n$i") for i in 1:size(a, 1)]
amat = NamedArray(Int.(a), (vars, vars), ("Rows", "Cols"));

dag = DAG("m_separation", amat)

m = maximize(ribbon_graph(amat, c=[:n1], m=[:n4]))

f = [:n1]
s = [:n4]
c = [:n3]

ms1 = m_separation(dag, f, s; debug=false)
ms2 = m_separation(dag, f, s; c=c)

@testset "m_separation" begin

  @test ms1 == true
  @test ms2 == false

end


using JuMP
import Ipopt
import NLopt

function P1(z)
    model = Model(Ipopt.Optimizer)
    JuMP.set_silent(model)
    JuMP.@variable(model, x1)
    JuMP.@variable(model, x2)
    JuMP.@NLexpression(model, nl_expr, z[1] * ((x1 - 3)^2 + (x2 - 1)^2) + z[2] * ((x1 - 1)^2 + (x2 - 1)^2))

    JuMP.@NLobjective(model, Min, nl_expr)
    JuMP.@constraint(model, c1, x1 + 2 * x2 <= 2)
    JuMP.@constraint(model, c2, x1 - 2 * x2 <= 2)
    JuMP.@constraint(model, c3, -x1 + 2 * x2 <= 2)
    JuMP.@constraint(model, c4, -x1 - 2 * x2 <= 2)
    JuMP.optimize!(model)
    return value.([x1, x2]), JuMP.objective_value(model)
end
function P2(v)
    model = Model(Ipopt.Optimizer)
    JuMP.set_silent(model)

    JuMP.@variable(model, z)
    JuMP.@variable(model, x1)
    JuMP.@variable(model, x2)
    JuMP.@objective(model, Min, z)
    JuMP.@constraint(model, c1, x1 + 2 * x2 <= 2)
    JuMP.@constraint(model, c2, x1 - 2 * x2 <= 2)
    JuMP.@constraint(model, c3, -x1 + 2 * x2 <= 2)
    JuMP.@constraint(model, c4, -x1 - 2 * x2 <= 2)
    # Gammax = Γ([x1, x2])
    # z[1] * ((x1 - 3)^2 + (x2 - 1)^2) + z[2] * ((x1 - 1)^2 + (x2 - 1)^2)
    JuMP.@NLconstraint(model, c5, (x1 - 3)^2 + (x2 - 1)^2 - z - v[1] <= 0)
    JuMP.@NLconstraint(model, c6, (x1 - 1)^2 + (x2 - 1)^2 - z - v[2] <= 0)
    JuMP.optimize!(model)
    return (value.([x1, x2]), value(z))
end
function dual_objective(u1, u2, u3, u4, w1, w2)
    model2 = Model(Ipopt.Optimizer)
    JuMP.set_silent(model2)
    JuMP.@variable(model2, x1)
    JuMP.@variable(model2, x2)
    JuMP.@NLexpression(model2, my_expr, w1 * (x1 - 3)^2 + (x2 - 1)^2 + w2 * (x1 - 1)^2 + (x2 - 1)^2)
    JuMP.@NLobjective(model2, Min,
        u1 * (x1 + 2 * x2 - 2) +
        u2 * (x1 - 2 * x2 - 2) +
        u3 * (-x1 + 2 * x2 - 2) +
        u4 * (-x1 - 2 * x2 - 2) +
        my_expr
    )
    JuMP.optimize!(model2)
    return JuMP.objective_value(model2)
end
# function D2(v)
#     model = Model(Ipopt.Optimizer)
#     @variable(model, u >= 0)
#     @variable(model, w1 >= 0)
#     @variable(model, w2 >= 0)
#     @NLobjective(
#         model,
#         Max,
#         (-(w1^2 + w2^2) / (4 * u) - u) + w1 * (1 - v[1]) + w2 * (1 - v[2])
#     )
#     @constraint(model, c1, w1 + w2 == 1)
#     optimize!(model)
#     return (value.([w1, w2]), value(u))
# end

# function D2(v)
#     model = Model(Ipopt.Optimizer)
#     @variable(model, u1 >= 0)
#     @variable(model, u2 >= 0)
#     @variable(model, u3 >= 0)
#     @variable(model, u4 >= 0)
#     @variable(model, w1 >= 0)
#     @variable(model, w2 >= 0)

#     @NLexpression(model, my_expr,
#         u3 * (u2 + u4 - u1 - u3 + (u1 + u2 - u3 - u4 - 2w2 - 6w1) / (2w1 + 2w2)) +
#         u1 * (u2 + u4 - u1 - u3 + (u3 + u4 - u1 - u2 + 2w2 + 6w1) / (2w1 + 2w2)) +
#         u4 * (-4 + u1 + u3 - u2 - u4 + (u1 + u2 - u3 - u4 - 2w2 - 6w1) / (2w1 + 2w2)) +
#         u2 * (-4 + u1 + u3 - u2 - u4 + (u3 + u4 - u1 - u2 + 2w2 + 6w1) / (2w1 + 2w2)) +
#         w1 * (-3 + (u3 + u4 - u1 - u2 + 2w2 + 6w1) / (2w1 + 2w2))^2 +
#         (1 / 2) * (-u1 + u2 - u3 + u4)^2 +
#         w2 * (-1 + (u3 + u4 - u1 - u2 + 2w2 + 6w1) / (2w1 + 2w2))^2)

#     @NLobjective(
#         model,
#         Max,
#         my_expr - w1 * v[1] - w2 * v[2]
#     )
#     @constraint(model, c1, w1 + w2 == 1)
#     optimize!(model)
#     return (value.([w1, w2]), value.([u1, u2, u3, u4]))
# end

function D2(v)
    model = Model(Ipopt.Optimizer)
    @variable(model, u1 >= 0)
    @variable(model, u2 >= 0)
    @variable(model, u3 >= 0)
    @variable(model, u4 >= 0)
    @variable(model, w1 >= 0)
    @variable(model, w2 >= 0)

    @NLexpression(model, my_expr,
        u3 * (u2 + u4 - u1 - u3 + (u1 + u2 - u3 - u4 - 2w2 - 6w1) / (2w1 + 2w2)) +
        u1 * (u2 + u4 - u1 - u3 + (u3 + u4 - u1 - u2 + 2w2 + 6w1) / (2w1 + 2w2)) +
        u4 * (-4 + u1 + u3 - u2 - u4 + (u1 + u2 - u3 - u4 - 2w2 - 6w1) / (2w1 + 2w2)) +
        u2 * (-4 + u1 + u3 - u2 - u4 + (u3 + u4 - u1 - u2 + 2w2 + 6w1) / (2w1 + 2w2)) +
        w1 * (-3 + (u3 + u4 - u1 - u2 + 2w2 + 6w1) / (2w1 + 2w2))^2 +
        (1 / 2) * (-u1 + u2 - u3 + u4)^2 +
        w2 * (-1 + (u3 + u4 - u1 - u2 + 2w2 + 6w1) / (2w1 + 2w2))^2)

    @NLobjective(
        model,
        Max,
        my_expr - w1 * v[1] - w2 * v[2]
    )
    @constraint(model, c1, w1 + w2 == 1)
    optimize!(model)
    return (value.([w1, w2]), value.([u1, u2, u3, u4]))
end
# function D2para(v)
#     model = Model(Ipopt.Optimizer)
#     @variable(model, u >= 0)
#     @variable(model, w1 >= 0)
#     @variable(model, w2 >= 0)
#     exprrr = ParameterD2()
#     @NLexpression(model,my_expr_obj,exprrr)
#     @NLobjective(
#         model,
#         Max,
#         exprrr - w1*v[1]-w2*v[2]
#     )
#     @constraint(model, c1, w1 + w2 == 1)
#     optimize!(model)
#     return (value.([w1, w2]), value(u))
# end

# function D2(v0::Vector)

#     function ExplicitFunction(uw::Vector, grad::Vector, v::Vector)
#         println("called my own OFunction")
#         if length(grad) > 0
#             grad[1] = 0
#             grad[2] = 0.5 / sqrt(x[2])
#         end
#         w1, w2, u1, u2, u3, u4 = uw
#         return u3 * (u2 + u4 - u1 - u3 + (u1 + u2 - u3 - u4 - 2w2 - 6w1) / (2w1 + 2w2)) +
#                u1 * (u2 + u4 - u1 - u3 + (u3 + u4 - u1 - u2 + 2w2 + 6w1) / (2w1 + 2w2)) +
#                u4 * (-4 + u1 + u3 - u2 - u4 + (u1 + u2 - u3 - u4 - 2w2 - 6w1) / (2w1 + 2w2)) +
#                u2 * (-4 + u1 + u3 - u2 - u4 + (u3 + u4 - u1 - u2 + 2w2 + 6w1) / (2w1 + 2w2)) +
#                w1 * (-3 + (u3 + u4 - u1 - u2 + 2w2 + 6w1) / (2w1 + 2w2))^2 +
#                (1 / 2) * (-u1 + u2 - u3 + u4)^2 +
#                w2 * (-1 + (u3 + u4 - u1 - u2 + 2w2 + 6w1) / (2w1 + 2w2))^2 - w1 * v[1] - w2 * v[2]
#     end
#     function ImplicitFunction(uw::Vector, grad::Vector, v::Vector)

#         if length(grad) > 0
#             grad[1] = 0
#             grad[2] = 0.5 / sqrt(x[2])
#         end
#         model2 = Model(Ipopt.Optimizer)
#         JuMP.set_silent(model2)
#         JuMP.@variable(model2, x1)
#         JuMP.@variable(model2, x2)
#         JuMP.@NLexpression(model2, my_expr, uw[1] * (x1 - 3)^2 + (x2 - 1)^2 + uw[2] * (x1 - 1)^2 + (x2 - 1)^2)
#         JuMP.@NLobjective(model2, Min,
#             uw[3] * (x1 + 2 * x2 - 2) +
#             uw[4] * (x1 - 2 * x2 - 2) +
#             uw[5] * (-x1 + 2 * x2 - 2) +
#             uw[6] * (-x1 - 2 * x2 - 2) +
#             my_expr
#         )
#         JuMP.optimize!(model2)
#         return JuMP.objective_value(model2) - uw[1] * v[1] - uw[2] * v[2]
#     end

#     function constraint(uw::Vector, _grad::Vector)
#         uw[1] + uw[2] - 1
#     end

#     opt = NLopt.Opt(:LN_COBYLA, 6) # LN_COBYLA LD_MMA
#     opt.lower_bounds = repeat([0], 6)
#     # opt.xtol_rel = 1e-2

#     opt.max_objective = (x, g) -> ImplicitFunction(x, g, v0)
#     NLopt.equality_constraint!(opt, constraint, 1e-8)

#     (minf, minx, ret) = NLopt.optimize(opt, repeat([0.9], 6))
#     numevals = opt.numevals # the number of function evaluations
#     w1, w2, u1, u2, u3, u4 = minx
#     return (value.([w1, w2]), value.([u1, u2, u3, u4]))
# end

function Γ(x)
    [(x[1] - 3)^2 + (x[2] - 1)^2; (x[1] - 1)^2 + (x[2] - 1)^2]
end


module Neural
    using Random
    using Plots
    include("functions.jl")

    mutable struct Layer
        W::AbstractArray{Float64} # weights
        b::Vector{Float64} # bias

        afun::Function     # activaion function
        dafun::Function    # derivative of the activation function

        z::Vector{Float64} # W*x+b
        a::Vector{Float64} # afun(z)

        dCda::Vector{Float64}
        dCdW::AbstractArray{Float64}
        dCdb::Vector{Float64}

        function Layer(inlen::Int64, outlen::Int64, afun::Function, dafun::Function)
            σ = 2/(inlen + outlen)

            new(rand(outlen, inlen) * σ, rand(outlen) * σ, afun, dafun, [zeros(outlen) for i in 1:3]..., zeros(outlen, inlen), zeros(outlen))
        end
    end

    mutable struct Net
        layers::Vector{Layer}
        C::Function
        dC::Function

        function Net(inlength::Int64, sizes::Vector{Int64}, activations::Vector; # activations can be eighter vector of symbols or vector of
                                                                                # Tuple{Function, Function} where the 2nd function is derivative of the 1st
                    C::Function = softmaxLikelihood, dC::Function = dSoftmaxLikelihood)

            layers = Vector{Layer}()
            ilen = inlength
            for (size, id) in zip(sizes, activations)
                push!(layers, Layer(ilen, size, getFunctions(id)...))
                ilen = size
            end

            new(layers, C, dC)
        end
    end

    @views function makechunks(X::AbstractArray, n::Integer)
        c = length(X) ÷ n
        return [X[1+c*k:(k == n-1 ? end : c*k+c)] for k = 0:n-1]
    end

    mean(x) = sum(x)/length(x)

    getFunctions(id::Symbol) = activations[id]

    getFunctions(f::Tuple{Function, Function}) = f

    function applyLayer!(layer::Layer, x::AbstractArray)
        layer.z = layer.W*x + layer.b
        layer.a = layer.afun(layer.z)

        return layer.a
    end

    function forward!(net::Net, x::AbstractArray)
        input = x
        map(layer -> input = applyLayer!(layer, input), net.layers)

        return input
    end

    function backward!(net::Net, x::AbstractArray, t::AbstractArray)
       lastLayer = net.layers[end]

       lastLayer.dCda =  lastLayer.dafun(lastLayer.z) * net.dC(lastLayer.a, t)
       lastLayer.dCdW += lastLayer.dCda * transpose(net.layers[end - 1].a)
       lastLayer.dCdb += lastLayer.dCda

       for i in (length(net.layers) - 1):-1:1
            layer = net.layers[i]
            next = net.layers[i + 1]
            prevout = i == 1 ? x : net.layers[i - 1].a

            layer.dCda = layer.dafun(layer.a) * transpose(next.W) * next.dCda
            layer.dCdW += layer.dCda * transpose(prevout)
            layer.dCdb += layer.dCda
       end
    end

    function updateLayer!(layer::Layer, α::Float64)
        layer.W -= α*layer.dCdW
        layer.b -= α*layer.dCdb

        layer.dCda = zeros(size(layer.dCda))
        layer.dCdb = zeros(size(layer.dCdb))
        layer.dCdW = zeros(size(layer.dCdW))
    end

    function identify(net::Net, x::AbstractArray)
        vec = forward!(net, x)
        id = argmax(vec)

        return id, vec
    end

    function accuracy(net::Net, x::AbstractArray, t::AbstractArray)
        identities = map(i -> identify(net, x[:, i])[1], 1:size(x, 2))
        cmp = [Int(identities[i] == argmax(t[:, i])) for i in 1:size(t, 2)]

        return sum(cmp)/size(t, 2)
    end

    function train!(net::Net, x::AbstractArray, t::AbstractArray, batchSize::Int64, epochs::Int64, α::Float64, αDecay::Float64 = 1.0)
    plotly()

    meanCosts = Vector{Float64}()
    accs = Vector{Float64}()

    testIndices = shuffle(1:size(x, 2))[1:2*batchSize]
    trainIndices = filter(i -> !(i in testIndices), 1:size(x, 2))

    for i in 1:epochs
        costs = Vector{Float64}()
        acc = 0
        indices = shuffle(trainIndices) # one column represents one data point

        for i in indices[1:batchSize]
            forward!(net, x[:, i])
            backward!(net, x[:, i], t[:, i])
            updateLayer!.(net.layers, α)

            output = forward!(net, x[:, i])
            push!(costs, net.C(output, t[:, i]))
        end
        acc = accuracy(net, x[:, testIndices], t[:, testIndices])
        meanCost = mean(costs)
        α /= αDecay

        push!(meanCosts, meanCost)
        push!(accs, acc);
        println("Epoch $i/$epochs:")
        println("\tMean cost: $meanCost")
        println("\tAccuracy: $(acc *100)%")
        println("\tLearning rate: $α")
    end

    plt = plot(1:epochs, meanCosts, xlabel = "Epoch", ylabel = "Mean cost", legend = false, line = [:scatter, :path], dpi = 600, size = (1440, 810), grid = false)
    plot!(plt, 1:epochs, accs, line = [:scatter, :path])
    display(plt)
    #savefig(plt, "plt.png")
end

end


using LinearAlgebra
using Distributions
using SpecialFunctions
using PDMats
using CSV
#using Profile
#using ProfileView
using TimerOutputs
include("kernel.jl")
include("quadrature.jl")
include("Derivatives.jl")
include("TDist.jl")
include("examples.jl")
include("transforms.jl")
include("settings.jl")
include("tensorgrid.jl")
include("priors.jl")

"""
Obtain Bayesian predictive density function p(z0|z) and cumulative density function P(z0|z)
by marginalizing out theta and lambda.

Arg \"quadtype\" refers to type of quadrature rule used to integrate out θ, options are
    1) Gaussian (does not use derivatives)
    2) Turan (uses higher derivatives)

Note: Gaussian quadrature is always used to integrate out λ
"""
function getBtgDensity(train::trainingData{A, B}, test::testingData{A}, rangeθ::B, rangeλ::B, transforms, quadtype = "Gaussian", priortype = "Uniform") where A<:Array{Float64, 2} where B<:Array{Float64, 1}
    if quadtype == "Turan"
        nodesWeightsθ = getTuranQuadratureData() #use 12 Gauss-Turan integration nodes and weights by default
    elseif quadtype == "Gaussian"
        nodesWeightsθ = getGaussQuadraturedata()
    else
        throw(ArgumentError("Quadrature rule not recognized"))
    end

    nodesWeightsθ = getGaussQuadraturedata()
    affineTransformNodes(nodesWeightsθ, rangeθ)
    #always use Gauss quadrature to marginalize out λ
    nodesWeightsλ = getGaussQuadraturedata()
    affineTransformNodes(nodesWeightsλ, rangeλ)
    priorθ = initialize_prior(rangeθ, priortype); priorλ = initialize_prior(rangeλ, priortype);
    (pdf, cdf) = getTensorGrid(train, test, priorθ, priorλ, nodesWeightsθ, nodesWeightsλ, transforms, quadtype)
end


porosity(image :: AbstractArray{Bool}) =
    count(iszero, image) / length(image)

function corrfn(image  :: AbstractArray{Bool},
                cutoff :: Integer)
    return Directional.s2(image, false; periodic = true, len = cutoff) |> mean
end

function likelihood_point(x  :: AbstractFloat,
                          xs :: AbstractVector{<:AbstractFloat})
    σ = var(xs)
    μ = mean(xs)

    return -log(sqrt(2π * σ)) - (x - μ)^2/(2σ)
end

function likelihood(cf, size,
                    porosity   :: AbstractFloat,
                    dist       :: Distribution;
                    ncloud     :: Integer = 100)
    cutoff = length(cf)
    cloud = mapreduce(vcat, 1:ncloud) do _
        corrfn(generate_image(size, porosity, dist), cutoff)'
    end

    cloud_points = (cloud[:, n] for n in 1:cutoff)
    res = sum(likelihood_point(x, xs) for (x, xs) in zip(cf, cloud_points))
    return isnan(res) ? -Inf : res
end

function likelihood(image  :: AbstractArray{Bool},
                    dist   :: Distribution;
                    ncloud :: Integer = 100,
                    cutoff :: Integer = typemax(Int))
    likelihood(corrfn(image, cutoff), size(image), porosity(image), dist;
               ncloud = ncloud)
end


@testset "DouglasRachford" begin
    # Though this seems a strange way, it is a way to compute the mid point
    M = Sphere(2)
    p = SnPoint([1., 0., 0.])
    q = SnPoint([0., 1., 0.])
    r = SnPoint([0., 0., 1.])
    start = SnPoint([0.,0.,1.])
    result = midPoint(M,p,q)
    F(x) = distance(M,x,p)^2 + distance(M,x,q)^2
    prox1 = (η,x) -> proxDistance(M,η,p,x)
    prox2 = (η,x) -> proxDistance(M,η,q,x)
    xHat = DouglasRachford(M,F,[prox1,prox2],start)
    @test F(start) > F(xHat)
    @test distance(M,xHat,midPoint(M,p,q)) ≈ 0
    # but we can also compute the riemannian center of mass (locally) on Sn
    # though also this is not that usweful, but easy to test that DR works
    F2(x) = distance(M,x,p)^2 + distance(M,x,q)^2 + distance(M,x,r)^2
    prox3 = (η,x) -> proxDistance(M,η,r,x)
    xHat2, drec2 = DouglasRachford(M,F2,[prox1,prox2,prox3],start;
    debug = [DebugCost(), DebugIterate(), DebugProximalParameter(),100],
    record = [RecordCost(), RecordProximalParameter()]
    )
    xCmp = SnPoint(1/sqrt(3)*ones(3))
    # since the default does not run that long -> rough estimate
    @test distance(M,xHat2,xCmp) ≈ 0 atol = 10^(-5)
end


push!(LOAD_PATH,"../src")

using Test
using TypesParticles
using TypesBasis
using BuildBasis
using BuildOverlap
using InputParser
using KineticEnergy

"""
		run_buildkineticenergy_test()

This is a function to test the user types in BuildOverlap.jl. In asserts overlap integral calculationfor two hydrogen atoms is the same as the output from Prof. James John's matlab code.

""" function run_buildkineticenergy_test()
    @testset "Kinetic Energy" begin
        system,basisset = getcalcsetup("../examples/HatreeFock.in")
        basisfunc,numelec = buildbasisfunc(system,basisset);

        Scalc = buildelecoverlap(numelec,basisfunc,basisset);
        kecalc = buildkineticenergy(numelec,basisfunc,basisset);
        ketest = [1.39568e0 0.25974e0 0.75367e0 0.22767e0;
                  0.25974e0 0.24192e0 0.22767e0 0.22327e0;
                  0.75367e0 0.22767e0 1.39568e0 0.25974e0;
                  0.22767e0 0.22327e0 0.25974e0 0.24192e0];

        tolerance = 1.0e-5;
        for i=1:4
            for j=1:4
                @test sign(kecalc[i,j]) == sign(ketest[i,j]);
                diff = abs(kecalc[i,j]-ketest[i,j]);
                @test diff <= tolerance;
	    end
        end #i

    end #testset
end #run_kineticenergy_test()


using Distributions
using Bijectors
using Transits: Kipping13Transform
using HypothesisTests

@testset "Kipping13" begin
    N = 10000
    @test length(Kipping13()) == 2
    samps = rand(rng, Kipping13(), N)
    u1, u2 = eachrow(samps)

    @test all(@.(u1 + u2 < 1)) # flux is always positive
    @test all(@.(u1 + 2 * u2 > 0)) # flux is strictly monotonic increasing
    @test all(@.(u1 > 0)) # u1 always positive (flux is always positive)
    @test loglikelihood(Kipping13(), samps) == 0.0

    # make sure qs are Uniform
    q1 = @. (u1 + u2)^2
    q2 = @. 0.5 * u1 / (u1 + u2)
    for q in (q1, q2)
        t = ApproximateOneSampleKSTest(q, Uniform(0, 1))
        @test t.δ < 0.05
    end

end

@testset "Kipping13Transform" begin
    N = 1000
    @test bijector(Kipping13()) isa Kipping13Transform

    𝔅 = Kipping13Transform()
    xs = rand(Kipping13(), N)
    for x in eachcol(xs)
        y = 𝔅(x)
        xi = inv(𝔅)(y)
        @test xi ≈ x
    end

    ys = 2 .* randn(2, N)
    @test all(eachcol(ys)) do y
        u1, u2 = inv(𝔅)(y)

        u1 + u2 ≤ 1 &&
        u1 + 2 * u2 ≥ 0 &&
        u1 ≥ 0
    end
end


"""
Describe a single layer operator from the surface to a line.

```math
<v, Su> = ∫_γ dx v(x) ∫_Γ dy \frac{e^{-ikR}}{4πR} u(y)
```
"""
mutable struct SingleLayerTrace{T} <: MaxwellOperator3D
    gamma::T
end

function quaddata(operator::SingleLayerTrace,
  localtestbasis::LagrangeRefSpace,
  localtrialbasis::LagrangeRefSpace,
  testelements,
  trialelements)

  tqd = quadpoints(localtestbasis,  testelements,  (10,))
  bqd = quadpoints(localtrialbasis, trialelements, (8,))

  #return QuadData(tqd, bqd)
  return (tpoints=tqd, bpoints=bqd)
end

# Use numerical quadrature for now
# Note: basis integral is over triangle, test over line
function quadrule(op::SingleLayerTrace, g::LagrangeRefSpace, f::LagrangeRefSpace, i, τ, j, σ, qd)
    DoubleQuadStrategy(
        qd.tpoints[1,i],
        qd.bpoints[1,j]
    )
end

integrand(op::SingleLayerTrace, kernel, g, τ, f, σ) = f[1]*g[1]*kernel.green


import Pkg
Pkg.activate("..")

using Distributions: Exponential, DiscreteUniform, Poisson, Normal, DiscreteNonParametric, truncated
using CrimsonSkyline
using Random: seed!
using Logging
using StatsBase: mean

seed!(2021)

# source: http://people.reed.edu/~jones/141/Coal.html
const coal_mining_data = [
    4, 5, 4, 1, 0, 4, 3, 4, 0, 6, 3, 3,
    4, 0, 2, 6, 3, 3, 5, 4, 5, 3, 1, 4, 4, 1, 5, 5, 3, 4, 2, 5, 2,
    2, 3, 4, 2, 1, 3, 2, 2, 1, 1, 1, 1, 3, 0, 0, 1, 0, 1, 1, 0, 0,
    3, 1, 0, 3, 2, 2, 0, 1, 1, 1, 0, 1, 0, 1, 0, 0, 0, 2, 1, 0, 0,
    0, 1, 1, 0, 2, 3, 3, 1, 1, 2, 1, 1, 1, 1, 2, 4, 2, 0, 0, 0, 1,
    4, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 1
]

function changepoint_model(t, data)
    rate_1 = sample(t, "rate_1", Exponential(5.0))
    rate_2 = sample(t, "rate_2", Exponential(5.0))
    changepoint = sample(t, "changepoint", DiscreteUniform(2, length(data) - 1))
    for (i, d) in enumerate(data)
        i < changepoint ? observe(t, "data $i", Poisson(rate_1), d) : observe(t, "data $i", Poisson(rate_2), d)
    end
end

rate_1_proposal(t0, t1, params...) = propose(t1, "rate_1", truncated(Normal(t0["rate_1"].value, 0.5), 0.01, Inf))
rate_2_proposal(t0, t1, params...) = propose(t1, "rate_2", truncated(Normal(t0["rate_2"].value, 0.5), 0.01, Inf))
function cp_proposal(t0, t1, data)
    ld = length(data) - 1
    pt = t0["changepoint"].value
    if pt == 2
        propose(t1, "changepoint", DiscreteNonParametric([2, 3], [0.5, 0.5]))
    elseif pt == ld
        propose(t1, "changepoint", DiscreteNonParametric([ld - 1, ld], [0.5, 0.5]))
    else
        propose(t1, "changepoint", DiscreteNonParametric([pt - 1, pt, pt + 1], [1/3.0, 1/3.0, 1/3.0]))
    end
end

function main()
    @time results = mh(
        changepoint_model, [rate_1_proposal, rate_2_proposal, cp_proposal];
        params = (coal_mining_data,), burn = 1000, thin = 100, num_iterations = 11000, inverse_verbosity = 1000
    )
    ests = hpdi(results, 0.9, ["rate_1", "rate_2", "changepoint"])
    @info "90% posterior HPDIs: $ests"
    plot_marginal(results, "changepoint", "plots/changepoint", "changepoint.png")
end

main()


function plotmarginaline(fluxobs, posterior, σ2)

  linemeansigma = marginaline(posterior, σ2)

  # project flux obs on line to find minimum, maximum t for plotting purposes

  D = round(Int, length(posterior)/2)

  ā, b̄ = mean(posterior)[1:D], mean(posterior)[D+1:2*D]

  proj(x) = dot(ā, x-b̄) / dot(ā, ā)

  tproj = [proj(fluxobs[:,i]) for i=1:size(fluxobs, 2)]

  tmin, tmax = minimum(tproj), maximum(tproj)

  tmin, tmax = sign(tmin) * abs(tmin * 1.2), sign(tmax) * abs(tmax * 1.2)

  trange = collect(LinRange(tmin, tmax, 100))

  predictions     = map(linemeansigma, trange)
  meanpredictions = reduce(hcat, [p[1] for p in predictions])
  covpredictions  = reduce(hcat, [diag(p[2]) for p in predictions])

  figure()

  # verification
  # plot3D(ā[1].*trange .+ b̄[1], ā[2].*trange .+ b̄[2], ā[3].*trange .+ b̄[3], "c-", label="mean fit")
  # map(x -> plot3D(x...,"k."), meanpredictions)

  plot3D(fluxobs[1,:], fluxobs[2,:], fluxobs[3,:], "bo", label="observed total flux")

  plot3D(vec(meanpredictions[1,:]), vec(meanpredictions[2,:]), vec(meanpredictions[3,:]), "k-", label="mean fit")

  plot3D(vec(meanpredictions[1,:]) .+ 2*vec(sqrt.(covpredictions[1,:])),
    vec(meanpredictions[2,:]),
    vec(meanpredictions[3,:]), "k--", label="std2")

  plot3D(vec(meanpredictions[1,:]),
    vec(meanpredictions[2,:] .+ 2*vec(sqrt.(covpredictions[2,:]))),
    vec(meanpredictions[3,:]), "k--", label="std2")

  plot3D(vec(meanpredictions[1,:]),
    vec(meanpredictions[2,:]),
    vec(meanpredictions[3,:]) .+ 2*vec(sqrt.(covpredictions[3,:])), "k--", label="std2")

end


using Yao
using YaoExtensions

# arXiv: 1704.05018, table S2
function hydrogen_hamiltonian()
    Z1 = put(2,1=>Z)
    Z2 = put(2,2=>Z)
    X1 = put(2,1=>X)
    X2 = put(2,2=>X)
    0.011280*Z1*Z2 + 0.397936*Z1 + 0.397936*Z2 + 0.180931*X1*X2
end

using Flux: Optimise
function train!(circ, hamiltonian; optimizer, niter::Int=100)
     params = parameters(circ)
     dispatch!(circ, :random)
     for i=1:niter
         _, grad = expect'(hamiltonian, zero_state(nqubits(circ)) => circ)
         Optimise.update!(optimizer, params, grad)
         dispatch!(circ, params)
         println("Energy = $(expect(hamiltonian, zero_state(nqubits(h)) |> circ) |> real)")
     end
     return expect(hamiltonian, zero_state(nqubits(h)) |> circ)
end

h = hydrogen_hamiltonian()
c = variational_circuit(2)
emin_vqe = train!(c, h; optimizer=Optimise.ADAM(0.1))

using LinearAlgebra
emin = eigvals(Matrix(mat(h)))[1]
@assert isapprox(emin, emin, atol=1e-1)


export Sector

#= Sector (for a base shape) =#

mutable struct Sector <: Shape{2,4}  # M = 2K
    c::SVector{2,Float64}  # center of circle
    r::Float64  # radius of circle
    ϕ₀::Float64  # center angle bisecting sector: -π ≤ ϕ₀ < π  (π excluded)
    ∆ϕ2::Float64  # "radius" in angle dimension: 0 ≤ ∆ϕ2 ≤ π (sector spans from ϕ₀ - ∆ϕ2 to ϕ₀ + ∆ϕ2)

    function Sector(c::AbstractVector{<:Real}, r::Real, ϕ::Real, ∆ϕ::Real)
        r≥0 || throw(ArgumentError("r = $r must be nonnegative."))
        -2π≤∆ϕ≤2π  || throw(ArgumentError("∆ϕ = $∆ϕ must be between -2π and 2π, inclusive."))

        ϕ₀ = rem(ϕ + ∆ϕ/2, 2π, RoundNearest)  # put ϕ₀ in [-π, π)
        ∆ϕ2 = abs(∆ϕ/2)

        return new(c, r, ϕ₀, ∆ϕ2)
    end
end

Base.:(==)(s1::Sector, s2::Sector) = s1.c==s2.c && s1.r==s2.r && s1.ϕ₀==s2.ϕ₀ && s1.∆ϕ2==s2.∆ϕ2
Base.isapprox(s1::Sector, s2::Sector) = s1.c≈s2.c && s1.r≈s2.r && s1.ϕ₀≈s2.ϕ₀ && s1.∆ϕ2≈s2.∆ϕ2
Base.hash(s::Sector, h::UInt) = hash(s.c, hash(s.r, hash(s.ϕ₀, hash(s.∆ϕ2, hash(:Sector, h)))))

distangle(ϕ::Real, ϕ₀:: Real) = rem(ϕ-ϕ₀, 2π, RoundNearest)  # ϕ measured from ϕ₀; result within [-π, π)

function level(x::SVector{2,<:Real}, s::Sector)
    d = x - s.c
    ld = norm(d)
    ϕ = ld==0 ? s.ϕ₀ : atan(d[2], d[1])  # angle to x with respect to c

    return max(ld/s.r, abs(distangle(ϕ, s.ϕ₀)) / s.∆ϕ2) - 1.0
end

function surfpt_nearby(x::SVector{2,<:Real}, s::Sector)
    # Basically mimic the same function for Prism, but proceeds in the (ρ,ϕ) domain.
    d = x - s.c
    ld = norm(d)

    # Calculate the closest point in the ρ dimension and outward normal direciton there.
    r2 = s.r / 2
    ρ = ld - r2  # positive if closer to arc; negative if closer to center
    d̂ = ld ≤ Base.rtoldefault(Float64) * r2  ? SVector(cos(s.ϕ₀),sin(s.ϕ₀)) : normalize(d)

    surfρ = ρ<0 ? 0.0 : s.r  # scalar: closest point to x between center and perimeter point
    noutρ = copysign(1.0,ρ) * d̂  # SVector{2}: outward direction normal at surfρ

    absρ = abs(ρ)
    abs∆ρ = abs(r2 - absρ)  # radial distance between x and either center or perimeter, whichever closer to x

    onbndρ = abs∆ρ ≤ Base.rtoldefault(Float64) * r2  # basically r2 ≈ ρ but faster
    isoutρ = (r2 < absρ) || onbndρ

    # Calculate the closest point in the ϕ dimension and outward normal direciton there.
    ϕ = ld==0 ? 0.0 : distangle(atan(d[2], d[1]), s.ϕ₀)  # positive if closer to end side; negative if closer to start side

    ϕtemp = ϕ<0 ? s.ϕ₀ - s.∆ϕ2 : s.ϕ₀ + s.∆ϕ2
    surfϕ = SVector(cos(ϕtemp), sin(ϕtemp))  # SVector{2}: closest point in ϕ dimension
    cosθ = d̂⋅surfϕ  # cosine of angle between d and surfϕ
    ldcosθ = ld*cosθ  # ldconθ .* surfϕ is surface point
    ldsinθ = ld*sqrt(1-cosθ^2)  # always positive

    ϕtemp += copysign(π/2,ϕ)
    noutϕ = SVector(cos(ϕtemp), sin(ϕtemp))  # SVector{2}: outward direction normal at surfϕ

    absϕ = abs(ϕ)
    abs∆ϕ = abs(s.∆ϕ2 - absϕ)  # angular distance between x and closer side of sector

    onbndϕ = abs∆ϕ ≤ Base.rtoldefault(Float64) * s.∆ϕ2  # basically ∆ϕ2 ≈ ϕ but faster
    isoutϕ = (s.∆ϕ2 < absϕ) || onbndϕ

    # Pick the surface point and outward direction normal depending on the location of x.
    if isoutρ && isoutϕ  # x outside in both ρ and ϕ dimensions
        surf = surfρ .* surfϕ  # one of two end points of arc
        nout = (onbndρ && onbndϕ) ? (noutρ + noutϕ) : (d - surf)
        nout = normalize(nout)
    elseif !isoutρ && isoutϕ  # x inside in ρ dimension, but outside in ϕ dimension
        (surf, nout) = (absϕ < s.∆ϕ2+π/2) ? (ldcosθ .* surfϕ, noutϕ) : (SVector(0.0,0.0), d̂)
    elseif isoutρ && !isoutϕ  # x outside in ρ dimension, but inside in ϕ dimension
        (surf, nout) = (surfρ .* d̂, noutρ)
    else  # !isoutρ && !isoutϕ: x strictly inside sector
        (surf, nout) = (abs∆ρ ≤ ldsinθ) ? (surfρ .* d̂, noutρ) : (ldcosθ .* surfϕ, noutϕ)
    end

    return surf+s.c, nout
end

function bounds(s::Sector)
    # Find the minimum and maximum coordinates among the center and two ends of the arc.
    ϕ = SVector(s.ϕ₀ - s.∆ϕ2, s.ϕ₀ + s.∆ϕ2)  # [start angle, end angle]
    v = s.r .* [SVector(0.0,0.0) [cos.(ϕ) sin.(ϕ)]']  # [center, start point of arc, end point of arc]

    # Consider using the code below once https://github.com/JuliaArrays/StaticArrays.jl/issues/498
    # is resolved:
    # l = minimum(v, dims=Val(2))[:,1]
    # u = maximum(v, dims=Val(2))[:,1]

    xs = v[1,:]
    ys = v[2,:]

    xmin = minimum(xs)
    xmax = maximum(xs)

    ymin = minimum(ys)
    ymax = maximum(ys)

    # Consider the extreme points on the arc in the Cartesian directions.
    abs(distangle(0, s.ϕ₀)) ≤ s.∆ϕ2 && (xmax = max(xmax,s.r))  # sector contains +x-direction from center
    abs(distangle(π/2, s.ϕ₀)) ≤ s.∆ϕ2 && (ymax = max(ymax, s.r))  # sector contains +y-direction from center
    abs(distangle(π, s.ϕ₀)) ≤ s.∆ϕ2  && (xmin = min(xmin, -s.r))  # sector contains -x-direction from center
    abs(distangle(3π/2, s.ϕ₀)) ≤ s.∆ϕ2  && (ymin = min(ymin, -s.r))  # sector contains -y-direction from center

    return (SVector(xmin,ymin)+s.c, SVector(xmax,ymax)+s.c)
end


using NPZ, Flux, Setfield, DifferentialEquations
using Statistics

using Revise
using GraphFlux

cora = npzread("download/cora/cora.npz")

g = GraphTuple(nodes=cora["x"] ./ sum(cora["x"], dims=1),
			senders=cora["edgelist"][:, 1],
			receivers=cora["edgelist"][:, 2]) |> gpu

y = Flux.onehotbatch(cora["y"], 1:7) |> Array |> gpu
trainidx = cora["trainidx"] |> gpu
validx = cora["validx"] |> gpu

modelf(dim) = Chain(
			NodeLayer(Dropout(0.5)),
			NodeLayer(Dense(1433, dim, relu)),
			NodeLayer(Dropout(0.1)),
			# NodeLayer(x -> vcat(x, zero(x))),
			GraphNeuralODE(GCN(Dense(zeros(dim, dim), zeros(dim)), identity), (0.0f0, 1.0f0), Tsit5()),
			# GraphNeuralODE(GCN(Dense(zeros(dim*2, dim*2), zeros(dim*2)), identity), (0.0f0, 1.0f0), Tsit5()),
			# GraphNeuralODE(GCN(dim*2, dim*2), (0.0f0, 1.0f0), Tsit5()),
			# g -> relu.(g.nodes[1:dim, :]),
			g -> relu.(g.nodes),
			Dense(dim, 7))

function training_run()
	m = modelf(64) |> gpu
	ps = Flux.params(m)

	bestloss = Inf
	bestmodel = nothing

	loss(g, y, idxs) = Flux.Losses.logitcrossentropy(m(g)[:, idxs], y[:, idxs])

	opt = Flux.Optimise.Optimiser(Flux.Optimise.WeightDecay(0.001), Flux.Optimise.ADAMW(0.01))

	Flux.@epochs 100 begin
		gs = gradient(ps) do
			losstrn = loss(g, y, trainidx)
			return losstrn
		end

		lossval = loss(g, y, validx)

		if lossval < bestloss
			@show bestloss = lossval

			acc = sum(Flux.onecold(m(g))[validx] .== Flux.onecold(y)[validx]) / length(y[validx])
			@show acc

			bestmodel = deepcopy(m)
		end
		Flux.update!(opt, ps, gs)
	end

	return bestmodel
end

bmod = training_run()
Flux.Losses.logitcrossentropy(bmod(g)[:, validx], y[:, validx])

sum(length, Flux.params(bmod))

# bmod[4].re(bmod[4].p).l.W
using BSON
bparams = Flux.params(bmod |> cpu)
BSON.@save "download/coraode.bson" bparams


using VortexDistributions
using JLD2
using GLMakie
using Interpolations
using NearestNeighbors
using LinearAlgebra
using Distances
## Include util functions
include("../util.jl")

## Load psi simulation
@load "data/box_rings.jld2" psi_ring1 psi_ring2 psi_ring3 X # make sure julia repl is in base directory
# @load "data/box_vorts.jld2" psi_chaos1 psi_chaos2 psi_knots1 psi_knots2 psi_knots3 psi_knots4 psi_tubes1 psi_tubes2 psi_ringtube X
# @load "data/slab_tubes.jld2" psi_tubes1 psi_tubes2 psi_tubes3 X

## Set psi
psi = psi_ring3;

## Plot iso surface
plot_iso(psi, X)

naivePlaquette(psi[:, :, 42], X, 0)
findvortices(Torus(psi[:, :, 42], X[1], X[2]))

## Find vortices
@time vorts_3d = findvortices3D_itp(psi, X, 2);
num_vorts = length(vcat(vorts_3d...)[:,1])
print("Number of vortices found: ", length(vcat(vorts_3d...)[:,1]))

## Scatter vortices found
scatterVortsOnIso(vcat(vorts_3d'...))

#############################

## Unique vortices
interp_depth = 2
rtol = 0.44 # Largest distance between two points

@time fils = uniqueVortices(rtol, psi, X, interp_depth);

length(fils) # number of filaments found

## Plotting unique filaments
plot_iso(psi, X)
scatterDemo(fils)

plot!()


function part1(lines)
    mask = Dict{Int, Int}()
    registers = Dict{Int, Int}()

    for line in lines
        m = match(r"^mask = ([01X]{36})$", line)
        if !isnothing(m)
            mask = map(filter(!=('X') ∘ last, collect(enumerate(reverse(m[1]))))) do (i, x)
                i => Int(x - '0')
            end |> Dict
            continue
        end
        m = match(r"^mem\[(\d+)\] = (\d+)$", line)
        if !isnothing(m)
            reg, bits = parse(Int, m[1]), digits(parse(Int, m[2]); base=2, pad=36)
            for (k, v) in mask
                bits[k] = v
            end
            registers[reg] = parse(Int, reverse(join(bits)); base=2)
        else
            error("neither mask nor mem: $line")
        end
    end

    sum(values(registers))
end

function apply(mask, bits, running, i, acc)
    if i == 0
        push!(acc, parse(Int, join(reverse(running)); base=2))
    elseif mask[i] == 'X'
        apply(mask, bits, [running; '0'], i - 1, acc)
        apply(mask, bits, [running; '1'], i - 1, acc)
    elseif mask[i] == '1'
        apply(mask, bits, [running; '1'], i - 1, acc)
    elseif mask[i] == '0'
        apply(mask, bits, [running; '0' + bits[i]], i - 1, acc)
    else
        error("unexpected char in mask at $i")
    end
end
apply(mask, bits) = apply(mask, bits, Char[], length(mask), Int[])

function part2(lines)
    mask = String[]
    registers = Dict{Int, Int}()

    for line in lines
        m = match(r"^mask = ([01X]{36})$", line)
        if !isnothing(m)
            mask = m[1]
            continue
        end

        m = match(r"^mem\[(\d+)\] = (\d+)$", line)
        if !isnothing(m)
            bits, value = reverse(digits(parse(Int, m[1]); base=2, pad=36)), parse(Int, m[2])
            foreach(reg -> registers[reg] = value, apply(mask, bits))
        else
            error("neither mask nor mem: $line")
        end
    end

    sum(values(registers))
end

function main()
    lines = readlines("input.txt")
    @info "Solution" part1(lines) part2(lines)
end

#  main()


"""
Kim, P.J., Chang, H.J., Song, D.S. and Choi, J.Y., 2007, June.
Fast support vector data description using k-means clustering.
In International Symposium on Neural Networks (pp. 506-514). Springer, Berlin, Heidelberg.

We dub the method K-Means subsamples for SVDD (KMSVDD)
Parameters:
on benchmark data k ∈ [5, 10, 20, 30, 40, 50]
init strategy and solver for SVDD solution
"""
struct KMSVDD <: Sampler
    k::Int
    init_strategy::SVDD.InitializationStrategy
    solver::OptimizerFactory
    function KMSVDD(k::Int, init_strategy::SVDD.InitializationStrategy, solver::OptimizerFactory)
        k > 0 || throw(ArgumentError("Invalid number of clusters $(k)."))
        return new(k, init_strategy, solver)
    end
end

function sample(sampler::KMSVDD, data::DataSet, labels::Vector{Symbol})::BitArray{1}
    clustering = kmeans(data, sampler.k)
    model = SVDD.VanillaSVDD(data)
    SVDD.initialize!(model, sampler.init_strategy)
    SVDD.set_adjust_K!(model, true)

    sample_mask = falses(size(data, 2))
    for i in 1:sampler.k
        cluster_indices = findall(clustering.assignments .== i)
        svs = train_and_find_support_vectors(model, data[:, cluster_indices], sampler.solver)
        sample_mask[cluster_indices[svs]] .= true
    end
    sv_mask = findall(sample_mask)
    svs = train_and_find_support_vectors(model, data[:, sv_mask], sampler.solver)
    sample_mask = falses(size(data, 2))
    sample_mask[sv_mask[svs]] .= true
    return sample_mask
end


RD.@autodiff struct Cartpole{T} <: RobotDynamics.ContinuousDynamics
    mc::T
    mp::T
    l::T
    g::T
end

Cartpole() = Cartpole{Float64}(1.0, 0.2, 0.5, 9.81)

function RobotDynamics.dynamics(model::Cartpole, x, u)
    mc = model.mc  # mass of the cart in kg (10)
    mp = model.mp   # mass of the pole (point mass at the end) in kg
    l = model.l   # length of the pole in m
    g = model.g  # gravity m/s^2

    q = x[ @SVector [1,2] ]
    qd = x[ @SVector [3,4] ]

    s = sin(q[2])
    c = cos(q[2])

    H = @SMatrix [mc+mp mp*l*c; mp*l*c mp*l^2]
    C = @SMatrix [0 -mp*qd[2]*l*s; 0 0]
    G = @SVector [0, mp*g*l*s]
    B = @SVector [1, 0]

    qdd = -H\(C*qd + G - B*u[1])
    return [qd; qdd]
end

function RobotDynamics.dynamics!(model::Cartpole, xdot, x, u)
    mc = model.mc  # mass of the cart in kg (10)
    mp = model.mp   # mass of the pole (point mass at the end) in kg
    l = model.l   # length of the pole in m
    g = model.g  # gravity m/s^2

    q = x[ @SVector [1,2] ]
    qd = x[ @SVector [3,4] ]

    s = sin(q[2])
    c = cos(q[2])

    H = @SMatrix [mc+mp mp*l*c; mp*l*c mp*l^2]
    C = @SMatrix [0 -mp*qd[2]*l*s; 0 0]
    G = @SVector [0, mp*g*l*s]
    B = @SVector [1, 0]

    qdd = -H\(C*qd + G - B*u[1])
    xdot[1] = qd[1]
    xdot[2] = qd[2]
    xdot[3] = qdd[1]
    xdot[4] = qdd[2]
    return nothing
end

function RD.jacobian!(model::Cartpole, J, xdot, x, u, t)
    mc = model.mc  # mass of the cart in kg (10)
    mp = model.mp   # mass of the pole (point mass at the end) in kg
    l = model.l   # length of the pole in m
    g = model.g  # gravity m/s^2

    q = x[ @SVector [1,2] ]
    qd = x[ @SVector [3,4] ]

    s = sin(q[2])
    c = cos(q[2])

    H = @SMatrix [mc+mp mp*l*c; mp*l*c mp*l^2]
    C = @SMatrix [0 -mp*qd[2]*l*s; 0 0]
    G = @SVector [0, mp*g*l*s]
    B = @SVector [1, 0]
    qdd = -H\(C*qd + G - B*u[1])

    ∂H∂z = SA[
        0 -mp*l*s*qdd[2] 0 0 0
        0 -mp*l*s*qdd[1] 0 0 0
    ]
    ∂C∂z = SA[
        0 -mp*l*c*qd[2]^2 0 -2*mp*l*qd[2]*s 0
        0 0 0 0 0
    ]
    ∂G∂z = SA[
       0 0 0 0 0
       0 mp*g*l*c 0 0 0
    ]
    ∂B∂z = SA[
        0 0 0 0 1
        0 0 0 0 0
    ]
    J[1:2,:] .= 0
    J[1,3] = 1.0
    J[2,4] = 1.0
    J[3:4,:] .= H\(-∂H∂z - ∂C∂z - ∂G∂z + ∂B∂z)
    return nothing
end

RobotDynamics.state_dim(::Cartpole) = 4
RobotDynamics.control_dim(::Cartpole) = 1


using LinearAlgebra
using Plots
using IJulia
using Printf

function meshgrid(x, y)
    x = reshape(x, 1, length(x))
    y = reshape(y, length(y), 1)
    X = repeat(x, length(y), 1)
    Y = repeat(y, 1, length(x))
    return X, Y
end

# choose y label to be either y or ϕ, depending on the geometry
function plot_basic_state(grid, background, geometry, file)

    if geometry == Cartesian()
        y = grid.y/1e3
        xlabel = "y (km)"
    elseif geometry == Spherical()
        y = rad2deg.(grid.ϕ)
        xlabel = "ϕ (°)"
    end

    solution = background[geometry]

     p1 = plot(
         y,
         solution.u,
         linewidth = 3,
         label = "u",
         title = "Velocity",
         ylabel = "(m/s)",
         xlims = (y[end], y[1])
         )

     p2 = plot(
         y,
         solution.η,
         linewidth = 3,
         label = "η",
         xlabel = xlabel,
         ylabel = "(m)",
         title = "Free-Surface",
         xlims = (y[end], y[1])
         )

     plt1 = plot(p1, p2, layout = (2,1))
     savefig(plt1, file)
end

function plot_growth_rates(ks, σ, phys, geometry, Nmodes, file)

    if geometry == Cartesian()
        k = ks[geometry]*1e3
        xlabel = "k (1/km)"
    elseif geometry == Spherical()
        k = ks[geometry]/1e3
        xlabel = "k (10^3/rad)"
    end

    titles = ["1st mode", "2nd mode"]
    plt = plot()
    for (index,value) in enumerate(1:Nmodes)
        σnormalized = [σ[(k, geometry)][index] for k in ks[geometry]]/phys[geometry].f₀
        plot!(plt,
        k,
        σnormalized,
        linewidth = 3,
        label=titles[index],
        xlabel=xlabel,
        ylabel="σ/f₀",
        title="Growth Rates",
        xlims=(k[1], k[end])
        )
    end
    savefig(plt,file)

end

#ticks = collect(0:0.2:1)
#ticklabels = [ @sprintf("%5.1f",x) for x in ticks ]
#plot(x,y)
#plot!(xticks=(ticks,ticklabels))

function  make_tick_labels(fs, Nfs)
    fs_max = maximum([real(fs); imag(fs)])
    fs_min = minimum([real(fs); imag(fs)])
    dfs = (fs_max - fs_min)/Nfs

    ticks = collect(fs_min:dfs:fs_max)
    ticklabels = [ @sprintf("%f",f) for f in ticks]

    return ticks, ticklabels
end

function plot_1D_streamfunction(k_index, k, phys, σmodes, geometry, mode_number, file)

    if geometry == Cartesian()
        y = phys[geometry].grid.y/1e3
        xlabel = "y (km)"
    elseif geometry == Spherical()
        y = rad2deg.(phys[geometry].grid.ϕ)
        xlabel = "ϕ (°)"
    end

    N = phys[geometry].grid.N
    Lx = 2*pi/k;
    Nx = N;
     x = LinRange(0, Lx, Nx+1)
  X, Y = meshgrid(x,y)

  uvec = σmodes[(k,geometry)][    1:N+1,   mode_number];
  vvec = σmodes[(k,geometry)][  N+2:2*N+2, mode_number];
  ηvec = σmodes[(k,geometry)][2*N+1:3*N+1, mode_number];

    p1 = plot(
      y,
      real(uvec),
      linewidth = 3,
      label = "real",
      title = "u",
      yticks=make_tick_labels(uvec, 3),
      xlims = (y[end], y[1])
      )
  plot!(
      p1,
      y,
      linewidth = 3,
      imag(uvec),
      label = "imag")
  p2 = plot(
      y,
      real(vvec),
      linewidth = 3,
      label = "real",
      title = "v",
      yticks=make_tick_labels(vvec, 3),
       xlims = (y[end], y[1])
      )
  plot!(
      p2,
      y,
      linewidth = 3,
      imag(vvec),
      label = "imag")
  p3 = plot(
      y,
      real(ηvec),
      linewidth = 3,
      label = "real",
      title = "η",
      yticks=make_tick_labels(ηvec, 3),
      xlabel = xlabel,
      xlims = (y[end], y[1])
      )
  plot!(
      p3,
      y,
      linewidth = 3,
      imag(ηvec),
      label = "imag")
  plt = plot(p1, p2, p3, layout = (3,1))
  savefig(plt, file)
  #display(plt)

end

function plot_2D_streamfunction(k_index, k, phys, σmodes, geometry, mode_number, file)

    N = phys[geometry].grid.N
    Nx = N;
    Lx = 2*pi/k;
    x = LinRange(0, Lx, Nx+1)

    if geometry == Cartesian()
        y = phys[geometry].grid.y/1e3
        ylabel = "y (km)"
        x = x./1e3
        kscaled = k*1e3
        xlabel = "x (km)"
    elseif geometry == Spherical()
        y = rad2deg.(phys[geometry].grid.ϕ)
        ylabel = "ϕ (°)"
        x = rad2deg.(x)
        kscaled = k * π / 180
        xlabel = "λ (°)"
     end

  X, Y = meshgrid(x,y)

  uvec = σmodes[(k,geometry)][    1:N+1,   mode_number];
  vvec = σmodes[(k,geometry)][  N+2:2*N+2, mode_number];
  ηvec = σmodes[(k,geometry)][2*N+1:3*N+1, mode_number];

    u =     repeat(real(uvec),1,Nx+1).*cos.(kscaled*X) -    repeat(imag(uvec),1,Nx+1).*sin.(kscaled*X);
    v = -k.*repeat(imag(vvec),1,Nx+1).*cos.(kscaled*X) - k.*repeat(real(vvec),1,Nx+1).*sin.(kscaled*X);
    η =     repeat(real(ηvec),1,Nx+1).*cos.(kscaled*X) -    repeat(imag(ηvec),1,Nx+1).*sin.(kscaled*X);

    kwargs = (
        xlabel = "x (km)",
        ylabel = ylabel,
          fill = true,
        levels = 20,
     linewidth = 0,
         color = :balance,
      colorbar = true,
          xlim = (x[1],   x[end]),
          ylim = (y[end], y[1])
    )

    #print("x ticks = ", make_tick_labels(x, 2), "\n")

    u_plt = contour(x, y[end:-1:1], u[:, end:-1:1], title="u"; kwargs...)
    v_plt = contour(x, y[end:-1:1], v[:, end:-1:1], title="v"; kwargs...)
    η_plt = contour(x, y[end:-1:1], η[:, end:-1:1], title="η"; kwargs...)

    plt1 = plot(u_plt, v_plt, η_plt, layout = (1,3))
    savefig(plt1, file)

end

function plot_2D_vorticity(phys, k_index, k, σmodes, geometry, mode_number, file)

    N = phys[geometry].grid.N
    Lx = 2*pi/k;
    Nx = N;
    x = LinRange(0, Lx, Nx+1)

    if geometry == Cartesian()
        y = phys[geometry].grid.y/1e3
        ylabel = "y (km)"
        Dy = phys[geometry].grid.D
        x = x./1e3
        kscaled = k*1e3
        xlabel = "x (km)"
    elseif geometry == Spherical()
        y = rad2deg.(phys[geometry].grid.ϕ)
        ylabel = "ϕ (°)"
        Dy = phys[geometry].grid.D/phys[geometry].grid.a
        x = rad2deg.(x)
        kscaled = k * π / 180
        xlabel = "λ (°)"
    end

    X, Y = meshgrid(x,y)

  uvec = σmodes[(k,geometry)][    1:N+1,   mode_number];
  vvec = σmodes[(k,geometry)][  N+2:2*N+2, mode_number];

    ζvec = im * k * vvec - Dy * uvec
    ζ    = repeat(real(ζvec),1,Nx+1).*cos.(kscaled*X) - repeat(imag(ζvec),1,Nx+1).*sin.(kscaled*X);

    kwargs = (
        xlabel = xlabel,
        ylabel = ylabel,
          fill = true,
        levels = 20,
     linewidth = 0,
         color = :balance,
      colorbar = true,
          xlim = (  x[1], x[end]),
         ylim = (y[end],  y[1])
    )

    plt = contour(x, y[end:-1:1], ζ[:, end:-1:1], title="ζ"; kwargs...)

    savefig(plt, file)

end

#string.(fieldnames(typeof(gridC)))


include("simulator.jl")

# ## Bayesian wrapper
using Gen, Genify, Distributions

trunc_normal(μ, σ, lb, ub) =
    Genify.WrappedDistribution(truncated(Normal(μ, σ), lb, ub))

# Set up fixed model parameters
function create_params(; β=0.5)
    Ns = [50; 50; 50]
    migration_rates = [0.95 0.04 0.01; 0.10 0.85 0.05; 0.01 0.04 0.95]
    β_det = β * ones(3)
    β_und = β_det ./ 10
    death_rate = 0.0
    Is = [10; 0; 0]
    return @dict(Ns, migration_rates, β_det, β_und, death_rate, Is)
end

# Define observation model
@gen function observe(model::ABM, noise::Float64=5.0)
    obs = zeros(length(nodes(model)) * 2)
    for city in nodes(model)
        infected = count(a.status == :I for a in get_node_agents(city, model))
        recovered = count(a.status == :R for a in get_node_agents(city, model))
        obs[[city*2-1, city*2]] = [infected, recovered]
        {:infected => city} ~ trunc_normal(infected, noise, 0, Inf)()
        {:recovered => city} ~ trunc_normal(recovered, noise, 0, Inf)()
    end
    for (i, agent_id) in enumerate(model.scheduler(model))
        agent = model.agents[agent_id]
        probs = [c == agent.pos ? 0.99 : 0.01/(model.C-1) for c in 1:model.C]
        {:location => i} ~ Gen.categorical(probs)
    end
    return obs
end

# Genify the Agents.step! method
abm_step! = genify(Agents.step!, AgentBasedModel, Any, Int, Bool; recurse=true)

# Define Unfold step
@gen function bayesian_sir_step!(t::Int, _::Nothing,
                                model::AgentBasedModel, noise::Float64)
    {:agents} ~ abm_step!(model, agent_step!, 1, true)
    {:obs} ~ observe(model, noise)
    return nothing
end

# Wrap SIR step model in Gen model with parameter uncertainty
@gen (static) function bayesian_sir(T::Int, noise::Float64)
    β ~ uniform_continuous(0, 2.0)
    params = create_params(β=β)
    model = model_initiation(; params...)
    {:step} ~ Unfold(bayesian_sir_step!)(T, nothing, model, noise)
    return model
end

load_generated_functions()

# Test run the model
params = create_params()
model = model_initiation(; params...)
step!(model, agent_step!, 1, true)
abm_step!(model, agent_step!, 1, true)


# Fallbacks
@inline diffusivity_mixing_length(m, i) = mixing_length(m, i)
@inline dissipation_length(m, i) = mixing_length(m, i)

#
# A "simple" mixing length
#

Base.@kwdef struct SimpleMixingLength{T} <: AbstractParameters
    Cᴸᵇ :: T = 3.9302
end

@inline function mixing_length(m::Model{<:SimpleMixingLength}, i)
    # Two mixing lengths:

    # 1. Based on distance from surface (and limited by grid spacing)
    @inbounds ℓᶻ = - m.grid.zc[i]

    # 2. Based on stratification
    ℓᵇ = nan2inf(m.mixing_length.Cᴸᵇ * sqrt_e(m, i) / oncell(sqrt_∂B∂z, m, i))

    # Take hard minimum between the two:
    ℓ = min(ℓᶻ, ℓᵇ)

    # Limit mixing length to half cell width
    ℓ = max(ℓ, Δf(m.grid, i) / 2)

    return ℓ
end

#
# Mixing length model due to Ignacio Lopez-Gomez + Clima
#

Base.@kwdef struct EquilibriumMixingLength{T} <: AbstractParameters
    Cᴸᵟ :: T = 1.0
    Cᴸʷ :: T = 1.688 # Limits to Von-Karman constant for stress-driven turbulence
    Cᴸᵇ :: T = 1.664
end

@inline function mixing_length(m::Model{<:EquilibriumMixingLength}, i)
    @inbounds e = m.solution.e[i] # TKE at cell i

    # "Smallest" diffusivity limited by grid spacing.
    ℓᵟ = m.mixing_length.Cᴸᵟ * Δf(m.grid, i)

    if e <= 0 # shortcut when TKE is zero (diffusivity is zero in this case regardless).
        return ℓᵟ
    else
        # For notational convenience
        Cᴸᵇ = m.mixing_length.Cᴸᵇ
        Cᴸʷ = m.mixing_length.Cᴸʷ
        Cᴰ  = m.tke_equation.Cᴰ

        # Extract buoyancy frequency and shear squared.
        N² = oncell_∂B∂z(m, i)
        S² = oncell(shear_squared, m, i)

        # Length scale associated with a production-buoyancy flux-dissipation balance
        # in the TKE budget. Valid only for a linear equation of state with a shared diffusivity
        # for salinity and temperature.
        #
        # Note that ω² = Cᴷu * S² - Cᴷc * N² can be negative, in which case this model predicts ℓᵀᴷᴱ=0.
        ω² = Cᴷu(m, i) * S² - Cᴷc(m, i) * N²
        ℓᵀᴷᴱ = maxsqrt(Cᴰ * e / ω²)

        # Length-scale limitation by strong stratification.
        ℓᵇ = N² <= 0 ? Inf : Cᴸᵇ * √(e / N²)

        # Near-wall length scale:
        ℓʷ = @inbounds - Cᴸʷ * m.grid.zc[i] * u★(m) / √e

        # Hard minimum for now
        ℓ = min(ℓᵀᴷᴱ, ℓᵇ, ℓʷ)

        return max(ℓ, ℓᵟ) # limits mixing length to be larger than grid spacing
    end
end


@testset "2 stages with $solver" for solver in lp_solvers
    include("prob5.2_data.jl")

    numScen = 2
    M1 = StructuredModel(num_scenarios=numScen)

    @variable M1 x1[1:n] >= 0
    @variable M1 v1[1:n] >= 0
    @constraints M1 begin
        x1 .== v1
    end
    @objective(M1, Min, dot(ic, v1))

    for s in 1:numScen
        M2 = StructuredModel(parent=M1, prob=p2[s], id=s)
        @variable(M2, y2[1:n, 1:m] >= 0)
        @constraints M2 begin
            demand[j=1:m], sum(y2[:,j]) == D2[j,s]
            ylim[i=1:n], sum(y2[i,:]) <= x1[i]
        end
        @objective(M2, Min, dot(C, y2 * T))
    end

    # With detectlb = true and MultiCut
    # Iteration | Status    |     LB    |      UB    | #FC | #OC |
    #     1     | Optimal   |      0.0  |     Inf    |  2  |  0  |
    #     2     | Optimal   |  22338    | 1196514.84 |  0  |  2  |
    #     3     | Optimal   |  89893.68 |  676982.10 |  0  |  2  |
    #     4     | Optimal   | 182717.07 |  353268.97 |  0  |  2  |
    #     5     | Optimal   | 303678.32 |  366193.45 |  0  |  2  |
    #     6     | Optimal   | 311405.77 |  356197.41 |  0  |  2  |
    #     7     | Optimal   | 319203.65 |  344833.83 |  0  |  2  |
    #     8     | Optimal   | 334922.82 |  368285.45 |  0  |  2  |
    #     9     | Optimal   | 339047.44 |  349576.21 |  0  |  1  |
    #    10     | Optimal   | 340315.52 |  340315.52 |  0  |  0  |
    # With detectlb = true and AveragedCut
    #     1     | Optimal   |      0.0  |     Inf    | #FC | #OC |
    #     2     | Optimal   |  22338    | 1196514.84 |  2  |  0  |
    #     3     | Optimal   |  73386.05 |  660474.47 |  0  |  1  |
    #     4     | Optimal   | 171778.69 |  356387.31 |  0  |  1  |
    #     5     | Optimal   | 297813.60 |  391523.18 |  0  |  1  |
    #     6     | Optimal   | 309298.85 |  353098.55 |  0  |  1  |
    #     7     | Optimal   | 314734.72 |  344892.19 |  0  |  1  |
    #     8     | Optimal   | 331094.09 |  376701.57 |  0  |  1  |
    #     9     | Optimal   | 336691.25 |  346276.39 |  0  |  1  |
    #    10     | Optimal   | 337198.76 |  346727.58 |  0  |  1  |
    #    11     | Optimal   | 339235.89 |  344363.19 |  0  |  1  | -> Pereira stops here
    #    12     | Optimal   | 339779.51 |  340949.37 |  0  |  1  |
    #    13     | Optimal   | 340315.52 |  340315.52 |  0  |  0  |
    function testniter(niter, K, maxncuts, cutgen, detectlb)
        if isa(cutgen, StructProg.MultiCutGenerator)
            @test niter == 10
        else
            if maxncuts == -1
                if K == -1
                    if detectlb
                        @test niter == 13
                    else
                        # 10 on Gurobi/CPLEX, 14 otherwise
                        @test niter == 10 || niter == 14
                    end
                else
                    # 9 on Gurobi/CPLEX, 11 otherwise
                    @test niter == 9 || niter == 11 # Pereira stops earlier
                end
            else
                if K == -1
                    if detectlb
                        # up to 15 on Mac OS, <= 14 on Linux
                        @test 13 <= niter <= 15
                    else
                        # 10 on Gurobi/CPLEX, between 14 and 15 otherwise
                        @test 10 <= niter <= 15
                    end
                else
                    # 9 on Gurobi/CPLEX, between 11 and 12 otherwise
                    @test 9 <= niter <= 12 # Pereira stops earlier
                end
            end
        end
    end
    fulltest(M1, 2, 340315.52, [5085,1311,3919,854], 334687.754566, 15869.996575, testniter, solver)
    # root = model2lattice(M1, 2, solver, cutmode)
    # sol = SDDP(root, 2, cutmode, :All, verbose)
    #
    # v11value = sol.sol[1:4]
    # @test sol.status == :Optimal
    # @test abs(sol.objval - 340315.52) < 0.1
    # @test norm(v11value - [5085,1311,3919,854]) < 0.1
    # SDDPclear(M1)
end