using Base.Iterators: partition
using Flux
using Flux.Optimise: update!
using Flux: logitbinarycrossentropy
using Images
using MLDatasets
using Statistics
using Parameters: @with_kw
using Printf
using Random

@with_kw struct HyperParams
    batch_size::Int = 128
    latent_dim::Int = 100
    epochs::Int = 20
    verbose_freq::Int = 1000
    output_x::Int = 6
    output_y::Int = 6
    lr_dscr::Float64 = 0.0002
    lr_gen::Float64 = 0.0002
end

function create_output_image(gen, fixed_noise, hparams)
    @eval Flux.istraining() = false
    fake_images = @. cpu(gen(fixed_noise))
    @eval Flux.istraining() = true
    image_array = permutedims(dropdims(reduce(vcat, reduce.(hcat, partition(fake_images, hparams.output_y))); dims=(3, 4)), (2, 1))
    image_array = @. Gray(image_array + 1f0) / 2f0
    return image_array
end

function Discriminator()
    return Chain(
            Conv((4, 4), 1 => 64; stride = 2, pad = 1),
            x->leakyrelu.(x, 0.2f0),
            Dropout(0.25),
            Conv((4, 4), 64 => 128; stride = 2, pad = 1),
            x->leakyrelu.(x, 0.2f0),
            Dropout(0.25),
            x->reshape(x, 7 * 7 * 128, :),
            Dense(7 * 7 * 128, 1))
end

function Generator()
    return Chain(
            Dense(hparams.latent_dim, 7 * 7 * 256),
            BatchNorm(7 * 7 * 256, relu),
            x->reshape(x, 7, 7, 256, :),
            ConvTranspose((5, 5), 256 => 128; stride = 1, pad = 2),
            BatchNorm(128, relu),
            ConvTranspose((4, 4), 128 => 64; stride = 2, pad = 1),
            BatchNorm(64, relu),
            ConvTranspose((4, 4), 64 => 1, tanh; stride = 2, pad = 1),
            )
end

# Loss functions
function discriminator_loss(real_output, fake_output)
    real_loss = mean(logitbinarycrossentropy.(real_output, 1f0))
    fake_loss = mean(logitbinarycrossentropy.(fake_output, 0f0))
    return real_loss + fake_loss
end

generator_loss(fake_output) = mean(logitbinarycrossentropy.(fake_output, 1f0))

function train_discriminator!(gen, dscr, x, opt_dscr, hparams)
    noise = randn!(similar(x, (hparams.latent_dim, hparams.batch_size)))
    fake_input = gen(noise)
    ps = Flux.params(dscr)
    # Taking gradient
    loss, back = Flux.pullback(ps) do
        discriminator_loss(dscr(x), dscr(fake_input))
    end
    grad = back(1f0)
    update!(opt_dscr, ps, grad)
    return loss
end

function train_generator!(gen, dscr, x, opt_gen, hparams)
    noise = randn!(similar(x, (hparams.latent_dim, hparams.batch_size)))
    ps = Flux.params(gen)
    # Taking gradient
    loss, back = Flux.pullback(ps) do
        generator_loss(dscr(gen(noise)))
    end
    grad = back(1f0)
    update!(opt_gen, ps, grad)
    return loss
end

function train(; kws...)
    # Model Parameters
    hparams = HyperParams(; kws...)

    # Load MNIST dataset
    images, _ = MLDatasets.MNIST.traindata(Float32)
    # Normalize to [-1, 1]
    image_tensor = reshape(@.(2f0 * images - 1f0), 28, 28, 1, :)
    # Partition into batches
    data = [image_tensor[:, :, :, r] |> gpu for r in partition(1:60000, hparams.batch_size)]

    fixed_noise = [randn(hparams.latent_dim, 1) |> gpu for _=1:hparams.output_x*hparams.output_y]

    # Discriminator
    dscr = Discriminator() |> gpu

    # Generator
    gen =  Generator() |> gpu

    # Optimizers
    opt_dscr = ADAM(hparams.lr_dscr)
    opt_gen = ADAM(hparams.lr_gen)

    # Training
    train_steps = 0
    for ep in 1:hparams.epochs
        @info "Epoch $ep"
        for x in data
            # Update discriminator and generator
            loss_dscr = train_discriminator!(gen, dscr, x, opt_dscr, hparams)
            loss_gen = train_generator!(gen, dscr, x, opt_gen, hparams)

            if train_steps % hparams.verbose_freq == 0
                @info("Train step $(train_steps), Discriminator loss = $(loss_dscr), Generator loss = $(loss_gen)")
                # Save generated fake image
                output_image = create_output_image(gen, fixed_noise, hparams)
                save(@sprintf("output/dcgan_steps_%06d.png", train_steps), output_image)
            end
            train_steps += 1
        end
    end

    output_image = create_output_image(gen, fixed_noise, hparams)
    save(@sprintf("output/dcgan_steps_%06d.png", train_steps), output_image)
end

cd(@__DIR__)
train()


using PowerSystems
using JuMP
using Ipopt
using PowerSimulations
const PS = PowerSimulations

ipopt_optimizer = JuMP.with_optimizer(Ipopt.Optimizer, tol=1e-6, print_level=0)

base_dir = dirname(dirname(pathof(PowerSystems)))
include(joinpath(base_dir,"data/data_5bus_pu.jl"))
sys5 = PowerSystem(nodes5, generators5, loads5_DA, branches5, nothing,  100.0);

@test try
    @info "testing copper plate"
    Net = PS.CopperPlatePowerModel
    m = Model(ipopt_optimizer);
    netinjection = PS.instantiate_network(Net, sys5);
    PS.constructdevice!(m, netinjection, ThermalGen, PS.ThermalDispatch, Net, sys5);
    PS.constructnetwork!(m, [(device=Line, formulation=PS.PiLine)], netinjection, Net, sys5)
    @objective(m, Min, m.obj_dict[:objective_function])
    JuMP.optimize!(m)
    isapprox(JuMP.objective_value(m), 3400, atol = 1000)
finally end

# Flow Models
@test try
    @info "testing net flow"
    Net = PS.StandardNetFlow
    m = Model(ipopt_optimizer);
    netinjection = PS.instantiate_network(Net, sys5);
    PS.constructdevice!(m, netinjection, ThermalGen, PS.ThermalDispatch, Net, sys5);
    PS.constructnetwork!(m, [(device=Line, formulation=PS.PiLine)], netinjection, Net, sys5)
    @objective(m, Min, m.obj_dict[:objective_function])
    JuMP.optimize!(m)
    isapprox(JuMP.objective_value(m), 3400, atol = 1000)
finally end

# Flow Models
@test try
    @info "testing PTDF 5-bus"
    Net = PS.StandardPTDF
    m = Model(ipopt_optimizer);
    ptdf,  A = PowerSystems.buildptdf(sys5.branches, sys5.buses)
    netinjection = PS.instantiate_network(Net, sys5);
    PS.constructdevice!(m, netinjection, ThermalGen, PS.ThermalDispatch, Net, sys5);
    PS.constructnetwork!(m, [(device=Line, formulation=PS.PiLine)], netinjection, Net, sys5, PTDF = ptdf)
    @objective(m, Min, m.obj_dict[:objective_function])
    JuMP.optimize!(m)
    isapprox(JuMP.objective_value(m), 3400, atol = 1000)
finally end

@test try
    @info "testing AngleDC-OPF 5-bus"
    Net = PS.DCAngleForm
    m = Model(ipopt_optimizer);
    netinjection = PS.instantiate_network(Net, sys5);
    PS.constructdevice!(m, netinjection, ThermalGen, PS.ThermalDispatch, Net, sys5);
    PS.constructnetwork!(m, [(device=Line, formulation=PS.PiLine)], netinjection, Net, sys5)
    @objective(m, Min, m.obj_dict[:objective_function])
    JuMP.optimize!(m)
    isapprox(JuMP.objective_value(m), 3400, atol = 1000)
finally end

@test try
    @info "testing ACP-OPF 5-bus"
    Net = PS.StandardAC
    m = Model(ipopt_optimizer);
    netinjection = PS.instantiate_network(Net, sys5);
    PS.constructdevice!(m, netinjection, ThermalGen, PS.ThermalDispatch, Net, sys5);
    PS.constructnetwork!(m, [(device=Line, formulation=PS.PiLine)], netinjection, Net, sys5)
    @objective(m, Min, m.obj_dict[:objective_function])
    JuMP.optimize!(m)
    isapprox(JuMP.objective_value(m), 3400, atol = 1000)
true finally end

@test try
    @info "testing ACP- QCWForm 5-bus"
    Net = PM.QCWRForm
    m = Model(ipopt_optimizer);
    netinjection = PS.instantiate_network(Net, sys5);
    PS.constructdevice!(m, netinjection, ThermalGen, PS.ThermalDispatch, Net, sys5);
    PS.constructnetwork!(m, [(device=Line, formulation=PS.PiLine)], netinjection, Net, sys5)
    @objective(m, Min, m.obj_dict[:objective_function])
    JuMP.optimize!(m)
    isapprox(JuMP.objective_value(m), 3400, atol = 1000)
true finally end

include(joinpath(base_dir,"data/data_14bus_pu.jl"))
sys14 = PowerSystem(nodes14, generators14, loads14, branches14, nothing,  100.0);

@test try
    @info "testing copper plate 14-bus"
    Net = PS.CopperPlatePowerModel
    m = Model(ipopt_optimizer);
    netinjection = PS.instantiate_network(Net, sys14);
    PS.constructdevice!(m, netinjection, ThermalGen, PS.ThermalDispatch, Net, sys14);
    PS.constructnetwork!(m, [(device=Line, formulation=PS.PiLine)], netinjection, Net, sys14)
    @objective(m, Min, m.obj_dict[:objective_function])
    JuMP.optimize!(m)
    isapprox(JuMP.objective_value(m), 1200, atol = 1000)
finally end

# Flow Models
@test try
    @info "testing net 14-bus"
    Net = PS.StandardNetFlow
    m = Model(ipopt_optimizer);
    netinjection = PS.instantiate_network(Net, sys14);
    PS.constructdevice!(m, netinjection, ThermalGen, PS.ThermalDispatch, Net, sys14);
    PS.constructnetwork!(m, [(device=Line, formulation=PS.PiLine)], netinjection, Net, sys14)
    @objective(m, Min, m.obj_dict[:objective_function])
    JuMP.optimize!(m)
    isapprox(JuMP.objective_value(m), 1200, atol = 1000)
finally end

# Flow Models
@test_skip try
    @info "testing PTDF 14-bus"
    Net = PS.StandardPTDF
    m = Model(ipopt_optimizer);
    ptdf,  A = PowerSystems.buildptdf(sys14.branches, sys14.buses)
    netinjection = PS.instantiate_network(Net, sys14);
    PS.constructdevice!(m, netinjection, ThermalGen, PS.ThermalDispatch, Net, sys14);
    PS.constructnetwork!(m, [(device=Line, formulation=PS.PiLine)], netinjection, Net, sys14, PTDF = ptdf)
    @objective(m, Min, m.obj_dict[:objective_function])
    JuMP.optimize!(m)
    isapprox(JuMP.objective_value(m), 1200, atol = 1000)
finally end

@test try
    @info "testing AngleDC-OPF 14-bus"
    Net = PS.DCAngleForm
    m = Model(ipopt_optimizer);
    netinjection = PS.instantiate_network(Net, sys14);
    PS.constructdevice!(m, netinjection, ThermalGen, PS.ThermalDispatch, Net, sys14);
    PS.constructnetwork!(m, [(device=Line, formulation=PS.PiLine)], netinjection, Net, sys14)
    @objective(m, Min, m.obj_dict[:objective_function])
    JuMP.optimize!(m)
    isapprox(JuMP.objective_value(m), 1200, atol = 1000)
finally end

@test try
    @info "testing ACP-OPF 14-bus"
    Net = PS.StandardAC
    m = Model(ipopt_optimizer);
    netinjection = PS.instantiate_network(Net, sys14);
    PS.constructdevice!(m, netinjection, ThermalGen, PS.ThermalDispatch, Net, sys14);
    PS.constructnetwork!(m, [(device=Line, formulation=PS.PiLine)], netinjection, Net, sys14)
    @objective(m, Min, m.obj_dict[:objective_function])
    JuMP.optimize!(m)
    isapprox(JuMP.objective_value(m), 1200, atol = 1000)
true finally end

@test try
    @info "testing ACP-QCWForm 14-bus"
    Net = PM.QCWRForm
    m = Model(ipopt_optimizer);
    netinjection = PS.instantiate_network(Net, sys14);
    PS.constructdevice!(m, netinjection, ThermalGen, PS.ThermalDispatch, Net, sys14);
    PS.constructnetwork!(m, [(device=Line, formulation=PS.PiLine)], netinjection, Net, sys14)
    @objective(m, Min, m.obj_dict[:objective_function])
    JuMP.optimize!(m)
    isapprox(JuMP.objective_value(m), 1200, atol = 1000)
true finally end


using Krylov
using LinearAlgebra
using SparseArrays
using CUDA

function square_preconditioned(n :: Int=10)
        A   = ones(n, n) + (n-1) * Matrix(I, n, n)
        b   = 10.0 * [1:n;]
        M⁻¹ = 1/n * Matrix(I, n, n)
        return A, b, M⁻¹
end

A, b, N = square_preconditioned(100)
spA = sparse(A)
spN = sparse(N)
y, status = Krylov.bicgstab(spA, copy(b), N=spN)
@show length(status.residuals)
@show y

cuspA = CUDA.CUSPARSE.CuSparseMatrixCSR(spA)
cuspN = CUDA.CUSPARSE.CuSparseMatrixCSR(spN)
cub = CuVector(b)
cuy, status = Krylov.bicgstab(cuspA, cub, N=cuspN)
@show cuy
@show length(status.residuals)


function solve(nums, use_growing=false)
    counter = zeros(Int, maximum(nums) + 1)
    grow = copy(counter)
    for num in nums
        counter[num + 1] += 1
    end
    s = 0
    while length(counter) > 1
        front, back = counter[1] + grow[1], counter[end] + grow[end]
        if front < back
            counter[2] += counter[1]
            popfirst!(counter)
            if use_growing
                popfirst!(grow)
                grow[1] = front
            end
            s += front
        else
            counter[end-1] += counter[end]
            pop!(counter)
            if use_growing
                pop!(grow)
                grow[end] = back
            end
            s += back
        end
    end
    return s
end

function main()
    nums = (x -> parse(Int, x)).(split(readlines()[1], ','))
    println(solve(nums, false))
    println(solve(nums, true))
end

main()


using Distributions

include("assigndist_normal.jl")
include("assigndist_uniform.jl")
include("assigndist_beta.jl")
include("assigndist_betaprime.jl")
include("assigndist_betabinomial.jl")
include("assigndist_binomial.jl")
include("assigndist_gamma.jl")
include("assigndist_frechet.jl")


module HeatEqCalcOld

function laplacian4!(v, u, dx)
    v[begin] = (u[end] + u[begin+1] - 2u[begin])/dx^2
    @. @views v[begin+1:end-1] = (u[begin:end-2] + u[begin+2:end] - 2u[begin+1:end-1])/dx^2
    v[end] = (u[end-1] + u[begin] - 2u[1end])/dx^2
    return
end

function heateq4(u0, dx, tmax, N=200)
    t = 0:dx:tmax
    dt = step(t)
    u = similar(u0, length(u0), length(t)+1)
    u[:, 1] = u0
    v = similar(u0)
    N = 100
    for i in 2:length(t)+1
        @. @views u[:, i] = u[:, i-1]
        for _ in 1:N
            @views laplacian4!(v, u[:, i], dx)
            @. @views u[:, i] += v*dt/N
        end
    end
    t, u
end

using StaticArrays

function calc_sol(
        n = 200,
        x = range(-π, π; length=n+1)[1:end-1],
        m = 200
    )
    dx = step(x)
    noise1, noise2 = 0.3randn(m), 0.3randn(m)
    f(x) = SVector{m}((1 .+ noise1)*sin(x) .- (1 .+ noise2)*cos(2x))
    u0 = f.(x)
    tmax = 1.0
    t, u = heateq4(u0, dx, tmax)
    (; t, x, u, m)
end

save_sol(sol) = write("sol.txt", repr(sol))

end # module


"""
Struct for MultiVariate Copula Process

		gaussianCopulaNVariateProcess=GaussianCopulaNVariateProcess(models::num1,λ::num2,p::num3,λ₊::num4,λ₋::num5) where {num1,num2,num3,num4,num5 <: Number}

Where:\n
		models  =	the processes.
		rho  = 	correlation matrix.
"""
mutable struct GaussianCopulaNVariateProcess{num3 <: Number, numtype <: Number} <: NDimensionalMonteCarloProcess{numtype}
    models::Tuple{Vararg{BaseProcess}}
    rho::Matrix{num3}
    function GaussianCopulaNVariateProcess(rho::Matrix{num3}, models::BaseProcess...) where {num3 <: Number}
        sz = size(rho)
        @assert sz[1] == sz[2]
        @assert length(models) == sz[1]
        @assert isposdef(rho)
        zero_typed = predict_output_type_zero_(models...) + zero(num3)
        return new{num3, typeof(zero_typed)}(models, rho)
    end
    function GaussianCopulaNVariateProcess(models::BaseProcess...)
        len_ = length(models)
        return GaussianCopulaNVariateProcess(Matrix{Float64}(I, len_, len_), models...)
    end
    function GaussianCopulaNVariateProcess(model1::BaseProcess, model2::BaseProcess, rho::num3) where {num3 <: Number}
        corr_matrix_ = [1 rho; rho 1]
        @assert isposdef(corr_matrix_)
        return GaussianCopulaNVariateProcess(corr_matrix_, model1, model2)
    end
end
support_type(z::GaussianCopulaNVariateProcess{num, num1}) where {num <: Number, num1 <: Number} = zero(num)

export GaussianCopulaNVariateProcess;

function simulate!(S_total, mcProcess::GaussianCopulaNVariateProcess, rfCurve::AbstractZeroRateCurve, mcBaseData::AbstractMonteCarloConfiguration, T::Number)
    Nsim = mcBaseData.Nsim
    Nstep = mcBaseData.Nstep
    @assert T > 0

    ####Simulation
    ## Simulate
    len_ = length(mcProcess.models)
    for i = 1:len_
        simulate!(S_total[i], mcProcess.models[i], rfCurve, mcBaseData, T)
    end
    rho = mcProcess.rho
    if (isdiag(rho))
        return
    end
    U_joint = Matrix{eltype(rho)}(undef, len_, Nsim)
    for j = 1:Nstep
        gausscopulagen2!(U_joint, rho, mcBaseData)

        for i = 1:len_
            @views tmp_ = S_total[i][:, j+1]
            sort!(tmp_)
            @views tmp_ .= Statistics.quantile(tmp_, U_joint[i, :]; sorted = true)
        end
    end

    ## Conclude
    return
end


function objective_min_redispatch_cost(pm::GenericPowerModel)
    PowerModels.check_cost_models(pm)

    gen_cost = Dict()
    for (n, nw_ref) in nws(pm)
        for (i,gen) in nw_ref[:gen]
            pg_delta = sum( var(pm, n, c, :pg_delta, i) for c in conductor_ids(pm, n) )
            if length(gen["cost"]) == 1
                gen_cost[(n,i)] = gen["cost"][1]
            elseif length(gen["cost"]) == 2
                gen_cost[(n,i)] = gen["cost"][1]*pg_delta + gen["cost"][2]
            elseif length(gen["cost"]) == 3
                gen_cost[(n,i)] = gen["cost"][1]*pg_delta^2 + gen["cost"][2]*pg_delta + gen["cost"][3]
            else
                gen_cost[(n,i)] = 0.0
            end
        end
    end
    return @objective(pm.model, Min,
        sum(
            sum( gen_cost[(n,i)] for (i,gen) in nw_ref[:gen] ) +
            sum(
                load["voll"]*sum( PowerModels.var(pm, n, c, :pl_delta, i) for c in PowerModels.conductor_ids(pm, n)) for (i,load) in PowerModels.ref(pm, :load))
        for (n, nw_ref) in PowerModels.nws(pm))
    )
end

function objective_total_risk(pm::GenericPowerModel, first_stage_network_id, second_stage_network_ids)
    constraint_first_stage_cost(pm::GenericPowerModel, first_stage_network_id)
    constraint_second_stage_risk(pm::GenericPowerModel, second_stage_network_ids)
    for (n, network) in pm.ref[:nw]
        constraint_dispatch_cost(pm::GenericPowerModel, n)
        constraint_redispatch_cost(pm::GenericPowerModel, n)
        constraint_loadshedding_cost(pm::GenericPowerModel, n)
    end
    first_stage_cost = PowerModels.var(pm, :first_stage_cost)
    second_stage_risk = PowerModels.var(pm, :second_stage_risk)

    return @objective(pm.model, Min, first_stage_cost + second_stage_risk)
end

function constraint_first_stage_cost(pm::GenericPowerModel, first_stage_network_id; cnd::Int = pm.ccnd)
    pg_delta = PowerModels.var(pm, first_stage_network_id, cnd, :pg_delta)
    pl_delta = PowerModels.var(pm, first_stage_network_id, cnd, :pl_delta)
    first_stage_cost = PowerModels.var(pm, :first_stage_cost)

    gen_cost_fs = Dict()
    n = first_stage_network_id
    for (i,gen) in pm.ref[:nw][n][:gen]
        pg_delta_fs = sum( var(pm, n, c, :pg_delta, i) for c in conductor_ids(pm, n) )
        if length(gen["cost"]) == 1
            gen_cost_fs[(n,i)] = 0
        elseif length(gen["cost"]) == 2
            gen_cost_fs[(n,i)] = gen["cost"][1]*pg_delta_fs + gen["cost"][2]
        elseif length(gen["cost"]) == 3
            gen_cost_fs[(n,i)] = gen["cost"][1]*pg_delta_fs^2 + gen["cost"][2]*pg_delta_fs + gen["cost"][3]
        else
            gen_cost_fs[(n,i)] = 0.0
        end
    end

    return  @constraint(pm.model, first_stage_cost ==
        sum(gen_cost_fs[(n,i)] for (i,gen) in PowerModels.ref(pm, :gen)) +
        sum(load["voll"]*pl_delta[i] for (i,load) in PowerModels.ref(pm, :load)))
end

function constraint_first_stage_cost(pm::GenericPowerModel{T}, first_stage_network_id; cnd::Int = pm.ccnd)  where T <: PowerModels.AbstractDCPForm
    pg_delta = PowerModels.var(pm, first_stage_network_id, cnd, :pg_delta)
    pl_delta = PowerModels.var(pm, first_stage_network_id, cnd, :pl_delta)
    first_stage_cost = PowerModels.var(pm, :first_stage_cost)

    gen_cost_fs = Dict()
    n = first_stage_network_id
    for (i,gen) in pm.ref[:nw][n][:gen]
        pg_delta_fs = sum( var(pm, n, c, :pg_delta, i) for c in conductor_ids(pm, n) )
        if length(gen["cost"]) == 1
            gen_cost_fs[(n,i)] = 0
        elseif length(gen["cost"]) == 2
            gen_cost_fs[(n,i)] = gen["cost"][1]*pg_delta_fs + gen["cost"][2]
        elseif length(gen["cost"]) == 3
            gen_cost_fs[(n,i)] = gen["cost"][2]*pg_delta_fs + gen["cost"][3]
        else
            gen_cost_fs[(n,i)] = 0.0
        end
    end

    return  @constraint(pm.model, first_stage_cost ==
        sum(gen_cost_fs[(n,i)] for (i,gen) in PowerModels.ref(pm, :gen)) +
        sum(load["voll"]*pl_delta[i] for (i,load) in PowerModels.ref(pm, :load)))
end

function constraint_second_stage_risk(pm::GenericPowerModel, second_stage_network_ids; cnd::Int = pm.ccnd)
    pg_delta = Dict(n => PowerModels.var(pm, n, cnd, :pg_delta) for (n, contingency_id) in second_stage_network_ids)
    pl_delta = Dict(n => PowerModels.var(pm, n, cnd, :pl_delta) for (n, contingency_id) in second_stage_network_ids)
    second_stage_risk = PowerModels.var(pm, :second_stage_risk)

    gen_cost_ss = Dict()
    for (cont_id, n)  in second_stage_network_ids
        for (i, gen) in pm.ref[:nw][n][:gen]
            pg_delta_ss = sum( var(pm, n, c, :pg_delta, i) for c in conductor_ids(pm, n) )
            if length(gen["cost"]) == 1
                gen_cost_ss[(n,i)] = 0
            elseif length(gen["cost"]) == 2
                gen_cost_ss[(n,i)] = gen["cost"][1]*pg_delta_ss + gen["cost"][2]
            elseif length(gen["cost"]) == 3
                gen_cost_ss[(n,i)] = gen["cost"][1]*pg_delta_ss^2 + gen["cost"][2]*pg_delta_ss + gen["cost"][3]
            else
                gen_cost_ss[(n,i)] = 0.0
            end
        end
    end
    return @constraint(pm.model, second_stage_risk ==
    sum(    PowerModels.ref(pm, n, :contingencies, cont_id)["prob"] *
            (
            sum(gen_cost_ss[(n,i)] for (i,gen) in PowerModels.ref(pm, n, :gen))  +
            sum(load["voll"]*pl_delta[n][i] for (i,load) in PowerModels.ref(pm, n, :load))
            )
            for (cont_id, n)  in second_stage_network_ids
    ))
end

function constraint_second_stage_risk(pm::GenericPowerModel{T}, second_stage_network_ids; cnd::Int = pm.ccnd) where T <: PowerModels.AbstractDCPForm
    pg_delta = Dict(n => PowerModels.var(pm, n, cnd, :pg_delta) for (n, contingency_id) in second_stage_network_ids)
    pl_delta = Dict(n => PowerModels.var(pm, n, cnd, :pl_delta) for (n, contingency_id) in second_stage_network_ids)
    second_stage_risk = PowerModels.var(pm, :second_stage_risk)

    gen_cost_ss = Dict()
    for (cont_id, n)  in second_stage_network_ids
        for (i, gen) in pm.ref[:nw][n][:gen]
            pg_delta_ss = sum( var(pm, n, c, :pg_delta, i) for c in conductor_ids(pm, n) )
            if length(gen["cost"]) == 1
                gen_cost_ss[(n,i)] = 0
            elseif length(gen["cost"]) == 2
                gen_cost_ss[(n,i)] = gen["cost"][1]*pg_delta_ss + gen["cost"][2]
            elseif length(gen["cost"]) == 3
                gen_cost_ss[(n,i)] = gen["cost"][2]*pg_delta_ss + gen["cost"][3]
            else
                gen_cost_ss[(n,i)] = 0.0
            end
        end
    end
    return @constraint(pm.model, second_stage_risk ==
    sum(    PowerModels.ref(pm, n, :contingencies, cont_id)["prob"] *
            (
            sum(gen_cost_ss[(n,i)] for (i,gen) in PowerModels.ref(pm, n, :gen))  +
            sum(load["voll"]*pl_delta[n][i] for (i,load) in PowerModels.ref(pm, n, :load))
            )
            for (cont_id, n)  in second_stage_network_ids
    ))
end
""
function constraint_dispatch_cost(pm::GenericPowerModel, n; cnd::Int = pm.ccnd)
    gen_cost = Dict()
    dispatch_cost = PowerModels.var(pm, :dispatch_cost)
    for (i,gen) in pm.ref[:nw][n][:gen]
        pg = sum( var(pm, n, c, :pg, i) for c in conductor_ids(pm, n) )
        if length(gen["cost"]) == 1
            gen_cost[(i)] = gen["cost"][1]
        elseif length(gen["cost"]) == 2
            gen_cost[(i)] = gen["cost"][1]*pg + gen["cost"][2]
        elseif length(gen["cost"]) == 3
            gen_cost[(i)] = gen["cost"][1]*pg^2 + gen["cost"][2]*pg + gen["cost"][3]
        else
            gen_cost[(i)] = 0.0
        end
    end
    return @constraint(pm.model, dispatch_cost[n] ==
            sum( gen_cost[(i)] for (i,gen) in pm.ref[:nw][n][:gen] )
    )
end

function constraint_dispatch_cost(pm::GenericPowerModel{T}, n; cnd::Int = pm.ccnd) where T <: PowerModels.AbstractDCPForm
    gen_cost = Dict()
    dispatch_cost = PowerModels.var(pm, :dispatch_cost)
    for (i,gen) in pm.ref[:nw][n][:gen]
        pg = sum( var(pm, n, c, :pg, i) for c in conductor_ids(pm, n) )
        if length(gen["cost"]) == 1
            gen_cost[(i)] = gen["cost"][1]
        elseif length(gen["cost"]) == 2
            gen_cost[(i)] = gen["cost"][1]*pg + gen["cost"][2]
        elseif length(gen["cost"]) == 3
            gen_cost[(i)] = gen["cost"][2]*pg + gen["cost"][3]
        else
            gen_cost[(i)] = 0.0
        end
    end
    return @constraint(pm.model, dispatch_cost[n] ==
            sum( gen_cost[(i)] for (i,gen) in pm.ref[:nw][n][:gen] )
    )
end

function constraint_redispatch_cost(pm::GenericPowerModel, n; cnd::Int = pm.ccnd)
    redispatch_cost = PowerModels.var(pm, :redispatch_cost)
    gen_cost = Dict()

    for (i,gen) in pm.ref[:nw][n][:gen]
        pg_delta = sum( var(pm, n, c, :pg_delta, i) for c in conductor_ids(pm, n) )
        if length(gen["cost"]) == 1
            gen_cost[(i)] = 0
        elseif length(gen["cost"]) == 2
            gen_cost[(i)] = gen["cost"][1]*pg_delta
        elseif length(gen["cost"]) == 3
            gen_cost[(i)] = gen["cost"][1]*pg_delta^2 + gen["cost"][2]*pg_delta
        else
            gen_cost[(i)] = 0.0
        end
    end

    return @constraint(pm.model, redispatch_cost[n] ==
            sum( gen_cost[(i)] for (i,gen) in pm.ref[:nw][n][:gen] )
    )
end

function constraint_redispatch_cost(pm::GenericPowerModel{T}, n; cnd::Int = pm.ccnd) where T <: PowerModels.AbstractDCPForm
    redispatch_cost = PowerModels.var(pm, :redispatch_cost)
    gen_cost = Dict()

    for (i,gen) in pm.ref[:nw][n][:gen]
        pg_delta = sum( var(pm, n, c, :pg_delta, i) for c in conductor_ids(pm, n) )
        if length(gen["cost"]) == 1
            gen_cost[(i)] = 0
        elseif length(gen["cost"]) == 2
            gen_cost[(i)] = gen["cost"][1]*pg_delta
        elseif length(gen["cost"]) == 3
            gen_cost[(i)] = gen["cost"][1]*pg_delta^2 + gen["cost"][2]*pg_delta
        else
            gen_cost[(i)] = 0.0
        end
    end

    return @constraint(pm.model, redispatch_cost[n] ==
            sum( gen_cost[(i)] for (i,gen) in pm.ref[:nw][n][:gen] )
    )
end

function constraint_loadshedding_cost(pm::GenericPowerModel, n; cnd::Int = pm.ccnd)
    pl_delta = PowerModels.var(pm, n, cnd, :pl_delta)
    loadshedding_cost = PowerModels.var(pm, :loadshedding_cost)

    return @constraint(pm.model, loadshedding_cost[n] ==
            sum(load["voll"]*pl_delta[i] for (i,load) in PowerModels.ref(pm, n,:load))
    )
end

# function objective_total_risk(pm::GenericPowerModel{T}, first_stage_network_id, second_stage_network_ids; cnd::Int = pm.ccnd) where T <: PowerModels.AbstractDCPForm
#
#     for (n, network) in pm.ref[:nw]
#         constraint_loadshedding_cost(pm::GenericPowerModel, n)
#     end
#     pg_delta = PowerModels.var(pm, first_stage_network_id, cnd, :pg_delta)
#     pl_delta = PowerModels.var(pm, first_stage_network_id, cnd, :pl_delta)
#
#     pg_delta_ss = Dict(n => PowerModels.var(pm, n, cnd, :pg_delta) for (n, contingency_id) in second_stage_network_ids)
#     pl_delta_ss = Dict(n => PowerModels.var(pm, n, cnd, :pl_delta) for (n, contingency_id) in second_stage_network_ids)
#
#     gen_cost_fs = Dict()
#     gen_cost_ss = Dict()
#     for n in first_stage_network_id
#         for (i,gen) in pm.ref[:nw][n][:gen]
#             pg_fs = sum( var(pm, n, c, :pg_delta, i) for c in conductor_ids(pm, n) )
#             if length(gen["cost"]) == 1
#                 gen_cost_fs[(i)] = 0
#             elseif length(gen["cost"]) == 2
#                 gen_cost_fs[(i)] = gen["cost"][1]*pg_fs
#             elseif length(gen["cost"]) == 3
#                 gen_cost_fs[(i)] = gen["cost"][2]*pg_fs
#             else
#                 gen_cost_fs[(i)] = 0.0
#             end
#         end
#     end
#     for (cont_id, n) in second_stage_network_ids
#         for (i,gen) in pm.ref[:nw][n][:gen]
#             pg_ss = sum( var(pm, n, c, :pg_delta, i) for c in conductor_ids(pm, n) )
#             if length(gen["cost"]) == 1
#                 gen_cost_ss[(n,i)] = 0
#             elseif length(gen["cost"]) == 2
#                 gen_cost_ss[(n,i)] = gen["cost"][1]*pg_ss
#             elseif length(gen["cost"]) == 3
#                 gen_cost_ss[(n,i)] = gen["cost"][2]*pg_ss
#             else
#                 gen_cost_ss[(n,i)] = 0.0
#             end
#         end
#     end
#
#     return @objective(pm.model, Min,
#             sum(gen_cost_fs[(i)] for (i,gen) in PowerModels.ref(pm, :gen)) +
#             sum(load["voll"]*pl_delta[i] for (i,load) in PowerModels.ref(pm, :load)) +
#             sum(    pm.ref[:nw][n][:contingencies][contingency_id]["prob"] *
#             (
#             sum(gen_cost_ss[(n,i)] for (i,gen) in pm.ref[:nw][n][:gen]) +
#             sum(load["voll"]*pl_delta_ss[n][i] for (i,load) in pm.ref[:nw][n][:load])
#             )
#             for (contingency_id, n) in second_stage_network_ids
#     ))
# end


using BenchmarkTools,
    Particles,
    Distributions,
    ConjugatePriors,
    StaticArrays,
    Compat,
    Random

using ConjugatePriors: NormalInverseWishart
using Compat.LinearAlgebra

Random.seed!(100)

truth = MixtureModel([MvNormal(-2 .* ones(2), Matrix(1.0I,2,2)),
                      MvNormal(2 .* ones(2), Matrix(1.0I,2,2))])
y = [rand(truth) for _ in 1:100]

prior = NormalInverseWishart(zeros(2), 0.1, Matrix(1.0I,2,2), 3.0)

function f(y, prior)
    stateprior = ChineseRestaurantProcess(1.)
    ps = ChenLiuParticles(100, prior, stateprior)
    filter!(ps, y, false)
end

f(y, prior)
@btime(f($y, $prior))

sprior = NormalInverseWishart(SVector{2}(zeros(2)), 0.1, cholesky(SMatrix{2,2}(1.0I)), 3.0)

f(y, sprior)
@btime(f($y, $sprior))

sy = SVector{2}.(y)
f(sy, sprior)
ps3 = @btime f($sy, $sprior)


#http://en.wikipedia.org/wiki/Stochastic_matrix

I = zeros(4,4);
[I[i,i] = 1 for i in 1:4];

f = open("m3.txt","r")
T = readdlm(f,',');
close(f);

Ep = [0 1 0 0]*inv(I - T)*[1,1,1,1];

println("Expected lifetime for the mouse is $Ep")


function solve()
    n = parse(Int, readline())
    ls = [parse(Int, x) for x in split(readline())]
    sort!(ls)

    ret = 0
    for i=1:n-2, j=i+1:n-1, k=j+1:n
        if ls[i] + ls[j] > ls[k] && ls[i] ≠ ls[j] && ls[j] ≠ ls[k]
            ret += 1
        end
    end
    ret
end

println(solve())


using LinearAlgebra: I

const POST_PRED_YVAL = -3.5
const POST_PRED_DIMENSION = 10
const POST_PRED_Y_OBSERVED = POST_PRED_YVAL * ones(POST_PRED_DIMENSION) / sqrt(POST_PRED_DIMENSION)
const POST_PRED_GROUND_TRUTH = pdf(
    MvNormal(POST_PRED_Y_OBSERVED / 2, I), -POST_PRED_Y_OBSERVED
)

@expectation function post_pred(y)
    x ~ MvNormal(zeros(length(y)), I)
    y ~ MvNormal(x, I)
    return pdf(MvNormal(-y, 0.5*I), x)
end


using Test: approx_full
using LinearAlgebra
using NamedDims
using NamedDims: dimnames
using Test

function baseline_tests(fact, identity)
    # A set of generic tests to ensure that our components don't accidentally reverse the
    # `:foo` and `:bar` labels for any components
    @testset "Baseline" begin
        names = (:foo, :bar)
        sz = (2, 3)
        data = [1.0 2 3; 4 5 6]
        nda = NamedDimsArray{names}(data)

        base_fact = fact(data)
        named_fact = fact(nda)

        # Check that the size is the same
        @test size(named_fact) == size(base_fact)

        # Check that the returned destructuring via iteration is the same
        @test parent.([named_fact...]) == [base_fact...]

        @test propertynames(named_fact) == propertynames(base_fact)

        # Test that all properties work as expected
        @testset "Property $P" for P in propertynames(named_fact)
            _base = getproperty(base_fact, P)
            _named = getproperty(named_fact, P)

            @test size(_base) == size(_named)

            # If our property is a NamedDimsArray make sure that the names make sense
            _named isa NamedDimsArray && @testset "Test name for dim $d" for d in 1:ndims(_named)
                # Don't think it make sense for an factorization to produce properties with
                # dimension sizes outside 1, 2 or 3
                @test size(_named, d) in (1, 2, 3)

                if size(_named, d) == 1
                    # Neither name makes sense here
                    @test dimnames(_named, d) == :_
                elseif size(_named, d) == 2
                    # Name must either be :foo or :_
                    @test dimnames(_named, d) in (:foo, :_)
                elseif size(_named, d) == 3
                    # Name must either be :bar or :_
                    @test dimnames(_named, d) in (:bar, :_)
                end
            end
        end

        @testset "Reconstructions" begin
            @test identity(base_fact) ≈ data
            @test identity(named_fact) ≈ nda
        end
    end
end

# LinearAlgebra
@testset "lu" begin
    baseline_tests(lu, F -> getindex(F.L * F.U, F.p, :))

    # Explicit `dimnames` tests for readability
    nda = NamedDimsArray{(:foo, :bar)}([1.0 2 3; 4 5 6])
    x = lu(nda)
    @test dimnames(x.L) == (:foo, :_)
    @test dimnames(x.U) == (:_, :bar)
    @test dimnames(x.p) == (:foo,)
    @test dimnames(x.P) == (:foo, :foo)

    # Opperations that should give back original dimnames
    @test dimnames(x.P * nda) == (:foo, :bar)
    @test dimnames(x.L * x.U) == (:foo, :bar)
    @test dimnames(nda[x.p, :]) == (:foo, :bar)
end

@testset "lq" begin
    baseline_tests(lq, S -> S.L * S.Q)

    # Explicit `dimnames` tests for readability
    nda = NamedDimsArray{(:foo, :bar)}([1.0 2 3; 4 5 6])
    x = lq(nda)
    @test size(x) == size(parent(x))
    @test dimnames(x.L) == (:foo, :_)
    @test dimnames(x.Q) == (:_, :bar)

    # Idenity opperations should give back original dimnames
    @test dimnames(x.L * x.Q) == (:foo, :bar)
end

@testset "svd" begin
    baseline_tests(svd, F -> F.U * Diagonal(F.S) * F.Vt)

    # Explicit `dimnames` tests for readability
    nda = NamedDimsArray{(:foo, :bar)}([1.0 2 3; 4 5 6])
    x = svd(nda)
    @test size(x) == size(parent(x))
    # Test based on visualization on wikipedia
    # https://en.wikipedia.org/wiki/File:Singular_value_decomposition_visualisation.svg
    @test dimnames(x.U) == (:foo, :_)
    @test dimnames(x.S) == (:_,)
    @test dimnames(x.V) == (:bar, :_)
    @test dimnames(x.Vt) == (:_, :bar)

    # Identity operation should give back original names
    @test dimnames(x.U * Diagonal(x.S) * x.Vt) == (:foo, :bar)
end

@testset "qr" begin
    baseline_tests(qr, F -> F.Q * F.R)
    baseline_tests(A -> qr(A, Val(true)), F -> F.Q * F.R * F.P')

    # Explicit `dimnames` tests for readability
    for pivot in (true, false)
        for data in ([1.0 2 3; 4 5 6], [big"1.0" 2; 3 4])
            nda = NamedDimsArray{(:foo, :bar)}(data)
            x = qr(nda, Val(pivot));
            @test size(x) == size(parent(x))
            @test dimnames(x.Q) == (:foo, :_)
            @test dimnames(x.R) == (:_, :bar)

            # Identity operation should give back original dimnames
            @test dimnames(x.Q * x.R) == (:foo, :bar)

            pivot && @testset "pivoted" begin
                @test parent(x) isa QRPivoted
                @test dimnames(x.p) == (:bar,)
                @test dimnames(x.P) == (:bar, :bar)

                # Identity operation should give back original dimnames
                @test dimnames(nda * x.P') == (:foo, :bar)
                @test dimnames(nda[:, x.p]) == (:foo, :bar)
                @test dimnames(x.Q * x.R * x.P') == (:foo, :bar)
            end
        end
    end
end

@testset "#164 factorization eltype not same as input eltype" begin
    # https://github.com/invenia/NamedDims.jl/issues/164
    nda = NamedDimsArray{(:foo, :bar)}([1 2 3; 4 5 6; 7 8 9])  # Int eltype
    @test qr(nda) isa NamedDims.NamedFactorization{(:foo, :bar),Float64}
end

@testset "LinearAlgebra.:ldiv " begin
    r1 = [2 3 5; 7 11 13; 17 19 23]
    r2 = r1[:, 1:2]
    b = [29, 31, 37]
    b_nda = NamedDimsArray{(:foo,)}(b)

    for A in (r1, r2)
        (m, n) = size(A)
        issquare = m == n
        fn = issquare ? (identity, triu, tril, Diagonal) : (identity,)
        for f in fn
            for B in (b_nda, b)
                nda = NamedDimsArray{(:foo, :bar)}(f(A))
                x = nda \ B
                @test parent(x) ≈ f(A) \ parent(B)
                # NOTE: Diagonal loses NamedDimness so specialcase
                f != Diagonal && @test dimnames(x) == (:bar,)
            end
        end
    end

    @test_throws DimensionMismatch (\)(
        NamedDimsArray{(:A, :B)}(r1), NamedDimsArray{(:NotA,)}(b)
    )
end


using PyPlot
using JLD
function plot_unstable_sens()
	X = load("../data/obj_erg_avg/cos4y_s1.jld")
	s1_arr = X["s1"]
	J_arr = X["J"]
	fig, ax = subplots(1,1)
	ax.plot(s1_arr, J_arr, ".", ms=10.0)
    ax.xaxis.set_tick_params(labelsize=32)
    ax.yaxis.set_tick_params(labelsize=32)
    ax.set_xlabel(L"$s_1$",fontsize=32)
    ax.set_ylabel(L"$\langle J\rangle$",fontsize=32)
	ax.grid(true)

	X = load("../data/unstable_sens/dJds1_test.jld")
	dJds = X["dJds"]
	s1 = X["s1"]

	X = load("../data/obj_erg_avg/cos4y_s1_sens.jld")
	J = X["J"]

	eps = 2e-2
	n = size(dJds)[1]
	J_pts = reshape([J .- eps*dJds J .+ eps*dJds], n,
					2)'
	s_pts = reshape([s1 .- eps s1 .+ eps], n,
					2)'

	ax.plot(s_pts[:,1:2:end-1], J_pts[:,1:2:end-1], "k",lw=6.0)
	ax.plot(s_pts[:,end], J_pts[:,end],"k",lw=6.0,label=L"$d_s\langle J\rangle$ from S3")
	ax.legend(fontsize=32)
	return dJds

end


# ---
# title: 55. Jump Game
# id: problem55
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Array, Greedy
# link: <https://leetcode.com/problems/jump-game/description/>
# hidden: true
# ---
#
# Given an array of non-negative integers, you are initially positioned at the
# first index of the array.
#
# Each element in the array represents your maximum jump length at that
# position.
#
# Determine if you are able to reach the last index.
#
#
#
# **Example 1:**
#
#
#
#     Input: nums = [2,3,1,1,4]
#     Output: true
#     Explanation: Jump 1 step from index 0 to 1, then 3 steps to the last index.
#
#
# **Example 2:**
#
#
#
#     Input: nums = [3,2,1,0,4]
#     Output: false
#     Explanation: You will always arrive at index 3 no matter what. Its maximum jump length is 0, which makes it impossible to reach the last index.
#
#
#
#
# **Constraints:**
#
#   * `1 <= nums.length <= 3 * 10^4`
#   * `0 <= nums[i][j] <= 10^5`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


# See estimate/posterior.jl for a test of the wrapper for dsgevecm_likelihood that
# takes a DSGEVECM as an input.
using DSGE, FileIO, Test, ModelConstructors
fp = dirname(@__FILE__)
@testset "Likelihood function of DSGE-VECM p(Y | θ, λ)" begin
    matdata = load(joinpath(fp, "../../reference/test_dsgevecm_likelihood.jld2"))

    # note nobs in the Matlab code is actually T in the Julia code (n_observations)
    # and lambdaT in the Matlab code is λ * T in the Julia code.
    @test matdata["lnpy"] ≈ DSGE.dsgevecm_likelihood(matdata["YYYY"],
                                                     matdata["XXYY"], matdata["XXXX"],
                                                     matdata["yyyyd"], matdata["xxyyd"],
                                                     matdata["xxxxd"], Int(matdata["nobs"]),
                                                     matdata["lambdaT"] / matdata["nobs"])
    @test matdata["lnpyinf"] ≈ DSGE.dsgevecm_likelihood(matdata["YYYY"],
                                                        matdata["XXYY"], matdata["XXXX"],
                                                        matdata["yyyyd"], matdata["xxyyd"],
                                                        matdata["xxxxd"], Int(matdata["nobs"]),
                                                        Inf)
    @test_throws DomainError matdata["lnpyinf"] ≈ DSGE.dsgevecm_likelihood(matdata["YYYY"],
                                                                           matdata["XXYY"], matdata["XXXX"],
                                                                           matdata["yyyyd"], matdata["xxyyd"],
                                                                           matdata["xxxxd"], Int(matdata["nobs"]),
                                                                           -1.)
end


using ChebyshevQuantum
using LinearAlgebra

function go(N)

    D = ChebyshevQuantum.Interp.getD(N);
    pts = ChebyshevQuantum.Interp.getCpts(N);
    v = eigvals( (D*D )[2:N,2:N]);
    v / (pi^2/4)
end


using Revise
using ADCME
using ADSeismic
using PyPlot
using DelimitedFiles
matplotlib.use("agg")

param = AcousticPropagatorParams(NX=150, NY=150, NSTEP=1000, DELTAT=1e-4,  DELTAX=1.0, DELTAY=1.0,
                                 vp_ref = 3000.0, Rcoef=0.001,
                                 USE_PML_XMIN=true, USE_PML_XMAX=true, USE_PML_YMIN=true, USE_PML_YMAX=true)

rc = Ricker(param, 15.0, 100.0, 1e10)
srci = [param.NX ÷ 2]
srcj = [param.NY ÷ 2]
srcv = reshape(rc, :, 1)
src = AcousticSource(srci, srcj, srcv)

c = 3000.0*ones(param.NX+2, param.NY+2)
model = AcousticPropagatorSolver(param, src, c)

sess = Session(); init(sess)
u = run(sess, model.u)

p = visualize_wavefield(u, param)
saveanim(p, "acoustic-wavefield.gif")


@testset "SLIC" begin
  Z = [ones(10,10) 2ones(10,10); 3ones(10,10) 4ones(10,10)]
  𝒮 = georef((Z=Z,))
  p = partition(𝒮, SLIC(4, 1.0))
  @test length(p) == 4
  @test all(nelements.(p) .== 100)
  p1, p2, p3, p4 = p
  @test mean(coordinates(centroid(p1, ind)) for ind in 1:nelements(p1)) == [5.0,5.0]
  @test mean(coordinates(centroid(p2, ind)) for ind in 1:nelements(p2)) == [15.0,5.0]
  @test mean(coordinates(centroid(p3, ind)) for ind in 1:nelements(p3)) == [5.0,15.0]
  @test mean(coordinates(centroid(p4, ind)) for ind in 1:nelements(p4)) == [15.0,15.0]

  C = cluster(𝒮, SLIC(4, 1.0))
  @test C.cluster == vec(Z')

  𝒮 = georef((z=[√(i^2+j^2) for i in 1:100, j in 1:100],))
  p = partition(𝒮, SLIC(50, 0.001))
  @test length(p) == 49

  if visualtests
    @test_reference "data/slic.png" plot(p)
  end

  # test SLIC with heterogeneous data
  Z = (a=rand(10), b=1:10, x=rand(10), y=rand(10))
  𝒮 = georef(Z, (:x, :y))
  C = cluster(𝒮, SLIC(2, 1.0))
  @test domain(C) == domain(𝒮)
  @test Set(C.cluster) ⊆ Set(1:2)

  # test SLIC for orphaned points
  a = [0,0,0,0,0,0,0,0,0,0]
  x = [0.4993029939801461, 0.14954882636793432, 0.23118957975519616, 0.6816610871344635, 0.6665309965318731, 0.691522274292691, 0.012495903053589608, 0.9831177095525963, 0.4445263730141056, 0.2175871587746574]
  y = [0.32721108209880256, 0.11427387079564899, 0.826401075107011, 0.6164294766961782, 0.6562529361193601, 0.43388375115444644, 0.7624847842129086, 0.1516623758764959, 0.07641616063237144, 0.8669098569279463]
  Z = (a=a, x=x, y=y)
  𝒮 = georef(Z, (:x, :y))
  C = cluster(𝒮, SLIC(2, 1.0))
  @test Set(C.cluster) ⊆ Set(1:2)

  # test SLIC with weights in attribute columns
  z1 = [√((i-0)^2+(j-0)^2) for i in 1:100, j in 1:100]
  z2 = [√((i-100)^2+(j-100)^2) for i in 1:100, j in 1:100]
  𝒮 = georef((z1=z1, z2=z2))
  w1 = Dict(:z1 => 10, :z2 => 0.1)
  w2 = Dict(:z1 => 0.1, :z2 => 10)
  p1 = partition(𝒮, SLIC(50, 0.001, weights=w1))
  p2 = partition(𝒮, SLIC(50, 0.001, weights=w2))
  @test length(p1) == 49
  @test length(p2) == 49

  if visualtests
    @test_reference "data/slic-w1.png" plot(p1)
    @test_reference "data/slic-w2.png" plot(p2)
  end
end


using NewtonMethods
using LinearAlgebra: Diagonal, I, Symmetric
using Roots: fzero

include("types.jl")
include("base.jl")
include("invert.jl")
include("cracks.jl")

println("Running test...")

# generate synthetic data, from "true model"
println("Make synthetic data...")

m0 = [0.1,0.05,0.,0.,0.0,deg2rad(28), deg2rad(39), deg2rad(58), deg2rad(90),deg2rad(58),deg2rad(90)]
ρ = 2.5
E₀ = 70
nu₀ = 0.25

C0 = cracks2moduli(E₀, nu₀, m0[1],m0[2],m0[3],m0[4],m0[5])

np=4
nsv=0
nsh=2

args = (np, nsv, nsh)

dobs = _gf(m0, E₀, nu₀, ρ, np, nsv, nsh)

p = [VMeasure(dobs[2k-1],exp(dobs[2k]),.01,0.01) for k in 1:np]
sv = [VMeasure(dobs[2np+2k-1],exp(dobs[2np+2k]),.01,0.01) for k in 1:nsv]
sh = [VMeasure(dobs[2np+2nsv+2k-1],exp(dobs[2np+2nsv+2k]),.01,0.01) for k in 1:nsh]

# a prior model parameters
m1 = [0.0,0.0,0.0,0.0,0.,deg2rad(25), deg2rad(40), deg2rad(55), deg2rad(90),deg2rad(90)]

C1 = cracks2moduli(E₀, nu₀, m1[1],m1[2],m1[3],m1[4],m1[5])

# run inversion
println("Run inversion to find crack parameters...")

m, anglep, anglesv, anglesh, cm = findcracks(p,sv,sh, ρ, E₀, nu₀, (m1[1], 1.0), (m1[2], 1.0), (m1[3], 1e-5), (m1[4], 1e-5), (m1[5],1e-5))

println()
println("Model parameters: input | inverted")
show(IOContext(stdout), "text/plain",[m0[1:5]  m])
println()
println()
dcalc = _gf(vcat(m, anglep, anglesv, anglesh), E₀, nu₀, ρ, np, nsv, nsh)
println("Observations: obs | calc")
show(IOContext(stdout), "text/plain",[dobs  dcalc])
println()

# run inversion moduli
println("Run inversion to find moduli...")

C, anglep, anglesv, anglesh, cm = findmoduli(p,sv,sh, ρ, C1[1,1], C1[1,3], C1[3,3], C1[5,5], C1[6,6], 100)

mc = vcat([C[1,1], C[1,3], C[3,3], C[5,5], C[6,6]], anglep, anglesv, anglesh)

println()
println("Model parameters: input | inverted")
show(IOContext(stdout), "text/plain",C0)
println("\n------")
show(IOContext(stdout), "text/plain",C)
println()
println()
dcalc = _gc(mc, ρ, np, nsv, nsh)
println("Observations: obs | calc")
show(IOContext(stdout), "text/plain",[dobs  dcalc])


# Prediction of the number of links from the number of species using the flexible links model
# and estimation of the probability of networks being stable according to their species richness

theme(:mute)
default(; frame=:box)

## Fit flexible links model using all food webs archived on Mangal

# Flexible links model and functions
@model FL(S,R) = begin
  N = length(S)
  # Number of trials
  F = S.*S .- (S.-1)
  # Parameters
  ϕ ~ Normal(3.0, 0.5)
  μ ~ Beta(3.0, 7.0)
  for i in 1:N
    R[i] ~ BetaBinomial(F[i], μ*exp(ϕ), (1-μ)*exp(ϕ))
  end
  return μ, ϕ
end

# Get the FL model
ls = DataFrame(CSV.read(joinpath("data", "mangal_foodwebs.csv"), DataFrame))
S = vec(ls[:,:S])
L = vec(ls[:,:L])
R = L .- (S.-1)
chain = sample(FL(S,R), HMC(0.01,10), 3000)

# Beta-binomial for the FL model
μ = mean(get_params(chain).μ)
ϕ = mean(get_params(chain).ϕ)

## Simulate counterfactuals

# Predict the number of links from the number of species
function prediction(chain, S)
  p = get_params(chain[200:end,:,:])
  i = rand(1:length(p.μ))
  μ, ϕ = p.μ[i], p.ϕ[i]
  return rand(BetaBinomial(S^2-(S-1), μ*exp(ϕ), (1-μ)*exp(ϕ))) + (S-1)
end

# Simulate the number of links for a range of species richness
sp = collect(5:1000) # numbers of species
nsp = length(sp)
nsim = 1000 # number of simulations

predicted_links = DataFrame()

for i in 1:nsim
  colname = "Lhat$i"
  Lhat = [prediction(chain, S) for S in sp]
  predicted_links[!, colname] = Lhat
end

## Summary statistics for the predicted numbers of links

measures = DataFrame()

# Species richnes
measures.S = sp

# Median predicted numbers of links
measures.L_median = median.(eachrow(predicted_links))

# 68% percentile intervals of predicted numbers of links
lower_links = quantile.(eachrow(predicted_links), 0.16)
upper_links = quantile.(eachrow(predicted_links), 0.84)
measures.L_PI = upper_links .- lower_links

## Estimation of the probability of the networks being stable (stability score)

# Integration over σ and get the area under the curve
function ∫(x,y)
  @assert length(x) == length(y)
  S = zero(Float64)
  for i in 2:length(x)
    S += (x[i]-x[i-1])*(y[i]+y[i-1])*0.5
  end
  return S
end

# This function is basically a wrapper around the code before
function stability_score(S, μ, ϕ)
    # This is the flexible link distribution
    D = BetaBinomial(S*S-(S-1), μ*exp(ϕ), (1-μ)*exp(ϕ))
    # We know that 0 < σ ≤ √(S/(S-1))
    Σ = LinRange(1e-5, sqrt(S/(S-1)), 100)
    # This is to store the values of maximal links before the network gets unstable
    V = zeros(Float64, length(Σ))
    for (i,σ) in enumerate(Σ)
        # This is the number of links for which the network is unstable at a given S and σ
        threshold = S/(σ*σ)-(S-1)
        # We can get the probability of the network being unstable with the CDF
        V[i] = cdf(D, threshold)
    end
    return ∫(Σ, V)
end

# Stability of networks of 5 to 1000 species
stab = zeros(Float64, length(sp))
p = Progress(length(sp))
Threads.@threads for i in 1:length(sp)
    stab[i] = stability_score(sp[i], μ, ϕ)
    next!(p)
end

measures.stab = stab

# Stability score as a function of species richness
plot(sp, stab, lab="", fill=(0, 0.1))
xaxis!(:log, extrema(sp), "Species richness")
yaxis!((0, 1), "Stability score")
savefig(joinpath("figures", "stability_score.png"))

# Write file
CSV.write(joinpath("data", "measures.csv"), measures)


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

"""
    Geometry{Dim,T}

A geometry embedded in a `Dim`-dimensional space with coordinates of type `T`.
"""
abstract type Geometry{Dim,T} end

"""
    embeddim(geometry)

Return the number of dimensions of the space where the `geometry` is embedded.
"""
embeddim(::Type{<:Geometry{Dim,T}}) where {Dim,T} = Dim
embeddim(g::Geometry) = embeddim(typeof(g))

"""
    paramdim(geometry)

Return the number of parametric dimensions of the `geometry`. For example, a
sphere embedded in 3D has 2 parametric dimension (polar and azimuthal angles).
"""
paramdim(g::Geometry) = paramdim(typeof(g))

"""
    coordtype(geometry)

Return the machine type of each coordinate used to describe the `geometry`.
"""
coordtype(::Type{<:Geometry{Dim,T}}) where {Dim,T} = T
coordtype(g::Geometry) = coordtype(typeof(g))

"""
    centroid(geometry)

Return the centroid of the `geometry`.
"""
centroid(g::Geometry) = center(g)

"""
    measure(geometry)

Return the measure of the `geometry`, i.e. the length, area, or volume.
"""
function measure end

"""
    boundary(geometry)

Return the boundary of the `geometry`.
"""
function boundary end

# ----------------
# IMPLEMENTATIONS
# ----------------
include("primitives.jl")
include("polytopes.jl")


module SpatialJackknife

using Statistics, LinearAlgebra
using NearestNeighbors, Distances

export jackknife, get_subvols

"""
This function gets dimensions of a given data array and returns an error if
they are not appropriate or consistent with the dimension of optional randoms.
"""
function get_dims(data::Array{Float64, 2},
                  randmask::Array{Float64, 2} = ones(0, 0))::Tuple{Int, Int}

    ndat, ndims = size(data)

    if ndat < ndims
        throw(ArgumentError("data array must have shape (ndat, ndims)"))
    end

    if !isempty(randmask)
        if ndims != size(randmask)[2]
            throw(ArgumentError("data array and mask array must have same number of dimensions"))
        end
    end

    ndat, ndims
end

"""
This function takes a dataset and returns an array of subvolume indices for
each datapoint. In this case, it takes a set of random mask points for
dividing arbitrary volumes. It computes the quantiles of the mask and then
finds the closest mask point to each data point to determine the subvolume
the data point belongs in.

The function expects a data array with shape (ndat, ndims) for Ndat samples in
ndims dimensions. The mask array should have shape (nmask, ndims) and then
there is an integere side_divs which specifies the number of times each
dimension is divided up to make side_divs^ndims subvolumes.
"""
function get_subvols(data::Array{Float64, 2},
                     randmask::Array{Float64, 2},
                     side_divs::Int,
                     metric::Metric = Euclidean())::Array{Int, 1}

    ndat, ndims = get_dims(data, randmask)
    nmask = size(randmask)[1]

    # now to figure out which subvolume each of the mask points is in by
    # getting the quantiles in each dimension
    maskvols = ones(Int, nmask)
    volnums = [1]  # start off with all points in the same volume
    for i in 1:ndims

        # creates a set of numbers to assign volumes for the ith dimension
        nuvolnums = [(a - 1) * side_divs^(ndims - i) for a in 1:(side_divs^i)] .+ 1

        for (j, volnum) in enumerate(volnums)
            dimmask = findall(maskvols .== volnum)
            dimquants = quantile(randmask[dimmask, i],
                                 range(0, stop = 1, length = side_divs + 1))
            dimquants[end] += 1

            for k in 2:side_divs
                submask = findall(dimquants[k] .<= randmask[dimmask, i] .< dimquants[k + 1])
                maskvols[dimmask[submask]] .= nuvolnums[side_divs * (j - 1) + k]
            end
        end

        volnums = nuvolnums
    end

    # and find the closest random to each data point to find the subvolume
    balltree = BallTree(collect(randmask'), metric)
    nninds, dists = knn(balltree, collect(data'), 1)

    subvols = maskvols[[n[1] for n in nninds]]
    subvols
end

"""
As with the other method for get_subvols, this function computes the subvolume
indices for given samples. It assumes a constant density throughout the
volume. The method computes the subvolumes based on a set of extrema for
sample values in each dimension and a number of volumes to divide the sample
on per side. If a single set of extrema are given, the assumption is that they are the same in each dimension. Alternatively, it can take a set of volume
edges in each dimension.
"""
function get_subvols(data::Array{Float64, 2};
                     side_divs::Int = 3,
                     edges::Array{Array{Float64, 1}, 1} = Array{Array{Float64, 1}, 1}(undef, 0))::Array{Int, 1}

    ndat, ndims = get_dims(data)

    if length(edges) == 1
        edges = [edges[1] for a in 1:ndims]
    elseif length(edges) != ndims
        throw(ArgumentError("edges argument must have either one array or as many arrays as dimensions in the data"))
    end

    if !all(length.(edges) .== length(edges[1]))
        throw(ArgumentError("expect the edges arrays to all have the same length since volume is a cube"))
    end

    if length(edges[1]) == 2
        edges = [collect(range(ed[1], stop = ed[2],
                               length = side_divs + 1)) for ed in edges]
    end

    # now to use the edges to assign subvolume indices
    subvols = ones(Int, ndat)
    diminds = ones(Int, ndims)
    dimfacs = [side_divs^(ndims - a) for a in 1:ndims]
    for i in 1:ndat
        for j in 1:ndims
            diminds[j] = maximum(findall(data[i, j] .>= edges[j]))
        end
        subvols[i] = sum(dimfacs .* (diminds .- 1)) + 1
    end

    subvols
end

"""
This function computes the jackknife mean and variance over the observations
computed from the data with obsfunc. They are computed by repeatedly applying
obsfunc to the data with one of the subvolumes removed. It therefore assumes
that obsfunc has normalised the observable for that remaining volume. The
function obsfunc must be callable with the form

obsfunc(data::Array{Float64, 2}, args...)

for optional args tuple and the value returned by obsfunc must be in the form
of an array of floats. If any of the returned values are NaNs, an error will be
raised. By default, the covariance matrix is computed for observables in more
than one dimension, but setting covar to false will result in only the diagonal
variances being returned.
"""
function jackknife(obsfunc::Function,
                   data::Array{Float64, 2},
                   subvolinds::Array{Int, 1},
                   args::Tuple = ();
                   covar::Bool = true,
                   pass_vol = false)::Tuple

    volset = BitSet(subvolinds)
    nvols = length(volset)
    meanfac = 1 / nvols
    varfac = (nvols - 1) / nvols

    # get the observables for the data with each volume left out
    local ivals = Array{Array{Float64, 1}, 1}(undef, nvols)
    for (i, vol) in enumerate(volset)
        if pass_vol
            argsvol = (args..., vol)
            ivals[i] = obsfunc(data[findall(subvolinds .!= vol), :],
                               argsvol...)
        else
            ivals[i] = obsfunc(data[findall(subvolinds .!= vol), :], args...)
        end
    end
    meanvals = sum(ivals) .* meanfac

    if any(isnan.(meanvals))
        throw(ErrorException("nans returned by obsfunc, perhaps the data is too sparse for the selected size/number of subvolumes"))
    end

    ndims = length(meanvals)

    # now to get the (co)variance
    covarmat = zeros(ndims, ndims)
    for i in 1:ndims
        for j in i:ndims
            if !covar && i != j
                continue
            end

            covarmat[i, j] = varfac * sum([(ivals[k][i] - meanvals[i]) *
                                           (ivals[k][j] - meanvals[j]) for k in 1:nvols])
            covarmat[j, i] = covarmat[i, j]
        end
    end

    if ndims == 1
        return meanvals, covarmat[1, 1]
    elseif !covar
        return meanvals, diag(covarmat)
    else
        return meanvals, covarmat
    end
end

end


module BOMBS

# Load dependencies
using BayesianOptimization
using BlackBoxOptim
using CSV
using Calculus
using CmdStan
using DataFrames
using Dates
using DiffEqBase
using DifferentialEquations
using Distributed
using Distributions
using GaussianProcesses
using GaussianMixtures
using JLD
using LinearAlgebra
using ODEInterfaceDiffEq
using OrdinaryDiffEq
using Plots
using Random
using Roots
using ScikitLearn.GridSearch
using ScikitLearnBase
using SharedArrays
# using Statistics
using StatsBase
using StatsPlots
using Sundials

# laod logo
include("logoBOMBS.jl")
export printLogo
export versionBOMBS

# load info function
include("InfoStructs.jl")
export infoAll

# Load section to generate the model scripts
include("ModelGen.jl")
export defModStruct
export checkStruct
export GenerateModel

# Load section to simulate the model
include("ModelSim.jl")
export simulateODEs
export defSimulStruct
export checkStructSimul
export fileStructInfo
export defSimulStructFiles
export extractSimulCSV
export plotSimsODE

# Load the section to generate pseudo-data
include("ModelPDat.jl")
export GenPseudoDat
export plotPseudoDatODE
export defPseudoDatStruct
export checkStructPseudoDat
# export fileStructInfo
export defPseudoDatStructFiles
export extractPseudoDatCSV
export PDatCSVGen

include("ModelMLE.jl")
export defMLEStruct
export SimToMle
export checkStructMLE
export selectObsSim_te
export restructInputs_te
export UVloglike
export MVloglike
export plotMLEResults
export defCrossValMLEStruct
export checkStructCrossValMLE
export plotCrossValMLEResults
export CrossValMLE
export finishMLEres
export MLEtheta

include("ModelStanInfer.jl")
export defBayInfStruct
export defBayInfDataStruct
export defBayInfDataFromFilesStruct
export defBasicStanSettingsStruct
export convertBoundTo2
export convertBoundToReal
export fitPriorSamps
export fitPriorSampsMultiNorm
export checkStructBayInf
export checkStructBayInfData
export checkStructBayInfDataFiles
export checkStructBayInfStanSettings
export genStanInitDict
export reparamDictStan
export genStanModel
export restructureDataInference
export getStanInferenceElements
export saveStanResults
export runStanInference
export plotStanResults
export StanInfer

include("ModelEntropyTheta.jl")
export genSamplesPrior
export computeH
export computeHgain

include("ModelOEDSelection.jl")
export defODEModelSelectStruct
export checkStructOEDMS
export BhattacharyyaDist
export EuclideanDist
export genOptimMSFuncts
export plotOEDMSResults
export settingsBayesOpt
export mainOEDMS

include("ModelOEDCalibration.jl")
export defODEModelCalibrStruct
export checkStructOEDMC
export genOptimMCFuncts
export plotOEDMCResults
export settingsBayesOptMC
export mainOEDMC

end


import StaticArrays: SVector

const Vertex = SVector{3, Float32}
const Normal = SVector{3, Float32}

struct Triangle
    normal::Normal
    a::Vertex
    b::Vertex
    c::Vertex
end

struct Plane
    normal::Normal
    point::Vertex
end

struct AABB
    xmin::Float32
    xmax::Float32
    ymin::Float32
    ymax::Float32
    zmin::Float32
    zmax::Float32
end

struct LineSegment
    a::Vertex
    b::Vertex
end

const Polygon = Vector{Vertex}


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

"""
    Box(min, max)

An axis-aligned box with `min` and `max` corners.
See https://en.wikipedia.org/wiki/Hyperrectangle.

## Example

```julia
Box(Point(0,0,0), Point(1,1,1)) # unit cube
```
"""
struct Box{Dim,T} <: Primitive{Dim,T}
  min::Point{Dim,T}
  max::Point{Dim,T}
end

Box(min::Tuple, max::Tuple) = Box(Point(min), Point(max))

paramdim(::Type{<:Box{Dim}}) where {Dim} = Dim

Base.minimum(b::Box) = b.min
Base.maximum(b::Box) = b.max
Base.extrema(b::Box) = b.min, b.max
center(b::Box) = Point((coordinates(b.max) + coordinates(b.min)) / 2)
measure(b::Box) = prod(b.max - b.min)
diagonal(b::Box) = norm(b.max - b.min)
sides(b::Box) = b.max - b.min

function Base.in(p::Point{Dim}, b::Box{Dim}) where {Dim}
  l, u = coordinates.((b.min, b.max))
  x = coordinates(p)
  for i in 1:Dim
    l[i] ≤ x[i] && x[i] ≤ u[i] || return false
  end
  true
end


using RandomMatrixDistributions
using Distributions
using Test

@testset "MatrixSamplers" begin

    n = 1000
    p = 400
    γ = p/n

    # test white Wishart
    for beta in [1, 4, 8]
        #println("Testing white Wishart with β = ", beta)

        λs = randeigvals(SpikedWishart(beta, n, p))

        @test length(λs) == p
        @test eltype(λs) <: Real

        @test minimum(λs) > 0

        #println("Testing Jacobi with β = ", beta)
        λs = randeigvals(Jacobi(beta, n, n, p))

        @test length(λs) == p
        @test eltype(λs) <: Real

        @test minimum(λs) > 0
    end

    #println("Testing Spiked Wishart with β = 1")
    s = 8
    W = SpikedWishart(1, n, p, [0.1, 2, 4, s])
    λs = randeigvals(W)/n

    @test length(λs) == p
    @test eltype(λs) <: Real

    # Check if the maximum spike is reasonable
    # This tests can fail in principle, but only with negligible probability
    #μ = (1 + s) * (1 + γ/s)
    #σ = (1 + s) * sqrt(γ * 2 * (1 - γ/s^2))

    spikedist = covspikedist(W)
    μ = mean(spikedist)[end]
    σ = sqrt(var(spikedist)[end])

    @test abs(maximum(λs) - μ)/σ < 4

    @test all(randeigstat(W, eigmax, 10) .> 0)

end


# Multithreaded multiplication using Polyester.jl @batch
using SparseMatricesCSR: nzrange

# sparse mat-vec multiplication
"""
    bmul!(y::AbstractVector, A::SparseMatrixCSR, x::AbstractVector, alpha::Number, beta::Number)
    bmul!(y::AbstractVector, A::SparseMatrixCSR, x::AbstractVector)

Evaluates `y = alpha*A*x + beta*y` (`y = A*x`)

In-place multithreaded version of sparse csr matrix - vector multiplication, using the threading provided by Polyester.jl
"""
function bmul!(y::AbstractVector, A::SparseMatrixCSR, x::AbstractVector, alpha::Number, beta::Number)

    A.n == size(x, 1) || throw(DimensionMismatch())
    A.m == size(y, 1) || throw(DimensionMismatch())

    o = getoffset(A)

    @batch minbatch = size(y, 1) ÷ matmul_num_threads[] for row in 1:size(y, 1)
        @inbounds begin
            accu = zero(eltype(y))
            for nz in nzrange(A, row)
                col = A.colval[nz] + o
                accu += A.nzval[nz]*x[col]
            end
            y[row] = alpha*accu + beta*y[row]

        end
    end

    return y

end

function bmul!(y::AbstractVector, A::SparseMatrixCSR, x::AbstractVector)

    return bmul!(y, A, x, true, false)

end

"""
    bmul(A::SparseMatrixCSR, x::AbstractVector)

Evaluates `A*x`.

Multithreaded version of sparse csr matrix - vector multiplication, using the threading provided by Polyester.jl
"""
function bmul(A::SparseMatrixCSR, x::AbstractVector)

    y = similar(x, promote_type(eltype(A), eltype(x)), size(A, 1))

    return bmul!(y, A, x, true, false)

end


function raindrops(number)
    # has 3 as a factor, add 'Pling' to the result.
    # has 5 as a factor, add 'Plang' to the result.
    # has 7 as a factor, add 'Plong' to the result.
    # does not have any of 3, 5, or 7 as a factor, the result should be the digits of the number.
    result = ""
    for (factor, string) in [(3, "Pling"), (5, "Plang"), (7, "Plong")]
        if number % factor == 0
            result *= string
        end
    end
    if result == ""
        return "$number"
    else
        return result
    end
end


using MultiModalMuSig
using Test

K = 2
α = 0.1
η = 0.1
X = Matrix{Int}[
    [
        1 5;
        2 8
    ],
    [
        1 2;
        2 5
    ]
]

@testset "constructor" begin
    model = MultiModalMuSig.LDA(K, α, η, X)
    @test model.K == 2
    @test model.D == 2
    @test model.N == [13, 7]
    @test model.V == 2

    @test size(model.λ) == (2, 2)
    @test all(model.λ .> 0) == true

    @test size(model.γ) == (2, 2)
    @test all(model.γ .> 0) == true

    @test sum(model.ϕ[1], dims=1) ≈ ones(2)'

    model = MultiModalMuSig.LDA(K, α, η, 3, X)
    @test model.V == 3
    @test size(model.λ) == (3, 2)
end

@testset "update_ϕ!" begin
    Elnθ = [0.5  -1.1; 2.3 -0.7]
    Elnβ = [-0.2 -0.9; -1.1 0.3]

    ϕ = Array{Float64}(undef, 2, 2)
    # K1 W1
    ϕ[1, 1] = exp(Elnθ[1, 1] + Elnβ[1, 1])
    # K1 W2
    ϕ[1, 2] = exp(Elnθ[1, 1] + Elnβ[2, 1])
    # K2 W1
    ϕ[2, 1] = exp(Elnθ[2, 1] + Elnβ[1, 2])
    # K2 W2
    ϕ[2, 2] = exp(Elnθ[2, 1] + Elnβ[2, 2])

    ϕ[:, 1] ./= sum(ϕ[:, 1])
    ϕ[:, 2] ./= sum(ϕ[:, 2])

    model = MultiModalMuSig.LDA(K, α, η, X)
    model.Elnθ .= Elnθ
    model.Elnβ .= Elnβ
    MultiModalMuSig.update_ϕ!(model)

    @test model.ϕ[1] ≈ ϕ
end

@testset "update_γ!" begin
    ϕ = [0.4 0.2; 0.6 0.8]

    γ = Array{Float64}(undef, 2)
    γ[1] = α + (ϕ[1, 1] * X[1][1, 2]) + (ϕ[1, 2] * X[1][2, 2])
    γ[2] = α + (ϕ[2, 1] * X[1][1, 2]) + (ϕ[2, 2] * X[1][2, 2])

    Elnθ = Array{Float64}(undef, 2)
    Elnθ[1] = digamma(γ[1]) - digamma(γ[1] + γ[2])
    Elnθ[2] = digamma(γ[2]) - digamma(γ[1] + γ[2])

    model = MultiModalMuSig.LDA(K, α, η, X)
    model.ϕ[1] = ϕ
    MultiModalMuSig.update_γ!(model)

    @test model.γ[:, 1] ≈ γ
    @test model.Elnθ[:, 1] ≈ Elnθ
end

@testset "update_λ!" begin
    ϕ = [[0.4 0.2; 0.6 0.8], [0.1 0.6; 0.9 0.4]]

    λ = Array{Float64}(undef, 2, 2)
    λ[1, 1] = η + (ϕ[1][1, 1] * X[1][1, 2]) + (ϕ[2][1, 1] * X[2][1, 2])
    λ[2, 1] = η + (ϕ[1][1, 2] * X[1][2, 2]) + (ϕ[2][1, 2] * X[2][2, 2])
    λ[1, 2] = η + (ϕ[1][2, 1] * X[1][1, 2]) + (ϕ[2][2, 1] * X[2][1, 2])
    λ[2, 2] = η + (ϕ[1][2, 2] * X[1][2, 2]) + (ϕ[2][2, 2] * X[2][2, 2])

    Elnβ = Array{Float64}(undef, 2, 2)
    Elnβ[1, 1] = digamma(λ[1, 1]) - digamma(λ[1, 1] + λ[2, 1])
    Elnβ[2, 1] = digamma(λ[2, 1]) - digamma(λ[1, 1] + λ[2, 1])
    Elnβ[1, 2] = digamma(λ[1, 2]) - digamma(λ[1, 2] + λ[2, 2])
    Elnβ[2, 2] = digamma(λ[2, 2]) - digamma(λ[1, 2] + λ[2, 2])

    model = MultiModalMuSig.LDA(K, α, η, X)
    model.ϕ = ϕ
    MultiModalMuSig.update_λ!(model)

    @test model.λ ≈ λ
    @test model.Elnβ ≈ Elnβ
end

@testset "calculate_elbo" begin
    model = MultiModalMuSig.LDA(K, α, η, X)

    @test typeof(MultiModalMuSig.calculate_ElnPβ(model)) == Float64
    @test typeof(MultiModalMuSig.calculate_ElnPθ(model)) == Float64
    @test typeof(MultiModalMuSig.calculate_ElnPZ(model)) == Float64
    @test typeof(MultiModalMuSig.calculate_ElnPX(model)) == Float64

    @test typeof(MultiModalMuSig.calculate_ElnQβ(model)) == Float64
    @test typeof(MultiModalMuSig.calculate_ElnQθ(model)) == Float64
    @test typeof(MultiModalMuSig.calculate_ElnQZ(model)) == Float64

    @test MultiModalMuSig.calculate_elbo(model) < 0.0
end


using PyCall
using DataFrames
using SMACOF
using Plots
using Statistics
skl = pyimport("sklearn.manifold")

max_iter = 200
function skl_smacof(Δ; anchors = nothing)
    Yhat = skl.smacof(Δ, n_init = 1, eps = eps(), max_iter = max_iter)[1]
    return align(Yhat, anchors)
end

function main(nlst = [10,20, 30, 50, 100, 200, 300, 500], nrepeats = 10)
    first_run = true

    df = DataFrame(n = Int[], method = Symbol[], mse = Float64[], time = Float64[])
    for n in nlst
        for i in 1:nrepeats
            X = 2 * rand(n, 2)

            Δ = dist(X)
            t = time()
            Y_jl = smacof(Δ, anchors = X, ε = eps(), itmax = max_iter, Xinit = rand(n, 2))
            t_jl = time() - t
            # push!(df, [n, :julia, mse(Y, X), time() - t])

            t = time()
            Y_sk = skl_smacof(Δ, anchors = X)
            t_sk = time() - t
            # push!(df, [n, :sklearn, mse(Yhat, X), time() - t])

            if first_run
                first_run = false
                continue
            end
            push!(df, [n, :julia, mse(Y_jl, X), t_jl])
            push!(df, [n, :sklearn, mse(Y_sk, X), t_sk])
        end
    end

    df_mean_times = combine(groupby(df, [:method, :n]), :time => mean => :mean)
    sort!(df_mean_times, :n)

    df_julia_times = df[df[!, :method] .== :julia, [:n, :time]]
    plot(legend = :outertopleft)
    scatter!(df_julia_times[!, :n], df_julia_times[!, :time], label = "julia", color = 1)
    plot!(df_mean_times[df_mean_times[!, :method] .== :julia, :n], df_mean_times[df_mean_times[!, :method] .== :julia, :mean], color = 1, label = nothing)

    df_sklearn_times = df[df[!, :method] .== :sklearn, [:n, :time]]
    scatter!(df_sklearn_times[!, :n], df_sklearn_times[!, :time], label = "sklearn", color = 2)
    plot!(df_mean_times[df_mean_times[!, :method] .== :sklearn, :n], df_mean_times[df_mean_times[!, :method] .== :sklearn, :mean], color = 2, label = nothing)

    title!("First to $max_iter iterations")
    xlabel!("sample size (n)")
    ylabel!("time (sec)")
    savefig(joinpath(dirname(@__FILE__), "results.png"))

end

main()


using jInv.InverseSolve
using jInv.Utils

n = 200
x = rand(n) + im*rand(n)
A = sprand(n,2*n,0.1)

Hx = A[1:2:end,1:2:end]*real(x) +
     A[1:2:end,2:2:end]*imag(x) +
     im*A[2:2:end,1:2:end]*real(x) +
     im*A[2:2:end,2:2:end]*imag(x)

@test all(abs.(HessMatVec(A,x)-Hx).<1e-12)


function getgenericHamiltoniantest( Id,
    Ix,
    Iy_no_im,
    Iz,
    ω0::Vector{T},
    J_vals::Vector{T},
    J_inds::Vector{Tuple{Int,Int}}) where T <: Real
#
    N = length(ω0)
    N_couplings = length(J_vals)
    @assert N_couplings == length(J_inds)

    # pair-wise terms.
    H1 = zeros(T, 2^N, 2^N)
    for i = 1:N_couplings

        j = J_inds[i][1]
        k = J_inds[i][2]

        H1 += getmultiIjIk(Id, Ix, Iy_no_im, Iz, j, k, N, 2*π*J_vals[i])
    end

    # single terms.
    ωIzs = collect( ω0[j] .* Iz for j = 1:N )
    H2 = Kronecker.kroneckersum(ωIzs...)

    return H1+H2, H1, H2
end

function getgenericHamiltoniantest( Id,
    Ix,
    Iy_no_im,
    Iz,
    ω0::Vector{T},
    J::Matrix{T}) where T <: Real

    N = length(ω0)
    @assert size(J,1) == size(J,2) == N

    # pair-wise terms.
    H1 = zeros(T, 2^N, 2^N)
    for j = 1:N
        for k = j+1:N
            #H += (2*π*J[j,k]) .* getmultiIjIk(Id, Ix, Iy_no_im, Iz, j, k, N)
            H1 += getmultiIjIk(Id, Ix, Iy_no_im, Iz, j, k, N, 2*π*J[j,k])

        end
    end

    # single terms.
    ωIzs = collect( ω0[j] .* Iz for j = 1:N )
    H2 = Kronecker.kroneckersum(ωIzs...)

    return H1+H2, H1, H2
end

# Only the upper-right off-diagonals of J is used.
# No secular dipole-dipole.
function getgenericHamiltonian( Id,
    Ix,
    Iy_no_im,
    Iz,
    ω0::Vector{T},
    J::Matrix{T}) where T <: Real

    N = length(ω0)
    @assert size(J,1) == size(J,2) == N

    # pair-wise terms.
    H = zeros(T, 2^N, 2^N)
    for j = 1:N
        for k = j+1:N
            #H += (2*π*J[j,k]) .* getmultiIjIk(Id, Ix, Iy_no_im, Iz, j, k, N)
            H += getmultiIjIk(Id, Ix, Iy_no_im, Iz, j, k, N, 2*π*J[j,k])

        end
    end

    # single terms.
    ωIzs = collect( ω0[j] .* Iz for j = 1:N )
    H += Kronecker.kroneckersum(ωIzs...)

    return H
end

# No secular dipole-dipole.
function getgenericHamiltonian0( Id,
    Ix,
    Iy_no_im,
    Iz,
    ω0::Vector{T},
    J::Matrix{T}) where T <: Real

    N = length(ω0)
    @assert size(J,1) == size(J,2) == N

    H = zeros(T, 2^N, 2^N)

    for j = 1:N
        H += ω0[j] .* getmultiIz(Id, j, N)

        for k = j+1:N
            H += (2*π*J[j,k]) .* getmultiIjIk0(Id, Ix, Iy_no_im, Iz, j, k, N)
        end
    end

    return H
end

# No secular dipole-dipole.
function getgenericHamiltonian( Id,
    Ix,
    Iy_no_im,
    Iz,
    ω0::Vector{T},
    J_vals::Vector{T},
    J_inds::Vector{Tuple{Int,Int}}) where T <: Real

    N = length(ω0)
    N_couplings = length(J_vals)
    @assert N_couplings == length(J_inds)

    # pair-wise terms.
    H = zeros(T, 2^N, 2^N)
    for i = 1:N_couplings

        j = J_inds[i][1]
        k = J_inds[i][2]

        H += getmultiIjIk(Id, Ix, Iy_no_im, Iz, j, k, N, 2*π*J_vals[i])
    end

    # single terms.
    ωIzs = collect( ω0[j] .* Iz for j = 1:N )
    H += Kronecker.kroneckersum(ωIzs...)

    return H
end

# has secular dipole-dipole. J here is not really J-coupling, but J - d.
function getgenericHamiltonian( Id,
    Ix,
    Iy_no_im,
    Iz,
    ω0::Vector{T},
    J_vals::Vector{T},
    J_inds::Vector{Tuple{Int,Int}},
    Di_vals::Vector{T},
    Di_inds::Vector{Tuple{Int,Int}}) where T <: Real

    N = length(ω0)
    N_couplings = length(Di_vals)
    @assert N_couplings == length(Di_inds)

    # Allocate base operators.
    z_operators = Vector{Matrix{T}}(undef, N)
    fill!(z_operators, Id)

    # pair-wise terms.
    H = getgenericHamiltonian(Id, Ix, Iy_no_im, Iz, ω0, J_vals, J_inds)
    for i = 1:N_couplings

        j = Di_inds[i][1]
        k = Di_inds[i][2]

        ### z.
        z_operators[j] = (3*Di_vals[i]) .* Iz
        z_operators[k] = Iz

        # get product operator.
        H_di = Kronecker.kronecker(z_operators...)

        H += H_di
    end

    return H
end

# v must be an eigen vector of Iz. This routine does not verify this condition.
function getzangularmomentum(v::Vector{T}, Iz; tol = 1e-7)::T where T

    b = Iz*v
    m = dot(b,v)/dot(v,v)

    ## This was meant to be a sanity-check, but when N > 4, the tol becomes larger than 1e-7.
    # The exclusion of these states mean there are insufficient number of splits,
    #   thus, the area ratio does not hold, unless weak coupling is "really satisifed"?
    # Solution: just disable this check, and assume v must be an eigenvector of Iz.
    # if norm(m .* v - b) > tol
    #     return -Inf
    # end

    return m
end

function getorderofcoherence(Iz, Q::Vector{Vector{T}}) where T

    M = length(Q)
    @assert size(Iz, 1) == M == size(Iz, 2)

    M_array = collect( getzangularmomentum(Q[r], Iz) for r = 1:length(Q))

    p = Matrix{T}(undef, M, M)
    fill!(p, NaN)
    for r = 1:M
        #Mr = getzangularmomentum(Q[r], Iz)
        Mr = M_array[r]

        for s = 1:M
            #Ms = getzangularmomentum(Q[s], Iz)
            Ms = M_array[s]

            p[r,s] = Mr-Ms
            #p[r,s] = round(Mr-Ms)
        end
    end

    return p, M_array
end

function normalizetoNspins!(A::Vector{T}, N_spins::Int) where T

    factor = N_spins/sum(A)
    for l = 1:length(A)
        A[l] = A[l] * factor
    end

    return nothing
end

function getaΩ( Iz_full,
    H::Matrix{T},
    Iys_no_im_full,
    Ix_full;
    tol = 1e-7) where T <: Real

    M = size(H,1)
    @assert size(H,2) == M == size(Iz_full,1) == size(Iz_full,2)

    ### eigen states.
    λ, Q_mat = eigen(H)
    # λ[6] = -λ[6]
    # Q_mat[:,6] = -Q_mat[:,6]
    Q = collect( Q_mat[:,m] for m = 1:M )

    ### pre-compute.
    IyQs = collect( Iys_no_im_full*Q[s] for s = 1:size(Q,1)) # no im since official formula for α is div by 2im.
    IxQs = collect( Ix_full*Q[r] for r = 1:size(Q,1))

    ### get order of coherence.
    #@time p, M_array = getorderofcoherence(Iz_full, Q) # TODO speed this up.
    #println("timing coherence")
    #@time p, M_array = getorderofcoherence(Iz_full, Q)
    p, M_array = getorderofcoherence(Iz_full, Q)

    a = Vector{T}(undef, M*M)
    F = Vector{T}(undef, M*M)
    inds = Vector{Tuple{Int,Int}}(undef, M*M)

    i_out = 0

    # for r = 1:M
    #     for s = 1:M
    for (r,s) in Base.Iterators.product(1:M,1:M)

        #if Utilities.isnumericallyclose(p[r,s], -one(T), tol)
        if isnumericallyclose(round(p[r,s]), -one(T), tol)
            #println("i_out = ", i_out)
            i_out += 1

            inds[i_out] = (r,s)

            ## signal amplitude. Section A.8.2 Spin Dynamics, Eq 2.1.131 Ernst.
            t1 = -dot(conj(Q[r]), IyQs[s])
            t2 = dot(conj(Q[s]), IxQs[r])

            # factor 2 instead of 2*im since Iys_no_im_full
            #   is divided by 2 instead of 2*im.
            a[i_out] = 2*(t1 * t2)

            ## signal frequency.
            F[i_out] = λ[s] - λ[r]
        end
    end

    resize!(inds, i_out)
    resize!(a, i_out)
    resize!(F, i_out)

    return a, F, p, λ, Q, inds, M_array
end

function getcfreq(z::Vector{T}) where T
    M = length(z)

    F = Matrix{T}(undef, M, M)
    for r = 1:M
        for s = 1:M
            F[r,s] = z[s] - z[r]
        end
    end

    return F
end

# Q is operator. Evalates <r| Q |s>.
function evaltrace( Q::Kronecker.KroneckerSum{T}, r, s)::T where T
    return dot(conj(r),Q*s)
end

function evaltrace( Q::Kronecker.KroneckerSum{T},
    V)::Matrix{T} where T

    M = length(V)

    out = Matrix{T}(undef, M, M)
    for r = 1:M
    for s = 1:M
    out[r,s] = evaltrace(Q, V[r], V[s])
    end
    end

    return out
end

# tp_basis[r,:] is the elementary state vector for each of the spins.
# +1 for state 1/2, -1 for state -1/2.
function gettensorproductbasis(Iz::Matrix{T}, Id, N_spins::Int) where T
    tp_basis = Matrix{Int}(undef, 2^N_spins, N_spins)

    for n = 1:N_spins

    Is = collect( Id for j = 1:N_spins )
    Is[n] = Iz
    I_full = Kronecker.kronecker(Is...)

        for i = 1:size(I_full,1)
            m = I_full[i,i]

            if m > 0
                tp_basis[i,n] = 1
            else
                tp_basis[i,n] = -1
            end
        end
    end

    return tp_basis
end

# for single (π/2)_x pulse. # equation A.32.
function evalamplitudeA32( Iys_no_im_full,
    Ix_full,
    V::Vector{Vector{T}})::Matrix{T} where T

    M = length(V)

    out = Matrix{T}(undef, M, M)
    for r = 1:M
    for s = 1:M

    t1 = -dot(conj(V[r]), Iys_no_im_full*V[s])
    t2 = dot(conj(V[s]), Ix_full*V[r])

    # factor 2 instead of 2*im since Iys_no_im_full
    #   is divided by 2 instead of 2*im.
    out[r,s] = 2*(t1 * t2)
    end
    end

    return out
end

# no phase. common peak width λ.
function evalabsorptionspectrum(x::T, λ::T, a::Array{T,D}, Ω::Array{T,D}) where {T,D}
    L = length(a)
    @assert length(Ω) == L

    out = zero(T)
    for l = 1:L
    denominator = λ^2 + (x-Ω[l])^2
    out += a[l]*λ/denominator
    end

    return out
end

function evalabsorptionspectrumarea(lower_limit::T,
    upper_limit::T,
    λ::T,
    a::Array{T,D},
    Ω::Array{T,D}) where {T,D}

    L = length(a)
    @assert length(Ω) == L

    out = zero(T)
    for l = 1:L

    numerator = Ω[l] - lower_limit
    Ia = -atan(numerator,λ)

    numerator = Ω[l] - upper_limit
    Ib = -atan(numerator,λ)

    out += a[l]*(Ib-Ia)
    end

    return out
end


function corners(box::IntervalBox{2,T}) where {T}
    xL = [box[1].lo, box[2].lo]
    xR = [box[1].hi, box[2].hi]
    return xL, xR
end

function split_box(box, numsplits)
    return mince(box, numsplits)
end

function min_diam(box)
    return minimum(diam.(box))
end

function interval_arithmetic_sign_search(
    func,
    initialbox,
    tol,
    perturbation,
    numsplits,
)

    rtol = tol * diam(initialbox)

    foundpos = foundneg = breachedtol = false
    queue = [initialbox]

    while !isempty(queue)
        if (foundpos && foundneg)
            break
        else
            box = popfirst!(queue)
            if min_diam(box) < rtol
                breachedtol = true
                break
            end
            funcrange = func(box)
            if inf(funcrange) > -perturbation
                foundpos = true
            elseif sup(funcrange) < perturbation
                foundneg = true
            else
                newboxes = split_box(box, numsplits)
                push!(queue, newboxes...)
            end
        end
    end

    if breachedtol
        return 2
    elseif foundpos && foundneg
        return 0
    elseif foundpos && !foundneg
        return +1
    elseif !foundpos && foundneg
        return -1
    else
        error("Unexpected scenario")
    end
end

function Base.sign(func, box; tol = 1e-3, perturbation = 0.0, numsplits = 2)
    return interval_arithmetic_sign_search(
        func,
        box,
        tol,
        perturbation,
        numsplits,
    )
end
"""
    sign(f, box)
return
- `+1` if `f` is uniformly positive on `int`
- `-1` if `f` is uniformly negative on `int`
- `0` if `f` has at least one zero crossing in `int` (f assumed continuous)
"""
function Base.sign(func, xL, xR; tol = 1e-3, perturbation = 0.0, numsplits = 2)
    box = IntervalBox(xL, xR)
    return sign(
        func,
        box;
        tol = tol,
        perturbation = perturbation,
        numsplits = numsplits,
    )
end

function extremal_coeffs_in_box(poly, xL, xR)

    max_coeff = -Inf
    min_coeff = Inf
    points = interpolation_points(basis(poly))
    dim, npoints = size(points)
    for i = 1:npoints
        p = view(points, :, i)
        if all(xL .<= p .<= xR)
            coeff = poly.coeffs[i]
            min_coeff = min(min_coeff, coeff)
            max_coeff = max(max_coeff, coeff)
        end
    end
    return max_coeff, min_coeff
end

function Base.sign(
    P::InterpolatingPolynomial{1,B},
    xL,
    xR;
    tol = 1e-3,
    perturbation = 0.0,
    numsplits = 2,
) where {B<:LagrangeTensorProductBasis}

    max_coeff, min_coeff = extremal_coeffs_in_box(P, xL, xR)
    if max_coeff > 0 && min_coeff < 0
        return 0
    else
        box = IntervalBox(xL, xR)
        return interval_arithmetic_sign_search(
            P,
            box,
            tol,
            perturbation,
            numsplits,
        )
    end
end

function (IP::InterpolatingPolynomial{1,B,T})(
    box::IntervalBox{2,S},
) where {B,T,S}
    @assert PolynomialBasis.dimension(IP) == 2
    return IP(box[1], box[2])
end

function PolynomialBasis.gradient(
    IP::InterpolatingPolynomial{1,B,T},
    box::IntervalBox{2,S},
) where {B,T,S}
    @assert PolynomialBasis.dimension(IP) == 2
    return gradient(IP, box[1], box[2])
end

function (IP::InterpolatingPolynomial{1,B,T})(
    box::IntervalBox{1,S},
) where {B<:PolynomialBasis.AbstractBasis{1},T,S}
    @assert PolynomialBasis.dimension(IP) == 1
    return IP(box[1])
end

function PolynomialBasis.gradient(
    IP::InterpolatingPolynomial{1,B,T},
    box::IntervalBox{1,S},
) where {B,T,S}
    @assert PolynomialBasis.dimension(IP) == 1
    return gradient(IP, box[1])
end


# If want to explore ternary plots and line fitting:

using JLD2, DataFrames, CSV, Dates, Query
using LinearAlgebra, GLM
using RCall
using Gadfly, Colors

include("/home/kristen/Documents/coda_utilities/coda_ops_transforms.jl")

#load in the data:
@load "/home/kristen/Documents/V6V8_analysis/analysis_products/oligo_df_zero_plus.jld2"
#should return dataframe oligo_df_plus

## Ternary Plots, courtesy of the lovely Ternary package in R:

R"require(Ternary)"
R"require(png)"

x136=closure(convert(Array,oligo_df_plus[!,[:Oligo1, :Oligo3,:Oligo6]])) #convert Dataframe columns to array
R"""
TernaryPlot()
TernaryPlot(atip = "O1-IC", btip = "O3-II/XV", ctip = "O6-I?", lab.col=c("#0000FF","#FF7F00","#000099"),tip.col=c("#0000FF","#FF7F00","#000099"))
TernaryPoints($x136, pch=16, cex=0.8, col='black')
"""

x245=closure(convert(Array,oligo_df_plus[!,[:csO2, :csO4,:csO5]]))
R"""
TernaryPlot()
TernaryPlot(atip = "O2-CB5",btip = "O4-III/IV", ctip = "O5-IE", lab.col=c("#CC334C","#009900","#00CCFF"),tip.col=c("#CC334C","#009900","#00CCFF"),grid.lines = 10, grid.minor.lines = 0, grid.col = "lightgrey")
TernaryPoints($x245, pch=16, cex=0.8, col='black')
"""

## Fancier year day colors:

x123=closure(convert(Array,oligo_df_plus[!,[:csO1,:csO2,:csO3]]))
yearday=dayofyear.(oligo_df_plus[!,:date])

R"""
TernaryPlot()
jet <- colorRampPalette(c("#00007F", "blue", "#007FFF", "cyan",
                     "#7FFF7F", "yellow", "#FF7F00", "red", "#7F0000"))
jmap <- jet(365)[$yearday] #first argument is max....I think...
TernaryPlot(atip = "O1-I", btip = "O2-CB5", ctip = "O3-II/XV", lab.col=c("#0000FF","#CC334C","#FF7F00"),tip.col=c("#0000FF","#CC334C","#FF7F00"),grid.lines = 10, grid.minor.lines = 0, grid.col = "lightgrey")
TernaryPoints($x123, pch=16, cex=0.8, col=jmap) #whoo hoo!!!!!
"""

## Fit a line through this data in the simplex

#Note, this is easiest to do with an ilr tranform, fitting the line with standard linear regression and then back transfrom to simplex:

#Line throuhg oligotype 1,2,3 data:
x123=closure(convert(Array,oligo_df_plus[!,[:csO1,:csO2,:csO3]]))
sbp=[-1 -1 1; 1 -1 -0] #to form basis
ilr_x, bal=ilr(x123,sbp) #ilr transform

#see what this transformed data looks like:
Gadfly.set_default_plot_size(35cm,25cm)
layer1=layer(x=ilr_x[:,2], y=ilr_x[:,1],Geom.point, style(default_color=RGB(0,0,1))) #okay, that seems to be working...
plt=plot(layer1)

#fit line:
T1=DataFrame(x = ilr_x[:,2], y = ilr_x[:,1])
test1 = lm(@formula(y ~ x), T1)
bf=zeros(2,1)
bf[2] = GLM.coef(test1)[2]
bf[1] = GLM.coef(test1)[1]

xf=collect(-2:0.1:8) #and see what this fit looks like:
yf=bf[1] .+ bf[2] .* xf
layer2=layer(x=xf, y=yf, Geom.line)
push!(plt,layer2) #check that line goes through points!

## Now plot this line back in the simplex and save:

ll=invilr([yf xf],bal) #back transform fitted line

R"""
TernaryPlot()
jet <- colorRampPalette(c("#00007F", "blue", "#007FFF", "cyan",
                     "#7FFF7F", "yellow", "#FF7F00", "red", "#7F0000"))
jmap <- jet(365)[$yearday] #first argument is max....I think...
TernaryPlot(atip = "O1-I", btip = "O2-CB5", ctip = "O3-II/XV", lab.col=c("#0000FF","#CC334C","#FF7F00"),tip.col=c("#0000FF","#CC334C","#FF7F00"),grid.lines = 10, grid.minor.lines = 0, grid.col = "lightgrey")
TernaryPoints($x123, pch=16, cex=0.8, col=jmap) #whoo hoo!!!!!
TernaryLines($ll, col='black')
"""

#and if wanted to save:
R"""
png(file="/home/kristen/Documents/V6V8_analysis/paper_V6V8/V6V8_Syn_CoDa_paper/figures/tern1.png")
TernaryPlot()
jet <- colorRampPalette(c("#00007F", "blue", "#007FFF", "cyan",
                     "#7FFF7F", "yellow", "#FF7F00", "red", "#7F0000"))
jmap <- jet(365)[$yearday] #first argument is max....I think...
TernaryPlot(atip = "O1-I", btip = "O2-CB5", ctip = "O3-II/XV", lab.col=c("#0000FF","#CC334C","#FF7F00"),tip.col=c("#0000FF","#CC334C","#FF7F00"),grid.lines = 10, grid.minor.lines = 0, grid.col = "lightgrey")
TernaryPoints($x123, pch=16, cex=0.8, col=jmap) #whoo hoo!!!!!
TernaryLines($ll, col='black')
dev.off()
"""

## Again for the remaining set:

x456=closure(convert(Array,oligo_df_plus[!,[:csO4,:csO5,:csO6]]))
sbp=[-1 1 -1; 1 0 -1]
ilr_x, bal=ilr(x456,sbp)
layer1=layer(x=ilr_x[:,2], y=ilr_x[:,1],Geom.point, style(default_color=RGB(0,0,1))) #okay, that seems to be working...
plt=plot(layer1)

#fit line just within date:
T1=DataFrame(x = ilr_x[:,2], y = ilr_x[:,1])
test1 = lm(@formula(y ~ x), T1)
bf=zeros(2,1)
bf[2] = GLM.coef(test1)[2]
bf[1] = GLM.coef(test1)[1]

xf=collect(-3:0.1:5)
yf=bf[1] .+ bf[2] .* xf
layer2=layer(x=xf, y=yf, Geom.line)
push!(plt,layer2) #check that line goes through points!

ll=invilr([yf xf],bal) #back transform fitted line...

R"""
png(file="/home/kristen/Documents/V6V8_analysis/paper_V6V8/V6V8_Syn_CoDa_paper/figures/tern2.png")
TernaryPlot()
jet <- colorRampPalette(c("#00007F", "blue", "#007FFF", "cyan",
                     "#7FFF7F", "yellow", "#FF7F00", "red", "#7F0000"))
jmap <- jet(365)[$yearday] #first argument is max....I think...
TernaryPlot(atip = "O4-III/IV", btip = "O5-IE", ctip = "O6-I",lab.col=c("#009900","#00CCFF","#000099"),tip.col=c("#009900","#00CCFF","#000099"),grid.lines = 10, grid.minor.lines = 0, grid.col = "lightgrey")
TernaryPoints($x456, pch=16, cex=0.8, col=jmap) #whoo hoo!!!!!
TernaryLines($ll, col='black')
dev.off()
"""

## to get the simplified relationship:

a=ll[:,1]./ll[:,2]
b=ll[:,2]./ll[:,3]

# a=ll[:,1]./ll[:,3]
# b=ll[:,3]./ll[:,2]

T1=DataFrame(x = log.(a), y = log.(b))

test2 = lm(@formula(y ~ x), T1)
bf=zeros(2,1)
bf[2] = GLM.coef(test2)[2]
bf[1] = GLM.coef(test2)[1]

plot(x=log.(a),y=log.(b),Geom.point)


using PyPlot, transD_GP, Random, LinearAlgebra

## model
ρbg        = 1000
zfixed     = [-1e5,   0,   (20:20:200)...]
rho        = [1e12,   ρbg, ρbg*ones(10)...]
ianom      = 5
rho[ianom] = 10.
δ          = 1e-4  # fd step size in log10

## Gauss Newton settings
use_fd          = false # in the one step for Gauss Newton
startfromcloser = false  # whether to start midway between anomaly and bg or plain bg in G-N step
λ²              = 5e-17 # model update damping

##  geometry and frequencies for VMD in frequency domain
nFreqsPerDecade     = 5
freqLowLimit        = 1e-3
freqHighLimit       = 1e5
freqs = 10 .^(log10(freqLowLimit):1/nFreqsPerDecade:log10(freqHighLimit))
rRx = 100.
zRx = -37.
zTx = -35.
nkᵣeval = 50
modelprimary = true
doconvramp = false
provideddt = true
Fvmd = transD_GP.AEM_VMD_HMD.HFieldDHT(
                      calcjacobian = true,
                      nmax = length(zfixed),
                      zTx    = zTx,
                      rRx    = rRx,
                      freqs  = freqs,
                      zRx    = zRx,
                      nkᵣeval = nkᵣeval,
                      modelprimary = modelprimary,
                      doconvramp = doconvramp,
                      provideddt = provideddt)
transD_GP.AEM_VMD_HMD.getfieldTD!(Fvmd, zfixed, rho)

## plot
times = Fvmd.interptimes
f, ax = plt.subplots(1,3, sharey=true, figsize=(8,6))
delz, delt = diff(zfixed[2:end]), diff(log10.(freqs))
tlabel = [(log10.(times) .- delt[1]/2)..., log10.(times)[end]+delt[end]/2]
zlab = [zfixed[2:end]..., zfixed[end]+delz[end]]
ax[1].step(log10.(rho[2:end]), zfixed[2:end])
im2 = ax[2].pcolormesh(tlabel, zlab, Fvmd.HTD_z_J_interp[2:end,:])
cb2 = colorbar(im2, ax=ax[2])
im3 = ax[3].pcolormesh(tlabel, zlab, log10.(abs.(Fvmd.HTD_z_J_interp[2:end,:])))
cb2 = colorbar(im3, ax=ax[3])
ax[1].set_xlabel("log10 Ω-m")
ax[1].invert_xaxis()
ax[2].invert_yaxis()
ax[2].set_xlabel("log10 s")
ax[3].set_xlabel("log10 s")
ax[1].set_ylabel("Depth m")
ax[2].set_title("∂f/∂(log10σ)")
ax[3].set_title("∂f/∂(log10σ)")

## approximate partials with central differences and compare
J = copy(Fvmd.HTD_z_J_interp)

function getJ_TD_noconv!(Fin, zin, ρ, Δ)
    # sensitivities computed in log10σ
    # same as in AEM_VMD_HMD code
    Jout = similar(Fin.HTD_z_J_interp)
    for i = 2:length(zin)
        modelfd = copy(ρ)
        modelfd[i] = 10 .^-(-log10(ρ[i]) .+ Δ/2)
        transD_GP.AEM_VMD_HMD.getfieldTD!(Fin, zin, modelfd)
        HTD_z1 = copy(Fin.HTD_z_interp)
        modelfd[i] = 10 .^-(-log10(ρ[i]) .- Δ/2)
        transD_GP.AEM_VMD_HMD.getfieldTD!(Fin, zin, modelfd)
        Jout[i,:] = (HTD_z1 - Fin.HTD_z_interp) / Δ
        end
    Jout
end
J_fd = getJ_TD_noconv!(Fvmd, zfixed, rho, δ)

## comparison plots
f, ax = plt.subplots(1,3, sharey=true, figsize=(10,6))
ax[1].step(log10.(rho[2:end]), zfixed[2:end])
im2 = ax[2].pcolormesh(tlabel, zlab, J[2:end,:])
cb2 = colorbar(im2, ax=ax[2])
im3 = ax[3].pcolormesh(tlabel, zlab, J_fd[2:end,:])
cb2 = colorbar(im3, ax=ax[3])
ax[1].set_xlabel("log10 Ω-m")
ax[1].invert_xaxis()
ax[2].invert_yaxis()
ax[2].set_xlabel("log10 Hz")
ax[3].set_xlabel("log10 s")
ax[1].set_ylabel("Depth m")
ax[2].set_title("∂f/∂(log10σ) analytic")
ax[3].set_title("∂f/∂(log10σ) FD")


import Base.Iterators: product

abstract type DesignType end

struct Opt <: DesignType
end
struct Rand <: DesignType
end

const Optimize = Opt()
const Randomize = Rand()

abstract type ModelType end

struct Stat <: DesignType
end
struct Dyn <: DesignType
end

const Static = Stat()
const Dynamic = Dyn()

"""
**Model**

Creates a model object containing a log likelihood function and prior distributions

- `prior`: a vector of distribution objects for each parameter. A uniform prior is used if no
prior is passed.
- `loglike`: a function that computes the log likelihood

Constructor

````julia
Model(args...; prior=nothing, loglike, kwargs...)
````
*Examples*

```julia
# Default uninform prior
loglike(θ, design, data) = ...

model = Model(;loglike)
```

```julia
# Custom priors
prior = [Beta(5,5),...]
loglike(θ, design, data) = ...

model = Model(;loglike, prior)
```
"""
mutable struct Model{F<:Function,P}
    prior::P
    loglike::F
end

function Model(args...; prior=nothing, loglike, kwargs...)
    return Model(prior, (x...)->loglike(x..., args...; kwargs...))
end

"""
    Optimizer(;task, model, grid_design, grid_parms, grid_response)

`Optimizer` constructs a model object for adaptive design optimization

- `model`: a model object
- `grid_design`:a grid of design parameters
- `grid_parms`: a grid of model parameters
- `grid_data`: a grid of data
- `log_like`: a three dimensional array of precomputed log likelihoods
- `marg_log_like`: a two dimensional array containing marginal log likelihoods for design and data
- `priors`: a multidimensional array of prior probabilities for parameters
- `log_post`: a one dimensional array of log posterior probabilities for parameters
- `entropy`: a two dimensional array of entropy values for parameter and design combinations
-
"""
mutable struct Optimizer{A,MT,M<:Model,T1,T2,T3,T4,T5,T6,T7,T8,T9,
    T10,T11}
    design_type::A
    model_type::MT
    model::M
    design_grid::T1
    parm_grid::T2
    data_grid::T3
    log_like::Array{Float64,3}
    marg_log_like::T4
    priors::T5
    log_post::Vector{Float64}
    entropy::T6
    marg_entropy::Vector{Float64}
    cond_entropy::Vector{Float64}
    mutual_info::Vector{Float64}
    best_design::T7
    parm_names::T8
    design_names::T9
    model_state::T10
    update_state!::T11
end

function Optimizer(args...; design_type=Optimize, model, design_list, parm_list, data_list,
    state_type=nothing, model_type=Static, update_state! = nothing, kwargs...)
    design_names,design_grid = to_grid(design_list)
    parm_names,parm_grid = to_grid(parm_list)
    _,data_grid = to_grid(data_list)
    dims = map(length, (parm_grid,design_grid,data_grid))
    model_state = create_state(model_type, state_type, dims, args...; kwargs...)
    log_like = loglikelihood(model, design_grid, parm_grid, data_grid, model_type, model_state)
    priors = prior_probs(model, parm_grid)
    post = priors[:]
    log_post = log.(post)
    entropy = compute_entropy(log_like)
    marg_log_like = marginal_log_like(log_post, log_like)
    marg_entropy = marginal_entropy(marg_log_like)
    cond_entropy = conditional_entropy(entropy, post)
    mutual_info = mutual_information(marg_entropy, cond_entropy)
    best_design = find_best_design(mutual_info, design_grid, design_names)
    return Optimizer(design_type, model_type, model, design_grid, parm_grid,
        data_grid, log_like,marg_log_like, priors, log_post, entropy, marg_entropy,
        cond_entropy, mutual_info, best_design, parm_names, design_names,
        model_state, update_state!)
end