## generation_tools.jl
export get_scan_pattern, ϕ_tuple, idx_to_dim

"""
    ϕ_tuple(t::NTuple)

Helper function to obtain the angle from a tuple of coordinates.
The azimuthal angle is calculated via the `atan`.
However, the order of the NTuple is (x,y).

Arguments:
+ `t`: an NTuple. Only the first two coordinates (x,y) are used and the azimuthal angle ϕ is returned
"""
ϕ_tuple(t::NTuple) = atan(t[2],t[1])

## conversion_tools.jl

"""
    idx_to_dim(idx_arr)  # this should be a view

Converts an N-dimensional array of NTuple to an N+1 dimensional array by orienting the (inner) tuple along the (outer) trailing+1 dimension.

Arguments:
+ `idx_arr`. The array of NTuple to convert

"""
function idx_to_dim(idx_arr::AbstractArray{T, N}) where {T, N}
    out_arr = Array{typeof(idx_arr[1][1])}(undef, size(idx_arr)..., length(idx_arr[begin]))

    # loop over array
    @inbounds for I in CartesianIndices(idx_arr)
        for (i, ind) in enumerate(axes(out_arr, N+1))
            out_arr[Tuple(I)..., ind] = idx_arr[I][i]
        end
    end
    return out_arr
end


# This test file test the auxillary functions in FractionalDiffEq.jl

using FractionalDiffEq
using Test

@testset "Test eliminator" begin
    @test isapprox(eliminator(3, 2), [1.0 0.0 0.0; 0.0 0.0 1.0]; atol=1e-5)
end

@testset "Test Riesz Matrix" begin
    @test isapprox(RieszMatrix(0.5, 3, 0.01), [-5.0 4.375 -0.3125; 4.375 -5.0 4.375; -0.3125 4.375 -5.0]; atol=0.00001)
end

@testset "Test ω function" begin
    @test isapprox(omega(3, 0.5), [1.0; -0.5; -0.125; -0.0625]; atol=1e-5)
    @test isapprox(omega(5, 0.5), [1.0; -0.5; -0.125; -0.0625; -0.0390625; -0.02734375]; atol=1e-5)
end

@testset "Test distributed order utils function" begin
    @test isapprox(DOB(x->6*x*(1-x), [0, 1], 0.1, 1, 0.1), [3.547014602981364]; atol=1e-5)
    @test isapprox(DOF(x->6*x*(1-x), [0, 1], 0.1, 1, 0.1), [-3.547014602981364]; atol=1e-5)
    @test isapprox(DORANORT(x->6*x*(1-x), [0, 1], 0.1, 1, 0.1), [2.120468497843435]; atol=1e-5)
end

#FIXME: Add array type tests
@testset "Test Mittag Leffler function" begin
    @test isapprox(mittleff(1, 2, 1), exp(1)-1; atol=1e-5)
    @test isapprox(mittleff(2, 2, 1), sinh(1); atol=1e-5)
    @test isapprox(mittleff(1, 1, 0), 1; atol=1e-5)
    @test isapprox(mittleff(1, 2, [1, 2, 3]), [1.718281828459045; 3.194528049465325; 6.361845641062556]; atol=1e-2)

    @test isapprox(mittleffderiv(1, 1, 1), 2.718281828459045; atol=1e-5)
    @test isapprox(mittleffderiv(1, 1), 2.718281828459045; atol=1e-5)

    @test isapprox(mittleffderiv(1, 1, 1/3), 1.39561242508609; atol=1e-5)
    @test isapprox(mittleffderiv(2, 2, 1), 0.18393972058572117; atol=1e-5)

    @test isapprox(mittleff(1, 2, 1, 1), 1.7182818284590438; atol=1e-5)
    @test isapprox(mittleff(2, 2, 1, 1), 1.1752011936438016; atol=1e-3)
    @test isapprox(mittleff(2.1, 2, 1, 1), 1.1527981788744044; atol=1e-4)

    # Also tests for Complex z:

    m = mittleff(1, 2, im)
    @test isapprox(real(m), 0.8414709848078965; atol=1e-5)
    @test isapprox(imag(m), 0.4596976941318611; atol=1e-5)
end

@testset "Test meshgrid" begin
    a, b = meshgrid(0:1, 0:1)
    isapprox(a, [0 1; 0 1]; atol=1e-2)
    isapprox(b, [0 0; 1 1]; atol=1e-2)
end

@testset "Test isFunction" begin
    @test isFunction(x->x)==true
    @test isFunction("Hello")==false
end

@testset "Test singleterm show method" begin
    # SingleTermFODEProblem
    singlefun(x, y) = 1-y
    singletermprob = SingleTermFODEProblem(singlefun, 1.8, 0, 5)
    @test_nowarn show(singletermprob)
end

@testset "Test Multiterms show method" begin
    # MultiTermsFODEProblem
    rightfun(x, y) = 172/125*cos(4/5*x)
    multitermsprob = MultiTermsFODEProblem([1, 1/16, 4/5, 3/2, 1/25, 6/5], [3, 2.5, 2, 1, 0.5, 0], rightfun, [0 0 0 0 0 0], (0, 10))
    multitermssol = solve(multitermsprob, 0.5, PIEX())

    @test_nowarn show(multitermsprob)
    @test_nowarn show(multitermssol)
end

@testset "Test FODESystem show method" begin
    # FODESystem
    h=0.5; tspan=(0, 1)
    alpha = [0.99, 0.99, 0.99]
    x0 = [1, 0, 1]
    function testf!(du, u, p, t)
        a, b, c = 10, 28, 8/3
        du[1] = a*(u[2]-u[1])
        du[2] = u[1]*(b-u[3])-u[2]
        du[3] = u[1]*u[2]-c*u[3]
    end
    fodesystemprob = FODESystem(testf!, alpha, x0, tspan)
    fodesystemsol = solve(fodesystemprob, h, GL())

    @test_nowarn show(fodesystemprob)
    @test_nowarn show(fodesystemsol)
end

@testset "Test FDDEProblem show method" begin
    # FDDEProblem
    function ϕ(x)
        if x == 0
            return 19.00001
        else
            return 19.0
        end
    end
    function delayf(t, y, ϕ)
        return 3.5*y*(1-ϕ/19)
    end
    fddeprob = FDDEProblem(delayf, ϕ, 0.97, 0.8, (0, 2))
    fddesol = solve(fddeprob, 0.5, DelayPI())

    @test_nowarn show(fddeprob)
    @test_nowarn show(fddesol)
end

@testset "Test FDDESystem show method" begin
    function EnzymeKinetics!(dy, y, ϕ, t)
        dy[1] = 10.5-y[1]/(1+0.0005*ϕ[4]^3)
        dy[2] = y[1]/(1+0.0005*ϕ[4]^3)-y[2]
        dy[3] = y[2]-y[3]
        dy[4] = y[3]-0.5*y[4]
    end
    q = [60, 10, 10, 20]; α = [0.95, 0.95, 0.95, 0.95]
    fddesystemprob = FDDESystem(EnzymeKinetics!, q, α, 4, 6)

    @test_nowarn show(fddesystemprob)
end

@testset "Test DODEProblem show method" begin
    h = 0.5; t = collect(0:h:1)
    dodefun(t)=0
    dodeprob = DODEProblem([1, 0.1], [x->6*x*(1-x), 0], (0, 1), dodefun, 1, t)
    dodesol = solve(dodeprob, h, DOMatrixDiscrete())

    @test_nowarn show(dodeprob)
    @test_nowarn show(dodesol)
end

@testset "Test FractionalDifferenceProblem" begin
    differencefun(x) = 0.5*x+1
    α=0.5;x0=1;
    T=1; h=0.1
    differenceprob = FractionalDifferenceProblem(differencefun, α, x0)
    differencesol = solve(differenceprob, T, h, PECEDifference())

    @test_nowarn show(differenceprob)
    @test_nowarn show(differencesol)
end

@testset "Test FIEProblem show method" begin
    fietspan = LinRange(-1, 1, 5)
    fieprob = FIEProblem([1, 1], [1, 0.5], 1, fietspan)
    fiesol = solve(fieprob, 20, SpectralUltraspherical())

    @test_nowarn show(fieprob)
    @test_nowarn show(fiesol)
end


#=
A file to hold the helper functions used for working with cubic polynomials
=#

"""A helper function; Returns true if any of the arguments is NaN"""
function any_nan{T}(x::T...)
	for n in x
		if isnan(n)
			return true
		end
	end

	#if the loop finished, then they are all not NAN so return false
	return false
end

function maxcubic{T}(x0::T, x1::T, a::T, b::T, c::T, d::T) #a+bx+cx^2+dx^3
	if any_nan(x0, x1, a, b, c, d)
		throw(ArgumentError("Maximum of a cubic with NaN parameters is not defined"))
	end

  if x0 == x1
    return (x0, @evalpoly(x0, a, b, c, d))
  end
  if d == 0 #then the max is at an endpoint. #FIXME: I don't think so
    endpt1_val = @evalpoly(x0, a, b, c, d)
    endpt2_val = @evalpoly(x1, a, b, c, d)
    return endpt2_val > endpt1_val ? (endpt2_val, x1) : (endpt1_val, x0)
  end

  #Otherwise, we use the quadratic formula on the derivative
	discriminant = (c)^2-(3d*b)
	if discriminant >= 0 #then the max may be  in [x0, x1]
		root1 = (-sqrt(discriminant)-c)/(3*d)
		root2 = (sqrt(discriminant)-c)/(3*d)
	end

	val1 = @evalpoly(x0, a, b, c, d)::T
	if discriminant >= 0
		val2 = @evalpoly(root1, a, b, c, d)::T
		val3 = @evalpoly(root2, a, b, c, d)::T
	else
		val2 = typemin(T)
		val3 = typemin(T)
	end
	val4 = @evalpoly(x1, a, b, c, d)::T

	if discriminant >= 0
		if root1 < x0 || root1 >x1
			val2 = typemin(T)
		end
		if root2 < x0 || root2 > x1
			val3 = typemin(T)
		end
	end

	theMax = max(val1, val2, val3, val4)

	maxPos = typemin(T)
	if val1 == theMax
		maxPos = x0
	elseif val2 == theMax
		maxPos = root1
	elseif val3 == theMax
		maxPos = root2
	elseif val4 == theMax
		maxPos = x1
	else
		println("Failed to find the max in maxcubic. This is probably a bug in Splines.jl")
		return nothing
	end

	return (theMax, maxPos)
end

function maxcubic{T}(x::T, a::T, b::T, c::T, d::T) #search on [0, x]
	return maxcubic(zero(T), x, a, b, c, d)
end

#FIXME: NaN if d = 0
function maxabscubic{T}(x0::T, x1::T, a::T, b::T, c::T, d::T) #a+bx+cx^2+dx^3
  @assert !isnan(x0)
  @assert !isnan(x1)
  @assert !isnan(a)
  @assert !isnan(b)
  @assert !isnan(c)
  @assert !isnan(d)
  #print("x0, x1, a, b, c, d: ", x0, a, b, c, d)
  #print("\n")
  if x0 == x1 #then there is no real max, but we do the sanest possible thing
    return (x0, @evalpoly(x0, a, b, c, d))
  end
  if d == 0 #then the max is at an endpoint. #FIXME: I don't think so
    endpt1_val = @evalpoly(x0, a, b, c, d)
    endpt2_val = @evalpoly(x1, a, b, c, d)
    return endpt2_val > endpt1_val ? (endpt2_val, x1) : (endpt1_val, x0)
  end

	#otherwise, we use the quadratic formula on the derivative
	discriminant = (c)^2-(3d*b)
	if discriminant >= 0 #then the max may be  in [x0, x1]
		root1 = (-sqrt(discriminant)-c)/(3*d)
		root2 = (sqrt(discriminant)-c)/(3*d)
	end

	val1 = abs(@evalpoly(x0, a, b, c, d)::T)
	if discriminant >= 0
		val2 = abs(@evalpoly(root1, a, b, c, d)::T)
		val3 = abs(@evalpoly(root2, a, b, c, d)::T)
	else
		val2 = typemin(T)
		val3 = typemin(T)
	end
	val4 = abs(@evalpoly(x1, a, b, c, d)::T)

	if discriminant >= 0
		if root1 < x0 || root1 >x1
			val2 = typemin(T)
		end
		if root2 < x0 || root2 > x1
			val3 = typemin(T)
		end
	end

	theMax = max(val1, val2, val3, val4)

  #	println(val1, val2, val3, val4, theMax)
	maxPos = typemin(T)
	if val1 == theMax
		maxPos = x0
	elseif val2 == theMax
		maxPos = root1
	elseif val3 == theMax
		maxPos = root2
	elseif val4 == theMax
		maxPos = x1
  #=	else
    println("Failed to find the max in maxabscubic. This is probably a bug in Splines.jl")
		return nothing=#
	end

	return (theMax, convert(T, maxPos))
end

function maxabscubic{T}(x::T, a::T, b::T, c::T, d::T) #search on [0, x]
	return maxcubic(zero(T), x, a, b, c, d)
end

function mincubic{T}(x0::T, x1::T, a::T, b::T, c::T, d::T) #a+bx+cx^2+dx^3
	if any_nan(x0, x1, a, b, c, d)
		throw(ArgumentError("Minimum of a cubic with NaN parameters is not defined"))
	end
  if x0 == x1
    return (x0, @evalpoly(x0, a, b, c, d))
  end
  if d == 0 #then the min is at an endpoint. #FIXME: I don't think so
    endpt1_val = @evalpoly(x0, a, b, c, d)
    endpt2_val = @evalpoly(x1, a, b, c, d)
    return endpt2_val < endpt1_val ? (endpt2_val, x1) : (endpt1_val, x0)
  end

	#otherwise, we use the quadratic formula on the derivative
	discriminant = (c)^2-(3*d*b)
	if discriminant >= 0
		root1 = (-sqrt(discriminant)-c)/(3*d)
		root2 = (sqrt(discriminant)-c)/(3*d)
	end

	val1 = @evalpoly(x0, a, b, c, d)
	if discriminant >= 0
		val2 = @evalpoly(root1, a, b, c, d)
		val3 = @evalpoly(root2, a, b, c, d)
	else
		val2 = typemax(T)
		val3 = typemax(T)
	end
	val4 = @evalpoly(x1, a, b, c, d)

	if discriminant >= 0
		if root1 < x0 || root1 > x1
			val2 = typemax(T)
		end

		if root2 < x0 || root2 > x1
			val3 = typemax(T)
		end
	end

	theMin = min(val1, val2, val3, val4)

	minPos = typemax(T)
	if val1 == theMin
		minPos = x0
	elseif val2 == theMin
	if !isdefined(:root1)
		println("val1: $val1 val2: $val2 val3: $val3 val4: $val4 theMin: $theMin")
	end
		minPos = root1
	elseif val3 == theMin
		minPos = root2
	elseif val4 == theMin
		minPos = x1
	else
		error("Failed to find the minimum. This is probably a bug in Splines.jl")
	end

	return (theMin, minPos)
end

function mincubic{T}(x::T, a::T, b::T, c::T, d::T) #search between 0 and x
	return mincubic(zero(T), x, a, b, c, d)
end

function integrate_cubic{T}(x0::T, x1::T, a::T, b::T, c::T, d::T) #a+bx+cx^2+dx^3
  return @evalpoly(x1, 0, a, b/2, c/3, d/4)-@evalpoly(x0, 0, a, b/2, c/3, d/4)
end


using ITensors
using OptimKit
using Zygote
using Random: seed!

include("circuit.jl")

function gate(::Val{:Ry}; θ)
  return [
    cos(θ / 2) -sin(θ / 2)
    sin(θ / 2) cos(θ / 2)
  ]
end

function gate(::Val{:CX})
  return [
    1 0 0 0
    0 1 0 0
    0 0 0 1
    0 0 1 0
  ]
end

function Rylayer(N, θ⃗)
  return [("Ry", (n,), (θ=θ⃗[n],)) for n in 1:N]
end

function CXlayer(N)
  return [("CX", (n, n + 1)) for n in 1:2:(N - 1)]
end

layer(N, θ⃗) = vcat(Rylayer(N, θ⃗), CXlayer(N))

# The variational circuit we want to optimize
function variational_circuit(N, nlayers, θ⃗)
  range = 1:N
  circuit = layer(N, θ⃗[range])
  for n in 1:(nlayers - 1)
    circuit = vcat(circuit, layer(N, θ⃗[range .+ n * N]))
  end
  return circuit
end

N = 4
s = siteinds("Qubit", N)

h = 1.3
ℋ = OpSum()
for j in 1:(N - 1)
  ℋ .+= -1, "Z", j, "Z", j + 1
end
for j in 1:N
  ℋ .+= h, "X", j
end
H = prod(MPO(ℋ, s))
ψ = prod(MPS(s, "0"))

N = 4
nlayers = 5

function loss(θ⃗)
  gates = variational_circuit(N, nlayers, θ⃗)
  U = buildcircuit(gates, s)
  return rayleigh_quotient(H, (U, ψ))
end

seed!(1234)
θ⃗₀ = 2π * rand(N * nlayers)
@show θ⃗₀
@show loss(θ⃗₀)
@show loss'(θ⃗₀)

fg(x) = (loss(x), convert(Vector, loss'(x)))
θ⃗ₒₚₜ, fₒₚₜ, gₒₚₜ, numfg, normgradhistory = optimize(fg, θ⃗₀, GradientDescent())
@show θ⃗ₒₚₜ
@show fₒₚₜ
@show gₒₚₜ
@show numfg
println("normgradhistory = ")
display(normgradhistory)

Uₒₚₜ = buildcircuit(variational_circuit(N, nlayers, θ⃗ₒₚₜ), s)
Uₒₚₜψ = apply(Uₒₚₜ, ψ)
@show (Uₒₚₜψ' * H * Uₒₚₜψ)[]

nothing


using SymbolicUtils, Test
using SymbolicUtils: Term, symtype
@testset "polyform" begin
    @syms a b c d
    @test expand(a * (b + -1 * c) + -1 * (b * a + -1 * c * a)) == 0
    @eqtest expand(sin(a+b)+sin(c+d)) == sin(a+b) + sin(c+d)
    @eqtest simplify(expand(sin((a+b)^2)^2)) == simplify(sin(a^2+2*(b*a)+b^2)^2)
    @test simplify(expand(sin((a+b)^2)^2 + cos((a+b)^2)^2)) == 1
    @syms x1::Real f(::Real)::Real

    # issue 193
    @test isequal(expand(f(x1 + 2.0)), f(2.0 + x1))
    @test symtype(expand(f(x1 + 2.0))) == Real

    # cleanup rules
    @test expand(Term{Number}(identity, 0)) == 0
    @test expand(Term{Number}(one, 0)) == 1
    @test expand(Term{Number}(zero, 0)) == 0
    @test expand(identity(a * b) - b * a) == 0
    @test expand(a * b - b * a) == 0
end


@testset "FixedHomotopy" begin
    PolyImpl.@polyvar x y a
    H = StraightLineHomotopy([y^2-x], [x^2+3.0y])

    H1 = FixedHomotopy(H, 1.0)
    H0 = FixedHomotopy(H, 0.0)
    HH = StraightLineHomotopy(H1, H0)

    @test typeof(H1)<:FixedHomotopy{Float64, typeof(H), Float64}

    @test evaluate(H1, [2.0, 1.0]) == evaluate(H, [2.0, 1.0], 1.0)
    @test H1([2.0, 1.0]) == H([2.0, 1.0], 1.0)
    @test H0([2.0, 1.0]) == H([2.0, 1.0], 0.0)
    @test HH([2.0, 1.0], 1.0) == H([2.0, 1.0], 1.0)
    @test HH([2.0, 1.0], 0.0) == H([2.0, 1.0], 0.0)

    J_H = differentiate(H)
    J_H1 = differentiate(H1)
    J_H0 = differentiate(H0)

    @test J_H0([1.0, 1.0]) == J_H([1.0, 1.0], 0.0)
    @test J_H1([1.0, 1.0]) == J_H([1.0, 1.0], 1.0)

    @test degrees(H1) == [2]

    H = FixedHomotopy(StraightLineHomotopy([y^2-x, 3x + y], [x^2+3y, 2y-3x]), 0.34)

    @test homogenized(H) == false
    K = homogenize(H)
    @test ishomogenous(H) == false
    @test homogenized(K)
    @test ishomogenous(K)
    @test nvariables(K) == 3
end


abstract type Input end

"""
    get_value(input::Input, time::Number)

Get input value at givent time.
"""
function get_value(input::Input, time::Number) end

function get_value(input_arr::Array{T,1}, t::Number) where T<:Input
    value = 0.0
    for input in input_arr
        value = value + get_value(input, t)
    end
    return value
end

"""
    ConstantInput(val)

Constant value input.
"""
struct ConstantInput <: Input
    val::Float64
end

get_value(input::ConstantInput, time::Number) = input.val

"""
    StepInput(val, tini, tfin)

Step input. Takes value val from tini to tfin and 0 otherwise.
"""
struct StepInput <: Input
    val::Float64
    tini::Float64
    tfin::Float64
end

get_value(input::StepInput, time::Number) = input.tini <= time < input.tfin ? input.val : 0.

"""
    DoubletInput(val, tini, tfin)

Symmetric doublet input with peak to peak amplitude val:
 * val/2 if tini <= t < (tini + tfin) / 2
 * -val/2 if (tini + tfin) / 2 <= t < tfin
"""
DoubletInput(val, tini, tfin) = vcat(
    StepInput(0.5*val, tini, 0.5*(tini+tfin)),
    StepInput(-0.5*val, 0.5*(tini+tfin), tfin)
)

"""
    RampInput(val, tini, tfin)

Ramp input. Takes value at tend increasing linearly from tini to tend.
"""
struct RampInput <: Input
    val::Float64
    tini::Float64
    tfin::Float64
end

function get_value(input::RampInput, time::Number)
    if input.tini <= time < input.tfin
        return input.val * (time - input.tini) / (input.tfin - input.tini)
    else
        return 0.
    end
end

"""
    SinusoidalInput(A, f, ϕ, tini, tfin)

Sinusoidal input. Takes a A/2 * sin(2πf (t-tini) + ϕ) value between tini and tfin.
Value is 0 otherwise.
"""
struct SinusoidalInput <: Input
    "Peak to peak amplitude"
    A::Float64
    "Frequency (1/s)"
    f::Float64
    "Phase (rad)"
    ϕ::Float64
    tini::Float64
    tfin::Float64
end

function get_value(input::SinusoidalInput, time::Number)
    if input.tini <= time < input.tfin
        return 0.5 * input.A * sin(2π * input.f * (time - input.tini) + input.ϕ)
    else
        return 0.
    end
end


using OrdinaryDiffEq, ParameterizedFunctions, Test, Random

f = @ode_def LotkaVolterra begin
  dx = 1.5x - x*y
  dy = -3y + x*y
end

prob = ODEProblem(f,big.([1.0;1.0]),(big(0.0),big(10.0)))

using DiffEqDevTools

dts = big(1/2).^(8:-1:4)
test_setup = Dict(:alg=>Vern9(),:reltol=>1e-25,:abstol=>1e-25)
sim1 = analyticless_test_convergence(dts,prob,Tsit5(),test_setup)
sim2 = analyticless_test_convergence(dts,prob,Vern9(),test_setup)

@test sim1.𝒪est[:final]-5 < 0.2
@test sim2.𝒪est[:final]-9 < 0.2

function f2(du,u,p,t)
  du .= 1.01u
end
function g2(du,u,p,t)
  du .= 1.01u
end
prob = SDEProblem(f2,g2,[1.0;1.0],(0.0,10.0))

using StochasticDiffEq

dts = (1/2).^(6:-1:3)
test_dt = 1/2^8
Random.seed!(100)
sim1 = analyticless_test_convergence(dts,prob,SRIW1(),test_dt,trajectories=400)
@test sim1.𝒪est[:final]-1.5 < 0.3


## R code 12.2
library(rethinking)
data(UCBadmit)
d <- UCBadmit
d$gid <- ifelse( d$applicant.gender=="male" , 1L , 2L )
dat <- list( A=d$admit , N=d$applications , gid=d$gid )
m12.1 <- ulam(
    alist(
        A ~ dbetabinom( N , pbar , theta ),
        logit(pbar) <- a[gid],
        a[gid] ~ dnorm( 0 , 1.5 ),
        transpars> theta <<- phi + 2.0,
        phi ~ dexp(1)
    ), data=dat , chains=4 )

## R code 12.6
library(rethinking)
data(Kline)
d <- Kline
d$P <- standardize( log(d$population) )
d$contact_id <- ifelse( d$contact=="high" , 2L , 1L )

dat2 <- list(
    T = d$total_tools,
    P = d$population,
    cid = d$contact_id )

m12.2 <- ulam(
    alist(
        T ~ dgampois( lambda , phi ),
        lambda <- exp(a[cid])*P^b[cid] / g,
        a[cid] ~ dnorm(1,1),
        b[cid] ~ dexp(1),
        g ~ dexp(1),
        phi ~ dexp(1)
    ), data=dat2 , chains=4 , log_lik=TRUE )

## R code 12.9
m12.3 <- ulam(
    alist(
        y ~ dzipois( p , lambda ),
        logit(p) <- ap,
        log(lambda) <- al,
        ap ~ dnorm( -1.5 , 1 ),
        al ~ dnorm( 1 , 0.5 )
    ) , data=list(y=y) , chains=4 )

## R code 12.11
m12.3_alt <- ulam(
    alist(
        y|y>0 ~ custom( log1m(p) + poisson_lpmf(y|lambda) ),
        y|y==0 ~ custom( log_mix( p , 0 , poisson_lpmf(0|lambda) ) ),
        logit(p) <- ap,
        log(lambda) <- al,
        ap ~ dnorm(-1.5,1),
        al ~ dnorm(1,0.5)
    ) , data=list(y=as.integer(y)) , chains=4 )

## R code 12.15
logit <- function(x) log(x/(1-x)) # convenience function
round( lco <- logit( cum_pr_k ) , 2 )

## R code 12.16
m12.4 <- ulam(
    alist(
        R ~ dordlogit( 0 , cutpoints ),
        cutpoints ~ dnorm( 0 , 1.5 )
    ) , data=list( R=d$response ), chains=4 , cores=4 )

## R code 12.17
m12.4q <- quap(
    alist(
        response ~ dordlogit( 0 , c(a1,a2,a3,a4,a5,a6) ),
        c(a1,a2,a3,a4,a5,a6) ~ dnorm( 0 , 1.5 )
    ) , data=d , start=list(a1=-2,a2=-1,a3=0,a4=1,a5=2,a6=2.5) )

## R code 12.24
dat <- list(
    R = d$response,
    A = d$action,
    I = d$intention,
    C = d$contact )
m12.5 <- ulam(
    alist(
        R ~ dordlogit( phi , cutpoints ),
        phi <- bA*A + bC*C + BI*I ,
        BI <- bI + bIA*A + bIC*C ,
        c(bA,bI,bC,bIA,bIC) ~ dnorm( 0 , 0.5 ),
        cutpoints ~ dnorm( 0 , 1.5 )
    ) , data=dat , chains=4 , cores=4 )

## R code 12.34
dat <- list(
    R = d$response ,
    action = d$action,
    intention = d$intention,
    contact = d$contact,
    E = as.integer( d$edu_new ),   # edu_new as an index
    alpha = rep( 2 , 7 ) )      # delta prior

m12.6 <- ulam(
    alist(
        R ~ ordered_logistic( phi , kappa ),
        phi <- bE*sum( delta_j[1:E] ) + bA*action + bI*intention + bC*contact,
        kappa ~ normal( 0 , 1.5 ),
        c(bA,bI,bC,bE) ~ normal( 0 , 1 ),
        vector[8]: delta_j <<- append_row( 0 , delta ),
        simplex[7]: delta ~ dirichlet( alpha )
    ), data=dat , chains=4 , cores=4 , coerce_int=FALSE )

## R code 12.37
dat$edu_norm <- normalize( d$edu_new )
m12.7 <- ulam(
    alist(
        R ~ ordered_logistic( mu , cutpoints ),
        mu <- bE*edu_norm + bA*action + bI*intention + bC*contact,
        c(bA,bI,bC,bE) ~ normal( 0 , 1 ),
        cutpoints ~ normal( 0 , 1.5 )
    ), data=dat , chains=4 , cores=4 )


"""
   probability_check(pomdp::POMDP)

Checks if the transition and observation function of the discrete `pomdp`
have probability mass that sums up to unity for all state-action pairs.
"""
function probability_check(pomdp::POMDP)
    obs_prob_consistency_check(pomdp)
    trans_prob_consistency_check(pomdp)
end

"""
    obs_prob_consistency_check(pomdp::POMDP)

Checks if the observation function of the discrete `pomdp`
has probability mass that sums up to unity for all state-action pairs.
"""
function obs_prob_consistency_check(pomdp::POMDP)
    # initalize space
    sspace = states(pomdp)
    aspace = actions(pomdp)
    ospace = observations(pomdp)
    # iterate through all s-a pairs
    for s in iterator(sspace)
        for a in iterator(aspace)
            obs = observation(pomdp, a, s)
            p = 0.0
            for o in iterator(ospace)
                p += pdf(obs, o)
            end
            @assert isapprox(p, 1.0) "Observation probability does not sum to unity for state: ", s, " action: ", a
        end
    end
end

"""
    trans_prob_consistency_check(pomdp::Union{MDP, POMDP})

Checks if the transition function of the discrete problem
has probability mass that sums up to unity for all state-action pairs.
"""
function trans_prob_consistency_check(pomdp::Union{MDP, POMDP})
    # initalize space
    sspace = states(pomdp)
    aspace = actions(pomdp)
    # iterate through all s-a pairs
    for s in iterator(sspace)
        for a in iterator(aspace)
            tran = transition(pomdp, s, a)
            p = 0.0
            for sp in iterator(sspace)
                p += pdf(tran, sp)
            end
            @assert isapprox(p, 1.0) "Transition probability does not sum to unity for state: ", s, " action: ", a
        end
    end
end


# Bell inequality

using Probably
using Statistics: mean
import Probably: apply!

° = π/180

rot(θ) = [cos(θ) -sin(θ)
          sin(θ) cos(θ)]

# Simulate measurement by a polaroid at angle θ to the x / |1⟩ axis.
# (Effectively we rotate the photon rather than the polaroid.)
function polaroid(x, θ)
  apply!(rot(θ), x)
  result = measure(x)
  apply!(rot(-θ), x)
  return result
end

# Create an EPR pair, measure both at a random angle, and see if they agree.
function experiment1()
  a, b = epr()
  θ = rand([-60°, 0°, 60°])
  !(polaroid(a, θ) ⊻ polaroid(b, θ))
end

# When measured at the same angle the qubits always agree.
mean(experiment1() for _ = 1:10^4)

# Create an EPR pair, measure both at different random angles, and see if they
# agree.
function experiment2()
  a, b = epr()
  θa, θb = rand([-60°, 0°, 60°], 2)
  !(polaroid(a, θa) ⊻ polaroid(b, θb))
end

# At different angles the qubits agree 50% of the time.
mean(experiment2() for _ = 1:10^4)

# A local hidden variable theory *cannot* give this prediction. Say qubit `a` is
# in a state that determines measurement outcomes at each angle [-60, 0, 60]
# (say [1, 0, 1]). We know that when the polaroids are at the same angle they
# always agree, so `b` must be in an equivalent state.

# Qubit state – [60° => 1, 0° => 0, 60° => 1]

# polaroid 1, polaroid 2, outcome (1 ⊻ 2)
#  0           0          1
#  60          60         1
# -60         -60         1
#  0           60         0 (× 2)
#  0          -60         0 (× 2)
#  60         -60         1 (× 2)

# In this hidden state the qubits agree 5/9s of the time. Other states like [1,
# 1, 1] and [1, 0, 0] will agree as much or more than this, so the probability
# of agreement is at least 5/9 overall, contradicting the result of our
# experiment.


#Test --- if ϵ (infectiousness of asymptomatics is zero) and external rates are zero then nothing should happen despite 1000 initial asymptomatics

function test_no_infecteds()
    #Load data completely susceptible Population
    u0,P,transport_matrix = KenyaCoV.model_ingredients_from_data("data/combined_population_estimates.csv",
                                                                 "data/optimal_transition_matrix.jld2",
                                                                "data/optimal_movement_matrix.jld2",
                                                                "data/flight_numbers.csv",
                                                                "data/projected_global_prevelance.csv")
    #You can modify other parameters directly
    P.τ = 0. #e.g. no treatment rate
    P.ϵ = 0.
    P.into_mom .= 0
    P.into_nai .= 0
    #Define initial conditions by modifying the completely susceptible population
    u0[30,3,1] += 1000#One asymptomatic in Nairobi

    #Go straight to solution using solver compiled in the KenyaCoV module
    sol_tl = KenyaCoV.solve_KenyaCoV_prob(u0,(0.,365.),P,1.)

    return ~any([sum(reshape(sol_tl[end],KenyaCoV.n,KenyaCoV.n_s,2)[i,7:8,1:2]) for i = 1:47] .!= 0)
end

function stays_nonnegative()
    u0,P,transport_matrix = KenyaCoV.model_ingredients_from_data("data/combined_population_estimates.csv",
                                                                 "data/optimal_transition_matrix.jld2",
                                                                 "data/optimal_movement_matrix.jld2",
                                                                 "data/flight_numbers.csv",
                                                                 "data/projected_global_prevelance.csv")
    P.τ = 0. #e.g. no treatment rate
    for (i,p) in enumerate(P.global_prev)
        P.global_prev[i] = 0.
    end
    # P.global_prev
    P.μ₁ = 10.
    P.ϵ = 1.
    P.β = 2.2*P.γ
    #Define initial conditions by modifying the completely susceptible population
    P.dt = 1.
    u0[30,4,1] += 10#One diseased in Nairobi
    prob = KenyaCoV.create_KenyaCoV_non_neg_prob(u0,(0.,365.),P)
    sol = solve(prob,FunctionMap(),dt = P.dt)
    ~any([any(u .< 0) for u in sol.u])
end
# function run_epidemic_infs_everywhere()
#     u0,P,transport_matrix = KenyaCoV.model_ingredients_from_data("data/combined_population_estimates.csv",
#                                                                  "data/optimal_transition_matrix.jld2",
#                                                                 "data/optimal_movement_matrix.jld2",
#                                                                 "data/flight_numbers.csv",
#                                                                 "data/projected_global_prevelance.csv")
#
#     u0[30,4,1] += 100#One diseased in Nairobi
#
#     #Create a JumpProblem which you can solve --- needs DifferentialEquations module for the solvers
#     jump_prob_tl = create_KenyaCoV_prob(u0,(0.,365.),P)
#     #Go straight to solution using solver compiled in the KenyaCoV module
#     sol_tl = solve_KenyaCoV_prob(u0,(0.,365.),P,1.)
# end


"""
    RotationFrameDiagonal(omega::Array{Float64},omega_drive::Float64)

# Argument:
- omega: transition frequencies
- omega_drive: driving frequency
- N: number of states

# Output:
- Diagonal matrix diag{0,omega[1]-omr,omega[2]-omr,...,omega[N-1]-omr}
"""
function RotationFrameDiagonal(omega::Array{Float64},omega_drive::Float64)
	diag_vec = 0.0
	for i = 1:length(omega)
		# potential troublemaker of auto-diff
		diag_vec = [diag_vec;diag_vec[end]+omega[i]-omega_drive]
	end
	HK = Diagonal(diag_vec)
    return HK
end

function RotationFrameDiagonal(omega,omega_drive)
	diag_vec = 0.0
	for i = 1:length(omega)
		# potential troublemaker of auto-diff
		diag_vec = [diag_vec;diag_vec[end]+omega[i]-omega_drive]
	end
	HK = Diagonal(diag_vec)
    return HK
end

"""
    RotationFrameDiagonal(gamma1::Array{Float64},gamma2::Array{Float64})

# Argument:
- gamma1: determine the decay operator in the Lindblad system
- gamma2: determine the dephasing operator in the Lindblad system

# Output:
- Ldecay: Ldecay[i][i+1] = sqrt(gamma1[i]), and zeros everywhere else
- Ldephase: Diagonal{0,sqrt(gamma2[1]),...)
"""
function RotationFrameLindblad(gamma1::Array{Float64},gamma2::Array{Float64})
	#L_decay = Bidiagonal(zeros(length(gamma1)+1),sqrt.(gamma1),:U)
	L_dephase = Diagonal([0;sqrt.(gamma2)])
	L_decay = [zeros(Float64,length(gamma1)) Diagonal(sqrt.(gamma1))]
	L_decay = [L_decay;zeros(Float64,1,length(gamma1)+1)]
	#L_decay = [0.0 sqrt(gamma1[1]) 0.0 0.0;
	#		   0.0 0.0      sqrt(gamma1[2]) 0.0;
	#		   0.0 0.0      0.0      sqrt(gamma1[3]);
	#		   0.0 0.0      0.0      0.0]
    return L_decay, L_dephase
end

function RotationFrameLindblad(gamma1,gamma2)
	#L_decay = Bidiagonal(zeros(length(gamma1)+1),sqrt.(gamma1),:U)
	L_dephase = Diagonal([0;sqrt.(gamma2)])
	L_decay = [zeros(Float64,length(gamma1)) Diagonal(sqrt.(gamma1))]
	L_decay = [L_decay;zeros(Float64,1,length(gamma1)+1)]
	#L_decay = [0.0 sqrt(gamma1[1]) 0.0 0.0;
	#		   0.0 0.0      sqrt(gamma1[2]) 0.0;
	#		   0.0 0.0      0.0      sqrt(gamma1[3]);
	#		   0.0 0.0      0.0      0.0]
    return L_decay, L_dephase
end

"""
    RotationFrameRamseyControl(abs_control::Float64,phase_control::Float64,N::Int64)

# Argument:
- abs_control: strength of the control ``\\Omega``
- phase_control: phase of the control signal ``\\theta``
- N: number of states

# Output:
- HCK,HCS: the real and imaginary part of the control operator, determined by
  ``p(a+a^\\dagger)+iq(a-a^\\dagger),``
  where ``a`` is the lowering operator, ``p=|\\Omega|\\cos(\\theta)`` and ``q=|\\Omega|\\sin(\\theta)``
"""
function RotationFrameRamseyControl(abs_control::Float64,phase_control::Float64,N::Int64)
    p = abs_control*cos(phase_control);
    q = abs_control*sin(phase_control);
	amat = [Matrix(Diagonal(sqrt.(collect(1:N-1))));zeros(Float64,1,N-1)]
	#amat = [Diagonal(sqrt.(collect(1:N-1)));zeros(Float64,1,N-1)]
	amat = [zeros(Float64,N) amat]
    #amat = Bidiagonal(zeros(N),sqrt.(collect(1:N-1)),:U) # trouble maker, maybe we can make a Bidiagonal adjoint
    HCK = p*(amat+transpose(amat))
    HCS = q*(amat-transpose(amat))
    return HCK,HCS
end
"""
    RotationFrameRamseyControl(N::Int64)

# Argument:
- abs_control: strength of the control ``\\Omega``
- phase_control: phase of the control signal ``\\theta``
- N: number of states

# Output:
- HCK,HCS: the real and imaginary part of the control operator, determined by
  ``(a+a^\\dagger)``
"""
function RotationFrameRamseyControl(N::Int64)
	amat = [Matrix(Diagonal(sqrt.(collect(1:N-1))));zeros(Float64,1,N-1)]
	amat = [zeros(Float64,N) amat]
    HCK = amat+transpose(amat)
	HCS = zeros(Float64,N,N)
    return HCK,HCS
end

"""
	RamseyForwardSolve(rho_u0::Array{Float64},rho_v0::Array{Float64},
						  omega::Array{Float64},omega_drive::Float64,
						  gamma1::Array{Float64},gamma2::Array{Float64},
						  InitialState::Int64,
						  TC::Float64,tdark_times::Array{Float64},
					   	  N_states::Int64=0;
					      initial_type="states",
					   	  method="exponential",
					      DiffEqKwargs...)

# Argument:
- rho_u0,rho_v0: initial states ``\\rho_{u_0}-i\\rho_{v_0}``
- omega: transition frequencies
- omega_drive: driving frequency
- gamma1: the decay parameter  of the Lindblad operators, equal to ``1/T_1`` with ``T_1`` being the relaxation time
- gamma2: the dephasing parameter  of the Lindblad operators, equal to ``1/T_2`` with ``T_2`` being the dephasing time
- InitialState: initial state of the density matrix
- TC: control time of the control signal
- tdark_times: the dark time in the Ramsey experiment
- N_states: number of states. If it is not specified, the default value will be 0, and the code will automatically detect how many states are considered.
- initial_type: type of the initial condition
  1. initial values are state vectors (default choice)
  2. initial values are density matrices saving in flattened vectors
- method: method to solve the Lindblad system
 1. "exponential": exponential time integrator (default choice)
 2. "DiffEqDefault": a default choice made by DifferentialEquations.jl
 3. Other solvers available in DifferentialEquations.jl, for example, method = Trapezoid()
- DiffEqKwargs: keyword arguments feed to the ``solve function" of DifferentialEquations.jl

# Output:
- ``\\rho_u,\\rho_v``: real and imaginary part of the density matrix ``\\rho=\\rho_u-i\\rho_v`` at sampled dark times, and they are saved as flattened vectors
"""
function RamseyForwardSolve(rho_u0::Array{Float64},rho_v0::Array{Float64},
				   		  omega::Array{Float64},omega_drive::Float64,
				   		  gamma1::Array{Float64},gamma2::Array{Float64},
						  InitialState::Int64,
				   		  TC::Float64,t_dark_times::Array{Float64},
				   		  N_states::Int64=0;
						  initial_type="states",
						  method="exponential",
						  DiffEqKwargs...)
    if(N_states==0)
		N_states = length(omega)+1
	end
	#################################
	# Step 1: assemble operators
	#################################
	# Hamiltonian without control
	HK_free = RotationFrameDiagonal(omega,omega_drive)
	# Hamiltonain with control
	if(method=="exponential")
		Ω = 0.5*pi/(TC*sqrt(InitialState+1.0))
		θ = 0.0
		HK_control,HS_control = RotationFrameRamseyControl(Ω,θ,N_states)
	else
		HK_control,HS_control = RotationFrameRamseyControl(N_states)
	end
	# Lindblad operator
	L1,L2 = RotationFrameLindblad(gamma1,gamma2)
	# Assemble the operators for the vectorized system
	# Half pi pulse
	LK_half_pi,LS_half_pi,LD = make_lindblad_operator(Matrix(HK_control+HK_free),Matrix(HS_control),(L1,L2))
	# Free evolution
	LK_free = make_hamiltonian_operator(Matrix(HK_free))
	#################################
	# Step 2: Forward solves
	#################################
	# initial conditions
	if(initial_type == "states")
		rho_u_initial = (rho_u0*transpose(rho_u0)+rho_v0*transpose(rho_v0))[:]
		rho_v_initial = (rho_v0*transpose(rho_u0)-rho_u0*transpose(rho_v0))[:]
		rho_vec0 = [rho_u_initial;rho_v_initial]
	elseif(initial_type == "density")
		rho_vec0 = [rho_u0;rho_v0];
	else
		println("Error! initial_type must be \"density\" or \"states\"")
		return
	end
	N_dark_times = length(t_dark_times)
	if(method=="exponential")
	#if(false)
		# Ramsey experiment
		Half_pi_operator = exp(0.5*TC*[LD+LS_half_pi -LK_half_pi; LK_half_pi LD+LS_half_pi])
		FreeEvolution = [ LD -LK_free;
					  	  LK_free  LD]
		# half-pi pusle
		rho_vec1 = Half_pi_operator*rho_vec0
		# free propagation for a dark time
		Free_operator = exp(t_dark_times[1]*FreeEvolution)
		rho_vec2 = Free_operator*rho_vec1
		# half-pi pulse
		rho_ramsey = Half_pi_operator*rho_vec2
		for i = 2:N_dark_times
			# free propagation for a dark time
			Free_operator = exp(t_dark_times[i]*FreeEvolution)
			rho_vec2 = Free_operator*rho_vec1
			# half-pi pulse
			rho_ramsey = [rho_ramsey Half_pi_operator*rho_vec2]
		end
	else
		function GaussianPulse(_time,_amplitude,_width)
			σ = _width/sqrt(2.0*pi)
			return _amplitude*exp(-(_time-0.5*TC)^2/(2.0*σ*σ))
		end
		width = TC/2.5
		GaussianAmplitude = 0.5*pi/(width*sqrt(InitialState+1.0))
		function ControlledODE!(_drho,_rho,_p,_t)
			# Half pi pulse
			_LK_half_pi,_LS_half_pi,_LD = make_lindblad_operator(HK_control*GaussianPulse(_t,GaussianAmplitude,width)+HK_free,
															     HS_control*GaussianPulse(_t,GaussianAmplitude,width),
															     (L1,L2))
			SystemOperator  = [_LD+_LS_half_pi -_LK_half_pi;
					  		   _LK_half_pi      _LD+_LS_half_pi]
			_drho .= SystemOperator*_rho
		end
		# DiffEq interface for the controlled problem
		controlled_problem = ODEProblem(ControlledODE!,rho_vec0,(0.0,0.5*TC),[1.0])
		# Free propagation part
		FreeEvolution = [ LD -LK_free;
						  LK_free  LD]
		# half-pi pusle
		if(method=="DiffEqDefault")
			sol = solve(controlled_problem;DiffEqKwargs...)
		else
			sol = solve(controlled_problem,method;DiffEqKwargs...)
		end
		rho_vec1 = sol[:,end]
		# free propagation for a dark time
		Free_operator = exp(t_dark_times[1]*FreeEvolution)
		rho_vec2 = Free_operator*rho_vec1
		# second half-pi pulse
		controlled_problem = remake(controlled_problem,u0=rho_vec2)
		if(method=="DiffEqDefault")
			sol = solve(controlled_problem;DiffEqKwargs...)
		else
			sol = solve(controlled_problem,method;DiffEqKwargs...)
		end
		rho_ramsey = sol[:,end]
		# free propagation for a dark time + second half-pi pulse
		n_dark_times = length(t_dark_times)
		for i = 2:n_dark_times
			# free propagation
			Free_operator = exp(t_dark_times[i]*FreeEvolution)
			rho_vec2 = Free_operator*rho_vec1
			# second half pi pulse
			controlled_problem = remake(controlled_problem,u0=rho_vec2)
			if(method=="DiffEqDefault")
				sol = solve(controlled_problem;DiffEqKwargs...)
			else
				sol = solve(controlled_problem,method;DiffEqKwargs...)
			end
			rho_ramsey = [ rho_ramsey sol[:,end] ]
		end
	end
	return rho_ramsey[1:N_states^2,:],rho_ramsey[N_states^2+1:end,:]
end

# Defined for the compatability of adjoint solves in Turing.jl to hold their data structure
function RamseyForwardSolve(rho_u0,rho_v0,
				   		  omega,omega_drive,
				   		  gamma1,gamma2,
						  InitialState,
				   		  TC,t_dark_times,
				   		  N_states::Int64=0;
						  initial_type="states",
						  method="exponential",
						  DiffEqKwargs...)
    if(N_states==0)
		N_states = length(omega)+1
	end
	#################################
	# Step 1: assemble operators
	#################################
	# Hamiltonian without control
	HK_free = RotationFrameDiagonal(omega,omega_drive)
	# Hamiltonain with control
	if(method=="exponential")
		Ω = 0.5*pi/(TC*sqrt(InitialState+1.0))
		θ = 0.0
		HK_control,HS_control = RotationFrameRamseyControl(Ω,θ,N_states)
	else
		HK_control,HS_control = RotationFrameRamseyControl(N_states)
	end
	# Lindblad operator
	L1,L2 = RotationFrameLindblad(gamma1,gamma2)
	# Assemble the operators for the vectorized system
	# Half pi pulse
	LK_half_pi,LS_half_pi,LD = make_lindblad_operator(Matrix(HK_control+HK_free),Matrix(HS_control),(L1,L2))
	# Free evolution
	LK_free = make_hamiltonian_operator(Matrix(HK_free))
	#################################
	# Step 2: Forward solves
	#################################
	# initial conditions
	if(initial_type == "states")
		rho_u_initial = (rho_u0*transpose(rho_u0)+rho_v0*transpose(rho_v0))[:]
		rho_v_initial = (rho_v0*transpose(rho_u0)-rho_u0*transpose(rho_v0))[:]
		rho_vec0 = [rho_u_initial;rho_v_initial]
	elseif(initial_type == "density")
		rho_vec0 = [rho_u0;rho_v0];
	else
		println("Error! initial_type must be \"density\" or \"states\"")
		return
	end
	N_dark_times = length(t_dark_times)
	if(method=="exponential")
	#if(false)
		# Ramsey experiment
		Half_pi_operator = exp(0.5*TC*[LD+LS_half_pi -LK_half_pi; LK_half_pi LD+LS_half_pi])
		FreeEvolution = [ LD -LK_free;
					  	  LK_free  LD]
		# half-pi pusle
		rho_vec1 = Half_pi_operator*rho_vec0
		# free propagation for a dark time
		Free_operator = exp(t_dark_times[1]*FreeEvolution)
		rho_vec2 = Free_operator*rho_vec1
		# half-pi pulse
		rho_ramsey = Half_pi_operator*rho_vec2
		for i = 2:N_dark_times
			# free propagation for a dark time
			Free_operator = exp(t_dark_times[i]*FreeEvolution)
			rho_vec2 = Free_operator*rho_vec1
			# half-pi pulse
			rho_ramsey = [rho_ramsey Half_pi_operator*rho_vec2]
		end
	else
		function GaussianPulse(_time,_amplitude,_width)
			σ = _width/sqrt(2.0*pi)
			return _amplitude*exp(-(_time-0.5*TC)^2/(2.0*σ*σ))
		end
		width = TC/2.5
		GaussianAmplitude = 0.5*pi/(width*sqrt(InitialState+1.0))
		function ControlledODE!(_drho,_rho,_p,_t)
			# Half pi pulse
			_LK_half_pi,_LS_half_pi,_LD = make_lindblad_operator(HK_control*GaussianPulse(_t,GaussianAmplitude,width)+HK_free,
															     HS_control*GaussianPulse(_t,GaussianAmplitude,width),
															     (L1,L2))
			SystemOperator  = [_LD+_LS_half_pi -_LK_half_pi;
					  		   _LK_half_pi      _LD+_LS_half_pi]
			_drho .= SystemOperator*_rho
		end
		# DiffEq interface for the controlled problem
		controlled_problem = ODEProblem(ControlledODE!,rho_vec0,(0.0,0.5*TC),[omega;gamma1;gamma2])
		# Free propagation part
		FreeEvolution = [ LD -LK_free;
						  LK_free  LD]
		# half-pi pusle
		if(method=="DiffEqDefault")
			sol = solve(controlled_problem;DiffEqKwargs...)
		else
			sol = solve(controlled_problem,method;DiffEqKwargs...)
		end
		rho_vec1 = sol[:,end]
		# free propagation for a dark time
		Free_operator = exp(t_dark_times[1]*FreeEvolution)
		rho_vec2 = Free_operator*rho_vec1
		# second half-pi pulse
		controlled_problem = remake(controlled_problem,u0=rho_vec2)
		if(method=="DiffEqDefault")
			sol = solve(controlled_problem;DiffEqKwargs...)
		else
			sol = solve(controlled_problem,method;DiffEqKwargs...)
		end
		rho_ramsey = sol[:,end]
		# free propagation for a dark time + second half-pi pulse
		n_dark_times = length(t_dark_times)
		for i = 2:n_dark_times
			# free propagation
			Free_operator = exp(t_dark_times[i]*FreeEvolution)
			rho_vec2 = Free_operator*rho_vec1
			# second half pi pulse
			controlled_problem = remake(controlled_problem,u0=rho_vec2)
			if(method=="DiffEqDefault")
				sol = solve(controlled_problem;DiffEqKwargs...)
			else
				sol = solve(controlled_problem,method;DiffEqKwargs...)
			end
			rho_ramsey = [ rho_ramsey sol[:,end] ]
		end
	end
	return rho_ramsey[1:N_states^2,:],rho_ramsey[N_states^2+1:end,:]
end
"""
    RamseyParityForwardSolve(rho_u0::Array{Float64},rho_v0::Array{Float64},
					   omega::Array{Float64},omega_drive:Float64,
					   charge_noise::Array{Float64},
					   gamma1::Array{Float64},gamma2::Array{Float64},
					   InitalState::Int64,
					   TC::Float64,tdark_times::Array{Float64},
					   N_states::Int64=0;
					   initial_type="states",
					   method="exponential",
					   DiffEqKwargs...)

# Argument:
- rho_u0,rho_v0: initial states ``\\rho_{u_0}-i\\rho_{v_0}``
- omega: transition frequencies
- omega_drive: driving frequency
- gamma1: the decay parameter  of the Lindblad operators, equal to ``1/T_1`` with ``T_1`` being the relaxation time
- gamma2: the dephasing parameter  of the Lindblad operators, equal to ``1/T_2`` with ``T_2`` being the dephasing time
- Initial State: initial state of the density matrix
- TC: control time of the control signal
- dark_times: the dark time in the Ramsey experiment
- N_states: number of states. If it is not specified, the default value will be 0, and the code will automatically detect how many states are considered.
- initial_type: type of the initial condition
  1. initial values are state vectors (default choice)
  2. initial values are density matrices saving in flattened vectors
- method: method to solve the Landlady system
 1. "exponential": exponential time integrator (default choice)
 2. "DiffEqDefault": a default choice made by DifferentialEquations.jl
 3. Other solvers available in DifferentialEquations.jl, for example, method = Trapezoid()
- DiffEqKwargs: keyword arguments feed to the ``solve function" of DifferentialEquations.jl

# Parity event
- Transition frequency will take equal probability to be omega``\\pm``0.5``\\times``charge_noise

# Output:
- ``\\rho_u,\\rho_v``: real and imaginary part of the density matrix ``\\rho=\\rho_u-i\\rho_v`` at sampled dark times, and they are saved as flattened vectors
"""
function RamseyParityForwardSolve(rho_u0::Array{Float64},rho_v0::Array{Float64},
				   		  omega::Array{Float64},omega_drive::Float64,
						  charge_noise::Array{Float64},
				   		  gamma1::Array{Float64},gamma2::Array{Float64},
						  InitialState::Int64,
				   		  TC::Float64,t_dark_times::Array{Float64},
				   		  N_states::Int64=0;
						  initial_type="states",
						  method="exponential",
						  DiffEqKwargs...)

	rho_u1, rho_v1 = RamseyForwardSolve(rho_u0,rho_v0,
					   		  omega-0.5.*charge_noise,omega_drive,
					   		  gamma1,gamma2,
							  InitialState,
					   		  TC,t_dark_times,
					   		  N_states;
							  initial_type,
							  method,
							  DiffEqKwargs...)
	rho_u2, rho_v2 = RamseyForwardSolve(rho_u0,rho_v0,
							  omega+0.5.*charge_noise,omega_drive,
							  gamma1,gamma2,
							  InitialState,
							  TC,t_dark_times,
							  N_states;
							  initial_type,
							  method,
							  DiffEqKwargs...)
	return 0.5.*(rho_u1+rho_u2),0.5.*(rho_v1+rho_v2)
end

"""
    EchoForwardSolve(rho_u0::Array{Float64},rho_v0::Array{Float64},
					   omega::Array{Float64},omega_drive::Float64,
					   gamma1::Array{Float64},gamma2::Array{Float64},
					   InitialState::Int64,
					   TC::Float64,tdark_times::Array{Float64},
					   N_states::Int64=0;
					   initial_type="states",
					   method="exponential",
					   DiffEqKwargs...)

# Argument:
- rho_u0,rho_v0: initial states ``\\rho_{u_0}-i\\rho_{v_0}``
- omega: transition frequencies
- omega_drive: driving frequency
- gamma1: the decay parameter  of the Lindblad operators, equal to ``1/T_1`` with ``T_1`` being the relaxation time
- gamma2: the dephasing parameter  of the Lindblad operators, equal to ``1/T_2`` with ``T_2`` being the dephasing time
- InitialState: initial state of the density matrix
- TC: control time of the control signal
- tdark_times: the dark time in the Ramsey experiment
- N_states: number of states. If it is not specified, thedefault value will be 0, and the code will automatically detect how many states are considered.
- initial_type: type of the initial condition
  1. initial values are state vectors (default choice)
  2. initial values are density matrices saving in flattened vectors
- method: method to solve the Lindblad system
 1. "exponential": exponential time integrator (default choice)
 2. "DiffEqDefault": a default choice made by DifferentialEquations.jl
 3. Other solvers available in DifferentialEquations.jl, for example, method = Trapezoid()
- DiffEqKwargs: keyword arguments feed to the ``solve function" of DifferentialEquations.jl

# Output:
- ``\\rho_u,\\rho_v``: real and imaginary part of the density matrix ``\\rho=\\rho_u-i\\rho_v`` at sampled dark times, and they are saved as flattened vectors
"""
function EchoForwardSolve(rho_u0::Array{Float64},rho_v0::Array{Float64},
				   omega::Array{Float64},omega_drive::Float64,
				   gamma1::Array{Float64},gamma2::Array{Float64},
				   InitialState::Int64,
				   TC::Float64,t_dark_times::Array{Float64},
				   N_states::Int64=0;
				   initial_type="states",
				   method="exponential",
				   DiffEqKwargs...)

	if(N_states==0)
		N_states = length(omega)+1
	end
	#################################
	# Step 1: assemble operators
	#################################
	# Hamiltonian without control
	HK_free = RotationFrameDiagonal(omega,omega_drive)
	# Hamiltonain with control
	if(method=="exponential")
		Ω = 0.5*pi/(TC*sqrt(InitialState+1.0))
		θ = 0.0
		HK_control,HS_control = RotationFrameRamseyControl(Ω,θ,N_states)
	else
		HK_control,HS_control = RotationFrameRamseyControl(N_states)
	end
	# Lindblad operator
	L1,L2 = RotationFrameLindblad(gamma1,gamma2)
	# Assemble the operators for the vectorized system
	# Half pi pulse
	LK_half_pi,LS_half_pi,LD = make_lindblad_operator(Matrix(HK_control+HK_free),Matrix(HS_control),(L1,L2))
	# Free evolution
	LK_free = make_hamiltonian_operator(Matrix(HK_free))
	#################################
	# Step 2: Forward solves
	#################################
	# initial conditions
	if(initial_type == "states")
		rho_u_initial = (rho_u0*transpose(rho_u0)+rho_v0*transpose(rho_v0))[:]
		rho_v_initial = (rho_v0*transpose(rho_u0)-rho_u0*transpose(rho_v0))[:]
		rho_vec0 = [rho_u_initial;rho_v_initial]
	elseif(initial_type == "density")
		rho_vec0 = [rho_u0;rho_v0];
	else
		println("Error! initial_type must be \"density\" or \"states\"")
		return
	end
	N_dark_times = length(t_dark_times)
	if(method=="exponential")
		# Assemble operators
		Half_pi_operator = exp(0.5*TC*[LD+LS_half_pi -LK_half_pi; LK_half_pi LD+LS_half_pi])
		Pi_operator = Half_pi_operator*Half_pi_operator
		FreeEvolution = [ LD -LK_free;
					      LK_free  LD]
					  # half-pi pusle
		rho_vec1 = Half_pi_operator*rho_vec0
		# free propagation for a dark time
		Free_operator = exp(t_dark_times[1]*FreeEvolution)
		rho_vec2 = Free_operator*rho_vec1
		# pi pulse
		rho_vec2 = Pi_operator*rho_vec2
		# free propagation for the same dark time
		rho_vec2 = Free_operator*rho_vec2
		# half-pi pulse
		rho_echo = Half_pi_operator*rho_vec2
		for i = 2:N_dark_times
			# free propagation for a dark time
			Free_operator = exp(t_dark_times[i]*FreeEvolution)
			rho_vec2 = Free_operator*rho_vec1
			# pi pulse
			rho_vec2 = Pi_operator*rho_vec2
			# free propagation for the same dark time
			rho_vec2 = Free_operator*rho_vec2
			# half-pi pulse
			rho_echo = [rho_echo Half_pi_operator*rho_vec2]
		end
	else
		function GaussianPulse(_time,_amplitude,_width)
			_tot_time = _width*2.5
			σ = _width/sqrt(2.0*pi)
			return _amplitude*exp(-(_time-0.5*_tot_time)^2/(2.0*σ*σ))
		end
		width = TC/2.5
		GaussianAmplitude = 0.5*pi/(width*sqrt(InitialState+1.0))
		function ControlledODE!(_drho,_rho,_p,_t)
			# Half pi pulse
			_LK_half_pi,_LS_half_pi,_LD = make_lindblad_operator(HK_control*GaussianPulse(_t,GaussianAmplitude,width)+HK_free,
															     HS_control*GaussianPulse(_t,GaussianAmplitude,width),
															     (L1,L2))
			SystemOperator  = [_LD+_LS_half_pi -_LK_half_pi;
					  		   _LK_half_pi      _LD+_LS_half_pi]
			_drho .= SystemOperator*_rho
		end
		# DiffEq interface for the controlled problem
		half_pi_pulse = ODEProblem(ControlledODE!,rho_vec0,(0.0,0.5*TC),[1.0])
		# Free propagation part
		FreeEvolution = [ LD -LK_free;
						  LK_free  LD]
		# half-pi pusle
		if(method=="DiffEqDefault")
			sol = solve(half_pi_pulse;DiffEqKwargs...)
		else
			sol = solve(half_pi_pulse,method;DiffEqKwargs...)
		end
		rho_vec1 = sol[:,end]
		# free propagation for a dark time
		Free_operator = exp(t_dark_times[1]*FreeEvolution)
		rho_vec2 = Free_operator*rho_vec1
		# pi pulse
		pi_pulse = ODEProblem(ControlledODE!,rho_vec2,(0.0,TC),[1.0])
		if(method=="DiffEqDefault")
			sol = solve(pi_pulse;DiffEqKwargs...)
		else
			sol = solve(pi_pulse,method;DiffEqKwargs...)
		end
		# free propagation for the same dark time
		rho_vec2 = Free_operator*sol[:,end]
		# half pi pulse
		half_pi_pulse = remake(half_pi_pulse,u0=rho_vec2)
		if(method=="DiffEqDefault")
			sol = solve(half_pi_pulse;DiffEqKwargs...)
		else
			sol = solve(half_pi_pulse,method;DiffEqKwargs...)
		end
		rho_echo = sol[:,end]
		# free propagation for a dark time + second half-pi pulse
		n_dark_times = length(t_dark_times)
		for i = 2:n_dark_times
			Free_operator = exp(t_dark_times[i]*FreeEvolution)
			# second half pi
			rho_vec2 = Free_operator*rho_vec1
			# pi pulse
			pi_pulse = remake(pi_pulse,u0=rho_vec2)
			if(method=="DiffEqDefault")
				sol = solve(pi_pulse;DiffEqKwargs...)
			else
				sol = solve(pi_pulse,method;DiffEqKwargs...)
			end
			# free propagation for the same dark time
			rho_vec2 = Free_operator*sol[:,end]
			# half pi pulse
			half_pi_pulse = remake(half_pi_pulse,u0=rho_vec2)
			if(method=="DiffEqDefault")
				sol = solve(half_pi_pulse;DiffEqKwargs...)
			else
				sol = solve(half_pi_pulse,method;DiffEqKwargs...)
			end
			rho_echo = [rho_echo sol[:,end]]
		end

	end

	return rho_echo[1:N_states^2,:],rho_echo[N_states^2+1:end,:]
end

"""
    EchoParityForwardSolve(rho_u0::Array{Float64},rho_v0::Array{Float64},
					   omega::Array{Float64},omega_drive::Float64,
					   charge_noise::Array{Float64},
					   gamma1::Array{Float64},gamma2::Array{Float64},
					   InitialState::Int64,
					   TC::Float64,tdark_times::Array{Float64},
					   N_states::Int64=0;
					   initial_type="states",
					   method="exponential",
					   DiffEqKwargs...)

# Argument:
- rho_u0,rho_v0: initial states ``\\rho_{u_0}-i\\rho_{v_0}``
- omega: transition frequencies
- omega_drive: driving frequency
- gamma1: the decay parameter  of the Lindblad operators, equal to ``1/T_1`` with ``T_1`` being the relaxation time
- gamma2: the dephasing parameter  of the Lindblad operators, equal to ``1/T_2`` with ``T_2`` being the dephasing time
- InitialState: initial state of the density matrix
- TC: control time of the control signal
- tdark_times: the dark time in the Ramsey experiment
- N_states: number of states. If it is not specified, the default value will be 0, and the code will automatically detect how many states are considered.
- initial_type: type of the initial condition
  1. initial values are state vectors (default choice)
  2. initial values are density matrices saving in flattened vectors
- method: method to solve the Lindblad system
 1. "exponential": exponential time integrator (default choice)
 2. "DiffEqDefault": a default choice made by DifferentialEquations.jl
 3. Other solvers available in DifferentialEquations.jl, for example, method = Trapezoid()
- DiffEqKwargs: keyword arguments feed to the ``solve function" of DifferentialEquations.jl

# Parity event
- Transition frequency will take equal probability to be omega``\\pm``0.5``\\times``charge_noise

# Output:
- ``\\rho_u,\\rho_v``: real and imaginary part of the density matrix ``\\rho=\\rho_u-i\\rho_v`` at sampled dark times, and they are saved as flattened vectors
"""
function EchoParityForwardSolve(rho_u0::Array{Float64},rho_v0::Array{Float64},
				   omega::Array{Float64},omega_drive::Float64,
				   charge_noise::Array{Float64},
				   gamma1::Array{Float64},gamma2::Array{Float64},
				   InitialState::Int64,
				   TC::Float64,t_dark_times::Array{Float64},
				   N_states::Int64=0;
				   initial_type="states",
				   method="exponential",
				   DiffEqKwargs...)

	rho_u1, rho_v1 = EchoForwardSolve(rho_u0,rho_v0,
							  omega-0.5.*charge_noise,omega_drive,
							  gamma1,gamma2,
							  InitialState,
							  TC,t_dark_times,
							  N_states;
							  initial_type,
							  method,
							  DiffEqKwargs...)
	rho_u2, rho_v2 = EchoForwardSolve(rho_u0,rho_v0,
							  omega+0.5.*charge_noise,omega_drive,
							  gamma1,gamma2,
							  InitialState,
							  TC,t_dark_times,
							  N_states;
							  initial_type,
							  method,
							  DiffEqKwargs...)
	return 0.5.*(rho_u1+rho_u2),0.5.*(rho_v1+rho_v2)
end

"""
    T1ForwardSolve(rho_u0::Array{Float64},rho_v0::Array{Float64},
					   omega::Array{Float64},omega_drive::Float64,
					   gamma1::Array{Float64},gamma2::Array{Float64},
					   InitialState::Int64,
					   TC::Float64,tdark_times::Array{Float64},
					   N_states::Int64=0;
					   initial_type="states",
					   method="exponential",
					   DiffEqKwargs...)

# Argument:
- rho_u0,rho_v0: initial states ``\\rho_{u_0}-i\\rho_{v_0}``
- omega: transition frequencies
- omega_drive: driving frequency
- gamma1: the decay parameter  of the Lindblad operators, equal to ``1/T_1`` with ``T_1`` being the relaxation time
- gamma2: the dephasing parameter  of the Lindblad operators, equal to ``1/T_2`` with ``T_2`` being the dephasing time
- InitialState: initial state of the density matrix
- TC: control time of the control signal
- tdark_times: the dark time in the Ramsey experiment
- N_states: number of states. If it is not specified, the default value will be 0, and the code will automatically detect how many states are considered.
- initial_type: type of the initial condition
  1. initial values are state vectors (default choice)
  2. initial values are density matrices saving in flattened vectors
- method: method to solve the Lindblad system
 1. "exponential": exponential time integrator (default choice)
 2. "DiffEqDefault": a default choice made by DifferentialEquations.jl
 3. Other solvers available in DifferentialEquations.jl, for example, method = Trapezoid()
- DiffEqKwargs: keyword arguments feed to the ``solve function" of DifferentialEquations.jl

# Output:
- ``\\rho_u,\\rho_v``: real and imaginary part of the density matrix ``\\rho=\\rho_u-i\\rho_v`` at sampled dark times, and they are saved as flattened vectors
"""
function T1ForwardSolve(rho_u0::Array{Float64},rho_v0::Array{Float64},
				   omega::Array{Float64},omega_drive::Float64,
				   gamma1::Array{Float64},gamma2::Array{Float64},
				   InitialState::Int64,
				   TC::Float64,t_dark_times::Array{Float64},
				   N_states::Int64=0;
				   initial_type="states",
				   method="exponential",
				   DiffEqKwargs...)

	if(N_states==0)
		N_states = length(omega)+1
	end
	#################################
	# Step 1: assemble operators
	#################################
	# Hamiltonian without control
	HK_free = RotationFrameDiagonal(omega,omega_drive)
	# Hamiltonain with control
	# Hamiltonain with control
	if(method=="exponential")
		Ω = 0.5*pi/(TC*sqrt(InitialState+1.0))
		θ = 0.0
		HK_control,HS_control = RotationFrameRamseyControl(Ω,θ,N_states)
	else
		HK_control,HS_control = RotationFrameRamseyControl(N_states)
	end
	# Lindblad operator
	L1,L2 = RotationFrameLindblad(gamma1,gamma2)
	# Assemble the operators for the vectorized system
	# Half pi pulse
	LK_half_pi,LS_half_pi,LD = make_lindblad_operator(Matrix(HK_control+HK_free),Matrix(HS_control),(L1,L2))
	# Free evolution
	LK_free = make_hamiltonian_operator(Matrix(HK_free))
	#################################
	# Step 2: Forward solves
	#################################
	# initial conditions
	if(initial_type == "states")
		rho_u_initial = (rho_u0*transpose(rho_u0)+rho_v0*transpose(rho_v0))[:]
		rho_v_initial = (rho_v0*transpose(rho_u0)-rho_u0*transpose(rho_v0))[:]
		rho_vec0 = [rho_u_initial;rho_v_initial]
	elseif(initial_type == "density")
		rho_vec0 = [rho_u0;rho_v0];
	else
		println("Error! initial_type must be \"density\" or \"states\"")
		return
	end
	N_dark_times = length(t_dark_times)
	if(method=="exponential")
		# Assemble operators
		Pi_operator = exp(TC*[LD+LS_half_pi -LK_half_pi; LK_half_pi LD+LS_half_pi])
		FreeEvolution = [ LD -LK_free;
					  	  LK_free  LD]
		# pi pusle
		rho_vec1 = Pi_operator*rho_vec0
		# free propagation for a dark time
		Free_operator = exp(t_dark_times[1]*FreeEvolution)
		rho_T1 = Free_operator*rho_vec1
		for i = 2:N_dark_times
			# free propagation for a dark time
			Free_operator = exp(t_dark_times[i]*FreeEvolution)
			rho_T1 = [rho_T1 Free_operator*rho_vec1]
		end
	else
		function GaussianPulse(_time,_amplitude,_width)
			σ = _width/sqrt(2.0*pi)
			return _amplitude*exp(-(_time-0.5*TC)^2/(2.0*σ*σ))
		end
		width = TC/2.5
		GaussianAmplitude = 0.5*pi/(width*sqrt(InitialState+1.0))
		function ControlledODE!(_drho,_rho,_p,_t)
			# Half pi pulse
			_LK_half_pi,_LS_half_pi,_LD = make_lindblad_operator(HK_control*GaussianPulse(_t,GaussianAmplitude,width)+HK_free,
																 HS_control*GaussianPulse(_t,GaussianAmplitude,width),
																 (L1,L2))
			SystemOperator  = [_LD+_LS_half_pi -_LK_half_pi;
							   _LK_half_pi      _LD+_LS_half_pi]
			_drho .= SystemOperator*_rho
		end
		# DiffEq interface for the controlled problem
		pi_pulse = ODEProblem(ControlledODE!,rho_vec0,(0.0,TC),[1.0])
		# Free propagation part
		FreeEvolution = [ LD -LK_free;
						  LK_free  LD]
		# pi pusle
		if(method=="DiffEqDefault")
			sol = solve(pi_pulse;DiffEqKwargs...)
		else
			sol = solve(pi_pulse,method;DiffEqKwargs...)
		end
		rho_vec1 = sol[:,end]
		Free_operator = exp(t_dark_times[1]*FreeEvolution)
		rho_T1 = Free_operator*rho_vec1
		for i = 2:N_dark_times
			# free propagation for a dark time
			Free_operator = exp(t_dark_times[i]*FreeEvolution)
			rho_T1 = [rho_T1 Free_operator*rho_vec1]
		end
	end
	return rho_T1[1:N_states^2,:],rho_T1[N_states^2+1:end,:]

end

"""
    T1ParityForwardSolve(rho_u0::Array{Float64},rho_v0::Array{Float64},
					   omega::Array{Float64},omega_drive::Float64,
					   charge_noise::Array{Float64},
					   gamma1::Array{Float64},gamma2::Array{Float64},
					   InitialState::Int64,
					   TC::Float64,tdark_times::Array{Float64},
					   N_states::Int64=0;
					   initial_type="states",
					   method="exponential",
					   DiffEqKwargs...)

# Argument:
- rho_u0,rho_v0: initial states ``\\rho_{u_0}-i\\rho_{v_0}``
- omega: transition frequencies
- omega_drive: driving frequency
- gamma1: the decay parameter  of the Lindblad operators, equal to ``1/T_1`` with ``T_1`` being the relaxation time
- gamma2: the dephasing parameter  of the Lindblad operators, equal to ``1/T_2`` with ``T_2`` being the dephasing time
- InitialState: initial state of the density matrix
- TC: control time of the control signal
- tdark_times: the dark time in the Ramsey experiment
- N_states: number of states. If it is not specified, the default value will be 0, and the code will automatically detect how many states are considered.
- initial_type: type of the initial condition
  1. initial values are state vectors (default choice)
  2. initial values are density matrices saving in flattened vectors
- method: method to solve the Lindblad system
 1. "exponential": exponential time integrator (default choice)
 2. "DiffEqDefault": a default choice made by DifferentialEquations.jl
 3. Other solvers available in DifferentialEquations.jl, for example, method = Trapezoid()
- DiffEqKwargs: keyword arguments feed to the ``solve function" of DifferentialEquations.jl

# Parity event
- Transition frequency will take equal probability to be omega``\\pm``0.5``\\times``charge_noise

# Output:
- ``\\rho_u,\\rho_v``: real and imaginary part of the density matrix ``\\rho=\\rho_u-i\\rho_v`` at sampled dark times, and they are saved as flattened vectors
"""
function T1ParityForwardSolve(rho_u0::Array{Float64},rho_v0::Array{Float64},
				   omega::Array{Float64},omega_drive::Float64,
				   charge_noise::Array{Float64},
				   gamma1::Array{Float64},gamma2::Array{Float64},
				   InitialState::Int64,
				   TC::Float64,t_dark_times::Array{Float64},
				   N_states::Int64=0;
				   initial_type="states",
				   method="exponential",
				   DiffEqKwargs...)

	rho_u1, rho_v1 = T1ForwardSolve(rho_u0,rho_v0,
							  omega-0.5.*charge_noise,omega_drive,
							  gamma1,gamma2,
							  InitialState,
							  TC,t_dark_times,
							  N_states;
							  initial_type,
							  method,
							  DiffEqKwargs...)
	rho_u2, rho_v2 = T1ForwardSolve(rho_u0,rho_v0,
							  omega+0.5.*charge_noise,omega_drive,
							  gamma1,gamma2,
							  InitialState,
							  TC,t_dark_times,
							  N_states;
							  initial_type,
							  method,
							  DiffEqKwargs...)
	return 0.5.*(rho_u1+rho_u2),0.5.*(rho_v1+rho_v2)
end


# NOTE: This script is used for generating data for the PPL comparisons.

import Pkg; Pkg.activate("../../../")

using Distributions
using JSON3
import Random

"""
Function to generate data for univariate Gaussian mixture model (GMM).

Args
====

- `n`: number of observations
- `mu`: Mixture component means
- `sig`: Standard deviations of each component
- `w`: Mixture component weights
- `seed`: Random number generator seed
"""
function generate_gmm_data(; n::Integer,
                           mu::Vector{Float64}, sig::Vector{Float64},
                           w::Vector{Float64},
                           seed=nothing) where T <: AbstractFloat
    if seed != nothing
        Random.seed!(seed)
    end

    # Number of mixture components
    num_components = length(mu)

    # Check that w is a proper probability vector
    @assert all(w .> 0) && sum(w) == 1

    # Generate random class labels
    rand_idx = wsample(1:num_components, w, n)

    # Generate random observations
    y = mu[rand_idx] + randn(n) .* sig[rand_idx]

    return Dict(:rand_idx => rand_idx,
                :mu => mu, :sig => sig, :w => w, :y => y)
end

function main(; seed=1, num_components=4,
              datasizes=[50, 100, 200, 400, 800, 1600],
              data_dir=joinpath(@__DIR__, "../data"))
  Random.seed!(seed)

  # simulate mu
  mu = let
      _mu = collect(1:num_components)
      centered_mu = _mu .- (num_components + 1) / 2
      jittered_mu = centered_mu .+ randn(num_components) * .1
      jittered_mu
  end

  # simulate sigma
  sig = rand(Uniform(.1, .2), num_components)

  # simulate weights
  w = let
      w_tmp = rand(Uniform(.8/num_components, 1.2/num_components),
                   num_components)
      w_tmp ./ sum(w_tmp)
  end

  # Create data sets and write to disk
  for n in datasizes
      # Generate GMM data.
      gmm_data = generate_gmm_data(n=n, mu=mu, sig=sig, w=w, seed=seed)

      # Write GMM data to json string.
      gmm_data_json = JSON3.write(gmm_data);

      # Write json string to disk.
      open("$(data_dir)/gmm-data-n$(n).json", "w") do io
        write(io, gmm_data_json)
      end
  end
end

# Generate data.
main()


import Distances.pairwise!
import Distances.pairwise
import Distances.colwise!
import Distances.colwise
import Distances.evaluate

## Generic distance evaluation
#define Abstract type Distance
abstract type Distance end
struct sql2 <: Distance end
struct l2 <: Distance end
struct l1 <: Distance end
evaluate(::sql2, a::AbstractVector{T}, b::AbstractVector{T}) where {T} = sql2dist(a,b)
evaluate(::l2, a::AbstractVector{T}, b::AbstractVector{T}) where {T} = l2dist(a,b)
evaluate(::l1, a::AbstractVector{T}, b::AbstractVector{T}) where {T} = l1dist(a,b)

#Columnwise distances
function colwise!(d::AbstractVector{<:AbstractFloat}, dm::Distance, x::AbstractMatrix{T}, y::AbstractMatrix{T}) where {T}
    @inbounds @views Threads.@threads for i ∈ axes(y,2)
        d[i] = evaluate(dm, x[:,i], y[:,i])
    end
end

function colwise!(d::AbstractVector{<:AbstractFloat}, dm::Distance, x::AbstractVector{T}, y::AbstractMatrix{T}) where {T}
    @inbounds @views Threads.@threads for i ∈ axes(y,2)
        d[i] = evaluate(dm,x,y[:,i])
    end
end

function colwise(dm::Distance, x::AbstractVecOrMat{T}, y::AbstractMatrix{T}) where {T}
    d = Vector{T}(undef,size(y,2))
    colwise!(d,dm,x,y)
    return d
end

#Pairwise distances
function pairwise!(d::AbstractVector{<:AbstractFloat}, dm::Distance, x::AbstractVector{T}, y::AbstractMatrix{T}) where {T}
    @inbounds @views Threads.@threads for i ∈ axes(y,2)
        d[i] = evaluate(dm, x, y[:,i])
    end
end

pairwise!(d::AbstractVector{<:AbstractFloat}, dm::Distance, x::AbstractMatrix{T}, y::AbstractVector{T}) where {T} = pairwise!(d,dm,y,x)

function pairwise!(d::AbstractMatrix{<:AbstractFloat}, dm::Distance, x::AbstractMatrix{T}, y::AbstractMatrix{T}) where {T}
    @inbounds @views Threads.@threads for i ∈ axes(x,2)
        for j ∈ axes(y,2)
            d[i,j] = evaluate(dm, x[:,i], y[:,j])
        end
    end
end

function pairwise(dm::Distance, x::AbstractVector{T}, y::AbstractMatrix{T}) where {T}
    d = Vector{T}(undef,size(y,2))
    pairwise!(d,dm,x,y)
    return d
end

pairwise(dm::Distance, x::AbstractMatrix{T}, y::AbstractVector{T}) where {T} = pairwise(dm,y,x)

function pairwise(dm::Distance, x::AbstractMatrix{T}, y::AbstractMatrix{T}) where {T}
    d = Matrix{T}(undef,size(x,2),size(y,2))
    pairwise!(d,dm,x,y)
    return d
end


using Revise
	using DiffEqOperators, Setfield, Parameters
	using PseudoArcLengthContinuation, LinearAlgebra, Plots, SparseArrays
	const PALC = PseudoArcLengthContinuation

heatmapsol(x) = heatmap(reshape(Array(x), Nx, Ny)', color=:viridis)

Nx = 151*1
	Ny = 100*1
	lx = 8pi
	ly = 2*2pi/sqrt(3)

function Laplacian2D(Nx, Ny, lx, ly, bc = :Neumann)
	hx = 2lx/Nx
	hy = 2ly/Ny
	D2x = CenteredDifference(2, 2, hx, Nx)
	D2y = CenteredDifference(2, 2, hy, Ny)
	if bc == :Dirichlet
		Qx = Dirichlet0BC(typeof(hx))
		Qy = Dirichlet0BC(typeof(hy))
	elseif bc == :Neumann
		Qx = Neumann0BC(hx)
		Qy = Neumann0BC(hy)
	elseif bc == :Periodic
		Qx = PeriodicBC(hx)
		Qy = PeriodicBC(hy)
	end
	# @show norm(D2x - Circulant(D2x[1,:]))
	A = kron(sparse(I, Ny, Ny), sparse(D2x * Qx)[1]) + kron(sparse(D2y * Qy)[1], sparse(I, Nx, Nx))
	return A, D2x
end

function F_sh(u, p)
	@unpack l, ν, L1 = p
	return -L1 * u .+ (l .* u .+ ν .* u.^2 .- u.^3)
end

function dF_sh(u, p)
	@unpack l, ν, L1 = p
	return -L1 .+ spdiagm(0 => l .+ 2 .* ν .* u .- 3 .* u.^2)
end

X = -lx .+ 2lx/(Nx) * collect(0:Nx-1)
Y = -ly .+ 2ly/(Ny) * collect(0:Ny-1)

sol0 = [(cos(x) .+ cos(x/2) * cos(sqrt(3) * y/2) ) for x in X, y in Y]
		sol0 .= sol0 .- minimum(vec(sol0))
		sol0 ./= maximum(vec(sol0))
		sol0 = sol0 .- 0.25
		sol0 .*= 1.7
		heatmap(sol0', color=:viridis)

Δ, D2x = Laplacian2D(Nx, Ny, lx, ly, :Neumann)
L1 = (I + Δ)^2
par = (l = -0.1, ν = 1.3, L1 = L1)

optnew = PALC.NewtonPar(verbose = true, tol = 1e-8, maxIter = 20)
# optnew = PALC.NewtonPar(verbose = true, tol = 1e-8, maxIter = 20, eigsolver = EigArpack(0.5, :LM))
	sol_hexa, hist, flag = @time PALC.newton(
				x ->  F_sh(x, par),
				u -> dF_sh(u, par),
				vec(sol0),
				optnew)
	println("--> norm(sol) = ", norm(sol_hexa, Inf64))
	heatmapsol(sol_hexa)

heatmapsol(0.2vec(sol_hexa) .* vec([exp(-(x+0lx)^2/25) for x in X, y in Y]))

###################################################################################################
# recherche de solutions
deflationOp = DeflationOperator(2.0, (x, y)->dot(x, y), 1.0, [sol_hexa])

optnew = @set optnew.maxIter = 250
outdef, _, flag, _ = @time PALC.newton(
				x -> F_sh(x, par),
				u -> dF_sh(u, par),
				# 0.4vec(sol_hexa) .* vec([exp(-1(x+1lx)^2/25) for x in X, y in Y]),
				0.4vec(sol_hexa) .* vec([1 .- exp(-1(x+0lx)^2/55) for x in X, y in Y]),

				optnew, deflationOp, normN = x -> norm(x, Inf))
	println("--> norm(sol) = ", norm(outdef))
	heatmapsol(outdef) |> display
	flag && push!(deflationOp, outdef)

heatmapsol(deflationOp[4])

heatmapsol(0.4vec(sol_hexa) .* vec([1 .- exp(-1(x+0lx)^2/55) for x in X, y in Y]))
###################################################################################################
normC = x -> norm(x, Inf)
optcont = ContinuationPar(dsmin = 0.0001, dsmax = 0.005, ds= -0.001, pMax = -0.095, pMin = -1.0, newtonOptions = setproperties(optnew; tol = 1e-9, maxIter = 15), maxSteps = 145, detectBifurcation = 2, nev = 40, detectFold = false, dsminBisection =1e-7, saveSolEveryNsteps = 4)
	optcont = @set optcont.newtonOptions.eigsolver = EigArpack(0.1, :LM)

	br, u1 = @time PALC.continuation(
		(x, p) ->  F_sh(x, @set par.l = p),
		(x, p) -> dF_sh(x, @set par.l = p),
		deflationOp[1], -0.1, optcont;
		plot = true, verbosity = 3,
		tangentAlgo = BorderedPred(),
		# linearAlgo = MatrixBLS(),
		plotSolution = (x; kwargs...) -> (heatmap!(X, Y, reshape(x, Nx, Ny)'; color=:viridis, label="", kwargs...);ylims!(-1,1,subplot=4);xlims!(-.5,.3,subplot=4)),
		printSolution = (x,p) -> norm(x),
		finaliseSolution = (z, tau, step, contResult) -> (Base.display(contResult.eig[end].eigenvals) ;true) ,
		normC = x -> norm(x, Inf))
###################################################################################################
using IncompleteLU
prec = ilu(L1 + I,τ=0.05)
prec = lu(L1 + I)
ls = GMRESIterativeSolvers(tol = 1e-5, N = Nx*Ny, Pl = prec)

sol_hexa, _, flag = @time PALC.newton(
			x ->  F_sh(x, par),
			u -> dF_sh(u, par),
			vec(sol0),
			@set optnew.linsolver = ls)
	println("--> norm(sol) = ", norm(sol_hexa, Inf64))
	heatmapsol(sol_hexa)


using Test
using ITensors
using Random

Random.seed!(1234)

@testset "Test contractions with scalar-like ITensors" begin
  i = Index(2, "i")
  j = Index(2, "j")
  k = Index(2, "k")
  α = Index(1, "α")

  is = (i, j, k)

  A = randomITensor(is..., dag(α))
  B = ITensor(2, α, α', α'')

  C = A * B
  @test C ≈ B[1, 1, 1] * A * ITensor(1, inds(B))

  C = emptyITensor(is..., α', α'')
  C .= A .* B
  @test C ≈ B[1, 1, 1] * A * ITensor(1, inds(B))

  C = emptyITensor(shuffle([(is..., α', α'')...])...)
  C .= A .* B
  @test C ≈ B[1, 1, 1] * A * ITensor(1, inds(B))
end

@testset "NaN in-place contraction bug regression test" begin
  BlasFloats = (Float32, Float64, ComplexF32, ComplexF64)
  @testset "Scalar contract, no permutation" for ElA in BlasFloats, ElB in BlasFloats
    i = Index(2, "i")
    j = Index(3, "j")
    α = Index(1, "α")
    A = randomITensor(ElA, i, j, α')
    B = randomITensor(ElB, dag(α)', α)
    R = ITensor(promote_type(ElA, ElB), i, j, α)

    R .= NaN
    @test any(isnan, R)

    R .= A .* B
    @test !any(isnan, R)
    @test array(R) ≈ array(A) * array(B)[]

    R .= NaN
    @test any(isnan, R)

    R .= B .* A
    @test !any(isnan, R)
    @test array(R) ≈ array(A) * array(B)[]
  end

  @testset "Scalar contraction, permutation" for ElA in BlasFloats, ElB in BlasFloats
    i = Index(2, "i")
    j = Index(3, "j")
    α = Index(1, "α")
    A = randomITensor(ElA, i, j, α')
    B = randomITensor(ElB, dag(α)', α)
    R = ITensor(promote_type(ElA, ElB), j, i, α)

    R .= NaN
    @test any(isnan, R)

    R .= A .* B
    @test !any(isnan, R)
    @test array(R) ≈ permutedims(array(A), (2, 1, 3)) * array(B)[]

    R .= NaN
    @test any(isnan, R)

    R .= B .* A
    @test !any(isnan, R)
    @test array(R) ≈ permutedims(array(A), (2, 1, 3)) * array(B)[]
  end

  @testset "General contraction, permutation" for ElA in BlasFloats, ElB in BlasFloats
    i = Index(2, "i")
    j = Index(3, "j")
    α = Index(2, "α")
    A = randomITensor(ElA, i, j, α')
    B = randomITensor(ElB, dag(α)', α)
    R = ITensor(promote_type(ElA, ElB), j, i, α)

    R .= NaN
    @test any(isnan, R)

    R .= A .* B
    @test !any(isnan, R)
    @test reshape(array(R), 6, 2) ≈
      reshape(permutedims(array(A), (2, 1, 3)), 6, 2) * array(B)

    R .= NaN
    @test any(isnan, R)

    R .= B .* A
    @test !any(isnan, R)
    @test reshape(array(R), 6, 2) ≈
      reshape(permutedims(array(A), (2, 1, 3)), 6, 2) * array(B)
  end
end


using Flux, Test
using Tracker: gradcheck

gradtest(f, xs::AbstractArray...) = gradcheck((xs...) -> sum(sin.(f(xs...))), xs...)
gradtest(f, dims...) = gradtest(f, rand.(Float64, dims)...)

@testset "Tracker" begin

@test gradtest(Flux.mse, rand(5,5), rand(5, 5))
@test gradtest(Flux.crossentropy, rand(5,5), rand(5, 5))

@test gradtest(x -> Flux.normalise(x), rand(4,3))
@test gradtest(x -> Flux.normalise(x, dims = 2), rand(3,4))

end


module BLS12381FieldTest

using Test

using TowerFields.BLS12381Field

const A2 = (BLS12381Field.Q - 1), (BLS12381Field.Q - 2)
const B2 = (BLS12381Field.Q - 2), (BLS12381Field.Q - 3)
const C2 = (BLS12381Field.Q - 3), (BLS12381Field.Q - 4)

const A6 = (A2, B2, C2)
const B6 = (B2, C2, A2)
const C6 = (C2, A2, B2)
const D6 = (C2, B2, A2)

const A12 = (A6, C6)
const B12 = (B6, D6)

@testset "add: FQ12" begin
    let + = BLS12381Field.:+
        @test A12 + B12 == (((4002409555221667393417789825735904156556882819939007885332058136124031650490837864442687629129015664037894272559784, 4002409555221667393417789825735904156556882819939007885332058136124031650490837864442687629129015664037894272559782), (4002409555221667393417789825735904156556882819939007885332058136124031650490837864442687629129015664037894272559782, 4002409555221667393417789825735904156556882819939007885332058136124031650490837864442687629129015664037894272559780), (4002409555221667393417789825735904156556882819939007885332058136124031650490837864442687629129015664037894272559783, 4002409555221667393417789825735904156556882819939007885332058136124031650490837864442687629129015664037894272559781)), ((4002409555221667393417789825735904156556882819939007885332058136124031650490837864442687629129015664037894272559781, 4002409555221667393417789825735904156556882819939007885332058136124031650490837864442687629129015664037894272559779), (4002409555221667393417789825735904156556882819939007885332058136124031650490837864442687629129015664037894272559784, 4002409555221667393417789825735904156556882819939007885332058136124031650490837864442687629129015664037894272559782), (4002409555221667393417789825735904156556882819939007885332058136124031650490837864442687629129015664037894272559784, 4002409555221667393417789825735904156556882819939007885332058136124031650490837864442687629129015664037894272559782)))
    end
end

@testset "sub: FQ12" begin
    let - = BLS12381Field.:-
        @test A12 - B12 == (((1, 1), (1, 1), (4002409555221667393417789825735904156556882819939007885332058136124031650490837864442687629129015664037894272559785, 4002409555221667393417789825735904156556882819939007885332058136124031650490837864442687629129015664037894272559785)), ((0, 0), (1, 1), (4002409555221667393417789825735904156556882819939007885332058136124031650490837864442687629129015664037894272559786, 4002409555221667393417789825735904156556882819939007885332058136124031650490837864442687629129015664037894272559786))) == (((1, 1), (1, 1), (4002409555221667393417789825735904156556882819939007885332058136124031650490837864442687629129015664037894272559785, 4002409555221667393417789825735904156556882819939007885332058136124031650490837864442687629129015664037894272559785)), ((0, 0), (1, 1), (4002409555221667393417789825735904156556882819939007885332058136124031650490837864442687629129015664037894272559786, 4002409555221667393417789825735904156556882819939007885332058136124031650490837864442687629129015664037894272559786)))
    end
end

@testset "mul: FQ12" begin
    let * = BLS12381Field.:*
        @test A12 * B12 == (((4002409555221667393417789825735904156556882819939007885332058136124031650490837864442687629129015664037894272559692, 46), (4002409555221667393417789825735904156556882819939007885332058136124031650490837864442687629129015664037894272559731, 59), (4002409555221667393417789825735904156556882819939007885332058136124031650490837864442687629129015664037894272559750, 68)), ((4002409555221667393417789825735904156556882819939007885332058136124031650490837864442687629129015664037894272559712, 53), (4002409555221667393417789825735904156556882819939007885332058136124031650490837864442687629129015664037894272559740, 63), (4002409555221667393417789825735904156556882819939007885332058136124031650490837864442687629129015664037894272559757, 72)))
    end
end

@testset "pow: FQ12" begin
    let ^ = BLS12381Field.:^,  * = BLS12381Field.:*
        @test A12 ^ 3 == (A12 * A12) * A12
    end
end

end


using Plots

include(string(pwd(), "/src/interpreter.jl"))
include(string(pwd(), "/src/gensys.jl"))
include(string(pwd(), "/galis_example.jl"))

@polyvar pi pi_e y y_e i v v_l ep

#variables = [pi pi_e y y_e v v_l ep]

#Equations of the model

eq1 = pi - bet*pi_e - kappa*y
eq2 = y + 1/sig*(i - pi_e) - y_e
eq3 = i - phi_pi*pi - phi_y*y - v_l
eq4 = v - rho_v*v_l - ep

#Dictionary with all the equations

model = Dict("IS"=>eq2,"Phillips"=>eq1,"Taylor"=>eq3,"Ar Shock"=>eq4)

modelgensys = model2gensys(model,[ep])

sol = gensys(modelgensys)

fris = irf(sol,15,0.25)

plot(4*fris[:,1],w = 3)
plot!(4*irfs_true[2:15,1], line = :dash,w = 3)

plot(fris[:,2],w = 3)
plot!(irfs_true[2:15,2], line = :dash,w = 3)

plot(4*fris[:,3],w = 3)
plot!(4*irfs_true[2:15,4], line = :dash,w = 3)

plot(fris[:,3],w = 3)
plot!(irfs_true[2:15,4], line = :dash,w = 3)


include("./params.jl")

function pprocess_gamma(lambda,k)
    # Example usage:
    # lambda = 20.0; # Spike-rate (Hz).
    # k = 2.3;       # Shape parameter.
    # x, rs = pprocess_gamma(lambda, k)
    #
    # Input argument
    # y:    Spike-rate.
    # k:    The shape parameter of the ISI distribution.
    #       * To use small k << 1, use small sampling time dt.
    #       * Maximum k allowed is ~10.
    #
    # Output argument
    # x:    Spike time.
    # rs:   Intensity function.

    #Initialize parameters

    x = [0.0]; # push!(x, num)
    rs = zeros(L);

    P = 0; R = 0;

    rs[1] = lambda;

    eta = - log(rand());
    for i = 2:L
        tau = i*dt - x[end]

        p = 1.0 / gamma(k) * (k * lambda).^k * tau.^(k-1) .* exp(-k*lambda*tau); # ISI
        P = P + p*dt;

        r = p / (1.0 - P); # Intensity
        R = R + r*dt;

        if eta <= R
            push!(x,i*dt) # Spike time

            P = 0.0; R = 0.0;
            eta = - log(rand());
        end
        rs[i] = r;
    end

    return x, rs
end


mutable struct TOProblemPar
    ndim::Int          # dimension: 2 or 3
    n::Int             # num of nodes in ground structure
    m::Int             # num of elements in ground structure
    X::Matrix{Float64} # n x ndim matrix, node coords
    T::Matrix{Int}     # m x 2 matrix, element connectivity
    B::Matrix{Float64} # m x (ndim x 2) matrix, compatability matrix, i.e. gamma_i, the direction cosine, i ∈ [m]

    # TODO: multiple load cases, Array{Matrix{Float64}}
    F::Matrix{Float64} # n x ndim matrix, load

    # TODO: multiple support cases, Array{Matrix{Int}}
    Bound::Matrix{Int} # n x ndim matrix, boundary, = 1 fixed, = 0 free

    dof_map::Matrix{Int} # m x (2xndim) matrix, element-dof (e_id, 1..2xndim)
end

mutable struct TOProblem
    m::Int             # num of nodes in ground structure
    c::Vector{Float64} # m x 1 vector, coeff of liear objective function

    τ::Float64 # stability load factor, .>= 1 to ensure global stability
    E::Float64 # Young's  modulus
    σ_c::Float64 # compression stress limit
    σ_t::Float64 # tension stress limit

    # TODO: multi-load cases
    # τ::Vector{Float64}
    l::Vector{Float64} # m x 1 vector, element length
    # TODO: stored as a sparse vector
    f::Vector{Float64} # nfree x 1 vector, external force mapped in free dof indices

    Γ::SparseMatrixCSC{Float64, Int} # m x nfree sp compatability matrix

    Δ::Array{SparseMatrixCSC{Float64, Int}, 1} # m x sp(nfree x nfree) array, Δ_i for geometric stiffness matrix
    eK::Array{SparseMatrixCSC{Float64, Int}, 1} # m x sp(nfree x nfree), K_i for element stiffness matrix in full dof index form

    x_init::Vector{Float64} # init x value
end


# Statistical model for analyzing Radial Velocity observations
# Assumes existance of data in arrays times, obs, sigma_obs
# Assumes existance of functions: TODO

export loglikelihood, logprior_noise, is_valid_noise

num_noise_param() = 1

if(!isdefined(:PDMats))        using PDMats       end

function is_valid_noise(p::Vector)
  param = extract_noise_param(p)
  @assert length(param) == num_noise_param()
  sigma_cor = param[1]
  tau_evolve   = 50.0 # param[2]
  lambda_p     = 0.5  # param[3]
  tau_rotate   = 20.0 # param[4]
  const min_jitter = 0.0
  const max_jitter = 99.0
  #if !(min_jitter<=sigma_sq_cor<=max_jitter) || !(0<tau_evolve) || !(0<lambda_p) || !(0<tau_rotate)
  if !(min_jitter<=sigma_cor<=max_jitter)
     return false
  end
  return true
end

function kernel_quaesiperiodic(dt::Float64,param::Vector)
  sigma_cor = param[1]
  tau_evolve   = 50.0 # param[2]
  lambda_p     = 0.5  # param[3]
  tau_rotate   = 20.0 # param[4]
  sigma_cor*sigma_cor*exp(-0.5*((dt/tau_evolve)^2 + (sin(pi*dt/tau_rotate)/lambda_p)^2) )
end

function loglikelihood(p::Vector)
  num_pl = num_planets(p)
  if num_pl <= 1 return 0.0 end
  if !is_valid(p) return -Inf end  # prempt model evaluation
  # Set t, o, and so to point to global arrays with observational data, while enforcing types
  t::Array{Float64,1} = times
  o::Array{Float64,1} = obs
  so::Array{Float64,1} = sigma_obs
  @assert length(times) == length(obs) == length(sigma_obs)
  gp_param = extract_noise_param(p)
  model_minus_data = map(i->calc_model_rv(p,t[i])-o[i],1:length(t))
  Sigma = PDMat( [ kernel_quaesiperiodic(t[i]-t[j],gp_param) + (i==j ? so[i]^2 : 0.0) for i in 1:length(t), j in 1:length(t)] )

  # chisq = model_minus_data'*(Sigma*model_minus_data)
  chisq = invquad(Sigma, model_minus_data)
  log_likelihood = -0.5*( chi_sq + logdet(Sigma) + length(t)*log(2pi) )
end

function logprior_noise(p::Vector)
  logp = zero(eltype(p))
  if !is_valid_noise(p)
     logp -= Inf
     return logp
  end
  param = extract_noise_param(p)
  @assert length(param) == num_noise_param()
  sigma_cor = sqrt(param[1])
  # tau_evolve   = param[2]
  # lambda_p     = param[3]
  # tau_rotate   = param[4]
  const Jitter0 = 1.0
  const min_jitter = 0.0
  const max_jitter = 1000.0
  const jitter_prior_norm = -log(Jitter0*log1p(max_jitter/Jitter0))
  logp += -log1p(sigma_cor/Jitter0) + jitter_prior_norm
  # TODO: Add prior for other GP parameters if they become variable
  return logp::eltype(p)
end