using Latexify
using ParameterizedFunctions
using DiffEqBiological
using Test

f = @ode_def TestAlignFeedback begin
    dx = y/c_1 - x
    dy = x^c_2 - y
end c_1 c_2

g = @ode_def TestAlign2 begin
    dx = p_x - d_x*x - d_x_y*y*x
    dy = p_y - d_y*y
    dz = p_z
end p_x d_x d_x_y p_y d_y p_z

@test latexify(f) ==
raw"\begin{align}
\frac{dx}{dt} =& \frac{y}{c_{1}} - x \\
\frac{dy}{dt} =& x^{c_{2}} - y
\end{align}
"

@test latexify(f; env=:align, field=:symjac) ==
raw"\begin{align}
\frac{dx}{dt} =& -1 \cdot 1 =& c_{1}^{-1} \\
\frac{dy}{dt} =& x^{-1 + c_{2}} \cdot c_{2} =& -1 \cdot 1
\end{align}
"

@test latexalign(["d$(x)/dt" for x = f.syms], f.funcs) ==
raw"\begin{align}
\frac{dx}{dt} =& \frac{y}{c_{1}} - x \\
\frac{dy}{dt} =& x^{c_{2}} - y
\end{align}
"

@test latexalign([f,g]) == "\\begin{align}\n\\frac{dx}{dt}  &=  \\frac{y}{c_{1}} - x  &  \\frac{dx}{dt}  &=  p_{x} - d_{x} \\cdot x - d_{x\\_y} \\cdot y \\cdot x  &  \\\\\n\\frac{dy}{dt}  &=  x^{c_{2}} - y  &  \\frac{dy}{dt}  &=  p_{y} - d_{y} \\cdot y  &  \\\\\n  &    &  \\frac{dz}{dt}  &=  p_{z} \\cdot 1  & \n\\end{align}\n"

@reaction_func hill2(x, v, k) = v*x^2/(k^2 + x^2)

rn = @reaction_network TestAlignMyRnType begin
  hill2(y, v_x, k_x), 0 --> x
  p_y, 0 --> y
  (d_x, d_y), (x, y) --> 0
  (r_b, r_u), x ↔ y
end v_x k_x p_y d_x d_y r_b r_u

@test latexalign(rn; clean=true).s ==
raw"\begin{align}
\frac{dx}{dt} =& \frac{v_{x} \cdot y^{2}}{k_{x}^{2} + y^{2}} - d_{x} \cdot x - r_{b} \cdot x + r_{u} \cdot y \\
\frac{dy}{dt} =& p_{y} - d_{y} \cdot y + r_{b} \cdot x - r_{u} \cdot y
\end{align}
"

@test latexify(((1.0, 2), (3, 4)); env=:align) ==
raw"\begin{align}
1.0 =& 3 \\
2 =& 4
\end{align}
"

# @test_throws MethodError latexify(rn; bad_kwarg="should error")


using Test
using QCMaterial

@testset "util.topylist" begin
    org_array = [1, 2, 3.0]
    pylist = topylist(org_array)
    @test all(org_array .== pylist)
end

@testset "util.doublefunc" begin
    org_array = [1, 2, 3]
    @test all(doublefunc(org_array) .== 2 .* org_array)
end

@testset "util.numerical_grad" begin
    # Some test for numerical_grad
    f(x) = x[1] + 2*x[2]
    deriv = numerical_grad(f, zeros(2))
    @assert deriv ≈ [1.0, 2.0]
end


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:light
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.4'
#       jupytext_version: 1.2.4
#   kernelspec:
#     display_name: Julia 1.1.0
#     language: julia
#     name: julia-1.1
# ---

# # This notebook
#
# _Notes:_ For documentation see <https://gaelforget.github.io/MeshArrays.jl/stable/>, <https://docs.juliadiffeq.org/latest/solvers/ode_solve.html> and <https://en.wikipedia.org/wiki/Displacement_(vector)>

# ## 1. import software

using IndividualDisplacements, MeshArrays, DifferentialEquations, Plots, Statistics
p=dirname(pathof(IndividualDisplacements))
include(joinpath(p,"PlotIndDisp.jl"))

# ## 2. Read gridded variables as `MeshArray`s

# Put grid variables in a dictionary.

# +
mygrid=gcmgrid("llc90_latlon/","ll",1,[(360,178)], [360 178], Float32, read, write)

GridVariables=Dict("XC" => read(mygrid.path*"XC.latlon.data",MeshArray(mygrid,Float32)),
"YC" => read(mygrid.path*"YC.latlon.data",MeshArray(mygrid,Float32)),
"XG" => read(mygrid.path*"XG.data",MeshArray(mygrid,Float32)),
"YG" => read(mygrid.path*"YG.data",MeshArray(mygrid,Float32)),
"DXC" => read(mygrid.path*"DXC.latlon.data",MeshArray(mygrid,Float32)),
"DYC" => read(mygrid.path*"DYC.latlon.data",MeshArray(mygrid,Float32)) );
# -

# Read velocity fields as `MeshArray`s.

# +
using MAT
file = matopen(mygrid.path*"uv_lonlat.mat")
u=read(file, "u")
v=read(file, "v")
close(file)

u=dropdims(mean(u,dims=3),dims=3)
v=dropdims(mean(v,dims=3),dims=3)

u=read(u,MeshArray(mygrid,Float32))
v=read(v,MeshArray(mygrid,Float32));

u[findall(isnan.(u))]=0.0
v[findall(isnan.(v))]=0.0

u0=u; u1=u;
v0=v; v1=v;
# -

# Put velocity fields and time range in a dictionary.

# +
t0=0.0; t1=86400*366*10.0; dt=3600;

u0=u0./GridVariables["DXC"]
u1=u1./GridVariables["DXC"]
v0=v0./GridVariables["DYC"]
v1=v1./GridVariables["DYC"]

uvt = Dict("u0" => u0, "u1" => u1, "v0" => v0, "v1" => v1, "t0" => t0, "t1" => t1, "dt" => dt) ;
# -

# Merge the two dictionaries and add masks

# +
uvetc=merge(uvt,GridVariables);

nr=50; kk=1;

mskW=read(mygrid.path*"hFacW.latlon.data",MeshArray(mygrid,Float32,nr))
mskW=1.0 .+ 0.0 * mask(mskW[:,kk],NaN,0.0)
mskS=read(mygrid.path*"hFacS.latlon.data",MeshArray(mygrid,Float32,nr))
mskS=1.0 .+ 0.0 * mask(mskS[:,kk],NaN,0.0)

msk=Dict("mskW" => mskW, "mskS" => mskS)

uvetc=merge(uvetc,msk);
# -

# ## 3. Visualize  gridded variables

#PyPlot.contourf(GridVariables["XG"].f[1], GridVariables["YC"].f[1], mskS.f[1].*v0.f[1])
#colorbar()
heatmap(v0.f[1])
#heatmap(mskS.f[1].*v0.f[1])
#PyPlot.contour(GridVariables["XG"].f[1], GridVariables["YC"].f[1], mskS.f[1].*v0.f[1])

# ## 4. Recompute displacements from gridded flow fields

# Initialize individual locations and define method aliases.

# +
uInit=[180.0,40.0] #uInit=[tmp[1,:lon];tmp[1,:lat]]./uvetc["dx"]
nSteps=Int32(uvt["t1"]/uvt["dt"]) #nSteps=Int32(tmp[end,:time]/3600)-2
du=fill(0.0,2);

comp_vel=IndividualDisplacements.VelComp
get_vel=IndividualDisplacements.VelCopy
# -

# Inspect how `comp_vel` behaves.

# +
if false
    comp_vel(du,uInit,uvetc,0.0); show(du)
    tmpdu=[uvt["u0"][1][uInit[1]+1,uInit[2]+1] uvt["v0"][1][uInit[1]+1,uInit[2]+1]]; show(tmpdu)
end

ii=uInit[1]-3:0.1:uInit[1]+3
jj=uInit[2]-3:0.1:uInit[2]+3
tmpu=zeros(size(ii))
tmpv=zeros(size(ii))
for i in eachindex(ii)
    comp_vel(du,[ii[i];jj[i]],uvetc,0.0)
    tmpu[i],tmpv[i]=du
end
#tmp=zeros(10,1)
#for comp_vel(du,[180.1;40.1],uvetc,0.0)
Plots.plot(tmpu)
Plots.plot!(tmpv)
# -

# ## 5. Solve through time using `DifferentialEquations.jl`

using DifferentialEquations
tspan = (0.0,nSteps*3600.0)
#prob = ODEProblem(get_vel,uInit,tspan,tmp)
prob = ODEProblem(comp_vel,uInit,tspan,uvetc)
sol = solve(prob,Tsit5(),reltol=1e-8,abstol=1e-8)
sol[1:4]

?Tsit5

#Plots.plot(sol[1,:],sol[2,:])
sol[:,end-4:end]

Plots.plot(sol[1,2:end]-sol[1,1:end-1])
Plots.plot!(sol[2,2:end]-sol[2,1:end-1])

Plots.plot(sol[1,:],sol[2,:])


#Pkg.add("UnstructuredGrids")
using UnstructuredGrids

# Generate a grid from its nodal coordinates, cell connectivities,
# and cell types

# Define coordinates
coords = zeros(2,9)
coords[:,1] = [1,1]
coords[:,2] = [3,1]
coords[:,3] = [4,1]
coords[:,4] = [1,2]
coords[:,5] = [2,2]
coords[:,6] = [1,3]
coords[:,7] = [2,3]
coords[:,8] = [3,3]
coords[:,9] = [4,3]

# Define connectivity
connect = [1,2,5,4,2,3,9,4,5,7,6,5,8,7,5,2,8,2,9,8]
offsets = [1,      5,    8,      12,   15,   18,   21]

# Define cell types
using UnstructuredGrids.RefCellGallery: SQUARE, TRIANGLE
refcells = [SQUARE, TRIANGLE]
types = [1,2,1,2,2,2]

# Generate the UGrid object
grid = UGrid(connect,offsets,types,refcells,coords)

# Export grid into vtk format
#writevtk(grid,"grid")

# Generate a global numbering for the edges (1d objects) of the grid,
# and find which edges are on the boundary of each cell
n=1
cell_to_edges = generate_cell_to_faces(n,grid)

#@show cell_to_edges.list
#@show cell_to_edges.ptrs

cell = 3
a = cell_to_edges.ptrs[cell]
b = cell_to_edges.ptrs[cell+1]-1
edges = cell_to_edges.list[a:b]
#@show edges

# Generate a toy structured grid of the unit cube
# with 2x3x2 quadrilateral cells
grid = UGrid(domain=(0,1,0,1,0,1),partition=(2,3,2))

# Get the connectivity of the cells of the grid (i.e., for each cell, the ids of its vertices)
cell_to_vertices = connections(grid)

# The cell connectivity is represented by a `Connections` object
@assert isa(cell_to_vertices, Connections)

# `Connections` is a struct representing a vector of vectors
# in compressed form. It is fully described by a vector containing
# the flatted underlying data and another vector containing
# the start and end of each of the sub-vectors.
# That is:

cell_to_vertices_data = list(cell_to_vertices)
cell_to_vertices_ptrs = ptrs(cell_to_vertices)

cell = 2
ibeg = cell_to_vertices_ptrs[cell]
iend = cell_to_vertices_ptrs[cell+1]-1
vertices = cell_to_vertices_data[ibeg:iend] # Vertices of cell 2
@assert vertices == [2, 3, 5, 6, 14, 15, 17, 18]

# Get the coordinates of the mesh vertices
vertex_to_coords = coordinates(grid)
@assert isa(vertex_to_coords, Array{Float64,2})

# First axis for space dimensions, second one for number of vertices
@assert size(vertex_to_coords) == (3,36)

# Find cells around each vertex
vertex_to_cells = generate_dual_connections(cell_to_vertices)

# The returned data is a vector of vectors represented by a Connections object
@assert isa(vertex_to_cells, Connections)

# Generate a global numbering for the edges (1d objects) of the grid,
# and find which edges are on the boundary of each cell
n=1
cell_to_edges = generate_cell_to_faces(n,grid)

# The returned data again a Connections object
@assert isa(cell_to_edges, Connections)

# Idem for faces (2d objects)
n=2
cell_to_faces = generate_cell_to_faces(n,grid)
@assert isa(cell_to_faces, Connections)


"""
    interpret_odds(odds::Number; rules::Union{String, Rules}="chen2010", logodds::Bool=false)

Interpret odds ratio with rules of thumb.

The odds ratio (OR) is probably one of the most widely used index of effect size.

# Arguments
- `odds`: Odds ratio value.
- `rules`: Can be 'chen2010' or custom set of [Rules](@ref).
- `logodds`: Are these log odds ratio (*e.g.*, extracted from a logistic regression)?

# Rules of thumb
- **Chen et al., (2010)**: Thee authors demonstrate that OR = 1.68, 3.47, and 6.71 are equivalent to Cohen's d = 0.2 (small), 0.5 (medium), and 0.8 (large) and Cohen's d < 0.2 when OR < 1.5, and Cohen's d > 0.8 when OR > 5.

# Examples
```jldoctest
interpret_odds(2, rules="chen2010")

# output

"small"

# References
- Chen, H., Cohen, P., & Chen, S. (2010). How big is a big odds ratio? Interpreting the magnitudes of odds ratios in epidemiological studies. Communications in Statistics—Simulation and Computation, 39(4), 860-864.
```
"""
function interpret_odds(odds::Number; rules::Union{String, Rules}="chen2010", logodds::Bool=false)

    if logodds == true
        odds = exp(abs(odds))
    end

    if isa(rules, String)
        if rules == "chen2010"
            rules = Rules(
                [1.68, 3.47, 6.71],
                ["very small", "small", "medium", "large"],
                true
                )
        else
            throw(ArgumentError("`rules` must be a String ('chen2010') or custom Rules."))
        end
    end

    interpretation = interpret(abs(odds), rules)
    return interpretation
end


# ------------------------------------------------------------------------------------------
# # Representing data in a computer
#
# The core of data science and machine learning is **data**: we are interested in extracting
# knowledge from data.
#
# But how exactly do computers represent data? Let's find out exactly what an "artificial
# intelligence" has at its disposal to learn from.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## Data is represented as arrays
#
# Let's take a look at some fruit. Using the `Images.jl` library, we can load in some
# images:
# ------------------------------------------------------------------------------------------

using Images

apple = load("data/10_100.jpg")

banana = load("data/104_100.jpg")

# ------------------------------------------------------------------------------------------
# Here we have images of apples and bananas. We would eventually like to build a program
# that can automatically distinguish between the two. However, the computer doesn't "see" an
# apple or a banana; instead, it just sees numbers.
#
# An image is encoded in something called an **array**, which is like a container that has
# boxes or slots for individual pieces of data:
#
# ![attachment:array_cartoon.png](data/array_cartoon.png)
#
# An array is a bunch of numbers in connected boxes; the figure above shows a 1-dimensional
# array. Our images are instead 2-dimensional arrays, or matrices, of numbers, arranged
# something like this:
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ![attachment:array2d.png](data/array2d.png)
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# For example, `apple` is an image, consisting of a 100x100 array of numbers:
# ------------------------------------------------------------------------------------------

typeof(apple)

typeof(a)

a = [ 1 2 3;4 5 6]

size(a)

# ------------------------------------------------------------------------------------------
# We can grab the datum stored in the box at row `i` and column `j` by *indexing* using
# square brackets: `[i, j]`. For example, let's get the pixel (piece of the image) in box
# $(40, 60)$, i.e. in the 40th row and 60th column of the image:
# ------------------------------------------------------------------------------------------

apple

dump(typeof(apple[40, 60]))

apple[18:20,29:31]

# ------------------------------------------------------------------------------------------
# We see that Julia displays a coloured box! Julia, via the `Colors.jl` package, is clever
# enough to display colours in a way that is useful to us humans!
#
# So, in fact, an image is a 2D array, in which each element of the array is an object (a
# collection of numbers) describing a coloured pixel.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## Colors as numbers
#
# How, then, are these colors actually stored? Computers store colors in RGB format, that is
# they store a value between 0 and 1 for each of three "channels": red, green, and blue.
# Here, 0 means none of that color and 1 means the brightest form of that color. The overall
# color is a combination of those three colors.
#
# For example, we can pull out the `red` value using the function `red` applied to the
# color. Since internally the actual value is stored in a special format, we choose to
# convert it to a standard floating-point number using the `Float64` function:
# ------------------------------------------------------------------------------------------

Float64(red(apple[40, 60]))

[ mean(float.(c.(img))) for c = [red,green,blue], img = [apple,banana] ]

using Plots

histogram(float.(green.(apple[:])),color="red",label="apple", normalize=true, nbins=25)
histogram!(float.(green.(banana[:])),color="yellow",label="banana",normalize=true, nbins=25)

apple

float(red(banana[50,20]))

banana[50,20]

pixel = apple[40, 60]

red_value   = Float64( red(pixel) )
green_value = Float64( green(pixel) )
blue_value  = Float64( blue(pixel) )

print("The RGB values are ($red_value, $green_value, $blue_value)")

# ------------------------------------------------------------------------------------------
# Since the red value is high while the others are low, this means that at pixel `(40, 60)`,
# the picture of the apple is very red. If we do the same at one of the corners of the
# picture, we get the following:
# ------------------------------------------------------------------------------------------

pixel = apple[1, 1]

red_value   = Float64( red(pixel) )
green_value = Float64( green(pixel) )
blue_value  = Float64( blue(pixel) )

print("The RGB values are ($red_value, $green_value, $blue_value)")

apple

# ------------------------------------------------------------------------------------------
# We see that every color is bright, which corresponds to white.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## Working on an image as a whole
#
# In Julia, to apply a function to the whole of an array, we place a `.` between the
# function name and the left parenthesis (`(`), so the following gives us the `red` value of
# every pixel in the image:
# ------------------------------------------------------------------------------------------

redpartofapple = Float64.(red.(apple))
mean(redpartofapple)

using Plots

gr()

histogram(redpartofapple[:],color=:red,label="redness in the apple")

# ------------------------------------------------------------------------------------------
# Note that we get a 2D array (matrix) back.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# Julia's [mathematical standard
# library](https://docs.julialang.org/en/stable/stdlib/math/#Mathematics-1) has many
# mathematical functions built in. One of them is the `mean` function, which computes the
# average value. If we apply this to our apple:
# ------------------------------------------------------------------------------------------

mean(Float64.(red.(apple)))

# ------------------------------------------------------------------------------------------
# we see that the value indicates that the average amount of red in the image is a value
# between the amount of red in the apple and the amount of red in the white background.
#
# *Somehow we need to teach a computer to use this information about a picture to recognize
# that there's an apple there!*
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## A quick riddle
#
# Here's a quick riddle. Let's check the average value of red in the image of the banana.
# ------------------------------------------------------------------------------------------

mean(Float64.(red.(banana)))

# ------------------------------------------------------------------------------------------
# Oh no, that's more red than our apple? This isn't a mistake and is actually true! Before
# you move onto the next exercise, examine the images of the apple and the banana very
# carefully and see if you can explain why this is expected.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 1
#
# What is the average value of blue in the banana?
#
# (To open a new box use <ESC>+b (b is for "below", what do you think a does?))
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 2
#
# Does the banana have more blue or more green?
# ------------------------------------------------------------------------------------------

.nbexports/introductory-tutorials/broader-topics-and-ecosystem/intro-to-ml/01. ML - Representing data in a computer.jl


using AdventOfCodeSolutions
using Test

function input(puzzle::Puzzle{2020, 7, n}) where n
    io = openInput(puzzle)
    split(read(io, String), "\n", keepempty=false)
end

function parseInput(input)
    rules = map(input) do rule
        (bag, containRules) = split(rule, " bags contain ")
        contains = map(split(containRules, ", ")) do rule
            m = match(r"^(\d+) (.*) bags?.?$", rule)

            if (isnothing(m))
                return nothing
            end

            return (parse(Int, m.captures[1]), string(m.captures[2]))
        end

        return (
            bag = bag,
            contains = filter(!isnothing, contains)
        )
    end

    return map(rule -> rule.bag => rule.contains, rules) |> Dict
end

function canContainShinyGold(rules, bag)
    if (!haskey(rules, bag))
        return false
    end
    # @show rules[bag]

    if "shiny gold" in last.(rules[bag])
        return true
    end

    return any(bag -> canContainShinyGold(rules, last(bag)), rules[bag])
end

function solve(::Puzzle{2020, 7, 1}, input)
    rules = parseInput(input)

    return count(
        ((bag, contains),) -> canContainShinyGold(rules, bag),
        rules
    )
end

function amountOfBagsIn(rules, bag)
    bags = rules[bag]

    if (length(bags) == 0)
        return 0
    end

    map(
        ((amount, name),) -> amount * (1 + amountOfBagsIn(rules, name)),
        bags
    ) |> sum
end

function solve(::Puzzle{2020, 7, 2}, input)
    rules = parseInput(input)
    amountOfBagsIn(rules, "shiny gold")
end

@testset "Day 7" begin
TEST_INPUT = split("""
light red bags contain 1 bright white bag, 2 muted yellow bags.
dark orange bags contain 3 bright white bags, 4 muted yellow bags.
bright white bags contain 1 shiny gold bag.
muted yellow bags contain 2 shiny gold bags, 9 faded blue bags.
shiny gold bags contain 1 dark olive bag, 2 vibrant plum bags.
dark olive bags contain 3 faded blue bags, 4 dotted black bags.
vibrant plum bags contain 5 faded blue bags, 6 dotted black bags.
faded blue bags contain no other bags.
dotted black bags contain no other bags.
""", '\n', keepempty=false)
    @test(solve(Puzzle(7, 1), TEST_INPUT) == 4)
    @test(solve(Puzzle(7, 1), input(Puzzle(7, 1))) == 144)

TEST_INPUT_2 = split("""
shiny gold bags contain 2 dark red bags.
dark red bags contain 2 dark orange bags.
dark orange bags contain 2 dark yellow bags.
dark yellow bags contain 2 dark green bags.
dark green bags contain 2 dark blue bags.
dark blue bags contain 2 dark violet bags.
dark violet bags contain no other bags.
""", '\n', keepempty=false)
    @test(solve(Puzzle(7, 2), TEST_INPUT_2) == 126)
    @test(solve(Puzzle(7, 2), input(Puzzle(7, 2))) == 5956)
end


function average_loss(m::Model, loss::Loss, inputs::Vector, weights::Vector)
    l, n = 0.0, 0
    n_inputs = length(inputs)
    for i in 1:n_inputs
        @inbounds x2, t2, x1, t1, y = inputs[i]
        @inbounds w = weights[i]
        p = predict(m, x2, t2, x1, t1)
        l += w * loss_value(loss, p, y)
        n += w
    end
    return l / n
end


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

using JuMP
using NLopt
using LinearAlgebra

function opt_solve(A::Matrix{Float64}, b::Vector{Float64}, x0::Vector{Float64}, lb::Vector{Float64}, ub::Vector{Float64})

    #mle = Model(optimizer_with_attributes(Ipopt.Optimizer, "print_level" => 0))
    mle = Model(NLopt.Optimizer)
    set_optimizer_attribute(mle, "algorithm", :LD_SLSQP) #LD_SLSQP

    x0[x0 .< 0] .= 0
    x0[x0 .> 1] .= 1

    @variable(mle, lb[i] <= x[i=1:length(x0)] <= ub[i])
    #@constraint(mle, sum(x) == 1)
    for n in 1:length(x0) set_start_value(x[n], x0[n]) end

    #@NLexpression(mle, mse ,  sum( (b .- sum(x[i] .* A[:,i] for i in 1:length(x) )).^2    ) )
    @expression(mle, mse ,  sum( (b .- sum(x[i] .* A[:,i] for i in 1:length(x) )).^2    ) )

    @NLobjective(mle, Min, mse)

    JuMP.optimize!(mle)
    return value.(mle[:x]) , exp(objective_value(mle))
end

function dolsq(A, b; method::String="default")
    if method == "default"
        x = A \ b
    elseif method == "pinv"
        x = pinv(A)*b
    elseif method == "qr"
        #Q,R = qr(A)
        #x = inv(R)*Q'*b
        qrA = qr(A)
        x = qrA\b
    end
    return x
end

function bvls(A, b, x_lsq, lb, ub, tol::Float64, max_iter::Int64, verbose::Int64, inverse_method::String)
    n_iter = 0
    m, n = size(A)

    x = x_lsq
    on_bound = zeros(n)

    mask = x .<= lb
    x[mask] = lb[mask]
    on_bound[mask] .= -1

    mask = x .>= ub
    x[mask] = ub[mask]
    on_bound[mask] .= 1

    free_set = on_bound .== 0
    active_set = .!free_set
    free_set = (1:size(free_set)[1])[free_set .!= 0]

    r = A * x - b
    cost = 0.5 * dot(r , r)
    initial_cost = cost
    g = A' * r

    cost_change = nothing
    step_norm = nothing
    iteration = 0

    while size(free_set)[1] > 0
        if verbose == 2
            optimality = compute_kkt_optimality(g, on_bound)
        end

        iteration += 1
        x_free_old = x[free_set]

        A_free = A[:, free_set]
        b_free = b - A * (x .* active_set)
        z = dolsq(A_free, b_free, method=inverse_method)

        lbv = z .< lb[free_set]
        ubv = z .> ub[free_set]

        v = lbv .| ubv

        if any(lbv)
            ind = free_set[lbv]
            x[ind] = lb[ind]
            active_set[ind] .= true
            on_bound[ind] .= -1
        end

        if any(ubv)
            ind = free_set[ubv]
            x[ind] = ub[ind]
            active_set[ind] .= true
            on_bound[ind] .= 1
        end

        ind = free_set[.!v]
        x[ind] = z[.!v]

        r = A * x -b
        cost_new = 0.5 * dot(r , r)
        cost_change = cost - cost_new
        cost = cost_new
        g = A' * r
        step_norm = sum((x[free_set] .- x_free_old).^2)

        if any(v)
            free_set = free_set[.!v]
        else
            break
        end
    end

    if isnothing(max_iter)
        max_iter = n
    end
    max_iter += iteration

    termination_status = nothing

    optimality = compute_kkt_optimality(g, on_bound)
    for iteration in iteration:max_iter
        if optimality < tol
            termination_status = 1
        end

        if !isnothing(termination_status)
            break
        end

        move_to_free = argmax(g .* on_bound)
        on_bound[move_to_free] = 0

        x_free = copy(x)
        x_free_old = copy(x)
        while true

            free_set = on_bound .== 0
            sum(free_set)
            active_set = .!free_set
            free_set = (1:size(free_set)[1])[free_set .!= 0]

            x_free = x[free_set]
            x_free_old = copy(x_free)
            lb_free = lb[free_set]
            ub_free = ub[free_set]

            A_free = A[:, free_set]
            b_free = b - A * (x .* active_set)
            z = dolsq(A_free, b_free, method=inverse_method)

            lbv = (1:size(free_set)[1])[ z .< lb_free]
            ubv = (1:size(free_set)[1])[ z .> ub_free]
            v = cat(lbv, ubv, dims=1)

            if size(v)[1] > 0
                alphas = cat(lb_free[lbv] - x_free[lbv], ub_free[ubv] - x_free[ubv],dims=1) ./ (z[v] - x_free[v])

                i = argmin(alphas)
                i_free = v[i]
                alpha = alphas[i]

                x_free .*= (1 .- alpha)
                x_free .+= (alpha .* z)
                x[free_set] = x_free

                vsize = size(lbv)
                if i <= size(lbv)[1]
                    on_bound[free_set[i_free]] = -1
                else
                    on_bound[free_set[i_free]] = 1
                end
            else
                x_free = z
                x[free_set] = x_free
                @goto start
            end
        end #while
        @label start
        step_norm = sum((x_free .- x_free_old).^2)

        r = A * x - b
        cost_new = 0.5 * dot(r , r)
        cost_change = cost - cost_new

        combo = tol * cost
        if cost_change < tol * cost
            termination_status = 2
        end
        cost = cost_new

        g = A' * r
        optimality = compute_kkt_optimality(g, on_bound)
    end #iteration

    if isnothing(termination_status)
        termination_status = 0
    end

    x[x .< 1e-5] .= 0
    return x, cost
end

function compute_kkt_optimality(g::Vector{Float64}, on_bound::Vector)
  g_kkt = g .* on_bound
  free_set = on_bound .== 0
  g_kkt[free_set] = broadcast(abs, g[free_set])

  return maximum(g_kkt)

end


using NBConvolution
using Distributions
using Base.Test

# compare the two methods to verify accuracy
d1 = FurmanNegativeBinomialConvolution([4.0, 1.0], [0.01, 0.5], 1000)
d2 = NegativeBinomialConvolution([4, 1], [0.01, 0.5])
@test abs(pdf(d1, 10) - pdf(d2, 10)) < 1e-10

d1 = FurmanNegativeBinomialConvolution([4.0, 10.0], [0.5, 0.5], 1000)
d2 = NegativeBinomialConvolution([4, 10], [0.5, 0.5])
@test abs(pdf(d1, 30) - pdf(d2, 30)) < 1e-10

d1 = FurmanNegativeBinomialConvolution([4.0, 10.0, 3.0, 6.0], [0.4, 0.5, 0.3, 0.6], 1000)
d2 = NegativeBinomialConvolution([4, 10, 3, 6], [0.4, 0.5, 0.3, 0.6])
@test abs(pdf(d1, 30) - pdf(d2, 30)) < 1e-10

# cdf
d2 = NegativeBinomialConvolution([4, 10], [0.7, 0.5])
@test abs(cdf(d2, 1000) - 1.0) < 1e-10

# non-integer rs
d3 = NegativeBinomialConvolution([4.2, 10], [0.7, 0.5])
@test abs(cdf(d3, 1000) - 1.0) < 1e-10
@test abs(pdf(d2, 10) - pdf(d3, 10)) > 1e-10
d4 = FurmanNegativeBinomialConvolution([4.2, 10.0], [0.7, 0.5], 1000)
@test abs(pdf(d3, 4) - pdf(d4, 4)) < 1e-10

# integrate with the distributions package
modelInner = MixtureModel([NegativeBinomialConvolution([3, 4], [0.1, 0.5]), NegativeBinomial(5, 0.4)], [0.6, 0.4])
pdf(modelInner, 0)
pdf(modelInner, 10)


using Test, OffsetArrays
using VanFoFy.Ellipticals: Theta, theta

@testset "θ-functions" begin
    ω1 = complex(1.0)
    ω3 = exp(1im)
    tol_check = 1e-15
    tol_compute = 1e-20
    θ = Theta(ω1, ω3, tol_compute)

    z = 0.63256929128615824176 + 0.47122375149242207160im

    th = OffsetArray{ComplexF64}(undef, 1:4, 0:4)

    for n in 0:4
        for k in 1:4
            th[k, n] = theta(θ, th_k=k, d_n=n, z=z)
        end
    end

    #####

    th_ma = [ 0.46214180304899833678 + 0.65513824162092381551im,
               0.96467023166692638429 + 0.11945667130003175272im,
               1.13835201868809736642 + 0.08123099419582718546im,
               0.86149287144921489500 - 0.08153295492698743884im]

    @test th[:, 0] ≈ th_ma atol=tol_check

    #####

    d1_th_ma = [ 0.97727706225681589769 + 0.08049915546539055744im,
             -0.42045304182205492196 - 0.63824580596766239552im,
              0.14241824316243509356 - 0.38564309515712238239im,
             -0.14363258751912619862 + 0.38620176200689986572im]

    @test th[:, 1] ≈ d1_th_ma atol=tol_check

    #####

    d2_th_ma = [-0.54549882845031761232 - 0.68893198268462155795im,
             -0.93946558497506382639 - 0.19739257822094024267im,
             -0.55247742739635853599 - 0.32311219018650595886im,
              0.55495918519936412129 + 0.32794356188507020189im]

    @test th[:, 2] ≈ d2_th_ma atol=tol_check

    #####

    d3_th_ma = [-1.05294508933080125267 + 0.15318331373277377347im,
              0.17025898335751795341 + 0.53691780124847588798im,
             -0.5623867730174600491 + 1.5392204494203911042im,
              0.5818162827245192532 - 1.5481591190168449923im]

    @test th[:, 3] ≈ d3_th_ma atol=tol_check

    #####

    d4_th_ma = [1.2954660922120757234 + 0.9931820896285708616im,
             0.71273193831970297909 + 0.89906624330941167631im,
             2.1950195882908139481 + 1.2634597310003502748im,
            -2.2347277131390164682 - 1.3407616781773902616im]

    @test th[:, 4] ≈ d4_th_ma atol=tol_check

    #####

    th1, th2, dth1, dth2 = theta(θ, z)

    @test th1  ≈ theta(θ, th_k=1, d_n=0, z=z) atol=tol_check
    @test th2  ≈ theta(θ, th_k=2, d_n=0, z=z) atol=tol_check
    @test dth1 ≈ theta(θ, th_k=1, d_n=1, z=z) atol=tol_check
    @test dth2 ≈ theta(θ, th_k=2, d_n=1, z=z) atol=tol_check
end


function likelihoods!(L::AbstractMatrix, hmm::AbstractHMM{Univariate}, observations)
    T, K = size(observations, 1), size(hmm, 1)
    @argcheck size(L) == (T, K)
    @inbounds for i in OneTo(K), t in OneTo(T)
        L[t,i] = pdf(hmm.B[i], observations[t])
    end
end

function loglikelihoods!(LL::AbstractMatrix, hmm::AbstractHMM{Univariate}, observations)
    T, K = size(observations, 1), size(hmm, 1)
    @argcheck size(LL) == (T, K)
    @inbounds for i in OneTo(K), t in OneTo(T)
        LL[t,i] = logpdf(hmm.B[i], observations[t])
    end
end

function likelihoods!(L::AbstractMatrix, hmm::AbstractHMM{Multivariate}, observations)
    T, K = size(observations, 1), size(hmm, 1)
    @argcheck size(L) == (T, K)
    @inbounds for i in OneTo(K), t in OneTo(T)
        L[t,i] = pdf(hmm.B[i], view(observations, t, :))
    end
end

function loglikelihoods!(LL::AbstractMatrix, hmm::AbstractHMM{Multivariate}, observations)
    T, K = size(observations, 1), size(hmm, 1)
    @argcheck size(LL) == (T, K)
    @inbounds for i in OneTo(K), t in OneTo(T)
        LL[t,i] = logpdf(hmm.B[i], view(observations, t, :))
    end
end


@doc doc"""
    TranslationAction(
        M::Manifold,
        Rn::TranslationGroup,
        AD::ActionDirection = LeftAction(),
    )

Space of actions of the [`TranslationGroup`](@ref) $\mathrm{T}(n)$ on a Euclidean-like
manifold `M`.

The left and right actions are equivalent.
"""
struct TranslationAction{TM<:Manifold,TRn<:TranslationGroup,TAD<:ActionDirection} <:
       AbstractGroupAction{TAD}
    M::TM
    Rn::TRn
end

function TranslationAction(
    M::Manifold,
    Rn::TranslationGroup,
    ::TAD = LeftAction(),
) where {TAD<:ActionDirection}
    return TranslationAction{typeof(M),typeof(Rn),TAD}(M, Rn)
end

function show(io::IO, A::TranslationAction)
    print(io, "TranslationAction($(A.M), $(A.Rn), $(direction(A)))")
end

base_group(A::TranslationAction) = A.Rn

g_manifold(A::TranslationAction) = A.M

function switch_direction(A::TranslationAction{TM,TRN,TAD}) where {TM,TRN,TAD}
    return TranslationAction(A.M, A.Rn, switch_direction(TAD()))
end

apply(A::TranslationAction, a, x) = x + a

apply!(A::TranslationAction{M,G}, y, a, x) where {M,G} = (y .= x .+ a)
apply!(A::TranslationAction{M,G}, y, e::Identity{G}, x) where {M,G} = copyto!(y, x)

inverse_apply(A::TranslationAction, a, x) = x - a

inverse_apply!(A::TranslationAction{M,G}, y, a, x) where {M,G} = (y .= x .- a)
inverse_apply!(A::TranslationAction{M,G}, y, e::Identity{G}, x) where {M,G} = copyto!(y, x)

apply_diff(A::TranslationAction, a, x, v) = v

apply_diff!(A::TranslationAction, vout, a, x, v) = copyto!(vout, v)

inverse_apply_diff(A::TranslationAction, a, x, v) = v

inverse_apply_diff!(A::TranslationAction, vout, a, x, v) = copyto!(vout, v)


### Mean

@mkapprule Mean :nargs => 1
@doap Mean(x::ListT) = Statistics.mean(margs(x))


include("../../common/sparse/src/sparse_module.jl")
include("../../../std/util/src/util_module.jl")

module linalg
using ..sparse
using ..util
using LinearAlgebra

include("norm.jl")
include("triangular.jl")
include("ldl.jl")
include("lu.jl")
include("tridiagonal.jl")

end # module


using ReservoirComputing

#model parameters
const approx_res_size = 30
const radius = 1.2
const activation = tanh
const degree = 6
const sigma = 0.1
const beta = 0.0
const alpha = 1.0
const nla_type = NLADefault()
const in_size = 3
const out_size = 3
W_in = ReservoirComputing.init_input_layer(approx_res_size, in_size, sigma)
W = ReservoirComputing.init_reservoir_givendeg(approx_res_size, radius, degree)
const extended_states = false
const h_steps = 2

const train_len = 50
const predict_len = 12
data = ones(Float64, in_size, 100)
train = data[:, 1:1+train_len-1]
test = data[:, train_len:train_len+predict_len-1]

#constructor 1
esn = ESN(approx_res_size,
    train,
    degree,
    radius,
    activation = activation,
    sigma = sigma,
    alpha = alpha,
    nla_type = nla_type,
    extended_states = extended_states)

#test constructor
@test isequal(Integer(floor(approx_res_size/in_size)*in_size), esn.res_size)
@test isequal(train, esn.train_data)
@test isequal(alpha, esn.alpha)
@test isequal(activation, esn.activation)
@test isequal(nla_type, esn.nla_type)
@test size(esn.W) == (esn.res_size, esn.res_size)
@test size(esn.W_in) == (esn.res_size, esn.in_size)
@test size(esn.states) == (esn.res_size, train_len)

#test train
W_out = ESNtrain(esn, beta)
@test size(W_out) == (out_size, esn.res_size)
#test predict
output = ESNpredict(esn, predict_len, W_out)
@test size(output) == (out_size, predict_len)
#test h steps predict
output = ESNpredict_h_steps(esn, predict_len, h_steps, test, W_out)
@test size(output) == (out_size, predict_len)

#constructor 2
esn = ESN(W,
    train,
    activation = activation,
    sigma = sigma,
    alpha = alpha,
    nla_type = nla_type,
    extended_states = extended_states)

#test constructor
@test isequal(approx_res_size, esn.res_size)
@test isequal(train, esn.train_data)
@test isequal(alpha, esn.alpha)
@test isequal(activation, esn.activation)
@test isequal(nla_type, esn.nla_type)
@test size(esn.W) == (esn.res_size, esn.res_size)
@test size(esn.W_in) == (esn.res_size, esn.in_size)
@test size(esn.states) == (esn.res_size, train_len)

#test train
W_out = ESNtrain(esn, beta)
@test size(W_out) == (out_size, esn.res_size)
#test predict
output = ESNpredict(esn, predict_len, W_out)
@test size(output) == (out_size, predict_len)

#constructor 3
esn = ESN(approx_res_size,
    train,
    degree,
    radius,
    W_in,
    activation = activation,
    alpha = alpha,
    nla_type = nla_type,
    extended_states = extended_states)

#test constructor
@test isequal(Integer(floor(approx_res_size/in_size)*in_size), esn.res_size)
@test isequal(train, esn.train_data)
@test isequal(alpha, esn.alpha)
@test isequal(activation, esn.activation)
@test isequal(nla_type, esn.nla_type)
@test size(esn.W) == (esn.res_size, esn.res_size)
@test size(esn.W_in) == (esn.res_size, esn.in_size)
@test size(esn.states) == (esn.res_size, train_len)

#test train
W_out = ESNtrain(esn, beta)
@test size(W_out) == (out_size, esn.res_size)
#test predict
output = ESNpredict(esn, predict_len, W_out)
@test size(output) == (out_size, predict_len)

#constructor 4
esn = ESN(W,
    train,
    W_in,
    activation = activation,
    alpha = alpha,
    nla_type = nla_type,
    extended_states = extended_states)

#test constructor
@test isequal(approx_res_size, esn.res_size)
@test isequal(train, esn.train_data)
@test isequal(alpha, esn.alpha)
@test isequal(activation, esn.activation)
@test isequal(nla_type, esn.nla_type)
@test size(esn.W) == (esn.res_size, esn.res_size)
@test size(esn.W_in) == (esn.res_size, esn.in_size)
@test size(esn.states) == (esn.res_size, train_len)

#test train
W_out = ESNtrain(esn, beta)
@test size(W_out) == (out_size, esn.res_size)
#test predict
output = ESNpredict(esn, predict_len, W_out)
@test size(output) == (out_size, predict_len)


using NLOptControl, JuMP, Parameters, VehicleModels
include("plots.jl");include("autonomousControl.jl");
c=defineCase("c1");

# initialize
n1 = NLOpt();s1=Settings(;MPC=false,save=true,simulate=false,format=:png);

# define
pa1 = VparaKB();  @unpack_VparaKB pa1 # vehicle parameters
X0 = [x0_,y0_,psi0_,u0_]; #TODO pull these states in
XF = [NaN,NaN,NaN,NaN];
XL = [x_min,y_min,psi_min,u_min];
XU = [x_max,y_max,psi_max,u_max];
CL = [sa_min,ax_min];
CU = [sa_max,ax_max];
n1 = define(n1,stateEquations=KinematicBicycle,numStates=4,numControls=2,X0=X0,XF=XF,XL=XL,XU=XU,CL=CL,CU=CU)

# build
n1 = configure(n1,Ni=2,Nck=[15,10];(:integrationMethod => :ps),(:integrationScheme => :lgrExplicit),(:finalTimeDV => false),(:tf => 4.0))
defineSolver(n1,solver=:IPOPT)
mdl1 = build(n1);

# addtional information
names = [:x,:y,:psi,:ux];
descriptions = ["X (m)","Y (m)","Yaw Angle (rad)","Longitudinal Velocity (m/s)"];
stateNames(n1,names,descriptions)
names = [:sr,:jx];
descriptions = ["Steering Angle (rad)","Longitudinal Acceleration (m/s^2)"];
controlNames(n1,names,descriptions);

#mXL=Any[false,false,false,false];mXU=Any[false,false,false,-1];  # set to false if you don't want to taper that side
#linearStateTolerances(n;mXL=mXL,mXU=mXU);

# setup OCP
params1 = [pa1];   # vehicle parameters
n1,r1 = OCPdef(mdl1,n1,s1,params1);
x_ref = 200; y_ref = 100; # define target TODO get from other model
@NLobjective(mdl1, Min, (r1.x[end,1]-x_ref)^2 + (r1.x[end,2]-y_ref)^2);

# solve
optimize(mdl1,n1,r1,s1)

# post process
r1.results_dir = string(main_dir,"/results_1","/test_1_",n1.solver,"/")
resultsDir(r1.results_dir);
using PrettyPlots, Plots
pyplot();
allPlots(n1,r1,s1,1)

idx=1;
pp=statePlot(n1,r1,s1,idx,1,2;(:lims=>false));
pp=obstaclePlot(n1,r1,s1,c,idx,pp;(:append=>true)); # add obstacles
pp=vehiclePlot(n1,r1,s1,c,idx,pp;(:append=>true)); # add the vehicle


Ints = Union{Array{Int8, 1}, Array{Int16, 1}, Array{Int32, 1}, Array{Int64, 1}, Array{Int128, 1}, Array{UInt8, 1}, Array{UInt16, 1}, Array{UInt32, 1}, Array{UInt64, 1}, Array{UInt128, 1}}

function count_sort!(data::Ints, output::Ints, c::Array{Int32, 1}, bits::Int32, offset::Int32)
	combs = 2^bits

	for i in 1:combs
		c[i] = 0
	end

	for i in data
		c[(i >> offset) & (combs - 1) + 1] += 1
	end

	for i in 2:combs
		c[i] += c[i - 1]
	end

	index = 0

	for i in length(data) : -1 : 1
		index = (data[i] >> offset) & (combs - 1) + 1
		output[c[index]] = data[i]
		c[index] -= 1
	end
end

function count_sort_sign!(data::Ints, output::Ints, c::Array{Int32, 1})
	c[1] = 0
	c[2] = 0

	bit = sizeof(data[1]) * 8

	for i in data
		c[~(i >> (bit - 1)) & 1 + 1] += 1
	end

	c[2] += c[1]

	index = 0

	for i in length(data) : -1 : 1
		index = ~(data[i] >> (bit - 1)) & 1 + 1
		output[c[index]] = data[i]
		c[index] -= 1
	end
end

function radix_sort!(data::Ints, log_bits::Int=3)
	if length(data) == 0
		return
	end
	T = typeof(data[1])
	mirror = zeros(T, length(data))
	bits::Int32 = 2^log_bits
	count = zeros(Int32, 2^bits)
	for i in 1:fld(sizeof(T) * 8, bits)
		count_sort!(data, mirror, count, bits, Int32(((i - 1) * 2) * bits))
		count_sort!(mirror, data, count, bits, Int32(((i - 1) * 2 + 1) * bits))
	end
	if !(T <: Unsigned)
		count_sort_sign!(data, mirror, count)
		for i in 1:length(data)
			data[i] = mirror[i]
		end
	end
end

function radix_sort!(data::Array{Float16, 1}, log_bits::Int)
	intdata::Array{Int16, 1} = unsafe_wrap(Array, convert(Ptr{Int16}, pointer(data)), length(data))

	for i in 1:length(intdata)
		if intdata[i] < 0
			intdata[i] = ~intdata[i] + Int16(-32768)
		end
	end

	radix_sort!(intdata, log_bits)

	for i in 1:length(intdata)
		if intdata[i] < 0
			intdata[i] = ~intdata[i] + Int16(-32768)
		end
	end
end

function radix_sort!(data::Array{Float32, 1}, log_bits::Int)
	intdata::Array{Int32, 1} = unsafe_wrap(Array, convert(Ptr{Int32}, pointer(data)), length(data))

	for i in 1:length(intdata)
		if intdata[i] < 0
			intdata[i] = ~intdata[i] + Int32(-2147483648)
		end
	end

	radix_sort!(intdata, log_bits)

	for i in 1:length(intdata)
		if intdata[i] < 0
			intdata[i] = ~intdata[i] + Int32(-2147483648)
		end
	end
end

function radix_sort!(data::Array{Float64, 1}, log_bits::Int)
	intdata::Array{Int64, 1} = unsafe_wrap(Array, convert(Ptr{Int64}, pointer(data)), length(data))

	for i in 1:length(intdata)
		if intdata[i] < 0
			intdata[i] = ~intdata[i] + Int64(-9223372036854775808)
		end
	end

	radix_sort!(intdata, log_bits)

	for i in 1:length(intdata)
		if intdata[i] < 0
			intdata[i] = ~intdata[i] + Int64(-9223372036854775808)
		end
	end
end


# ## Heat

# ## Model

module heat #jl

using ReachabilityAnalysis, MAT, SparseArrays
using ReachabilityModels: @relpath

file = matopen(@relpath "heat.mat")

## system matrix
A = read(file, "A")

## input matrix
B = sparse([67], [1], [1.0], size(A, 1), 1)

## state domain
X = Universe(200)

## input domain
U = BallInf([0.0], 0.5)

function model(X0)
    S = @system(x' = Ax + Bu, x ∈ X, u ∈ U)
    return IVP(S, X0)
end

end  #jl

# ## References


first_3_of_6(v::SVector{6,T}) where {T} = SVector{3,T}(v[1], v[2], v[3])
last_3_of_6(v::SVector{6,T}) where {T} = SVector{3,T}(v[4], v[5], v[6])

function getTop(a::SMatrix{4,4,T,16}) where {T}
    return @SMatrix [
    a[1]  a[5]  a[9]  a[13];
    a[2]  a[6]  a[10] a[14];
    a[3]  a[7]  a[11] a[15]
    ]
end

unPad(a::SVector{4,T}) where {T} = SVector{3,T}(a[1], a[2], a[3])
unPad(a::SMatrix{1,4,T,4}) where {T} = SVector{3,T}(a[1], a[2], a[3])
onePad(a::SVector{3,T}) where {T} = SVector{4,T}(a[1], a[2], a[3], one(T))
zeroPad(a::SVector{3,T}) where {T} = SVector{4,T}(a[1], a[2], a[3], zero(T))

@inline mul_then_un_pad(m::SMatrix{4,4,T1,16}, v::SVector{4,T2}) where {T1,T2} = unPad(m * v)
@inline one_pad_then_mul(m::SMatrix{4,4,T1,16}, v::SVector{3,T2}) where {T1,T2} = m * onePad(v)

function weightPoly(p1::SVector{N,T}, p2::SVector{N,T}, w1::T, w2::T) where {N,T}
  sum_weight = w1 - w2
  c1 = w1 / sum_weight
  c2 = w2 / sum_weight
  return c1 * p2 - c2 * p1
end

function make_pd_gains(time_settle::Float64, damping::Float64=1.0)
    w_n = 2 * pi / time_settle
    p = w_n^2
    d = 2 * damping * w_n
    return p, d
end

"""
Creates a random positive definite SMatrix for a size.
"""
function rand_pd(n::Int64)
    USV = svd(randn(SMatrix{n,n,Float64,n^2}))
    return Hermitian(USV.U * Diagonal(SVector(USV.S)) * USV.U')
end


# Even though shapes are specified with real geometric parameters, the size of voxels can be
# complex when the underlying space is stretched by complex factors.  So, we allow complex
# `vxl` below.  Even in such a case, r₀ and nout are still real.  (Maybe this condition
# should be relaxed later.)

export volfrac

const X, Y, Z = 1, 2, 3
const XYZ, YZX, ZXY = (X,Y,Z), (Y,Z,X), (Z,X,Y)
const UVW = YZX, ZXY, XYZ
const N, P = 1, 2  # negative, positive
const NP = (N, P)

corner(vxl::NTuple{2,SVector{3,<:Number}}, sx::Integer, sy::Integer, sz::Integer) =
    SVector(vxl[sx][X], vxl[sy][Y], vxl[sz][Z])

# Calculate the bit array that indicates corner contained-ness.
#
# The bit array is 8 bits (= 1 byte).  The kth bit is 1 if the kth corner is contained in
# the half space defined by the plane, an 0 otherwise.  The 8 corners of the voxel are
# indexed 1 through 8 in the order of (---), (+--), (-+-), (++-), (--+), (+-+), (-++), (+++),
# where, e.g., (+--) indicates the corner at the +x, -y, -z corner.
function corner_bits(vxl::NTuple{2,SVector{3,<:Number}},  # two ends of solid diagonal of voxel
                     nout::SVector{3,<:Real}, # unit outward normal of plane
                     nr₀::Real)  # equation of plane: nout⋅(r - r₀) = 0, or nout⋅r = nr₀
    cbits = 0x00
    bit = 0x01
    n_on = 0
    for sz = NP, sy = NP, sx = NP
        r = corner(vxl, sx, sy, sz)
        nr = nout⋅r
        if nr ≤ nr₀  # corner is contained and not on boundary (nout is outward normal)
            cbits |= bit
            n_on += nr==nr₀  # corner is on boundary
        end
        bit <<= 1
    end

    return cbits, n_on
end

# Determine if the corner contained-ness bit array corresponds to the case where the plane
# crosses one of the three sets of four parallel edges of the voxel.
function isquadsect(cbits::UInt8)
    return cbits==0x0F || cbits==~0x0F || cbits==0x33 || cbits==~0x33 || cbits==0x55 || cbits==~0x55

    # Equivalent to
    # n = count_ones(cbits)
    # m = count_ones(cbits & 0x0F)
    # return n==4 && iseven(m)
end

# For the cases where the plane crosses a set of four parallel edges of the voxel, determine
# which direction those edges lie.
function edgedir_quadsect(cbits::UInt8)
    if cbits==0x0F || cbits==~0x0F
        dir = Z
    elseif cbits==0x33 || cbits==~0x33
        dir = Y
    else
        @assert cbits==0x55 || cbits==~0x55
        dir = X
    end

    return dir
end

# Return the volume fraction when the plane croses a set of four parallel edges.
function rvol_quadsect(vxl::NTuple{2,SVector{3,<:Number}}, nout::SVector{3,<:Real}, nr₀, cbits::UInt8)
    w = edgedir_quadsect(cbits)
    ∆w = vxl[P][w] - vxl[N][w]

    u, v, _w = UVW[w]
    nu, nv, nw = nout[u], nout[v], nout[w]
    mean_cepts = 4nr₀
    for sv = NP, su = NP
        mean_cepts -= nu*vxl[su][u] + nv*vxl[sv][v]
    end
    mean_cepts /=  nw * 4∆w

    sw = nw>0 ? N : P  # nw cannot be 0
    return abs(mean_cepts - vxl[sw][w]/∆w)
end

# Calculate the volume fraction for the most general cases, by cutting out corners from a
# triangular pyramid.
#
# Assume count_ones(cbits) ≤ 4.  Othewise, call this function with flipped nout, nr₀,
# cbits.
function rvol_gensect(vxl::NTuple{2,SVector{3,<:Number}}, nout::SVector{3,<:Real}, nr₀::Real, cbits::UInt8)
    s = (nout.<0) .+ 1
    c = corner(vxl, s[X], s[Y], s[Z])  # corner coordinates
    ∆ = vxl[P] - vxl[N]  # vxl edges
    nc = nout .* c
    rmax, rmid, rmin =  abs.(((nr₀-sum(nc)) .+ nc) ./ nout - c) ./ ∆ # (lengths from corner to intercetps) / (voxel edges)

    # nx, ny, nz = nout
    # sx, sy, sz = ((nx≥0 ? N : P), (ny≥0 ? N : P), (nz≥0 ? N : P))  # signs of corner
    # cx, cy, cz = vxl[nX][sx], vxl[nY][sy], vxl[nZ][sz]  # corner coordinates
    # ∆x, ∆y, ∆z = vxl[nX][nP]-vxl[nX][nN], vxl[nY][nP]-vxl[nY][nN], vxl[nZ][nP]-vxl[nZ][nN]  # voxel edges
    # nxcx, nycy, nzcz = nx*cx, ny*cy, nz*cz
    # rmax, rmid, rmin =  # (lengths from corner to intercetps) / (voxel edges)
    #     abs((nr₀-nycy-nzcz)/nx-cx)/∆x, abs((nr₀-nzcz-nxcx)/ny-cy)/∆y, abs((nr₀-nxcx-nycy)/nz-cz)/∆z

    # Sort rmax, rmin, rmin properly.
    if rmax < rmid; rmax, rmid = rmid, rmax; end
    if rmid < rmin; rmid, rmin = rmin, rmid; end
    if rmax < rmid; rmax, rmid = rmid, rmax; end

    # Calculate the volume of the triangular pyramid, and cut off appropriate corners.
    tmax = 1 - 1/rmax
    rvol_core = 1 + tmax + tmax^2

    # Below, if rmax == Inf, rmid and rmin must be ≤ 1 in exact arithmetic and subtraction
    # must not occur, but they can be > 1 in floating-point arithmetic and lead to large
    # subtraction.  Prevent this by subtracting only for rmax ≠ Inf.
    isfin_rmax = !isinf(rmax)
    if rmid > 1 && isfin_rmax
        tmid = 1 - 1/rmid
        rvol_core -= rmax * tmid^3
    end
    if rmin > 1 && isfin_rmax
        tmin = 1 - 1/rmin
        rvol_core -= rmax * tmin^3
    end

    return rvol_core * rmin * rmid / 6
end

"""
    volfrac(voxel, nout, r₀)

Returns the volume fraction `rvol = vol(voxel ⋂ half-space) / vol(voxel)` that indicates how
much portion of the volume of a voxel is included the half-space defined by a plane.  The
result is between 0.0 and 1.0, inclusive.  Inside and outside of the half space is
determined by `nout`, which is the outward normal.  (If a point is on the `nout` side of the
plane, then it is outside.)

The voxel needs to be aligned with the Cartesian axes, and is described by `v =
([xn,yn,zn], [xp,yp,zp])` that specifies the x-, y-, z-boundaries of the voxel.  The
half-space is described by the boundary plane.  The boundary plane is described by the
outward normal vector `nout = [nx, ny, nz]` and a point `r₀ = [rx, ry, rz]` on the plane.
`nout` does not have to be normalized.
"""
function volfrac(vxl::NTuple{2,SVector{3,<:Number}}, nout::SVector{3,<:Real}, r₀::SVector{3,<:Real})
    nr₀ = nout⋅r₀
    cbits, n_on = corner_bits(vxl, nout, nr₀)
    n_in = count_ones(cbits)  # number of corners contained

    if n_in == 8  # voxel is inside half-space
        rvol = 1.
    elseif n_in - n_on == 0  # voxel is outside half-space
        rvol = 0.
    elseif isquadsect(cbits) # plane crosses a set of four parallel edges of voxel
        rvol = rvol_quadsect(vxl, nout, nr₀, cbits)
    elseif n_in ≤ 4 # general cases with n_in ≤ 4
        rvol = rvol_gensect(vxl, nout, nr₀, cbits)
    else  # general cases with n_in ≥ 5
        @assert n_in ≥ 5
        rvol = 1. - rvol_gensect(vxl, .-nout, -nr₀, ~cbits)
    end

    return rvol
end

# volfrac() for a 2D voxel (= pixel).
# - Turn the nout and r₀ into 3D vectors on the xy-plane.
# - Turn the pixel into a 3D voxel whose base is the pixel and the height is 1.
# - Pass these 3D objects to volfrac() for 3D.
volfrac(vxl::NTuple{2,SVector{2,<:Number}}, nout::SVector{2,<:Real}, r₀::SVector{2,<:Real}) =
    volfrac((SVector(vxl[N][1], vxl[N][2], 0), SVector(vxl[P][1], vxl[P][2], 1)),
            SVector(nout[1], nout[2], 0),
            SVector(r₀[1], r₀[2], 0))

# volfrac() for 1D voxel (= line segment).
# - Turn the nout and r₀ into 3D vectors along the z-axis
# - Turn the line segment into a 3D voxel whose height is the line segment and the base is a 1×1 pixel.
# - Pass these 3D objects to volfrac() for 3D.
#
# volfrac() for 1D can be implemented from scratch without calling volfrac() for 3D, but
# let's keep it this way for now.
volfrac(vxl::NTuple{2,SVector{1,<:Number}}, nout::SVector{1,<:Real}, r₀::SVector{1,<:Real}) =
    volfrac((SVector(0, 0, vxl[N][1]), SVector(1, 1, vxl[P][1])),
            SVector(0, 0, nout[1]),
            SVector(0, 0, r₀[1]))


using PtFEM, Test

data = Dict(
  # Plane(ndim, nst, nxe, nye, nip, direction, finite_element(nod, nodof), axisymmetric)
  :struc_el => Plane(2, 3, 3, 2, 4, :x, Quadrilateral(8, 2), false),
  :properties => [1.0e6 0.3;],
  :x_coords => [0.0, 3.0, 6.0],
  :y_coords => [0.0, -3.0, -6.0, -9.0],
  :support => [
    ( 1, [0 1]),
    ( 6, [0 1]),
    ( 9, [0 1]),
    (14, [0 1]),
    (17, [0 1]),
    (22, [0 1]),
    (25, [0 0]),
    (26, [0 0]),
    (27, [0 0]),
    (28, [0 0]),
    ( 5, [0 1]),
    ( 8, [0 1]),
    (13, [0 1]),
    (16, [0 1]),
    (21, [0 1]),
    (24, [0 1]),
    (29, [0 0])
    ],
  :loaded_nodes => [
    (1, [0.0 -0.5]),
    (2, [0.0 -2.0]),
    (3, [0.0 -0.5])
    ]
)

@time m, dis_df, fm_df = p51(data);

@test m.loads[0:36] ≈ [0.0,-5.310672021656476e-6,-4.211153422117119e-7,
  -5.04123153829146e-6,-7.221882550664876e-7,-3.3428571428571473e-6,
  -4.211153422117125e-7,-1.6444827474228391e-6,-1.3750422640578195e-6,
  -4.288483723850779e-6,3.774202388609178e-7,-2.7857142857142914e-6,
  -1.2829448475778024e-6,-3.2431230529276257e-6,2.7081477113201545e-7,
  -2.8731660139934707e-6,3.6699957275470236e-7,-2.2285714285714337e-6,
  2.708147711320155e-7,-1.5839768431493952e-6,-1.2140198042152427e-6,
  -2.217378390101892e-6,2.996313701739242e-7,-1.6714285714285747e-6,
  -1.125478752755258e-6,-1.3787735959184822e-6,1.3704853665438675e-7,
  -1.298707004971293e-6,1.912334941440869e-7,-1.1142857142857151e-6,
  1.3704853665438667e-7,-9.298644236001403e-7,-8.497978326529501e-7,
  -6.453529188458675e-7,1.1219944448620244e-7,-5.571428571428575e-7,
  -4.68932795439848e-7]


function epsilon_dominates_clear(u, v, epsilon)
  veps = v .+ epsilon
  all(u .<= veps) && any(u .< veps)
end

function epsilon_dominates_fast(u, v, epsilon)
  pareto_dominates_fast(u, v .+ epsilon)
end

function pareto_dominates_clear(u, v)
  all(u .<= v) && any(u .< v)
end

function pareto_dominates_fast(u, v)
  one_smaller = false
  for i in 1:length(u)
    if u[i] > v[i]
      return false
    elseif u[i] < v[i]
      one_smaller = true
    end
  end
  return one_smaller
end

# Returns two booleans where the first indicates e-box dominance and the 2nd indicates
# epsilon progress
function epsilon_box_dominates_and_epsilon_progress(u, v, epsilon)
  uindex = floor(u ./ epsilon)
  vindex = floor(v ./ epsilon)
  inner_epsilon_box_dominates_and_epsilon_progress(
    u, uindex, (epsilon .* uindex),
    v, vindex, (epsilon .* vindex))
end

epsilon_box_dominates(u, v, epsilon) = epsilon_box_dominates_and_epsilon_progress(u, v, epsilon)[1]
epsilon_box_progress(u, v, epsilon) = epsilon_box_dominates_and_epsilon_progress(u, v, epsilon)[2]

# Inner boolean expression when the intermediate results have already been calculated.
# This is to speed up processing when there is a whole archive of solutions to compare
# a new solution to.
function inner_epsilon_box_dominates_and_epsilon_progress(
  u, uindex, uindextimesepsilon,
  v, vindex, vindextimesepsilon)

  if pareto_dominates_fast(uindex, vindex)
    return (true, true)
  elseif all(uindex .== vindex)
    smaller_distance = norm(u .- uindextimesepsilon) < norm(v .- vindextimesepsilon)
    return (smaller_distance, false)
  else
    return (false, false)
  end

end

# An EpsilonBoxArchive keeps a epsilon Pareto set of the best, non-dominated solutions
# found so far for a multi-objective search/optimization problem.
mutable struct EpsilonBoxArchive
  epsilon
  numsolutions::Int
  solutions::Matrix{Float64}       # solutions in this archive, one solution per column

  # For speedup we cache the boxes and leftpoints of all solutions. They have the same index as in
  # the solutions above.
  boxes::Matrix{Float64}      # the epsilon-boxes corresponding to each solution (same index as solutions)
  leftpoints::Matrix{Float64} # Position of "lower left" point of box (same index as solutions)

  function EpsilonBoxArchive(dimension::Integer, epsilon = 0.01, size = 100)
    new(
      epsilon,
      0,
      zeros(Float64, dimension, size),
      zeros(Float64, dimension, size),
      zeros(Float64, dimension, size)
      )
  end
end

add_if_dominates!(a::EpsilonBoxArchive)


function train_layer_1_sigmoid(neurons,vars,obs,n_per_part,inds_all,x,y)

    # a_all=zeros(neurons,vars+1)
    a_all=Vector{Vector{Float64}}()
    ooops=0
    for i=1:neurons # num_ranges^vars
        # if i in 1:1000:neurons println(Dates.format(now(), "HH:MM:SS")," Calculating weights w for ",i," neuron, of ",neurons," Total.") end
        # if n_per_part[i+1]-n_per_part[i]!=0

            try
                atmp=[ones(n_per_part[i+1]-n_per_part[i]) x[inds_all[n_per_part[i]+1:n_per_part[i+1]],:]]\isigm1.(y[inds_all[n_per_part[i]+1:n_per_part[i+1]]])
                if sum(isnan.(atmp))==0
                    # a_all[i,:].=atmp
                    err=[ones(n_per_part[i+1]-n_per_part[i]) x[inds_all[n_per_part[i]+1:n_per_part[i+1]],:]]*atmp.-isigm1.(y[inds_all[n_per_part[i]+1:n_per_part[i+1]]])
                    # println(i," ",mean(abs.(err)))
                else
                    ooops+=1
                    # println("local inverse problem, neuron=",i," total ooops=",ooops)
                end
                # # println(a_all[i,:])

                obs1=length(inds_all[n_per_part[i]+1:n_per_part[i+1]])
                # println(i," ",obs1)
                if sum(isnan.(atmp))==0 && sum(isinf.(atmp))==0 push!(a_all,atmp) else push!(a_all,zeros(vars+1)); ooops+=1 end

            catch ex
                ooops+=1
                println("problem neuron=",i," total ooops=",ooops," ex=",ex)
                # push!(a_all,zeros(length(inds_all[n_per_part[i]+1:n_per_part[i+1]])))
                push!(a_all,zeros(vars+1))
                # continue
            end
        # end
        # println("calculating weights for neuron ",i," of ",neurons)
    end
    # println("ooops=",ooops)
    layer1=zeros(obs,neurons)
    for i=1:neurons
        if i in 1:100:neurons println(Dates.format(now(), "HH:MM:SS")," Applying sigmoid on ",i," neuron, of ",neurons," Total.") end
        # layer1[:,i].=sigm1.([ones(obs) x]*a_all[i,:])
        # println(size(a_all[i]))
        layer1[:,i].=sigm1.([ones(obs) x]*a_all[i])
        # x[i_so,:]*a_all[i,:]
    end
    # println("calculating linear terms...")
    # a_layer1=[sigm1.(layer1) ones(obs)]\y
    # println(Dates.format(now(), "HH:MM:SS")," Solving system (",obs," by ",neurons+1,") for weights v in the output Layer")
    a_layer1=[layer1 ones(obs)]\y
    # println(Dates.format(now(), "HH:MM:SS")," All weights computed.")
    # predl1=[sigm1.(layer1) ones(obs)]*a_layer1

    return a_all,a_layer1,layer1

end


module Day9
export day9
import ..Day1

function readInput(inputfile="inputs/input09.txt")
    [parse(Int, line) for line in eachline(inputfile)]
end

function day9(inputfile="inputs/input09.txt")
    values = readInput(inputfile)
    part1 = day9Part1(values)
    part2 = day9Part2(values, part1)
    (part1, part2)
end

function day9Part1(values)
    preamble = 25
    for (i, val) in enumerate(values[27:end])
        front,back = Day1.doublePointerSearch(sort(values[i:i+25]), val)
        if front == back # failure
            return val
        end
    end
end

function day9Part2(values, part1)
    lower = 1
    upper = 2
    s = values[lower]+values[upper]
    while s !== part1 && upper <= length(values)
        if s > part1
            lower += 1
        else
            upper += 1
        end
        s = sum(@view values[lower:upper])
    end
    if upper > length(values)
        error("[Day09 Part2]No contiguous range in input summing to $part1 found!")
    end
    return sum(extrema(@view values[lower:upper]))
end

end #module


using Test
using NSGAII

include("./cornercases.jl")

@testset "Domination" begin
    L=[[2,1,1],[1,2,2],[1,0,0]]
    fitness(x::FloatIndividual)=0
    cfg=Cambrian.get_config("./test.yaml";n_population= 3,n_offsprings=5,d_fitness=3)
    e=NSGA2Evolution{FloatIndividual}(cfg,fitness)
    for i in 1:3
        copyto!(e.population[i].fitness,L[i])
    end
    @test dominates(e,e.population[1],e.population[2])==false
    @test dominates(e,e.population[2],e.population[1])==false
    @test dominates(e,e.population[1],e.population[3])==true
    @test dominates(e,e.population[3],e.population[1])==false
end

@testset "fastNonDominatedSort" begin
    fitness1,ranks1=pop1D()
    fitness2,ranks2=pop3D()
    fitness(x::FloatIndividual)=0
    cfg=Cambrian.get_config("./test.yaml";n_population= 20,n_offsprings=5,d_fitness=1)
    e=NSGA2Evolution{FloatIndividual}(cfg,fitness)
    for i in 1:20
        copyto!(e.population[i].fitness,fitness1[i])
    end
    oldPop1D=copy(e.population)
    fastNonDominatedSort!(e)
    @test length(e.population)==length(oldPop1D)
    for x in oldPop1D
        @test x in e.population
    end
    for i in 1:e.config.n_population
        @test ranks1[i]==e.rank[objectid(e.population[i])]
    end

    cfg=Cambrian.get_config("./test.yaml";n_population= 16,n_offsprings=5,d_fitness=3)
    e=NSGA2Evolution{FloatIndividual}(cfg,fitness)
    for i in 1:16
        copyto!(e.population[i].fitness,fitness2[i])
    end
    oldPop3D=copy(e.population)
    fastNonDominatedSort!(e)
    @test length(e.population)==length(oldPop3D)
    for x in oldPop3D
        @test x in e.population
    end
    for i in 1:e.config.n_population
        @test ranks2[i]==e.rank[objectid(e.population[i])]
    end
end


using Markowitz
using LinearAlgebra
import Plots

assets = [ "Bonds - US Government"
           "Bonds - US Corporate"
           "Bonds - International"
           "Bonds - High Yield"
           "Bonds - Bank Loans"
           "Bonds - Emerging USD Debt"
           "Bonds - Emerging Local Debt"
           "Alternative - Emerging FX"
           "Alternative - Commodities"
           "Alternative - REITs"
           "Stocks - US Large"
           "Stocks - US Small"
           "Stocks - International"
           "Stocks - Emerging" ]

V = [ 133.0  39.0  36.0  -17.0  -28.0   31.0   -5.0   -6.0  -34.0  -10.0  -46.0  -68.0  -52.0  -63.0
       39.0  16.0  17.0   10.0    0.0   22.0   13.0    7.0    1.0   19.0    3.0    0.0    5.0    9.0
       36.0  17.0  38.0   19.0    3.0   34.0   45.0   29.0   37.0   48.0   20.0   21.0   42.0   50.0
      -17.0  10.0  19.0   98.0   69.0   65.0   73.0   43.0   68.0  151.0   99.0  130.0  121.0  165.0
      -28.0   0.0   3.0   69.0   84.0   34.0   42.0   25.0   55.0   99.0   64.0   93.0   87.0  119.0
       31.0  22.0  34.0   65.0   34.0   93.0   83.0   48.0   57.0  115.0   76.0   85.0   99.0  145.0
       -5.0  13.0  45.0   73.0   42.0   83.0  154.0   87.0  111.0  168.0  113.0  143.0  165.0  225.0
       -6.0   7.0  29.0   43.0   25.0   48.0   87.0   57.0   75.0   95.0   71.0   89.0  104.0  142.0
      -34.0   1.0  37.0   68.0   55.0   57.0  111.0   75.0  286.0  117.0  101.0  125.0  164.0  239.0
      -10.0  19.0  48.0  151.0   99.0  115.0  168.0   95.0  117.0  491.0  231.0  327.0  259.0  322.0
      -46.0   3.0  20.0   99.0   64.0   76.0  113.0   71.0  101.0  231.0  214.0  256.0  217.0  269.0
      -68.0   0.0  21.0  130.0   93.0   85.0  143.0   89.0  125.0  327.0  256.0  380.0  265.0  342.0
      -52.0   5.0  42.0  121.0   87.0   99.0  165.0  104.0  164.0  259.0  217.0  265.0  297.0  359.0
      -63.0   9.0  50.0  165.0  119.0  145.0  225.0  142.0  239.0  322.0  269.0  342.0  359.0  556.0 ]

E = [ 0.1 0.7 0.8 2.3 2.2 1.9 5.6 5.6 2.2 1.3 0.7 -0.1 4.1 7.2 ]

class = vec([ :FI :FI :FI :FI :FI :FI :FI :ALT :ALT :ALT :EQ :EQ :EQ :EQ ])
colors = map(c-> c==:FI ? "blue" : (c==:ALT ? "green" : "red"),class)
alphas = 0.5 .+ 0.5*[ (1:7)/7 ; (1:3)/3 ; (1:4)/4 ]

function plot_frontier()
    m = f.problem
    unused = vec(sum(f.weights, dims=1).==0)
    xrange = [0, max(maximum(sqrt.(diag(m.V))),maximum(f.vol))]*1.1
    yrange = [minimum(m.E),max(maximum(m.E),maximum(f.ret))]+(maximum(m.E)-minimum(m.E))*[-1,1]*0.1
    dense = smooth(f)
    p1 = Plots.plot(dense[:,1], dense[:,2], legend=false, label="", color="orange", width=2)
    for i = 1:length(m.E)
        Plots.scatter!([sqrt.(diag(m.V)[i])], [m.E[i]], label=m.names[i],
                       markersize=4-1.5*unused[i], markercolor=colors[i], markeralpha=alphas[i],
                       markerstrokecolor=colors[i], markerstrokealpha=alphas[i])
    end
    Plots.scatter!(f.vol, f.ret, legend=false, label="", color="orange", markerstrokecolor="orange", markersize=1)
    Plots.xlabel!("Volatility")
    Plots.ylabel!("Return")
    Plots.xlims!(Tuple(xrange))
    Plots.ylims!(Tuple(yrange))
    weights = f.weights[:,.!unused]
    indices = [1]
    for i in 2:size(weights,1)
        w0 = weights[i-1,:]
        w1 = weights[i,:]
        if sum((sign.(w0).*sign.(w1)).==-1)>0
            for j in 1:length(w0)
                if sign.(w0[j]) != sign.(w1[j])
                    indices = [indices ; i-1+abs(w0[j]/(w1[j]-w0[j]))]
                end
            end
        end
        indices = [indices ; i]
    end
    extweights = zeros(0,size(weights,2))
    extret = []
    for idx in sort(indices)
        if rem(idx,1) == 0
            extweights = [ extweights ; weights[[Int(idx)],:] ]
            extret = [ extret ; f.ret[Int(idx)] ]
        else
            w0 = weights[[Int(floor(idx))],:]
            w1 = weights[[Int(ceil(idx))],:]
            r0 = f.ret[[Int(floor(idx))]]
            r1 = f.ret[[Int(ceil(idx))]]
            extweights = [ extweights ; w0+rem(idx,1)*(w1-w0) ]
            extret = [ extret ; r0+rem(idx,1)*(r1-r0) ]
        end
    end
    wrange = [ minimum(sum((extweights.<0) .* (extweights), dims=2)), maximum(sum((extweights.>0) .* (extweights), dims=2)) ]
    wrange = [ wrange[1], wrange[1] + 2.5*(wrange[2]-wrange[1])]
    ylong = [ zeros(size(extweights,1),1) cumsum((extweights.>0) .* (extweights), dims=2) ]
    yshort = [ zeros(size(extweights,1),1) cumsum((extweights.<0) .* (extweights), dims=2) ]
    p2 = Plots.plot(legend=:right, background_color_legend="#ffffff00", legendfont=Plots.font(5), grid=nothing)
    Plots.xticks!([0,0.5,1])
    for i = 1:(size(ylong,2)-1)
        Plots.plot!(Plots.Shape([ vec(ylong[:,i]); reverse(vec(ylong[:,i+1])) ], [ extret; reverse(extret)]),
                    color=colors[.!unused][i], alpha=alphas[.!unused][i], label=m.names[.!unused][i])
    end
    for i = 1:(size(yshort,2)-1)
        Plots.plot!(Plots.Shape([ vec(yshort[:,i]); reverse(vec(yshort[:,i+1])) ], [ extret; reverse(extret)]),
                    color=colors[.!unused][i],alpha=alphas[.!unused][i], label="")
    end
    Plots.xlims!(Tuple(wrange))
    Plots.ylims!(Tuple(yrange))
    Plots.plot(p1, p2, size=(800, 400))
end

m = markowitz(E, V, names=assets)
unit_sum(m) # total weight = 100%
f=frontier(m)
plot_frontier()
optimal(f) # volatility, return and weights for the minimum variance portofolio
optimal(f,4) # volatility, return and weights for the optimal portofolio with return = 4

m=markowitz(E, V, names=assets, lower=0.05, upper=0.15) # min 5%, max 15% per position
unit_sum(m)
f=frontier(m)
plot_frontier()

m=markowitz(E, V, names=assets, # asset bounds by class: stocks -10/30, bonds 0/20, alt. 0/10
            lower = -0.1 * (class .== :EQ),
            upper = 0.3 * (class .== :EQ) + 0.2 * (class .== :FI) + 0.1 * (class .== :ALT))
unit_sum(m)
add_constraint(m, 1 * (class .== :EQ), '>', 0.3) # net equity exposure between 30% and 60%
add_constraint(m, 1 * (class .== :EQ), '<', 0.6)
add_constraint(m, [1 1 0 0 0 0 0 0 0 0 0 0 0 0], '=', 0.25) # US govt + Investment Grade = 25%
f=frontier(m)
plot_frontier()


# #############################################################################
# #############################################################################
#PSS clustering methods

#This file contains functions to quantify and analyze the degree of clustering of points on a Poincare Surface of Section, which would indicate the presence of a nearby periodic mode.

#Finding the order of periodicity of a series of bounces from their positions on the PSS is a constrained clustering problem, where the data is already ordered such that the clusters are sampled periodically, so points from the same cluster are separated by uniform intervals equal to the number of clusters.

#   clustcount,clustsepmean = findcluster(bncpts::Array{Float64,2};threshhold::Float64=0.1,samplesize::Int64=30,maxtestinterval::Int64=0)
#   This function estimates the number of clusters in the sequence <bncpts> of (theta,chi) bouncepoint locations on the Poincare Surface of Section, but averaging the pairwise distances of <samplesize> bouncepoints separated by a fixed interval in the sequence, and finding the highest interval not exceeding <maxtestinterval> (no limit if set to 0) for which there is a sharp drop in average distance to a proportion below <threshhold>.
#   The number <clustcount> of clusters and the associated mean separation <clustsepmean> of points within each cluster are returned. If no clear indication of clustering is found (no drop in mean intra-cluster separation below threshhold), then <clustcount> is set to 1.

# #############################################################################
# #############################################################################
#Initiate

#Dependencies
#require("util.jl")
#require("boundary.jl")
#require("refractive_index.jl")
#require("ray_solver.jl")

# #############################################################################
# #############################################################################
#Distance definitions on the Poincare Surface of Section

#Convert incidence angle chi to value that grows approximately linearly with
#the periodicity number of whispering gallery modes in a circular cavity. Returns
#the number of sides of a regular polygon-shaped mode if given the appropriate
#incidence angle (half the interior angle).
function chi2whisporder(chi::Float64)
    return 2*pi/(pi-2*chi)::Float64
end

#Metric on the PSS
#Treats (theta,chi) pairs as positions on a half-infinite cylinder embedded in R^3,
#theta being the angular position and chi2whisporder(chi) being the vertical position.
function pssmetric(theta1::Float64,theta2::Float64,chi1::Float64,chi2::Float64)
    return hypot(2*sin(0.5*(theta1-theta2)),
                 chi2whisporder(chi1)-chi2whisporder(chi2))::Float64
end

#Distance function between 2 points on the PSS
#Column vectors
function pssdist{N1,N2}(p1::Array{Float64,N1},p2::Array{Float64,N2})
    return pssmetric(p1[1],p2[1],p1[2],p2[2])::Float64
end

# #############################################################################
# #############################################################################
#Algorithm to estimate number of clusters

#This algorithm relies on the following properties of clusters in the specific context of k-periodic bounce points:

#   - points an interval of k apart in the bounce sequence are separated by a very
#     small PSS distance.
#   - due to the fractal nature of the PSS, the clusters of bounce points themselves
#     appear in clusters of a divisor of k.

#This means that one can progressively test increasing intervals until a sharp drop in PSS distance indicates a periodic cluster, and subsequently increment the test intervals by the number of clusters, since k should be some multiple of it. The details of the algorithm are as follows:

#   1) Distance between pairs of points a fixed interval apart in the bounce sequence
#      is estimated by the mean of such pairwise distances from all such pairs
#      available in the bounce data. This approximates the distance between their
#      corresponding clusters, so a sudden low value indicates that each pair belongs
#      to the same cluster.
#   2) The cluster count is initially estimated to be 1
#   3) The cluster distance for various intervals of bounce points is computed, with
#      the interval incremented by the cluster count, since that should be a divisor
#      of the "actual" order of periodicity.
#   4) If a sharp drop is detected below a threshhold proportion of the
#      typical cluster distances, then it is likely that there is a clustering
#      present. The cluster count is now estimated to be the current test interval.
#   5) Repeat (3) and (4) until the test intervals is too high relative to the total
#      number of bounce points for meaningful clusters to form. If the cluster count
#      remains 1, then it is inconclusive whether there is any clustering present at
#      all.

#Mean PSS distance of pairs of bouncepoints a specified interval apart
function periodic_pssdist(bncpts::Array{Float64,2},order::Int64,samplesize::Int64=0)
    const maxsamplesize::Int64 = size(bncpts,1)-order
    if samplesize == 0 || samplesize > maxsamplesize
        samplesize = maxsamplesize
    end
    return mean([pssdist(bncpts[i,:],bncpts[i+order,:]) for i=1:samplesize])
end

#Estimate the number of clusters of bouncepoints
function findcluster(bncpts::Array{Float64,2};threshhold::Float64=0.1,samplesize::Int64=30,maxtestinterval::Int64=0)
    #Invalid maximum test interval means no limit
    if maxtestinterval <= 0
        maxtestinterval = size(bncpts,1) - samplesize - 1
    end

    #Initiate
    clustcount::Int64 = 1
    testinterval::Int64 = 2
    clustsepsum::Float64 = periodic_pssdist(bncpts,1,samplesize)
    clustsepcount::Int64 = 1

    while(testinterval <= maxtestinterval)
        #Get mean PSS distance for test interval
        testdist::Float64 = periodic_pssdist(bncpts,testinterval,samplesize)

        #Test against threshhold
        if testdist < threshhold*clustsepsum/clustsepcount
            #Sharp drop detected
            #cluster count likely to be multiple of test interval
            clustcount = testinterval
            #Reset intra-cluster distance data
            clustsepsum = testdist
            clustsepcount = 1
        else
            #No significant drop detected
            #Collect intra-cluster separation data
            clustsepsum += testdist
            clustsepcount += 1
        end

        #Increment test interval by the current cluster count estimate
        testinterval += clustcount
    end

    #Return best estimate of number of clusters, and mean separation of points
    #within each cluster
    return clustcount, clustsepsum/clustsepcount
end


module UMVETests

using LinearAlgebra
using Subsystems
using Test

@testset "umve estimator" begin
    A = Matrix{Float64}(I, 2, 2)
    B = reshape([1., 0.], :, 1)
    C = I
    Aij = [0.1 0; 0 0]
    W = Diagonal([0.1, 0.1])
    V = W

    @testset "handles empty neighbourhoods" begin
        s = LdtiSubsystem(1, A, B, C)
        @testset "with svd" begin
            o = UMVE(s.A, s.B, s.C, s.E, W, V)
            Rk = Subsystems.R(o.A, o.C, o.Pi, o.W, o.V)
            M∅ = Subsystems.M(o.G, o.C, Rk)
            Kk = Subsystems.K(o.A, o.C, o.Pi, o.W, Rk)
            Ā∅ = Subsystems.barA(Kk, o.C, o.G, M∅)
            L̄∅ = Subsystems.barL(Kk, o.C, o.G, M∅)

            @test M∅ == zeros(0,2)
            @test Ā∅ == (I - Kk*o.C)
            @test L̄∅ == Kk
        end
        @testset "without svd" begin
            o = UMVE(s.A, s.B, s.C, s.E, W, V, false)
            Rk = Subsystems.R(o.A, o.C, o.Pi, o.W, o.V)
            M∅ = Subsystems.M(o.G, o.C, Rk)
            Kk = Subsystems.K(o.A, o.C, o.Pi, o.W, Rk)
            Ā∅ = Subsystems.barA(Kk, o.C, o.G, M∅)
            L̄∅ = Subsystems.barL(Kk, o.C, o.G, M∅)

            @test M∅ == zeros(0,2)
            @test Ā∅ == (I - Kk*o.C)
            @test L̄∅ == Kk
        end
    end

    @testset "Constructor" begin
        @testset "builds correctly if exists" begin
            N = 4
            LSS = [LdtiSubsystem(i, A, B, C) for i in 1:N]
            # make ring topology
            [addNeighbour(LSS[i], LSS[i%N + 1], Aij) for i in 1:N]

            @testset "with svd" begin
                O = [UMVE(s.A, s.B, s.C, s.E, W, V) for s in LSS]
                for o in O
                    @test o.G ≈ reshape([1. 0], :, 1)
                end
            end
            @testset "without svd" begin
                O = [UMVE(s.A, s.B, s.C, s.E, W, V, false) for s in LSS]
                for o in O
                    @test o.G ≈ Matrix{eltype(o.G)}(I, o.nxi, o.nxi)
                end
            end
        end
        @testset "can build with no neighbours" begin
            s = LdtiSubsystem(1, A, B, C)
            @testset "with svd" begin
                o = UMVE(s.A, s.B, s.C, s.E, W, V)
                @test o.G == zeros(s.nx, 0)
            end
            @testset "without svd" begin
                o = UMVE(s.A, s.B, s.C, s.E, W, V, false)
                @test o.G == zeros(s.nx, 0)
                @test o.Ebar == zeros(s.nx, 0)
            end
        end
    end

end

end


# test_convergence.jl

function get_problem(ns)
    n = m = ns
    L = laplace2d(n, m)
    A = V_cycle(L, (n, m))
    pts = 1/n:1/n:1-1/n
    b = Float64[h(x,y) for x in pts, y in pts][:]
    return (A,b)
end

@testset "Convergence for example from Briggs et. al." begin
    for n in [16 32 64 128 256 512 1024]
        (A, b) = get_problem(n)
        A.grids[1].b .= b # copy rhs
        for item in Base.Iterators.take(A,15) end # iterate
        @test A.resnorm[end] < 1/n^2
        log(A,2)
    end
end


mutable struct Variable
    data::Array
    grad::Array
    Variable(data) = new(data)
end

abstract type _Function end
function _function(f::_Function, forward::Function, input::Variable)::Variable
    x = input.data
    y::Array = forward(x)
    output = Variable(y)
    f.input = input
    return output
end

mutable struct Square <: _Function
    input::Variable
    Square() = new()
end
(s::Square)(input::Variable)::Variable = _function(s, x -> x .^ 2, input)
backward(s::Square, gy::Array) = 2 * s.input.data .* gy

mutable struct Exp <: _Function
    input::Variable
    Exp() = new()
end
(s::Exp)(input::Variable)::Variable = _function(s, x -> exp.(x), input)
backward(e::Exp, gy::Array) = exp.(e.input.data) .* gy


function solve_re{T<:FloatingPoint}(model::Blanchard_Kahn_Form{T},cutoff::T)

  nx    = copy(model.nx)
  ny    = copy(model.ny)
  a     = copy(model.a)
  c     = copy(model.c)
  sigma = copy(model.sigma)

  n = nx+ny

  (s,q) = schur(complex(a))
  q = q'  # So now q*a*q' = s

  # Reorder the eigenvalues so that those with modulus greater than "cutoff" reside at the bottom.

  (s,q) = reorder_complex_schur(s,q)

  # Calculate the number of unstable eigenvalues

  grc = 0
  for i = 1:n

    if norm(s[i,i]) > cutoff
      grc += 1
    end

  end

  soln_type = "determinate"
  if grc < ny
    soln_type = "indeterminate"
  elseif grc > ny
    soln_type = "explosive"
  end

  # Construct the rational expectations equilibrium by eliminating the unstable dynamics

  q11 = q[1:nx,1:nx]
  q12 = q[1:nx,(nx+1):n]
  q21 = q[(nx+1):n,1:nx]
  q22 = q[(nx+1):n,(nx+1):n]
  s11 = s[1:nx,1:nx]

  p = real((q11-(q12/q22)*q21)\s11*(q11-(q12/q22)*q21))
  k = c[1:nx,:]
  f = -real(q22\q21)

  soln = Blanchard_Kahn_Soln(p,k,f,sigma,grc,soln_type)

  return soln

end

function solve_re{T<:FloatingPoint}(model::Blanchard_Kahn_Form{T},cutoff::T,tol::T)

  nx    = copy(model.nx)
  ny    = copy(model.ny)
  a     = copy(model.a)
  c     = copy(model.c)
  sigma = copy(model.sigma)

  n = nx+ny

  s = eigvals(a)

  grc = 0
  for i = 1:n

    if norm(s[i]) > cutoff
      grc += 1
    end

  end

  soln_type = "determinate"
  if grc < ny
    soln_type = "indeterminate"
  elseif grc > ny
    soln_type = "explosive"
  end

  # Construct the rational expectations equilibrium

  a11 = a[1:nx,1:nx]
  a12 = a[1:nx,(nx+1):n]
  a21 = a[(nx+1):n,1:nx]
  a22 = a[(nx+1):n,(nx+1):n]

  f0 = -a22\a21
  f = Array(T,ny,nx)
  len = Inf
  while len > tol

    f = (a22-f0*a12)\(f0*a11-a21)
    len = maxabs(f0-f)
    f0 = copy(f)

  end

  p = a11+a12*f
  k = c[1:nx,:]

  soln = Blanchard_Kahn_Soln(p,k,f,sigma,grc,soln_type)

  return soln

end

function solve_re{T<:FloatingPoint}(model::Klein_Form{T},cutoff::T)

  nx    = copy(model.nx)
  ny    = copy(model.ny)
  a     = copy(model.a)
  b     = copy(model.b)
  c     = copy(model.c)
  sigma = copy(model.sigma)

  n = nx+ny

  (s,t,q,z) = schur(complex(a),complex(b))
  q = q'  # So now q*a*z = s and q*b*z = t

  # Reorder the generalized eigenvalues so that those with modulus greater than "cutoff" reside at the bottom.

  (s,t,q,z) = reorder_generalized_complex_schur(s,t,q,z)

  # Calculate the number of unstable eigenvalues

  grc = 0
  for i = 1:n

    if norm(t[i,i]) < 2*eps(T)*norm(t) || norm(s[i,i]/t[i,i]) > cutoff
      grc += 1
    end

  end

  soln_type = "determinate"
  if grc < ny
    soln_type = "indeterminate"
  elseif grc > ny
    soln_type = "explosive"
  end

  # Construct the rational expectations equilibrium by eliminating the unstable dynamics

  z11 = z[1:nx,1:nx]
  z21 = z[(nx+1):n,1:nx]
  t11 = t[1:nx,1:nx]
  s11 = s[1:nx,1:nx]

  p = real((z11/t11)*(s11/z11))
  k = c[1:nx,:]
  f = real(z21/z11)

  soln = Klein_Soln(p,k,f,sigma,grc,soln_type)

  return soln

end

function solve_re{T<:FloatingPoint}(model::Structural_Form{T},cutoff::T)

  a     = copy(model.a)
  a1    = copy(model.a1)
  b     = copy(model.b)
  c     = copy(model.c)
  sigma = copy(model.sigma)

  nx = size(a1,1)
  ny = size(b,1)

  n = nx+ny

  a = [zeros(nx,nx) eye(ny); -a1 a]
  b = [eye(nx) zeros(nx,ny); zeros(ny,nx) b]
  c = [zeros(nx,size(c,2)); -c]

  (s,t,q,z) = schur(complex(a),complex(b))
  q = q'  # So now q*a*z = s and q*b*z = t

  # Reorder the generalized eigenvalues so that those with modulus greater than "cutoff" reside at the bottom.

  (s,t,q,z) = reorder_generalized_complex_schur(s,t,q,z)

  # Calculate the number of unstable eigenvalues

  grc = 0
  for i = 1:n

    if norm(t[i,i]) < 2*eps(T)*norm(t) || norm(s[i,i]/t[i,i]) > cutoff
      grc += 1
    end

  end

  soln_type = "determinate"
  if grc < ny
    soln_type = "indeterminate"
  elseif grc > ny
    soln_type = "explosive"
  end

  # Construct the rational expectations equilibrium by eliminating the unstable dynamics

  z11 = z[1:nx,1:nx]
  z12 = z[1:nx,(nx+1):n]
  z21 = z[(nx+1):n,1:nx]
  z22 = z[(nx+1):n,(nx+1):n]
  t11 = t[1:nx,1:nx]
  s11 = s[1:nx,1:nx]
  s12 = s[1:nx,(nx+1):n]
  s22 = s[(nx+1):n,(nx+1):n]

  qc  = -q*c
  qcs = qc[1:nx,:]       # The part of qc that is in the stable block
  qcu = qc[(nx+1):n,:]   # The part of qc that is in the unstable block

  p = real((z11/t11)*(s11/z11))
  k = real((z11/t11)*(s12/s22)*qcu-(z11/t11)*(s11/z11)*(z12/s22)*qcu-(z11/t11)*qcs)

  soln = Structural_Soln(p,k,sigma,grc,soln_type)

  return soln

end

function solve_re{T<:FloatingPoint}(model::Structural_Form{T},cutoff::T,tol::T)

  a     = copy(model.a)
  a1    = copy(model.a1)
  b     = copy(model.b)
  c     = copy(model.c)
  sigma = copy(model.sigma)

  p0 = copy(a1)
  p  = copy(p0)

  len = Inf
  while len > tol

    p = (a-b*p0)\a1
    len = maxabs(p-p0)
    p0 = copy(p)

  end
  k = (a-b*p0)\c

  nx = size(a1,1)
  ny = size(b,1)
  aa = [zeros(nx,nx) eye(ny); -a1 a]
  bb = [eye(nx) zeros(nx,ny); zeros(ny,nx) b]

  lambda = eigvals(bb,aa)

  grc = 0
  for i = 1:(nx+ny)

    if norm(lambda[i]) > cutoff
      grc += 1
    end

  end

  soln_type = "determinate"
  if grc < ny
    soln_type = "indeterminate"
  elseif grc > ny
    soln_type = "explosive"
  end

  soln = Structural_Soln(p,k,sigma,grc,soln_type)

  return soln

end

function solve_re{T<:FloatingPoint}(model::Sims_Form{T},cutoff::T)

  # This function is modeled on a routine written by
  # Christopher A. Sims (http://www.princeton.edu/~sims/)

  gamma0 = copy(model.gamma0)
  gamma1 = copy(model.gamma1)
  c      = copy(model.c)
  psi    = copy(model.psi)
  pi     = copy(model.pi)
  sigma  = copy(model.sigma)

  n  = size(gamma0,1)

  (s,t,q,z) = schur(complex(gamma1),complex(gamma0))
  q = q'  # So now q*gamma1*z = s and q*gamma0*z = t

  # Reorder the generalized eigenvalues so that those with modulus greater than "cutoff" reside at the bottom.

  (s,t,q,z) = reorder_generalized_complex_schur(s,t,q,z)

  # Calculate the number of unstable eigenvalues

  grc = 0
  for i = 1:n

    if norm(t[i,i]) < 2*eps(T)*norm(t) || norm(s[i,i]/t[i,i]) > cutoff
      grc += 1
    end

  end

  # Define some ingredients that are needed for any solution

  q1 = q[1:(n-grc),:]
  q2 = q[(n-grc+1):n,:]
  z1 = z[:,1:(n-grc)]'
  z2 = z[:,(n-grc+1):n]'
  t2 = t[(n-grc+1):n,(n-grc+1):n]
  s2 = s[(n-grc+1):n,(n-grc+1):n]
  q2pi  = q2*pi
  q2psi = q2*psi

  # Check for the existance of a solution

  (pi_rows,pi_cols)   = size(pi)
  (psi_rows,psi_cols) = size(psi)

  if grc == 0
    q2pi    = zeros(0,pi_cols)
    uq2pi   = zeros(0,0)
    dq2pi   = zeros(0,0)
    vq2pi   = zeros(pi_cols,0)

    q2psi   = zeros(0,psi_cols)
    uq2psi  = zeros(0,0)
    dq2psi  = zeros(0,0)
    vq2psi  = zeros(psi_cols,0)

    non_zero_index_psi = []
  else
    (uq2pi,dq2pi,vq2pi) = svd(q2pi)
    non_zero_index_pi = find(dq2pi .> 2*eps(T)*norm(q2pi))
    uq2pi = uq2pi[:,non_zero_index_pi]
    vq2pi = vq2pi[:,non_zero_index_pi]
    dq2pi = diagm(dq2pi[non_zero_index_pi])

    (uq2psi,dq2psi,vq2psi) = svd(q2psi)
    non_zero_index_psi = find(dq2psi .> 2*eps(T)*norm(q2psi))
    uq2psi = uq2psi[:,non_zero_index_psi]
    vq2psi = vq2psi[:,non_zero_index_psi]
    dq2psi = diagm(dq2psi[non_zero_index_psi])
  end

  if isempty(non_zero_index_psi)
    iid_exist = true
  else
    iid_exist = (norm(uq2psi - uq2pi*uq2pi'*uq2psi) < 2*eps(T)*norm(uq2psi))
  end

  if iid_exist == false
    general_exist = false
  else
    t2iq2psi = t2\q2psi
    p        = t2\s2
    h        = t2iq2psi
    for i = 1:n

      h = [p*h t2iq2psi]

    end
    (uh,dh,vh) = svd(h)
    non_zero_index_dh = find(dh .> 2*eps(T)*norm(h))
    uh = uh[:,non_zero_index_dh]
    vh = vh[:,non_zero_index_dh]
    dh = diagm(dh[non_zero_index_dh])
    general_exist = (norm(uh-uq2pi*uq2pi'*uh) < 2*eps(T)*norm(uh,1))
  end

  # Check for uniqueness

  if grc == n
    q1pi  = zeros(0,pi_cols)
    uq1pi = zeros(0,0)
    vq1pi = zeros(pi_cols,0)
    dq1pi = zeros(0,0)
  else
    q1pi = q1*pi
    (uq1pi,dq1pi,vq1pi) = svd(q1pi)
    non_zero_index_dq1pi = find(dq1pi .> 2*eps(T)*norm(q1pi))
    uq1pi = uq1pi[:,non_zero_index_dq1pi]
    vq1pi = vq1pi[:,non_zero_index_dq1pi]
    dq1pi = diagm(dq1pi[non_zero_index_dq1pi])
  end

  if isempty(vq1pi)
    unique = true
  else
    resid = vq1pi-vq2pi*vq2pi'vq1pi
    (u1,d1,v1) = svd(resid)
    unique = ((sumabs(d1) > 2*eps(T)*n) != true)
  end
  tmat = [eye(n-grc) -(uq2pi*(dq2pi\vq2pi')*vq1pi*dq1pi*uq1pi')']
  g0   = [tmat*t; zeros(grc,(n-grc)) eye(grc)]
  g1   = [tmat*s; zeros(grc,n)]

  g0i    = inv(g0)
  g1     = real(z*g0i*g1*z')
  c      = real(z*g0i*[tmat*q*c; (t[(n-grc+1):n,(n-grc+1):n]-s[(n-grc+1):n,(n-grc+1):n])\q2*c])
  impact = real(z*g0i*[tmat*q*psi; zeros(grc,size(psi,2))])
  fmat   = real(s[(n-grc+1):n,(n-grc+1):n]\t[(n-grc+1):n,(n-grc+1):n])
  fwt    = -real(s[(n-grc+1):n,(n-grc+1):n]\q2*psi)
  ywt    = real(z*g0i[:,(n-grc+1):n])

  soln = Sims_Soln(g1,c,impact,ywt,fmat,fwt,sigma,grc,iid_exist,general_exist,unique)

  return soln

end

function solve_re{T<:FloatingPoint}(model::Gomme_Klein_Form,cutoff::T)

  nx     = copy(model.nx)
  ny     = copy(model.ny)
  deriv1 = copy(model.derivs)
  deriv2 = copy(model.hessians)
  eta    = copy(model.eta)
  sigma  = copy(model.sigma)

  n = nx+ny

  # Reorder the second derivatives to put them in the order that Gomme and Klein use
  # Probably do away with this in a subsequent update

  hes = [deriv2[(n+1):2*n,:]; deriv2[1:n,:]]
  for i = 2:n

    hes = [hes; deriv2[2*n*(i-1)+n+1:2*n*i,:]; deriv2[2*n*(i-1)+1:2*n*(i-1)+n,:]]

  end
  hes = [hes[:,(n+1):2*n] hes[:,1:n]]

  # Construct partitioned first derivative matrices

  fx  = deriv1[1:n,1:nx]
  fy  = deriv1[1:n,(nx+1):n]
  fxp = deriv1[1:n,(n+1):(n+nx)]
  fyp = deriv1[1:n,(n+nx+1):(2*n)]

  # Solve the linearized first-order system

  a =  [fx fy]
  b = -[fxp fyp]
  c =  [eta; zeros(ny,size(eta,2))]

  first_order_model = Klein_Form(nx,ny,a,b,c,model.sigma)
  first_order_soln  = solve_re(first_order_model,cutoff)

  hx = first_order_soln.p
  gx = first_order_soln.f
  grc = first_order_soln.grc

  soln_type = "determinate"
  if grc < ny
    soln_type = "indeterminate"
  elseif grc > ny
    soln_type = "explosive"
  end

  # Set up the LP problem needed to construct the matrices on the second-order states

  m = [hx; gx*hx; eye(nx); gx]
  q = kron(eye(n),m')*hes*m
  b1 = kron(fxp,eye(nx))
  b2 = kron(fyp,eye(nx))
  b4 = kron(fy,eye(nx))
  c1 = kron(eye(ny),hx')
  c2 = kron(gx,eye(nx))

  qq = [kron(eye(nx),b4)+kron(hx',b2*c1) kron(eye(nx),b1+b2*c2)]
  xx = -qq\vec(q)
  gxx = reshape(xx[1:nx^2*ny],nx*ny,nx)
  hxx = reshape(xx[nx^2*ny+1:n*nx*nx],nx^2,nx)

  # Set up the LP problem needed to construct the intercepts, which contain the volatility effects

  qq = [fxp+fyp*gx fyp+fy]
  nn = [eye(nx); gx; zeros(n,nx)]
  q = fyp*tracem(kron(eye(ny),eta*(sigma')*eta')*gxx)+tracem(kron(eye(n),nn')*hes*nn*eta*(sigma')*eta')
  ss = -qq\q
  ssh = ss[1:nx]
  ssg = ss[nx+1:n]

  soln = Gomme_Klein_Soln(ssh,hx,hxx,ssg,gx,gxx,sigma,grc,soln_type)

  return soln

end

function solve_re{T<:FloatingPoint}(model::Lombardo_Sutherland_Form,cutoff::T)

  nx     = copy(model.nx)
  ny     = copy(model.ny)
  deriv1 = copy(model.derivs)
  deriv2 = copy(model.hessians)
  eta    = copy(model.eta)
  sigma  = copy(model.sigma)

  n  = nx+ny
  ns = size(eta,2)

  # Solve the first order system

  a1 = -deriv1[1:n,(n+1):2*n]
  a2 = deriv1[1:n,1:n]
  c  = [eta; zeros(ny,size(eta,2))]

  (s,t,q,z) = schur(complex(a2),complex(a1))
  q = q'  # So now q*a*z = s and q*b*z = t

  # Reorder the generalized eigenvalues so that those with modulus greater than "cutoff" reside at the bottom.

  (s,t,q,z) = reorder_generalized_complex_schur(s,t,q,z)

  # Calculate the number of unstable eigenvalues

  grc = 0
  for i = 1:n

    if norm(t[i,i]) < 2*eps(T)*norm(t) || norm(s[i,i]/t[i,i]) > cutoff
      grc += 1
    end

  end

  soln_type = "determinate"
  if grc < ny
    soln_type = "indeterminate"
  elseif grc > ny
    soln_type = "explosive"
  end

  # Construct the rational expectations equilibrium by eliminating the unstable dynamics

  z11 = z[1:nx,1:nx]
  z12 = z[1:nx,nx+1:n]
  z21 = z[(nx+1):n,1:nx]
  z22 = z[(nx+1):n,(nx+1):n]

  t11 = t[1:nx,1:nx]
  t12 = t[1:nx,nx+1:n]
  t22 = t[nx+1:n,nx+1:n]

  s11 = s[1:nx,1:nx]
  s12 = s[1:nx,nx+1:n]
  s22 = s[nx+1:n,nx+1:n]

  phi   = real((z11/t11)*(s11/z11))
  gamma = c[1:nx,:]
  omega = [eye(nx); real(z21/z11)]  # Defines [x; y] as a function of [x]

  soln_type = "determinate"
  if grc < ny
    soln_type = "indeterminate"
  elseif grc > ny
    soln_type = "explosive"
  end

  # Set up the matrices for the second order terms

  for i = 1:n

    for  j = 1:2*n

      deriv2[(i-1)*2*n+j,j] = deriv2[(i-1)*2*n+j,j]/2

    end

    deriv2[(i-1)*2*n+1:i*2*n,:] = 2*deriv2[(i-1)*2*n+1:i*2*n,:]

  end

  v2vh_n = vec_to_vech(n)

  a4 = v2vh_n*vec(deriv2[1:n,1:n])
  a5 = v2vh_n*vec(deriv2[(n+1):2*n,(n+1):2*n])
  for i = 2:n

    a4 = [a4 v2vh_n*vec(deriv2[(i-1)*2*n+1:(i-1)*2*n+n,1:n])]
    a5 = [a5 v2vh_n*vec(deriv2[(i-1)*2*n+n+1:i*2*n,(n+1):2*n])]

  end

  a4 = a4'
  a5 = a5'

  vh2v_nx = vech_to_vec(nx)
  v2vh_nx = vec_to_vech(nx)
  vh2v_ns = vech_to_vec(ns)
  v2vh_ns = vec_to_vech(ns)
  p       = permutation(nx,ns)

  r           = v2vh_n*kron(omega,omega)*vh2v_nx            # n*(n+1)/2   times nx*(nx+1)/2
  phi_tilda   = v2vh_nx*kron(phi,phi)*vh2v_nx               # nx*(nx+1)/2 times nx*(nx+1)/2
  gamma_tilda = v2vh_nx*kron(gamma,gamma)*vh2v_ns           # nx*(nx+1)/2 times ns*(ns+1)/2
  psi_tilda   = v2vh_nx*(kron(phi,gamma)+kron(gamma,phi)*p) # nx*(nx+1)/2 times nx*ns

  g = a4*r+a5*r*phi_tilda                                   # n times nx*(nx+1)/2
  h = a5*r*gamma_tilda                                      # n times ns*(ns+1)/2

  sig = v2vh_ns*vec(sigma)                                  # ns*(ns+1)/2 times ns*ns

  # The next step is to resolve the system using the q and z matrices obtained earlier

  g = q*g
  h = q*h

  gs = g[1:nx,:]
  gu = g[nx+1:n,:]

  hs = h[1:nx,:]
  hu = h[nx+1:n,:]

  m = reshape((kron(phi_tilda',t22)-kron(eye(int(nx*(nx+1)/2)),s22))\vec(gu),ny,int(nx*(nx+1)/2))

  hx  = real(z11*(t11\s11)/z11)
  hxx = real(-(z11*(t11\s11)/z11)*z12*m + z11*(t11\(s12*m-t12*m*phi_tilda+gs)) + z12*m*phi_tilda)

  gx  = real(z21/z11)
  gxx = real((z22-z21*(z11\z12))*m)

  vbar = (eye(int(nx*(nx+1)/2))-phi_tilda)\gamma_tilda
  m3   = (eye(ny)-s22\t22)\(s22\(gu*vbar+hu))
  m2   = reshape((eye(int(nx*(nx+1)/2)*ny)-kron(phi_tilda',s22\t22))\(vec(s22\gu)),ny,int(nx*(nx+1)/2))

  p4 = -real(z22'\(m3-m2*vbar))
  p3 = -z22'\m2
  p2 = -z22'\z12'

  r4 = (t11*z21'+t12*z22')*p4
  r3 = (t11*z21'+t12*z22')*p3
  r2 = t11*z11'+t12*z12'+(t11*z21'+t12*z22')*p2

  d4 = (s11*z21'+s12*z22')*p4+hs

  ssg = p4*sig
  ssh = real(r2\(d4-r4-r3*gamma_tilda))*sig

  soln = Lombardo_Sutherland_Soln(ssh,hx,hxx,ssg,gx,gxx,phi_tilda,gamma_tilda,psi_tilda,sigma,grc,soln_type)

  return soln

end


## U updates

"""
    updateU!(M::HePPCATModel,Y,method)

Update the factor eigenvectors `M.U` with respect to data `Y` using `method`.
"""
function updateU! end

# Update method: Minorize Maximize (with linear minorizer)
"""
    updateU!(M::HePPCATModel,Y,::MinorizeMaximize)

Minorize maximize update of `M.U` using a linear minorizer.
"""
function updateU!(M::HePPCATModel,Y,::MinorizeMaximize)
    M.U .= polar(gradF(M.U,M.λ,M.v,Y))
    return M
end

# Update method: Projected Gradient Ascent
"""
    updateU!(M::HePPCATModel,Y,pga::ProjectedGradientAscent)

Projected gradient ascent update of `M.U`.
Supported step size types:
`Number` (i.e., constant step size),
[`InverseLipschitz`](@ref),
and
[`ArmijoSearch`](@ref) (experimental).
"""
function updateU!(M::HePPCATModel,Y,pga::ProjectedGradientAscent{<:Number})
    if pga.stepsize == Inf
        M.U .= polar(gradF(M.U,M.λ,M.v,Y))
    else
        M.U .= polar(M.U + pga.stepsize*gradF(M.U,M.λ,M.v,Y))
    end
    return M
end
function updateU!(M::HePPCATModel,Y,pga::ProjectedGradientAscent{<:InverseLipschitz})
    return updateU!(M,Y,ProjectedGradientAscent(inv(pga.stepsize.bound(M,Y))))
end
function updateU!(M::HePPCATModel,Y,pga::ProjectedGradientAscent{<:ArmijoSearch})
    params = pga.stepsize

    dFdU = gradF(M.U,M.λ,M.v,Y)
    F0, FΔ = F(M.U,M.λ,M.v,Y), params.tol * norm(dFdU)^2
    m = findfirst(IdentityUnitRange(0:params.maxsearches-1)) do m
        Δ = params.contraction^m * params.stepsize
        F(polar(M.U + Δ*dFdU),M.λ,M.v,Y) >= F0 + Δ * FΔ
    end
    if isnothing(m)
        @warn "Exceeded maximum line search iterations. Accuracy not guaranteed."
        m = params.maxsearches
    end

    Δ = params.contraction^m * params.stepsize
    M.U .= polar(M.U + Δ*dFdU)
    return M
end

# Update method: Stiefel Gradient Ascent
"""
    updateU!(M::HePPCATModel,Y,sga::StiefelGradientAscent)

Stiefel gradient ascent update of `M.U`.
Supported step size types:
[`ArmijoSearch`](@ref).
"""
function updateU!(M::HePPCATModel,Y,sga::StiefelGradientAscent{<:ArmijoSearch})
    params = sga.stepsize

    dFdU = gradF(M.U,M.λ,M.v,Y)
    ∇F = dFdU - M.U*(dFdU'M.U)

    F0, FΔ = F(M.U,M.λ,M.v,Y), params.tol * norm(∇F)^2
    m = findfirst(IdentityUnitRange(0:params.maxsearches-1)) do m
        Δ = params.contraction^m * params.stepsize
        F(geodesic(M.U,∇F,Δ),M.λ,M.v,Y) >= F0 + Δ * FΔ
    end
    if isnothing(m)
        @warn "Exceeded maximum line search iterations. Accuracy not guaranteed."
        m = params.maxsearches
    end

    Δ = params.contraction^m * params.stepsize
    M.U .= geodesic(M.U,∇F,Δ)
    return M
end

# Objective with respect to U
"""
    F(U,λ,v,Y)

Objective for the optimization problem w.r.t `U`.
"""
F(U,λ,v,Y) = 1/2*sum(norm(sqrt(Diagonal(λ./vl./(λ.+vl)))*U'*Yl)^2 for (Yl,vl) in zip(Y,v))

"""
    gradF(U,λ,v,Y)

Euclidean gradient of objective for the optimization problem w.r.t `U`.
"""
gradF(U,λ,v,Y) = sum(Yl * Yl' * U * Diagonal(λ./vl./(λ.+vl)) for (Yl,vl) in zip(Y,v))

"""
    LipBoundU1(M::HePPCATModel,Y)

Lipschitz bound of objective w.r.t. `U` at `M` with data `Y`.
"""
function LipBoundU1(M::HePPCATModel,Y)
    L, λmax = length(M.v), maximum(M.λ)
    return sum(norm(Y[l])^2*λmax/M.v[l]/(λmax+M.v[l]) for l in 1:L)
end

"""
    LipBoundU2(M::HePPCATModel,Y)

Lipschitz bound of objective w.r.t. `U` at `M` with data `Y`.
"""
function LipBoundU2(M::HePPCATModel,Y)
    L, λmax = length(M.v), maximum(M.λ)
    return sum(opnorm(Y[l])^2*λmax/M.v[l]/(λmax+M.v[l]) for l in 1:L)
end

# Utilities
"""
    polar(A)

Polar factor of `A`.
"""
function polar(A)
    F = svd(A)
    return F.U*F.Vt
end

"""
    skew(A)

Skew-symmetrize `A` by computing `(A-A')/2`.
"""
skew(A) = (A-A')/2

"""
    geodesic(U,X,t)

Geodesic step of size `t` from `U` in direction of `X` along the Stiefel manifold.
"""
function geodesic(U,X,t)
    k = size(U,2)

    A = skew(U'X)
    Q,R = qr(X - U*(U'X))

    MN = exp(t*[A -R'; R zeros(k,k)])[:,1:k]
    M, N = MN[1:k,:], MN[k+1:end,:]

    return U*M + Matrix(Q)*N
end