text:
# BFGS solver for DMD problem
mutable struct DMD_BFGS_Options{T<:AbstractFloat}
	itm::Integer
	tol::T
	ptf::Integer
	# pre-allocate variables
	ar  ::Vector{T}
	nar ::Vector{T}
	gar ::Vector{T}
	ngar::Vector{T}
	par ::Vector{T}
	sar ::Vector{T}
	yar ::Vector{T}
	Har ::Matrix{T}
end

# constructor for BFGS options
function DMD_BFGS_Options(σ, params; itm=100, tol=1e-5, ptf=10)
	T = typeof(real(params.X[1]));
	tol = T(tol);
	k = params.k;

	ar   = zeros(T, 2*k);
	nar  = zeros(T, 2*k);
	gar  = zeros(T, 2*k);
	ngar = zeros(T, 2*k);
	par  = zeros(T, 2*k);
	sar  = zeros(T, 2*k);
	yar  = zeros(T, 2*k);
	Har  = diagm(0 => fill(σ, 2*k));

	return DMD_BFGS_Options(itm, tol, ptf, ar, nar,
		gar, ngar, par, sar, yar, Har)
end

# BFGS solver
function solveDMD_withBFGS(params, opts)
	# load all the variables
	T    = typeof(real(params.X[1]));
	itm  = opts.itm;
	tol  = opts.tol;
	ptf  = opts.ptf;

	ar   = params.ar;
	nar  = opts.nar;
	gar  = opts.gar;
	ngar = opts.ngar;
	par  = opts.par;
	sar  = opts.sar;
	yar  = opts.yar;
	Har  = opts.Har;

	# first gradient step
	nobj = aBFunc(params);
	aBGrad(ngar, params);
	err = norm(ngar, Inf);
	noi = 0;

	obj_his = zeros(T, itm);
	err_his = zeros(T, itm);

	while err ≥ tol
		# calculate the direction
		BLAS.gemv!('N', T(-1.0), Har, ngar, T(0.0), par);
		# decend line search
		η = 1.0;
		copyto!(nar, ar); BLAS.axpy!(η, par, ar);
		obj = aBFunc(params);
		while obj ≥ nobj
			η *= 0.5;
			η < 1e-10 && break;
			copyto!(ar, nar); BLAS.axpy!(η, par, ar);
			obj = aBFunc(params);
		end

		# update gradient
		aBGrad(gar, params);
		# update differences
		sar .= ar - nar; copyto!(nar, ar);
		yar .= gar - ngar; copyto!(ngar, gar);

		# update Hessian approximation
		ρ = T(1.0)/dot(yar, sar);
		μ = sum(abs2, sar);
		ρ = min(T(10.0), ρ*μ)/μ;
		if ρ > 0.0
			BLAS.gemv!('N', T(1.0), Har, yar, T(0.0), par);
			β = dot(yar, par)*ρ^2 + ρ;
			BLAS.ger!(-ρ, sar, par, Har);
			BLAS.ger!(-ρ, par, sar, Har);
			BLAS.ger!( β, sar, sar, Har);
		end

		# update information
		nobj = obj;
		err = norm(par, Inf);
		noi = noi + 1;

		obj_his[noi] = obj;
		err_his[noi] = err;

		# print information
		noi % ptf == 0 && @printf("iter %4d, obj %1.2e, err %1.2e\n",
			noi, obj, err);
		noi ≥ itm && break;
	end

	return obj_his[1:noi], err_his[1:noi]
end


text:
module IMFData

using HTTP
using JSON
using Dates
using DataFrames
import Base.show

const API_URL = "http://dataservices.imf.org/REST/SDMX_JSON.svc"

export
    IMFSeries,
    get_imf_datasets,
    get_imf_datastructure,
    get_ifs_data

abstract type IMFSeries end

include("get_imf_datasets.jl")
include("get_imf_datastructure.jl")
include("get_ifs_data.jl")

end # module

ext: jl


text:
#!/usr/bin/env julia

module QJuliaCUDAFields

#load path to qjulia home directory
push!(LOAD_PATH, joinpath(@__DIR__, "../..", "core"))
push!(LOAD_PATH, @__DIR__)

import QJuliaRegisters
import QJuliaEnums
import QJuliaGrid
import QJuliaFields

using CUDAnative
using CuArrays

const complex_length = 2

function CreateGenericField(fdesc::QJuliaGrid.QJuliaLatticeFieldDescr_qj)

  generic_field = QJuliaFields.QJuliaLatticeField_qj(fdesc)

  if fdesc.grid.location == QJuliaEnums.QJULIA_CUDA_FIELD_LOCATION
    # Get reference to field parameters:
    fdesc = generic_field.field_desc
    # Set the field array total elements:
    tot_elems = complex_length*fdesc.siteSubset*(fdesc.geom == QJuliaEnums.QJULIA_SCALAR_GEOMETRY ? fdesc.nSpin : fdesc.nColor)*fdesc.nColor*fdesc.grid.grid_volumeCB

    generic_field.v = CuArray{fdesc.grid.register_type, 2}(undef, tot_elems, Int(fdesc.geom)*fdesc.nBlock)

    #address_space = AS.Global or AS.Generic etc.
    #dev_ptr = CUDAnative.DevicePtr{fdesc.grid.register_type,address_space}(tot_elems*Int(fdesc.geom)*fdesc.nBlock)
    #generic_field.v = CuDeviceArray{fdesc.grid.register_type,2,address_space}((tot_elems, Int(fdesc.geom)*fdesc.nBlock), dev_ptr)
  end

  return generic_field
end

end # module QJuliaCUDAFields

max_stars_repo_name: alexstrel/QJulia

max_issues_repo_name: alexstrel/QJulia

max_forks_repo_name: alexstrel/QJulia


text:
using PlantBiophysics
using Documenter
using PlantBiophysics
using DataFrames
using CSV
using Plots

DocMeta.setdocmeta!(PlantBiophysics, :DocTestSetup, :(using PlantBiophysics,Plots,DataFrames,CSV); recursive = true)

makedocs(;
    modules = [PlantBiophysics],
    authors = "remi.vezy <VEZY@users.noreply.github.com> and contributors",
    repo = "https://github.com/VEZY/PlantBiophysics.jl/blob/{commit}{path}#L{line}",
    sitename = "PlantBiophysics.jl",
    format = Documenter.HTML(;
        prettyurls = get(ENV, "CI", "false") == "true",
        canonical = "https://VEZY.github.io/PlantBiophysics.jl",
        assets = String[],
    ),
    pages = [
        "Home" => "index.md",
        "Getting started" => "get_started.md",
        "Models" => Any[
            "Photosynthesis" => "photosynthesis.md",
            "Stomatal conductance" => "gs.md",
            "Energy balance" => "energy_balance.md",
            "Light interception" => "light.md"
            ],
        "Details" => Any[
            "Package design" => "concepts.md",
            "Implement a model" => "implement_a_model.md",
            "Uncertainty propagation" => "uncertainty_propagation.md",
        ],
        "API" => "functions.md"
    ],
)

deploydocs(;
    repo = "github.com/VEZY/PlantBiophysics.jl.git",
)


text:
# ---
# title: 283. Move Zeroes
# id: problem283
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Array, Two Pointers
# link: <https://leetcode.com/problems/move-zeroes/description/>
# hidden: true
# ---
#
# Given an array `nums`, write a function to move all `0`'s to the end of it
# while maintaining the relative order of the non-zero elements.
#
# **Example:**
#
#
#
#     Input: [0,1,0,3,12]
#     Output: [1,3,12,0,0]
#
# **Note** :
#
#   1. You must do this **in-place** without making a copy of the array.
#   2. Minimize the total number of operations.
#
#
## @lc code=start
using LeetCode

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


text:
# Test functions related to the reduced clique graph and the merging strategy that uses it
using COSMO, Test, LinearAlgebra, Random

# We use the example graph from the paper in Habib, Stacho - Reduced clique graphs of chordal graphs(2011) - Fig. 1
snd = [Set([4, 5]); Set([1, 4, 6]); Set([1, 7]); Set([1, 8]); Set([1, 3, 4]); Set([1,2,3]); Set([2, 3, 9]); Set([3, 4, 11]); Set([3, 10])]
sep = [Set([1, 3]); Set([1, 4]); Set([2, 3]); Set([3, 4]); Set([1]); Set([3]); Set([4])]

rows, cols = COSMO.compute_reduced_clique_graph!(sep, snd)

edges = Tuple{Int64, Int64}[]
for k = 1:length(rows)
    push!(edges, (rows[k], cols[k]))
end
# compare this to the known reference solution
edges_ref = [(2, 1), (5, 1), (8, 1), (9, 8), (9, 5), (9, 7), (7, 6), (6, 4), (5, 4), (4, 2), (4, 3), (3, 2), (5, 3), (6, 3), (9, 6), (8, 5), (5, 2), (6, 5)]
permissible_ref = falses(length(edges_ref))
permissible_ref[[7, 11, 16, 17, 18]] .= true
@testset "Reduced clique graph" begin

    edges_contained = true
    for edge in edges
        if edge ∉ edges_ref
        edges_contained = false
        break
    end
    end
    @test edges_contained

    # let's make a supernodetree structure to test merging related functions
    N = 11
    t = COSMO.SuperNodeTree(deepcopy(snd), N)
    t.sep = sep
    COSMO.initialise!(t, t.strategy)

    # check whether an edge is permissible
    permissible = falses(length(edges))

    for (k, edge) in enumerate(edges)
        permissible[k] = COSMO.ispermissible(edge, t.strategy.adjacency_table, t.snd)
    end
    permissible_edges_ref = edges_ref[permissible_ref]
    for permissible_edge in edges[permissible]
        @test permissible_edge in permissible_edges_ref
    end

    cand, weight = COSMO.traverse(t, t.strategy)
    # merge the following cliques
    cand = [5, 2];

    @test !COSMO.evaluate(t, t.strategy, cand)
    COSMO.merge_two_cliques!(t, cand, t.strategy)
    @test isempty(t.snd[2])
    @test t.snd[5] == Set([1,3,4,6])

    # update the graph information
    COSMO.update!(t.strategy, t, cand, true)

    @test !haskey(t.strategy.adjacency_table, 2)
    @test mapreduce(x -> 2 ∈ x, +,  t.strategy.adjacency_table) == 0 #check that all 2's are deleted

    # consider a new merge example: 7 and 6
    N = 11
    t = COSMO.SuperNodeTree(deepcopy(snd), N)
    t.sep = sep
    COSMO.initialise!(t, t.strategy)
    cand = COSMO.traverse(t, t.strategy)
    cand = [7, 6];

    @test !COSMO.evaluate(t, t.strategy, cand)
    COSMO.merge_two_cliques!(t, cand, t.strategy)
    @test isempty(t.snd[6])
    @test t.snd[7] == Set([1,2,3,9])
    COSMO.update!(t.strategy, t, cand, true)
    @test !haskey(t.strategy.adjacency_table, 6)
    @test mapreduce(x -> 6 ∈ x, +,  t.strategy.adjacency_table) == 0 #check that all 6's are deleted

    # Check the recomputation step
    COSMO.clique_intersections!(t.strategy.edges, t.snd)
    COSMO.kruskal!(t.strategy.edges, t.num)
    COSMO.determine_parent_cliques!(t.snd_par, t.snd_child, t.snd, t.post, t.strategy.edges)
    t.snd_post = COSMO.post_order(t.snd_par, t.snd_child, t.num)
    t.sep = [Set{Int64}() for i=1:length(t.snd)]
    COSMO.split_cliques!(t.snd, t.sep, t.snd_par, t.snd_post, t.num)

    @test check_clique_tree(t.snd, t.sep, t.snd_par, t.snd_child)
end #@testset


text:
using SurveyDataWeighting
using Test

target_populations = [ 50.0, 20.0, 230.0, 35.0 ]

data = [
 1.0 1.0 0.0 0.0 ;
 0.0 1.0 0.0 0.0 ;
 1.0 0.0 2.0 0.0 ;
 0.0 0.0 6.0 1.0 ;
 1.0 0.0 4.0 1.0 ;
 1.0 1.0 0.0 0.0 ;
 1.0 0.0 5.0 0.0 ;
 0.0 0.0 6.0 1.0 ;
 0.0 1.0 0.0 0.0 ;
 0.0 0.0 3.0 1.0 ;
 1.0 0.0 2.0 0.0 ;
 1.0 1.0 0.0 1.0 ;
 1.0 0.0 3.0 1.0 ;
 1.0 0.0 4.0 0.0 ;
 0.0 0.0 5.0 0.0 ;
 0.0 1.0 0.0 1.0 ;
 1.0 0.0 2.0 1.0 ;
 0.0 0.0 6.0 0.0 ;
 1.0 0.0 4.0 1.0 ;
 0.0 1.0 0.0 0.0  ]

initial_weights = [
   3.0,
   3.0,
   5.0,
   4.0,
   2.0,
   5.0,
   5.0,
   4.0,
   3.0,
   3.0,
   5.0,
   4.0,
   4.0,
   3.0,
   5.0,
   3.0,
   4.0,
   5.0,
   4.0,
   3.0 ]

nrows = size( data, 1 )
ncols = size( data, 2 )
sp = size( target_populations, 1 )

@testset "Reproduce the basic test case in Creedy NEW ZEALAND TREASURY WORKING PAPER 03/17" begin

   @test ncols == size( target_populations, 1 )
   @test nrows == size( initial_weights, 1 )

   println( "target popns $target_populations" )

   # a = target_populations - (data'*initial_weights)
   # print( a )

   wchi = do_chi_square_reweighting( data, initial_weights, target_populations )
   println( "direct chi-square results $wchi")
   weighted_popn_chi = (wchi' * data)'
   println( "wchisq; got $weighted_popn_chi")
   @test weighted_popn_chi ≈ target_populations
   lower_multiple = 0.20 # any smaller min and d_and_s_constrained fails on this dataset
   upper_multiple = 2.19
   for m in instances( DistanceFunctionType )
      println( "on method $m")
      rw = do_reweighting(
            data               = data,
            initial_weights    = initial_weights,
            target_populations = target_populations,
            functiontype       = m,
            lower_multiple     = lower_multiple,
            upper_multiple     = upper_multiple,
            tolx               = 0.000001,
            tolf               = 0.000001 )
      println( "results for method $m = $rw" )
      weights = rw.weights
      weighted_popn = (weights' * data)'
      println( "weighted_popn = $weighted_popn" )
      @test weighted_popn ≈ target_populations
      if m != chi_square
         for w in weights # check: what's the 1-liner for this?
            @test w > 0.0
         end
      else
         @test weights ≈ wchi # chisq the direct way should match chisq the iterative way
      end
      if m in [constrained_chi_square, d_and_s_constrained ]
         # check the constrainted methods keep things inside ll and ul
         for r in 1:nrows
            @test weights[r] <= initial_weights[r]*upper_multiple
            @test weights[r] >= initial_weights[r]*lower_multiple
         end
      end
   end # meth loop
end # tests

include( "large_scale_test.jl")

max_stars_repo_name: grahamstark/SurveyDataWeighting

max_issues_repo_name: grahamstark/SurveyDataWeighting

max_forks_repo_name: grahamstark/SurveyDataWeighting


text:
export random_position, get_name, sigmoid, activ_dict

function rand_weight()
    rand() * 2.0 - 1.0 # initial weight is uniformly distributed between -1 and 1
end

function random_position(origin::Float64, destination::Float64)
    n_min = maximum([0, minimum([origin, destination])])
    n_max = minimum([1, maximum([origin, destination])])
    n = n_min + rand() * (n_max - n_min)
    if n <= 0 || n >= 1
        throw("Invalid neuron position: [$origin, $destination] => [$n_min, $n_max] => $n")
    end
    n
end

## Activation functions

function sigmoid(x::Float64)
    1 / (1 + exp(-5 * x))
end

function ReLU(x::Float64)
    if x < 0
        0
    else
        x
    end
end

function identity_activ(x::Float64)
    x
end

function gauss(x::Float64)
    exp(-5.0 * x^2)
end

function tanh_activ(x::Float64)
    tanh(2.5 * x)
end

activ_dict = Dict(
    "sigmoid" => sigmoid,
    "ReLU" => ReLU,
    "sin" => sin,
    "cos" => cos,
    "tanh" => tanh_activ,
    "abs" => abs,
    "identity" => identity_activ,
    "gauss" => gauss,
)


text:
# -------- Auxiliary Functions ----------------------------------------------- #

# Set an optimizer for all parameters of a model
function set_optimizer!(model, opt = Knet.Adam; kwargs...)

  for param in Knet.params(model)

    # Feature heads that are not used have length 0. They do not get an
    # optimizer.
    if length(param) > 0

      # If opt is given, overwrite every optimizer in param.opt
      if !isnothing(opt)

        param.opt = opt(; kwargs...)

      # If opt is not given, overwrite only non-existing param.opt with Adam
      elseif isnothing(opt) && isnothing(param.opt)

        param.opt = Knet.Adam(; kwargs...)

      end

    end

  end

end

# A single training step, the loss is returned
function train_step!(l :: Loss, model, cache :: DataCache)

  tape = Knet.@diff sum(loss(l, model, cache))

  for param in Knet.params(model)
    Knet.update!(Knet.value(param), Knet.grad(tape, param), param.opt)
  end

  Knet.value(tape)

end

# -------- Training on Datasets ---------------------------------------------- #

"""
    train!(model/player, dataset; <keyword arguments>)

Train `model`, or the training model of `player`, on `dataset`.

# Arguments
' `loss = Loss()`: Loss used for training.
- `epochs = 10`: Number of iterations through `dataset`.
- `batchsize = 50`: Batchsize for the update steps.
- `callback_epoch`: Function called after each epoch.
- `callback_step`: Function called after each update step.
- `optimizer = nothing`: Optimizer initialized for each weight in `model`
- `kwargs...`: Keyword arguments for `optimizer`

# Examples
```julia
set = record_self(MCTSPlayer(), 10, game = TicTacToe())
model = NeuralModel(TicTacToe, @chain TicTacToe Dense(50))
loss = Loss(policy = 0.15, regularization=1e-4)
train!(model, set, loss = loss, epochs = 15)
```

"""
function train!( player  :: Union{Player, Model}
               , trainset :: DataSet
               ; loss = Loss()
               , epochs = 10
               , batchsize = 50
               , callback_step = (_) -> nothing
               , callback_epoch = (_) -> nothing
               , optimizer = nothing
               , kwargs... )

  # Get basic info about the player's model
  model = training_model(player)
  gpu = on_gpu(model)

  # Check if features can be enabled
  use_features = check_features(loss, model, trainset)

  # Set/overwrite optimizers
  set_optimizer!(model, optimizer; kwargs...)

  # Generate the batch iterator
  batches = Batches( trainset
                   , batchsize
                   , shuffle = true
                   , partial = false
                   , gpu = gpu )

  # Iterate through the trainset
  for j in 1:epochs

    for (i, cache) in enumerate(batches)

      train_step!(loss, model, cache)
      callback_step(i)

    end

    callback_epoch(j)

  end

end

# -------- Training by Playing ----------------------------------------------- #

function _train!( player :: Player{G}
                , gen_data :: Function
                ; loss = Loss()
                , epochs = 10
                , playings = 20
                , iterations = 10
                , batchsize = 50
                , testfrac = 0.1
                , optimizer = Knet.Adam
                , replays = 0
                , quiet = false
                , callback_epoch = (_) -> nothing
                , callback_iter = (_) -> nothing
                , kwargs...
                ) where {G <: Game}

  @assert playings > 0 "Number of playings for training must be positive"

  # If we do not print results, it is not necessary to test
  quiet || testfrac < 0 && (testfrac = 0.)

  # Get number of playings used for training
  train_playings = ceil(Int, playings * (1-testfrac))

  # Print the loss header if not quiet
  !quiet && print_loss_header(loss, check_features(loss, player))

  # Set the optimizer for the player's model
  set_optimizer!(training_model(player), optimizer; kwargs...)

  # Save the last `replay` generated datasets and use them for training
  replay_buffer = []

  for i in 1:epochs

    # Generate callback functions for the progress meter
    step, finish = stepper("# Playing...", playings)
    cb = quiet ? () -> nothing : step

    # Generate new datasets for this generation
    datasets = gen_data(cb, playings)

    # Create the training set by merging with the last `playing` generations of
    # datasets that are stored in the replay_buffer
    trainset = merge(datasets[1:train_playings]...)
    for set in replay_buffer
      trainset = merge(trainset, set...)
    end

    # Create the testing set by collecting the remaining datasets generated
    # in this epoche
    if testfrac > 0
      testset = merge(datasets[train_playings+1:playings]...)
    else
      testset = DataSet{G}()
    end

    # Clear the progress bar (if it was printed)
    finish()

    # Prepare the next progress bar
    steps = iterations * div(length(trainset), batchsize)
    step, finish = stepper("# Learning...", steps)
    cb = (_) -> quiet ? nothing : step()

    # The Knet allocator works better for training but worse for playing
    switch_knet_allocator()

    # Train the player with the generated dataset
    train!( player
          , trainset
          , loss = loss
          , epochs = iterations
          , batchsize = batchsize
          , callback_step = cb
          , callback_epoch = callback_iter )

    # Clear the progress bar
    finish()

    # Calculate and print loss for this epoch if not quiet
    !quiet && print_loss(loss, player, i, trainset, testset)

    # Undo the changes in the Knet allocator
    switch_knet_allocator()

    # Bring the full dataset from this epoche in the replay buffer
    push!(replay_buffer, datasets)

    # Remove the oldest item from the replay buffer if it is to large
    length(replay_buffer) >= replays+1 && deleteat!(replay_buffer, 1)

    # Call callback function
    callback_epoch(i)

  end

end

"""
    train_self!(player; <keyword arguments>)

Train an MCTS `player` via playing against itself.

The training model of `player` learns to predict the MCTS policy, the value of
game states, and possible features (if `player` supports the same features as
the keyword argument `loss`). Note that only players with NeuralModel-based
training models can be trained currently.

# Arguments
- `loss = Loss()`: Loss used for training.
- `epochs = 10`: Number of epochs.
- `playings = 20`: Games played per `epoch` for training-set generation.
- `iterations = 10`: Number of training epochs per training-set.
- `batchsize = 50`: Batchsize during training from the training-set.
- `testfrac = 0.1`: Fraction of `playings` used to create test-sets.
- `branching = 0.`: Random branching probability.
- `augment = true`: Whether to use augmentation on the created data sets.
- `replays = 0`: Add datasets from last `replay` epochs to the trainset.
- `quiet = false`: Whether to suppress logging of training progress.
- `callback_epoch`: Function called after every epoch.
- `callback_iter`: Function called after every iteration.
- `distributed = false`: Shall workers be used for self-playings?
- `tickets = nothing`: Number of tickets if distributed.
- `optimizer = Adam`: Optimizer for each weight in the training model.
- `kwargs...`: Keyword arguments for `optimizer`

# Examples
```julia
# Self-train a simple neural network model for 5 epochs
model = NeuralModel(TicTacToe, @chain TicTacToe Conv(64, relu) Dense(32, relu))
player = MCTSPlayer(model, power = 50, temperature = 0.75, exploration = 2.)
train_self!(player, epochs = 5, playings = 100)
```
"""
function train_self!( player :: MCTSPlayer
                    ; loss = Loss()
                    , prepare = prepare(steps = 0)
                    , branch = branch(prob = 0.)
                    , augment = true
                    , distributed = false
                    , tickets = nothing
                    , kwargs... )

  # Only use player-enabled features for recording if they are compatible with
  # the loss
  features = check_features(loss, player) ? Jtac.features(player) : Feature[]

  # Function to generate datasets through selfplays
  gen_data = (cb, n) -> record_self( player
                                   , n
                                   , prepare = prepare
                                   , branch = branch
                                   , augment = augment
                                   , features = features
                                   , merge = false
                                   , callback = cb
                                   , distributed = distributed
                                   , tickets = tickets )

  _train!(player, gen_data; loss = loss, kwargs...)

end

"""
    train_against!(player, enemy; <keyword arguments>)

Train an MCTS `player` under `loss` through playing against `enemy`.

The training model of `player` learns to predict the MCTS policy, the value of
game states, and possible features (if `player` supports the same features as
the keyword argument `loss`) when playing against `enemy`. If the enemy is too
good, this may turn out to be a bad learning mode: a player that loses all the
time will produce very pessimistic value predictions for each single game state,
which will harm the MCTS algorithm for improved policies. Note that only players
with NeuralModel-based training models can be trained currently.

# Arguments
- `loss = Loss()`: Loss used for training.
- `start`: Function that (randomly) yields -1 or 1 to fix the starting player.
- `epochs = 10`: Number of epochs.
- `playings = 20`: Games played per `epoch` for training-set generation.
- `iterations = 10`: Number of training epochs per training-set.
- `batchsize = 50`: Batchsize during training from the training-set.
- `testfrac = 0.1`: Fraction of `playings` used to create test-sets.
- `branching = 0.`: Random branching probability.
- `augment = true`: Whether to use augmentation on the created data sets.
- `replays = 0`: Add datasets from last `replay` epochs to the trainset.
- `quiet = false`: Whether to suppress logging of training progress.
- `callback_epoch`: Function called after every epoch.
- `callback_iter`: Function called after every iteration.
- `distributed = false`: Shall workers be used for self-playings?
- `tickets = nothing`: Number of tickets if distributed.
- `optimizer = Adam`: Optimizer for each weight in the training model.
- `kwargs...`: Keyword arguments for `optimizer`

# Examples
```julia
# Train a simple neural network model against an MCTS player for 5 epochs
model = NeuralModel(TicTacToe, @chain TicTacToe Conv(64, relu) Dense(32, relu))
player = MCTSPlayer(model, power = 50, temperature = 0.75, exploration = 2.)
enemy = MCTSPlayer(power = 250)
train_against!(player, enemy, epochs = 5, playings = 100)
```
"""
function train_against!( player :: MCTSPlayer
                       , enemy
                       ; loss = Loss()
                       , start :: Function = () -> rand([-1, 1])
                       , branching = 0.
                       , augment = true
                       , distributed = false
                       , tickets = nothing
                       , kwargs... )

  features = check_features(loss, player) ? features(player) : Feature[]

  gen_data = (cb, n) -> record_against( player
                                      , enemy
                                      , n
                                      , start = start
                                      , augment = augment
                                      , branching = branching
                                      , features = features
                                      , merge = false
                                      , callback = cb
                                      , distributed = distributed
                                      , tickets = tickets )

  _train!(player, gen_data; loss = loss, kwargs...)

end

"""
    train_from_model!(pupil, teacher [, players]; <keyword arguments>)

Train a `pupil` by letting it approximate the predictions of
`teacher`'s model on game states created by `players`.

In each epoch, two random `players` are chosen (by default, these are the pupil
and the teacher) to generate a set of game states by playing the game. These
games are used to create a training set by applying the teachers training model.
Note that only players with NeuralModel-based training models can be trained
currently.

# Arguments
- `loss = Loss()`: Loss used for training.
- `epochs = 10`: Number of epochs.
- `playings = 20`: Games played per `epoch` for training-set generation.
- `iterations = 10`: Number of training epochs per training-set.
- `batchsize = 50`: Batchsize during training from the training-set.
- `testfrac = 0.1`: Fraction of `playings` used to create test-sets.
- `branching = 0.`: Random branching probability.
- `augment = true`: Whether to use augmentation on the created data sets.
- `replays = 0`: Add datasets from last `replay` epochs to the trainset.
- `quiet = false`: Whether to suppress logging of training progress.
- `callback_epoch`: Function called after every epoch.
- `callback_iter`: Function called after every iteration.
- `optimizer = Adam`: Optimizer for each weight in the training model.
- `kwargs...`: Keyword arguments for `optimizer`
"""
function train_from_model!( pupil :: Player{G}
                          , teacher :: Union{Model{H}, Player{H}}
                          , players = [pupil]
                          ; loss = Loss()
                          , branching = 0.
                          , augment = true
                          , kwargs...
                          ) where {H <: Game, G <: H}

  # Complement the branching options (before, during, steps)
  branching = branch_options(branching)

  # Check if pupil and teacher are compatible featurewise
  use_features = features(pupil) == features(teacher) != Feature[]

  # Function that generates datasets
  gen_data = (cb, n) -> begin

    # TODO: Async?

    # Construct a list of games by letting two players compete
    datasets = asyncmap(1:n) do _

      # Get two players at random
      p1, p2 = rand(players, 2)

      # Watch them playing
      games = pvp_games(p1, p2, game = branch_root(G(), branching))[1:end-1]

      # Force the players to 'make errors' sometimes
      branchgames = branch(games, branching)
      branches = map(g -> pvp_games(p1, p2, g), branchgames)

      # Let the teacher model look at all games to generate a dataset
      ds = record_model( training_model(teacher)
                       , vcat(games, branches...)
                       , use_features = use_features
                       , augment = augment )

      # Give the signal that one playing is complete
      cb()

      ds

    end

    datasets

  end

  _train!(pupil, gen_data; loss = loss, kwargs...)

end

# -------- Training With Contests -------------------------------------------- #

"""
    with_contest(train_function, player, <arguments>; <keyword arguments>)

Train `player` according to `train_function` while conducting regular contests
during the training process.

The argument `train_function` must be one of `train_self!`, `train_against!`, or
`train_from!`. The function `with_contest` will always at least print the
results of a contest before training and after training.

# Arguments
The non-keyword arguments supported are the non-keyword arguments that are also
accepted / needed by `train_function`. As keyword arguments, all keyword
arguments of `train_function` are supported. Additionally, the following
arguments can be used to adapt the behavior of `with_contest`:

- `interval = 10`: After how many epochs do contests take place.
- `length = 250`: Number of playings during one contest.
- `opponents`: List of opponents for the competition. To this list the
   two players `current` and `initial` are added, which are IntuitionPlayers
   with the model (or a copy of the initial model) of `player`.
- `temperature`: Temperature of the intuition players `current` and `initial`.
- `cache = 0`: Number of games to be cached. If `cache > 0`, a number
   of `cache` games between all players that are independent of `player`'s model
   are conducted and cached for all competitions that follow.

# Examples
```julia

G = TicTacToe
loss = Loss(policy = 0.25)
opponents = [MCTSPlayer(power = 50), MCTSPlayer(power = 500)]

model = NeuralModel(G, @chain G Conv(100, relu) Dense(32, relu))
player = MCTSPlayer(model, power = 25)

with_contest( train_self!
            , player
            , loss = loss
            , cache = 1000
            , length = 500
            , opponents = opponents
            , interval = 5
            , epochs = 20
            , playings = 150 )
```

"""
function with_contest( trainf!     # the training function
                     , player :: MCTSPlayer
                     , args...
                     ; loss = Loss()
                     , opponents = Player[]
                     , length :: Int = 250
                     , cache :: Int = 0
                     , interval :: Int = 10
                     , temperature = player.temperature
                     , distributed = false
                     , epochs = 10
                     , kwargs... )

  # One can set length == 0 to access trainf! without contests
  if length <= 0

    @info "Contests are disabled."
    return trainf!( player
                  , args...
                  ; loss = loss
                  , epochs = epochs
                  , distributed = distributed
                  , kwargs... )

  end

  # Rename the length keyword argument
  len = length
  length = Base.length

  # List of all players that will compete in the contest
  players = [

    IntuitionPlayer( player.model
                   , temperature = temperature
                   , name = "current" );
    IntuitionPlayer( copy(player.model)
                   , temperature = temperature
                   , name = "initial" );
    opponents

  ]

  # Creating a cache of contest games for all players that do not depend on the
  # model that will be trained
  if cache > 0

    # Get the training model
    model = training_model(player)

    # Reorder players such that the active (i.e., learning) players come last
    # and passive ones first
    active = filter(p -> training_model(p) == model, players)
    passive = filter(p -> training_model(p) != model, players)
    players = [passive; active]

    # Get the respective numbers of players
    n  = length(players)
    np = length(passive)
    na = length(active)

    # Get the indices of all active models
    aidx = collect(1:na) .+ np

    # Get the progress-meter going
    step, finish = stepper("# Caching...", cache)

    # Create the cache for games between passive players
    cache_results = zeros(Int, n, n, 3)
    cache_results[1:np,1:np,:] = compete( passive
                                        , cache
                                        , distributed = distributed
                                        , callback = step )

    # Remove the progress bar
    finish()

    #  Leave a message that confirms caching.
    print(gray_crayon)
    println("# Cached $cache matches by $n players")

  else

    # If caching is not activated, regard each player as active
    aidx = 1:length(players)
    cache_results = zeros(Int, length(players), length(players), 3)

  end

  # Create the callback function
  cb = epoch -> begin
    if epoch % interval == 0

      step, finish = stepper("# Contest...", len)

      results = compete( players
                       , len
                       , aidx
                       , distributed = distributed
                       , callback = step )

      finish()

      print_ranking(Ranking(players, results + cache_results))

      0 < epoch < epochs && print_loss_header(loss, check_features(loss, player))

    end
  end

  # Run contest before training
  cb(0)

  # Train with a contest every interval epochs
  trainf!( player
         , args...
         ; loss = loss
         , callback_epoch = cb
         , epochs = epochs
         , distributed = distributed
         , kwargs... )

  # Run contest after training, if it was not run in trainf! already
  epochs % interval == 0 ? nothing : cb(epochs)

end

max_stars_repo_name: roSievers/jtac


text:
function SphericalCoordinates(arg0::Vector3D)
    return SphericalCoordinates((Vector3D,), arg0)
end

function SphericalCoordinates(arg0::jdouble, arg1::jdouble, arg2::jdouble)
    return SphericalCoordinates((jdouble, jdouble, jdouble), arg0, arg1, arg2)
end

function get_cartesian(obj::SphericalCoordinates)
    return jcall(obj, "getCartesian", Vector3D, ())
end

function get_phi(obj::SphericalCoordinates)
    return jcall(obj, "getPhi", jdouble, ())
end

function get_r(obj::SphericalCoordinates)
    return jcall(obj, "getR", jdouble, ())
end

function get_theta(obj::SphericalCoordinates)
    return jcall(obj, "getTheta", jdouble, ())
end

function to_cartesian_gradient(obj::SphericalCoordinates, arg0::Vector{jdouble})
    return jcall(obj, "toCartesianGradient", Vector{jdouble}, (Vector{jdouble},), arg0)
end

function to_cartesian_hessian(obj::SphericalCoordinates, arg0::Vector{Vector{jdouble}}, arg1::Vector{jdouble})
    return jcall(obj, "toCartesianHessian", Vector{Vector{jdouble}}, (Vector{Vector{jdouble}}, Vector{jdouble}), arg0, arg1)
end


text:
#!/usr/bin/julia

# Copyright (C) 2021, 2022 Matheus Fernandes Bigolin <mfrdrbigolin@disroot.org>
# SPDX-License-Identifier: MIT

"Day Fifteen, Chiton."

include("Utils.jl")

# N.b.: external package, used for the PriorityQueue data structure.
using DataStructures

using DelimitedFiles
using .Utils

"Find the `cavern` path with the lowest risk; then, return this total risk."
function solve(cavern :: Array{Int8, 2}) :: Int
    sizepaths = size(cavern)

    unvisited = [(i,j) for j in 1:sizepaths[2] for i in 1:sizepaths[1]]

    distance = PriorityQueue(zip(unvisited, fill(Inf, length(cavern))))
    distance[(1,1)] = 0

    current = (1,1)

    # Dijkstra’s algorithm.
    while current != sizepaths
        i, j = current
        neighbors = Array{Tuple{Int64, Int64}}([])

        for idx in [(i - 1, j), (i + 1, j), (i, j - 1), (i, j + 1)]
            idxs = CartesianIndices(cavern)
            if CartesianIndex(idx) in idxs && idx in keys(distance)
                push!(neighbors, idx)
            end
        end

        for neighbor in neighbors
            rel_dist = cavern[neighbor...] + peek(distance)[2]
            if rel_dist < distance[neighbor]
                distance[neighbor] = rel_dist
            end
        end

        dequeue!(distance)

        current = peek(distance)[1]
    end

    return distance[sizepaths]
end

"Enlarge `cavern` to its actual dimension."
function enlarge(cavern :: Array{Int8, 2}) :: Array{Int8, 2}
    pattern = reshape([Int8(i + j) for i in 0:4 for j in 0:4], (5,5))

    enlarged_pattern = repeat(pattern, inner=size(cavern))
    repeated_cavern = repeat(cavern, 5, 5)

    return [Int8(mod1(i, 9)) for i in repeated_cavern + enlarged_pattern]
end

function main()
    usage_and_exit(length(ARGS) != 1)

    data = collect.(readdlm(first(ARGS), String))

    cavern = hcat([parse.(Int8, path) for path in data]...)
    large_cavern = enlarge(cavern)

    println(solve(cavern))
    println(solve(large_cavern))
end

if abspath(PROGRAM_FILE) == @__FILE__
    main()
end


text:
# tests value choice points, i.e. using choose(DataType) construct

# TODO to avoid repetition, would like to construct tests for types using metaprogramming
# HOWEVER, this works with Expressions constructed without interpolation, but as soon as
# interpolation is performed:
# (a) an explicit macroexpand is required before the eval
# (b) even then, either (DataGenerators.)choose is not recognised or, if instead Autotest.eval is explicitly using, then (DataGenerators.)gen is not recognised - a scoping issue/bug?
# And it does not appear to be possible to manipulate code as a string as an alternative (using parse, then eval)
#
# NOTE: For this reason, currently only the most common types are tested, and even within these, not all
# combinations of parameters are tested

@generator VCBoolGen begin # prefix VC (for Value Choice points) to avoid type name clashes
    start() = choose(Bool)
end

@generator VCIntGen begin
    start() = choose(Int,-1000,2000)
end

@generator VCIntNonLiteralGen begin
    start() = choose(Int,-10*100,2*1000)
end

@generator VCIntVariableGen begin
    start() = choose(Int,-1000,2000)
    x() = -1000
    y() = 2000
end

@generator VCIntNoMaxGen begin
    start() = choose(Int,-1000)
end

@generator VCIntNoRangeGen begin
    start() = choose(Int)
end

@generator VCInt64Gen begin
    start() = choose(Int64,-1000,2000)
end

@generator VCInt32Gen begin
    start() = choose(Int32,-1000,2000)
end

@generator VCUIntNoRangeGen begin
    start() = choose(UInt)
end

@generator VCUIntGen begin
    start() = choose(UInt,4,420)
end

@generator VCFloat64Gen begin
    start() = choose(Float64,-4.2,32.7)
end

@generator VCFloat64NoMaxGen begin
    start() = choose(Float64,-4.2)
end

@generator VCFloat64NoRangeGen begin
    start() = choose(Float64,-4.2)
end

@generator VCFloat64NonLiteralGen begin
    start() = choose(Float64,-2.1*2,32.0+0.7)
end

@generator VCFloat64VariableGen begin
    start() = choose(Float64,x(),y())
    x() = -4.2
    y() = 32.7
end

@generator VCFloat32Gen begin
    start() = choose(Float32,-4.2,32.7)
end

@testset "value choice points" begin

	@testset "choose(Bool)" begin

	    gn = VCBoolGen()

		@mtestset reps=Main.REPS alpha=Main.ALPHA begin
		    td = choose(gn)
		    @test typeof(td) == Bool
		    @mtest_values_vary td
		end

	end

	@testset "choose(Int)" begin

		gn = VCIntGen()

		@mtestset reps=Main.REPS alpha=Main.ALPHA begin
		    td = choose(gn)
		    @test typeof(td) == Int
		    @test -1000 <= td <= 2000
		    @mtest_values_vary td
		end

	end

	@testset "choose(Int) using range defined by non-literals" begin

		gn = VCIntNonLiteralGen()

		@mtestset reps=Main.REPS alpha=Main.ALPHA begin
		    td = choose(gn)
		    @test typeof(td) == Int
		    @test -1000 <= td <= 2000
		    @mtest_values_vary td
		end

	end

	@testset "choose(Int) using range defined by variables" begin

		gn = VCIntVariableGen()

		@mtestset reps=Main.REPS alpha=Main.ALPHA begin
		    td = choose(gn)
		    @test typeof(td) == Int
		    @test -1000 <= td <= 2000
		    @mtest_values_vary td
		end

	end

	@testset "choose(Int) with no max specified" begin

		gn = VCIntNoMaxGen()

		@mtestset reps=Main.REPS alpha=Main.ALPHA begin
		    td = choose(gn)
		    @test typeof(td) == Int
		    @test -1000 <= td
		    @mtest_values_vary td
		end

	end

	@testset "choose(Int) with no min nor max specified" begin

		gn = VCIntNoRangeGen()

		@mtestset reps=Main.REPS alpha=Main.ALPHA begin
		    td = choose(gn)
		    @test typeof(td) == Int
		    @mtest_values_vary td
		end

	end

	@testset "choose(Int64)" begin

		gn = VCInt64Gen()

		@mtestset reps=Main.REPS alpha=Main.ALPHA begin
		    td = choose(gn)
		    @test typeof(td) == Int
		    @test -1000 <= td <= 2000
		    @mtest_values_vary td
		end

	end

	@testset "choose(Int32)" begin

		gn = VCInt32Gen()

		@mtestset reps=Main.REPS alpha=Main.ALPHA begin
		    td = choose(gn)
		    @test typeof(td) == Int32
		    @test -1000 <= td <= 2000
		    @mtest_values_vary td
		end

	end

	@testset "choose(UInt) with no min nor max specified" begin

		gn = VCUIntNoRangeGen()

		@mtestset reps=Main.REPS alpha=Main.ALPHA begin
		    td = choose(gn)
		    @test typeof(td) == UInt
		    @mtest_values_vary td
		end

	end

	@testset "choose(UInt)" begin

		gn = VCUIntGen()

		@mtestset reps=Main.REPS alpha=Main.ALPHA begin
		    td = choose(gn)
		    @test typeof(td) == UInt
		    @test 4 <= td <= 420
		    @mtest_values_vary td
		end

	end

	@testset "choose(Float64)" begin

		gn = VCFloat64Gen()

		@mtestset reps=Main.REPS alpha=Main.ALPHA begin
		    td = choose(gn)
		    @test typeof(td) == Float64
		    @test -4.2 <= td <= 32.7
		    @mtest_values_vary td
		end

	end

	@testset "choose(Float64) with no max specified" begin

		gn = VCFloat64NoMaxGen()

		@mtestset reps=Main.REPS alpha=Main.ALPHA begin
		    td = choose(gn)
		    @test typeof(td) == Float64
		    @test -4.2 <= td
		    @mtest_values_vary td
		end

	end

	@testset "choose(Float64) with no min nor max specified" begin

		gn = VCFloat64NoRangeGen()

		@mtestset reps=Main.REPS alpha=Main.ALPHA begin
		    td = choose(gn)
		    @test typeof(td) == Float64
		    @mtest_values_vary td
		end

	end

	@testset "choose(Float64) using defined by non-literals" begin

		gn = VCFloat64NonLiteralGen()

		@mtestset reps=Main.REPS alpha=Main.ALPHA begin
		    td = choose(gn)
		    @test typeof(td) == Float64
		    @test -4.2 <= td <= 32.7
		    @mtest_values_vary td
		end

	end

	@testset "choose(Float64) using range defined by variables" begin

		gn = VCFloat64VariableGen()

		@mtestset reps=Main.REPS alpha=Main.ALPHA begin
		    td = choose(gn)
		    @test typeof(td) == Float64
		    @test -4.2 <= td <= 32.7
		    @mtest_values_vary td
		end

	end

	@testset "choose(Float32)" begin

		gn = VCFloat32Gen()

		@mtestset reps=Main.REPS alpha=Main.ALPHA begin
		    td = choose(gn)
		    @test typeof(td) == Float32
		    @test -4.2 <= td <= 32.7
		    @mtest_values_vary td
		end

	end

end


text:
#generic high order function to apply any function to a vector column of a df
function ApplyColOper(df, colName::String, f)
    return df[!, colName] |> (f);
end


text:
using PeaceVote
using PeaceVote: Vote, Option

setnamespace(@__MODULE__)

import Pkg
Pkg.develop(Pkg.PackageSpec(url = "https://github.com/PeaceFounder/Community.jl"))
Pkg.resolve()

import Community

#import Pkg.Types.UUID
#mycom = UUID("3fc59347-698f-470f-bc28-6120569cf229")
uuid = PeaceVote.uuid("Community")

# Cleanup from the previous tests
dirs = [PeaceVote.keydir(uuid), PeaceVote.datadir(uuid), PeaceVote.communitydir(uuid)]
for dir in dirs
    isdir(dir) && rm(dir,recursive=true)
end

# Start a community server.
import PeaceFounder

maintainer = PeaceVote.Signer(uuid,"maintainer")
server = PeaceVote.Signer(uuid,"server")

ballotconfig = Community.BallotConfig(2000,2001,3,server.id,(uuid,server.id)) # self referencing
braidchainconfig = PeaceFounder.BraidChainConfig([(uuid,maintainer.id)],maintainer.id,2002,2003,2004,ballotconfig)
systemconfig = Community.SystemConfig(2001,2005,braidchainconfig)

Community.save(systemconfig) # Only necessary for testing.

task = @async Community.serve(server)

### The participatory element

# In one way or another members obtain a valid certificates for the community
certificates = []
for i in 1:3
    account = "account$i"
    member = PeaceVote.KeyChain(uuid,account)
    identification = PeaceVote.ID("$i","today",member.member.id)
    cert = PeaceVote.Certificate(identification,maintainer)
    push!(certificates,cert)
end

# Each member then participates.
@sync for (cert,i) in zip(certificates,1:3)
    @async begin
        account = "account$i"
        keychain = PeaceVote.KeyChain(uuid,account)
        PeaceVote.register(uuid,cert)
    end
end

@sync for i in 1:3
    @async begin
        account = "account$i"
        keychain = PeaceVote.KeyChain(uuid,account)
        PeaceVote.braid!(keychain)
    end
end

# Someone puts a proposal
pmember = PeaceVote.KeyChain(uuid,"account2")
PeaceVote.propose("Found peace for a change?",["yes","no","maybe"],pmember);

@sync for i in 1:3
    @async begin
        account = "account$i"
        keychain = PeaceVote.KeyChain(uuid,account)
        PeaceVote.braid!(keychain)
    end
end

pmember = PeaceVote.KeyChain(uuid,"account1")
PeaceVote.propose("Democracy for everyone?",["yes","no"],pmember);

sleep(1)

braidchain = PeaceVote.braidchain(uuid)
proposals = PeaceVote.proposals(braidchain)

@sync for i in 1:3
    @async begin
        account = "account$i"
        keychain = PeaceVote.KeyChain(uuid,account)

        option = PeaceVote.Option(proposals[1],rand(1:3))
        PeaceVote.vote(option,keychain,braidchain)

        option = PeaceVote.Option(proposals[2],rand(1:2))
        PeaceVote.vote(option,keychain,braidchain)
    end
end

### Methods for analyzing the braidchain.

sleep(1)

braidchain = PeaceVote.braidchain(uuid)
members = PeaceVote.members(braidchain)

for proposal in PeaceVote.proposals(braidchain)
    tally = PeaceVote.count(uuid,proposal,braidchain)
    @show proposal,tally
end


text:
include("run_scf_and_compare.jl")
include("testcases.jl")

function run_silicon_lda(T ;Ecut=5, grid_size=15, spin_polarization=:none, kwargs...)
    # These values were computed using ABINIT with the same kpoints as testcases.jl
    # and Ecut = 25
    ref_lda = [
        [-0.178566465714968, 0.261882541175914, 0.261882541178847, 0.261882541181782,
          0.354070367072414, 0.354070367076363, 0.354070367080310, 0.376871160884678],
        [-0.127794342370963, 0.064395861472044, 0.224958824747686, 0.224958824750934,
          0.321313617512188, 0.388442495007398, 0.388442495010722, 0.542078732298094],
        [-0.108449612789883, 0.077125812982728, 0.172380374761464, 0.172380374766260,
          0.283802499666810, 0.329872296009131, 0.525606867582028, 0.525606867585921],
        [-0.058089253154566, 0.012364292440522, 0.097350168867990, 0.183765652148129,
          0.314593174568090, 0.470869435132365, 0.496966579772700, 0.517009645871194],
    ]
    ref_etot = -7.911817522631488

    fft_size = fill(grid_size, 3)
    Si = ElementPsp(silicon.atnum, psp=load_psp(silicon.atnum, functional="lda", family="hgh"))
    model = model_DFT(Array{T}(silicon.lattice), [Si => silicon.positions], [:lda_x, :lda_c_vwn],
                      spin_polarization=spin_polarization)
    basis = PlaneWaveBasis(model, Ecut, silicon.kcoords, silicon.ksymops; fft_size=fft_size)

    spin_polarization == :collinear && (ref_lda = vcat(ref_lda, ref_lda))
    run_scf_and_compare(T, basis, ref_lda, ref_etot;
                        ρ=guess_density(basis),
                        kwargs...)
end

@testset "Silicon LDA (small, Float64)" begin
    run_silicon_lda(Float64, Ecut=7, test_tol=0.03, n_ignored=0, grid_size=17, scf_tol=1e-5,
                    n_ep_extra=0)
end

if !isdefined(Main, :FAST_TESTS) || !FAST_TESTS
    @testset "Silicon LDA (large, Float64)" begin
        run_silicon_lda(Float64, Ecut=25, test_tol=1e-5, n_ignored=0,
                        grid_size=33, scf_tol=1e-7, n_ep_extra=0)
    end
end

@testset "Silicon LDA (small, Float32)" begin
    run_silicon_lda(Float32, Ecut=7, test_tol=0.03, n_ignored=1, grid_size=19, scf_tol=1e-4,
                    n_ep_extra=1)
end

@testset "Silicon LDA (small, collinear spin)" begin
    run_silicon_lda(Float64, Ecut=7, test_tol=0.03, n_ignored=0, grid_size=17,
                    scf_tol=1e-5, n_ep_extra=0, spin_polarization=:collinear)
end

if !isdefined(Main, :FAST_TESTS) || !FAST_TESTS
    @testset "Silicon LDA (large, collinear spin)" begin
        run_silicon_lda(Float64, Ecut=25, test_tol=1e-5, n_ignored=0,
                        grid_size=33, scf_tol=1e-7, n_ep_extra=0, spin_polarization=:collinear)
    end
end


text:
#=

    This is a port of the R2 model SurveyResultUnbias.cs

    With the datafiles
    - survey_result_unbias_population.csv
    - survey_result_unbias_gender.csv
    - survey_result_unbias_person_gender.csv
    - survey_result_unbias_answer.csv

    Note: there is only 5 observations in the answer file, but reading the files makes it look
          more general...

    Output of the R2 model:
    ```
    <0>[0] Mean: 0.35699               bias[0]
    <0>[0] Variance: 7.87941e-006

    <0>[1] Mean: 0.573151              bias[1]
    <0>[1] Variance: 4.22538e-005

    <1>[0] Mean: 1121.3                mean[0]
    <1>[0] Variance: 77.7373

    <1>[1] Mean: 1557.82               mean[1]
    <1>[1] Variance: 312.151

    <2>[0] Mean: 720.986               variance[0]
    <2>[0] Variance: 6.41627

    <2>[1] Mean: 664.841               variance[1]
    <2>[1] Variance: 7.56368

    <3>[0] Mean: 1125.37               votes[0]
    <3>[0] Variance: 28.9169

    <3>[1] Mean: 1570.39               votes[1]
    <3>[1] Variance: 53.9271
    ```

    The output of this model is a little different than the R2 output:
    Summary Statistics
    parameters        mean       std   naive_se      mcse       ess      rhat   ess_per_sec
        Symbol     Float64   Float64    Float64   Float64   Float64   Float64       Float64

         bias[1]      0.8220    0.0000     0.0000    0.0000    3.2352    1.4114        0.2795
         bias[2]      0.3806    0.0000     0.0000    0.0000    5.6540    1.0970        0.4885
        votes[1]   2582.0475    0.0000     0.0000    0.0000    3.2164    1.4135        0.2779
        votes[2]   1034.6060    0.0000     0.0000    0.0000    5.6039    1.0987        0.4842

    pop_mean, pop_variance for gender 0 and 1:
    2-element Vector{Vector{Float64}}:
    [2582.0475011949065, 1034.6060354928363]
    [4.736500784568854e-6, 2.4354395705793703e-6]

=#

using Turing, StatsPlots, Distributions, StatsBase
using CSV
include("jl_utils.jl")

@model function survey_result_unbias(population, person_gender, answer)
    pop_len = length(population)
    answer_len = length(answer)

    bias ~ filldist(Beta(1,1),pop_len)

    pop_mean = population .*  bias
    pop_variance = (pop_mean .- pop_mean .* bias)
    votes ~ MvNormal(pop_mean, pop_variance)

    ansBias =  Vector{Real}(undef, answer_len)
    for i in 1:answer_len
        ansBias[i] = bias[person_gender[i]+1] # gender is 0 and 1 in the person_gender file. Hence the + 1.
        answer[i] ~ Bernoulli(ansBias[i])
    end

    return [pop_mean, pop_variance]
end

# 3141 2718
population = parse.(Int64,split(readline("survey_result_unbias_population.csv"),","))
# 0 1
# gender = parse.(Int64,split(readline("survey_result_unbias_gender.csv"),","))
# 0 1 0 0 1
person_gender = parse.(Int64,split(readline("survey_result_unbias_person_gender.csv"),","))
# 1 0 1 1 1
answer = parse.(Int64,split(readline("survey_result_unbias_answer.csv"),","))

model = survey_result_unbias(population, person_gender, answer)

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
# chns = sample(model, PG(5), 10_000)
# chns = sample(model, SMC(), 1_000)
# chns = sample(model, IS(), 1_000)

chns = sample(model, NUTS(), 1_000)
# chns = sample(model, HMC(0.1,6), 10_000)

display(chns)

println("pop_mean, pop_variance for gender 0 and 1:")
chains_params = Turing.MCMCChains.get_sections(chns, :parameters)
genq = generated_quantities(model, chains_params)
display(mean(genq))
# show_var_dist_pct(genq)

max_stars_repo_name: tias/hakank

max_issues_repo_name: tias/hakank

max_forks_repo_name: tias/hakank


text:
import PencilFFTs

using Oceananigans.Solvers: copy_real_component!
using Oceananigans.Distributed: rank2index

import Oceananigans.Solvers: poisson_eigenvalues, solve!
import Oceananigans.Architectures: architecture

struct DistributedFFTBasedPoissonSolver{P, F, L, λ, S}
              plan :: P
       global_grid :: F
        local_grid :: L
       eigenvalues :: λ
           storage :: S
end

architecture(solver::DistributedFFTBasedPoissonSolver) =
    architecture(solver.global_grid)

function DistributedFFTBasedPoissonSolver(global_grid, local_grid)

    arch = architecture(local_grid)
    topo = (TX, TY, TZ) = topology(global_grid)

    λx = poisson_eigenvalues(global_grid.Nx, global_grid.Lx, 1, TX())
    λy = poisson_eigenvalues(global_grid.Ny, global_grid.Ly, 2, TY())
    λz = poisson_eigenvalues(global_grid.Nz, global_grid.Lz, 3, TZ())

    arch.ranks[1] == arch.ranks[3] == 1 || @warn "Must have Rx == Rz == 1 for distributed fft solver"

    Rx, Ry, Rz = arch.ranks

    # PencilFFT performs a permutation y -> x.
    # x will be the "distributed direction" when  s = b / (λx + λy + λz)
    # we have to permute (Rx, Ry, Rz) with (Ry, Rx, Rz)
    I, J, K = rank2index(arch.local_rank, Ry, Rx, Rz)

    perm_Nx = global_grid.Nx ÷ Ry

    i₁ = (I-1) * perm_Nx + 1
    i₂ = I * perm_Nx
    λx = λx[i₁:i₂, :, :]

    eigenvalues = (; λx, λy, λz)

    transform = PencilFFTs.Transforms.FFT!()
    proc_dims = (arch.ranks[2], arch.ranks[3])
    plan = PencilFFTs.PencilFFTPlan(size(global_grid), transform, proc_dims, MPI.COMM_WORLD)
    storage = PencilFFTs.allocate_input(plan)

    return DistributedFFTBasedPoissonSolver(plan, global_grid, local_grid, eigenvalues, storage)
end

function solve!(x, solver::DistributedFFTBasedPoissonSolver)
    arch = architecture(solver)
    λx, λy, λz = solver.eigenvalues

    # Apply forward transforms.
    solver.plan * solver.storage

    # Solve the discrete Poisson equation, storing the solution
    # temporarily in xc and later extracting the real part into
    # the solution, x.
    xc = b = solver.storage[2]
    @. xc = - b / (λx + λy + λz)

    # Setting DC component of the solution (the mean) to be zero. This is also
    # necessary because the source term to the Poisson equation has zero mean
    # and so the DC component comes out to be ∞.
    if MPI.Comm_rank(MPI.COMM_WORLD) == 0
        xc[1, 1, 1] = 0
    end

    # Apply backward transforms.
    solver.plan \ solver.storage
    xc_transposed = first(solver.storage)

    copy_event = launch!(arch, solver.local_grid, :xyz, copy_real_component!, x, xc_transposed, dependencies=device_event(arch))
    wait(device(arch), copy_event)

    return x
end


text:
module Network_ADER
# Include Polynomial reconstruction stuff
include("compute_constants/polynomial_reconstruction.jl")

end


text:
# Export Knet functionalities
using Knet

Parameter = Knet.param

export @diff, grad, value, update!, params  # gradient
export relu, sigm, invx, elu, selu          # activation
export KnetArray                            # datatype


text:
# Configuration Parameters
"""
    sdad_setup!(;db_user::AbstractString = "",
                 db_pwd::AbstractString = "",
                 github_login::AbstractString = "",
                 github_token::AbstractString = "",
                 inserver::Bool = false)
This function writes the configurations to `confs/config.simple`
# Arguments
- `db_user`: Your username for the `postgis_1` [database](http://sdad.policy-analytics.net:8080/?pgsql=postgis_1&db=oss&ns=universe) (i.e., your UVA computing ID).
- `db_pwd`: The password for the `postgis_1` [database](http://sdad.policy-analytics.net:8080/?pgsql=postgis_1&db=oss&ns=universe).
- `github_login`: Your GitHub login (handle).
- `github_token`: A 40 alphanumeric characters string. Obtain a GitHub personal access token [here](https://help.github.com/en/articles/creating-a-personal-access-token-for-the-command-line) (only select read-access is required).
- `inserver`: Whether the project is being used within the server (container) or not.
# Examples
## Good!
```
julia> sdad_setup!(db_user = "jbs3hp",
                   db_pwd = "MyVerySafePwd",
                   github_login = "Nosferican",
                   github_token = "0ipg0jvonteb54lv7j6cbgwn2snq3d3ac1pthxvz",
                   inserver = true)
😃
```
## Incomplete!
```
julia> sdad_setup!(db_user = "jbs3hp",
                   db_pwd = "MyVerySafePwd",
                   github_login = "Nosferican")
Warning: github_login has not been defined.
😞
```
# Updating the configuration file.
```
julia> using GitHubAPI
julia> sdad_setup!(db_user = "jbs3hp",
                   db_pwd = "MyVerySafePwd",
                   github_login = "Nosferican",
                   github_token = "0ipg0jvonteb54lv7j6cbgwn2snq3d3ac1pthxvz",
                   inserver = true)
julia> exit()
> julia
julia> using GitHubAPI # config up-to-date
```
!!! note
    Updating the configuration file requires a restart of the session.
"""
function sdad_setup!(;db_user::AbstractString = "",
                      db_pwd::AbstractString = "",
                      github_login::AbstractString = "",
                      github_token::AbstractString = "",
                      inserver::Bool = false)
    args = ["db_user", "db_pwd", "github_login", "github_token", "inserver"]
    isdir(joinpath(dirname(@__FILE__), "..", "confs")) || mkdir(joinpath(dirname(@__FILE__), "..", "confs"))
    isfile(joinpath(dirname(@__FILE__), "..", "confs", "config.simple")) ||
        touch(joinpath(dirname(@__FILE__), "..", "confs", "config.simple"))
    conf = ConfParse(joinpath(dirname(@__FILE__), "..", "confs", "config.simple"),
                     "simple")
    parse_conf!(conf)
    for (key, val) ∈ zip(args, [db_user, db_pwd, github_login, github_token, inserver])
        isempty(val) || commit!(conf, key, val)
    end
    save!(conf)
    notdefined = filter(key -> !haskey(conf, key),
                        args)
    foreach(key -> @warn("$key has not been defined."),
            notdefined)
    if isinteractive()
        if isempty(notdefined)
            println("😃")
        else
            println("😞")
        end
    end
end
isfile(joinpath(dirname(@__FILE__), "..", "confs", "config.simple")) ||
    sdad_setup!(db_user = get(ENV, "db_user", ""),
                db_pwd = get(ENV, "db_pwd", ""),
                github_login = get(ENV, "github_login", ""),
                github_token = get(ENV, "github_token", ""),
                inserver = get(ENV, "inserver", false))
const conf = ConfParse(joinpath(dirname(@__FILE__), "..", "confs", "config.simple"),
                       "simple");
parse_conf!(conf);
"""
    db_user::String
The username for the [database](http://sdad.policy-analytics.net:8080/?pgsql=postgis_1&db=oss&ns=universe) (i.e., your UVA computing ID).
"""
const db_user = haskey(conf, "db_user") ? retrieve(conf, "db_user") : "";
"""
    db_pwd::String
The password for the [database](http://sdad.policy-analytics.net:8080/?pgsql=postgis_1&db=oss&ns=universe).
"""
const db_pwd = haskey(conf, "db_pwd") ? retrieve(conf, "db_pwd") : "";
"""
    github_login::String
Your GitHub handle.
"""
const github_login = haskey(conf, "github_login") ? retrieve(conf, "github_login") : "";
"""
    github_header = Dict("User-Agent" => github_login)
Header for the GitHub API.
"""
const github_header = Dict("User-Agent" => github_login);
const github_token = haskey(conf, "github_token") ? retrieve(conf, "github_token") : "";
"""
    db_host = "postgis_1" | "sdad.policy-analytics.net"
Host for the [database](http://sdad.policy-analytics.net:8080/?pgsql=postgis_1&db=oss&ns=universe).
"""
const db_host = (haskey(conf, "inserver") &&
    retrieve(conf, "inserver") == "true") ? "postgis_1" : "sdad.policy-analytics.net";
"""
    db_port = 5432 | 5434
Port for the `postgis_1` in the [database](http://sdad.policy-analytics.net:8080/?pgsql=postgis_1&db=oss&ns=universe).
"""
const db_port = (haskey(conf, "inserver") &&
    retrieve(conf, "inserver") == "true") ? 5432 : 5434;
"""
    dbname = "oss"
Database for the Open-Source Software project.
"""
const dbname = "oss";
"""
    response_dtf = dateformat"d u y H:M:S Z"
HTTP responses require this datetime format.
"""
const response_dtf = dateformat"d u y H:M:S Z";
"""
    postgis_dtf = dateformat"yyyy-mm-dd HH:MM:SSzzzz"
HTTP responses require this datetime format.
"""
const postgis_dtf = dateformat"yyyy-mm-dd HH:MM:SSzzzz";
"""
    github_dtf = "yyyy-mm-ddTHH:MM:SSzzzz"
GitHub zoned datetime format.
"""
const github_dtf = "yyyy-mm-ddTHH:MM:SSzzzz";
# until = today() |>
#     yearmonthday |>
#     (dt -> ZonedDateTime(dt..., localzone()))
#     (td -> ZonedDateTime(year(td), month(td), day(td), localzone()));
"""
    until::ZonedDateTime
Until when should the scrapper query data. Currently at `"2019-08-15T00:00:00-04:00"`.
"""
const until = ZonedDateTime("2019-08-15T00:00:00-04:00",
                            github_dtf)


text:
module MultBoxes

include("./../../constraints/geometric/constants.jl")
using JuMP
using .Constants

export cons_ds_mult_boxes

function cons_ds_mult_boxes(m)

    @expression(m, mb_sum_of_boxes_below_box_i[i=1:length(boxes)],
        sum([m[:if_on_top][i, j] for j in 1:length(boxes) if i!=j])
    )

    @expression(m, ds_mb_sum_of_boxes_below_obj[i=1:length(boxes)],
        length(boxes) - m[:mb_sum_of_boxes_below_box_i][i]
    )

    return m
end

end


text:
import Oceananigans: short_show
using Oceananigans.Grids: Face, Bounded, topology
using Oceananigans.Fields: location

"""
    struct FieldSlicer{I, J, K, W}

Slices fields along indices with or without halo regions as specified.
"""
struct FieldSlicer{I, J, K}
    i :: I
    j :: J
    k :: K
    with_halos :: Bool
end

"""
    FieldSlicer(i=Colon(), j=Colon(), k=Colon(), with_halos=false)

Returns `FieldSlicer` that slices a field prior to output or time-averaging.

The keyword arguments `i, j, k` prescribe an `Integer` index, `UnitRange`,
or `StepRange` of indices in `x, y, z`, respectively.

The default for `i`, `j`, and `k` is `Colon()` which indicates "all indices".

The keyword `with_halos` denotes whether halo data is saved or not. Halo regions are
sliced off output for `UnitRange`, `StepRange`, and `Colon` index specifications.
"""
FieldSlicer(; i=Colon(), j=Colon(), k=Colon(), with_halos=false) =
    FieldSlicer(num2int(i), num2int(j), num2int(k), with_halos)

# To support FieldSlicer construction with all number types (but we always want to
# index with an integer).
num2int(i) = i
num2int(i::Number) = Int(i)

#####
##### Slice of life, err... data
#####

# Integer slice
parent_slice_indices(loc, topo, N, H, i::Int, with_halos) = UnitRange(i + H, i + H)

# Colon slicing
parent_slice_indices(loc, topo, N, H,
                     ::Colon, with_halos) = with_halos ? UnitRange(1, N+2H) : UnitRange(H+1, H+N)

parent_slice_indices(::Type{Face}, ::Type{Bounded}, N, H,
                     ::Colon, with_halos) = with_halos ? UnitRange(1, N+1+2H) : UnitRange(H+1, H+N+1)

# Slicing along reduced dimensions
parent_slice_indices(::Type{Nothing}, loc, N, H, ::Colon, with_halos) = 1:1
parent_slice_indices(::Type{Nothing}, loc, N, H, ::AbstractRange, with_halos) = 1:1
parent_slice_indices(::Type{Nothing}, loc, N, H, ::Int, with_halos) = 1:1

# Safe slice ranges without halos
right_parent_index_without_halos(loc, topo, N, H, right) = min(H + N, right + H)
right_parent_index_without_halos(::Type{Face}, ::Type{Bounded}, N, H, right) = min(H + N + 1, right + H)

function parent_slice_index_endpoints(loc, topo, N, H, rng, with_halos)
    if with_halos
        left = rng[1] + H
        right = rng[end] + H
    else
        left = max(H, rng[1] + H)
        right = right_parent_index_without_halos(loc, topo, N, H, rng[end])
    end

    return left, right
end

function parent_slice_indices(loc, topo, N, H, rng::UnitRange, with_halos)
    left, right = parent_slice_index_endpoints(loc, topo, N, H, rng, with_halos)
    return UnitRange(left, right)
end

function parent_slice_indices(loc, topo, N, H, rng::StepRange, with_halos)
    left, right = parent_slice_index_endpoints(loc, topo, N, H, rng, with_halos)
    return StepRange(left, step(rng), right)
end

"""
    slice(slicer, field)

Returns a view over parent(field) associated with slice.i, slice.j, slice.k.
"""
function slice_parent(slicer, field)

    # Unpack
    Nx, Ny, Nz = field.grid.Nx, field.grid.Ny, field.grid.Nz
    Hx, Hy, Hz = field.grid.Hx, field.grid.Hy, field.grid.Hz
    Lx, Ly, Lz = location(field)
    Tx, Ty, Tz = topology(field)

    x_data_range = slicer.i
    y_data_range = slicer.j
    z_data_range = slicer.k

    # Convert slicer indices to parent indices, and managing halos
    x_parent_range = parent_slice_indices(Lx, Tx, Nx, Hx, x_data_range, slicer.with_halos)
    y_parent_range = parent_slice_indices(Ly, Ty, Ny, Hy, y_data_range, slicer.with_halos)
    z_parent_range = parent_slice_indices(Lz, Tz, Nz, Hz, z_data_range, slicer.with_halos)

    return field.data.parent[x_parent_range, y_parent_range, z_parent_range]
end

slice_parent(::Nothing, field) = parent(field)

show_index(i) = string(i)
show_index(i::Colon) = ":"

short_show(fs::FieldSlicer) = string("FieldSlicer(",
                                     "$(show_index(fs.i)), ",
                                     "$(show_index(fs.j)), ",
                                     "$(show_index(fs.k)), with_halos=$(fs.with_halos))")


text:
# https://github.com/JuliaLang/julia/blob/5e048f3e537387fd388a5360ce6163c8f7c61ccf/base/compiler/abstractinterpretation.jl#L31-L196

function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(f), argtypes::Vector{Any}, @nospecialize(atype), sv::InferenceState,
                                  max_methods::Int = InferenceParams(interp).MAX_METHODS)
    if sv.params.unoptimize_throw_blocks && sv.currpc in sv.throw_blocks
        return CallMeta(Any, false)
    end
    valid_worlds = WorldRange()
    atype_params = unwrap_unionall(atype).parameters
    splitunions = 1 < unionsplitcost(atype_params) <= InferenceParams(interp).MAX_UNION_SPLITTING
    mts = Core.MethodTable[]
    fullmatch = Bool[]
    if splitunions
        splitsigs = switchtupleunion(atype)
        applicable = Any[]
        infos = MethodMatchInfo[]
        for sig_n in splitsigs
            mt = ccall(:jl_method_table_for, Any, (Any,), sig_n)
            if mt === nothing
                add_remark!(interp, sv, "Could not identify method table for call")
                return CallMeta(Any, false)
            end
            mt = mt::Core.MethodTable
            matches = findall(sig_n, method_table(interp); limit=max_methods)
            if matches === missing
                add_remark!(interp, sv, "For one of the union split cases, too many methods matched")
                return CallMeta(Any, false)
            end
            push!(infos, MethodMatchInfo(matches))
            append!(applicable, matches)
            valid_worlds = intersect(valid_worlds, matches.valid_worlds)
            thisfullmatch = _any(match->(match::MethodMatch).fully_covers, matches)
            found = false
            for (i, mt′) in enumerate(mts)
                if mt′ === mt
                    fullmatch[i] &= thisfullmatch
                    found = true
                    break
                end
            end
            if !found
                push!(mts, mt)
                push!(fullmatch, thisfullmatch)
            end
        end
        info = UnionSplitInfo(infos)
    else
        mt = ccall(:jl_method_table_for, Any, (Any,), atype)
        if mt === nothing
            add_remark!(interp, sv, "Could not identify method table for call")
            return CallMeta(Any, false)
        end
        mt = mt::Core.MethodTable
        matches = findall(atype, method_table(interp, sv); limit=max_methods)
        if matches === missing
            # this means too many methods matched
            # (assume this will always be true, so we don't compute / update valid age in this case)
            add_remark!(interp, sv, "Too many methods matched")
            return CallMeta(Any, false)
        end
        push!(mts, mt)
        push!(fullmatch, _any(match->(match::MethodMatch).fully_covers, matches))
        info = MethodMatchInfo(matches)
        applicable = matches.matches
        valid_worlds = matches.valid_worlds
    end
    update_valid_age!(sv, valid_worlds)
    applicable = applicable::Array{Any,1}
    napplicable = length(applicable)
    rettype = Bottom
    edgecycle = false
    edges = Any[]
    nonbot = 0  # the index of the only non-Bottom inference result if > 0
    seen = 0    # number of signatures actually inferred
    istoplevel = sv.linfo.def isa Module
    multiple_matches = napplicable > 1

    if f !== nothing && napplicable == 1 && is_method_pure(applicable[1]::MethodMatch)
        val = pure_eval_call(f, argtypes)
        if val !== false
            # TODO: add some sort of edge(s)
            return CallMeta(val, MethodResultPure())
        end
    end

    for i in 1:napplicable
        match = applicable[i]::MethodMatch
        method = match.method
        sig = match.spec_types
        if istoplevel && !isdispatchtuple(sig)
            # only infer concrete call sites in top-level expressions
            add_remark!(interp, sv, "Refusing to infer non-concrete call site in top-level expression")
            rettype = Any
            break
        end
        sigtuple = unwrap_unionall(sig)::DataType
        splitunions = false
        this_rt = Bottom
        # TODO: splitunions = 1 < unionsplitcost(sigtuple.parameters) * napplicable <= InferenceParams(interp).MAX_UNION_SPLITTING
        # currently this triggers a bug in inference recursion detection
        if splitunions
            splitsigs = switchtupleunion(sig)
            for sig_n in splitsigs
                rt, edgecycle1, edge = abstract_call_method(interp, method, sig_n, svec(), multiple_matches, sv)
                if edge !== nothing
                    push!(edges, edge)
                end
                edgecycle |= edgecycle1::Bool
                this_rt = tmerge(this_rt, rt)
                this_rt === Any && break
            end
        else
            this_rt, edgecycle1, edge = abstract_call_method(interp, method, sig, match.sparams, multiple_matches, sv)
            edgecycle |= edgecycle1::Bool
            if edge !== nothing
                push!(edges, edge)
            end
        end
        if this_rt !== Bottom
            if nonbot === 0
                nonbot = i
            else
                nonbot = -1
            end
        end
        seen += 1
        rettype = tmerge(rettype, this_rt)
        rettype === Any && break
    end
    # try constant propagation if only 1 method is inferred to non-Bottom
    # this is in preparation for inlining, or improving the return result
    is_unused = call_result_unused(sv)
    if nonbot > 0 && seen == napplicable && (!edgecycle || !is_unused) &&
            is_improvable(rettype) && InferenceParams(interp).ipo_constant_propagation
        # if there's a possibility we could constant-propagate a better result
        # (hopefully without doing too much work), try to do that now
        # TODO: it feels like this could be better integrated into abstract_call_method / typeinf_edge
        const_rettype = abstract_call_method_with_const_args(interp, rettype, f, argtypes, applicable[nonbot]::MethodMatch, sv, edgecycle)
        if const_rettype ⊑ rettype
            # use the better result, if it's a refinement of rettype
            rettype = const_rettype
        end
    end
    if is_unused && !(rettype === Bottom)
        add_remark!(interp, sv, "Call result type was widened because the return value is unused")
        # We're mainly only here because the optimizer might want this code,
        # but we ourselves locally don't typically care about it locally
        # (beyond checking if it always throws).
        # So avoid adding an edge, since we don't want to bother attempting
        # to improve our result even if it does change (to always throw),
        # and avoid keeping track of a more complex result type.
        rettype = Any
    end
    if !(rettype === Any) # adding a new method couldn't refine (widen) this type
        for edge in edges
            add_backedge!(edge::MethodInstance, sv)
        end
        for (thisfullmatch, mt) in zip(fullmatch, mts)
            if !thisfullmatch
                # also need an edge to the method table in case something gets
                # added that did not intersect with any existing method
                add_mt_backedge!(mt, atype, sv)
            end
        end
    end
    #print("=> ", rettype, "\n")
    return CallMeta(rettype, info)
end


text:
using Reproject
using Test
using PythonCall
using FITSIO, WCS
using SHA: sha256

rp = pyimport("reproject")
astropy = pyimport("astropy")

include("parsers.jl")
include("utils.jl")
include("core.jl")


text:
# Training a normalizing flow for 2D Rosenbrock example
# Authors: Ali Siahkoohi, alisk@gatech.edu
#          Gabrio Rizzuti, rizzuti.gabrio@gatech.edu
#          Philipp Witte, pwitte3@gatech.edu
# Date: September 2020
"""
Example:
julia train_warm-start_unsupervised_hint.jl
"""

using DrWatson
@quickactivate :FastApproximateInference

using InvertibleNetworks
using Random
using Statistics
using ArgParse
using BSON
using Printf
using Flux: gpu, cpu
using Flux.Data: DataLoader
using Flux.Optimise: Optimiser, ExpDecay, update!, ADAM

# Random seed
Random.seed!(19)

s = ArgParseSettings()
@add_arg_table! s begin
    "--max_epoch"
        help = "Maximum # of epochs"
        arg_type = Int
        default = 5
    "--lr"
        help = "starting learning rate"
        arg_type = Float32
        default = 1f-3
    "--lr_step"
        help = "lr scheduler applied at evey lr_step epoch"
        arg_type = Int
        default = 12
    "--batchsize"
        help = "batch size"
        arg_type = Int
        default = 64
    "--n_hidden"
        help = "# hidden channels"
        arg_type = Int
        default = 64
    "--depth"
        help = "depth of the network"
        arg_type = Int
        default = 8
    "--sigma"
        help = "noise std"
        arg_type = Float32
        default = 4f-1
    "--sim_name"
        help = "simulation name"
        arg_type = String
        default = "warm-start-unsupervised-hint-rosenbrock"
end
parsed_args = parse_args(s)

max_epoch = parsed_args["max_epoch"]
lr = parsed_args["lr"]
lr_step = parsed_args["lr_step"]
batchsize = parsed_args["batchsize"]
n_hidden = parsed_args["n_hidden"]
depth = parsed_args["depth"]
sigma = parsed_args["sigma"]
sim_name = parsed_args["sim_name"]

# Define network
nx = 1
ny = 1
n_in = 2

CH = NetworkConditionalHINT(n_in, n_hidden, depth)

# Loading the experiment—only network weights and training loss
pretrained_args = Dict(
    "max_epoch" => 25,
    "lr"        => 1f-3,
    "lr_step"   => 1,
    "batchsize" => 64,
    "n_hidden"  => 64,
    "depth"     => 8,
    "sigma"     => 4f-1,
    "sim_name"  => "supervised-hint-rosenbrock"
)
Params, fval = load_experiment(pretrained_args)
put_params!(CH, Params)

# Prior distribution
RB_dist = RosenbrockDistribution(0f0, 5f-1, ones(Float32, 2, 1))

# Forward operator
A = SquareCSop(n_in; s=3f0)

# Observed data
X_true = zeros(Float32, 1, 1, 2, 1)
X_true[1, 1, 1, 1] = 2f0
X_true[1, 1, 2, 1] = 4f0
Y_obs = A*X_true + sigma*randn(Float32, nx, ny, n_in, 1)

# Corresponding data latent variable
Zy_obs = CH.forward_Y(Y_obs)

# reversed network
CHrev = reverse(CH)

# Training
# Batch extractor
ntrain = 1000
Z_train = randn(Float32, nx, ny, n_in, ntrain)

num_batches = cld(ntrain, batchsize)
train_loader = DataLoader(Z_train, batchsize=batchsize, shuffle=true)

# Optimizer
opt = Optimiser(ExpDecay(lr, .9f0, num_batches*lr_step, 0f0), ADAM(lr))

# Training log keeper
fval = zeros(Float32, num_batches*max_epoch)

for j=1:max_epoch
    for (itr, Zx) in enumerate(train_loader)

        fval[(j-1)*num_batches + itr] = loss_unsupervised(
            CHrev, Y_obs, Zx, Zy_obs, A, sigma, RB_dist
        )

        @printf(
            "Epoch: [%d/%d] | Itr: [%d/%d] | Objective: %4.2f\r",
            j, max_epoch, itr, num_batches, fval[(j-1)*num_batches + itr]
        )

        # Update params
        for p in get_params(CHrev)
            update!(opt, p.data, p.grad)
        end
        clear_grad!(CHrev)
    end
end

# Saving parameters and logs
CHrev = CHrev |> cpu
Params = get_params(CHrev)
save_dict = @strdict max_epoch lr lr_step batchsize n_hidden depth sigma sim_name Y_obs Zy_obs X_true Z_train A.A pretrained_args Params fval
@tagsave(datadir(sim_name, savename(save_dict, "bson")), save_dict; safe=true)


text:
#=
JuMP helpers for constructing extended formulations for vector and matrix
spectral/eigenvalue cones
=#

abstract type SpectralExtender end

# vector spectral formulations
abstract type VecSpecExt <: SpectralExtender end

# vector negative geomean formulations
struct VecNegRtdet <: VecSpecExt end
struct VecNegRtdetEFExp <: VecSpecExt end
struct VecNegRtdetEFPow <: VecSpecExt end

const VecNegRtdetAll = Union{VecNegRtdet, VecNegRtdetEFExp, VecNegRtdetEFPow}

# vector separable spectral primal formulations
struct VecLogCone <: VecSpecExt end
struct VecNegLog <: VecSpecExt end
struct VecNegLogEF <: VecSpecExt end
struct VecNegEntropy <: VecSpecExt end
struct VecNegEntropyEF <: VecSpecExt end
struct VecNegSqrt <: VecSpecExt end
struct VecNegSqrtEF <: VecSpecExt end
struct VecNegPower01 <: VecSpecExt p::Real end
struct VecNegPower01EF <: VecSpecExt p::Real end
struct VecPower12 <: VecSpecExt p::Real end
struct VecPower12EF <: VecSpecExt p::Real end

const VecSepSpecPrim = Union{VecLogCone, VecNegLog, VecNegEntropy, VecNegSqrt,
    VecNegPower01, VecPower12}
const VecSepSpecPrimEF = Union{VecNegLogEF, VecNegEntropyEF, VecNegSqrtEF,
    VecNegPower01EF, VecPower12EF}
const VecSepSpecPrimAll = Union{VecSepSpecPrim, VecSepSpecPrimEF}

# vector separable spectral dual formulations
struct VecNegEntropyConj <: VecSpecExt end
struct VecNegEntropyConjEF <: VecSpecExt end
struct VecNegSqrtConj <: VecSpecExt end
struct VecNegSqrtConjEF <: VecSpecExt end
struct VecNegPower01Conj <: VecSpecExt p::Real end
struct VecNegPower01ConjEF <: VecSpecExt p::Real end
struct VecPower12Conj <: VecSpecExt p::Real end
struct VecPower12ConjEF <: VecSpecExt p::Real end

const VecSepSpecDual = Union{VecNegEntropyConj, VecNegSqrtConj,
    VecNegPower01Conj, VecPower12Conj}
const VecSepSpecDualEF = Union{VecNegEntropyConjEF, VecNegSqrtConjEF,
    VecNegPower01ConjEF, VecPower12ConjEF}
const VecSepSpecDualAll = Union{VecSepSpecDual, VecSepSpecDualEF}

const VecSepSpec = Union{VecSepSpecPrim, VecSepSpecDual}
const VecSepSpecEF = Union{VecSepSpecPrimEF, VecSepSpecDualEF}

# matrix spectral formulations
abstract type MatSpecExt <: SpectralExtender end

# matrix negative geomean formulations
struct MatNegRtdet <: MatSpecExt end
struct MatNegRtdetEFExp <: MatSpecExt end
struct MatNegRtdetEFPow <: MatSpecExt end

const MatNegRtdetAll = Union{MatNegRtdet, MatNegRtdetEFExp, MatNegRtdetEFPow}

# matrix separable spectral primal formulations
struct MatLogdetCone <: MatSpecExt end
struct MatNegLog <: MatSpecExt end
struct MatNegLogEigOrd <: MatSpecExt end
struct MatNegLogDirect <: MatSpecExt end
struct MatNegEntropy <: MatSpecExt end
struct MatNegEntropyEigOrd <: MatSpecExt end
struct MatNegSqrt <: MatSpecExt end
struct MatNegSqrtEigOrd <: MatSpecExt end
struct MatNegPower01 <: MatSpecExt p::Real end
struct MatNegPower01EigOrd <: MatSpecExt p::Real end
struct MatPower12 <: MatSpecExt p::Real end
struct MatPower12EigOrd <: MatSpecExt p::Real end

const MatSepSpecPrim = Union{MatLogdetCone, MatNegLog, MatNegEntropy,
    MatNegSqrt, MatNegPower01, MatPower12}
const MatSepSpecPrimEF = Union{MatNegLogEigOrd, MatNegLogDirect,
    MatNegEntropyEigOrd, MatNegSqrtEigOrd, MatNegPower01EigOrd,
    MatPower12EigOrd}
const MatSepSpecPrimAll = Union{MatSepSpecPrim, MatSepSpecPrimEF}

# matrix separable spectral dual formulations
struct MatNegEntropyConj <: MatSpecExt end
struct MatNegEntropyConjEigOrd <: MatSpecExt end
struct MatNegSqrtConj <: MatSpecExt end
struct MatNegSqrtConjEigOrd <: MatSpecExt end
struct MatNegSqrtConjDirect <: MatSpecExt end
struct MatNegPower01Conj <: MatSpecExt p::Real end
struct MatNegPower01ConjEigOrd <: MatSpecExt p::Real end
struct MatPower12Conj <: MatSpecExt p::Real end
struct MatPower12ConjEigOrd <: MatSpecExt p::Real end

const MatSepSpecDual = Union{MatNegEntropyConj, MatNegSqrtConj,
    MatNegPower01Conj, MatPower12Conj}
const MatSepSpecDualEF = Union{MatNegEntropyConjEigOrd, MatNegSqrtConjEigOrd,
    MatNegSqrtConjDirect, MatNegPower01ConjEigOrd, MatPower12ConjEigOrd}
const MatSepSpecDualAll = Union{MatSepSpecDual, MatSepSpecDualEF}

const MatSepSpec = Union{MatSepSpecPrim, MatSepSpecDual}
const MatSepSpecEigOrd = Union{MatNegLogEigOrd, MatNegEntropyEigOrd,
    MatNegSqrtEigOrd, MatNegPower01EigOrd, MatPower12EigOrd,
    MatNegEntropyConjEigOrd, MatNegSqrtConjEigOrd, MatNegPower01ConjEigOrd,
    MatPower12ConjEigOrd}

const NegRtdetAll = Union{VecNegRtdetAll, MatNegRtdetAll}
const SepSpecPrimAll = Union{VecSepSpecPrimAll, MatSepSpecPrimAll}
const SepSpecDualAll = Union{VecSepSpecDualAll, MatSepSpecDualAll}
const SepSpecAll = Union{SepSpecPrimAll, SepSpecDualAll}

# helpers
get_vec_ef(::MatNegRtdetEFExp) = VecNegRtdetEFExp()
get_vec_ef(::MatNegRtdetEFPow) = VecNegRtdetEFPow()
get_vec_ef(::MatNegLogEigOrd) = VecNegLogEF()
get_vec_ef(::MatNegEntropyEigOrd) = VecNegEntropyEF()
get_vec_ef(::MatNegSqrtEigOrd) = VecNegSqrtEF()
get_vec_ef(ext::MatNegPower01EigOrd) = VecNegPower01EF(ext.p)
get_vec_ef(ext::MatPower12EigOrd) = VecPower12EF(ext.p)
get_vec_ef(::MatNegEntropyConjEigOrd) = VecNegEntropyConjEF()
get_vec_ef(::MatNegSqrtConjEigOrd) = VecNegSqrtConjEF()
get_vec_ef(ext::MatNegPower01ConjEigOrd) = VecNegPower01ConjEF(ext.p)
get_vec_ef(ext::MatPower12ConjEigOrd) = VecPower12ConjEF(ext.p)

get_ssf(::Union{VecLogCone, VecNegLog, VecNegLogEF, MatLogdetCone, MatNegLog,
    MatNegLogEigOrd, MatNegLogDirect}) =
    Cones.NegLogSSF()
get_ssf(::Union{VecNegEntropy, VecNegEntropyEF, VecNegEntropyConj,
    VecNegEntropyConjEF, MatNegEntropy, MatNegEntropyEigOrd, MatNegEntropyConj,
    MatNegEntropyConjEigOrd}) =
    Cones.NegEntropySSF()
get_ssf(::Union{VecNegSqrt, VecNegSqrtEF, VecNegSqrtConj, VecNegSqrtConjEF,
    MatNegSqrt, MatNegSqrtEigOrd, MatNegSqrtConj, MatNegSqrtConjEigOrd,
    MatNegSqrtConjDirect}) =
    Cones.NegSqrtSSF()
get_ssf(ext::Union{VecNegPower01, VecNegPower01EF, VecNegPower01Conj,
    VecNegPower01ConjEF, MatNegPower01, MatNegPower01EigOrd, MatNegPower01Conj,
    MatNegPower01ConjEigOrd}) =
    Cones.NegPower01SSF(ext.p)
get_ssf(ext::Union{VecPower12, VecPower12EF, VecPower12Conj, VecPower12ConjEF,
    MatPower12, MatPower12EigOrd, MatPower12Conj, MatPower12ConjEigOrd}) =
    Cones.Power12SSF(ext.p)

nat_name(::NegRtdetAll) = "NegRtdet"
get_name_ssf(ext::SepSpecAll) = string(nameof(typeof(get_ssf(ext))))
nat_name(ext::SepSpecPrimAll) = get_name_ssf(ext)
nat_name(ext::SepSpecDualAll) = get_name_ssf(ext) * "Conj"

is_domain_pos(::SpectralExtender) = true
is_domain_pos(ext::SepSpecDualAll) = Cones.h_conj_dom_pos(get_ssf(ext))

get_val(x::Vector, ::NegRtdetAll) = -exp(sum(log, x) / length(x))
get_val(x::Vector, ext::SepSpecPrimAll) = Cones.h_val(x, get_ssf(ext))
get_val(x::Vector, ext::SepSpecDualAll) = Cones.h_conj(x, get_ssf(ext))

pos_only(x::Vector{T}, minval::T = eps(T)) where {T <: Real} =
    [(x_i < minval ? minval : x_i) for x_i in x]

#=
homogenizes separable spectral vector/matrix constraints
=#
function add_homog_spectral(
    ext::SepSpecAll,
    d::Int,
    aff::Vector{JuMP.AffExpr},
    model::JuMP.Model,
    )
    # epigraph and perspective variables are swapped if dual cone is used
    aff_new = vcat(swap_if_dual(aff[1], 1.0, ext)..., aff[2:end])
    add_spectral(ext, d, aff_new, model)
    return
end

#=
VecNegRtdet
(u, w > 0) : u > -(Πᵢ wᵢ)^(1/d)
=#
function add_homog_spectral(
    ::VecNegRtdet,
    d::Int,
    aff::Vector{JuMP.AffExpr},
    model::JuMP.Model,
    )
    @assert 1 + d == length(aff)
    aff_new = vcat(-aff[1], aff[2:end])
    JuMP.@constraint(model, aff_new in Hypatia.HypoGeoMeanCone{Float64}(1 + d))
    return
end

#=
VecNegRtdetEFExp
negative geometric mean -> exponential cone EF
(u, w > 0) : u > -(Πᵢ wᵢ)^(1/d)
↔ ∃ y > 0, x ∈ ℝᵈ : e'x < 0, (-xᵢ, y - u, wᵢ) ∈ MOI.ExponentialCone, ∀i
=#
function add_homog_spectral(
    ::VecNegRtdetEFExp,
    d::Int,
    aff::Vector{JuMP.AffExpr},
    model::JuMP.Model,
    )
    @assert 1 + d == length(aff)

    y = JuMP.@variable(model)
    JuMP.@constraint(model, y >= 0)
    x = JuMP.@variable(model, [1:d])
    JuMP.@constraint(model, sum(x) <= 0)

    for i in 1:d
        aff_new = vcat(-x[i], y - aff[1], aff[1 + i])
        JuMP.@constraint(model, aff_new in MOI.ExponentialCone())
    end

    return
end

#=
VecNegRtdetEFPow
negative geometric mean -> power cone EF
(u, w > 0) : u > -(Πᵢ wᵢ)^(1/d)
↔ ∃ y > 0, x ∈ ℝᵈ⁻² : (w₁, w₂, x₁) ∈ MOI.PowerCone(1/2),
(w[i], x[i-2], x[i-1]) ∈ MOI.PowerCone(1/i), ∀i = 3..d-1
(w[d], x[d-2], y - u) ∈ MOI.PowerCone(1/d)
=#
function add_homog_spectral(
    ::VecNegRtdetEFPow,
    d::Int,
    aff::Vector{JuMP.AffExpr},
    model::JuMP.Model,
    )
    @assert d > 2 # EF does not work for d = 2
    @assert 1 + d == length(aff)
    w = aff[2:end]

    y = JuMP.@variable(model)
    JuMP.@constraint(model, y >= 0)
    x = JuMP.@variable(model, [1:(d - 2)])

    JuMP.@constraint(model, vcat(w[1], w[2], x[1]) in MOI.PowerCone(1.0 / 2))
    for i in 3:(d - 1)
        aff_new_i = vcat(w[i], x[i - 2], x[i - 1])
        JuMP.@constraint(model, aff_new_i in MOI.PowerCone(1.0 / i))
    end
    aff_new = vcat(w[end], x[end], y - aff[1])
    JuMP.@constraint(model, aff_new in MOI.PowerCone(1.0 / d))

    return
end

#=
VecLogCone
=#
function add_spectral(
    ::VecLogCone,
    d::Int,
    aff::Vector{JuMP.AffExpr},
    model::JuMP.Model,
    )
    @assert 2 + d == length(aff)
    aff_new = vcat(-aff[1], aff[2], aff[3:end])
    JuMP.@constraint(model, aff_new in
        Hypatia.HypoPerLogCone{Float64}(length(aff)))
    return
end

#=
VecSepSpec
=#
function add_spectral(
    ext::VecSepSpec,
    d::Int,
    aff::Vector{JuMP.AffExpr},
    model::JuMP.Model,
    )
    @assert 2 + d == length(aff)
    is_dual = (ext isa VecSepSpecDualAll)
    cone = Hypatia.EpiPerSepSpectralCone{Float64}(get_ssf(ext),
        Cones.VectorCSqr{Float64}, d, is_dual)
    JuMP.@constraint(model, aff in cone)
    return
end

#=
VecSepSpecEF
=#
function add_spectral(
    ext::VecSepSpecEF,
    d::Int,
    aff::Vector{JuMP.AffExpr},
    model::JuMP.Model,
    )
    @assert 2 + d == length(aff)
    # epigraph and perspective variables are swapped if dual cone is used
    (epi, per) = swap_if_dual(aff[1], aff[2], ext)

    # linear constraint
    x = JuMP.@variable(model, [1:d])
    JuMP.@constraint(model, epi >= sum(x))

    # 3-dim constraints
    for i in 1:d
        extend_atom(ext, (1.0 * x[i], per, aff[2 + i]), model)
    end

    # perspective nonnegative NOTE could be redundant with atom constraints
    JuMP.@constraint(model, per >= 0)

    return
end

#=
MatNegRtdet
=#
function add_homog_spectral(
    ::MatNegRtdet,
    d::Int,
    aff::Vector{JuMP.AffExpr},
    model::JuMP.Model,
    )
    aff_new = vcat(-aff[1], aff[2:end])
    R = get_R(d, length(aff) - 1)
    JuMP.@constraint(model, aff_new in
        Hypatia.HypoRootdetTriCone{Float64, R}(length(aff)))
    return
end

#=
MatNegRtdetEF
(u, W ≻ 0) : u > -rootdet(W)
↔ ∃ upper triangular U : [W U'; U Diag(δ)] ≻ 0, δ = diag(U),
(u, δ) ∈ NegRtdetean
=#
function add_homog_spectral(
    ext::Union{MatNegRtdetEFExp, MatNegRtdetEFPow},
    d::Int,
    aff::Vector{JuMP.AffExpr},
    model::JuMP.Model,
    )
    R = get_R(d, length(aff) - 1)
    δ = extend_det(R, d, aff[2:end], model)
    add_homog_spectral(get_vec_ef(ext), d, vcat(aff[1], δ), model)
    return
end

#=
MatLogdetCone
=#
function add_spectral(
    ::MatLogdetCone,
    d::Int,
    aff::Vector{JuMP.AffExpr},
    model::JuMP.Model,
    )
    aff_new = vcat(-aff[1], aff[2], aff[3:end])
    R = get_R(d, length(aff) - 2)
    JuMP.@constraint(model, aff_new in
        Hypatia.HypoPerLogdetTriCone{Float64, R}(length(aff)))
    return
end

#=
MatSepSpec
=#
function add_spectral(
    ext::MatSepSpec,
    d::Int,
    aff::Vector{JuMP.AffExpr},
    model::JuMP.Model,
    )
    R = get_R(d, length(aff) - 2)
    is_dual = (ext isa MatSepSpecDualAll)
    cone = Hypatia.EpiPerSepSpectralCone{Float64}(get_ssf(ext),
        Cones.MatrixCSqr{Float64, R}, d, is_dual)
    JuMP.@constraint(model, aff in cone)
    return
end

#=
MatSepSpecEigOrd
uses eigenvalue ordering constraints from Ben-Tal & Nemirovski
=#
function add_spectral(
    ext::MatSepSpecEigOrd,
    d::Int,
    aff::Vector{JuMP.AffExpr},
    model::JuMP.Model,
    )
    λ = JuMP.@variable(model, [1:d])
    JuMP.@constraint(model, [i in 1:(d - 1)], λ[i] >= λ[i + 1])

    vec_aff = vcat(aff[1], aff[2], λ)
    add_spectral(get_vec_ef(ext), d, vec_aff, model)

    w = aff[3:end]
    R = get_R(d, length(w))

    if R <: Real
        W = get_smat(d, w)
        trW = tr(W)
        for i in 1:(d - 1)
            Z = JuMP.@variable(model, [1:d, 1:d], PSD)
            s = JuMP.@variable(model)
            JuMP.@constraint(model, sum(λ[1:i]) - i * s - tr(Z) >= 0)
            mat = Symmetric(Z - W + s * Matrix(I, d, d), :U)
            JuMP.@SDconstraint(model, mat >= 0)
        end
    else
        (Wr, Wi) = get_smat_complex(d, w)
        trW = tr(Wr)
        for i in 1:(d - 1)
            Zr = JuMP.@variable(model, [1:d, 1:d], Symmetric)
            Zia = JuMP.@variable(model, [1:(d - 1), 1:(d - 1)], Symmetric)
            Zi = make_skewsymm(d, 1.0 * Zia)
            Z2 = Symmetric(hvcat((2, 2), Zr, Zi', Zi, Zr), :U)
            JuMP.@SDconstraint(model, Z2 >= 0)

            s = JuMP.@variable(model)
            JuMP.@constraint(model, sum(λ[1:i]) - i * s - tr(Zr) >= 0)

            Mr = Symmetric(Zr - Wr + s * Matrix(I, d, d), :U)
            Mi = Zi - Wi
            mat = Symmetric(hvcat((2, 2), Mr, Mi', Mi, Mr), :U)
            JuMP.@SDconstraint(model, mat >= 0)
        end
    end
    JuMP.@constraint(model, trW == sum(λ))

    return
end

#=
MatNegLogDirect
(u, v > 0, W ≻ 0) : u > -v logdet(W / v)
↔ ∃ upper triangular U : [W U'; U Diag(δ)] ≻ 0, δ = diag(U),
(u, v, δ) ∈ NegLogVector
=#
function add_spectral(
    ::MatNegLogDirect,
    d::Int,
    aff::Vector{JuMP.AffExpr},
    model::JuMP.Model,
    )
    R = get_R(d, length(aff) - 2)
    δ = extend_det(R, d, aff[3:end], model)
    vec_aff = vcat(aff[1], aff[2], δ)
    add_spectral(VecNegLogEF(), d, vec_aff, model)
    return
end

#=
MatNegSqrtConjDirect
note this is a dual cone EF, so u and v are swapped
(u > 0, v > 0, W ≻ 0) : v > 1/4 * u tr(inv(W / u))
↔ ∃ Z : [Z uI; uI W] ≻ 0, 4 * v > tr(Z), u > 0 (use Schur complement)
=#
function add_spectral(
    ::MatNegSqrtConjDirect,
    d::Int,
    aff::Vector{JuMP.AffExpr},
    model::JuMP.Model,
    )
    w = aff[3:end]
    R = get_R(d, length(w))

    JuMP.@constraint(model, aff[1] >= 0)
    Z = JuMP.@variable(model, [1:d, 1:d], Symmetric)
    JuMP.@constraint(model, 4 * aff[2] >= tr(Z))

    uI = aff[1] * Matrix(I, d, d)
    if R <: Real
        W = get_smat(d, w)
        mat = Symmetric(hvcat((2, 2), Z, uI, uI, W), :U)
    else
        (Wr, Wi) = get_smat_complex(d, w)
        Zia = JuMP.@variable(model, [1:(d - 1), 1:(d - 1)], Symmetric)
        Zi = make_skewsymm(d, 1.0 * Zia)

        Mr = Symmetric(hvcat((2, 2), Z, uI, uI, Wr), :U)
        Mi = hvcat((2, 2), Zi, zeros(d, d), zeros(d, d), Wi)
        mat = Symmetric(hvcat((2, 2), Mr, Mi', Mi, Mr), :U)
    end
    JuMP.@SDconstraint(model, mat >= 0)

    return
end

#=
helpers
=#

# get epigraph and perspective variables, swapping if using dual cone
swap_if_dual(x, y, ext::SepSpecDualAll) = (y, x)
swap_if_dual(x, y, ext::SepSpecAll) = (x, y)

# get real or complex type from length of affine expression
function get_R(d::Int, w_len::Int)
    if w_len == d^2
        return ComplexF64
    else
        @assert w_len == Cones.svec_length(d)
        return Float64
    end
end

# get upper triangle of symmetric matrix W from vectorized w
function get_smat(d::Int, w::Vector{JuMP.AffExpr})
    @assert Cones.svec_length(d) == length(w)
    W = zeros(JuMP.AffExpr, d, d)
    Cones.svec_to_smat!(W, w, sqrt(2))
    return W
end

# get real and imaginary matrix parts from vectorized w
function get_smat_complex(d::Int, w::Vector)
    @assert d^2 == length(w)
    rt2 = sqrt(2)
    Wr = zeros(eltype(w), d, d)
    Wi = zero(Wr)
    idx = 1
    for j in 1:d
        for i in 1:(j - 1)
            Wr[i, j] = Wr[j, i] = w[idx] / rt2
            Wi[i, j] = w[idx + 1] / rt2
            Wi[j, i] = -Wi[i, j]
            idx += 2
        end
        Wr[j, j] = w[idx]
        idx += 1
    end
    return (Wr, Wi)
end

# make a skew-symmetric matrix
function make_skewsymm(d::Int, U::Symmetric)
    Z = zeros(eltype(U), d, d)
    Z[1:(d - 1), 2:d] = UpperTriangular(U)
    @assert istriu(Z)
    Z .-= Z'
    @assert all(iszero, diag(Z))
    return Z
end

# construct the real matrix part of the det EF, return δ like eigenvalues
function extend_det(
    R::Type{Float64},
    d::Int,
    w::Vector{JuMP.AffExpr},
    model::JuMP.Model,
    )
    u = JuMP.@variable(model, [1:length(w)])
    U = get_smat(d, 1.0 * u)
    δ = diag(U)

    W = get_smat(d, w)
    mat = Symmetric(hvcat((2, 2), W, U', U, Diagonal(δ)), :U)
    JuMP.@SDconstraint(model, mat >= 0)

    return δ
end

# construct the complex matrix part of the det EF, return δ like eigenvalues
function extend_det(
    R::Type{ComplexF64},
    d::Int,
    w::Vector{JuMP.AffExpr},
    model::JuMP.Model,
    )
    (Wr, Wi) = get_smat_complex(d, w)

    Ur = UpperTriangular(JuMP.@variable(model, [1:d, 1:d], Symmetric))
    δ = diag(Ur)
    Uia = JuMP.@variable(model, [1:(d - 1), 1:(d - 1)], Symmetric)
    Ui = zero(Wi)
    Ui[1:(d - 1), 2:d] = UpperTriangular(Uia)

    Mr = Symmetric(hvcat((2, 2), Wr, Ur', Ur, Diagonal(δ)), :U)
    Mi = hvcat((2, 2), Wi, -Ui', Ui, zeros(d, d))
    mat = Symmetric(hvcat((2, 2), Mr, Mi', Mi, Mr), :U)
    JuMP.@SDconstraint(model, mat >= 0)

    return δ
end

#=
NegLogSSF:
(x, y > 0, z > 0) : x > y * -log(z / y)
↔ y * exp(-x / y) < z
↔ (-x, y, z) ∈ MOI.ExponentialCone
conjugate does not provide additional modeling power
=#
function extend_atom(
    ::VecNegLogEF,
    aff::NTuple{3, JuMP.AffExpr},
    model::JuMP.Model,
    )
    aff_new = vcat(-aff[1], aff[2], aff[3])
    JuMP.@constraint(model, aff_new in MOI.ExponentialCone())
    return
end

#=
NegEntropySSF:
(x, y > 0, z > 0) : x > z * log(z / y) = -z * log(y / z)
↔ z * exp(-x / z) < y
↔ (-x, z, y) ∈ MOI.ExponentialCone
or for the conjugate: (z can be negative)
(x > 0, y > 0, z) : x > y * exp(-z / y - 1) = y * exp((-z - y) / y)
↔ (-z - y, y, x) ∈ MOI.ExponentialCone
=#
function extend_atom(
    ::VecNegEntropyEF,
    aff::NTuple{3, JuMP.AffExpr},
    model::JuMP.Model,
    )
    aff_new = vcat(-aff[1], aff[3], aff[2])
    JuMP.@constraint(model, aff_new in MOI.ExponentialCone())
    return
end

function extend_atom(
    ::VecNegEntropyConjEF,
    aff::NTuple{3, JuMP.AffExpr},
    model::JuMP.Model,
    )
    aff_new = vcat(-aff[3] - aff[2], aff[2], aff[1])
    JuMP.@constraint(model, aff_new in MOI.ExponentialCone())
    return
end

#=
NegSqrtSSF:
(x, y > 0, z > 0) : x > -y * sqrt(z / y) = -sqrt(y * z)
↔ ∃ θ > 0 : (x - θ)^2 < y * z
↔ (y, z, sqrt(2) * (x - θ)) ∈ JuMP.RotatedSecondOrderCone, θ ∈ ℝ₊
or for the conjugate:
(x > 0, y > 0, z > 0) : x > 1/4 * y * inv(z / y) = 1/4 * y^2 / z
↔ 2 * x * z > y^2 / 2
↔ (x, z, y / sqrt(2)) ∈ JuMP.RotatedSecondOrderCone
=#
function extend_atom(
    ::VecNegSqrtEF,
    aff::NTuple{3, JuMP.AffExpr},
    model::JuMP.Model,
    )
    θ = JuMP.@variable(model)
    JuMP.@constraint(model, θ >= 0)
    aff_new = vcat(aff[2], aff[3], sqrt(2) * (aff[1] - θ))
    JuMP.@constraint(model, aff_new in JuMP.RotatedSecondOrderCone())
    return
end

function extend_atom(
    ::VecNegSqrtConjEF,
    aff::NTuple{3, JuMP.AffExpr},
    model::JuMP.Model,
    )
    aff_new = vcat(aff[1], aff[3], aff[2] / sqrt(2))
    JuMP.@constraint(model, aff_new in JuMP.RotatedSecondOrderCone())
    return
end

#=
NegPower01SSF(p) for p ∈ (0, 1)
(x, y > 0, z > 0) : x > -y * (z / y)^p = -y^(1-p) * z^p
↔ ∃ θ > 0 : z^p * y^(1-p) > |x - θ|
↔ (z, y, x - θ) ∈ MOI.PowerCone(p), θ ∈ ℝ₊
or for the conjugate:
let q = p / (p - 1) < 0, so 1 - p = 1 / (1 - q)
let c = (1 - p) * p^-q > 0
(x > 0, y > 0, z > 0) : x > c * y * (z / y)^q = c * z^q * y^(1-q)
↔ x^(1/(1-q)) > c^(1/(1-q)) * z^(q/(1-q)) * y
↔ z^(q/(q-1)) * x^(1/(1-q)) > c^(1/(1-q)) * y
↔ z^p * x^(1-p) > |b * y|, y > 0
where b = c^(1/(1-q)) = p^p * (1 - p)^(1-p) > 0
↔ (z, x, b * y) ∈ MOI.PowerCone(p), y ∈ ℝ₊
=#
function extend_atom(
    ext::VecNegPower01EF,
    aff::NTuple{3, JuMP.AffExpr},
    model::JuMP.Model,
    )
    p = ext.p
    @assert 0 < p < 1
    θ = JuMP.@variable(model)
    JuMP.@constraint(model, θ >= 0)
    aff_new = vcat(aff[3], aff[2], aff[1] - θ)
    JuMP.@constraint(model, aff_new in MOI.PowerCone(p))
    return
end

function extend_atom(
    ext::VecNegPower01ConjEF,
    aff::NTuple{3, JuMP.AffExpr},
    model::JuMP.Model,
    )
    p = ext.p
    @assert 0 < p < 1
    q = p / (p - 1)
    b = p^p * (1 - p)^(1 - p)
    aff_new = vcat(aff[3], aff[1], b * aff[2])
    JuMP.@constraint(model, aff_new in MOI.PowerCone(p))
    return
end

#=
Power12SSF(p) for p ∈ (1, 2]
(x > 0, y > 0, z > 0) : x > y * (z / y)^p = y^(1-p) * z^p
↔ x^(1/p) * y^(1-1/p) > |z|, z > 0
↔ (x, y, z) ∈ MOI.PowerCone(1/p), z ∈ ℝ₊
or for the conjugate: (z can be negative)
let z₋ = {0 if z ≥ 0, or -z if z < 0}, q = p / (p - 1) > 2
(x > 0, y > 0, z) : x > (p - 1) * y * (z₋ / y / p)^q
= (p - 1) / p^q * z₋^q * y^(1-q)
↔ x^(1/q) > c * z₋ * y^(1/q-1)
↔ x^(1/q) * y^(1-1/q) > c * z₋
↔ ∃ θ > 0 : x^(1/q) * y^(1-1/q) > |c * (z - θ)|
where c = (p - 1)^(1/q) / p > 0
↔ (x, y, c * (z - θ)) ∈ MOI.PowerCone(1/q), θ ∈ ℝ₊
=#
function extend_atom(
    ext::VecPower12EF,
    aff::NTuple{3, JuMP.AffExpr},
    model::JuMP.Model,
    )
    p = ext.p
    @assert 1 < p <= 2
    JuMP.@constraint(model, aff[3] >= 0)
    JuMP.@constraint(model, collect(aff) in MOI.PowerCone(inv(p)))
    return
end

function extend_atom(
    ext::VecPower12ConjEF,
    aff::NTuple{3, JuMP.AffExpr},
    model::JuMP.Model,
    )
    p = ext.p
    @assert 1 < p <= 2
    q = p / (p - 1)
    c = (p - 1)^inv(q) / p
    θ = JuMP.@variable(model)
    JuMP.@constraint(model, θ >= 0)
    aff_new = vcat(aff[1], aff[2], c * (aff[3] - θ))
    JuMP.@constraint(model, aff_new in MOI.PowerCone(inv(q)))
    return
end