##############################################################################
#
# In this test problem we consider a quadratic penalty method.
# This example features an algorithm with the 3 steps:
# the penalization - the unconstrained min - the 1d min
#
# Note that there is no optimization of the evaluations here.
# The penalization gives an approximation of the gradients, multipliers...
#
# Note the use of a structure for the algorithmic parameters which is
# forwarded to all the 3 steps. If a parameter is not mentioned, then the default
# entry in the algorithm will be taken.
#
#############################################################################

#include("uncons.jl")

##############################################################################
#
# Quadratic penalty algorithm
# fill_in! used instead of update! (works but usually more costly in evaluations)
# subproblems are solved via Newton method
#
#############################################################################
function penalty(stp :: NLPStopping; rho0 = 1.0, rho_min = 1e-10,
                                     rho_update = 0.5, prms = nothing)

 #algorithm's parameters
 rho = rho0

 #First call to the stopping
 #Becareful here, some entries must be filled in first.
 fill_in!(stp, stp.current_state.x)
 OK = start!(stp)

 #prepare the subproblem stopping:
 sub_nlp_at_x = NLPAtX(stp.current_state.x)
 sub_pb  = ADNLPModel(x -> obj(stp.pb, x)
                      + 1/rho * norm(max.(cons(stp.pb, x) - stp.pb.meta.ucon, 0.0))^2
                      + 1/rho * norm(max.(- cons(stp.pb, x) + stp.pb.meta.lcon, 0.0))^2,  x0)
 sub_stp = NLPStopping(sub_pb, unconstrained_check,
                               sub_nlp_at_x, main_stp = stp)

 #main loop
 while !OK

  #solve the subproblem
  reinit!(sub_stp)
  sub_stp.meta.atol = min(rho, sub_stp.meta.atol)
  global_newton(sub_stp, prms)

  #Update all the entries of the State
  fill_in!(stp, sub_stp.current_state.x)

  #Either stop! is true OR the penalty parameter is too small
  if rho < rho_min stp.meta.suboptimal = true end
  OK = stop!(stp)

  @show stp.meta.nb_of_stop, OK, rho

  #update the penalty parameter if necessary
  if !OK
   rho = rho * rho_update
   sub_stp.pb  = ADNLPModel(x -> obj(stp.pb, x)
                            + 1/rho * norm(max.(cons(stp.pb, x) - stp.pb.meta.ucon, 0.0))^2
                            + 1/rho * norm(max.(- cons(stp.pb, x) + stp.pb.meta.lcon, 0.0))^2,  x0)
  end
 end

 return stp
end

##############################################################################
#
# Quadratic penalty algorithm: buffer function
#
#############################################################################
function penalty(stp :: NLPStopping, prms)

 #extract required values in the prms file
 r0 = :rho0       ∈ fieldnames(typeof(prms)) ? prms.rho0       : 1.0
 rm = :rho_min    ∈ fieldnames(typeof(prms)) ? prms.rho_min    : 1e-10
 ru = :rho_update ∈ fieldnames(typeof(prms)) ? prms.rho_update : 0.5

 return penalty(stp, rho0 = r0, rho_min = rm, ru = 0.5, prms = prms)
end

##############################################################################
#
# Algorithmic parameters structure
#
#############################################################################

mutable struct Param

    #parameters for the penalty
    rho0       :: Float64 #initial value of the penalty parameter
    rho_min    :: Float64 #smallest possible parameter
    rho_update :: Float64 #update of the penalty parameter

    #parameters of the unconstrained minimization
    armijo_prm  :: Float64 #Armijo parameter
    wolfe_prm   :: Float64 #Wolfe parameter
    onedsolve   :: Function #1D solver
    ls_func     :: Function

    #parameters of the 1d minimization
    back_update :: Float64 #backtracking update

    function Param(;rho0        :: Float64 = 1.0,
                    rho_min     :: Float64 = sqrt(eps(Float64)),
                    rho_update  :: Float64 = 0.5,
                    armijo_prm  :: Float64 = 0.01,
                    wolfe_prm   :: Float64 = 0.99,
                    onedsolve   :: Function = backtracking_ls,
                    ls_func     :: Function = (x,y)-> armijo(x,y, τ₀ = armijo_prm),
                    back_update :: Float64 = 0.5)
        return new(rho0, rho_min, rho_update,
                   armijo_prm, wolfe_prm, onedsolve, ls_func,
                   back_update)
    end
end


path = joinpath(dirname(@__FILE__), "..", "..")
push!(Base.LOAD_PATH, path)
datapath = joinpath(dirname(@__FILE__), "data")

using ProxSDP

sympackedlen(n) = div(n*(n+1), 2)
sympackeddim(n) = div(isqrt(1+8n) - 1, 2)
function ivech!(out::AbstractMatrix{T}, v::AbstractVector{T}) where T
    n = sympackeddim(length(v))
    n1, n2 = size(out)
    @assert n == n1 == n2
    c = 0
    for j in 1:n, i in 1:j
        c += 1
        out[i,j] = v[c]
    end
    return out
end
function ivech(v::AbstractVector{T}) where T
    n = sympackeddim(length(v))
    out = zeros(n, n)
    ivech!(out, v)

    return out
end

include("moitest.jl")


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    SimpleKriging(γ, μ)
    SimpleKriging(data, γ, μ)

Simple Kriging with variogram model `γ` and constant mean `μ`.

Optionally, pass the geospatial `data` to the [`fit`](@ref) function.

### Notes

* Simple Kriging requires stationary variograms
"""
struct SimpleKriging{G<:Variogram,V} <: KrigingEstimator
  # input fields
  γ::G
  μ::V

  function SimpleKriging{G,V}(γ, μ) where {G<:Variogram,V}
    @assert isstationary(γ) "Simple Kriging requires stationary variogram"
    new(γ, μ)
  end
end

SimpleKriging(γ, μ) = SimpleKriging{typeof(γ),typeof(μ)}(γ, μ)

SimpleKriging(data, γ, μ) = GeoStatsBase.fit(SimpleKriging(γ, μ), data)

nconstraints(::SimpleKriging) = 0

set_constraints_lhs!(::SimpleKriging, LHS::AbstractMatrix, domain) = nothing

factorize(::SimpleKriging, LHS::AbstractMatrix) = cholesky(Symmetric(LHS), check=false)

set_constraints_rhs!(::FittedKriging{<:SimpleKriging}, pₒ) = nothing

function combine(fitted::FittedKriging{<:SimpleKriging},
                 weights::KrigingWeights, z::AbstractVector)
  γ = fitted.estimator.γ
  μ = fitted.estimator.μ
  b = fitted.state.RHS
  λ = weights.λ
  y = z .- μ

  μ + y⋅λ, sill(γ) - b⋅λ
end


@testset "condition 1" begin
    t = trace()
    n_datapoints = 2
    program!(t, n_datapoints)
    @test t[(:data, 1)].interpretation == NONSTANDARD
    cp = condition(program!, Dict((:data, 1) => 4))
    t, _ = cp(t, n_datapoints)
    @test t[(:data, 1)].interpretation == CONDITIONED
    @test t[(:data, 1)].value == 4
    @info "(:data, 1) = $(t[(:data, 1)])"
end

function program2!(t :: Trace, data)
    loc = sample(t, :loc, Normal())
    scale = sample(t, :scale, Gamma(2.0, 2.0))
    log_rate = sample(t, :log_rate, Normal(loc, scale))
    obs = Array{Int64, 1}(undef, length(data))
    for (i, d) in enumerate(data)
        obs[i] = observe(t, (:data, i), Poisson(exp(log_rate)), d)
    end
    obs
end

@testset "observing data" begin
    data = [1, 1, 2, 1, 5]
    t = trace()
    obs = program2!(t, data)
    logprob!(t)
    @test obs == data
    # log likelihood is always >= log joint density
    log_joint = logprob(t)
    log_likelihood = loglikelihood(t)
    @info "log p(x, z) = $log_joint"
    @info "log p(x | z) = $log_likelihood"
    @test log_likelihood >= log_joint

    # copy the old trace and demonstrate modification
    t_copy = deepcopy(t)
    obs = program2!(t, data)
    @info ":loc node in copied original: $(t_copy[:loc])"
    @info ":loc node in modified original: $(t[:loc])"
    @test t[:loc].value != t_copy[:loc].value
end

@testset "replay" begin
    data = [1, 1, 2, 1, 5]
    # trace the program
    t = trace()
    obs = program2!(t, data)
    logprob!(t)
    # modify a trace element
    t2 = deepcopy(t)
    t2[:scale] = node(2.0, :scale, t[:scale].dist, false, NONSTANDARD)
    @test t[:scale].value != t2[:scale].value
    @test t[:scale].dist == t2[:scale].dist
    # replay the program using the new trace
    t_new, replayed_f = replay(program2!, t2)
    r_obs = replayed_f(data)
    logprob!(t_new)
    @test r_obs == data
    @test t[:scale].dist == t_new[:scale].dist
    @test t[:log_rate].dist != t_new[:log_rate].dist
    @test t_new[:scale].value == 2.0
    @info "Old scale node: $(t[:scale])"
    @info "New scale node: $(t_new[:scale])"
    @info "Old log_rate node: $(t[:log_rate])"
    @info "New log_rate node: $(t_new[:log_rate])"
end

@testset "block" begin
    data = [1, 1, 2, 1, 5]
    t = trace()
    obs = program2!(t, data)
    logprob!(t)
    @test :loc in keys(t)

    t_new, blocked_f = block(program2!, t, (:loc,))
    r_obs = blocked_f(data)
    logprob!(t_new)
    @test !(:loc in keys(t_new))
end

function plate_program!(t :: Trace, data)
    loc = sample(t, :loc, Normal())
    scale = sample(t, :scale, Gamma(2.0, 2.0))
    log_rate = sample(t, :log_rate, Normal(loc, scale))
    plate(t, observe, :data, Poisson(exp(log_rate)), data)
end

@testset "plate observe" begin
    data = [1, 2, 3]
    t = trace()
    plate_program!(t, data)
    @info "Trace: $t"
    @test typeof(t[:data].value) == Vector{Int64}
    d = Poisson(exp(t[:log_rate].value))
    @test isapprox(t[:data].logprob_sum, sum(logpdf(d, i) for i in data))
end

@testset "plate latent" begin
    t = trace()
    plate(t, sample, "address", Normal(), 10)
    @info "Trace: $t"
    @test isapprox(t["address"].logprob_sum, sum(logpdf(Normal(), r) for r in t["address"].value))
end

@testset "plate mh inference" begin
    data = [1, 3, 1, 5, 4]
    results = mh(plate_program!; params = (data,))
    @info "Single results trace: $(results.traces[end])"
    @test typeof(results.traces[end][:data].value) == Vector{Int64}
end

@testset "plate replay" begin
    t = trace()
    f = (t,) -> plate(t, sample, "address", Normal(), 3)
    plate(t, sample, "address", Normal(), 3)
    @info "Original trace: $t"
    (new_t, g) = replay(f, t)
    @info "Replayed trace before function exec: $new_t"
    g()
    @info "Replayed trace after function exec: $new_t"
end

function expanded_plate_model(t, data)
    hl = sample(t, "hl", Normal())
    locs = plate(t, sample, "locs", Normal(hl, 1.0), length(data))
    plate(t, observe, "data", Normal, data, (locs, ones(length(data))))
end

@testset "expanded obs plate 1" begin
    t = trace()
    data = randn(3)
    expanded_plate_model(t, data)
    @info t
end

locs_proposal(t0, t1, params...) = propose(t1, "locs", MvNormal(t0["locs"].value, 1.0))
hl_proposal(t0, t1, params...) = propose(t1, "hl", Normal(t0["hl"].value, 1.0))

@testset "expanded obs plate 2" begin
    data = randn(10)
    results = mh(expanded_plate_model, [locs_proposal, hl_proposal]; params = (data,))
end


function bench()
    print("( ")
    for p in PrimesPaged() p > 100 && break; print(p, " ") end
    println(")")
    countPrimesTo(Prime(100)) # warm up JIT
#=
    println(@time let count = 0
                      for p in PrimesPaged()
                          p > 1000000000 && break
                          count += 1
                      end; count end) # much slower counting over iteration
=#
    println(@time countPrimesTo(Prime(1000000000)))
end
bench()


#!/usr/bin/env julia
using Dierckx
using Compat
using Base.Test

# Answers 'ans' are from scipy.interpolate,
# generated with genanswers.py script.

# -----------------------------------------------------------------------------
# Spline1D

x = [1., 2., 3.]
y = [0., 2., 4.]
spl = Spline1D(x, y; k=1, s=length(x))

yi = evaluate(spl, [1.0, 1.5, 2.0])
@test_approx_eq(yi, [0.0, 1.0, 2.0])
@test_approx_eq(evaluate(spl, 1.5), 1.0)
@test_approx_eq(get_knots(spl), [1., 3.])
@test_approx_eq(get_coeffs(spl), [0., 4.])
@test_approx_eq_eps(get_residual(spl), 0.0, 1.e-30)

# test call on v0.4+
if VERSION >= v"0.4.0"
    @test_approx_eq(spl([1.0, 1.5, 2.0]), [0.0, 1.0, 2.0])
    @test_approx_eq(spl(1.5), 1.0)
end

# test that a copy is returned by get_knots()
knots = get_knots(spl)
knots[1] = 1000.
@test_approx_eq(get_knots(spl), [1., 3.])

# test ported from scipy.interpolate testing this bug:
# http://mail.scipy.org/pipermail/scipy-dev/2008-March/008507.html
x = [-1., -0.65016502, -0.58856235, -0.26903553, -0.17370892,
     -0.10011001, 0., 0.10011001, 0.17370892, 0.26903553, 0.58856235,
     0.65016502, 1.]
y = [1.,0.62928599, 0.5797223, 0.39965815, 0.36322694, 0.3508061,
     0.35214793, 0.3508061, 0.36322694, 0.39965815, 0.5797223,
     0.62928599, 1.]
w = [1.00000000e+12, 6.88875973e+02, 4.89314737e+02, 4.26864807e+02,
     6.07746770e+02, 4.51341444e+02, 3.17480210e+02, 4.51341444e+02,
     6.07746770e+02, 4.26864807e+02, 4.89314737e+02, 6.88875973e+02,
     1.00000000e+12]
spl = Spline1D(x, y; w=w, s=@compat(Float64(length(x))))
desired = [0.35100374, 0.51715855, 0.87789547, 0.98719344]
actual = evaluate(spl, [0.1, 0.5, 0.9, 0.99])
@test_approx_eq_eps(actual, desired, 5e-4)

# tests for out-of-range
x = [0.0:4.0;]
y = x.^3

xp = linspace(-8.0, 13.0, 100)
xp_zeros = Float64[(0. <= xi <= 4.) ? xi : 0.0 for xi in xp]
xp_clip = Float64[(0. <= xi <= 4.) ? xi : (xi<0.0)? 0.0 : 4. for xi in xp]

spl = Spline1D(x, y)
t = get_knots(spl)[2: end-1]  # knots, excluding those at endpoints
spl2 = Spline1D(x, y, t)

@test_approx_eq(evaluate(spl, xp), xp_clip.^3)
@test_approx_eq(evaluate(spl2, xp), xp_clip.^3)

# test other bc's
spl = Spline1D(x, y; bc="extrapolate")
@test_approx_eq(evaluate(spl, xp), xp.^3)
spl = Spline1D(x, y; bc="zero")
@test_approx_eq(evaluate(spl, xp), xp_zeros.^3)
spl = Spline1D(x, y; bc="error")
@test_throws ErrorException evaluate(spl, xp)

# test unknown bc
@test_throws ErrorException Spline1D(x, y; bc="unknown")

# test derivative
x = linspace(0, 1, 70)
y = x.^3
spl = Spline1D(x, y)
xt = [0.3, 0.4, 0.5]
@test_approx_eq(derivative(spl, xt), 3xt.^2)

# test integral
x = linspace(0, 10, 70)
y = x.^2
spl = Spline1D(x, y)
@test_approx_eq(integrate(spl, 1.0, 5.0), 5.^3/3 - 1/3)

# test roots
x = linspace(0, 10, 70)
y = (x-4).^2-1
spl = Spline1D(x, y)
@test_approx_eq(roots(spl), [3, 5])

# -----------------------------------------------------------------------------
# Spline2D

# test linear
x = [1., 1., 1., 2., 2., 2., 3., 3., 3.]
y = [1., 2., 3., 1., 2., 3., 1., 2., 3.]
z = [0., 0., 0., 2., 2., 2., 4., 4., 4.]
spl = Spline2D(x, y, z; kx=1, ky=1, s=length(x))
tx, ty = get_knots(spl)
@test_approx_eq(tx, [1., 3.])
@test_approx_eq(ty, [1., 3.])
@test_approx_eq_eps(get_residual(spl), 0.0, 1e-16)
@test_approx_eq(evaluate(spl, 2.0, 1.5), 2.0)
@test_approx_eq(evalgrid(spl, [1.,1.5,2.], [1.,1.5]), [0. 0.; 1. 1.; 2. 2.])

# test 1-d grid arrays
@test_approx_eq(evalgrid(spl, [2.0], [1.5])[1, 1], 2.0)

# In this setting, lwrk2 is too small in the default run.
x = linspace(-2, 2, 80)
y = linspace(-2, 2, 80)
z = x .+ y
spl = Spline2D(x, y, z; s=length(x))
@test_approx_eq(evaluate(spl, 1.0, 1.0), 2.0)

# In this setting lwrk2 is too small multiple times!
# Eventually an error about s being too small is thrown.
srand(0)
x = rand(100)
y = rand(100)
z = sin(x) .* sin(y)
@test_throws ErrorException Spline2D(x, y, z; kx=1, ky=1, s=0.0)

# test grid input creation
x = [0.5, 2., 3., 4., 5.5, 8.]
y = [0.5, 2., 3., 4.]
z = [1. 2. 1. 2.;  # shape is (nx, ny)
     1. 2. 1. 2.;
     1. 2. 3. 2.;
     1. 2. 2. 2.;
     1. 2. 1. 2.;
     1. 2. 3. 1.]
spl = Spline2D(x, y, z)

# element-wise output
xi = [1., 1.5, 2.3, 4.5, 3.3, 3.2, 3.]
yi = [1., 2.3, 5.3, 0.5, 3.3, 1.2, 3.]
ans = [2.94429906542,
       1.25537598131,
       2.00063588785,
       1.0,
       2.93952664,
       1.06482509358,
       3.0]
zi = evaluate(spl, xi, yi)
@test_approx_eq(zi, ans)

# test call on v0.4+
if VERSION >= v"0.4.0"
    zi = spl(xi, yi)
    @test_approx_eq(zi, ans)
end

# grid output
xi = [1., 1.5, 2.3, 4.5]
yi = [1., 2.3, 5.3]
ans = [2.94429906542  1.16946130841  1.99831775701;
       2.80393858478  1.25537598131  1.99873831776;
       1.67143209613  1.94853338542  2.00063588785;
       1.89392523364  1.8126946729  2.01042056075]
zi = evalgrid(spl, xi, yi)
@test_approx_eq(zi, ans)

println("All tests passed.")


abstract type AbstractDist{E} end

np_float(::AbstractDist{E}) where {E} = E
np_fixed(::AbstractDist{E}) where {E} = !E

include("dist_tanh.jl")

struct EqualSpacing{E} <: AbstractDist{E}
    np::Int
    dsmax::Float64
end

function EqualSpacing(np::Int=0; dsmax::Real=-1)
    E = np < 2
    if E
        dsmax <= 0 && "`dsmax`(> 0) must be specified if `np` < 2 (auto determin)"
    else
        dsmax <= 0 || @warn "`dsmax` is ignored since `np` has been specified"
    end
    return EqualSpacing{E}(np, dsmax)
end

isvalid(d::EqualSpacing{true}) = d.dsmax > 0
isvalid(d::EqualSpacing{false}) = d.np > 1

function nodes(d::EqualSpacing{false}; len::L=1.0) where {L<:Real}
    return collect(LinRange(0, len, d.np))
end

function nodes(::EqualSpacing, np::Int; len::L=1.0) where {L<:Real}
    return collect(LinRange(0, len, np))
end

function nodes(d::EqualSpacing{true}; len::L=1.0) where {L<:Real}
    np = ceil(Int, len / d.dsmax) + 1
    return nodes(d, np; len=len)
end


## ExoplanetsSysSim/src/transit_detection_model.jl
## (c) 2015 Eric B. Ford

# Several functions below based on https://github.com/christopherburke/KeplerPORTs/blob/master/KeplerPORTs_utils.py
# That follows the procedure outlined in Burke et al.(2015).
# However we don't currently interpolate the mesthreshold to the relevant duration
import SpecialFunctions.lgamma

function real_log_choose(m::Float64, n::Float64)::Float64
  lgamma(m+1)-lgamma(n+1)-lgamma(m-n+1.0)
end

function real_binom(k::Float64, BigM::Float64, f::Float64)::Float64
    F1 = real_log_choose(BigM,k)
    F2 = k*log(f)
    F3 = (BigM-k)*log(1.0-f)
    x = exp(F1+F2+F3)
    return x
end

"""
    kepler_window_function_binomial_model(exp_num_transits_no_gaps, duty_cycle; min_transits = 3.0)

Binomial window function model for use with Kepler taken from Burke et al. (2015).

# Arguments:
- `exp_num_transits_no_gaps::Float64`: Expected number of transits assuming no gaps in the observation
- `duty_cycle::Float64`: Kepler duty cycle (i.e. fraction of observation time with usable data)
- `min_transits::Float64 = 3.0`: Minimum number of observed transits required for consideration as planet candidate

# Returns:
Probability of detecting at least min_transits given the provided properties of the Kepler target star and planet.
"""
function kepler_window_function_binomial_model(exp_num_transits_no_gaps::Float64, duty_cycle::Float64; min_transits::Float64 = 3.0)::Float64
  if exp_num_transits_no_gaps < min_transits
     return 0.0
  else
     return max(1.0 - real_binom(min_transits,exp_num_transits_no_gaps,duty_cycle), 0.0)
  end
end

"""
    kepler_window_function_binomial_model(t, exp_num_transits_no_gaps, period, duration; min_transits = 3.0)

Binomial window function model for use with Kepler taken from Burke et al. (2015).

# Arguments:
- `t::KeplerTarget`: Kepler target object
- `exp_num_transits_no_gaps::Float64`: Expected number of transits assuming no gaps in the observation
- `period::Float64`: Orbital period of planet (in days)
- `duration::Float64`: Transit duration of planet (in days)
- `min_transits::Float64 = 3.0`: Minimum number of observed transits required for consideration as planet candidate

# Returns:
Probability of detecting at least min_transits given the provided properties of the Kepler target star and planet.
"""
function kepler_window_function_binomial_model(t::KeplerTarget, exp_num_transits_no_gaps::Float64, period::Float64, duration::Float64; min_transits::Float64 = 3.0)::Float64
   kepler_window_function_binomial_model(exp_num_transits_no_gaps, t.duty_cycle, min_transits=min_transits)
end

"""
    kepler_window_function_dr25_model(t, exp_num_transits_no_gaps, period, duration)

Window function model for use with Kepler that were fit per target for DR25 (see KSCI-19101-002).

# Arguments:
- `t::KeplerTarget`: Kepler target object
- `exp_num_transits_no_gaps::Float64`: Expected number of transits assuming no gaps in the observation
- `period::Float64`: Orbital period of planet (in days)
- `duration::Float64`: Transit duration of planet (in days)

# Returns:
Probability of detecting at least 3 transits given the provided properties of the Kepler target star and planet.
"""
function kepler_window_function_dr25_model(t::KeplerTarget, exp_num_transits_no_gaps::Float64, period::Float64, duration::Float64)::Float64
   ExoplanetsSysSim.WindowFunction.eval_window_function(t.window_function_id, Duration=duration, Period=period)
end

# WARNING: Hard-coded function variable for computational efficiency,
#          replace with appropriate window function model.
#kepler_window_function = kepler_window_function_binomial_model
kepler_window_function = kepler_window_function_dr25_model

"""
    frac_depth_to_tps_depth(frac_depth)

Converts fractional transit depth to depth used by Kepler Transiting Planet Search (TPS) module. (See Christiansen et al. (2015))
NOTE: This assumes a linear limbdarkening coefficient of 0.6

# Arguments:
- `frac_depth::Float64`: Fractional transit depth of planet at transit center.

# Returns:
Kepler TPS depth.
"""
function frac_depth_to_tps_depth(frac_depth::Float64)
    alp = 1.0874
    bet = 1.0187
    REALDEPTH2TPSSQUARE = 1.0  # WARNING: Waiting for this to be confirmed
    k = sqrt(frac_depth)
    tps_depth = min( (alp-bet*k) * frac_depth* REALDEPTH2TPSSQUARE, 1.0)   # NOTE: I added the max based on common sense
    return tps_depth::Float64
end

"""
    detection_efficiency_theory(mes, expected_num_transits; min_pdet_nonzero = 0.0)

Detection efficiency model assuming perfect theoretical 7.1-sigma error function response.

# Arguments:
- `mes::Float64`: Estimated multiple event statistic (or signal-to-noise ratio) for planet
- `expected_num_transits::Float64`: Expected number of planet transits (assuming no observation gaps)
- `min_pdet_nonzero::Float64 = 0.0`: Minimum probability of detection (if transiting) to be treated as identifiable

# Results:
Probability of detection (if transiting) for planet
"""
function detection_efficiency_theory(mes::Float64, expected_num_transits::Float64; min_pdet_nonzero::Float64 = 0.0)
   muoffset =  0.0
   sig =  1.0
   mesthresh = 7.1
   mes *= 1.003
   if mes > (9.0 - mesthresh - muoffset)
      return 1.0
   else
      pdet = 0.5 + 0.5*erf((mes - mesthresh - muoffset) / sqrt(2.0*sig*sig))
      pdet = pdet >= min_pdet_nonzero ? pdet : 0.0
      return pdet
   end
end

"""
    detection_efficiency_fressin2013(mes, expected_num_transits)

Detection efficiency model using the linear ramp from Fressin (2013).

# Arguments:
- `mes::Float64`: Estimated multiple event statistic (or signal-to-noise ratio) for planet
- `expected_num_transits::Float64`: Expected number of planet transits (assuming no observation gaps)

# Results:
Probability of detection (if transiting) for planet
"""
function detection_efficiency_fressin2013(mes::Float64, expected_num_transits::Float64)
    mesmin =  6.0
    mesmax =  16.0
   if mes <= mesmin
      return 0.0
   elseif mes >= mesmax
      return 1.0
   else
     return (mes - mesmin) / (mesmax - mesmin)
  end
end

"""
    detection_efficiency_christiansen2015(mes, expected_num_transits; mes_threshold = 7.1, min_pdet_nonzero = 0.0)

Detection efficiency model using a Gamma function taken from Christiansen (2015).

# Arguments:
- `mes::Float64`: Estimated multiple event statistic (or signal-to-noise ratio) for planet
- `expected_num_transits::Float64`: Expected number of planet transits (assuming no observation gaps)
- `mes_threshold::Float64 = 7.1`: Minimum multiple event statistic value for signal to have been considered a potential planet candidate signal
- `min_pdet_nonzero::Float64 = 0.0`: Minimum probability of detection (if transiting) to be treated as identifiable

# Results:
Probability of detection (if transiting) for planet
"""
function detection_efficiency_christiansen2015(mes::Float64, expected_num_transits::Float64; mes_threshold::Float64 = 7.1, min_pdet_nonzero::Float64 = 0.0)
    a =  4.65  # from code for detection_efficiency(...) at https://github.com/christopherburke/KeplerPORTs/blob/master/KeplerPORTs_utils.py
   # b =  1.05
   # a =  4.35 # from arxiv abstract. Informal testing showed it didn't matter
    b =  0.98
   mes *= 1.003
   usemes = max(0.0,mes - 7.1 - (mes_threshold - 7.1))
   pdet = cdf(Gamma(a,b), usemes)
   pdet = pdet >= min_pdet_nonzero ? pdet : 0.0
   return pdet
end

"""
    detection_efficiency_dr25_simple(mes, expected_num_transits; min_pdet_nonzero = 0.0)

Detection efficiency model using a Gamma function fit over a FGK sample of DR25 targets, taken from KSCI-19110-001.

# Arguments:
- `mes::Float64`: Estimated multiple event statistic (or signal-to-noise ratio) for planet
- `expected_num_transits::Float64`: Expected number of planet transits (assuming no observation gaps)
- `min_pdet_nonzero::Float64 = 0.0`: Minimum probability of detection (if transiting) to be treated as identifiable

# Results:
Probability of detection (if transiting) for planet
"""
function detection_efficiency_dr25_simple(mes::Float64, expected_num_transits::Float64; min_pdet_nonzero::Float64 = 0.0)::Float64
   a = 30.87  # from pg 16 of https://exoplanetarchive.ipac.caltech.edu/docs/KSCI-19110-001.pdf
   b = 0.271
   c = 0.940
   mes *= 1.003
   dist = Gamma(a,b)
   pdet::Float64 = c*cdf(dist, mes)::Float64
   pdet = pdet >= min_pdet_nonzero ? pdet : 0.0
   return pdet
end

"""
    get_param_for_detection_and_vetting_efficiency_depending_on_num_transits(num_tr)

Determine appropriate parameters to use in Gamma function model for detection efficiency and vetting of Kepler planet candidates.  Fit using the simulated transit injection tests on Kepler DR25 targets.

# Arguments:
- `num_tr::Integer`: (Expected) number of observed transits for planet

# Returns:
alpha, beta, and C, the three parameters to be used for the Gamma detection efficiency function
"""
function get_param_for_detection_and_vetting_efficiency_depending_on_num_transits(num_tr::Integer)
    if num_tr <= 3
        return (33.3884, 0.264472, 0.699093)
    elseif num_tr <= 4
        return  (32.886, 0.269577, 0.768366)
    elseif num_tr <= 5
        return (31.5196, 0.282741, 0.833673)
    elseif num_tr <= 6
        return (30.9919, 0.286979, 0.859865)
    elseif num_tr <= 9
        return (30.1906, 0.294688, 0.875042)
    elseif num_tr <= 18
        return (31.6342, 0.279425, 0.886144)
    elseif num_tr <= 36
        return (32.6448, 0.268898, 0.889724)
    else
        return (27.8185, 0.32432, 0.945075)
    end
end

"""
    detection_and_vetting_efficiency_model_v1(mes, expected_num_transits; min_pdet_nonzero = 0.0)

Detection efficiency model using a Gamma function whose parameters were determined using a fit over the simulated transit injection tests of a sample of FGK DR25 targets (see Hsu et al. (2019)).
# WARNING: Combined detection and vetting efficiency model - do NOT include additional vetting efficiency

# Arguments:
- `mes::Float64`: Estimated multiple event statistic (or signal-to-noise ratio) for planet
- `expected_num_transits::Float64`: Expected number of planet transits (assuming no observation gaps)
- `min_pdet_nonzero::Float64 = 0.0`: Minimum probability of detection (if transiting) to be treated as identifiable

# Returns:
Probability of detection (if transiting) for planet
"""
function detection_and_vetting_efficiency_model_v1(mes::Float64, expected_num_transits::Float64; min_pdet_nonzero::Float64 = 0.0)::Float64
    mes *= 1.003
    num_transit_int = convert(Int64,floor(expected_num_transits))
    num_transit_int += rand() < expected_num_transits-num_transit_int ? 1 : 0
    a, b, c = get_param_for_detection_and_vetting_efficiency_depending_on_num_transits(num_transit_int)
    dist = Gamma(a,b)
    pdet::Float64 = c*cdf(dist, mes)::Float64
    pdet = pdet >= min_pdet_nonzero ? pdet : 0.0
    return pdet
end

# WARNING: Hardcoded choice of transit detection efficiency here for speed and so as to not have it hardcoded in multiple places
#detection_efficiency_model = detection_efficiency_christiansen2015
# detection_efficiency_model = detection_efficiency_dr25_simple
detection_efficiency_model = detection_and_vetting_efficiency_model_v1

"""
    vetting_efficiency_none(R_p, P)

Vetting efficiency model assuming all identified planet candidate signals are true positives

# Arguments:
- `R_p::Real`: Planet radius
- `P::Real`: Planet orbital period

# Returns:
1
"""
function vetting_efficiency_none(R_p::Real, P::Real)
    return 1.0
end

"""
    vetting_efficiency_dr25_mulders(R_p, P)

Vetting efficiency model from Mulders et al. (2018) using parameters fit for Gaia DR2 stellar properties applied to Kepler DR25 targets independent of reliability score (through e-mail communication).

# Arguments:
- `R_p::Real`: Planet radius (in R_sol)
- `P::Real`: Planet orbital period (in days)

# Returns:
Probability of planet with given properties to be vetted as a planet candidate.
"""
function vetting_efficiency_dr25_mulders(R_p::Real, P::Real)
    c = 0.93
    a_R = -0.03
    P_break = 205.
    a_P = 0.00
    b_P = -0.24
    pvet = c*(R_p/earth_radius)^a_R
    if P < P_break
        pvet *= (P/P_break)^a_P
    else
        pvet *= (P/P_break)^b_P
    end
    return pvet
end

"""
    vetting_efficiency_dr25_mulders_score_cut(R_p, P)

Vetting efficiency model from Mulders et al. (2018) using parameters fit for Gaia DR2 stellar properties applied to Kepler DR25 targets selected via reliability score (arXiv 1805.08211).

# Arguments:
- `R_p::Real`: Planet radius (in R_sol)
- `P::Real`: Planet orbital period (in days)

# Returns:
Probability of planet with given properties to be vetted as a planet candidate.
"""
function vetting_efficiency_dr25_mulders_score_cut(R_p::Real, P::Real)
    c = 0.63
    a_R = 0.19
    P_break = 53.
    a_P = -0.07
    b_P = -0.39
    pvet = c*(R_p/earth_radius)^a_R
    if P < P_break
        pvet *= (P/P_break)^a_P
    else
        pvet *= (P/P_break)^b_P
    end
    return pvet
end

# WARNING: Hardcoded choice of planet vetting efficiency here for speed and so as to not have it hardcoded in multiple places
vetting_efficiency = vetting_efficiency_none

"""
    interpolate_cdpp_to_duration_use_target_cdpp(t, duration)

Linearly interpolated CDPP using list of provided CDPP at different transit durations stored in Kepler target object (DEPRECATED).

# Arguments:
- `t::KeplerTarget`: Kepler target object
- `duration::Float64`: Transit duration (in days)

# Results:
Estimated CDPP at given transit duration for Kepler target.
"""
function interpolate_cdpp_to_duration_use_target_cdpp(t::KeplerTarget, duration::Float64)::Float64
   duration_in_hours = duration *24.0
   dur_idx = searchsortedlast(cdpp_durations,duration_in_hours)   # cdpp_durations is defined in constants.jl
   if dur_idx <= 0
      cdpp = t.cdpp[1,1]
   elseif dur_idx==length(cdpp_durations) && (duration_in_hours >= cdpp_durations[end]) # Should be 15 cdpp_durations.
      cdpp = t.cdpp[length(cdpp_durations),1]
   else
      w = ((duration_in_hours)-cdpp_durations[dur_idx]) / (cdpp_durations[dur_idx+1]-cdpp_durations[dur_idx])
      cdpp = w*t.cdpp[dur_idx+1,1] + (1-w)*t.cdpp[dur_idx,1]
   end
   return cdpp
end

"""
    interpolate_cdpp_to_duration_lookup_cdpp(t, duration)

Linearly interpolated CDPP using list of provided CDPP at different transit durations stored in stellar catalog dataframe.

# Arguments:
- `t::KeplerTarget`: Kepler target object
- `duration::Float64`: Transit duration (in days)

# Results:
Estimated CDPP at given transit duration for Kepler target.
"""
function interpolate_cdpp_to_duration_lookup_cdpp(t::KeplerTarget, duration::Float64)::Float64
   duration_in_hours = duration *24.0
   dur_idx = searchsortedlast(cdpp_durations,duration_in_hours)   # cdpp_durations is defined in constants.jl
   get_cdpp(i::Integer) = 1.0e-6*sqrt(cdpp_durations[i]/24.0/LC_duration)*star_table(t,duration_symbols[i])

   if dur_idx <= 0
      cdpp = get_cdpp(1)
   elseif dur_idx==length(cdpp_durations) && (duration_in_hours >= cdpp_durations[end]) # Should be 15 cdpp_durations.
      cdpp = get_cdpp(length(cdpp_durations))
   else
      w = ((duration_in_hours)-cdpp_durations[dur_idx]) / (cdpp_durations[dur_idx+1]-cdpp_durations[dur_idx])
      cdpp = w*get_cdpp(dur_idx+1) + (1-w)*get_cdpp(dur_idx)
   end
   return cdpp
end

#interpolate_cdpp_to_duration = interpolate_cdpp_to_duration_use_target_cdpp
interpolate_cdpp_to_duration = interpolate_cdpp_to_duration_lookup_cdpp

"""
    calc_snr_if_transit_cdpp(t, depth, duration, cdpp, sim_param; num_transit = 1)

Calculate the expected multiple event statistic (signal-to-noise ratio) for planet around Kepler target star using CDPP.

# Arguments:
- `t::KeplerTarget`: Kepler target object
- `depth::Float64`: (Fractional) transit depth
- `duration::Float64`: Transit duration (in days)
- `cdpp::Float64`: CDPP for target star given transit duration
- `sim_param::SimParam`: Simulation parameters
- `num_transit::Float64 = 1`: Expected number of transits observed

# Returns:
(Expected) multiple event statistic (signal-to-noise ratio)
"""
function calc_snr_if_transit_cdpp(t::KeplerTarget, depth::Float64, duration::Float64, cdpp::Float64, sim_param::SimParam; num_transit::Float64 = 1)
  # depth_tps = frac_depth_to_tps_depth(depth)                  # TODO SCI:  WARNING: Hardcoded this conversion.  Remove once depth calculated using limb darkening model
  # snr = depth_tps*sqrt(num_transit*duration*LC_rate)/cdpp     # WARNING: Assumes measurement uncertainties are uncorrelated & CDPP based on LC
  snr = depth*sqrt(num_transit*duration*LC_rate)/cdpp     # WARNING: Assumes measurement uncertainties are uncorrelated & CDPP based on LC
end

"""
    calc_snr_if_transit(t, depth, duration, osd, sim_param; num_transit = 1)

Calculate the expected multiple event statistic (signal-to-noise ratio) for planet around Kepler target star using 1-sigma depth function (OSD).
# NOTE: Assumes OSD functions have already been read in for all relevant Kepler targets.

# Arguments:
- `t::KeplerTarget`: Kepler target object
- `depth::Real`: (Fractional) transit depth
- `duration::Real`: Transit duration (in days)
- `osd::Real`: OSD for target star given transit duration and period
- `sim_param::SimParam`: Simulation parameters
- `num_transit::Real = 1`: Expected number of transits observed

# Returns:
(Expected) multiple event statistic (signal-to-noise ratio)
"""
function calc_snr_if_transit(t::KeplerTarget, depth::Real, duration::Real, osd::Real, sim_param::SimParam; num_transit::Real = 1)
  # depth_tps = frac_depth_to_tps_depth(depth)                  # WARNING: Hardcoded this conversion
  # snr = depth_tps/osd*1.0e6 # osd is in ppm
  snr = depth/osd*1.0e6 # osd is in ppm
end

"""
    calc_snr_if_transit_central(t, s, p, sim_param)

Calculate the expected multiple event statistic (signal-to-noise ratio) for planet around Kepler target star using 1-sigma depth function (OSD) at transit center.
# NOTE: Assumes OSD functions have already been read in for all relevant Kepler targets.

# Arguments:
- `t::KeplerTarget`: Kepler target object
- `s::Integer`: Star index within Kepler target
- `p::Integer`: Planet index within Kepler target
- `sim_param::SimParam`: Simulation parameters

# Returns:
(Expected) multiple event statistic (signal-to-noise ratio) at transit center
"""
function calc_snr_if_transit_central(t::KeplerTarget, s::Integer, p::Integer, sim_param::SimParam)
  period = t.sys[s].orbit[p].P
  depth = calc_transit_depth(t,s,p)
  duration_central = calc_transit_duration_eff_central(t,s,p)
  kepid = StellarTable.star_table(t.sys[s].star.id, :kepid)
  osd_duration_central = get_durations_searched_Kepler(period,duration_central)	#tests if durations are included in Kepler's observations for a certain planet period. If not, returns nearest possible duration
  osd_central = WindowFunction.interp_OSD_from_table(kepid, period, osd_duration_central)
  if osd_duration_central > duration_central				#use a correcting factor if this duration is lower than the minimum searched for this period.
      osd_central = osd_central*osd_duration_central/duration_central
  end
  num_transit = calc_expected_num_transits(t,s,p,sim_param)
  calc_snr_if_transit(t,depth,duration_central,osd_central, sim_param,num_transit=num_transit)
end

"""
    calc_snr_if_transit_central_cdpp(t, s, p, sim_param)

Calculate the expected multiple event statistic (signal-to-noise ratio) for planet around Kepler target star using CDPP at transit center.

# Arguments:
- `t::KeplerTarget`: Kepler target object
- `s::Integer`: Star index within Kepler target
- `p::Integer`: Planet index within Kepler target
- `sim_param::SimParam`: Simulation parameters

# Returns:
(Expected) multiple event statistic (signal-to-noise ratio) at transit center
"""
function calc_snr_if_transit_central_cdpp(t::KeplerTarget, s::Integer, p::Integer, sim_param::SimParam)
  depth = calc_transit_depth(t,s,p)
  duration_central = calc_transit_duration_eff_central(t,s,p)
  cdpp = interpolate_cdpp_to_duration(t, duration_central)
  num_transit = calc_expected_num_transits(t,s,p,sim_param)
  calc_snr_if_transit_cdpp(t,depth,duration_central,cdpp, sim_param,num_transit=num_transit)
end

"""
    calc_prob_detect_if_transit(t, snr, period, duration, sim_param; num_transit = 1)
    calc_prob_detect_if_transit(t, depth, period, duration, osd, sim_param; num_transit = 1)

Calculate probability of detecting planet signal (if planet transits) using 1-sigma depth (OSD) function.
# NOTE: Assumes OSD functions already read in.

# Arguments:
- `t::KeplerTarget`: Kepler target object
- `snr::Float64`: Expected multiple event statistic (signal-to-noise ratio)
- `depth::Float64`: (Fractional) transit depth
- `period::Float64`: Orbital period (in days)
- `duration::Float64`: Transit duration (in days)
- `osd::Float64`: OSD for target star given transit duration and period
- `sim_param::SimParam`: Simulation parameters
- `num_transit::Float64 = 1`: Expected number of transits observed

# Returns:
Probability of detecting planet (if it transits)
"""
function calc_prob_detect_if_transit(t::KeplerTarget, snr::Float64, period::Float64, duration::Float64, sim_param::SimParam; num_transit::Float64 = 1)
  min_pdet_nonzero = 1.0e-4                                                # TODO OPT: Consider raising threshold to prevent a plethora of planets that are very unlikely to be detected due to using 0.0 or other small value here
  wf = kepler_window_function(t, num_transit, period, duration)
  return wf*detection_efficiency_model(snr, num_transit, min_pdet_nonzero=min_pdet_nonzero)
end

function calc_prob_detect_if_transit(t::KeplerTarget, depth::Float64, period::Float64, duration::Float64, osd::Float64, sim_param::SimParam; num_transit::Float64 = 1)
  snr = calc_snr_if_transit(t,depth,duration,osd, sim_param, num_transit=num_transit)
  return calc_prob_detect_if_transit(t, snr, period, duration, sim_param, num_transit=num_transit)
end

"""
    calc_prob_detect_if_transit_cdpp(t, depth, period, duration, cdpp, sim_param; num_transit = 1)

Calculate probability of detecting planet signal (if planet transits) using CDPP.

# Arguments:
- `t::KeplerTarget`: Kepler target object
- `depth::Float64`: (Fractional) transit depth
- `period::Float64`: Orbital period (in days)
- `duration::Float64`: Transit duration (in days)
- `cdpp::Float64`: CDPP for target star given transit duration
- `sim_param::SimParam`: Simulation parameters
- `num_transit::Float64 = 1`: Expected number of transits observed

# Returns:
Probability of detecting planet (if it transits)
"""
function calc_prob_detect_if_transit_cdpp(t::KeplerTarget, depth::Float64, period::Float64, duration::Float64, cdpp::Float64, sim_param::SimParam; num_transit::Float64 = 1)
  snr = calc_snr_if_transit_cdpp(t,depth,duration,cdpp, sim_param, num_transit=num_transit)
  return calc_prob_detect_if_transit(t, snr, period, duration, sim_param, num_transit=num_transit)
end

"""
    calc_prob_detect_if_transit_central(t, s, p, sim_param)

Calculate probability of detecting planet signal (if planet transits) at transit center using 1-sigma depth (OSD) functions.
# NOTE: Assumes OSD functions already read in.

# Arguments:
- `t::KeplerTarget`: Kepler target object
- `s::Integer`: Star index within Kepler target
- `p::Integer`: Planet index within Kepler target
- `sim_param::SimParam`: Simulation parameters

# Returns:
Probability of detecting planet (if it transits) at transit center
"""
function calc_prob_detect_if_transit_central(t::KeplerTarget, s::Integer, p::Integer, sim_param::SimParam)
  period = t.sys[s].orbit[p].P
  depth = calc_transit_depth(t,s,p)
  duration_central = calc_transit_duration_eff_central(t,s,p)
  kepid = StellarTable.star_table(t.sys[s].star.id, :kepid)
  osd_duration_central = get_durations_searched_Kepler(period,duration_central)	#tests if durations are included in Kepler's observations for a certain planet period. If not, returns nearest possible duration
  osd_central = WindowFunction.interp_OSD_from_table(kepid, period, osd_duration_central)
  if osd_duration_central > duration_central				#use a correcting factor if this duration is lower than the minimum searched for this period.
      osd_central = osd_central*osd_duration_central/duration_central
  end
  ntr = calc_expected_num_transits(t,s,p,sim_param)
  calc_prob_detect_if_transit(t,depth,period,duration_central,osd_central, sim_param, num_transit=ntr)
end

"""
    calc_prob_detect_if_transit_central_cdpp(t, s, p, sim_param)

Calculate probability of detecting planet signal (if planet transits) at transit center using CDPP

# Arguments:
- `t::KeplerTarget`: Kepler target object
- `s::Integer`: Star index within Kepler target
- `p::Integer`: Planet index within Kepler target
- `sim_param::SimParam`: Simulation parameters

# Returns:
Probability of detecting planet (if it transits) at transit center
"""
function calc_prob_detect_if_transit_central_cdpp(t::KeplerTarget, s::Integer, p::Integer, sim_param::SimParam)
  period = t.sys[s].orbit[p].P
  depth = calc_transit_depth(t,s,p)
  duration_central = calc_transit_duration_eff_central(t,s,p)
  cdpp = interpolate_cdpp_to_duration(t, duration_central)
  ntr = calc_expected_num_transits(t,s,p,sim_param)
  calc_prob_detect_if_transit_cdpp(t,depth,period,duration_central,cdpp, sim_param, num_transit=ntr)
end

"""
    calc_prob_detect_if_transit_with_actual_b(t, s, p, sim_param)

Calculate probability of detecting planet signal (if planet transits) at transit center using 1-sigma depth (OSD) functions and the impact parameter of the provided orbit.
# NOTE: Assumes OSD functions already read in.

# Arguments:
- `t::KeplerTarget`: Kepler target object
- `s::Integer`: Star index within Kepler target
- `p::Integer`: Planet index within Kepler target
- `sim_param::SimParam`: Simulation parameters

# Returns:
Probability of detecting planet (if it transits) given the impact parameter of the planet's orbit
"""
function calc_prob_detect_if_transit_with_actual_b(t::KeplerTarget, s::Integer, p::Integer, sim_param::SimParam)
  period = t.sys[s].orbit[p].P
  size_ratio = t.sys[s].planet[p].radius/t.sys[s].star.radius
  depth = calc_transit_depth(t,s,p)
  duration = calc_transit_duration_eff(t,s,p)
  b = calc_impact_parameter(t.sys[s],p)
  snr_correction = calc_depth_correction_for_grazing_transit(b,size_ratio)
  depth *= snr_correction
  kepid = StellarTable.star_table(t.sys[s].star.id, :kepid)
  osd_duration = get_durations_searched_Kepler(period,duration)	#tests if durations are included in Kepler's observations for a certain planet period. If not, returns nearest possible duration
  osd = WindowFunction.interp_OSD_from_table(kepid, period, osd_duration)
  if osd_duration > duration				#use a correcting factor if this duration is lower than the minimum searched for this period.
      osd = osd*osd_duration/duration
  end
  ntr = calc_expected_num_transits(t,s,p,sim_param)
  calc_prob_detect_if_transit(t,depth,period,duration,osd, sim_param, num_transit=ntr)
end

"""
    calc_prob_detect_if_transit_with_actual_b_cdpp(t, s, p, sim_param)

Calculate probability of detecting planet signal (if planet transits) at transit center using CDPP and the impact parameter of the provided orbit.

# Arguments:
- `t::KeplerTarget`: Kepler target object
- `s::Integer`: Star index within Kepler target
- `p::Integer`: Planet index within Kepler target
- `sim_param::SimParam`: Simulation parameters

# Returns:
Probability of detecting planet (if it transits) given the impact parameter of the planet's orbit
"""
function calc_prob_detect_if_transit_with_actual_b_cdpp(t::KeplerTarget, s::Integer, p::Integer, sim_param::SimParam)
  period = t.sys[s].orbit[p].P
  size_ratio = t.sys[s].planet[p].radius/t.sys[s].star.radius
  depth = calc_transit_depth(t,s,p)
  duration = calc_transit_duration_eff(t,s,p)
  b = calc_impact_parameter(t.sys[s],p)
  snr_correction = calc_depth_correction_for_grazing_transit(b,size_ratio)
  depth *= snr_correction
  cdpp = interpolate_cdpp_to_duration(t, duration)
  ntr = calc_expected_num_transits(t,s,p,sim_param)
  calc_prob_detect_if_transit_cdpp(t,depth,period,duration,cdpp, sim_param, num_transit=ntr)
end

"""
    calc_ave_prob_detect_if_transit_from_snr(t, snr_central, period, duration_central, size_ratio, osd_central, sim_param; num_transit = 1)

Calculate probability of detecting planet signal (if planet transits) using 1-sigma depth (OSD) function and averaged over impact parameters b~U[0,1).
# NOTE: Assumes OSD functions already read in.

# Arguments:
- `t::KeplerTarget`: Kepler target object
- `snr_central::Float64`: Expected multiple event statistic (signal-to-noise ratio) at transit center
- `period::Float64`: Orbital period (in days)
- `duration_central::Float64`: Transit duration (in days) at transit center
- `size_ratio::Float64`: Ratio of planet-to-star radii
- `osd_central::Float64`: OSD for target star given transit duration at transit center and period
- `sim_param::SimParam`: Simulation parameters
- `num_transit::Float64 = 1`: Expected number of transits observed

# Returns:
Probability of detecting planet (if it transits) averaged over impact parameter
"""
function calc_ave_prob_detect_if_transit_from_snr(t::KeplerTarget, snr_central::Float64, period::Float64, duration_central::Float64, size_ratio::Float64, osd_central::Float64, sim_param::SimParam; num_transit::Float64 = 1)
  min_pdet_nonzero = 1.0e-4
  wf = kepler_window_function(t, num_transit, period, duration_central)

  detection_efficiency_central = detection_efficiency_model(snr_central, num_transit, min_pdet_nonzero=min_pdet_nonzero)
  if wf*detection_efficiency_central <= min_pdet_nonzero
     return 0.
  end

  # Breaking integral into two sections [0,1-b_boundary) and [1-b_boundary,1], so need at least 5 points to evaluate integral via trapezoid rule
  num_impact_param_low_b =  7                            # Number of points to evaluate integral over [0,1-b_boundary) via trapezoid rule
  num_impact_param_high_b = 5 # (size_ratio<=0.05) ? 5 : 11  # Number of points to evaluate integral over [1-b_boudnary,1) via trapezoid rule.  If using 2*size_ratio for bondary for small planets, then keep this odd, so one point lands on 1-size_ratio.
  @assert(num_impact_param_low_b >= 5)
  @assert(num_impact_param_high_b >= 3)
  num_impact_param = num_impact_param_low_b+num_impact_param_high_b-1 # One point is shared
  b_boundary = (size_ratio <= 0.15) ? 2*size_ratio : min(max(0.3,size_ratio),0.5)
  b = Array{Float64}(undef,num_impact_param)
  weight = Array{Float64}(undef,num_impact_param)
  b[1:num_impact_param_low_b] = range(0.0,stop=1-b_boundary,length=num_impact_param_low_b)
  b[num_impact_param_low_b:num_impact_param] .= range(1-b_boundary,stop=1.0,length=num_impact_param_high_b)
  weight[1:num_impact_param_low_b] .= (1-b_boundary)/(num_impact_param_low_b-1)  # Points for first integral
  weight[1] *= 0.5                        # Lower endpoint of first integral
  weight[num_impact_param_low_b] *= 0.5   # Upper endpoint of first integral
  weight[num_impact_param_low_b] += 0.5*(b_boundary)/(num_impact_param_high_b-1) # Also lower endpoint of second integral
  weight[(num_impact_param_low_b+1):num_impact_param] .= b_boundary/(num_impact_param_high_b-1)
  weight[num_impact_param] *= 0.5         # Upper endpoint of second integral
  @assert isapprox(sum(weight),1.0)

  function integrand(b::Float64)::Float64
     depth_factor = calc_depth_correction_for_grazing_transit(b,size_ratio)
     duration_factor = calc_transit_duration_eff_factor_for_impact_parameter_b(b,size_ratio)
     kepid = StellarTable.star_table(t.sys[1].star.id, :kepid)
     osd_duration = get_durations_searched_Kepler(period,duration_central*duration_factor)	#tests if durations are included in Kepler's observations for a certain planet period. If not, returns nearest possible duration
     osd = WindowFunction.interp_OSD_from_table(kepid, period, osd_duration)
     if osd_duration > duration_central*duration_factor				#use a correcting factor if this duration is lower than the minimum searched for this period.
	osd = osd*osd_duration/(duration_central*duration_factor)
     end
     snr_factor = depth_factor*(osd_central/osd)
     detection_efficiency_model(snr_central*snr_factor, num_transit, min_pdet_nonzero=min_pdet_nonzero)
  end

  ave_detection_efficiency = sum(weight .* map(integrand,b)::Vector{Float64} )

  return wf*ave_detection_efficiency
end

"""
    calc_ave_prob_detect_if_transit_from_snr_cdpp(t, snr_central, period, duration_central, size_ratio, osd_central, sim_param; num_transit = 1)

Calculate probability of detecting planet signal (if planet transits) using CDPP and averaged over impact parameters b~U[0,1).

# Arguments:
- `t::KeplerTarget`: Kepler target object
- `snr_central::Float64`: Expected multiple event statistic (signal-to-noise ratio) at transit center
- `period::Float64`: Orbital period (in days)
- `duration_central::Float64`: Transit duration (in days) at transit center
- `size_ratio::Float64`: Ratio of planet-to-star radii
- `cdpp_central::Float64`: CDPP for target star given transit duration at transit center
- `sim_param::SimParam`: Simulation parameters
- `num_transit::Float64 = 1`: Expected number of transits observed

# Returns:
Probability of detecting planet (if it transits) averaged over impact parameter
"""
function calc_ave_prob_detect_if_transit_from_snr_cdpp(t::KeplerTarget, snr_central::Float64, period::Float64, duration_central::Float64, size_ratio::Float64, cdpp_central::Float64, sim_param::SimParam; num_transit::Float64 = 1)
  min_pdet_nonzero = 1.0e-4
  wf = kepler_window_function(t, num_transit, period, duration_central)

  detection_efficiency_central = detection_efficiency_model(snr_central, num_transit, min_pdet_nonzero=min_pdet_nonzero)
  if wf*detection_efficiency_central <= min_pdet_nonzero
     return 0.
  end

  # Breaking integral into two sections [0,1-b_boundary) and [1-b_boundary,1], so need at least 5 points to evaluate integral via trapezoid rule
  num_impact_param_low_b =  7                            # Number of points to evaluate integral over [0,1-b_boundary) via trapezoid rule
  num_impact_param_high_b = 5 # (size_ratio<=0.05) ? 5 : 11  # Number of points to evaluate integral over [1-b_boudnary,1) via trapezoid rule.  If using 2*size_ratio for bondary for small planets, then keep this odd, so one point lands on 1-size_ratio.
  @assert(num_impact_param_low_b >= 5)
  @assert(num_impact_param_high_b >= 3)
  num_impact_param = num_impact_param_low_b+num_impact_param_high_b-1 # One point is shared
  b_boundary = (size_ratio <= 0.15) ? 2*size_ratio : min(max(0.3,size_ratio),0.5)
  b = Array{Float64}(undef,num_impact_param)
  weight = Array{Float64}(undef,num_impact_param)
  b[1:num_impact_param_low_b] = range(0.0,stop=1-b_boundary,length=num_impact_param_low_b)
  b[num_impact_param_low_b:num_impact_param] .= range(1-b_boundary,stop=1.0,length=num_impact_param_high_b)
  weight[1:num_impact_param_low_b] .= (1-b_boundary)/(num_impact_param_low_b-1)  # Points for first integral
  weight[1] *= 0.5                        # Lower endpoint of first integral
  weight[num_impact_param_low_b] *= 0.5   # Upper endpoint of first integral
  weight[num_impact_param_low_b] += 0.5*(b_boundary)/(num_impact_param_high_b-1) # Also lower endpoint of second integral
  weight[(num_impact_param_low_b+1):num_impact_param] .= b_boundary/(num_impact_param_high_b-1)
  weight[num_impact_param] *= 0.5         # Upper endpoint of second integral
  @assert isapprox(sum(weight),1.0)

  function integrand(b::Float64)::Float64
     depth_factor = calc_depth_correction_for_grazing_transit(b,size_ratio)
     duration_factor = calc_transit_duration_eff_factor_for_impact_parameter_b(b,size_ratio)
     cdpp = interpolate_cdpp_to_duration(t,duration_central*duration_factor)
     snr_factor = depth_factor*sqrt(duration_factor)*(cdpp_central/cdpp)
     detection_efficiency_model(snr_central*snr_factor, num_transit, min_pdet_nonzero=min_pdet_nonzero)
  end

  ave_detection_efficiency = sum(weight .* map(integrand,b)::Vector{Float64} )

  return wf*ave_detection_efficiency
end

"""
    calc_ave_prob_detect_if_transit_cdpp(t, depth, period, duration_central, size_ratio, sim_param; num_transit = 1)
    calc_ave_prob_detect_if_transit_cdpp(t, s, p, sim_param)

Calculate probability of detecting planet signal (if planet transits) using CDPP and averaged over impact parameters b~U[0,1).

# Arguments:
- `t::KeplerTarget`: Kepler target object
- `depth::Float64`: (Fractional) transit depth
- `period::Float64`: Orbital period (in days)
- `duration_central::Float64`: Transit duration (in days) at transit center
- `size_ratio::Float64`: Ratio of planet-to-star radii
- `s::Integer`: Star index within Kepler target
- `p::Integer`: Planet index within Kepler target
- `sim_param::SimParam`: Simulation parameters
- `num_transit::Float64 = 1`: Expected number of transits observed

# Returns:
Probability of detecting planet (if it transits) averaged over impact parameter
"""
function calc_ave_prob_detect_if_transit_cdpp(t::KeplerTarget, depth::Float64, period::Float64, duration_central::Float64, size_ratio::Float64, sim_param::SimParam; num_transit::Float64 = 1)
  cdpp_central = interpolate_cdpp_to_duration(t, duration_central)
  snr_central = calc_snr_if_transit(t,depth,duration_central,cdpp_central, sim_param, num_transit=num_transit)
  return calc_ave_prob_detect_if_transit_from_snr_cdpp(t, snr_central, period, duration_central, size_ratio, cdpp_central, sim_param, num_transit=num_transit)
end

function calc_ave_prob_detect_if_transit_cdpp(t::KeplerTarget, s::Integer, p::Integer, sim_param::SimParam)
  size_ratio = t.sys[s].planet[p].radius/t.sys[s].star.radius
  depth = calc_transit_depth(t,s,p)
  period = t.sys[s].orbit[p].P
  duration_central = calc_transit_duration_eff_central(t,s,p)
  ntr = calc_expected_num_transits(t,s,p,sim_param)
  calc_ave_prob_detect_if_transit_cdpp(t,depth,period,duration_central, size_ratio, sim_param, num_transit=ntr)
end


function free_energy(xs::Vector)
  I1 = pcubature(θs -> ρ(θs, xs)*sin(θs[2])*u(θs),
                 _LOWER_BOUNDS_, _UPPER_BOUNDS_; reltol=RELTOL_INT, maxevals=MAXEVALS_INT);
  I2 = pcubature(θs -> ρ(θs, xs)*sin(θs[2])*log(ρ(θs, xs)),
                 _LOWER_BOUNDS_, _UPPER_BOUNDS_; reltol=RELTOL_INT, maxevals=MAXEVALS_INT);
  return I1[1] + kT*I2[1] - kT*N*log(N);
end

function free_energy_smallomega(x::Vector)
  c, λ, α = x[:];
  I1 = pcubature(θs -> c*(1 - ω(θs) + α*cos(θs[1])*sin(θs[2]))*exp(λ*cos(θs[2]))*sin(θs[2])*u(θs),
                 _LOWER_BOUNDS_, _UPPER_BOUNDS_; reltol=RELTOL_INT, maxevals=MAXEVALS_INT);
  I2 = pcubature(θs -> c*(1 - ω(θs) + α*cos(θs[1])*sin(θs[2]))*exp(λ*cos(θs[2]))*sin(θs[2])*log(ρ(θs, xs)),
                 _LOWER_BOUNDS_, _UPPER_BOUNDS_; reltol=RELTOL_INT, maxevals=MAXEVALS_INT);
  return I1[1] + kT*I2[1] - kT*N*log(N);
end

function free_energy_smalllambda(x::Vector)
  c, λ, α = x[:];
  I1 = pcubature(θs -> c*(1 + λ*cos(θs[2]) + α*cos(θs[1])*sin(θs[2]))*exp(-ω(θs))*sin(θs[2])*u(θs),
                 _LOWER_BOUNDS_, _UPPER_BOUNDS_; reltol=RELTOL_INT, maxevals=MAXEVALS_INT);
                 I2 = pcubature(θs -> c*(1 + λ*cos(θs[2]) + α*cos(θs[1])*sin(θs[2]))*exp(-ω(θs))*sin(θs[2])*log(ρ(θs, xs)),
                 _LOWER_BOUNDS_, _UPPER_BOUNDS_; reltol=RELTOL_INT, maxevals=MAXEVALS_INT);
  return I1[1] + kT*I2[1] - kT*N*log(N);
end


using DifferentialEquations, ModelingToolkit
using HoJMTK.Mechanical1D

k = 1.0
m = 1.0
b = 1.0
v = 1.0
@named spring = Spring(k=k)
@named mass = Mass(m=m)
@named damper = Damper(b=b)
@named source = ConstantVelocity(v=v)
@named fixed = Fixed()

# mech1d_eqs = [
#           connect(source.p1, damper.p1)
#           connect(damper.p2, mass.p2)
#           connect(mass.p2, source.p2, fixed.g)
#          ]
mech1d_eqs = [
          connect(source.p1, mass.p1)
          connect(mass.p2, damper.p1, spring.p1)
          connect(damper.p2, spring.p2, fixed.g)
          connect(source.p2, fixed.g)
         ]

@variables t
@named _mech1d_model = ODESystem(mech1d_eqs, t)
@named mech1d_model = compose(_mech1d_model, [spring, mass, damper, source, fixed])
sys = structural_simplify(mech1d_model)
u0 = [
    mass.p1 => 0.0
     ]
prob = ODAEProblem(sys, u0, (0, 10.0))
sol = solve(prob, Tsit5())

using Plots
plot(sol)


const SIDE_COEF_SQUARE = 1 / 2
const SIDE_COEF_CIRCLE = 1 / sqrt(pi)
const SIDE_COEF_DIAMOND = sqrt(2) / 2
const SIDE_COEF_TRIANGLE = sqrt(2 / (3 * sqrt(3 / 2)))
const SIDE_COEF_PENTAGON = sqrt(2 / (5 * sin(1.2566)))
const SIDE_COEF_HEXAGON = sqrt(2 / (6 * sin(1.0472)))
const SIDE_COEF_OCTAGON = sqrt(2 / (8 * sin(0.7854)))

function _polygon_radius(a::Real, n::Integer, r::Real, start::Real=-pi / 2)
    a -= start
    while true
        if a > 2pi / n
            a -= 2pi / n
        elseif a < 0
            a += 2pi / n
        else
            break
        end
    end
    return r * cos(pi / n) / cos(pi / n - a)
end

"""
    node_radius(::Type{Val{:shape}}, s::Real, a::Real)

Get the distance between the center and the edge of a node at angle `a`.
The node is of shape `:shape` and size `s`.
"""
node_radius(_, s::Real, a::Real) = s / 2

node_radius(::Type{Val{:square}}, s::Real, a::Real) =
    s * SIDE_COEF_SQUARE / max(abs(sin(a)), abs(cos(a)))

node_radius(::Type{Val{:circle}}, s::Real, a::Real) = s * SIDE_COEF_CIRCLE

node_radius(::Type{Val{:diamond}}, s::Real, a::Real) =
    s * SIDE_COEF_DIAMOND / (1 + abs(sin(a) / cos(a))) / abs(cos(a))

node_radius(::Type{Val{:triangle}}, s::Real, a::Real) =
    _polygon_radius(a, 3, s * SIDE_COEF_TRIANGLE)

node_radius(::Type{Val{:pentagon}}, s::Real, a::Real) =
    _polygon_radius(a, 5, s * SIDE_COEF_PENTAGON)

node_radius(::Type{Val{:hexagon}}, s::Real, a::Real) =
    _polygon_radius(a, 6, s * SIDE_COEF_HEXAGON)

node_radius(::Type{Val{:octagon}}, s::Real, a::Real) =
    _polygon_radius(a, 8, s * SIDE_COEF_OCTAGON, pi / 8)

function _polygon_path(context::CairoContext, x, y, n, r, start=-pi / 2)
    a = start
    move_to(context, x + r * cos(a), y + r * sin(a))
    for i = 1:n - 1
        a += 2pi / n
        line_to(context, x + r * cos(a), y + r * sin(a))
    end
    close_path(context)
end

"""
    outline_node!(::Type{Val{:shape}}, context, x, y, s)

Create a path on a Cairo `context` that corresponds to `:shape` of size `s`.
"""
outline_node!(_, context, x, y, r) = nothing

outline_node!(::Type{Val{:circle}}, context, x, y, r) = circle(context, x, y, r * SIDE_COEF_CIRCLE)

outline_node!(::Type{Val{:square}}, context, x, y, r) = rectangle(context,
    x - r * SIDE_COEF_SQUARE, y - r * SIDE_COEF_SQUARE, r, r)

outline_node!(::Type{Val{:diamond}}, context, x, y, side) = begin
    move_to(context, x - side * SIDE_COEF_DIAMOND, y)
    line_to(context, x, y - side * SIDE_COEF_DIAMOND)
    line_to(context, x + side * SIDE_COEF_DIAMOND, y)
    line_to(context, x, y + side * SIDE_COEF_DIAMOND)
    close_path(context)
end

outline_node!(::Type{Val{:triangle}}, context, x, y, side) =
    _polygon_path(context, x, y, 3, side * SIDE_COEF_TRIANGLE)

outline_node!(::Type{Val{:pentagon}}, context, x, y, side) =
    _polygon_path(context, x, y, 5, side * SIDE_COEF_PENTAGON)

outline_node!(::Type{Val{:hexagon}}, context, x, y, side) =
    _polygon_path(context, x, y, 6, side * SIDE_COEF_HEXAGON)

outline_node!(::Type{Val{:octagon}}, context, x, y, side) =
    _polygon_path(context, x, y, 8, side * SIDE_COEF_OCTAGON, pi / 8)


using SIAN

println("Setting up the problem")

ode = @ODEmodel(
  x'(t) = lm - d * x(t) - beta * x(t) * v(t),
  y'(t) = beta * x(t) * v(t) - a * y(t),
  v'(t) = k * y(t) - u * v(t),
  w'(t) = c * z(t) * y(t) * w(t) - c * q * y(t) * w(t) - b * w(t),
  z'(t) = c * q * y(t) * w(t) - h * z(t),
  y1(t) = w(t),
  y2(t) = z(t)
)

res = identifiability_ode(ode, get_parameters(ode); p=0.99, p_mod=2^29 - 3, nthrds=1)

println(res)


# # Arbitrary Precision
#
# COSMO allows you to solve problems with arbitrary floating-point precision, e.g. by using `BigFloat` problem data. To do this, the desired floating point type has to be consistent across the model `COSMO.Model{<: AbstractFloat}`, the input data and the (optional) settings object `COSMO.Settings{<: AbstractFloat}`.
# As an example, assume we want to solve the following quadratic program:
#
# $$
# \begin{array}{ll} \text{minimize} &  1/2 x^\top P x + q^\top x \\
# \text{subject to} &  l \leq A x \leq u
# \end{array}
# $$
# where $P = \begin{bmatrix} 4 & 1 \\ 1 & 2\end{bmatrix}$, $q = [1, 1]^\top$, $A = \begin{bmatrix} 1 & 1 \\ 1 & 0 \\ 0 & 1\end{bmatrix}$ and $l= [1,0, 0]^\top$, $u=[1, 0.7, 0.7]^\top$. We start by creating the model with the desired precision:
using COSMO, LinearAlgebra, SparseArrays
model = COSMO.Model{BigFloat}()

#-
# Next, we define the problem data as `BigFloat` arrays and create the constraint:
q = BigFloat[1; 1.]
P = sparse(BigFloat[4. 1; 1 2])
A = BigFloat[1. 1; 1 0; 0 1]
l = BigFloat[1.; 0; 0]
u = BigFloat[1; 0.7; 0.7]

constraint = COSMO.Constraint(A, zeros(BigFloat, 3), COSMO.Box(l, u))

# Notice that the constraint type parameter is dependent on the input data. The same is true for the constraint set `Box`. Next, we define the settings
settings = COSMO.Settings{BigFloat}(verbose = true, kkt_solver = QdldlKKTSolver)

# and assemble and solve the problem:
assemble!(model, P, q, constraint, settings = settings)
result = COSMO.optimize!(model);

# Moreover, notice that when no type parameter is specified, all objects default to `Float64`:
model = COSMO.Model()

# Two limitations to arbitrary precision:
# - Since we call the LAPACK function `syevr` for eigenvalue decompositions, we currently only support solving problems with PSD constraints in `Float32` and `Float64`.
# - We suggest to use the pure Julia QDLDL linear system solver (`kkt_solver = QdldlKKTSolver`) when working with arbitrary precision types as some of the other available solvers don't support all available precisions.

#md # !!! note
#md #     `JuMP` does not currently support arbitrary precision. However, if you want to use `COSMO` directly with `MathOptInterface`, you can use: `COSMO.Optimizer{<: AbstractFloat}` as your optimizer. Again, the problem data precision of your MathOptInterface-model has to agree with the optimizer's precision.


#=
# 4. Multiple stationary bodies
Adding multiple bodies to a problem is easy, using the concepts of a `BodyList`
and `RigidTransformList`.
=#
using ViscousFlow
#-
using Plots

#=
In this example, we will set up a problem with three cylinders arranged in a
formation in a free stream.
=#
Re = 200 # Reynolds number
U = 1.0 # Free stream velocity
U∞ = (U,0.0);

#=
Set up the domain, grid spacing, and time step size
=#
xlim = (-2.0,4.0)
ylim = (-2.0,2.0)
Δx, Δt = setstepsizes(Re,gridRe=4)

#=
### Set up bodies
We start by initializing a `BodyList` and an associated `RigidTransformList`.
Each member of the `RigidTransformList` will be used to place the respective body in
the correct position and orientation.
=#
bl = BodyList()
tl = RigidTransformList()

# Place the first body at (-1,0)
push!(bl,Circle(0.5,1.5Δx))
push!(tl,RigidTransform((-1.,0.),0.));

# Place the second body at (1,-1)
push!(bl,Circle(0.5,1.5Δx))
push!(tl,RigidTransform((1.,-1.),0.));

# and place the third body at (1,1)
push!(bl,Circle(0.5,1.5Δx))
push!(tl,RigidTransform((1.,1.),0.))

# Perform the actual transformation. Note that this operation works `in-place`:
tl(bl)

# #### Plot the initial configuration of the bodies
# Just to check they are in the right places
plot(bl,xlim=xlim,ylim=ylim)

#=
### Construct the system structure
We construct the system with the same syntax as for a single body:
=#
sys = NavierStokes(Re,Δx,xlim,ylim,Δt,bl,freestream = U∞)
#-
u0 = newstate(sys)
tspan = (0.0,10.0)
integrator = init(u0,tspan,sys)
#=
### Solve
Here, we run it for only a little while, just to demonstrate:
=#
@time step!(integrator,0.5)
#=
### Examine
Let's make an animation
=#
sol = integrator.sol;
@gif for (u,t) in zip(sol.u,sol.t)
    plot(vorticity(u,sys,t),sys,clim=(-10,10),levels=range(-10,10,length=30), color = :RdBu)
end every 5

# Now we will examine the force on each body
fx1, fy1 = force(sol,sys,1)
fx2, fy2 = force(sol,sys,2)
fx3, fy3 = force(sol,sys,3);
#-
plt = plot(layout = (2,1), size = (600, 400))
plot!(plt[1],sol.t,2*fx1,xlim=(0,20),ylim=(0,4),xlabel="Convective time",ylabel="\$C_D\$",label="Lead body",title="Drag force")
plot!(plt[2],sol.t,2*fy1,xlim=(0,20),ylim=(-2,2),xlabel="Convective time",ylabel="\$C_L\$",label="Lead body",title="Side force")
plot!(plt[1],sol.t,2*fx2,xlim=(0,20),ylim=(0,4),xlabel="Convective time",ylabel="\$C_D\$",label="Trailing body",title="Drag force")
plot!(plt[2],sol.t,2*fy2,xlim=(0,20),ylim=(-2,2),xlabel="Convective time",ylabel="\$C_L\$",label="Trailing body",title="Side force")
#-
println("Mean drag coefficient on lead body = ", GridUtilities.mean(2*fx1))
#-
println("Mean drag coefficient on trailing body = ", GridUtilities.mean(2*fx2))


# include("./utils.jl")

abstract type Smoother end
mutable struct GaussianProcessSmoother <: Smoother
    smoother_function::Function
    kernel
    mean_function
    GaussianProcessSmoother(x, y, z) = GaussianProcessSmoother(x, y, z)
    GaussianProcessSmoother() = new()
    function GaussianProcessSmoother(mean_function, kernel)
        gp = new()
        gp.smoother_function = GP
        gp.mean_function = mean_function
        gp.kernel = kernel
        return gp
    end
end

function fit(smoother::GaussianProcessSmoother, x, y)
    model = smoother.smoother_function(
                        x,
                        y,
                        smoother.mean_function,
                        smoother.kernel
                        )
    return model
end

function predict(model::GPBase, smoothed_t)
    smoothed_x, _ = predict_y(model, smoothed_t)
    return smoothed_x
end

function smoothing(x::AbstractVector, smoother::Smoother;step=0.01)
    t = 1.0:length(x)
    standardized_x = standardize(x)
    model = fit(smoother, t, standardized_x)
    lower_bound, upper_bound = minimum(t), maximum(t)
    smoothed_t = range(lower_bound, stop=upper_bound, step=step)
    n_points = length(smoothed_t)
    smoothed_x = predict(model, smoothed_t)
    smoothed_x = reverse_standardization(smoothed_x, x)
    return smoothed_x
end
function smoothing(x::AbstractVector;step=0.01)
    smoother =  GaussianProcessSmoother(MeanZero(), SE(1.0, 0.0))
    smoothed_x = smoothing(x, smoother, step=step)
    return smoothed_x
end

function smoothing(x::Matrix{T}, smoothers::Vector{<:Smoother};step=0.01) where {T <: Number}
    n_row, n_col = size(x, 1), size(x, 2)
    @assert n_col == length(smoothers)
    t = 1.0:n_row
    lower_bound, upper_bound = minimum(t), maximum(t)
    smoothed_t = range(lower_bound, stop=upper_bound, step=step)
    n_points = length(smoothed_t)
    new_x = zeros(eltype(x), n_points, n_col)

    for i in 1:n_col
        smoothed_x = smoothing(x[:,i], smoothers[i], step=step)
        new_x[:, i] = smoothed_x
    end
    return new_x
end
function smoothing(x::Matrix{T}, smoother::Smoother;step=0.01) where {T <: Number}
    smoother_x = smoother
    smoother_y = deepcopy(smoother)
    smoother_z = deepcopy(smoother)
    smoothers = [smoother_x,smoother_y,smoother_z]
    return smoothing(x, smoothers;step=step)
end
function smoothing(x::Matrix{T};step=0.01) where {T <: Number}
    n_row, n_col = size(x, 1), size(x, 2)
    t = 1.0:n_row
    lower_bound, upper_bound = minimum(t), maximum(t)
    smoothed_t = range(lower_bound, stop=upper_bound, step=step)
    n_points = length(smoothed_t)
    new_x = zeros(eltype(x), n_points, n_col)
    for i in 1:n_col
        smoothed_x = smoothing(x[:,i], step=step)
        new_x[:, i] = smoothed_x
    end
    return new_x
end


# # Polarizability using automatic differentiation
#
# Simple example for computing properties using (forward-mode)
# automatic differentation.
# For a more classical approach and more details about computing polarizabilities,
# see [Polarizability by linear response](@ref).

using DFTK
using LinearAlgebra
using ForwardDiff

## Construct PlaneWaveBasis given a particular electric field strength
## Again we take the example of a Helium atom.
function make_basis(ε::T; a=10., Ecut=30) where T
    lattice=T(a) * I(3)  # lattice is a cube of ``a`` Bohrs
    He = ElementPsp(:He, psp=load_psp("hgh/lda/He-q2"))
    atoms = [He => [[1/2; 1/2; 1/2]]]  # Helium at the center of the box

    model = model_DFT(lattice, atoms, [:lda_x, :lda_c_vwn];
                      extra_terms=[ExternalFromReal(r -> -ε * (r[1] - a/2))],
                      symmetries=false)
    PlaneWaveBasis(model; Ecut, kgrid=[1, 1, 1])  # No k-point sampling on isolated system
end

## dipole moment of a given density (assuming the current geometry)
function dipole(basis, ρ)
    @assert isdiag(basis.model.lattice)
    a  = basis.model.lattice[1, 1]
    rr = [a * (r[1] - 1/2) for r in r_vectors(basis)]
    sum(rr .* ρ) * basis.dvol
end

## Function to compute the dipole for a given field strength
function compute_dipole(ε; tol=1e-8, kwargs...)
    scfres = self_consistent_field(make_basis(ε; kwargs...), tol=tol)
    dipole(scfres.basis, scfres.ρ)
end;

# With this in place we can compute the polarizability from finite differences
# (just like in the previous example):
polarizability_fd = let
    ε = 0.01
    (compute_dipole(ε) - compute_dipole(0.0)) / ε
end

# ## Forward-mode implicit differentiation
#
# Right now DFTK has no out-of-the-box support for implicit differentiation through the SCF.
# However one can easily work around this as follows. We keep both a non-dual basis
# and a basis including duals for easy bookkeeping (but redundant computation ...).

function self_consistent_field_dual(basis::PlaneWaveBasis, basis_dual::PlaneWaveBasis{T};
                                    kwargs...) where T <: ForwardDiff.Dual
    scfres = self_consistent_field(basis; kwargs...)
    ψ, occupation = DFTK.select_occupied_orbitals(basis, scfres.ψ, scfres.occupation)

    ## promote everything eagerly to Dual numbers
    occupation_dual = [T.(occupation[1])]
    ψ_dual = [Complex.(T.(real(ψ[1])), T.(imag(ψ[1])))]
    ρ_dual = compute_density(basis_dual, ψ_dual, occupation_dual)

    _, δH = energy_hamiltonian(basis_dual, ψ_dual, occupation_dual; ρ=ρ_dual)
    δHψ = δH * ψ_dual
    δHψ = [ForwardDiff.partials.(δHψ[1], 1)]
    δψ = DFTK.solve_ΩplusK(basis, ψ, -δHψ, occupation)
    δρ = DFTK.compute_δρ(basis, ψ, δψ, occupation)
    ρ = ForwardDiff.value.(ρ_dual)
    ψ, ρ, δψ, δρ
end;

# This function is now used in the following to provide a dual version
# for the compute_dipole function:
function compute_dipole(ε::ForwardDiff.Dual; tol=1e-8, kwargs...)
    T = ForwardDiff.tagtype(ε)
    basis = make_basis(ForwardDiff.value(ε); kwargs...)
    basis_dual = make_basis(ε; kwargs...)
    ψ, ρ, δψ, δρ = self_consistent_field_dual(basis, basis_dual; tol)
    ρ_dual = ForwardDiff.Dual{T}.(ρ, δρ)
    dipole(basis_dual, ρ_dual)
end;

# This setup allows to compute the polarizability via automatic differentiation:
polarizability = ForwardDiff.derivative(compute_dipole, 0.0)
println()
println("Polarizability via ForwardDiff:       $polarizability")
println("Polarizability via finite difference: $polarizability_fd")


using Test

# see https://discourse.julialang.org/t/writing-tests-in-vs-code-workflow-autocomplete-and-tooltips/57488
# see https://github.com/julia-vscode/julia-vscode/issues/800
if isdefined(@__MODULE__, :LanguageServer)
    include("../src/Classic.jl")
    using .Classic
else
    # invoked during test
    using Classic
end

cc = Classic

function run_astar()
    maze = cc.new_maze(10, 10, (1,1), (10,10), 0.2)

    # create functions for bfs()
    is_goal(pt)     = cc.is_goal(maze.goal, pt)
    next_points(pt) = cc.next_points(maze.grid, pt)
    heuristic(pt)   = cc.manhattan_distance(pt, maze.goal)

    node = cc.a_star((1,1), is_goal, next_points, heuristic)

    if isnothing(node)
        println("Path not found")
        display(maze.grid) # print the unsovled maze
    else
        cc.mark_path!(maze.grid, node, start=maze.start, goal=maze.goal)
        display(maze.grid) # print the maze and the path
    end
end

function run_bfs()
    maze = cc.new_maze(10, 10, (1,1), (10,10), 0.2)
    #maze = cc.new_maze(1)
    #maze = cc.new_maze(2)

    # create functions for bfs()
    is_goal(pt)     = cc.is_goal(maze.goal, pt)
    next_points(pt) = cc.next_points(maze.grid, pt)

    node = cc.bfs((1,1), is_goal, next_points)

    if isnothing(node)
        println("Path not found")
        display(maze.grid) # print the unsovled maze
    else
        cc.mark_path!(maze.grid, node, start=maze.start, goal=maze.goal)
        display(maze.grid) # print the maze and the path
    end
end

function run_cell()
    x = cc.EMPTY
    y = cc.START

    println("x = $(string(x)), y = $(string(y))")
end

"""
    run_dfs()

Use depth-first search to solve a maze.
"""
function run_dfs()
    maze = cc.new_maze() #new_maze(10,10, [1,1], [10,10], 0.2)
    #@show maze
    #display(maze.grid)

    #node = ch2.Node()
    #locations = ch2.successors(maze.grid, (2,2))
    #@show locations

    # pass this to dfs(..., successors)
    successors(here) = cc.successors(maze.grid, here)
    # pass this to dfs(..., goal_test, ...)
    is_goal(here) = cc.is_goal(maze.goal, here)

    node = cc.dfs((1,1), is_goal, successors)
    if isnothing(node)
        println("Path not found")
        display(maze.grid)
    else
        cc.mark_path!(maze.grid, node, start=maze.start, goal=maze.goal)
        #@show maze
        display(maze.grid)
    end
    #@show node
end

function test_bfs()
    for grid_ind = 1:2
        maze = new_maze(new_grid(grid_ind))

        # create functions for bfs()
        is_goal(pt)     = cc.is_goal(maze.goal, pt)
        next_points(pt) = cc.next_points(maze.grid, pt)

        node = bfs(maze.start, is_goal, next_points)
        @test !isnothing(node)

        mark_path!(maze.grid, node, start=maze.start, goal=maze.goal)
        #display(maze.grid)
        sol = new_maze(solution_grid(grid_ind))
        @test sol.grid == maze.grid
    end
end

function test_convert()
    # Char -> Cell
    x = 'O'
    @test Classic.EMPTY == Base.convert(Cell, x) # EMPTY not exported so must use Classic.EMPTY
    x = 'X'
    @test Classic.BLOCKED == Base.convert(Cell, x)
    x = 'S'
    @test Classic.START == Base.convert(Cell, x)
    x = 'G'
    @test Classic.GOAL == Base.convert(Cell, x)
    x = '+'
    @test Classic.PATH == Base.convert(Cell, x)

    x = fill('X', (30,30))
    y = fill(Classic.BLOCKED, (30,30))
    @test y == convert(Matrix{Cell}, x)

    x = fill('X', (30,30))
    x[16:end,16:end] .= '+'
    y = fill(Classic.BLOCKED, (30,30))
    y[16:end,16:end] .= Classic.PATH
    @test y == convert(Matrix{Cell}, x)

    # Cell -> Int
    x = Classic.EMPTY
    @test 0 == Base.convert(Int64, x)
    x = Classic.BLOCKED
    @test 1 == Base.convert(Int64, x)
    x = Classic.START
    @test 2 == Base.convert(Int64, x)
    x = Classic.GOAL
    @test 3 == Base.convert(Int64, x)
    x = Classic.PATH
    @test 4 == Base.convert(Int64, x)

    x = fill(Classic.BLOCKED, (30,30))
    y = fill(1, (30,30))
    @test y == convert(Matrix{Int64}, x)
end

function test_manhattan_distance()
    goal = (12,13)
    pt = (1,1)
    @test 23 == cc.manhattan_distance(pt, goal)

    goal = (1,1)
    pt = (12,13)
    @test 23 == cc.manhattan_distance(pt, goal)

    goal = (1,1)
    pt = (-12,-13)
    @test 27 == cc.manhattan_distance(pt, goal)
end

#run_astar()
#run_cell()
#run_dfs()
#run_bfs()
@testset "Maze" begin
    test_convert()
end

@testset "Breadth-first search" begin
    test_bfs()
end

@testset "A-star search" begin
    test_manhattan_distance()
end

nothing

#A = test_manhattan_distance()
#A # Test.Pass or not


export optimal_traditional_de

"""
    optimal_traditional_de(s, method = "afnn", dmethod = "mi_min"; kwargs...) → 𝒟, τ, E

Produce an optimal delay embedding `𝒟` of the given timeseries `s` by
using the traditional approach of first finding an optimal (and constant) delay
time using [`estimate_delay`](@ref) with the given `dmethod`, and then an optimal
embedding dimension, by calculating an appropriate statistic for each dimension `d ∈ 1:dmax`.
Return the embedding `𝒟`, the optimal delay time `τ`
(the optimal embedding dimension `d` is just `size(𝒟, 2)`) and the actual
statistic `E` used to estimate optimal `d`.

Notice that `E` is a function of the embedding dimension, which ranges from 1 to `dmax`.

For calculating `E` to estimate the dimension we use the given `method` which can be:

* `"afnn"` (default) is Cao's "Averaged False Nearest Neighbors" method[^Cao1997],
    which gives a ratio of distances between nearest neighbors.
* `"ifnn"` is the "Improved False Nearest Neighbors" from Hegger & Kantz[^Hegger1999],
    which gives the fraction of false nearest neighbors.
* `"fnn"` is Kennel's "False Nearest Neighbors" method[^Kennel1992], which gives
    the number of points that cease to be "nearest neighbors" when the dimension
    increases.
* `"f1nn"` is Krakovská's "False First Nearest Neighbors" method[^Krakovská2015],
    which gives the ratio of pairs of points that cease to be "nearest neighbors"
    when the dimension increases.

For more details, see individual methods: [`delay_afnn`](@ref), [`delay_ifnn`](@ref),
[`delay_fnn`](@ref), [`delay_f1nn`](@ref).

!!! warn "Careful in automated methods"
    While this method is automated if you want to be **really sure** of the results,
    you should directly calculate the statistic and plot its values versus the
    dimensions.

## Keyword Arguments
The keywords
```julia
τs = 1:100, dmax = 10
```
denote which delay times and embedding dimensions `ds ∈ 1:dmax` to consider when calculating
optimal embedding. The keywords
```julia
slope_thres = 0.05, stoch_thres = 0.1, fnn_thres = 0.05
```
are specific to this function, see Description below.
All remaining keywords are propagated to the low level functions:
```
w, rtol, atol, τs, metric, r
```

## Description
We estimate the optimal embedding dimension based on the given delay time gained
from `dmethod` as follows: For Cao's method the optimal dimension is reached,
when the slope of the `E₁`-statistic (output from `"afnn"`) falls below the
threshold `slope_thres` and the according stochastic test turns
out to be false, i.e. if the `E₂`-statistic's first value is `< 1 - stoch_thres`.

For all the other methods we return the optimal embedding dimension
when the corresponding FNN-statistic (output from `"ifnn"`, `"fnn"` or `"f1nn"`)
falls below the fnn-threshold `fnn_thres` AND the slope of the
statistic falls below the threshold `slope_thres`. Note that with noise
contaminated time series, one might need to adjust `fnn_thres` according to the
noise level.

See also the file `test/compare_different_dimension_estimations.jl` for a comparison.

[^Cao1997]: Liangyue Cao, [Physica D, pp. 43-50 (1997)](https://www.sciencedirect.com/science/article/pii/S0167278997001188?via%3Dihub)

[^Kennel1992]: M. Kennel *et al.*, [Phys. Review A **45**(6), (1992)](https://journals.aps.org/pra/abstract/10.1103/PhysRevA.45.3403).

[^Krakovská2015]: Anna Krakovská *et al.*, [J. Complex Sys. 932750 (2015)](https://doi.org/10.1155/2015/932750)

[^Hegger1999]: Hegger & Kantz, [Improved false nearest neighbor method to detect determinism in time series data. Physical Review E 60, 4970](https://doi.org/10.1103/PhysRevE.60.4970).
"""
function optimal_traditional_de(s::AbstractVector, dimensionmethod::String = "afnn",
        delaymethod::String= "mi_min";
        fnn_thres::Real = 0.05, slope_thres::Real = .05, dmax::Int = 10, w::Int=1,
        rtol=10.0, atol=2.0, τs = 1:100, metric = Euclidean(), r::Real=2.0,
        stoch_thres = 0.1,
    )

    @assert dimensionmethod ∈ ("afnn", "fnn", "ifnn", "f1nn")
    τ = estimate_delay(s, delaymethod, τs)
    ds = 1:dmax
    γs = ds .- 1 # TODO: This must be updated to dimension in 2.0

    if dimensionmethod=="afnn"
        dimension_statistic = delay_afnn(s, τ, ds; metric, w)
        Y, τ = find_cao_optimal(s, τ, dimension_statistic, slope_thres)
        flag = is_stochastic(s, τ, 1:3, stoch_thres)
        flag && println("Stochastic signal, valid embedding NOT achieved ⨉.")
    elseif dimensionmethod=="fnn"
        dimension_statistic = delay_fnn(s, τ, ds; rtol, atol)
        Y, τ = find_fnn_optimal(s, τ, dimension_statistic, fnn_thres, slope_thres)
    elseif dimensionmethod=="ifnn"
        dimension_statistic = delay_ifnn(s, τ, ds; r, w, metric)
        Y, τ = find_fnn_optimal(s, τ, dimension_statistic, fnn_thres, slope_thres)
    elseif dimensionmethod=="f1nn"
        dimension_statistic = delay_f1nn(s, τ, ds; metric)
        Y, τ = find_fnn_optimal(s, τ, dimension_statistic, fnn_thres, slope_thres)
    end
    return Y, τ, dimension_statistic
end

"""
Helper function for selecting the appropriate embedding dimension from the
statistic when using Kennel's, Hegger's or Krakovskas's method.
"""
function find_fnn_optimal(s::AbstractVector{T}, τ::Int, rat::Vector, fnn_thres::Real,
                                                    slope_thres::Real) where {T}
    @assert length(rat) > 1
    y = abs.(diff(rat))
    flag = false
    m = 0
    for i = 2:length(rat)
        if rat[i] ≤ fnn_thres && y[i-1] ≤ slope_thres
            m = i
            break
        elseif rat[i] > rat[i-1]
            flag = true
            m = i-1
            break
        end
    end
    if m == length(rat) || m == 0
        println("Sufficiently small FNNs NOT reached. "*
                "Valid embedding NOT achieved ⨉.")
    elseif flag
        println("Algorithm stopped due to increasing FNNs. "*
        "Double-check the FNN-statistic.")
    else
        println("Algorithm stopped due to sufficiently small FNNs. "*
                "Valid embedding achieved ✓.")
    end
    Y = embed(s, max(1, m), τ) # you can embed in 1 dimension in latest version
    return Y, τ
end

"""
Helper function for selecting the appropriate embedding dimension from the
statistic when using Cao's method.
"""
function find_cao_optimal(s::AbstractVector{T}, τ::Int, rat::Vector, thres::Real) where {T}
    m = 0
    y = abs.(diff(rat))
    for i = 1:length(rat)-1
        if y[i] ≤ thres && rat[i] > 0.5
            m = i
            break
        end
    end
    if m == 0
        println("NO convergence of E₁-statistic. "*
                "Valid embedding NOT achieved ⨉.")
    else
        Y =  m > 1 ? embed(s, m, τ) : s
        println("Algorithm stopped due to convergence of E₁-statistic. "*
                "Valid embedding achieved ✓.")
    end
    Y = embed(s, max(1, m), τ) # you can embed in 1 dimension in latest version
    return Y, τ
end

"""
Helper function for Cao's method, whether to decide if the input signal is
stochastic or not.
"""
function is_stochastic(s, τ, ds, stoch_thres) where {T}
    E2 = stochastic_indicator(s, τ, ds)
    E2[1] > 1 - stoch_thres
end


@testset "heapify!" begin
    @testset for order = (Forward,Reverse)
        @testset for n = 0
            v = []
            heapify!(v, order)
            @test v == []
        end
        @testset for n = 1
            v = [0]
            heapify!(v, order)
            @test v == [0]
        end
        @testset for n = 2:10
            Random.seed!(42)
            for _ = 1:10
                v = rand(1:100,n)
                heapify!(v,order)
                @test isheap(v,order)
            end
        end
    end
end

@testset "bubble!" begin
    @testset for order = (Forward,Reverse)
        @testset for n = 0
            v = []
            i = bubble!((v,0),1,0,order) == []
            @test v == []
        end
        @testset for n = 1
            v = [0]
            i = bubble!((v,1),1,1,order)
            @test i == 1
            @test v == [1]
        end
        @testset for n = 2:10
            Random.seed!(42)
            for _ = 1:10
                v = rand(1:100,n)
                heapify!(v,order)
                vi = rand(1:100)
                i = bubble!((v,n),vi,rand(1:n),order)
                @test v[i] == vi
                @test isheap(v,order)
            end
        end
    end
end


using Base.Test
import InfiniteVector

@testset "RangeInfiniteVector tests" begin
    for rep = 1:10
        a, b = sort(rand(-100:100, 2))
        vals = rand(length(a:b))
        def = rand()
        v = InfiniteVector.RangeInfiniteVector(a:b, def, vals)

        for i = a:b
            @test v[i] == vals[i - a + 1]
        end

        for r = 1 : 10
            i = rand(Int)
            if i < a || i > b
                @test v[i] == def
            else
                @test v[i] == vals[i - a + 1]
            end
        end

        i = rand(a:b)
        newv = rand()
        v[i] = newv
        @test v[i] == newv

        @test_throws BoundsError v[a - 1] = 1.
        @test_throws BoundsError v[b + 10] = 1.
    end
end


# Sum of the largest k components

# export SumLargest

# TODO: SumLargest(r) is (the postcomposition of) the conjugate of
#  (1) ind{0 <= x <= 1 : sum(x) = r},
# where instead IndSimplex(r) corresponds to
#  (2) ind{0 <= x : sum(x) = r}.
# Therefore SumLargest(r) now is only correct when r = 1.
# To make SumLargest correct we should extend IndSimplex to allow for that
# additional bound in its definition, by adding a second argument to the
# constructor. Then in the following line we should replace IndSimplex(k)
# with IndSimplex(k, 1.0). Note that (1) is proper only if x ∈ Rⁿ for n ⩾ r.

"""
    SumLargest(k::Integer=1, λ::Real=1.0)

Returns the function `g(x) = λ⋅sum(x_[1], ..., x_[k])`, for an integer k ⩾ 1 and `λ ⩾ 0`.
"""
SumLargest(k::I=1, lambda::R=1.0) where {I <: Integer, R <: Real} = Postcompose(Conjugate(IndSimplex(k)), lambda)

function (f::Conjugate{IndSimplex{I}})(x::AbstractArray{S}) where {I <: Integer, S <: Real}
    if f.f.a == 1
        return maximum(x)
    end
    v = zero(S)
    if ndims(x) == 1
        p = partialsortperm(x, 1:f.f.a, rev=true)
        v = sum(x[p])
    else
        p = partialsortperm(x[:], 1:f.f.a, rev=true)
        v = sum(x[p])
    end
    return v
end

fun_name(f::Postcompose{Conjugate{IndSimplex{I}}, R}) where {I <: Integer, R <: Real} = "sum of k largest components"
fun_expr(f::Postcompose{Conjugate{IndSimplex{I}}, R}) where {I <: Integer, R <: Real} = "x ↦ λ⋅sum(x_[1], ..., x_[k])"
fun_params(f::Postcompose{Conjugate{IndSimplex{I}}, R}) where {I <: Integer, R <: Real} = "k = $(f.f.f.a), λ = $(f.a)"


function gnomesort!(arr::Vector)
    i, j = 1, 2
    while i < length(arr)
        if arr[i] ≤ arr[i+1]
            i = j
            j += 1
        else
            arr[i], arr[i+1] = arr[i+1], arr[i]
            i -= 1
            if i == 0
                i = j
                j += 1
            end
        end
    end
    return arr
end

v = rand(-10:10, 10)
println("# unordered: $v\n -> ordered: ", bogosort!(v))