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

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

# 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
   if any(P.<=zero(P)) 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 make_fischer_information_matrix(model::Array{T,2}, Sigma::SigmaT) where {T<:Real, SigmaT<:AbstractPDMat{T}}
   FIM = Xt_invA_X(Sigma,model)
   make_matrix_posdef(FIM)
end

function calc_chisq(data::Vector{T}, model::Array{T,2}, Sigma::SigmaT; FIM::PDMat{T} = make_fischer_information_matrix(model,Sigma), 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} = make_fischer_information_matrix(model,Sigma), calclogprior::Function = x->(0.0,0.0), report_log_MAP::Bool = false ) 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]
  if report_log_MAP     return (param_bf_linalg, sigma_param,log_mode)  end
  LaplaceApprox = 0.5*nparam*log(2pi)-0.5*logdet(FIM)+log_mode + log_jacobian
  return (param_bf_linalg,sigma_param,LaplaceApprox)
end

function compute_max_aposterior_Approx(data::Vector{T}, model::Array{T,2}, Sigma::SigmaT; 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
  return (log_mode, logprior_mode, loglikelihood_mode, log_jacobian)
end

function compute_Laplace_Approx(data::Vector{T}, model::Array{T,2}, Sigma::SigmaT; FIM::PDMat{T} = make_fischer_information_matrix(model,Sigma), calclogprior::Function = x->(0.0,0.0), report_log_MAP::Bool = false ) where {T<:Real, SigmaT<:AbstractPDMat{T} }
  @assert length(data) == size(model,1) == size(Sigma,1)
  nparam = size(model,2)
  nobs = length(data)
#=
  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
=#
  (log_mode, logprior_mode, loglikelihood_mode, log_jacobian) = compute_max_aposterior_Approx(data,model,Sigma,calclogprior=calclogprior)
  if report_log_MAP return log_mode  end

  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)
  j = 1
  while j<=nsamples
    theta = bf .+ sigma*randn(nparam)
	for i in 1:npl
       Ks[j,i] = sqrt.(theta[1+namp_per_pl*(i-1)].^2+theta[2+namp_per_pl*(i-1)].^2)
	   hs[j,i] = theta[3+namp_per_pl*(i-1)] ./ Ks[j,i]
	   ks[j,i] = theta[4+namp_per_pl*(i-1)] ./ Ks[j,i]
	   es[j,i] = sqrt.(hs[j,i].^2.+ks[j,i].^2)
	end
	if all(es[j,:].<1.0)
	     j +=1
	end
  end
  return (Ks,es,hs,ks)
end

function generate_frac_valid_e(bf::Vector, sigma, nsamples::Integer = 100)
  nparam = length(bf)
  const namp_per_pl = 4
  npl = convert(Int64,floor((nparam-1)//namp_per_pl))
  es = generate_Ke_samples(bf,sigma,nsamples)[2]
  num_pass = 0
  for i in 1:size(es,1)
     if all(es[i,:].<1.0) num_pass += 1 end
  end
  num_pass/nsamples
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 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)
    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

function calc_map_no_planet(times::Vector{T}, vels::Vector{T}, sigma_obs::Vector{T}) where T<:Real
  function calc_map_fixed_jitter(sigmaj::Real)
    Sigma = Sigma0 + ScalMat(length(times),sigmaj)
    evid = compute_Laplace_Approx(vels,model,Sigma,report_log_MAP=true)
	evid += logprior_jitter(sigmaj)
    -evid
  end
  model = ones(length(times),1)
  Sigma0 = PDMat(make_Sigma(times,sigma_obs,0.))
  result = optimize(calc_map_fixed_jitter,sigma_j_min,sigma_j_max,rel_tol=1e-6)
  return -result.minimum
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)), report_log_MAP::Bool=false ) 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 = make_fischer_information_matrix(X,Sigma)
     evid = compute_Laplace_Approx(vels,X,Sigma,FIM=FIM,calclogprior=calclogprior_linear,report_log_MAP=report_log_MAP)
 	 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 = make_fischer_information_matrix(X,Sigma)
   logp = compute_Laplace_Approx(vels,X,Sigma,FIM=FIM,calclogprior=calclogprior_linear,report_log_MAP=report_log_MAP)
   logp += logprior_jitter(sigmaj)+ logprior_periods(periods)
   if !report_log_MAP
      logp += 0.5*(log(2pi)-log(0.5*trace((Sigma \ full(dSdj))^2)))
   end
   #println("# logp = ",logp)
   return logp
end

function make_Sigma_cache(num_Sigmas::Integer, Sigma0::SigmaT ) where {T<:Real, SigmaT<:AbstractPDMat{T}}
   @assert 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(size(Sigma0,1),s),sigmaj_list);
   return (sigmaj_list, weights, SigmaList)
end

function make_Sigma_cache(times::Vector{T}, sigma_obs::Vector{T}, num_Sigmas::Integer ) where T<:Real
   @assert num_Sigmas >= 1
   Sigma0 = PDMat(make_Sigma(times,sigma_obs,0.0))
   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);
   return (sigmaj_list, weights, SigmaList)
end

function compute_marginalized_likelihood_const_sigmalist(times::Vector{T}, vels::Vector{T}, sigma_obs::Vector{T}; Sigma0 = PDMat(make_Sigma(times,sigma_obs,0.0)), num_Sigmas::Integer = 20, Sigma_cache = make_Sigma_cache(num_Sigmas,Sigma0) ) where T<:Real
   X = ones(length(times),1)
   const nparam = size(X,2)
   sigmaj_list = Sigma_cache[1]
   weights = Sigma_cache[2]
   SigmaList = Sigma_cache[3]

	   marginalized_likelihood_fix_periods_sigmaj = map(k->compute_marginalized_likelihood_fix_periods_sigmaj(times, vels, sigma_obs, Float64[], sigmaj_list[k], Sigma=SigmaList[k],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

function compute_marginalized_likelihood_fix_periods_sigmalist(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)), num_Sigmas::Integer = 20, Sigma_cache = make_Sigma_cache(num_Sigmas,Sigma0) ) where T<:Real
   const nparam = size(X,2)
   sigmaj_list = Sigma_cache[1]
   weights = Sigma_cache[2]
   SigmaList = Sigma_cache[3]

	   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

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), FIM = make_fischer_information_matrix(X,Sigma) ) 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)

   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

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(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 = 20, verbose::Integer = 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)
  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);
	 =#
	 (sigmaj_list,weights,SigmaList) = make_Sigma_cache(num_Sigmas,Sigma0)
  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
	   logp = compute_marginalized_likelihood_fix_periods_sigmalist(times, vels, sigma_obs, periods, calclogprior_linear=calclogprior_linear, make_design_matrix=make_design_matrix, X=X, Sigma0=Sigma0,  num_Sigmas=num_Sigmas, Sigma_cache=(sigmaj_list,weights,SigmaList) )
	#=
	   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 verbose >=2 && 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
	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

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, sigmaj::Real = -1.0, num_Sigmas::Integer = 20, verbose::Integer = 0 ) 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_opt = 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_opt) # PDMat(make_Sigma(times,sigma_obs,sigmaj_opt))
  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);
	 =#
	 (sigmaj_list,weights,SigmaList) = make_Sigma_cache(num_Sigmas,Sigma0)
  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
	   if sigmaj < 0
   	      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)
	   else
	      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)
	   end
    else
		  logp = compute_marginalized_likelihood_fix_periods_sigmalist(times, vels, sigma_obs, periods, calclogprior_linear=calclogprior_linear, make_design_matrix=make_design_matrix, X=X, Sigma0=Sigma0,  num_Sigmas=num_Sigmas, Sigma_cache=(sigmaj_list,weights,SigmaList) )
       #=
	   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 verbose >=2 && 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
	period_grid[i] = P
	power_grid[i] = logp

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

# Optimize periods
function compute_marginalized_likelihood_near_periods_laplace(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, report_log_MAP::Bool=false ) 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,report_log_MAP=report_log_MAP)
 	 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,report_log_MAP=report_log_MAP)
 	 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,report_log_MAP=report_log_MAP)
 	 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)
   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,report_log_MAP=report_log_MAP)
   logp += logprior_jitter(sigmaj)+ logprior_periods(periods)

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

   if !report_log_MAP
      # 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))
   else
      log_int_over_sigma_j = 0.
   end

   # Laplace approx to integrate over period
   sigma_P = zeros(length(periods))
   log_int_over_periods = 0
   duration = maximum(times)-minimum(times)
   param_per_pl = (calclogprior_linear == logprior_sinusoids) ? 2 : 4
   for i in 1:length(best_period)
     periods = copy(best_period)

     amp = sqrt(param_bf[1+(i-1)*param_per_pl]^2+param_bf[2+(i-1)*param_per_pl]^2)
     eps = periods[i]/(40*2pi*duration*amp)
	 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(abs(1/d2logpdp2))
     #println("# d^2 logp/dP_",i,"^2 = ",d2logpdp2/log(10))
     log_int_over_periods += 0.5*(log(2pi)-log(abs(d2logpdp2)))
   end

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

   logp += log_int_over_periods # Approx as diagonal
   println("# marginal over P at ",best_period," and sigma_j at ",sigmaj," (laplace): ",logp/log(10))
   if calclogprior_linear == logprior_epicycles
      frac_valid_e = generate_frac_valid_e(param_bf,sigma_param)
      println("# frac w/ valid e = ", frac_valid_e)
      logp += log(frac_valid_e)
   end
   return (param_bf,sigma_param,logp,best_period,sigma_P)
end

function compute_marginalized_likelihood_near_periods_gauss(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, num_Sigmas::Integer = 20 ) 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-4)
	  periods = [result2.minimizer]
   else
	  result2 = optimize(optimize_period_helper,periods,f_tol=1e-4) # abs(result.minimum))
	  periods = result2.minimizer
   end
   #println("# Optimize period: ", result2)
   best_period = copy(periods)
   X = make_design_matrix(times,periods)
  (param_bf,sigma_param,logp) = compute_Laplace_Approx_param_and_integral(vels,X,Sigma,calclogprior=calclogprior_linear)

    (sigmaj_list,weights,SigmaList) = make_Sigma_cache(num_Sigmas,Sigma0)
     logp = compute_marginalized_likelihood_fix_periods_sigmalist(times, vels, sigma_obs, periods, calclogprior_linear=calclogprior_linear, make_design_matrix=make_design_matrix, X=X, Sigma0=Sigma0,  num_Sigmas=num_Sigmas, Sigma_cache=(sigmaj_list,weights,SigmaList) )

   # Laplace approx to integrate over period
   sigma_P = zeros(length(periods))
   log_int_over_periods = 0
   duration = maximum(times)-minimum(times)
   param_per_pl = (calclogprior_linear == logprior_sinusoids) ? 2 : 4
   for i in 1:length(best_period)
     periods = copy(best_period)
	 amp = sqrt(param_bf[1+(i-1)*param_per_pl]^2+param_bf[2+(i-1)*param_per_pl]^2)
     eps = periods[i]/(40*2pi*duration*amp)
	 periods[i] += eps

     X = make_design_matrix(times,periods)
	 logp_hi = compute_marginalized_likelihood_fix_periods_sigmalist(times, vels, sigma_obs, periods, calclogprior_linear=calclogprior_linear, make_design_matrix=make_design_matrix, X=X, Sigma0=Sigma0,  num_Sigmas=num_Sigmas, Sigma_cache=(sigmaj_list,weights,SigmaList) )

	 periods[i] -= 2*eps
     X = make_design_matrix(times,periods)
	 logp_lo = compute_marginalized_likelihood_fix_periods_sigmalist(times, vels, sigma_obs, periods, calclogprior_linear=calclogprior_linear, make_design_matrix=make_design_matrix, X=X, Sigma0=Sigma0,  num_Sigmas=num_Sigmas, Sigma_cache=(sigmaj_list,weights,SigmaList) )

     d2logpdp2 = (logp_hi-2*logp+logp_lo)/eps^2
     sigma_P[i] = sqrt(abs(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," (gauss): ",logp/log(10))
   if calclogprior_linear == logprior_epicycles
      frac_valid_e = generate_frac_valid_e(param_bf,sigma_param)
      println("# frac w/ valid e = ", frac_valid_e)
      logp += log(frac_valid_e)
   end
   return (param_bf,sigma_param,logp,best_period,sigma_P)
end

function find_n_peaks_in_pgram(period::Vector{Float64}, power::Vector{Float64}; num_peaks::Integer = 1, exclude_period_factor::Real = 1.2, exclude_peaks::Vector{Float64} = Float64[] )
  @assert num_peaks>= 1
  peak_periods = zeros(num_peaks)
  peak_logps = zeros(num_peaks)
  idx_active = trues(length(period))
  for i in 1:length(exclude_peaks)
     for j in 1:length(period)
        if exclude_peaks[i]/exclude_period_factor <= period[j] <= exclude_peaks[i]*exclude_period_factor
           idx_active[j] = false
        end # if near an excluded peak
     end # for over periods
  end # for over excluded peaks
  peaks_found = 1
  idx_peak = findmax(power[idx_active])[2]
  peak_periods[peaks_found] = period[idx_active][idx_peak]
  peak_logps[peaks_found] = power[idx_active][idx_peak]
  while peaks_found < num_peaks
     for j in 1:length(period)
        if peak_periods[peaks_found]/exclude_period_factor <= period[j] <= peak_periods[peaks_found]*exclude_period_factor
           idx_active[j] = false
        end # blackout periods near most recent peak
     end # for over periods

     peaks_found += 1
	 idx_peak = findmax(power[idx_active])[2]
     peak_periods[peaks_found] = period[idx_active][idx_peak]
	 peak_logps[peaks_found] = power[idx_active][idx_peak]

	 end # while more peaks to be found
  (peak_periods,peak_logps)
end

function compute_evidences_laplace_circ{T}(times::Vector{T}, vels::Vector{T}, sigma_obs::Vector{T}; num_Sigmas::Integer = 20, samples_per_peak::Integer = 1, Pminsearch::Vector{Float64} = 10.0*ones(3), report_log_MAP::Bool=false, max_num_pl::Integer=3 )
	 #(sigmaj_list,weights,SigmaList) = make_Sigma_cache(num_Sigmas,Sigma0)

runtimes = zeros(max_num_pl+1)
tic()
log10_evidences = zeros(max_num_pl+1)
res0l = compute_marginalized_likelihood_fix_periods(times,vels,sigma_obs,Float64[],report_log_MAP=report_log_MAP)
res0g = compute_marginalized_likelihood_const_sigmalist(times,vels,sigma_obs,num_Sigmas=num_Sigmas)
println("# Evidence for npl=0: ",res0l/log(10),", ", res0g/log(10))
if report_log_MAP
  log10_evidences[1] = res0l/log(10)
else
  log10_evidences[1] = res0g/log(10)
end
runtimes[1] = toq()
if max_num_pl<=0
  return (log10_evidences,[],[res0l],runtimes)
end

tic()
@time res1b = brute_force_over_periods_1d(times,vels,sigma_obs,Pminsearch[1],Pmax,samples_per_peak=samples_per_peak,num_Sigmas=num_Sigmas)
P1 = find_n_peaks_in_pgram(res1b[1],res1b[2])[1][1]

res1l = compute_marginalized_likelihood_near_periods_laplace(times,vels,sigma_obs, [P1],report_log_MAP=report_log_MAP )
res1g = compute_marginalized_likelihood_near_periods_gauss(times,vels,sigma_obs, [P1],num_Sigmas=num_Sigmas )
if report_log_MAP
   log10_evidences[2] = res1l[3]/log(10)
else
   log10_evidences[2] = res1g[3] > res1b[3] ? res1g[3]/log(10) : res1b[3]/log(10)
end
P1 = res1l[4][1]
runtimes[2] = toq()

if max_num_pl<=1
  return (log10_evidences,[P1],[res1l],runtimes)
end

tic()
@time res2b = brute_force_over_periods_multiplanet_1d(times,vels,sigma_obs,[P1],Pminsearch[2],Pmax,samples_per_peak=samples_per_peak,num_Sigmas=num_Sigmas)
println("# About to search for second planet, excluding peak near P=",P1)
P2 = find_n_peaks_in_pgram(res2b[1],res2b[2],exclude_peaks=[P1])[1][1]

res2l = compute_marginalized_likelihood_near_periods_laplace(times,vels,sigma_obs, [P1,P2],report_log_MAP=report_log_MAP )
res2g = compute_marginalized_likelihood_near_periods_gauss(times,vels,sigma_obs, [P1,P2],num_Sigmas=num_Sigmas  )
if report_log_MAP
   log10_evidences[3] = res2l[3]/log(10)
else
   log10_evidences[3] = res2g[3] > res2b[3] ? res2g[3]/log(10) : res2b[3]/log(10)
end
P1 = res2l[4][1]
P2 = res2l[4][2]
runtimes[3] = toq()

if max_num_pl<=2
  return (log10_evidences,[P1,P2],[res1l,res2l],runtimes)
end

tic()
@time res3b = brute_force_over_periods_multiplanet_1d(times,vels,sigma_obs,[P1,P2],Pminsearch[3],Pmax,samples_per_peak=samples_per_peak,num_Sigmas=num_Sigmas)
P3 = find_n_peaks_in_pgram(res3b[1],res3b[2],num_peaks=10,exclude_peaks=[P1,P2])[1][1]

res3l = compute_marginalized_likelihood_near_periods_laplace(times,vels,sigma_obs, [P1,P2,P3],report_log_MAP=report_log_MAP )
res3g = compute_marginalized_likelihood_near_periods_gauss(times,vels,sigma_obs, [P1,P2,P3],num_Sigmas=num_Sigmas  )
if report_log_MAP
   log10_evidences[4] = res3l[3]/log(10)
else
   log10_evidences[4] = res3g[3] > res3b[3] ? res3g[3]/log(10) : res3b[3]/log(10)
end
runtimes[4] = toq()

  return (log10_evidences,[P1,P2,P3],[res1l,res2l,res3l],runtimes)
end

function make_kepler_param_from_linear_output(output, npl::Integer)
  evidence_linear = output[1][npl]
  periods = output[2][1:npl]
  param_linear = output[4][npl]
  sigmaj = output[6][npl]
  chisq_linear = output[10][npl]
  const nparam_per_pl = 5
  const nparam_non_pl = 2
  const nparam_linear_per_pl = 4
  const nparam = npl*nparam_per_pl + nparam_non_pl
  const tref = 300.0
  param = Array{Float64}(nparam)
  for i in 1:npl
      P = periods[i]
	  Ksinl = param_linear[1+(i-1)*nparam_linear_per_pl]
	  Kcosl = param_linear[2+(i-1)*nparam_linear_per_pl]
	  K = sqrt(Ksinl^2+Kcosl^2)
	  wpM = atan2(Ksinl,Kcosl)
	  h = K>0 ? param_linear[3+(i-1)*nparam_linear_per_pl]/K : 0
	  k = K>0 ? param_linear[4+(i-1)*nparam_linear_per_pl]/K : 0
	  w = atan2(h,k)
      param[1+(i-1)*nparam_per_pl] = P
	  param[2+(i-1)*nparam_per_pl] = K > 0 ? K : 0.5
	  param[3+(i-1)*nparam_per_pl] = 1e-8 # h/100
	  param[4+(i-1)*nparam_per_pl] = 1e-4 # k/100
	  param[5+(i-1)*nparam_per_pl] = mod2pi(wpM+2pi*tref/P)
  end
  param[1+npl*nparam_per_pl] = param_linear[1+npl*nparam_linear_per_pl]
  param[2+npl*nparam_per_pl] = sigmaj
  param
end

function optimize_phase{T}(param, times::Vector{T}, vels::Vector{T}, Sigma, plid::Integer)
  const nparam_per_pl = 5
  const nparam_non_pl = 2
  npl = convert(Int64,floor((length(param)-1)//nparam_per_pl))
  function optimize_phase_helper(x::Real)
    param[5+(plid-1)*nparam_per_pl] = x
    -loglikelihood_fixed_jitter(param,times,vels,Sigma)
  end
  result = optimize(optimize_phase_helper,0.0,2pi)
  param[5+(plid-1)*nparam_per_pl] = mod2pi(result.minimizer)
  param
end

function optimize_omega{T}(param, times::Vector{T}, vels::Vector{T}, Sigma, plid::Integer)
  const nparam_per_pl = 5
  const nparam_non_pl = 2
  npl = convert(Int64,floor((length(param)-1)//nparam_per_pl))
  function optimize_omega_helper(x::Vector)
    param[3+(plid-1)*nparam_per_pl:4+(plid-1)*nparam_per_pl] = x
    -loglikelihood_fixed_jitter(param,times,vels,Sigma)
  end
  result = optimize(optimize_omega_helper, param[3+(plid-1)*nparam_per_pl:4+(plid-1)*nparam_per_pl] )
  if sum(result.minimizer.^2)<1
     param[3+(plid-1)*nparam_per_pl:4+(plid-1)*nparam_per_pl] = result.minimizer
  else
     param[3+(plid-1)*nparam_per_pl] = 0.0 #1e-6*param[3+(plid-1)*nparam_per_pl]
	 param[4+(plid-1)*nparam_per_pl] = 0.0 #1e-6*param[4+(plid-1)*nparam_per_pl]
  end
  param
end

function optimize_phases{T}(param, times::Vector{T}, vels::Vector{T}, Sigma)
  const nparam_per_pl = 5
  const nparam_non_pl = 2
  npl = convert(Int64,floor((length(param)-1)//nparam_per_pl))
  for plid in 1:npl
     optimize_phase(param,times,vels,Sigma,plid)
	 println("# param (after phase ",plid," ) = ", param)
	 #optimize_omega(param,times,vels,Sigma,plid)
	 #println("# param = (after omega ",plid," = ", param)
  end
  param
end

function analyze_dataset(filename::String; num_Sigmas::Integer = 20, samples_per_peak::Real = 1.0, Pminsearch::Vector{Float64} = 10.0*ones(3))
  # Read dataset
  data = readdlm(filename);
  times = data[:,1];
  vels = data[:,2];
  sigma_obs = data[:,3];
  nobs = length(vels);

   (evidence_circ_list, P_best, result_laplace_circ_list, runtimes_circ) = compute_evidences_laplace_circ(times,vels,sigma_obs,num_Sigmas=num_Sigmas, Pminsearch=Pminsearch)
   runtimes_kepler = zeros(runtimes_circ)

   param_bf_list = Vector{Vector{Float64}}(length(P_best))
   resid_list = Vector{Vector{Float64}}(length(P_best))
   rms_list = Vector{Float64}(length(P_best))
   chisq_list = Vector{Float64}(length(P_best))
   sigmaj_list = Vector{Float64}(length(P_best))
   sigma_P_list = Vector{Vector{Float64}}(length(P_best))
   sigma_sigmaj_list = Vector{Float64}(length(P_best))
   sigma_param_list = Vector(length(P_best))
   frac_valid_list = Vector{Float64}(length(P_best))
   evidence_kepler_list = Vector{Float64}(length(P_best))
   kepler_fit_list = Vector(length(P_best))
   for np in 1:length(P_best)  # Loop over number of planets to include
      sigma_P_list[np] = result_laplace_circ_list[3][5]
	  tic()
	  X = make_design_matrix_epicycle(times,P_best[1:np])
      sigmaj_opt = calc_opt_sigmaj(times, vels, sigma_obs, X, calclogprior_linear=logprior_epicycles).minimizer
      Sigma = PDMat(make_Sigma(times,sigma_obs,sigmaj_opt))
      param_bf_list[np] = calc_best_fit(vels,X,Sigma)
      predict = X*param_bf_list[np]
      # Estimate Uncertainties
      FIM = make_fischer_information_matrix(X,Sigma)
      sigma_param = inv(FIM).chol[:U]
      frac_valid_e = generate_frac_valid_e(param_bf_list[np],sigma_param)
      # Uncertainty on sigmaj
      dSdj = make_dSigmadsigmaj(times,sigma_obs,sigmaj_opt)
      sigma_sigmaj = 1/sqrt(abs(0.5*trace((Sigma \ full(dSdj))^2)))
      println("# np=", np," sigma_sigmaj = ",sigma_sigmaj,"  frac w/ valid e = ", frac_valid_e)
      evidence_kepler_list[np] = compute_marginalized_likelihood_near_periods_gauss(times,vels,sigma_obs, P_best[1:np],num_Sigmas=num_Sigmas, calclogprior_linear=logprior_epicycles, make_design_matrix=make_design_matrix_epicycle )[3]/log(10)

      resid_list[np] = vels.-predict
      rms_list[np] = std(resid_list[np])
      chisq_list[np] = invquad(Sigma,resid_list[np])
	  sigma_param_list[np] = sigma_param
	  sigmaj_list[np] = sigmaj_opt
	  sigma_sigmaj_list[np] = sigma_sigmaj
	  frac_valid_list[np] = frac_valid_e
	  runtimes_kepler[np] += toq()
  end

   result_linear = (evidence_circ_list, P_best,sigma_P_list, param_bf_list, sigma_param_list, sigmaj_list, sigma_sigmaj_list, resid_list, rms_list, chisq_list, frac_valid_list, runtimes_circ, evidence_kepler_list, runtimes_kepler)

   #=
   Sigma0 = PDMat(make_Sigma(times,sigma_obs,0.0));
   for npl in 1:length(P_best)
     tic()
     param_tmp = make_kepler_param_from_linear_output(result_linear,npl)
     SigmaBig = PDMat(make_Sigma(times,sigma_obs,sigmaj_list[npl]));
     optimize_phases(param_tmp,times,vels,SigmaBig)
     kepler_fit_list[npl] = laplace_approximation_kepler_model_jitter_separate(param_tmp,times,vels,Sigma0)
     evidence_kepler_list[npl] = kepler_fit_list[npl][1]/log(10)
	 println("# logp npl=",npl," (laplace): ",evidence_kepler_list[npl])
     flush(STDOUT)
     evidence_kepler_list[npl] = laplace_approximation_kepler_model_jitter_separate_gauss(param_tmp,times,vels,Sigma0)[3]/log(10)
	 println("# logp npl=",npl," (gauss): ",evidence_kepler_list[npl])
	 flush(STDOUT)
     runtimes_kepler[np] += toq()
  end

   return (evidence_circ_list, P_best,sigma_P_list, param_bf_list, sigma_param_list, sigmaj_list, sigma_sigmaj_list, resid_list, rms_list, chisq_list, frac_valid_list,evidence_kepler_list,runtimes_kepler, kepler_fit_list,runtimes_circ,runtimes_kepler)
   =#
end

function write_evidences_file(filename::String,output; use_method::Integer = 1)
   if use_method==1
      evidence = output[1]
	  runtime = convert(Array{Int64,1},ceil.(output[12]))
   elseif use_method==2
	  evidence = vcat(output[1][1],output[13])
	  runtime = convert(Array{Int64,1},ceil.(output[14]))
   end
   output_circ_data = Array{Any}(4,8)
   output_circ_data[1,1] = 0
   output_circ_data[:,1] = runtime
   output_circ_data[:,2] = collect(0:3)
   for i in 3:8
      output_circ_data[:,i] = evidence
   end
   writedlm(filename, output_circ_data, ", ")
end

function make_param_string(paramid::String, npl::Integer, plid::Integer,bf::Real, sigma::Real; runtime::Integer=0)
	   p50 = bf
	   p02 = bf-2*sigma
	   p16 = bf-sigma
	   p84 = bf+sigma
	   p98 = bf+2*sigma
	   string(runtime,',',paramid,',',npl,',',plid,',',bf,',',p50,',',p02,',',p16,',',p84,',',p98)
end

function write_parameters_file(filename::String,output; use_method::Integer = 1, input_filename::String="../data/rvs_0001.txt", num_pl::Integer=3)
    file = open(filename,"w")
	evidence = output[1]
	runtime = convert(Array{Int64,1},ceil.(output[12]))[num_pl]
    println(file,make_param_string("J",num_pl,0,output[6][num_pl],output[7][num_pl],runtime=runtime))
	(Ks,es,hs,ks) = generate_Ke_samples(output[4][num_pl],output[5][num_pl],1000)
	for plid = 1:num_pl
	   bf = output[2][plid]
	   sigma = output[3][num_pl][plid]
	   output_string = make_param_string("P",num_pl,plid,bf,sigma,runtime=runtime)
	   println(file,output_string)
	   bf = mean(Ks[:,plid])
	   sigma = std(Ks[:,plid])
	   output_string = make_param_string("K",num_pl,plid,bf,sigma,runtime=runtime)
	   println(file,output_string)
	 end
	close(file)
end

#=
for num_pl in 1:3
   write_parameters_file(string("params_",num_pl,"_000",1,".txt"),output1,input_filename="../data/rvs_0001.txt",num_pl=num_pl)
   write_parameters_file(string("params_",num_pl,"_000",2,".txt"),output2,input_filename="../data/rvs_0002.txt",num_pl=num_pl)
   write_parameters_file(string("params_",num_pl,"_000",3,".txt"),output3,input_filename="../data/rvs_0003.txt",num_pl=num_pl)
   write_parameters_file(string("params_",num_pl,"_000",4,".txt"),output4,input_filename="../data/rvs_0004.txt",num_pl=num_pl)
   write_parameters_file(string("params_",num_pl,"_000",5,".txt"),output5,input_filename="../data/rvs_0005.txt",num_pl=num_pl)
   write_parameters_file(string("params_",num_pl,"_000",6,".txt"),output6,input_filename="../data/rvs_0006.txt",num_pl=num_pl)
end
=#

function analyze_dataset(filename::String; num_Sigmas::Integer = 20, samples_per_peak::Real = 1.0, Pminsearch::Vector{Float64} = 10.0*ones(3))
  # Read dataset
  data = readdlm(filename);
  times = data[:,1];
  vels = data[:,2];
  sigma_obs = data[:,3];
  nobs = length(vels);

   (evidence_circ_list, P_best, result_laplace_circ_list, runtimes_circ) = compute_evidences_laplace_circ(times,vels,sigma_obs,num_Sigmas=num_Sigmas, Pminsearch=Pminsearch)
    write_parameters_file("tmp1.out",ans, input_filename=filename, num_pl=1)
	write_parameters_file("tmp2.out",ans, input_filename=filename, num_pl=2)
	write_parameters_file("tmp3.out",ans, input_filename=filename, num_pl=3)
end

function write_evidences_file(filename::String,output; use_method::Integer = 1, input_filename::String="../data/rvs_0001.txt")
   if use_method==1
      evidence = output[1]
	  runtime = convert(Array{Int64,1},ceil.(output[12]))
   elseif use_method==2
	  evidence = vcat(output[1][1],output[13])
	  runtime = convert(Array{Int64,1},ceil.(output[14]))
   elseif use_method==3
       evidence = zeros(4)
	   runtime = zeros(4)
	   data = readdlm(input_filename)
       times = data[:,1];
       vels = data[:,2];
       sigma_obs = data[:,3];
	   Sigma0 = PDMat(make_Sigma(times,sigma_obs,0.0));
	   evidence[1] = NaN
       for np in 1:3
	      tic()
		  param = make_kepler_param_from_linear_output(output,np)
	      SigmaBig = PDMat(make_Sigma(times,sigma_obs,param[end]));
		  optimize_phases(param,times,vels,SigmaBig)
	      param2 = laplace_approximation_kepler_model_jitter_separate(param,times,vels,Sigma0,opt_jitter=false)[2].minimizer
		  evidence[1+np] = laplace_approximation_kepler_model_jitter_separate(param2,times,vels,Sigma0,opt_jitter=true)[1]/log(10)
		  runtime[1+np] = toc()
	   end
   elseif use_method==4
       evidence = zeros(4)
	   runtime = zeros(4)
	   data = readdlm(input_filename)
       times = data[:,1];
       vels = data[:,2];
       sigma_obs = data[:,3];
	   Sigma0 = PDMat(make_Sigma(times,sigma_obs,0.0));
	   evidence[1] = NaN
       for np in 0:3
	      tic()
		  periods = output[2][1:np]
		  if np == 1
		     param = vcat(output[4][np][[1,2,5]],output[6])
		  elseif np == 2
		     idx = vcat(output[4][np][[1,2,5,6,9]],output[6])
		  elseif np == 3
		     idx = vcat(output[4][np][[1,2,5,6,9,10,13]],output[6])
		  end
          if np == 0
 		     evidence[1+np] = calc_map_no_planet(times,vels,sigma_obs)/log(10)
		  else
             evidence[1+np] = compute_marginalized_likelihood_near_periods_laplace(times, vels, sigma_obs, periods,  calclogprior_linear=logprior_sinusoids, make_design_matrix=make_design_matrix_circ, report_log_MAP=true)[3]/log(10) # Actually logMAP
		  end
		  runtime[1+np] = toc()
	   end
   elseif use_method==5
       evidence = zeros(4)
	   runtime = zeros(4)
	   data = readdlm(input_filename)
       times = data[:,1];
       vels = data[:,2];
       sigma_obs = data[:,3];
	   Sigma0 = PDMat(make_Sigma(times,sigma_obs,0.0));
	   evidence[1] = NaN
       for np in 0:3
	      tic()
		  periods = output[2][1:np]
		  if np == 1
		     param = vcat(output[4][np][[1,2,5]],output[6])
		  elseif np == 2
		     idx = vcat(output[4][np][[1,2,5,6,9]],output[6])
		  elseif np == 3
		     idx = vcat(output[4][np][[1,2,5,6,9,10,13]],output[6])
		  end
          if np == 0
 		     evidence[1+np] = calc_map_no_planet(times,vels,sigma_obs)/log(10)
		  else
             evidence[1+np] = compute_marginalized_likelihood_near_periods_laplace(times, vels, sigma_obs, periods,  calclogprior_linear=logprior_epicycles, make_design_matrix=make_design_matrix_epicycle, report_log_MAP=true)[3]/log(10) # Actually logMAP
	      end
		  runtime[1+np] = toc()
	   end
   elseif use_method==6
       evidence = zeros(4)
	   runtime = zeros(4)
	   data = readdlm(input_filename)
       times = data[:,1];
       vels = data[:,2];
       sigma_obs = data[:,3];
	   Sigma0 = PDMat(make_Sigma(times,sigma_obs,0.0));
	   evidence[1] = NaN
       for np in 0:3
	      tic()
		  if np == 0
 		     evidence[1+np] = calc_map_no_planet(times,vels,sigma_obs)/log(10)
		  else
		  param = make_kepler_param_from_linear_output(output,np)
	      SigmaBig = PDMat(make_Sigma(times,sigma_obs,param[end]));
		  optimize_phases(param,times,vels,SigmaBig)
	      #param2 = laplace_approximation_kepler_model_jitter_separate(param,times,vels,Sigma0,opt_jitter=false,report_log_MAP=true)[2].minimizer
		  #evidence[1+np] = laplace_approximation_kepler_model_jitter_separate(param2,times,vels,Sigma0,opt_jitter=true,report_log_MAP=true)[1]/log(10)
		  evidence[1+np] = laplace_approximation_kepler_model_jitter_separate(param,times,vels,Sigma0,opt_jitter=false,report_log_MAP=true)[1]/log(10)
		  end
		  runtime[1+np] = toc()
	   end
   end
   output_circ_data = Array{Any}(4,8)
   output_circ_data[1,1] = 0
   output_circ_data[:,1] = runtime
   output_circ_data[:,2] = collect(0:3)
   for i in 3:8
      output_circ_data[:,i] = evidence
   end
   writedlm(filename, output_circ_data, ", ")
end

#=
include("calc_evidence_laplace_lowe.jl")
output1 = analyze_dataset("../data/rvs_0001.txt",num_Sigmas = 4,samples_per_peak=1,Pminsearch=[30.0,10.0,10.])
write_evidences_file("laplace_linearized_circ/evidences_0001.txt",output1)
write_evidences_file("laplace_linearized_ecc/evidences_0001.txt",output1,use_method=2)
@save "output_kepler_1.jld"

output1 = analyze_dataset("../data/rvs_0001.txt",num_Sigmas = 40,samples_per_peak=4,Pminsearch=[30.0,10.0,10.0])

output1 = analyze_dataset("../data/rvs_0001.txt",num_Sigmas = 4,samples_per_peak=1,Pminsearch=[10.0,10.0,1.2])
output2 = analyze_dataset("../data/rvs_0002.txt",num_Sigmas = 4,samples_per_peak=1,Pminsearch=[10.0,10.0,1.2])
output3 = analyze_dataset("../data/rvs_0003.txt",num_Sigmas = 4,samples_per_peak=1,Pminsearch=[10.0,10.0,1.2])
output4 = analyze_dataset("../data/rvs_0004.txt",num_Sigmas = 4,samples_per_peak=1,Pminsearch=[10.0,10.0,1.2])
output5 = analyze_dataset("../data/rvs_0005.txt",num_Sigmas = 4,samples_per_peak=1,Pminsearch=[10.0,10.0,1.2])
output6 = analyze_dataset("../data/rvs_0006.txt",num_Sigmas = 4,samples_per_peak=1,Pminsearch=[10.0,10.0,1.2])
@save "output_kepler_4.jld"
=#
#=
include("calc_evidence_laplace_lowe.jl")
output1 = analyze_dataset("../data/rvs_0001.txt",num_Sigmas = 40,samples_per_peak=4,Pminsearch=[10.0,10.,1.2])
write_evidences_file("laplace_linearized_circ/evidences_0001.txt",output1)
write_evidences_file("laplace_linearized_ecc/evidences_0001.txt",output1,use_method=2)
@save "output_kepler_40_4_10.jld"
output2 = analyze_dataset("../data/rvs_0002.txt",num_Sigmas = 40,samples_per_peak=4,Pminsearch=[10.0,10.,1.2])
write_evidences_file("laplace_linearized_circ/evidences_0002.txt",output2)
write_evidences_file("laplace_linearized_ecc/evidences_0002.txt",output2,use_method=2)
output3 = analyze_dataset("../data/rvs_0003.txt",num_Sigmas = 40,samples_per_peak=4,Pminsearch=[10.0,10.,1.2])
write_evidences_file("laplace_linearized_circ/evidences_0003.txt",output3)
write_evidences_file("laplace_linearized_ecc/evidences_0003.txt",output3,use_method=2)
output4 = analyze_dataset("../data/rvs_0004.txt",num_Sigmas = 40,samples_per_peak=4,Pminsearch=[10.0,10.,1.2])
write_evidences_file("laplace_linearized_circ/evidences_0004.txt",output4)
write_evidences_file("laplace_linearized_ecc/evidences_0004.txt",output4,use_method=2)
output5 = analyze_dataset("../data/rvs_0005.txt",num_Sigmas = 40,samples_per_peak=4,Pminsearch=[10.0,10.,1.2])
write_evidences_file("laplace_linearized_circ/evidences_0005.txt",output5)
write_evidences_file("laplace_linearized_ecc/evidences_0005.txt",output5,use_method=2)
output6 = analyze_dataset("../data/rvs_0006.txt",num_Sigmas = 40,samples_per_peak=4,Pminsearch=[10.0,10.,1.2])
write_evidences_file("laplace_linearized_circ/evidences_0006.txt",output6)
write_evidences_file("laplace_linearized_ecc/evidences_0006.txt",output6,use_method=2)
@save "output_kepler_40_4_10.jld"
=#

#=
data = readdlm("../data/rvs_0001.txt")
times = data[:,1];
vels = data[:,2];
sigma_obs = data[:,3];
output1 = analyze_dataset("../data/rvs_0001.txt",num_Sigmas = 40,Pminsearch=[30.0,10.0,10.0])
param_tmp = make_kepler_param_from_linear_output(output1,1)
SigmaBig = PDMat(make_Sigma(times,sigma_obs,param_tmp[end]));
optimize_phases(param_tmp,times,vels,SigmaBig)
res1 = laplace_approximation_kepler_model_jitter_separate(param_tmp,times,vels,Sigma0)
res1[1]/log(10)

param_tmp = make_kepler_param_from_linear_output(output1,2)
SigmaBig = PDMat(make_Sigma(times,sigma_obs,param_tmp[end]));
optimize_phases(param_tmp,times,vels,SigmaBig)
res2 = laplace_approximation_kepler_model_jitter_separate(param_tmp,times,vels,Sigma0)
res2[1]/log(10)

param_tmp = make_kepler_param_from_linear_output(output1,3)
SigmaBig = PDMat(make_Sigma(times,sigma_obs,param_tmp[end]));
optimize_phases(param_tmp,times,vels,SigmaBig)
res3 = laplace_approximation_kepler_model_jitter_separate(param_tmp,times,vels,Sigma0)
res3[1]/log(10)

=#

#=
param1a = [42.0,5.0,0.01,0.01,1.,0.]
res = laplace_approximation_kepler_model(param1a,times,vels,SigmaBig,loglikelihood_fixed_jitter)
param1a = res[2].minimizer;
laplace_approximation_kepler_model_jitter_separate(param1a,times,vels,Sigma0)[1]/log(10)
=#

function compute_marginalized_likelihood_cuba(times::Vector{T}, vels::Vector{T}, sigma_obs::Vector{T}, min_param::Vector{T}, max_param::Vector{T}; num_Sigmas::Integer = 20 ) where T<:Real
  @assert length(min_param)==length(max_param)
  log_normalization = 0.0
  Sigma0 = PDMat(make_Sigma(times,sigma_obs,0.))
  Sigma_cache = make_Sigma_cache(num_Sigmas,Sigma0)
  function log_integrand_cuba(x::Vector)
    param = (max_param.-min_param).*x.+min_param

	logp = loglikelihood_marinalize_jitter(param,times,vels,Sigma_cache=Sigma_cache)
	#println("# param= ",param," logp= ",logp)
	logp
  end
  function integrand_cuba(x::Vector, output::Vector)
    output[1] = exp(log_integrand_cuba(x)-log_normalization)
  end

  log_normalization = log_integrand_cuba(0.5*ones(min_param))
  result = suave(integrand_cuba,length(min_param),1,reltol=1e-3,maxevals=10000)
  #result = log(result.integral[1]) + log(prod(max_param.-min_param))
  #result+log_normalization
end

function compute_marginalized_likelihood_cuba_near(times::Vector{T}, vels::Vector{T}, sigma_obs::Vector{T}, param::Vector{T}; num_Sigmas::Integer = 20 ) where T<:Real
  const nparam_per_pl = 5
  const nparam_non_pl = 2
  npl = convert(Int64,floor((length(param)-1)//nparam_per_pl))
  param_min = similar(param); param_max = similar(param);
  for p in 1:npl
    param_min[1+(p-1)*nparam_per_pl] = param[1+(p-1)*nparam_per_pl]*0.98
    param_max[1+(p-1)*nparam_per_pl] = param[1+(p-1)*nparam_per_pl]*1.02
	param_min[2+(p-1)*nparam_per_pl] = param[2+(p-1)*nparam_per_pl]*0.5
    param_max[2+(p-1)*nparam_per_pl] = param[2+(p-1)*nparam_per_pl]*2.00
	param_min[3+(p-1)*nparam_per_pl] = -1.0
    param_max[3+(p-1)*nparam_per_pl] = 1.0
	param_min[4+(p-1)*nparam_per_pl] = -1.0
    param_max[4+(p-1)*nparam_per_pl] = 1.0
	param_min[5+(p-1)*nparam_per_pl] = 0.0
    param_max[5+(p-1)*nparam_per_pl] = 2pi
  end
  param_min[1+npl*nparam_per_pl] = param[1+npl*nparam_per_pl]-2.0
  param_max[1+npl*nparam_per_pl] = param[1+npl*nparam_per_pl]+2.0
  param_min[1+npl*nparam_per_pl] = 0.0
  param_max[1+npl*nparam_per_pl] = 2*param[2+npl*nparam_per_pl]
  res = compute_marginalized_likelihood_cuba(times,vels,sigma_obs,param_min,param_max,num_Sigmas=num_Sigmas)
end

#param_tmp = make_kepler_param_from_linear_output(output1,1)
#compute_marginalized_likelihood_cuba_near(times,vels,sigma_obs,param_tmp,num_Sigmas=3)

#=
include("calc_evidence_laplace_lowe.jl")
using JLD
@load "output_kepler_40_4_10.jld"
@load "output_case_4.jld"

write_evidences_file("map_linearized_circ/evidences_0001.txt",output1,use_method=4,input_filename="../data/rvs_0001.txt")
write_evidences_file("map_linearized_circ/evidences_0002.txt",output2,use_method=4,input_filename="../data/rvs_0002.txt")
write_evidences_file("map_linearized_circ/evidences_0003.txt",output3,use_method=4,input_filename="../data/rvs_0003.txt")
write_evidences_file("map_linearized_circ/evidences_0004.txt",output4,use_method=4,input_filename="../data/rvs_0004.txt")
write_evidences_file("map_linearized_circ/evidences_0005.txt",output5,use_method=4,input_filename="../data/rvs_0005.txt")
write_evidences_file("map_linearized_circ/evidences_0006.txt",output6,use_method=4,input_filename="../data/rvs_0006.txt")
@save "output_method4.jld"

write_evidences_file("map_linearized_ecc/evidences_0001.txt",output1,use_method=5,input_filename="../data/rvs_0001.txt")
write_evidences_file("map_linearized_ecc/evidences_0002.txt",output2,use_method=5,input_filename="../data/rvs_0002.txt")
write_evidences_file("map_linearized_ecc/evidences_0003.txt",output3,use_method=5,input_filename="../data/rvs_0003.txt")
write_evidences_file("map_linearized_ecc/evidences_0004.txt",output4,use_method=5,input_filename="../data/rvs_0004.txt")
write_evidences_file("map_linearized_ecc/evidences_0005.txt",output5,use_method=5,input_filename="../data/rvs_0005.txt")
write_evidences_file("map_linearized_ecc/evidences_0006.txt",output6,use_method=5,input_filename="../data/rvs_0006.txt")
@save "output_method5.jld"

write_evidences_file("map_keplerian/evidences_0001.txt",output1,use_method=6,input_filename="../data/rvs_0001.txt")
write_evidences_file("map_keplerian/evidences_0002.txt",output2,use_method=6,input_filename="../data/rvs_0002.txt")
write_evidences_file("map_keplerian/evidences_0003.txt",output3,use_method=6,input_filename="../data/rvs_0003.txt")
write_evidences_file("map_keplerian/evidences_0004.txt",output4,use_method=6,input_filename="../data/rvs_0004.txt")
write_evidences_file("map_keplerian/evidences_0005.txt",output5,use_method=6,input_filename="../data/rvs_0005.txt")
write_evidences_file("map_keplerian/evidences_0006.txt",output6,use_method=6,input_filename="../data/rvs_0006.txt")
@save "output_method6.jld"

write_evidences_file("laplace_keplerian/evidences_0001.txt",output1,use_method=3,input_filename="../data/rvs_0001.txt")
write_evidences_file("laplace_keplerian/evidences_0002.txt",output2,use_method=3,input_filename="../data/rvs_0002.txt")
write_evidences_file("laplace_keplerian/evidences_0003.txt",output3,use_method=3,input_filename="../data/rvs_0003.txt")
write_evidences_file("laplace_keplerian/evidences_0004.txt",output4,use_method=3,input_filename="../data/rvs_0004.txt")
write_evidences_file("laplace_keplerian/evidences_0005.txt",output5,use_method=3,input_filename="../data/rvs_0005.txt")
write_evidences_file("laplace_keplerian/evidences_0006.txt",output6,use_method=3,input_filename="../data/rvs_0006.txt")
@save "output_method3.jld"

=#

function direct_integration(data::Vector{T}, model::Array{T,2}, Sigma::SigmaT; FIM::PDMat{T} = make_fischer_information_matrix(model,Sigma), calclogprior::Function = x->(0.0,0.0), report_log_MAP::Bool = false ) 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]
  if report_log_MAP     return (param_bf_linalg, sigma_param,log_mode)  end
  LaplaceApprox = 0.5*nparam*log(2pi)-0.5*logdet(FIM)+log_mode + log_jacobian
  return (param_bf_linalg,sigma_param,LaplaceApprox)
end

function compute_marginalized_likelihood_const_gaussgk(times::Vector{T}, vels::Vector{T}, sigma_obs::Vector{T}; Sigma0 = PDMat(make_Sigma(times,sigma_obs,0.0)) ) where T<:Real
    X = ones(length(times),1)
    const nparam = size(X,2)
    function log_integrand(x)
        Sigma = Sigma0 + ScalMat(length(times),x)
		compute_marginalized_likelihood_fix_periods_sigmaj(times, vels, sigma_obs, Float64[], x, Sigma=Sigma,X=X)+logprior_jitter(x)
	end
    int_norm = log_integrand(0.0)
	function integrand(x)
	   exp(log_integrand(x)-int_norm)
	end
	integral = QuadGK.quadgk(integrand,0.0,2.0,99.0,reltol=1e-14)

    (log(integral[1]) + int_norm, log(integral[2]/integral[1]) )
end

function loglikelihood_zero_pl( param::Vector{T}, times::Vector{T}, vels::Vector{T}, sigma_obs::Vector{T}; calclogprior::Function = x->(0.0,0.0) ) where T<:Real
  @assert length(vels) == length(sigma_obs) # size(Sigma,1)
  nobs = length(vels)

  predict = ones(nobs)*param[1]
  delta = vels.-predict
  Sigma = make_Sigma(times,sigma_obs,param[2])
  chisq = invquad(Sigma,delta)
  loglikelihood_mode = 0.5*(-chisq-logdet(Sigma)-nobs*log(2pi))
end

function logposterior_zero_pl( param::Vector{T}, times::Vector{T}, vels::Vector{T}, sigma_obs::Vector{T}) where T<:Real
  logprior_offset(param[1])+logprior_jitter(param[2])+loglikelihood_zero_pl(param,times,vels,sigma_obs)
end

function loglikelihood_zero_pl( param::Vector{T}, data::Vector{T}, Sigma::SigmaT; calclogprior::Function = x->(0.0,0.0) ) where {T<:Real, SigmaT<:AbstractPDMat{T} }
  @assert length(data) == size(Sigma,1)
  nobs = length(data)

  predict = ones(nobs)*param[1]
  delta = data.-predict
  Sigma = make_Sigma(times,sigma_obs,param[2])
  chisq = invquad(Sigma,delta)

  loglikelihood_mode = 0.5*(-chisq-logdet(Sigma)-nobs*log(2pi))
end

function logposterior_zero_pl( param::Vector{T}, data::Vector{T}, Sigma::SigmaT ) where {T<:Real, SigmaT<:AbstractPDMat{T} }
  logprior_offset(param[1])+logprior_jitter(param[2])+loglikelihood_zero_pl(param,data,Sigma)
end

function cuba_integrand(x::Vector{T}, output::Vector{T}) where T<:Real
  param_cuba = similar(x)
  param_cuba[1] = -1000+2000*x[1]
  param_cuba[2] = 0+99*x[2]
  output[1] = exp(logposterior_zero_pl(param_cuba,times,vels,sigma_obs)-int_norm)

end

#=
int_norm = logposterior_zero_pl([0.0,2.0],times,vels,sigma_obs)
diovenne(cuba_integrand,2)
=#

function logposterior_zero_pl_marginalize_offset( sigmaj::T, times::Vector{T}, vels::Vector{T}, sigma_obs::Vector{T}; xmax::T=Cmax, xmin::T=-xmax) where T<:Real
  #@assert length(vels) == length(times) = length(sigma_obs)
  param = zeros(2)

  function log_integrand(x::T)
    param[1] = x
    logprior_offset(param[1])+logprior_jitter(param[2])+loglikelihood_zero_pl(param,times,vels,sigma_obs)
  end

  param[2] = sigmaj
  int_norm = log_integrand(0.)
  function integrand(x::T)
     exp(log_integrand(x)-int_norm)
  end

  log(QuadGK.quadgk(integrand,xmin,xmax)[1])+int_norm
end

#=
data = readdlm("../data/rvs_0001.txt");
times = data[:,1];
vels = data[:,2];
sigma_obs = data[:,3];
Sigma0 = make_Sigma(times,sigma_obs,0.0)
Sigma2 = make_Sigma(times,sigma_obs,2.0)
=#

#=
include("calc_evidence_laplace_lowe.jl")
zero_pl_tests =
	[("rvs_0001.txt", [-0.728509750583953064E+00, 0.208717785347551521E+01]),
	("rvs_0002.txt", [ 0.653665406348636679E+01, 0.171886843613598783E+01]),
	("rvs_0003.txt", [-0.869753669161968901E+01, 0.994907921829712771E+00]),
	("rvs_0004.txt", [-0.658082070168154587E+01, 0.682875011391089748E+00]),
	("rvs_0005.txt", [-0.486556982931438142E+01, 0.930406316677828249E+00]),
	("rvs_0006.txt", [-0.108329221460122653E+02, 0.115487852990739137E+01])]

filename_out = "tmp.out"
file_out = open(filename_out,"w")
for test in zero_pl_tests
  filename = joinpath("..","data",test[1])
  data = readdlm(filename)
  times = data[:,1];
  vels = data[:,2];
  sigma_obs = data[:,3];
  lnlikelihood = loglikelihood_zero_pl(test[2],times,vels,sigma_obs)
  println(file_out,test, ":  ", lnlikelihood)
  test2 = copy(test[2])
  test2[1] = 0.0
  lnlikelihood = loglikelihood_zero_pl(test2,times,vels,sigma_obs)
  println(file_out,(test[1],test2), ":  ", lnlikelihood)
end
close(file_out)

=#

#loglikelihood_zero_pl([-0.639837062706153,0.0],times,vels,sigma_obs)


#
# This example assumes below packages are available in your environment
#

using Pkg, DrWatson

#@quickactivate "StatisticalRethinkingStan"
using StatsModelComparisons, StanSample
using StatsFuns, CSV, Random
using ParetoSmooth
using StatisticalRethinking

ProjDir = @__DIR__

df = CSV.read(joinpath(ProjDir, "cars.csv"), DataFrame)

cars_stan = "
data {
    int N;
    vector[N] speed;
    vector[N] dist;
}
parameters {
    real a;
    real b;
    real sigma;
}
transformed parameters{
    vector[N] mu;
    mu = a + b * speed;
}
model {
    a ~ normal(0, 100);
    b ~ normal(0, 10);
    sigma ~ exponential(1);
    dist ~ normal(mu, sigma)    ;
}
generated quantities {
    vector[N] log_lik;
    for (i in 1:N)
        log_lik[i] = normal_lpdf(dist[i] | mu[i], sigma);
}
"

Random.seed!(1)
cars_stan_model = SampleModel("cars.model", cars_stan)
data = (N = size(df, 1), speed = df.Speed, dist = df.Dist)
rc = stan_sample(cars_stan_model; data)
println()

if success(rc)
    #stan_summary(cars_stan_model, true)
    nt_cars = read_samples(cars_stan_model);
end

log_lik = nt_cars.log_lik';
n_sam, n_obs = size(log_lik)
lppd = reshape(logsumexp(log_lik .- log(n_sam); dims=1), n_obs);

pwaic = [var(log_lik[:, i]) for i in 1:n_obs];
@show -2(sum(lppd) - sum(pwaic))
println()

@show waic(log_lik)
println()

#Random.seed!(123)
loo, loos, pk = psisloo(log_lik);
@show -2loo

#Random.seed!(123)
ll = reshape(nt_cars.log_lik, 50, 1000, 4);
psis_ll = psis(ll);

lwp = deepcopy(ll);
lwp += psis_ll.weights;
lwpt = Matrix(reshape(lwp, 50, 4000)');
loos = reshape(logsumexp(lwpt; dims=1), size(lwpt, 2));

@show loo = sum(loos)
@show 2loo

pk_plot(pk)
savefig(joinpath(ProjDir, "pk.png"))

pk_plot(psis_ll.pareto_k)
savefig(joinpath(ProjDir, "pareto_k.png"))


using Revise # reduce recompilation time
using Plots
# using Documenter
using LinearAlgebra
using SparseArrays
using BenchmarkTools
using UnPack

push!(LOAD_PATH, "./examples/EntropyStableEuler.jl/src")
using EntropyStableEuler
using EntropyStableEuler.Fluxes1D
const γ = 1.4

function wavespeed_1D(rho,rhou,E)
    p = pfun_nd(rho,(rhou,),E)
    cvel = @. sqrt(γ*p/rho)
    return @. abs(rhou/rho) + cvel
end
unorm(U) = sum(map((x->x.^2),U))
function pfun_nd(rho, rhoU, E)
    rhoUnorm2 = unorm(rhoU)./rho
    return @. (γ-1)*(E - .5*rhoUnorm2)
end

K = 200
T = 6.0

# Bl = -0.5
# Br = 0.5
Bl = 0.0
Br = 9.0
VX = LinRange(Bl,Br,K+1)
# Assume uniform intervals
m_i = (Br-Bl)/K # mass
C = spdiagm(-1 => -1/2*ones(K-1), 1 => 1/2*ones(K-1))# stiffness
C[1,end] = -1/2
C[end,1] = 1/2

x = zeros(K,1)
for i = 1:K
    x[i] = (VX[i+1]+VX[i])/2
end
"""Initial condition"""
rho_x(x) = (x <= 0.0) ? 1.0 : 0.125
u_x(x) = 0.0
p_x(x) = (x <= 0.0) ? 1.0 : 0.1
# rho_x(x) = (x <= 0.0) ? 1.0 : 0.5
# u_x(x) = 1.0
# p_x(x) = (x <= 0.0) ? 1.0 : 1.0
rho_x(x) = (x <= 3.0) ? 1.0 : 0.001
u_x(x) = 0.0
p_x(x) = (x <= 3.0) ? 0.1 : 1e-7

rho = @. rho_x(x)
u = @. u_x(x)
p = @. p_x(x)
Q = primitive_to_conservative(rho,u,p)

function rhs_IDP(Q,K,m_i,C)
    p = pfun_nd.(Q[1],Q[2],Q[3])
    flux = zero.(Q)
    @. flux[1] = Q[2]
    @. flux[2] = Q[2]^2/Q[1]+p
    @. flux[3] = Q[3]*Q[2]/Q[1]+p*Q[2]/Q[1]

    dfdx = (zeros(K),zeros(K),zeros(K))
    for i = 1:K
        for c = 1:3
            dfdx[c][i] += C[i,mod1(i-1,K)]*(flux[c][mod1(i-1,K)]+flux[c][i])+
                          C[i,mod1(i+1,K)]*(flux[c][mod1(i+1,K)]+flux[c][i])
        end
    end

    visc = (zeros(K),zeros(K),zeros(K))
    for i = 1:K
        wavespd_curr = abs(wavespeed_1D(Q[1][i],Q[2][i],Q[3][i]))
        wavespd_R = abs(wavespeed_1D(Q[1][mod1(i+1,K)],Q[2][mod1(i+1,K)],Q[3][mod1(i+1,K)]))
        wavespd_L = abs(wavespeed_1D(Q[1][mod1(i-1,K)],Q[2][mod1(i-1,K)],Q[3][mod1(i-1,K)]))
        wavespd_curr_m = abs(wavespeed_1D(Q[1][i],-Q[2][i],Q[3][i]))
        wavespd_R_m = abs(wavespeed_1D(Q[1][mod1(i+1,K)],-Q[2][mod1(i+1,K)],Q[3][mod1(i+1,K)]))
        wavespd_L_m = abs(wavespeed_1D(Q[1][mod1(i-1,K)],-Q[2][mod1(i-1,K)],Q[3][mod1(i-1,K)]))
        dL = 1/2*max(wavespd_curr,wavespd_L)
        dR = 1/2*max(wavespd_curr,wavespd_R)
        for c = 1:3
            visc[c][i] = dL*(Q[c][mod1(i-1,K)]-Q[c][i]) + dR*(Q[c][mod1(i+1,K)]-Q[c][i])
        end
    end

    return -1/m_i.*(dfdx.-visc)
end

gr(size=(300,300),ylims=(0,1.2),legend=false,markerstrokewidth=0,markersize=1)
plt = plot()

dt = 0.0001
t = 0.0
Nsteps = Int(T/dt)
Q = collect(Q)
@gif for i = 1:Nsteps
    rhsQ = rhs_IDP(Q,K,m_i,C)
    global @. Q =  Q + dt*rhsQ

    if i%10==0 || i==Nsteps
        println("Number of time steps $i out of $Nsteps")

        scatter(x,Q[1])
        #for k = 1:K
            #xp = LinRange(VX[k], VX[k+1], 3)
            #scatter!(xp,Q[1][k]*ones(length(xp)))
        #end
    end
end every 1000

# p = pfun_nd.(Q[1],Q[2],Q[3])
# for k = 1:K
#     xp = LinRange(VX[k], VX[k+1], 3)
#     scatter!(xp,Q[1][k]*ones(length(xp)))

# end
# plt


using FractalTools
using Makie

# Contruct a tesslation
tridln = Node(TriDelaunay([0., 0.], [1., 0], [0.5, 1.], addboundarypoints=true))

# Definea an observables
msh = lift(tridln) do val
    tomesh(val.delaunay)
end
vtx = lift(msh) do val
    val.position
end

# Test point
tpnt = Node(getpoint(tridln[]))
tripts = lift(tpnt) do val
    tidx = simplices(tridln[])[locate(tridln[], val), :]
    map(val -> Point(val...), eachrow(points(tridln[])[tidx, :]))
end

# Plot tesselation
fig = Figure(resolution=(500, 500), scale_plot = false, camera = Makie.campixel!)
ax = fig[1, 1] =  Axis(fig)
mesh!(ax, msh, color=:lightblue)
wireframe!(ax, msh, color=:black)
scatter!(ax, vtx, color=:black)
scatter!(ax, tripts, color=:red, markersize=10)
scatter!(ax, tpnt, color=:orange, markersize=10)

# Add mouse button interaction
on(ax.scene.events.mousebuttons) do buttons
   if ispressed(ax.scene, Mouse.left)
        area = pixelarea(ax.scene)[]
        mp = ax.scene.events.mouseposition[]
        mp = Point2f0(mp) .- minimum(area)
        pos = to_world(ax.scene, mp)
        @info "$pos added"
        tpnt[] = pos
   end
end

# Add button interaction.
on(ax.scene.events.keyboardbuttons) do button
    if ispressed(button, Keyboard.a)
        pnt = addpoint!(tridln[])
        @info "$pnt is added"
    end
    if ispressed(button, Keyboard.f)
        finegrain!(tridln[], 10)
        @info "delaunay is fine grained"
    end
    tridln[] = tridln[]   # Trigger change
end

# Display scene
display(fig)


using Distributions
using StaticArrays
import Printf.@sprintf
import Distributions.RealInterval

abstract type AbstractUncertainValue end

Broadcast.broadcastable(d::AbstractUncertainValue) = Ref(d)

##############################
# Support of an observation
##############################
import Distributions.support

"""
    support(o::AbstractUncertainValue) -> IntervalArithmetic.Interval

Find the support of the uncertain observation.
"""
function support(o::AbstractUncertainValue)
    supp = support(o.distribution)
    lowerbound = supp.lb
    upperbound = supp.ub
    @interval(lowerbound, upperbound)
end

export support

##############################
# Intersection of the supports
##############################
"""
    support_overlap(o1::AbstractUncertainValue, o2::AbstractUncertainValue)

Compute the overlap in the supports of two uncertain observations.
"""
function support_overlap(uval1::AbstractUncertainValue, uval2::AbstractUncertainValue)
    s1 = support(uval1)
    s2 = support(uval2)

    if s1 isa RealInterval
        s1 = interval(s1...)
    end

    if s2 isa RealInterval
        s2 = interval(s2...)
    end
    s1 ∩ s2
end

export support_overlap

##################################################################
# Intersection of two UncertainValues as a mixture model
##################################################################
import Base.intersect, Base.∩

"""
    intersect(o1::AbstractUncertainValue, o2::AbstractUncertainValue)

Compute the intersection between two uncertain observations probabilistically.
The intersection is represented as a mixture model of the distributions
furnishing the observations.
"""
function intersect(o1::AbstractUncertainValue, o2::AbstractUncertainValue)
    # Create a mixture model representing the intersection
    if support_overlap(o1, o2) == ∅
        throw(DomainError((o1, o2), "intersect(o1, o2) == ∅. Cannot compute mixture model."))
    end

    MixtureModel([o1.distribution, o2.distribution])
end

export intersect

###################
# Pretty printing
###################

function summarise(o::AbstractUncertainValue)
    _type = typeof(o)
    "$_type"
end
Base.show(io::IO, q::AbstractUncertainValue) = print(io, summarise(q))

dimension(usv::AbstractUncertainValue) = 1

Base.max(uv::AbstractUncertainValue) = maximum(uv.distribution)
Base.min(uv::AbstractUncertainValue) = minimum(uv.distribution)
Base.minimum(uv::AbstractUncertainValue) = minimum(uv.distribution)
Base.maximum(uv::AbstractUncertainValue) = maximum(uv.distribution)

Base.maximum(uvs::AbstractVector{AbstractUncertainValue}) =
    maximum([maximum(uv) for uv in uvs])
Base.minimum(uvs::AbstractVector{AbstractUncertainValue}) =
    minimum([minimum(uv) for uv in uvs])

export
AbstractUncertainValue


module MyPackage
   # Declare dependencies
   using ForwardDiff, DelimitedFiles

   # Fuctions to export for external use
   export my_f, my_f_der

   # define your function here
   my_f(x)=3x^2-5x+1;

   # automated differentiation using ForwardDiff Package.
   my_f_der(t)=ForwardDiff.derivative(x->my_f(x),t)

end


using LinearAlgebra
using TensorOperations

####################################################################
#
# construct \Omega matrix
#

function Omega_uu(N::Int)
    id = Matrix{Float64}(I, N, N)
    zs = zeros(N,N)
    return [zs -id; id zs]
end

function Omega_dd(N::Int)
    id = Matrix{Float64}(I, N, N)
    zs = zeros(N,N)
    return [zs id; -id zs]
end

function change_coordinate(X::Array{<:Number},S::AbstractArray{<:Number,2})
    size_X =  size(X)
    N = div(size_X[1], 2)
    rank_X = ndims(X)
    X_prime = zeros(Complex{Float64},size_X...)
    indeces_set = Iterators.product(ntuple(i->1:2N,rank_X)...)

    for inds_after in indeces_set
        for inds_before in indeces_set
            Xprime_element = X[inds_before...];
            for i in 1:rank_X
                Xprime_element = Xprime_element * S[inds_before[i],inds_after[i]]
            end;
            X_prime[inds_after...] += Xprime_element
        end;
    end;
    return X_prime
end

############################################################
#symmetrize tensors
function symmetrize_tensor(X::Array{<:Number})
    rank_X = ndims(X)
    if rank_X ==1; return X; end;

    #construct permutation_group
    permutation_group = Set([])
    inds_for_perms = Iterators.product(ntuple(i->1:(rank_X-i+1),rank_X)...)
    size_perm_set = length(inds_for_perms)
    for perm in inds_for_perms
        perm = collect(perm)
        vec_perm  = zeros(Int64,rank_X) # permutated [1:N]
        #generate an element of the permutation group from perm
        for i in 1:rank_X
            j = 1
            while true
                if perm[i] == 1 && vec_perm[j] == 0; vec_perm[j]=i; break;
                elseif vec_perm[j] != 0; j+=1;
                else j+=1; perm[i]-=1
                end;
            end;
        end;
        push!(permutation_group,vec_perm)
    end;

    indeces_set = Iterators.product(ntuple(i->1:2N,rank_X)...)
    X_sym = zeros(eltype(X),size(X)...)
    for inds in indeces_set
        for perm in permutation_group
            inds_perm = Tuple(permute!([inds...],perm)) #permute inds by perm
            X_sym[inds...] += 1.0 / size_perm_set * X[inds_perm...]
        end;
    end;
    return X_sym
end;

############################################################
# get tensors of the effecitive Hamiltonian including Lindblad terms
#
function get_effective_Hamiltonian_tensors!(A′::Array{Complex{Float64},1},B′::Array{Complex{Float64},2},C′::Array{Complex{Float64},3},D′::Array{Complex{Float64},4},A::Array{Float64,1},B::Array{Float64,2},C::Array{Float64,3},D::Array{Float64,4},vec_EF::Array{Array{Complex{Float64}},2},μ₁::Array{Float64,1},μ₂::Array{Float64,2})

    N = Int(length(A)/2)
    Ω_dd = Omega_dd(N)
    A′ .= A; B′ .= B; C′ .= C; D′ .= D;

    for i in 1:size(vec_EF)[1]
        E = vec_EF[i,1];
        F = vec_EF[i,2];
        real_E = real.(E);
        real_F = real.(F);
        exp_Re_L = dot(real_E,μ₁) + tr(real_F*μ₂)
        @tensor begin
            A′[i]       += im * exp_Re_L * E[i]
            B′[i,j]     += im * exp_Re_L * F[i,j]
            B′[i,j]     -= im *  real_E[i] * E[j]
            C′[i,j,k]   -= im *  real_E[i] * F[j,k]
            C′[i,j,k]   -= im *  real_F[i,j] *E[k]
            D′[i,j,k,l] -= im *  real_F[i,j] * F[k,l]
        end
    end

    return (A′,symmetrize_tensor(B′),symmetrize_tensor(C′),symmetrize_tensor(D′))
end

############################################################
# get raw 2rd-order moments from central moments
#
function get_2nd_raw_moment(μ::Array{<:Number,1},μ_dd::Array{<:Number,2})

    @tensor begin
        μ′_dd[i,j]  := μ_dd[i,j] + μ[i] * μ[j]
    end

    return symmetrize_tensor(μ′_dd)
end
############################################################
# get raw 3rd-order moments from central moments
#
function get_3rd_raw_moment(μ::Array{<:Number,1},μ_dd::Array{<:Number,2},μ_ddd::Array{<:Number,3})

    @tensor begin
        μ′_ddd[i,j,k]  := μ_ddd[i,j,k] + 3.0 * μ_dd[i,j] * μ[k] + μ[i] * μ[j] * μ[k]
    end

    return symmetrize_tensor(μ′_ddd)
end

############################################################
# get raw 4th-order moments from central moments
#
function get_4th_raw_moment(μ::Array{<:Number,1},μ_dd::Array{<:Number,2},μ_ddd::Array{<:Number,3},μ_dddd::Array{<:Number,4})

    @tensor begin
        μ′_dddd[i,j,k,l]  := μ_ddd[i,j,k,l] + 4. * μ_ddd[i,j,k] * μ[l] + 6.0 * μ_dd[i,j] * μ[k] * μ[l] + μ[i] * μ[j] * μ[k]
    end

    return symmetrize_tensor(μ′_dddd)
end

#################################################################################################
#  get symmetrized z coordinate tensors from a coordinate tensors
#
#

function get_sym_tensors_with_z_coord(array_of_tensors::Array{Array{Complex{Float64}},1})

    ret = Array{Array{Float64}}(undef,length(array_of_tensors))
    for i in 1:length(ret)
        ret[i] = real.(symmetrize_tensor(converta2z(array_of_tensors[i])))
    end
    return Tuple(ret)
end


using CompScienceMeshes
using BEAST

fn = joinpath(dirname(@__FILE__),"assets","torus.msh")
m = CompScienceMeshes.read_gmsh_mesh(fn)

X = raviartthomas(m)
Y = buffachristiansen(m)

N = BEAST.NCross()
K = Maxwell3D.doublelayer(gamma=0.0)

verts = skeleton(m,0)
edges = skeleton(m,1)
faces = skeleton(m,2)
Λ = connectivity(verts, edges)
Σ = connectivity(faces, edges)

M = assemble(K+0.5N,Y,X)

using LinearAlgebra
s = svdvals(M)

using Plots
plot(log10.(s))

x = nullspace(M, 1.1*s[end])

norm(x)
norm(M*x)
norm(Σ'*Λ)


using NLPModels
using SolverTools

abstract type TR1DModel{T} end

abstract type QuadraticModel{T} <: TR1DModel{T} end

"Quadratic model solution"
function solveModel(M:: QuadraticModel{T}, t :: T) where T
    return -M.g/M.H
end

"Quadratic model solution in region  Δn ≤ t ≤ Δp "
function solveModel(M:: QuadraticModel{T}, t :: T, Δp :: T, Δn :: T) where T
    dN = -M.g/M.H

    if (Δm(M, Δp, Δn) < 0.0) #| (Δn==0.0)
        d=Δp
        @debug "dp"
    else
        d=Δn
        @debug "dn"
    end

    if  (dN < Δp) & (dN > Δn) & (Δm(M, d, dN) > 0.0)
        d=dN
        @debug "dN = ", d
    end

    return d
end

"Quadratic model evaluation"
function m(M:: QuadraticModel{T}, t :: T) where T
    return M.f + M.g * t + T(0.5)*M.H * t^2
end

"Quadratic model predicted reduction"
function Δm(M:: QuadraticModel{T}, t :: T) where T
    return M.g * t  +  T(0.5) * M.H * t^2
end

"Quadratic model predicted reduction given two points"
function Δm(M:: QuadraticModel{T}, t1 :: T, t2 :: T) where T
    #q1 = M.g * t1  +  T(0.5) * M.H * t1^2
    #q2 = M.g * t2  +  T(0.5) * M.H * t2^2
    #diff = q1 - q2
    diff = (M.g + T(0.5) * M.H * (t1+t2)) * (t1 - t2)
    return diff
end

"Taylor quadratic polynomial model"
mutable struct Taylor2{T} <: QuadraticModel{T}
    f :: T
    g :: T
    H :: T
end
Taylor2{T}() where T = Taylor2{T}(0,0,0)

function initModel(M:: Taylor2{T}, ϕ :: OneDModel, t :: T, ϕt :: T, dϕt :: T) where T
    M.f = ϕt
    M.g = dϕt
    M.H = hess(ϕ, t)
    return M
end

function updateModel!(M:: Taylor2{T}, ϕ :: OneDModel, t :: T, ϕt :: T, dϕt :: T) where T
    M.f = ϕt
    M.g  = dϕt
    M.H = hess(ϕ, t)
    return M
end

function initH_N(ϕ :: LineModel, t :: Float64, bidon)
    return hess(ϕ, t)
end

function updateH_N(ϕ :: LineModel, t :: Float64, dϕ :: Float64)
    return hess(ϕ, t)
end

"Secant quadratic model"
mutable struct secant{T} <: QuadraticModel{T}
    f :: T
    g :: T
    H :: T
end
secant{T}() where T = secant{T}(0,0,0)

function initModel(M:: secant{T}, ϕ :: OneDModel, t :: T, ϕt :: T, dϕt :: T) where T
    M.f = ϕt
    M.g = dϕt
    M.H = initH_s!(ϕ, t, dϕt)
    return M
end

function updateModel!(M:: secant{T}, ϕ :: OneDModel, t :: T, ϕt :: T, dϕt :: T) where T
    M.f = ϕt
    M.g = dϕt
    M.H = updateH_s!(ϕ, t, dϕt)
    return M
end

let (tpred, dϕpred) = (NaN, NaN)

    global function initH_s!(ϕ :: OneDModel, t :: Float64, dϕ :: Float64)
        (tpred, dϕpred) = (t, dϕ)
        return 1.0
    end

    global function updateH_s!(ϕ :: OneDModel, t :: Float64, dϕt :: Float64)
        s = t - tpred
        y = dϕt - dϕpred
        H = y/s

        tpred, dϕpred = t, dϕt

        return H
    end
end


using D4M

n = 2 .^ [6 6 7 8 9 10 11 12 13 14]

K = 8

MaxGB = 2
MaxGF = 4 * 1.6

m = K .* n
assoc_gbytes = zeros(1,length(n))
assoc_flops = zeros(1,length(n))
assoc_gflops = zeros(1,length(n))
assoc_time = zeros(1,length(n))

for i = 1:length(n)

    ii = round(Int, floor(rand(m[i]) .* n[i]) +1)
    jj = round(Int, floor(rand(m[i]) .* n[i]) +1)
    A = Assoc(ii,jj,1.0)

    ii = round(Int, floor(rand(m[i]) .* n[i]) +1)
    jj = round(Int, floor(rand(m[i]) .* n[i]) +1)
    B = Assoc(ii,jj,1.0)

    tic()
        C = CatValMul(A,B)
    assoc_time[i] = toq()
    assoc_flops[i] = 2*sum(C)
    ii, jj, vv = find(C)
    assoc_gbytes[i] = assoc_gbytes[i] + (length(ii) + length(jj)) + 8 .* m[i] ./ 1e9
    assoc_gflops[i] = assoc_flops[i] ./ assoc_time[i] ./ 1e9
    print("Time: ", assoc_time[i])
    print(", GFlops: ", assoc_gflops[i])
    println(", GBytes: ", assoc_gbytes[i])
end


@testset "Maximum likelihood" begin
    b = 1.0
    M0 = 0.5
    mags = rand(GRclassic(b,M0), 100000)
    filter!(x->rand() < cdf(Normal(0.8,0.2),x), mags)
    ΔM = 0.1

    M0 = 1.25 - ΔM/2
    mags = mags[mags .>= M0]

    @test round.( bvalue(fit_mle(GRclassic,mags,M0)), digits=3) == 0.995
    @test round.( fit_stderr(fit_mle(GRclassic,mags,M0), mags), digits=3) == 0.007

    @test round.( bvalue(fit_binned(GRclassic,mags,M0,ΔM)), digits=3) == 0.999
    @test round.( fit_stderrbinned(fit_binned(GRclassic,mags,M0,ΔM), mags, ΔM), digits=3) == 0.007

end


using ModelingToolkit
using Base.Test

struct Lorenz <: AbstractComponent
  x::Variable
  y::Variable
  z::Variable
  σ::Variable
  ρ::Variable
  β::Variable
  eqs::Vector{Expression}
end
function generate_lorenz_eqs(t,x,y,z,σ,ρ,β)
  D = Differential(t)
  [D*x ~ σ*(y-x)
   D*y ~ x*(ρ-z)-y
   D*z ~ x*y - β*z]
end
Lorenz(t) = Lorenz(first(@DVar(x(t))),first(@DVar(y(t))),first(@DVar(z(t))),first(@Param(σ)),first(@Param(ρ)),first(@Param(β)),generate_lorenz_eqs(t,x,y,z,σ,ρ,β))

@IVar t
lz1 = Lorenz(t)
lz2 = Lorenz(t)
Expression[lz1.x ~ lz2.x
           lz1
           lz2]


const tolerance_deg = 0.00001

"""
    normalize(value::Float64[, lower::Float64 = 0.0, upper::Float64 = 360.0])

Normalize a `value` to stay within the `lower` and `upper` limit.

Example:
normalize(370.0, lower=0.0, upper=360.0)
10.0
"""
function normalize(value::Float64, lower::Float64 = 0.0, upper::Float64 = 360.0)
    result = value
    if result > upper || value == lower
        value = lower + abs(value + upper) % (abs(lower) + abs(upper))
    end
    if result < lower || value == upper
        value = upper - abs(value - lower) % (abs(lower) + abs(upper))
    end
    value == upper ? result = lower : result = value
    return result
end

"""
    normalize(point::Point)

Normalize the point to keep the longitude between -180 and 180 degrees and the
latitude between -90 and 90.

Source: gist.github.com/missinglink/ > wrap.js
"""
function normalize(point::Point)
    ϕ, λ = _normalize(point.ϕ, point.λ)
    return Point(ϕ, λ)
end

function _normalize(pointϕ::Float64, pointλ::Float64)
    ϕ = pointϕ
    λ = pointλ
    quadrant = floor(Int64, abs(ϕ)/90) % 4
    pole = ϕ > 0 ? 90.0 : -90.0
    offset = ϕ % 90.0

    if quadrant == 0
        ϕ = offset
    elseif quadrant == 1
        ϕ = pole - offset
        λ += 180.0
    elseif quadrant == 2
        ϕ = -offset
        λ += 180.0
    elseif quadrant == 3
        ϕ = - pole + offset
    end

    if λ > 180.0 || λ < -180.0
        λ -= floor((λ + 180.0) / 360) * 360
    end

    return ϕ, λ
end


function meshgrid(vx::AbstractVector{T}, vy::AbstractVector{T}) where {T}
    m, n = length(vy), length(vx)
    gx = reshape(repeat(vx, inner = m, outer = 1), m, n)
    gy = reshape(repeat(vy, inner = 1, outer = n), m, n)
    return gx, gy
end

function plotedges(edges)::Tuple{Array{Float64,1}, Array{Float64,1}}
    x = Array{Float64,1}()
    y = Array{Float64,1}()
    for edge in edges
        push!(x, getx(geta(edge)))
        push!(x, getx(getb(edge)))
        push!(x, NaN)
        push!(y, gety(geta(edge)))
        push!(y, gety(getb(edge)))
        push!(y, NaN)
    end
    return x, y
end


# Testing of functions in norms.jl

using NumericExtensions
using Base.Test

x1 = randn(6)
y1 = randn(6)

x2 = randn(5, 6)
y2 = randn(5, 6)

x3 = randn(3, 4, 5)
y3 = randn(3, 4, 5)

xs = Any[x1, x2, x3]
ys = Any[y1, y2, y3]
ps = Any[1, 2, 3, Inf]

tdims = Any[
    Any[1],   # N = 1
    Any[1, 2, (1, 2)], # N = 2
    Any[1, 2, 3, (1, 2), (1, 3), (2, 3), (1, 2, 3)] # N = 3
]

safe_vnorm(x, p) = p == 1 ? sum(abs(x)) :
                   p == 2 ? sqrt(sum(abs2(x))) :
                   p == Inf ? maximum(abs(x)) :
                   sum(abs(x) .^ p) .^ inv(p)

safe_vnorm(x, p, reg) = p == 1 ? sum(abs(x), reg) :
                        p == 2 ? sqrt(sum(abs2(x), reg)) :
                        p == Inf ? maximum(abs(x), reg) :
                        sum(abs(x) .^ p, reg) .^ inv(p)

# vnorm

for x in xs, p in ps
    @test_approx_eq vnorm(x, p) safe_vnorm(x, p)
end

for x in xs, reg in tdims[ndims(x)], p in ps
    @test_approx_eq vnorm(x, p, reg) safe_vnorm(x, p, reg)
end

# vnormdiff

for (x, y) in zip(xs, ys), p in ps
    @test_approx_eq vnormdiff(x, y, p) safe_vnorm(x - y, p)
end

for (x, y) in zip(xs, ys), reg in tdims[ndims(x)], p in ps
    @test_approx_eq vnormdiff(x, y, p, reg) safe_vnorm(x - y, p, reg)
end

# normalize

for x in xs, p in ps
    r = x ./ vnorm(x, p)
    @test_approx_eq normalize(x, p) r

    xc = copy(x)
    normalize!(xc, p)
    @test_approx_eq xc r
end

for x in xs, reg in tdims[ndims(x)], p in ps
    r = x ./ vnorm(x, p, reg)
    @test_approx_eq normalize(x, p, reg) r

    xc = copy(x)
    normalize!(xc, p, reg)
    @test_approx_eq xc r
end


"""
`TransformedKernel(k::Kernel,t::Transform)`
Return a kernel where inputs are pretransformed by `t` : `k(t(x),t(x'))`
"""
struct TransformedKernel{Tk<:Kernel,Tr<:Transform} <: Kernel
    kernel::Tk
    transform::Tr
end

"""
```julia
    transform(k::BaseKernel, t::Transform) (1)
    transform(k::BaseKernel, ρ::Real) (2)
    transform(k::BaseKernel, ρ::AbstractVector) (3)
```
(1) Create a TransformedKernel with transform `t` and kernel `k`
(2) Same as (1) with a `ScaleTransform` with scale `ρ`
(3) Same as (1) with an `ARDTransform` with scales `ρ`
"""
transform

transform(k::BaseKernel, t::Transform) = TransformedKernel(k, t)

transform(k::BaseKernel, ρ::Real) = TransformedKernel(k, ScaleTransform(ρ))

transform(k::BaseKernel,ρ::AbstractVector) = TransformedKernel(k, ARDTransform(ρ))

kernel(κ) = κ.kernel

kappa(κ::TransformedKernel, x) = kappa(κ.kernel, x)

metric(κ::TransformedKernel) = metric(κ.kernel)

Base.show(io::IO,κ::TransformedKernel) = printshifted(io,κ,0)

function printshifted(io::IO,κ::TransformedKernel,shift::Int)
    printshifted(io,κ.kernel,shift)
    print(io,"\n"*("\t"^(shift+1))*"- $(κ.transform)")
end


# Plot the loss function on a random 1D line
# Authors: Ali Siahkoohi, alisk@gatech.edu
# Date: December 2020

using DrWatson
@quickactivate :ExtendedConv

using Flux
using PyPlot
using Seaborn
using Printf
using LinearAlgebra
using Optim
using JLD2
using Random
using ProgressMeter
seed1 = 21
seed2 = 12

Random.seed!(seed1)

matplotlib.use("Agg")
set_style("whitegrid")
rc("font", family="serif", size=10)
font_prop = matplotlib.font_manager.FontProperties(
    family="serif",
    style="normal",
    size=10
)
sfmt=matplotlib.ticker.ScalarFormatter(useMathText=true)
sfmt.set_powerlimits((0, 0))

sim_name = "2D-visualization"

nx = 16
ny = 16
nc = 1
batchsize = 1
k = 3
nlayers = 4
lambda = .04f0
f_tol = 1f-6
iterations = 100

G = Chain([Conv((k, k), nc => nc, sigmoid, pad=(1, 1), stride=(1, 1))  for i=1:nlayers]...)

w = deepcopy(Flux.params(G))
w_true = cat([w[i] for i=1:2:2*nlayers]..., dims=3)[:, :, :, 1]

C = extended_conv_op(nx, ny; nc=nc, batchsize=1, k=k, nlayers=nlayers, w=w_true)

z_std = 1f-2
z = z_std*randn(Float32, nx, ny, nc, batchsize)
y = G(z)

G_ext(x) = reshape(C(reshape(x, :)), nx, ny, nc, 1)[end:-1:1, end:-1:1]
perm_y(y, idx) = reshape(y[end:-1:1, end:-1:1, :, idx], :)

αs = range(-4f0, 4f0, step=4f-1)
βs = range(-4f0, 4f0, step=4f-1)
n_grid = length(αs)

Random.seed!(seed2)

opt_var_proj = Optim.Options(
    iterations=iterations,
    f_tol=f_tol,
    show_trace=false,
    store_trace=false,
    show_every=1
)

num_figs = 4
p = Progress(num_figs*n_grid^2)

for fig_idx = 1:num_figs
    v1 = filter_normalization(w)
    v2 = filter_normalization(w)

    loss = zeros(Float32, length(αs), length(βs))
    loss_ext = zeros(Float32, length(αs), length(βs))

    for (i, α) in enumerate(αs)
        for (j, β) in enumerate(βs)
            Base.flush(Base.stdout)

            idx = rand(1:batchsize)

            ŵ = w + α*v1 + β*v2
            Flux.loadparams!(G, ŵ)

            ŵ0 = cat([ŵ[i] for i=1:2:2*nlayers]..., dims=3)[:, :, :, 1]
            C = extended_conv_op(nx, ny; nc=nc, batchsize=1, k=k, nlayers=nlayers, w=ŵ0)
            w_ex_0 = cat(
                [randn(Float32, size(C[i].mask)).*C[i].mask for i=1:length(C)]...,
                dims=3
            )

            loss_ext[i, j] = var_proj_optim(
                reshape(z[:, :, :, idx:idx], :), perm_y(y, idx), C, w_ex_0, ŵ0, lambda,
                opt_var_proj
            )[1]

            loss[i, j] = 5f-1*norm(G(z[:, :, :, idx:idx]) - y[:, :, :, idx:idx])^2

            ProgressMeter.next!(
                p;
                showvalues = [
                    ("Figure number", fig_idx),
                    (:α, α),
                    (:β, β),
                    ("Objectibe", loss[i, j]),
                    ("Extended objective", loss_ext[i, j])
                ]
            )

        end
    end

    save_dict = @strdict loss loss_ext lambda n_grid f_tol iterations batchsize z_std nlayers nx ny k v1 v2 seed1 seed2 w z y αs βs
    save_path = plotsdir(sim_name, savename(save_dict; digits=12))

    @tagsave(
        datadir(sim_name, savename(save_dict, "jld2"; digits=12)),
        save_dict;
        safe=true
    )

    loss = (loss .- minimum(loss))./maximum(loss)
    loss_ext = (loss_ext .- minimum(loss_ext))./maximum(loss_ext)

    extent = [minimum(αs), maximum(αs), maximum(βs), minimum(βs)]

    fig = figure("conventional", figsize=(4, 4))
    contour(αs, βs, loss, 25, cmap=Seaborn.ColorMap("mako"), linewidths=1.5,
        vmin=0.0, vmax=1f0)
    title("Loss landscape")
    colorbar(fraction=0.057, pad=0.01, format=sfmt)
    PyPlot.scatter(
        βs[argmin(loss)[2]], αs[argmin(loss)[1]], s=30.0,
        color="#b80000", marker="*"
    )
    xlabel(L"$\beta$")
    ylabel(L"$\alpha$")
    safesave(joinpath(save_path, "loss.png"), fig)
    close(fig)

    fig = figure("conventional", figsize=(4, 4))
    contour(αs, βs, loss_ext, 25, cmap=Seaborn.ColorMap("mako"), linewidths=1.5,
        vmin=0.0, vmax=1f0)
    title("Extended loss landscape")
    colorbar(fraction=0.057, pad=0.01, format=sfmt)
    PyPlot.scatter(
        βs[argmin(loss_ext)[2]], αs[argmin(loss_ext)[1]], s=30.0,
        color="#b80000", marker="*"
    )
    xlabel(L"$\beta$")
    ylabel(L"$\alpha$")
    safesave(joinpath(save_path, "loss_ext.png"), fig)
    close(fig)

    fig = figure("conventional", figsize=(4, 4))
    contourf(αs, βs, loss, 25, cmap=Seaborn.ColorMap("mako"), linewidths=1.5,
        vmin=-1f-1, vmax=1f0)
    title("Loss landscape")
    cb = colorbar(fraction=0.057, pad=0.01, format=sfmt)
    PyPlot.scatter(
        βs[argmin(loss)[2]], αs[argmin(loss)[1]], s=30.0,
        color="#b80000", marker="*"
    )
    xlabel(L"$\beta$")
    ylabel(L"$\alpha$")
    safesave(joinpath(save_path, "loss_contourf.png"), fig)
    close(fig)

    fig = figure("conventional", figsize=(4, 4))
    contourf(αs, βs, loss_ext, 25, cmap=Seaborn.ColorMap("mako"), linewidths=1.5,
        vmin=-1f-1, vmax=1f0)
    title("Extended loss landscape")
    colorbar(fraction=0.057, pad=0.01, format=sfmt)
    PyPlot.scatter(
        βs[argmin(loss_ext)[2]], αs[argmin(loss_ext)[1]], s=30.0,
        color="#b80000", marker="*"
    )
    xlabel(L"$\beta$")
    ylabel(L"$\alpha$")
    safesave(joinpath(save_path, "loss_ext_contourf.png"), fig)
    close(fig)

    fig = figure("trace", figsize=(3*1.5, 7.5*1.5))
    plt.plot(αs, loss[:, argmin(loss)[2]], color="b", ls="--")
    plt.plot(αs, loss_ext[:, argmin(loss_ext)[2]], color="b")

    plt.plot(βs, loss[argmin(loss)[1], :], color="k", ls="--")
    plt.plot(βs, loss_ext[argmin(loss_ext)[1], :], color="k")

    plt.plot(αs, diag(loss), color="r", ls="--")
    plt.plot(αs, diag(loss_ext), color="r")
    safesave(joinpath(save_path, "trace.png"), fig)
    close(fig)

end