#
# models.jl --
#
# Models of Point Spread Functions (PSF) for Julia.
#

const J1 = besselj1
const AIRY_FWHM = 1.02899397
const AIRY_FIRST_ZERO  = 1.22
const AIRY_SECOND_ZERO = 2.23
const AIRY_THIRD_ZERO  = 3.24
const AIRY_FOURTH_ZERO = 4.24
const AIRY_FIFTH_ZERO  = 5.24

"""

Instances of sub-types of `PointSpreadFunctions.AbstractPSF{N}` are used to
store the parameters of various *Point Spread Functions* (PSF) models, `N`
is the number of parameters.  Let `P` be such an instance, then:

    P([T,] r)
    P([T,] x, y)

yield an approximation of the PSF at a distance `r` or at a position
`(x,y)` relative to the position of the point source.

Use index notation to retrieve the parameters as an `N`-tuple or a specific
parameters:

    P[:]  # yields all parameters
    P[i]  # yields i-th parameter

The PSF is normalized such that its peak value (usually at the center) is
equal to one.  The rationale is that normalization by the peak value does
not depend on the dimensionality while the integral of the PSF does depend
on the dimensionality.

See also [`AiryPSF`](@ref), [`CauchyPSF`](@ref), [`GaussianPSF`](@ref),
[`MoffatPSF`](@ref).

""" AbstractPSF

abstract type AbstractPSF{N} end

@inline getindex(P::AbstractPSF, ::Colon) = parameters(P)
@inline @propagate_inbounds getindex(P::AbstractPSF, i::Integer) =
    getindex(parameters(P), i)

"""
    getfwhm(P)

yields the full width at half maximum (FWHM) of point spread function `P`.
Computing the result may require some calculations, so it is better to save
the result if this value must be used several times.

""" getfwhm

"""
    check_structure(psf)

throws an exception if the parameters of the point spread function `psf`
are invalid.  Call [`isvalid(psf)`](@ref) to check the validity of the
parameters without throwing exceptions.

"""
check_structure(psf::AbstractPSF) =
    isvalid(psf) || throw_bad_argument("PSF has invalid parameter(s)")

"""
    isvalid(psf) -> bool

yields whether the parameters of the point spread function `psf` are
correct.  The [`check_structure`](@ref) method is similar but throws for
invalid parameters.

"""
isvalid(::AbstractPSF) = true

"""
    P = AiryPSF(lim [, ρ=0])

defines the *Point Spread Function* (PSF) `P` of a circular pupil.
Argument `lim` is the distance in the focal plane corresponding to the
diffraction limit.  Assuming the distance is given in angular units
(radians), `lim = λ/D` with `λ` the wavelength and `D` the pupil diameter
(both expressed in the same units).  Optional argument `ρ` is the ratio of
the size of the central obscuration to that of the pupil (necessarily `0 ≤
ρ < 1`).  If not specified, `ρ = 0` is assumed (unobstructed pupil).  Then:

    P([T,] r)
    P([T,] x, y)

yield the PSF at distance `r` or position `(x,y)` relative to the position of
the source and normalized so that the peak value of the PSF is one.  Optional
argument `T` is the floating-point type of the result.  Arguments `lim`, `r`,
`x` and `y` are in the same units.

The peak-normalized Airy PSF at distance `r` is given by:

    P(r) = ((a/(c*r))*(J1(c*r) - ρ*J1(ρ*c*r)))^2

with `a = 2/(1 - ρ^2)`, `c = π/lim` and `J1` the Bessel function of the
first kind.

See also [`CauchyPSF`](@ref), [`GaussianPSF`](@ref), [`MoffatPSF`](@ref).

""" AiryPSF

struct AiryPSF <: AbstractPSF{2}
    prm::NTuple{2,Float64}
    a::Float64
    c::Float64
    AiryPSF(lim::Real, rho::Real = 0.0) =
        new((lim, rho),
            (rho == 0 ? 2.0 : 2.0/((1.0 - rho)*(1.0 + rho))),
            π/lim)
end

@inline isvalid(P::AiryPSF) = _check_airy_lim(P) && _check_airy_rho(P)

function check_structure(P::AiryPSF)
    _check_airy_lim(P) ||
        throw_bad_argument("bad diffraction limit value")
    _check_airy_rho(P) ||
        throw_bad_argument("out of range central obscuration ratio")
end

@inline _check_airy_lim(P::AiryPSF) = finite_and_positive(_get_lim(P))
@inline _check_airy_rho(P::AiryPSF) = (ρ = _get_rho(P);
                                       isfinite(ρ) && 0.0 ≤ ρ < 1.0)

Base.show(io::IO, P::AiryPSF) =
    print(io, "AiryPSF(", P[1], ",", P[2], ")")

@inline parameters(P::AiryPSF) = getfield(P, :prm)
@inline obstrucated(P::AiryPSF) =  _get_rho(P) == 0

@inline _get_lim(P::AiryPSF) = @inbounds P[1]
@inline _get_rho(P::AiryPSF) = @inbounds P[2]
@inline _get_a(P::AiryPSF) = getfield(P, :a)
@inline _get_c(P::AiryPSF) = getfield(P, :c)

function getfwhm(P::AiryPSF; kwds...)
    l, ρ, a, c =  _get_lim(P), _get_rho(P), _get_a(P), _get_c(P)
    v = 1/sqrt(2)
    f(r) = v - (ρ == 0 ? _circular_amplitude(a, c, r) :
                _annular_amplitude(a, c, ρ, r))
    return fzero(f, 0.52*l, 0.35*l; kwds...)[1]*2
end

@inline _circular_amplitude(a::T, c::T, r::T) where {T<:AbstractFloat} =
    r == zero(T) ? one(T) : (u = c*r; (a/u)*J1(u))

@inline _annular_amplitude(a::T, c::T, ρ::T, r::T) where {T<:AbstractFloat} =
    r == zero(T) ? one(T) : (u = c*r; (ρ == zero(T) ?
                                       (a/u)*J1(u) :
                                       (a/u)*(J1(u) - ρ*J1(ρ*u))))

# No central obscuration.
@inline _airy(a::T, c::T, r::T) where {T<:AbstractFloat} =
    r == zero(T) ? one(T) : (u = c*r; ((a/u)*J1(u))^2)

# With central obscuration.
@inline _airy(a::T, c::T, ρ::T, r::T) where {T<:AbstractFloat} =
    r == zero(T) ? one(T) : (u = c*r; (ρ == zero(T) ?
                                       (a/u)*J1(u) :
                                       (a/u)*(J1(u) - ρ*J1(ρ*u)))^2)

(P::AiryPSF)(::Type{T}, r::Real) where {T<:AbstractFloat} =
    _airy(to_float(T, _get_a(P)),
          to_float(T, _get_c(P)),
          to_float(T, _get_rho(P)),
          to_float(T, r))

(P::AiryPSF)(::Type{T}, x::Real, y::Real) where {T<:AbstractFloat} =
    P(T, hypot(to_float(T, x), to_float(T, y)))

function (P::AiryPSF)(::Type{T},
                      r::AbstractArray{<:Real}) where {T<:AbstractFloat}
    a = to_float(T, _get_a(P))
    c = to_float(T, _get_c(P))
    ρ = to_float(T, _get_rho(P))
    ρ == zero(ρ) ?
        map(r -> _airy(a, c,    to_float(T, r)), r) :
        map(r -> _airy(a, c, ρ, to_float(T, r)), r)
end

function (P::AiryPSF)(::Type{T},
                      x::AbstractArray{<:Real},
                      y::AbstractArray{<:Real}) where {T<:AbstractFloat}
    a = to_float(T, _get_a(P))
    c = to_float(T, _get_c(P))
    ρ = to_float(T, _get_rho(P))
    ρ == zero(ρ) ?
        broadcast((x, y) -> _airy(a, c,    hypot(to_float(T, x), to_float(T, y))),
                  x, y) :
        broadcast((x, y) -> _airy(a, c, ρ, hypot(to_float(T, x), to_float(T, y))),
                  x, y)
end

"""
    findzeros([T=Float64,] P, n; kwds...)

yield the `n`-th zero of *Point Spread Function* (PSF) `P`.  If `n` is a
vector of integers, a vector with the corresponding zeros is returned.
For example:

    findzeros(P, 1:3)

yields the first 3 zeros of `P`.

Optional argument `T` can be used to specify the floating-point type of the
result.  The keywords of `OptimPack.fzero` can be specified.

"""
findzeros(P::AiryPSF, args...; kwds...) =
    findzeros(Float64, P::AiryPSF, args...; kwds...)

function findzeros(::Type{T},
                   P::AiryPSF,
                   n::AbstractVector{I};
                   kwds...) where {T<:AbstractFloat,I<:Integer}
    result = Array{T}(undef, length(n))
    for i in eachindex(result, n)
        result[i] = findzeros(T, P, n[i])
    end
    return result
end

# FIXME: the initial bracket has to be improved for obstruction ratio around
# 0.5, in fact a method like BRADI should be used
function findzeros(::Type{T}, P::AiryPSF, n::Integer;
                   kwds...) where {T<:AbstractFloat}
    @assert n ≥ 1
    l, ρ, a, c = T(_get_lim(P)), T(_get_rho(P)), T(_get_a(P)), T(_get_c(P))
    g = (n + T(0.25) - T(0.49)*ρ)*l  # guess for the n-th zero
    h = T(0.7)*l                     # half-width of search interval
    obstrucated(P) ?
        fzero(r -> _annular_amplitude(a, c, ρ, r), g - h, g + h; kwds...)[1] :
        fzero(r -> _circular_amplitude(a, c, r),   g - h, g + h; kwds...)[1]
end

"""
    P = CauchyPSF(fwhm)

defines a Cauchy *Point Spread Function* (PSF) `P` (or Lorentzian) of full
width at half maximum `fwhm` which can be used as:

    P([T,] r)
    P([T,] x, y)

to compute the PSF at distance `r` or position `(x,y)` relative to the
position of the source and normalized so that the peak value of the PSF is
one.  Optional argument `T` is the floating-point type of the result.
Arguments `fwhm`, `r`, `x` and `y` are in the same units.

The peak-normalized Cauchy PSF at distance `r` is given by:

    P(r) = q/(q + r^2)

with `q = (fwhm/2)^2`.

See also [`AiryPSF`](@ref), [`GaussianPSF`](@ref), [`MoffatPSF`](@ref).

""" CauchyPSF

struct CauchyPSF <: AbstractPSF{1}
    fwhm::Float64
    q::Float64
    CauchyPSF(fwhm::Real) = new(fwhm, (fwhm/2)^2)
end

@inline isvalid(P::CauchyPSF) = finite_and_positive(getfwhm(P))

check_structure(P::CauchyPSF) =
    finite_and_positive(getfwhm(P)) || throw_bad_argument("bad FWHM value")

Base.show(io::IO, P::CauchyPSF) =
    print(io, "CauchyPSF(", P[1], ")")

@inline getfwhm(P::CauchyPSF) = getfield(P, :fwhm)
@inline parameters(P::CauchyPSF) = (getfwhm(P),)

@inline _get_q(P::CauchyPSF) = getfield(P, :q)

@inline _cauchy(q::T, r::T) where {T<:AbstractFloat} = q/(r*r + q)

@inline _cauchy(q::T, x::T, y::T) where {T<:AbstractFloat} = q/(x*x + y*y + q)

(P::CauchyPSF)(::Type{T}, r::Real) where {T<:AbstractFloat} =
    _cauchy(to_float(T, _get_q(P)),
            to_float(T, r))

(P::CauchyPSF)(::Type{T}, x::Real, y::Real) where {T<:AbstractFloat} =
    _cauchy(to_float(T, _get_q(P)),
            to_float(T, x),
            to_float(T, y))

function (P::CauchyPSF)(::Type{T},
                        r::AbstractArray{<:Real}) where {T<:AbstractFloat}
    q = to_float(T, _get_q(P))
    map(r -> _cauchy(q, to_float(T, r)), r)
end

function (P::CauchyPSF)(::Type{T},
                        x::AbstractArray{<:Real},
                        y::AbstractArray{<:Real}) where {T<:AbstractFloat}

    q = to_float(T, _get_q(P))
    broadcast((x, y) -> _cauchy(q, to_float(T, x), to_float(T, y)), x, y)
end

"""
    P = GaussianPSF(fwhm)

defines a Gaussian *Point Spread Function* (PSF) `P` of full width at half
maximum `fwhm` which can be used as:

    P([T,] r)
    P([T,] x, y)

to compute the PSF at distance `r` or position `(x,y)` relative to the
position of the source and normalized so that the peak value of the PSF is
one.  Optional argument `T` is the floating-point type of the result.
Arguments `fwhm`, `r`, `x` and `y` are in the same units.

The peak-normalized Gaussian PSF at distance `r` is given by:

    P(r) = exp(q*r^2)

with `q = -4*log(2)/fwhm^2`.

See also [`AiryPSF`](@ref), [`CauchyPSF`](@ref), [`MoffatPSF`](@ref).

""" GaussianPSF

struct GaussianPSF <: AbstractPSF{1}
    fwhm::Float64
    q::Float64
    GaussianPSF(fwhm::Real) = new(fwhm, -log(2.0)*(2.0/fwhm)^2)
end

@inline isvalid(P::GaussianPSF) = finite_and_positive(getfwhm(P))

check_structure(P::GaussianPSF) =
    finite_and_positive(getfwhm(P)) || throw_bad_argument("bad FWHM value")

Base.show(io::IO, P::GaussianPSF) =
    print(io, "GaussianPSF(", P[1], ")")

@inline getfwhm(P::GaussianPSF) = getfield(P, :fwhm)
@inline parameters(P::GaussianPSF) = (getfwhm(P),)

@inline _get_q(P::GaussianPSF) = getfield(P, :q)

@inline _gauss(q::T, r::T) where {T<:AbstractFloat} = exp(r*r*q)

@inline _gauss(q::T, x::T, y::T) where {T<:AbstractFloat} = exp((x*x + y*y)*q)

(P::GaussianPSF)(::Type{T}, r::Real) where {T<:AbstractFloat} =
    _gauss(to_float(T, _get_q(P)),
           to_float(T, r))

(P::GaussianPSF)(::Type{T}, x::Real, y::Real) where {T<:AbstractFloat} =
    _gauss(to_float(T, _get_q(P)),
           to_float(T, x),
           to_float(T, y))

function (P::GaussianPSF)(::Type{T},
                        r::AbstractArray{<:Real}) where {T<:AbstractFloat}
    q = to_float(T, _get_q(P))
    map(r -> _gauss(q, to_float(T, r)), r)
end

function (P::GaussianPSF)(::Type{T},
                        x::AbstractArray{<:Real},
                        y::AbstractArray{<:Real}) where {T<:AbstractFloat}

    if axes(x) == axes(y)
        q = to_float(T, _get_q(P))
        broadcast((x, y) -> _gauss(q, to_float(T, x), to_float(T, y)), x, y)
    else
        P(T,x).*P(T,y)
    end
end

"""
    P = MoffatPSF(fwhm, beta)

defines a Moffat PSF `P` of full width at half maximum `fwhm` and exponent
`beta` which can be used as:

    P([T,] r)
    P([T,] x, y)

to compute the PSF at distance `r` or position `(x,y)` relative to the
position of the source and normalized so that the peak value of the PSF is
one.  Optional argument `T` is the floating-point type of the result.
Arguments `fwhm`, `r`, `x` and `y` are in the same units.

The peak-normalized Moffat PSF at distance `r` is given by:

    P(r) = (1 + q*r^2)^(-beta)

with `q = (2^(1/beta) - 1)*(2/fwhm)^2`.

See also [`AiryPSF`](@ref), [`CauchyPSF`](@ref), [`GaussianPSF`](@ref).

""" MoffatPSF

struct MoffatPSF <: AbstractPSF{2}
    prm::NTuple{2,Float64} # (fwhm, beta)
    p::Float64 # (2^(1/beta) - 1)*(2/fwhm)^2
    q::Float64 # -beta
    MoffatPSF(fwhm::Real, beta::Real) =
        new((fwhm, beta), -beta, (2.0^(1.0/beta) - 1.0)*(2.0/fwhm)^2)
end

@inline isvalid(P::MoffatPSF) =
    finite_and_positive(_get_fwhm(P)) && finite_and_positive(_get_beta(P))

function check_structure(P::MoffatPSF)
    finite_and_positive(_get_fwhm(P)) ||
        throw_bad_argument("bad FWHM value")
    finite_and_positive(_get_beta(P)) ||
        throw_bad_argument("bad value for exponent `beta`")
end

Base.show(io::IO, P::MoffatPSF) =
    print(io, "MoffatPSF(", P[1], ",", P[2], ")")

@inline getfwhm(P::MoffatPSF) = _get_fwhm(P)
@inline parameters(P::MoffatPSF) =  getfield(P, :prm)

@inline _get_fwhm(P::MoffatPSF) = @inbounds P[1]
@inline _get_beta(P::MoffatPSF) = @inbounds P[2]
@inline _get_p(P::MoffatPSF) = getfield(P, :p)
@inline _get_q(P::MoffatPSF) = getfield(P, :q)

@inline _moffat(p::T, q::T, r::T) where {T<:AbstractFloat} =
    (r*r*q + one(T))^p

@inline _moffat(p::T, q::T, x::T, y::T) where {T<:AbstractFloat} =
    ((x*x + y*y)*q + one(T))^p

(P::MoffatPSF)(::Type{T}, r::Real) where {T<:AbstractFloat} =
    _moffat(to_float(T, _get_p(P)),
            to_float(T, _get_q(P)),
            to_float(T, r))

(P::MoffatPSF)(::Type{T}, x::Real, y::Real) where {T<:AbstractFloat} =
    _moffat(to_float(T, _get_p(P)),
            to_float(T, _get_q(P)),
            to_float(T, x),
            to_float(T, y))

function (P::MoffatPSF)(::Type{T},
                        r::AbstractArray{<:Real}) where {T<:AbstractFloat}
    p = to_float(T, _get_p(P))
    q = to_float(T, _get_q(P))
    map(r -> _moffat(p, q, to_float(T, r)), r)
end

function (P::MoffatPSF)(::Type{T},
                        x::AbstractArray{<:Real},
                        y::AbstractArray{<:Real}) where {T<:AbstractFloat}

    p = to_float(T, _get_p(P))
    q = to_float(T, _get_q(P))
    broadcast((x, y) -> _moffat(p, q, to_float(T, x), to_float(T, y)), x, y)
end

# Define methods common to all sub-types of AbstractPSF.
for T in (:AiryPSF, :CauchyPSF, :GaussianPSF, :MoffatPSF)
    @eval begin
        (P::$T)(r::Real) = P(Float64, r)
        (P::$T)(r::AbstractArray) = P(Float64, r)
        (P::$T)(x::Real, y::Real) = P(Float64, x, y)
        (P::$T)(x::AbstractArray, y::AbstractArray) = P(Float64, x, y)
    end
end

@noinline throw_bad_argument(args...) = throw_bad_argument(string(args...))
@noinline throw_bad_argument(mesg::String) = throw(ArgumentError(mesg))

@inline finite_and_positive(x::Real) = (isfinite(x) && x > 0)
@inline finite_and_nonnegative(x::Real) = (isfinite(x) && x ≥ 0)

"""
    to_float(T, x)

lazily yields `x` converted to floating-point type `T`.

"""
to_float(::Type{T}, x::T   ) where {T<:AbstractFloat} = x
to_float(::Type{T}, x::Real) where {T<:AbstractFloat} = T(x)

"""
    to_int(x)

lazily yields `x` converted to integer type `int`.

"""
to_int(x::Int) = x
to_int(x::Integer) = Int(x)


using SententialDecisionDiagrams

# set up vtree and manager
vtree = SententialDecisionDiagrams.read_vtree("$(@__DIR__)/input/simple.vtree")
sdd = SententialDecisionDiagrams.sdd_manager(vtree)

println("Created an SententialDecisionDiagrams with $(SententialDecisionDiagrams.var_count(sdd)) variables")
root = SententialDecisionDiagrams.read_cnf("$(@__DIR__)/input/simple.cnf", sdd, compiler_options=SententialDecisionDiagrams.CompilerOptions(vtree_search_mode=-1))
# For DNF functions use `read_dnf_file`

# Model Counting
wmc = SententialDecisionDiagrams.wmc_manager(root, log_mode=true)
w = SententialDecisionDiagrams.propagate(wmc)
println("Model count: $(convert(Int32,exp(w)))")

# Weighted Model Counting
lits = [SententialDecisionDiagrams.literal(i,sdd) for i in 1:SententialDecisionDiagrams.var_count(sdd)]

# Positive literal weight
SententialDecisionDiagrams.set_literal_weight(lits[1], log(0.5), wmc)
# Negative literal weight
SententialDecisionDiagrams.set_literal_weight(~lits[1], log(0.5), wmc)

w = SententialDecisionDiagrams.propagate(wmc)
println("Weighted model count: $(exp(w))")

# Visualize SententialDecisionDiagrams and VTREE
println("saving sdd and vtree ... ")
SententialDecisionDiagrams.dot("$(@__DIR__)/output/simple-vtree.dot", vtree)
SententialDecisionDiagrams.dot("$(@__DIR__)/output/simple-sdd-cnf.dot", root)


loaddict(filename::AbstractString) = Dict(lval => rval for (rval, lval) in split.(eachline(filename), " -> "))
function prettyprint(dict)
    cache = Dict{AbstractString, UInt16}()
    foreach(sort!(collect(keys(dict)))) do key
        println(key, ": ", bobbycalculate(dict, key, cache))
    end
end

bobbycalculate(filename::AbstractString, var) = bobbycalculate(loaddict(filename), var)
function bobbycalculate(dict::Dict, var::AbstractString, cache = Dict{AbstractString, UInt16}())
    get!(cache, var) do
        str = dict[var]
        if all(isdigit, str)
            parse(UInt16, str)
        elseif startswith(str, "NOT")
            ~bobbycalculate(dict, str[5:end], cache)
        elseif str in keys(dict)
            bobbycalculate(dict, str, cache)
        else
            l, i, r = match(r"(.+) (AND|OR|[LR]SHIFT) (.+)", str)
            lval = all(isdigit, l) ? parse(UInt16, l) : bobbycalculate(dict, l, cache)
            instr = if i == "AND"
                (&)
            elseif i == "OR"
                (|)
            elseif i == "LSHIFT"
                (<<)
            elseif i == "RSHIFT"
                (>>)
            end
            rval = all(isdigit, r) ? parse(UInt16, r) : bobbycalculate(dict, r, cache)
            instr(lval, rval)
        end
    end
end

if basename(pwd()) == "aoc"
    cd("2015/7")
end
part1() = bobbycalculate(loaddict("input.txt"), "a")
part2() = bobbycalculate(loaddict("input.txt"), "a", Dict("b"=>0xb3f1))


using PDMats  # Positive Definite Matrices
using QuadGK  # 1-d integration
using Cuba    # multi-d integration
using Optim   # Optimizing over non-linear parameters

# Read dataset
data = readdlm("../data/rvs_0001.txt");
times = data[:,1]
vels = data[:,2]
sigma_obs = data[:,3]
nobs = length(vels)

# Set priors
const Cmax = 1000.0
function logprior_offset(param)
  -log(2*Cmax)
end

const sigma_j_min = 0.0
const sigma_j_max = 99.0
const sigma_j_0 = 1.0
function logprior_jitter(sigma_j::Real)
  -log1p(sigma_j/sigma_j_0)-log(log1p(sigma_j_max/sigma_j_0)) # -log(sigma_j_0)
end

const Pmin = 1.0
const Pmax = 1.0e4
function logprior_period(P::Real)
   -log(P)-log(log(Pmax/Pmin))
end
function logprior_periods(P::Vector{T}) where T<:Real
   if length(P) == 0  return 0 end
   -sum(log.(P))-length(P)*log(log(Pmax/Pmin))
end

const Kmin = 0.0
const Kmax = 999.0
const K0 = 1.0
function logprior_amplitude(K::Real)
  logp = -log1p(K/K0)-log(log1p(Kmax/K0)) # -log(K0)
end

function logprior_sinusoids(param)
  const nparam = length(param)
  const namp_per_pl = 2
  npl = convert(Int64,floor((nparam-1)//namp_per_pl))
  logp = 0
  logj = 0
  for i in 0:(npl-1)
    K = sqrt(param[1+namp_per_pl*i]^2+param[2+namp_per_pl*i]^2)
    logp += logprior_amplitude(K)
	logj -= log(K)
  end
  logp -= npl*log(2pi)
  C = param[nparam]
  if !(-Cmax<=C<=Cmax) logp -=Inf
  else  logp -= log(2*Cmax) end
  return (logp,logj)
end

function logprior_epicycles(param)
  nparam = length(param)
  namp_per_pl = 4
  npl = convert(Int64,floor((nparam-1)//namp_per_pl))
  logp = 0
  logj = 0
  for i in 0:(npl-1)
    K  = sqrt(param[1+namp_per_pl*i]^2+param[2+namp_per_pl*i]^2)
	eK = sqrt(param[3+namp_per_pl*i]^2+param[4+namp_per_pl*i]^2)
	if !(0.0<=eK<K) logp -= Inf end
    logp += logprior_amplitude(K)
	logj -= 2*log(K)+log(eK)
  end
  logp -= 2*npl*log(2pi)
  C = param[nparam]
  if !(-Cmax<=C<=Cmax) logp -=Inf
  else  logp -= log(2*Cmax) end
  return (logp,logj)
end

#sincos from LombScargle.jl
# `sincos` is not yet in Julia, but it's available through the math library.  The following
# definitions are from Yichao Yu's code at https://github.com/nacs-lab/yyc-data and are
# likely to be later added to Julia (so they can be removed here).  See also
# https://discourse.julialang.org/t/poor-performance-of-the-cis-function/3402.
if !isdefined(Base, :sincos)
    @inline function sincos(v::Float64)
        Base.llvmcall("""
        %f = bitcast i8 *%1 to void (double, double *, double *)*
        %pres = alloca [2 x double]
        %p1 = getelementptr inbounds [2 x double], [2 x double]* %pres, i64 0, i64 0
        %p2 = getelementptr inbounds [2 x double], [2 x double]* %pres, i64 0, i64 1
        call void %f(double %0, double *nocapture noalias %p1, double *nocapture noalias %p2)
        %res = load [2 x double], [2 x double]* %pres
        ret [2 x double] %res
        """, Tuple{Float64,Float64}, Tuple{Float64,Ptr{Void}}, v,
                      cglobal((:sincos, Base.libm_name)))
    end
    @inline function sincos(v::Float32)
        Base.llvmcall("""
        %f = bitcast i8 *%1 to void (float, float *, float *)*
        %pres = alloca [2 x float]
        %p1 = getelementptr inbounds [2 x float], [2 x float]* %pres, i64 0, i64 0
        %p2 = getelementptr inbounds [2 x float], [2 x float]* %pres, i64 0, i64 1
        call void %f(float %0, float *nocapture noalias %p1, float *nocapture noalias %p2)
        %res = load [2 x float], [2 x float]* %pres
        ret [2 x float] %res
            """, Tuple{Float32,Float32}, Tuple{Float32,Ptr{Void}}, v,
                      cglobal((:sincosf, Base.libm_name)))
    end
    @inline function sincos(v)
        @fastmath (sin(v), cos(v))
    end
end

# Fit planet models
function make_design_matrix_circ(times::Vector{T},periods::Vector{T}) where T<:Real
  const namp_per_pl = 2
  nparam = 1+namp_per_pl*length(periods)
  freqs = 2pi./periods
  X = Array{T}(length(times),nparam)
  for i in 1:length(periods)
     for j in 1:length(times)
       (X[j,1+(i-1)*namp_per_pl],X[j,2+(i-1)*namp_per_pl]) = sincos(freqs[i]*times[j])
	 end
  end
  X[:,nparam] = one(T)
  return X
end

function make_design_matrix_circ_old(times::Vector{T},periods::Vector{T}) where T<:Real
  const namp_per_pl = 2
  nparam = 1+namp_per_pl*length(periods)
  freqs = 2pi./periods
  X = Array{T}(length(times),nparam)
  for i in 1:length(periods)
     X[:,1+(i-1)*namp_per_pl] = sin.(freqs[i]*times)
	 X[:,2+(i-1)*namp_per_pl] = cos.(freqs[i]*times)
  end
  X[:,nparam] = one(T)
  return X
end

function update_design_matrix_circ(X::Array{T,2}, times::Vector{T},periods::Vector{T}, pl_to_update::Integer) where T<:Real
  const namp_per_pl = 2
  nparam = 1+namp_per_pl*length(periods)
  freq = 2pi./periods[pl_to_update]
  @assert size(X) == (length(times),nparam)
  i = pl_to_update
  idx_s = 1+(i-1)*namp_per_pl
  idx_c = 2+(i-1)*namp_per_pl
  for j in 1:length(times)
     (X[j,idx_s],X[j,idx_c]) = sincos(freq*times[j])
  end
  return X
end

function make_design_matrix_epicycle(times::Vector{T},periods::Vector{T}) where T<:Real
  const namp_per_pl = 4
  nparam = 1+namp_per_pl*length(periods)
  freqs = 2pi./periods
  X = Array{T}(length(times),nparam)
  for i in 1:length(periods)
     s = sin.(freqs[i]*times)
	 c = cos.(freqs[i]*times)
	 X[:,1+(i-1)*namp_per_pl] = s
	 X[:,2+(i-1)*namp_per_pl] = c
	 X[:,3+(i-1)*namp_per_pl] = 2*s.*c
	 X[:,4+(i-1)*namp_per_pl] = c.^2.-s.^2
  end
  X[:,nparam] = one(T)
  return X
end

function update_design_matrix_epicycle(X::Array{T,2}, times::Vector{T},periods::Vector{T}, pl_to_update::Integer) where T<:Real
  const namp_per_pl = 4
  nparam = 1+namp_per_pl*length(periods)
  freq = 2pi./periods[pl_to_update]
  @assert size(X) == (length(times),nparam)
  i = pl_to_update
  idx_s = 1+(i-1)*namp_per_pl
  for j in 1:length(times)
	 (s,c) = sincos(freq*times[j])
     X[j,idx_s  ] = s
	 X[j,idx_s+1] = c
	 X[j,idx_s+2] = 2*s.*c
	 X[j,idx_s+3] = c.^2.-s.^2
  end
  return X
end

function calc_best_fit(data::Vector{T}, model::Array{T,2}, Sigma::PDMat{T}) where T<:Real
  inv_chol_Sigma = inv(Sigma.chol[:L])
  param_bestfit = (inv_chol_Sigma*model) \ (inv_chol_Sigma*data)
end

function calc_best_fit(data::Vector{T}, model::Array{T,2}, Sigma::PDiagMat{T}) where T<:Real
  inv_chol_Sigma = sqrt.(Sigma.inv_diag)
  param_bestfit = (inv_chol_Sigma.*model) \ (inv_chol_Sigma.*data)
end

function calc_chisq(data::Vector{T}, model::Array{T,2}, Sigma::SigmaT; FIM::PDMat{T} = PDMat(Xt_invA_X(Sigma,model)), calclogprior::Function = x->(0.0,0.0) ) where {T<:Real, SigmaT<:AbstractPDMat{T} }
  @assert length(data) == size(model,1) == size(Sigma,1)
  nobs = length(data)
  nparam = size(model,2)

  param_bf_linalg = calc_best_fit(data,model,Sigma)
  predict = model*param_bf_linalg
  delta = data.-predict
  chisq = invquad(Sigma,delta) # sum(delta'*Sigma^-1*delta)
end

function compute_Laplace_Approx_param_and_integral(data::Vector{T}, model::Array{T,2}, Sigma::SigmaT; FIM::PDMat{T} = PDMat(Xt_invA_X(Sigma,model)), calclogprior::Function = x->(0.0,0.0) ) where {T<:Real, SigmaT<:AbstractPDMat{T} }
  @assert length(data) == size(model,1) == size(Sigma,1)
  nobs = length(data)
  nparam = size(model,2)

  param_bf_linalg = calc_best_fit(data,model,Sigma)
  predict = model*param_bf_linalg
  delta = data.-predict
  chisq = invquad(Sigma,delta) # sum(delta'*Sigma^-1*delta)

  loglikelihood_mode = 0.5*(-chisq-logdet(Sigma)-nobs*log(2pi))
  (logprior_mode, log_jacobian) = calclogprior(param_bf_linalg)
  log_mode = logprior_mode + loglikelihood_mode

  sigma_param = inv(FIM).chol[:U]
  LaplaceApprox = 0.5*nparam*log(2pi)-0.5*logdet(FIM)+log_mode + log_jacobian
  return (param_bf_linalg,sigma_param,LaplaceApprox)
end

function compute_Laplace_Approx(data::Vector{T}, model::Array{T,2}, Sigma::SigmaT; FIM::PDMat{T} = PDMat(Xt_invA_X(Sigma,model)), calclogprior::Function = x->(0.0,0.0) ) where {T<:Real, SigmaT<:AbstractPDMat{T} }
  @assert length(data) == size(model,1) == size(Sigma,1)
  nobs = length(data)
  nparam = size(model,2)

  param_bf_linalg = calc_best_fit(data,model,Sigma)
  predict = model*param_bf_linalg
  delta = data.-predict
  chisq = invquad(Sigma,delta) # sum(delta'*Sigma^-1*delta)

  loglikelihood_mode = 0.5*(-chisq-logdet(Sigma)-nobs*log(2pi))
  (logprior_mode, log_jacobian) = calclogprior(param_bf_linalg)
  log_mode = logprior_mode + loglikelihood_mode

  LaplaceApprox = 0.5*nparam*log(2pi)-0.5*logdet(FIM)+log_mode + log_jacobian
  return LaplaceApprox
end

function generate_K_samples(bf::Vector, sigma, nsamples::Integer = 1)
  nparam = length(bf)
  const namp_per_pl = 2
  theta = bf .+ sigma*randn(nparam,nsamples)
  #theta = param_bf'.+randn(nsamples,nparam)*sigma_param
  npl = convert(Int64,floor((nparam-1)//namp_per_pl))

  Ks = zeros(nsamples,npl)
  for i in 1:npl
    Ks[:,i] = sqrt.(theta[1+namp_per_pl*(i-1),:].^2+theta[2+namp_per_pl*(i-1),:].^2)
  end
  return Ks
end

function generate_Ke_samples(bf::Vector, sigma, nsamples::Integer = 1)
  nparam = length(bf)
  const namp_per_pl = 4
  theta = bf .+ sigma*randn(nparam,nsamples)
  npl = convert(Int64,floor((nparam-1)//namp_per_pl))

  Ks = zeros(nsamples,npl)
  es = zeros(nsamples,npl)
  hs = zeros(nsamples,npl)
  ks = zeros(nsamples,npl)
  for i in 1:npl
    Ks[:,i] = sqrt.(theta[1+namp_per_pl*(i-1),:].^2+theta[2+namp_per_pl*(i-1),:].^2)
	hs[:,i] = theta[3+namp_per_pl*(i-1),:] ./ Ks[:,i]
	ks[:,i] = theta[4+namp_per_pl*(i-1),:] ./ Ks[:,i]
	es[:,i] = sqrt.(hs[:,i].^2.+ks[:,i].^2)
  end
  return (Ks,es,hs,ks)
end

# Now with covariances
function kernel_quaesiperiodic(dt::Float64)
  alpha_sq     = 3.0
  tau_evolve   = 50.0
  lambda_p     = 0.5
  tau_rotate   = 20.0
  alpha_sq*exp(-0.5*((dt/tau_evolve)^2 + (sin(pi*dt/tau_rotate)/lambda_p)^2) )
end

function make_Sigma(t::Vector{Float64}, sigma_obs::Vector{Float64}, sigma_j::Float64)
  PDMat( [ kernel_quaesiperiodic(t[i]-t[j]) + (i==j ? sigma_obs[i]^2 + sigma_j^2 : 0.0) for i in 1:length(t), j in 1:length(t)] )
end

function make_dSigmadsigmaj(t::Vector{Float64}, sigma_obs::Vector{Float64}, sigma_j::Float64)
  #PDMat( [ 0.0 + (i==j ? 2*sigma_j : 0.0) for i in 1:length(t), j in 1:length(t)] )
  PDiagMat( 2*sigma_j*ones(length(t)) )
  ScalMat(length(t),2*sigma_j)
end

#=
function compute_log_evidence(sigma_j::Real)
  Sigma = make_Sigma(times,sigma_obs,sigma_j)
  log_prior_outside = logprior_jitter(sigma_j) + logprior_periods(periods)
  log_likelihood = compute_Laplace_Approx(vels,X,Sigma,calclogprior=logprior_epicycles)
  log_prior_outside+log_likelihood
end

function compute_evidence(sigma_j::Real; log_normalization::Real=0.0)
  exp(compute_log_evidence(sigma_j)-log_normalization)
end
=#

#=
periods = [42.0]
X = make_design_matrix_epicycle(times,periods);
Sigma = PDMat(make_Sigma(times,sigma_obs,sigma_j_0))
FIM = PDMat(Xt_invA_X(Sigma,X))
(param_bf,sigma_param,evid) = compute_Laplace_Approx_param_and_integral(vels,X,Sigma,FIM=FIM,calclogprior=logprior_epicycles)

(Ks,es) = generate_Ke_samples(param_bf,sigma_param,1000)
=#

#=
lognorm = compute_log_evidence(sigma_j_0)  # can help to prevent underflow
integrand(x) = compute_evidence(x,log_normalization=lognorm)
result = QuadGK.quadgk(integrand,sigma_j_min,sigma_j_max).*(exp(lognorm)/log(10))
=#

function calc_opt_sigmaj(times::Vector{T}, vels::Vector{T}, sigma_obs::Vector{T}, model::Array{T,2}; calclogprior_linear::Function = logprior_sinusoids) where T<:Real
  function calc_laplace_approx_arg_jitter(sigmaj::Real)
    Sigma = Sigma0 + ScalMat(length(times),sigmaj) # PDMat(make_Sigma(times,sigma_obs,sigmaj))
    FIM = PDMat(Xt_invA_X(Sigma,model))
    evid = compute_Laplace_Approx(vels,model,Sigma,FIM=FIM,calclogprior=calclogprior_linear)
	evid += logprior_jitter(sigmaj)
    -evid
  end
  Sigma0 = PDMat(make_Sigma(times,sigma_obs,0.))
  result = optimize(calc_laplace_approx_arg_jitter,sigma_j_min,sigma_j_max,rel_tol=1e-6)
end

#=
periods = [42.0]
X = make_design_matrix_sinusoids(times,periods);
res = calc_opt_sigmaj(times,vels,sigma_obs,X)
dSdj = make_dSigmadsigmaj(times,sigma_obs,res.minimizer)
(compute_log_evidence(res.minimizer)+0.5*(log(2pi)-log(0.5*trace((Sigma \ full(dSdj))^2))))/log(10)
=#

#=
function make_matrix_pd(A::Array{Float64,2}; epsabs::Float64 = 0.0, epsfac::Float64 = 1.0e-6)
  @assert(size(A,1)==size(A,2))
  B = copy(A)
  itt = 1
  while !isposdef(B)
	eigvalB,eigvecB = eig(B)
	neweigval = (epsabs == 0.0) ? epsfac*minimum(eigvalB[eigvalB.>0]) : epsabs
	eigvalB[eigvalB.<0] = neweigval
	B = eigvecB *diagm(eigvalB)*eigvecB'
#	println(itt,": ",B)
	itt +=1
	if itt>size(A,1)
	  error("There's a problem in make_matrix_pd.\n")
  	  break
	end
  end
  return B
end
=#

# Integrate over period
function compute_marginalized_likelihood_fix_periods(times::Vector{T}, vels::Vector{T}, sigma_obs::Vector{T}, periods::Vector{T}; calclogprior_linear::Function = logprior_sinusoids, make_design_matrix::Function = make_design_matrix_circ, X = make_design_matrix(times,periods), Sigma0 = PDMat(make_Sigma(times,sigma_obs,0.0)) ) where T<:Real
   const nparam = size(X,2)

   function optimize_jitter_helper(sigmaj::Real)
     Sigma = Sigma0 + ScalMat(length(times),sigmaj) # PDMat(make_Sigma(times,sigma_obs,sigmaj))
	 FIM = Xt_invA_X(Sigma,X)
	 if !isposdef(FIM)
	   FIM += diagm(1e-3*ones(nparam))
	 end
	 FIM = PDMat(FIM)
     evid = compute_Laplace_Approx(vels,X,Sigma,FIM=FIM,calclogprior=calclogprior_linear)
 	 evid += logprior_jitter(sigmaj)
	 -evid
   end
   result = optimize(optimize_jitter_helper,sigma_j_min,sigma_j_max,rel_tol=1e-6)
   #println(result)
   sigmaj = result.minimizer
   #Sigma = PDMat(make_Sigma(times,sigma_obs,sigmaj))
   Sigma = Sigma0 + ScalMat(length(times),sigmaj) # PDMat(make_Sigma(times,sigma_obs,sigmaj))
   dSdj = make_dSigmadsigmaj(times,sigma_obs,sigmaj)
   FIM = Xt_invA_X(Sigma,X)
	 if !isposdef(FIM)
	   FIM += diagm(1e-3*ones(nparam))
	 end
   FIM = PDMat(FIM)

   logp = compute_Laplace_Approx(vels,X,Sigma,FIM=FIM,calclogprior=calclogprior_linear)
   logp += logprior_jitter(sigmaj)+ logprior_periods(periods)
   logp += 0.5*(log(2pi)-log(0.5*trace((Sigma \ full(dSdj))^2)))
   #println("# logp = ",logp)
   return logp
end

function compute_marginalized_likelihood_fix_periods_sigmaj(times::Vector{T}, vels::Vector{T}, sigma_obs::Vector{T}, periods::Vector{T}, sigmaj::T; calclogprior_linear::Function = logprior_sinusoids, make_design_matrix::Function = make_design_matrix_circ, Sigma=PDMat(make_Sigma(times,sigma_obs,sigmaj)), X= make_design_matrix(times,periods) ) where T<:Real
   #X = make_design_matrix(times,periods)
   const nparam = size(X,2)
   #Sigma = PDMat(make_Sigma(times,sigma_obs,sigmaj))
   #dSdj = make_dSigmadsigmaj(times,sigma_obs,sigmaj)
   FIM = Xt_invA_X(Sigma,X)
	 if !isposdef(FIM)
	   FIM += diagm(1e-3*ones(nparam))
	 end
   FIM = PDMat(FIM)

   logp = compute_Laplace_Approx(vels,X,Sigma,FIM=FIM,calclogprior=calclogprior_linear)
   #logp += logprior_jitter(sigmaj)+ logprior_periods(periods)
   logp += logprior_periods(periods)
   #logp += 0.5*(log(2pi)-log(0.5*trace((Sigma \ full(dSdj))^2)))
   #println("# logp = ",logp)
   return logp
end

function compute_marginalized_likelihood(times::Vector{T}, vels::Vector{T}, sigma_obs::Vector{T}, minperiods::Vector{T}, maxperiods::Vector{T}; calclogprior_linear::Function = logprior_sinusoids, make_design_matrix::Function = make_design_matrix_circ  ) where T<:Real
  @assert length(minperiods)==length(maxperiods)
  function helper(periods::Vector{T})
     logp = compute_marginalized_likelihood_fix_periods(times,vels,sigma_obs,periods,calclogprior_linear=calclogprior_linear,make_design_matrix=make_design_matrix)
	 (logp == -Inf) ? 0. : exp(logp)
  end
  function helper_1d(period::T)
    println("# Evaluating P= ",period)
	flush(STDOUT)
    logp = compute_marginalized_likelihood_fix_periods(times,vels,sigma_obs,[period],calclogprior_linear=calclogprior_linear,make_design_matrix=make_design_matrix)
	println("#      ", logp)
	flush(STDOUT)
	return (logp == -Inf) ? 0. : exp(logp)
  end

  if length(minperiods)!=1
     error("# Sorry, we haven't implemented multiple periods yet")
  end
  result = QuadGK.quadgk(helper_1d,minperiods[1],maxperiods[1],reltol=1e-3)
end

compute_marginalized_likelihood(times,vels,sigma_obs,[11.9],[12.1])./log(10)

#=
compute_marginalized_likelihood(times,vels,sigma_obs,[11.9],[12.1],calclogprior_linear=logprior_epicycles,make_design_matrix=make_design_matrix_epicycle)./log(10)

compute_marginalized_likelihood(times,vels,sigma_obs,[10.],[2400.])./log(10)
=#

function logsumlogs(x::T,y::T) where T<:Real
  (x<=y) ? y + log1p(exp(x-y)) : x + log1p(exp(y-x))
end

#=
function brute_force_over_periods_1d_old(times::Vector{T}, vels::Vector{T}, sigma_obs::Vector{T}, Plo::Real, Phi::Real; samples_per_peak::Real = 4.0, calclogprior_linear::Function = logprior_sinusoids, make_design_matrix::Function = make_design_matrix_circ ) where T<:Real
  min_freq = 1/Phi
  max_freq = 1/Plo
  duration = maximum(times)-minimum(times)
  delta_freq = 1/(samples_per_peak * duration)
  log_delta_freq = log(delta_freq)
  range = min_freq:delta_freq:max_freq
  logintegral = -Inf
  for f in range
    P = 1/f
	logp = compute_marginalized_likelihood_fix_periods(times,vels,sigma_obs,[P], calclogprior_linear=calclogprior_linear, make_design_matrix=make_design_matrix)
	logp -= log(f)
	if logintegral == -Inf
	  logintegral = logp+log_delta_freq
	else
	  logintegral = logsumlogs(logintegral,logp+log_delta_freq)
	end
  end
  logintegral
end

function brute_force_over_periods_1d(times::Vector{T}, vels::Vector{T}, sigma_obs::Vector{T}, Plo::Real, Phi::Real; samples_per_peak::Real = 1.0, calclogprior_linear::Function = logprior_sinusoids, make_design_matrix::Function = make_design_matrix_circ, sigmaj::Real = -1.0 ) where T<:Real
  min_freq = 1/Phi
  max_freq = 1/Plo
  duration = maximum(times)-minimum(times)
  mean_vel = sum(vels./sigma_obs.^2)/sum(sigma_obs.^(-2))
  rms_vel = std((vels.-mean_vel)./sigma_obs)
  num_periods = convert(Int64,ceil(samples_per_peak*2pi*duration*rms_vel*(max_freq-min_freq)))
  println("# num_periods = ",num_periods)
  deltaP_const = 1/(samples_per_peak*2pi*duration*rms_vel)
  P = Plo
  periods = Array(Float64,2*num_periods)
  powers  = Array(Float64,2*num_periods)
  logintegral = -Inf
  i = 0
  while P<=Phi
    deltaP = deltaP_const*P^2
	if i >0 P += deltaP end
    i += 1
	if sigmaj < 0
	   logp = compute_marginalized_likelihood_fix_periods(times,vels,sigma_obs,[P], calclogprior_linear=calclogprior_linear, make_design_matrix=make_design_matrix)
	else
	   logp = compute_marginalized_likelihood_fix_periods_sigmaj(times,vels,sigma_obs,[P],sigmaj, calclogprior_linear=calclogprior_linear, make_design_matrix=make_design_matrix)

	   logp += logprior_jitter(sigmaj)
	end
	logp += log(deltaP/P)
    if i%10 == 0
	  println("#i=",i," P=",P," dP=",deltaP," logp=",logp," logI=",logintegral)
	end
	if logintegral == -Inf
	  logintegral = logp
	else
	  logintegral = logsumlogs(logintegral,logp)
	end
	@assert i<=length(periods)
	periods[i] = P
	powers[i] = logp

  end
  resize!(periods,i)
  resize!(powers,i)
  println("# Last P = ", P)
  return (periods,powers,logintegral)
end

res = brute_force_over_periods_1d(times,vels,sigma_obs,20.0,10000.0,samples_per_peak=1,sigmaj = 1.5)

result = QuadGK.quadgk(x->exp(brute_force_over_periods_1d(times,vels,sigma_obs,20.0,10000.0,samples_per_peak=1,sigmaj=x)),0.0,10.0,reltol=1e-2)

res = brute_force_over_periods_1d(times,vels,sigma_obs,20.0,10000.0,samples_per_peak=1)

=#

using FastGaussQuadrature
num_Sigmas = 40
nodes, weights = gausslegendre( num_Sigmas );
weights *= 0.5;
nodes = 0.5+0.5*nodes;
sigmaj_list = sigma_j_0*(exp.(nodes*log1p(sigma_j_max/sigma_j_0)).-1);
Sigma0 = PDMat(make_Sigma(times,sigma_obs,0.));
SigmaList = map(sigmaj->Sigma0+ScalMat(length(times),sigmaj),sigmaj_list);

marginalized_likelihood_fix_periods_sigmaj = map(i->compute_marginalized_likelihood_fix_periods_sigmaj(times, vels, sigma_obs, [12.1], sigmaj_list[i], Sigma=SigmaList[i]),1:length(sigmaj_list))
int_norm = maximum(marginalized_likelihood_fix_periods_sigmaj);
integral = log(dot( weights, exp.(marginalized_likelihood_fix_periods_sigmaj-int_norm) ))+int_norm
compute_marginalized_likelihood_fix_periods(times, vels, sigma_obs, [12.1] )

function brute_force_over_periods_1d_new(times::Vector{T}, vels::Vector{T}, sigma_obs::Vector{T}, Plo::Real, Phi::Real; samples_per_peak::Real = 1.0, calclogprior_linear::Function = logprior_sinusoids, make_design_matrix::Function = make_design_matrix_circ, sigmaj::Real = -1.0, num_Sigmas::Integer = 40 ) where T<:Real
  min_freq = 1/Phi
  max_freq = 1/Plo
  duration = maximum(times)-minimum(times)
  mean_vel = sum(vels./sigma_obs.^2)/sum(sigma_obs.^(-2))
  rms_vel = std((vels.-mean_vel)./sigma_obs)
  num_periods = convert(Int64,ceil(samples_per_peak*2pi*duration*rms_vel*(max_freq-min_freq)))
  println("# num_periods = ",num_periods)
  periods = zeros(1)
  sigma_eps = 1e-2
  Sigma0 = PDMat(make_Sigma(times,sigma_obs,0.+sigma_eps));
  if sigmaj >= 0.0
     Sigma = Sigma0 + ScalMat(length(times),sigmaj)
  else
     Sigma = Sigma0 + ScalMat(length(times),rms_vel)
  end

  if num_Sigmas > 1
     nodes, weights = gausslegendre( num_Sigmas );
     weights *= 0.5;
     nodes = 0.5+0.5*nodes;
     sigmaj_list = sigma_j_0*(exp.(nodes*log1p(sigma_j_max/sigma_j_0)).-1);
     SigmaList = map(s->Sigma0+ScalMat(length(times),s),sigmaj_list);
  end

  deltaP_const = 1/(samples_per_peak*2pi*duration*rms_vel)
  P = Plo
  period_grid = Array{Float64}(2*num_periods)
  power_grid  = Array{Float64}(2*num_periods)
  logintegral = -Inf
  i = 0
  while P<=Phi
    deltaP = deltaP_const*P^2
	if i >0 P += deltaP end
	periods[1] = P
    i += 1
	X = make_design_matrix(times,periods)
    if num_Sigmas >1
	   marginalized_likelihood_fix_periods_sigmaj = map(k->compute_marginalized_likelihood_fix_periods_sigmaj(times, vels, sigma_obs, periods, sigmaj_list[k], Sigma=SigmaList[k],calclogprior_linear=calclogprior_linear, make_design_matrix=make_design_matrix, X=X),1:length(sigmaj_list))
       int_norm = maximum(marginalized_likelihood_fix_periods_sigmaj);
       logp = log(dot( weights, exp.(marginalized_likelihood_fix_periods_sigmaj-int_norm) ))+int_norm

    else
	   if sigmaj < 0
	      logp = compute_marginalized_likelihood_fix_periods(times,vels,sigma_obs,periods,Sigma0=Sigma0, calclogprior_linear=calclogprior_linear, make_design_matrix=make_design_matrix,X=X)
	   else
	      logp = compute_marginalized_likelihood_fix_periods_sigmaj(times,vels,sigma_obs,periods,sigmaj, Sigma=Sigma, calclogprior_linear=calclogprior_linear, make_design_matrix=make_design_matrix,X=X)
	      logp += logprior_jitter(sigmaj)
	   end

    end
	logp += log(deltaP/P)
    if i%10 == 0
	  println("#i=",i," P=",P," dP=",deltaP," logp=",logp," logI=",logintegral)
	end
	if logintegral == -Inf
	  logintegral = logp
	else
	  logintegral = logsumlogs(logintegral,logp)
	end
	@assert i<=length(period_grid)
	period_grid[i] = P
	power_grid[i] = logp

  end
  resize!(period_grid,i)
  resize!(power_grid,i)
  println("# Last P = ", period_grid[end])
  return (period_grid,power_grid,logintegral)
end

res = brute_force_over_periods_1d_new(times,vels,sigma_obs,20.0,10000.0,samples_per_peak=1,num_Sigmas=1)

function brute_force_over_periods_multiplanet_1d_old(times::Vector{T}, vels::Vector{T}, sigma_obs::Vector{T}, periods_fixed::Vector{T}, Plo::Real, Phi::Real; samples_per_peak::Real = 4.0, calclogprior_linear::Function = logprior_sinusoids, make_design_matrix::Function = make_design_matrix_circ, update_design_matrix::Function = update_design_matrix_circ ) where T<:Real
  periods = zeros(length(periods_fixed)+1)
  periods[1:length(periods_fixed)] = periods_fixed
  periods[end] = (Plo+Phi)/2
  X = make_design_matrix(times,periods)
  min_freq = 1/Phi
  max_freq = 1/Plo
  duration = maximum(times)-minimum(times)
  delta_freq = 1/(samples_per_peak * duration)
  log_delta_freq = log(delta_freq)
  range = min_freq:delta_freq:max_freq
  logintegral = -Inf
  for f in range
    periods[end] =  1/f
	update_design_matrix(X,times,periods,length(periods))
	logp = compute_marginalized_likelihood_fix_periods(times,vels,sigma_obs,periods, calclogprior_linear=calclogprior_linear, X=X)
	logp -= log(f)
	if logintegral == -Inf
	  logintegral = logp+log_delta_freq
	else
	  logintegral = logsumlogs(logintegral,logp+log_delta_freq)
	end
  end
  logintegral
end

function brute_force_over_periods_multiplanet_1d(times::Vector{T}, vels::Vector{T}, sigma_obs::Vector{T}, periods_fixed::Vector{T}, Plo::Real, Phi::Real; samples_per_peak::Real = 2.0, calclogprior_linear::Function = logprior_sinusoids, make_design_matrix::Function = make_design_matrix_circ, update_design_matrix::Function = update_design_matrix_circ, num_Sigmas::Integer = 40 ) where T<:Real
  periods = zeros(length(periods_fixed)+1)
  periods[1:length(periods_fixed)] = periods_fixed

  # Compute how much signal is remaining so can choose appropriate grid density
  X = make_design_matrix(times,periods_fixed)
  sigmaj = calc_opt_sigmaj(times, vels, sigma_obs, X, calclogprior_linear=calclogprior_linear).minimizer
  Sigma0 = PDMat(make_Sigma(times,sigma_obs,0.))
  Sigma = Sigma0 + ScalMat(length(times),sigmaj) # PDMat(make_Sigma(times,sigma_obs,sigmaj))
  param_bf_linalg = calc_best_fit(vels,X,Sigma)
  predict = X*param_bf_linalg

  chisq = invquad(Sigma0,vels.-predict)
  rms_vel = std((vels.-predict)./sigma_obs)
  println("# RMS resid (uncor) = ", rms_vel)
  #rms_vel_eff = sqrt(chisq/length(times))
  #println("# RMS resid (corel) = ", rms_vel_eff)

  periods[end] = (Plo+Phi)/2
  X = make_design_matrix(times,periods)

  if num_Sigmas > 1
     nodes, weights = gausslegendre( num_Sigmas );
     weights *= 0.5;
     nodes = 0.5+0.5*nodes;
     sigmaj_list = sigma_j_0*(exp.(nodes*log1p(sigma_j_max/sigma_j_0)).-1);
     Sigma0 = PDMat(make_Sigma(times,sigma_obs,0.));
     SigmaList = map(sigmaj->Sigma0+ScalMat(length(times),sigmaj),sigmaj_list);
  end

  # Now compute periodogram
  min_freq = 1/Phi
  max_freq = 1/Plo
  duration = maximum(times)-minimum(times)
  num_periods = convert(Int64,ceil(samples_per_peak*2pi*duration*rms_vel*(max_freq-min_freq)))
  println("# num_periods = ",num_periods)
  period_grid = Array{Float64}(2*num_periods)
  power_grid  = Array{Float64}(2*num_periods)
  deltaP_const = 1/(samples_per_peak*2pi*duration*rms_vel)
  P = Plo
  logintegral = -Inf
  i = 0
  while P<=Phi
    deltaP = deltaP_const*P^2
	if i >0 P += deltaP end
    i += 1
    periods[end] =  P
	#println("# size(X) = ",size(X), " len(times)= ", length(times), " len(periods) = ", length(periods))
	update_design_matrix(X,times,periods,length(periods))

    if num_Sigmas == 1
	   #logp = compute_marginalized_likelihood_fix_periods(times,vels,sigma_obs,periods, calclogprior_linear=calclogprior_linear, X=X, Sigma0=Sigma0) # make_design_matrix=make_design_matrix)
	   logp = compute_marginalized_likelihood_fix_periods_sigmaj(times,vels,sigma_obs,periods,sigmaj, calclogprior_linear=calclogprior_linear, X=X, Sigma=Sigma) # make_design_matrix=make_design_matrix)
	else
	   marginalized_likelihood_fix_periods_sigmaj = map(k->compute_marginalized_likelihood_fix_periods_sigmaj(times, vels, sigma_obs, periods, sigmaj_list[k], Sigma=SigmaList[k],calclogprior_linear=calclogprior_linear, make_design_matrix=make_design_matrix, X=X),1:length(sigmaj_list))
       int_norm = maximum(marginalized_likelihood_fix_periods_sigmaj);
       logp = log(dot( weights, exp.(marginalized_likelihood_fix_periods_sigmaj-int_norm) ))+int_norm
	end

	logp += log(deltaP/P)
    if i%100 == 0
	  println("# i=",i," P=",P," dP=",deltaP," logp=",logp," logI=",logintegral)
	end
	if logintegral == -Inf
	  logintegral = logp
	else
	  logintegral = logsumlogs(logintegral,logp)
	end
	period_grid[i] = P
	power_grid[i] = logp

  end
  resize!(period_grid,i)
  resize!(power_grid,i)
  println("# Last P = ", P)
  (period_grid,power_grid,logintegral)
end

# Optimzie periods
function compute_marginalized_likelihood_near_periods(times::Vector{T}, vels::Vector{T}, sigma_obs::Vector{T}, periods::Vector{T}; calclogprior_linear::Function = logprior_sinusoids, make_design_matrix::Function = make_design_matrix_circ ) where T<:Real
   X = make_design_matrix(times,periods)

   function optimize_jitter_helper(sigmaj::Real)
     #println("# making Sigma")
	 Sigma = Sigma0 + ScalMat(length(times),sigmaj) # PDMat(make_Sigma(times,sigma_obs,sigmaj))
	 #println("# making FIM")
     FIM = PDMat(Xt_invA_X(Sigma,X))
     #println("# approximating integral")
	 evid = compute_Laplace_Approx(vels,X,Sigma,FIM=FIM,calclogprior=calclogprior_linear)
 	 evid += logprior_jitter(sigmaj)
	 #println("# ",evid/log(10))
     -evid
   end
   Sigma0 = PDMat(make_Sigma(times,sigma_obs,0.))
   result = optimize(optimize_jitter_helper,sigma_j_min,sigma_j_max,rel_tol=1e-6)
   #println("# Optimize jitter: ", result)
   sigmaj = result.minimizer
   Sigma = Sigma0 + ScalMat(length(times),sigmaj) # PDMat(make_Sigma(times,sigma_obs,sigmaj))

   function optimize_period_helper(P::T)
     X = make_design_matrix(times,[P])
     FIM = PDMat(Xt_invA_X(Sigma,X))
     evid = compute_Laplace_Approx(vels,X,Sigma,FIM=FIM,calclogprior=calclogprior_linear)
 	 evid += logprior_jitter(sigmaj)
     -evid
   end
   function optimize_period_helper(periods::Vector{T})
     X = make_design_matrix(times,periods)
     FIM = PDMat(Xt_invA_X(Sigma,X))
     evid = compute_Laplace_Approx(vels,X,Sigma,FIM=FIM,calclogprior=calclogprior_linear)
 	 evid += logprior_jitter(sigmaj)
     -evid
   end

   if(length(periods)==1)
      result2 = optimize(optimize_period_helper,periods[1]*0.98,periods[1]*1.02,rel_tol=1e-6)
	  periods = [result2.minimizer]
   else
	  result2 = optimize(optimize_period_helper,periods,f_tol=1e-3) # abs(result.minimum))
	  periods = result2.minimizer
   end
   #println("# Optimize period: ", result2)
   best_period = copy(periods)
   #=
   X = make_design_matrix(times,periods)
   result = optimize(optimize_jitter_helper,sigma_j_min,sigma_j_max,rel_tol=1e-6)
   #println("# Optimize jitter: ", result)
   sigmaj = result.minimizer
   #Sigma = PDMat(make_Sigma(times,sigma_obs,sigmaj))
   Sigma = Sigma0 + ScalMat(length(times),sigmaj) # PDMat(make_Sigma(times,sigma_obs,sigmaj))
   result2 = optimize(optimize_period_helper,periods[1]*0.98,periods[1]*1.02,rel_tol=1e-6)
   println("# Optimize period: ", result2)
   best_period = result2.minimizer
   =#

   X = make_design_matrix(times,periods)
   FIM = PDMat(Xt_invA_X(Sigma,X))

   (param_bf,sigma_param,logp) = compute_Laplace_Approx_param_and_integral(vels,X,Sigma,FIM=FIM,calclogprior=calclogprior_linear)
   logp += logprior_jitter(sigmaj)+ logprior_periods(periods)

   #println("# marginal at best sigma_j & P: ",logp/log(10))

   # Laplace approx to integrate over jitter parameter
   dSdj = make_dSigmadsigmaj(times,sigma_obs,sigmaj)
   log_int_over_sigma_j = 0.5*(log(2pi)-log(0.5*trace((Sigma \ full(dSdj))^2)))
   logp += log_int_over_sigma_j

   #println("# marginal over sigma_j at best P: ",logp/log(10))

   # Laplace approx to integrate over period (one peak only)
   sigma_P = zeros(length(periods))
   log_int_over_periods = 0
   for i in 1:length(best_period)
     periods = copy(best_period)
     eps = 1e-5*periods[i]
	 periods[i] += eps

     X = make_design_matrix(times,periods)
     logp_hi = compute_Laplace_Approx(vels,X,Sigma,calclogprior=calclogprior_linear)
     logp_hi += logprior_jitter(sigmaj)+ logprior_periods(periods)
     logp_hi += log_int_over_sigma_j  # Approximates this term as const
	 periods[i] -= 2*eps
     X = make_design_matrix(times,periods)
     logp_lo = compute_Laplace_Approx(vels,X,Sigma,calclogprior=calclogprior_linear)
     logp_lo += logprior_jitter(sigmaj)+ logprior_periods(periods)
     logp_lo += log_int_over_sigma_j # Approximates this term as const
     d2logpdp2 = (logp_hi-2*logp+logp_lo)/eps^2
     sigma_P[i] = sqrt(-1/d2logpdp2)
     println("# d^2 logp/dP_",i,"^2 = ",d2logpdp2/log(10))
     log_int_over_periods += 0.5*(log(2pi)-log(abs(d2logpdp2)))
   end
   logp += log_int_over_periods # Approx as diagonal
   println("# marginal over P at ",best_period," and sigma_j at ",sigmaj," : ",logp/log(10))

   return (param_bf,sigma_param,logp,best_period,sigma_P)
end

res0 = compute_marginalized_likelihood_fix_periods(times,vels,sigma_obs,Float64[])/log(10)
res1  = compute_marginalized_likelihood_near_periods(times,vels,sigma_obs,[42.0])[3] /log(10)
res2  = compute_marginalized_likelihood_near_periods(times,vels,sigma_obs,[42.0,12.1])[3]/log(10)

res1e = compute_marginalized_likelihood_near_periods(times,vels,sigma_obs,[42.0],calclogprior_linear=logprior_epicycles,make_design_matrix=make_design_matrix_epicycle)

res2e = compute_marginalized_likelihood_near_periods(times,vels,sigma_obs,[42.0,12.0],calclogprior_linear=logprior_epicycles,make_design_matrix=make_design_matrix_epicycle)

res2b =  brute_force_over_periods_multiplanet_1d(times,vels,sigma_obs,[42.0777],10.0,2400.0,samples_per_peak=8)

function compute_marginalized_likelihood_cuba(times::Vector{T}, vels::Vector{T}, sigma_obs::Vector{T}, minperiods::Vector{T}, maxperiods::Vector{T}; calclogprior_linear::Function = logprior_sinusoids, make_design_matrix::Function = make_design_matrix_circ  ) where T<:Real
  @assert length(minperiods)==length(maxperiods)
  log_normalization = 0.0
  function helper_1d(period::T) # Note, this exp's unlike multi-d helper
    #println("# Evaluating P= ",period)
	#flush(STDOUT)
    logp = compute_marginalized_likelihood_fix_periods(times,vels,sigma_obs,[period],calclogprior_linear=calclogprior_linear,make_design_matrix=make_design_matrix)
	#println("#      ", logp)
	#flush(STDOUT)
	return (logp == -Inf) ? 0. : exp(logp-log_normalization)
  end
  function helper(periods::Vector{T})
    #println("# Evaluating P= ",periods)
	#flush(STDOUT)
     logp = compute_marginalized_likelihood_fix_periods(times,vels,sigma_obs,periods,calclogprior_linear=calclogprior_linear,make_design_matrix=make_design_matrix)
	#println("#      ", logp)
	#flush(STDOUT)
	return logp
	# (logp == -Inf) ? 0. : exp(logp-log_normalization)
  end
  function integrand_cuba(x::Vector, output::Vector)
    @assert length(output) == 1
	#println("# x=",x)
    periods = (maxperiods.-minperiods).*x.+minperiods
    output[1] = exp(helper(periods)-log_normalization)
	println("# P= ",periods," logp= ",output[1])
  end

  log_normalization = helper(0.5*(minperiods.+maxperiods))
  if length(minperiods)==1
     result = QuadGK.quadgk(helper_1d,minperiods[1],maxperiods[1],reltol=1e-3)
	 result = log(result[1])
     #result = divonne(integrand_cuba,1,1,reltol=1e-3,maxevals=100)*prod(maxperiods.-minperiods) # Problem?
  else

     result = divonne(integrand_cuba,length(minperiods),1,reltol=1e-3,maxevals=1000)
	 result = log(result.integral[1]) + log(prod(maxperiods.-minperiods))
  end
  result+log_normalization
end

res = compute_marginalized_likelihood_cuba(times,vels,sigma_obs,[41.9,11.9],[42.3,12.2])

#=
@time vegas(integrand_cuba,2,1,reltol=1e-3,maxevals=1e5)
@time suave(integrand_cuba,2,1,reltol=1e-3,maxevals=1e5)
@time divonne(integrand_cuba,2,1,reltol=1e-3,maxevals=1e5)
@time cuhre(integrand_cuba,2,1,reltol=1e-3,minevals=1e3,maxevals=1e5)
=#

#=
# Brute force 2-d integration
using Cubature

function log_target(x::Vector)
  offset = x[1]
  sigma_j = x[2]
  logprior = logprior_offset([offset]) + logprior_jitter(sigma_j)
  Sigma = make_Sigma(times,sigma_obs,sigma_j)
  delta = vels.-offset
  chisq = invquad(Sigma,delta)
  loglikelihood = 0.5*(-chisq -logdet(Sigma)-nobs*log(2pi))
  logprior + loglikelihood
end

function integrand(x::Vector; log_normalization::Real=0.0)
  exp(log_target(x)-log_normalization)
end
@time hcubature(integrand, [-Cmax/100,sigma_j_min],[Cmax/100,sigma_j_max],reltol=1e-2)
#@time pcubature(integrand, [-Cmax/100,sigma_j_min],[Cmax/100,sigma_j_max],reltol=1e-2)
=#
#=
using Cuba
function integrand_cuba(x::Vector, output::Vector; log_normalization::Real=0.0)
  xx = zeros(2)
  xx[1] = x[1]*2*Cmax/100-Cmax/100
  xx[2] = x[2]*(sigma_j_max-sigma_j_min)+sigma_j_min
  output[1] = exp(log_target(xx)-log_normalization)*(2*Cmax/100)*(sigma_j_max-sigma_j_min)
end
@time vegas(integrand_cuba,2,1,reltol=1e-3,maxevals=1e5)
@time suave(integrand_cuba,2,1,reltol=1e-3,maxevals=1e5)
@time divonne(integrand_cuba,2,1,reltol=1e-3,maxevals=1e5)
@time cuhre(integrand_cuba,2,1,reltol=1e-3,minevals=1e3,maxevals=1e5)
=#
#=
# Brute Force:  Sequence of Nested 1-d integrations
function compute_log_target_fixed_offset(offset::Real)
  delta = vels.-offset
  function integrand(sigma_j::Real)
    logprior = logprior_offset([offset]) + logprior_jitter(sigma_j)
    Sigma = Sigma0 + ScalMat(length(times),sigmaj) # PDMat(make_Sigma(times,sigma_obs,sigmaj))
    chisq = invquad(Sigma,delta)
    loglikelihood = 0.5*(-chisq -logdet(Sigma)-nobs*log(2pi))
    logtarget = logprior + loglikelihood
	exp(logtarget) # +log(1e210))
  end
  Sigma0 = make_Sigma(times,sigma_obs,0.)
  result = QuadGK.quadgk(integrand,sigma_j_min,sigma_j_0,sigma_j_max,reltol=1e-2)[1]
end
@time result = QuadGK.quadgk(x->begin p=compute_log_target_fixed_offset(x); println(x," ",p); p end,-Cmax/100,Cmax/100,reltol=1e-2)./log(10)

# Brute Force: 2-d integration
function log_target(x::Vector)
  offset = x[1]
  sigma_j = x[2]
  logprior = logprior_offset([offset]) + logprior_jitter(sigma_j)
  Sigma = make_Sigma(times,sigma_obs,sigma_j)
  delta = vels.-offset
  chisq = invquad(Sigma,delta)
  loglikelihood = 0.5*(-chisq -logdet(Sigma)-nobs*log(2pi))
  logprior + loglikelihood
end
=#

#=
# Fit one planet model
periods = [42.0]
nparam = 1+2*length(periods)
freqs = 2pi./periods
X = [sin(freqs[1]*times) cos(freqs[1]*times) ones(times)];
Sigma = PDMat(make_Sigma(times,sigma_obs,sigma_j_0))
FIM = PDMat(Xt_invA_X(Sigma,X))
(param_bf,sigma_param,evid) = compute_Laplace_Approx_param_and_integral(vels,X,Sigma,FIM=FIM,calclogprior=logprior_sinusoids)

function calc_opt_sigmaj(data::Vector{T}, model::Array{T,2}) where T<:Real
  function calc_laplace_approx_arg_jitter(sigmaj::Real)
	Sigma = Sigma0 + ScalMat(length(times),sigmaj) # PDMat(make_Sigma(times,sigma_obs,sigmaj))
    FIM = PDMat(Xt_invA_X(Sigma,X))
    evid = compute_Laplace_Approx(vels,X,Sigma,FIM=FIM,calclogprior=logprior_sinusoids)
    evid += logprior_jitter(sigmaj)
	-evid
  end
  Sigma0 = make_Sigma(times,sigma_obs,0.)
  result = optimize(calc_laplace_approx_arg_jitter,sigma_j_min,sigma_j_max,rel_tol=1e-4)
end
res = calc_opt_sigmaj(vels,X)
Sigma = PDMat(make_Sigma(times,sigma_obs,res.minimizer))
dSigmadsigmaj = make_dSigmadsigmaj(times,sigma_obs,res.minimizer)
FIM = PDMat(Xt_invA_X(Sigma,X))
(param_bf,sigma_param,evid) = compute_Laplace_Approx_param_and_integral(vels,X,Sigma,FIM=FIM,calclogprior=logprior_sinusoids)
evid += 0.5*(log(2pi)-log(0.5*trace((Sigma \ full(dSdj))^2)))

Ks = generate_K_samples(param_bf,sigma_param,1000);
m = mean(Ks,1);
sig = std(Ks,1,mean=m);
m-sig,m,m+sig

function profile_sigmaj(data::Vector{T}, model::Array{T,2}) where T<:Real
  function calc_laplace_approx_arg_jitter(sigmaj::Real) #param::Vector{T})
    #sigmaj = param[1]
	Sigma = PDMat(make_Sigma(times,sigma_obs,sigmaj))
    FIM = PDMat(Xt_invA_X(Sigma,X))
    evid = compute_Laplace_Approx(vels,X,Sigma,FIM=FIM,calclogprior=logprior_sinusoids)
	evid += logprior_jitter(sigmaj)
    evid
  end
  #result = optimize(calc_laplace_approx_arg_jitter,[sigma_j_0])
  result = optimize(calc_laplace_approx_arg_jitter,sigma_j_min,sigma_j_max,rel_tol=1e-6)
end
res = calc_opt_sigmaj(vels,X)

# Fit two planet model
periods = [42.0,12.0]
X = make_design_matrix(times,periods)
Sigma = PDMat(make_Sigma(times,sigma_obs,0.68442348767604010984))
FIM = PDMat(Xt_invA_X(Sigma,X))
(param_bf,sigma_param,evid) = compute_Laplace_Approx_param_and_integral(vels,X,Sigma,FIM=FIM,calclogprior=logprior_sinusoids)

Ks = generate_K_samples(param_bf,sigma_param,1000)
m = mean(Ks,1)
sig = std(Ks,1,mean=m)
m-sig,m+sig

=#


using Revise, BenchmarkTools
using DifferentialEquations, Optim, DiffEqFlux, DiffEqSensitivity, Plots
using Zygote, ForwardDiff, ReverseDiff
using LinearAlgebra, RecursiveArrayTools, Measurements, Statistics
using DataFrames
using BlackBoxOptim
# -------------------------------------------------------------------------

function run_model(model,p,data,densities,constants,t) # version of run function with multiple models
    # f = (du,u,p,t) -> model(du,u,10 .^ p,t,densities,constants)
    f = (du,u,p,t) -> model(du,u,p,t,densities,constants)
    tmp_prob = ODEProblem(f,u₀,tspan,p)
    tmp_sol = solve(tmp_prob,AutoVern7(KenCarp4()),saveat=t, abstol=1e-8,reltol=1e-8,verbose=false)
    # tmp_sol = solve(tmp_prob,Vern7(),saveat=t, abstol=1e-8,reltol=1e-8,verbose=false)
    tmp_sol
end

function loss(p,data,model,dens,cons,t,unique_t)
    Δy_all = [] # keep track of errors for all curves (change this to array of arrays!)
    # Δy_all = Zygote.Buffer([[0.,0.],[0.,0.]],length(data))
    for i in 1:length(data)
      # run model with densities and parameter/dens/cons combinations
      tmp_sol = run_model(model,p,data[i],dens[i],cons[i],unique_t[i])

      if tmp_sol.retcode != :Success #if there is an error return infitite loss
        Δy_sol = fill(Inf, length(data[i]))
      else
        Δy_sol = find_diffs(tmp_sol,t[i],unique_t[i],data[i])
      end
      Δy_all = vcat(Δy_all,Δy_sol)
      # Δy_all[i] = Δy_sol
    end
    # Δy = copy(Δy_all
    # Losses = [sum(Δ .^ 2) for Δ in Δy]
    # Losses = [sum( (Δy[i] .^ 2) .* t[i] .^2 ) for i in 1:length(Δy)]
    t_all = reduce(vcat,t)

    # Losses = Δy_all .^2 .* t_all .^2
    Losses = Δy_all .^2 .* t_all
    L = sum(Losses) # calculate error across all differences from all curves

    return L
end

function find_diffs(tmp_sol,t,unique_t,data) # solutions use unique time points,
# need to evaluate the solution at each experimental time point as well as
# calculating the difference between fit and experiment

    Σ_sol = sum(Array(tmp_sol),dims=1)
    # Σ_sol = sum(tmp_sol,dims=1)
    # Δy = Zygote.Buffer(t,length(t))
    Δy = []
    # Δy = zeros(length(t))
    for (i, tᵢ) in enumerate(t)
        ind = isequal.(unique_t,tᵢ)
        yᵢ_exp = Σ_sol[ind][1]
        yᵢ_obs = data[i]
        Δyᵢ = yᵢ_obs .- yᵢ_exp
        # Δy[i] = Δyᵢ
        # push!(Δy,Δyᵢ)
        Δy=vcat(Δy,Δyᵢ)
    end
    return copy(Δy) #Δy[1:end] # copy(Δy)
end

# function callback(p,l,tmp_sol,data,model,dens,cons,t)
function callback(p,l,data,model,dens,cons,t,unique_t)
  # @show l, p
  @show l
  # Σ_sol = run_model(model,p,data,dens,cons,unique_t)

  # fig = plot(unique_t,Σ_sol')
  # scatter!(fig,t,data,legend=:bottomright)
  # display(fig)

  false
end

function scan_param_space(p,data,model,dens,cons,t,unique_t,lb,ub)
      #create rand mesh
      n = 5000
      Δbounds = ub - lb
      # ps = rand(n,size(p,1)) .* 10 .- 8
      ps = rand(n,size(p,1)) .* Δbounds  .+ lb
      # ps = rand(n,size(p,1)) .* 100
      errors = zeros(n)
      for i in 1:n
          # error, a = loss(ps[i,:],data,model,dens,cons,t)
          error = loss(ps[i,:],data,model,dens,cons,t,unique_t)
          errors[i] = error
          # @show error
          if mod(i,100) == 0
            @show i, error
          end
      end
      min_ind = argmin(errors)
      min_error = minimum(errors)
      @show findmin(errors)
      p_best = ps[min_ind,:]
      @show p_best
      if size(p,1) == 2
          fig = scatter(ps[:,1],ps[:,2],log10.(errors),
          zcolor = log10.(errors),camera=[0,90])
          scatter!([p_best[1]],[p_best[2]],[log10(min_error)],markercolor=:green,markersize=5)
          display(fig)
      end
      return p_best
end

function optimization(model,data,dens,cons,options,t)
# fit function - given model, data, densities, options get best parameters
  unique_t = [unique(times) for times in t]
  loss_in = (p) -> loss(p,data,model,dens,cons,t,unique_t)
  callback_in = (p,l) -> callback(p,l,data,model,dens,cons,t,unique_t)

  # @unpack lb,ub = options
  lb,ub = options

  # rates = [0.41751687626687767, -0.23866855433293688, 0.06689410353314687, -0.5510289313898404, -0.09606981254493874, -0.1155499030032302, -3.9999999999999964, -1.7908513448187287]
  # hes_zyg = Zygote.hessian(loss_in,rates)
  @show loss_in(pinit)
  # start = scan_param_space(pinit,data,model,dens,cons,t,unique_t,lb,ub) # !! change behavior so that it doesnt need pinit
  # start = DiffEqFlux.sciml_train(loss_in,start,ADAM(0.001),cb=callback_in,maxiters=50)
  # start = start.minimizer
  # @show start
  # res = optimize(loss_in,start,BFGS(),Optim.Options(f_tol = 1e-8))
  # res = optimize(loss_in,start,NelderMead(),Optim.Options(f_tol = 1e-8))
  res_bbo = bboptimize(loss_in; SearchRange = (lb, ub),
                                NumDimensions = length(pinit),
                                Method = :adaptive_de_rand_1_bin_radiuslimited,
                                # Method = :separable_nes,
                                NThreads=Threads.nthreads()-1,
                                MaxSteps = 10000)#,

                                # # MaxFuncEvals = 10000)#,
                                # TargetFitness=0.05)#,
                                # NThreads=Threads.nthreads()-1)
  p₀_bbo = best_candidate(res_bbo)
  @show p₀_bbo
  # p₀_bbo = pinit
  # bbo_minimizer = [0.829897, -3.39453, -3.83109, 0.215222, 0.848485, -3.08371, 0.606198, -0.906295]
  # res = optimize(loss_in,lb,ub,p₀_bbo, Fminbox(NelderMead()),
  res = optimize(loss_in,lb,ub,p₀_bbo, Fminbox(BFGS()), autodiff = :forward,
                                Optim.Options(show_trace=true,
                                              f_tol = 1e-5,
                                              outer_iterations = 10))

  # res = optimize(loss_in,p₀_bbo, LBFGS(),
  #               Optim.Options(show_trace=true,
  #                             f_tol = 1e-4))
  # res = optimize(loss_in,lb,ub,bbo_minimizer, Fminbox(NelderMead()) , Optim.Options(show_trace=true))

  # SCIML_TRAIN otpimizers
  # res = DiffEqFlux.sciml_train(loss_in,p₀_bbo,ADAM(),cb=callback_in,maxiters=100)
  # res = DiffEqFlux.sciml_train(loss_in,start,BBO(),cb=callback_in,maxiters=100)

  rates = res.minimizer
  l = res.minimum

  # bootstrap_rates = bootstrapping(rates,data,model,dens,cons,t)
  # bootstrap_rates = []

  @show model
  # @show start
  @show l
  @show rates
  # @show best_fitness(res_bbo), best_candidate(res_bbo)
  # hes_for = ReverseDiff.hessian(loss_in,rates)
  # # hes_for = ForwardDiff.hessian(loss_in,rates)
  # hes_zyg = Zygote.hessian(loss_in,rates)
  # @show hes_for ≈ hes_zyg
  hes_zyg = []
  return rates , hes_zyg, l
end

# function bootstrapping(rates,data,model,dens,cons,t)
#   #create sample
#   data = n
#   B_size = 10
#   # rand_mat = rand(1:length(data),(length(data),B_size))
#
#   rand_mat = [rand(1:length(data),length(data)) for i in 1:B_size]
#
#   bootstrap_rates = zeros(size(rates))
#   # bootstrap_rates = []
#   for i in 1:B_size
#     resampled_data = data[rand_mat[i]]
#     resampled_t = t[rand_mat[i]]
#     sort_ind = sortperm(resampled_t)
#     resampled_data = resampled_data[sort_ind]
#     resampled_t = resampled_t[sort_ind]
#     unique_t = unique(resampled_t)
#
#     # loss_in = (p) -> loss_2(p,resampled_data,model,dens,cons,resampled_t,unique_t)
#     # callback_in = (p,l) -> callback(p,l,resampled_data,model,dens,cons,resampled_t,unique_t)
#     loss_in = (p) -> loss(p,data,model,dens,cons,t,unique(t))
#     callback_in = (p,l) -> callback(p,l,data,model,dens,cons,t,unique(t))
#
#     # @show loss_in(rates)
#     res = optimize(loss_in,-9, 2, rates,Fminbox(NelderMead()))
#     # @time res = DiffEqFlux.sciml_train(loss_in,rates,BFGS(),
#     #                                 lower_bounds = -9,
#     #                                 upper_bounds = 2,
#     #                                 cb = callback_in,
#     #                                 maxiters=50,
#     #                                 f_tol = 1e-4)
#
#     # @show res.minimizer
#     bootstrap_rates= hcat(bootstrap_rates,res.minimizer)
#     # push!(bootstrap_rates,res.minimizer)
#     # Σ_sol = run_model_2(model,p,resampled_data,dens,cons,unique_t)
#   end
#
#   return bootstrap_rates[:,2:end]
# end
#


module LineSpectraEstimation

#TODO change benchmark to JuMP?

using BenchmarkTools
using StructuredOptimization
using Convex
using MathProgBase
using ECOS
using SCS
using AbstractOperators
using PyPlot

function set_up(;opt="")

    srand(39)

    fs = 16e3
    Nt = 2^7 #time samples
    y0 = zeros(Nt)
    SNR = 20

    s = 4 #super-resolution factor
    f_s  = linspace(0,  fs,       s*Nt+1)[1:end-1]     # super resolution frequency axis
    f_s2 = linspace(0,fs/2,div(s*Nt,2)+1)              # super resolution frequency axis (up to Nyquist)
    t  = 0:1/fs:(Nt-1)/fs                              # time axis
    f  = linspace(0,fs,Nt+1)[1:end-1]                  # frequency axis
    f2 = linspace(0,fs/2,div(Nt,2)+1)                  # frequency axis
    K = 14                                             # number of sinusoids
    Nf = div(s*Nt,4)+1
    fk = f_s2[randperm(Nf)[1:K]]                   # sinusoids frequencies
    ak = 0.1*randn(K)+0.7                              # amplitude

    for i in eachindex(fk) y0 .+= ak[i].*sin.(2*π*fk[i].*t) end
    y = y0.+10^(-SNR/10)*sqrt(var(y0))*randn(length(y0))

	xzp = rfft([y;zeros((s-1)*length(y))])
	IDFTm = [exp(im*2*pi*k*n/(s*Nt))  for k =0:s*Nt-1, n=0:s*Nt-1] #Inverse Fourier Matrix
	S = [speye(Nt) spzeros(Nt,(s-1)*Nt)] # selection matrix
	Fm = full(1/(s*Nt).*S*IDFTm)
    alpha = 0.001
	lambda_max_m = norm(Fm'*y, Inf)
	lambda_m = alpha*lambda_max_m

    F = 1/(s*Nt)*DFT(s*Nt)'[1:Nt] # Abstract Operator
	lambda_max = norm(F'*y, Inf)
	lambda = alpha*lambda_max

	Fc = [real(Fm) -imag(Fm); imag(Fm) real(Fm)]
	# needed in cvxpy (currently does not support complex variables)

	if opt == "MatrixFree"
		setup = K,  F, Fc, lambda, lambda_m
	else
		setup = K, Fm, Fc, lambda, lambda_m
	end
	return setup, t, f, f_s2, fk, ak, s, Nt, fs, xzp, y

end

function run_demo()

    tol = 1e-8
    verb = 1
    solver = PANOC

	setup, t, f, fs, fk, ak, s, Nt, Fs, xzp, y = set_up(opt = "MatrixFree")
	slv = solver(tol = tol, verbose = verb)
	x = init_variable(s*Nt,slv)

	println("Solving LASSO (Abstract Operator)")
	@time solve_problem!(slv, x, y, setup...)
	x1 = copy(~x)

	println("Refine solution by solving non-convex problem (Abstract Operator)")
	@time solve_problem_ncvx!(slv, x, y, setup...)
	x0 = copy(~x)

	slv = solver(tol = tol, verbose = verb)
	setup, t, f, fs, fk, ak, s, Nt, Fs, xzp, y = set_up()
	x0m = init_variable(s*Nt,slv)

	println("Solving LASSO (Matrix Operator)")
	@time solve_problem!(slv, x0m, y, setup...)

	println("Refine solution by solving non-convex problem (Matrix Operator)")
	@time solve_problem_ncvx!(slv, x0m, y, setup...)

	x1 = x1[1:div(s*Nt,2)+1]
	x0 = x0[1:div(s*Nt,2)+1]
	return t, f, fs, fk, ak, s, Nt, Fs, xzp, y, x1, x0
end

function run_demo_Convex()
	slv = ECOSSolver()
	setup, t, f, fs, fk, ak, s, Nt, Fs, xzp, y = set_up()
	x0m = init_variable(s*Nt,slv)

	println("Solving LASSO with Convex.jl (Matrix Operator)")
	@time x0m, = solve_problem!(slv, x0m, y, setup...)

	x1 = x0m.value
	x1 = x1[1:div(s*Nt,2)+1]
	return t, f, fs, fk, ak, s, Nt, Fs, xzp, y, x1, zeros(x1)
end

init_variable(N,slv::S) where {S <: StructuredOptimization.ForwardBackwardSolver} = StructuredOptimization.Variable(Complex{Float64}, N)
init_variable(N,slv::S) where {S <: MathProgBase.SolverInterface.AbstractMathProgSolver} = Convex.ComplexVariable(N)
init_variable(N,slv::S) where {S <: AbstractString} = cvx.Variable(N), cvx.Variable(N)

#StructuredOptimization Matrix Free
function solve_problem!(slv::S, x0, y, K, F::A, Fc, lambda, lambda_m) where {S <: StructuredOptimization.ForwardBackwardSolver, A <: AbstractOperator}
    it, = @minimize ls(F*x0-y)+lambda*norm(x0,1) with slv
	return x0, it
end

function solve_problem_ncvx!(slv::S, x0, y, K, F::A, Fc, lambda, lambda_m) where {S <: StructuredOptimization.ForwardBackwardSolver, A <: AbstractOperator}
	it, = @minimize ls(F*x0-y) st norm(x0,0) <= 2*K with slv
	return x0, it
end

#StructuredOptimization non-Matrix Free
function solve_problem!(slv::S, x0, y, K, F::A, Fc, lambda, lambda_m) where {S <: StructuredOptimization.ForwardBackwardSolver, A <: AbstractMatrix}
    it, = @minimize ls(F*x0-complex(y))+lambda_m*norm(x0,1) with slv
	return x0, it
end

function solve_problem_ncvx!(slv::S, x0, y, K, F::A, Fc, lambda, lambda_m) where {S <: StructuredOptimization.ForwardBackwardSolver, A <: AbstractMatrix}
    it, = @minimize ls(F*x0-complex(y)) st norm(x0,0) <= 2*K with slv
	return x0, it
end

#Convex
function solve_problem!(slv::S, x0, y, K, F, Fc, lambda, lambda_m) where {S <: MathProgBase.SolverInterface.AbstractMathProgSolver}
	problem = minimize(0.5*norm(F*x0-y,2)^2+lambda_m*norm(x0,1))
	Convex.solve!(problem,slv)
	return x0, 0
end

function benchmark(;verb = 0, samples = 5, seconds = 100, tol = 1e-8, maxit = 20000)

	suite = BenchmarkGroup()

	solvers = ["PANOC", "ZeroFPR", "FPG", "PG",]
	slv_opt = ["(verbose = $verb, tol = $tol, maxit = $maxit)",
               "(verbose = $verb, tol = $tol, maxit = $maxit)",
               "(verbose = $verb, tol = $tol, maxit = $maxit)",
               "(verbose = $verb, tol = $tol, maxit = $maxit)",]

	its = Dict([(sol,0.) for sol in solvers])
	for i in eachindex(solvers)

		setup, t, f, fs, fk, ak, s, Nt, Fs, xzp, y = set_up()
		solver = eval(parse(solvers[i]*slv_opt[i]))
		x0 = init_variable(Nt*s,solver)

		suite[solvers[i]] =
		@benchmarkable((x0,it) = solve_problem!(solver, x0, y, setup...),
			       setup = (
					it = 0;
					x0 = deepcopy($x0);
					setup = deepcopy($setup);
					y = $y;
					solver = deepcopy($solver) ),
			       teardown = (
					  $its[$solvers[$i]] = it;
					  ),
			       evals = 1, samples = samples, seconds = seconds)
	end

	results = run(suite, verbose = (verb != 0))
	println("LineSpectraEstimation its")
	println(its)
	return results
end

function benchmarkMatrixFree(;verb = 0, samples = 5, seconds = 100, tol = 1e-8, maxit = 20000)

	suite = BenchmarkGroup()

	solvers = ["PANOC", "ZeroFPR", "FPG", "PG"]
	slv_opt = ["(verbose = $verb, tol = $tol, maxit = $maxit)",
               "(verbose = $verb, tol = $tol, maxit = $maxit)",
               "(verbose = $verb, tol = $tol, maxit = $maxit)",
               "(verbose = $verb, tol = $tol, maxit = $maxit)"]

	its = Dict([(sol,0.) for sol in solvers])
	for i in eachindex(solvers)

		setup, t, f, fs, fk, ak, s, Nt, Fs, xzp, y = set_up(opt = "MatrixFree")
		solver = eval(parse(solvers[i]*slv_opt[i]))
		x0 = init_variable(Nt*s,solver)

		suite[solvers[i]] =
		@benchmarkable((x0,it) = solve_problem!(solver, x0, y, setup...),
			       setup = (
					it = 0;
					x0 = deepcopy($x0);
					setup = deepcopy($setup);
					y = $y;
					solver = deepcopy($solver) ),
			       teardown = (
					  $its[$solvers[$i]] = it;
					  ),
			       evals = 1, samples = samples, seconds = seconds)
	end

	results = run(suite, verbose = (verb != 0))
	println("LineSpectraEstimation (MatrixFree) its")
	println(its)
	return results
end

function show_results(t, f, fs, fk, ak, s, Nt, Fs, xzp, y, x1, x0)

	figure()
	plot(fs,abs.(xzp./Nt ), label = "dft zero pad.")
	plot(f,        abs.(fft(y)./Nt)       , label = "dft")
	plot(fk,       abs.(ak)/2     , "r*", label = "true amp.")
    plot(fs,abs.(x1./(s*Nt)), "k*", label = "LASSO")
    plot(fs,abs.(x0./(s*Nt)), "go", label = "IndBallL0")
	xlim([0;Fs/4])
	legend()

end

end


using Flux: Dense, BatchNorm, Chain, relu, sigmoid

# A linear layer followed by batch normalization.
function DenseBNReLU(in:: Integer, out:: Integer; bias = false)
    dense = Dense(in, out; bias=bias)
    batchnorm = BatchNorm(out, relu)
    Chain(dense, batchnorm)
end

"""
    MLPEncoder(in, latent, hidden; bias)

A MLP encoder with variable number of hidden layers.
"""
function MLPEncoder(in:: Integer, latent:: Integer, hidden:: AbstractVector{T}; bias:: Bool = false) where {T <: Integer}
    hidden = append!([in], hidden)
    layers = [DenseBNReLU(hidden[i], hidden[i + 1], bias=bias) for i=1:length(hidden) - 1]
    Chain(layers..., Dense(hidden[end], latent; bias=bias))
end

"""
    MLPDecoder(in, latent, hidden; bias)

A MLP decoder with variable number of hidden layers.
"""
function MLPDecoder(in:: Integer, latent:: Integer, hidden:: AbstractVector{T}; bias:: Bool = false) where {T <: Integer}
    hidden = append!([latent], hidden)
    layers = [DenseBNReLU(hidden[i], hidden[i + 1], bias=bias) for i=1:length(hidden) - 1]
    Chain(layers..., Dense(hidden[end], in, sigmoid; bias=bias))
end

"""
    MLPAutoEncoder(in, latent, hidden; bias)

A MLP auto-encoder with variable number of hidden layers.
"""
function MLPAutoEncoder(in:: Integer, latent:: Integer, hidden:: AbstractVector{T}; bias:: Bool = false) where {T <: Integer}
    encoder = MLPEncoder(in, latent, hidden; bias)
    decoder = MLPDecoder(in, latent, reverse(hidden); bias)
    encoder, decoder
end


using Test
using QuantumOptics
using Random, SparseArrays, LinearAlgebra

@testset "manybody" begin

Random.seed!(0)

D(op1::AbstractOperator, op2::AbstractOperator) = abs(tracedistance_nh(dense(op1), dense(op2)))
D(x1::StateVector, x2::StateVector) = norm(x2-x1)

# Test state creation
Nmodes = 5
b = GenericBasis(Nmodes)
@test bosonstates(b, 1) == fermionstates(b, 1)
@test bosonstates(b, 2) != fermionstates(b, 2)
@test length(bosonstates(b, 1)) == Nmodes
@test bosonstates(b, 1) ∪ bosonstates(b, 2) == bosonstates(b, [1, 2])
@test fermionstates(b, 1) ∪ fermionstates(b, 2) == fermionstates(b, [1, 2])

@test ManyBodyBasis(b, bosonstates(b, 1)) == ManyBodyBasis(b, fermionstates(b, 1))
@test ManyBodyBasis(b, bosonstates(b, 2)) != ManyBodyBasis(b, fermionstates(b, 2))

# Test basisstate
b_mb = ManyBodyBasis(b, bosonstates(b, 2))
psi_mb = basisstate(b_mb, [2, 0, 0, 0, 0])
op = dm(basisstate(b, 1))
@test onebodyexpect(op, psi_mb) ≈ 2
psi_mb = basisstate(b_mb, [1, 0, 1, 0, 0])
@test onebodyexpect(op, psi_mb) ≈ 1
@test_throws ArgumentError basisstate(b_mb, [1, 0, 0, 0, 0])

# Test creation operator
b_mb = ManyBodyBasis(b, bosonstates(b, [0, 1, 2]))
vac = basisstate(b_mb, [0, 0, 0, 0, 0])
at1 = create(b_mb, 1)
at2 = create(b_mb, 2)
at3 = create(b_mb, 3)
at4 = create(b_mb, 4)
at5 = create(b_mb, 5)
@test 1e-12 > D(at1*vac, basisstate(b_mb, [1, 0, 0, 0, 0]))
@test 1e-12 > D(at2*vac, basisstate(b_mb, [0, 1, 0, 0, 0]))
@test 1e-12 > D(at3*vac, basisstate(b_mb, [0, 0, 1, 0, 0]))
@test 1e-12 > D(at4*vac, basisstate(b_mb, [0, 0, 0, 1, 0]))
@test 1e-12 > D(at5*vac, basisstate(b_mb, [0, 0, 0, 0, 1]))

@test 1e-12 > D(at1*at1*vac, sqrt(2)*basisstate(b_mb, [2, 0, 0, 0, 0]))
@test 1e-12 > D(at1*at2*vac, basisstate(b_mb, [1, 1, 0, 0, 0]))
@test 1e-12 > norm(at1*basisstate(b_mb, [2, 0, 0, 0, 0]))

# Test annihilation operator
b_mb = ManyBodyBasis(b, bosonstates(b, [0, 1, 2]))
vac = basisstate(b_mb, [0, 0, 0, 0, 0])
a1 = destroy(b_mb, 1)
a2 = destroy(b_mb, 2)
a3 = destroy(b_mb, 3)
a4 = destroy(b_mb, 4)
a5 = destroy(b_mb, 5)
@test 1e-12 > D(a1*basisstate(b_mb, [1, 0, 0, 0, 0]), vac)
@test 1e-12 > D(a2*basisstate(b_mb, [0, 1, 0, 0, 0]), vac)
@test 1e-12 > D(a3*basisstate(b_mb, [0, 0, 1, 0, 0]), vac)
@test 1e-12 > D(a4*basisstate(b_mb, [0, 0, 0, 1, 0]), vac)
@test 1e-12 > D(a5*basisstate(b_mb, [0, 0, 0, 0, 1]), vac)

@test 1e-12 > D(a1*a1*basisstate(b_mb, [2, 0, 0, 0, 0]), sqrt(2)*vac)
@test 1e-12 > D(a1*a2*basisstate(b_mb, [1, 1, 0, 0, 0]), vac)
@test 1e-12 > norm(a1*vac)

# Test number operator
b = GenericBasis(3)
b_mb = ManyBodyBasis(b, bosonstates(b, [0, 1, 2, 3, 4]))
vac = basisstate(b_mb, [0, 0, 0])
n = number(b_mb)
n1 = number(b_mb, 1)
n2 = number(b_mb, 2)
n3 = number(b_mb, 3)

@test 1 ≈ expect(n2, basisstate(b_mb, [0, 1, 0]))
@test 2 ≈ expect(n3, basisstate(b_mb, [0, 0, 2]))
@test 0 ≈ expect(n1+n2+n3, vac)
@test 4 ≈ expect(n, basisstate(b_mb, [1, 3, 0]))

psi = randstate(b_mb)
rho = randoperator(b_mb)

n1_ = dm(basisstate(b, 1))
n2_ = dm(basisstate(b, 2))
n3_ = dm(basisstate(b, 3))

@test expect(n1, psi) ≈ onebodyexpect(n1_, psi)
@test expect(n2, psi) ≈ onebodyexpect(n2_, psi)
@test expect(n3, psi) ≈ onebodyexpect(n3_, psi)
@test expect(n1, rho) ≈ onebodyexpect(n1_, rho)
@test expect(n2, rho) ≈ onebodyexpect(n2_, rho)
@test expect(n3, rho) ≈ onebodyexpect(n3_, rho)

# Test transition operator
b = NLevelBasis(4)
b_mb = ManyBodyBasis(b, bosonstates(b, [0, 1, 2, 3]))
t23 = transition(b_mb, 2, 3)
t32 = transition(b_mb, 3, 2)

@test t23 == dagger(t32)
@test t32*basisstate(b_mb, [0, 1, 0, 0]) == basisstate(b_mb, [0, 0, 1, 0])
@test 1e-12 > D(t32*basisstate(b_mb, [0, 2, 1, 0]), 2*basisstate(b_mb, [0, 1, 2, 0]))
@test 1e-12 > D(number(b_mb, 2), transition(b_mb, 2, 2))

# Single particle operator in second quantization
b_single = GenericBasis(Nmodes)
b = ManyBodyBasis(b_single, bosonstates(b_single, [1, 2]))
x = randoperator(b_single)
y = randoperator(b_single)

X = manybodyoperator(b, x)
Y = manybodyoperator(b, y)

@test 1e-12 > D(X + Y, manybodyoperator(b, x + y))

x_ = sparse(x)
y_ = sparse(y)
@test 1e-12 > D(X, manybodyoperator(b, x_))
@test 1e-12 > D(Y, manybodyoperator(b, y_))
@test 1e-12 > D(X + Y, manybodyoperator(b, x_ + y_))

# Particle-particle interaction operator in second quantization
x = randoperator(b_single ⊗ b_single)
y = randoperator(b_single ⊗ b_single)

X = manybodyoperator(b, x)
Y = manybodyoperator(b, y)

@test 1e-12 > D(X + Y, manybodyoperator(b, x + y))

x_ = sparse(x)
y_ = sparse(y)
@test 1e-12 > D(X, manybodyoperator(b, x_))
@test 1e-12 > D(Y, manybodyoperator(b, y_))
@test 1e-12 > D(X + Y, manybodyoperator(b, x_ + y_))

# Test one-body expect
x = randoperator(b_single)
y = randoperator(b_single)
X = manybodyoperator(b, x)
Y = manybodyoperator(b, y)

psi = randstate(b)

@test onebodyexpect(x, psi) ≈ expect(X, dm(psi))
@test onebodyexpect(x, Y) ≈ expect(X, Y)
@test onebodyexpect(sparse(x), psi) ≈ expect(X, dm(psi))
@test onebodyexpect(sparse(x), Y) ≈ expect(X, Y)
@test onebodyexpect(x, [psi, Y]) == [onebodyexpect(x, psi), onebodyexpect(x, Y)]

@test_throws ArgumentError manybodyoperator(b_mb, x)
@test_throws ArgumentError onebodyexpect(X, psi)
@test_throws ArgumentError onebodyexpect(X, dm(psi))

end # testset


import LinearAlgebra.norm

# deprecate these soon
norm_inf(x::AbstractExpr) = maximum(abs(x))
norm_1(x::AbstractExpr) = sum(abs(x))
norm_fro(x::AbstractExpr) = norm2(vec(x))

"""
    norm(x::AbstractExpr, p::Real=2)

Computes the `p`-norm `‖x‖ₚ = (∑ᵢ |xᵢ|^p)^(1/p)` of a vector expression `x`.

This function uses specialized methods for `p=1, 2, Inf`. For `p > 1` otherwise,
this function uses the procedure documented at
[`rational_to_socp.pdf`](https://github.com/jump-dev/Convex.jl/raw/master/docs/supplementary/rational_to_socp.pdf),
based on the paper "Second-order cone programming" by F. Alizadeh and D. Goldfarb,
Mathematical Programming, Series B, 95:3-51, 2001.

!!! warning
    For versions of Convex.jl prior to v0.14.0, `norm` on a matrix expression returned
    the operator norm ([`opnorm`](@ref)), which matches Julia v0.6 behavior. This functionality
    was deprecated since Convex.jl v0.8.0, and has been removed. In the future,
    `norm(x, p)` will return `‖vec(x)‖ₚ`, matching the behavior of [`norm`](@ref)
    for numeric matrices.
"""
function LinearAlgebra.norm(x::AbstractExpr, p::Real=2)
    if length(size(x)) <= 1 || minimum(size(x))==1
        # x is a vector
        if p == 1
            return norm_1(x)
        elseif p == 2
            return norm2(x)
        elseif p == Inf
            return norm_inf(x)
        elseif p > 1
            # TODO: allow tolerance in the rationalize step
            return rationalnorm(x, rationalize(Int, float(p)))
        else
            error("vector p-norms not defined for p < 1")
        end
    else
        error("In Convex.jl v0.13 and below, `norm(x, p)` meant `opnorm(x, p)` (but was deprecated since v0.8.0). In the future, `norm(x,p)` for matrices will be equivalent to `norm(vec(x),p)`. This is currently an error to ensure you update your code!")
    end
end


using PointPatternStatistics
using Test
include("reference_implementation.jl")
@testset "PointPatternStatistics.jl" begin
    window = (x=(-1.3,1), y=(0,1.6))
    randu((a, b)) = a + rand() * (b-a)
    randw(window) = (randu(window.x), randu(window.y))
    xy = [randw(window) for _ in 1:300]
    r = range(0, 0.25, length=512)
    @testset "iterate" begin
        @test collect(PointPattern(xy, window)) == xy
    end

    @testset "pcf" begin
        est = pcf(PointPattern(xy, window), r)
        ref = pcfref(xy, window, r)
        @test sum(abs2, (est - ref)[2:end]) < 0.0001
    end

    @testset "K" begin
        est = Kest(PointPattern(xy, window), r)
        ref = Kref(xy, window, r)
        @test est ≈ ref
    end

    @testset "F" begin
        include("F.jl")
    end
    @testset "G" begin
        include("G.jl")
    end
    #include("envelope.jl")
    @testset "SSI" begin
        include("SSI.jl")
    end
end


using CoexistenceHoles
using Test

@testset "is_GLVlocallystable.jl" begin
    A1 = [-1     0       0       -0.148;
         0      -1      0       0;
         0      0.0997  -1      0;
         0.14   0       -0.05   -1];
    r1 = [0.1, -0.3, 2.1, -1];

    A2 = [-1     0       0       -0.148;
         0      -1      0       0;
         0      0.0997  -1      0;
         0.14   0       -0.05   -1];
    r2 = [0.13, 0.1, 0.1, 0.2];

    @test !is_GLVlocallystable(A1,r1)
    @test is_GLVlocallystable(A2, r2)

end

@testset "is_GLVpermanent.jl" begin
    A1 = [-1     0       0       -0.148;
         0      -1      0       0;
         0      0.0997  -1      0;
         0.14   0       -0.05   -1];
    r1 = [0.1, -0.3, 2.1, -1];

    A2 = [-1     0       0       -0.148;
         0      -1      0       0;
         0      0.0997  -1      0;
         0.14   0       -0.05   -1];
    r2 = [0.13, 0.1, 0.1, 0.2];

    @test !is_GLVpermanent(A1, r1)
    @test is_GLVpermanent(A2, r2)
end


using DistanceTransforms
using Test

@testset ExtendedTestSet "mean_hausdorff" begin
    @testset ExtendedTestSet "mean_hausdorff" begin
    x = [1 1 1 0; 1 1 1 0; 1 1 1 0; 1 1 1 0]
    y = [1 1 1 0; 1 1 1 0; 1 1 1 0; 1 1 1 0]

    @test mean_hausdorff(x, y) ≈ 0
    end
end

@testset ExtendedTestSet "mean_dice" begin
    @testset ExtendedTestSet "mean_dice" begin
    x = [1 1 1 0; 1 1 1 0; 1 1 1 0; 1 1 1 0]
    y = [1 1 1 0; 1 1 1 0; 1 1 1 0; 1 1 1 0]

    @test mean_dice(x, y) ≈ 1
    end
end


@doc doc"""
`Q` state Potts model with energy $E = -\sum_{i,j} \delta_{\sigma_i, \sigma_j}$,
where $\sigma_i$ takes an integer value from $1$ to $Q$ and $\delta$ is a Kronecker's delta.
Order parameter (total magnetization) is defined as
\begin{equation}
    M = \frac{Q-1}{Q}N_1 - \frac{1}{Q}(N-N_1),
\end{equation}
where $N$ is the number of sites and $N_1$ is the number of $\sigma=1$ spins.
"""
mutable struct Potts <: Model
    lat :: Lattice
    Q :: Int
    spins :: Matrix{Int}
    rng :: Random.MersenneTwister

    function Potts(lat::Lattice, Q::Integer, rng::Random.AbstractRNG)
        if Q < 2
            error("Q should be 2 or more: Q = $Q")
        end
        spins = rand(rng, 1:Q, 1, numsites(lat))
        return new(lat, Q, spins, rng)
    end
end

Potts(lat::Lattice, Q::Integer) = Potts(lat, Q, Random.seed!(Random.MersenneTwister(0)))
Potts(lat::Lattice, Q::Integer, seed) = Potts(lat, Q, Random.seed!(Random.MersenneTwister(0), seed...))

@doc doc"""
    Potts(param)

Generates `Potts` using `param["Lattice"]`, `param["Q"]`, and `param["Seed"]` (if defined).
Each spin will be initialized randomly and independently.
"""
function Potts(param::Parameter)
    lat = generatelattice(param)
    Q = param["Q"]
    if "Seed" in keys(param)
        return Potts(lat, Q, param["Seed"])
    else
        return Potts(lat, Q)
    end
end

include("update.jl")
include("estimator.jl")
include("postproc.jl")


include("../src/geom_module.jl")

module dtest
using Random
using BenchmarkTools
using ..geom
using PyPlot
PyPlot.pygui(true)

# x_ = vec([0 1 0.3 1.5 1.7 -0.9 -1])
# y_ = vec([ 0 0.2 1 1.5 3 2.55 0.65])

N = 1000
R = 10.0
# # R_x = 30
# # R_y = 3
# #
# # phi = LinRange(0, 2*π, N)
# # x_ = R_x .* map(cos, phi)
# # y_ = R_y .* map(sin, phi)
#
r_ = R .* Random.rand(N)
phi_ = (2*π) .* Random.rand(N)
x_ = r_ .* cos.(phi_)
y_ = r_ .* sin.(phi_)

point_array = map(geom.MutablePoint, x_, y_)

dmesh = geom.build_delaunay_mesh!(point_array)

point_array2 = copy(point_array)
rnb_idx = Random.randperm(length(point_array2))

(p_x, p_y) = geom.discretize_for_plot(point_array)
(init_tri_x, init_tri_y) = geom.discretize_for_plot(
  [dmesh.vertex_array[1], dmesh.vertex_array[2], dmesh.vertex_array[3]], true)

tri_vec = geom.triangle_vector_for_plot(dmesh)

PyPlot.figure()
PyPlot.grid()
PyPlot.plot(p_x, p_y,
    color = "k",
    alpha = 0.7,
    linestyle = "none",
    marker = ".",
    markersize = 10)
PyPlot.plot(init_tri_x, init_tri_y,
  marker = "o", linewidth = 2)

for kk = 1:length(tri_vec)
    (p_x, p_y) = geom.discretize_for_plot(tri_vec[kk])
    PyPlot.plot(p_x, p_y,
      marker = "o",
      linewidth = 1,
      color = "r",
      alpha = 0.7,)
end

end # module dtest


# ---
# title: 407. Trapping Rain Water II
# id: problem407
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Heap, Breadth-first Search
# link: <https://leetcode.com/problems/trapping-rain-water-ii/description/>
# hidden: true
# ---
#
# Given an `m x n` matrix of positive integers representing the height of each
# unit cell in a 2D elevation map, compute the volume of water it is able to
# trap after raining.
#
# **Example:**
#
#
#
#     Given the following 3x6 height map:
#     [
#       [1,4,3,1,3,2],
#       [3,2,1,3,2,4],
#       [2,3,3,2,3,1]
#     ]
#
#     Return 4.
#
#
# ![](https://assets.leetcode.com/uploads/2018/10/13/rainwater_empty.png)
#
# The above image represents the elevation map
# `[[1,4,3,1,3,2],[3,2,1,3,2,4],[2,3,3,2,3,1]]` before the rain.
#
#
#
# ![](https://assets.leetcode.com/uploads/2018/10/13/rainwater_fill.png)
#
# After the rain, water is trapped between the blocks. The total volume of water
# trapped is 4.
#
#
#
# **Constraints:**
#
#   * `1 <= m, n <= 110`
#   * `0 <= heightMap[i][j] <= 20000`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


#! /usr/bin/julia

# Rosetta Code, Knapsack problem/Continuous

immutable KPCSupply{S<:String, T<:Real}
    item::S
    weight::T
    value::T
    uvalue::T
end
Base.isless(a::KPCSupply, b::KPCSupply) = a.uvalue<b.uvalue

function KPCSupply{S<:String, T<:Real}(item::S, weight::T, value::T)
    KPCSupply(item, weight, value, value/weight)
end

function solve{S<:String, T<:Real}(store::Array{KPCSupply{S,T},1},
                                   capacity::T)
    ksack = KPCSupply{S,T}[]
    kweight = zero(T)
    for s in sort(store, rev=true)
        if kweight + s.weight <= capacity
            kweight += s.weight
            push!(ksack, s)
        else
            w = capacity-kweight
            v = w*s.uvalue
            push!(ksack, KPCSupply(s.item, w, v, s.uvalue))
            break
        end
    end
    return ksack
end

store = [KPCSupply("beef", 38//10, 36//1),
         KPCSupply("pork", 54//10, 43//1),
         KPCSupply("ham", 36//10, 90//1),
         KPCSupply("greaves", 24//10, 45//1),
         KPCSupply("flitch", 4//1, 30//1),
         KPCSupply("brawn", 25//10, 56//1),
         KPCSupply("welt", 37//10, 67//1),
         KPCSupply("salami", 3//1, 95//1),
         KPCSupply("sausage", 59//10, 98//1)]

sack = solve(store, 15//1)

println("The store contains:")
println("        Item Weight  Unit Price")
for s in store
    println(@sprintf("%12s %4.1f    %6.2f",
                     s.item, float(s.weight), float(s.uvalue)))
end

println()
println("The thief should take:")
println("        Item Weight  Value")
for s in sack
    println(@sprintf("%12s %4.1f    %6.2f",
                     s.item, float(s.weight), float(s.value)))
end
w = mapreduce(x->x.weight, +, sack)
v = mapreduce(x->x.value, +, sack)
println(@sprintf("%12s %4.1f    %6.2f",
                 "Total", float(w), float(v)))


using DelayDiffEq
using Test

# Test parameterized delayed logistic equation

# delayed logistic equation
f_inplace(du, u, h, p, t) = (du[1] = p[1] * u[1] * (1 - h(p, t-1; idxs = 1)))
f_scalar(u, h, p, t) = p[1] * u * (1 - h(p, t-1))

# simple history function
h(p, t; idxs=nothing) = 0.1

@testset for inplace in (true, false)
  # define problem
  # we specify order_discontinuity_t0 = 1 to indicate that the discontinuity at
  # t = 0 is of first order
  prob = DDEProblem(inplace ? f_inplace : f_scalar,
                    inplace ? [0.1] : 0.1,
                    h, (0.0, 50.0), [0.3];
                    constant_lags = [1],
                    order_discontinuity_t0 = 1)

  # solve problem with initial parameter:
  sol1 = solve(prob, MethodOfSteps(Tsit5()))
  @test length(sol1) == 21
  @test first(sol1(12)) ≈ 0.884 atol = 1e-4
  @test first(sol1[end]) ≈ 1 atol = 1e-5

  # solve problem with updated parameter
  prob.p[1] = 1.4
  sol2 = solve(prob, MethodOfSteps(Tsit5()))
  @test length(sol2) == 47
  @test first(sol2(12)) ≈ 1.125 atol = 5e-4
  @test first(sol2[end]) ≈ 0.994 atol = 2e-5
end


using CUDA
using IndexFunArrays

function f(arr, x)
    arr .= arr.^2 .+ arr.^3 .+ sqrt.(x)
    return arr
end

function f2(arr)
    arr .= arr.^2 .+ arr.^3 .+ sqrt.(arr .^2)
    return arr
end

function f3(arr)
    arr = arr.^2 .+ arr.^3 .+ sqrt.(arr .^2)
    return arr
end

function test(s)
    arr = randn(Float32, s)
    arr_c = CuArray(arr)
    x = rr2(Float32, s)
    x_c = CuArray(rr(Float32, s))

    @info "f on CPU"
    @time f(arr, x);
    @time f(arr, x);
    @info "f2 on CPU"
    @time f2(arr);
    @time f2(arr);
    @info "f on GPU"
    CUDA.@time f(arr_c, x);
    CUDA.@time f(arr_c, x);
    @info "f2 on GPU"
    CUDA.@time f2(arr_c);
    CUDA.@time f2(arr_c);
    @info "f3 on GPU"
    CUDA.@time f3(arr_c);
    CUDA.@time f3(arr_c);

    return
end

function MandelbrotIterations(c, Niter=500)
    x = c
    res = 0
    for n in 1:Niter
        x = x.^2 .+ c
        #if abs(x) > 4.0 && (res == 0)
        #    res = n
        #end
        res += (res == 0)*(abs(x)> 4.0)*n  # ensure that only the first breakage counts but avoid if
    end
    return (res % 127)
end

function MandelbrotIterations(idx,start,scale, Niter=100)
    c = Complex(start[1].+idx[1].*scale[1],start[2].+idx[2].*scale[2])
    return MandelbrotIterations(c..., Niter)
end

function testMandelbrotArray(s=(1024,1024), Niter=100)
    arr = zeros(Int64, s)
    scale = (3.0,3.0)./s;
    start=(-2.3,-1.5);
    @info "Test MandelBrot Array CPU"
    @time arr .= MandelbrotIterations.(CartesianIndices(arr), [scale], [start])
    @time arr .= MandelbrotIterations.(CartesianIndices(arr), [scale], [start])
end

function testMandelbrotIFA(s=(1024,1024), Niter=100)
    arr = zeros(Int64, s)
    scale = (3.0,3.0)./s;
    start=(-2.3,-1.5);
    Mandelbrot = IndexFunArray(x -> MandelbrotIterations(x,start, scale,Niter), size(arr));
    @info "Test MandelBrot IndexFunArray CPU"
    @time arr .= Mandelbrot;
    @time arr .= Mandelbrot;

    #arr_c = CuArray(arr)
    #CUDA.@time arr_c .= Mandelbrot
    #CUDA.@time arr_c .= Mandelbrot
    #CUDA.@time arr_c .= MandelbrotIterations.(arr_c,start,scale,Niter)
    return arr
end

#using Napari
#napari.view_image(testMandelbrotIFA())


function euler_step(q,p,dt,HS)
    Q = q+dt*HS.Minv*p
    P = p+dt*HS.F(q)
    return (dt,Q,P)
end

function i_euler_step(q,p,dt,HS;tol = 1e-6,max_iter=100,no_warning=false)
    Q=q;P=p
    for i in 1:max_iter
        Qn = q+dt*HS.Minv*P
        Pn = p+dt*HS.F(Q)
        if norm(Qn-Q)<=tol && norm(Pn-P)<=tol
            return (dt,Qn,Pn)
        end
        Q=Qn
        P=Pn
    end
    if !no_warning
        println("WARNING: Implicit euler step did not reach tol.")
    end
    return (dt,Q,P)
end

function midpoint_step(q,p,dt,HS)
    Q = q+dt*HS.Minv*p
    P = p+dt*HS.F(q)
    Q = q+dt*HS.Minv*((p+P)/2)
    P = p+dt*HS.F((q+Q)/2)
    return (dt,Q,P)
end

function i_midpoint_step(q,p,dt,HS;tol = 1e-6,max_iter=100,no_warning=false)
    Q=q;P=p
    for i in 1:max_iter
        Qn = q+dt*HS.Minv*((p+P)/2)
        Pn = p+dt*HS.F((q+Q)/2)
        if norm(Qn-Q)<=tol && norm(Pn-P)<=tol
            return (dt,Qn,Pn)
        end
        Q=Qn
        P=Pn
    end
    if !no_warning
        println("WARNING: Implicit midpoint step did not reach tol.")
    end
    return (dt,Q,P)
end

function theta_step(q,p,dt,HS;theta=0.5, kwargs...)
    _,Qe,Pe = euler_step(q,p,dt,HS)
    _,Qi,Pi = i_euler_step(q,p,dt,HS;kwargs...)
    Q,P = theta.*(Qe,Pe).+(1-theta).*(Qi,Pi)
    return (dt,Q,P)
end

function s_euler_step(q,p,dt,HS)
    P = p+dt*HS.F(q)
    Q = q+dt*HS.Minv*P
    return (dt,Q,P)
end

function sa_euler_step(q,p,dt,HS)
    Q = q+dt*HS.Minv*p
    P = p+dt*HS.F(Q)
    return (dt,Q,P)
end

function verlet_step(q,p,dt,HS)
    Q = q+dt*HS.Minv*p+dt^2/2*HS.Minv*HS.F(q)
    P = p+dt/2*(HS.F(q)+HS.F(Q))
    return (dt,Q,P)
end

function q_split_verlet_step(q,p,dt,HS)
    Q = q + dt/2*HS.Minv*p
    P = p + dt*(HS.F(Q))
    Q = Q + dt/2*HS.Minv*P
    return (dt,Q,P)
end

function p_split_verlet_step(q,p,dt,HS)
    P = p + dt/2*(HS.F(q))
    Q = q + dt*HS.Minv*P
    P = P + dt/2*(HS.F(Q))
    return (dt,Q,P)
end

function puls_projected_step(q,p,dt,HS;step=verlet_step,kwargs...)
    (dt,Q,P) = verlet_step(q,p,dt,HS;kwargs...);
    E = getEnergy(HS,HS.q0,HS.p0)
    if P'*HS.Minv*P<0.001 || U(Q)>E
        gamma = 1
    else
        gamma = sqrt(2*(E-HS.U(Q))/(P'*HS.Minv*P))
    end
    return (dt,Q,gamma*P)
end

function rk4_step(q,p,dt,HS)
    k1q = dt*HS.Minv*p
    k1p = dt*HS.F(q)

    k2q = dt*HS.Minv*(p+k1p/2)
    k2p = dt*HS.F(q+k1q/2)

    k3q = dt*HS.Minv*(p+k2p/2)
    k3p = dt*HS.F(q+k2q/2)

    k4q = dt*HS.Minv*(p+k3p)
    k4p = dt*HS.F(q+k3q)

    Q = q + 1/6*(k1q + 2*k2q + 2*k3q + k4q)
    P = p + 1/6*(k1p + 2*k2p + 2*k3p + k4p)
    return (dt,Q,P)
end

function projected_step(q,p,dt,HS;tol=1e-5,sigma=0.5,max_iter=100,step=rk4_step,no_warning=false,kwargs...)
    dt,Q,P = step(q,p,dt,HS,kwargs...)
    E0 = getEnergy(HS,HS.q0,HS.p0)
    E=getEnergy(HS,Q,P)-E0
    i=0;
    while abs(E)>tol
        if i > max_iter
            if !no_warning
                println("WARNING: projected step did not reach tol")
            end
            return (dt,Q,P)
        end
        dq = HS.F(Q)*E
        dp = -HS.Minv*P*E
        s=1/sigma
        diff = 0
        E_new = E
        j = 0
        while diff>=0
            if j>=max_iter
                if !no_warning
                    println("WARNING: projected step could not find sigma small enough!")
                end
                return (dt,Q,P)
            end
            s*=sigma
            E_new = getEnergy(HS,(Q+s*dq),(P+s*dp))-E0
            diff = E_new^2-E^2
            j+=1
        end
        E=E_new
        Q=Q+s*dq
        P=P+s*dp
        i+=1
    end
    return (dt,Q,P)
end

println("Finished loading HamiltonSteps")