text:
stress_scale = 1.0e5
strain_scale = 1

force_scales = [2.0]
force_scale = 2.0
porder = 2
T = 0.2
NT = 200
Δt = T/NT

nntype = "stiffmat"
include("nnutil.jl")

#use only rho
testtype = "PlaneStressIncompressibleRivlinSaunders"
#https://classes.engineering.wustl.edu/2009/spring/mase5513/abaqus/docs/v6.6/books/bmk/default.htm?startat=ch01s01ach08.html
prop = Dict("name"=> testtype, "rho"=> 0.8000, "C1"=>0.1863e6, "C2"=>0.00979e6)

# ! define H0
# Trained with nx, ny = 10, 5
#H0 = [1.04167e6  2.08333e5  0.0
#      2.08333e5  1.04167e6  0.0
#      0.0        0.0        4.16667e5]/stress_scale

H0 = [1.0406424793819175e6 209077.08366547766         0.0
      209077.08366547766   1.0411467691352057e6       0.0
      0.0                  0.0                   419057.32049008965]/stress_scale

n_data = [100, 101, 102, 103, 104, 200, 201, 202, 203, 204]

# DNS computaional domain
fiber_size = 2
# nx_f, ny_f = 40*fiber_size, 20*fiber_size
nx_f, ny_f = 10*fiber_size, 5*fiber_size
Lx, Ly = 0.1, 0.05

nx, ny = 10, 5
# number of subelements in one element in each directions
sx_f, sy_f = div(nx_f,nx), div(ny_f,ny)

ndofs = 2
fine_to_coarse = zeros(Int64, ndofs*(nx*porder+1)*(ny*porder+1))
for idof = 1:ndofs
    for iy = 1:ny*porder+1
        for ix = 1:nx*porder+1
            fine_to_coarse[ix + (iy - 1)*(nx*porder+1) + (idof-1)*(nx*porder+1)*(ny*porder+1)] =
            1 + (ix - 1) * sx_f + (iy - 1) * (nx_f*porder + 1) * sy_f + (nx_f*porder + 1)*(ny_f*porder + 1)*(idof - 1)
        end
    end
end

function approximate_stress(tid, force_scale, method)
    nodes, EBC, g, gt, FBC, fext, ft, npoints, node_to_point = BoundaryCondition(tid, nx, ny, porder, Lx, Ly; force_scale=force_scale )
    domain = Domain(nodes, elements, ndofs, EBC, g, FBC, fext)
    setGeometryPoints!(domain, npoints, node_to_point)

    state = zeros(domain.neqs)
    ∂u = zeros(domain.neqs)
    globdat = GlobalData(state,zeros(domain.neqs), zeros(domain.neqs),∂u, domain.neqs, gt, ft)
    assembleMassMatrix!(globdat, domain)

    # full_state_history, full_fext_history = read_data("$(@__DIR__)/Data/order$porder/$(tid)_$(force_scale)_$(fiber_size).dat")
    full_state_history, full_fext_history = read_data("$(@__DIR__)/Data/order$porder/$(tid)_$(force_scale)_$(fiber_size).dat")

    #update state history and fext_history on the homogenized domain
    state_history = [x[fine_to_coarse] for x in full_state_history]
    fext_history = []
    setNeumannBoundary!(domain, FBC, fext)
    for i = 1:NT
        globdat.time = Δt*i
        updateDomainStateBoundary!(domain, globdat)
        push!(fext_history, domain.fext[:])
    end

    E_all, S_all = LSfittingStress(domain, globdat, hcat(state_history...), hcat(fext_history...), Δt, method)

end

nodes, _, _, _, _, _, _,_,_ = BoundaryCondition(n_data[1], nx, ny, porder, Lx, Ly; force_scale=force_scale)
elements = []
for j = 1:ny
    for i = 1:nx
        n = (nx*porder+1)*(j-1)*porder + (i-1)porder+1
        #element (i,j)
        if porder == 1
            #   4 ---- 3
            #
            #   1 ---- 2

            elnodes = [n, n + 1, n + 1 + (nx + 1), n + (nx + 1)]
        elseif porder == 2
            #   4 --7-- 3
            #   8   9   6
            #   1 --5-- 2
            elnodes = [n, n + 2, n + 2 + 2*(2*nx+1),  n + 2*(2*nx+1), n+1, n + 2 + (2*nx+1), n + 1 + 2*(2*nx+1), n + (2*nx+1), n+1+(2*nx+1)]
        else
            error("polynomial order error, porder= ", porder)
        end
        coords = nodes[elnodes,:]
        push!(elements,SmallStrainContinuum(coords,elnodes, prop, 3))
    end
end

#################### Gradient test
# tid = n_data[1]
# method = "Constant"
# nodes, EBC, g, gt, FBC, fext, ft, npoints, node_to_point = BoundaryCondition(tid, nx, ny, porder; force_scale=force_scale )
# domain = Domain(nodes, elements, ndofs, EBC, g, FBC, fext)
# setGeometryPoints!(domain, npoints, node_to_point)
# function f(S_comp)

#     stress = reshape(S_comp, 3, div(size(S_comp,1),3))'|>Array

#     LSfittingStressHelper(domain, stress, method)
# end
# nS_comp = (method == "Constant") ? domain.neles : domain.npoints;
# S_comp = rand(Float64, 3*nS_comp)
# gradtest(f, S_comp)
#########################

#################### S_all compare with real S_all
# method = "Linear"
# tid = n_data[1]
# E_all, S_all = approximate_stress(tid, force_scale, method)
# @load "Data/order$porder/domain$(tid)_$(force_scale)_1.jld2" domain
# S_all_real = domain.history["stress"]
# E_all_real = domain.history["strain"]
# for it = 1:NT
#     @show "E_all[", it, "] error is ", norm(E_all[it+1,:,:] - E_all_real[it]), " / ", norm(E_all_real[it])
#     @show "S_all[", it, "] error is ", norm(S_all[it+1,:,:] - S_all_real[it]), " / ", norm(S_all_real[it])
# end
#########################

loss = constant(0.0)

method = "Linear"
for tid in n_data
    E_all, S_all = approximate_stress(tid, force_scale, method)
    ngp = size(E_all,2)

    X = zeros(Float64, (NT-1)*ngp, 9)
    Y = zeros(Float64, (NT-1)*ngp, 3)
    k = 1
    for it = 2:NT
        for igp = 1:ngp
            X[k,:] = [E_all[it,igp,:]; E_all[it-1,igp,:]; S_all[it-1,igp,:]]#ε, ε0, σ0
            Y[k,:] = S_all[it,igp,:]
            k = k + 1
        end
    end

    x = constant(X)
    y = nn(x[:,1:3], x[:,4:6], x[:,7:9])

    global loss
    loss += mean((y - Y)^2) #/stress_scale^2
end

sess = Session(); init(sess)
@show run(sess, loss)
# ADCME.load(sess, "Data/$(nntype)/NNPreLSfit.mat")
for i = 1:5
    BFGS!(sess, loss, 1000)
    ADCME.save(sess, "Data/$(nntype)/NNPreLSfit_$(idx)_$(H_function)_$(i).mat")
end


text:
import Base: show, display, getindex, read

###########################
#    Types
###########################

abstract type AtomFileType end

immutable FrameHeader
    frame_start::Int64
    frame_size::Int64

    num_atoms::Int64
    atom_data_start::Int64
    atom_data_size::Int64

    num_bonds::Int64
    bond_data_start::Int64
    bond_data_size::Int64
end

FrameHeader(frame_start::Int64, frame_size::Int64, num_atoms::Int64, atom_data_start::Int64, atom_data_size::Int64) = FrameHeader(frame_start, frame_size, num_atoms, atom_data_start, atom_data_size, 0, 0, 0)

immutable AtomFile{F <: AtomFileType}
    filename::AbstractString
    frame_headers::Vector{FrameHeader}
    frames::Vector{Frame}

    AtomFile{F}(filename::AbstractString) where F <: AtomFileType = new(filename, Vector{FrameHeader}(), Vector{Frame}())
end

function AtomFile(filename::AbstractString)
    base, ext = splitext(filename)
    for file_type in subtypes(AtomFileType)
        if ext ∈ file_exts(file_type)
            return AtomFile{file_type}(filename)
        end
    end

    throw(ArgumentError("Unknown file type $ext"))
end

AtomFile{F <: AtomFileType}(::Type{F}, filename::AbstractString) = AtomFile{F}(filename)

type Chunk
    data::Vector{UInt8}
    i::Int64

    Chunk(num_bytes::Int64) = new(Vector{UInt8}(num_bytes), 1)
    Chunk(data::Vector{UInt8}) = new(data, 1)
end

################################
#    File reading functions
################################

function read_byte(chunk::Chunk)
    b = chunk.data[chunk.i]
    chunk.i += 1
    return b
end

preview_byte(chunk::Chunk) = chunk.data[chunk.i]

next_byte(chunk::Chunk) = chunk.i += 1

read_newline(chunk::Chunk) = while read_byte(chunk) != 0x0a end

function read_newlines(chunk::Chunk, num_newlines::Int)
    s = 0
    while s < num_newlines
        if read_byte(chunk) == 0x0a
            s += 1
        end
    end
end

function read_float{T}(::Type{T}, chunk::Chunk)
    a = T(0)
    b = T(0)
    p = 1

    c = preview_byte(chunk)
    if c == 0x2d
        p = -1
        next_byte(chunk)
    end

    c == 0x2b && next_byte(chunk)

    while true
        c = preview_byte(chunk)
        (c < 0x30 || c > 0x39) && break
        next_byte(chunk)
        a = 10a + (c - 0x30)
    end

    c != 0x2e && return p*a
    next_byte(chunk)

    i = 1
    while true
        c = preview_byte(chunk)
        (c < 0x30 || c > 0x39) && break
        next_byte(chunk)
        a = 10a + (c - 0x30)
        i *= 10
    end

    c != 0x45 && c != 0x65 && return (p*a/i)
    next_byte(chunk)

    b = read_int(chunk)
    return (p*a/i) * T(10)^b
end

function read_int(chunk::Chunk)
    n = 0
    p = 1

    c = preview_byte(chunk)
    if c == 0x2d
        p = -1
        next_byte(chunk)
    end

    c == 0x2b && next_byte(chunk)

    while true
        c = preview_byte(chunk)
        (c < 0x30 || c > 0x39) && break
        next_byte(chunk)
        n = 10n + (c - 0x30)
    end

    return p*n
end

function read_symbol(chunk::Chunk, atom_types::Matrix{UInt8}, i::Int)
    n = size(atom_types, 1)

    j = 1
    while j <= n
        c = preview_byte(chunk)
        (c == 0x20 || c == 0x09) && break
        atom_types[j, i] = c
        next_byte(chunk)
        j += 1
    end

    while j <= n
        atom_types[j, i] = 0x20
        j += 1
    end
end

function read_symbol(chunk::Chunk)
    i1 = chunk.i
    while true
        c = preview_byte(chunk)
        if (c == 0x20 || c == 0x09)
            return ccall(:jl_symbol_n, Ref{Symbol}, (Ptr{UInt8}, Int32), pointer(chunk.data, i1), Int32(chunk.i - i1))
        end
        next_byte(chunk)
    end
end

function read_nonwhitespace(chunk::Chunk)
    while true
        c = preview_byte(chunk)
        (c == 0x20 || c == 0x09 || c == 0x23 || c == 0x0a) && return
        next_byte(chunk)
    end
end

function read_whitespace(chunk::Chunk)
    while true
        c = preview_byte(chunk)
        c != 0x20 && c != 0x09 && c != 0x23 && return
        next_byte(chunk)
    end
end

function read_line(chunk::Chunk)
    i = chunk.i
    read_newline(chunk)
    #ascii(pointer(chunk.data, i), chunk.i - i - 1)
    String(chunk.data[i:chunk.i - 1])
end

function read_through_regex(chunk::Chunk, r::Regex)
    while true
        s = read_line(chunk)
        ismatch(r, s) && return match(r, s)
    end
end

function write_atoms{F <: AtomFileType}(file::AtomFile{F}, frames::Vector{Frame})
    filename = file.filename
    io = open(filename, "w")

    for frame in frames
        write_atoms(F, io, frame)
    end

    close(io)
end

write_atoms{F <: AtomFileType}(file::AtomFile{F}, frame::Frame) = write_atoms(file, [frame])
write_atoms(filename::AbstractString, frames::Vector{Frame}) = write_atoms(AtomFile(filename), frames)
write_atoms(filename::AbstractString, frame::Frame) = write_atoms(AtomFile(filename), frame)

function read_atoms(filename::AbstractString, end_frames::Int = 0)
    file = AtomFile(filename)
    num_frames = length(read_frame_headers(file))
    end_frames == 0 && (end_frames = num_frames)
    frames = read_atoms(file, frames = num_frames-end_frames+1:num_frames)
    end_frames == 1 ? frames[1] : frames
end

function read_atoms{F <: AtomFileType}(file::AtomFile{F}; frames::AbstractVector{Int} = eachindex(file.frame_headers))
    max_buffer_size = 1*1024*1024*1024

    length(file.frame_headers) == 0 && read_frame_headers(file)

    if length(frames) == 0
        frames = eachindex(file.frame_headers)
    end

    filename = file.filename
    frame_headers = file.frame_headers[frames]
    frame_data = file.frames[frames]

    num_frames = length(frame_headers)
    max_frame_size = maximum(f -> f.frame_size, frame_headers)

    max_buffer_size >= max_frame_size || error("Enlarge buffer size.")

    num_chunks = floor(Int, max_buffer_size / max_frame_size)
    num_chunks = min(num_chunks, num_frames)

    chunks = [Chunk(max_frame_size) for _ in 1:num_chunks]

    atom_indices = cumsum([f.num_atoms for f in frame_headers])
    total_atoms = atom_indices[end]
    atom_indices .-= atom_indices[1] - 1

    atoms = Atoms(F, total_atoms)

    ios = open(filename, "r")

    for (i, f) in enumerate(frame_headers)
        atom_start = atom_indices[i]
        atom_end = atom_start + f.num_atoms - 1
        atom_range = atom_start:atom_end

        frame_data[i][:atoms] = atoms[atom_range]
        j = mod1(i, num_chunks)
        # Seek indexing is 0-based
        seek(ios, f.atom_data_start - 1)
        chunks[j].i = 1
        readbytes!(ios, chunks[j].data, f.atom_data_size, all = false)
        read_atoms(Float64, F, chunks[j], atoms, atom_range)

        if frame_headers[i].num_bonds > 0
            bond_range = 1:f.num_bonds
            bonds = Matrix{Int}(2, f.num_bonds)
            frame_data[i][:bonds] = bonds

            seek(ios, f.bond_data_start - 1)
            chunks[j].i = 1
            readbytes!(ios, chunks[j].data, f.bond_data_size, all = false)
            read_bonds(F, chunks[j], bonds, bond_range)
        end
    end

    close(ios)

    return frame_data
end

function read_frame_headers{F <: AtomFileType}(file::AtomFile{F})
    file_size = filesize(file.filename)
    chunk = Chunk(file_size)

    ios = open(file.filename, "r")
    readbytes!(ios, chunk.data, file_size, all = true)
    close(ios)

    empty!(file.frame_headers)
    empty!(file.frames)
    while chunk.i < file_size
        frame_header, frame = read_frame_header(Float64, F, chunk)
        push!(file.frame_headers, frame_header)
        push!(file.frames, frame)
    end

    return file.frames
end

ext: jl

max_stars_repo_name: NickMcNutt/Alpine

max_issues_repo_name: NickMcNutt/Alpine

max_forks_repo_name: NickMcNutt/Alpine


text:
#=using DataStructures

abstract type AbstractFactor end

struct Nominal <: AbstractFactor
    levels::Set{String}
end

struct Ordinal <: AbstractFactor
    levels::OrderedSet{String}
end

struct Discrete <: AbstractFactor
    levels::OrderedSet{Int}
end

struct Continuous <: AbstractFactor
    levels::Vector{Float64}
end

struct Run
    setup::Int
    start::DateTime
    comment::String
    pois::Vector{POI}
end

struct Experiment{N}
    description::String
    factors::NTuple{N, String}
    levels::NTuple{N, AbstractFactor}
    setups::Vector{NTuple{N, Int}}
    runs::Vector{Run}
end=#

struct Run
    setup::Vector{Union{Missing, String}}
    start::DateTime
    comment::String
    pois::Dict{Symbol, POI}
end

struct Experiment
    description::String
    runs::Vector{Run}
end


text:
using Test, IonSim
using IonSim.PhysicalConstants: m_ca40
using Suppressor

@suppress_err begin

@testset "ion_configurations -- LinearChain" begin
    C = Ca40()
    lc = LinearChain(
            ions=[C, C, C, C], com_frequencies=(x=5, y=5, z=1),
            vibrational_modes=(y=[1], z=[4])
        )
    @test ions(lc) == lc.ions
    # test get_vibrational_modes, which should return an array of the selected
    # VibrationalModes in the linear chain
    vms = lc.vibrational_modes
    @test get_vibrational_modes(lc) == [vms.x..., vms.y..., vms.z...]

    # make sure ion numbers are updated
    @test [I.number for I in lc.ions] == [1, 2, 3, 4]

    # cmode -> pre-evaluated equilibrium positions for four ion chain
    cmode = [-0.06392393573, -0.0202155287427, 0.0202155287427, 0.0639239357]
    if lc.ions[1].position > 0
        cmode .*= -1
    end
    @test [I.position for I in lc.ions] ≈ cmode rtol=1e-6

    # should get warning if same ion is input multiple times to ion kwarg
    warning = "Some ions point to the same thing. Making copies."
    @test_logs (:warn, warning) LinearChain(
        ions=[C, C], com_frequencies=(x=4,y=4,z=1), vibrational_modes=(x=[],y=[],z=[1,2])
    )
    # and copies should be made of the repeated ions, so that they are no longer the same
    chain1 = LinearChain(
        ions=[C, C], com_frequencies=(x=4,y=4,z=1), vibrational_modes=(x=[],y=[],z=[1,2])
    )
    @test !(chain1.ions[1] ≡ chain1.ions[2])

    # Make sure there are no errors with print/show
    @suppress print(lc); show(lc)
end

@testset "ion_configurations -- general" begin
    # test calculation of equilibrium positions for a linear chain of equal mass ions against
    # known values
    posL = [-2.8708, -2.10003, -1.4504, -0.85378, -0.2821]
    pos = [posL; -1 .* reverse(posL)]
    @test any(isapprox.(linear_equilibrium_positions(10), pos, rtol=1e-4))

    # and test calculation of characterstic length scale for linear chain, equal mass
    @test characteristic_length_scale(m_ca40, 1e6) ≈ 4.449042804354e-6

    # and do the same for Anm, which computes the normal modes
    @test_throws AssertionError Anm(2, (x=0.5, y=0.5, z=1), (x=1, y=0, z=0))
    cst = [-0.2132, 0.6742, -0.6742, 0.2132]
    freq, mode = Anm(4, (x=2, y=2, z=1), (x=0, y=0, z=1))[end]
    @test freq ≈ √9.308 rtol=1e-4
    if mode[1] > 0
        cst = -1 .* cst
    end
    @test any(isapprox.(mode, cst, rtol=1e-4))

    # test '_sparsify', which should drop small values from an array
    x = [1e-6, 1e-5]
    IonSim._sparsify!(x, 2e-6)
    @test any(x .== [0, 1e-5])
end
end  # end suppress


text:
# basically the same as map but maps over each array in an array of arrays
function maparray(a::Any, f::Function)

	[f.(elem) for elem in a]

end

# find the product of a sparse and dense matrix, faster than native julia multiplication when the density is v low (less than 1%)
function sparsedensemult(sparsemat::SparseMatrixCSC, densemat)

	outmat = sparse([],[],[],size(sparsemat,1),size(sparsemat,2))
	(i,j,v) = findnz(sparsemat)

	for it = 1:length(i)

		broadcast(*, outmat[i[it],j[it]], v[it],densemat[i[it],j[it]])

	end

	return outmat

end

function Δ(td::Number)

	td == 0 ? 1 : 0
end

function Δ(td::Array)

	collect( td .== 0 )

end


text:
# sample

using MAT

vars = matread("c1p8.mat")

stim = vars["stim"]
ρ = vars["rho"]

max_stars_repo_name: voutasaurus/comp-neuro

max_issues_repo_name: voutasaurus/comp-neuro

max_forks_repo_name: voutasaurus/comp-neuro


text:
#Program using Agatz et al. method to make a TSP-D using a TSP solution
#See "Optimization Approaches for the Traveling SalesmanProblem with Drone" - Niels AgatzPaul BoumanMarie Schmidt
#######################
#Constantes

#######################
include("tsp.jl")

#=
structure of a TSP-D solution :
tsp : TSP solution used to create the TSP-D
drone : nodes where the drones are assigned
z : value of the objective function for this solution
=#
mutable struct Solution
	tsp::Array{Integer}
	drone::Array{Integer}
	combined::Array{Integer}
	z::Integer
	Solution(tsp::Array{Integer},drone::Array{Integer},combined::Array{Integer},z::Integer) = new(sol,drone,combined,z)
	Solution(tsp::Array{Integer},z::Integer) = new(sol,Array{Integer}(),Array{Integer}(),z)
end

struct Constants
	drone_speed::Integer
	truck_speed::Integer
	drone_range::Integer
end

function evaluate!(sol::Solution)

end
#=
function that assign the drone to the node
=#
function greedy_partitionning!(sol::Solution,con::Constants)
	simple = sol.tsp
	util = zeros(Integer, simple.length)
	for node in util
		node = cout(node-1 % length(sol.tsp),node,con.truck_speed) + cout(node,node+1 % length(sol.tsp),con.truck_speed) - max(cout(node-1% length(sol.tsp),node,con.drone_speed)+cout(node,node+1% length(sol.tsp),con.drone_speed),cout(node-1% length(sol.tsp),node+1% length(sol.tsp),con.truck_speed))
	end

	while !isempty(simple)
		select = whereis(max(util))
		makefly!(sol select)
		if cout(select-1% length(simple),select,1) <= con.drone_range && cout(select,select+1,% length(simple),1) <= con.drone_range
				push!(sol.drone,select)
				push!(sol.combine,select+1%length(sol.tsp))
				push!(sol.combine,select-1%length(sol.tsp))
		else

		pop!(simple,select)

	end

end

#=
function that assign a drone to a node (cf the article p.18)
PRE: node != 1 && node != length(sol.tsp)
=#
function make_fly!(sol::Solution,node::Integer)

	push!(sol.drone,node)
	push!(sol.combined,node+1)
	push!(sol.combined,node-1)

	#we save the performing of the operation
	sol.z -= cout(node-1% length(sol.tsp),node,con.truck_speed) + cout(node,node+1,% length(sol.tsp)con.truck_speed) - max(cout(node-1% length(sol.tsp),node,con.drone_speed)+cout(node,node+1% length(sol.tsp),con.drone_speed),cout(node-1% length(sol.tsp),node+1% length(sol.tsp),con.truck_speed))

end

#=

=#
function push_left!()

end

#=

=#
function push_right!()

end

max_stars_repo_name: Emmanuel-kun/TSPD

max_issues_repo_name: Emmanuel-kun/TSPD

max_forks_repo_name: Emmanuel-kun/TSPD


text:
"""
    getPhi(Zchain,Beta,W,niter,nchains,k)
get the topic/word posterior for `niter` samples
# Examples
```
getPhi(Zchain,Beta,W,niter,nchains,k)
```
"""
function getPhi(Zchain,Beta,W,niter,nchains,k,nlex;verbose=false)
    println("niter:",niter)
    Phi = SharedArray(zeros(niter,nlex,k,nchains))
    println("Phi size: ",size(Phi))
    for j in 1:nchains
        Threads.@threads for i in 1:niter
            verbose ? println("getting chain $j, iter $i") : nothing
            nw = permutedims(getZ(Zchain[i,:,j],k)*W)
            Phi[i,:,:,j] .= (nw .+ Beta) ./ sum(nw .+ Beta)
        end
    end
    Phi
end

"""
    getTheta(Zchain,J,doc_id,niter,nchains,k)
get the thetas for `niter` samples of document `doc_id`
# Examples
```
getTheta(Zchain,J,doc_id,niter,nchains,k)
```
"""
function getTheta!(Thetas,Zchain,J,Alpha,doc_id,niter,chain_id,k;verbose=false)
    @sync @distributed for i = 1:niter
        thetaIter!(Thetas,Zchain[i,:,chain_id],J,doc_id,Alpha,i,chain_id,k)
    end
end

function thetaIter!(Theta,Z,J,doc_id,Alpha,iter,chain_id,k)
    nj = getZ(Z,k)*J
    Theta[iter,:,chain_id,doc_id] .= vec((nj[:,doc_id] .+ Alpha) ./ sum(nj[:,doc_id] .+ Alpha))
end

"""
    getTheta(Zchain,J,doc_id,niter,nchains,k)
get the thetas for `niter` samples of document `doc_id`
# Examples
```
getTheta(Zchain,J,doc_id,niter,nchains,k)
```
"""
function getTheta(Zchain,J,Alpha,doc_id,niter,nchains,k;verbose=false)
    Theta = SharedArray(zeros(niter,k,nchains))
    for j in 1:nchains
        @sync @distributed for i = 1:niter
            thetaIterOld!(Theta,Zchain,J,doc_id,Alpha,i,j,k)
        end
    end
    Theta
end

function thetaIterOld!(Theta,Zchain,J,doc_id,Alpha,i,j,k)
    nj = getZ(Zchain[i,:,j],k)*J
    Theta[i,:,j] .= vec((nj[:,doc_id] .+ Alpha) ./ sum(nj[:,doc_id] .+ Alpha))
end

"""
    getThetaHat(thetas,ndocs,ntopic)
Get the estimate of theta by averaging samples across chains
# Examples
```
getThetaHat(thetas,ndocs,ntopic)
```
"""
function getThetaHat(thetas,ndocs,ntopic)
    thetahat0 = mean(thetas,dims=[1,3])
    thetahat = zeros(ndocs,ntopic)
    for i in 1:size(thetahat0)[4]
        thetahat[i,:] .= vec(thetahat0[:,:,:,i])
    end
    thetahat
end

"""
    getPhiHat(phis,nlex,ntopic)
Get the estimate of phi by averaging samples across chains
# Examples
```
getPhiHat(phis,nlex,ntopic)
```
"""
function getPhiHat(phis,nlex,ntopic)
    phihat0 = mean(phis,dims=[1,4])
    phihat = zeros(nlex,ntopic)
    for i in 1:size(phihat0)[3]
        phihat[:,i] .= phihat0[1,:,i,1]
    end
    phihat
end

"""
    getHPI(x; alpha=0.05)
Get the highest-probability interval
# Examples
```
getHPI(rand(Beta(6,2),10000))
```
"""
function getHPI(x; alpha=0.05)
    # this is MCMCChains._hpd @ master "2ff0a26beda7223552e47d25d948d6c7e44baa95"
    n = length(x)
    m = max(1, ceil(Int, alpha * n))

    y = sort(x)
    a = y[1:m]
    b = y[(n - m + 1):n]
    _, i = findmin(b - a)

    return [a[i], b[i]]
end

"""
    getETI(x; alpha=0.05)
Get the equal-tailed interval (quantiles)
# Examples
```
getETI(rand(Beta(6,2),10000))
```
"""
function getETI(x; alpha=0.05)
    a=quantile!(x,alpha/2)
    b=quantile!(x,1-alpha/2)
    return [a,b]
end


text:
include("constants.jl")

#λν = c
fre2wl(ν)    = c / ν #cm
wl2fre(λ)    = c / λ #Hz

#Distance
m2cm(m)      = m   * 100
mm2cm(mm)    = mm  * 0.1
mm2m(mm)     = mm  * 0.001
mm2um(mm)    = mm  * 1000
cm2m(cm)     = cm  * 0.01
cm2mm(cm)    = cm  * 10

#Frequency
hz2ghz(hz)   = hz  / 1e9
hz2mhz(hz)   = hz  / 1e6
ghz2hz(ghz)  = ghz * 1e9
mhz2hz(mhz)  = mhz * 1e6

#Angle
arcs2dg(arcs) = arcs * 0.00028

export fre2wl, wl2fre,
       m2cm,
       mm2cm, mm2m, mm2um,
       cm2m, cm2mm,
       hz2ghz, hz2mhz,
       ghz2hz, mhz2hz,
       arcs2dg


text:
"""
Linear-Stability-Calculators: Shallow Water Model
=================================================
"""

include("src/shallow_water_setup.jl")

             ks = Dict()
ks[Cartesian()] = collect(2e-5:2e-5:2.5e-3)
ks[Spherical()] = collect(2e-5:2e-5:2.5e-3) * 6.3781e6 * cos(π/4)

        Uj = 1.0                 # Jet Parameters
   Ljscale = 20
    Nmodes = 2
         σ = Dict()
         ω = Dict()
    σmodes = Dict()
      phys = Dict()
background = Dict()

geometrys = (Cartesian(), Spherical())

for geometry in geometrys

     @info string("Setting Geometry = ", geometry)

                    grid = Grid(geometry, N=250)
          phys[geometry] = Physics(geometry; grid=grid)
    background[geometry] = Bickley_Jet(geometry; phys=phys[geometry], Uj, Ljscale)

     plot_basic_state(
          grid,
          background,
          geometry,
          string("basic_state_",typeof(geometry),".png")
          )

     for (index, k) in enumerate(ks[geometry])
          A = build_matrix(geometry; k, background, phys)
          σ[(k,geometry)], ω[(k,geometry)], σmodes[(k,geometry)] = compute_spectrum(A, k, phys, Nmodes);
          @info string("     k = ", k, " and max growth = ", maximum(σ[(k,geometry)]), "\n")
     end

     plot_growth_rates(
          ks,
          σ,
          phys,
          geometry,
          Nmodes,
          string("growth_rates_",typeof(geometry),"_Bickley_Jet.png")
          )

     mode_number = 1;
           σ_max = maximum( σ[(k,geometry)][mode_number] for k in ks[geometry]);
         k_index = sortperm([σ[(k,geometry)][mode_number] for k in ks[geometry]],rev=true)[1];
               k = ks[geometry][k_index];

     print("σ_max = ",σ_max, " with k = ", k, " at k_index =", k_index, "\n")
     plot_1D_streamfunction(
          k_index,
          k,
          phys,
          σmodes,
          geometry,
          mode_number,
          string("modes_1D_streamfunction_",typeof(geometry),".png")
          )

     plot_2D_streamfunction(
          k_index,
          k,
          phys,
          σmodes,
          geometry,
          mode_number,
          string("modes_2D_streamfunction_",typeof(geometry),".png")
          )

     plot_2D_vorticity(
          phys,
          k_index,
          k,
          σmodes,
          geometry,
          mode_number,
          string("modes_2D_vorticity_",typeof(geometry),".png")
          )

end


text:
include("juliacomputing-repos.jl")
include("julialang-repos.jl")
include("miscellaneous-gpu-related-repos.jl")
include("miscellaneous-julia-related-repos.jl")
include("miscellaneous-ml-related-repos.jl")
include("unregistered-packages.jl")

const ADDITIONAL_REPOS = convert(
    Vector{MirrorUpdater.Types.SrcDestPair},
    vcat(
        JULIACOMPUTING_REPOS,
        JULIALANG_REPOS,
        MISCELLANEOUS_GPU_RELATED_REPOS,
        MISCELLANEOUS_JULIA_RELATED_REPOS,
        MISCELLANEOUS_ML_RELATED_REPOS,
        UNREGISTERED_PACKAGES,
        ),
    )


text:
using DataStructs
using Test

tests = [
         "list",
        ]

if length(ARGS) > 0
    tests = ARGS
end
@testset "DataStructs.jl" begin
    # Write your tests here.
for t in tests
    fp = joinpath(dirname(@__FILE__), "test_$t.jl")
    println("$fp ...")
    include(fp)
end

end


text:
################################################################################
#
#   METHODS FOR DUMPING A LATTICE TO BLENDER INPUT FILE
#   (input file can be used in Blender with the provided AddOn)
#
#   STRUCTURE OF THE FILE
#
#   1) SAVING LATTICE DATA TO TEXT FILE FOR BLENDER USE
#
################################################################################

# function to save lattice data to a file for blender to read
"""
    saveBlenderFile(
            lattice::Lattice,
         [  filename_output::String="AUTO"
          ; visualize_periodic::Bool=false,
            site_radius::Float64=-0.2,
            bond_thickness::Float64=-0.1,
            print_used_options::Bool=false,
            site_labels::String="OFF",
            colorcode_sites::Dict = Dict(0 => [255,255,255]),
            colorcode_bonds::Dict = Dict("0" => [0,0,0]),
            colorcode_bonds_automation::String = "OFF" ]
        )

Save the `Lattice` object `lattice` to a file so that Blender can be used for visualizing the lattice.
The filename can be specified optionally as well as some other plotting options.

Note: 2D lattices will be extended with z=0 to 3D.

# Options

Additional options include the following:

- `filename_output::String` the filename of the output file. If not chosen explicitly, automatically sets itself based on the lattice filename.

- `colorcode_bonds::Dict` a dictonary in which every bond interaction strength (as string) is mapped to a color
- `colorcode_bonds_automation::String` determines if the bond colorcode defined by `colorcode_bonds` should be redefined
  based on bond-strength. Possible options are
    - `"OFF"`   no redefinition
    - `"GREY"`  redefinition with random grey colors
    - `"COLOR"` redefinition with random colors
- `bond_thickness::Float64` gives the thickness of bonds in the plot. For positive numbers, the unit is px, for negative numbers,
  the thickness is given in terms of shortest bond length.
- `visualize_periodic::Bool` determines wheather periodic connections should be drawn as dashed lines

- `colorcode_sites::Dict` a dictonary in which every site index (as integer) is mapped to a color
- `site_radius::Float64` gives the radius of sites in the plot. For positive numbers, the unit is px, for negative numbers,
  the radius is given in terms of shortest bond length.
- `site_labels::String` (NOT IMPLEMENTED SO FAR) options for on-site labels, possible options are
    - `"OFF"`   no labels
    - `"POSITION INDEX"` indices given by the `positions_indices` field inside the `Lattice` object
    - `"LATTICE INDEX"` the indices with which the site can be adressed in the `Lattice` object

- `print_used_options::Bool` determines if the options that are chosen should be printed as well

# Blender AddOn usage

To plot the resulting file in Blender, proceed as follows:

#### Installing the AddOn (only once)

1) Open Blender
2) Open User Preferences (Ctrl + Alt + U)
3) Go to "Add-ons"
4) Choose "Install from file" (bottom line)
5) Choose the provided script, located in ".../LatticePhysics.jl/blender/LatticePhysics_blender_addon.py"
6) Enable the AddOn by ticking the box next to the newly appeared label.
7) Save user Preferences to load AddOn on startup of Blender

#### Plotting lattices

1) Within a Blender project, hit spacebar
2) Type "LatticePhysics" and select the AddOn from the list
3) Choose the file to import
4) Optionally adjust some settings for the poly resolution
5) Use "Import LatticePhysics dump file"

# Examples

```julia-repl
julia> saveBlenderFile(lattice)
...

julia> saveBlenderFile(lattice, "myfile.txt")
...

julia> saveBlenderFile(lattice, site_radius=-0.2, bond_thickness=-0.1)
...

```
"""
function saveBlenderFile(
            lattice::Lattice,
            filename_output::String="AUTO";
            visualize_periodic::Bool=false,
            site_radius::Float64=-0.2,
            bond_thickness::Float64=-0.1,
            print_used_options::Bool=false,
            site_labels::String="OFF",
    		colorcode_sites::Dict = Dict(0 => [255,255,255]),
    		colorcode_bonds::Dict = Dict("0" => [0,0,0]),
            colorcode_bonds_automation::String = "OFF"
        )

    ##########---------------------------------
    # STEP 1 #  INITIALIZATION OF PARAMTERS
    ##########---------------------------------

    # Maybe print what is being done
    if print_used_options
        println("Dumping lattice object with filename \"$(lattice.filename)\" to blender file")
    end

    # define the FILENAME of the output file if it is set to "AUTO"
    if filename_output=="AUTO"
        filename_output = "$(lattice.filename[1:end-4])_blender_dump.txt"
        filename_output = replace(filename_output, FOLDER_LATTICES, "")
        # maybe print
        if print_used_options
            println("- Changing Blender filename to \"$(filename_output)\"")
        end
    elseif print_used_options
        println("- Using given SVG filename \"$(filename_output)\"")
    end

    ############---------------------------------
    # STEP 2.1 #  LOADING OF CONNECTION DATA
    ############---------------------------------

	# connections to plot (all)
	connections     = deepcopy(lattice.connections)
	positions	    = deepcopy(lattice.positions)
    indices_to_plot = deepcopy(lattice.positions_indices)

    # Function to determine if a connection should be plotted
    # 0 yes, no periodic
    # 1 yes, periodic
    # -1 no
    function shallBePlotted(c::Array{Any,1})
        # check if the indices are okay (exclude double bonds)'
        if c[1] > c[2]
            return -1
        end
        # check the sum of all wraps
        if sum(abs.(c[4])) > 0
            if visualize_periodic
                return 1
            else
                return -1
            end
        else
            return 0
        end
    end

    # print plotting of connections
    if print_used_options
        if visualize_periodic
            println("- Periodic connections will be plotted, based on visualize_periodic=$(visualize_periodic)")
        else
            println("- Periodic connections will not be plotted, based on visualize_periodic=$(visualize_periodic)")
        end
    end

    ############-------------------------------------
    # STEP 2.2 #  DETERMINATION OF SITE & BOND DATA
    ############-------------------------------------

    ##############--------------
    # STEP 2.2.1 #  SITE DATA
    ##############--------------

    # check if site radius has to be chosen
    if site_radius < 0
        # site radius is given in relative units, first maybe print
        if print_used_options
            println("- Site radius is given in relative units to be site_radius=$(site_radius)")
        end
        # then, calculate the minimum length of connections
        min_length = 999999999
        for c in connections
            # check if the connection is NOT PERIODIC and NOT FROM SITE TO ITSELF
            if sum(abs.(c[4]))==0 && c[1]!=c[2]
                # determine the length of the position on the canvas
                c_vector = positions[Int(c[1])] .- positions[Int(c[2])]
                c_length = sqrt(sum(c_vector .* c_vector))
                # check if smaller then minimum length
                min_length = min(min_length, c_length)
            end
        end
        # set the radius of a site to be -site_radius*min_length
        site_radius = -site_radius*min_length
        # maybe print
        if print_used_options
            println("- Site radius changed to absolute units to be site_radius=$(site_radius) as minimal bond length is $(min_length)")
        end
    else
        # site radius is given in absolute units, maybe print
        if print_used_options
            println("- Site radius is given in absolute units to be site_radius=$(site_radius)")
        end
    end

    # repair color dictonary by making sure default elements exist
    colorcode_sites[0]   = get(colorcode_sites, 0, [255,255,255])
    # maybe print
    if print_used_options
        println("- Site colors passed with dictonary:")
        for k in keys(colorcode_sites)
            println("    - $(k) -> $(colorcode_sites[k])")
        end
    end

    ##############--------------
    # STEP 2.2.2 #  BOND DATA
    ##############--------------

    # check if bond thickness has to be chosen
    if bond_thickness < 0
        # bond thickness is given in relative units, first maybe print
        if print_used_options
            println("- Bond thickness is given in relative units to be bond_thickness=$(bond_thickness)")
        end
        # then, calculate the minimum length of connections
        min_length = 999999999
        for c in connections
            # check if the connection is NOT PERIODIC and NOT FROM SITE TO ITSELF
            if sum(abs.(c[4]))==0 && c[1]!=c[2]
                # determine the length of the position on the canvas
                c_vector = positions[Int(c[1])] .- positions[Int(c[2])]
                c_length = sqrt(sum(c_vector .* c_vector))
                # check if smaller then minimum length
                min_length = min(min_length, c_length)
            end
        end
        # set the thickness of a bond to be -bond_thickness*min_length
        bond_thickness = -bond_thickness*min_length
        # maybe print
        if print_used_options
            println("- Bond thickness changed to absolute units to be bond_thickness=$(bond_thickness) as minimal bond length is $(min_length)")
        end
    else
        # bond thickness is given in absolute units, maybe print
        if print_used_options
            println("- Bond thickness is given in absolute units to be bond_thickness=$(bond_thickness)")
        end
    end

    # HANDLING BOND COLORS

    # get list of all connections strengths
    cs_list = getConnectionStrengthList(lattice)
    # parse all elements to strings
    cs_list = String["$(element)" for element in cs_list]
    # sort the list
    sort!(cs_list)

    # BOND COLOR DICTONARY
    # maybe overwrite dictonary
    if colorcode_bonds_automation == "GREY"
        # get the color code list
        cc_list = LatticePhysics.getRandomGreys(length(cs_list))
        # put in a new dictonary
        colorcode_bonds = Dict()
        # insert all pairs
        for i in 1:size(cc_list, 1)
            colorcode_bonds[cs_list[i]] = cc_list[i]
        end
        # maybe print
        if print_used_options
            println("- Using random GREY for bonds, $(length(cc_list)) greys used")
        end
    elseif colorcode_bonds_automation == "COLOR"
        # get the color code list
        cc_list = LatticePhysics.getRandomColors(length(cs_list))
        # put in a new dictonary
        colorcode_bonds = Dict()
        # insert all pairs
        for i in 1:size(cc_list, 1)
            colorcode_bonds[cs_list[i]] = cc_list[i]
        end
        # maybe print
        if print_used_options
            println("- Using random COLOR for bonds, $(length(cc_list)) colors used")
        end
    elseif colorcode_bonds_automation == "OFF"
        # maybe print
        if print_used_options
            println("- No redefinition of bond colors, using provided dictonary:")
            for k in keys(colorcode_bonds)
                println("    - \"$(k)\" -> $(colorcode_bonds[k])")
            end
        end
    else
        # maybe print
        if print_used_options
            println("- Could not identify redefinition of bond colors based on colorcode_bonds_automation=\"$(colorcode_bonds_automation)\", using provided dictonary:")
            for k in keys(colorcode_bonds)
                println("    - \"$(k)\" -> $(colorcode_bonds[k])")
            end
        end
    end

	# repair color dictonary by making sure default elements exist
	colorcode_bonds["0"] = get(colorcode_bonds, "0", [0,0,0])

    ##########---------------------
    # STEP 3 #  BLENDER FILE CREATION
    ##########---------------------

    # open the file in write mode
    f = open(filename_output, "w")

    # write all materials (bond / site color directories)
    for k in keys(colorcode_sites)
        # write start of the line
        write(f, "material:\t")
        # write the name
        write(f, "site_material_$(k)\t")
        # write the color
        site_color = colorcode_sites[k]
        write(f, "$(site_color[1]), $(site_color[2]), $(site_color[3])")
        # write line break
        write(f, "\n")
    end
    for k in keys(colorcode_bonds)
        # write start of the line
        write(f, "material:\t")
        # write the name
        write(f, "bond_material_$(k)\t")
        # write the color
        connection_color = colorcode_bonds[k]
        write(f, "$(connection_color[1]), $(connection_color[2]), $(connection_color[3])")
        # write line break
        write(f, "\n")
    end

    # write all sites
    for (i,p) in enumerate(positions)
        # write start of the line
        write(f, "site:\t$(i), ")
        # write the position
        if length(p) == 2
            write(f, "$(p[1]), $(p[2]), 0.0")
        elseif length(p) == 3
            write(f, "$(p[1]), $(p[2]), $(p[3])")
        else
            println("site has not correct amount of dimensions!")
        end
        # write the radius
        write(f, ", $(site_radius)")
        # write the color
        if indices_to_plot[i] in keys(colorcode_sites)
            write(f, "\tsite_material_$(indices_to_plot[i])")
        else
            write(f, "\tsite_material_$(0)")
        end
        # write line break
        write(f, "\n")
    end

    # write all bonds
	for (i,c) in enumerate(connections)
        # non periodic or periodic and wanted to be plotted
		if shallBePlotted(c) >= 0
            # write start of the line
            write(f, "bond:\t$(i), ")
            # write the position from where
            p = positions[Int(c[1])]
            if length(p) == 2
                write(f, "$(p[1]), $(p[2]), 0.0")
            elseif length(p) == 3
                write(f, "$(p[1]), $(p[2]), $(p[3])")
            else
                println("from site has not correct amount of dimensions!")
            end
            # write the position to where
            p = positions[Int(c[2])]
            if length(p) == 2
                write(f, ", $(p[1]), $(p[2]), 0.0")
            elseif length(p) == 3
                write(f, ", $(p[1]), $(p[2]), $(p[3])")
            else
                println("from site has not correct amount of dimensions!")
            end
            # write the thickness
            write(f, ", $(bond_thickness)")
            # write the color
            if string(c[3]) in keys(colorcode_bonds)
                write(f, "\tbond_material_$(string(c[3]))")
            else
                write(f, "\tbond_material_0")
            end
            # write line break
            write(f, "\n")
        end
	end

    # close the file
    close(f)

    ##########---------------------
    # STEP 4 #  FILE FINISH
    ##########---------------------

    # return the filename
    return filename_output

end
export saveBlenderFile


text:
# VSCode users:
# Do not run next line of code with Ctrl-Enter!
# Run next line of code with Alt-Enter (run line), Shift-Enter (run cell), or Ctrl-F5 (run file)
mypath = dirname(@__FILE__)

## ##########################################################################
# Initialization

using ModelBaseEcon
using StateSpaceEcon
using TimeSeriesEcon

using Test
using Plots
using Random
using Distributions

# Fix the random see for reproducibility
Random.seed!(1234);

nothing
## ##########################################################################
#### Part 1: The model

### Loading the model

# The model is described in its own dedicated module, which is contained in its
# own file, `SW07.jl`. We can load the module with `using SW07`; the model itself
# is a global variable called `model` within that module.

unique!(push!(LOAD_PATH, mypath))
using SW07
m = SW07.model

## ##########################################################################
### Examining the model

# This model is too big to fit all of its details in the REPL window, so only
# summary information is displayed. We can see the entire model with `fullprint`.
fullprint(m)

# We can also examine individual components using the commands `parameters`,
# `variables`, `shocks`, `equations`.
parameters(m)

## ##########################################################################
### Setting the model parameters

# We must not change any part of the model in the active Julia session except for
# the model parameters and steady state constraints.
# If we want to add variables, shocks, or equations, we
# must do so in the model module file and restart Julia to load the new model.

# When it comes to the model parameters, we can access them by their names from
# the model object using dot notation.
m.crr # read a parameter value
m.cgy = 0.5187 # modify a parameter value

# !!! note
#     In this model the values of the parameters have been set according to the
#     [replication data](http://doi.org/10.3886/E116269V1).

## ##########################################################################
### Model flags and options

# In addition to model parameters, which are values that appear in the model
# equations, the model object also holds two other sets of parameters, namely
# flags and options.

# Flags are (usually boolean) values which characterize the type of model we have.
# For example, a linear model should have its `linear` flag set to `true`.
# Typically, this is done in the model file before calling `@initialize`.
m.flags

# Options are values that adjust the operations of the algorithms. For example, we
# have `tol` and `maxiter`, which set the desired accuracy and maximum number of
# iterations for the iterative solvers. These can be adjusted as needed at any
# time. Another useful option is `verbose`, which controls the level of verbosity
# of the different commands.

# Many functions in `StateSpaceEcon` have optional arguments of the same name as a
# model option. When the argument is not explicitly given in the function call,
# these functions will use the value from the model option of the same name.
m.verbose = true
m.options

## ##########################################################################
#### Part 2: The Steady state solution

# The steady state is a special solution of the dynamic system that remains
# constant over time. It is important on its own, but also it can be useful in
# several ways. For example, linearizing the model requires a particular solution
# about which to linearize, and the steady state is typically used for this
# purpose.

# In addition to the steady state, we also consider another kind of special
# solution which grows linearly in time. If we know that the steady state solution
# is constant (i.e., its slope is zero), we can set the model flag `ssZeroSlope`
# to `true`. This is not required; however in a large model it might help the
# steady state solver converge faster to the solution.

# The model object `m` stores information about the steady state. This includes
# the steady state solution itself, as well as a (possibly empty) set of additional
# constraints that apply only to the steady state. This information can be
# accessed via `m.sstate`.
m.sstate

## ##########################################################################
### Steady state constraints

# Sometimes the steady state is not unique, and we can use steady state
# constraints to specify the particular steady state we want. Also, if the model
# is non-linear, these constraints can be used to help the steady state solver
# converge. Steady state constraints can be added with the `@steadystate`
# macro. The constraint can be as simple as giving a specific value; we can also
# write an equation with multiple variables. We're allowed to use model parameters
# in these equations as well.

@steadystate m a = 5
m.sstate

# We can clean up the constraints by emptying the constraints container.

empty!(m.sstate.constraints)
m.sstate

# !!! note "Important note"
#     Steady state constraints that are always valid can be pre-defined in the model file.
#     In that case, all calls to the `@steadystate` macro must be made after calling
#     `@initialize`.

## ##########################################################################
### Solving for the steady state

# The steady state solution is stored within the model object. Before solving, we
# have to specify an initial condition. If the model is linear, this makes no
# difference, but in a non-linear model a good or a bad initial guess might be the
# difference between success and failure of the steady state solver.

# We specify the initial guess by calling `clear_sstate!`. This call
# removes any previously stored solution, sets the initial guess, and runs the
# pre-solve pass of the steady state solver. The initial guess can be given with
# the `lvl` and `slp` arguments; if not provided, an initial guess is chosen
# automatically.

# Once that's done, we call `sssolve!` to find the steady state. This
# function returns a `Vector{Float64}` containing the steady state solution, and
# it also writes that solution into the model object. The vector is of length
# `2*nvariables(m)` and contains the level and the slope for each variable.

clear_sstate!(m)
sssolve!(m);

# If in doubt, we can use `check_sstate` to make sure the steady state solution
# stored in the model object indeed satisfies the steady state system of equations.
# This function returns the number of equations that are not satisfied.
# A value of 0 is what we want to see. In verbose mode, it also lists the
# problematic equations and their residuals.

@test check_sstate(m) == 0

## ##########################################################################
### Examining the steady state

# We can access the steady state solution via `m.sstate` using the dot notation.
m.sstate.dc

# We can also assign new values to the steady state solution, but we should be
# careful to make sure it remains a valid steady state solution.
m.sstate.dc.level = 0.43121
@test check_sstate(m) > 0

# As the code above shows, a wrong steady state solution (based on the specified
# precision in the `tol` option) will result in one or more equation not being
# satisfied. Let's put back the correct value.
m.sstate.dc.level = 0.4312
@test check_sstate(m) == 0

# We can examine the entire steady state solution with `printsstate`.
printsstate(m)

## ##########################################################################
#### Part 3: Impulse response

### Simulation plan

# Before we can simulate the model, we have to decide on the length of the
# simulation and what data is available for each period, i.e., what values are
# known (exogenous). This is done with an object of type `Plan`.

# To create a plan, all we need is the model object and a range for the
# simulation.
sim_rng = 2000Q1:2039Q4
p = Plan(m, sim_rng)

# The plan shows us the list of exogenous values (variable or shocks) for each
# period or sub-range of the simulation. By default, all shocks are exogenous and
# all variables are endogenous.

# We also see that the range of the plan has been extended before and after the
# simulation range. This is necessary because we need to set initial and final
# conditions. The number of periods for initial conditions is equal to the largest
# lag in the model. Similarly, final conditions have to be imposed over as many
# periods as the largest lead.

p.range          # the full range of the plan
init_rng = first(p.range):first(sim_rng)-1   # the range for initial conditions
final_rng = last(sim_rng)+1:last(p.range)     # the range for final conditions
@test length(init_rng) == m.maxlag
@test length(final_rng) == m.maxlead

## ##########################################################################
### Exogenous data

# We have to provide the data for the simulation. We start with all zeros and fill
# in the external data, which must include initial conditions for all variable and
# shocks, exogenous values (according to the plan), and possibly final conditions.

#### Initial conditions

# In this example, we want to simulate an impulse response, so it makes sense to
# start from the steady state, so that is what we set as the initial condition. We
# leave the initial conditions for the shocks at 0.
exog = zerodata(m, p);
for var in variables(m)
    exog[init_rng, var] = m.sstate[var].level
end
exog

# !!! tip "Pro tip"
#     The above works because the steady state is stationary, i.e., all slopes
#     are zero. If we had a model with linear growth steady state, we could do
#     something like the following (see `@rec`):
    for var in variables(m)
        ss = m.sstate[var]
        exog[init_rng, var] = ss.level
        if ss.slope != 0
            # recursively update by adding the slope
            @rec init_rng[2:end] exog[t, var] = exog[t - 1, var] + ss.slope
        end
    end

#### Final conditions

# For the final conditions we can use the steady state again, because we expect
# that the economy will eventually return to it if the simulation is sufficiently
# long past the last shock. We can do this by assigning the values of the steady
# state to the final periods after the simulation, similarly to what we did with the
# initial conditions.

# Alternatively, we can specify that we want to use the steady state in the call
# to simulate by passing `fctype=fclevel`. Yet another possibility is to set the
# final condition so that the solution slope matches the slope of the steady state
# by setting `fctype=fcslope`. In both cases, we do not need to set anything in
# the exogenous data array because those values would be ignored.

# !!! tip "Pro tip"
#     In the Smets and Wouters 2007 model the two ways of using the steady state
#     for final conditions (level or slope) are equivalent, because the steady
#     state here is stationary and unique. In models where the steady state has
#     non-zero slope, or the steady state has zero slope but the level is not
#     unique, we should use `fctype=fcslope`.

#### A quick sanity check

# If we were to run a simulation where the economy started in the steady state and
# there were no shocks at all, we'd expect that the economy would remain in steady
# state forever.

ss = simulate(m, exog, p; fctype=fcslope);

# The simulated data, `ss`, should equal (up to the accuracy of the solution) the
# steady state data. Similar to `zerodata`, we can use `steadystatedata`
# to create a data set containing the steady state solution.

@test ss ≈ steadystatedata(m, p)

#### Exogenous data

# All shocks are exogenous by default. All we have left to do is to set the value
# of the shock.

# Let's say that we want to shock `epinf` for the first four quarters by `0.1`.
exog[sim_rng[1:4], :epinf] = 0.1;
exog[shocks(m)]

#### Running the simulation

# We call `simulate`, providing the model, the exogenous data, and the
# plan. We also specify the type of final condition we want to impose.
irf = simulate(m, exog, p, fctype=fcslope);

# We can now take a look at how some of the observable variables in the model have
# responded to this shock. We use `plot` from the `Plots` package to for that. We
# specify the variables we want to plot using `vars` and the names of the
# datasets being plotted (for the legend) in the `names` option.
plot(ss, irf,
     vars=(:pinfobs, :dy, :labobs, :robs),
     names=("SS", "IRF"),
     legend=[true false false false],
     size=(600, 400)
    );

png_fname = joinpath(mypath, "irf.png")
rm(png_fname, force=true)
savefig(png_fname)
@test isfile(png_fname)

## ##########################################################################
#### Part 4: Stochastic shocks simulation

# Now let's run a simulation with stochastic shocks. We will have random shocks
# over two years and then have no shocks for several years afterwards to allow
# time for the economy to return to its steady state.

sim_rng = 2000Q1:2049Q4      # simulate 50 years starting 2000
shk_rng = 2004Q1 .+ (0:7)    # shock 8 quarters starting in 2004
p = Plan(m, sim_rng)
init_rng = first(p.range):first(sim_rng) - 1
final_rng = last(sim_rng) + 1:last(p.range)
exog = zerodata(m, p);
for v in variables(m)
    exog[init_rng, v] = m.sstate[v].level
end

# The distributions of the shocks are assumed normal with mean zero and standard
# deviations that have been estimated in the [replication data](http://doi.org/10.3886/E116269V1).
# We use packages `Distributions` and `Random` to draw the
# necessary random values.

shk_dist = (ea = Normal(0.0, 0.4618),
            eb = Normal(0.0, 1.8513),
            eg = Normal(0.0, 0.6090),
            eqs = Normal(0.0, 0.6017),
            em = Normal(0.0, 0.2397),
            epinf = Normal(0.0, 0.1455),
            ew = Normal(0.0, 0.2089));

for (shk, dist) in pairs(shk_dist)
    exog[shk_rng, shk] = rand(dist, length(shk_rng))
end
exog[shk_rng, shocks(m)]

# Now we are ready to simulate. We can set the shocks to be anticipated or
# unanticipated by setting the `anticipate` parameter in `simulate`.

sim_a = simulate(m, exog, p; fctype=fcslope, anticipate=true);
sim_u = simulate(m, exog, p; fctype=fcslope, anticipate=false);

# As before, we can review the responses of the observed variables to these shocks
# using `plot`.

observed = (:dy, :dc, :dinve, :labobs, :pinfobs, :dw, :robs);
ss = steadystatedata(m, p);
plot(ss, sim_a, sim_u,
     vars=observed,
     names=("SS", "Anticipated", "Unanticipated"),
     legend=[true (false for i = 1:6)...],
     linewidth=1.5,
     size=(900, 600)
    );

png_fname = joinpath(mypath, "stoch_shk.png")
rm(png_fname, force=true)
savefig(png_fname)
@test isfile(png_fname)

# We see that when the shocks are anticipated, the variables start to react to
# them right away; in the unanticipated case, there is no movement until the
# shocks actually hit.

## ##########################################################################
#### Part 5: Backing out historical shocks

# Now let's pretend that the simulated values are historical data and that we do
# not know the magnitude of the shocks. We can treat the observed (simulated)
# values of the variables as known by making them exogenous. At the same time we
# will make the shocks endogenous, so that we can solve for their values during
# the simulation.

# The "historic" (simulated, assumed observed) range is from the first period of
# the simulation until the last shock in the previous exercise.
hist_rng = first(sim_rng):last(shk_rng)

# We use `exogenize!` and `endogenize!` to set up a plan in which
# observed variables are exogenous and shocks are endogenous throughout the
# historic range.
endogenize!(p, shocks(m), hist_rng);
exogenize!(p, observed, hist_rng);
p

# As we can see above, the plan now reflects our intentions.

# Finally, we need to set up the exogenous data. This time we do not specify the
# shocks; instead, we assign the known data for the observed variables for the
# historic range. We start with initial conditions.
exog = zerodata(m, p);
for v in variables(m)
    exog[init_rng, v] = m.sstate[v].level
end

# We take the observed data from the simulation above. We show the anticipated
# version first.
for v in observed
    exog[hist_rng, v] = sim_a[v]
end
back_a = simulate(m, exog, p, fctype=fcslope, anticipate=true);

# Now we show the unanticipated case.
for v in observed
    exog[hist_rng, v] = sim_u[v]
end
back_u = simulate(m, exog, p, fctype=fcslope, anticipate=false);

# If we did everything correctly, the shocks we recovered must match exactly the
# shocks we used when we simulated the "historical" data.
@test sim_a[shocks(m)] ≈ back_a[shocks(m)]
@test sim_u[shocks(m)] ≈ back_u[shocks(m)]

# Moreover, we must have the unobserved variables match as well. In fact, all the
# data must match over the entire simulation range.
@test sim_a ≈ back_a
@test sim_u ≈ back_u

### References

# [Smets, F., Wouters, R., 2007. Shocks and frictions in US business cycles: A bayesian DSGE approach. The American Economic Review 97(3), 586–606.](https://www.aeaweb.org/articles?id=10.1257/aer.97.3.586)


text:
Kallen(x,y,z) = x^2+y^2+z^2 - 2x*y - 2y*z - 2z*x
KallenFact(a²,b²,c²) = ((a,b,c)=sqrt.((a²,b²,c²)); (a-(b+c))*(a-(b-c))*(a+(b+c))*(a+(b-c)))
Kibble(σs,msq) = Kallen(Kallen(msq[4],msq[1],σs[1]),
                        Kallen(msq[4],msq[2],σs[2]),
                        Kallen(msq[4],msq[3],σs[3]))
#
"""
    σkofi(k,z,σi,msq)

Computes invariant σk = (pi + pj)² from
the scattering angle z=cosθjk in the rest from of (j,k),
given the mass of the system m(j,k)² = σi

Explicit forms: `σ3of1`, `σ1of2`, `σ2of3`.

See also `σkofj(k,z,σj,msq)`
"""
function σkofi(k,z,σi,msq)
    i,j,_ = ijk(k)
    #
    s = msq[4]
    EE4σ = (σi+msq[j]-msq[k])*(s-σi-msq[i])
    p²q²4σ = Kallen(σi,msq[j],msq[k])*Kallen(s,σi,msq[i])
    p²q²4σ = (p²q²4σ < 0) ? 0.0 : p²q²4σ # for numerical errors
    σk = msq[i]+msq[j] + (EE4σ + sqrt(p²q²4σ)*z) / (2σi)
    return σk
end
σ3of1(z,σ1,msq) = σkofi(3,z,σ1,msq)
σ1of2(z,σ2,msq) = σkofi(1,z,σ2,msq)
σ2of3(z,σ3,msq) = σkofi(2,z,σ3,msq)

"""
    σkofj(k,z,σj,msq)

Computes invariant σk = (pi + pj)² from
the scattering angle z=cosθki in the rest from of (k,i),
given the mass of the system m(k,i)² = σj

Explicit forms: `σ3of2`, `σ1of3`, `σ2of1`.

See also `σkofi(k,z,σi,msq)`
"""
function σkofj(k,z,σj,msq)
    σi = σkofi(k,z,σj,msq)
    sum(msq)-σi-σj
end
σ3of2(z,σ2,msq) = σkofj(3,z,σ2,msq)
σ1of3(z,σ3,msq) = σkofj(1,z,σ3,msq)
σ2of1(z,σ1,msq) = σkofj(2,z,σ1,msq)

# Scattering angle

"""
    cosθij(k,σs,msq)

Isobar decay angle for the chain-k, i.e.
an angle of between vectors pi and -pk in the (ij) rest frame.

Explicit forms: `cosθ23`, `cosθ31`, `cosθ12`.
"""
function cosθij(k,σs,msq)
    (i,j) = ij_from_k(k)
    #
    s = msq[4]
    EE4σ = (σs[k]+msq[i]-msq[j])*(s-σs[k]-msq[k])
    pp4σ = sqrt(Kallen(σs[k],msq[i],msq[j])*Kallen(s,σs[k],msq[k]))
    rest = σs[j]-msq[k]-msq[i]
    return (2σs[k]*rest-EE4σ)/pp4σ
end
cosθ23(σs,msq) = cosθij(1,σs,msq)
cosθ31(σs,msq) = cosθij(2,σs,msq)
cosθ12(σs,msq) = cosθij(3,σs,msq)

max_stars_repo_name: mmikhasenko/ThreeBodyDecay

max_issues_repo_name: mmikhasenko/ThreeBodyDecay

max_forks_repo_name: mmikhasenko/ThreeBodyDecay


text:
###############################################################################
# TASK ########################################################################
###### The `task` field should not be accessed outside this section. ##########

## Affine Constraints #########################################################
##################### lc ≤ Ax ≤ uc ############################################

function allocateconstraints(m::MosekModel, N::Int)
    numcon = getnumcon(m.task)
    alloced = ensurefree(m.c_block,N)
    id = newblock(m.c_block, N)

    if alloced > 0
        appendcons(m.task, alloced)
    end
    return id
end

function getconboundlist(t::Mosek.Task, subj::Vector{Int32})
    n = length(subj)
    bk = Vector{Boundkey}(undef,n)
    bl = Vector{Float64}(undef,n)
    bu = Vector{Float64}(undef,n)
    for i in 1:n
        bki,bli,bui = getconbound(t,subj[i])
        bk[i] = bki
        bl[i] = bli
        bu[i] = bui
    end
    bk,bl,bu
end

# `putbaraij` and `putbarcj` need the whole matrix as a sum of sparse mat at once
function split_scalar_matrix(m::MosekModel, terms::Vector{MOI.ScalarAffineTerm{Float64}},
                             set_sd::Function)
    cols = Int32[]
    values = Float64[]
    sd_row  = Vector{Int32}[Int32[] for i in 1:length(m.sd_dim)]
    sd_col  = Vector{Int32}[Int32[] for i in 1:length(m.sd_dim)]
    sd_coef = Vector{Float64}[Float64[] for i in 1:length(m.sd_dim)]
	function add(col::ColumnIndex, coefficient::Float64)
		push!(cols, col.value)
		push!(values, coefficient)
	end
	function add(mat::MatrixIndex, coefficient::Float64)
        coef = mat.row == mat.column ? coefficient : coefficient / 2
        push!(sd_row[mat.matrix], mat.row)
        push!(sd_col[mat.matrix], mat.column)
        push!(sd_coef[mat.matrix], coef)
    end
	for term in terms
        add(mosek_index(m, term.variable_index), term.coefficient)
    end
    for j in 1:length(m.sd_dim)
        if !isempty(sd_row[j])
            id = appendsparsesymmat(m.task, m.sd_dim[j], sd_row[j],
                                    sd_col[j], sd_coef[j])
            set_sd(j, [id], [1.0])
        end
    end
    return cols, values
end

function set_row(task::Mosek.MSKtask, row::Int32, cols::ColumnIndices,
                 values::Vector{Float64})
    putarow(task, row, cols.values, values)
end
function set_row(m::MosekModel, row::Int32,
                 terms::Vector{MOI.ScalarAffineTerm{Float64}})
    cols, values = split_scalar_matrix(m, terms,
        (j, ids, coefs) -> putbaraij(m.task, row, j, ids, coefs))
    set_row(m.task, row, ColumnIndices(cols), values)
end

function set_coefficients(task::Mosek.MSKtask, rows::Vector{Int32},
                          cols::ColumnIndices, values::Vector{Float64})
    putaijlist(task, rows, cols.values, values)
end

function set_coefficients(task::Mosek.MSKtask, rows::Vector{Int32},
                          col::ColumnIndex, values::Vector{Float64})
    n = length(rows)
    @assert n == length(values)
    set_coefficient(task, rows, ColumnIndices(fill(col.value, n)), values)
end
function set_coefficients(m::MosekModel, rows::Vector{Int32},
                          vi::MOI.VariableIndex, values::Vector{Float64})
    set_coefficient(m.task, rows, mosek_index(m, vi), values)
end

function set_coefficient(task::Mosek.MSKtask, row::Int32, col::ColumnIndex,
                         value::Float64)
    putaij(task, row, col.value, value)
end
function set_coefficient(m::MosekModel, row::Int32, vi::MOI.VariableIndex,
                         value::Float64)
    set_coefficient(m.task, row, mosek_index(m, vi), value)
end

bound_key(::Type{MOI.GreaterThan{Float64}}) = MSK_BK_LO
bound_key(::Type{MOI.LessThan{Float64}})    = MSK_BK_UP
bound_key(::Type{MOI.EqualTo{Float64}})     = MSK_BK_FX
bound_key(::Type{MOI.Interval{Float64}})    = MSK_BK_RA
add_bound(m::MosekModel, row::Int32, dom::MOI.GreaterThan{Float64}) = putconbound(m.task, row, bound_key(typeof(dom)), dom.lower, dom.lower)
add_bound(m::MosekModel, row::Int32, dom::MOI.LessThan{Float64})    = putconbound(m.task, row, bound_key(typeof(dom)), dom.upper, dom.upper)
add_bound(m::MosekModel, row::Int32, dom::MOI.EqualTo{Float64})     = putconbound(m.task, row, bound_key(typeof(dom)), dom.value, dom.value)
add_bound(m::MosekModel, row::Int32, dom::MOI.Interval{Float64})    = putconbound(m.task, row, bound_key(typeof(dom)), dom.lower, dom.upper)
function bounds_to_set(::Type{S}, bk, bl, bu) where S
    if S == MOI.GreaterThan{Float64}
        return S(bl)
    elseif S == MOI.LessThan{Float64}
        return S(bu)
    elseif S == MOI.EqualTo{Float64}
        @assert bl == bu
        return S(bu)
    else
        @assert S == MOI.Interval{Float64}
        return S(bl, bu)
    end
end
function get_bound(m::MosekModel,
                   ci::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}, S}) where {S}
    bounds_to_set(S, getconbound(m.task, row(m, ci))...)
end

## Variable Constraints #######################################################
####################### lx ≤ x ≤ u ############################################
#######################      x ∈ K ############################################

function getvarboundlist(t::Mosek.Task, subj::Vector{Int32})
    n = length(subj)
    bk = Vector{Boundkey}(undef,n)
    bl = Vector{Float64}(undef,n)
    bu = Vector{Float64}(undef,n)
    for i in 1:n
        bki,bli,bui = getvarbound(t, subj[i])
        bk[i] = bki
        bl[i] = bli
        bu[i] = bui
    end
    bk,bl,bu
end

function delete_variable_constraint(m::MosekModel, col::ColumnIndex,
                                    ::Type{<:Union{MOI.Interval, MOI.EqualTo}})
    putvarbound(m.task, col.value, MSK_BK_FR, 0.0, 0.0)
end
function delete_variable_constraint(m::MosekModel, col::ColumnIndex,
                                    ::Type{MOI.Integer})
    putvartype(m.task, col.value, MSK_VAR_TYPE_CONT)
end
function delete_variable_constraint(m::MosekModel, col::ColumnIndex,
                                    ::Type{MOI.ZeroOne})
    putvartype(m.task, col.value, MSK_VAR_TYPE_CONT)
    putvarbound(m.task, col.value, MSK_BK_FR, 0.0, 0.0)
end
function delete_variable_constraint(m::MosekModel, col::ColumnIndex,
                                    ::Type{MOI.LessThan{Float64}})
    bk, lo, up = getvarbound(m.task, col.value)
    if bk == MSK_BK_UP
        bk = MSK_BK_FR
    else
        @assert bk == MSK_BK_RA
        bk = MSK_BK_LO
    end
    putvarbound(m.task, col.value, bk, lo, 0.0)
end
function delete_variable_constraint(m::MosekModel, col::ColumnIndex,
                                    ::Type{MOI.GreaterThan{Float64}})
    bk, lo, up = getvarbound(m.task, col.value)
    if bk == MSK_BK_LO
        bk = MSK_BK_FR
    else
        @assert bk == MSK_BK_RA
        bk = MSK_BK_UP
    end
    putvarbound(m.task, col.value, bk, 0.0, up)
end
function add_variable_constraint(m::MosekModel, col::ColumnIndex, dom::MOI.Interval)
    putvarbound(m.task, col.value, MSK_BK_RA, dom.lower, dom.upper)
end
function add_variable_constraint(m::MosekModel, col::ColumnIndex, dom::MOI.EqualTo)
    putvarbound(m.task, col.value, MSK_BK_FX, dom.value, dom.value)
end
function add_variable_constraint(m::MosekModel, col::ColumnIndex, ::MOI.Integer)
    putvartype(m.task, col.value, MSK_VAR_TYPE_INT)
end
function add_variable_constraint(m::MosekModel, col::ColumnIndex, ::MOI.ZeroOne)
    putvartype(m.task, col.value, MSK_VAR_TYPE_INT)
    putvarbound(m.task, col.value, MSK_BK_RA, 0.0, 1.0)
end
function add_variable_constraint(m::MosekModel, col::ColumnIndex, dom::MOI.LessThan)
    bk, lo, up = getvarbound(m.task, col.value)
    if bk == MSK_BK_FR
        bk = MSK_BK_UP
    else
        @assert bk == MSK_BK_LO
        bk = MSK_BK_RA
    end
    putvarbound(m.task, col.value, bk, lo, dom.upper)
end
function add_variable_constraint(m::MosekModel, col::ColumnIndex,
                                 dom::MOI.GreaterThan)
    bk, lo, up = getvarbound(m.task, col.value)
    if bk == MSK_BK_FR
        bk = MSK_BK_LO
    else
        @assert bk == MSK_BK_UP
        bk = MSK_BK_RA
    end
    putvarbound(m.task, col.value, bk, dom.lower, up)
end
function get_variable_constraint(m::MosekModel,
                                 col::ColumnIndex,
                                 ci::MOI.ConstraintIndex{MOI.SingleVariable, S}) where S
    return bounds_to_set(S, getvarbound(m.task, col.value)...)
end
function get_variable_constraint(m::MosekModel, vi::MOI.VariableIndex,
                                 ci::MOI.ConstraintIndex)
    return get_variable_constraint(m, mosek_index(m, vi), ci)
end

cone_parameter(dom :: MOI.PowerCone{Float64})     = dom.exponent
cone_parameter(dom :: MOI.DualPowerCone{Float64}) = dom.exponent
cone_parameter(dom :: C) where C <: MOI.AbstractSet = 0.0

function add_cone(m::MosekModel, cols::ColumnIndices, set)
    appendcone(m.task, cone_type(typeof(set)), cone_parameter(set), cols.values)
    id = getnumcone(m.task)
    if DEBUG
        putconename(m.task, id, "$id")
    end
    return id
end

## Name #######################################################################
###############################################################################

function set_row_name(task::Mosek.MSKtask, row::Int32, name::String)
    putconname(task, row, name)
end

function set_row_name(m::MosekModel,
                      c::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}},
                      name::AbstractString)
    set_row_name(m.task, row(m, c), name)
end
function set_row_name(m::MosekModel, c::MOI.ConstraintIndex,
                      name::AbstractString)
    # Fallback for `SingleVariable` and `VectorOfVariables`.
    m.con_to_name[c] = name
end

function delete_name(m::MosekModel, ci::MOI.ConstraintIndex)
    name = MOI.get(m, MOI.ConstraintName(), ci)
    if !isempty(name)
        cis = m.constrnames[name]
        deleteat!(cis, findfirst(isequal(ci), cis))
    end
end

###############################################################################
# INDEXING ####################################################################
###############################################################################

function row(m::MosekModel,
             c::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}})::Int32
    return getindex(m.c_block, c.value)
end
function columns(m::MosekModel, ci::MOI.ConstraintIndex{MOI.VectorOfVariables})
    return ColumnIndices(getcone(m.task, ci.value)[4])
end

const VectorCone = Union{MOI.SecondOrderCone,
                         MOI.RotatedSecondOrderCone,
                         MOI.PowerCone,
                         MOI.DualPowerCone,
                         MOI.ExponentialCone,
                         MOI.DualExponentialCone}

# Two `SingleVariable`-in-`S` cannot be set to the same variable if
# the two constraints
# * both set a lower bound, or
# * both set an upper bound, or
# * both set it to integer.
# The `incompatible_mask` are computed according to these rules.
flag(::Type{MOI.EqualTo{Float64}}) = 0x1
incompatible_mask(::Type{MOI.EqualTo{Float64}}) = 0x2f
flag(::Type{MOI.GreaterThan{Float64}}) = 0x2
incompatible_mask(::Type{MOI.GreaterThan{Float64}}) = 0x2b
flag(::Type{MOI.LessThan{Float64}}) = 0x4
incompatible_mask(::Type{MOI.LessThan{Float64}}) = 0x2d
flag(::Type{MOI.Interval{Float64}}) = 0x8
incompatible_mask(::Type{MOI.Interval{Float64}}) = 0x2f
flag(::Type{MOI.Integer}) = 0x10
incompatible_mask(::Type{MOI.Integer}) = 0x30
flag(::Type{MOI.ZeroOne}) = 0x20
incompatible_mask(::Type{MOI.ZeroOne}) = 0x3f
flag(::Type{<:VectorCone}) = 0x40
incompatible_mask(::Type{<:VectorCone}) = 0x40

function set_flag(model::MosekModel, vi::MOI.VariableIndex, S::Type)
    model.x_constraints[vi.value] |= flag(S)
end
function unset_flag(model::MosekModel, vi::MOI.VariableIndex, S::Type)
    model.x_constraints[vi.value] &= ~flag(S)
end
function has_flag(model::MosekModel, vi::MOI.VariableIndex, S::Type)
    return !iszero(model.x_constraints[vi.value] & flag(S))
end

###############################################################################
# MOI #########################################################################
###############################################################################

const ScalarLinearDomain = Union{MOI.LessThan{Float64},
                                 MOI.GreaterThan{Float64},
                                 MOI.EqualTo{Float64},
                                 MOI.Interval{Float64}}
const ScalarIntegerDomain = Union{MOI.ZeroOne, MOI.Integer}

## Add ########################################################################
###############################################################################

MOI.supports_constraint(m::MosekModel, ::Type{<:Union{MOI.SingleVariable, MOI.ScalarAffineFunction}}, ::Type{<:ScalarLinearDomain}) = true
MOI.supports_constraint(m::MosekModel, ::Type{MOI.VectorOfVariables}, ::Type{<:Union{VectorCone, MOI.PositiveSemidefiniteConeTriangle}}) = true
MOI.supports_constraint(m::MosekModel, ::Type{MOI.SingleVariable}, ::Type{<:ScalarIntegerDomain}) = true

## Affine Constraints #########################################################
##################### lc ≤ Ax ≤ uc ############################################

function MOI.add_constraint(m  ::MosekModel,
                            axb::MOI.ScalarAffineFunction{Float64},
                            dom::D) where {D <: MOI.AbstractScalarSet}

    if !iszero(axb.constant)
        throw(MOI.ScalarFunctionConstantNotZero{Float64, typeof(axb), D}(axb.constant))
    end

    # Duplicate indices not supported
    axb = MOIU.canonical(axb)

    N = 1
    conid = allocateconstraints(m, N)
    ci = MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}, D}(conid)
    r = row(m, ci)
    set_row(m, r, axb.terms)

    add_bound(m, r, dom)

    return ci
end

## Variable Constraints #######################################################
####################### lx ≤ x ≤ u ############################################
#######################      x ∈ K ############################################

# We allow following. Each variable can have
# - at most most upper and one lower bound
# - belong to at most one non-semidefinite cone
# - any number of semidefinite cones, which are implemented as ordinary constraints
# This is when things get a bit funky; By default a variable has no
# bounds, i.e. "free". Adding a `GreaterThan`
# constraint causes it to have a defined lower bound but no upper
# bound, allowing a `LessThan` constraint to be
# added later. Adding a `Interval` constraint defines both upper and
# lower bounds.

cone_type(::Type{MOI.ExponentialCone})        = MSK_CT_PEXP
cone_type(::Type{MOI.DualExponentialCone})    = MSK_CT_DEXP
cone_type(::Type{MOI.PowerCone{Float64}})     = MSK_CT_PPOW
cone_type(::Type{MOI.DualPowerCone{Float64}}) = MSK_CT_DPOW
cone_type(::Type{MOI.SecondOrderCone})        = MSK_CT_QUAD
cone_type(::Type{MOI.RotatedSecondOrderCone}) = MSK_CT_RQUAD

function MOI.add_constraint(
    m   :: MosekModel,
    xs  :: MOI.SingleVariable,
    dom :: D) where {D <: MOI.AbstractScalarSet}

    msk_idx = mosek_index(m, xs.variable)
    if !(msk_idx isa ColumnIndex)
        error("Cannot add $D constraint on a matrix variable")
    end

    if !iszero(incompatible_mask(D) & m.x_constraints[xs.variable.value])
        error("Cannot put multiple bound sets of the same type on a variable")
    end

    set_flag(m, xs.variable, D)

    add_variable_constraint(m, msk_idx, dom)

    return MOI.ConstraintIndex{MOI.SingleVariable, D}(xs.variable.value)
end

function MOI.add_constraint(m::MosekModel, xs::MOI.VectorOfVariables,
                            dom::D) where {D<:VectorCone}
    if any(vi -> is_matrix(m, vi), xs.variables)
        error("Cannot add $D constraint on a matrix variable")
    end
    cols = ColumnIndices(reorder(columns(m, xs.variables).values, D))

    if !all(vi -> iszero(incompatible_mask(D) & m.x_constraints[vi.value]), xs.variables)
        error("Cannot multiple bound sets of the same type to a variable")
    end

    N = MOI.dimension(dom)
    id = add_cone(m, cols, dom)

    return MOI.ConstraintIndex{MOI.VectorOfVariables, D}(id)
end

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

function MOI.add_constraint(m  ::MosekModel,
                            xs ::MOI.VectorOfVariables,
                            dom::MOI.PositiveSemidefiniteConeTriangle)
    N = dom.side_dimension
    if N < 2
        error("Invalid dimension for semidefinite constraint, got $N which is ",
              "smaller than the minimum dimension 2.")
    end
    appendbarvars(m.task, [Int32(N)])
    push!(m.sd_dim, N)
    id = length(m.sd_dim)
    sd_row  = Dict{Int, Vector{Int32}}()
    sd_col  = Dict{Int, Vector{Int32}}()
    sd_coef = Dict{Int, Vector{Float64}}()
    obj_row = Int32[]
    obj_col = Int32[]
    obj_coef = Float64[]
    k = 0
    for i in 1:N
        for j in 1:i
            k += 1
            vi = xs.variables[k]
            m.x_sd[vi.value] = MatrixIndex(id, i, j)

            if variable_type(m, vi) == MatrixVariable
                error("Variable $vi cannot be part of two matrix variables.")
            elseif variable_type(m, vi) == ScalarVariable
                # The variable has already been used, we need replace the
                # coefficients in the `A` matrix by matrix coefficients
                nnz, rows, vals = getacol(m.task, column(m, vi).value)
                @assert nnz == length(rows) == length(vals)
                for ii in 1:nnz
                    if !haskey(sd_row, rows[ii])
                        sd_row[rows[ii]] = Int32[]
                        sd_col[rows[ii]] = Int32[]
                        sd_coef[rows[ii]] = Float64[]
                    end
                    push!(sd_row[rows[ii]], i)
                    push!(sd_col[rows[ii]], j)
                    push!(sd_coef[rows[ii]], i == j ? vals[ii] : vals[ii] / 2)
                end
                cj = getcj(m.task, column(m, vi).value)
                if !iszero(cj)
                    push!(obj_row, i)
                    push!(obj_col, j)
                    push!(obj_coef, i == j ? cj : cj / 2)
                end
                MOI.delete(m, vi)
            end
            m.x_type[vi.value] = MatrixVariable
        end
    end
    for row in eachindex(sd_row)
        sid = appendsparsesymmat(m.task, N, sd_row[row], sd_col[row],
                                sd_coef[row])
        putbaraij(m.task, row, id, [sid], [1.0])
    end
    if !isempty(obj_row)
        sid = appendsparsesymmat(m.task, N, obj_row, obj_col, obj_coef)
        putbarcj(m.task, id, [sid], [1.0])
    end
    @assert k == length(xs.variables)

    conref = MOI.ConstraintIndex{MOI.VectorOfVariables, MOI.PositiveSemidefiniteConeTriangle}(id)

    return conref
end

## Get ########################################################################
###############################################################################

_variable(ci::MOI.ConstraintIndex{MOI.SingleVariable}) = MOI.VariableIndex(ci.value)
function MOI.get(m::MosekModel, ::MOI.ConstraintFunction,
                 ci::MOI.ConstraintIndex{MOI.SingleVariable}) where S <: ScalarLinearDomain
    return MOI.SingleVariable(_variable(ci))
end
function MOI.get(m::MosekModel, ::MOI.ConstraintSet,
                 ci::MOI.ConstraintIndex{MOI.SingleVariable, S}) where S <: ScalarIntegerDomain
    return S()
end
function MOI.get(m::MosekModel, ::MOI.ConstraintSet,
                 ci::MOI.ConstraintIndex{MOI.SingleVariable, S}) where S <: ScalarLinearDomain
    sv = MOI.get(m, MOI.ConstraintFunction(), ci)
    return get_variable_constraint(m, sv.variable, ci)
end

function MOI.get(m::MosekModel, ::MOI.ConstraintFunction,
                 ci::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64},
                                         <:ScalarLinearDomain})
    nnz, cols, vals = getarow(m.task, row(m, ci))
    @assert nnz == length(cols) == length(vals)
    terms = MOI.ScalarAffineTerm{Float64}[
        MOI.ScalarAffineTerm(vals[i], index_of_column(m, cols[i])) for i in 1:nnz]
    # TODO add matrix terms
    return MOI.ScalarAffineFunction(terms, 0.0)
end
function MOI.get(m::MosekModel, ::MOI.ConstraintSet,
                 ci::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64},
                                         S}) where S
    return get_bound(m, ci)
end

## Modify #####################################################################
###############################################################################

### SET
chgbound(bl::Float64,bu::Float64,k::Float64,dom :: MOI.LessThan{Float64})    = bl,dom.upper-k
chgbound(bl::Float64,bu::Float64,k::Float64,dom :: MOI.GreaterThan{Float64}) = dom.lower-k,bu
chgbound(bl::Float64,bu::Float64,k::Float64,dom :: MOI.EqualTo{Float64})     = dom.value-k,dom.value-k
chgbound(bl::Float64,bu::Float64,k::Float64,dom :: MOI.Interval{Float64})    = dom.lower-k,dom.upper-k

function MOI.set(m::MosekModel,
                 ::MOI.ConstraintSet,
                 ci::MOI.ConstraintIndex{F,D},
                 dom::D) where {F<:MOI.SingleVariable,
                                D<:ScalarLinearDomain}
    col = column(m, _variable(ci))
    bk, bl, bu = getvarbound(m.task, col.value)
    bl, bu = chgbound(bl, bu, 0.0, dom)
    putvarbound(m.task, col.value, bk, bl, bu)
end

function MOI.set(m::MosekModel,
                 ::MOI.ConstraintSet,
                 cref::MOI.ConstraintIndex{F,D},
                 dom::D) where { F    <: MOI.ScalarAffineFunction,
                                 D    <: ScalarLinearDomain }
    cid = ref2id(cref)
    i = getindex(m.c_block,cid) # since we are in a scalar domain
    bk, bl, bu = getconbound(m.task,i)
    bl, bu = chgbound(bl,bu,0.0,dom)
    putconbound(m.task,i,bk,bl,bu)
end

### MODIFY
function MOI.modify(m   ::MosekModel,
                    c   ::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64},D},
                    func::MOI.ScalarConstantChange{Float64}) where {D <: MOI.AbstractSet}
    if !iszero(func.new_constant)
        throw(MOI.ScalarFunctionConstantNotZero{Float64, MOI.ScalarAffineFunction{Float64}, D}(func.new_constant))
    end
end

function MOI.modify(m   ::MosekModel,
                    c   ::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}},
                    func::MOI.ScalarCoefficientChange{Float64})
    set_coefficient(m, row(m, c), func.variable, func.new_coefficient)
end

### TRANSFORM
function MOI.transform(m::MosekModel,
                       cref::MOI.ConstraintIndex{F,D},
                       newdom::D) where {F <: MOI.AbstractFunction,
                                         D <: MOI.AbstractSet}
    MOI.modify(m,cref,newdom)
    return cref
end

function MOI.transform(m::MosekModel,
                       cref::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64},D1},
                       newdom::D2) where {D1 <: ScalarLinearDomain,
                                          D2 <: ScalarLinearDomain}
    F = MOI.ScalarAffineFunction{Float64}

    cid = ref2id(cref)

    r = row(m, cref)
    add_bound(m, r, newdom)

    newcref = MOI.ConstraintIndex{F, D2}(cid)
    return newcref
end

## Delete #####################################################################
###############################################################################

function MOI.is_valid(model::MosekModel,
                      ci::MOI.ConstraintIndex{<:MOI.ScalarAffineFunction{Float64},
                                              S}) where S<:ScalarLinearDomain
    return allocated(model.c_block, ci.value) && getconbound(model.task, row(model, ci))[1] == bound_key(S)
end
function MOI.delete(
    m::MosekModel,
    cref::MOI.ConstraintIndex{F,D}) where {F <: MOI.ScalarAffineFunction{Float64},
                                           D <: ScalarLinearDomain}
    if !MOI.is_valid(m, cref)
        throw(MOI.InvalidIndex(cref))
    end

    delete_name(m, cref)

    subi = getindexes(m.c_block, cref.value)

    n = length(subi)
    subi_i32 = convert(Vector{Int32}, subi)
    ptr = fill(Int64(0), n)
    putarowlist(m.task,subi_i32,ptr,ptr,Int32[],Float64[])
    b = fill(0.0,n)
    putconboundlist(m.task,subi_i32,fill(MSK_BK_FX,n),b,b)

    deleteblock(m.c_block, cref.value)
end

function MOI.is_valid(model::MosekModel,
                      ci::MOI.ConstraintIndex{MOI.SingleVariable,
                                              S}) where S<:Union{ScalarLinearDomain,
                                                                 ScalarIntegerDomain}
    return allocated(model.x_block, ci.value) && has_flag(model, _variable(ci), S)
end
function MOI.delete(
    m::MosekModel,
    ci::MOI.ConstraintIndex{MOI.SingleVariable, S}) where S<:Union{ScalarLinearDomain,
                                                                   ScalarIntegerDomain}
    if !MOI.is_valid(m, ci)
        throw(MOI.InvalidIndex(ci))
    end
    delete_name(m, ci)
    vi = _variable(ci)
    unset_flag(m, vi, S)
    delete_variable_constraint(m, column(m, vi), S)
end
function MOI.is_valid(model::MosekModel,
                      ci::MOI.ConstraintIndex{MOI.VectorOfVariables,
                                              S}) where S<:VectorCone
    return 1 ≤ ci.value ≤ getnumcone(model.task) &&
        getconeinfo(model.task, ci.value)[1] == cone_type(S)
end
function MOI.is_valid(model::MosekModel,
                      ci::MOI.ConstraintIndex{MOI.VectorOfVariables,
                                              MOI.PositiveSemidefiniteConeTriangle})
    # TODO add supports for deletion
    return 1 ≤ ci.value ≤ length(model.sd_dim)
end

## List #######################################################################
###############################################################################
function MOI.get(m::MosekModel, ::MOI.ListOfConstraintAttributesSet)
    set = MOI.AbstractConstraintAttribute[]
    if !isempty(m.constrnames)
        push!(set, MOI.ConstraintName())
    end
    # TODO add VariablePrimalStart when get is implemented on it
    return set
end

## Name #######################################################################
###############################################################################
function MOI.supports(::MosekModel, ::MOI.ConstraintName,
                      ::Type{<:MOI.ConstraintIndex})
    return true
end
function MOI.set(m::MosekModel, ::MOI.ConstraintName, ci::MOI.ConstraintIndex,
                 name ::AbstractString)
    delete_name(m, ci)
    if !haskey(m.constrnames, name)
        m.constrnames[name] = MOI.ConstraintIndex[]
    end
    push!(m.constrnames[name], ci)
    set_row_name(m, ci, name)
end
function MOI.get(m::MosekModel, ::MOI.ConstraintName,
                 ci::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}})
    # All rows should have same name so we take the first one
    return getconname(m.task, getindexes(m.c_block, ref2id(ci))[1])
end
function MOI.get(m::MosekModel, ::MOI.ConstraintName, ci::MOI.ConstraintIndex)
    return get(m.con_to_name, ci, "")
end
function MOI.get(m::MosekModel, CI::Type{<:MOI.ConstraintIndex}, name::String)
    #asgn, row = getconnameindex(m.task, name)
    #if iszero(asgn)
    #    return nothing
    #else
    #    # TODO how to recover function and constraint type ?
    #end
    if !haskey(m.constrnames, name)
        return nothing
    end
    cis = m.constrnames[name]
    if isempty(cis)
        return nothing
    end
    if !isone(length(cis))
        error("Multiple constraints have the name $name.")
    end
    ci = first(cis)
    if !(ci isa CI)
        return nothing
    end
    return ci
end


text:
import Optim: retract!, project_tangent!

# * Sphere

metric_norm(x, g::UniformScaling) = norm(x)*√(g.λ)
metric_norm(x, g) = √(dot(x,g,x))

metric_normalize!(x, ::UniformScaling{Bool}) = normalize!(x)
function metric_normalize!(x, g)
    x ./= metric_norm(x,g)
end

"""Spherical manifold {|x| = 1}."""
struct MetricSphere{M} <: Manifold
    "Metric used for inner products, i.e. `⟨a|b⟩_g = gₘₙaᵐbⁿ"
    g::M
end
MetricSphere() = MetricSphere(I)
Optim.retract!(S::MetricSphere, x) = metric_normalize!(x, S.g)
Optim.project_tangent!(S::MetricSphere,g,x) = (g .-= real(dot(x, S.g, g)).*x)

# * Stiefel

matfun(f::Function, ee) = ee.vectors*Diagonal(f.(ee.values))*ee.vectors'

# For orthogonal bases, e.g. FE-DVR
löwdin_transform(g::UniformScaling{Bool}) = g,g
# For quasi-orthogonal bases, e.g. finite-differences
löwdin_transform(g::UniformScaling) = inv(√(g.λ))*I,√(g.λ)*I
löwdin_transform(g::Diagonal) = inv(√(g)),√(g)
# For non-orthogonal bases, e.g. B-splines
function löwdin_transform(g::AbstractMatrix)
    ee = eigen(g)
    matfun(λ -> inv(√(λ)), ee), matfun(λ -> √(λ), ee)
end
function löwdin_transform(g::BandedMatrix)
    m,n = size(g)
    m == n || throw(DimensionMismatch("Löwdin transform only valid for square matrices"))
    l,u = bandwidths(g)
    if l == u == 0
        löwdin_transform(Diagonal(g))
    else
        löwdin_transform(Symmetric(g))
    end
end

löwdin_transform!(v, ::UniformScaling{Bool}) = v
löwdin_transform!(v, S⁻ᴴ) = (v .= S⁻ᴴ*v)

abstract type MetricStiefel <: Stiefel end
struct MetricStiefel_CholQR{L} <: MetricStiefel
    "Löwdin transform to orthogonal basis"
    S⁻ᴴ::L
    "Löwdin transform to non-orthogonal basis"
    Sᴴ::L
end
struct MetricStiefel_SVD{L} <: MetricStiefel
    "Löwdin transform to orthogonal basis"
    S⁻ᴴ::L
    "Löwdin transform to non-orthogonal basis"
    Sᴴ::L
end
function MetricStiefel(g, retraction=:SVD)
    S⁻ᴴ,Sᴴ = löwdin_transform(g)
    if retraction == :CholQR
        MetricStiefel_CholQR(S⁻ᴴ,Sᴴ)
    elseif retraction == :SVD
        MetricStiefel_SVD(S⁻ᴴ,Sᴴ)
    else
        throw(ArgumentError("Unknown retraction $retraction"))
    end
end
function retract!(St::MetricStiefel_SVD, X)
    löwdin_transform!(X, St.Sᴴ)
    U,S,V = svd(X)
    X .= U*V'
    löwdin_transform!(X, St.S⁻ᴴ)
end
function retract!(St::MetricStiefel_CholQR, X)
    löwdin_transform!(X, St.Sᴴ)
    overlap = X'X
    X .= X/cholesky(overlap).L
    löwdin_transform!(X, St.S⁻ᴴ)
end
# For functions depending only on the subspace spanned by X, we always
# have G = A*X for some A, and so X'G = G'X, and Stiefel == Grassmann
# Edelman et al. have G .-= X*G'X (2.53), corresponding to a different
# metric ("canonical metric"). We follow Absil et al. here and use the
# metric inherited from Nxn matrices.
function project_tangent!(St::MetricStiefel, G, X)
    löwdin_transform!(X, St.Sᴴ)
    löwdin_transform!(G, St.Sᴴ)
    XG = X'G
    G .-= X*((XG .+ XG')./2)
    löwdin_transform!(X, St.S⁻ᴴ)
    löwdin_transform!(G, St.S⁻ᴴ)
end

# * Sphere & Stiefel combo

struct SphereStiefelCombo{Sp,St,M,F} <: Manifold
    "Indices of columns that only need normalization"
    normal_indidices::Vector{Int}
    "Manifold used to normalize columns"
    sphere::Sp
    "Sets of indices of columns that should be mutually orthogonalized"
    ortho_indidices::Vector{Vector{Int}}
    "Manifold used to orthogonalize sets of columns"
    stiefel::St
    "Temporary storage for orthgonalization"
    ortho_X::M
    "Temporary storage for orthgonalization"
    ortho_G::M
    f::F
end

function copytotmp!(dst, src, ortho_indices)
    for (i,j) in enumerate(ortho_indices)
        copyto!(view(dst, :, i), view(src, :, j))
    end
end

function copyfromtmp!(dst, src, ortho_indices)
    for (i,j) in enumerate(ortho_indices)
        copyto!(view(dst, :, j), view(src, :, i))
    end
end

function retract!(ssc::SphereStiefelCombo, X)
    retract!(ssc.sphere, view(X, :, ssc.normal_indidices))
    for o in ssc.ortho_indidices
        copytotmp!(ssc.ortho_X, X, o)
        retract!(ssc.stiefel, view(ssc.ortho_X, :, 1:length(o)))
        copyfromtmp!(X, ssc.ortho_X, o)
    end
    X
end

function project_tangent!(ssc::SphereStiefelCombo, G, X)
    project_tangent!(ssc.sphere, view(G, :, ssc.normal_indidices), view(X, :, ssc.normal_indidices))
    for o in ssc.ortho_indidices
        sel = 1:length(o)
        copytotmp!(ssc.ortho_G, G, o)
        copytotmp!(ssc.ortho_X, X, o)
        project_tangent!(ssc.stiefel, view(ssc.ortho_G, :, sel), view(ssc.ortho_X, :, sel))
        copyfromtmp!(G, ssc.ortho_G, o)
        copyfromtmp!(X, ssc.ortho_X, o)
    end
    G
end

# * Manifold setup

function setup_manifold(f::FockProblem)
    fock = f.fock
    P = orbitals(fock.quantum_system)
    m,n = size(P)

    normal_indices = sort(reduce(vcat, filter(s -> length(s)==1, fock.symmetries), init=Int[]))
    ortho_indices = filter(s -> length(s)≠1, fock.symmetries)
    n_ortho = length.(ortho_indices)

    length(normal_indices) + sum(n_ortho) == n ||
        throw(DimensionMismatch("Number of orbitals does not agree with amount of symmetry specified"))

    sphere = Optim.PowerManifold(MetricSphere(fock.S), (m,), (length(normal_indices),))
    stiefel = MetricStiefel(fock.S)

    nn = isempty(n_ortho) ? 0 : maximum(n_ortho)

    T = eltype(P)
    ortho_X = Matrix{T}(undef, m, nn)
    ortho_G = Matrix{T}(undef, m, nn)

    SphereStiefelCombo(normal_indices, sphere,
                       ortho_indices, stiefel,
                       ortho_X, ortho_G, f)
end


text:
using Test, SpecialFunctions, Statistics
import TruncatedNormal as TruncNorm

@test isnan(TruncNorm.tnvar(1, 0))
@test isnan(TruncNorm.tnvar(+Inf, 0))
@test isnan(TruncNorm.tnvar(0, -Inf))
@test isnan(TruncNorm.tnvar(NaN, NaN))
@test isnan(TruncNorm.tnvar(0, NaN))
@test isnan(TruncNorm.tnvar(NaN, 0))

@test TruncNorm.tnvar(-Inf, +Inf) == 1
@test TruncNorm.tnvar(+Inf, +Inf) == 0
@test TruncNorm.tnvar(-Inf, -Inf) == 0
@test TruncNorm.tnvar(0, +Inf) ≈ 1 - 2/π
@test TruncNorm.tnvar(-Inf, 0) ≈ 1 - 2/π

for x = -10:10
    @test TruncNorm.tnvar(x, +Inf) ≈ 1 + √(2/π) * x / erfcx(+x / √2) - 2/π / erfcx(+x / √2)^2
    @test TruncNorm.tnvar(-Inf, x) ≈ 1 - √(2/π) * x / erfcx(-x / √2) - 2/π / erfcx(-x / √2)^2
end

@test TruncNorm.tnvar(50, 70) ≈ 0.0003990431868038995479099272265360593305365

@test TruncNorm.tnvar(-Inf, Inf, 0, 1) == 1

for x = -10:10
    @test TruncNorm.tnvar(x, +Inf, 0, 1) ≈ TruncNorm.tnvar(x, +Inf)
    @test TruncNorm.tnvar(-Inf, x, 0, 1) ≈ TruncNorm.tnvar(-Inf, x)
end

@test TruncNorm.tnvar(50, 70, -2, 3) ≈ 0.029373438107168350377591231295634273607812172191712
@test TruncNorm.tnvar(-100.0, 0.0, 0.0, 2.0986317998643735) ≈ 1.6004193412141677189841357987638847137391508803335
@test TruncNorm.tnvar(0.0, 0.9, 0.0, 0.07132755843183151) ≈ 0.0018487407287725028827020557707636415445504260892486
@test TruncNorm.tnvar(-100.0, 100.0, 0.0, 17.185261847875548) ≈ 295.333163899557735486302841237124507431445
@test TruncNorm.tnvar(-100.0, 0.5, 0.0, 0.47383322897860064) ≈ 0.145041095812679283837328561547251019229612
@test TruncNorm.tnvar(-100.0, 100.0, 0.0, 17.185261847875548) ≈ 295.333163899557735486302841237124507431445
@test TruncNorm.tnvar(-10000, 10000, 0, 1) ≈ 1

# https://github.com/JuliaStats/Distributions.jl/issues/827
@test TruncNorm.tnvar(999000, 1e6) ≥ 0
@test TruncNorm.tnvar(0, 1000, 1000000, 1) ≥ 0

@test_broken TruncNorm.tnvar(1e6, Inf) ≈ 9.99999999994000000000050000000e-13
@test_broken TruncNorm.tnvar(999000, 1e6) ≈ 1.00200300399898194688784897455e-12
@test_broken TruncNorm.tnvar(0, 1000, 1000000, 1) ≈ 1.00200300399898194688784897455909681215272616880064850605775930845141e-12


text:
module HTgnss
using Debugger
using Geodesy
using Dates
using TimesDates
using LeapSeconds
using Parameters
using Printf
using HTrg

include("GNSS/Systems.jl")
include("Time/Systems.jl")
include("Position/Systems.jl")

# println(GNSSBase.ObsID)
# ObsID = GNSSBase.ObsID
include("RINEX/Rinex.jl")

greet() = print("Hello World!")

end # module