# For converting a number of molecules to molars, given the cell volume
# Arguments:  molecules = number of molecules
#             volume = volume in metres cubed
# Returns:    Float64, the molar density

function moleculesToMolars(molecules::Float64, volume::Float64)
  N_av = 6.02214085774e+23 # avagadro's number
  volume_litres = 1000 * volume
  molars = molecules/(N_av * volume_litres)
  return molars
end

function moleculesToMolars(molecules::Array{Float64, 1}, volume::Float64)
  N_av = 6.02214085774e+23 # avagadro's number
  volume_litres = 1000 * volume
  molars = molecules/(N_av * volume_litres)
  return molars
end

function moleculesToMolars(molecules::Array{Float64, 2}, volume::Float64)
  N_av = 6.02214085774e+23 # avagadro's number
  volume_litres = 1000 * volume
  molars = molecules/(N_av * volume_litres)
  return molars
end


###############################################################################
#
# Big leaf canopy model
#
###############################################################################
const _K_D         = 0.7
const _ALPAR       = 0.8
const _ALNIR       = 0.2
const _ALPAR_SQ    = sqrt(_ALPAR)
const _ALNIR_SQ    = sqrt(_ALNIR)
const _ALPAR_SQ_KD = _ALPAR_SQ * _K_D
const _ALNIR_SQ_KD = _ALNIR_SQ * _K_D

"""
    big_leaf_partition(
                lai::FT,
                zenith::FT,
                r_all::FT,
                r_dir::FT = FT(0.8)
    ) where {FT<:AbstractFloat}

Partition the big-leaf canopy into sunlit and shaded layers, given
- `lai` Leaf area index
- `zenith` Zenith angle in degree
- `r_all` Total radiation in `[W m⁻²]`
- `r_dir` Direct radiation partition in `r_all`

The function returns
- `ratio` ratio of sunlit leaves out of all leaves
- `q_slm` Mean sunlit layer PAR
- `q_shm` Mean shaded layer PAR
- `e_sl` Mean sunlit layer absorbed total energy
- `e_sh` Mean shaded layer absorbed total energy
"""
function big_leaf_partition(
            lai::FT,
            zenith::FT,
            r_all::FT,
            r_dir::FT = FT(0.8)
) where {FT <:AbstractFloat}
    # 1. assume 50%-50% PAR and NIR
    par_tot = r_all / 2 / FT(0.235);
    q_ob    = par_tot * r_dir;
    q_od    = par_tot * (1 - r_dir);

    # 2. calculate the LAI
    shape = FT(1.0);
    shapa = FT(2.0);

    # 3. calculate the mean sunlit layer PAR
    # For vertical leaves k_be = 2.0 * tand(zenith) / pi
    k_be  = sqrt( shape^2+tand(zenith)^2 ) /
              ( shape + FT(1.774) * (shape+FT(1.182))^FT(-0.733) );
    q_btp = q_ob * exp( -FT(_ALPAR_SQ) * k_be * lai );
    q_btn = q_ob * exp( -FT(_ALNIR_SQ) * k_be * lai );
    q_b   = q_ob * exp( -k_be*lai );
    q_scp = q_btp - q_b;
    q_scn = q_btn - q_b;

    # 4. mean par in sunlit and shade layers
    q_dm  = q_od * ( 1 - exp(-FT(_ALPAR_SQ_KD)) ) / ( FT(_ALPAR_SQ_KD) * lai );
    r_dm  = q_od * ( 1 - exp(-FT(_ALNIR_SQ_KD)) ) / ( FT(_ALNIR_SQ_KD) * lai );
    q_shm = q_dm + q_scp/shapa;
    r_shm = r_dm + q_scn/shapa;
    q_slm = k_be*q_ob + q_dm + q_scp/shapa;
    r_slm = k_be*q_ob + r_dm + q_scn/shapa;

    # 5. calculate lasi_sl and lai_sh
    lai_sl = ( 1 - exp(-k_be*lai) ) / k_be;
    lai_sh = lai - lai_sl;
    ratio  = lai_sl / lai;

    # 6. calculate the radiation based on PAR and NIR
    ep_sl = q_slm * FT(_ALPAR) * FT(0.235);
    ep_sh = q_shm * FT(_ALPAR) * FT(0.235);
    en_sl = r_slm * FT(_ALNIR) * FT(0.235);
    en_sh = r_shm * FT(_ALNIR) * FT(0.235);

    # e_sl and e_sh are leaf-area-based mean radiation
    e_sl = ep_sl + en_sl;
    e_sh = ep_sh + en_sh;

    # return values
    return ratio, q_slm, q_shm, e_sl, e_sh
end


mutable struct Sphere{T} <: Intersectable{T}
    pos::SVector{3,T}
    radius::T
end

#=
\[d=-[\mathbf{l}\cdot(\mathbf{o}-\mathbf{c})]\pm\sqrt{[\mathbf{l}\cdot(\mathbf{o}-\mathbf{c})]^2-||\mathbf{o}-\mathbf{c}||^2+r^2}\]
=#

function extents(sphere::Sphere{T}) where T
    e = √(3one(T))
    [(pos[i]-e,pos[i]+e) for i=1:3]
end

function Base.intersect(ray::Ray{T}, sphere::Sphere{T}) where T
    oc = ray.pos - sphere.pos
    loc = dot(ray.dir, oc)
    loc^2 - dot(oc, oc) + sphere.radius^2 >= 0
end

function Base.intersect!(i::Intersection{T}, sphere::Sphere{T}) where T
    oc = i.ray.pos - sphere.pos
    loc = dot(i.ray.dir, oc)
    D = loc^2 - dot(oc, oc) + sphere.radius^2
    (D < 0) && return
    t = D != 0 ? -loc + sign(loc)*sqrt(D) : -loc
    (t > i.t) && return
    i.o = sphere
    i.t = t
end

normal(sphere::Sphere{T}, p::SVector{3,T}, ::Intersection{T}) where T =
    normalize(p-sphere.pos)

aabb(s::Sphere{T}) where T = AABB{T}(s.pos-s.radius,s.pos+s.radius)

export Sphere, extents, normal, aabb


##### Function to use the photoevaporation boundary in Carrera et al 2018 to separate planets:

function photoevap_boundary_Carrera2018(R::Float64, P::Float64)
    #R is the planet radius in Earth radii, P is the period in days
    #This function returns 1 if the planet is above the boundary, and 0 if the planet is below the boundary as defined by Eq. 5 in Carrera et al 2018
    Rtrans = 2.6*P^(-0.1467)
    if R > Rtrans
        return 1
    elseif R < Rtrans
        return 0
    end
end


using ECOS, SCS, Gurobi
using JuMP

mysolver = ECOSSolver()
#mysolver = SCSSolver()
#mysolver = GurobiSolver()

items  = [:Gold, :Silver, :Bronze]
values = Dict(:Gold => 5.0,  :Silver => 3.0,  :Bronze => 1.0)
weight = Dict(:Gold => 2.0,  :Silver => 1.5,  :Bronze => 0.3)

m = Model(solver=mysolver)
@defVar(m, 0 <= take_var[items] <= 1)  # Define a variable for each item
@setObjective(m, Max, sum{ values[item] * take_var[item], item in items})
@addConstraint(m, sum{ weight[item] * take_var[item], item in items} <= 3)
@time JuMP.solve(m)

println("Optimal objective: ",getObjectiveValue(m))
println("Optimal vars: ",getValue(take_var))

# take
# [  Gold] = 0.9999999680446406
# [Silver] = 0.46666670881026834
# [Bronze] = 0.9999999633898735


import Pkg; Pkg.add("DataFramesMeta")
using Random
using LinearAlgebra
using Statistics
using Optim
using DataFrames
using DataFramesMeta
using CSV
using HTTP
using GLM
# read in function to create state transitions for dynamic model
#The function didn't work for me for some reason - I tried it a couple of times but couldn't resolve the error so I'm copying the code from the file you uploaded which creates the grid#
function create_grids()

    function xgrid(theta,xval)
        N      = length(xval)
        xub    = vcat(xval[2:N],Inf)
        xtran1  = zeros(N,N)
        xtran1c = zeros(N,N)
        lcdf   = zeros(N)
        for i=1:length(xval)
            xtran1[:,i]   = (xub[i] .>= xval).*(1 .- exp.(-theta*(xub[i] .- xval)) .- lcdf)
            lcdf        .+= xtran1[:,i]
            xtran1c[:,i] .+= lcdf
        end
        return xtran1, xtran1c
    end

    zval   = collect([.25:.01:1.25]...)
    zbin   = length(zval)
    xval   = collect([0:.125:25]...)
    xbin   = length(xval)
    tbin   = xbin*zbin
    xtran  = zeros(tbin,xbin)
    xtranc = zeros(xbin,xbin,xbin)
    for z=1:zbin
        xtran[(z-1)*xbin+1:z*xbin,:],xtranc[:,:,z] = xgrid(zval[z],xval)
    end

    return zval,zbin,xval,xbin,xtran
end

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 1: reshaping the data
#:::::::::::::::::::::::::::::::::::::::::::::::::::
# load in the data
using DataFrames
using CSV
using HTTP
url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2021/master/ProblemSets/PS5-ddc/busdataBeta0.csv"
df = CSV.read(HTTP.get(url).body, DataFrame)
# create bus id variable
df = @transform(df, bus_id = 1:size(df,1))

#---------------------------------------------------
# reshape from wide to long (must do this twice be-
# cause DataFrames.stack() requires doing it one
# variable at a time)
#---------------------------------------------------
# first reshape the decision variable
dfy = @select(df, :bus_id,:Y1,:Y2,:Y3,:Y4,:Y5,:Y6,:Y7,:Y8,:Y9,:Y10,:Y11,:Y12,:Y13,:Y14,:Y15,:Y16,:Y17,:Y18,:Y19,:Y20,:RouteUsage,:Branded)
dfy_long = DataFrames.stack(dfy, Not([:bus_id,:RouteUsage,:Branded]))
rename!(dfy_long, :value => :Y)
dfy_long = @transform(dfy_long, time = kron(collect([1:20]...),ones(size(df,1))))
select!(dfy_long, Not(:variable))

# next reshape the odometer variable
dfx = @select(df, :bus_id,:Odo1,:Odo2,:Odo3,:Odo4,:Odo5,:Odo6,:Odo7,:Odo8,:Odo9,:Odo10,:Odo11,:Odo12,:Odo13,:Odo14,:Odo15,:Odo16,:Odo17,:Odo18,:Odo19,:Odo20)
dfx_long = DataFrames.stack(dfx, Not([:bus_id]))
rename!(dfx_long, :value => :Odometer)
dfx_long = @transform(dfx_long, time = kron(collect([1:20]...),ones(size(df,1))))
select!(dfx_long, Not(:variable))

# join reshaped df's back together
df_long = leftjoin(dfy_long, dfx_long, on = [:bus_id,:time])
sort!(df_long,[:bus_id,:time])

#:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 2: estimate a static version of the model using GLM
#:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
theta_hat = glm(@formula(Y ~ Odometer + Branded ), df_long, Binomial(), LogitLink())
println(theta_hat)
#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 3a,b,c: read in data for dynamic model
#:::::::::::::::::::::::::::::::::::::::::::::::::::
using DataFrames
using CSV
using HTTP
url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2021/master/ProblemSets/PS5-ddc/busdata.csv"
df = CSV.read(HTTP.get(url).body, DataFrame)
Y = Matrix(df[:,[:Y1,:Y2,:Y3,:Y4,:Y5,:Y6,:Y7,:Y8,:Y9,:Y10,:Y11,:Y12,:Y13,:Y14,:Y15,:Y16,:Y17,:Y18,:Y19,:Y20]])
# [ convert other data frame columns to matrices]
zval,zbin,xval,xbin,xtran = create_grids()
@views @inbounds function myfun(theta)   #3f?#
N=size(xtran,1) #
T=size(Y,2) #
FV=zeros(N,2,T+1) #
#Now write four nested for loops over each of the possible states#
    for t in T:-1:1 #Loop backwards over t from T +1 to 1#
    for b in 0:1#Loop over the two possible brand states {0,1}#
    for z in 1:zbin#Loop over the possible permanent route usage states (i.e. from 1 to zbin)#
    for x in 1:xbin#Loop over the possible odometer states (i.e. from 1 to xbin#
    rowtm=x+(z-1)*xbin #Create an object that marks the row of the transition matrix that we need to - This will be x + (z-1)*xbin (where x indexes the mileage bin and z indexes the route usage bin), given how the xtran matrix was constructed in the create grids() function.#
    end
    #conditional value function#
v1=theta[1]+theta[2]*xval[x]+theta[3]*b+ beta*xtran[row,:]'*FV[(z-1)*xbin+1:z*xbin,b+1,t+1]
v0=beta*xtran[1+(z-1)*xbin,:]'*FV[(z-1)*xbin+1:z*xbin,b+1,t+1]
FV[row,b+1,t]=beta*log(exp(v1)+exp(v0))
beta=0.9
#3d#
loglikelihood=0
v_0=beta*xtran[1+(:Zst-1)*xbin,:]'*FV[(:Zst-1)*xbin+1:Zst*xbin,b+1,t+1]
v_1=theta[1]+theta[2]*xval[:xst]+theta[3]*b+ beta*xtran[row,:]'*FV[(:zst-1)*xbin+1:z*xbin,b+1,t+1]
#3e#
return=-loglikelihood


#
function permutedims(A::AbstractArray, perm)
    dest = similar(A, genperm(axes(A), perm))
    permutedims!(dest, A, perm)
end

"""
    permutedims(m::AbstractMatrix)
Permute the dimensions of the matrix `m`, by flipping the elements across the diagonal of
the matrix. Differs from `LinearAlgebra`'s [`transpose`](@ref) in that the
operation is not recursive.
# Examples
```jldoctest; setup = :(using LinearAlgebra)
julia> a = [1 2; 3 4];
julia> b = [5 6; 7 8];
julia> c = [9 10; 11 12];
julia> d = [13 14; 15 16];
julia> X = [[a] [b]; [c] [d]]
2×2 Matrix{Matrix{Int64}}:
 [1 2; 3 4]     [5 6; 7 8]
 [9 10; 11 12]  [13 14; 15 16]
julia> permutedims(X)
2×2 Matrix{Matrix{Int64}}:
 [1 2; 3 4]  [9 10; 11 12]
 [5 6; 7 8]  [13 14; 15 16]
julia> transpose(X)
2×2 transpose(::Matrix{Matrix{Int64}}) with eltype Transpose{Int64, Matrix{Int64}}:
 [1 3; 2 4]  [9 11; 10 12]
 [5 7; 6 8]  [13 15; 14 16]
```
"""
permutedims(A::AbstractMatrix) = permutedims(A, (2,1))

"""
    permutedims(v::AbstractVector)
Reshape vector `v` into a `1 × length(v)` row matrix.
Differs from `LinearAlgebra`'s [`transpose`](@ref) in that
the operation is not recursive.
# Examples
```jldoctest; setup = :(using LinearAlgebra)

 """

permutedims([1, 2, 3, 4])
"""1×4 Matrix{Int64}:
 1  2  3  4"""

 V = [[[1 2; 3 4]]; [[5 6; 7 8]]]

 """2-element Vector{Matrix{Int64}}:
 [1 2; 3 4]
 [5 6; 7 8]
"""

permutedims(V)
"""1×2 Matrix{Matrix{Int64}}:
 [1 2; 3 4]  [5 6; 7 8]"""

transpose(V)

"""1×2 transpose(::Vector{Matrix{Int64}}) with eltype Transpose{Int64, Matrix{Int64}}:
 [1 3; 2 4]  [5 7; 6 8]"""


using KendallTau
using Test

X = Float64[1 0; 2 1; 3 0; 4 1; 5 10]
Y = Float64[5 5 6; 3 4 1; 4 0 4; 2 6 1; 5 7 10]
Xm = [1 0; missing 1; 2 1; 3 0; 4 1; 5 10]
Ym = [5 5 6; 1 2 3; 3 4 1; 4 0 4; 2 6 1; 5 7 10]

x1 = X[:,1]
x2 = X[:,2]
y = Y[:,1]

#= corspearman

@test corspearman(x1, y) ≈ -0.102597835208515
@test corspearman(x2, y) ≈ -0.081110710565381

@test corspearman(X, y) ≈ [-0.102597835208515, -0.081110710565381]
@test corspearman(y, X) ≈ [-0.102597835208515 -0.081110710565381]

c11 = corspearman(x1, x1)
c12 = corspearman(x1, x2)
c22 = corspearman(x2, x2)
@test c11 ≈ 1.0
@test c22 ≈ 1.0
@test corspearman(X, X) ≈ [c11 c12; c12 c22]
@test corspearman(X)    ≈ [c11 c12; c12 c22]

@test corspearman(X, Y) ==
     [corspearman(X[:,i], Y[:,j]) for i in axes(X, 2), j in axes(Y, 2)] =#

# corkendall

# Check error, handling of NaN, Inf etc
@test_throws DimensionMismatch("Vectors must have same length") corkendall([1,2,3,4], [1,2,3])
@test isnan(corkendall([1,2], [3,NaN]))
@test isnan(corkendall([1,1,1], [1,2,3]))
@test corkendall([-Inf,-0.0,Inf], [1,2,3]) == 1.0

# Test, with exact equality, some known results.
# RealVector, RealVector
@test corkendall(x1, y) == -1 / sqrt(90)
@test corkendall(x2, y) == -1 / sqrt(72)
# RealMatrix, RealVector
@test corkendall(X, y)  == [-1 / sqrt(90), -1 / sqrt(72)]
# RealVector, RealMatrix
@test corkendall(y, X)  == [-1 / sqrt(90) -1 / sqrt(72)]

# n = 78_000 tests for overflow errors on 32 bit
# Testing for overflow errors on 64bit would require n be too large for practicality
# This also tests merge_sort! since n is (much) greater than SMALL_THRESHOLD.
n = 78_000
# Test with many repeats
@test corkendall(repeat(x1, n), repeat(y, n)) ≈ -1 / sqrt(90)
@test corkendall(repeat(x2, n), repeat(y, n)) ≈ -1 / sqrt(72)
@test corkendall(repeat(X, n), repeat(y, n))  ≈ [-1 / sqrt(90), -1 / sqrt(72)]
@test corkendall(repeat(y, n), repeat(X, n))  ≈ [-1 / sqrt(90) -1 / sqrt(72)]
@test corkendall(repeat([0,1,1,0], n), repeat([1,0,1,0], n)) == 0.0

# Test with no repeats, note testing for exact equality
@test corkendall(collect(1:n), collect(1:n))          == 1.0
@test corkendall(collect(1:n), reverse(collect(1:n))) == -1.0

# All elements identical should yield NaN
@test isnan(corkendall(repeat([1], n), collect(1:n)))

# Test handling of missings
@test corkendall(Xm, Xm, skipmissing=:pairwise) == corkendall(X, X)
@test corkendall(Xm, Xm, skipmissing=:complete) == corkendall(X, X)
@test corkendall(Xm, Ym, skipmissing=:complete) == corkendall(X, Y)
@test corkendall(Xm, Ym, skipmissing=:pairwise) ≈  [-1/√90 0.4 1/√90;
                                                  -2/√154 7/√165 -1/√154]

# Test not-correct values of skipmissing
@test_throws ArgumentError corkendall(Xm)
@test_throws ArgumentError corkendall(X, skipmissing=:foo)

c11 = corkendall(x1, x1)
c12 = corkendall(x1, x2)
c22 = corkendall(x2, x2)

# RealMatrix, RealMatrix
@test corkendall(X, X) ≈ [c11 c12; c12 c22]
# RealMatrix
@test corkendall(X)    ≈ [c11 c12; c12 c22]

@test c11 == 1.0
@test c22 == 1.0
@test c12 == 3 / sqrt(20)

# Finished testing for overflow, so redefine n for speedier tests
n = 100

@test corkendall(repeat(X, n), repeat(X, n)) ≈ [c11 c12; c12 c22]
@test corkendall(repeat(X, n))               ≈ [c11 c12; c12 c22]

# All eight three-element permutations
z = [1  1  1;
     1  1  2;
     1  2  2;
     1  2  2;
     1  2  1;
     2  1  2;
     1  1  2;
     2  2  2]

@test corkendall(z)         == [1 0 1 / 3; 0 1 0; 1 / 3 0 1]
@test corkendall(z, z)      == [1 0 1 / 3; 0 1 0; 1 / 3 0 1]
@test corkendall(z[:,1], z) == [1 0 1 / 3]
@test corkendall(z, z[:,1]) == [1; 0; 1 / 3]

z = float(z)
@test corkendall(z)         == [1 0 1 / 3; 0 1 0; 1 / 3 0 1]
@test corkendall(z, z)      == [1 0 1 / 3; 0 1 0; 1 / 3 0 1]
@test corkendall(z[:,1], z) == [1 0 1 / 3]
@test corkendall(z, z[:,1]) == [1; 0; 1 / 3]

w = repeat(z, n)
@test corkendall(w)         == [1 0 1 / 3; 0 1 0; 1 / 3 0 1]
@test corkendall(w, w)      == [1 0 1 / 3; 0 1 0; 1 / 3 0 1]
@test corkendall(w[:,1], w) == [1 0 1 / 3]
@test corkendall(w, w[:,1]) == [1; 0; 1 / 3]

KendallTau.midpoint(1,10)        == 5
KendallTau.midpoint(1,widen(10)) == 5

# NaN handling

Xnan = copy(X)
Xnan[1,1] = NaN
Ynan = copy(Y)
Ynan[2,1] = NaN

# for f in (corspearman, corkendall)
for f in (corkendall,)
    @test isnan(f([1.0, NaN, 2.0], [2.0, 1.0, 3.4]))
    @test all(isnan, f([1.0, NaN], [1 2; 3 4]))
    @test all(isnan, f([1 2; 3 4], [1.0, NaN]))
    @test isequal(f([1 NaN; NaN 4]), [1 NaN; NaN 1])
    @test all(isnan, f([1 NaN; NaN 4], [1 NaN; NaN 4]))
    @test all(isnan, f([1 NaN; NaN 4], [NaN 1; NaN 4]))

    @test isequal(f(Xnan, Ynan),
                   [f(Xnan[:,i], Ynan[:,j]) for i in axes(Xnan, 2), j in axes(Ynan, 2)])
    @test isequal(f(Xnan),
                   [i == j ? 1.0 : f(Xnan[:,i], Xnan[:,j])
                    for i in axes(Xnan, 2), j in axes(Xnan, 2)])
     for k in 1:2
          @test isequal(f(Xnan[:,k], Ynan),
                        [f(Xnan[:,k], Ynan[:,j]) for i in 1:1, j in axes(Ynan, 2)])
          # TODO: fix corkendall (PR#659)
          if f === corkendall
               @test isequal(f(Xnan, Ynan[:,k]),
                             [f(Xnan[:,i], Ynan[:,k]) for i in axes(Xnan, 2)])
          else
               @test isequal(f(Xnan, Ynan[:,k]),
                             [f(Xnan[:,i], Ynan[:,k]) for i in axes(Xnan, 2), j in 1:1])
          end
     end

     # Wrong dimensions
     @test_throws DimensionMismatch f([1], [1, 2])
     @test_throws DimensionMismatch f([1], [1 2; 3 4])
     @test_throws DimensionMismatch f([1 2; 3 4], [1])
     @test_throws ArgumentError f([1 2; 3 4:4 6], [1 2; 3 4])

end


function AUC_ROC(Ytrue, scores)
     perm = sortperm(scores)

     roc_y = Ytrue[perm]
     stack_x = cumsum(roc_y)/sum(roc_y)
     stack_y = cumsum(! roc_y)/sum(! roc_y)
     N = length(roc_y)

     auc = sum( (stack_x[2:N] - stack_x[1:N-1]) .* stack_y[2:N] )
     return auc
end


module NLP1

using JuMP, BARON, Test

m = Model(BARON.Optimizer)
ub = [6,4]
@variable(m, 0 ≤ x[i=1:2] ≤ ub[i])

@constraint(m, x[1]*x[2] ≤ 4)

@objective(m, Min, -x[1] - x[2])

optimize!(m)

@test isapprox(value(x[1]), 6, rtol=1e-6)
@test isapprox(value(x[2]), 2/3, rtol=2e-6)
@test isapprox(objective_value(m), -20/3, rtol=1e-6)

end # module


using SolverPeeker
using Test
using JuMP
using MathOptInterface
const MOI = MathOptInterface

# Model with linear constraints and objective
model1 = Model()
@variable(model1, x >= 5, Int)
@variable(model1, y <= 5)
@variable(model1, 0 <= z <= 50)
@objective(model1, Max, 3x+5y-2z+2)
@constraint(model1, x+2y <= 45)
@constraint(model1, z - y >= 4)
@constraint(model1, 0 <= 2z <= 7)
@constraint(model1, x+y+2z <= 55)

function fill_small_test_model!(model::JuMP.Model)
    # The model does not need to make sense, just use many different features.
    @variable(model, a[1:5] >= 0, Int)
    @variable(model, b[6:10], Bin)
    @variable(model, c[1:3] == 0)
    @variable(model, 10 <= d[1:3] <= 20)
    @constraint(model, con1, sum(a) + sum(b) <= 5)
    @constraint(model, con2, sum(b) >= 3)
    @constraint(model, con3, sum(d[1:2]) >= 5)
    @constraint(model, con4, sum(d) <= (sum(c) + 10))
    @objective(model, Max, sum(a) - sum(b) + sum(d))
    return model
end

@testset "lprep.jl" begin
    #lpmod = lpmodel(model1)
    #show(lpmodel)
    lprep = LPRep(model1)
    show(lprep.A)
    print("\n\n")
    #show(typeof(model1))
    #show(lprep.obj_constant)
    print("\n\n")
    #show(typeof(lpmodel.sense))
    #cis_i = MOI.get(lpmodel, MOI.ListOfConstraintIndices{MathOptInterface.SingleVariable,MathOptInterface.Integer}())
end


# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
# ==============================================================================
#
#    Coordinate transformations related with the geodetic and geocentric
#    coordinates.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# References
# ==============================================================================
#
#   [1] Vallado, D. A (2013). Fundamentals of Astrodynamics and Applications.
#       Microcosm Press, Hawthorn, CA, USA.
#
#   [2] ESA Navipedia: http://www.navipedia.net/
#
#   [3] mu-blox ag (1999). Datum Transformations of GPS Positions. Application
#       Note.
#
#   [4] ISO TC 20/SC 14 N (2011). Geomagnetic Reference Models.
#
#   [5] Borkowski, K. M (1987). Transformation of geocentric to geodetic
#       coordinates without approximations. Astrophysics and Space Science, vol.
#       139, pp. 1-4.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

export ecef_to_geodetic, geodetic_to_ecef
export geocentric_to_geodetic, geodetic_to_geocentric

"""
    ecef_to_geodetic(r_e::AbstractVector)

Convert the vector `r_e` [m] represented in the Earth-Centered, Earth-Fixed
(ECEF) reference frame into Geodetic coordinates (WGS-84).

!!! info
    The algorithm is based in **[3]**.

# Returns

* Latitude [rad].
* Longitude [rad].
* Altitude [m].

# Reference

- **[3]**: mu-blox ag (1999). Datum Transformations of GPS Positions.
    Application Note.
"""
function ecef_to_geodetic(r_e::AbstractVector)
    # Auxiliary variables.
    X = r_e[1]
    Y = r_e[2]
    Z = r_e[3]

    # Auxiliary variables.
    p = √(X^2 + Y^2)
    θ = atan(Z * a_wgs84, p * b_wgs84)
    e_wgs84² = e_wgs84^2
    sin_θ, cos_θ = sincos(θ)

    # Compute Geodetic.
    lon = atan(Y, X)
    lat = atan(
        Z + el_wgs84^2 * b_wgs84 * sin_θ^3,
        p -   e_wgs84² * a_wgs84 * cos_θ^3
    )

    sin_lat, cos_lat = sincos(lat)

    N = a_wgs84/√(1 - e_wgs84² * sin_lat^2)

    # Avoid singularity if we are near the poles (~ 1 deg according to [1,
    # p.172]). Note that `cosd(1) = -0.01745240643728351`.
    if !(-0.01745240643728351 < cos_lat < 0.01745240643728351)
        h = p/cos_lat - N
    else
        h = Z/sin_lat - N*(1 - e_wgs84²)
    end

    return lat, lon, h
end

"""
    geodetic_to_ecef(lat::Number, lon::Number, h::Number)

Convert the latitude `lat` [rad], longitude `lon` [rad], and altitude `h` \\[m]
(WGS-84) into a vector represented on the Earth-Centered, Earth-Fixed (ECEF)
reference frame.

!!! info
    The algorithm is based in **[3]**.

# Reference

- **[3]**: mu-blox ag (1999). Datum Transformations of GPS Positions.
    Application Note.
"""
function geodetic_to_ecef(lat::Number, lon::Number, h::Number)
    # Auxiliary variables.
    sin_lat, cos_lat = sincos(lat)
    sin_lon, cos_lon = sincos(lon)

    # Radius of curvature [m].
    N = a_wgs84/√(1 - e_wgs84^2 * sin_lat^2)

    # Compute the position in ECEF frame.
    return SVector(
        (                        N + h) * cos_lat * cos_lon,
        (                        N + h) * cos_lat * sin_lon,
        ((b_wgs84 / a_wgs84)^2 * N + h) * sin_lat
    )
end

"""
    geocentric_to_geodetic(ϕ_gc::Number, r::Number)

Compute the geodetic latitude and altitude (WGS-84) from the geocentric latitude
`ϕ_gc` (-π/2, π/2) [rad] and radius `r` [m]. Notice that the longitude is the
same in both geocentric and geodetic coordinates.

!!! info
    The algorithm is based in **[5]**.

# Returns

* Geodetic latitude [rad].
* Altitude above the reference ellipsoid (WGS-84) [m].

# References

- **[5]** Borkowski, K. M (1987). Transformation of geocentric to geodetic
    coordinates without approximations. Astrophysics and Space Science, vol.
    139, pp. 1-4.
"""
function geocentric_to_geodetic(ϕ_gc::Number, r::Number)
    # Obtain the `z` component and the equatorial component `re`.
    sin_ϕ_gc, cos_ϕ_gc = sincos(ϕ_gc)
    re = r * cos_ϕ_gc
    z  = r * sin_ϕ_gc
    sign_z = z == 0 ? +1 : sign(z)

    # Auxiliary variables.
    a  = a_wgs84
    a² = a^2
    b  = sign_z * b_wgs84
    b² = b^2

    # Compute the parameters.
    E  = (b * z - (a² - b²)) / (a * re)
    E² = E^2
    F  = (b * z + (a² - b²)) / (a * re)
    P  = (4 / 3) * (E * F + 1)
    Q  = 2 * (E² - F^2)
    D  = P^3 + Q^2

    if D ≥ 0
        aux = √D
        v = (aux - Q)^(1/3) - (Q + aux)^(1/3)
    else
        aux = √(-P)
        v = 2 * aux * cos(acos(Q / (P * aux)) / 3)
    end

    G = (√(E² + v) + E)/2

    # NOTE: Reference [5] appears to have an error in Eq. (13), where we must
    # have G^2 instead of G inside the square root. The correct version can be
    # seen here:
    #
    #   https://www.astro.uni.torun.pl/~kb/Papers/geod/Geod-BG.htm
    #
    t = √(G^2 + (F - v * G) / (2 * G - E)) - G

    # Compute the geodetic latitude and altitude.
    ϕ_gd = atan(a * (1 - t^2)/(2b * t))
    sin_ϕ_gd, cos_ϕ_gd = sincos(ϕ_gd)
    h = (re - a * t) * cos_ϕ_gd + (z - b) * sin_ϕ_gd

    return ϕ_gd, h
end

"""
    geodetic_to_geocentric(ϕ_gd::Number, h::Number)

Compute the geocentric latitude and radius from the geodetic latitude `ϕ_gd`
(-π/2, π/2) [rad] and height above the reference ellipsoid `h` \\[m] (WGS-84).
Notice that the longitude is the same in both geocentric and geodetic
coordinates.

!!! info
    The algorithm is based in **[4]**(p. 3).

# Returns

* Geocentric latitude [rad].
* Radius from the center of the Earth [m].

# References

- **[4]** ISO TC 20/SC 14 N (2011). Geomagnetic Reference Models.
"""
function geodetic_to_geocentric(ϕ_gd::Number, h::Number)
    # Auxiliary variables to decrease computational burden.
    sin_ϕ_gd, cos_ϕ_gd = sincos(ϕ_gd)
    sin²_ϕ_gd = sin_ϕ_gd^2
    e²_wgs84  = e_wgs84^2

    # Radius of curvature in the prime vertical [m].
    N    = a_wgs84/√(1 - e²_wgs84 * sin²_ϕ_gd )

    # Compute the geocentric latitude and radius from the Earth center.
    ρ    = (N + h) * cos_ϕ_gd
    z    = (N * (1 - e²_wgs84) + h) * sin_ϕ_gd
    r    = √(ρ^2 + z^2)
    ϕ_gc = asin(z / r)

    return ϕ_gc, r
end


"""
A Tiler is an iterator that, for each iteration, returns a tuple of:

- the `x`/`y` point of the center of each tile in a set of tiles that divide up a rectangular space such as a page into rows and columns (relative to current 0/0)

- the number of the tile

    tiles = Tiler(areawidth, areaheight, nrows, ncols, margin=20)

where `width`, `height` is the dimensions of the area to be tiled, `nrows`/`ncols`
is the number of rows and columns required, and `margin` is applied to all four
edges of the area before the function calculates the tile sizes required.

    tiles = Tiler(1000, 800, 4, 5, margin=20)
    for (pos, n) in tiles
    # the point pos is the center of the tile
    end

You can access the calculated tile width and height like this:

    tiles = Tiler(1000, 800, 4, 5, margin=20)
    for (pos, n) in tiles
      ellipse(pos.x, pos.y, tiles.tilewidth, tiles.tileheight, :fill)
    end
"""
type Tiler
  width::Real
  height::Real
  tilewidth::Real
  tileheight::Real
  nrows::Int
  ncols::Int
  margin::Real
  function Tiler(areawidth, areaheight, nrows::Int, ncols::Int; margin=10)
      tilewidth  = (areawidth - 2margin)/ncols
      tileheight = (areaheight - 2margin)/nrows
      new(areawidth, areaheight, tilewidth, tileheight, nrows, ncols, margin)
  end
end

function Base.start(pt::Tiler)
# return the initial state
  x = -(pt.width/2)  + pt.margin + (pt.tilewidth/2)
  y = -(pt.height/2) + pt.margin + (pt.tileheight/2)
  return (Point(x, y), 1)
end

function Base.next(pt::Tiler, state)
  # Returns the item and the next state
  # state[1] is the Point
  x = state[1].x
  y = state[1].y
  # state[2] is the tilenumber
  tilenumber = state[2]
  x1 = x + pt.tilewidth
  y1 = y
  if x1 > (pt.width/2) - pt.margin
    y1 += pt.tileheight
    x1 = -(pt.width/2) + pt.margin + (pt.tilewidth/2)
  end
  return ((Point(x, y), tilenumber), (Point(x1, y1), tilenumber + 1))
end

function Base.done(pt::Tiler, state)
  # Tests if there are any items remaining
  state[2] > (pt.nrows * pt.ncols)
end


using MultipleViewGeometry, Test
using StaticArrays, Random

# Vectors are scaled so that the last component is unity.
v = [2,2,2]
@test 𝑛(v) == [1.0, 1.0, 1.0]

v = Point2DH([2,2,2])
@test 𝑛(v) == Point2DH([1.0, 1.0, 1.0])

# Vectors which represent points at infinity are unchanged.
v = [2,2,0]
@test 𝑛(v) == [2,2,0]

v = Point2DH([2,2,0])
@test 𝑛(v) == Point2DH([2,2,0])

v = [1, 2, 3]
@test vec2antisym(v) == @SMatrix [ 0 -3  2;
                                   3  0 -1;
                                  -2  1  0]

 v = Point2DH([1, 2, 3])
 @test vec2antisym(v) == @SMatrix [ 0 -3  2;
                                    3  0 -1;
                                   -2  1  0]


mg!(st::SolveState;kw...) = iterate!(st,Vcycle!;mxiter=32,kw...)
function mg(A::AbstractMatrix,b::AbstractVector;kw...)
    x=zero(b)
    return x,mg!(mg_state(A,x,b);kw...)
end
mg_state(A,x,b) = fill_children!(SolveState(A,x,residual(A,x,b)))

@inline Vcycle!(st::SolveState;kw...) = Vcycle!(st,st.child;kw...)
@inline Vcycle!(fine::SolveState,coarse::Nothing;kw...) = GS!(fine;kw...)
function Vcycle!(fine::SolveState,coarse::SolveState;
                 precond!::Function=(st;kw...)->GS!(st,inner=0),smooth!::Function=GS!,kw...)
    # set up & solve coarse recursively
    precond!(fine;kw...)
    restrict!(coarse.r,fine.r)
    fill!(coarse.x,0.)
    Vcycle!(coarse;precond!,smooth!,kw...)
    # correct & solve fine
    prolongate!(fine.ϵ,coarse.x)
    increment!(fine)
    smooth!(fine;kw...)
end

restrict!(a,b) = @loop a[I] = sum(@inbounds(b[J]) for J ∈ up(I))
prolongate!(a,b) = @loop a[I] = b[down(I)]

@inline up(I::CartesianIndex{N},a=0) where N = (2I-2oneunit(I)):(2I-oneunit(I)-δ(a,N))
@inline down(I) = CartesianIndex((I+2oneunit(I)).I .÷2)

function fill_children!(st::SolveState)
    if isdivisible(st) && isnothing(st.child)
        st.child = create_child(st.A.L,eltype(st.x))
        fill_children!(st.child)
    end
    return st
end

function create_child(b::AbstractArray{T,N},xT::Type) where {T,N}
    dims = 1 .+ Base.front(size(b)) .÷ 2
    a = zeros(T,dims...,N-1)
    x = FieldVector(zeros(xT,dims))
    for i ∈ 1:N-1, I ∈ x.R
        a[I,i] = 0.5sum(@inbounds(b[J,i]) for J ∈ up(I,i))
    end
    SolveState(Poisson(a),x,zero(x))
end

@inline isdivisible(N::Int) = mod(N,2)==0 && N>2
@inline isdivisible(st::SolveState) = all(size(st.x.R) .|> isdivisible)


using FastTransforms, Compat
using Compat.Test

import FastTransforms: normalizecolumns!, maxcolnorm

@testset "Slow plan" begin
    N = round.([Int],10 .^ range(1,stop=2.5,length=10))

    t = zeros(length(N))
    err = zeros(length(N))

    Nr = 2

    j = 1
    for n in N
        nrms = zeros(Nr)
        for kk = 1:Nr
            A = sphrandn(Float64, n+1, n+1)
            normalizecolumns!(A)
            Ac = copy(A)
            B = zero(A)
            SP = SlowSphericalHarmonicPlan(A)
            mul!(B, SP, A)
            fill!(A, 0.0)
            t[j] += @elapsed At_mul_B!(A, SP, B)
            nrms[kk] = maxcolnorm(A - Ac)
        end
        t[j] /= Nr
        err[j] = Statistics.mean(nrms)
        println("At a bandlimit of ",n,", the maximum 2-norm in the columns over ",Nr," trials is: ",err[j])
        @test err[j] < 10eps()*sqrt(n)
        j+=1
    end
end


#=
An example of Type unions
=#

mutable struct Point2D
    x::Float64
    y::Float64
end

mutable struct Vector2D
    x::Float64
    y::Float64
end

VecOrPoint2D = Union{Vector2D,Point2D}

import Base.+

function +(u::VecOrPoint2D, v::VecOrPoint2D)
    @show typeof(u),typeof(v)
    if typeof(u) == Vector2D && typeof(v) == Vector2D
        Vector2D(u.x+v.x,u.y+v.y)
    else
        Point2D(u.x+v.x,u.y+v.y)
    end
end

p = Point2D(2,4)
q = Point2D(4,3)
v = Vector2D(3,2)
w = Vector2D(-1,2)

@show typeof(p+q)
@show typeof(p+v)
@show typeof(w+p)
@show typeof(v+w)


using .Clustering
using .Plots
using .Colors
using .LinearAlgebra

function clustercovariance(Σ²)
    issymmetric(Σ²) || (Σ² = cov(Σ²'))
    Dr = 1.0.-abs.(Σ²)
    if !issymmetric(Dr)
        @warn "Correlation distance matrix is not symmetric, so not clustering"
    end
    Clustering.hclust(Dr; linkage=:average, branchorder=:optimal)
end

"""
    covarianceimage(f, F; [palette=[:cornflowerblue, :crimson, :forestgreen], colormode=:top, colorbargrad=:binary, donames=true, kwargs...])
Plot the covariance matrix of the columns of an `Array`, coloring elements by their contribution to each of the top 3 principal components.
Either provide as positional arguments a vector `f` of N row names and an N×_ matrix `F`, or just a `Catch22.FeatureArray`.

# Keyword Arguments
- `palette`: a vector containing a color for each principal component.
- `colorbargrad`: the color gradient for the colorbar.
- `donames`: override `f` and display no tick labels.
- `colormode`: how to color the covariance matrix. `:raw` gives no coloring by principal components, `:top` is a combination of the top three PC colors and `:all` is a combination of all PC colors, where PCN = :black if N > length(palette).
- `docluster`: whether to perform clustering on the covariance matrix.
- `kwargs...`: other `Plots` attributes.
"""
covarianceimage
@userplot CovarianceImage
Plots.@recipe function f(g::CovarianceImage; palette=[:cornflowerblue, :crimson, :forestgreen], colormode=:top, colorbargrad=:binary, donames=true, docluster=true, verbose=true, dendrogram=false)
    if g.args[1] isa AbstractFeatureArray || g.args[1] isa AbstractDimArray
        f, Σ² = string.(getdim(g.args[1], 1)), g.args[1] |> Array
    elseif length(g.args) == 2 && g.args[2] isa AbstractMatrix
        f, Σ² = g.args[1], g.args[2] |> Array
    else
        @error "Incorrect arguments; give row names and a matrix or an annotated DimArray"
    end

    issymmetric(Σ²) || (Σ² = cov(Σ²'))

    any(diag(Σ²) .== 0) && @warn "Covariance matrix is not positive definite, which may cause an error"
    linecolor --> nothing
    if docluster == true
        idxs = clustercovariance(Σ²).order
    elseif docluster isa Union{AbstractVector, Tuple}
        idxs = docluster # Precomputed indices
    else
        idxs = 1:size(Σ², 1)
    end
    Σ̂² = Σ²[idxs, idxs]
    A = abs.(Σ̂²)./max(abs.(Σ̂²)...)
    f̂ = f[idxs]
    N = min(length(palette), size(Σ̂², 1))
    if colormode == :raw # * Don't color by PC's
        H = abs.(Σ̂²)
        colorbar --> true
    else
        λ = (eigvals∘Symmetric∘Array)(Σ̂²)
        λi = sortperm(abs.(λ), rev=true)
        λ = λ[λi]
        P = (eigvecs∘Symmetric∘Array)(Σ̂²)[:, λi] # Now sorted by decreasing eigenvalue norm
        vidxs = sortperm(abs.(P[:, 1]), rev=true)
        verbose && isnothing(printstyled("Feature weights:\n", color=:red, bold=true)) && display(vcat(hcat("Feature", ["PC$i" for i ∈ 1:N]...) , hcat(f̂[vidxs], round.(P[vidxs, 1:N], sigdigits=3))))
        P = abs.(P)
        if colormode == :top # * Color by the number of PC's given by the length of the color palette
            P = P[:, 1:N]
            P̂ = P.^2.0./sum(P.^2.0, dims=2)
            # Square the loadings, since they are added in quadrature. Maybe not a completely faithful representation of the PC proportions, but should get the job done.
            𝑓′ = parse.(Colors.XYZ, palette[1:N]);
        elseif colormode == :all # * Color by all PC's. This can end up very brown
            Σ̂′² = Diagonal(abs.(λ))
            P̂ = P.^2.0./sum(P.^2.0, dims=2)
            p = fill(:black, size(P, 2))
            p[1:N] = palette[1:N]
            𝑓′ = parse.(Colors.XYZ, p);
            [𝑓′[i] = Σ̂′²[i, i]*𝑓′[i] for i ∈ 1:length(𝑓′)]
        end
        𝑓 = Vector{eltype(𝑓′)}(undef, size(P̂, 1))
        try # Load colors by PC weights
            𝑓 = P̂*𝑓′
        catch
            # Equivalent but slower
            @info "Iterating to load covariances"
            for ii ∈ 1:length(𝑓)
                𝑓[ii] = sum([P̂[ii, jj]*𝑓′[jj] for jj ∈ 1:length(𝑓′)])
            end
        end

        H = Array{Colors.XYZA}(undef, size(Σ̂²))
        for (i, j) ∈ Tuple.(CartesianIndices(H)) # Apply the correlations as transparencies
            J = (𝑓[i] + 𝑓[j])/2
            H[i, j] = Colors.XYZA(J.x, J.y, J.z, A[i, j])
        end
        H = convert.((Colors.RGBA,), H)
    end
    @series begin
        seriestype := :heatmap
        if colormode == :raw
            seriescolor --> colorbargrad
        end
        if backend() == Plots.GRBackend() # For some reason, GR does heatmaps differently
            xs = 0.5:1:size(H, 1)+0.5
        else
            xs = 1:size(H, 1)
        end
        (xs, xs, H)
    end
    # Plot dummy data and set attributes
    @series begin
        if dendrogram
            colorbar := false
        else
            colorbar --> true
        end
        seriestype := :scatter
        markersize := 0.0
        label := nothing
        legend := :none
        marker_z := [0, max(abs.(Σ²)...)]
        markercolor := colorbargrad
        (zeros(2), zeros(2))
    end
    for i ∈ 1:N
        @series begin
            seriestype := :shape
            if colormode != :raw
                label := "PC$i"
                legend --> :topright
                line_width := 20
            else
                label := nothing
                legend := nothing
            end
            colorbar_title --> "\$|\\Sigma^2|\$"
            colorbar_titlefontsize := 14
            xticks := :none
            size --> (800, 400)
            yflip --> true
            lims := (0.5, size(H, 1)+0.5)
            aspect_ratio := :equal
            legendfontsize := 8
            if donames
                yticks := (1:size(H, 1), replace.(string.(f̂), (r"_" => s"\\_",)))
            else
                yticks := nothing
            end
            grid := :none
            framestyle := :box
            seriescolor := palette[i]
            (Shape([0.0;], [0.0;]))
        end
    end
    if dendrogram
        @series begin
            inset_subplots := (1, bbox(0.895, 0.0, 0.1, 1.0))
            subplot := 2
            orientation := :horizontal
            yticks := nothing
            showaxis := false
            colorbar := nothing
            axis := nothing
            useheight := false
            yflip := true
            x := clustercovariance(Σ²)
        end
    end
end


using OrdinaryDiffEq
using DiffEqCallbacks
using Surrogates
using Conda
using PyCall
using Random
using Plots
using BenchmarkTools;

env = Conda.ROOTENV
Conda.pip_interop(true, env)
Conda.pip("install", "mogp-emulator");

random = pyimport("random")
np = pyimport("numpy")
mogp = pyimport("mogp_emulator");

Random.seed!(123)
random.seed(123)
np.random.seed(123);

function sir_ode!(du,u,p,t)
    (S,I,R) = u
    (β,γ) = p
    @inbounds begin
        du[1] = -β*S*I
        du[2] = β*S*I - γ*I
        du[3] = γ*I
    end
    nothing
end;

tmax = 10000.0
tspan = (0.0,tmax)
δt = 1.0;

n_train = 50 # Number of training samples
n_test = 1000; # Number of test samples

# Parameters are β, γ
lb = [0.00005, 0.1]
ub = [0.001, 1.0];

N = 1000.0
u0 = [990.0,10.0,0.0]
p = [0.0005,0.25]
prob_ode = ODEProblem(sir_ode!,u0,tspan,p);

sampler = LatinHypercubeSample();

θ = Surrogates.sample(n_train,lb,ub,sampler);

logit = (x) -> log(x/(1-x))
invlogit = (x) -> exp(x)/(exp(x)+1.0)
cb_ss = TerminateSteadyState()
logit_final_size = function(z)
  prob = remake(prob_ode;p=z)
  sol = solve(prob, ROS34PW3(),callback=cb_ss)
  fsp = sol[end][3]/N
  logit(fsp)
end;

lfs = logit_final_size.(θ);

gp = mogp.GaussianProcess(θ, lfs, nugget="fit");

gp = mogp.fit_GP_MAP(gp, n_tries=100);

lfs_train_pred = gp.predict(θ);

scatter(invlogit.(lfs),
        invlogit.(lfs_train_pred["mean"]),
        xlabel = "Model final size",
        ylabel = "Surrogate final size",
        legend = false,
        title = "Training set")

θ_test = sample(n_test,lb,ub,sampler)
lfs_test = logit_final_size.(θ_test)
lfs_test_pred = gp.predict(θ_test);

scatter(invlogit.(lfs_test),
        invlogit.(lfs_test_pred["mean"]),
        xlabel = "Model final size",
        ylabel = "Surrogate final size",
        legend = false,
        title = "Test set")

β_grid = collect(lb[1]:0.00001:ub[1])
θ_eval = [[βᵢ,0.25] for βᵢ in β_grid]
lfs_eval = gp.predict(θ_eval)
fs_eval = invlogit.(lfs_eval["mean"])
fs_eval_uc = invlogit.(lfs_eval["mean"] .+ 1.96 .* sqrt.(lfs_eval["unc"]))
fs_eval_lc = invlogit.(lfs_eval["mean"] .- 1.96 .* sqrt.(lfs_eval["unc"]))
plot(β_grid,
     fs_eval,
     xlabel = "Infectivity parameter, β",
     ylabel = "Final size",
     label = "Model")
plot!(β_grid,
      invlogit.(logit_final_size.(θ_eval)),
      ribbon = (fs_eval .- fs_eval_lc, fs_eval_uc - fs_eval),
      label = "Surrogate",
      legend = :right)

γ_grid = collect(lb[2]:0.001:ub[2])
θ_eval = [[0.001,γᵢ] for γᵢ in γ_grid]
lfs_eval = gp.predict(θ_eval)
fs_eval = invlogit.(lfs_eval["mean"])
fs_eval_uc = invlogit.(lfs_eval["mean"] .+ 1.96 .* sqrt.(lfs_eval["unc"]))
fs_eval_lc = invlogit.(lfs_eval["mean"] .- 1.96 .* sqrt.(lfs_eval["unc"]))
plot(γ_grid,
     fs_eval,
     xlabel = "Recovery rate, γ",
     ylabel = "Final size",
     label = "Model")
plot!(γ_grid,
      invlogit.(logit_final_size.(θ_eval)),
      ribbon = (fs_eval .- fs_eval_lc, fs_eval_uc - fs_eval),
      label = "Surrogate")

obs = logit_final_size(p)
invlogit(obs)

hm = mogp.HistoryMatching(gp=gp,
                          obs=obs,
                          coords=np.array(θ_test),
                          threshold=3.0);

nroy_points = hm.get_NROY() .+ 1
length(nroy_points),n_test

x = [θᵢ[1] for θᵢ in θ_test]
y = [θᵢ[2] for θᵢ in θ_test]
l = @layout [a b]
pl1 = histogram(x[nroy_points],legend=false,xlim=(lb[1],ub[1]),bins=lb[1]:0.00005:ub[1],title="NROY values for β")
vline!(pl1,[p[1]])
pl2 = histogram(y[nroy_points],legend=false,xlim=(lb[2],ub[2]),bins=lb[2]:0.05:ub[2],title="NROY values for γ")
vline!(pl2,[p[2]])
plot(pl1, pl2, layout = l)

@benchmark logit_final_size(p)

@benchmark gp.predict(p)


"""
    tsz(νGHz; units = "cmb", Tcmb = 2.7255)

Spectrum of the thermal Sunyaev-Zeldovich effect.

*Reference*: Equation (V) in Appendix of Zeldovich, Sunyaev, Astrophys. Space Sci. 4, 301 (1969)

# Arguments
- `νGHz::Real`: frequency in units of GHz.

# Optional keyword arguments
- `units::String="cmb"`: units of the spectrum. For Rayleigh-Jeans temperature (brightness temperature) units, `units = "rj"`. The default is the CMB units.
- `Tcmb::Real=2.7255`: present-day temperature of the CMB in units of Kelvin.
"""
function tsz(νGHz::Real; units = "cmb", Tcmb = 2.7255)
    hp = 6.62607015e-34
    kB = 1.380649e-23
    x = hp * νGHz * 1e9 / kB / Tcmb
    g = (exp(x) - 1)^2 / x^2 / exp(x)
    tsz = x * coth(x / 2) - 4
    if units == "rj"
        return tsz / g
    else
        return tsz
    end
end

"""
    dust1(νGHz; Td = 19.6, βd = 1.6, νd = 353, units = "cmb", Tcmb = 2.7255)

Spectrum of 1-component modified black-body thermal dust emission. The output is normalized to unity at `νGHz = νd`.

# Arguments
- `νGHz::Real`: frequency in units of GHz.

# Optional keyword arguments
- `Td::Real=19.6`: dust temperature  in units of Kelvin.
- `βd::Real=1.6`: dust emissivity index.
- `νd::Real=353`: frequency at which the output is normalized to unity.
- `units::String="cmb"`: units of the spectrum. For Rayleigh-Jeans temperature (brightness temperature) units, `units = "rj"`. The default is the CMB units.
- `Tcmb::Real=2.7255`: present-day temperature of the CMB in units of Kelvin.
"""
function dust1(
    νGHz::Real;
    Td = 19.6,
    βd = 1.6,
    νd = 353,
    units = "cmb",
    Tcmb = 2.7255,
)
    hp = 6.62607015e-34
    kB = 1.380649e-23
    x(ν) = hp * ν * 1e9 / kB / Tcmb
    g(ν) = (exp(x(ν)) - 1)^2 / x(ν)^2 / exp(x(ν))
    xd(ν) = hp * ν * 1e9 / kB / Td
    mbb_rj = (νGHz / νd)^(βd + 1) * (exp(xd(νd)) - 1) / (exp(xd(νGHz)) - 1)
    if units == "rj"
        return mbb_rj
    else
        return g(νGHz) / g(νd) * mbb_rj
    end
end

"""
    synch(νGHz; βs = -3, νs = 23, Cs = 0, νC = 40, units = "cmb", Tcmb = 2.7255)

Spectrum of synchrotron emission. The output is normalized to unity at `νGHz = νs`.

# Arguments
- `νGHz::Real`: frequency in units of GHz.

# Optional keyword arguments
- `βs::Real=-3`: synchrotron power-law index.
- `νs::Real=23`: frequency at which the output is normalized to unity.
- `Cs::Real=0`: curvature of the synchrotron spectrum.
- `νC::Real=40`: pivot frequency for curvature of the synchrotron spectrum.
- `units::String="cmb"`: units of the spectrum. For Rayleigh-Jeans temperature (brightness temperature) units, `units = "rj"`. The default is the CMB units.
- `Tcmb::Real=2.7255`: present-day temperature of the CMB in units of Kelvin.
"""
function synch(
    νGHz::Real;
    βs = -3,
    νs = 23,
    Cs = 0,
    νC = 40,
    units = "cmb",
    Tcmb = 2.7255,
)
    hp = 6.62607015e-34
    kB = 1.380649e-23
    x(ν) = hp * ν * 1e9 / kB / Tcmb
    g(ν) = (exp(x(ν)) - 1)^2 / x(ν)^2 / exp(x(ν))
    synch_rj = (νGHz / νs)^(βs + Cs * log(νGHz / νC))
    if units == "rj"
        return synch_rj
    else
        return g(νGHz) / g(νs) * synch_rj
    end
end


"""
sample_theta()

Description of sample_theta
------------------------------
Returns θ in radians Uniformly distributed. θ is the polar angle. See Distributions.Uniform

"""
function sample_theta()
    return acos(1 - 2 * rand(Uniform())) ## polar angle in radians
end

"""
    sample_phi()

Description of sample_phi
------------------------------
Returns ϕ in radians Uniformly distributed. phi is the azimuthal angle. See Distributions.Uniform

Sampling equation is:
    + First θ1 is sampled uniformly from 2π uniform pdf.
    + Second angle θdif (the angle between the two electrons) is sampled from pdf = 0.5 - 0.5x. This pdf represents the 1 - cos(θ) angular distribution.
        It is sampled through inverse CDF method. Integrating pdf and finding the inverse function we get ``` x = 1 - 2\\sqrt{1 - \\gamma}```
    + θ2 is determined as θ2 = θdif + θ1.

"""
function sample_phi()
    return 2*π*rand(Uniform()) ## azimuthal angle in radians
end

"""
    sample_theta_dif()

Description of sample_theta_dif
------------------------------
Returns θdif in radians distrtibuted according to pdf = 0.5 - 0.5x (where x = cos(θdif)). θdif is the angle between the two electrons.

Sampling equation is:
    + θdif (the angle between the two electrons) is sampled from pdf = 0.5 - 0.5x. This pdf represents the 1 - cos(θ) angular distribution.
        It is sampled through inverse CDF method. Integrating pdf and finding the inverse function we get ``` x = 1 - 2\\sqrt{1 - \\gamma}```
    + returned angle is the acos(x)

"""
function sample_theta_dif()
    a = 1.
    b = -2.
    c = -3 + 4*rand(Uniform())
    cosθdif = solve_quadratic(a,b,c)[1]
    return acos(cosθdif)
end

"""
#### function ```sample_energies(df::DataFrame, thickness = 0.001)```

    Description of ```sample_energies```
    ------------------------------
Returns a tuple of energies (E1, E2) given by the distribution specified in dataframe ```df```.

```df``` must have fields: ```:E1, :minE, :maxE, :a, :b, :cdf (given by get_cdf!(df))```

Sampling is done in the following way.
E2 is sampled using inverse CDF method - where cdf is obtained from ```get_cdf!(df)```.
E1 is dependent on E2 as it can only be sampled in the given row where E2 was chosen.
E1 is sampled as a random uniform number between (row of E1, next row).

It takes roughly 100s to sample 1e5 energies.

"""
function sample_energies(df::DataFrame, thickness = 0.001)
    gamma = rand(Uniform(0.0, maximum(df.cdf))) # random uniform number from 0 to maximum of cdf column (it is not exactly 1.0 due to discrete numbers used in PC)

    r_id = findfirst(x -> x .>= gamma,df.cdf)

    r_id > 1 ? Pᵢ₋₁ = df[r_id-1, 8] : Pᵢ₋₁ = 0

    df_Pᵢ= df[r_id, :]

    a = 0.5* df_Pᵢ.a * thickness
    b = df_Pᵢ.b * thickness
    c = Pᵢ₋₁ - gamma - a*df_Pᵢ.minE^2 - b*df_Pᵢ.minE

    x1, x2 = solve_quadratic(a, b, c)

    if (df_Pᵢ.minE < x1) && (df_Pᵢ.maxE > x1)
        return rand(Uniform(df_Pᵢ.E1, df_Pᵢ.E1+thickness)), x1  #E1 is sampled as uniform number between steps,
                                                                #E2 is sampled by CDF
    elseif (df_Pᵢ.minE < x2) && (df_Pᵢ.maxE > x2)
        return rand(Uniform(df_Pᵢ.E1, df_Pᵢ.E1+thickness)), x2
    else
        sample_energies(df)
    end

end

function solve_quadratic(a, b, c)
    d  = sqrt(b^2 - 4*a*c)
    return (-b - d) / (2*a), (-b + d) / (2*a)
end


@with_kw struct LQR
    γ::Float64 = 1.0
end

discount(mdp::LQR) = mdp.γ

generate_start_state(mdp::LQR) = rand(Normal(0.3,0.1))

function transition(mdp::LQR, s::Float64, a::Float64)
    # NOTE: Truncated to prevent going off to infinity with poor policies
    return truncated(Normal(s + a, 0.1), -10.0, 10.0)
end
reward(mdp::LQR, s::Float64, a::Float64) = -s^2

function MDP(mdp::LQR; γ::Float64=mdp.γ)
    return MDP(
            γ,
            nothing, # continuous state
            nothing, # continuous action space
            (s,a, s′=nothing) -> begin
                S′ = transition(mdp, s, a)
                if s′ == nothing
                    return S′
                end
                return pdf(S′, s′)
            end,
            (s,a) -> reward(mdp, s, a),
            (s, a)->begin
                s′ = rand(transition(mdp,s,a))
                r = reward(mdp, s, a)
                return (s′, r)
            end
        )
end

const LINEAR_1D_GAUSSIAN_POLICY = (θ::Vector{Float64}, s::Float64) -> begin
            λ = θ[1]
            σ = abs(θ[2]) + 0.00001
            return Normal(λ*s, σ)
        end
const LINEAR_1D_GAUSSIAN_POLICY_∇ = (θ::Vector{Float64}, a::Float64, s::Float64) -> begin
            λ = θ[1]
            σ = θ[2] # note: standard dev would be abs(θ[2])
            ∇₁ = s*(a - λ*s)*exp(-(a-λ*s)^2/(2*σ^2)) / (sqrt(2π)*abs(σ)^3)
            ∇₂ = sign(σ)*exp(-(a-λ*s)^2/(2*σ^2)) * ((a-λ*s)^2 - σ^2)/(sqrt(2π)*σ^4)
            if abs(σ) < 1e-6
                ∇₁ = 0.0
                ∇₂ = 0.0
            end
            return [∇₁, ∇₂]
        end
const LINEAR_1D_GAUSSIAN_POLICY_∇logπ = (θ::Vector{Float64}, a::Float64, s::Float64) -> begin
            λ = θ[1]
            σ = θ[2]
            ∇₁ = (a*s - λ*s^2) / (σ^2)
            ∇₂ = (a^2-2a*λ*s - σ^2 + λ^2*s^2)/(σ^3)
            if abs(σ) < 1e-6
                ∇₁ = 0.0
                ∇₂ = 0.0
            end
            return [∇₁, ∇₂]
        end
const LINEAR_1D_GAUSSIAN_POLICY_KLDIV = (θ::Vector{Float64}, θ′::Vector{Float64}, s::Float64) -> begin
            μ₁ = θ[1]*s
            μ₂ = θ′[1]*s
            σ₁ = abs(θ[2])
            σ₂ = abs(θ′[2])
            return log(σ₂/σ₁) + (σ₁^2 + (μ₁ - μ₂)^2)/(2*σ₂^2) - 1/2
        end
const LINEAR_1D_GAUSSIAN_POLICY_H_KL = (θ::Vector{Float64}, θ′::Vector{Float64}, s::Float64) -> begin
            F = zeros(Float64,2,2)
            λ₁ = θ[1]
            λ₂ = θ′[1]
            σ₁ = θ[2]
            σ₂ = θ′[2]
            F[1,1] = s^2/(σ₂^2)
            F[2,1] = F[1,2] = 2*s^2*(λ₁ - λ₂)/(σ₂^3)
            F[2,2] = 3*((λ₁ - λ₂)^2*s^2+σ₁^2)/(σ₂^4) - 1/(σ₂)^2
            return F
        end


# This file is a part of Julia. License is MIT: https://julialang.org/license

__precompile__(true)

"""
Arnoldi and Lanczos iteration for computing eigenvalues
"""
module IterativeEigenSolvers

using Base.LinAlg: BlasFloat, BlasInt, SVD, checksquare

export eigs, svds

include("arpack.jl")

using .ARPACK

## eigs
"""
    eigs(A; nev=6, ncv=max(20,2*nev+1), which=:LM, tol=0.0, maxiter=300, sigma=nothing, ritzvec=true, v0=zeros((0,))) -> (d,[v,],nconv,niter,nmult,resid)

Computes eigenvalues `d` of `A` using implicitly restarted Lanczos or Arnoldi iterations for real symmetric or
general nonsymmetric matrices respectively.

The following keyword arguments are supported:

* `nev`: Number of eigenvalues
* `ncv`: Number of Krylov vectors used in the computation; should satisfy `nev+1 <= ncv <= n`
  for real symmetric problems and `nev+2 <= ncv <= n` for other problems, where `n` is the
  size of the input matrix `A`. The default is `ncv = max(20,2*nev+1)`. Note that these
  restrictions limit the input matrix `A` to be of dimension at least 2.
* `which`: type of eigenvalues to compute. See the note below.

| `which` | type of eigenvalues                                                                                |
|:--------|:--------------------------------------------------------------------------------------------------------------------------|
| `:LM`   | eigenvalues of largest magnitude (default)                                                                                |
| `:SM`   | eigenvalues of smallest magnitude                                                                                |
| `:LR`   | eigenvalues of largest real part                                                                                |
| `:SR`   | eigenvalues of smallest real part                                                                                |
| `:LI`   | eigenvalues of largest imaginary part (nonsymmetric or complex `A` only)                                                  |
| `:SI`   | eigenvalues of smallest imaginary part (nonsymmetric or complex `A` only)                                                 |
| `:BE`   | compute half of the eigenvalues from each end of the spectrum, biased in favor of the high end. (real symmetric `A` only) |

* `tol`: parameter defining the relative tolerance for convergence of Ritz values (eigenvalue estimates).
     A Ritz value ``θ`` is considered converged when its associated residual
     is less than or equal to the product of `tol` and ``max(ɛ^{2/3}, |θ|)``,
     where `ɛ = eps(real(eltype(A)))/2` is LAPACK's machine epsilon.
     The residual associated with ``θ`` and its corresponding Ritz vector ``v``
     is defined as the norm ``||Av - vθ||``.
     The specified value of `tol` should be positive; otherwise, it is ignored
     and ``ɛ`` is used instead.
     Default: ``ɛ``.

* `maxiter`: Maximum number of iterations (default = 300)
* `sigma`: Specifies the level shift used in inverse iteration. If `nothing` (default),
  defaults to ordinary (forward) iterations. Otherwise, find eigenvalues close to `sigma`
  using shift and invert iterations.
* `ritzvec`: Returns the Ritz vectors `v` (eigenvectors) if `true`
* `v0`: starting vector from which to start the iterations

`eigs` returns the `nev` requested eigenvalues in `d`, the corresponding Ritz vectors `v`
(only if `ritzvec=true`), the number of converged eigenvalues `nconv`, the number of
iterations `niter` and the number of matrix vector multiplications `nmult`, as well as the
final residual vector `resid`.

# Examples
```jldoctest
julia> A = Diagonal(1:4);

julia> λ, ϕ = eigs(A, nev = 2);

julia> λ
2-element Array{Float64,1}:
 4.0
 3.0
```

!!! note
    The `sigma` and `which` keywords interact: the description of eigenvalues
    searched for by `which` do *not* necessarily refer to the eigenvalues of
    `A`, but rather the linear operator constructed by the specification of the
    iteration mode implied by `sigma`.

    | `sigma`         | iteration mode                   | `which` refers to eigenvalues of |
    |:----------------|:---------------------------------|:---------------------------------|
    | `nothing`       | ordinary (forward)               | ``A``                            |
    | real or complex | inverse with level shift `sigma` | ``(A - \\sigma I )^{-1}``        |

!!! note
    Although `tol` has a default value, the best choice depends strongly on the
    matrix `A`. We recommend that users _always_ specify a value for `tol`
    which suits their specific needs.

    For details of how the errors in the computed eigenvalues are estimated, see:

    * B. N. Parlett, "The Symmetric Eigenvalue Problem", SIAM: Philadelphia, 2/e
      (1998), Ch. 13.2, "Accessing Accuracy in Lanczos Problems", pp. 290-292 ff.
    * R. B. Lehoucq and D. C. Sorensen, "Deflation Techniques for an Implicitly
      Restarted Arnoldi Iteration", SIAM Journal on Matrix Analysis and
      Applications (1996), 17(4), 789–821.  doi:10.1137/S0895479895281484
"""
eigs(A; kwargs...) = eigs(A, I; kwargs...)
eigs(A::AbstractMatrix{<:BlasFloat}, ::UniformScaling; kwargs...) = _eigs(A, I; kwargs...)

eigs(A::AbstractMatrix{T}, B::AbstractMatrix{T}; kwargs...) where {T<:BlasFloat} = _eigs(A, B; kwargs...)
eigs(A::AbstractMatrix{BigFloat}, B::AbstractMatrix...; kwargs...) = throw(MethodError(eigs, Any[A,B,kwargs...]))
eigs(A::AbstractMatrix{BigFloat}, B::UniformScaling; kwargs...) = throw(MethodError(eigs, Any[A,B,kwargs...]))
function eigs(A::AbstractMatrix{T}, ::UniformScaling; kwargs...) where T
    Tnew = typeof(zero(T)/sqrt(one(T)))
    eigs(convert(AbstractMatrix{Tnew}, A), I; kwargs...)
end
function eigs(A::AbstractMatrix, B::AbstractMatrix; kwargs...)
    T = promote_type(eltype(A), eltype(B))
    Tnew = typeof(zero(T)/sqrt(one(T)))
    eigs(convert(AbstractMatrix{Tnew}, A), convert(AbstractMatrix{Tnew}, B); kwargs...)
end
"""
    eigs(A, B; nev=6, ncv=max(20,2*nev+1), which=:LM, tol=0.0, maxiter=300, sigma=nothing, ritzvec=true, v0=zeros((0,))) -> (d,[v,],nconv,niter,nmult,resid)

Computes generalized eigenvalues `d` of `A` and `B` using implicitly restarted Lanczos or Arnoldi iterations for
real symmetric or general nonsymmetric matrices respectively.

The following keyword arguments are supported:

* `nev`: Number of eigenvalues
* `ncv`: Number of Krylov vectors used in the computation; should satisfy `nev+1 <= ncv <= n`
  for real symmetric problems and `nev+2 <= ncv <= n` for other problems, where `n` is the
  size of the input matrices `A` and `B`. The default is `ncv = max(20,2*nev+1)`. Note that
  these restrictions limit the input matrix `A` to be of dimension at least 2.
* `which`: type of eigenvalues to compute. See the note below.

| `which` | type of eigenvalues                                                                                |
|:--------|:--------------------------------------------------------------------------------------------------------------------------|
| `:LM`   | eigenvalues of largest magnitude (default)                                                                                |
| `:SM`   | eigenvalues of smallest magnitude                                                                                |
| `:LR`   | eigenvalues of largest real part                                                                                |
| `:SR`   | eigenvalues of smallest real part                                                                                |
| `:LI`   | eigenvalues of largest imaginary part (nonsymmetric or complex `A` only)                                                  |
| `:SI`   | eigenvalues of smallest imaginary part (nonsymmetric or complex `A` only)                                                 |
| `:BE`   | compute half of the eigenvalues from each end of the spectrum, biased in favor of the high end. (real symmetric `A` only) |

* `tol`: relative tolerance used in the convergence criterion for eigenvalues, similar to
     `tol` in the [`eigs(A)`](@ref) method for the ordinary eigenvalue
     problem, but effectively for the eigenvalues of ``B^{-1} A`` instead of ``A``.
     See the documentation for the ordinary eigenvalue problem in
     [`eigs(A)`](@ref) and the accompanying note about `tol`.
* `maxiter`: Maximum number of iterations (default = 300)
* `sigma`: Specifies the level shift used in inverse iteration. If `nothing` (default),
  defaults to ordinary (forward) iterations. Otherwise, find eigenvalues close to `sigma`
  using shift and invert iterations.
* `ritzvec`: Returns the Ritz vectors `v` (eigenvectors) if `true`
* `v0`: starting vector from which to start the iterations

`eigs` returns the `nev` requested eigenvalues in `d`, the corresponding Ritz vectors `v`
(only if `ritzvec=true`), the number of converged eigenvalues `nconv`, the number of
iterations `niter` and the number of matrix vector multiplications `nmult`, as well as the
final residual vector `resid`.

# Examples
```jldoctest
julia> A = sparse(1.0I, 4, 4); B = Diagonal(1:4);

julia> λ, ϕ = eigs(A, B, nev = 2);

julia> λ
2-element Array{Float64,1}:
 1.0
 0.4999999999999999
```

!!! note
    The `sigma` and `which` keywords interact: the description of eigenvalues searched for by
    `which` do *not* necessarily refer to the eigenvalue problem ``Av = Bv\\lambda``, but rather
    the linear operator constructed by the specification of the iteration mode implied by `sigma`.

    | `sigma`         | iteration mode                   | `which` refers to the problem      |
    |:----------------|:---------------------------------|:-----------------------------------|
    | `nothing`       | ordinary (forward)               | ``Av = Bv\\lambda``                |
    | real or complex | inverse with level shift `sigma` | ``(A - \\sigma B )^{-1}B = v\\nu`` |
"""
eigs(A, B; kwargs...) = _eigs(A, B; kwargs...)
function _eigs(A, B;
               nev::Integer=6, ncv::Integer=max(20,2*nev+1), which=:LM,
               tol=0.0, maxiter::Integer=300, sigma=nothing, v0::Vector=zeros(eltype(A),(0,)),
               ritzvec::Bool=true)
    n = checksquare(A)

    T = eltype(A)
    iscmplx = T <: Complex
    isgeneral = B !== I
    sym = !iscmplx && issymmetric(A) && issymmetric(B)
    nevmax = sym ? n-1 : n-2
    if nevmax <= 0
        throw(ArgumentError("input matrix A is too small. Use eigfact instead."))
    end
    if nev > nevmax
        warn("Adjusting nev from $nev to $nevmax")
        nev = nevmax
    end
    if nev <= 0
        throw(ArgumentError("requested number of eigenvalues (nev) must be ≥ 1, got $nev"))
    end
    ncvmin = nev + (sym ? 1 : 2)
    if ncv < ncvmin
        warn("Adjusting ncv from $ncv to $ncvmin")
        ncv = ncvmin
    end
    ncv = BlasInt(min(ncv, n))
    bmat = isgeneral ? "G" : "I"
    isshift = sigma !== nothing

    if isa(which,AbstractString)
        warn("Use symbols instead of strings for specifying which eigenvalues to compute")
        which=Symbol(which)
    end
    if (which != :LM && which != :SM && which != :LR && which != :SR &&
        which != :LI && which != :SI && which != :BE)
        throw(ArgumentError("which must be :LM, :SM, :LR, :SR, :LI, :SI, or :BE, got $(repr(which))"))
    end
    if which == :BE && !sym
        throw(ArgumentError("which=:BE only possible for real symmetric problem"))
    end
    isshift && which == :SM && warn("use of :SM in shift-and-invert mode is not recommended, use :LM to find eigenvalues closest to sigma")

    if which==:SM && !isshift # transform into shift-and-invert method with sigma = 0
        isshift=true
        sigma=zero(T)
        which=:LM
    end

    if sigma !== nothing && !iscmplx && isa(sigma,Complex)
        throw(ArgumentError("complex shifts for real problems are not yet supported"))
    end
    sigma = isshift ? convert(T,sigma) : zero(T)

    if !isempty(v0)
        if length(v0) != n
            throw(DimensionMismatch())
        end
        if eltype(v0) != T
            throw(ArgumentError("starting vector must have element type $T, got $(eltype(v0))"))
        end
    end

    whichstr = "LM"
    if which == :BE
        whichstr = "BE"
    end
    if which == :LR
        whichstr = (!sym ? "LR" : "LA")
    end
    if which == :SR
        whichstr = (!sym ? "SR" : "SA")
    end
    if which == :LI
        if !sym
            whichstr = "LI"
        else
            throw(ArgumentError("largest imaginary is meaningless for symmetric eigenvalue problems"))
        end
    end
    if which == :SI
        if !sym
            whichstr = "SI"
        else
            throw(ArgumentError("smallest imaginary is meaningless for symmetric eigenvalue problems"))
        end
    end

    # Refer to ex-*.doc files in ARPACK/DOCUMENTS for calling sequence
    matvecA!(y, x) = A_mul_B!(y, A, x)
    if !isgeneral           # Standard problem
        matvecB = x -> x
        if !isshift         #    Regular mode
            mode       = 1
            solveSI = x->x
        else                #    Shift-invert mode
            mode       = 3
            F = factorize(A - UniformScaling(sigma))
            solveSI = x -> F \ x
        end
    else                    # Generalized eigenproblem
        matvecB = x -> B * x
        if !isshift         #    Regular inverse mode
            mode       = 2
            F = factorize(B)
            solveSI = x -> F \ x
        else                #    Shift-invert mode
            mode       = 3
            F = factorize(A - sigma*B)
            solveSI = x -> F \ x
        end
    end

    # Compute the Ritz values and Ritz vectors
    (resid, v, ldv, iparam, ipntr, workd, workl, lworkl, rwork, TOL) =
       ARPACK.aupd_wrapper(T, matvecA!, matvecB, solveSI, n, sym, iscmplx, bmat, nev, ncv, whichstr, tol, maxiter, mode, v0)

    # Postprocessing to get eigenvalues and eigenvectors
    output = ARPACK.eupd_wrapper(T, n, sym, iscmplx, bmat, nev, whichstr, ritzvec, TOL,
                                 resid, ncv, v, ldv, sigma, iparam, ipntr, workd, workl, lworkl, rwork)

    # Issue 10495, 10701: Check that all eigenvalues are converged
    nev = length(output[1])
    nconv = output[ritzvec ? 3 : 2]
    nev ≤ nconv || warn("not all wanted Ritz pairs converged. Requested: $nev, converged: $nconv")

    return output
end

## svds
struct SVDAugmented{T,S} <: AbstractArray{T, 2}
    X::S
    SVDAugmented{T,S}(X::AbstractMatrix) where {T,S} = new(X)
end

function SVDAugmented(A::AbstractMatrix{T}) where T
    Tnew = typeof(zero(T)/sqrt(one(T)))
    Anew = convert(AbstractMatrix{Tnew}, A)
    SVDAugmented{Tnew,typeof(Anew)}(Anew)
end

function Base.A_mul_B!(y::StridedVector{T}, A::SVDAugmented{T}, x::StridedVector{T}) where T
    m, mn = size(A.X, 1), length(x)
    A_mul_B!( view(y, 1:m), A.X, view(x, m + 1:mn)) # left singular vector
    Ac_mul_B!(view(y, m + 1:mn), A.X, view(x, 1:m)) # right singular vector
    return y
end
Base.size(A::SVDAugmented)  = ((+)(size(A.X)...), (+)(size(A.X)...))
Base.ishermitian(A::SVDAugmented) = true

struct AtA_or_AAt{T,S} <: AbstractArray{T, 2}
    A::S
    buffer::Vector{T}
end

function AtA_or_AAt(A::AbstractMatrix{T}) where T
    Tnew = typeof(zero(T)/sqrt(one(T)))
    Anew = convert(AbstractMatrix{Tnew}, A)
    AtA_or_AAt{Tnew,typeof(Anew)}(Anew, Vector{Tnew}(uninitialized, max(size(A)...)))
end

function Base.A_mul_B!(y::StridedVector{T}, A::AtA_or_AAt{T}, x::StridedVector{T}) where T
    if size(A.A, 1) >= size(A.A, 2)
        A_mul_B!(A.buffer, A.A, x)
        return Ac_mul_B!(y, A.A, A.buffer)
    else
        Ac_mul_B!(A.buffer, A.A, x)
        return A_mul_B!(y, A.A, A.buffer)
    end
end
Base.size(A::AtA_or_AAt) = ntuple(i -> min(size(A.A)...), Val(2))
Base.ishermitian(s::AtA_or_AAt) = true

svds(A::AbstractMatrix{<:BlasFloat}; kwargs...) = _svds(A; kwargs...)
svds(A::AbstractMatrix{BigFloat}; kwargs...) = throw(MethodError(svds, Any[A, kwargs...]))
function svds(A::AbstractMatrix{T}; kwargs...) where T
    Tnew = typeof(zero(T)/sqrt(one(T)))
    svds(convert(AbstractMatrix{Tnew}, A); kwargs...)
end

"""
    svds(A; nsv=6, ritzvec=true, tol=0.0, maxiter=1000, ncv=2*nsv, v0=zeros((0,))) -> (SVD([left_sv,] s, [right_sv,]), nconv, niter, nmult, resid)

Computes the largest singular values `s` of `A` using implicitly restarted Lanczos
iterations derived from [`eigs`](@ref).

**Inputs**

* `A`: Linear operator whose singular values are desired. `A` may be represented as a
  subtype of `AbstractArray`, e.g., a sparse matrix, or any other type supporting the four
  methods `size(A)`, `eltype(A)`, `A * vector`, and `A' * vector`.
* `nsv`: Number of singular values. Default: 6.
* `ritzvec`: If `true`, return the left and right singular vectors `left_sv` and `right_sv`.
   If `false`, omit the singular vectors. Default: `true`.
* `tol`: tolerance, see [`eigs`](@ref).
* `maxiter`: Maximum number of iterations, see [`eigs`](@ref). Default: 1000.
* `ncv`: Maximum size of the Krylov subspace, see [`eigs`](@ref) (there called `nev`). Default: `2*nsv`.
* `v0`: Initial guess for the first Krylov vector. It may have length `min(size(A)...)`, or 0.

**Outputs**

* `svd`: An `SVD` object containing the left singular vectors, the requested values, and the
  right singular vectors. If `ritzvec = false`, the left and right singular vectors will be
  empty.
* `nconv`: Number of converged singular values.
* `niter`: Number of iterations.
* `nmult`: Number of matrix--vector products used.
* `resid`: Final residual vector.

# Examples
```jldoctest
julia> A = Diagonal(1:4);

julia> s = svds(A, nsv = 2)[1];

julia> s[:S]
2-element Array{Float64,1}:
 4.0
 2.9999999999999996
```

!!! note "Implementation"
    `svds(A)` is formally equivalent to calling [`eigs`](@ref) to perform implicitly restarted
    Lanczos tridiagonalization on the Hermitian matrix ``A^\\prime A`` or ``AA^\\prime`` such
    that the size is smallest.
"""
svds(A; kwargs...) = _svds(A; kwargs...)
function _svds(X; nsv::Int = 6, ritzvec::Bool = true, tol::Float64 = 0.0, maxiter::Int = 1000, ncv::Int = 2*nsv, v0::Vector=zeros(eltype(X),(0,)))
    if nsv < 1
        throw(ArgumentError("number of singular values (nsv) must be ≥ 1, got $nsv"))
    end
    if nsv > minimum(size(X))
        throw(ArgumentError("number of singular values (nsv) must be ≤ $(minimum(size(X))), got $nsv"))
    end
    m, n = size(X)
    otype = eltype(X)
    if length(v0) ∉ [0,n]
        throw(DimensionMismatch("length of v0, the guess for the starting right Krylov vector, must be 0, or $n, got $(length(v0))"))
    end
    ex    = eigs(AtA_or_AAt(X), I; which = :LM, ritzvec = ritzvec, nev = nsv, tol = tol, maxiter = maxiter, v0=v0)
    # ind   = [1:2:ncv;]
    # sval  = abs.(ex[1][ind])

    svals = sqrt.(real.(ex[1]))
    if ritzvec
        # calculating singular vectors
        # left_sv  = sqrt(2) * ex[2][ 1:size(X,1),     ind ] .* sign.(ex[1][ind]')
        if size(X, 1) >= size(X, 2)
            V = ex[2]
            U = qr(scale!(X*V, inv.(svals)))[1]
        else
            U = ex[2]
            V = qr(scale!(X'U, inv.(svals)))[1]
        end

        # right_sv = sqrt(2) * ex[2][ size(X,1)+1:end, ind ]
        return (SVD(U, svals, V'), ex[3], ex[4], ex[5], ex[6])
    else
        #The sort is necessary to work around #10329
        return (SVD(zeros(eltype(svals), n, 0),
                    svals,
                    zeros(eltype(svals), 0, m)),
                    ex[2], ex[3], ex[4], ex[5])
    end
end

end # module


function jump_sensorloc(solver, seed, n; verbose = false, test = false)

    m, x_true, a, d, d_bar = sensorloc_data(seed, n)

    model = Model(with_optimizer(solver))

    # Build SDP problem
    @variable(model, X[1:n+2, 1:n+2], PSD)

    # Constraint with distances from anchors to sensors
    for j in 1:n
        for k in 1:m
            # e = zeros(n, 1)
            # e[j] = -1.0
            # v = vcat(a[k], e)
            # V = v * v'
            # @constraint(model, sum(V .* X) == d_bar[k, j]^2)
            @constraint(model, X[1,1]*a[k][1]*a[k][1] + X[2,2]*a[k][2]*a[k][2]
                                 - 2 * X[1, j+2] * a[k][1]
                                 - 2 * X[2, j+2] * a[k][2]
                                 + X[j+2, j+2]
                                == d_bar[k, j]^2)
        end
    end

    # Constraint with distances from sensors to sensors
    count, count_all = 0, 0
    for i in 1:n
        for j in 1:i - 1
            count_all += 1
            if rand() > 0.9
                count += 1
                # e = zeros(n, 1)
                # e[i] = 1.0
                # e[j] = -1.0
                # v = vcat(zeros(2, 1), e)
                # V = v * v'
                # @constraint(model, sum(V .* X) == d[i, j]^2)
                @constraint(model, X[i+2,i+2] + X[j+2,j+2] - 2*X[i+2,j+2] == d[i, j]^2)
            end
        end
    end
    if verbose
        @show count_all, count
    end

    @constraint(model, X[1, 1] == 1.0)
    @constraint(model, X[1, 2] == 0.0)
    @constraint(model, X[2, 1] == 0.0)
    @constraint(model, X[2, 2] == 1.0)

    # Feasibility objective function
    # L = eye(n+1)
    # @objective(model, Min, sum(L[i, j] * X[i, j] for i in 1:n+1, j in 1:n+1))
    @objective(model, Min, 0.0 * X[1, 1] + 0.0 * X[2, 2])

    teste = @time optimize!(model)

    XX = value.(X)

    verbose && sensorloc_eval(n, m, x_true, XX)

    objval = objective_value(model)
    stime = MOI.get(model, MOI.SolveTime())

    # @show tp = typeof(model.moi_backend.optimizer.model.optimizer)
    # @show fieldnames(tp)
    rank = -1
    try
        @show rank = model.moi_backend.optimizer.model.optimizer.sol.final_rank
    catch
    end
    status = 0
    if JuMP.termination_status(model) == MOI.OPTIMAL
        status = 1
    end
    return (objval, stime, rank, status)
    # return (objval, stime)
end


A = randn(3,5)
b = randn(3)

x_pg = Variable(5)
prob_pg = problem(ls(A*x_pg - b) + 1e-3*norm(x_pg, 1))
sol_pg = solve(prob_pg, StructuredOptimization.PG())

x_zfpr = Variable(5)
prob_zfpr = problem(ls(A*x_zfpr - b) + 1e-3*norm(x_zfpr, 1))
sol_zfpr = solve(prob_zfpr, StructuredOptimization.ZeroFPR())

x_pnc = Variable(5)
prob_pnc = problem(ls(A*x_pnc - b) + 1e-3*norm(x_pnc, 1))
sol_pnc = solve(prob_pnc, StructuredOptimization.PANOC())


function next(w, η)
	return (w + η*sin(2π*w) + 1) % 1
end
function dnext(w, η)
	return 1.0 + 2π*η*cos(2π*w)
end
function d2next(w, η)
	return -2π*2π*η*sin(2π*w)
end


function ljpotential(xij, qi, qj, ti, tj, mu, eta, kappa)

A = mu[1];
B = mu[2];

r2 = xij[1,:].*xij[1,:] .+ xij[2,:].*xij[2,:] .+ xij[3,:].*xij[3,:];
#r1 = sqrt.(r2);
r4 = r2.*r2;
r6 = r2.*r4;
r12 = r6.*r6;
eij = (A./r12 - B./r6);

r8 = r6.*r2;
r14 = r12.*r2;
fij = xij*1.0;
for i = 1:3
    fij[i,:] = xij[i,:].*((6*B)./r8 - (12*A)./r14);
end
# r7 = r6.*r2;
# r13 = r12.*r1;
# dr = xij./r1;
# fij = dr.*((6*B)./r7 - (12*A)./r13);

return eij, fij

end


using Test
using IntervalArithmetic

@testset "atan2" begin include("atan2.jl") end
@testset "c-xsc" begin include("c-xsc.jl") end
@testset "fi_lib" begin include("fi_lib.jl") end
@testset "IEEE 1788" begin
    @testset "ieee1788-constructors" begin include("ieee1788-constructors.jl") end
    @testset "ieee1788-exceptions" begin include("ieee1788-exceptions.jl") end
    @testset "libieeep1788_bool" begin include("libieeep1788_bool.jl") end
    @testset "libieeep1788_cancel" begin include("libieeep1788_cancel.jl") end
    @testset "libieeep1788_class" begin include("libieeep1788_class.jl") end
    @testset "libieeep1788_elem" begin include("libieeep1788_elem.jl") end
    @testset "libieeep1788_num" begin include("libieeep1788_num.jl") end
    @testset "libieeep1788_overlap" begin include("libieeep1788_overlap.jl") end
    @testset "libieeep1788_rec_bool" begin include("libieeep1788_rec_bool.jl") end
    @testset "libieeep1788_reduction" begin include("libieeep1788_reduction.jl") end
    @testset "libieeep1788_set" begin include("libieeep1788_set.jl") end
    @testset "libieeep1788_rev" begin include("libieeep1788_rev.jl") end
    @testset "libieeep1788_mul_rev" begin include("libieeep1788_mul_rev.jl") end
end # IEEE 1788
@testset "mpfi" begin include("mpfi.jl") end
@testset "abs_rev" begin include("abs_rev.jl") end
@testset "pow_rev" begin include("pow_rev.jl") end


# indicator of the free cone

export IndFree

"""
**Indicator of the free cone**

    IndFree()

Returns the indicator function of the whole space, or "free cone", *i.e.*,
a function which is identically zero.
"""
struct IndFree <: ProximableFunction end

is_separable(f::IndFree) = true
is_convex(f::IndFree) = true
is_affine(f::IndFree) = true
is_cone(f::IndFree) = true
is_smooth(f::IndFree) = true
is_quadratic(f::IndFree) = true

const Zero = IndFree

function (f::IndFree)(x::AbstractArray{T}) where {R, T <: RealOrComplex{R}}
  return zero(R)
end

function prox!(y::AbstractArray{T}, f::IndFree, x::AbstractArray{T}, gamma::R=one(R)) where {R, T <: RealOrComplex{R}}
  y .= x
  return zero(R)
end

prox!(y::AbstractArray{T}, f::IndFree, x::AbstractArray{T}, gamma::AbstractArray{T}) where {R, T <: RealOrComplex{R}} = prox!(y, f, x)

function gradient!(y::AbstractArray{T}, f::IndFree, x::AbstractArray{T}) where {R, T <: RealOrComplex{R}}
  y .= zero(T)
  return zero(R)
end

fun_name(f::IndFree) = "indicator of the free cone"
fun_dom(f::IndFree) = "AbstractArray{Real}, AbstractArray{Complex}"
fun_expr(f::IndFree) = "x ↦ 0"
fun_params(f::IndFree) = "none"

function prox_naive(f::IndFree, x::AbstractArray, gamma=1.0)
  return x, 0.0
end


#Methods to transform a PEP to a companion linearized form and solve the corresponding linearized pencil
using SparseArrays

export companion
export polyeig #Wrapper around the solver for a linearized PEP pencil

"""
    E,A = companion(nep::PEP);

Linearizes a  polynomial eigenvalue problem (PEP) a to the companion form, as in the paper by Mehrmann and Voss.
More precisely, for a k-th degree PEP with n-by-n coefficient matrices,
this returns matrices E and A, both kn-by-kn, corresponding to the linearized problem
```math
Ax = λEx
```

# Example
```julia-repl
julia> pep = nep_gallery("pep0");
julia> E,A = companion(pep);
julia> λ, V = eigen(A,E);
julia> minimum(svd(compute_Mder(pep,λ[1])).S)
2.4845217786736996e-12
```

# References
* V. Mehrmann and H. Voss, Non-linear eigenvalue problems, a challenge for modern eigenvalue methods, GAMM‐Mitteilungen (2004)
"""
    function companion(pep::PEP)

        n = size(pep,1) #Size of monomial coefficient matrices
        d = size(pep.A,1)-1 #Degree of pep
        T = eltype(pep.A[1]) #Deduce the type of elements in the PEP matrices

        ##### Check sparsity of the problem and allocate memory accordingly #####
        if issparse(pep)
            E = spzeros(T, d*n, d*n)
            A = spzeros(T, n*d, n*d)

            #(d-1)n-by-(d-1)n matrix (Used to construct both E and A)
            Iblock = sparse(kron(Matrix{T}(I, d-1, d-1), Matrix{T}(I, n, n)))
        else
            E = zeros(T, d*n, d*n)
            A = zeros(T, n*d, n*d)

            Iblock = kron(Matrix{T}(I, d-1, d-1), Matrix{T}(I, n, n))
        end

        E[1:n,1:n] = pep.A[d+1];#Fill block (1,1)
        E[n+1:d*n,n+1:d*n] = Iblock;#Fill all blocks on the diagonal with eye(n)

        #####Construct A #####

        #First row block of A
        for i=1:d
           A[1:n,(i-1)*n+1:i*n] = pep.A[d-i+1];
        end
        #Lower part of A
        A[n+1:d*n,1:(d-1)*n] = T(-1.0)*Iblock

        return E,-A

    end

"""
    λ,v = polyeig([eltype],nep::PEP,[eigsolvertype,])

Linearizes a  polynomial eigenvalue problem (PEP) a to the companion form
and solves the corresponding linear eigenvalue problem; see [`companion`](@ref).
The `eigsolvertype` is optinal can be used to specify how the linear problem
is solved; see [`eig_solve`](@ref), and [`EigSolver`](@ref).

# Example
```julia-repl
julia> pep = nep_gallery("pep0");
julia> λ,V = polyeig(pep);
julia> minimum(svd(compute_Mder(pep,λ[1])).S)
1.2050763381899922e-14
julia> norm(compute_Mlincomb(pep,λ[2],vec(V[:,2])))
1.0245470569036458e-12
```
"""
    polyeig(pep::PEP,vargs...)=polyeig(ComplexF64,pep,vargs...)
    function polyeig(::Type{T},pep::PEP,eigsolvertype::Type=DefaultEigSolver) where T

        #Linearize to Ax = λEx
        E,A = companion(pep);

        solver::EigSolver = eigsolvertype(A,E);
        D,V = eig_solve(solver,target=one(T),nev=size(A,1));

        D = Vector{T}(D)
        V = Matrix{T}(V)

        return D,V[1:size(pep,1),:]
    end

"""
    λ,v = polyeig([eltype],nep::ChebPEP)

Computes a companion linearization for the NEP represented
in a Chebyshev basis, and returns eigenpairs.

# Example
```julia
julia> using LinearAlgebra
julia> nep=nep_gallery("dep0");
julia> chebpep=ChebPEP(nep,9,-3,1,cosine_formula_cutoff=5);
julia> (λv,V)=polyeig(chebpep);
julia> ii=argmin(abs.(λv));
julia> λ=λv[ii];
julia> v=V[:,ii];
julia> norm(compute_Mlincomb(chebpep,λ,v))
1.3543968603949142e-14
julia> # Actually, it's not a bad approx to the original NEP either
julia> norm(compute_Mlincomb(nep,λ,v))
4.326355966047557e-6
```

See also [`ChebPEP`](@ref).

# References:

* Amiraslani, A., Corless, R. M. & Lancaster, P. "Linearization of matrix polynomials expressed in poly-nomial bases" IMA J. Numer. Anal.,29 (2009): 141–157.
* Effenberger and Kressner. "Chebyshev interpolation for nonlinear eigenvalue problems." BIT Numerical Mathematics 52.4 (2012): 933-951.

"""
polyeig(pep::ChebPEP,vargs...)=polyeig(ComplexF64,pep,vargs...)
function polyeig(::Type{T}, pep::ChebPEP{TT,Ftype}) where {T,Ftype,TT}
    k=pep.k
    Fk=get_Av(pep);

    n=size(pep,1);
    L0=zeros(T,n*(k-1),n*(k-1));
    L1=zeros(T,n*(k-1),n*(k-1));
    II=Matrix{Ftype}(I,n,n);

    for j=1:(k-2)
        L0[((j-1)*n) .+ (1:n), j*n .+ (1:n)]=II
        L0[j*n .+ (1:n), ((j-1)*n) .+ (1:n)]=II
    end

    for j=1:k-1
        L0[((k-2)*n) .+ (1:n), (n*(j-1)).+ (1:n)]=-Fk[j];
    end

    L0[((k-2)*n) .+ (1:n), (n*(k-3)) .+ (1:n)] += Fk[k];

    for j=1:k-2
        factor=2;
        if (j==1)
            factor=1
        end
        L1[((j-1)*n) .+ (1:n), ((j-1)*n) .+ (1:n)]=factor*II
    end
    L1[((k-2)*n) .+ (1:n),
       ((k-2)*n) .+ (1:n)]=2*Fk[k]

    LL=eigen(L0,L1);

    a=pep.a; b=pep.b;
    return ((b-a)*(LL.values .+ 1)/2 .+ a,
            hcat(map(i -> normalize(LL.vectors[1:n,i]), 1:(n*(k-1)))...))

end


"""
    bbo_floor_int_searchrange(L::Int)

Generate a BlackBoxOptim SearchRange for optimizing integers on a
float scale when we are going to use floor to get back an Int.
"""
bbo_floor_int_searchrange(L::Int) = (0.499999999-L, L+0.499999999)

@inline floortoints(λ::Vector{N}) where {N<:Number} = floor.(Int, λ)
@inline floortoints(λ::Vector{Int}) = λ

"""
    nnzeroasint(λ::Vector{N})

Count number of nonzero elements of a vector when floored to integers.
"""
nnzeroasint(λ::Vector{N}) where {N<:Number} = sum(floortoints(λ) .!= 0)

"""
    nonzeroindices(λ::Vector{N}, sorted = false)

Return the indices of the nonzero elements of a vector when they are floored to integers.
"""
function nonzeroindices(λ::Vector{N}, sorted = true) where {N<:Number}
    idxs = findall(v -> v != 0, λ)
    return sorted ? idxs[sortperm(λ[idxs], rev=true)] : idxs
end

"""
    sign_predict(λ::Vector, X)

Predict class, indicated with +1 or -1, i.e. sign-based, by multiplying
a vector `λ` interpreted as integer coefficients and a design matrix `X`.
"""
sign_predict(λ::Vector{Int}, X::AbstractMatrix) = Int[(v<0 ? -1 : 1) for v in (X * λ)]
sign_predict(λ::Vector{N}, X::AbstractMatrix) where {N<:Number} =
    sign_predict(floortoints(λ), X)

"""
    sign_accuracy(yhat, y)

Accuracy with which yhat predicts the (signed) class labels in y.
"""
@inline function sign_accuracy(yhat::Vector{N1}, y::Vector{N2}) where {N1<:Number,N2<:Number}
    @assert length(yhat) == length(y)
    sum((yhat .* y) .> 0)/float(length(yhat))
end

"""
    sign_accuracy(λ, X, y)

Accuracy with which a coefficient vector `λ` predicts `y` given design matrix `X`.
"""
sign_accuracy(λ::Vector{N}, X::AbstractMatrix{N1}, y::Vector{N2}) where {N<:Number,N1<:Number,N2<:Number} =
    sign_accuracy(sign_predict(λ, X), y)

"""
    sign_misclassification_rate(yhat, y)

Misclassification rate for the prediction yhat of the (signed) class labels in y.
"""
@inline sign_misclassification_rate(yhat, y) = 1.0 - sign_accuracy(yhat, y)

"""
    aspct(v::Float64)

Convert to a percentage value.
"""
@inline aspct(v::Float64) = round(100.0*v, digits=3)

as_sign_labels(y::Vector{N}) where {N<:Number} = Int[(c>0 ? 1 : -1) for c in y]
function as_sign_labels(y::AbstractVector)
    uniquevalues = sort(unique(y))
    @assert 1 <= length(uniquevalues) <= 2
    Int[(c==uniquevalues[1] ? 1 : -1) for c in y]
end

function uptobutexcluding(P::Int, excl::AbstractVector)
    setdiff(collect(1:P), excl)
end


# ---
# title: 15. 3Sum
# id: problem15
# author: zhwang
# date: 2022-03-15
# difficulty: Medium
# categories: Array, Two Pointers
# link: <https://leetcode.com/problems/3sum/description/>
# hidden: true
# ---
#
# Given an array `nums` of _n_ integers, are there elements _a_ , _b_ , _c_ in
# `nums` such that _a_ \+ _b_ \+ _c_ = 0? Find all unique triplets in the array
# which gives the sum of zero.
#
# Notice that the solution set must not contain duplicate triplets.
#
#
#
# **Example 1:**
#
#
#
#     Input: nums = [-1,0,1,2,-1,-4]
#     Output: [[-1,-1,2],[-1,0,1]]
#
#
# **Example 2:**
#
#
#
#     Input: nums = []
#     Output: []
#
#
# **Example 3:**
#
#
#
#     Input: nums = [0]
#     Output: []
#
#
#
#
# **Constraints:**
#
#   * `0 <= nums.length <= 3000`
#   * `-105 <= nums[i] <= 105`
#
#
## @lc code=start
using LeetCode

three_sum(nums::Vector{Int}) = three_sum!(copy(nums))
function three_sum!(nums::Vector{Int})::Vector{NTuple{3, Int}}
    sort!(nums)
    n, res, pre = length(nums), NTuple{3, Int}[], NaN
    for (i, first) in enumerate(nums)
        first == pre && continue
        pre, left, right = first, i + 1, n
        while left < right
            total = nums[left] + nums[right]
            if total < -first
                left += 1
            elseif total > -first
                right -= 1
            else
                push!(res, (first, nums[left], nums[right]))
                pos = findfirst(!=(nums[left]), @view(nums[left+1:right]))
                isnothing(pos) ? break : left += pos
                pos = findfirst(!=(nums[right]), @view(nums[right-1:-1:left]))
                isnothing(pos) ? break : right -= pos
            end
        end
    end
    res
end

## @lc code=end

## @lc test=start
@testset "15.3sum.jl" begin
    @test three_sum([-1,0,1,2,-1,-4]) == [(-1,-1,2),(-1,0,1)]
    @test three_sum([0]) == three_sum!(Int[]) == NTuple{3, Int}[]
    @test three_sum([1, 1, -2]) == [(-2, 1, 1)]
    @test three_sum([0, 0, 0, 0]) == [(0, 0, 0)]
end
## @lc test=end

## @lc info=start
# link: [solution to 15]()
## @lc info=end


#!/usr/bin/julia
# FizzBuzz using basic division rules
# Author: @joaogui1

function div3(n)
    s = string(n)
    while length(s) > 1
        s = string(sum([parse(Int, c) for c in s]))
    end
    return s in ["0", "3", "6", "9"]
end

function div5(n)
    s = string(n)
    return s[end] in "05"
end

function fizzbuzz(n)
    for i = 1:n
        print(i, ": ")
        by3, by5 = div3(i), div5(i)

        if by3 && by5
            println("FizzBuzz")
        elseif by3
            println("Fizz")
        elseif by5
            println("Buzz")
        else
            println(i)
        end
    end
end

fizzbuzz(100)


export PointOperation

export apply_operation
export domaintype
export isidentity, istranslation, ispoint

import LinearAlgebra
import LinearAlgebraX

"""
    PointOperation{S<:Real}

Represents point symmetry operation

# Fields
* `matrix`: rotation/mirror/inversion matrix
"""
struct PointOperation{S<:Real} <:AbstractSpaceSymmetryOperation{S}
    matrix::Matrix{S}

    @doc """
        PointOperation(matrix::AbstractMatrix{S}) where {S<:Real}

    Construct a point operation with `matrix`.
    """
    function PointOperation(matrix::AbstractMatrix{S}) where {S<:Real}
        return PointOperation{S}(matrix)
    end

    function PointOperation{S}(matrix::AbstractMatrix{S2}) where {S<:Real, S2<:Real}
        dim = size(matrix, 1)
        size(matrix, 2) != dim && throw(DimensionMismatch("matrix not square"))
        return new{S}(matrix)
    end
end

function Base.one(x::PointOperation{S}) where {S}
    return PointOperation(Matrix{S}(LinearAlgebra.I, size(x.matrix)))
end

# function Base.convert(::Type{PointOperation{S}}, obj::IdentityOperation{S}) where S
#     dim = dimension(obj)
#     return PointOperation{S}(Matrix(LinearAlgebra.I, dim, dim))
# end

function Base.convert(::Type{PointOperation{S}}, matrix::AbstractMatrix{S}) where S
    return PointOperation{S}(matrix)
end

# function Base.promote_rule(::Type{PointOperation{S}}, ::Type{IdentityOperation{S}}) where S
#     return PointOperation{S}
# end

## properties

"""
    isidentity(arg::PointOperation)

Check whether `arg` is identity.
"""
isidentity(arg::PointOperation) = isone(arg.matrix)

"""
    istranslation(arg::PointOperation)

Check whether `arg` is a translation operation, i.e. identity.
"""
istranslation(arg::PointOperation) = isone(arg.matrix)

"""
    ispoint(arg::PointOperation)

Check whether `arg` is a point operation. Always `true`.
"""
ispoint(arg::PointOperation) = true

"""
    dimension(arg::PointOperation)

Return spatial dimension of `arg`.
"""
dimension(arg::PointOperation) = size(arg.matrix, 1)

Base.hash(arg::PointOperation{S}, h::UInt) where S = hash(PointOperation{S}, hash(arg.matrix, h))

## operators
function Base.:(==)(lhs::PointOperation, rhs::PointOperation)
    return lhs.matrix == rhs.matrix
end
# function Base.:(==)(pop::PointOperation, ::IdentityOperation)
#     return isidentity(pop)
# end
# function Base.:(==)(::IdentityOperation, pop::PointOperation)
#     return isidentity(pop)
# end
function Base.:(==)(pop::PointOperation, top::TranslationOperation)
    return isidentity(pop) && isidentity(top)
end
function Base.:(==)(top::TranslationOperation, pop::PointOperation)
    return isidentity(pop) && isidentity(top)
end

function Base.:(*)(lhs::PointOperation, rhs::PointOperation)
    return PointOperation(lhs.matrix * rhs.matrix)
end

function Base.:(^)(lhs::PointOperation{S}, rhs::Integer) where S
    if rhs >= 0
        return PointOperation(lhs.matrix^rhs)
    else
        lhs_inv_matrix = Matrix{S}(LinearAlgebraX.invx(lhs.matrix))
        return PointOperation{S}(lhs_inv_matrix^(-rhs))
    end
end

function Base.inv(lhs::PointOperation{S}) where {S<:Real}
    return PointOperation{S}(LinearAlgebraX.invx(lhs.matrix))
end

## apply
"""
    apply_operation(symop::PointOperation{S}, coordinate::AbstractArray{S}) where {S}

Apply point operation to the coordinates.
"""
function apply_operation(symop::PointOperation{S}, coord::AbstractVector{S}) where S
    return symop.matrix * coord
end

function (symop::PointOperation{S})(coord::AbstractVector{S}) where S
    return symop.matrix * coord
end


# # Sphere sound-soft scattering

# In this example we solve the scattering by a sphere. First, we need to load
# the modules
using LinearAlgebra
using WaveProp
using WaveProp.Nystrom
using WaveProp.Geometry
using WriteVTK

# We now specify which *pde* we want to solve, as well as the geometry:
dim = 3; k=2π
pde = Helmholtz(dim=dim,k=k)
Ω,mesh   = WaveProp.IO.gmsh_sphere(dim=2,h=0.1)
WaveProp.Mesh.compute_quadrature!(mesh,dim=2,order=1,need_normal=true)

# We are interested in solving a sound-soft scattering problem. We will take an
# incident field `uᵢ` as a plane wave, and solve for the scattered field using a
# combined field integral equation:
uᵢ       = (x)   -> exp(im*k*dot((1,0,0),x))
rhs      = γ₀((x)-> -uᵢ(x),mesh)
S,D      = single_double_layer(pde,mesh)
L        = I/2 + D - im*k*S
ϕ        = L\rhs

# Now that the density ϕ has been found, we can represent the solution
𝒮        = SingleLayerPotential(pde,mesh)
𝒟        = DoubleLayerPotential(pde,mesh)
u(x)     = uᵢ(x) + 𝒟[ϕ](x) - im*k*𝒮[ϕ](x)

# and output it on a plane for visualization

## output the mesh
fname = joinpath(@__DIR__,"mesh")
vtkfile = WaveProp.IO.vtk_mesh_file(mesh,Ω,fname) |> vtk_save

## compute the solution on a plane
rect,out_mesh  =  WaveProp.IO.gmsh_rectangle(origin=(-5,-5,0),dx=10,dy=10,h=0.1,dim=2)
vals           = [u(x) for x in WaveProp.Mesh.nodes(out_mesh)]
fname = joinpath(@__DIR__,"output")
vtkfile = WaveProp.IO.vtk_mesh_file(out_mesh,rect,fname)
vtkfile["total field (real part)"] = real(vals)
vtkfile["total field (imag part)"] = imag(vals)
vtk_save(vtkfile)


vec_infeasible = @stochastic_model begin
    @stage 1 begin
        @parameters begin
            lb = zeros(2)
            c  = [3., 2.]
        end
        @decision(model, x[i in 1:2] >= lb[i])
        @objective(model, Min, dot(c, x))
    end
    @stage 2 begin
        @parameters begin
            q = [-15., -12.]
            T = [-1. 0;
                  0  -1.]
            W = [3. 2.;
                 2. 5.]
        end
        @uncertain ξ₁ ξ₂
        ub = [ξ₁, ξ₂]
        lb = 0.8 * ub
        @variable(model, lb[i] <= y[i in 1:2] <= ub[i])
        @objective(model, Min, dot(q, y))
        @constraint(model, T * x + W * y in MOI.Nonpositives(2))
    end
end

ξ₁ = Scenario(ξ₁ = 6., ξ₂ = 8., probability = 0.5)
ξ₂ = Scenario(ξ₁ = 4., ξ₂ = 4., probability = 0.5)

vec_infeasible_res = SPResult([27.2,41.6], 36.4, 9.2, 27.2, Inf, 9.2, Inf)
push!(problems, (vec_infeasible, [ξ₁,ξ₂], vec_infeasible_res, "Infeasible"))


# This file plots two dimensional interpolation with irregular data

using FractalTools
using GeometryBasics
using Makie

# Construct interpolation data
f = FractalTools.ackley
α = 10
ngon = [BigFloat.([-1., -1.]), BigFloat.([1., -1]), BigFloat.([0., 1.])] * α
npts = 100
pts = getdata(f, ngon, npts)

# Construct interpolant
freevar = 0.001
interp = interpolate(pts, Interp2D(freevar))

# Define error function
erf(x, y) = abs(f(x, y) - interp(x, y))

# Construct test data
tpts = getdata(ngon, 7 * npts)

# Plots
fig = Figure()
ls1 = fig[1, 1] = LScene(fig)
ls2 = fig[1, 2] = LScene(fig)
trisurf!(ls1, tpts, f, meshcolor3=:red)
trisurf!(ls1, tpts, interp, meshcolor3=:blue)
scatter!(ls1, tpts)
trisurf!(ls2, tpts, erf, meshcolor3=:green)


# See [Chapter 9, in Finite Element Procedures, K-J Bathe].
function forced_oscillator(; X0=(zeros(2), zeros(2)), X=Universe(2))
    # Mx'' + Kx = R
    M = [2 0; 0 1.]
    K = [6 -2; -2 4.]
    C = zeros(2, 2)
    R = [0, 10.]
    example_9_1_Bathe = SecondOrderAffineContinuousSystem(M, C, K, R)
    prob = InitialValueProblem(example_9_1_Bathe, X0)
    tspan = (0.0, 20.0)
    return prob, tspan
end

# analytic solution (see Example 9.7 in Bathe)
function forced_oscillator_solution()
    A = [1/√3  (1/2)*√(2/3);
         1/√3      -√(2/3)]
    x₁(t) = (5 / √3) * (1 - cos(t*√2))
    x₂(t) = (2 * √(2/3)) * (-1 + cos(t*√5))
    U(t) = A * [x₁(t), x₂(t)]
end


"""
    EGPnormal(σ, ξ, κ)

*EGPnormal* corresponds to the extended GP model developed by Gamet and Jalbert (2020), with the truncated normal distribution of mean 1 and variance 1/κ^2.

It is a three parameters family: κ controls the shape of the lower tail, σ is a scale parameter, and ξ controls the rate of upper tail decay.

```julia
EGPnormal(σ, ξ, κ)   # EGP of Gamet and Jalbert (2020) with scale parameter σ, rate of upper tail decay ξ and shape of the lower tail κ.

params(d)           # Get the parameters, i.e. (σ, ξ, κ)
```

Reference :

* Unpublished for now
"""
struct EGPnormal{T<:Real} <: ContinuousUnivariateDistribution
    σ::T    # scale parameter
    ξ::T    # rate of upper tail decay
    κ::T    # shape of the lower tail

    function EGPnormal{T}(σ::T, ξ::T, κ::T) where {T <: Real}
        new{T}(σ, ξ, κ)
    end
end

function EGPnormal(σ::T, ξ::T, κ::T; check_args=true) where {T <: Real}
    check_args && @check_args(EGPnormal, σ > zero(σ) && κ >= zero(κ))
    return EGPnormal{T}(σ, ξ, κ)
end

EGPnormal(σ::Real, ξ::Real, κ::Real) = EGPnormal(promote(σ, ξ, κ)...)
EGPnormal(σ::Integer, ξ::Integer, κ::Integer) = EGPnormal(float(σ), float(ξ), float(κ))

minimum(d::EGPnormal) = 0.0
maximum(d::EGPnormal) = Inf * (d.ξ >= 0) - ( d.σ / d.ξ ) * (d.ξ < 0)
insupport(d::EGPnormal, x::Real) = minimum(d) <= x <= maximum(d)

#### Parameters

scale(d::EGPnormal) = d.σ
decay(d::EGPnormal) = d.ξ    # noms à vérifier...
shape(d::EGPnormal) = d.κ

params(d::EGPnormal) = (d.σ, d.ξ, d.κ)
partype(::EGPnormal{T}) where {T} = T

#### Evaluation

function pdf(d::EGPnormal{T}, x::Real) where T<:Real

    (σ, ξ, κ) = params(d)
    η = 1

    if κ == 0

        return Distributions.pdf.(GeneralizedPareto(0, σ, ξ), x)

    else

        return κ/(Distributions.cdf(Normal(), κ*(1-η))-Distributions.cdf(Normal(), -κ*η))*Distributions.pdf(Normal(), κ*(Distributions.cdf(GeneralizedPareto(0, σ, ξ),x)-η))*Distributions.pdf(GeneralizedPareto(0, σ, ξ),x)

    end
end

logpdf(d::EGPnormal, x::Real) = log(pdf(d, x))

function cdf(d::EGPnormal{T}, x::Real) where T<:Real

    (σ, ξ, κ) = params(d)

    return 2/Distributions.erf(κ/sqrt(2)) * (Distributions.cdf(Normal(), κ*(Distributions.cdf(GeneralizedPareto(σ, ξ), x)-1)) - 1 + Distributions.cdf(Normal(), κ))

end

logcdf(d::EGPnormal, x::Real) = log(cdf(d, x))

function quantile(d::EGPnormal{T}, p::Real) where T<:Real
    @assert zero(p)<p<one(p) "the quantile level should be between 0 and 1."

    (σ, ξ, κ) = params(d)

    return Distributions.quantile(GeneralizedPareto(σ,ξ), 1 + 1/κ * Distributions.quantile(Normal(), 0.5*(1 - (1-p)*Distributions.erf(κ/sqrt(2)))))

end

#### Sampling

function rand(rng::AbstractRNG, d::EGPnormal)
    # Generate a Float64 random number uniformly in (0,1].
    u = 1 - rand(rng)

    if u == 1

        return maximum(d)

    else

        return quantile(d, u)

    end
end

sampler(d::EGPnormal) = d


function genKnap(num_items::Int64, range::Int64 = 1000; type::Symbol = :Uncorrelated, seed::Int64 = 42, num_tests::Int64 = 1000)::KnapData
    Random.seed!(seed)

    sum = 0
    items = KnapItem[]
    for i in 1:num_items

        weight = rand(1:range)
        sum += weight

        if (type == :Uncorrelated)
            profit = rand(1:range)
        elseif (type == :WeakCorrelated)
            profit = rand(1:(range ÷ 5 + 1)) + weight - range ÷ 10
            if (profit <= 0) profit = 1 end
        elseif (type == :StrongCorrelated)
            profit = weight + 10
        elseif (type == :SubsetSum)
            profit = weight
        else
            error("Wrong profit correlation type! Use :Uncorrelated, :WeakCorrelated, :StrongCorrelated or :SubsetSum.")
        end

        push!(items, KnapItem(weight, profit))
    end

    capacity = (seed * sum) ÷ (num_tests + 1)
    if (capacity <= range) capacity = range + 1 end

    return KnapData(capacity, items)
end


# a simple LP example
#
#   maximize x + y
#
#   s.t. 50 x + 24 y <= 2400
#        30 x + 33 y <= 2100
#        x >= 45, y >= 5
#
#   solution: x = 45, y = 6.25, objv = 51.25

using Xpress

#method = getparam(env, "Method")
#println("method = $method")

model = Xpress.Model("lp_01", :maximize)

# add variables
add_cvar!(model, 1.0, 45., Inf)  # x
add_cvar!(model, 1.0,  5., Inf)  # y
#update_model!(model)

# add constraints
add_constr!(model, [50., 24.], '<', 2400.)
add_constr!(model, [30., 33.], '<', 2100.)
#update_model!(model)

println(model)

# perform optimization
optimize(model)

# show results
info = get_optiminfo(model)
println()
println(info)

sol = get_solution(model)
println("soln = $(sol)")

objv = get_objval(model)
println("objv = $(objv)")

gc()  # test finalizers