# Normal distribution

import StatsFuns
export Normal

import Base: eltype

# Note: `@measure Normal Lebesgue` expands to (cleaned up)

# quote
#     struct Normal{P, X} <: AbstractMeasure{X}
#         par::P
#     end

#     function Normal(nt::NamedTuple)
#         P = typeof(nt)
#         return Normal{P, eltype(Normal{P})}
#     end

#     Normal(; kwargs...) = Normal((; kwargs...))

#     (baseMeasure(μ::Normal{P, X}) where {P, X}) = Lebesgue{X}

#     (:≪(::Normal{P, X}, ::Lebesgue{X}) where {P, X}) = true

# end
# @measure Normal Lebesgue

@measure Normal(μ,σ) ≃ (1/sqrt2π) * Lebesgue(X)

function logdensity(d::Normal{P} , x::X) where {P <: NamedTuple{(:μ, :σ)}, X}
    return - (log(2) + log(π)) / 2 - log(d.par.σ)  - (x - d.par.μ)^2 / (2 * d.par.σ^2)
end

# Standard normal

function logdensity(d::Normal{EmptyNamedTuple,X} , ::Lebesgue{X}) where {X <: Real}
    return - (log(2) + log(π)) / 2  - x^2 / 2
end

Normal() = Normal(NamedTuple())

# @implement HasDensity{Normal{P,X},X} where {X, P <: NamedTuple{(:μ, :σ)}} begin
#     baseMeasure(d) = Lebesgue(X)
#     logdensity(d, x) = - (log(2) + log(π)) / 2 - log(d.par.σ)  - (x - d.par.μ)^2 / (2 * d.par.σ^2)
# end

# # μ, τ

# eltype(::Type{Normal}, ::Type{NamedTuple{(:μ, :τ), Tuple{A, B}}}) where {A,B} = promote_type(A,B)

# @implement HasDensity{Normal{P,X},X} where {X, P <: NamedTuple{(:μ, :τ)}} begin
#     baseMeasure(d) = Lebesgue(X)
#     logdensity(d, x) = - (log(2) + log(π) - log(d.par.τ)  + d.par.τ * (x - d.par.μ)^2) / 2
# end

# # σ

# eltype(::Type{Normal}, ::Type{NamedTuple{(:σ,), Tuple{B}}}) where {B} = B

# @implement HasDensity{Normal{P,X},X} where {X, P <: NamedTuple{(:σ,)}} begin
#     baseMeasure(d) = Lebesgue(X)
#     logdensity(d, x) = - (log(2) + log(π)) / 2 - log(d.par.σ)  - x^2 / (2 * d.par.σ^2)
# end

# # τ

# eltype(::Type{Normal}, ::Type{NamedTuple{(:τ,), Tuple{B}}}) where {B} = B

# @implement HasDensity{Normal{P,X},X} where {X, P <: NamedTuple{(:τ,)}} begin
#     baseMeasure(d) = Lebesgue(X)
#     logdensity(d, x) = - (log(2) + log(π) - log(d.par.τ)  + d.par.τ * x^2) / 2
# end

# @implement HasRand{Normal{P,X},X} where {X, P <: NamedTuple{(:μ, :σ)}} begin
#     rand(d) = Distributions.rand(Normal(d.par.μ,d.par.σ))
# end

# @implement IsMeasurable{Normal{P,X},S,X} where {X <: Real, S <: Interval{X}, P <: NamedTuple{(:μ, :σ)}} begin
#     measure(d, i) = StatsFuns.normcdf(d.par.μ, d.par.σ, i.last) - StatsFuns.normcdf(d.par.μ, d.par.σ, i.first)
# end

# logdensity(Normal(σ=0.5), 4.0)
#  == logdensity(Normal(μ=0.0, σ=0.5), 4.0)
#  == logdensity(Normal(μ=0.0, τ=4.0), 4.0)
#  == logdensity(Normal(τ=4.0), 4.0)


#== # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
#
#   Compute the beta angle of a satellite.
#
#   The algorithm was based on:
#
#       Mortari, D., Wilkins, M. P., and Bruccoleri, C.  On Sun-Synchronous
#       Orbits and Associated Constellations
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # ==#

export beta_angle

"""
    beta_angle(JD₀::Number, a::Number, e::Number, i::Number, RAAN::Number, Δt::Integer, pert::Symbol = :J2)

Compute the beta angle of an orbit with semi-major axis `a` [m], eccentricity
`e`, inclination `i` [rad], and initial right ascension of the ascending node
`RAAN` [rad]. The orbit epoch, which is also the day in which the analysis will
begin, is `JD₀` [Julian Day]. The analysis will be performed for each day during
`Δt` days.

The argument `pert` can be used to select the perturbation terms that must be
used when propagating the right ascencion of the ascending node. The possible
values are:

* `:J0`: Consider a Keplerian orbit.
* `:J2`: Consider the perturbation terms up to J2.
* `:J4`: Consider the perturbation terms J2, J4, and J2².

If `pert` is omitted, then it defaults to `:J2`.

# Returns

An array with two columns. The first one contains the days of the analysis and
the second one contains the beta angle [rad] for each day.

"""
function beta_angle(JD₀::Number, a::Number, e::Number, i::Number, RAAN::Number,
                    Δt::Integer, pert::Symbol = :J2)

    # Vector of the days in which the beta angle will be computed.
    days = collect(0:1:Δt-1)

    # Output vector.
    β = Vector{Float64}(undef,Δt)

    # RAAN rotation rate [rad/day].
    δΩ = 86400dRAAN(a, e, i, pert)

    # Loop
    @inbounds @views for d in days
        # Compute the RAAN at the day d.
        Ω = RAAN + δΩ*d

        # Compute the versor N represented in the Inertial ref. frame.
        Dio = angle_to_dcm(-i, -Ω, 0, :XZX)
        N_i = Dio[:,3]  # -> Dio*[0;0;1]

        # Compute the Sun position at noon (UT) represented in the Inertial ref.
        # frame.
        S_i = sun_position_i(JD₀+d)
        S_i = S_i/norm(S_i)

        # Compute the beta angle, which is the angle between the Sun vector and
        # the orbit plane.
        β[d+1] = 90 - acosd(dot(N_i,S_i))
    end

    return [days β]
end


module PAM

export pam

function build_phase(D,k)
    N = size(D,1)
    total_dists = [sum(D[:,j]) for j in 1:N]

    # initialize medoids with index of object with shortest distance to all others
    medoids = Int[findmin(total_dists)[2]]
    for j in 1:k-1
        TD = Vector{Float64}(undef,N)
        for a in 1:N
            td = 0.0
            for i in 1:N
                td += reduce(min, (D[i,m] for m in medoids), init=D[a,i])
            end
            TD[a] = td
        end
        push!(medoids, findmin(TD)[2])
    end

    return medoids
end

function swap_phase(D,k,M)
    M = copy(M)
    Mⱼ = similar(M)
    # Perform clustering
    assignments = Int[findmin(view(D, i,M))[2] for i in axes(D,1)]
    cumulative = similar(assignments)
    while true
        # Find minimum sum for each cluster (i.e. find the best medoid)
        for i in 1:k
            cluster = assignments .== i
            cumsum!(cumulative, cluster)
            D_slice = view(D, cluster, cluster)
            distances = sum(@view D_slice[:,i] for i in 1:last(cumulative))
            smallest_distance_idx = findmin(distances)[2]
            Mⱼ[i] = findfirst(==(smallest_distance_idx), cumulative)::Int
        end
        if sort(M) == sort(Mⱼ)
            return (medoids=M,assignments=assignments)
        else
            M,Mⱼ = Mⱼ,M
        end
    end
end

function pam(D,k)
    M = build_phase(D,k)
    return swap_phase(D,k,M)
end

end


# m6.7s.jl
# Clip-06-06-09s.jl

using Pkg, DrWatson

using MonteCarloMeasurements
using StanSample
using StatisticalRethinking

N = 100
df = DataFrame(
    :h0 => rand(Normal(10,2 ), N),
    :treatment => vcat(zeros(Int, Int(N/2)), ones(Int, Int(N/2)))
);
df[!, :fungus] = [rand(Binomial(1, 0.5 - 0.4 * df[i, :treatment]), 1)[1] for i in 1:N]
df[!, :h1] = [df[i, :h0] + rand(Normal(5 - 3 * df[i, :fungus]), 1)[1] for i in 1:N]

stan6_7 = "
data {
    int <lower=1> N;
    vector[N] h0;
    vector[N] h1;
    vector[N] treatment;
    vector[N] fungus;
}
parameters{
    real a;
    real bt;
    real bf;
    real<lower=0> sigma;
}
model {
    vector[N] mu;
    vector[N] p;
    a ~ lognormal(0, 0.2);
    bt ~ normal(0, 0.5);
    bf ~ normal(0, 0.5);
    sigma ~ exponential(1);
    for ( i in 1:N ) {
    p[i] = a + bt*treatment[i] + bf*fungus[i];
    mu[i] = h0[i] * p[i];
    }
    h1 ~ normal(mu, sigma);
}
"

m6_7_data = Dict(
    :N => nrow(df),
    :h0 => df[:, :h0],
    :h1 => df[:, :h1],
    :fungus => df[:, :fungus],
    :treatment => df[:, :treatment]
)

m6_7s = SampleModel("m6.7", stan6_7)
rc6_7s = stan_sample(m6_7s; data=m6_7_data)

if success(rc6_7s)
    part6_7s = read_samples(m6_7s, :particles);
    part6_7s |> display
end

# End of m6.7s.jl


randexp() = Base.Random.randmtzig_exprnd()

immutable ExponentialSampler <: Sampleable{Univariate,Continuous}
	scale::Float64
end

rand(s::ExponentialSampler) = s.scale * randexp()


#
using diffMesh, Plots

ng = 10;
np = 10*ng;

xp = Array(range(-1,1,length=np))'

pou = intervalPOU(ng,x -> @. x^2)

plt = plot()
plt = plot!(xp',pou(xp)',width=3)


using Distributions
using Base.Test

v = 7.0
S = eye(2)
S[1, 2] = S[2, 1] = 0.5

W = Wishart(v,S)
IW = InverseWishart(v,S)

for d in [W,IW]
    @test size(d) == size(rand(d))
    @test length(d) == length(rand(d))
    @test typeof(d)(params(d)...) == d
end

@test_approx_eq_eps mean(rand(W,100000)) mean(W) 0.1
@test_approx_eq_eps mean(rand(IW,100000)) mean(IW) 0.1

v = 3.0

@test_approx_eq_eps pdf(Wishart(v,S), S) 0.04507168 1e-8
@test_approx_eq_eps pdf(Wishart(v,S), inv(S)) 0.01327698 1e-8
@test_approx_eq_eps pdf(Wishart(v,inv(S)),S) 0.0148086 1e-8
@test_approx_eq_eps pdf(Wishart(v,inv(S)),inv(S)) 0.01901462 1e-8

@test_approx_eq_eps pdf(InverseWishart(v,S), S) 0.04507168 1e-8
@test_approx_eq_eps pdf(InverseWishart(v,S), inv(S)) 0.006247377 1e-8
@test_approx_eq_eps pdf(InverseWishart(v,inv(S)),S)  0.03147137 1e-8
@test_approx_eq_eps pdf(InverseWishart(v,inv(S)),inv(S)) 0.01901462 1e-8


using Clp: ClpSolver
using SparseArrays: spzeros

@testset "Minimum-cost flow" begin

    # bipartite oriented + source & sink
    # source 5, sink 6, v1 {1, 2} v2 {3, 4}
    g = lg.DiGraph(6)
    lg.add_edge!(g, 5, 1)
    lg.add_edge!(g, 5, 2)
    lg.add_edge!(g, 3, 6)
    lg.add_edge!(g, 4, 6)
    lg.add_edge!(g, 1, 3)
    lg.add_edge!(g, 1, 4)
    lg.add_edge!(g, 2, 3)
    lg.add_edge!(g, 2, 4)
    w = zeros(6,6)
    w[1,3] = 10.
    w[1,4] = 5.
    w[2,3] = 2.
    w[2,4] = 2.
    # v2 -> sink have demand of one
    demand = spzeros(6,6)
    demand[3,6] = 1
    demand[4,6] = 1
    capacity = ones(6,6)

    flow = mincost_flow(g, spzeros(lg.nv(g)),capacity, w, ClpSolver(), edge_demand=demand,source_nodes=[5],sink_nodes= [6])
    @test flow[5,1] == 1
    @test flow[5,2] == 1
    @test flow[3,6] == 1
    @test flow[4,6] == 1
    @test flow[1,4] == 1
    @test flow[2,4] == 0
    @test flow[2,3] == 1
    @test flow[1,3] == 0
    @test sum(diag(flow)) == 0

    # flow conservation property
    for n1 = 1:4
        @test sum(flow[n1,:]) ≈ sum(flow[:,n1])
    end

    #same test with exact flows
    flow = mincost_flow(g, spzeros(lg.nv(g)),capacity, w, ClpSolver(),edge_demand_exact=true, edge_demand=demand,source_nodes=[5],sink_nodes= [6])
    @test flow[5,1] == 1
    @test flow[5,2] == 1
    @test flow[3,6] == 1
    @test flow[4,6] == 1
    @test flow[1,4] == 1
    @test flow[2,4] == 0
    @test flow[2,3] == 1
    @test flow[1,3] == 0
    @test sum(diag(flow)) == 0

    # flow conservation property
    for n1 = 1:4
        @test sum(flow[n1,:]) ≈ sum(flow[:,n1])
    end

    #--- Same test without edge_demands & positive cost should result in null flow
    flow = mincost_flow(g, spzeros(lg.nv(g)),capacity, w, ClpSolver(),source_nodes=[5],sink_nodes= [6])
    @test flow[5,1] == 0
    @test flow[5,2] == 0
    @test flow[3,6] == 0
    @test flow[4,6] == 0
    @test flow[1,4] == 0
    @test flow[2,4] == 0
    @test flow[2,3] == 0
    @test flow[1,3] == 0
    @test sum(diag(flow)) == 0

    # flow conservation property
    for n1 = 1:4
        @test sum(flow[n1,:]) ≈ sum(flow[:,n1])
    end

    #run sam test again with nodal demands
    node_demand=spzeros(6)
    node_demand[5]=2
    node_demand[6]=-2
    #--- Same test without edge_demands
    flow = mincost_flow(g, node_demand,capacity, w, ClpSolver())
    @test flow[5,1] == 1
    @test flow[5,2] == 1
    @test flow[3,6] == 1
    @test flow[4,6] == 1
    @test flow[1,4] == 1
    @test flow[2,4] == 0
    @test flow[2,3] == 1
    @test flow[1,3] == 0
    @test sum(diag(flow)) == 0

    # flow conservation property
    for n1 = 1:4
        @test sum(flow[n1,:]) ≈ sum(flow[:,n1])
    end

    # no demand => null flow
    d2 = spzeros(6,6)
    flow = mincost_flow(g, spzeros(lg.nv(g)),capacity, w, ClpSolver(), edge_demand=d2,source_nodes=[5],sink_nodes= [6])
    for idx in 1:6
        for jdx in 1:6
            @test flow[idx,jdx] ≈ 0.0
        end
    end

    # graph without sink or source
    g = lg.DiGraph(6)
    # create circuit
    for s in 1:5
        lg.add_edge!(g, s, s+1)
    end
    lg.add_edge!(g, 6, 1)
    lg.add_edge!(g, 2, 5) # shortcut

    capacity = ones(6,6)
    demand = spzeros(6,6)
    demand[1,2] = 1
    costs = ones(6,6)
    flow = mincost_flow(g,spzeros(lg.nv(g)), capacity, costs, ClpSolver(), edge_demand=demand)
    active_flows = [(1,2), (2,5), (5,6),(6,1)]
    for s in 1:6
        for t in 1:6
            if (s,t) in active_flows
                @test flow[s,t] ≈ 1.
            else
                @test flow[s,t] ≈ 0.
            end
        end
    end
    # flow conservation property
    for n1 = 1:6
        @test sum(flow[n1,:]) ≈ sum(flow[:,n1])
    end
    # higher short-circuit cost
    costs[2,5] = 10.
    flow = mincost_flow(g,spzeros(lg.nv(g)), capacity, costs, ClpSolver(), edge_demand=demand)
    active_flows = [(1,2),(2,3),(3,4),(4,5),(5,6),(6,1)]
    for s in 1:6
        for t in 1:6
            if (s,t) in active_flows
                @test flow[s,t] ≈ 1.
            else
                @test flow[s,t] ≈ 0.
            end
        end
    end
        # flow conservation property
    for n1 = 1:6
        @test sum(flow[n1,:]) ≈ sum(flow[:,n1])
    end
end


include("problem.jl")
include("options.jl")

include("debug_options.jl")
include("record_options.jl")

include("cost_plan.jl")
include("gradient_plan.jl")
include("conjugate_gradient_plan.jl")
include("quasi_newton_plan.jl")
include("hessian_plan.jl")
include("proximal_plan.jl")
include("subgradient_plan.jl")
include("particle_swarm_plan.jl")
include("primal_dual_plan.jl")
include("stochastic_gradient_plan.jl")

include("stopping_criterion.jl")
include("stepsize.jl")


#Demo 9: Integration/differentiation
#-------------------------------------------------------------------------------

using CMDimCircuits
CMDimCircuits.@using_CData()

#Get a demo display:
include(CMDimCircuits.demoplotcfgscript)
#Normally use something like:
#CMDimData.@includepkg EasyPlotInspect

#==Constants
===============================================================================#
vvst = cons(:a, labels = set(yaxis="Amplitude (V)", xaxis="Time (s)"))
sigvscosref = cons(:a, labels = set(yaxis="Signal Amplitude (V)", xaxis="Ref cos(t) Value"))
noline = cons(:a, line = set(style=:none))

#==Input data
===============================================================================#
osr = 50 #(fundamental) oversampling ratio
ncycles = 5

#==Computations
===============================================================================#
t = DataF1(0:(2pi/osr):(ncycles*(2pi)))
y = cos(t)
dydt = deriv(y)
∫ydt = iinteg(y)
x2x=xcross1(y - ∫ydt, n=2)
x2y=ycross1(y, ∫ydt, n=2)
println("crossing#2: ($x2x, $x2y)")

#Generate parameter sweeps:
sweeplist = PSweep[
	PSweep("phi", pi*[0, .25, .5, .75])
]

𝜙 = parameter(DataHR, sweeplist, "phi")
ydel = cos(t+𝜙)
yref = y

#==Generate plot
===============================================================================#
plot1 = push!(cons(:plot, vvst, title="Integration/Differentiation"),
	cons(:wfrm, y, label="sig"),
	cons(:wfrm, dydt, label="dy/dt"),
	cons(:wfrm, ∫ydt, label="integ(y)dt"),
	cons(:wfrm, ycross(y, dydt), noline, glyph=set(shape=:o, color=:red), label="crossings"),
)
plot2 = push!(cons(:plot, sigvscosref, title="Lissajous Curves: yvsx(cos(t+ϕ), cos(t))"),
	cons(:wfrm, yvsx(ydel,yref)),
)

pcoll = push!(cons(:plotcoll, title="Integration/Differentiation/yvsx"), plot1, plot2)
	pcoll.ncolumns = 1

#==Display results in pcoll
===============================================================================#
EasyPlot.displaygui(pcoll)

#==Return pcoll to user (call evalfile(...))
===============================================================================#
pcoll


using HPAT
using MPI
using DocOpt
ParallelAccelerator.ParallelIR.PIRHoistParfors(1)
ParallelAccelerator.CGen.setRecreateLoops(true)
ParallelAccelerator.CGen.setRecreateConds(true)
HPAT.DistributedPass.DistOptimize(true)
#HPAT.DistributedPass.set_debug_level(3)

@acc hpat function linear_regression(iterations, N)
    D = 10  # Number of features
    p = 4
    labels = rand(p,N)
    points = rand(D,N)
    w = zeros(p,D)
    alphaN = 0.01/N

    for i in 1:iterations
       w -= alphaN*((w*points)-labels)*points'
    end
    w
end

function main()
    doc = """linear regression statistical method.

Usage:
  linear_regression.jl -h | --help
  linear_regression.jl [--iterations=<iterations>] [--instances=<instances>]

Options:
  -h --help                  Show this screen.
  --iterations=<iterations>  Specify number of iterations; defaults to 20.
  --instances=<instances>    Specify number of instances; defaults to 10^7.
"""
    arguments = docopt(doc)
    iterations = 20
    instances = 10^7

    if (arguments["--iterations"] != nothing)
        iterations = parse(Int, arguments["--iterations"])
    end

    if (arguments["--instances"] != nothing)
        instances = parse(Int, arguments["--instances"])
    end

    W = linear_regression(iterations, instances)

    if MPI.Comm_rank(MPI.COMM_WORLD)==0 println("result = ", W) end
end

main()


using Test
using ModelingToolkit, OrdinaryDiffEq

# Basic electric components
const t = Sym{ModelingToolkit.Parameter{Real}}(:t)
function Pin(;name)
    @variables v(t) i(t)
    ODESystem(Equation[], t, [v, i], [], name=name, defaults=[v=>1.0, i=>1.0])
end

function Ground(name)
    @named g = Pin()
    eqs = [g.v ~ 0]
    ODESystem(eqs, t, [], [], systems=[g], name=name)
end

function ConstantVoltage(name; V = 1.0)
    val = V
    @named p = Pin()
    @named n = Pin()
    @parameters V
    eqs = [
           V ~ p.v - n.v
           0 ~ p.i + n.i
          ]
    ODESystem(eqs, t, [], [V], systems=[p, n], defaults=Dict(V => val), name=name)
end

function Resistor(name; R = 1.0)
    val = R
    @named p = Pin()
    @named n = Pin()
    @variables v(t)
    @parameters R
    eqs = [
           v ~ p.v - n.v
           0 ~ p.i + n.i
           v ~ p.i * R
          ]
    ODESystem(eqs, t, [v], [R], systems=[p, n], defaults=Dict(R => val), name=name)
end

function Capacitor(name; C = 1.0)
    val = C
    @named p = Pin()
    @named n = Pin()
    @variables v(t)
    @parameters C
    D = Differential(t)
    eqs = [
           v ~ p.v - n.v
           0 ~ p.i + n.i
           D(v) ~ p.i / C
          ]
    ODESystem(eqs, t, [v], [C], systems=[p, n], defaults=Dict(C => val), name=name)
end

R = 1.0
C = 1.0
V = 1.0
resistor = Resistor(:resistor, R=R)
capacitor = Capacitor(:capacitor, C=C)
source = ConstantVoltage(:source, V=V)
ground = Ground(:ground)

function connect(ps...)
    eqs = [
           0 ~ sum(p->p.i, ps) # KCL
          ]
    # KVL
    for i in 1:length(ps)-1
        push!(eqs, ps[i].v ~ ps[i+1].v)
    end

    return eqs
end
rc_eqs = [
          connect(source.p, resistor.p)
          connect(resistor.n, capacitor.p)
          connect(capacitor.n, source.n, ground.g)
         ]

rc_model = ODESystem(rc_eqs, t, systems=[resistor, capacitor, source, ground], name=:rc)
sys = structural_simplify(rc_model)
@test !isempty(ModelingToolkit.defaults(sys))
u0 = [
      capacitor.v => 0.0
      capacitor.p.i => 0.0
      resistor.v => 0.0
     ]
prob = ODEProblem(sys, u0, (0, 10.0))
sol = solve(prob, Rodas4())

@test sol[resistor.p.i] == sol[capacitor.p.i]
@test sol[resistor.n.i] == -sol[capacitor.p.i]
@test sol[capacitor.n.i] == -sol[capacitor.p.i]
@test iszero(sol[ground.g.i])
@test iszero(sol[ground.g.v])
@test sol[resistor.v] == sol[source.p.v] - sol[capacitor.p.v]

prob = ODAEProblem(sys, u0, (0, 10.0))
sol = solve(prob, Tsit5())

@test sol[resistor.p.i] == sol[capacitor.p.i]
@test sol[resistor.n.i] == -sol[capacitor.p.i]
@test sol[capacitor.n.i] == -sol[capacitor.p.i]
@test iszero(sol[ground.g.i])
@test iszero(sol[ground.g.v])
@test sol[resistor.v] == sol[source.p.v] - sol[capacitor.p.v]
#using Plots
#plot(sol)


"""
   Solve linear least square problem by LSQR method, the method solve
   min ||Ax-b||₂² if damp = 0 or
   min ||Ax-b||₂² + damp² ||x||₂² otherwise;
   A can be an explicit matrix or a function(linear operator)
   A(x, 1; params...) forward operator
   A(x, 2; params...) adjoint operator; the input and output are standard vector
   atol: the stopping tolerance for the residue norm
   conlim: the limite of the estimated condition number
   params is a dictionary contain the parameters for forward modeling operator

   x: final solution
   rnorm: sqrt(||Ax-b||₂² + damp² ||x||₂²)
   Anorm: the Frobenous norm of A
   Acond: the estimated condition number of A
   diagInv: diag(Dk * Dk'), where Dk is n*k matrix of search directions after k iterations,
          it satisfy Dk' (A'A + damp²* I) Dk = I, so the var is the estimated diagonals
          of the inverse matrix.
"""
function LSQR(A, b::Vector{Tv}; params=[], damp=0.0,
              atol=1e-9, btol=1e-9, conlim=1e8, maxIter=100,
              print_flag=false, diagInv_flag=false) where {Tv<:Number}

    # determine whether A is a matrix or operator
    if isa(A, Matrix)
       (m, n) = size(A)
       explicitA = true
    elseif isa(A, Function)
       m = length(b)
       explicitA = false
    else
       error("A either to be a matrix or operator")
    end

    msg = ["the exact solution is x = 0"
           "Ax-b is small enough given atol and btol"
           "the estimated cond(Abar) has exceeded conlim"
           "the max iterations has been reached"]

    # initialize some intermediate variables
    itn   = 0     # count number of iterations
    istop = 0     # the indicater for stop reason
    ctol  = conlim > 0 ? 1.0/conlim : 0.0
    Anorm = 0.0   # the estimated norm of vcat(A, damp I)
    Acond = 0.0   # the estimated condition number of vcat(A, damp I)
    xnorm = 0.0   # the norm of computed x
    dampsq = damp * damp;
    ddnorm = 0.0; res2 = 0.0; xxnorm = 0.0;
    z      = 0.0; cs2  =-1.0; sn2    = 0.0;

    # set up the first vectors for lanczos bidiagonalization
    # beta * u = b; alpha * v = A' * u
    u    = copy(b)
    beta = norm(u)
    alpha= 0.0

    if beta > 0.0
       u = u / beta
       if explicitA
          v = A' * u
       else
          v = A(u, 2; params...)
       end
       n = length(v)
       x = zeros(eltype(b), n)
       alpha = norm(v)
    end

    if alpha > 0.0
       v = v / alpha
       w = copy(v)
    end

    # estimate the diagonal element of inverse (A'A + dmap^2 I)^(-1)
    if diagInv_flag
       diagInv = zeros(eltype(b), n)
    end

    Arnorm = alpha * beta
    if Arnorm == 0
       println(msg[1])
       return
    end

    rhobar = alpha; phibar = beta ; bnorm  = beta ;
    rnorm  = beta ; r1norm = rnorm; r2norm = rnorm;

    # some header for print info
    test1 = 1.0; test2 = alpha / beta;
    header = "Itn               r1norm               r2norm               norm A               cond A"
    if print_flag
       println("")
       println(header)
       @printf("%3g %20.10e %20.10e %20.10e %20.10e\n", itn, r1norm, r2norm, test1, test2)
    end

     while itn < maxIter

           itn = itn + 1

           # do next step lanczos bidiagonalization
           if explicitA
              u = A * v              - alpha * u
           else
              u = A(v, 1; params...) - alpha * u
           end
           beta = norm(u)

           if beta > 0.0
              u = u / beta
              Anorm = norm([Anorm; alpha; beta; damp])
              if explicitA
                 v = A' * u             - beta * v
              else
                 v = A(u, 2; params...) - beta * v
              end

              alpha = norm(v)
              if alpha > 0.0
                 v = v / alpha
              end
           end

           # plane rotation to eliminate damping parameter
           rhobar1 =   norm([rhobar; damp])
           cs1     =   rhobar  / rhobar1
           sn1     =   damp    / rhobar1
           psi     =   sn1     * phibar
           phibar  =   cs1     * phibar

           # plane rotation
           rho     =   norm([rhobar1; beta])
           cs      =   rhobar1 / rho
           sn      =   beta    / rho
           theta   =   sn      * alpha
           rhobar  = - cs      * alpha
           phi     =   cs      * phibar
           phibar  =   sn      * phibar
           tau     =   sn      * phi

           # update x and w
           t1      =   phi     / rho
           t2      = - theta   / rho
           dk      =   w       / rho
           x       =   x +  t1 * w
           w       =   v +  t2 * w
           ddnorm  =   ddnorm  + norm(dk)^2

           # estimate the diagonal element of Inverse matrix
           if diagInv_flag
              diagInv = diagInv + dk .* dk
           end

           # plane rotation
           delta  =   sn2    * rho
           gambar = - cs2    * rho
           rhs    =   phi    - delta * z
           zbar   =   rhs    / gambar
           xnorm  =   sqrt(xxnorm + zbar^2)
           gamma  =   norm([gambar; theta])
           cs2    =   gambar / gamma
           sn2    =   theta  / gamma
           z      =   rhs    / gamma
           xxnorm =   xxnorm + z^2

           # test for convergence
           Acond  = Anorm * sqrt(ddnorm)
           res1   = phibar^2
           res2   = res2 + psi^2
           rnorm  = sqrt(res1 + res2)
           Arnorm = alpha * abs(tau)

           r1sq   = rnorm^2 - dampsq * xxnorm
           r1norm = sqrt(abs(r1sq)); r1sq = r1sq > 0.0 ? r1sq : -r1sq;
           r2norm = rnorm

           test1  = rnorm  / bnorm
           test2  = Arnorm / (Anorm * rnorm)
           test3  = 1.     / Acond
           t1     = test1  / (1   + Anorm * xnorm / bnorm)
           rtol   = btol   + atol * Anorm * xnorm / bnorm

           istop  = itn        >= maxIter ? 7 : 0
           istop  = 1. + test3 <= 1.      ? 6 : 0
           istop  = 1. + test2 <= 1.      ? 5 : 0
           istop  = 1. + t1    <= 1.      ? 4 : 0
           istop  = test3      <= ctol    ? 3 : 0
           istop  = test2      <= atol    ? 2 : 0
           istop  = test1      <= rtol    ? 1 : 0

           if istop > 0
              break
           end

           if print_flag
              @printf("%3g %20.10e %20.10e %20.10e %20.10e\n", itn, r1norm, r2norm, test1, test2)
           end

     end

     if diagInv_flag
        optout = Dict(:r1norm=>r1norm, :r2norm=>r2norm,
                      :Anorm=>Anorm  , :Acond=>Acond,
                      :Arnorm=>Arnorm, :diagInv=>diagInv)
     else
        optout = Dict(:r1norm=>r1norm, :r2norm=>r2norm,
                      :Anorm=>Anorm  , :Acond=>Acond,
                      :Arnorm=>Arnorm                 )
     end

     return x, optout
end

# the LSQR method for LSRTM
function LSQR(A::Function, b::Ts; params=[], path="NULL", damp=0.0,
              atol=1e-9, btol=1e-9, conlim=1e8, maxIter=100,
              print_flag=false, save_process_flag=false, diagInv_flag=false) where {Ts<:Union{String, Vector{String}}}

    # set the work space
    if path == "NULL"
       path = pwd()
    end

    # number of observation files
    if typeof(b) == String
       num_file = 1
       u = join([path "/u.bin"])
       q = join([path "/q.bin"])

    elseif typeof(b) == Vector{String}
       num_file = length(b)
       path_tmp = joinpath(path, "mu")
       if !isdir(path_tmp)
          mkdir(path_tmp)
          if !isdir(path_tmp)
             error("could not create directory")
          end
       end
       u = Vector{String}(num_file)
       q = Vector{String}(num_file)
       for i = 1 : num_file
           u[i] = join([path_tmp "/u" "_" "$i" ".bin"])
           q[i] = join([path_tmp "/q" "_" "$i" ".bin"])
       end
    end

    v = join([path "/v.bin"])
    x = join([path "/x.bin"])
    w = join([path "/w.bin"])
    p = join([path "/p.bin"])
    dk= join([path "/dk.bin"])

    msg = ["the exact solution is x = 0"
           "Ax-b is small enough given atol and btol"
           "the estimated cond(Abar) has exceeded conlim"
           "the max iterations has been reached"]

    # initialize some intermediate variables
    itn   = 0     # count number of iterations
    istop = 0     # the indicater for stop reason
    ctol  = conlim > 0 ? 1.0/conlim : 0.0
    Anorm = 0.0   # the estimated norm of vcat(A, damp I)
    Acond = 0.0   # the estimated condition number of vcat(A, damp I)
    xnorm = 0.0   # the norm of computed x
    dampsq = damp * damp
    ddnorm = 0.0; res2 = 0.0; xxnorm = 0.0;
    z      = 0.0; cs2  =-1.0; sn2    = 0.0;

    # set up the first vectors for lanczos bidiagonalization
    # beta * u = b; alpha * v = A' * u
    cp(b, u, remove_destination=true)
    beta = norm(u)
    alpha= 0.0

    if beta > 0.0
       scaling!(u, beta)

       if typeof(params) <: Dict{Symbol, Any}
          params[:path_m] = v
       elseif typeof(params) == Vector{Dict}
         for i = 1 : length(params)
             params[i][:path_m] = v
         end
       end
       v = A(u, 2; params=params)

       # x = zeros(v)
       (h, d) = read_USdata(v)
       write_USdata(x, h, zeros(d))
       alpha = norm(v)
    end

    if alpha > 0.0
       scaling!(v, alpha)
       cp(v, w, remove_destination=true)
    end

    # estimate the diagonal element of inverse (A'A + dmap^2 I)^(-1)
    if diagInv_flag
       diagInv = join([path "/diagInv.bin"])
       (h, d) = read_USdata(v)
       write_USdata(diagInv, h, zeros(d))
    end

    Arnorm = alpha * beta
    if Arnorm == 0
       println(msg[1])
       return
    end

    rhobar = alpha; phibar = beta ; bnorm  = beta ;
    rnorm  = beta ; r1norm = rnorm; r2norm = rnorm;

    # some header for print info
    test1 = 1.0; test2 = alpha / beta;
    header = "Itn               r1norm               r2norm               norm A               cond A"
    if print_flag
       println("")
       println(header)
       @printf("%3g %20.10e %20.10e %20.10e %20.10e\n", itn, r1norm, r2norm, test1, test2)
    end

     while itn < maxIter

           itn = itn + 1

           # do next step lanczos bidiagonalization; forward
           if typeof(params) <: Dict{Symbol, Any}
              params[:path_fwd] = q
           elseif typeof(params) == Vector{Dict}
              for i = 1 : length(params)
                  params[i][:path_fwd] = q[i]
              end
           end
           q = A(v, 1; params=params)
           x_plus_alpha_q!(q, -alpha, u; iflag=2)
           beta = norm(u)

           if beta > 0.0
              scaling!(u, beta)
              Anorm = norm([Anorm; alpha; beta; damp])

              # adjoint
              if typeof(params) <: Dict{Symbol, Any}
                 params[:path_m] = p
              elseif typeof(params) == Vector{Dict}
                 for i = 1 : length(params)
                     params[i][:path_m] = p
                 end
              end
              p = A(u, 2; params=params)
              x_plus_alpha_q!(p, -beta, v; iflag=2)

              alpha = norm(v)
              if alpha > 0.0
                 scaling!(v, alpha)
              end
           end

           # plane rotation to eliminate damping parameter
           rhobar1 =   norm([rhobar; damp])
           cs1     =   rhobar  / rhobar1
           sn1     =   damp    / rhobar1
           psi     =   sn1     * phibar
           phibar  =   cs1     * phibar

           # plane rotation
           rho     =   norm([rhobar1; beta])
           cs      =   rhobar1 / rho
           sn      =   beta    / rho
           theta   =   sn      * alpha
           rhobar  = - cs      * alpha
           phi     =   cs      * phibar
           phibar  =   sn      * phibar
           tau     =   sn      * phi

           # update x and w
           t1      =   phi     / rho
           t2      = - theta   / rho
           scaling!(dk, w, rho)
           x_plus_alpha_q!(x, t1, w; iflag=1)
           x_plus_alpha_q!(v, t2, w; iflag=2)
           ddnorm  =   ddnorm  + norm(dk)^2

           # save the result of each iteration
           if save_process_flag
              path_process = join([path "/iter" "_" "$itn" ".bin"])
              cp(x, path_process, remove_destination=true)
           end

           # estimate the diagonal element of Inverse matrix
           if diagInv_flag
              diagInv = diagInv + dk .* dk
           end

           # plane rotation
           delta  =   sn2    * rho
           gambar = - cs2    * rho
           rhs    =   phi    - delta * z
           zbar   =   rhs    / gambar
           xnorm  =   sqrt(xxnorm + zbar^2)
           gamma  =   norm([gambar; theta])
           cs2    =   gambar / gamma
           sn2    =   theta  / gamma
           z      =   rhs    / gamma
           xxnorm =   xxnorm + z^2

           # test for convergence
           Acond  = Anorm * sqrt(ddnorm)
           res1   = phibar^2
           res2   = res2 + psi^2
           rnorm  = sqrt(res1 + res2)
           Arnorm = alpha * abs(tau)

           r1sq   = rnorm^2 - dampsq * xxnorm
           r1norm = sqrt(abs(r1sq)); r1sq = r1sq > 0.0 ? r1sq : -r1sq;
           r2norm = rnorm

           test1  = rnorm  / bnorm
           test2  = Arnorm / (Anorm * rnorm)
           test3  = 1.     / Acond
           t1     = test1  / (1   + Anorm * xnorm / bnorm)
           rtol   = btol   + atol * Anorm * xnorm / bnorm

           istop  = itn        >= maxIter ? 7 : 0
           istop  = 1. + test3 <= 1.      ? 6 : 0
           istop  = 1. + test2 <= 1.      ? 5 : 0
           istop  = 1. + t1    <= 1.      ? 4 : 0
           istop  = test3      <= ctol    ? 3 : 0
           istop  = test2      <= atol    ? 2 : 0
           istop  = test1      <= rtol    ? 1 : 0

           if istop > 0
              break
           end

           if print_flag
              @printf("%3g %20.10e %20.10e %20.10e %20.10e\n", itn, r1norm, r2norm, test1, test2)
           end

     end

     if diagInv_flag
        optout = Dict(:r1norm=>r1norm, :r2norm=>r2norm,
                      :Anorm=>Anorm  , :Acond=>Acond,
                      :Arnorm=>Arnorm, :diagInv=>diagInv)
     else
        optout = Dict(:r1norm=>r1norm, :r2norm=>r2norm,
                      :Anorm=>Anorm  , :Acond=>Acond,
                      :Arnorm=>Arnorm                 )
     end

     return x, optout
end

# ==========================support functions for CGLS==========================
function cp(ps::Ts, pd::Ts; remove_destination=false) where {Ts<:Vector{String}}
    if length(ps) != length(pd)
       error("number of file does not match")
    end
    for i = 1 : length(ps)
        cp(ps[i], pd[i], remove_destination=remove_destination)
    end
    return nothing
end

function norm(s::Ts) where {Ts<:Union{String, Vector{String}}}

    if typeof(s) == String
       (hdr, d) = read_USdata(s)
       return vecnorm(d)
    elseif typeof(s) == Vector{String}
       r = 0.0
       for i = 1 : length(s)
           (hdr, d) = read_USdata(s[i])
           tmp = vecnorm(d)
           r = r + tmp * tmp
       end
       return sqrt(r)
    end
end

function x_plus_alpha_q!(x::Ts, alpha::Tv, q::Ts; iflag=1) where {Ts<:Union{String, Vector{String}}, Tv<:AbstractFloat}

    if typeof(x) == String
       (hdr_x, dx) = read_USdata(x)
       (hdr_q, dq) = read_USdata(q)
       for i = 1 : length(dx)
           dx[i] = dx[i] + alpha * dq[i]
       end

       if iflag == 1
          write_USdata(x, hdr_x, dx)
       elseif iflag == 2
          write_USdata(q, hdr_x, dx)
       end

    elseif typeof(x) == Vector{String}
       if length(x) != length(q)
          error("number of file does not match")
       end
       for i = 1 : length(x)
           (hdr_x, dx) = read_USdata(x[i])
           (hdr_q, dq) = read_USdata(q[i])
           for j = 1 : length(dx)
               dx[j] = dx[j] + alpha * dq[j]
           end

           if iflag == 1
              write_USdata(x[i], hdr_x, dx)
           elseif iflag == 2
              write_USdata(q[i], hdr_x, dx)
           end
       end
    end

    return nothing
end

function scaling!(u::Ts, beta::Tv) where {Ts<:Union{String, Vector{String}}, Tv<:AbstractFloat}

    beta_inv = 1.0 / beta

    if typeof(u) == String
       (hdr, d) = read_USdata(u)
       for i = 1 : length(d)
           d[i] = d[i] * beta_inv
       end
       write_USdata(u, hdr, d)

    elseif typeof(u) == Vector{String}

       for i = 1 : length(u)
           (hdr, d) = read_USdata(u[i])

           for j = 1 : length(d)
               d[j] = d[j] * beta_inv
           end

           write_USdata(u[i], hdr, d)
       end
    end
    return nothing
end

function scaling!(a::Ts, u::Ts, beta::Tv) where {Ts<:Union{String, Vector{String}}, Tv<:AbstractFloat}

    beta_inv = 1.0 / beta

    if typeof(u) == String
       (hdr, d) = read_USdata(u)
       for i = 1 : length(d)
           d[i] = d[i] * beta_inv
       end
       write_USdata(a, hdr, d)

    elseif typeof(u) == Vector{String}

       for i = 1 : length(u)
           (hdr, d) = read_USdata(u[i])

           for j = 1 : length(d)
               d[j] = d[j] * beta_inv
           end

           write_USdata(a[i], hdr, d)
       end
    end
    return nothing
end


# # Conditional Maximum Likelihood for the Rasch Model
#
#-
#md # !!! tip
#md #     This example is also available as a Jupyter notebook:
#md #     [``rasch.ipynb``](@__NBVIEWER_ROOT_URL__examples/generated/rasch.ipynb)
#-
using Optim, Random #hide
#
# The Rasch model is used in psychometrics as a model for
# assessment data such as student responses to a standardized
# test. Let $X_{pi}$ be the response accuracy of student $p$
# to item $i$ where $X_{pi}=1$ if the item was answered correctly
# and $X_{pi}=0$ otherwise for $p=1,\ldots,n$ and $i=1,\ldots,m$.
# The model for this accuracy is
# ```math
#   P(\mathbf{X}_{p}=\mathbf{x}_{p}|\xi_p, \mathbf\epsilon) = \prod_{i=1}^m \dfrac{(\xi_p \epsilon_j)^{x_{pi}}}{1 + \xi_p\epsilon_i}
# ```
# where $\xi_p > 0$ the latent ability of person $p$ and $\epsilon_i > 0$
# is the difficulty of item $i$.

# We simulate data from this model:

Random.seed!(123)
n = 1000
m = 5
theta = randn(n)
delta = randn(m)
r = zeros(n)
s = zeros(m)

for i in 1:n
  p = exp.(theta[i] .- delta) ./ (1.0 .+ exp.(theta[i] .- delta))
  for j in 1:m
    if rand() < p[j] ##correct
      r[i] += 1
      s[j] += 1
    end
  end
end
f = [sum(r.==j) for j in 1:m];

# Since the number of parameters increases
# with sample size standard maximum likelihood will not provide us
# consistent estimates. Instead we consider the conditional likelihood.
# It can be shown that the Rasch model is an exponential family model and
# that the sum score $r_p = \sum_{i} x_{pi}$ is the sufficient statistic for
# $\xi_p$. If we condition on the sum score we should be able to eliminate
# $\xi_p$. Indeed, with a bit of algebra we can show
# ```math
# P(\mathbf{X}_p = \mathbf{x}_p | r_p, \mathbf\epsilon) = \dfrac{\prod_{i=1}^m \epsilon_i^{x{ij}}}{\gamma_{r_i}(\mathbf\epsilon)}
# ```
# where $\gamma_r(\mathbf\epsilon)$ is the elementary symmetric function of order $r$
# ```math
# \gamma_r(\mathbf\epsilon) = \sum_{\mathbf{y} : \mathbf{1}^\intercal \mathbf{y} = r} \prod_{j=1}^m \epsilon_j^{y_j}
# ```
# where the sum is over all possible answer configurations that give a sum
# score of $r$. Algorithms to efficiently compute $\gamma$ and its
# derivatives are available in the literature (see eg Baker (1996) for a review
# and Biscarri (2018) for a more modern approach)

function esf_sum!(S::AbstractArray{T,1}, x::AbstractArray{T,1}) where T <: Real
  n = length(x)
  fill!(S,zero(T))
  S[1] = one(T)
  @inbounds for col in 1:n
    for r in 1:col
      row = col - r + 1
      S[row+1] = S[row+1] + x[col] * S[row]
    end
  end
end

function esf_ext!(S::AbstractArray{T,1}, H::AbstractArray{T,3}, x::AbstractArray{T,1}) where T <: Real
  n = length(x)
  esf_sum!(S, x)
  H[:,:,1] .= zero(T)
  H[:,:,2] .= one(T)

  @inbounds for i in 3:n+1
    for j in 1:n
      H[j,j,i] = S[i-1] - x[j] * H[j,j,i-1]
      for k in j+1:n
        H[k,j,i] = S[i-1] - ((x[j]+x[k])*H[k,j,i-1] + x[j]*x[k]*H[k,j,i-2])
        H[j,k,i] = H[k,j,i]
      end
    end
  end
end

# The objective function we want to minimize is the negative log conditional
# likelihood
# ```math
# \begin{aligned}
# \log{L_C(\mathbf\epsilon|\mathbf{r})} &= \sum_{p=1}^n \sum_{i=1}^m x_{pi} \log{\epsilon_i} - \log{\gamma_{r_p}(\mathbf\epsilon)}\\
#   &= \sum_{i=1}^m s_i \log{\epsilon_i} - \sum_{r=1}^m f_r \log{\gamma_r(\mathbf\epsilon)}
# \end{aligned}
# ```
ϵ = ones(Float64, m)
β0 = zeros(Float64, m)
last_β = fill(NaN, m)
S = zeros(Float64, m+1)
H = zeros(Float64, m, m, m+1)

function calculate_common!(x, last_x)
  if x != last_x
    copyto!(last_x, x)
    ϵ .= exp.(-x)
    esf_ext!(S, H, ϵ)
  end
end
function neglogLC(β)
  calculate_common!(β, last_β)
  return -s'log.(ϵ) + f'log.(S[2:end])
end

# Parameter estimation is usually performed with respect to the unconstrained parameter
# $\beta_i = -\log{\epsilon_i}$. Taking the derivative with respect to $\beta_i$
# (and applying the chain rule) one obtains
# ```math
#   \dfrac{\partial\log L_C(\mathbf\epsilon|\mathbf{r})}{\partial \beta_i} = -s_i + \epsilon_i\sum_{r=1}^m \dfrac{f_r \gamma_{r-1}^{(j)}}{\gamma_r}
# ```
# where $\gamma_{r-1}^{(i)} = \partial \gamma_{r}(\mathbf\epsilon)/\partial\epsilon_i$.

function g!(storage, β)
  calculate_common!(β, last_β)
  for j in 1:m
    storage[j] = s[j]
    for l in 1:m
      storage[j] -= ϵ[j] * f[l] * (H[j,j,l+1] / S[l+1])
    end
  end
end

# Similarly the Hessian matrix can be computed
# ```math
#   \dfrac{\partial^2 \log L_C(\mathbf\epsilon|\mathbf{r})}{\partial \beta_i\partial\beta_j} = \begin{dcases} -\epsilon_i \sum_{r=1}^m \dfrac{f_r\gamma_{r-1}^{(i)}}{\gamma_r}\left(1 - \dfrac{\gamma_{r-1}^{(i)}}{\gamma_r}\right) & \text{if $i=j$}\\
#     -\epsilon_i\epsilon_j\sum_{r=1}^m \dfrac{f_r \gamma_{r-2}^{(i,j)}}{\gamma_r} - \dfrac{f_r\gamma_{r-1}^{(i)}\gamma_{r-1}^{(j)}}{\gamma_r^2} &\text{if $i\neq j$}
#    \end{dcases}
# ```
# where $\gamma_{r-2}^{(i,j)} = \partial^2 \gamma_{r}(\mathbf\epsilon)/\partial\epsilon_i\partial\epsilon_j$.

function h!(storage, β)
  calculate_common!(β, last_β)
  for j in 1:m
    for k in 1:m
      storage[k,j] = 0.0
      for l in 1:m
        if j == k
          storage[j,j] += f[l] * (ϵ[j]*H[j,j,l+1] / S[l+1]) *
            (1 - ϵ[j]*H[j,j,l+1] / S[l+1])
        elseif k > j
          storage[k,j] += ϵ[j] * ϵ[k] * f[l] *
            ((H[k,j,l] / S[l+1]) - (H[j,j,l+1] * H[k,k,l+1]) / S[l+1] ^ 2)
        else #k < j
          storage[k,j] += ϵ[j] * ϵ[k] * f[l] *
            ((H[j,k,l] / S[l+1]) - (H[j,j,l+1] * H[k,k,l+1]) / S[l+1] ^ 2)
        end
      end
    end
  end
end

# The estimates of the item parameters are then obtained via standard optimization
# algorithms (either Newton-Raphson or L-BFGS). One last issue is that the model is
# not identifiable (multiplying the $\xi_p$ by a constant and dividing the $\epsilon_i$
# by the same constant results in the same likelihood). Therefore some kind of constraint
# must be imposed when estimating the parameters. Typically either $\epsilon_1 = 0$ or
# $\prod_{i=1}^m \epsilon_i = 1$ (which is equivalent to $\sum_{i=1}^m \beta_i = 0$).

con_c!(c, x) = (c[1] = sum(x); c)
function con_jacobian!(J, x)
  J[1,:] .= ones(length(x))
end
function con_h!(h, x, λ)
    for i in 1:size(h)[1]
        for j in 1:size(h)[2]
            h[i,j] += (i == j) ? λ[1] : 0.0
        end
    end
end
lx = Float64[]; ux = Float64[]
lc = [0.0]; uc = [0.0]
df = TwiceDifferentiable(neglogLC, g!, h!, β0)
dfc = TwiceDifferentiableConstraints(con_c!, con_jacobian!, con_h!, lx, ux, lc, uc)
res = optimize(df, dfc, β0, IPNewton())

# Compare the estimate to the truth
delta_hat = res.minimizer
[delta delta_hat]


# Test creation and evaluation of SPNs

@testset "Defining and evaluating discrete SPNs" begin
    import SPNetworks: ncircuits, scopes
    import SPNetworks: plogpdf

    @testset "Simple DAG SPN" begin
        SPN = SumProductNetwork(
            [
                SumNode([2,3,4],[0.2,0.5,0.3]),        # 1
                ProductNode([5,7]),                    # 2
                ProductNode([5,8]),                    # 3
                ProductNode([6,8]),                    # 4
                CategoricalDistribution(1, [0.6,0.4]), # 5
                CategoricalDistribution(1, [0.1,0.9]), # 6
                CategoricalDistribution(2, [0.3,0.7]), # 7
                CategoricalDistribution(2, [0.8,0.2])  # 8
            ]
        )
        @test length(SPN) == 8
        @test ncircuits(SPN) == 3
        @test nparams(SPN) == 11
        sc = scope(SPN)
        # println("Scope: ", map(string,sc))
        @test (length(sc) == 2) && (1 in sc) && (2 in sc)
        @testset "Evaluate node scopes" begin
            scl = scopes(SPN)
            @test length(scl) == 8
            @test (length(scl[1]) == 2) && (1 in scl[1]) && (2 in scl[1])
            @test (length(scl[2]) == 2) && (1 in scl[2]) && (2 in scl[2])
            @test (length(scl[3]) == 2) && (1 in scl[3]) && (2 in scl[3])
            @test (length(scl[4]) == 2) && (1 in scl[4]) && (2 in scl[4])
            @test (length(scl[5]) == 1) && (1 in scl[5])
            @test (length(scl[6]) == 1) && (1 in scl[6])
            @test (length(scl[7]) == 1) && (2 in scl[7])
            @test (length(scl[8]) == 1) && (2 in scl[8])
        end
        results = [0.3, 0.15, 0.4, 0.15]
        @testset "Evaluation at $a,$b" for a=1:2, b=1:2
            @test SPN(a,b) ≈ results[2*(a-1) + b]
            @test logpdf(SPN,[a,b]) ≈ log(results[2*(a-1) + b])
            @test plogpdf(SPN,[a,b]) ≈ log(results[2*(a-1) + b])
        end
        @testset "Marginalization" begin
            @test SPN(1,NaN) ≈ 0.45
            @test SPN(NaN,2) ≈ 0.3
            @test logpdf(SPN,[NaN,2]) ≈ log(0.3)
            @test plogpdf(SPN,[NaN,2]) ≈ log(0.3)
        end
        @testset "Sampling" begin
            x = rand(SPN)
            @test length(x) == 2
            #N = 10000
            N = 1000
            data = rand(SPN,N)
            counts = [0 0; 0 0]
            for i = 1:N
                counts[Int(data[i,1]), Int(data[i,2])] += 1
            end
            @testset "Verifing empirical estimates at $a,$b" for a in 1:2, b in 1:2
                ref = SPN([a,b])
                # println("S($a,$b) = $ref ≈ ", counts[a,b]/N)
                @test ref ≈ counts[a,b]/N atol=0.1
            end
        end
    end # end of DAG SPN testset
    @testset "DAG SPN encoding HMM" begin
        HMM = SumProductNetwork(
            [
            SumNode([2,3],[0.3,0.7]),               # 1
            ProductNode([4,5]),                     # 2
            ProductNode([6,7]),                     # 3
            CategoricalDistribution(1,[0.3,0.7]),   # 4 (D1)
            SumNode([8,9],[0.5,0.5]),               # 5
            SumNode([8,9],[0.2,0.8]),               # 6
            CategoricalDistribution(1,[0.8,0.2]),   # 7 (D2)
            ProductNode([10,11]),                   # 8
            ProductNode([12,13]),                   # 9
            CategoricalDistribution(2,[0.4,0.6]),   # 10 (D3)
            SumNode([14,15],[0.6,0.4]),             # 11
            SumNode([14,15],[0.4,0.6]),             # 12
            CategoricalDistribution(2,[0.25,0.75]), # 13 (D4)
            CategoricalDistribution(3,[0.9,0.1]),   # 14 (D5)
            CategoricalDistribution(3,[0.42,0.58])  # 15 (D6)
            ]
        )
        @test length(HMM) == 15
        @test ncircuits(HMM) == 8
        @test length(scope(HMM)) == 3
        @testset "Evaluate node scopes" begin
            scl = scopes(HMM)
            @test length(scl) == 15
            @test (length(scl[1]) == 3) && (1 in scl[1]) && (2 in scl[1]) && (3 in scl[1])
            @test (length(scl[5]) == 2) && (2 in scl[5]) && (3 in scl[5])
            @test (length(scl[6]) == 2) && (2 in scl[6]) && (3 in scl[6])
            @test (length(scl[8]) == 2) && (2 in scl[8]) && (3 in scl[8])
            @test (length(scl[9]) == 2) && (2 in scl[9]) && (3 in scl[9])
        end
        results = [0.11989139999999997,0.06615860000000003,0.29298060000000004,
        0.1709694,0.0708666,0.03658340000000001,0.1561014,0.08644860000000001]
        @testset "Evaluating HMM at $a,$b,$c" for a=1:2, b=1:2, c=1:2
            v = HMM([a,b,c])
            @test v ≈ results[4*(a-1) + 2*(b-1) + c]
            @test plogpdf(HMM, [a,b,c]) ≈ log(results[4*(a-1) + 2*(b-1) + c])
        end
        # println("HMM() ≈ $(HMM([NaN,NaN,NaN]))")
        @test logpdf(HMM,[NaN,NaN,NaN]) ≈ 0.0
        @test plogpdf(HMM,[NaN,NaN,NaN]) ≈ 0.0
        # println("HMM(X1=1) ≈ $(HMM([1,NaN,NaN]))")
        @test logpdf(HMM,[1,NaN,NaN]) ≈ log(0.65)
        @test plogpdf(HMM,[1,NaN,NaN]) ≈ log(0.65)
        x = rand(HMM)
        @test length(x) == 3
        x[3] = NaN
        # project/prune network
        prunedHMM = SPNetworks.project(HMM, Set([1]), x)
        @testset "Projection" for a = 1:2
            x[1] = a
            @test HMM(x) ≈ prunedHMM(x)
        end
        x = [1.0, 1.0, 2.0]
        prunedHMM = SPNetworks.project(HMM, Set([1]), x)
        @testset "Projection" for a = 1:2
            x[1] = a
            @test HMM(x) ≈ prunedHMM(x)
        end
    end # end of HMM testset
    @testset "Selective SPN" begin
        selSPN = SumProductNetwork(
            [
            SumNode([2,3],[0.4,0.6]),              # 1
            ProductNode([4,5,6]),                  # 2
            ProductNode([7,8,9]),                  # 3
            IndicatorFunction(1,2.0),              # 4
            CategoricalDistribution(2,[0.3,0.7]),  # 5
            CategoricalDistribution(3,[0.4,0.6]),  # 6
            CategoricalDistribution(2,[0.8,0.2]),  # 7
            CategoricalDistribution(3,[0.9,0.1]),  # 8
            IndicatorFunction(1,1.0)               # 9
            ]
        )
        @test length(selSPN) == 9
        @test ncircuits(selSPN) == 2
        @test length(scope(selSPN)) == 3
        results = [0.432,0.048,0.108,0.012,0.048,0.072,0.112,0.168]
        @testset "Evaluating Sel SPN at $a,$b,$c" for a=1:2, b=1:2, c=1:2
            v = selSPN(Float64[a,b,c])
            # println("SEL($a,$b,$c) = $v")
            @test v ≈ results[4*(a-1) + 2*(b-1) + c]
            @test plogpdf(selSPN, [a,b,c]) ≈ log(results[4*(a-1) + 2*(b-1) + c])
        end
        # println("- Selective SPN")
        value = selSPN([1,NaN,NaN])
        # println("S2(A=1) = $value")
        @test value ≈ 0.6
        @test logpdf(selSPN,[2,NaN,NaN]) ≈ log(0.4)
        @test plogpdf(selSPN,[2,NaN,NaN]) ≈ log(0.4)
        # println()
        x = rand(selSPN)
        @test length(x) == 3
    end # end of selective SPN testset
    @testset "SPN encoding PSDD" begin
        # taken from https://github.com/UCLA-StarAI/Circuit-Model-Zoo/blob/master/psdds/little_4var.psdd
        PSDD = SumProductNetwork(
            [
                ProductNode([2,3]),                    # 1 (10)
                SumNode([7,6,5,4],exp.([-1.6094379124341003,-1.2039728043259361,-0.916290731874155,-2.3025850929940455])),     # 2 (9)
                SumNode([11,10,9,8],exp.([-2.3025850929940455,-2.3025850929940455,-2.3025850929940455,-0.35667494393873245])), # 3 (8)
                ProductNode([14,12]), # 4
                ProductNode([14,13]), # 5
                ProductNode([15,12]), # 6
                ProductNode([15,13]), # 7
                ProductNode([18,16]), # 8
                ProductNode([18,17]), # 9
                ProductNode([19,16]), # 10
                ProductNode([19,17]), # 11
                IndicatorFunction(4, 1.0), # 12
                IndicatorFunction(4, 2.0), # 13
                IndicatorFunction(3, 1.0), # 14
                IndicatorFunction(3, 2.0), # 15
                IndicatorFunction(2, 1.0), # 16
                IndicatorFunction(2, 2.0), # 17
                IndicatorFunction(1, 1.0), # 18
                IndicatorFunction(1, 2.0)  # 19
            ]
        )
        @test length(PSDD) == 19
        @test ncircuits(PSDD) == 16
        @test length(scope(PSDD)) == 4
        results = [0.07, 0.27999999999999997, 0.20999999999999996, 0.14, 0.010000000000000005, 0.04000000000000001, 0.029999999999999995, 0.02000000000000001, 0.010000000000000005, 0.04000000000000001, 0.029999999999999995, 0.02000000000000001, 0.010000000000000005, 0.04000000000000001, 0.029999999999999995, 0.02000000000000001]
        @testset "Evaluating PSDD at $a,$b,$c,$d" for a=1:2, b=1:2, c=1:2, d=1:2
            v = PSDD([a,b,c,d])
            # println("PSDD($a,$b,$c,$d) = $v")
            @test v ≈ results[8*(a-1) + 4*(b-1) + 2*(c-1) + d]
            @test plogpdf(PSDD, [a,b,c,d]) ≈ log(results[8*(a-1) + 4*(b-1) + 2*(c-1) + d])
        end
        @testset "Sampling" begin
            x = rand(PSDD)
            @test length(x) == 4
        end
    end # end of PSDD test
end # END of discrete testset


real_checked(A::AbstractArray) = A
function real_checked(A::AbstractArray{Complex{T}}) where T
    epsT = eps(real(eltype(A)))
    rtol = 1000epsT
    atol = 100epsT

    if any(abs(imag(x)) > rtol * abs(x) && abs(imag(x)) > atol for x in A)
        relerror = imag(A) ./ abs.(A)
        relerror[map(x -> abs(imag(x)) < atol, A)] .= 0
        @warn "Large imaginary part" norm(relerror)
    end

    real(A)
end

"""
Zero all elements of a Fourier space array which have no complex-conjugate partner
and may thus lead to an imaginary component in real space (after an iFFT).
"""
function clear_without_conjugate!(A::AbstractArray{T, 3}) where {T <: Complex}
    fft_size = size(A)
    # For even-length grids the element with largest negative integer coordinate
    # in G_vectors has no complex-conjugated partner. In order to avoid having
    # a spurious imaginary component in real-space, we set these elements to zero
    iseven(fft_size[1]) && (A[div(end, 2) + 1, :, :] .= 0)
    iseven(fft_size[2]) && (A[:, div(end, 2) + 1, :] .= 0)
    iseven(fft_size[3]) && (A[:, :, div(end, 2) + 1] .= 0)

    A
end


#POTENTIAL ISSUE - line 22 may throw an Inexact Error for large problems

# Parameters:
#	N::Int
#	- N is the size of P
#	P::Array{Int, 1}
#	- P is a Partition
#Return Value:
#	NST::Int
#	- the number of Standard Tableau of the Young Diagram of P
function degree(N::Int, P::Array{Int, 1})
	if N > 20
		NST = 2432902008176640000.0
		for i = 1:length(P)
			for j = 1:P[i]
				NST /= hook_length(P, i, j)
			end
		end
		for n = 21:N
			NST *= n
		end
		NST = round(Int,NST)
		return NST
	else
		NST = factorial(N)
		NST /= hook_product(P)
		NST = round(Int,NST)
		return NST
	end
end

# Parameters:
#	P::Array{Int, 1}
#	- P is a Partition
# Return Values:
#	HP::Int
#	- the Hookproduct of the Young Diagram of P
function hook_product(P::Array{Int, 1})
	HP = 1
	for i = 1:length(P)
		for j = 1:P[i]
			HP *= hook_length(P, i, j)
		end
	end
	return HP
end

# Parameters:
#	P::Array{Int, 1}
#	- P is a Partition
#	i::Int
#	- the row index
#	j::Int
#	- the column index
# Return Value
#	HL::Int
#	- the HookLength of the box in the ith row and jth colomn of the Young Diagram of P
function hook_length(P::Array{Int, 1}, i::Int, j::Int)
	R = P[i] - j
	B = 0
	i += 1
	while i <= length(P) && j <= P[i]
		i += 1
		B += 1
	end
	HL = 1 + R + B
	return HL
end


using SpecialFunctions

"""
    black(F,K,s)

The Black function of standardised arguments:

    black(F,K,s) = F ⋅ Φ(log(F/K)/s + s/2) - K ⋅ Φ(log(F/K)/s - s/2)

This is the Black-Scholes price of a call option with spot price `F`, strike `K` and volatility `s`, assuming the discount rate `r` and carry rate `q` are zero, and the time to maturity is one.

This is monotonic increasing in `s`, with

    black(K,F,0) == max(F-K, 0)
    black(K,F,∞) == F

The corresponding put price is:

    black(K,F,s)

and so put-call parity implies:

    black(F,K,s) + K = black(K,F,s) + F

`black(F,K,s)` is the expectation of `max(exp(Y)-K,0)` where `Y` has a `Normal(log(F)-s^2/2, s)` distribution.

This can be used to compute the complementary CDF of the inverse Gaussian distribution with parameters `(λ,μ)` at `x` via,

    P[X > x] = black(1, exp(2λ/μ), 2*sqrt(λ/x))

This is largely based on the method of Jäckel (2015).

References

* Jäckel, P. (2015), Let's Be Rational. Wilmott, 2015(75): 40–53. doi:10.1002/wilm.10395

"""
function black(F,K,s)
    F > zero(F) && K > zero(K) && s ≥ 0 || throw(DomainError())
    if s == 0
        max(F-K,zero(K))
    else
        if K >= F
            x = -log(K/F)
            h = x/s
            t = s/2
            F * _black(h,t,K/F)
        else
            # use put-call parity
            x = -log(F/K)
            h = x/s
            t = s/2
            (F-K) + K * _black(h,t,F/K)
        end
    end
end

"""
    _black(h, t, K = exp(-2*h*t))

The black kernel function:

    Φ(h + t) - K*Φ(h - t)

where `K = exp(-2*h*t)` (this is provided as an extra argument, as it is typically known a priori.

Assumes `h <= 0`, `t >= 0` , which implies `K >= 1`.
"""
_black(h, t, K = exp(-2*h*t)) = _black_normcdf(h, t, K) # generic fallback

function _black(h::Float64, t::Float64, K::Float64=exp(-2*h*t))
    if h > -10 && t < 0.21022410381342865 # 2*eps()^(1/16)
        _black_small_t(h, t)
    elseif t > 0.85 - h
        _black_normcdf(h, t, K)
    else
        _black_erfcx(h, t)
    end
end

"""
    _black_normcdf(h, t, K=exp(-2*h*t))

Computes the black kernel function directly using the normal CDF.
"""
_black_normcdf(h, t, K=exp(-2*h*t)) = Φ(h+t) - K*Φ(h-t)

"""
    _black_erfcx(h, t)

Compute the black kernel function using the `erfcx` function:

    b  =  Φ(h+t) - exp(-2⋅h⋅t) ⋅ Φ(h-t)

          exp(-(h²+t²)/2)
       =  ---------------  ·  [ Φ(h+t)/ϕ(h+t) - Φ(h-t)/ϕ(h-t) ]
              √(2π)

       =  exp(-(h²+t²)/2) ⋅ [erfcx(-(h+t)/√2) - erfcx(-(h-t)/√2)]/2

"""
_black_erfcx(h,t) =
    exp(-(h+t)^2/2)*(erfcx(-invsqrt2*(h+t)) - erfcx(-invsqrt2*(h-t)))/2

"""
    _black_small_t(h, t)

Compute the black kernel function using Taylor series around `t=0`:

    b  =  Φ(h+t) - exp(-2⋅h⋅t) ⋅ Φ(h-t)

          exp(-(h²+t²)/2)
       =  ---------------  ·  [ Φ(h+t)/ϕ(h+t) - Φ(h-t)/ϕ(h-t) ]
              √(2π)

       =  exp(-(h²+t²)/2) ⋅ [erfcx(-(h+t)/√2) - erfcx(-(h-t)/√2)]/2

Using

    Y(h) = Φ(h)/φ(h) = √(π/2)·erfcx(-h/√2)
    a    = 1+h·Y(h)

Note that due to `h < 0`, and `h·Y(h) -> -1` (from above) as `h -> -∞`, we also have that `a > 0` and `a -> 0` as `h -> -∞`.
"""
function _black_small_t(h,t)
    a = 1+h*sqrthalfπ*erfcx(-invsqrt2*h)
    t2 = t*t
    h2 = h*h
    expansion = 2*t*(@horner(t2,
        a,
        @horner(h2,-1+3*a,a)/6,
        @horner(h2,-7+15*a,-1+10*a,a)/120,
        @horner(h2,-57+105*a,-18+105*a,-1+21*a,a)/5040,
        @horner(h2,-561+945*a,-285+1260*a,-33+378*a,-1+36*a,a)/362880,
        @horner(h2,-6555+10395*a,-4680+17325*a,-840+6930*a,-52+990*a,-1+55*a,a)/39916800,
        @horner(h2,-89055+135135*a,-82845+270270*a,-20370+135135*a,-1926+25740*a,-75+2145*a,-1+78*a,a)/6227020800))

    return invsqrt2π*exp(-(h+t)^2/2)*expansion
end

"""
    _cblack(h, t, K=exp(-2*h*t))

The complement of the black kernel function

    1 - _black(h, t, K)

"""
_cblack(h, t, K=exp(-2*h*t)) = Φ(-(h+t)) + K*Φ(h-t)

"""
    _black_vega(h, t)

The vega corresponding to the black kernel function:

    Φ(h + t) - K*Φ(h - t)

where `h = -log(K)/s` and `t = s/2`.
"""
_black_vega(h,t) = ϕ(h+t)


# draw histgram of sensor data

module HistogramMean
    using Plots, DataFrames, CSV, Statistics
    pyplot()

    function main(is_test=false)
        data_path = joinpath(split(@__FILE__, "src")[1], "data/sensor_data_200.txt")
        df_200_mm = CSV.read(data_path, DataFrame,
                             header=["date", "time", "ir", "lidar"],
                             delim=' ')

        a = sum(df_200_mm.lidar)
        println("Sum of sensor data = $a")

        b = length(df_200_mm.lidar)
        println("Length of sensor data = $b")

        mean_1 = a / b
        mean_2 = mean(df_200_mm.lidar)
        println("Mean(Sum / Length) = $mean_1")
        println("Mean(mean()) = $mean_2")

        bin_min_max = maximum(df_200_mm.lidar) - minimum(df_200_mm.lidar)
        histogram(df_200_mm.lidar, bins=bin_min_max, color=:orange,
                  label="histogram")
        plot!([mean_1], st=:vline, color=:red,
              ylim=(0, 5000), linewidth=5, label="mean")

        if is_test == false
            save_path = joinpath(split(@__FILE__, "src")[1], "src/prob_stats/freq_dist/histogram/histogram_mean.png")
            savefig(save_path)
        end
    end
end


"""
Regression Modelling Splines from Frank Harrell's RMS book

These methods use truncated power bases to implement regression splines.

Frank E. Harrell, Jr.. 2006. Regression Modeling Strategies. Springer-Verlag, Berlin, Heidelberg.
"""
module RMSSplines

__precompile__()

include("spline_interface.jl")

using StatsModels
using StatsBase

"""Compute Truncated power basis"""
@inline U₊(v) = v > 0 ? v : zero(v)
@inline P(v, k) = v > 0 ? v ^ k : zero(v)
@inline P₃(v) = P(v, 3)

function rms_linear_spline(x::AbstractVector, knots::AbstractVector)
    k = length(knots)

    M = Array{Float64}(undef, k, length(x))

    for i = 1:k
        M[i, :] = U₊.(x .- knots[i])
    end

    M
end

@inline _rcs(x::Missing, j::Int, t::AbstractVector) = missing
@inline _rcs(x::Int, j::Int, t::AbstractVector) = _rcs(float(x), j, t)

"""Compute Truncated Power Basis for Natural/Restricted Cubic Spline """
@inline function _rcs(x, j::Int, t::AbstractVector)
    k = length(t)

    d = (t[k] - t[k-1])
    τ = (t[k] - t[1]) ^ 2

    (
     P₃(x - t[j])
     - P₃(x - t[k - 1]) * (t[k] - t[j]) / d
     + P₃(x - t[k]) * (t[k-1] - t[j]) / d
    ) / τ
end

const RealOrMissing = Union{<:Real, Missing}

@inline function rms_restricted_cubic_spline!(out::AbstractVector, x::Missing, knots::AbstractVector)
    out .= missing
    out
end

@inline function rms_restricted_cubic_spline!(out::AbstractVector, x::Real, knots::AbstractVector)
    k = length(knots)

    out[1] = x
    @simd for j in 2:k-1
        out[j] = _rcs(x, j - 1, knots)
    end

    out
end

function rms_restricted_cubic_spline(x::AbstractVector{<:Real}, knots::AbstractVector{<:Real})
    k = length(knots)
    M = Array{Float64}(undef, length(x), k-1)

    for i in eachindex(x)
        row = view(M, i, :)
        rms_restricted_cubic_spline!(row, x[i], knots)
    end

    M
end

function rms_restricted_cubic_spline(x::AbstractVector{<:RealOrMissing}, knots::AbstractVector{<:Real})
    k = length(knots)

    R = eltype(x) <: Real ? Float64 : Union{Float64, Missing}
    M = Array{R}(undef, length(x), k-1)

    for i in eachindex(x)
        row = view(M, i, :)
        rms_restricted_cubic_spline!(row, x[i], knots)
    end

    M
end

"""RMS quantiles for regression spline knots"""
function rms_quantiles_for_knots(k::Int)
    if k <= 3
        outer = 0.1
    elseif k <= 6
        outer = 0.05
    else
        outer = 0.025
    end

    collect(range(outer, 1-outer, length=k))
end

function rms_generate_knots(x::AbstractVector, knots=3; unique_xs=false)
    if length(x) < knots
        throw("Not enough data to generate knots")
    end

    q = rms_quantiles_for_knots(knots)

    if unique_xs
        quantile(x |> skipmissing |> collect, q)
    else
        quantile(x |> skipmissing |> unique |> collect, q)
    end
end

rcs(x::AbstractVector, knots::Int) = rcs(x, rms_generate_knots(x, knots))
rcs(x::AbstractVector, knots::AbstractVector) = rms_restricted_cubic_spline(x, knots)

################################################################################
"""Define Formula interface for Restricted Cubic Splines"""

struct RCSTerm{T, K} <: AbstractSplineTerm
    term::T
    knots::K
end

shorthand(::Type{<:RCSTerm}) = "rcs"
spline_term(v::RCSTerm) = v.term
spline_knots(v::RCSTerm) = v.knots
spline_deg(v::RCSTerm) = 3

rcs(x::Symbol, knots::Int=3) = RCSTerm(term(x), term(knots))
rcs(t::AbstractTerm, knots::Int=3) = RCSTerm(t, term(knots))

function StatsModels.apply_schema(t::FunctionTerm{typeof(rcs)}, sch::StatsModels.Schema, Mod::Type{<:SplineTermContext})
    apply_schema(RCSTerm(t.args_parsed...), sch, Mod)
end

function StatsModels.apply_schema(t::RCSTerm, sch::StatsModels.Schema, Mod::Type{<:SplineTermContext})
    term = apply_schema(t.term, sch, Mod)

    isa(t.knots, ConstantTerm) || throw("Knots must be a constant term")
    t.knots.n > 2 || throw("Restricted cubic spline requires at least 3 knots")

    RCSTerm(term, t.knots)
end

function StatsModels.modelcols(t::RCSTerm, d::NamedTuple)
    col = modelcols(t.term, d)
    rcs(col, t.knots.n)
end

StatsModels.width(t::RCSTerm) = t.knots.n - 1

export rcs, rms_generate_knots, rms_quantiles_for_knots

end


using  Test
import RecursiveFactorization
import LinearAlgebra
using  LinearAlgebra: norm
using  Random

Random.seed!(12)

const baselu = LinearAlgebra.lu
const mylu = RecursiveFactorization.lu

function testlu(A, MF, BF)
    @test MF.info == BF.info
    @test norm(MF.L*MF.U - A[MF.p, :], Inf) < sqrt(eps(real(first(A))))
    nothing
end

@testset "Test LU factorization" begin
    for p in (Val(true), Val(false)), T in (Float64, Float32, ComplexF64, ComplexF32, Real)
        siz = (50, 100)
        if isconcretetype(T)
            A = rand(T, siz...)
        else
            _A = rand(50, 100)
            A  = Matrix{T}(undef, siz...)
            copyto!(A, _A)
        end
        MF = mylu(A, p)
        BF = baselu(A, p)
        testlu(A, MF, BF)
        for i in 50:7:100 # test `MF.info`
            A[:, i] .= 0
            MF = mylu(A, p, check=false)
            BF = baselu(A, p, check=false)
            testlu(A, MF, BF)
        end
    end
end


using PolyChaosODE, BenchmarkTools, DifferentialEquations, Distributions, PolyChaos, Plots

### This package ###

x0 = 2.0
μ, σ = -0.5, 0.05
tend = 10.0
L = 6

function exp_ode!(du, u, p, x, t)
    du[1] = (p[1] + p[2] * x[1]) * u[1]
end

prob_dim = 1
vars = (Normal(0.0, sqrt(2) / 2),)
num_polys = (L + 1,)
stoch_galerkin_ode = StochGalerkinODE(exp_ode!, prob_dim, vars; num_polys = num_polys)

x0_ = [x0]
tspan = (0, tend)
# sqrt(2) factor in the second term because we are using the Physicist's Hermite polynomials instead of the Probabilist's
p = [-0.5, 0.05 * sqrt(2)]
sol = stoch_galerkin_ode(x0_, tspan, p; alg = VCABM())

t_step = 0.01
interval_t = 0.0:t_step:tend
std_this =
    (
        x -> sqrt.(x)
    ).(
        compute_expectation_and_diag_variance(
            stoch_galerkin_ode,
            interval_t,
            [sol(t) for t in interval_t],
        )[2],
    )

@btime stoch_galerkin_ode($x0_, $tspan, $p; alg = VCABM())
@btime (
    x -> sqrt.(x)
).(
    compute_expectation_and_diag_variance(
        stoch_galerkin_ode,
        interval_t,
        [sol(t) for t in interval_t],
    )[2],
)

### PolyChaos (based on https://timueh.github.io/PolyChaos.jl/stable/random_ode/)###

Nrec = 40
opq = GaussOrthoPoly(L; Nrec = Nrec, addQuadrature = true)

a = [convert2affinePCE(μ, σ, opq); zeros(Float64, L - 1)]
xinit = [x0; zeros(Float64, L)]

const t2_ = Tensor(2, opq);
const t3_ = Tensor(3, opq);

function ODEgalerkin(du, u, p, t)
    du[:] = [
        sum(p[j+1] * u[k+1] * t3_.get([j, k, m]) / t2_.get([m, m]) for j = 0:L for k = 0:L) for m = 0:L
    ]
end

probgalerkin = ODEProblem(ODEgalerkin, xinit, tspan, a)
solgalerkin = solve(probgalerkin; alg = VCABM())

t, x = solgalerkin.t, solgalerkin.u;
mean_poly_chaos = [mean(x_, opq) for x_ in x]
std_poly_chaos = [std(x_, opq) for x_ in x]

@btime solve(probgalerkin; alg = VCABM())
@btime [mean(x_, opq) for x_ in x]
@btime [std(x_, opq) for x_ in x]

plot(t, mean_pce .- std_pce)
plot!(t, mean_pce)
plot!(t, mean_pce .+ std_pce)


@add_problem socp function socp_dual_norm_2_atom(
    handle_problem!,
    ::Val{test},
    atol,
    rtol,
    ::Type{T},
) where {T,test}
    x = Variable(2, 1)
    A = [1 2; 2 1; 3 4]
    b = [2; 3; 4]
    p = minimize(norm2(A * x + b); numeric_type = T)

    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 0.64888 atol = atol rtol = rtol
        @test evaluate(norm2(A * x + b)) ≈ 0.64888 atol = atol rtol = rtol
    end

    x = Variable(2, 1)
    A = [1 2; 2 1; 3 4]
    b = [2; 3; 4]
    lambda = 1
    p = minimize(norm2(A * x + b) + lambda * norm2(x), x >= 1; numeric_type = T)

    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 14.9049 atol = atol rtol = rtol
        @test evaluate(norm2(A * x + b) + lambda * norm2(x)) ≈ 14.9049 atol =
            atol rtol = rtol
        @test p.constraints[1].dual ≈ [4.4134, 5.1546] atol = atol rtol = rtol
    end

    x = Variable(2)

    p = minimize(
        norm2([x[1] + 2x[2] + 2; 2x[1] + x[2] + 3; 3x[1] + 4x[2] + 4]) +
        lambda * norm2(x),
        x >= 1;
        numeric_type = T,
    )

    if test
        @test vexity(p) == ConvexVexity()
    end

    handle_problem!(p)
    if test
        @test p.optval ≈ 14.9049 atol = atol rtol = rtol
        @test evaluate(norm2(A * x + b) + lambda * norm2(x)) ≈ 14.9049 atol =
            atol rtol = rtol
        @test p.constraints[1].dual ≈ [4.4134, 5.1546] atol = atol rtol = rtol
    end

    x = Variable(2, 1)
    A = [1 2; 2 1; 3 4]
    b = [2; 3; 4]
    lambda = 1
    p = minimize(
        norm2(A * x + b) + lambda * norm_1(x),
        x >= 1;
        numeric_type = T,
    )

    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 15.4907 atol = atol rtol = rtol
        @test evaluate(norm2(A * x + b) + lambda * norm_1(x)) ≈ 15.4907 atol =
            atol rtol = rtol
        @test p.constraints[1].dual ≈ [4.7062, 5.4475] atol = atol rtol = rtol
    end
end

@add_problem socp function socp_dual_frobenius_norm_atom(
    handle_problem!,
    ::Val{test},
    atol,
    rtol,
    ::Type{T},
) where {T,test}
    m = Variable(4, 5)
    c = [m[3, 3] == 4, m >= 1]
    p = minimize(norm(vec(m), 2), c; numeric_type = T)

    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ sqrt(35) atol = atol rtol = rtol
        @test evaluate(norm(vec(m), 2)) ≈ sqrt(35) atol = atol rtol = rtol
        @test p.constraints[1].dual ≈ 0.6761 atol = atol rtol = rtol
        dual = 0.1690 .* ones(4, 5)
        dual[3, 3] = 0
        @test p.constraints[2].dual ≈ dual atol = atol rtol = rtol
    end
end

@add_problem socp function socp_quad_over_lin_atom(
    handle_problem!,
    ::Val{test},
    atol,
    rtol,
    ::Type{T},
) where {T,test}
    x = Variable(3)
    A = [2 -3 5; -2 9 -3; 5 -8 3]
    b = [-3; 9; 5]
    c = [3, 2, 4]
    d = -3
    p = minimize(quadoverlin(A * x + b, dot(c, x) + d); numeric_type = T)

    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 17.7831 atol = atol rtol = rtol
        @test evaluate(quadoverlin(A * x + b, dot(c, x) + d)) ≈ 17.7831 atol =
            atol rtol = rtol
    end
end

@add_problem socp function socp_sum_squares_atom(
    handle_problem!,
    ::Val{test},
    atol,
    rtol,
    ::Type{T},
) where {T,test}
    x = Variable(2, 1)
    A = [1 2; 2 1; 3 4]
    b = [2; 3; 4]
    p = minimize(sumsquares(A * x + b); numeric_type = T)

    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 0.42105 atol = atol rtol = rtol
        @test (evaluate(sumsquares(A * x + b)))[1] ≈ 0.42105 atol = atol rtol =
            rtol
    end
end

@add_problem socp function socp_square_atom(
    handle_problem!,
    ::Val{test},
    atol,
    rtol,
    ::Type{T},
) where {T,test}
    x = Variable(2, 1)
    A = [1 2; 2 1; 3 4]
    b = [2; 3; 4]
    p = minimize(sum(square(A * x + b)); numeric_type = T)

    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 0.42105 atol = atol rtol = rtol
        @test evaluate(sum(square(A * x + b))) ≈ 0.42105 atol = atol rtol = rtol
    end

    x = Variable(2, 1)
    A = [1 2; 2 1; 3 4]
    b = [2; 3; 4]
    expr = A * x + b
    p = minimize(sum(dot(^)(expr, 2)); numeric_type = T) # elementwise ^
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 0.42105 atol = atol rtol = rtol
        @test evaluate(sum(broadcast(^, expr, 2))) ≈ 0.42105 atol = atol rtol =
            rtol
    end

    p = minimize(sum(dot(*)(expr, expr)); numeric_type = T) # elementwise *
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 0.42105 atol = atol rtol = rtol
        @test evaluate(sum((dot(*))(expr, expr))) ≈ 0.42105 atol = atol rtol =
            rtol
    end
end

@add_problem socp function socp_inv_pos_atom(
    handle_problem!,
    ::Val{test},
    atol,
    rtol,
    ::Type{T},
) where {T,test}
    x = Variable(4)
    p = minimize(
        sum(invpos(x)),
        invpos(x) < 2,
        x > 1,
        x == 2,
        2 == x;
        numeric_type = T,
    )

    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 2 atol = atol rtol = rtol
        @test evaluate(sum(invpos(x))) ≈ 2 atol = atol rtol = rtol
    end

    x = Variable(3)
    p = minimize(sum(dot(/)([3, 6, 9], x)), x <= 3; numeric_type = T)

    handle_problem!(p)
    if test
        @test evaluate(x) ≈ fill(3.0, (3, 1)) atol = atol rtol = rtol
        @test p.optval ≈ 6 atol = atol rtol = rtol
        @test evaluate(sum((dot(/))([3, 6, 9], x))) ≈ 6 atol = atol rtol = rtol
    end

    x = Variable()
    p = minimize(sum([3, 6, 9] / x), x <= 3; numeric_type = T)

    handle_problem!(p)
    if test
        @test evaluate(x) ≈ 3 atol = atol rtol = rtol
        @test p.optval ≈ 6 atol = atol rtol = rtol
        @test evaluate(sum([3, 6, 9] / x)) ≈ 6 atol = atol rtol = rtol
    end
end

@add_problem socp function socp_geo_mean_atom(
    handle_problem!,
    ::Val{test},
    atol,
    rtol,
    ::Type{T},
) where {T,test}
    x = Variable(2)
    y = Variable(2)
    p = minimize(geomean(x, y), x >= 1, y >= 2; numeric_type = T)

    # not DCP compliant
    if test
        @test vexity(p) == ConcaveVexity()
    end
    p = maximize(geomean(x, y), 1 < x, x < 2, y < 2; numeric_type = T)

    # Just gave it a vector as an objective, not okay
    if test
        @test_throws Exception handle_problem!(p)
    end

    p = maximize(sum(geomean(x, y)), 1 < x, x < 2, y < 2; numeric_type = T)

    handle_problem!(p)
    if test
        @test p.optval ≈ 4 atol = atol rtol = rtol
        @test evaluate(sum(geomean(x, y))) ≈ 4 atol = atol rtol = rtol
    end
end

@add_problem socp function socp_sqrt_atom(
    handle_problem!,
    ::Val{test},
    atol,
    rtol,
    ::Type{T},
) where {T,test}
    x = Variable()
    return p = maximize(sqrt(x), 1 >= x; numeric_type = T)
end

@add_problem socp function socp_quad_form_atom(
    handle_problem!,
    ::Val{test},
    atol,
    rtol,
    ::Type{T},
) where {T,test}
    x = Variable(3, 1)
    A = [0.8608 0.3131 0.5458; 0.3131 0.8584 0.5836; 0.5458 0.5836 1.5422]
    p = minimize(quadform(x, A), [x >= 1]; numeric_type = T)

    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 6.1464 atol = atol rtol = rtol
        @test evaluate(quadform(x, A)) ≈ 6.1464 atol = atol rtol = rtol
    end

    x = Variable(3, 1)
    A =
        -1.0 *
        [0.8608 0.3131 0.5458; 0.3131 0.8584 0.5836; 0.5458 0.5836 1.5422]
    c = [3 2 4]
    p = maximize(c * x, [quadform(x, A) >= -1]; numeric_type = T)

    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 3.7713 atol = atol rtol = rtol
        @test evaluate(quadform(x, A)) ≈ -1 atol = atol rtol = rtol
    end

    # https://github.com/jump-dev/Convex.jl/issues/398
    n = 3
    b = randn(n)
    x = Variable(n)
    H = Semidefinite(n)
    Hval = randn(n, n)
    Hval .= Hval' * Hval + 10 * diagm(0 => ones(n)) # symmetric positive definite
    fix!(H, Hval)
    p = minimize(x'b + quadform(x, H), [x >= 0]; numeric_type = T)
    handle_problem!(p)

    p2 = minimize(x'b + quadform(x, Hval), [x >= 0]; numeric_type = T)
    handle_problem!(p2)

    if test
        @test p.optval ≈ p2.optval atol = atol rtol = rtol
        @test evaluate(H) ≈ Hval atol = atol rtol = rtol
    end

    # https://github.com/jump-dev/Convex.jl/pull/444
    x = Variable(3)
    a, b, c, d, e, f = rand(6)
    M = [
        2 a-b*im c-d*im
        a+b*im 2 e-f*im
        c+d*im e+f*im 2
    ]
    y = rand(3)
    p = minimize(quadform(x - y, M); numeric_type = T)
    handle_problem!(p)
    if test
        @test evaluate(x) ≈ y atol = atol rtol = rtol
    end

    # https://github.com/jump-dev/Convex.jl/issues/446
    x = Variable(3)
    A = zeros(3, 3)
    A[1, 1] = 1.0
    p = minimize(
        quadform(x, A) - 2 * x[2] - x[3],
        norm(x, 1) <= 1;
        numeric_type = T,
    )
    handle_problem!(p)
    if test
        @test evaluate(x) ≈ [0.0, 1.0, 0.0] atol = atol rtol = rtol
    end
end

@add_problem socp function socp_huber_atom(
    handle_problem!,
    ::Val{test},
    atol,
    rtol,
    ::Type{T},
) where {T,test}
    x = Variable(3)
    p = minimize(sum(huber(x, 1)), x >= 2; numeric_type = T)

    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)
    if test
        @test p.optval ≈ 9 atol = atol rtol = rtol
        @test evaluate(sum(huber(x, 1))) ≈ 9 atol = atol rtol = rtol
    end
end

@add_problem socp function socp_rational_norm_atom(
    handle_problem!,
    ::Val{test},
    atol,
    rtol,
    ::Type{T},
) where {T,test}
    A = [1 2 3; -1 2 3]
    b = A * ones(3)
    x = Variable(3)
    p = minimize(norm(x, 4.5), [A * x == b]; numeric_type = T)

    if test
        @test vexity(p) == ConvexVexity()
    end
    # Solution is approximately x = [1, .93138, 1.04575]
    handle_problem!(p)
    if test
        @test p.optval ≈ 1.2717 atol = atol rtol = rtol
        @test evaluate(norm(x, 4.5)) ≈ 1.2717 atol = atol rtol = rtol
    end
end

@add_problem socp function socp_rational_norm_dual_norm(
    handle_problem!,
    ::Val{test},
    atol,
    rtol,
    ::Type{T},
) where {T,test}
    v = [0.463339, 0.0216084, -2.07914, 0.99581, 0.889391]
    x = Variable(5)
    q = 1.379  # q norm constraint that generates many inequalities
    qs = q / (q - 1)  # Conjugate to q
    p = minimize(x' * v; numeric_type = T)

    p.constraints += (norm(x, q) <= 1)
    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p) # Solution is -norm(v, q / (q - 1))
    if test
        @test p.optval ≈ -2.144087 atol = atol rtol = rtol
        @test sum(evaluate(x' * v)) ≈ -2.144087 atol = atol rtol = rtol
        @test evaluate(norm(x, q)) ≈ 1 atol = atol rtol = rtol
        @test sum(evaluate(x' * v)) ≈ -(sum(abs.(v) .^ qs)^(1 / qs)) atol = atol rtol =
            rtol
    end
end

@add_problem socp function socp_rational_norm_atom_sum(
    handle_problem!,
    ::Val{test},
    atol,
    rtol,
    ::Type{T},
) where {T,test}
    A = [
        -0.719255 -0.229089
        -1.33632 -1.37121
        0.703447 -1.4482
    ]
    b = [-1.82041, -1.67516, -0.866884]
    q = 1.5
    xvar = Variable(2)
    p = minimize(
        0.5 * sumsquares(xvar) + norm(A * xvar - b, q);
        numeric_type = T,
    )

    if test
        @test vexity(p) == ConvexVexity()
    end
    handle_problem!(p)

    if test
        # Compute gradient, check it is zero(ish)
        x_opt = evaluate(xvar)
        margins = A * x_opt - b
        qs = q / (q - 1)  # Conjugate
        denom = sum(abs.(margins) .^ q)^(1 / qs)
        g = x_opt + A' * (abs.(margins) .^ (q - 1) .* sign.(margins)) / denom
        @test p.optval ≈ 1.7227 atol = atol rtol = rtol
        @test norm(g, 2)^2 ≈ 0 atol = atol rtol = rtol
    end
end

@add_problem socp function socp_norm_consistent_with_Base_for_matrix_variables(
    handle_problem!,
    ::Val{test},
    atol,
    rtol,
    ::Type{T},
) where {T,test}
    A = randn(4, 4)
    x = Variable(4, 4)
    set_value!(x, A)
    # Matrix norm
    if test
        @test evaluate(opnorm(x)) ≈ opnorm(A) atol = atol rtol = rtol
        @test evaluate(opnorm(x, 1)) ≈ opnorm(A, 1) atol = atol rtol = rtol
        @test evaluate(opnorm(x, 2)) ≈ opnorm(A, 2) atol = atol rtol = rtol
        @test evaluate(opnorm(x, Inf)) ≈ opnorm(A, Inf) atol = atol rtol = rtol
    end
    # Vector norm
    # TODO: Once the deprecation for norm on matrices is removed, remove the `vec` calls
    if test
        @test evaluate(norm(vec(x), 1)) ≈ norm(vec(A), 1) atol = atol rtol =
            rtol
        @test evaluate(norm(vec(x), 2)) ≈ norm(vec(A), 2) atol = atol rtol =
            rtol
        @test evaluate(norm(vec(x), 7)) ≈ norm(vec(A), 7) atol = atol rtol =
            rtol
        @test evaluate(norm(vec(x), Inf)) ≈ norm(vec(A), Inf) atol = atol rtol =
            rtol
    end
end

@add_problem socp function socp_fix_and_free_addition(
    handle_problem!,
    ::Val{test},
    atol,
    rtol,
    ::Type{T},
) where {T,test}
    x = Variable()
    y = Variable()

    p = minimize(x + y, x >= 0, y >= 0; numeric_type = T)

    handle_problem!(p)
    if test
        @test p.optval ≈ 0 atol = atol rtol = rtol
    end

    set_value!(y, 4)
    fix!(y)
    handle_problem!(p)
    if test
        @test p.optval ≈ 4 atol = atol rtol = rtol
    end

    free!(y)
    handle_problem!(p)
    if test
        @test p.optval ≈ 0 atol = atol rtol = rtol
    end
end

@add_problem socp function socp_fix_multiplication(
    handle_problem!,
    ::Val{test},
    atol,
    rtol,
    ::Type{T},
) where {T,test}
    a = [1, 2, 3, 2, 1]
    x = Variable(length(a))
    gamma = Variable(Positive())
    fix!(gamma, 0.7)

    p = minimize(
        norm(x - a) + gamma * norm(x[1:end-1] - x[2:end]);
        numeric_type = T,
    )

    handle_problem!(p)
    if test
        o1 = p.optval
        # x should be very close to a
        @test o1 ≈ 0.7 * norm(a[1:end-1] - a[2:end]) atol = atol rtol = rtol
    end
    # increase regularization
    fix!(gamma, 1.0)
    handle_problem!(p)

    if test
        o2 = p.optval
        # x should be very close to mean(a)
        @test o2 ≈ norm(a .- mean(a)) atol = atol rtol = rtol
    end

    if test
        @test o1 <= o2
    end
end

@add_problem socp function socp_dual_minimal_norm_solutions(
    handle_problem!,
    ::Val{test},
    atol,
    rtol,
    ::Type{T},
) where {T,test}
    x = Variable(2)
    A = [1 2; 2 4]
    b = [3, 6]
    p = minimize(norm(x, 1), A * x == b; numeric_type = T)

    if test
        @test vexity(p) == ConvexVexity()
    end

    handle_problem!(p)
    if test
        @test p.optval ≈ 1.5 atol = atol rtol = rtol
        @test evaluate(x) ≈ [0, 1.5] atol = atol rtol = rtol
        @test evaluate(norm(x, 1)) ≈ p.optval atol = atol rtol = rtol
        @test dot(b, p.constraints[1].dual) ≈ p.optval atol = atol rtol = rtol
    end

    x = Variable(2)
    A = [1 2; 2 4]
    b = [3, 6]
    p = minimize(norm(x, 2), A * x == b; numeric_type = T)

    test && @test vexity(p) == ConvexVexity()

    handle_problem!(p)

    if test
        @test p.optval ≈ 3 / sqrt(5) atol = atol rtol = rtol
        @test evaluate(x) ≈ [3 / 5, 6 / 5] atol = atol rtol = rtol
        @test evaluate(norm(x, 2)) ≈ p.optval atol = atol rtol = rtol
        @test dot(b, p.constraints[1].dual) ≈ p.optval atol = atol rtol = rtol
    end

    x = Variable(2)
    A = [1 2; 2 4]
    b = [3, 6]
    p = minimize(norm(x, Inf), A * x == b; numeric_type = T)

    test && @test vexity(p) == ConvexVexity()

    handle_problem!(p)

    if test
        @test p.optval ≈ 1.0 atol = atol rtol = rtol
        @test evaluate(x) ≈ [1, 1] atol = atol rtol = rtol
        @test evaluate(norm(x, Inf)) ≈ p.optval atol = atol rtol = rtol
        @test dot(b, p.constraints[1].dual) ≈ p.optval atol = atol rtol = rtol
    end
end


using Test
using HyperbolicPlane

p = HPoint(0)
q = HPoint(.2)
r = HPoint(0-0.2im)
@test dist(p,q) == dist(r)

p[:color] = "red"
@test p[:color] == "red"

P = RandomHPoint()
Q = RandomHPoint()
S = P+Q
L = HLine(P,Q)
@test issubset(S,L)
M = midpoint(S)
@test in(P,L)
@test in(P,S)
@test issubset(P+M,S)

@test length(S) == dist(P,Q)

x1 = length(S)
x2 = length(P+M) + length(M+Q)
@test abs(x1-x2) < 100*eps(1.)

T = P+Q+M
@test area(T) < 100*eps(1.0)

L = HLine(1,3)
LL = HLine(2,pi)
P = meet(L,LL)
@test in(P,L)
@test in(P,LL)

TT = reflect_across(T,L)
TT = reflect_across(TT,L)
@test T == TT

S = RandomHSegment()
L = HLine(S)
SS = S'
LL = HLine(SS)
@test L == LL'

SS = -S
LL = HLine(SS)
@test L == -LL

a,b,c,d = [RandomHPoint() for t=1:4]
X = HPolygon(a,b,c,d,HPoint(im*eps(1.0)))
Y = HPolygon(d,c,b,a,HPoint(0))
@test X==Y

C = HContainer()
P = RandomHPolygon(5)
for k=1:10
    add_object!(C,P)
end

T = RandomHTriangle()
add_object!(C,T)
TT = T'
add_object!(C,TT')
@test length(C) == 2

a,b,c = endpoints(T)
@test a+c+b == T

P = RandomHPolygon(3)
T = HTriangle(P)
aP = sort(angles(P))
aT = angles(T)
@test sum(abs.(aP-aT)) < 100*eps(1.0)

T = HPoint(1.0,2.0) + HPoint(3.0,-1) + HPoint(2,0)
P = HPolygon(T)
@test perimeter(T) == perimeter(P)

# point on line stuff
L = RandomHLine()
P = point_on_line(L)
@test in(P,L)

PP = points_on_line(L,2)
P,Q = PP
@test in(P,L) && in(Q,L)

R = RandomHRay()
V = get_vertex(R)
P = point_on_ray(R)
@test in(P,R)
A = midpoint(V,P)
@test in(A,R)

# Polygon area
X = equilateral(6,1)
T = X.plist[1] + X.plist[2] + HPoint()
AX = area(X)
AT = area(T)
@test abs(AX-6AT) <= equality_threshold() * eps(1.0)


# from FastTransforms
function backend_fft_pow2!{T}(x::Vector{T})
    n,big2,bigpi=length(x),2one(T),T(pi)
    nn,j=n÷2,1
    @inbounds for i=1:2:n-1
        if j>i
            x[j], x[i] = x[i], x[j]
            x[j+1], x[i+1] = x[i+1], x[j+1]
        end
        m = nn
        while m ≥ 2 && j > m
            j -= m
            m = m÷2
        end
        j += m
    end
    logn = 2
    while logn < n
        piθ=-bigpi*big2/logn
        wtemp = sin(piθ/2)
        wpr, wpi = -2wtemp^2, sin(piθ)
        wr, wi = one(T), zero(T)
        for m=1:2:logn-1
            @inbounds for i=m:2logn:n
                j=i+logn
                mixr, mixi = wr*x[j]-wi*x[j+1], wr*x[j+1]+wi*x[j]
                x[j], x[j+1] = x[i]-mixr, x[i+1]-mixi
                x[i], x[i+1] = x[i]+mixr, x[i+1]+mixi
            end
            wr = (wtemp=wr)*wpr-wi*wpi+wr
            wi = wi*wpr+wtemp*wpi+wi
        end
        logn = logn << 1
    end
    return x
end

function fft_pow2{T}(x::Vector{T})
    y = FastTransforms.interlace(real(x),imag(x))
    backend_fft_pow2!(y)
@compat    return complex.(y[1:2:end],y[2:2:end])
end

## TAYLOR

taylorinterpolationerror(r,Ninterp::Int,M,R) = M*(r/R)^Ninterp/(1-r/R)

function taylortransform{T}(vals::Vector{Complex{T}},r::T,M::T,R::T)
  Ninterp = length(vals)
  cfs = fft_pow2(vals)/Ninterp
  err= erroradjustment(Complex{T},taylorinterpolationerror(r,Ninterp,M,R))
  powr = one(T)
  @inbounds for i = 1:Ninterp
    cfs[i] *= powr
    cfs[i] += err
    err /= R
    powr /= r
  end
  cfs
end

## CHEBYSHEV

function chebyshevtransform{T}(a::AbstractArray{Complex{T}})
  N = length(a) #big(length(a))
    ispow2(N) || error("N not a power of 2")
    c = fft_pow2([a; flipdim(a,1)])
    d = c[1:N]
 @compat    d .*= exp.((-im*convert(T,pi)).*(0:N-1)./(2*N))
     d[1] = d[1] / 2 #sqrt(big(2))
     scale!(inv(convert(T,N)), d)
end
@compat chebyshevtransform{T<:Real}(a::AbstractArray{T}) = real(chebyshevtransform(complex(a)))

function cos_vec!(A::Vector)
  for i in eachindex(A)
    A[i] = cos(A[i])
  end
  A
end
chebyshevpoints(N_interp,T) = cos_vec!(collect(1:2:2N_interp-1)*(T(pi)/2N_interp));

chebyshev_shifted_sq(x,k) = (-1)^k * cos(2k*asin(x))
chebyshev_shifted(x,k) = chebyshev_shifted_sq(sqrt(x),k)

function chebyshev_shifted_by_squaring{T}(x::T,k)
  tot=k;
  s0sq = 4x*(1-x)
  cn = 2x-1; sn_div = one(T)
  ctot = one(T); stot_div = zero(T)
  while tot>0
    if isodd(tot)
      ctot,stot_div = ctot*cn - stot_div*sn_div*s0sq, stot_div*cn - ctot*sn_div
    end
    sn_div *= 2cn; cn *= 2cn; cn-=1;
    tot ÷=2
  end
  ctot
end

function chebyshev_shifted_multiel{T}(x::T,K)
  Tk = Array{T}(K)
  mul = 4x-2
  K > 0 && (Tk[1] = one(typeof(x)))
  K > 1 && (Tk[2] = 2x-1)
  @inbounds for kk = 1:div(K-1,2)
    Tk[2kk+1] = 2Tk[kk+1]^2-1
    Tk[2kk+2] = mul*Tk[2kk+1]-Tk[2kk]
  end
  isodd(K) && (Tk[K] = 2Tk[div(K,2)+1]^2-1)
  Tk
end

function fejerweights(n,T::DataType=Float64)
  v = Array{Complex{T}}(n)
  v[1] = 2
  # twid = exp(-im*T(pi)/n) # conjugating here to get ifft
  # twidp = copy(twid)
  @inbounds for k = 1:div(n-1,2)
    twidp = exp(-im*k*T(pi)/n)
    v[k+1] = 2twidp/(1-4k^2)
    # twidp *= twid
  end
  mod(n,2)==0 && (v[div(n,2)+1]=0)
  @inbounds for k = 1:div(n-1,2)
    v[n-k+1] = conj(v[k+1])
  end
   real(fft_pow2(v))/n
end
 # fejerweights{T<:Real}(n,::Type{Complex{T}}) = fejerweights(n,T)

fejererror(N,M,ρ) = 2M/(ρ-1)*ρ^(-N)


test = """
0,9 -> 5,9
8,0 -> 0,8
9,4 -> 3,4
2,2 -> 2,1
7,0 -> 7,4
6,4 -> 2,0
0,9 -> 2,9
3,4 -> 1,4
0,0 -> 8,8
5,5 -> 8,2
"""

function read_input(io)
    lines = readlines(io)
    numbers = hcat(split.(lines, r" -> |,")...) |> permutedims
    return parse.(Int, numbers) .+ 1
end

function board(values::Matrix; diagonal::Bool = false)
    m = maximum(values)
    M = fill(0, m, m)
    for (x1, y1, x2, y2) in eachrow(values)
        if x1 == x2 || y1 == y2
            (x1, x2) = extrema((x1, x2))
            (y1, y2) = extrema((y1, y2))
            M[y1:y2, x1:x2] .+= 1
        elseif diagonal
            for (x, y) in zip(_range(x1, x2), _range(y1, y2))
                M[y, x] += 1
            end
        end
    end
    return M
end

_range(a, b) = range(a, b; step = sign(b - a))

function count_overlaps(values::Matrix; diagonal::Bool = false)
    return count(>(1), board(values; diagonal))
end

let T = read_input(IOBuffer(test))
    @assert count_overlaps(T) == 5
    @assert count_overlaps(T; diagonal = true) == 12
end

M = read_input("day05.in")
println("Q1: ", count_overlaps(M))
println("Q2: ", count_overlaps(M; diagonal = true))


module FastInverseSqrt

export rsqrt

# See https://en.wikipedia.org/wiki/Fast_inverse_square_root#Overview_of_the_code & https://www.youtube.com/watch?v=p8u_k2LIZyo
function rsqrt(x::Float32)
    xₛ = x
    int32 = reinterpret(UInt32, xₛ)
    int32 = 0x5f3759df - int32 >> 1
    xₛ = reinterpret(Float32, int32)
    xₛ *= 1.5f0 - x * 0.5f0 * xₛ^2
    return xₛ
end
function rsqrt(x::Float64)
    xₛ = x
    int64 = reinterpret(UInt64, xₛ)
    int64 = 0x5fe6eb50c7b537a9 - int64 >> 1  # See https://stackoverflow.com/a/11644533/3260253
    xₛ = reinterpret(Float64, int64)
    xₛ *= 1.5 - x * 0.5 * xₛ^2
    return xₛ
end
rsqrt(x::Real) = rsqrt(float(x))

end


 #mtdna_model = @reaction_network mtdna begin
 # b, wildtype --> 2wildtype
 # d, wildtype  --> null
 # bmut, mut --> 2mut
 # dmut, mut --> null
 # m, wildtype --> wildtype + mut
 # # can we have wildtype --> 2mut
#end b d bmut dmut m

using Random
using Colors
using Plots
theme(:ggplot2)

function update(t, vals)
  wt = vals[1];
  mut = vals[2];
  haz = [ifelse(sum(vals)<target,b*wt,0.0), d*wt, ifelse(sum(vals)<target,bmut*mut,0.0), dmut*mut, m*wt];
  # ^ specifying the reactions. 1st = wt replication,  b, wildtype --> 2wildtype
  # 2nd = wt degredation, d, wildtype  --> null
  # 3rd = mut replication, bmut, mut --> 2mut
  # 4th = mut degredation, dmut, mut --> null
  # 5th = mutaion, m, wildtype --> wildtype + mut
  hazfrac = cumsum(haz/sum(haz));
  index = minimum(findall(hazfrac .- Random.rand() .> 0));
  tnew = t + (1.0/sum(haz))*(Random.randexp());
  valsnew = [wt + delta_wt[index], mut + delta_mut[index]]
  return((tnew,valsnew));
end

b = 0.001875*365
d = 0.001875*365*0.86 # Note, slightly lower than b, so that when total is low, population grows
bmut = 0.001875*365
dmut = 0.001875*365*0.86 # Note, slightly lower than b, so that when total is low, population grows
m = (5e-6)*365
target = 100.0

vals0 = [100, 0]
tmax = 100.0

delta_wt = [1,-1,0,0,0];
delta_mut = [0,0,1,-1,1];

simres = [vals0]
times = [0.0]

while (times[end] < tmax) & (sum(simres[end]) > 0)
 global simres, times
 (tnew, valsnew) = update(times[end],simres[end])
 push!(simres,valsnew)
 push!(times,tnew)
end

sol = (t = times, u = simres)

plot(sol.t, [sol.u[i][1] for i in 1:length(sol.t)], labels="Wildtype", color= "blue", lw=1.0)
plot!(sol.t, [sol.u[i][2] for i in 1:length(sol.t)], labels="Mutant", color="red", lw=1.0)
plot!(sol.t, [sol.u[i][1]+sol.u[i][2] for i in 1:length(sol.t)], legend=:left, labels="Total", color="white", lw=0.2, title="Low Initial Mutation Load Where b<bmut", xaxis="Time (Years)", yaxis="Population Size")

plot!(sol.t, [sol.u[i][1] for i in 1:length(sol.t)], labels="", color= "blue",lw=0.4)
plot!(sol.t, [sol.u[i][2] for i in 1:length(sol.t)], labels="", color="red", lw=0.4)
plot!(sol.t, [sol.u[i][1]+sol.u[i][2] for i in 1:length(sol.t)], legend=:left,layout=4, labels="", color="white", lw=0.2)

function sim_gillespie(;b,d,bmut,dmut,m,target,vals0,tmax)
    function update(t, vals)
      wt = vals[1];
      mut = vals[2];
      haz = [ifelse(sum(vals)<target,b*wt,0.0), d*wt, ifelse(sum(vals)<target,bmut*mut,0.0), dmut*mut, m*wt];
      # ^ specifying the reactions. 1st = wt replication,  b, wildtype --> 2wildtype
      # 2nd = wt degredation, d, wildtype  --> null
      # 3rd = mut replication, bmut, mut --> 2mut
      # 4th = mut degredation, dmut, mut --> null
      # 5th = mutaion, m, wildtype --> wildtype + mut
      hazfrac = cumsum(haz/sum(haz));
      index = minimum(findall(hazfrac .- Random.rand() .> 0));
      tnew = t + (1.0/sum(haz))*(Random.randexp());
      valsnew = [wt + delta_wt[index], mut + delta_mut[index]]
      return((tnew,valsnew));
    end
    delta_wt = [1,-1,0,0,0];
    delta_mut = [0,0,1,-1,1];

    simres = [vals0]
    times = [0.0]

    while (times[end] < tmax) & (sum(simres[end]) > 0)
     (tnew, valsnew) = update(times[end],simres[end])
     push!(simres,valsnew)
     push!(times,tnew)
    end
    sol = (t = times, u = simres)
    return(sol)
end

sol = sim_gillespie(;b = 0.001875*365,d = 0.001875*365*0.86,bmut = 0.001875*365,dmut = 0.001875*365*0.86,m = (5e-6)*365, target = 100.0,vals0 = [100, 0],tmax = 100.0)

nsims=150
transparency = 0.4
thick = 0.2
solutions =[sim_gillespie(;b = 0.001875*365,d = 0.001875*365*0.86,bmut = 0.001875*365*1.0,dmut = 0.001875*365*0.86,m = (5e-7)*365, target = 100.0,vals0 = [100, 0],tmax = 100.0) for i in 1:nsims]
panelA = plot(solutions[1].t,[solutions[1].u[i][1]+solutions[1].u[i][2] for i in 1:length(solutions[1].t)], labels="", color=RGBA(0,0,0,transparency), lw=thick)
for j in 2:nsims
    plot!(solutions[j].t,[solutions[j].u[i][1]+solutions[j].u[i][2] for i in 1:length(solutions[j].t)], labels="", color=RGBA(0,0,0,transparency), lw=thick, legend=false, xaxis="Time (Years)", yaxis="mtDNA copy number", title="High initial pathogenic variant load where b=bmut", title_location=:left)
end
panelB = plot(solutions[1].t,[solutions[1].u[i][2]/(solutions[1].u[i][1]+solutions[1].u[i][2]) for i in 1:length(solutions[1].t)], labels="", color=RGBA(1,0,0,transparency), lw=thick)
for j in 2:nsims
    plot!(solutions[j].t,[solutions[j].u[i][2]/(solutions[j].u[i][1]+solutions[j].u[i][2]) for i in 1:length(solutions[j].t)], labels="", color=RGBA(1,0,0,transparency), lw=thick, legend=false, xaxis="Time (Years)", yaxis="Pathogenic variant load")
    #plot!(rand(100), lw=0)
end

plot(panelA,panelB,layout=(1,2), fmt=:pdf)

savefig("D:\\himl.pdf")

nsims=1
transparency = 1.0
thick = 0.7
solutions =[sim_gillespie(;b = 0.001875*365,d = 0.001875*365*0.86,bmut = 0.001875*365*1.1,dmut = 0.001875*365*0.86,m = (5e-7)*365, target = 100.0,vals0 = [100, 0],tmax = 100.0) for i in 1:nsims]
panelC = plot(solutions[1].t,[solutions[1].u[i][1]+solutions[1].u[i][2] for i in 1:length(solutions[1].t)], labels="mtDNA copy number", color=RGBA(0,0,0,transparency), lw=thick, legend=false,  xaxis="Time (Years)", yaxis="Population Size", title="Copy Number of Cell")
panelB = plot(solutions[1].t,[solutions[1].u[i][2] for i in 1:length(solutions[1].t)], color=RGBA(0.5,0,0.5,transparency), lw=thick, legend=false,  xaxis="Time (Years)", yaxis="Population Size", title="Mutant mtDNA")
panelA = plot(solutions[1].t,[solutions[1].u[i][1] for i in 1:length(solutions[1].t)], color=RGBA(0,0,1,transparency), lw=thick, legend=false,  xaxis="Time (Years)", yaxis="Population Size", title="Wild-type mtDNA")
panelD = plot(solutions[1].t,[solutions[1].u[i][2]/(solutions[1].u[i][1]+solutions[1].u[i][2]) for i in 1:length(solutions[1].t)], title="Pathogenic variant load", color=RGBA(1,0,0,transparency), lw=thick, legend=false, xaxis="Time (Years)", yaxis="Fraction Value")

plot(panelA,panelB,panelC,panelD,layout=(2,2), fmt=:pdf)

savefig("D:\\4panel.pdf")

# What happens if a cells is born with no mutations?

# What happens if a cell inherits a lot of mutations?  70%?

# What happens in both of the scenarios above if bmut > b?