text:
#   Copyright 2022 Velexi Corporation
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

"""
jltest/utils.jl defines utility functions to support testing.
"""

# --- Exports

export find_tests, run_tests

# --- Imports

# Standard library
using Logging
using Test
using Test: AbstractTestSet

# External packages
using OrderedCollections: OrderedDict
using Suppressor: @capture_err

# --- Private utility functions

"""
    reemit_log_msg(messages::AbstractString)

Remove missing dependency warnings from log messages and reemit them.
"""
function reemit_log_msg(messages::AbstractString)

    # --- Preparations

    # Get lines of message
    lines = split(messages, '\n')

    # --- Extract, reformat, and re-emit messages

    message_finished = true
    message_lines = nothing
    for i in 1:length(lines)

        # --- Extract next log message

        if message_finished
            message_lines = [lines[i]]

            # Check for multi-line message
            if startswith(lines[i], "┌ ")
                message_finished = false
            end

        else
            push!(message_lines, lines[i])

            # Check for end of multi-line message
            if startswith(lines[i], "└ ")
                message_finished = true
            end
        end

        if message_finished

            # --- Skip missing dependency warning

            if occursin(
                r"^┌ Warning: Package TestTools does not have \S* in its dependencies:",
                message_lines[1],
            )
                continue
            end

            # --- Preparations

            # Get log level
            log_level = Logging.Warn
            if occursin("Info", message_lines[1])
                log_level = Logging.Info
            elseif occursin("Debug", message_lines[1])
                log_level = Logging.Debug
            end

            # --- Reformat message

            # Reformat first line
            if startswith(message_lines[1], "┌ Warning: ")
                message_lines[1] = replace(message_lines[1], "┌ Warning: " => "")
            elseif startswith(message_lines[1], "[ Warning: ")
                message_lines[1] = replace(message_lines[1], "[ Warning: " => "")
            elseif startswith(message_lines[1], "┌ Info: ")
                message_lines[1] = replace(message_lines[1], "┌ Info: " => "")
            elseif startswith(message_lines[1], "[ Info: ")
                message_lines[1] = replace(message_lines[1], "[ Info: " => "")
            elseif startswith(message_lines[1], "┌ Debug: ")
                message_lines[1] = replace(message_lines[1], "┌ Debug: " => "")
            elseif startswith(message_lines[1], "[ Debug: ")
                message_lines[1] = replace(message_lines[1], "[ Debug: " => "")
            end

            # Reformat middle lines
            for i in 2:(length(message_lines) - 1)
                message_lines[i] = replace(message_lines[i], "│ " => "")
            end

            # Reformat last line of message
            message_lines[end] = replace(message_lines[end], "└ " => "")

            # --- Emit message

            message = join(message_lines, '\n')
            @logmsg log_level message _module = nothing _file = nothing _group = nothing
        end
    end
end

"""
    handle_test_exception(error::Base.error)

Handle exception thrown when running a test file. If `error` is a missing package
dependency, construct a LoadError to throw outside of the catch block. Otherwise, rethrow
the error.
"""
function handle_test_exception(error::Exception)
    if error isa LoadError && error.error isa ArgumentError
        match_results = match(
            r".*Package TestTools does not have (.*?) in its dependencies:.*",
            error.error.msg,
        )

        if !isnothing(match_results)
            message = strip(
                """
                The test environment is missing $(match_results[1]) from its dependencies.
                Error occurred at $(error.file):$(error.line)
                """
            )
            return ErrorException(message)
        end
    end

    # Not a missing dependency error, so rethrow it.
    rethrow()

    return nothing
end

"""
    run_all_tests(test_files::Vector{<:AbstractString})

Run all tests contained in `test_files`.
"""
function run_all_tests(test_files::Vector{<:AbstractString})

    # Get current directory
    cwd = pwd()

    # Run tests files
    if !isempty(test_files)
        for test_file in test_files
            # Restore current directory before each test file is run
            cd(cwd)

            # Construct module name
            module_name = splitext(relpath(test_file, cwd))[1]

            # Run test, capturing log messages
            println()
            print(module_name, ": ")
            mod = gensym(module_name)
            missing_dependency_error = nothing
            log_msg = strip(@capture_err begin
                try
                    @eval module $mod
                    Base.include($mod, abspath($test_file))
                    end
                catch error
                    missing_dependency_error = handle_test_exception(error)
                end
            end)

            if !isnothing(missing_dependency_error)
                throw(missing_dependency_error)
            end

            # Suppress warnings about missing TestTools dependencies
            if !isempty(log_msg)
                reemit_log_msg(log_msg)
            end
        end
    end
end

# --- Functions/Methods

"""
    run_tests(tests::Vector; <keyword arguments>)
    run_tests(tests::AbstractString; <keyword arguments>)

Run tests in the list of files or modules provided in `tests`. If `tests` is an empty list
or an empty string, an `ArgumentError` is thrown. File names in `tests` may be specified
with or without the `.jl` extension.

# Keyword Arguments

* `desc::AbstractString`: description to use for test set used to group `tests`.
  Default: the default description set by `@testset`

* `test_set_type::Type`: type of test set to use to group tests. When `test_set_type`
  is set to `nothing`, the tests are run individually.
  Default: `EnhancedTestSet{DefaultTestSet}`

  !!! note
      When the `JLTEST_FAIL_FAST` environment variable is set to "true", the `test_set_type`
      argument is overridden and set to `EnhancedTestSet{Test.FallbackTestSet}` so that
      tests are run in fail-fast mode.

* `recursive::Bool`: flag indicating whether or not to run tests found in subdirectories
  of directories in `tests`. Default: `true`

* `exclude_runtests::Bool`: flag indicating whether or not to exclude files named
  `runtests.jl` from the list of test files that are run. Default: `true`
"""
function run_tests(
    tests::Vector;
    desc::AbstractString="",
    test_set_type::Union{Type{<:AbstractTestSet},Nothing}=EnhancedTestSet{DefaultTestSet},
    recursive::Bool=true,
    exclude_runtests::Bool=true,
)
    # --- Check arguments

    # Ensure that `tests` is not empty
    if isempty(tests)
        throw(ArgumentError("`tests` may not be empty"))
    end

    # Ensure that `tests` contains strings
    tests = convert(Vector{String}, tests)

    # Enable fail-fast if JLTEST_FAIL_FAST environment variable is set to "true"
    if lowercase(get(ENV, "JLTEST_FAIL_FAST", "false")) == "true"
        test_set_type = EnhancedTestSet{Test.FallbackTestSet}
    end

    # --- Preparations

    test_files = Vector{String}()

    for test in tests
        if isdir(test)
            # Find tests contained in the directory
            test_files = vcat(
                test_files,
                find_tests(test; recursive=recursive, exclude_runtests=exclude_runtests),
            )
        else
            # Get the file name
            file_name = test
            if !endswith(file_name, ".jl")
                # `test` is a module, so append ".jl"
                file_name = join([file_name, ".jl"])
            end

            push!(test_files, file_name)
        end
    end

    # --- Run tests

    if isnothing(test_set_type)
        run_all_tests(test_files)
    else
        if isempty(desc)
            @testset test_set_type begin
                run_all_tests(test_files)
            end
        else
            @testset test_set_type "$desc" begin
                run_all_tests(test_files)
            end
        end
    end

    return nothing
end

# run_tests(tests::AbstractString) method that converts the argument to a Vector{String}
function run_tests(
    test::AbstractString;
    desc::AbstractString="",
    test_set_type::Union{Type{<:AbstractTestSet},Nothing}=EnhancedTestSet{DefaultTestSet},
    recursive::Bool=true,
)
    # --- Check arguments

    # Ensure that `tests` is not an empty string
    if isempty(test)
        throw(ArgumentError("`test` may not be empty"))
    end

    # --- Run tests

    run_tests(Vector{String}([test]); desc=desc, test_set_type=test_set_type)

    return nothing
end

"""
    find_tests(dir::AbstractString), <keyword arguments>)::Vector{String}

Construct a list of absolute paths to Julia test files contained in `dir`.

# Keyword Arguments

* `recursive::Bool`: flag indicating whether or not tests found in subdirectories of `dir`
  should be included in results. Default: `true`

* `exclude_runtests::Bool`: flag indicating whether or not to exclude files named
  `runtests.jl` from the list of test files. Default: `true`
"""
function find_tests(
    dir::AbstractString; recursive::Bool=true, exclude_runtests::Bool=true
)::Vector{String}

    # --- Check arguments

    # Ensure `dir` is an absolute path
    dir = abspath(dir)

    # --- Find test files

    # Find test files in `dir`
    files = filter(f -> endswith(f, ".jl"), readdir(dir))
    if exclude_runtests
        files = filter(f -> f != "runtests.jl", files)
    end
    # TODO: filter `tests` to exclude files that do not contain tests

    tests = [joinpath(dir, file) for file in files]

    # Recursively search directories
    if recursive
        dirs = filter(f -> isdir(f), [joinpath(dir, f) for f in readdir(dir)])
        for dir in dirs
            tests = vcat(tests, find_tests(dir; exclude_runtests=exclude_runtests))
        end
    end

    return tests
end


text:
ENV["JULIA_CXX_RTTI"]=1
using Cxx

abstract type ROSTYPE end

ext: jl


text:
module SocioHydrology

export SocioFEHM

include("PathFile.jl");
import SupportFunc: updaterec2, getrisk2, pidcontrol;
import SocialDynamics: socdyncoupad, socdyncoupexp;
import BounWrite: setboun, timelimit;

"""

    SocioFEHM(; conc_risk = 0, exp_adv = 0, exprate_advper = 0.03 , simulationtime = 20, cthresh = 30, pid = [0, 0, 0],  measure_period = 1, predetect_time = 0, startpumptime = 10, init_urate = 0.99, popul= 10000, agelimit = 70, risk_percentile = .90,  adtargetfrac = 0.5, FEHMprefix = "cm", outputfile = "default_output_file.txt")
[requires names of variables]

Simulate a coupled consistency-conformity social dynamics with FEHM for contaminated aquifer.

Arguments:

- `conc_risk` : concentration control (0), or risk control (1).
- `exp_adv` : expert influence (0), or media influence (1).
- `exprate_advper` : fraction of experts in population (if `exp_adv = 0`), or advertising period (if `exprate_advper = 1`).
- `simulationtime` : Total time of coupled simulation (in years).
- `cthresh` : concentration threshold (control reference). Risk threshold `rthresh` is calculated in equivalence.
- `pid` : vector `[pK, dK, iK]` of proportional-derivative-integral control coefficients.
- `measure_period` : time interval between two successive measurement events (years).
- `predetect_time` : time before contaminants are detected in acquifer (in years).
- `startpumptime` : Time (from beginning of FEHM simulation) after which pumping starts (years)
- `init_urate` : fraction of the population that initially uses the water.
- `popul` : population size.
- `agelimit` : age at which people die, and get replaced by a new-born (used for risk control).
- `risk_percentile`: Percentile of people kept below risk threshold.
- `adtargetfrac` : fraction of population reached by media influence (used for `exp_adv = 1`).
- `FEHMprefix` : FEHM prefix of files for aquifer simulation, placed in `./contaminantion-model/`.
- `outputfile` : Default suffix for saving data in text files; set to `""` to prevent saving.

Returns:

return times, prate, adorbel, adornumexp, ucount

- `times`: Times in years since of FEHM simulation.
- `adornumexp`: number of employed experts if `exp_adv = 0`, or period between sucessive ads if `exp_adv = 1`
- `ucount`: number of water users.
- `conchist` : aquifer concentration history

"""

function SocioFEHM(; conc_risk = 0, exp_adv = 0, exprate_advper = 0.03 , simulationtime = 20, cthresh = 30, pid = [0, 0, 0],  measure_period = 1, predetect_time = 0, startpumptime = 10, init_urate = 0.99, popul= 10000, agelimit = 70, risk_percentile = .90,  adtargetfrac = 0.5, FEHMprefix = "cm", outputfile = "default_output_file.txt")

	### Fixed Variables ###
	period 		= 1; 		# Number of social dynamics simulations for 1 year of FEHM simulations
	initialrisk 	= 0;		# Initial health risk (per thousand)
	useday 		= 2660; 	# Water use per day per person in Liters, estimated as USGS to be 80-100 gallons - 340L
	wdensity 	= 1; 		# Water density  (Kg/L)
	interfreq 	= 15; 		# Calculated such that 1000 steps make around 50 years. (DAYS)
	location 	= 3612;		# Position of well in FEHM (3612 for homogeneous)
	pcons 		= 0.35; 	# Consistency rate
	numatr		= 2; 		# Belief - Usage:  attributes
	Dt 		= 1;	 	# Integration/derivative time; TO BE FIXED IF DECISION TIME IS NOT = 1

	### RISK DATA ###
	# Reference: http://pubs.acs.org/doi/abs/10.1021/es400316c
	# (Paper has few mistakes in units)

	irbw 		= 0.022;	# Ingestion rate of water per unit body weight (L kg/day)
	ef 		= 350;		# Standard exposure frequency (day/year)
	mfi 		= 1;		# Metabolized fraction of contaminant
	cpf 		= 0.42; 	# Cancer potency factor (Kg day/mg)

	#### INITIAL CONDITIONS ####
	socatribute 	= zeros(Bool, popul, numatr); 	# Belief(False - bel. not cont. ; True - bel. cont.) , Usage (False - using ; True - not using),
	socconc 	= zeros(Float64, popul, 1);  	# Total usage concentration
	socage		= zeros(Int, popul, 1); 	# Age
	soc 		= hcat(socatribute, socconc, socage);

	times 		= Float32[];		# Time of pumping (years)
	prate 		= Float32[];		# pumping rate function of time
	conchist 	= Float32[];		# pumped concentration as a function of time
	riskhist 	= Float32[];		# Control risk as a function of time
	riskarray 	= zeros(Float32, popul, 1); # Risk of all population at each iteration
	adfreqmat 	= Int64[];		# ad frequency function of time
	adflagmat	= Bool[];		# Ads or no ads per year
	expb  		= Bool[];		# expert belief function of time
	num_expert	= Int32[];		# Number of experts function of time
	soc_rec		= Array{Bool}[];
	ucount 		= Int32[];		# Number of people using the water
	simstep 	= round(Int, simulationtime/period);	# simluation time per iteration
	adcounter 	= 0;			# Counts years before next ad
	adf		= exprate_advper;	# ad frequency.
	adflag		= 0; 			# Flag whether to make ads (that year) or not.
	ca 		= 1;
	belief		= false;
	expmax_advper 	= convert(Bool, exp_adv) ? exprate_advper : convert(Int32, float(exprate_advper) * popul);
	cc 		= 0;		# Initial concentration
	numexp		= expmax_advper;

	risk = initialrisk;
	rthresh = getrisk2([1,1,0,0]', cthresh, agelimit, irbw, ef, mfi, cpf)[1]; # Risk threshold calculated in equivalence to concentration threshold
	rsig 	= rthresh/3;
	sig	= cthresh/3;	# Normal distribution std (concentration control) used for pid control

	# PID Parameters
	pK 		= pid[1];		# Proportional coefficient
	dK 		= pid[2];		# Derivative coefficient
	iK 		= pid[3];		# Integral coefficient

	# uniform distribution of ages
	for i = 1:popul
		soc[i,4] = floor(Int, rand() * agelimit);
	end

	# Specify the initial users
	soc[end - round(Int, (1-init_urate)*popul) : end, 2] = true;

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

	println("Simulation Parameters:");
	println(convert(Bool, exp_adv) ? "Media Influence" : "Expert Influence");
	println(convert(Bool, conc_risk) ? "Risk Control" : "Concentration Control");
	println("____________________________________________________");

	for i = 1:simstep

	# DECISION MAKING (CONTROL)
		if(exp_adv == 1)

			adcounter += period;
			if(mod(i, measure_period) == 0 && i > predetect_time && i>1)
				println("Making Decision");
				if(adcounter >= adf)
					if(conc_risk == 1)
						belief, adf  = pidcontrol(rthresh, pK, dK, iK, riskhist, Dt, rsig, exp_adv, expmax_advper);
					elseif(conc_risk == 0)
						belief, adf  = pidcontrol(cthresh, pK, dK, iK, conchist, Dt, sig, exp_adv, expmax_advper);
					end
					adcounter = 0;
					adflag = true;
				else
					adflag = false;
				end
			end
			append!(adflagmat, [adflag]);
			append!(adfreqmat, [adf]);

			println(string("Advertising this year:\t ", convert(Bool, adflag) ? "yes!" : "no"));
			println(string("Media Belief:\t ", belief ? "It's contaminated!" : "it's not contaminated"));
			println(string("Ads counter/freq:\t ", adcounter, " out of ", adf));

		else

			if(mod(i, measure_period) == 0 && i > predetect_time && i > 1)
				println("Making Decision");
				if(conc_risk==1)
					belief, numexp = pidcontrol(rthresh, pK, dK, iK, riskhist, Dt, rsig, exp_adv, expmax_advper)
				elseif(conc_risk==0)
					belief, numexp = pidcontrol(cthresh, pK, dK, iK, conchist, Dt, sig, exp_adv, expmax_advper)
				end
			end
			append!(expb, [belief]);
			append!(num_expert, [numexp]);

			println(string("expert belief:\t\t ", belief ? "it's contaminated" : "it's not contaminated"));
			println(string("number of experts:\t ", numexp));
		end

			println(string("year:\t\t\t ", i, " of ", simulationtime));
			println(string("risk:\t\t\t ", Float16(risk), " to ", Float16(rthresh)));

		# Initialize FEHM files
		usage = popul - round(Int, sum(soc,1)[2]);
		println(string("number of users:\t ", usage, " out of ", popul));
		append!(ucount, [usage]);
		append!(prate, [usage*useday/24/3600*wdensity]);
		append!(times, [startpumptime + period*(i-1)]);
		setboun(times, prate, location, FEHMprefix);
		timelimit(round(Int, times[end]*365.25/8), location, FEHMprefix);

		#### RUN FEHM CODE ####
		run(`sh ${srcdir}fehmrun.sh $fehmdir $FEHMprefix`);
		steps = round(Int, period*365.25/interfreq);

		#### SOCIAL DYNAMICS ####
		if(exp_adv == 1)
			soc[:,1:2] = socdyncoupad(popul, adflag, belief, adtargetfrac, soc[:,1:2], pcons, steps);
		else
			soc[:,1:2] = socdyncoupexp(popul, numexp, belief, soc[:,1:2], pcons, steps);
		end

		#### UPDATE SOCIAL RECORDS ####
		soc = updaterec2(soc, cc, period, agelimit);
		push!(soc_rec, soc[:,1:2]);

		#### RISK ####
		population_risk = getrisk2(soc, cc, agelimit, irbw, ef, mfi, cpf);
		riskarray = hcat(riskarray, population_risk);

		if(length(population_risk) >= 2)
			saferank = round(Int, (1 - risk_percentile) * length(population_risk) ); 	# rank based on percentile
			risk = sort(population_risk, rev=true)[saferank]; 			# risk associated with that rank
		else
			risk = 0;
		end
		append!(riskhist, [risk]);

		# Read contamination level from (FEHM trc) file
		file = readdlm(string(fehmdir, FEHMprefix, ".trc"));
		cc = Float32( file[end,1] );
		append!(conchist, [cc]);
		println(string("Concentration:\t\t ", cc));
		println(string("_______________________________________________"));
	end

	if(exp_adv==1)
		adorbel = adflagmat;
		adornumexp = adfreqmat;
	else
		adorbel = expb;
		adornumexp = num_expert;
	end

	stmat = hcat(times, prate, adorbel, adornumexp, ucount);

	if(outputfile != "")
		writedlm(string(fehmdirresult, "hist_", outputfile), stmat); 	# years, pumping_rate, ads_or_expertbelief, adfrequency_or_expnumber, num_users
		writedlm(string(fehmdirresult, "risk_", outputfile), riskarray);# Risk for all individuals for all years.
		writedlm(string(fehmdirresult, "socrec_", outputfile), soc_rec);# Belief and Usage bits for all individuals for all years.
		#writedlm(string(fehmdirresult, "soc_", outputfile), soc);	# Last year's: belief, usage, concentration accumulation, age
	end

return times, adornumexp, ucount, conchist, riskhist

end

end

max_stars_repo_name: josephbakarji/socio-hydro-abm

max_issues_repo_name: josephbakarji/socio-hydro-abm

max_forks_repo_name: josephbakarji/socio-hydro-abm


text:
using EconPDEs

Base.@kwdef struct BoltonChenWangModel
  r::Float64 = 0.06
  δ::Float64 = 0.10
  A::Float64 = 0.18
  σ::Float64 = 0.09
  θ::Float64 = 1.5
  λ::Float64 = 0.01
  l::Float64 = 0.9
  γ::Float64 = 0.06
  ϕ::Float64 = 0.01
end

function (m::BoltonChenWangModel)(state::NamedTuple, y::NamedTuple)
  (; r, δ, A, σ, θ, λ, l, γ, ϕ) = m
  (; w) = state
  (; v, vw_up, vw_down, vww) = y
  vw = vw_up
  iter = 0
  @label start
  i = 1 / θ * (v / vw - w - 1)
  μw = (r - λ) * w + A - i - θ * i^2 / 2 - (i - δ) * w
  if (iter == 0) & (μw <= 0)
    iter += 1
    vw = vw_down
    @goto start
  end
  vt = - ((i - δ) * (v - vw * w) + ((r - λ) * w + A - i - θ * i^2 / 2) * vw + σ^2 / 2 * vww - r * v)
  return (; vt), (; v, vw, vww, w)
end

m = BoltonChenWangModel()
stategrid = OrderedDict(:w => range(0.0, 0.3, length = 100))
yend = OrderedDict(:v => stategrid[:w])
result = pdesolve(m, stategrid, yend; bc = OrderedDict(:vw => (1.5, 1.0)))
@assert result.residual_norm <= 1e-5

#========================================================================================

Iterate on boundary conditions until the solution satisfies the conditions given in Bolton Chen Wang (2009)

========================================================================================#
using NLsolve
function f(m, x, stategrid, y)
  result = pdesolve(m, stategrid, y; bc = OrderedDict(:vw => (x[1], x[2])), verbose = false)
  v, vw, w = result.optional[:v], result.optional[:vw], result.optional[:w]
  out = zeros(2)
  mi = argmin(abs.(1 + m.γ .- vw))
  out[1] =  v[mi] - m.ϕ - (1 + m.γ) * w[mi] - v[1]
  wi = argmin(abs.(1 .- vw))
  i = 1 / m.θ * (v[wi] - w[wi] - 1)
  out[2] = m.r * v[wi] - (i - m.δ) * (v[wi] - w[wi]) - ((m.r - m.λ) * w[wi] + m.A - i - m.θ * i^2 / 2)
  return out
end

newsol = nlsolve(x -> f(m, x, stategrid, yend), [1.0,  1.0])
result = pdesolve(m, stategrid, yend; bc = OrderedDict(:vw => (newsol.zero[1], newsol.zero[2])))


text:
export Hankel

immutable Hankel{T} <: AbstractArray{T, 2}
    c :: Vector{T}
end
Hankel{T}(c::Vector{T}) = length(c) % 2 == 1 ? Hankel{T}(c) : throw(ArgumentError(""))

#XXX Inefficient but works
getindex(H::Hankel, i, j) = getindex(full(H), i, j)
isassigned(H::Hankel, i, j) = isassigned(full(H), i, j)

size(H::Hankel, r::Int) = (r==1 || r==2) ? 1 + div(length(H.c),2) :
    throw(ArgumentError("Invalid dimension $r"))
size(H::Hankel) = size(H,1), size(H,2)

function full{T}(H::Hankel{T})
    n=size(H, 1)
    M=Array(T, n, n)
    for i=1:n
        M[:,i] = H.c[i:i+n-1]
    end
    M
end


text:
cfg_control = NeuroEvolution.get_config("../cfg/test.yaml")
cfg["n_out"]=1

ind = NEATIndiv(cfg)
ind_control = NEATIndiv(ind)
@test test_identical(ind, ind_control)

@testset "Mutation Weight" begin
    cfg = NeuroEvolution.get_config("../cfg/test.yaml")
    cfg["n_out"]=1
    cfg["p_mut_weights"]=1
    ind_mut = mutate_weight(ind, cfg)
    @test test_identical(ind, ind_control)
    @test !test_identical(ind, ind_mut)
    test_process(ind_mut, cfg)
end

@testset "Mutation Add connection" begin
    cfg = NeuroEvolution.get_config("../cfg/test.yaml")
    cfg["n_out"]=1
    ind_mut = mutate_connect(ind, cfg)
    @test test_identical(ind, ind_control)
    @test !test_identical(ind, ind_mut)
    @test length(ind.genes) < length(ind_mut.genes)
    @test cfg["innovation_max"] > cfg_control["innovation_max"]
    test_process(ind_mut, cfg)
end

@testset "New neuron position" begin
    # Not recurrent
    i = 1.0.*rand(-10:0, 100)
    j = 1.0.*rand(1:10, 100)
    n = random_position.(i, j)
    @test all(n .< j)
    @test all(n .> i)
    @test all(n .> 0.0)
    @test all(n .< 1.0)

    # Recurrent
    i = rand(1:10, 100)
    j = rand(-10:0, 100)
    n = random_position.(1.0*i, 1.0*j)
    @test all(n .< i)
    @test all(n .> j)
    @test all(n .> 0.0)
    @test all(n .< 1.0)
end

@testset "Mutation Add neuron" begin
    cfg = NeuroEvolution.get_config("../cfg/test.yaml")
    cfg["n_out"]=1
    ind_mut = mutate_neuron(ind, cfg)
    @test test_identical(ind, ind_control)
    @test !test_identical(ind, ind_mut)
    @test length(ind.neuron_pos) < length(ind_mut.neuron_pos)
    test_process(ind_mut, cfg)
end

@testset "Mutation Enable" begin
    cfg = NeuroEvolution.get_config("../cfg/test.yaml")
    cfg["n_out"]=1
    ind_mut = mutate_enabled(ind, cfg)
    @test test_identical(ind, ind_control)
    @test !test_identical(ind, ind_mut)
    enabled = true
    for g in values(ind_mut.genes)
        enabled &= g.enabled
    end
    @test ! enabled
    test_process(ind_mut, cfg)
end

@testset "Mutation" begin
    cfg = NeuroEvolution.get_config("../cfg/test.yaml")
    cfg["n_out"]=1
    cfg["p_mutate_enabled"] = 1 # Ensure there is a mutation
    cfg["p_mut_weights"]=1 # Ensure weights mutation really happens
    ind_mut = mutate(ind, cfg)
    @test test_identical(ind, ind_control)
    @test !test_identical(ind, ind_mut)
    test_process(ind_mut, cfg)
end


text: last_col() = df->names(df)[end]


text:
#*******************************************************/
#* Copyright(c) 2019 by Artelys                        */
#* This source code is subject to the terms of the     */
#* MIT Expat License (see LICENSE.md)                  */
#*******************************************************/

#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
#  This example demonstrates how to load a MPS file with Knitro MPS
#  reader.
#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

using KNITRO

# Absolute path to MPS file.
mps_file = joinpath(dirname(@__FILE__), "..", "examples", "lp.mps")

# Create a new Knitro solver instance.
kc = KNITRO.KN_new()

# Illustrate how to override default options by reading from
# the knitro.opt file.
options = joinpath(dirname(@__FILE__), "..", "examples", "knitro.opt")
KNITRO.KN_load_param_file(kc, options)

# Import MPS file inside Knitro.
err = KNITRO.KN_load_mps_file(kc, mps_file)

# Solve the problem.
# Return status codes are defined in "kn_defines.jl" and described
# in the Knitro manual.
nStatus = KNITRO.KN_solve(kc)

println("Knitro converged with final status = ", nStatus)

# An example of obtaining solution information.
nStatus, objSol, x, lambda_ =  KNITRO.KN_get_solution(kc)

# Delete the Knitro solver instance.
KNITRO.KN_free(kc)

@testset "Exemple lp1" begin
    @test nStatus == 0
    @test objSol ≈ 250 / 3 atol=1e-5
end


text:
include("utils.jl")

function test_proj_zero()
    Random.seed!(0)
    n = 100
    for _ in 1:10
        x = randn(n)
        @assert zeros(n) ≈ ConeProgramDiff._proj(x, MOI.Zeros(n))
        @assert x ≈ ConeProgramDiff._proj(x, MOI.dual_set(MOI.Zeros(n)))
    end
    return true
end

function test_proj_pos()
    Random.seed!(0)
    n = 100
    for _ in 1:10
        x = randn(n)
        @assert max.(x,zeros(n)) ≈ ConeProgramDiff._proj(x, MOI.Nonnegatives(n))
        @assert max.(x,zeros(n)) ≈ ConeProgramDiff._proj(x, MOI.dual_set(MOI.Nonnegatives(n)))
    end
    return true
end

function test_proj_soc()
    Random.seed!(0)
    n = 100
    for _ in 1:10
        x = randn(n)
        model = Model()
        set_optimizer(model, optimizer_with_attributes(SCS.Optimizer, "eps" => 1e-8, "max_iters" => 10000, "verbose" => 0))
        @variable(model, z[1:n])
        @variable(model, t)
        @objective(model, Min, t)
        @constraint(model, sum((x-z).^2) <= t)
        @constraint(model, z in MOI.SecondOrderCone(n))
        optimize!(model)

        p = ConeProgramDiff._proj(x, MOI.SecondOrderCone(n))
        @assert p ≈ value.(z)
        @assert p ≈ ConeProgramDiff._proj(x, MOI.dual_set(MOI.SecondOrderCone(n)))
    end
    return true
end

# TODO: Fix PSD cone -- Triangle vs square matrix
function test_proj_psd()
    Random.seed!(0)
    n = 10
    for _ in 1
        x = randn(n,n)
        x = x + x'
        x_vec = x[LinearAlgebra.tril(trues(size(x)))']
        model = Model()
        set_optimizer(model, optimizer_with_attributes(SCS.Optimizer, "eps" => 1e-8, "max_iters" => 10000, "verbose" => 0))
        @variable(model, z[1:n,1:n])
        @variable(model, t)
        @objective(model, Min, t)
        @constraint(model, sum((x-z).^2) <= t)
        @constraint(model, z in PSDCone())
        optimize!(model)

        z_star = value.(z)[LinearAlgebra.tril(trues(size(x)))']
        p = ConeProgramDiff._proj(x_vec, MOI.PositiveSemidefiniteConeTriangle(n))
        if !isapprox(p, z_star)
            error("Error in PSD Cone Projection:\np=$p\nz=$z_star")
        end
        @assert p ≈ ConeProgramDiff._proj(x_vec, MOI.dual_set(MOI.PositiveSemidefiniteConeTriangle(n)))
    end
    return true
end

# Test Exponential Cone
function det_case_exp_cone(v; dual=false)
    v = dual ? -v : v
    if ConeProgramDiff.in_exp_cone(v)
        return 1
    elseif ConeProgramDiff.in_exp_cone_dual(v)
        return 2
    elseif v[1] <= 0 && v[2] <= 0 #TODO: threshold here??
        return 3
    else
        return 4
    end
end

function test_proj_exp(tol)
    function _test_proj_exp_cone_help(x, tol; dual=false)
        cone = dual ? MOI.DualExponentialCone() : MOI.ExponentialCone()
        model = Model()
        set_optimizer(model, optimizer_with_attributes(SCS.Optimizer, "eps" => 1e-8, "max_iters" => 10000, "verbose" => 0))
        @variable(model, z[1:3])
        @variable(model, t)
        @objective(model, Min, t)
        @constraint(model, sum((x-z).^2) <= t)
        @constraint(model, z in cone)
        optimize!(model)
        z_star = value.(z)
        @assert isapprox(ConeProgramDiff._proj_exp_cone(x,  dual=dual), z_star, atol=tol)
    end

    Random.seed!(0)
    n = 3
    case_p = zeros(4)
    case_d = zeros(4)
    for _ in 1:100
        # x = ConeProgramDiff._proj_exp_cone(x) + randn(3)
        x = randn(3)
        # println(x)
        case_p[det_case_exp_cone(x; dual=false)] += 1
        _test_proj_exp_cone_help(x, tol; dual=false)

        case_d[det_case_exp_cone(x; dual=true)] += 1
        _test_proj_exp_cone_help(x, tol; dual=true)
    end
    @assert all(case_p .> 0) && all(case_d .> 0)
    return true
end

function det_case_pow_cone(v, α; dual=false)
    v = dual ? -v : v
    x, y, z = v

    if ConeProgramDiff.in_pow_cone(x, y, z, α)
        # if in power cone
        return 1
    elseif ConeProgramDiff.in_pow_cone_polar(x, y, z, α)
        # if in polar cone Ko = -K*
        return 2
    elseif abs(z) <= 1e-6
        # if not in K, Ko and z = 0
        return 3
    else
        return 4
    end
end

function test_proj_pow_cone(tol)
    function _test_proj_pow_cone_help(x, α, tol; dual=false)
        cone = dual ? MOI.DualPowerCone(α) : MOI.PowerCone(α)
        model = Model()
        set_optimizer(model, optimizer_with_attributes(
            SCS.Optimizer, "eps" => 1e-8, "max_iters" => 10000,
            "verbose" => 0, "scale" => 10.0))
        @variable(model, z[1:3])
        @variable(model, t)
        @objective(model, Min, t)
        @constraint(model, sum((x-z).^2) <= t)
        @constraint(model, z in cone)
        optimize!(model)
        z_star = value.(z)
        xp = ConeProgramDiff._proj_pow_cone(x, α, dual=dual)
        if !isapprox(xp, z_star, atol=tol)
            error("x = $x,\nz* = $z_star,\nxp = $xp,\ncase = $(det_case_pow_cone(x, α; dual=dual))")
        end
    end
    Random.seed!(0)
    n = 3
    case_p = zeros(4)
    case_d = zeros(4)
    for _ in 1:100
        x = randn(3)
        α = rand()*0.6 + 0.2
        # α = rand()

        # Need to get some into case 3
        if rand(1:10) == 1
            x[3] = 0
        end
        case_p[det_case_pow_cone(x, α; dual=false)] += 1
        _test_proj_pow_cone_help(x, α, tol; dual=false)

        case_d[det_case_pow_cone(x, α; dual=true)] += 1
        _test_proj_pow_cone_help(x, α, tol; dual=true)
    end

    @assert all(case_p .> 0) && all(case_d .> 0)
    return true
end

function test_project_onto_cone()
    function _test_proj_all_cones(x, cone_prod; dual=false)
        cones = dual ? [MOI.dual_set(c) for c in cone_prod] : cone_prod
        xp = project_onto_cone(x, cones)
        offset = 0
        for cone in cones
            inds = 1+offset:offset+MOI.dimension(cone)
            @assert xp[inds] ≈ ConeProgramDiff._proj(x[inds], cone)
            offset += MOI.dimension(cone)
        end
    end

    for _ in 1:10
        cone_prod = get_random_product_cone(10)
        N = sum([MOI.dimension(c) for c in cone_prod])
        local x = randn(N)
        _test_proj_all_cones(x, cone_prod; dual=false)
        _test_proj_all_cones(x, cone_prod; dual=true)
    end
    return true
end

# Test sets
# ------------------------------------------------------------------------
@testset "Projections Test (Zero, Pos, SOC, PSD)" begin
    @test test_proj_zero()
    @test test_proj_pos()
    @test test_proj_soc()
    @test test_proj_psd()
end

@testset "Projection Test: Exp Cone" begin
    @test test_proj_exp(1e-7)
end

@testset "Projection Test: Pow Cone" begin
    @test test_proj_pow_cone(1e-6)
end

@testset "Projection: project_onto_cone()" begin
    test_project_onto_cone()
end


text:
@symbol_func function c(cur_reactor::AbstractReactor)
  cur_c = G_CR(cur_reactor)

  cur_c *= safe_symbol(cur_reactor, :R_0)

  cur_c += G_CO(cur_reactor)

  cur_c
end


text:
function _tensor2matrix(tensor, axis)
    # internal function to move the axis dim to the first index, then reshape the tensor into a matrix
    dim = length(size(tensor))
    n1 = size(tensor)[axis]
    partialsize = deleteat!(collect(size(tensor)), axis) # the size of the tensor except the axis-th dimension
    n2 = reduce(*, partialsize)

    if axis == 1 #no need to permutate the axis
        return reshape(tensor, (n1, n2)), partialsize
    elseif axis == 2 && dim == 2 #for matrix, simply transpose, no copy is created
        return transpose(tensor), partialsize
    else
        permu = [i for i = 1:dim]
        permu[1], permu[axis] = axis, 1
        partialsize = collect(size(tensor)[permu][2:end])
        ntensor = permutedims(tensor, permu) # permutate the axis-th and the 1st dim, a copy of the tensor is created
        # ntensor = nocopy ? PermutedDimsArray(tensor, permu) : permutedims(tensor, permu) # permutate the axis-th and the 1st dim
        ntensor = reshape(ntensor, (n1, n2)) # no copy is created
        return ntensor, partialsize
    end
end

function _matrix2tensor(mat, partialsize, axis)
    # internal function to reshape matrix to a tensor, then swap the first index with the axis-th dimension
    @assert size(mat)[2] == reduce(*, partialsize) # total number of elements of mat and the tensor must match
    tsize = vcat(size(mat)[1], partialsize)
    tensor = reshape(mat, Tuple(tsize))
    dim = length(partialsize) + 1

    if axis == 1
        return tensor
    elseif axis == 2 && dim == 2
        return transpose(tensor) #transpose do not create copy
    else
        permu = [i for i = 1:dim]
        permu[1], permu[axis] = axis, 1
        return permutedims(tensor, permu) # permutate the axis-th and the 1st dim, a copy of the tensor is created
        # ntensor = nocopy ? PermutedDimsArray(tensor, permu) : permutedims(tensor, permu) # permutate the axis-th and the 1st dim
        # return ntensor
    end
end

function _weightedLeastSqureFit(Gτ, error, kernel)
    """
    the current algorithm with weight is only accurate up to 1e-9
    """
    # solve the linear equation: (Kᵀ⋅W⋅K) A = (Kᵀ⋅W) Gτ, where A is the spectral density to calculate, W=1/error^2
    # assume Gτ: (Nτ, N), kernel: (Nω, Nτ), weight: (Nτ, N)
    # B: (Nω, Nω),  C: (Nω, N)
    if isnothing(error)
        B = kernel
        C = Gτ
    else
        @assert size(error) == size(Gτ)
        for i = 1:size(error)[1]
            error[i, :] /= sum(error[i, :]) / length(error[i, :])
        end
        # W = Diagonal(weight)
        # B = transpose(kernel) * W * kernel
        # C = transpose(kernel) * W * Gτ
        w = 1.0 ./ error
        B = w .* kernel
        # B = Diagonal(w) * kernel
        C = w .* Gτ
    end
    # ker, ipiv, info = LAPACK.getrf!(B) # LU factorization
    # coeff = LAPACK.getrs!('N', ker, ipiv, C) # LU linear solvor for green=kernel*coeff
    coeff = B \ C #solve C = B * coeff
    return coeff
end

"""
function tau2dlr(dlrGrid::DLRGrid, green, τGrid = dlrGrid.τ; error = nothing, axis = 1)

    imaginary-time domain to DLR representation

#Members:
- `dlrGrid`: DLRGrid struct.
- `green` : green's function in imaginary-time domain.
- `τGrid` : the imaginary-time grid that Green's function is defined on.
- `error` : error the Green's function.
- `axis`: the imaginary-time axis in the data `green`.
"""
function tau2dlr(dlrGrid::DLRGrid, green, τGrid = dlrGrid.τ; error = nothing, axis = 1)
    @assert length(size(green)) >= axis "dimension of the Green's function should be larger than axis!"
    @assert size(green)[axis] == length(τGrid)
    ωGrid = dlrGrid.ω

    kernel = Spectral.kernelT(dlrGrid.isFermi, dlrGrid.symmetry, τGrid, ωGrid, dlrGrid.β, true)
    typ = promote_type(eltype(kernel), eltype(green))
    kernel = convert.(typ, kernel)
    green = convert.(typ, green)

    g, partialsize = _tensor2matrix(green, axis)

    if isnothing(error) == false
        @assert size(error) == size(green)
        error, psize = _tensor2matrix(error, axis)
    end
    coeff = _weightedLeastSqureFit(g, error, kernel)
    return _matrix2tensor(coeff, partialsize, axis)
end

"""
function dlr2tau(dlrGrid::DLRGrid, dlrcoeff, τGrid = dlrGrid.τ; axis = 1)

    DLR representation to imaginary-time representation

#Members:
- `dlrGrid` : DLRGrid
- `dlrcoeff` : DLR coefficients
- `τGrid` : expected fine imaginary-time grids
- `axis`: imaginary-time axis in the data `dlrcoeff`
"""
function dlr2tau(dlrGrid::DLRGrid, dlrcoeff, τGrid = dlrGrid.τ; axis = 1)
    @assert length(size(dlrcoeff)) >= axis "dimension of the dlr coefficients should be larger than axis!"
    @assert size(dlrcoeff)[axis] == size(dlrGrid)

    β = dlrGrid.β
    ωGrid = dlrGrid.ω

    kernel = Spectral.kernelT(dlrGrid.isFermi, dlrGrid.symmetry, τGrid, ωGrid, β, true)

    coeff, partialsize = _tensor2matrix(dlrcoeff, axis)

    G = kernel * coeff # tensor dot product: \sum_i kernel[..., i]*coeff[i, ...]

    return _matrix2tensor(G, partialsize, axis)
end

"""
function matfreq2dlr(dlrGrid::DLRGrid, green, nGrid = dlrGrid.n; error = nothing, axis = 1)

    Matsubara-frequency representation to DLR representation

#Members:
- `dlrGrid`: DLRGrid struct.
- `green` : green's function in Matsubara-frequency domain
- `nGrid` : the n grid that Green's function is defined on.
- `error` : error the Green's function.
- `axis`: the Matsubara-frequency axis in the data `green`
"""
function matfreq2dlr(dlrGrid::DLRGrid, green, nGrid = dlrGrid.n; error = nothing, axis = 1)
    @assert length(size(green)) >= axis "dimension of the Green's function should be larger than axis!"
    @assert size(green)[axis] == length(nGrid)
    @assert eltype(nGrid) <: Integer
    ωGrid = dlrGrid.ω

    kernel = Spectral.kernelΩ(dlrGrid.isFermi, dlrGrid.symmetry, nGrid, ωGrid, dlrGrid.β, true)
    typ = promote_type(eltype(kernel), eltype(green))
    kernel = convert.(typ, kernel)
    green = convert.(typ, green)

    g, partialsize = _tensor2matrix(green, axis)

    if isnothing(error) == false
        @assert size(error) == size(green)
        error, psize = _tensor2matrix(error, axis)
    end
    coeff = _weightedLeastSqureFit(g, error, kernel)
    return _matrix2tensor(coeff, partialsize, axis)
end

"""
function dlr2matfreq(dlrGrid::DLRGrid, dlrcoeff, nGrid = dlrGrid.n; axis = 1)

    DLR representation to Matsubara-frequency representation

#Members:
- `dlrGrid` : DLRGrid
- `dlrcoeff` : DLR coefficients
- `nGrid` : expected fine Matsubara-freqeuncy grids (integer)
- `axis`: Matsubara-frequency axis in the data `dlrcoeff`
"""
function dlr2matfreq(dlrGrid::DLRGrid, dlrcoeff, nGrid = dlrGrid.n; axis = 1)
    @assert length(size(dlrcoeff)) >= axis "dimension of the dlr coefficients should be larger than axis!"
    @assert size(dlrcoeff)[axis] == size(dlrGrid)
    @assert eltype(nGrid) <: Integer
    ωGrid = dlrGrid.ω

    kernel = Spectral.kernelΩ(dlrGrid.isFermi, dlrGrid.symmetry, nGrid, ωGrid, dlrGrid.β, true)

    coeff, partialsize = _tensor2matrix(dlrcoeff, axis)

    G = kernel * coeff # tensor dot product: \sum_i kernel[..., i]*coeff[i, ...]

    return _matrix2tensor(G, partialsize, axis)
end

"""
function tau2matfreq(dlrGrid, green, nNewGrid = dlrGrid.n, τGrid = dlrGrid.τ; error = nothing, axis = 1)

    Fourier transform from imaginary-time to Matsubara-frequency using the DLR representation

#Members:
- `dlrGrid` : DLRGrid
- `green` : green's function in imaginary-time domain
- `nNewGrid` : expected fine Matsubara-freqeuncy grids (integer)
- `τGrid` : the imaginary-time grid that Green's function is defined on.
- `error` : error the Green's function.
- `axis`: the imaginary-time axis in the data `green`
"""
function tau2matfreq(dlrGrid, green, nNewGrid = dlrGrid.n, τGrid = dlrGrid.τ; error = nothing, axis = 1)
    coeff = tau2dlr(dlrGrid, green, τGrid; error = error, axis = axis)
    return dlr2matfreq(dlrGrid, coeff, nNewGrid, axis = axis)
end

"""
function matfreq2tau(dlrGrid, green, τNewGrid = dlrGrid.τ, nGrid = dlrGrid.n; error = nothing, axis = 1)

    Fourier transform from Matsubara-frequency to imaginary-time using the DLR representation

#Members:
- `dlrGrid` : DLRGrid
- `green` : green's function in Matsubara-freqeuncy repsentation
- `τNewGrid` : expected fine imaginary-time grids
- `nGrid` : the n grid that Green's function is defined on.
- `error` : error the Green's function.
- `axis`: Matsubara-frequency axis in the data `green`
"""
function matfreq2tau(dlrGrid, green, τNewGrid = dlrGrid.τ, nGrid = dlrGrid.n; error = nothing, axis = 1)
    coeff = matfreq2dlr(dlrGrid, green, nGrid; error = error, axis = axis)
    return dlr2tau(dlrGrid, coeff, τNewGrid, axis = axis)
end

"""
function tau2tau(dlrGrid, green, τNewGrid, τGrid = dlrGrid.τ; weight = nothing, axis = 1)

    Interpolation from the old imaginary-time grid to a new grid using the DLR representation

#Members:
- `dlrGrid` : DLRGrid
- `green` : green's function in imaginary-time domain
- `τNewGrid` : expected fine imaginary-time grids
- `τGrid` : the imaginary-time grid that Green's function is defined on.
- `error` : error the Green's function.
- `axis`: the imaginary-time axis in the data `green`
"""
function tau2tau(dlrGrid, green, τNewGrid, τGrid = dlrGrid.τ; error = nothing, axis = 1)
    coeff = tau2dlr(dlrGrid, green, τGrid; error = error, axis = axis)
    return dlr2tau(dlrGrid, coeff, τNewGrid, axis = axis)
end

"""
function matfreq2matfreq(dlrGrid, green, nNewGrid, nGrid = dlrGrid.n; error = nothing, axis = 1)

    Fourier transform from Matsubara-frequency to imaginary-time using the DLR representation

#Members:
- `dlrGrid` : DLRGrid
- `green` : green's function in Matsubara-freqeuncy repsentation
- `nNewGrid` : expected fine Matsubara-freqeuncy grids (integer)
- `nGrid` : the n grid that Green's function is defined on.
- `error` : error the Green's function.
- `axis`: Matsubara-frequency axis in the data `green`
"""
function matfreq2matfreq(dlrGrid, green, nNewGrid, nGrid = dlrGrid.n; error = nothing, axis = 1)
    coeff = matfreq2dlr(dlrGrid, green, nGrid; error = error, axis = axis)
    return dlr2matfreq(dlrGrid, coeff, nNewGrid, axis = axis)
end

# function convolution(dlrGrid, green1, green2; axis = 1)

# end


text:
#=
Copyright 2021 BlackRock, Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
=#

@testset "Membership" begin
    S = Interval(2.0, 4.0)
    @test 3.0 ∈ S
    @test 2.0 ∈ S
    @test 4.0 ∈ S
    @test !(5.0 ∈ S)
end

@testset "Inclusion" begin
    @test Interval(2.0, 4.0) ⊆ Interval(1.0, 5.0)
    @test !(Interval(1.0, 5.0) ⊆ Interval(2.0, 4.0))
    @test !(Interval(2.0, 5.0) ⊆ Interval(1.0, 4.0))
end

@testset "Is empty" begin
    @test !isempty(Interval(2.0, 4.0))
    @test isempty(Interval(4.0, 2.0))
    @test !isempty(Interval(4.0, 4.0))
end

@testset "Equals" begin
    @test Interval(2.0, 4.0) == Interval(2.0, 4.0)
    @test !(Interval(2.0, 4.0) == Interval(3.0, 5.0))
    @test Interval(2.0, 4.0) != Interval(3.0, 5.0)
    @test !(Interval(2.0, 4.0) != Interval(2.0, 4.0))
end

@testset "Approx. Equals" begin
    ϵ = 1e-12
    @test !(Interval(2 + ϵ, 4 + ϵ) ≈ Interval(2.0, 4.0))
    ϵ /= 10
    @test Interval(2 + ϵ, 4 + ϵ) ≈ Interval(2.0, 4.0)
    @test !(Interval(2 + ϵ, 4 + ϵ) ≈ Interval(3.0, 5.0))
end

@testset "Intersection" begin
    @test Interval(2.0, 4.0) ∩ Interval(1.0, 5.0) ≈ Interval(2.0, 4.0)
    @test Interval(2.0, 5.0) ∩ Interval(1.0, 4.0) ≈ Interval(2.0, 4.0)
    @test Interval(1.0, 4.0) ∩ Interval(2.0, 5.0) ≈ Interval(2.0, 4.0)
    @test isempty(Interval(1.0, 2.0) ∩ Interval(4.0, 5))
end

@testset "Less than" begin
    @test Interval(2.0, 4.0) < Interval(5.0, 7.0)
    @test !(Interval(2.0, 4.0) < Interval(1.0, 3.0))
end

@testset "Greater than" begin
    @test Interval(5.0, 7.0) > Interval(2.0, 4.0)
    @test !(Interval(1.0, 3.0) > Interval(2.0, 4.0))
end


text:
using LightXML
"""
Verision Comment
v0.1     hf, firstversion
v0.2     hf, use exiftool instead of tiffcommit

ref: http://web.mit.edu/jhawk/mnt/cgs/Image-ExifTool-6.99/html/exiftool_pod.html
    https://exiftool.org/TagNames/EXIF.html
    https://docs.openmicroscopy.org/ome-model/5.6.3/ome-tiff/specification.html
    https://micro-manager.org/wiki/Micro-Manager_File_Formats
"""

"""
Create and return XML string with xyzt
"""
function generatexml(SizeX::Integer, SizeY::Integer, SizeZ::Integer, SizeT::Integer)
    tiffXML = XMLDocument();
    OME = create_root(tiffXML, "OME");
    xmlns="http://www.openmicroscopy.org/Schemas/OME/2016-06"
    xmlns_xsi="http://www.w3.org/2001/XMLSchema-instance"
    Creator="OME Bio-Formats 5.2.2"
    UUID="urn:uuid:a5ae8c1b-ac04-4544-97c1-bbdd0bdf8629" # use fixed value now
    xsi_schemaLocation="http://www.openmicroscopy.org/Schemas/OME/2016-06 http://www.openmicroscopy.org/Schemas/OME/2016-06/ome.xsd"
    OME_attributes = Dict("xmlns"=>xmlns, "xmlns:xsi"=>xmlns_xsi, "Creator"=>Creator, "UUID"=>UUID, "xsi:schemaLocation"=>xsi_schemaLocation)
    set_attributes(OME, OME_attributes)

    name = "test.tiff"
    Image= new_child(OME, "Image")
    set_attributes(Image, Dict("ID"=>"Image:0", "Name"=>name))

    Pixels = new_child(Image, "Pixels");
    SizeC=1;
    #SizeT=10;  #SizeX=1300; #SizeY=1900; #SizeZ=2;
    PhysicalSizeX=PhysicalSizeY=0.108; PhysicalSizeZ=0.5;
    TimeIncrement=600;
    Pixels_attributes = Dict("ID"=>"Pixels:0","Type"=>"uint16", "BigEndian"=>"false",
    "DimensionOrder"=>"XYZCT", "SizeC"=>1,"SizeX"=>SizeX, "SizeY"=>SizeY, "SizeZ"=>SizeZ, "SizeT"=>SizeT,
    "PhysicalSizeX"=>PhysicalSizeX, "PhysicalSizeY"=>PhysicalSizeY, "PhysicalSizeZ"=>PhysicalSizeZ,
    "PhysicalSizeXUnit"=>"µm", "PhysicalSizeYUnit"=>"µm", "PhysicalSizeZUnit"=>"µm",
    "TimeIncrement"=>TimeIncrement, "TimeIncrementUnit"=>"s");
    set_attributes(Pixels, Pixels_attributes)

    Channel_ = new_child(Pixels, "Channel");
    set_attributes(Channel_, Dict("ID"=>"Channel:0:0", "SamplesPerPixel"=>"1"));
    LightPath =new_child(Channel_, "LightPath");

    for i in 1:SizeC*SizeZ*SizeT
        new_child(Pixels, "TiffData")
    end
	"""
	for c in 0:SizeC-1
    	for t in 0:SizeT-1
        	for z in 0:SizeZ-1
               TiffData = new_child(Pixels, "TiffData")
               TiffData_attributes = Dict("FirstC"=>c, "FirstT"=>t, "FirstZ"=>z,
										  "IFD"=>t*SizeZ+z,"PlaneCount"=>1)
                set_attributes(TiffData, TiffData_attributes)
            end
        end
    end
	"""
    tiffXML;
end

"""
Embed OME-XML to tiff with XYZT informance
"""
function embedxml(SizeX::Integer, SizeY::Integer, SizeZ::Integer, SizeT::Integer, img_name::String)
    tiffxml = generatexml(SizeX, SizeY, SizeZ, SizeT);
    print("Embedding OME-XML")
    #run(`/home/hf/Bin/bftools/tiffcomment -set $tiffxml $img_name`);
    if Sys.which("exiftool") == Nothing
        print("Please install exiftool")
    else
		run(`exiftool -ImageDescription=$tiffxml $img_name -overwrite_original`, wait=false)
		#run(pipeline(`echo $tiffxml`, `exiftool -ImageDescription\<= $img_name -overwrite_original`), wait=false)
		#run(`ulimit -s 65536 ; exiftool -ImageDescription\<= $img_name -overwrite_original`, wait=true)
		# run in backgroud to continue next computing work
        # I assume that run computing takes more time to embed xml
        #run(`rm $img_name"_original"`)
    end
    println("")
    nothing;
end


text:
@params struct FunctionEvaluations
    obj::Int
    constr
end
function FunctionEvaluations(optimizer)
    obj_fevals = getfevals(optimizer.obj)
    constr_fevals = getfevals.(optimizer.constr)
    return FunctionEvaluations(obj_fevals, constr_fevals)
end

Base.zero(t::FunctionEvaluations) = FunctionEvaluations(zero(t.obj), zero.(t.constr))
for f in (:-, :+)
    @eval begin
        function Base.$(f)(f1::FunctionEvaluations, f2::FunctionEvaluations)
            obj_fevals = $(f)(f1.obj, f2.obj)
            constr_fevals = $(f).(f1.constr, f2.constr)
            return FunctionEvaluations(obj_fevals, constr_fevals)
        end
    end
end
Base.broadcastable(f::FunctionEvaluations) = Ref(f)
function Base.all(f, fevals::FunctionEvaluations)
    f(fevals.obj) && all(f, fevals.constr)
end
function Base.any(f, fevals::FunctionEvaluations)
    f(fevals.obj) || any(f, fevals.constr)
end

@params mutable struct SIMPResult{T, TF <: FunctionEvaluations}
    topology::AbstractVector{T}
    objval::T
    fevals::TF
    convstate
    penalty_trace::AbstractVector{Pair{T, TF}}
    nsubproblems::Int
end
GPUUtils.whichdevice(s::SIMPResult) = whichdevice(s.topology)

function NewSIMPResult(::Type{T}, optimizer, ncells) where {T}
    FunctionEvaluations(optimizer)
    fevals = zero(FunctionEvaluations(optimizer))
    SIMPResult(fill(T(NaN), ncells), T(NaN), fevals, MMA.ConvergenceState(), Pair{T, typeof(fevals)}[], 0)
end

@params mutable struct SIMP{T} <: AbstractSIMP
    optimizer
    penalty
    prev_penalty
    result::SIMPResult{T}
    tracing::Bool
end
GPUUtils.whichdevice(s::SIMP) = whichdevice(s.optimizer)

function SIMP(optimizer, p::T; tracing=true) where T
    penalty = getpenalty(optimizer)
    penalty = @set penalty.p = p
    prev_penalty = @set penalty.p = NaN
    ncells = getncells(optimizer.obj.f.problem)
    result = NewSIMPResult(T, optimizer, ncells)

    return SIMP(optimizer, penalty, prev_penalty, result, tracing)
end

Utilities.getpenalty(s::AbstractSIMP) = s.penalty
function Utilities.setpenalty!(s::AbstractSIMP, p::Number)
    penalty = s.penalty
    s.prev_penalty = penalty
    s.penalty = @set penalty.p = p
    setpenalty!(s.optimizer, p)
end

function (s::SIMP{T, TO})(x0=s.optimizer.obj.f.solver.vars, prev_fevals = getfevals(s.optimizer)) where {T, TO<:MMAOptimizer}
    setpenalty!(s.optimizer, s.penalty.p)
    mma_results = s.optimizer(x0)
    update_result!(s, mma_results, prev_fevals)
    return s.result
end

function (s::SIMP{T, TO})(workspace::MMA.Workspace, prev_fevals = getfevals(s.optimizer)) where {T, TO <: MMAOptimizer}
    mma_results = s.optimizer(workspace)
    update_result!(s, mma_results, prev_fevals)
    return s.result
end

function get_topologies(problem, trace::TopOptTrace)
    @unpack black, white, varind = problem
    @unpack x_hist = trace
    nel = length(black)
    topologies = Vector{Float64}[]
    topology = zeros(T, nel)
    for i in 1:length(x_hist)
        update_topology!(topology, black, white, x_hist[i], varind)
        push!(topologies, copy(topology))
    end
    return topologies
end

function update_result!(s::SIMP{T}, mma_results, prev_fevals) where T
    # Postprocessing
    @unpack result, optimizer = s
    @unpack obj = optimizer
    @unpack problem = obj.f
    @unpack black, white, varind = problem
    @unpack x_hist = obj.f.topopt_trace
    nel = getncells(problem)

    update_topology!(result.topology, black, white, mma_results.minimizer, varind)
    result.objval = mma_results.minimum

    new_fevals = getfevals(optimizer)
    extra_fevals = new_fevals - prev_fevals
    result.fevals += extra_fevals
    if s.tracing
        push!(result.penalty_trace, (getpenalty(s).p => extra_fevals))
    end
    if all(x -> x > 0, extra_fevals)
        result.nsubproblems += 1
    end
    result.convstate = deepcopy(mma_results.convstate)
    return result
end

function update_topology!(topology, black, white, x, varind)
    for i in 1:length(black)
        if black[i]
            topology[i] = 1
        elseif white[i]
            topology[i] = 0
        else
            topology[i] = x[varind[i]]
        end
    end

    return
end


text:
using DashBootstrapComponents, DashHtmlComponents

modal = html_div([
    dbc_button("Small modal", id="open-sm", className="mr-1", n_clicks=0),
    dbc_button("Large modal", id="open-lg", className="mr-1", n_clicks=0),
    dbc_button("Extra large modal", id="open-xl", n_clicks=0),
    dbc_modal(
        [
            dbc_modalheader("Header"),
            dbc_modalbody("A small modal."),
            dbc_modalfooter(
                dbc_button("Close", id="close-sm", className="ml-auto", n_clicks=0),
            ),
        ],
        id="modal-sm",
        size="sm",
        is_open=false,
    ),
    dbc_modal(
        [
            dbc_modalheader("Header"),
            dbc_modalbody("A large modal."),
            dbc_modalfooter(
                dbc_button("Close", id="close-lg", className="ml-auto", n_clicks=0),
            ),
        ],
        id="modal-lg",
        size="lg",
        is_open=false,
    ),
    dbc_modal(
        [
            dbc_modalheader("Header"),
            dbc_modalbody("An extra large modal."),
            dbc_modalfooter(
                dbc_button("Close", id="close-xl", className="ml-auto", n_clicks=0),
            ),
        ],
        id="modal-xl",
        size="xl",
        is_open=false,
    ),
]);

function toggle_modal(n1, n2, is_open)
    return n1 > 0 || n2 > 0 ? is_open == 0 : is_open
end;

callback!(
    app,
    Output("modal-sm", "is_open"),
    Input("open-sm", "n_clicks"),
    Input("close-sm", "n_clicks"),
    State("modal-sm", "is_open"),
) do n1, n2, is_open
    return toggle_modal(n1, n2, is_open)
end;

callback!(
    app,
    Output("modal-lg", "is_open"),
    Input("open-lg", "n_clicks"),
    Input("close-lg", "n_clicks"),
    State("modal-lg", "is_open"),
) do n1, n2, is_open
    return toggle_modal(n1, n2, is_open)
end;

callback!(
    app,
    Output("modal-xl", "is_open"),
    Input("open-xl", "n_clicks"),
    Input("close-xl", "n_clicks"),
    State("modal-xl", "is_open"),
) do n1, n2, is_open
    return toggle_modal(n1, n2, is_open)
end;

max_stars_repo_name: harisbal/dash-bootstrap-components

max_issues_repo_name: harisbal/dash-bootstrap-components

max_forks_repo_name: harisbal/dash-bootstrap-components


text:
using Test
using EBayes
using Random
using MLJ

Random.seed!(0)
ffh_sim = EBayes.FriedmanFayHerriotSimulation(1000, 5 , 0.5, 1.0)
ffh_sample = rand(ffh_sim)

lin_mlj = MLJ.@load LinearRegressor pkg=GLM
lin_mach = fit!(machine(lin_mlj, ffh_sample.X, response(ffh_sample.ss)), verbosity=0)
lin_preds = predict_mean(lin_mach)

mse_linpred = mean((lin_preds .- ffh_sample.true_μs).^2)
mse_mle = mean((response(ffh_sample.ss) - ffh_sample.true_μs).^2)

ebcf_lin = EBayesCrossFit(lin_mlj)
ebcf_fit = fit(ebcf_lin, ffh_sample.X, ffh_sample.ss)

ebcf_preds = predict(ebcf_fit)

mse_ebcf = mean((ebcf_preds .- ffh_sample.true_μs).^2)
@test mse_ebcf < mse_linpred
@test mse_ebcf < mse_mle


text:
"""
    kron_id(n, k)

Compute the k-fold Kronecker product of the identity: `I ⊗ I ⊗ ... ⊗ I`, k times.

### Input

- `n` -- integer representing the order (dimension of the identity)
- `k` -- integer representing the power

### Output

A `Diagonal` matrix with element `1` in the diagonal and order ``n^k``.

### Examples

```jldoctest
julia> kron_id(2, 2)
4×4 Diagonal{Float64,Array{Float64,1}}:
 1.0   ⋅    ⋅    ⋅
  ⋅   1.0   ⋅    ⋅
  ⋅    ⋅   1.0   ⋅
  ⋅    ⋅    ⋅   1.0
```
"""
function kron_id(n::Int, k::Int)
    return Diagonal(ones(n^k))
end

"""
    kron_sandwich(F, n, k1, k2)

Compute `A ⊗ F ⊗ C` where `A = I^{⊗ k1}` and `C = I^{⊗ k2}` where `I` is the
identity matrix of order `n`, the `k1`-fold and `k2`-fold Kronecker product of the
identity and `F` in between.

### Input

- `F`  -- matrix
- `n`  -- integer, dimension of the identity
- `k1` -- nonnegative integer
- `k2` -- nonnegative integer

### Output

The kronecker product `I^{⊗ k1} ⊗ F ⊗ I^{⊗ k2}`, represented as a sparse matrix
if `F` is sparse.

### Examples

```jldoctest
julia> F = sparse([0 1; -1 0.])
2×2 SparseMatrixCSC{Float64,Int64} with 2 stored entries:
  [2, 1]  =  -1.0
  [1, 2]  =  1.0

julia> Q = kron_sandwich(F, 2, 2, 2);

julia> size(Q)
(32, 32)

julia> I2 = I(2)
IdentityMultiple{Float64} of value 1.0 and order 2

julia> Q == reduce(kron, [I2, I2, F, I2, I2])
true
```
"""
function kron_sandwich(F::AbstractMatrix, n::Int, k1::Int, k2::Int)
    # compute A ⊗ F
    if k1 == 0
        AF = F
    else
        A = kron_id(n, k1)
        AF = kron(A, F)
    end

    # compute A ⊗ F ⊗ C
    if k2 == 0
        AFC = AF
    else
        C = kron_id(n, k2)
        AFC = kron(AF, C)
    end

    return AFC
end

"""
    kron_sum(F::AbstractMatrix, k::Int)

Compute the Kronecker sum of order k defined as:

```math
    F ⊕ F ⊕ ... ⊕ F := F ⊗ I ⊗ ... ⊗ I + I ⊗ F ⊗ I ⊗ ... ⊗ I + ... + I ⊗ ... ⊗ I ⊗ F
```
where each term has `k` products and there are a total of `k` summands
and `I` is the identity matrix of order `n`.

### Examples

It holds that:

- ``k = 1: F``
- ``k = 2: F ⊗ I + I ⊗ F``
- ``k = 3: F ⊗ I ⊗ I + I ⊗ F ⊗ I + I ⊗ I ⊗ F``

```jldoctest
julia> F = sparse([0 1; -1 0.])
2×2 SparseMatrixCSC{Float64,Int64} with 2 stored entries:
  [2, 1]  =  -1.0
  [1, 2]  =  1.0

julia> kron_sum(F, 1) == F
true

julia> I2 = I(2)
IdentityMultiple{Float64} of value 1.0 and order 2

julia> kron_sum(F, 2) == kron(F, I2) + kron(I2, F)
true

julia> kron_sum(F, 3) == sum(reduce(kron, X) for X in [[F, I2, I2], [I2, F, I2], [I2, I2, F]])
true
```
"""
function kron_sum(F::AbstractMatrix, k::Int)
    if k < 1
        throw(ArgumentError("expected k ≥ 1, got $k"))
    elseif k == 1
        return F
    end

    n = size(F, 1) # leading dimension
    k1 = 0 # terms on the left of F
    k2 = k-1 # terms on the right of F
    A = kron_sandwich(F, n, k1, k2)  # I^(⊗k1) ⊗ F ⊗ I^(⊗k2)
    for i in 2:k
        k1 += 1
        k2 -= 1
        B = kron_sandwich(F, n, k1, k2)
        A += B
    end
    return A
end

"""
    kron_pow(x::Vector{<:AbstractVariable}, pow::Int)

Compute the higher order concrete Kronecker power: `x ⊗ x ⊗ ... ⊗ x`, `pow` times
for a vector of symbolic monomials.

### Input

- `x`   -- polynomial variable
- `pow` -- integer

### Output

Vector of multivariate monomial corresponding to `x^{⊗ pow}`.

### Examples

```jldoctest
julia> using DynamicPolynomials

julia> @polyvar x[1:2]
(PolyVar{true}[x₁, x₂],)

julia> x
2-element Vector{PolyVar{true}}:
 x₁
 x₂

julia> kron_pow(x, 2)
4-element Vector{Monomial{true}}:
 x₁²
 x₁x₂
 x₁x₂
 x₂²
```
"""
function kron_pow(x::Vector{<:AbstractVariable}, pow::Int)
    @assert pow > 0 "expected positive power, got $pow"
    if pow == 1
        return x
    else
        return kron(x, kron_pow(x, pow-1))
    end
end

kron_pow(x::NTuple, pow::Int) = kron_pow([xi for xi in x], pow)


text:
using StructuredMechModels
using Test

using BatchedRoutines
using Zygote
using TrajectoryOptimization
using Flux

function ngradient(f, xs::AbstractArray...)
    grads = zero.(xs)
    for (x, Δ) in zip(xs, grads), i in 1:length(x)
        δ = sqrt(eps())
        tmp = x[i]
        x[i] = tmp - δ/2
        y1 = f(xs...)
        x[i] = tmp + δ/2
        y2 = f(xs...)
        x[i] = tmp
        Δ[i] = (y2-y1)/δ
    end
    return grads
end

gradcheck(f, xs...) =
    all(isapprox.(ngradient(f, xs...),
                  gradient(f, xs...), rtol = 1e-5, atol = 1e-5))

gradtest(f, xs::AbstractArray...) = gradcheck((xs...) -> sum(sin.(f(xs...))), xs...)

@testset "StructuredMechModels.jl" begin

    @testset "batchmul" begin
        A = rand(4, 2, 3)
        B = rand(2, 4, 3)
        C = StructuredMechModels.batchmul(A, B)
        C_f = batched_gemm('N', 'N', A, B)
        C_s = StructuredMechModels.batchmul_slow(A, B)
        @test all(isapprox(C, C_s, rtol=1e-5,atol=1e-5))
        @test all(isapprox(C_f, C_s, rtol=1e-5,atol=1e-5))
    end

    @testset "grads" begin
        x = rand(4, 5)
        u = rand(1, 5)
        target = rand(4, 5)
        nn = StructuredMechModels.DeLan(2, 1, [1, 1, 0, 0])
        g = Zygote.gradient(() -> Flux.mse(nn(x, u), target), params(nn))
        @test gradtest(nn, x, u)
    end

    @testset "pytorch" begin
        @testset "$mtype" for mtype in (:Naive, :DeLan, :ControlAffine)
            @testset "$act" for act in ("ReLU", "SoftPlus", "Tanh")
                nn_torch = getfield(StructuredMechModels.Torch, mtype)(2, 1, [1, 1, 0, 0], activation=act)
                nn = getfield(StructuredMechModels, mtype)(2, 1, [1, 1, 0, 0], activation=getfield(StructuredMechModels, Symbol(act))())
                Flux.loadparams!(nn, StructuredMechModels.Torch.params(nn_torch))
                x = rand(4, 8);
                u = rand(1, 8);
                @test all(isapprox.(nn(x, u), nn_torch(x, u)))
            end
        end
    end

    @testset "model" begin
        @testset "$mtype" for mtype in (:Naive, :DeLan, :ControlAffine)
            @testset "$act" for act in ("ReLU", "SoftPlus", "Tanh")
                nn = getfield(StructuredMechModels, mtype)(2, 1, [1, 1, 0, 0], activation=getfield(StructuredMechModels, Symbol(act))())
                model = Model(nn)
                x = rand(4, 8);
                u = rand(1, 8);
                @test all(isapprox.(nn(x, u), model.params.f(x, u)))

                @testset "jac" begin
                    model_d = TrajectoryOptimization.discretize_model(model, :rk4)
                    x = rand(4)
                    u = rand(1)
                    Z = TrajectoryOptimization.PartedMatrix(model_d)
                    @test sum(TrajectoryOptimization.jacobian!(Z, model_d, x, u, 0.1)) != 0.0
                end
            end
        end
    end
end


text:
using JLD2
using Statistics

mutable struct ISRatio <: LoggerKeyData
    isRatio::Array{Any}
    log_interval::Int

    function ISRatio(logger_init_info)
        num_logged_steps = fld(logger_init_info[LoggerInitKey.TOTAL_STEPS], logger_init_info[LoggerInitKey.INTERVAL])
        println("The ISRatio logger might be pretty inefficient and should probably only be used for Auto debug. If you are planning on running a sweep, it might be worth not logging step sizes.")
        new([], logger_init_info[LoggerInitKey.INTERVAL])
    end
end

function lg_step!(self::ISRatio, env, agent, s, a, s_next, r, is_terminal, cur_step_in_episode, cur_step_total)
    if rem(cur_step_total, self.log_interval) == 0
        ind = fld(cur_step_total, self.log_interval)
        learner = agent.demon_learner
        lu = update(learner)
        @assert lu isa ETB ||  lu isa PriorTB "Update is not ETB or PriorTB Please set it to ETB or PriorTB  (or add my emphasis type here) or remove LoggerKey.IS_RATIO from logger keys"

        rho = lu.rho_logging
        push!(self.isRatio, copy(rho))
    end
end

function lg_episode_end!(self::ISRatio, cur_step_in_episode, cur_step_total)
end

function save_log(self::ISRatio, save_dict::Dict)
    # Throwing away the first info.
    save_dict[:is_ratio] = cat(self.isRatio[2:end]..., dims=2)
end

max_stars_repo_name: MatthewMcLeod/curiosity

max_issues_repo_name: MatthewMcLeod/curiosity

max_forks_repo_name: MatthewMcLeod/curiosity


text:
#=
Various routines to discretize AR(1) processes

@author : Spencer Lyon <spencer.lyon@nyu.edu>

@date : 2014-04-10 23:55:05

References
----------

https://lectures.quantecon.org/jl/finite_markov.html
=#

import Optim
import NLopt
import Distributions: pdf, Normal, quantile

std_norm_cdf(x::T) where {T <: Real} = 0.5 * erfc(-x/sqrt(2))
std_norm_cdf(x::Array{T}) where {T <: Real} = 0.5 .* erfc(-x./sqrt(2))

@doc doc"""
Tauchen's (1996) method for approximating AR(1) process with finite markov chain

The process follows

```math
    y_t = \mu + \rho y_{t-1} + \epsilon_t
```

where ``\epsilon_t \sim N (0, \sigma^2)``

##### Arguments

- `N::Integer`: Number of points in markov process
- `ρ::Real` : Persistence parameter in AR(1) process
- `σ::Real` : Standard deviation of random component of AR(1) process
- `μ::Real(0.0)` : Mean of AR(1) process
- `n_std::Integer(3)` : The number of standard deviations to each side the process
  should span

##### Returns

- `mc::MarkovChain` : Markov chain holding the state values and transition matrix

"""
function tauchen(N::Integer, ρ::T1, σ::T2, μ=zero(promote_type(T1, T2)), n_std::Integer=3) where {T1 <: Real, T2 <: Real}
    # Get discretized space
    a_bar = n_std * sqrt(σ^2 / (1 - ρ^2))
    y = range(-a_bar, stop=a_bar, length=N)
    d = y[2] - y[1]

    # Get transition probabilities
    Π = zeros(promote_type(T1, T2), N, N)
    for row = 1:N
        # Do end points first
        Π[row, 1] = std_norm_cdf((y[1] - ρ*y[row] + d/2) / σ)
        Π[row, N] = 1 - std_norm_cdf((y[N] - ρ*y[row] - d/2) / σ)

        # fill in the middle columns
        for col = 2:N-1
            Π[row, col] = (std_norm_cdf((y[col] - ρ*y[row] + d/2) / σ) -
                           std_norm_cdf((y[col] - ρ*y[row] - d/2) / σ))
        end
    end

    # NOTE: I need to shift this vector after finding probabilities
    #       because when finding the probabilities I use a function
    #       std_norm_cdf that assumes its input argument is distributed
    #       N(0, 1). After adding the mean E[y] is no longer 0, so
    #       I would be passing elements with the wrong distribution.
    #
    #       It is ok to do after the fact because adding this constant to each
    #       term effectively shifts the entire distribution. Because the
    #       normal distribution is symmetric and we just care about relative
    #       distances between points, the probabilities will be the same.
    #
    #       I could have shifted it before, but then I would need to evaluate
    #       the cdf with a function that allows the distribution of input
    #       arguments to be [μ/(1 - ρ), 1] instead of [0, 1]

    yy = y .+ μ / (1 - ρ) # center process around its mean (wbar / (1 - rho)) in new variable

    # renormalize. In some test cases the rows sum to something that is 2e-15
    # away from 1.0, which caused problems in the MarkovChain constructor
    Π = Π./sum(Π, dims = 2)

    MarkovChain(Π, yy)
end

@doc doc"""
Rouwenhorst's method to approximate AR(1) processes.

The process follows

```math
    y_t = \mu + \rho y_{t-1} + \epsilon_t
```

where ``\epsilon_t \sim N (0, \sigma^2)``

##### Arguments
- `N::Integer` : Number of points in markov process
- `ρ::Real` : Persistence parameter in AR(1) process
- `σ::Real` : Standard deviation of random component of AR(1) process
- `μ::Real(0.0)` :  Mean of AR(1) process

##### Returns

- `mc::MarkovChain{Float64}` : Markov chain holding the state values and transition matrix

"""
function rouwenhorst(N::Integer, ρ::Real, σ::Real, μ::Real=0.0)
    σ_y = σ / sqrt(1-ρ^2)
    p  = (1+ρ)/2
    Θ = [p 1-p; 1-p p]
    ψ = sqrt(N-1) * σ_y
    m = μ / (1 - ρ)

    state_values, p = _rouwenhorst(p, p, m, ψ, N)
    MarkovChain(p, state_values)
end

function _rouwenhorst(p::Real, q::Real, m::Real, Δ::Real, n::Integer)
    if n == 2
        return [m-Δ, m+Δ],  [p 1-p; 1-q q]
    else
        _, θ_nm1 = _rouwenhorst(p, q, m, Δ, n-1)
        θN = p    *[θ_nm1 zeros(n-1, 1); zeros(1, n)] +
             (1-p)*[zeros(n-1, 1) θ_nm1; zeros(1, n)] +
             q    *[zeros(1, n); zeros(n-1, 1) θ_nm1] +
             (1-q)*[zeros(1, n); θ_nm1 zeros(n-1, 1)]

        θN[2:end-1, :] ./= 2

        return range(m-Δ, stop=m+Δ, length=n), θN
    end
end

# These are to help me order types other than vectors
@inline _emcd_lt(a::T, b::T) where {T} = isless(a, b)
# @inline _emcd_lt(a::Vector{T}, b::Vector{T}) where {T} = Base.lt(Base.Order.Lexicographic, a, b)

@doc doc"""
Accepts the simulation of a discrete state Markov chain and estimates
the transition probabilities

Let ``S = s_1, s_2, \ldots, s_N`` with ``s_1 < s_2 < \ldots < s_N`` be the discrete
states of a Markov chain. Furthermore, let ``P`` be the corresponding
stochastic transition matrix.

Given a history of observations, ``\{X\}_{t=0}^{T}`` with ``x_t \in S \forall t``,
we would like to estimate the transition probabilities in ``P`` with ``p_{ij}``
as the ith row and jth column of ``P``. For ``x_t = s_i`` and ``x_{t-1} = s_j``,
let ``P(x_t | x_{t-1})`` be defined as ``p_{i,j}`` element of the stochastic
matrix. The likelihood function is then given by

```math
  L(\{X\}^t; P) = \text{Prob}(x_1) \prod_{t=2}^{T} P(x_t | x_{t-1})
```

The maximum likelihood estimate is then just given by the number of times
a transition from ``s_i`` to ``s_j`` is observed divided by the number of times
``s_i`` was observed.

Note: Because of the estimation procedure used, only states that are observed
in the history appear in the estimated Markov chain... It can't divine whether
there are unobserved states in the original Markov chain.

For more info, refer to:

- http://www.stat.cmu.edu/~cshalizi/462/lectures/06/markov-mle.pdf
- https://stats.stackexchange.com/questions/47685/calculating-log-likelihood-for-given-mle-markov-chains

##### Arguments

- `X::Vector{T}` : Simulated history of Markov states

##### Returns

- `mc::MarkovChain{T}` : A Markov chain holding the state values and
  transition matrix

"""
function estimate_mc_discrete(X::Vector{T}, states::Vector{T}) where T
    # Get length of simulation
    capT = length(X)

    # Make sure all of the passed in states appear in X... If not
    # throw an error
    if any(!in(x, X) for x in states)
        error("One of the states does not appear in history X")
    end

    # Count states and store in dictionary
    nstates = length(states)
    d = Dict{T, Int}(zip(states, 1:nstates))

    # Counter matrix and dictionary mapping i -> states
    cm = zeros(nstates, nstates)

    # Compute conditional probabilities for each state
    state_i = d[X[1]]
    for t in 1:capT-1
        # Find next period's state
        state_j = d[X[t+1]]
        cm[state_i, state_j] += 1.0

        # Tomorrow's state is j
        state_i = state_j
    end

    # Compute probabilities using counted elements
    P = cm ./ sum(cm, dims = 2)

    return MarkovChain(P, states)
end

function estimate_mc_discrete(X::Vector{T}) where T
    # Get unique states and sort them
    states = sort!(unique(X); lt=_emcd_lt)

    return estimate_mc_discrete(X, states)
end

@doc doc"""

types specifying the method for `discrete_var`

"""
abstract type VAREstimationMethod end
struct Even <: VAREstimationMethod end
struct Quantile <: VAREstimationMethod end
struct Quadrature <: VAREstimationMethod end

@doc doc"""

Compute a finite-state Markov chain approximation to a VAR(1) process of the form

```math
    y_{t+1} = b + By_{t} + \Psi^{\frac{1}{2}}\epsilon_{t+1}
```

where ``\epsilon_{t+1}`` is an vector of independent standard normal
innovations of length `M`

```julia
P, X = discrete_var(b, B, Psi, Nm, n_moments, method, n_sigmas)
```

##### Arguments

- `b::Union{Real, AbstractVector}` : constant vector of length `M`.
                                     `M=1` corresponds scalar case
- `B::Union{Real, AbstractMatrix}` : `M x M` matrix of impact coefficients
- `Psi::Union{Real, AbstractMatrix}` : `M x M` variance-covariance matrix of
                                       the innovations
    - `discrete_var` only accepts non-singular variance-covariance matrices, `Psi`.
- `Nm::Integer > 3` : Desired number of discrete points in each dimension

##### Optional

- `n_moments::Integer` : Desired number of moments to match. The default is 2.
- `method::VAREstimationMethod` : Specify the method used to determine the grid
                                  points. Accepted inputs are `Even()`, `Quantile()`,
                                  or `Quadrature()`. Please see the paper for more details.
- `n_sigmas::Real` : If the `Even()` option is specified, `n_sigmas` is used to
                     determine the number of unconditional standard deviations
                     used to set the endpoints of the grid. The default is
                     `sqrt(Nm-1)`.

##### Returns

- `P` : `Nm^M x Nm^M` probability transition matrix. Each row
        corresponds to a discrete conditional probability
        distribution over the state M-tuples in `X`
- `X` : `M x Nm^M` matrix of states. Each column corresponds to an
        M-tuple of values which correspond to the state associated
        with each row of `P`

##### NOTES

- discrete_var only constructs tensor product grids where each dimension
  contains the same number of points. For this reason it is recommended
  that this code not be used for problems of more than about 4 or 5
  dimensions due to curse of dimensionality issues.
- Future updates will allow for singular variance-covariance matrices and
  sparse grid specifications.

##### Reference

- Farmer, L. E., & Toda, A. A. (2017).
  "Discretizing nonlinear, non‐Gaussian Markov processes with exact conditional moments,"
  Quantitative Economics, 8(2), 651-683.

"""
function discrete_var(b::Union{Real, AbstractVector},
                      B::Union{Real, AbstractMatrix},
                      Psi::Union{Real, AbstractMatrix},
                      Nm::Integer,
                      n_moments::Integer=2,
                      method::VAREstimationMethod=Even(),
                      n_sigmas::Real=sqrt(Nm-1))
    # b = zeros(2)
    # A = [0.9809 0.0028; 0.041 0.9648]
    # Sigma = [7.569e-5 0.0; 0.0 0.00068644]
    # N = 9
    # n_moments = nMoments
    # method = Quantile()
    # b, B, Psi, Nm = (zeros(2), A, Sigma, N, nMoments, Quantile())
    M, M_ = size(B, 1), size(B, 2)

    # Check size restrictions on matrices
    M == M_ || throw(ArgumentError("B must be a scalar or square matrix"))
    M == length(b) || throw(ArgumentError("b must have the same number of rows as B"))

    #% Check that Psi is a valid covariance matrix
    isposdef(Psi) || throw(ArgumentError("Psi must be a positive definite matrix"))

    # Check that Nm is a valid number of grid points
    Nm >= 3 || throw(ArgumentError("Nm must be a positive interger greater than 3"))

    # Check that n_moments is a valid number
    if n_moments < 1 || !(n_moments % 2 == 0 || n_moments == 1)
        error("n_moments must be either 1 or a positive even integer")
    end

    # warning about persistency
    warn_persistency(B, method)

    # Compute polynomial moments of standard normal distribution
    gaussian_moment = zeros(n_moments)
    c = 1
    for k=1:floor(Int, n_moments/2)
        c = (2*k-1)*c
        gaussian_moment[2*k] = c
    end
    # Compute standardized VAR(1) representation
    # (zero mean and diagonal covariance matrix)
    A, C, mu, Sigma = standardize_var(b, B, Psi, M)

    # Construct 1-D grids
    y1D, y1Dhelper = construct_1D_grid(Sigma, Nm, M, n_sigmas, method)

    # Construct all possible combinations of elements of the 1-D grids
    D = allcomb3(y1D')'

    # Construct finite-state Markov chain approximation
    # conditional mean of the VAR process at each grid point
    cond_mean = A*D
    # probability transition matrix
    P = ones(Nm^M, Nm^M)
    # normalizing constant for maximum entropy computations
    scaling_factor = y1D[:, end]
    # used to store some intermediate calculations
    temp = Matrix{Float64}(undef, Nm, M)
    # store optimized values of lambda (2 moments) to improve initial guesses
    lambda_bar = zeros(2*M, Nm^M)
    # small positive constant for numerical stability
    kappa = 1e-8

    for ii = 1:(Nm^M)

        # Construct prior guesses for maximum entropy optimizations
        q = construct_prior_guess(cond_mean[:, ii], Nm, y1D, y1Dhelper, method)

        # Make sure all elements of the prior are stricly positive
        q[q.<kappa] .= kappa

        for jj = 1:M
            # Try to use intelligent initial guesses
            if ii == 1
                lambda_guess = zeros(2)
            else
                lambda_guess = lambda_bar[(jj-1)*2+1:jj*2, ii-1]
            end

            # Maximum entropy optimization
            if n_moments == 1 # match only 1 moment
                temp[:, jj], _, _ = discrete_approximation(y1D[jj, :],
                    X -> (X'.-cond_mean[jj, ii])/scaling_factor[jj],
                    [0.0], q[jj, :], [0.0])
            else # match 2 moments first
                p, lambda, moment_error = discrete_approximation(y1D[jj, :],
                    X -> polynomial_moment(X, cond_mean[jj, ii], scaling_factor[jj], 2),
                    [0; 1]./(scaling_factor[jj].^(1:2)), q[jj, :], lambda_guess)
                if !(norm(moment_error) < 1e-5) # if 2 moments fail, just match 1 moment
                    @warn("Failed to match first 2 moments. Just matching 1.")
                    temp[:, jj], _, _ = discrete_approximation(y1D[jj, :],
                    X -> (X'.-cond_mean[jj, ii])/scaling_factor[jj],
                        [0.0], q[jj, :], [0.0])
                    lambda_bar[(jj-1)*2+1:jj*2, ii] = zeros(2,1)
                elseif n_moments == 2
                    lambda_bar[(jj-1)*2+1:jj*2, ii] = lambda
                    temp[:, jj] = p
                else # solve maximum entropy problem sequentially from low order moments
                    lambda_bar[(jj-1)*2+1:jj*2, ii] = lambda
                    for mm = 4:2:n_moments
                        lambda_guess = vcat(lambda, 0.0, 0.0) # add 0 to previous lambda
                        pnew, lambda, moment_error = discrete_approximation(y1D[jj,:],
                            X -> polynomial_moment(X, cond_mean[jj,ii],
                                                    scaling_factor[jj], mm),
                            gaussian_moment[1:mm]./(scaling_factor[jj].^(1:mm)),
                            q[jj, :], lambda_guess)
                        if !(norm(moment_error) < 1e-5)
                            @warn(
                            "Failed to match first $mm moments.  Just matching $(mm-2).")
                            break
                        else
                            p = pnew
                        end
                    end
                    temp[:, jj] = p
                end
            end
        end
        P[ii, :] .= vec(prod(allcomb3(temp), dims = 2))
    end

    X = C*D .+ mu # map grids back to original space

    M != 1 || (return MarkovChain(P, vec(X)))
    return MarkovChain(P, [X[:, i] for i in 1:Nm^M])
end

"""
check persistency when `method` is `Quadrature` and give warning if needed

##### Arguments
- `B::Union{Real, AbstractMatrix}` : impact coefficient
- `method::VAREstimationMethod` : method for grid making

##### Returns
- `nothing`

"""
function warn_persistency(B::AbstractMatrix, ::Quadrature)
    if any(eigen(B).values .> 0.9)
        @warn("The quadrature method may perform poorly for highly persistent processes.")
    end
    return nothing
end
warn_persistency(B::Real, method::Quadrature) = warn_persistency(fill(B, 1, 1), method)
warn_persistency(::Union{Real, AbstractMatrix}, ::VAREstimationMethod) = nothing

"""

return standerdized AR(1) representation

##### Arguments

- `b::Real` : constant term
- `B::Real` : impact coefficient
- `Psi::Real` : variance of innovation
- `M::Integer == 1` : must be one since the function is for AR(1)

##### Returns

- `A::Real` : impact coefficient of standardized AR(1) process
- `C::Real` : standard deviation of the innovation
- `mu::Real` : mean of the standardized AR(1) process
- `Sigma::Real` : variance of the standardized AR(1) process

"""
function standardize_var(b::Real, B::Real, Psi::Real, M::Integer)
    C = sqrt(Psi)
    A = B
    mu = [b/(1-B)] # mean of the process
    Sigma = 1/(1-B^2) #
    return A, C, mu, Sigma
end

"""

return standerdized VAR(1) representation

##### Arguments

- `b::AbstractVector` : `M x 1` constant term vector
- `B::AbstractMatrix` : `M x M` matrix of impact coefficients
- `Psi::AbstractMatrix` : `M x M` variance-covariance matrix of innovations
- `M::Intger` : number of variables of the VAR(1) model

##### Returns

- `A::Matirx` : impact coefficients of standardized VAR(1) process
- `C::AbstractMatrix` : variance-covariance matrix of standardized model innovations
- `mu::AbstractVector` : mean of the standardized VAR(1) process
- `Sigma::AbstractMatrix` : variance-covariance matrix of the standardized VAR(1) process

"""
function standardize_var(b::AbstractVector, B::AbstractMatrix,
                         Psi::AbstractMatrix, M::Integer)
    C1 = cholesky(Psi).L
    mu = ((I - B)\ I)*b
    A1 = C1\(B*C1)
    # unconditional variance
    Sigma1 = reshape(((I-kron(A1,A1))\I)*vec(Matrix(I, M, M)),M,M)
    U, _ = min_var_trace(Sigma1)
    A = U'*A1*U
    Sigma = U'*Sigma1*U
    C = C1*U
    return A, C, mu, Sigma
end

"""
construct prior guess for evenly spaced grid method

##### Arguments

- `cond_mean::AbstractVector` : conditional Mean of each variable
- `Nm::Integer` : number of grid points
- `y1D::AbstractMatrix` : grid of variable
- `::AbstractMatrix` : bounds of each grid bin
- `method::Even` : method for grid making

"""
construct_prior_guess(cond_mean::AbstractVector, Nm::Integer,
                      y1D::AbstractMatrix, ::Nothing, method::Even) =
    pdf.(Normal.(repeat(cond_mean, 1, Nm), 1), y1D)

"""
construct prior guess for quantile grid method

##### Arguments

- `cond_mean::AbstractVector` : conditional Mean of each variable
- `Nm::Integer` : number of grid points
- `::AbstractMatrix` : grid of variable
- `y1Dbounds::AbstractMatrix` : bounds of each grid bin
- `method::Quantile` : method for grid making

"""
construct_prior_guess(cond_mean::AbstractVector, Nm::Integer,
                      ::AbstractMatrix, y1Dbounds::AbstractMatrix, method::Quantile) =
    cdf.(Normal.(repeat(cond_mean, 1, Nm), 1), y1Dbounds[:, 2:end]) -
           cdf.(Normal.(repeat(cond_mean, 1, Nm), 1), y1Dbounds[:, 1:end-1])

"""
construct prior guess for quadrature grid method

##### Arguments

- `cond_mean::AbstractVector` : conditional Mean of each variable
- `Nm::Integer` : number of grid points
- `y1D::AbstractMatrix` : grid of variable
- `weights::AbstractVector` : weights of grid `y1D`
- `method::Quadrature` : method for grid making

"""
construct_prior_guess(cond_mean::AbstractVector, Nm::Integer,
                      y1D::AbstractMatrix, weights::AbstractVector, method::Quadrature) =
    (pdf.(Normal.(repeat(cond_mean, 1, Nm), 1), y1D) ./ pdf.(Ref(Normal()), y1D)).*weights'
"""

construct one-dimensional evenly spaced grid of states

##### Argument

- `Sigma::ScalarOrArray` : variance-covariance matrix of the standardized process
- `Nm::Integer` : number of grid points
- `M::Integer` : number of variables (`M=1` corresponds to AR(1))
- `n_sigmas::Real` : number of standard error determining end points of grid
- `method::Even` : method for grid making

##### Return

- `y1D` : `M x Nm` matrix of variable grid
- `nothing` : `nothing` of type `Void`

"""
function construct_1D_grid(Sigma::Union{Real, AbstractMatrix}, Nm::Integer,
                           M::Integer, n_sigmas::Real, method::Even)
    min_sigmas = sqrt(minimum(eigen(Sigma).values))
    y1Drow = collect(range(-min_sigmas*n_sigmas, stop=min_sigmas*n_sigmas, length=Nm))'
    y1D = repeat(y1Drow, M, 1)
    return y1D, nothing
end

"""

construct one-dimensional quantile grid of states

##### Argument

- `Sigma::AbstractMatrix` : variance-covariance matrix of the standardized process
- `Nm::Integer` : number of grid points
- `M::Integer` : number of variables (`M=1` corresponds to AR(1))
- `n_sigmas::Real` : number of standard error determining end points of grid
- `method::Quntile` : method for grid making

##### Return

- `y1D` : `M x Nm` matrix of variable grid
- `y1Dbounds` : bounds of each grid bin

"""
function construct_1D_grid(Sigma::AbstractMatrix, Nm::Integer,
                           M::Integer, n_sigmas::Real, method::Quantile)
    sigmas = sqrt.(diag(Sigma))
    y1D = quantile.(Normal.(0, sigmas), (2*(1:Nm)'.-1)/(2*Nm))
    y1Dbounds = hcat(fill(-Inf, M, 1),
                     quantile.(Normal.(0, sigmas), ((1:(Nm-1))')/Nm),
                     fill(Inf, M, 1))
    return y1D, y1Dbounds
end
construct_1D_grid(Sigma::Real, Nm::Integer,
                  M::Integer, n_sigmas::Real, method::Quantile) =
    construct_1D_grid(fill(Sigma, 1, 1), Nm, M, n_sigmas, method)

"""

construct one-dimensional quadrature grid of states

##### Argument

- `::ScalarOrArray` : not used
- `Nm::Integer` : number of grid points
- `M::Integer` : number of variables (`M=1` corresponds to AR(1))
- `n_sigmas::Real` : not used
- `method::Quadrature` : method for grid making

##### Return

- `y1D` : `M x Nm` matrix of variable grid
- `weights` : weights on each grid

"""
function construct_1D_grid(::ScalarOrArray, Nm::Integer,
                           M::Integer, ::Real, method::Quadrature)
    nodes, weights = qnwnorm(Nm, 0, 1)
    y1D = repeat(nodes', M, 1)
    return y1D, weights
end

"""

Return combinations of each column of matrix `A`.
It is simiplifying `allcomb2` by using `gridmake` from QuantEcon

##### Arguments

- `A::AbstractMatrix` : `N x M` Matrix

##### Returns

- `N^M x M` Matrix, combination of each row of `A`.

###### Example
```julia-repl
julia> allcomb3([1 4 7;
                 2 5 8;
                 3 6 9]) # numerical input
27×3 Array{Int64,2}:
 1  4  7
 1  4  8
 1  4  9
 1  5  7
 1  5  8
 1  5  9
 1  6  7
 1  6  8
 1  6  9
 2  4  7
 ⋮
 2  6  9
 3  4  7
 3  4  8
 3  4  9
 3  5  7
 3  5  8
 3  5  9
 3  6  7
 3  6  8
 3  6  9
```

"""
allcomb3(A::AbstractMatrix) =
    reverse(gridmake(reverse([A[:, i] for i in 1:size(A, 2)], dims = 1)...), dims = 2)

"""
Compute a discrete state approximation to a distribution with known moments,
using the maximum entropy procedure proposed in Tanaka and Toda (2013)

```julia
p, lambda_bar, moment_error = discrete_approximation(D, T, TBar, q, lambda0)
```

##### Arguments

- `D::AbstractVector` : vector of grid points of length `N`.
                        N is the number of points at which an approximation
                        is to be constructed.
- `T::Function` : A function that accepts a single `AbstractVector` of length `N`
                  and returns an `L x N` matrix of moments evaluated
                  at each grid point, where L is the number of moments to be
                  matched.
- `TBar::AbstractVector` : length `L` vector of moments of the underlying distribution
                           which should be matched

##### Optional

- `q::AbstractVector` : length `N` vector of prior weights for each point in D.
                        The default is for each point to have an equal weight.
- `lambda0::AbstractVector` : length `L` vector of initial guesses for the dual problem
                              variables. The default is a vector of zeros.

##### Returns

- `p` : (1 x N) vector of probabilties assigned to each grid point in `D`.
- `lambda_bar` : length `L` vector of dual problem variables which solve the
                maximum entropy problem
- `moment_error` : vector of errors in moments (defined by moments of
                  discretization minus actual moments) of length `L`

"""
function discrete_approximation(D::AbstractVector, T::Function, TBar::AbstractVector,
                                  q::AbstractVector=ones(N)/N, # Default prior weights
                                  lambda0::AbstractVector=zeros(Tbar))

    # Input error checking
    N = length(D)

    Tx = T(D)
    L, N2 = size(Tx)

    if N2 != N || length(TBar) != L || length(lambda0) != L || length(q) != N2
        error("Dimension mismatch")
    end

    # Compute maximum entropy discrete distribution
    options = Optim.Options(f_tol=1e-16, x_tol=1e-16)
    obj(lambda) = entropy_obj(lambda, Tx, TBar, q)
    grad!(grad, lambda) = entropy_grad!(grad, lambda, Tx, TBar, q)
    hess!(hess, lambda) = entropy_hess!(hess, lambda, Tx, TBar, q)
    res = Optim.optimize(obj, grad!, hess!, lambda0, Optim.Newton(), options)
    # Sometimes the algorithm fails to converge if the initial guess is too far
    # away from the truth. If this occurs, the program tries an initial guess
    # of all zeros.
    if !Optim.converged(res) && all(lambda0 .!= 0.0)
        @warn("Failed to find a solution from provided initial guess. Trying new initial guess.")
        res = Optim.optimize(obj, grad!, hess!, zeros(lambda0), Optim.Newton(), options)
        # check convergence
        Optim.converged(res) || error("Failed to find a solution.")
    end
    # Compute final probability weights and moment errors
    lambda_bar = Optim.minimizer(res)
    minimum_value = Optim.minimum(res)

    Tdiff = Tx .- TBar
    p = (q'.*exp.(lambda_bar'*Tdiff))/minimum_value
    grad = similar(lambda0)
    grad!(grad, Optim.minimizer(res))
    moment_error = grad/minimum_value
    return p, lambda_bar, moment_error
end
# discrete_approximation(D::AbstractVector, T::Function, TBar::Real,
#                        q::AbstractVector=ones(N)/N, lambda0::Real=0) =
#     discrete_approximation(D, T, [TBar], q, [lambda0])

"""

Compute the moment defining function used in discrete_approximation

```julia
T = polynomial_moment(X, mu, scaling_factor, mMoments)
```

##### Arguments:

- `X::AbstractVector` : length `N` vector of grid points
- `mu::Real` : location parameter (conditional mean)
- `scaling_factor::Real` : scaling factor for numerical stability.
                          (typically largest grid point)
- `n_moments::Integer` : number of polynomial moments

##### Return

- `T` : moment defining function used in `discrete_approximation`

"""
function polynomial_moment(X::AbstractVector, mu::Real,
                           scaling_factor::Real, n_moments::Integer)
    # Check that scaling factor is positive
    scaling_factor>0 || error("scaling_factor must be a positive number")
    Y = (X.-mu)/scaling_factor # standardized grid
    T = Y'.^collect(1:n_moments)
end

"""

Compute the maximum entropy objective function used in `discrete_approximation`

```julia
obj = entropy_obj(lambda, Tx, TBar, q)
```

##### Arguments

- `lambda::AbstractVector` : length `L` vector of values of the dual problem variables
- `Tx::AbstractMatrix` : `L x N` matrix of moments evaluated at the grid points
                         specified in discrete_approximation
- `TBar::AbstractVector` : length `L` vector of moments of the underlying distribution
                           which should be matched
- `q::AbstractVector` : length `N` vector of prior weights for each point in the grid.

##### Returns

- `obj` : scalar value of objective function evaluated at `lambda`

"""
function entropy_obj(lambda::AbstractVector, Tx::AbstractMatrix,
                     TBar::AbstractVector, q::AbstractVector)
    # Compute objective function
    Tdiff = Tx .- TBar
    temp = q' .* exp.(lambda'*Tdiff)
    obj = sum(temp)
    return obj
end

"""
Compute gradient of objective function

##### Returns

- `grad` : length `L` gradient vector of the objective function evaluated at `lambda`

"""
function entropy_grad!(grad::AbstractVector, lambda::AbstractVector,
                       Tx::AbstractMatrix, TBar::AbstractVector, q::AbstractVector)
    Tdiff = Tx .- TBar
    temp = q'.*exp.(lambda'*Tdiff)
    temp2 = temp.*Tdiff
    grad .= vec(sum(temp2, dims = 2))
end

"""
Compute hessian of objective function

##### Returns

- `hess` : `L x L` hessian matrix of the objective function evaluated at `lambda`

"""
function entropy_hess!(hess::AbstractMatrix, lambda::AbstractVector,
                       Tx::AbstractMatrix, TBar::AbstractVector, q::AbstractVector)
    Tdiff = Tx .- TBar
    temp = q'.*exp.(lambda'*Tdiff)
    temp2 = temp.*Tdiff
    hess .= temp2*Tdiff'
end

@doc doc"""

find a unitary matrix `U` such that the diagonal components of `U'AU` is as
close to a multiple of identity matrix as possible

##### Arguments

- `A::AbstractMatrix` : square matrix

##### Returns

- `U` : unitary matrix
- `fval` : minimum value

"""
function min_var_trace(A::AbstractMatrix)

    ==(size(A)...) || throw(ArgumentError("input matrix must be square"))

    K = size(A, 1) # size of A
    d = tr(A)/K # diagonal of U'*A*U should be closest to d
    function obj(X, grad)
        X = reshape(X, K, K)
        return (norm(diag(X'*A*X) .- d))
    end
    function unitary_constraint(res, X, grad)
        X = reshape(X, K, K)
        res .= vec(X'*X - Matrix(I, K, K))
    end

    opt = NLopt.Opt(:LN_COBYLA, K^2)
    NLopt.min_objective!(opt, obj)
    NLopt.equality_constraint!(opt, unitary_constraint, zeros(K^2))
    fval, U_vec, ret = NLopt.optimize(opt, vec(Matrix(I, K, K)))

    return reshape(U_vec, K, K), fval
end


text:
using DelayDiffEq, DiffEqDevTools, DiffEqProblemLibrary, Plots
using DiffEqProblemLibrary.DDEProblemLibrary: importddeproblems; importddeproblems()
import DiffEqProblemLibrary.DDEProblemLibrary: prob_dde_wheldon
gr()

sol = solve(prob_dde_wheldon, MethodOfSteps(Vern9(); fpsolve = NLFunctional(; max_iter = 1000)); reltol=1e-14, abstol=1e-14)
test_sol = TestSolution(sol)
plot(sol)

abstols = 1.0 ./ 10.0 .^ (4:7)
reltols = 1.0 ./ 10.0 .^ (1:4)

setups = [Dict(:alg=>MethodOfSteps(BS3())),
          Dict(:alg=>MethodOfSteps(Tsit5())),
          Dict(:alg=>MethodOfSteps(RK4())),
          Dict(:alg=>MethodOfSteps(DP5())),
          Dict(:alg=>MethodOfSteps(OwrenZen3())),
          Dict(:alg=>MethodOfSteps(OwrenZen4())),
          Dict(:alg=>MethodOfSteps(OwrenZen5()))]
names = ["BS3", "Tsit5", "RK4", "DP5", "OwrenZen3", "OwrenZen4", "OwrenZen5"]
wp = WorkPrecisionSet(prob_dde_wheldon,abstols,reltols,setups;names=names,
                      appxsol=test_sol,maxiters=Int(1e5),error_estimate=:final)
plot(wp)

abstols = 1.0 ./ 10.0 .^ (4:7)
reltols = 1.0 ./ 10.0 .^ (1:4)

setups = [Dict(:alg=>MethodOfSteps(BS3())),
          Dict(:alg=>MethodOfSteps(Tsit5())),
          Dict(:alg=>MethodOfSteps(RK4())),
          Dict(:alg=>MethodOfSteps(DP5())),
          Dict(:alg=>MethodOfSteps(OwrenZen3())),
          Dict(:alg=>MethodOfSteps(OwrenZen4())),
          Dict(:alg=>MethodOfSteps(OwrenZen5()))]
names = ["BS3", "Tsit5", "RK4", "DP5", "OwrenZen3", "OwrenZen4", "OwrenZen5"]
wp = WorkPrecisionSet(prob_dde_wheldon,abstols,reltols,setups;names=names,
                      appxsol=test_sol,maxiters=Int(1e5),error_estimate=:L2)
plot(wp)

abstols = 1.0 ./ 10.0 .^ (8:11)
reltols = 1.0 ./ 10.0 .^ (5:8)

setups = [Dict(:alg=>MethodOfSteps(BS3())),
          Dict(:alg=>MethodOfSteps(Tsit5())),
          Dict(:alg=>MethodOfSteps(RK4())),
          Dict(:alg=>MethodOfSteps(DP5())),
          Dict(:alg=>MethodOfSteps(OwrenZen3())),
          Dict(:alg=>MethodOfSteps(OwrenZen4())),
          Dict(:alg=>MethodOfSteps(OwrenZen5()))]
names = ["BS3", "Tsit5", "RK4", "DP5", "OwrenZen3", "OwrenZen4", "OwrenZen5"]
wp = WorkPrecisionSet(prob_dde_wheldon,abstols,reltols,setups;names=names,
                      appxsol=test_sol,maxiters=Int(1e5),error_estimate=:final)
plot(wp)

abstols = 1.0 ./ 10.0 .^ (8:11)
reltols = 1.0 ./ 10.0 .^ (5:8)

setups = [Dict(:alg=>MethodOfSteps(BS3())),
          Dict(:alg=>MethodOfSteps(Tsit5())),
          Dict(:alg=>MethodOfSteps(RK4())),
          Dict(:alg=>MethodOfSteps(DP5())),
          Dict(:alg=>MethodOfSteps(OwrenZen3())),
          Dict(:alg=>MethodOfSteps(OwrenZen4())),
          Dict(:alg=>MethodOfSteps(OwrenZen5()))]
names = ["BS3", "Tsit5", "RK4", "DP5", "OwrenZen3", "OwrenZen4", "OwrenZen5"]
wp = WorkPrecisionSet(prob_dde_wheldon,abstols,reltols,setups;names=names,
                      appxsol=test_sol,maxiters=Int(1e5),error_estimate=:L2)
plot(wp)

abstols = 1.0 ./ 10.0 .^ (4:7)
reltols = 1.0 ./ 10.0 .^ (1:4)

setups = [Dict(:alg=>MethodOfSteps(Vern6())),
          Dict(:alg=>MethodOfSteps(Vern7())),
          Dict(:alg=>MethodOfSteps(Vern8())),
          Dict(:alg=>MethodOfSteps(Vern9())),
          Dict(:alg=>MethodOfSteps(OwrenZen4()))]
names = ["Vern6", "Vern7", "Vern8", "Vern9", "OwrenZen4"]
wp = WorkPrecisionSet(prob_dde_wheldon,abstols,reltols,setups;names=names,
                      appxsol=test_sol,maxiters=Int(1e5),error_estimate=:final)
plot(wp)

abstols = 1.0 ./ 10.0 .^ (4:7)
reltols = 1.0 ./ 10.0 .^ (1:4)

setups = [Dict(:alg=>MethodOfSteps(Vern6())),
          Dict(:alg=>MethodOfSteps(Vern7())),
          Dict(:alg=>MethodOfSteps(Vern8())),
          Dict(:alg=>MethodOfSteps(Vern9())),
          Dict(:alg=>MethodOfSteps(OwrenZen4()))]
names = ["Vern6", "Vern7", "Vern8", "Vern9", "OwrenZen4"]
wp = WorkPrecisionSet(prob_dde_wheldon,abstols,reltols,setups;names=names,
                      appxsol=test_sol,maxiters=Int(1e5),error_estimate=:L2)
plot(wp)

abstols = 1.0 ./ 10.0 .^ (8:11)
reltols = 1.0 ./ 10.0 .^ (5:8)

setups = [Dict(:alg=>MethodOfSteps(Vern6())),
          Dict(:alg=>MethodOfSteps(Vern7())),
          Dict(:alg=>MethodOfSteps(Vern8())),
          Dict(:alg=>MethodOfSteps(Vern9())),
          Dict(:alg=>MethodOfSteps(OwrenZen4()))]
names = ["Vern6", "Vern7", "Vern8", "Vern9", "OwrenZen4"]
wp = WorkPrecisionSet(prob_dde_wheldon,abstols,reltols,setups;names=names,
                      appxsol=test_sol,maxiters=Int(1e5),error_estimate=:final)
plot(wp)

abstols = 1.0 ./ 10.0 .^ (8:11)
reltols = 1.0 ./ 10.0 .^ (5:8)

setups = [Dict(:alg=>MethodOfSteps(Vern6())),
          Dict(:alg=>MethodOfSteps(Vern7())),
          Dict(:alg=>MethodOfSteps(Vern8())),
          Dict(:alg=>MethodOfSteps(Vern9())),
          Dict(:alg=>MethodOfSteps(OwrenZen4()))]
names = ["Vern6", "Vern7", "Vern8", "Vern9", "OwrenZen4"]
wp = WorkPrecisionSet(prob_dde_wheldon,abstols,reltols,setups;names=names,
                      appxsol=test_sol,maxiters=Int(1e5),error_estimate=:L2)
plot(wp)

using DiffEqBenchmarks
DiffEqBenchmarks.bench_footer(WEAVE_ARGS[:folder],WEAVE_ARGS[:file])