"""
This file constructs the functions to find lower and upper bound of optimal
set of sourcing countries according to Jia's algorithm.
"""

## Define module and things to be exported
module JiaAlgorithm

export lowerbound_setup, lowerbound, upperbound_setup, upperbound, optimalset

## Load packages
using LinearAlgebra, Random, Distributions, Statistics, DataFrames, StatsBase
using Combinatorics

## Initialize lower bound and source potential matrix
function lowerbound_setup(N, ξ, my_exp)
    J_lb = zeros(1, N)  # the lower bound is empty set
    source_start_lb = (J_lb * ξ).^my_exp
    temp = repeat(J_lb, N, 1) + I
    temp2 = ones(size(temp,1), size(temp,2))
    check_matrix_lb = min.(temp, temp2)
    source_check_lb = (check_matrix_lb * ξ).^my_exp

    return source_start_lb, source_check_lb
end

## Jia's lower bound algorithm
function lowerbound(source_start, source_check, ϕ_σ_B, fc, N, ξ, my_exp, firm)

    # Start iteration for Marginal Benefit (MB)
    k = 1
    Z_start  = zeros(1, N)

    while k <= N

    # compute MB
        if k > 1
            source_start = (Z_start * ξ).^my_exp
            source_check = (Z_start * ξ .+ ξ .* (1 .- Z_start')).^my_exp
        end

        source_potential_start = ϕ_σ_B[firm] .* source_start
        source_potential_new_vec = ϕ_σ_B[firm] .* source_check

        # generate matrix with 1 if MB positive and update set of sourcing countries
        MB = (source_potential_new_vec' - fc[firm, :]' .- source_potential_start .> 0)
        Z_new = min.(Z_start + MB, ones(size(MB,1), size(MB,2)))

        if Z_start == Z_new
            return Z_new
            @goto end_lb_algorithm
        end

        k += 1
        Z_start = Z_new
        return Z_new
    end

    @label end_lb_algorithm
    return Z_new, k
end

## Initialize upper bound and source potential matrix
function upperbound_setup(N, ξ, my_exp)
    J_ub = ones(1, N)  # the upper bound is full set
    source_start_ub = (J_ub * ξ).^my_exp
    temp = repeat(J_ub, N, 1) - I
    temp2 = zeros(size(temp,1), size(temp,2))
    check_matrix_ub = max.(temp, temp2)
    source_check_ub = (check_matrix_ub * ξ).^my_exp

    return source_start_ub, source_check_ub
end

## Jia's upper bound algorithm
function upperbound(source_start, source_check, ϕ_σ_B, fc, N, ξ, my_exp, firm)

    # Start iteration for Marginal Benefit (MB)
    k = 1
    Z_start  = ones(1, N)

    while k <= N

    # compute MB
        if k > 1
            source_start = (Z_start * ξ).^my_exp
            source_check = (Z_start * ξ .- ξ .* (Z_start')).^my_exp
        end

        source_potential_start = ϕ_σ_B[firm] .* source_start
        source_potential_new_vec = ϕ_σ_B[firm] .* source_check

        # generate matrix with 1 if MB positive and update set of sourcing countries
        MB = (source_potential_start .- source_potential_new_vec' - fc[firm, :]'  .< 0)
        Z_new = max.(Z_start - MB, zeros(size(MB,1), size(MB,2)))

        if Z_start == Z_new
            return Z_new
            @goto end_ub_algorithm
        end

        k += 1
        Z_start = Z_new
        return Z_new
    end

    @label end_ub_algorithm
    return Z_new, k
end

## Check if Jia's algorithm produce the same lower and upper bound
function optimalset(Z, gap_bounds, firm, Z_lb, Z_ub, S, N, num_rand_checks, rand_check_matrix, fc, ξ, my_exp, ϕ_σ_B)

    if Z_lb == Z_ub
        print("")
        Z[firm,:] = Z_lb

    # I don't know how to do lines 64-88 of gmm_objective.m: the case when lower and
    # upper bounds are not too different

    else
        #print("WARNING! The sourcing strategy may not be solved correctly")
        #print("for firm number")
        #println("$firm")
        Z_check = repeat(Z_lb, num_rand_checks, 1)
        ind_diffZ = zeros(1,N)
        for i in 1:N
            if Z_ub[i] == Z_lb[i]
                ind_diffZ[i] = 0.0
            else
                ind_diffZ[i] = 1.0
            end
        end
        gap_bounds[firm] = sum(ind_diffZ)
        K_top = Int(sum(ind_diffZ))
        K_diff = [i[2] for i in findall(x->x!=0.0, ind_diffZ)]
        for K in 1:K_top
            Z_check[:,K_diff[K]] = (rand_check_matrix[:,K_diff[K]] .> 0.5)
        end

        # Use the check and both bounds to find new set of sourcing countries
        Z_check = [Z_check; Z_lb; Z_ub]
        fc_payments = sum(Z_check .* fc[firm,:]', dims=2)
        source_potential_shock_mat_check = (Z_check * ξ).^my_exp
        source_potential_vec = ϕ_σ_B[firm] .* source_potential_shock_mat_check
        tot_profit = source_potential_vec - fc_payments
        loc_firm_best = findmax(tot_profit, dims = 1)[2]

        # replace Z as the set of locations from Z_check that maximize profits
        Z[firm,:] = Z_check[loc_firm_best[1][1],:]
    end

    return Z
end

end     # end module


@testset "n-Queens" begin

    N = 8
    P = 100
    generatePositions = ()->collect(1:N)[randperm(N)]

    # Vector of N cols filled with numbers from 1:N specifying row position
    function nqueens(queens::Vector{Int})
        n = length(queens)
        fitness = 0
        for i=1:(n-1)
            for j=(i+1):n
                k = abs(queens[i] - queens[j])
                if (j-i) == k || k == 0
                    fitness += 1
                end
                # println("$(i),$(queens[i]) <=> $(j),$(queens[j]) : $(fitness)")
            end
        end
        return fitness
    end
    @test nqueens([2,4,1,3]) == 0
    @test nqueens([3,1,2]) == 1

    # Testing: GA solution with various mutations
    for muts in [inversion, insertion, swap2, scramble, shifting]
        result = Evolutionary.optimize(nqueens,
            generatePositions,
            GA(
                populationSize = P,
                selection = susinv,
                crossover = PMX,
                mutation = muts
            ));
        # show(result)
        println("GA:PMX:$(string(muts))(N=$(N), P=$(P)) => F: $(minimum(result)), C: $(Evolutionary.iterations(result))")
        @test nqueens(Evolutionary.minimizer(result)) == 0
        @test Evolutionary.iterations(result) < 100 # Test early stopping
    end

    # Testing: ES
    for muts in [inversion, insertion, swap2, scramble, shifting]
        for sel in [:plus, :comma]
            μ = 20
            result = Evolutionary.optimize(nqueens, generatePositions,
                ES(mutation = mutationwrapper(muts), μ=μ, ρ=1, λ=P, selection=sel),
                Evolutionary.Options(show_trace=false))
            println("($μ$(sel == :plus ? "+" : ",")$(P))-ES:$(string(muts)) => F: $(minimum(result)), C: $(Evolutionary.iterations(result))")
            # show(result)
            @test nqueens(Evolutionary.minimizer(result)) == 0
        end
    end

end


import sbp
using MAT
using LinearAlgebra
include("../common/wave_equation.jl")
include("../common/bound.jl")
include("../common/discretization.jl")
disc = Discretization
pde=WaveEquation

datafile=ARGS[1]
dR = parse(Float64, ARGS[2]) * ones(1,1)
nsamples = 4
nrefinements = 6
nx = 20
T = 1.2
dt = 0.25 / (nx - 1)
nt = T / dt
order = 4

mms = pde.manufactured_solution

function data()
        return Dict("p" => [], "v" => [] ,
                    "Fp" => [], "Fv" => [])
end

function discretize(nx, dRi)
    ops = disc.OperatorsType(nx, order)
    M = disc.discretization(pde.A(), ops)
    Rl = -ones(1,1)
    Rr = -ones(1,1)
    Pl = disc.projection_matrix(pde.A(pde.LEFT), Rl, dRi)
    Pr = disc.projection_matrix(pde.A(pde.RIGHT), Rr, dRi)
    Sl, Slg = pde.sat(pde.A(pde.LEFT), pde.LEFT, ops, Pl)
    Sr, Srg = pde.sat(pde.A(pde.RIGHT), pde.RIGHT, ops, Pr)
    gl = t -> pack(mms(ops.x[1], t))
    gr = t -> pack(mms(ops.x[end], t))
    f = (u,t) -> (M + Sl + Sr) * u - Slg * gl(t) - Srg * gr(t)
    return f, ops
end

function pack(mms_)
   p = mms_["p"]
   v = mms_["v"]
   return [p; v]
end

function initialize(ops, nx)
   return pack(mms(ops.x, 0))
end

function error(d, u, ops, t)
    mms = pde.manufactured_solution(ops.x,  t)
    F_mms = pde.functionals(t)

    soln = Dict(sbp.MMS.extract_vars(u, pde.sizes(ops.nx)))
    norms = Dict("p" => ops.Hx, "v" => ops.Hx)
    err = sbp.MMS.l2_errors(soln, mms, norms)
    F_err = sbp.MMS.functional_errors(soln, F_mms, norms)
    append!(d["p"], err["p"])
    append!(d["v"], err["v"])
    append!(d["Fp"], F_err["p"])
    append!(d["Fv"], F_err["v"])
end

function info(i, nx, nt, dt, ops, err)
        println("Grid = $i, nx = $nx, p = ", err["p"][end],
                " v = ", err["v"][end],
                ", rate p = ", sbp.MMS.log2_convergence_rates(err["p"])[end],
                ", rate v = ", sbp.MMS.log2_convergence_rates(err["v"])[end],
                ", rate Fp = ", sbp.MMS.log2_convergence_rates(err["Fp"])[end],
                ", rate Fv = ", sbp.MMS.log2_convergence_rates(err["Fv"])[end]
                )
end

errors = data()
rates = data()

x = disc.convergence_study(nx, nt, dt, nrefinements, errors, rates,
                           initialize, x -> discretize(x, dR), error,
                           info)
println("")
matwrite(datafile,
         Dict("dR" => dR,
              "x" => x,
              "p" => errors["p"],
              "v" => errors["v"],
              "Fp" => errors["Fp"],
              "Fv" => errors["Fv"],
              "q_p" => rates["p"],
              "q_v" => rates["v"],
              "q_Fp" => rates["Fp"],
              "q_Fv" => rates["Fv"]
               ))


###################################
## Example code
## for Bayesian Poisson Mixture Model

push!(LOAD_PATH,".")
using PyPlot, PyCall
import PoissonMixtureModel

"""
Visualize data & estimation in 2D space.
"""
function visualize_2D(X::Matrix{Float64}, S::Matrix{Float64}, S_est::Matrix{Float64}, text)
    cmp = get_cmap("jet")

    K1 = size(S, 1)
    K2 = size(S_est, 1)
    col1 = [pycall(cmp.o, PyAny, Int(round(val)))[1:3] for val in linspace(0,255,K1)]
    col2 = [pycall(cmp.o, PyAny, Int(round(val)))[1:3] for val in linspace(0,255,K2)]

    f, (ax1, ax2) = subplots(1,2,num=text)
    f[:clf]()
    f, (ax1, ax2) = subplots(1,2,num=text)

    for k in 1 : K1
        ax1[:scatter](X[1, S[k,:].==1], X[2, S[k,:].==1], color=col1[k])
    end
    ax1[:set_title]("truth")

    for k in 1 : K2
        ax2[:scatter](X[1, S_est[k,:].==1], X[2, S_est[k,:].==1], color=col2[k])
    end

    ax2[:set_title]("estimation")
end

function draw_hist(ax, X, S, label)
    counts, bins, patches = ax[:hist](X', 20)
    for i in 1 : length(patches)
        if counts[i] > 0
            S_tmp = S[:,bins[i] .<= X[1,:] .<= bins[i+1]]
            S_sum = sum(S_tmp, 2) / sum(S_tmp)
            patches[i][:set_facecolor]((S_sum[1], 0, S_sum[2]))
        end
    end
    ax[:set_title](label)
end

"""
Visualize data & estimation using 1D histogram.
"""
function visualize_1D(X::Matrix{Float64}, S::Matrix{Float64}, S_est::Matrix{Float64})
    # separated figures
    f1, ax1 = subplots(1,1,num="observation")
    f2, ax2 = subplots(1,1,num="estimation")
    f1[:clf]()
    f2[:clf]()
    _, ax1 = subplots(1,1,num="observation")
    _, ax2 = subplots(1,1,num="estimation")
    ax1[:hist](X', 20)
    ax1[:set_title]("observation")
    draw_hist(ax2, X, S_est, "estimation")
end

"""
Run a test script for 1D data clustering.
"""
function test_1D()
    ## set model
    D = 1 # data dimension, must be 1.
    K = 2 #  number of mixture components, must be 2.
    alpha = 100.0 * ones(K)
    cmp = [PoissonMixtureModel.Gam(1.0*ones(D), 0.01) for i in 1 : K]
    bpmm = PoissonMixtureModel.BPMM(D, K, alpha, cmp)

    ## generate data
    N = 1000
    pmm = PoissonMixtureModel.sample_PMM(bpmm)
    X, S = PoissonMixtureModel.sample_data(pmm, N)

    ## inference
    max_iter = 100
    tic()
    S_est, post_bpmm, VB = PoissonMixtureModel.learn_VI(X, bpmm, max_iter)
    #S_est, post_bpmm, VB = PoissonMixtureModel.learn_GS(X, bpmm, max_iter)
    #S_est, post_bpmm, VB = PoissonMixtureModel.learn_CGS(X, bpmm, max_iter)
    toc()

    ## plot
    visualize_1D(X, S, S_est)

    figure("ELBO")
    clf()
    plot(VB)
    ylabel("ELBO")
    xlabel("iterations")
    show()
end

"""
Run a test script for 2D data clustering.
"""
function test_2D()
    ## set model
    D = 2 # data dimension, must be 2.
    K = 8 # number of mixture components
    #K = 5

    alpha = 100.0 * ones(K)
    cmp = [PoissonMixtureModel.Gam(1.0*ones(D), 0.01) for i in 1 : K]
    bpmm = PoissonMixtureModel.BPMM(D, K, alpha, cmp)

    ## generate data
    N = 300
    pmm = PoissonMixtureModel.sample_PMM(bpmm)
    X, S = PoissonMixtureModel.sample_data(pmm, N)

    ## inference
    max_iter = 100
    tic()
    S_est, post_bpmm, VB = PoissonMixtureModel.learn_VI(X, bpmm, max_iter)
    #S_est, post_bpmm, VB = PoissonMixtureModel.learn_GS(X, bpmm, max_iter)
    #S_est, post_bpmm, VB = PoissonMixtureModel.learn_CGS(X, bpmm, max_iter)
    toc()

    ## plot
    visualize_2D(X, S, PoissonMixtureModel.winner_takes_all(S_est), "2D plot")

    # VB check
    figure("ELBO")
    clf()
    plot(VB)
    ylabel("ELBO")
    xlabel("iterations")
    show()
end

test_1D()
#test_2D()


using RefractiveIndexDatabase
using BenchmarkTools

suite = BenchmarkGroup()
suite["formula"] = BenchmarkGroup()
suite["tabulated"] = BenchmarkGroup()

let material, l, f
    material = get_material("other", "Mg-LiTaO3", "Moutzouris-o")
    l = range(material.n.λrange..., length=10_000)
    f = RefractiveIndex(material)
    suite["formula"]["MgLiTaO3"] = @benchmarkable $f.($l)
end

let material, l, f
    material = get_material("main", "Ag", "Johnson")
    l = range(extrema(material.λ)..., length=10_000)
    f = RefractiveIndex(material)
    suite["tabulated"]["Ag"] = @benchmarkable $f.($l)
end

tune!(suite)
results = run(suite; verbose=true, seconds = 1)

for (suite, group) in results
    @show suite
    for (benchmark, result) in group
        @show benchmark
        display(result)
    end
end

# # If a cache of tuned parameters already exists, use it, otherwise, tune and cache
# # the benchmark parameters. Reusing cached parameters is faster and more reliable
# # than re-tuning `suite` every time the file is included.
# paramspath = joinpath(dirname(@__FILE__), "params.json")
#
# if isfile(paramspath)
#     loadparams!(suite, BenchmarkTools.load(paramspath)[1], :evals)
# else
#     tune!(suite)
#     BenchmarkTools.save(paramspath, params(suite))
# end


# Chapter 17. Multiple Dispatch
# https://benlauwens.github.io/ThinkJulia.jl/latest/book.html#chap17

using Printf

# Methods
# https://benlauwens.github.io/ThinkJulia.jl/latest/book.html#_methods
function printtime(mytime::MyTime)
    @printf("%02d:%02d:%02d\n", mytime.hour, mytime.minute, mytime.second)
end

# Exercise 17-1

function timetosecond(time::MyTime)
    seconds = time.hour*3600 + time.minute*60 + time.second
end

function secondtotime(totalsecond::Int64)
    (minute, second) = divrem(totalsecond, 60)
    (hour, minute) = divrem(minute, 60)
    return MyTime(hour, minute, second)
end

# Constructors
# https://benlauwens.github.io/ThinkJulia.jl/latest/book.html#constructor
function MyTime(mytime::MyTime)
    MyTime(mytime.hour, mytime.minute, mytime.second)
end

function Base.show(io::IO, time::MyTime)
    @printf(io, "%02d:%02d:%02d", time.hour, time.minute, time.second)
end

import Base.+
function +(t1::MyTime, t2::MyTime)
    seconds = timetosecond(t1) + timetosecond(t2)
    secondtotime(seconds)
end


using Test
#using AbstractPlotting
import Calculus
import ForwardDiff
#using GLMakie; GLMakie.activate!()
#using MakieLayout
import NonlinearEquations
import PyPlot
import Random
import SparseArrays
using StaticArrays

ForwardDiff_gradient(x...) = ForwardDiff.gradient(x...)

function kr(psi, alpha, N)
	if psi < 0
		m = (N - 1) / N
		denom = 1 + abs(alpha * psi) ^ N
		numer = 1 - abs(alpha * psi) ^ (N - 1) * denom ^ (-m)
		return numer ^ 2 / denom ^ (m / 2)
	else
		return one(psi)
	end
end

kr(x::AbstractArray) = kr(x[1], x[2], x[3])

function Calculus.differentiate(x::Calculus.SymbolParameter{:kr}, args, wrt)
	chain_part = map(x->Calculus.simplify(Calculus.differentiate(x, wrt)), args)
	if chain_part == [0, 0, 0]
		return :(0)
	else
		return :(sum(ForwardDiff_gradient(kr, SA[$(args...)]) .* SA[$(chain_part...)]))
	end
end

function hm(x, y)
	return 2 / (1 / x + 1 / y)
end

hm(x::AbstractArray) = hm(x[1], x[2])

function Calculus.differentiate(x::Calculus.SymbolParameter{:hm}, args, wrt)
	chain_part = map(x->Calculus.simplify(Calculus.differentiate(x, wrt)), args)
	if chain_part == [0, 0]
		return :(0)
	else
		return :(sum(ForwardDiff_gradient(hm, SA[$(args...)]) .* SA[$(chain_part...)]))
	end
end

function Calculus.differentiate(x::Calculus.SymbolParameter{:abs}, args, wrt)
	if length(args) > 1
		error("Too many arguments passed to abs()")
	end
	arg = args[1]
	return :(ifelse($arg > 0, 1, -1) * $(Calculus.differentiate(arg, wrt)))
end

function effective_saturation(alpha::Number, psi::Number, N::Number)
	m = (N - 1) / N
	if psi < 0
		return (1 + abs(alpha * psi) ^ N) ^ (-m)
	else
		return 1
	end
end

function saturation(psi::Array, coords::Array)
	Se = similar(psi)
	for i = 1:length(psi)
		if inclay(coords[1, i], coords[2, i])
			_, alpha, N, residual_saturation = params[:clay]
		else
			_, alpha, N, residual_saturation = params[:claysilt]
		end
		Se[i] = residual_saturation + effective_saturation(alpha, psi[i], N) * (1 - residual_saturation)
	end
	return Se
end

function callback(psi, residuals, J, i, saturation=saturation(psi, coords))
	if mod(i, 10) == 0
		fig, axs = PyPlot.subplots(2, 2, figsize=(16, 9))
		psi = reshape(psi, ns[2], ns[1])
		saturation = reshape(saturation, ns[2], ns[1])
		residuals = reshape(residuals, ns[2], ns[1])
		img = axs[1].imshow(map(x->x > 0 ? x : 0, psi), origin="lower", interpolation="nearest", extent=[mins[1], maxs[1], mins[2], maxs[2]])
		fig.colorbar(img, ax=axs[1])
		img = axs[2].imshow(saturation, origin="lower", interpolation="nearest", extent=[mins[1], maxs[1], mins[2], maxs[2]])
		fig.colorbar(img, ax=axs[2])
		goodindices = filter(i->coords[2, i] == maxs[2], 1:size(coords, 2))
		img = axs[3].imshow(map(x->x < 0 ? x : 0, psi), origin="lower", interpolation="nearest", extent=[mins[1], maxs[1], mins[2], maxs[2]])
		fig.colorbar(img, ax=axs[3])
		img = axs[4].imshow(residuals, origin="lower", interpolation="nearest", extent=[mins[1], maxs[1], mins[2], maxs[2]])
		fig.colorbar(img, ax=axs[4])
		display(fig)
		println()
		println()
		PyPlot.close(fig)
	end
end

#=
function animatetransient(soln)
	hend = soln(tspan[end])
	scene, layout = layoutscene(resolution=(600, 600))
	t_observable = Node(0.0)
	h_observable = lift(t->reshape(soln(t), ns[2], ns[1])', t_observable)
	saturation_observable = lift(t->reshape(saturation(soln(t), coords), ns[2], ns[1])', t_observable)
	ax1 = layout[1, 1] = LAxis(scene; title="Pressure Head")
	ax2 = layout[2, 1] = LAxis(scene; title="Saturation")
	xs = range(mins[1]; stop=maxs[1], length=ns[1])
	ys = range(mins[2]; stop=maxs[2], length=ns[2])
	heatmap!(ax1, xs, ys, h_observable; colorrange=extrema(hend))
	heatmap!(ax2, xs, ys, saturation_observable, colorrange=(0, 1))
	for ax in [ax1, ax2]
		xlims!(ax, mins[1], maxs[1])
		ylims!(ax, mins[2], maxs[2])
	end
	layout[1, 2] = LColorbar(scene; limits=extrema(hend), width=30)
	layout[2, 2] = LColorbar(scene; limits=(0, 1), width=30)
	#timeplot = layout[2, :] = LAxis(scene; title="time", xgridvisible=false, ygridvisible=false, yticklabelsvisible=false, yticksvisible=false)
	#timeplot.height=20
	#lines!(timeplot, t_observable, [0, 1], linewidth=5)
	#ylims!(timeplot, 0, 1)
	#xlims!(timeplot, 0, tspan[end])
	#display(scene)
	record(scene, "animation_$(ns[1]).mp4", 0:0.2:62; framerate=60) do t
		t_observable[] = t
	end
end

function plottransient(soln)
	hend = soln(tspan[end])
	scene, layout = layoutscene(resolution=(1200, 600))
	t_slider = layout[3, 1:3] = LSlider(scene; range=range(0, tspan[end]; length=500), startvalue=0)
	t_observable = lift(t->[t, t], t_slider.attributes[:value])
	h_observable = lift(t->reshape(soln(t), ns[2], ns[1])', t_slider.attributes[:value])
	saturation_observable = lift(t->reshape(saturation(soln(t), coords), ns[2], ns[1])', t_slider.attributes[:value])
	ax1 = layout[1, 1] = LAxis(scene; title="Pressure Head")
	ax2 = layout[1, 3] = LAxis(scene; title="Saturation")
	xs = range(mins[1]; stop=maxs[1], length=ns[1])
	ys = range(mins[2]; stop=maxs[2], length=ns[2])
	heatmap!(ax1, xs, ys, h_observable; colorrange=extrema(hend))
	heatmap!(ax2, xs, ys, saturation_observable, colorrange=(0, 1))
	for ax in [ax1, ax2]
		xlims!(ax, mins[1], maxs[1])
		ylims!(ax, mins[2], maxs[2])
	end
	layout[1, 2] = LColorbar(scene; limits=extrema(hend), width=30)
	layout[1, 4] = LColorbar(scene; limits=(0, 1), width=30)
	timeplot = layout[2, :] = LAxis(scene; title="time", xgridvisible=false, ygridvisible=false, yticklabelsvisible=false, yticksvisible=false)
	timeplot.height=20
	lines!(timeplot, t_observable, [0, 1], linewidth=5)
	ylims!(timeplot, 0, 1)
	xlims!(timeplot, 0, tspan[end])
	display(scene)
end
=#

function i2i1i2(i)
	i2 = 1 + mod(i - 1, ns[2])
	i1 = 1 + div(i - i2, ns[2])
	return i1, i2
end

function updateneighborarrays!(a1, a2, pair, value)
	n1, n2 = pair
	if i2i1i2(n1)[1] == i2i1i2(n2)[1]#the neighbors have the same x coordinate, update a2
		i1, i2_1 = i2i1i2(n1)
		i2_2 = i2i1i2(n2)[2]
		@assert i2_1 < i2_2
		a2[i2_1, i1] = value
	else#the neighbors have the same y coordinate, update a1
		i1_1, i2 = i2i1i2(n1)
		i1_2 = i2i1i2(n2)[1]
		@assert i1_1 < i1_2
		a1[i2, i1_1] = value
	end
end

function breakupgradient(grad)
	grad1 = zeros(ns[2], ns[1] - 1)
	grad2 = zeros(ns[2] - 1, ns[1])
	for (v, pair) in zip(grad, neighbors)
		updateneighborarrays!(grad1, grad2, pair, v)
	end
	return grad1, grad2
end

function plotgradient(grad, obsnode; vmin=minimum(grad), vmax=maximum(grad))
	grad1, grad2 = breakupgradient(grad)
	fig, axs = PyPlot.subplots(1, 2, figsize=(16, 9))
	img = axs[1].imshow(grad1, extent=[mins[1], maxs[1], mins[2], maxs[2]], origin="lower", interpolation="nearest", vmin=vmin, vmax=vmax)
	fig.colorbar(img, ax=axs[1])
	img = axs[2].imshow(grad2, extent=[mins[1], maxs[1], mins[2], maxs[2]], origin="lower", interpolation="nearest", vmin=vmin, vmax=vmax)
	fig.colorbar(img, ax=axs[2])
	display(fig)
	println()
	println()
	PyPlot.close(fig)
end


function digitalroot(n)
    ns = string(n)
    sum = 0
    for i in ns
      sum += (Int(i) - Int('0'))
    end
    if(sum > 9)
      return digitalroot(sum)
    else
      return sum
    end
 end


### A Pluto.jl notebook ###
# v0.16.0

using Markdown
using InteractiveUtils

# ╔═╡ 3d1e8af8-16ef-11ec-18d8-e7fe83a40de4
begin
    using Pkg
    Pkg.activate("..")
    using Javis, Colors, Plots
end

# ╔═╡ 985f6ae4-c559-4756-b9d7-98696c7cbf7f
begin
    function ground(args...)
        background("white")
        sethue("black")
    end
end

# ╔═╡ c25e5278-b7fa-4ec4-abcc-08617fa08872
begin
    mutable struct HarmonicSum
        α::Float64
        values::Vector{Point}
    end

    function HarmonicSum(α::Float64, start_pos::Point; scaling = 20, n_points = 100)
        sequence = map(k -> 1 / k^α, 1:n_points) .* scaling
        return HarmonicSum(α, map(1:n_points) do k
            start_pos + (0, -reduce(+, sequence[1:k]))
        end)
    end
end

# ╔═╡ ef565322-cc4d-4f9d-8558-c722d490c611
function show_sum(S::HarmonicSum, firstframe, lastframe; color = colorant"red")
    points = S.values
    n_frames = lastframe - firstframe + 1
    pointiter = range(1, length(points) - 1, step = (length(points) - 1) ÷ n_frames)
    for idx in 1:n_frames
        Object(
            (firstframe + idx):lastframe,
            @JShape begin
                setline(3)
                sethue(color)
                line(points[pointiter[idx]], points[pointiter[idx + 1]], :stroke)
            end
        )
    end
end

# ╔═╡ 80489189-55c9-46c3-a665-cb58b8f0152d
function dash_sum(S::HarmonicSum, firstframe, lastframe; color = colorant"red")
    points = S.values
    n_frames = lastframe - firstframe + 1
    pointiter = range(1, length(points) - 1, step = (length(points) - 1) ÷ n_frames)
    for idx in 1:n_frames
        Object(
            (firstframe + idx):lastframe,
            @JShape begin
                sethue(color)
                line(
                    points[pointiter[idx + 1]] - (2, 0),
                    points[pointiter[idx + 1]] + (2, 0),
                    :stroke,
                )
            end
        )
    end
end

# ╔═╡ f9c60c0d-9150-42e7-98ed-83d971366292
begin
    function ζ(;
        n_sums = 100,
        height = 500,
        width = 500,
        firstframe = 1,
        lastframe = 100,
        margin = 10,
        llim = -200,
        rlim = 200,
    )

        n_frames = lastframe - firstframe
        colorpalette = range(colorant"red", colorant"blue", length = n_sums)
        Object(
            @JShape begin
                tickline(
                    Point(-200, 200 - margin),
                    Point(200, 200 - margin),
                    startnumber = 0.0,
                    finishnumber = 2.0,
                )
            end
        )
        Object(JLine(Point(0, -height ÷ 2), Point(0, 200 - 20), color = "black"))

        Xpoints = range(llim, rlim, length = n_sums)
        sums = map(enumerate(Xpoints)) do (idx, x)
            s = HarmonicSum(1.0 + x / rlim, Point(x, 200), n_points = 10000)
            dash_sum(s, firstframe, lastframe, color = colorpalette[idx])
            s
        end
    end
    series_video = Video(500, 500)
    Background(1:250, (args...) -> begin
        background("white")
        sethue("black")
    end)

    mypoint = Point(range(-200, 200, length = 100)[50], 200)
    l1 = @JLayer 1:250 begin
        s = HarmonicSum(1.0 + mypoint.x / 200, mypoint, n_points = 10000)
        dash_sum(s, 1, 250)
    end
    act!(l1, Action(50:100, anim_translate(mypoint, Point(200, 0))))
    act!(l1, Action(150:200, anim_translate(Point(200, 0), mypoint)))

    l2 = @JLayer 101:250 begin
        ζ(n_sums = 100, firstframe = 1, lastframe = 150)
    end

    act!(l2, Action(1:50, appear(:fade)))

    render(series_video, pathname = "../output/series_sum.gif")
end

# ╔═╡ Cell order:
# ╟─3d1e8af8-16ef-11ec-18d8-e7fe83a40de4
# ╟─985f6ae4-c559-4756-b9d7-98696c7cbf7f
# ╟─c25e5278-b7fa-4ec4-abcc-08617fa08872
# ╟─ef565322-cc4d-4f9d-8558-c722d490c611
# ╟─80489189-55c9-46c3-a665-cb58b8f0152d
# ╟─f9c60c0d-9150-42e7-98ed-83d971366292


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

"""
    Segment(p1, p2)

An oriented line segment with end points `p1`, `p2`.
The segment can be called as `s(t)` with `t` between
`0` and `1` to interpolate linearly between its endpoints.

See also [`Line`](@ref).
"""
struct Segment{Dim,T,V<:AbstractVector{Point{Dim,T}}} <: Polytope{1,Dim,T}
  vertices::V
end

isconvex(::Type{<:Segment}) = true

nvertices(::Type{<:Segment}) = 2
nvertices(s::Segment) = nvertices(typeof(s))

Base.minimum(s::Segment) = s.vertices[1]
Base.maximum(s::Segment) = s.vertices[2]
Base.extrema(s::Segment) = s.vertices[1], s.vertices[2]
measure(s::Segment) = norm(s.vertices[2] - s.vertices[1])
Base.length(s::Segment) = measure(s)

function (s::Segment)(t)
  if t < 0 || t > 1
    throw(DomainError(t, "s(t) is not defined for t outside [0, 1]."))
  end
  a, b = s.vertices
  a + t * (b - a)
end

function Base.in(p::Point{Dim,T}, s::Segment{Dim,T}) where {Dim,T}
  a, b = s.vertices
  if Dim == 2
    _0 = zero(T)
  elseif Dim == 3
    _0 = Vec{3,T}(0, 0, 0)
  else
    throw(ErrorException("not implemented"))
  end
  # given collinear points (a, b, p), the point p intersects
  # segment ab if and only if vectors satisfy 0 ≤ ap ⋅ ab ≤ ||ab||²
  iscollinear = isapprox((b - a) × (p - a), _0, atol=atol(T)^2)
  iscollinear && zero(T) ≤ (b - a) ⋅ (p - a) ≤ (b - a) ⋅ (b - a)
end


# outputs I want
# true parameter value list
#
# TODO: how to compute the posterior distribution of parameters in ABC
function BayesianCalibration(N_experiments,N_samples,alpha,
    algo_list,algo_parameter_list;
    ode_model = solve_ode,
    initial_state = [99.0;1.0;0.0],
    time_window=(0,10.0),
    add_noise = false,
    true_p_dist=(Gamma(2,1),Gamma(1,1)),
    save_param_output=false)

    # initial pop
    s0 = sum(initial_state)

    # num of algorithms
    N_algo = length(algo_list)

    # N param
    N_param = length(true_p_dist)

    # list of param values
    true_param = rand.(true_p_dist,N_experiments)

    # Calibration value
    res = zeros(N_experiments,N_algo)

    # algo output
    if save_param_output
        param_output = zeros(N_experiments,N_samples,N_param,N_algo)
    end
    # algo time
    cputime_output = zeros(N_experiments, N_algo)

    # ess
    ess_output = zeros(N_experiments,N_algo)

    # ess per unit time
    ess_output_time = copy(ess_output)

    # data generator
    function data_generator(p)
        y = ode_model(initial_state,time_window,p)
        if add_noise
            y = y + rand(LogNormal(0,0.5),size(y))
            y = y ./ sum(y,dims=1) * s0
        end
        y
    end

    for i in 1:N_experiments
        # generate data
        # firt sample true SIR parameters
        # true_p = (true_param[1][i],true_param[2][i])
        true_p_vec=zeros(N_param)
        for k in 1:N_param
            true_p_vec[k] = true_param[k][i]
        end
        true_p = Tuple(true_p_vec)
        # generate data
        y = data_generator(true_p)

        # iterate over algos
        for algo_index in 1:N_algo
            # perform algo
            output, cpu_time = algo_list[algo_index](y,data_generator,algo_parameter_list[algo_index],N_samples)
            # check whether parameters are contained in the 95% posterior dist
            # given independence both have to be contained

            res[i,algo_index] = all(map(ind -> check_in_interval(true_p[ind],compute_posterior(vec(output[:,ind]),alpha)),1:length(true_p)))

            # save output
            if save_param_output
                param_output[i,:,:,algo_index] = output
            end
            cputime_output[i,algo_index] =  cpu_time
            # use R's package mcmcse to compute ESS
            ess_output[i,algo_index]  = try rcopy(R"mean(mcmcse::ess($output))") catch; 0 end
            ess_output_time[i,algo_index]  = ess_output[i,algo_index] / cpu_time
        end

        if !save_param_output
            param_output = nothing
        end

        if i % 10 == 0
            @show(i)
        end
    end

    return(true_param = true_param,
           calibration = res,
           simulated_param = param_output,
           cpu_time = cputime_output,
           ess = ess_output,
           ess_time = ess_output_time,
           algo_param_list = algo_parameter_list)
end


module Factors

import Base.Iterators.product
const FactorTable = Dict{NamedTuple,Float64}

struct Variable
    name::Symbol
    m::Int # number of possible values
end

struct Factor
    vars::Vector{Variable}
    table::FactorTable
end

variablenames(ϕ::Factor) = [var.name for var in ϕ.vars]

function namedtuple(n::Array{Symbol,1}, v::Tuple)
    return NamedTuple{Tuple(n)}(v)
end

function assignments(vars::AbstractVector{Variable})
    n = [var.name for var in vars]
    #n = Tuple([var.name for var in vars])
    #return [NamedTuple{n}(v) for v in product((1:v.m for v in vars)...)]
    return [namedtuple(n, v) for v in product((1:v.m for v in vars)...)]
end

X = Variable(:x, 2)
Y = Variable(:y, 2)
Z = Variable(:z, 2)

ϕ = Factor([X, Y, Z], FactorTable(
    (x=1, y=1, z=1) => 0.08,
    (x=1, y=1, z=2) => 0.31,
    (x=1, y=2, z=1) => 0.09,
    (x=1, y=2, z=2) => 0.37,
    (x=2, y=1, z=1) => 0.01,
    (x=2, y=1, z=2) => 0.05,
    (x=2, y=2, z=1) => 0.02,
    (x=2, y=2, z=2) => 0.07,
))

display(ϕ.table)

display(ϕ.vars)

display(variablenames(ϕ))

"""
    normalize!(ϕ::Factor)
Normalize a factor ϕ, which divides all the entries in the factor by the
same scalar so they sum to 1.
"""
function normalize!(ϕ::Factor)

    z = sum(p for (a,p) in ϕ.table)
    for (a,p) in ϕ.table
        ϕ.table[a] = p/z
    end
    return ϕ
end

k = assignments(ϕ.vars)
@show k
@show k[2,1,1]
@show k[2,1,1].x
@show k[2,1,1].y
@show k[2,1,1].z

println(k)
println(typeof(k))
println(size(k))

sum([2,3,2])

end


# define commonly used symbols
const σ0 = @SMatrix [1 + 0im 0; 0 1 + 0im]
const σx = @SMatrix [0 1 + 0im; 1 + 0im 0]
const σy = @SMatrix [0 1im; -1im 0]
const σz = @SMatrix [1 + 0im 0; 0 -1 + 0im]


"""
 Holds the tableau of a stochastic implicit split partitioned Runge-Kutta method.
 qdrift, pdrift1, pdrift2 hold the RK coefficients for the drift parts,
 and qdiff, pdiff1, pdiff2 hold the RK coefficients for the diffusion part of the SDE.
"""
struct TableauSISPRK{T} <: AbstractTableauIRK{T}
    name::Symbol
    s::Int
    qdrift ::CoefficientsRK{T}
    qdiff  ::CoefficientsRK{T}
    pdrift1::CoefficientsRK{T}
    pdrift2::CoefficientsRK{T}
    pdiff1 ::CoefficientsRK{T}
    pdiff2 ::CoefficientsRK{T}

    # Order of the tableau is not included, because unlike in the deterministic
    # setting, it depends on the properties of the noise (e.g., the dimension of
    # the Wiener process and the commutativity properties of the diffusion matrix)
    #
    # Orders stored in qdrift and qdiff are understood as the classical orders of these methods.

    function TableauSISPRK{T}(name, s, qdrift, qdiff, pdrift1, pdrift2, pdiff1, pdiff2) where {T}
        # THE COMMENTED OUT PART WAS FOR TableauFIRK. MAY IMPLEMENT SOMETHING
        # SIMILAR FOR TableauSFIRK LATER.

        # if (q.s > 1 && istrilstrict(q.a)) || (q.s==1 && q.a[1,1] == 0)
        #     warn("Initializing TableauFIRK with explicit tableau ", q.name, ".\n",
        #          "You might want to use TableauERK instead.")
        # elseif q.s > 1 && istril(q.a)
        #     warn("Initializing TableauFIRK with diagonally implicit tableau ", q.name, ".\n",
        #          "You might want to use TableauDIRK instead.")
        # end

        @assert s == qdrift.s == qdiff.s == pdrift1.s == pdrift2.s == pdiff1.s == pdiff2.s

        new(name, s, qdrift, qdiff, pdrift1, pdrift2, pdiff1, pdiff2)
    end
end

function TableauSISPRK(name::Symbol, qdrift::CoefficientsRK{T},
                                     qdiff::CoefficientsRK{T},
                                     pdrift1::CoefficientsRK{T}, pdrift2::CoefficientsRK{T},
                                     pdiff1::CoefficientsRK{T}, pdiff2::CoefficientsRK{T}) where {T}
    TableauSISPRK{T}(name, qdrift.s, qdrift, qdiff, pdrift1, pdrift2, pdiff1, pdiff2)
end

function TableauSISPRK(name::Symbol, qorder_drift::Int, qa_drift::Matrix{T}, qb_drift::Vector{T}, qc_drift::Vector{T},
                                      qorder_diff::Int, qa_diff::Matrix{T}, qb_diff::Vector{T}, qc_diff::Vector{T},
                                      porder1_drift::Int, pa1_drift::Matrix{T}, pb1_drift::Vector{T}, pc1_drift::Vector{T},
                                      porder2_drift::Int, pa2_drift::Matrix{T}, pb2_drift::Vector{T}, pc2_drift::Vector{T},
                                      porder1_diff::Int, pa1_diff::Matrix{T}, pb1_diff::Vector{T}, pc1_diff::Vector{T},
                                      porder2_diff::Int, pa2_diff::Matrix{T}, pb2_diff::Vector{T}, pc2_diff::Vector{T}) where {T}
    TableauSISPRK{T}(name, length(qc_drift), CoefficientsRK(name, qorder_drift, qa_drift, qb_drift, qc_drift),
                                              CoefficientsRK(name, qorder_diff, qa_diff, qb_diff, qc_diff),
                                              CoefficientsRK(name, porder1_drift, pa1_drift, pb1_drift, pc1_drift),
                                              CoefficientsRK(name, porder2_drift, pa2_drift, pb2_drift, pc2_drift),
                                              CoefficientsRK(name, porder1_diff, pa1_diff, pb1_diff, pc1_diff),
                                              CoefficientsRK(name, porder2_diff, pa2_diff, pb2_diff, pc2_diff))
end

# TODO function readTableauSFIRKFromFile(dir::AbstractString, name::AbstractString)

"""
Parameters for right-hand side function of implicit Runge-Kutta methods.
  A - if positive, the upper bound of the Wiener process increments; if A=0.0, no truncation
"""
mutable struct ParametersSISPRK{DT, TT, ET <: SPSDE{DT,TT}, D, M, S} <: Parameters{DT,TT}
    equ::ET
    tab::TableauSISPRK{TT}
    Δt::TT
    ΔW::Vector{DT}
    ΔZ::Vector{DT}
    A::DT

    t::TT
    q::Vector{DT}
    p::Vector{DT}
end

function ParametersSISPRK(equ::ET, tab::TableauSISPRK{TT}, Δt::TT, ΔW::Vector{DT}, ΔZ::Vector{DT}, A::DT) where {DT, TT, ET <: SPSDE{DT,TT}}
    @assert equ.m == length(ΔW) == length(ΔZ)
    ParametersSISPRK{DT, TT, ET, equ.d, equ.m, tab.s}(equ, tab, Δt, ΔW, ΔZ, A, 0, zeros(DT, equ.d), zeros(DT, equ.d))
end

"""
Structure for holding the internal stages Q, the values of the drift vector
and the diffusion matrix evaluated at the internal stages V=v(Q), B=B(Q),
and the increments Y = Δt*a_drift*v(Q) + a_diff*B(Q)*ΔW
"""
struct NonlinearFunctionCacheSISPRK{DT}
    Q ::Vector{Vector{DT}}
    P ::Vector{Vector{DT}}
    V ::Vector{Vector{DT}}
    F1::Vector{Vector{DT}}
    F2::Vector{Vector{DT}}
    B ::Vector{Matrix{DT}}
    G1::Vector{Matrix{DT}}
    G2::Vector{Matrix{DT}}
    Y ::Vector{Vector{DT}}
    Z ::Vector{Vector{DT}}

    v ::Vector{DT}
    f1::Vector{DT}
    f2::Vector{DT}
    b ::Matrix{DT}
    g1::Matrix{DT}
    g2::Matrix{DT}
    y ::Vector{DT}
    z ::Vector{DT}

    function NonlinearFunctionCacheSISPRK{DT}(D, M, S) where {DT}

        # create internal stage vectors
        Q  = create_internal_stage_vector(DT, D, S)
        P  = create_internal_stage_vector(DT, D, S)
        V  = create_internal_stage_vector(DT, D, S)
        F1 = create_internal_stage_vector(DT, D, S)
        F2 = create_internal_stage_vector(DT, D, S)
        B  = create_internal_stage_vector(DT, D, M, S)
        G1 = create_internal_stage_vector(DT, D, M, S)
        G2 = create_internal_stage_vector(DT, D, M, S)
        Y  = create_internal_stage_vector(DT, D, S)
        Z  = create_internal_stage_vector(DT, D, S)

        # create velocity and update vector
        v  = zeros(DT,D)
        f1 = zeros(DT,D)
        f2 = zeros(DT,D)
        b  = zeros(DT,D,M)
        g1 = zeros(DT,D,M)
        g2 = zeros(DT,D,M)
        y  = zeros(DT,D)
        z  = zeros(DT,D)

        new(Q, P, V, F1, F2, B, G1, G2, Y, Z, v, f1, f2, b, g1, g2, y, z)
    end
end

"""
Unpacks the data stored in x = (Y[1][1], Y[1][2], ... Y[1][D], Y[2][1], ..., Z[1][1], Z[1][2], ... Z[1][D], Z[2][1], ...)
into Y, Z::Vector{Vector}, calculates the internal stages Q, P, the values of the RHS
of the SDE ( vi(Q,P), fi(Q,P), Bi(Q,P) and Gi(Q,P) ), and assigns them to V[i], F[i], B[i] and G[i].
Unlike for FIRK, here
Y = Δt a_drift v(Q,P) + a_diff B(Q,P) ΔW,
Z = Δt ̃a1_drift f1(Q,P) + Δt ̃a2_drift f2(Q,P) + ̃a1_diff G1(Q,P) ΔW + ̃a2_diff G2(Q,P) ΔW.
"""
function compute_stages!(x::Vector{ST}, Q::Vector{Vector{ST}}, P::Vector{Vector{ST}},
                                        V::Vector{Vector{ST}}, F1::Vector{Vector{ST}}, F2::Vector{Vector{ST}},
                                        B::Vector{Matrix{ST}}, G1::Vector{Matrix{ST}}, G2::Vector{Matrix{ST}},
                                        Y::Vector{Vector{ST}}, Z::Vector{Vector{ST}},
                                        params::ParametersSISPRK{DT,TT,ET,D,M,S}) where {ST,DT,TT,ET,D,M,S}

    local tqᵢ::TT       #times for the q internal stages
    local tpᵢ::TT       #times for the p internal stages

    @assert S == length(Q) == length(V) == length(B) == length(P) == length(F1) == length(F2) == length(G1) == length(G2)

    # copy x to Y, Z and calculate Q, P
    for i in 1:S
        @assert D == length(Q[i]) == length(V[i]) == size(B[i],1) == length(P[i]) == length(F1[i]) == length(F2[i]) == size(G1[i],1) == size(G2[i],1)
        @assert M == size(B[i],2) == size(G1[i],2) == size(G2[i],2)

        for k in 1:D
            Y[i][k] = x[D*(  i-1)+k]
            Z[i][k] = x[D*(S+i-1)+k]
        end

        Q[i] .= params.q .+ Y[i]
        P[i] .= params.p .+ Z[i]
    end

    # compute Vi = vi(Q) and Bi=Bi(Q)
    for i in 1:S
        tqᵢ = params.t + params.Δt * params.tab.qdrift.c[i]
        # Not sure what about pdrift2.c[i] --- should there be another time point?
        tpᵢ = params.t + params.Δt * params.tab.pdrift1.c[i]
        # calculates v(t,Q,P), f1(t,Q,P), f2(t,Q,P) and assigns to V, F1, F2
        params.equ.v(tqᵢ, Q[i], P[i], V[i])
        params.equ.f1(tpᵢ, Q[i], P[i], F1[i])
        params.equ.f2(tpᵢ, Q[i], P[i], F2[i])
        # calculates B(t,Q,P), G1(t,Q,P), G2(t,Q,P) and assigns to the matrices B[i], G1[i], G2[i]
        params.equ.B(tqᵢ, Q[i], P[i], B[i])
        params.equ.G1(tpᵢ, Q[i], P[i], G1[i])
        params.equ.G2(tpᵢ, Q[i], P[i], G2[i])
    end
end

"Compute stages of stochastic implicit split partitioned Runge-Kutta methods."
@generated function function_stages!(x::Vector{ST}, b::Vector{ST}, params::ParametersSISPRK{DT,TT,ET,D,M,S}) where {ST,DT,TT,ET,D,M,S}

    cache = NonlinearFunctionCacheSISPRK{ST}(D, M, S)

    quote
        compute_stages!(x, $cache.Q, $cache.P, $cache.V, $cache.F1, $cache.F2, $cache.B, $cache.G1, $cache.G2, $cache.Y, $cache.Z, params)

        local y1::ST
        local y2::ST
        local z1::ST
        local z2::ST

        # compute b = - (Y-AV)
        for i in 1:S
            for k in 1:D
                y1 = 0
                y2 = 0
                z1 = 0
                z2 = 0
                for j in 1:S
                    y1 += params.tab.qdrift.a[i,j]  * $cache.V[j][k]  * params.Δt + params.tab.qdiff.a[i,j]  * dot($cache.B[j][k,:], params.ΔW)
                    y2 += params.tab.qdrift.â[i,j]  * $cache.V[j][k]  * params.Δt + params.tab.qdiff.â[i,j]  * dot($cache.B[j][k,:], params.ΔW)
                    z1 += params.tab.pdrift1.a[i,j] * $cache.F1[j][k] * params.Δt + params.tab.pdrift2.a[i,j] * $cache.F2[j][k] * params.Δt
                        + params.tab.pdiff1.a[i,j]  * dot($cache.G1[j][k,:], params.ΔW) + params.tab.pdiff2.a[i,j] * dot($cache.G2[j][k,:], params.ΔW)
                    z2 += params.tab.pdrift1.â[i,j] * $cache.F1[j][k] * params.Δt + params.tab.pdrift2.â[i,j] * $cache.F2[j][k] * params.Δt
                        + params.tab.pdiff1.â[i,j]  * dot($cache.G1[j][k,:], params.ΔW) + params.tab.pdiff2.â[i,j] * dot($cache.G2[j][k,:], params.ΔW)
                end
                b[D*(  i-1)+k] = - $cache.Y[i][k] + (y1 + y2)
                b[D*(S+i-1)+k] = - $cache.Z[i][k] + (z1 + z2)
            end
        end
    end
end

"Stochastic implicit partitioned Runge-Kutta integrator."
struct IntegratorSISPRK{DT, TT, PT <: ParametersSISPRK{DT,TT},
                                ST <: NonlinearSolver{DT}, N} <: StochasticIntegrator{DT,TT}
    params::PT
    solver::ST
    # InitialGuessPSDE not implemented for SFIPRK
    #iguess::IT
    fcache::NonlinearFunctionCacheSISPRK{DT}

    q::Matrix{Vector{TwicePrecision{DT}}}
    p::Matrix{Vector{TwicePrecision{DT}}}
end

# K - the integer in the bound A = √(2 K Δt |log Δt|) due to Milstein & Tretyakov; K=0 no truncation
function IntegratorSISPRK(equation::SPSDE{DT,TT,VT,F1T,F2T,BT,G1T,G2T,N}, tableau::TableauSISPRK{TT}, Δt::TT; K::Int=0) where {DT,TT,VT,F1T,F2T,BT,G1T,G2T,N}
    D = equation.d
    M = equation.m
    NS= equation.ns
    NI= equation.n
    S = tableau.s

    # create params
    K==0 ? A = 0.0 : A = sqrt( 2*K*Δt*abs(log(Δt)) )
    params = ParametersSISPRK(equation, tableau, Δt, zeros(DT,M), zeros(DT,M), A)

    # create solver
    solver = create_nonlinear_solver(DT, 2*D*S, params)

    # Not implementing InitialGuessSDE
    # create initial guess
    #iguess = InitialGuessODE(get_config(:ig_interpolation), equation, Δt)

    # create cache for internal stage vectors and update vectors
    fcache = NonlinearFunctionCacheSISPRK{DT}(D, M, S)

    # create solution vectors
    q = create_solution_vector(DT, D, NS, NI)
    p = create_solution_vector(DT, D, NS, NI)

    # create integrator
    IntegratorSISPRK{DT, TT, typeof(params), typeof(solver), N}(params, solver, fcache, q, p)
end

equation(integrator::IntegratorSISPRK) = integrator.params.equ
timestep(integrator::IntegratorSISPRK) = integrator.params.Δt
tableau(integrator::IntegratorSISPRK) = integrator.params.tab
Base.eltype(integrator::IntegratorSISPRK{DT, TT, PT, ST, N}) where {DT, TT, PT, ST, N} = DT

function initialize!(int::IntegratorSISPRK{DT,TT}, sol::SolutionPSDE, k::Int, m::Int) where {DT,TT}
    check_solution_dimension_asserts(sol, k, m)

    # copy the m-th initial condition for the k-th sample path
    get_initial_conditions!(sol, int.q[k,m], int.p[k,m], k, m)

    # Not implementing InitialGuessSDE
    # # initialise initial guess
    # initialize!(int.iguess, m, sol.t[0], int.q[m])
end

"""
This function computes initial guesses for Y, Z and assigns them to int.solver.x
The prediction is calculated using an explicit integrator.
"""
function initial_guess!(int::IntegratorSISPRK{DT,TT}) where {DT,TT}

    # NOT IMPLEMENTING InitialGuessSDE

    # SIMPLE SOLUTION
    # The simplest initial guess for Y, Z is 0
    int.solver.x .= zeros(eltype(int), 2*int.params.tab.s*ndims(int))

    # FIX NEEDED !!!
    # APPARENTLY PREDICTION USING AN EXPLICIT METHOD DOESNT WORK WELL HERE
    # PROBABLY BECAUSE THE pdrift2.c[i] POINTS ARE NOT INCLUDED

    # # USING AN EXPLICIT INTEGRATOR TO COMPUTE AN INITIAL GUESS
    # # Below we use the R2 method of Burrage & Burrage to calculate
    # # the internal stages at the times c[1]...c[s].
    # # This approach seems to give very good approximations if the time step
    # # and magnitude of noise are not too large. If the noise intensity is too big,
    # # one may have to perform a few iterations of the explicit method with a smaller
    # # time step, use a higher-order explicit method (e.g. CL or G5), or use
    # # the simple solution above.
    #
    # local tV1  = zeros(DT,int.params.equ.d)
    # local tV2  = zeros(DT,int.params.equ.d)
    # local tF1  = zeros(DT,int.params.equ.d)
    # local tF2  = zeros(DT,int.params.equ.d)
    # local tmpF = zeros(DT,int.params.equ.d)
    # local tB1  = zeros(DT,int.params.equ.d, int.params.equ.m)
    # local tB2  = zeros(DT,int.params.equ.d, int.params.equ.m)
    # local tG1  = zeros(DT,int.params.equ.d, int.params.equ.m)
    # local tG2  = zeros(DT,int.params.equ.d, int.params.equ.m)
    # local tmpG = zeros(DT,int.params.equ.d, int.params.equ.m)
    # local Q    = zeros(DT,int.params.equ.d)
    # local P    = zeros(DT,int.params.equ.d)
    # local t2::TT
    # local Δt_local::TT
    # local ΔW_local = zeros(DT,int.params.equ.m)
    #
    # # When calling this function, int.params should contain the data:
    # # int.params.q - the q solution at the previous time step
    # # int.params.p - the p solution at the previous time step
    # # int.params.t - the time of the previous step
    # # int.params.ΔW- the increment of the Brownian motion for the current step
    #
    # #Evaluating the functions v and B at t,q - same for all stages
    # int.params.equ.v(int.params.t, int.params.q, int.params.p, tV1)
    # int.params.equ.B(int.params.t, int.params.q, int.params.p, tB1)
    # int.params.equ.f1(int.params.t, int.params.q, int.params.p, tF1)
    # int.params.equ.f2(int.params.t, int.params.q, int.params.p, tmpF)
    # #tF1 = tF1 + tmpF
    # int.params.equ.G1(int.params.t, int.params.q, int.params.p, tG1)
    # int.params.equ.G2(int.params.t, int.params.q, int.params.p, tmpG)
    # #tG1 = tG1 + tmpG
    #
    #
    # # Calculating the positions q at the points qdrift.c[i]
    # # if qdrift.c==pdrift.c, then also calculating the momenta
    # for i in 1:int.params.tab.s
    #
    #     # Taking the c[i] from the qdrift tableau.
    #     Δt_local  = int.params.tab.qdrift.c[i]*int.params.Δt
    #     ΔW_local .= int.params.tab.qdrift.c[i]*int.params.ΔW
    #
    #     Q = int.params.q + 2./3. * Δt_local * tV1 + 2./3. * tB1 * ΔW_local
    #     P = int.params.p + 2./3. * Δt_local * tF1 + 2./3. * tG1 * ΔW_local
    #
    #     t2 = int.params.t + 2./3.*Δt_local
    #
    #     int.params.equ.v(t2, Q, P, tV2)
    #     int.params.equ.B(t2, Q, P, tB2)
    #     int.params.equ.f1(t2, Q, P, tF2)
    #     int.params.equ.f2(t2, Q, P, tmpF)
    #     #tF2 = tF2 + tmpF
    #     int.params.equ.G1(t2, Q, P, tG2)
    #     int.params.equ.G2(t2, Q, P, tmpG)
    #     #tG2 = tG2 + tmpG
    #
    #     #Calculating the Y's and assigning them to the array int.solver.x as initial guesses
    #     for j in 1:int.params.equ.d
    #         int.solver.x[(i-1)*int.params.equ.d+j] =  Δt_local*(1./4.*tV1[j] + 3./4.*tV2[j]) + dot( (1./4.*tB1[j,:] + 3./4.*tB2[j,:]), ΔW_local )
    #     end
    #
    #     # if the collocation points are the same for both q and p parts
    #     # not sure what about pdrift2.c
    #     if int.params.tab.qdrift.c==int.params.tab.pdrift1.c
    #         for j in 1:int.params.equ.d
    #             int.solver.x[(int.params.tab.s+i-1)*int.params.equ.d+j] =  Δt_local*(1./4.*tF1[j] + 3./4.*tF2[j]) + dot( (1./4.*tG1[j,:] + 3./4.*tG2[j,:]), ΔW_local )
    #         end
    #     end
    # end
    #
    #
    # # If qdrift.c != pdrift1.c, then calculating the momenta p at the points pdrift1.c[i]
    # # not sure what about pdrift2.c
    # if int.params.tab.qdrift.c != int.params.tab.pdrift1.c
    #
    #     for i in 1:int.params.tab.s
    #
    #         # Taking the c[i] from the pdrift1 tableau.
    #         Δt_local  = int.params.tab.pdrift1.c[i]*int.params.Δt
    #         ΔW_local .= int.params.tab.pdrift1.c[i]*int.params.ΔW
    #
    #         Q = int.params.q + 2./3. * Δt_local * tV1 + 2./3. * tB1 * ΔW_local
    #         P = int.params.p + 2./3. * Δt_local * tF1 + 2./3. * tG1 * ΔW_local
    #
    #         t2 = int.params.t + 2./3.*Δt_local
    #
    #         int.params.equ.v(t2, Q, P, tV2)
    #         int.params.equ.B(t2, Q, P, tB2)
    #         int.params.equ.f1(t2, Q, P, tF2)
    #         int.params.equ.f2(t2, Q, P, tmpF)
    #         #tF2 = tF2 + tmpF
    #         int.params.equ.G1(t2, Q, P, tG2)
    #         int.params.equ.G2(t2, Q, P, tmpG)
    #         #tG2 = tG2 + tmpG
    #
    #         # Calculating the Z's and assigning them to the array int.solver.x as initial guesses
    #         # The guesses for the Y's have already been written to x above
    #         for j in 1:int.params.equ.d
    #             int.solver.x[(int.params.tab.s+i-1)*int.params.equ.d+j] =  Δt_local*(1./4.*tF1[j] + 3./4.*tF2[j]) + dot( (1./4.*tG1[j,:] + 3./4.*tG2[j,:]), ΔW_local )
    #         end
    #     end
    # end

end

"""
Integrate PSDE with a stochastic implicit partitioned Runge-Kutta integrator.
 Integrating the k-th sample path for the m-th initial condition
"""
function integrate_step!(int::IntegratorSISPRK{DT,TT}, sol::SolutionPSDE{DT,TT,NQ,NW}, k::Int, m::Int, n::Int) where {DT,TT,NQ,NW}
    check_solution_dimension_asserts(sol, k, m, n)

    # set time for nonlinear solver
    int.params.t  = sol.t[0] + (n-1)*int.params.Δt
    int.params.q .= int.q[k, m]
    int.params.p .= int.p[k, m]

    # copy the increments of the Brownian Process
    if NW==1
        #1D Brownian motion, 1 sample path
        int.params.ΔW[1] = sol.W.ΔW[n-1]
        int.params.ΔZ[1] = sol.W.ΔZ[n-1]
    elseif NW==2
        #Multidimensional Brownian motion, 1 sample path
        int.params.ΔW .= sol.W.ΔW[n-1]
        int.params.ΔZ .= sol.W.ΔZ[n-1]
    elseif NW==3
        #1D or Multidimensional Brownian motion, k-th sample path
        int.params.ΔW .= sol.W.ΔW[n-1,k]
        int.params.ΔZ .= sol.W.ΔZ[n-1,k]
    end

    # truncate the increments ΔW with A
    if int.params.A>0
        for i in eachindex(int.params.ΔW)
            if int.params.ΔW[i]<-int.params.A
                int.params.ΔW[i] = -int.params.A
            elseif int.params.ΔW[i]>int.params.A
                int.params.ΔW[i] = int.params.A
            end
        end
    end

    # compute initial guess and assign to int.solver.x
    initial_guess!(int)

    # call nonlinear solver
    solve!(int.solver)

    # print solver status
    print_solver_status(int.solver.status, int.solver.params)

    # check if solution contains NaNs or error bounds are violated
    check_solver_status(int.solver.status, int.solver.params)

    # compute the drift vector field and the diffusion matrix at internal stages
    compute_stages!(int.solver.x, int.fcache.Q, int.fcache.P, int.fcache.V, int.fcache.F1, int.fcache.F2, int.fcache.B, int.fcache.G1, int.fcache.G2, int.fcache.Y, int.fcache.Z, int.params)

    # compute final update
    update_solution!(int.q[k,m], int.p[k,m], int.fcache.V, int.fcache.F1, int.fcache.F2, int.fcache.B, int.fcache.G1, int.fcache.G2,
                    int.params.tab.qdrift.b, int.params.tab.qdrift.b̂, int.params.tab.qdiff.b, int.params.tab.qdiff.b̂,
                    int.params.tab.pdrift1.b, int.params.tab.pdrift1.b̂, int.params.tab.pdrift2.b, int.params.tab.pdrift2.b̂,
                    int.params.tab.pdiff1.b, int.params.tab.pdiff1.b̂, int.params.tab.pdiff2.b, int.params.tab.pdiff2.b̂,
                    int.params.Δt, int.params.ΔW)

    # # NOT IMPLEMENTING InitialGuessSDE
    # # # copy solution to initial guess
    # # update!(int.iguess, m, sol.t[0] + n*int.params.Δt, int.q[m])

    # take care of periodic solutions
    cut_periodic_solution!(int.q[k,m], int.params.equ.periodicity)

    # # copy to solution
    set_solution!(sol, int.q[k,m], int.p[k,m], n, k, m)
end


"""
    val_interp = gaugeinterp(gauge::VisClaw.Gauge, time_interp; varname=:eta::Symbol)

interp gauge values on a non-uniform time
"""
function gaugeinterp(gauge::VisClaw.Gauge, time_interp; varname=:eta::Symbol)
    val = getfield(gauge,varname)
    itp = Interpolations.LinearInterpolation(gauge.time, val, extrapolation_bc=Interpolations.Flat())
    val_interp = itp(time_interp)
    return  val_interp
end


function splice_splines(left_spline::Schumaker, right_spline::Schumaker, splice_point::Real)
    end_in_left_spline = searchsortedlast(left_spline.IntStarts_, splice_point)
    left_starts        = left_spline.IntStarts_[1:end_in_left_spline]
    left_coefficients  = left_spline.coefficient_matrix_[1:end_in_left_spline,:]

    start_in_right     = searchsortedlast(right_spline.IntStarts_, splice_point)
    right_starts       = right_spline.IntStarts_[start_in_right:length(right_spline.IntStarts_)]
    right_coefficients = right_spline.coefficient_matrix_[start_in_right:length(right_spline.IntStarts_),:]

    # As the splice_point is probably not an interval start in the right_spline we need to adjust.
    G                       = splice_point - right_starts[1]
    a                       = right_coefficients[1,1]
    b                       = right_coefficients[1,2]
    c                       = right_coefficients[1,3]
    right_starts[1]         = splice_point
    right_coefficients[1,2] = 2*G*a + b
    right_coefficients[1,3] = a*(G^2) + b*G + c

    IntStarts = vcat(left_starts, right_starts)
    Coeffs    = vcat(left_coefficients, right_coefficients)
    return Schumaker(IntStarts, Coeffs)
end


# # Sums of Hermitian squares

#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/Sums of Hermitian squares.ipynb)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/Sums of Hermitian squares.ipynb)
# **Contributed by**: Benoît Legat

using Test #src

using SumOfSquares

using DynamicPolynomials
@ncpolyvar x y
p = (x + 1.0im * y) * (x - im * y)

import CSDP
model = Model(CSDP.Optimizer)
add_bridge(model, SumOfSquares.Bridges.Constraint.EmptyBridge)
add_bridge(model, SumOfSquares.Bridges.Constraint.PositiveSemidefinite2x2Bridge)
add_bridge(model, SumOfSquares.COI.Bridges.Variable.HermitianToSymmetricPSDBridge)
add_bridge(model, SumOfSquares.COI.Bridges.Constraint.SplitZeroBridge)
MOI.Bridges.add_bridge(backend(model).optimizer, PolyJuMP.ZeroPolynomialBridge{Complex{Float64}})
MOI.Bridges.add_bridge(backend(model).optimizer, SumOfSquares.Bridges.Constraint.SOSPolynomialBridge{Complex{Float64}})
cone = NonnegPolyInnerCone{SumOfSquares.COI.HermitianPositiveSemidefiniteConeTriangle}()
certificate = SumOfSquares.Certificate.MaxDegree(cone, MonomialBasis, 2)
c = SumOfSquares.add_constraint(model, p, cone, ideal_certificate = certificate)
optimize!(model)
dec = sos_decomposition(c, 1e-6) #src
@test length(dec.ps) == 1 #src
@test sign(real(first(coefficients(dec.ps[1])))) * dec.ps[1] ≈ x - im * y atol=1e-6 rtol=1e-6 #src
sos_decomposition(c, 1e-6) #!src


using Test
using CuboidGrids
Lar = CuboidGrids

@testset "LarGrid-0-1 Tests" begin
   @testset "grid_0" begin
      @test Lar.grid_0(1) == [0  1]
      @test Lar.grid_0(2) == [0  1  2]
      @test Lar.grid_0(3) == [0  1  2  3]
   end

   @testset "Grid_1 Tests" begin
      @test repr(Lar.grid_1(1)) == "[0; 1]"
      @test Lar.grid_1(2) == [0 1; 1 2]
      @test Lar.grid_1(3) == [0 1 2; 1 2 3]
   end

   @testset "Grid Tests" begin
      @test repr(Lar.larGrid(1)(0)) == "[0 1]"
      @test repr(Lar.larGrid(1)(1)) == "[0; 1]"
      @test Lar.larGrid(2)(0) == [0  1  2]
      @test Lar.larGrid(2)(1) == [0 1; 1 2]
      @test Lar.larGrid(3)(0) == [0  1  2  3]
      @test Lar.larGrid(3)(1) == [0  1  2; 1  2  3]
   end
end

@testset "LarVertProd" begin
      vertexDomain(n) = hcat([k for k in 0:n-1]...)

      @testset "LarVertProd 1D" begin
            shape = [3]
	      vertLists = [vertexDomain(k+1) for k in shape]
            @test typeof(Lar.larVertProd(vertLists))==Array{Int64,2}
            @test size(Lar.larVertProd(vertLists))==(1, 4)
            @test Lar.larVertProd(vertLists)[:,1]==[0]
            @test Lar.larVertProd(vertLists)[:,4]==[3]
      end

      @testset "LarVertProd 2D" begin
            shape = [3,2]
	      vertLists = [vertexDomain(k+1) for k in shape]
            @test typeof(Lar.larVertProd(vertLists))==Array{Int64,2}
            @test size(Lar.larVertProd(vertLists))==(2, 12)
            @test Lar.larVertProd(vertLists)[:,1]==[0;0]
            @test Lar.larVertProd(vertLists)[:,12]==[3;2]
      end

      @testset "LarVertProd 3D" begin
            shape = [3,2,1]
	      vertLists = [vertexDomain(k+1) for k in shape]
            @test typeof(Lar.larVertProd(vertLists))==Array{Int64,2}
            @test size(Lar.larVertProd(vertLists))==(3, 24)
            @test Lar.larVertProd(vertLists)[:,1]==[0;0;0]
            @test Lar.larVertProd(vertLists)[:,24]==[3;2;1]
      end
end

@testset "Index2addr Tests" begin
   @testset "Shape 1D Tests" begin
      @test Lar.index2addr([10])([0])==1
      @test Lar.index2addr([10])([9])==10
      @test [Lar.index2addr([10])([index]) for index in collect(0:9)]==collect(1:10)
   end

   @testset "Shape 2d Tests" begin
      aa = Lar.cart([[0;1;2],[0;1]])
      bb = Lar.cart([collect(0:9),collect(0:1)])
      cc = Lar.cart([collect(0:2),collect(0:2)])
      dd = "Tuple{Int64,Int64}[(0,0),(0,1),(1,0),(1,1),(2,0),(2,1)]"
      ee = "Tuple{Int64,Int64}[(0, 0), (0, 1), (1, 0), (1, 1), (2, 0), (2, 1)]"
      ff = repr( Lar.cart([collect(0:2),collect(0:1)]) )
      @test [ Lar.index2addr([3,2])(collect(index)) for index in aa ]==collect(1:6)
      @test [ Lar.index2addr([10,2])(collect(index)) for index in bb ] == collect(1:20)
      @test [ Lar.index2addr([3,3])(collect(index)) for index in cc] == collect(1:9)
   end

   @testset "Shape 3d Tests" begin
      aaa = Lar.cart([collect(0:3),collect(0:2),collect(0:1)])
      bbb = Lar.cart([[0;1;2],[0;1],[0;1]])
      ccc = Lar.cart([[0;1],[0;1],[0;1;2]])
      @test Lar.index2addr([3,2,1])([0,0,0]) == 1
      @test [Lar.index2addr([4,3,2])(collect(index))
      	for index in aaa ] == collect(1:24)
      @test [Lar.index2addr([3,2,2])(collect(index))
      	for index in bbb ] == collect(1:12)
      @test [Lar.index2addr([2,2,3])(collect(index))
      	for index in ccc ] == collect(1:12)
   end
end

@testset "LarCellProd Tests" begin
	 @testset "3D shape" begin
		@test typeof(Lar.larGridSkeleton([3,2,1])(3)) ==
		   		Array{Array{Int64,1},1}
		@test typeof(Lar.larGridSkeleton([3,2,1])(2)) ==
		   		Array{Array{Int64,1},1}
		@test typeof(Lar.larGridSkeleton([3,2,1])(1)) ==
		   		Array{Array{Int64,1},1}
		@test typeof(Lar.larGridSkeleton([3,2,1])(0)) ==
		   		Array{Array{Int64,1},1}
		@test typeof(Lar.larGridSkeleton([1,1,1])(3)) ==
		   		Array{Array{Int64,1},1}
		@test typeof(Lar.larGridSkeleton([1,1,1])(2)) ==
		   		Array{Array{Int64,1},1}
		@test typeof(Lar.larGridSkeleton([1,1,1])(1)) ==
		   		Array{Array{Int64,1},1}
		@test typeof(Lar.larGridSkeleton([1,1,1])(0)) ==
		   		Array{Array{Int64,1},1}
	 end
	 @testset "2D shape" begin
		@test typeof(Lar.larGridSkeleton([3,2])(2)) ==
		   		Array{Array{Int64,1},1}
		@test typeof(Lar.larGridSkeleton([3,2])(1)) ==
		   		Array{Array{Int64,1},1}
		@test typeof(Lar.larGridSkeleton([3,2])(0)) ==
		   		Array{Array{Int64,1},1}
	 end
	 @testset "1D shape" begin
		@test typeof(Lar.larGridSkeleton([3])(1)) ==
		   		Array{Array{Int64,1},1}
		@test typeof(Lar.larGridSkeleton([3])(0)) ==
		   		Array{Array{Int64,1},1}
	 end
end

@testset "FilterByOrder Tests" begin
      term = "000"
      bit = '0'
      theTerm = Vector{Char}(term)
      @test typeof(theTerm) == Array{Char,1}
      @test parse(Int8,bit) == 0
      @test [parse(Int8,bit) for bit in theTerm] == zeros(3)
      out = hcat([[parse(Int8,bit) for bit in term] for term in
      	Lar.binaryRange(3)]...)
      @test typeof(out) == Array{Int8,2}
      @test size(out) == (3,8)
      @test repr(out) == "Int8[0 0 0 0 1 1 1 1; 0 0 1 1 0 0 1 1; 0 1 0 1 0 1 0 1]"

      @testset "$n" for n in 1:4
            data = [Lar.filterByOrder(n)[k] for (k,el) in enumerate(Lar.filterByOrder(n))]
            @test sum(map(length,data)) == 2^n
      end

      @testset "$n,$k" for n in 1:4, k in 0:n
            @test length(Lar.filterByOrder(n)[k+1]) == binomial(n,k)
      end
end

@testset "LarImageVerts Tests" "$shape" for shape in [[3,2,1],[3,2],[10,10,10]]
          @test size(Lar.larImageVerts(shape)) == (length(shape),prod(shape .+ 1))
       end

@testset "LarCuboids Tests" begin
      @testset "Simple shape" begin
            shape = [3,2,1]
            cubes = Lar.cuboidGrid(shape,true)
            verts, cells = cubes
            @test typeof(verts) == Array{Float64,2}
            VV,EV,FV,CV = cells
            @testset "$basis" for basis in [VV,EV,FV,CV]
              @test typeof(basis) == Array{Array{Int64,1},1}
            end
      end

      @testset "$shape" for shape in [[3,2,1],[1,1,1],[3,3,10]]
            endShape = shape .+ 1
            cubes = Lar.cuboidGrid(shape,true);
            verts, cells = cubes;
            VV,EV,FV,CV = cells;
            @test size(Lar.cuboidGrid(shape)[1],2) == prod(collect(endShape))
            @test length(Lar.cuboidGrid(shape)[2]) == prod(shape)
            @test typeof(verts) == Array{Float64,2}
	      @test typeof(VV) == Array{Array{Int64,1},1}
	      @test typeof(EV) == Array{Array{Int64,1},1}
	      @test typeof(FV) == Array{Array{Int64,1},1}
	      @test typeof(FV) == Array{Array{Int64,1},1}
      end

      @testset "$shape" for shape in [[3,2],[1,1],[3,10]]
            endShape = shape .+ 1
            @test size(Lar.cuboidGrid(shape)[1],2) == prod(collect(endShape))
            @test length(Lar.cuboidGrid(shape)[2]) == prod(shape)
            cubes = Lar.cuboidGrid(shape,true)
            verts, cells = cubes
            @test typeof(verts) == Array{Float64,2}
            VV,EV,FV = cells
	      @test typeof(VV) == Array{Array{Int64,1},1}
	      @test typeof(EV) == Array{Array{Int64,1},1}
	      @test typeof(FV) == Array{Array{Int64,1},1}
      end

      @testset "$shape" for shape in [[3],[1],[10]]
            endShape = shape .+ 1
            @test size(Lar.cuboidGrid(shape)[1],2) == prod(collect(endShape))
            @test length(Lar.cuboidGrid(shape)[2]) == prod(shape)
            cubes = Lar.cuboidGrid(shape,true)
            verts, cells = cubes
            @test typeof(verts) == Array{Float64,2}
            VV,EV = cells
	      @test typeof(VV) == Array{Array{Int64,1},1}
	      @test typeof(EV) == Array{Array{Int64,1},1}
      end
end

@testset "LarModelProduct Tests" begin
      geom_0 = hcat([[x] for x=0.:10]...)
      topol_0 = [[i,i+1] for i=1:9]
      geom_1 = hcat([[0.],[1.],[2.]]...)
      topol_1 = [[1,2],[2,3]]
      model_0 = (geom_0, topol_0)
      model_1 = (geom_1, topol_1)
      model_2 = Lar.larModelProduct(model_0, model_1)
      model_3 = Lar.larModelProduct(model_2, model_1)

      @testset "Result size" begin
            @test size(model_2[1],2) == size(model_0[1],2) * size(model_1[1],2)
            @test length(model_2[2]) == length(model_0[2]) * length(model_1[2])
            @test size(model_3[1],2) == size(model_2[1],2) * size(model_1[1],2)
            @test length(model_3[2]) == length(model_2[2]) * length(model_1[2])
      end

      @testset "Model 0 type" begin
            m0_V,m0_EV = model_0
            @test typeof(m0_V) == Array{Float64,2}
            @test typeof(m0_EV) == Array{Array{Int64,1},1}
      end

      @testset "Model 1 type" begin
            m1_V,m1_EV = model_1
            @test typeof(m1_V) == Array{Float64,2}
            @test typeof(m1_EV) == Array{Array{Int64,1},1}
      end

      @testset "Model 2 type" begin
            m2_V,m2_EV = model_2
            @test typeof(m2_V) == Array{Float64,2}
            @test typeof(m2_EV) == Array{Array{Int64,1},1}
      end

      @testset "Model 3 type" begin
            m3_V,m3_EV = model_3
            @test typeof(m3_V) == Array{Float64,2}
            @test typeof(m3_EV) == Array{Array{Int64,1},1}
      end

      @testset "Simple product" begin
            modelOne = ([0 0 1 1; 0 1 0 1], Array{Int64,1}[[1, 2, 3, 4]])
            modelTwo = ([0 1], Array{Int64,1}[[1, 2]])
            @test Lar.larModelProduct(modelOne, modelTwo) == Lar.larModelProduct(modelTwo, modelOne)
            @test Lar.larModelProduct(modelTwo, modelTwo)[1] == [0  0  1  1; 0  1  0  1]
            @test Lar.larModelProduct(modelOne,modelOne)[2][1] == Any[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
      end
end


using ITensors

examples_dir = joinpath(ITensors.examples_dir(), "src")
include(joinpath(examples_dir, "trg.jl"))
include(joinpath(examples_dir, "2d_classical_ising.jl"))

function run(; maxdim::Int,
               nsweeps::Int = 20,
               outputlevel::Int = 0,
               cutoff::Float64 = 0.0,
               β::Float64 = 1.001 * βc)
  # Make Ising model MPO
  s = Index(2)
  sₕ = addtags(s, "horiz")
  sᵥ = addtags(s, "vert")
  T = ising_mpo(sₕ, sᵥ, β)
  κ, T = trg(T; χmax = maxdim, cutoff = cutoff, nsteps = nsweeps)
  if outputlevel > 0
    @show nsweeps
    @show cutoff
    @show κ, exp(-β * ising_free_energy(β))
  end
  return ITensors.maxdim(T)
end


function a_ideal(model::SAFT, z, v, T)
    x = z/sum(z)
    return sum(x[i]*log(z[i]/v) for i in model.components)-1
end


using JuMP
using SCIP
using Clp;
using Cbc
model = Model(SCIP.Optimizer)
# set_optimizer_attribute(model, "msg_lev", GLPK.GLP_MSG_ON)
# set_optimizer_attribute(model, "maxNodes", 5)
month_num = 8
product_num = 1
month = 1:month_num
product = 1:product_num
# paramerters
q = [3000] # product fix cost
p = [60] # cost per product
d = [200, 200, 200, 200, 200, 200, 500, 500]' # demand of product by month
s_init = [100] # initial stock
h = [3] # stock cost per product month

@variable(model, x[i in product, j in month] >= 0)
@variable(model, s[i in product, j in month] >= 0)
@variable(model, y[i in product, j in month], Bin)

# 库存平衡
@constraint(model ,dem_sat1[i in product], s_init[i] + x[i, 1] == d[i, 1] + s[i, 1])
@constraint(model, dem_sat[i in product, j in month[2:end]], s[i, j - 1] + x[i, j] == d[i, j] + s[i, j] )
# 最大产量约束
@constraint(model, vub[i in product, j in month], x[i, j] <= sum(d[i, k]  for k in j:month_num) * y[i, j])
# 目标函数
@expression(model ,cost, sum(p[i] * x[i, j] + q[i] * y[i, j] for i in product, j in month) + sum(h[i] * s[i, j] for i in product, j in month[1:end - 1]) + sum(h[i] * s[i, month[end]] / 2 for i in product))
@objective(model, Min, cost)
optimize!(model)
# 显示结果
cost_val = value.(cost)
x_val = value.(x)
s_val = value.(s)
@show cost_val;
@show x_val;
@show s_val;

model2 = Model(SCIP.Optimizer);
x_val[1,1] = 600
x_val[1,4] = 500
order = 1:8;
orderquantity = [100, 200, 200, 200, 200, 200, 500, 500];
ordertime = [1,2,3,4,5,6,7,8];
@variable(model2, wx[t in month, p in order] >= 0);
@variable(model2, 0 <= wy[t in month, p in order] <= 1, Bin);
@constraint(model2, dem_sato[p in order], sum(wx[t, p] for t in 1:ordertime[p])  == orderquantity[p]);
# @constraint(model2, xlimit[t in month, p in order], wx[t, p] <= x_val[1,t])
@constraint(model2, xlimit[t in month, p in order], wx[t, p] <= x_val[1,t] * wy[t, p]);
@constraint(model2, xeq[t in month], sum(wx[t, p] for p in order) == x_val[1, t]);
@objective(model2, Min, sum(wy[t, p] * ordertime[p] for t in month, p in order));
optimize!(model2);
@show value.(wx);
@show value.(wy);


let
  i = Index(2)
  A = randomITensor(i, i', i'')
  for T in (Float64, ComplexF64)
    A = randomITensor(T, i, i')
    B = randomITensor(T, i', i'')
    C = A * B
    U, S, V = svd(A, i)
    A, B = factorize(A, i)
  end
end


const Point{T} = AbstractVector{T} where T<:Real
const Points = AbstractVector{<:Point}
const Sigma = AbstractVector{<:Integer}  # Encoded the vertex by the ids of its generators.
const Vertex = Tuple{<:Sigma, <: Point}
const Vertices = Dict{<:Sigma, <:Point}

voronoi(x; kwargs...) = voronoi(vecvec(x); kwargs...)

""" construct the voronoi diagram from `x` through breadth-first search """
function voronoi(xs::Points, searcher = Raycast(xs))
    sig, r = descent(xs, searcher)
    verts, boundary = explore(sig, r, xs, searcher)
    return verts::Vertices, xs, boundary
end

voronoi_random(x, args...; kwargs...) = voronoi_random(vecvec(x), args...; kwargs...)

""" construct a (partial) voronoi diagram from `x` through a random walk """
function voronoi_random(xs::Points, iter=1000; maxstuck=typemax(Int))
    searcher = Raycast(xs)
    sig, r = descent(xs, searcher)
    verts = walk(sig, r, iter, xs, searcher, maxstuck)
    return verts::Vertices, xs
end

vecvec(x::Matrix) = map(SVector{size(x,1)}, eachcol(x))
vecvec(x::Vector{<:SVector}) = x

""" starting at given points, run the ray shooting descent to find vertices """
function descent(xs::Points, searcher, start = 1) :: Vertex
    sig = [start]
    r = xs[start]
    d = length(r)

    for k in d:-1:1  # find an additional generator for each dimension
        u = randray(xs[sig])
        (tau, t) = raycast(sig, r, u, xs, searcher)
        if t == Inf
            u = -u
            (tau, t) = raycast(sig, r, u, xs, searcher)
        end
        if t == Inf
            error("Could not find a vertex in both directions of current point." *
                "Consider increasing search range (tmax)")
        end
        sig = tau
        r = r + t*u
    end

    return (sig, r)
end

""" starting at vertices, walk nsteps along the voronoi graph to find new vertices """
function walk(sig::Sigma, r::Point, nsteps::Int, xs::Points, searcher, maxstuck::Int) :: Vertices
    verts = Dict(sig => r)
    nonew = 0
    prog = Progress(maxstuck, 1, "Voronoi walk")
    progmax = 0

    for s in 1:nsteps
        nonew += 1
        progmax = max(progmax, nonew)
        ProgressMeter.update!(prog, progmax)
        randdir = rand(1:length(sig))
        sig, r = walkray(sig, r, xs, searcher, randdir)
        r == Inf && continue  # we encountered an infinite ray
        get!(verts, sig) do
            nonew = 0
            return r
        end
        nonew > maxstuck && break
    end

    return verts
end

""" find the vertex connected to `v` by moving away from its `i`-th generator """
function walkray(sig::Sigma, r::Point, xs::Points, searcher, i)
    sig_del = deleteat(sig, i)
    u = randray(xs[sig_del])
    if (u' * (xs[sig[i]] - xs[sig_del[1]])) > 0
        u = -u
    end
    sig′, t = raycast(sig_del, r, u, xs, searcher)
    if t < Inf
        r′ = r + t*u
        return sig′, r′
    else
        return sig, r  # if the vertex has an unbounded ray, return the same vertex
    end
end

""" BFS of vertices starting from `S0` """
function explore(sig, r, xs::Points, searcher) # :: Vertices
    verts = Dict(sig=>r)
    queue = copy(verts)
    edgecount = Dict{Vector{Int64}, Int}()
    boundary = []

    cache = true  # Caches if both points along an edge are known. Trades memory for runtime.
    hit = 0
    miss = 0
    new = 0
    unbounded = 0

    while length(queue) > 0
        (sig,r) = pop!(queue)
        for i in 1:length(sig)

            if cache && get(edgecount, deleteat(sig, i), 0) == 2
                hit += 1
                continue
            end

            sig′, r′ = walkray(sig, r, xs, searcher, i)

            if sig′ == sig
                push!(boundary, sig)
                unbounded += 1
                continue
            end

            if !haskey(verts, sig′)
                push!(queue, sig′ => r′)
                push!(verts, sig′ => r′)
                if cache
                    for j in 1:length(sig′)
                        edge = deleteat(sig′, j)
                        edgecount[edge] = get(edgecount, edge, 0) + 1
                    end
                end
                new += 1
            else
                miss += 1
            end
        end

    end

    #@show hit, miss, new, unbounded
    return verts, boundary
end

deleteat(sig, i) = deleteat!(copy(sig), i)

""" generate a random ray orthogonal to the subspace spanned by the given points """
function randray(xs::Points)
    k = length(xs)
    d = length(xs[1])
    v = similar(xs, k-1)

    # Gram Schmidt
    for i in 1:k-1
        v[i] = xs[i] .- xs[k]
        for j in 1:(i-1)
            v[i] = v[i] .- dot(v[i], v[j]) .* v[j]
        end
        v[i] = normalize(v[i])
    end

    u = randn(d)
    for i in 1:k-1
        u = u - dot(u, v[i]) * v[i]
    end
    u = normalize(u)
    return u
end


module LaplaceBeltramiCubedSphereTestsSeq

  using Test
  using Gridap
  using GridapGeosciences
  using Plots

  include("../ConvergenceAnalysisTools.jl")
  include("../LaplaceBeltramiCubedSphereTests.jl")

  n_values=generate_n_values(2)
  hs=[2.0/n for n in n_values]
  model_args_series=[(n,) for n in n_values]

  @time ahs1,ak1errors,as1=convergence_study(solve_laplace_beltrami,hs,model_args_series,1,8)
  @test round(as1,digits=1) ≈ 1.0

  @time ahs2,ak2errors,as2=convergence_study(solve_laplace_beltrami,hs,model_args_series,2,12)
  @test round(as2,digits=1) ≈ 2.0

end


module ReusableResourceAllocation

using JuMP
using Gurobi

using LinearAlgebra
using MathOptInterface

export reusable_resource_allocation

function reusable_resource_allocation(
		initial_supply::Array{<:Real,1},
		supply::Array{<:Real,2},
		demand::Array{<:Real,2},
		adj_matrix::BitArray{2};
		obj_dir::Symbol=:shortage,
		send_new_only::Bool=false,
		sendrecieve_switch_time::Int=0,
		min_send_amt::Real=0,
		smoothness_penalty::Real=0,
		setup_cost::Real=0,
		sent_penalty::Real=0,
		verbose::Bool=false,
)
	N, T = size(supply)
	@assert(size(initial_supply, 1) == N)
	@assert(size(demand, 1) == N)
	@assert(size(demand, 2) == T)
	@assert(size(adj_matrix, 1) == N)
	@assert(size(adj_matrix, 2) == N)
	@assert(obj_dir in [:shortage, :overflow])

	model = Model(Gurobi.Optimizer)
	if !verbose set_silent(model) end

	@variable(model, sent[1:N,1:N,1:T])
	@variable(model, obj_dummy[1:N,1:T] >= 0)

	if min_send_amt <= 0
		@constraint(model, sent .>= 0)
	else
		@constraint(model, [i=1:N,j=1:N,t=1:T], sent[i,j,t] in MOI.Semicontinuous(Float64(min_send_amt), Inf))
	end

	objective = @expression(model, sum(obj_dummy))
	if sent_penalty > 0
		add_to_expression!(objective, sent_penalty*sum(sent))
	end
	if smoothness_penalty > 0
		@variable(model, smoothness_dummy[i=1:N,j=1:N,t=1:T-1] >= 0)
		@constraint(model, [t=1:T-1],  (sent[:,:,t] - sent[:,:,t+1]) .<= smoothness_dummy[:,:,t])
		@constraint(model, [t=1:T-1], -(sent[:,:,t] - sent[:,:,t+1]) .<= smoothness_dummy[:,:,t])

		add_to_expression!(objective, smoothness_penalty * sum(smoothness_dummy))
		add_to_expression!(objective, smoothness_penalty * sum(sent[:,:,1]))
	end
	if setup_cost > 0
		@variable(model, setup_dummy[i=1:N,j=i+1:N], Bin)
		@constraint(model, [i=1:N,j=i+1:N], [1-setup_dummy[i,j], sum(sent[i,j,:])+sum(sent[j,i,:])] in MOI.SOS1([1.0, 1.0]))
		add_to_expression!(objective, setup_cost*sum(setup_dummy))
	end
	@objective(model, Min, objective)

	if send_new_only
		@constraint(model, [t=1:T],
			sum(sent[:,:,t], dims=2) .<= max.(0, supply[:,t])
		)
	else
		@constraint(model, [i=1:N,t=1:T],
			sum(sent[i,:,t]) <=
				initial_supply[i]
				+ sum(supply[i,1:t])
				- sum(sent[i,:,1:t-1])
				+ sum(sent[:,i,1:t-1])
		)
	end

	for i = 1:N
		for j = 1:N
			if ~adj_matrix[i,j]
				@constraint(model, sum(sent[i,j,:]) .== 0)
			end
		end
	end

	if sendrecieve_switch_time > 0
		@constraint(model, [i=1:N,t=1:T-1],
			[sum(sent[:,i,t]), sum(sent[i,:,t:min(t+sendrecieve_switch_time,T)])] in MOI.SOS1([1.0, 1.0])
		)
		@constraint(model, [i=1:N,t=1:T-1],
			[sum(sent[:,i,t:min(t+sendrecieve_switch_time,T)]), sum(sent[i,:,t])] in MOI.SOS1([1.0, 1.0])
		)
	end

	flip_sign = (obj_dir == :shortage) ? 1 : -1
	z1, z2 = (obj_dir == :shortage) ? (0, -1) : (-1, 0)
	@constraint(model, [i=1:N,t=1:T],
		obj_dummy[i,t] >= flip_sign * (
			demand[i,t] - (
				initial_supply[i]
				+ sum(supply[i,1:t])
				- sum(sent[i,:,1:t+z1])
				+ sum(sent[:,i,1:t+z2])
			)
		)
	)

	optimize!(model)
	return model
end

end;


using ApproxFun, DualNumbers

import EquilibriumMeasures: emmoment

Vp = x -> -im*x^2

# Vp = x -> im
a,b = -1.1+0.5im,1+im
a,b = 1,im
N = (a,b) -> (f = Fun(Vp, Segment(a,b)); [f.coefficients[1],emmoment(f)-1])

J=[epsilon.(N(dual(a,1),b)) epsilon.(N(a,dual(b,1)))]
a,b = [a,b]-J\N(a,b)

plot(Segment(a,b))


###################
# Dynamics Models #
###################

export fderiv_earth_orbit
"""
Compute the state derivative.

Arguments:
- `epc::Epoch`: Current epoch
- `x::Array{<:Real, 1}`: Satellite state vector
- `mass::Real`: Satellite mass [kg]
- `area_drag`: Velocity-facing area affected by drag. [m^2]
- `coef_drag`: Coefficient of drag [dimensionless]
- `area_srp`: Velocity-facing area affected by drag. [m^2]
- `coef_srp`: Coefficient of reflectivity [dimensionless]
- `n_grav::Integer`: Gravity model degree (Default: 20)
- `m_grav::Integer`: Gravity model order (Default: 20)
- `drag::Bool`: Include cannonball atomospheric drag in force model (Default: `true`)
- `srp::Bool`: Include flat-plate solar radiation pressure in force model (Default: `true`)
- `moon::Bool`: Include thridbody lunar gravity in force model (Default: `true`)
- `sun::Bool`: Include thirdbody solar in force model (Default: `true`)
- `relativity::Bool`: Include relativistic effects in force model (Default: `true`)

Returns:
- `dx::Array{<:Float64, 1}`: Satellite state derivative, velocity and accelerations [m; m/s]
"""
function fderiv_earth_orbit(epc::Epoch, x::Array{<:Real} ;
             mass::Real=1.0, area_drag::Real=1.0, coef_drag::Real=2.3,
             area_srp::Real=1.0, coef_srp::Real=1.8,
             n_grav::Integer=20, m_grav::Integer=20,
             drag::Bool=true, srp::Bool=true, moon::Bool=true, sun::Bool=true,
             relativity::Bool=true)

    # Extract position and velocity
    r = x[1:3]
    v = x[4:6]

    # Compute ECI to ECEF Transformation -> IAU2010 Theory
    PN = bias_precession_nutation(epc)
    E  = earth_rotation(epc)
    W  = polar_motion(epc)
    R  = W * E * PN

    # Compute geolocation
    geod = sECEFtoGEOD(R*r)

    # Compute sun and moon position
    r_sun  = sun_position(epc)
    r_moon = moon_position(epc)

    # Compute acceleration
    a = zeros(Float64, 3)

    a += accel_gravity(x, R, n_grav, m_grav)

    # Drag
    if drag
        # Use PN*x[1:3] to compute the satellite position in the true-of-date inertial frame
        rho = density_nrlmsise00(epc, geod)
        a  += accel_drag(x, rho, mass, area_drag, coef_drag, Array{Real, 2}(PN))
    end

    # Solar Radiation Pressure
    if srp
        nu = eclipse_conical(x, r_sun)
        a += nu*accel_srp(x, r_sun, mass, area_srp, coef_srp)
    end

    # Thirdbody
    if sun
        a += accel_thirdbody_sun(x, r_sun)
    end

    if moon
        a += accel_thirdbody_moon(x, r_moon)
    end

    # Relativity
    if relativity
        a += accel_relativity(x)
    end

    return vcat(v, a)
end

##############
# Propagator #
##############

export EarthInertialState
"""
Satellite orbit state vector using an Earth inertial state representation and
dynamics model.

If an initial state transition matrix is provided it will be used for propagation
if there is no state transition matrix then only the state will be propagated.

Attributes:
- `rk4::RK4` Internal numerical integrator used for state propagation
- `dt::Real` Default propagation time step. All steps will be this size unless
the state vector is requested to propagate to a time smaller than this step size,
which it will do.
- `epc::Epoch` Epoch of state
- `x::Array{Float64, 1}` State vector. Earth-centered inertial Cartesian state.
- `phi::Union{Nothing, Array{Float64, 2}}` State transition matrix, or the matrix
of partial derivatives of the state at the current time with respect to the
start of propagation.

The following force model parametters can be set as keyword arguments

Parameters:
- `mass::Real`: Satellite mass [kg]
- `area_drag`: Velocity-facing area affected by drag. [m^2]
- `coef_drag`: Coefficient of drag [dimensionless]
- `area_srp`: Velocity-facing area affected by drag. [m^2]
- `coef_srp`: Coefficient of reflectivity [dimensionless]
- `n_grav::Integer`: Gravity model degree (Default: 20)
- `m_grav::Integer`: Gravity model order (Default: 20)
- `drag::Bool`: Include cannonball atomospheric drag in force model (Default: `true`)
- `srp::Bool`: Include flat-plate solar radiation pressure in force model (Default: `true`)
- `moon::Bool`: Include thridbody lunar gravity in force model (Default: `true`)
- `sun::Bool`: Include thirdbody solar in force model (Default: `true`)
- `relativity::Bool`: Include relativistic effects in force model (Default: `true`)

"""
mutable struct EarthInertialState
    rk4::RK4
    dt::Real
    epc::Epoch
    x::Array{Float64, 1}
    phi::Union{Nothing, Array{Float64, 2}}
end

function EarthInertialState(epc::Epoch, x::Array{<:Real, 1}, phi::Union{Nothing, Array{Float64, 2}}=nothing; dt::Real=30.0, kwargs...)
    rk4 = RK4(fderiv_earth_orbit; kwargs...)
    return EarthInertialState(rk4, dt, epc, x, phi)
end

export step!
"""
Step dynamics state the specified time.

Arguments:
- `state::EarthInertialState` State vector to propagate
- `dt::Real` Step size. If no value is provided will use default state stepsize
"""
function step!(state::EarthInertialState, dt::Real=0.0)
    # Set step size
    if dt == 0.0
        dt = state.h
    end

    if state.phi == nothing
        # If no state transition matrix is provided integrate
        # without propagating variational equations
        state.x = istep(state.rk4, state.epc, dt, state.x)
    else
        state.x, state.phi = istep(state.rk4, state.epc, dt, state.x, state.phi)
    end

    # Advance state propagator time
    state.epc += dt
end

export stepto!
"""
Step dynamics state to the specified time. Will take as many internal steps as
required to advance the propagator to the correct time. No internal step will
exceed the size specified in the state propagator.

Arguments:
- `state::EarthInertialState` State vector to propagate
- `time::Union{Real, Epoch}` Time to propagate internal state too. Can be either
a real number to advance the state by or the Epoch
"""
function stepto!(state::EarthInertialState, time::Union{Real, Epoch}=0.0)
    if typeof(time) <: Real
        time = state.epc + time
    end

    if time < state.epc
        @warn "Propagation time is before current state epoch. Will integrate dyanmics in reverse."
    end

    # Propagate state
    while state.epc < time
        # Set propagation state size
        dt = min(state.dt, time - state.epc)

        # Propagate state
        step!(state, dt)
    end
end

export sim!
"""
Simulate the sttate dynamics to the specified time. Takes same inputs as `stepto!`
but instead of just updating the state vector to the specified time, this
function also returns the timesteps, state values, and state transition
matrices for each time step.

Arguments:
- `state::EarthInertialState` State vector to propagate
- `time::Union{Real, Epoch}` Time to propagate internal state too. Can be either
a real number to advance the state by or the Epoch

Returns:
- `t::Array{Float64, 1}` Elapsed time as a scalar from the initial simulation epoch
- `epc::Array{Epoch, 1}` Epoch at each timestep
- `x::Array{Float64, 2}` State vectors at each time step. Time is along second axis
- `Phi::Array{Float64, 2}` Stacked array of state transition matrices
"""
function sim!(state::EarthInertialState, time::Union{Real, Epoch}=0.0)
    if typeof(time) <: Real
        time = state.epc + time
    end

    if time < state.epc
        @warn "Propagation time is before current state epoch. Will integrate dyanmics in reverse."
    end

    # Number of state variables
    n = length(state.x)

    # Compute number of time steps in propagation
    n_steps = ceil(Int, (time - state.epc)/state.dt) + 1

    # Initialize containers to store output
    t   = zeros(Float64, n_steps)
    epc = Array{Epoch}(undef, n_steps)
    x   = zeros(Float64, n, n_steps)
    A   = zeros(Float64, n*n_steps, n)

    # Save initial state
    idx = 1

    epc[idx]  = state.epc
    t[idx]    = epc[idx] - epc[1]
    x[:, idx] = state.x

    if state.phi != nothing
        A[(1+n*(idx-1)):(n*idx), :] = state.phi
    end

    # Propagate state
    while state.epc < time
        # Increase storage index
        idx += 1

        # Set propagation state size
        dt = min(state.dt, time - state.epc)

        # Propagate state
        step!(state, dt)

        # Save output
        epc[idx]  = state.epc
        t[idx]    = epc[idx] - epc[1]
        x[:, idx] = state.x
        if state.phi != nothing
            A[(1+n*(idx-1)):(n*idx), :] = state.phi
        end
    end

    # Output
    if state.phi != nothing
        # Return with state transition matrix if used
        return t, epc, x, A
    end

    # Default return without STM
    return t, epc, x
end

export reinit!
"""
Reinitialize State transition matrix to identity at the current time step.

Used to reinitialize the state transition matrix value to identity
"""
function reinit!(state::EarthInertialState)
    state.phi = diagm(0 => ones(Float64, length(state.x)))
end


# This script actually solves the WachterDisasterRisk model with a risk-adjusted linearization
# and times the methods, if desired
using BenchmarkTools, RiskAdjustedLinearizations, Test
include(joinpath(dirname(@__FILE__), "..", "..", "examples", "rbc_cc", "rbc_cc.jl"))

# Set up
m_rbc_cc = RBCCampbellCochraneHabits()
m = rbc_cc(m_rbc_cc, 0)

# Solve!
solve!(m; algorithm = :relaxation)

sssout = JLD2.jldopen(joinpath(dirname(@__FILE__), "..", "reference", "rbccc_sss_iterative_output.jld2"), "r")

@test isapprox(sssout["z_rss"], m.z, atol=1e-4)
@test isapprox(sssout["y_rss"], m.y, atol=1e-4)
@test isapprox(sssout["Psi_rss"], m.Ψ, atol=1e-4)


export weightednorm
export entropyintegral

using LinearAlgebra: norm, dot

weightednorm(dg, q, p = 2; componentwise=false) =
  weightednorm(q, Val(p), dg.MJ, componentwise)

function weightednorm(q, ::Val{p}, MJ, componentwise) where {p}
  n = map(components(q)) do f
    sum(MJ .* abs.(f) .^ p) ^ (1 // p)
  end
  componentwise ? n : norm(n, p)
end

function weightednorm(q, ::Val{Inf}, _, componentwise)
  n = map(components(q)) do f
    maximum(abs.(f))
  end
  componentwise ? n : norm(n, Inf)
end

entropyintegral(dg, q) = sum(dg.MJ .* entropy.(Ref(dg.law), q, dg.auxstate))
entropyproduct(dg, p, q) = sum(dg.MJ .* dot.(entropyvariables.(Ref(dg.law), p, dg.auxstate), q))


function ftriple!(f::Vector{Float64},x::Vector{Float64},beta::Vector{Float64},A0::Vector{Float64},
    rho::Float64,c0::Vector{Float64},W::Vector{Float64})
    # f = zeros(8);
    # conservation of mass
    f[1] = x[1] .- x[3] .- x[5] .- x[7];

    # total P
    f[2] = (beta[1].*(sqrt.(x[2]) .- sqrt.(A0[1])) + 0.5.*rho.*(x[1]/x[2]).^2) .-
        (beta[2].*(sqrt.(x[4]) .- sqrt.(A0[2])) + 0.5.*rho.*(x[3]/x[4]).^2);
    f[3] = (beta[1].*(sqrt.(x[2]) .- sqrt.(A0[1])) + 0.5.*rho.*(x[1]/x[2]).^2) .-
        (beta[3].*(sqrt.(x[6]) .- sqrt.(A0[3])) + 0.5.*rho.*(x[5]/x[6]).^2);
    f[4] = (beta[1].*(sqrt.(x[2]) .- sqrt.(A0[1])) + 0.5.*rho.*(x[1]/x[2]).^2) .-
        (beta[4].*(sqrt.(x[8]) .- sqrt.(A0[4])) + 0.5.*rho.*(x[7]/x[8]).^2);

    # Riemann invariants
    f[5] = (x[1]./x[2]) .+ 4.*(sqrt.(0.5.*beta[1]./rho).*x[2].^0.25 .- c0[1]) .- W[1];
    f[6] = (x[3]./x[4]) .- 4.*(sqrt.(0.5.*beta[2]./rho).*x[4].^0.25 .- c0[2]) .- W[2];
    f[7] = (x[5]./x[6]) .- 4.*(sqrt.(0.5.*beta[3]./rho).*x[6].^0.25 .- c0[3]) .- W[3];
    f[8] = (x[7]./x[8]) .- 4.*(sqrt.(0.5.*beta[4]./rho).*x[8].^0.25 .- c0[4]) .- W[4];

    # return f
end


# ------------------------------------------------------------------------------------------
# # Building a single neural network layer using `Flux.jl`
#
# In this notebook, we'll move beyond binary classification. We'll try to distinguish
# between three fruits now, instead of two. We'll do this using **multiple** neurons
# arranged in a **single layer**.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## Read in and process data
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# We can start by loading the necessary packages and getting our data into working order
# with similar code we used at the beginning of the previous notebooks, except that now we
# will combine three different apple data sets, and will add in some grapes to the fruit
# salad!
# ------------------------------------------------------------------------------------------

# Load apple data in with `readdlm` for each file
apples1, applecolnames1 = readdlm("data/Apple_Golden_1.dat", '\t', header = true)
apples2, applecolnames2 = readdlm("data/Apple_Golden_2.dat", '\t', header = true)
apples3, applecolnames3 = readdlm("data/Apple_Golden_3.dat", '\t', header = true)

# Check that the column names are the same for each apple file
println( applecolnames1 == applecolnames2 == applecolnames3)

# ------------------------------------------------------------------------------------------
# Since each apple file has columns with the same headers, we know we can concatenate these
# columns from the different files together:
# ------------------------------------------------------------------------------------------

apples = vcat(apples1, apples2, apples3)

# ------------------------------------------------------------------------------------------
# And now let's build an array called `x_apples` that stores data from the `red` and `blue`
# columns of `apples`. From `applecolnames1`, we can see that these are the 3rd and 5th
# columns of `apples`:
# ------------------------------------------------------------------------------------------

applecolnames1[3], applecolnames1[5]

length(apples[:, 1])

x_apples  = [ [apples[i, 3], apples[i, 5]] for i in 1:length(apples[:, 3]) ]

# ------------------------------------------------------------------------------------------
# Similarly, let's create arrays called `x_bananas` and `x_grapes`:
# ------------------------------------------------------------------------------------------

# Load data from *.dat files
bananas, bananacolnames = readdlm("data/Banana.dat", '\t', header = true)
grapes1, grapecolnames1 = readdlm("data/Grape_White.dat", '\t', header = true)
grapes2, grapecolnames2 = readdlm("data/Grape_White_2.dat", '\t', header = true)

# Concatenate data from two grape files together
grapes = vcat(grapes1, grapes2)

# Check that column 3 and column 5 refer to the "red" and "blue" columns from each file
println("All column headers are the same: ", bananacolnames == grapecolnames1 == grapecolnames2 == applecolnames1)

# Build x_bananas and x_grapes from bananas and grapes
x_bananas  = [ [bananas[i, 3], bananas[i, 5]] for i in 1:length(bananas[:, 3]) ]
x_grapes = [ [grapes[i, 3], grapes[i, 5]] for i in 1:length(grapes[:, 3]) ]

# ------------------------------------------------------------------------------------------
# ## One-hot vectors
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# Now we wish to classify *three* different types of fruit. It is not clear how to encode
# these three types using a single output variable; indeed, in general this is not possible.
#
# Instead, we have the idea of encoding $n$ output types from the classification into
# *vectors of length $n$*, called "one-hot vectors":
#
# $$
# \textrm{apple} = \begin{pmatrix} 1 \\ 0 \\ 0 \end{pmatrix};
# \quad
# \textrm{banana} = \begin{pmatrix} 0 \\ 1 \\ 0 \end{pmatrix};
# \quad
# \textrm{grape} = \begin{pmatrix} 0 \\ 0 \\ 1 \end{pmatrix}.
# $$
#
# The term "one-hot" refers to the fact that each vector has a single $1$, and is $0$
# otherwise.
#
# Effectively, the first neuron will learn whether or not (1 or 0) the data corresponds to
# an apple, the second whether or not (1 or 0) it corresponds to a banana, etc.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# `Flux.jl` provides an efficient representation for one-hot vectors, using advanced
# features of Julia so that it does not actually store these vectors, which would be a waste
# of memory; instead `Flux` just records in which position the non-zero element is. To us,
# however, it looks like all the information is being stored:
# ------------------------------------------------------------------------------------------

using Flux: onehot

onehot(1, 1:3)

# ------------------------------------------------------------------------------------------
# #### Exercise 1
#
# Make an array `labels` that gives the labels (1, 2 or 3) of each data point. Then use
# `onehot` to encode the information about the labels as a vector of `OneHotVector`s.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## Single layer in Flux
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# Let's suppose that there are two pieces of input data, as in the previous single neuron
# notebook. Then the network has 2 inputs and 3 outputs:
# ------------------------------------------------------------------------------------------

include("draw_neural_net.jl")
draw_network([2, 3])

# ------------------------------------------------------------------------------------------
# `Flux` allows us to express this again in a simple way:
# ------------------------------------------------------------------------------------------

using Flux

model = Dense(2, 3, σ)

# ------------------------------------------------------------------------------------------
# #### Exercise 2
#
# Now what do the weights inside `model` look like? How does this compare to the diagram of
# the network layer above?
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## Training the model
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# Despite the fact that the model is now more complicated than the single neuron from the
# previous notebook, the beauty of `Flux.jl` is that the rest of the training process
# **looks exactly the same**!
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 3
#
# Implement training for this model.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 4
#
# Visualize the result of the learning for each neuron. Since each neuron is sigmoidal, we
# can get a good idea of the function by just plotting a single contour level where the
# function takes the value 0.5, using the `contour` function with keyword argument
# `levels=[0.5, 0.501]`.
# ------------------------------------------------------------------------------------------

plot()

contour!(0:0.01:1, 0:0.01:1, (x,y)->model([x,y]).data[1], levels=[0.5, 0.501], color = cgrad([:blue, :blue]))
contour!(0:0.01:1, 0:0.01:1, (x,y)->model([x,y]).data[2], levels=[0.5,0.501], color = cgrad([:green, :green]))
contour!(0:0.01:1, 0:0.01:1, (x,y)->model([x,y]).data[3], levels=[0.5,0.501], color = cgrad([:red, :red]))

scatter!(first.(x_apples), last.(x_apples), m=:cross, label="apples")
scatter!(first.(x_bananas), last.(x_bananas), m=:circle, label="bananas")
scatter!(first.(x_grapes), last.(x_grapes), m=:square, label="grapes")

# ------------------------------------------------------------------------------------------
# #### Exercise 5
#
# Interpret the results by checking which fruit each neuron was supposed to learn and what
# it managed to achieve.
# ------------------------------------------------------------------------------------------

.nbexports/introductory-tutorials/broader-topics-and-ecosystem/intro-to-ml/16. Tools - Using Flux to build a single layer neural net.jl


@testset "1601.maximum-number-of-achievable-transfer-requests.jl" begin
    @test maximum_requests(
        20,
        [
            [7, 11],
            [6, 4],
            [12, 15],
            [8, 3],
            [13, 12],
            [14, 19],
            [5, 16],
            [5, 6],
            [15, 0],
            [6, 5],
            [6, 0],
            [12, 15],
            [18, 8],
            [15, 5],
            [3, 16],
            [18, 13],
        ],
    ) == 2
    @test maximum_requests(5, [[0, 1], [1, 0], [0, 1], [1, 2], [2, 0], [3, 4]]) == 5
    @test maximum_requests(3, [[0, 0], [1, 2], [2, 1]]) == 3
    @test maximum_requests(4, [[0, 3], [3, 1], [1, 2], [2, 0]]) == 4
    @test maximum_requests(
        20,
        [
            [0, 7],
            [2, 0],
            [2, 15],
            [5, 6],
            [6, 16],
            [6, 16],
            [6, 17],
            [7, 14],
            [11, 2],
            [13, 2],
            [14, 11],
            [15, 6],
            [16, 5],
            [16, 6],
            [17, 13],
            [19, 10],
        ],
    ) == 15
end