module UncValue

import Base: +, -, *, /, \, ÷, ^, ==, !=, <, >, <=, >=
import Printf: @printf

struct Value{T <: Real}
    x::T
    σ::T

    Value{T}(x, σ) where {T <: Real} = σ < 0 ? error("Negative uncertainty has no meaning.") : new{T}(x, σ)
end
function Value(x::T, σ::U) where {T <: Real, U <: Real}
    xT, σT = promote(x, σ)
    @warn "Value was initialize with different types, promoting $(typeof(x)) & $(typeof(σ)) to $(typeof(xT))."
    Value{typeof(xT)}(xT, σT)
end
Value(x::T, σ::T) where {T <: Real} = Value{T}(x, σ)
Value(x::T) where {T <: Real} = Value{T}(x, zero(x))

export Value, val, unc, set_unc

## MATHEMATICAL OPERATORS

+(a::Value, b::T) where {T <: Real} = Value(a.x + b, a.σ)
+(a::T, b::Value) where {T <: Real} = Value(a + b.x, b.σ)
+(a::Value, b::Value) = Value(a.x + b.x, hypot(a.σ, b.σ))

-(a::Value) = Value(-a.x, a.σ)
-(a::Value, b::T) where {T <: Real} = Value(a.x - b, a.σ)
-(a::T, b::Value) where {T <: Real} = Value(a - b.x, b.σ)
-(a::Value, b::Value) = Value(a.x - b.x, hypot(a.σ, b.σ))

*(a::Value, b::T) where {T <: Real} = Value(a.x * b, a.σ * abs(b))
*(a::T, b::Value) where {T <: Real} = Value(a * b.x, b.σ * abs(a))
*(a::Value, b::Value) = Value(a.x * b.x, hypot(a.σ * b.x, a.x * b.σ))

/(a::Value, b::T) where {T <: Real} = Value(a.x / b, a.σ / abs(b))
/(a::T, b::Value) where {T <: Real} = Value(a / b.x, abs(a) * b.σ / b.x^2)
/(a::Value, b::Value) = Value(a.x / b.x, hypot(a.σ / b.x, b.σ * a.x / b.x^2))

\(a::Value, b::T) where {T <: Real} = Value(a.x \ b, a.x^2 \ b * a.σ)
\(a::T, b::Value) where {T <: Real} = Value(a \ b.x, a \ b.σ)
\(a::Value, b::Value) = Value(a.x \ b.x, hypot( a.x^2 \ b.x * a.σ, a.x \ b.σ))

÷(a::Value, b::T) where {T <: Real} = Value(a.x ÷ b, a.σ ÷ b)
÷(a::T, b::Value) where {T <: Real} = Value(a ÷ b.x, b.σ * a ÷ b.x^2)
÷(a::Value, b::Value) = Value(a.x ÷ b.x, round(typeof(a.x), hypot(a.σ ÷ b.x, b.σ * a.x ÷ b.x^2)))

^(a::Value, b::T) where {T <: Real} = Value(a.x ^ b, abs(b * a.x^(b-1)) * a.σ)
^(a::T, b::Value) where {T <: Real} = Value(a ^ b.x, abs(a^b.x * log(abs(a))) * b.σ)
^(a::Value, b::Value) where {T <: Real} =
    Value(a.x ^ b.x, hypot(b.x * a.x^(b.x-1) * a.σ, a.x^b.x * log(abs(a.x)) * b.σ))

==(a::Value, b::T) where {T <: Real} = a.x == b
==(a::T, b::Value) where {T <: Real} = a == b.x
==(a::Value, b::Value) = a.x == b.x

!=(a::Value, b::T) where {T <: Real} = a.x != b
!=(a::T, b::Value) where {T <: Real} = a != b.x
!=(a::Value, b::Value) = a.x != b.x

>(a::Value, b::T) where {T <: Real} = a.x > b
>(a::T, b::Value) where {T <: Real} = a > b.x
>(a::Value, b::Value) = a.x > b.x

<(a::Value, b::T) where {T <: Real} = a.x < b
<(a::T, b::Value) where {T <: Real} = a < b.x
<(a::Value, b::Value) = a.x < b.x

>=(a::Value, b::T) where {T <: Real} = a.x >= b
>=(a::T, b::Value) where {T <: Real} = a >= b.x
>=(a::Value, b::Value) = a.x >= b.x

<=(a::Value, b::T) where {T <: Real} = a.x <= b
<=(a::T, b::Value) where {T <: Real} = a <= b.x
<=(a::Value, b::Value) = a.x <= b.x

## MATHEMATICAL FUNCTIONS
Base.abs(a::Value) = Value(abs(a.x), a.σ)
Base.abs2(a::Value) = Value(abs2(a.x), 2 * abs(a.x) * a.σ)

Base.sin(a::Value) = Value(sin(a.x), abs(cos(a.x)) * a.σ)
Base.sind(a::Value) = Value(sind(a.x), abs(cosd(a.x)) * a.σ)
Base.sinh(a::Value) = Value(sinh(a.x), cosh(a.x) * a.σ)
Base.Math.sinpi(a::Value) = Value(sinpi(a.x), abs(π * cospi(a.x)) * a.σ)

Base.cos(a::Value) = Value(cos(a.x), abs(sin(a.x)) * a.σ)
Base.cosd(a::Value) = Value(cosd(a.x), abs(sind(a.x)) * a.σ)
Base.cosh(a::Value) = Value(cosh(a.x), abs(sinh(a.x)) * a.σ)
Base.Math.cospi(a::Value) = Value(cospi(a.x), abs(π * sinpi(a.x)) * a.σ)

Base.tan(a::Value) = Value(tan(a.x),  a.σ * sec(a.x)^2)
Base.tand(a::Value) = Value(tand(a.x),  a.σ * secd(a.x)^2)
Base.tanh(a::Value) = Value(tanh(a.x),  a.σ * sech(a.x)^2)

Base.asin(a::Value) = Value(asin(a.x),  a.σ / sqrt(1-a.x^2))
Base.Math.asind(a::Value) = Value(asind(a.x),  a.σ / sqrt(1-a.x^2))
Base.asinh(a::Value) = Value(asinh(a.x),  a.σ / sqrt(1 + a.x^2))

Base.acos(a::Value) = Value(acos(a.x),  a.σ / sqrt(1-a.x^2))
Base.Math.acosd(a::Value) = Value(acosd(a.x),  a.σ / sqrt(1-a.x^2))
Base.acosh(a::Value) = Value(acosh(a.x), a.σ / sqrt(a.x^2-1))

Base.atan(a::Value) = Value(atan(a.x), a.σ / (1 + a.x^2))
Base.atan(a::Value, b::T) where {T <: Real} = Value(atan(a.x, b), abs(b) * a.σ / (a.x^2 + b^2))
Base.atan(a::T, b::Value) where {T <: Real} = Value(atan(a, b.x), abs(a) * b.σ / (a^2 + b.x^2))
Base.atan(a::Value, b::Value) =
    Value(atan(a.x, b.x), hypot(b.x * a.σ, a.x * b.σ) / (a.x^2 + b.x^2))
Base.Math.atand(a::Value) = Value(atand(a.x),  a.σ / (1 + a.x^2))
Base.Math.atand(a::Value, b::T) where {T <: Real} = Value(atand(a.x, b), abs(b) * a.σ / (a.x^2 + b^2))
Base.Math.atand(a::T, b::Value) where {T <: Real} = Value(atand(a, b.x), abs(a) * b.σ / (a^2 + b.x^2))
Base.Math.atand(a::Value, b::Value) =
    Value(atand(a.x, b.x), hypot(b.x * a.σ, a.x * b.σ) / (a.x^2 + b.x^2))
Base.atanh(a::Value) = Value(atanh(a.x),  a.σ / (1 - a.x^2))

Base.Math.sec(a::Value) = Value(sec(a.x), abs(sec(a.x) * tan(a.x)) * a.σ)
Base.Math.secd(a::Value) = Value(secd(a.x), abs(secd(a.x) * tand(a.x)) * a.σ)
Base.Math.sech(a::Value) = Value(sech(a.x), abs(tanh(a.x) * sech(a.x)) * a.σ)

Base.Math.csc(a::Value) = Value(csc(a.x), abs(cot(a.x) * csc(a.x)) * a.σ)
Base.Math.cscd(a::Value) = Value(cscd(a.x), abs(cotd(a.x) * cscd(a.x)) * a.σ)
Base.Math.csch(a::Value) = Value(csch(a.x), abs(coth(a.x) * csch(a.x)) * a.σ)

Base.Math.cot(a::Value) = Value(cot(a.x), csc(a.x)^2 * a.σ)
Base.Math.cotd(a::Value) = Value(cotd(a.x), cscd(a.x)^2 * a.σ)
Base.Math.coth(a::Value) = Value(coth(a.x), csch(a.x)^2 * a.σ)

Base.Math.asec(a::Value) = Value(asec(a.x), a.σ / sqrt(a.x^2 - 1))
Base.Math.asecd(a::Value) = Value(asecd(a.x), a.σ / sqrt(a.x^2 - 1))
Base.Math.asech(a::Value) = Value(asech(a.x), a.σ / sqrt(1 - a.x^2))

Base.Math.acsc(a::Value) = Value(acsc(a.x), a.σ / sqrt(a.x^2 - 1))
Base.Math.acscd(a::Value) = Value(acscd(a.x), a.σ / sqrt(a.x^2 - 1))
Base.Math.acsch(a::Value) = Value(acsch(a.x), a.σ / sqrt(a.x^2 + 1))

Base.Math.acot(a::Value) = Value(acot(a.x), a.σ / (a.x^2 + 1))
Base.Math.acotd(a::Value) = Value(acotd(a.x), a.σ / (a.x^2 + 1))
Base.Math.acoth(a::Value) = Value(acoth(a.x), a.σ / abs(1 - a.x^2))

Base.Math.sinc(a::Value) = Value(sinc(a.x), a.σ * abs(cosc(a.x)))
Base.Math.cosc(a::Value) = Value(cosc(a.x), a.σ * abs((2 * a.x * cos(a.x) + (a.x^2 - 2) * sin(a.x)) / a.x^3))

Base.Math.deg2rad(a::Value) = Value(deg2rad(a.x), deg2rad(a.σ))
Base.Math.rad2deg(a::Value) = Value(rad2deg(a.x), rad2deg(a.σ))

Base.exp(a::Value) = Value(exp(a.x), exp(a.x) * a.σ)
Base.exp2(a::Value) = Value(exp2(a.x), exp2(a.x) * a.σ * log(2))
Base.exp10(a::Value) = Value(exp10(a.x), exp10(a.x) * a.σ * log(10))

Base.log(a::Value) = Value(log(a.x), a.σ / abs(a.x))
Base.log(b::T, a::Value) where {T <: Real} = Value(log(b, a.x), a.σ / abs(a.x * log(b)))
Base.log2(a::Value) = Value(log2(a.x), a.σ / abs(a.x * log(2)))
Base.log10(a::Value) = Value(log10(a.x), a.σ / abs(a.x * log(10)))
Base.log1p(a::Value) = Value(log1p(a.x), a.σ / abs(1 + a.x))

Base.sqrt(a::Value) = Value(sqrt(a.x), a.σ / (2 * sqrt(a.x)))
Base.hypot(a::Value, b::T) where {T <: Real} = Value(hypot(a.x, b), abs(a.x) * a.σ / hypot(a.x, b))
Base.hypot(a::T, b::Value) where {T <: Real} = Value(hypot(a, b.x), abs(b.x) * b.σ / hypot(a, b.x))
function Base.hypot(a::Value, b::Value)
    h = hypot(a.x, b.x)
    Value(h, hypot(a.x * a.σ, b.x * b.σ) / h)
end

Base.cbrt(a::Value) = Value(cbrt(a.x), a.σ / (3 * cbrt(a.x)^2))

## OTHER FUNCTIONS

Base.sign(a::Value) = sign(a.x)
Base.signbit(a::Value) = signbit(a.x)

Base.inv(a::Value) = Value(inv(a.x), a.σ * inv(a.x)^2)

Base.fld(a::Value, b::T) where {T <: Real} = Value(fld(a.x, b), fld(a.σ, b))
Base.fld(a::T, b::Value) where {T <: Real} = Value(fld(a, b.x), fld(b.σ * a, b.x^2))
Base.fld(a::Value, b::Value) = Value(fld(a.x, b.x), floor(typeof(a.x), hypot(fld(a.σ, b.x), fld(b.σ * a.x, b.x^2))))

Base.cld(a::Value, b::T) where {T <: Real} = Value(cld(a.x, b), cld(a.σ, b))
Base.cld(a::T, b::Value) where {T <: Real} = Value(cld(a, b.x), cld(b.σ * a, b.x^2))
Base.cld(a::Value, b::Value) = Value(cld(a.x, b.x), ceil(typeof(a.x), hypot(cld(a.σ, b.x), cld(b.σ * a.x, b.x^2))))

Base.cmp(a::Value, b::T) where {T <: Real} = cmp(a.x, b)
Base.cmp(a::T, b::Value) where {T <: Real} = cmp(a, b.x)
Base.cmp(a::Value, b::Value) = cmp(a.x, b.x)

Base.isless(a::Value, b::T) where {T <: Real} = cmp(a.x, b) < 0
Base.isless(a::T, b::Value) where {T <: Real} = cmp(a, b.x) < 0
Base.isless(a::Value, b::Value) = cmp(a.x, b.x) < 0

function Base.clamp(x::T, lh::Value) where {T <: Real}
    if x > lh.x + lh.σ
        return lh.x + lh.σ
    elseif x < lh.x - lh.σ
        return lh.x - lh.σ
    end

    x
end

function Base.clamp(rh::Value, lh::Value) where {T <: Real}
    if rh.x > lh.x + lh.σ
        return Value(lh.x + lh.σ, rh.σ)
    elseif rh.x < lh.x - lh.σ
        return Value(lh.x - lh.σ, rh.σ)
    end

    rh
end

#check how to return also the uncertainty!
Base.max(a::Value) = a.x + a.σ
Base.max(a::Value, b::T) where {T <: Real} = cmp(a.x, b) <= 0 ? b : a
Base.max(a::T, b::Value) where {T <: Real} = cmp(a, b.x) <= 0 ? b : a
Base.max(a::Value, b::Value)= cmp(a.x, b.x) <= 0 ? b : a

Base.min(a::Value) = a.x - a.σ
Base.min(a::Value, b::T) where {T <: Real} = cmp(a.x, b) <= 0 ? a : b
Base.min(a::T, b::Value) where {T <: Real} = cmp(a, b.x) <= 0 ? a : b
Base.min(a::Value, b::Value)= cmp(a.x, b.x) <= 0 ? a : b

Base.isapprox(a::Value, b::T; significance::Int = 1) where {T <: Real} =
    a.x - significance * a.σ <= b <= a.x + significance * a.σ
Base.isapprox(a::T, b::Value; significance::Int = 1) where {T <: Real} =
    b.x - significance * b.σ <= a <= b.x + significance * b.σ
Base.isapprox(a::Value, b::Value; significance::Int = 1) =
    (b.x - significance * b.σ <= a.x <= b.x + significance * b.σ) |
    (a.x - significance * a.σ <= b.x <= a.x + significance * a.σ)

Base.length(a::Value) = length(a.x)

Base.iterate(a::Value) = (a, nothing)
Base.iterate(a::Value, b::Nothing) = (a, 2)

## PRINTED RESULTS# https://docs.julialang.org/en/v1/base/math/
# precision(a::T) where {T <: Real} = floor(Int, log10(a)) - 1
Base.precision(a::Value) = floor(Int, log10(a.σ)) - 1

function Base.show(io::IO, a::Value)
    p = precision(a)
    fix_x = round(a.x * 10.0^(-p); digits = 2, base = 10) / 10
    fix_ux = round(a.σ * 10.0^(-p); digits = 2, base = 10) / 10
    p += 1
    if p == 0
        @printf(io, "%.1f ± %.1f\n", fix_x, fix_ux)
    else
        @printf(io, "(%.1f ± %.1f)·10^%d\n", fix_x, fix_ux, p)
    end
end

## NEW FUNCTIONS
val(item::Number) = item
val(item::Value) = item.x
function val(A::Array)
    type = typeof(A[1])
    if type <: Value
        type = typeof(A[1].x)
    end

    B = Array{type, ndims(A)}(undef, size(A))
    for index in eachindex(A)
        B[index] = val(A[index])
    end

    B
end

unc(item::Number) = zero(item)
unc(item::Value) = item.σ
function unc(A::Array)
    type = typeof(A[1])
    if type <: Value
        type = typeof(A[1].x)
    end

    B = Array{type, ndims(A)}(undef, size(A))
    for index in eachindex(A)
        B[index] = unc(A[index])
    end

    B
end

function set_unc(A::Array{T, N}, unc::Real) where {T, N}
    type = typeof(A[1])
    if type <: Value
        type = typeof(A[1].x)
    end

    B = Array{Value{type}, ndims(A)}(undef, size(A))
    for index in eachindex(A)
        B[index] = Value(val(A[index]), unc)
    end

    B
end

function set_unc(A::Array{T, N}, unc::Array{T2, N}) where {T, T2 <: Real, N}
    type = typeof(A[1])
    if type <: Value
        type = typeof(A[1].x)
    end

    B = Array{Value{type}, ndims(A)}(undef, size(A))
    for index in eachindex(A)
        B[index] = Value(val(A[index]), unc[index])
    end

    B
end

end


function hessian_replicates(dfns, d2fns, parms, pts, datav)
  hess = zeros(3,3)
  _S   = Symmetric([matern(x, y, parms) for x in pts, y in pts])
  S    = cholesky(_S)
  for j in 1:3
    dfj = dfns[j]
    Sj  = Symmetric([dfj(x, y, parms) for x in pts, y in pts])
    for k in j:3
      # first derivative matrix:
      dfk   = dfns[k]
      Sk    = Symmetric([dfk(x, y, parms)  for x in pts, y in pts])
      # second derivative matrix:
      dfjk  = d2fns[j][k-j+1]
      Sjk   = Symmetric([dfjk(x, y, parms) for x in pts, y in pts])
      # Compute the complicated derivative of the qform (not efficient or
      # thoughtful, so don't use this code for real somewhere):
      dqf   = -(S\(Sk*(S\(Sj*inv(_S)))) )
      dqf  += S\(Sjk*inv(_S))
      dqf  -= S\(Sj*(S\(Sk*inv(_S))))
      # compute the term:
      term  = (tr(S\Sjk) - tr(S\(Sj*(S\Sk))))*length(datav)/2
      term -= sum(z->dot(z, dqf, z), datav)/2
      hess[j,k] = term
      hess[k,j] = term
    end
  end
  hess
end


using AnalyticalEOR, Roots

kr_ww = RelPerms(swr=0.2,
                sor=0.2,
                krw0=0.5,
                kro0=1.0,
                nw=3.0,
                no=2.0)

kr_ow = RelPerms(swr=0.2,
                sor=0.2,
                krw0=0.245,
                kro0=0.54,
                nw=2.111,
                no=3.706)

μw = 1.0
μo = 3.0

si = 0.55
sj = 0.80

D = 0.2

pf = solve_polymerflood(si, sj, kr_ow, kr_ww, D, μw, μo)
tracer = solve_tracer(pf)

plot_sw_profile(pf, tracer, 0.2)

animate_sw_profile(pf, tracer)


"A measure associated with a parameterization "
struct ParameterizationMeasure{P,T} <: DomainIntegrals.Weight{T}
    param   ::  P
end

ParameterizationMeasure(param) = ParameterizationMeasure{typeof(param),eltype(domain(param))}(param)

DomainIntegrals.unsafe_weightfun(p::ParameterizationMeasure, t) = gradientnorm(p.param, t)
DomainIntegrals.support(p::ParameterizationMeasure) = domain(p.param)

export BoundaryIntegralOperator
"A boundary integral operator with a parameterized boundary."
struct BoundaryIntegralOperator{K,P,T} <: IntegralOperator
    kernel      ::  K
    boundary    ::  Domain{T}
    param       ::  P
end

BasisFunctions.name(op::BoundaryIntegralOperator) = "Boundary integral operator"

kernel(op::BoundaryIntegralOperator) = op.kernel
boundary(op::BoundaryIntegralOperator) = op.boundary
parameterization(op::BoundaryIntegralOperator) = op.param
parameterdomain(op::BoundaryIntegralOperator) = domain(parameterization(op))

eval_kernel(op::BoundaryIntegralOperator, x, y) =
    kernel(op)(x, y, parameterization(op))

measure(op::BoundaryIntegralOperator) = ParameterizationMeasure(parameterization(op))

singularity(op::BoundaryIntegralOperator) = singularity(kernel(op))

export op_apply
"Apply the integral operator to the given density function and evaluate at the point t."
function op_apply(op::BoundaryIntegralOperator, density, t)
    μ = measure(op)
    integrand = y -> eval_kernel(op, t, y) * density(y) * unsafe_weightfun(μ, y)
    DomainIntegrals.integral(QuadAdaptive(), integrand, parameterdomain(op), LogSingPoint(t))
end


# output = vcat(input1, input2, ...)
immutable Concat{T} <: Transformation
    n::Int
    input::CatNode{T,1}
    output::OutputNode{T,1}
end
Concat(n::Int) = Concat(Float64, n)

# sort of tricky... we do the concatenating in the input node, so instead
# of copying val/∇ through, just make them point to the same arrays
function Concat{T}(::Type{T}, n::Int)
    input = InputNode(:cat, T, n)
    output = OutputNode(T, n)
    output.val = input.val
    output.∇ = input.∇
    Concat(n, input, output)
end

transform!(t::Concat) = t.output.val
grad!(t::Concat) = t.input.∇

# ---------------------------------------------------------

# output = input1 .* input2 .* ...
immutable Gate{T} <: Transformation
    n::Int
    input::ProdNode{T,1}
    output::OutputNode{T,1}
end
Gate(n::Int) = Gate(Float64, n)

# sort of tricky... we do the gating in the input node, so instead
# of copying val/∇ through, just make them point to the same arrays
function Gate{T}(::Type{T}, n::Int)
    input = InputNode(:*, T, n)
    output = OutputNode(T, n)
    output.val = input.val
    output.∇ = input.∇
    Gate(n, input, output)
end

transform!(t::Gate) = t.output.val
grad!(t::Gate) = t.input.∇


#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/build_literate/tutorial.ipynb)

# # Quick Start Tutorial

# The core of generalized Bloch model is implemented in the function [`apply_hamiltonian_gbloch!(∂m∂t, m, mfun, p, t)`](@ref), which calculates the derivative `∂m/∂t` for a given magnetization vector `m` and stores it in-place in the the variable `∂m∂t`. The function interface is written in a way that we can directly feed it into a differential equation solver of the [DifferentialEquations.jl](https://diffeq.sciml.ai/stable/) package.

# For this example, we need the following packages:
using MRIgeneralizedBloch
using DifferentialEquations
using Plots
plotlyjs(bg = RGBA(31/255,36/255,36/255,1.0), ticks=:native); #hide #!nb

# We simulate the dynamics of a coupled spin system with the following parameters:
m0s = 0.15
R1f = 0.3 # 1/s
R2f = 15 # 1/s
R1s = 2 # 1/s
T2s = 10e-6 # s
Rx = 30; # 1/s

# and the thermal equilibrium of the magnetization `m = [xf; yf; zf; zs; 1]`:
m0 = [0; 0; 1-m0s; m0s; 1];

# during a rectangular RF-pulse with the flip angle and pulse duration
α = π
TRF = 100e-6; # s

# Further, we assume a perfectly calibrated, on-resonant RF-pulse:
B1 = 1
ω0 = 0; # rad/s

# as well as a [super-Lorentzian lineshape](http://dx.doi.org/10.1002/mrm.1910330404). We interpolate the corresponding Green's function to improve performance:
G = interpolate_greens_function(greens_superlorentzian, 0, TRF / T2s);

# The generalized Bloch model is a so-called integro-differential equation where the derivative ``∂m/∂t`` at the time ``t_1`` does not just depend on ``m(t_1)``, but on ``m(t)`` for ``t \in [0, t_1]``. This is solved with a [delay differential equation (DDE) solver](https://diffeq.sciml.ai/stable/tutorials/dde_example/) that stores an interpolated *history function* `mfun(p, t)`, which we use in the [`apply_hamiltonian_gbloch!`](@ref) function to evaluate the integral. This history function has to be initialized with `mfun(p, 0) = m0`. Here, we use a slightly more complicated initialization that allows us to index the history function in [`apply_hamiltonian_gbloch!`](@ref), which improves performance:
mfun(p, t; idxs=nothing) = typeof(idxs) <: Number ? m0[idxs] : m0;

# With this, we are ready to formulate and solve the differential equation:
param = (α/TRF, B1, ω0, m0s, R1f, R2f, Rx, R1s, T2s, G) # defined by apply_hamiltonian_gbloch!
prob = DDEProblem(apply_hamiltonian_gbloch!, m0, mfun, (0, TRF), param)
sol = solve(prob)

# The plot function is implemented for such solution objects and we can plot the solution simply with
p = plot(sol, labels=["xᶠ" "yᶠ" "zᶠ" "zˢ" "1"], xlabel="t [s]", ylabel="m(t)")
#md Main.HTMLPlot(p) #hide

# More details on the interface, including the linear approximation of the generalized Bloch model can found in the following scripts that replicate all simulations, data analyses, and figures of the generalized Bloch paper.


using JuMP
using SumOfSquares

export getlyap, soslyap, soslyapb, soslyapbs

# Storing the Lyapunov
function setlyap!(s, lyap::Lyapunov)
    d = lyap.d
    lyaps = getlyaps(s)
    if length(lyaps) < d
        sizehint!(lyaps, d)
        while length(lyaps) < d
            push!(lyaps, nothing)
        end
    end
    lyaps[d] = lyap
end
function getlyap(s::AbstractSwitchedSystem, d::Int; kws...)
    lyaps = getlyaps(s)
    if d > length(lyaps) || lyaps[d] === nothing
        soslyapb(s, d; cached=true, kws...)
    end
    return lyaps[d]
end

function getsoslyapinitub(s::AbstractDiscreteSwitchedSystem, d::Integer)
    #_, sosub = pradiusb(s, 2*d)
    #sosub
    return Inf
end
#function getsoslyapinitub(s::AbstractContinuousSwitchedSystem, d::Integer)
#    return Inf
#end

function getsoslyapinit(s, d)
    lyaps = getlyaps(s)
    if d <= length(lyaps) && lyaps[d] !== nothing
        lyap = lyaps[d]
        return lyap.soslb, lyap.dual, lyap.sosub, lyap.primal
    else
        # The SOS ub is greater than the JSR hence also greater than any of its lower bound.
        # Hence getlb(s) can be used as an initial lowerbound
        return getlb(s), nothing, getsoslyapinitub(s, d), nothing
    end
end

function superset(s::HybridSystems.AbstractHybridSystem, mode, d)
    q = GramMatrix(SOSDecomposition(variables(s, mode).^d))
    return SetProg.Sets.PolySet(2d, q)
end

# For values of γ far from 1.0, it is better to divide A_i's by γ,
# it results in a problem that is better conditioned.
# This is clearly visible in [Example 5.4, PJ08] for which the JSR is ≈ 8.9
# So we ask the user to give a scaled system, i.e. do `p(A/γx) ≥ p(x)` rather
# than do `p(Ax) ≥ γ p(x)`
"""
    soslyap(s::AbstractSwitchedSystem, d; optimizer_constructor=nothing)

Find Sum-of-Squares Lyapunov functions; i.e. solves [(5), PJ08]
or gives moment matrices certifying the infeasibility of the problem.
Use [`ScaledHybridSystem`](@ref) to use a different growth rate than 1.

[PJ08] P. Parrilo and A. Jadbabaie.
*Approximation of the joint spectral radius using sum of squares*.
Linear Algebra and its Applications, Elsevier, **2008**, 428, 2385-2402
"""
function soslyap(s::HybridSystems.AbstractHybridSystem, d; optimizer_constructor=nothing)
    sets = HybridSystems.state_property(s, PolynomialLyapunov{Float64})
    infeasibility_certificates = HybridSystems.transition_property(s, MeasureLyapunov{Float64})
    set_variables = PolySet[PolySet(symmetric=true, degree=2d, superset=superset(s, mode, d))
                            for mode in states(s)]
    status = invariant_sets!(sets, 1:nstates(s), s, optimizer_constructor, set_variables,
                             volume_heuristic = nothing,
                             infeasibility_certificates = infeasibility_certificates,
                             verbose=0)
    if isinfeasible(status, true)
        @assert !isfeasible(status, true)
        if status[3] == MOI.NO_SOLUTION
            return status, nothing, nothing
        else
            return status, nothing, infeasibility_certificates
        end
    elseif isfeasible(status, true)
        status, sets, nothing
    else
        @assert !isdecided(status, true)
        status, nothing, nothing
    end
end

function increaselb(s::AbstractDiscreteSwitchedSystem, lb, step)
    lb *= step
end

soschecktol(soslb, sosub) = sosub - soslb
soschecktol(s::AbstractDiscreteSwitchedSystem, soslb, sosub) = soschecktol(log(soslb), log(sosub))
tol_diff_str(s::AbstractDiscreteSwitchedSystem) = "Log-diff   "
#soschecktol(s::AbstractContinuousSwitchedSystem, soslb, sosub) = soschecktol(soslb, sosub)

sosshift(s::AbstractDiscreteSwitchedSystem, b, shift) = exp(log(b) + shift)
function sosshift(s::AbstractDiscreteSwitchedSystem, b, shift, scaling)
    return sosshift(s, b / scaling, shift) * scaling
end
#sosshift(s::AbstractContinuousSwitchedSystem, b, shift) = b + shift

function sosmid(soslb, sosub, step)
    if isfinite(soslb) && isfinite(sosub)
        mid = (soslb + sosub) / 2
    elseif isfinite(soslb)
        mid = soslb + step
    elseif isfinite(sosub)
        mid = sosub - step
    else
        mid = 0
    end
end
usestep(soslb, sosub) = isfinite(soslb) ⊻ isfinite(sosub)
sosmid(s::AbstractDiscreteSwitchedSystem, soslb, sosub, step) = exp(sosmid(log(soslb), log(sosub), step))
usestep(s::AbstractDiscreteSwitchedSystem, soslb, sosub) = usestep(log(soslb), log(sosub))
#sosmid(s::AbstractContinuousSwitchedSystem, soslb, sosub, step) = sosmid(soslb, sosub, step)
function sosmid(s::AbstractDiscreteSwitchedSystem, soslb, sosub, step, scaling)
    sosmid(s, soslb / scaling, sosub / scaling, step) * scaling
end

function soslb2lb(s::AbstractDiscreteSwitchedSystem, soslb, d)
    n = maximum(statedim.(s, states(s)))
    η = min(ρA(s), binomial(n+d-1, d))
    soslb / η^(1/(2*d))
end
#soslb2lb(s::AbstractContinuousSwitchedSystem, soslb, d) = -Inf

function showbs(s, soslb, sosub, tol, verbose, ok::Bool)
    if verbose >= 2 || (ok && verbose >= 1)
        println("Lower bound: $soslb")
        println("Upper bound: $sosub")
        println("$(tol_diff_str(s)): $(soschecktol(s, soslb, sosub)) $(ok ? '≤' : '>') $tol")
    end
end

function showmid(γ, status, verbose)
    if verbose >= 3
        println("  Trial value of γ: $γ")
        println("Termination status: $(status[1])")
        println("     Primal status: $(status[2])")
        println("       Dual status: $(status[3])")
        println("        Raw status: $(status[4])")
        if !isdecided(status, true)
            problem_status = "Unknown"
        elseif isfeasible(status, true)
            problem_status = "Feasible"
        else
            @assert isinfeasible(status, true)
            problem_status = "Infeasible"
        end
        println("    Problem status: $problem_status")
    end
end

# Binary Search

"""
    soslyapbs(s::AbstractSwitchedSystem, d::Integer,
              soslb, dual,
              sosub, primal;
              verbose=0, tol=1e-5, step=0.5, scaling=quickub(s),
              ranktols=tol, disttols=tol, kws...)

Find the smallest `γ` such that [`soslyap`](@ref) is feasible.
"""
function soslyapbs(s::AbstractSwitchedSystem, d::Integer,
                   soslb, dual,
                   sosub, primal;
                   verbose=0, tol=1e-5, step=0.5, scaling=quickub(s),
                   ranktols=tol, disttols=tol, kws...)
    _lyap(γ) = soslyap(ScaledHybridSystem(s, γ), d; kws...)
    while soschecktol(s, soslb, sosub) > tol
        showbs(s, soslb, sosub, tol, verbose, false)
        mid = sosmid(s, soslb, sosub, step, scaling)
        status, curprimal, curdual = _lyap(mid)
        showmid(mid, status, verbose)
        if !isdecided(status, true)
            if usestep(s, soslb, sosub)
                step *= 2
                continue
            end
            # If mid-tol/2 and mid+tol/2 also Stall, there would be an interval of length tol of Stall -> impossible to satisfy requirements
            # the distance between soslb and mid is at least tol/2.
            # Sometimes, mid is far from soslb and is at a point where the solver Stall even if it is far from the optimum point.
            # In that case, it is better to take (mid + soslb)/2
            midlb = min(sosmid(s, soslb, mid, step, scaling),
                        sosshift(s, mid, -tol/2, scaling))
            # If mid-tol/2 is too close to soslb, we would not make progress!
            # So we ensure we make a progress of at least tol/8. If dual is nothing, then that would still be progress to find a dual
            if dual !== nothing
                midlb = max(midlb, sosshift(s, soslb, tol/8))
            end
            statuslb, curprimallb, curduallb = _lyap(midlb)
            showmid(midlb, statuslb, verbose)
            if isdecided(statuslb, true)
                mid = midlb
                status = statuslb
                curprimal = curprimallb
                curdual = curduallb
            else
                midub = max(sosmid(s, mid, sosub, step, scaling),
                            sosshift(s, mid, tol/2, scaling))
                if primal !== nothing
                    midub = min(midub, sosshift(s, sosub, -tol/8))
                end
                statusub, curprimalub, curdualub = _lyap(midub)
                showmid(midub, statusub, verbose)
                if isdecided(statusub, true)
                    mid = midub
                    status = statusub
                    curprimal = curprimalub
                    curdual = curdualub
                end
            end
        end
        if isinfeasible(status, true)
            dual = curdual
            if dual !== nothing
                sosextractcycle(s, dual, d, ranktols=ranktols, disttols=disttols)
            end
            soslb = mid
        elseif isfeasible(status, true)
            if !(curprimal === nothing) # FIXME remove
                primal = curprimal
            end
            sosub = mid
        else
            @warn("Solver returned with status : $statuslb for γ=$midlb, $status for γ=$mid and $statusub for γ=$midub. Stopping bisection with $(soschecktol(s, soslb, sosub)) > $tol (= tol)")
            break
        end
    end
    if soschecktol(s, soslb, sosub) ≤ tol # it is not guaranteed because of the break
        showbs(s, soslb, sosub, tol, verbose, true)
    end
    soslb, dual, sosub, primal
end

"""
    soslyapbs(s::AbstractSwitchedSystem, d::Integer,
              soslb, dual,
              sosub, primal;
              verbose=0, tol=1e-5, step=0.5, scaling=quickub(s),
              ranktols=tol, disttols=tol, kws...)

Find upper bounds to the (constrained) Joint Spectral Radius [(5), PJ08].
Lower bounds a obtained using guarantees in [LPJ19].

[LPJ19] B. Legat, P. A. Parrilo and R. M. Jungers
*An entropy-based bound for the computational complexity of a switched system*.
IEEE Transactions on Automatic Control, **2019**.

[PJ08] P. Parrilo and A. Jadbabaie.
*Approximation of the joint spectral radius using sum of squares*.
Linear Algebra and its Applications, Elsevier, **2008**, 428, 2385-2402
"""
function soslyapb(s::AbstractSwitchedSystem, d::Integer; optimizer_constructor=nothing, tol=1e-5, cached=true, kws...)
    soslb, dual, sosub, primal = soslyapbs(s::AbstractSwitchedSystem, d::Integer, getsoslyapinit(s, d)...; optimizer_constructor=optimizer_constructor, tol=tol, kws...)
    _lyap(γ) = soslyap(ScaledHybridSystem(s, γ), d, optimizer_constructor=optimizer_constructor)
    if cached
        if primal === nothing
            if isfinite(sosub)
                status, primal, _ = _lyap(sosub)
                @assert isfeasible(status, true)
                @assert primal !== nothing
            else
                error("Bisection ended with infinite sosub=$sosub")
            end
        end
        if dual === nothing
            if isfinite(soslb)
                status, _, dual = _lyap(soslb)
                if !isinfeasible(status, true)
                    soslb = sosshift(s, soslb, -tol)
                    status, _, dual = _lyap(soslb)
                    if !isinfeasible(status, true)
                        @warn("We ignore getlb and start from scratch. tol was probably set too small and soslb is too close to the JSR so soslb-tol is too close to the JSR")
                        soslb = 0. # FIXME fix for continuous
                    end
                    soslb, dual, sosub, primal = soslyapbs(s::AbstractSwitchedSystem, d::Integer, soslb, dual, sosub, primal; optimizer_constructor=optimizer_constructor, tol=tol, kws...)
                    @assert dual !== nothing
                end
            else
                error("Bisection ended with infinite soslb=$soslb")
            end
        end
        setlyap!(s, Lyapunov(d, soslb, dual, sosub, primal))
    end
    ub = sosub
    lb = soslb2lb(s, soslb, d)
    updateb!(s, lb, ub)
end


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

"""
    nnpars(lpars, data, target)

Pass local anisotropies `lpars` from `data` to `target` object, via nearest
neighbors. Useful to pass parameters from domain to samples.
"""
nnpars(lpars, data, target) = slice(lpars, grid2hd_ids(target, data))

"""
    idwpars(localaniso, searcher, domain; power=2, metric=Euclidean())

Interpolate `LocalAnisotropy` data into a `domain`, using the
local neighbors returned from the `searcher`. The interpolation can be inverse
distance weighted by given `power` and `metric` or a simply averaged if
`power=0`. The ellipses/ellipsoids rotations are interpolated using
(weighted) average of quaternions as described by Markley et al (2007). The
interpolated magnitude is the weighted median of the neighbors magnitude.

## Example

```julia
searcher = KNearestSearch(data, 10)
idw3_into_grid = idwpars(localaniso, searcher, grid, power=3)
```

## Reference

Markley, F.L., et al. (2007). [Averaging quaternions](https://ntrs.nasa.gov/archive/nasa/casi.ntrs.nasa.gov/20070017872.pdf)
"""
idwpars(lpars, searcher::NeighborSearchMethod, domain; power=2.0, metric=Euclidean()) =
	interpolate(lpars, searcher, domain, power=power)

"""
    smoothpars(localaniso, searcher; power=0, metric=Euclidean())

Smooth `LocalAnisotropy`, using the local neighbors returned from the `searcher`.
The interpolation can be inverse distance weighted by given `power` and `metric`
or a simply averaged if `power=0`. The ellipses/ellipsoids rotations are
interpolated using (weighted) average of quaternions as described by Markley
et al. (2007). The interpolated magnitude is the weighted median of the
neighbors magnitude.

## Example

```julia
searcher = KNearestSearch(data, 10)
averaged_inplace = smoothpars(localaniso, searcher)
```

## Reference

Markley, F.L., et al. (2007). [Averaging quaternions](https://ntrs.nasa.gov/archive/nasa/casi.ntrs.nasa.gov/20070017872.pdf)
"""
smoothpars(lpars, searcher::NeighborSearchMethod; power=0.0, metric=Euclidean()) =
	interpolate(lpars, searcher, power=power)

function interpolate(lpars, searcher::NeighborSearchMethod, domain=nothing;
	power::Real=0, metric=Euclidean())
	D = searcher.domain
	N = embeddim(D)
	len = domain==nothing ? nelements(D) : nelements(domain)
	@assert nelements(D)==nvals(lpars) "searcher domain must match number of local anisotropies"

    quat = Array{Quaternion}(undef,len)
	mag  = lpars.magnitude
    m    = domain==nothing && power==0 ? mag : Array{Float64}(undef,N,len)

    Threads.@threads for i in 1:len
        ic  = domain==nothing ? centroid(D,i) : centroid(domain,i)
		icoords  = coordinates(ic)
		neighids = search(ic, searcher)

		if length(neighids) == 0
			throw(ErrorException("zero neighbors at some location; adjust searcher"))
			# or accept missing in LocalAnisotropy?
		elseif power==0.0
			quat[i] = quatavg(rotation(lpars,neighids))
			if domain!=nothing
				mi      = magnitude(lpars,neighids)
				m[:,i] .= mapreduce(x->quantile(view(mi,x,:),0.5), vcat, 1:N)
			end
		else
			xcoords = coords(D, neighids)
			prewgts = 1 ./ (eps() .+ Distances.colwise(metric, icoords, xcoords)) .^ power
			weights = prewgts ./ sum(prewgts)
			quat[i] = quatavg(rotation(lpars,neighids), weights)
			mi      = magnitude(lpars,neighids)
			wgts    = Weights(weights)
			m[:,i] .= mapreduce(x->quantile(view(mi,x,:),wgts,0.5), vcat, 1:N)
			# mean instead median for magnitude?
		end
    end

    LocalAnisotropy(quat, m)
end

quat2vector(q) = [q.q0; q.q1; q.q2; q.q3]
tensor(q) = q' .* q
wtensor(q) = q[1] .* (q[2]' .* q[2])

function quatavg(qarr, warr=[])
	length(qarr) == 1 && (return qarr)
	#@assert in(length(warr),[n,0])
	Q = mapreduce(quat2vector, hcat, qarr)
    n = size(Q,2)
    W = length(warr) > 0

    A = W ? mapreduce(wtensor, +, zip(warr, eachcol(Q))) : mapreduce(tensor, +, eachcol(Q))

    # scale
    Nw = W ? sum(warr) : n
    A ./= Nw

    # compute eigenvalues and -vectors
    T = eigen(Symmetric(SMatrix{4,4}(A)))
    V = view(T.vectors[:, sortperm(T.values,rev=true)],:,1)
    Quaternion(V...)
end


module CampoMinado end

export campoMinado
export ehValido

function ehValido(campo ::String)
 length(campo) == 0 || ( ismatch(r"((X{1}\s{1}|O{1}\s{1})\n{1})+", campo) && checarTamanhoDeLinhasDo(campo))
end

function checarTamanhoDeLinhasDo(campo::String)
  strs = split(campo, '\n')

  if length(strs) != 0
    w = map (x -> mapreduce(z-> 1, +, x) , strs[1:end-1])

    if length(w) != 0
        m = maximum(w);
        return all(z ->  m == z, w)
      end
    return false
  end
end

function paraArray(campo::String)
    map( x -> reduce(push!, Char[], x), split (filter(c -> c != ' ', campo), '\n')[1:end-1])
  end

function campoMinado(campo ::String)
  if ehValido(campo)
    r = ""
    campoarr = paraArray(campo)

    valor = function(arr, x, y)
      if x > 0 && y > 0 && x < length(arr) && y < length(arr[x])
        if(arr[x,y] = 'X')
          1
        end
      end
      0
    end

    for i = [1:length(campoarr)]
       for j = [1:length(campoarr[i])]
          if campoarr[i][j] != 'X'
             contar =  valor(campoarr,i-1,j-1) + valor(campoarr,i,j-1) + valor(campoarr,i + 1,j - 1)
                    +  valor(campoarr,i-1,j)   + valor(campoarr,i+1,j)
                    +  valor(campoarr,i-1,j)   + valor(campoarr,i,j+1)   + valor(campoarr,i + 1, j + 1)
             insert!(campoarr[i], j, char(contar) )
          end
      end
    end
   return campoarr
  end
  throw(ArgumentError("Campo minado com formato incorreto!"))
end


"""
# Brent
## Constructor
```julia
    Brent(;)
```

## Description
Also known as the Brent-Dekker algorith, `Brent` is a univariate optimization
algorithm for minimizing functions on some interval `[a,b]`. The method uses bisection
to find a zero of the gradient. If the original interval contains a minimum,
bisection will reliably find the solution, but can be slow. To this end `Brent`
combines bisection with the secant method and inverse quadratic interpolation to
accelerate convergence.

## References
R. P. Brent (2002) Algorithms for Minimization Without Derivatives. Dover edition.
"""
struct Brent <: UnivariateOptimizer end

Base.summary(::Brent) = "Brent's Method"

function optimize(
        f, x_lower::T, x_upper::T,
        mo::Brent;
        rel_tol::T = sqrt(eps(T)),
        abs_tol::T = eps(T),
        iterations::Integer = 1_000,
        store_trace::Bool = false,
        show_trace::Bool = false,
        callback = nothing,
        show_every = 1,
        extended_trace::Bool = false) where T <: AbstractFloat
    t0 = time()
    options = (store_trace=store_trace, show_trace=show_trace, show_every=show_every, callback=callback)
    if x_lower > x_upper
        error("x_lower must be less than x_upper")
    end

    # Save for later
    initial_lower = x_lower
    initial_upper = x_upper

    golden_ratio::T = T(1)/2 * (3 - sqrt(T(5.0)))

    new_minimizer = x_lower + golden_ratio*(x_upper-x_lower)
    new_minimum = f(new_minimizer)
    best_bound = "initial"
    f_calls = 1 # Number of calls to f
    step = zero(T)
    old_step = zero(T)

    old_minimizer = new_minimizer
    old_old_minimizer = new_minimizer

    old_minimum = new_minimum
    old_old_minimum = new_minimum

    iteration = 0
    converged = false

    # Trace the history of states visited
    tr = OptimizationTrace{T, typeof(mo)}()
    tracing = store_trace || show_trace || extended_trace || callback != nothing
    stopped_by_callback = false
    if tracing
        # update trace; callbacks can stop routine early by returning true
        state = (new_minimizer=new_minimizer,
                 x_lower=x_lower,
                 x_upper=x_upper,
                 best_bound=best_bound,
                 new_minimum=new_minimum)
        stopped_by_callback = trace!(tr, nothing, state, iteration, mo, options, time()-t0)
    end

    while iteration < iterations && !stopped_by_callback

        p = zero(T)
        q = zero(T)

        x_tol = rel_tol * abs(new_minimizer) + abs_tol

        x_midpoint = (x_upper+x_lower)/2

        if abs(new_minimizer - x_midpoint) <= 2*x_tol - (x_upper-x_lower)/2
            converged = true
            break
        end

        iteration += 1

        if abs(old_step) > x_tol
            # Compute parabola interpolation
            # new_minimizer + p/q is the optimum of the parabola
            # Also, q is guaranteed to be positive

            r = (new_minimizer - old_minimizer) * (new_minimum - old_old_minimum)
            q = (new_minimizer - old_old_minimizer) * (new_minimum - old_minimum)
            p = (new_minimizer - old_old_minimizer) * q - (new_minimizer - old_minimizer) * r
            q = 2(q - r)

            if q > 0
                p = -p
            else
                q = -q
            end
        end

        if abs(p) < abs(q*old_step/2) && p < q*(x_upper-new_minimizer) && p < q*(new_minimizer-x_lower)
            old_step = step
            step = p/q

            # The function must not be evaluated too close to x_upper or x_lower
            x_temp = new_minimizer + step
            if ((x_temp - x_lower) < 2*x_tol || (x_upper - x_temp) < 2*x_tol)
                step = (new_minimizer < x_midpoint) ? x_tol : -x_tol
            end
        else
            old_step = (new_minimizer < x_midpoint) ? x_upper - new_minimizer : x_lower - new_minimizer
            step = golden_ratio * old_step
        end

        # The function must not be evaluated too close to new_minimizer
        if abs(step) >= x_tol
            new_x = new_minimizer + step
        else
            new_x = new_minimizer + ((step > 0) ? x_tol : -x_tol)
        end

        new_f = f(new_x)
        f_calls += 1

        if new_f < new_minimum
            if new_x < new_minimizer
                x_upper = new_minimizer
                best_bound = "upper"
            else
                x_lower = new_minimizer
                best_bound = "lower"
            end
            old_old_minimizer = old_minimizer
            old_old_minimum = old_minimum
            old_minimizer = new_minimizer
            old_minimum = new_minimum
            new_minimizer = new_x
            new_minimum = new_f
        else
            if new_x < new_minimizer
                x_lower = new_x
            else
                x_upper = new_x
            end
            if new_f <= old_minimum || old_minimizer == new_minimizer
                old_old_minimizer = old_minimizer
                old_old_minimum = old_minimum
                old_minimizer = new_x
                old_minimum = new_f
            elseif new_f <= old_old_minimum || old_old_minimizer == new_minimizer || old_old_minimizer == old_minimizer
                old_old_minimizer = new_x
                old_old_minimum = new_f
            end
        end
        if tracing
            # update trace; callbacks can stop routine early by returning true
            state = (new_minimizer=new_minimizer,
                     x_lower=x_lower,
                     x_upper=x_upper,
                     best_bound=best_bound,
                     new_minimum=new_minimum)
            stopped_by_callback = trace!(tr, nothing, state, iteration, mo, options, time()-t0)
        end
    end

    return UnivariateOptimizationResults(mo,
                                         initial_lower,
                                         initial_upper,
                                         new_minimizer,
                                         new_minimum,
                                         iteration,
                                         iteration == iterations,
                                         converged,
                                         rel_tol,
                                         abs_tol,
                                         tr,
                                         f_calls)
end

function trace!(tr, d, state, iteration, method::Brent, options, curr_time=time())
    dt = Dict()
    dt["time"] = curr_time
    dt["minimizer"] = state.new_minimizer
    dt["x_lower"] = state.x_lower
    dt["x_upper"] = state.x_upper
    dt["best bound"] = state.best_bound
    T = eltype(state.new_minimum)

    update!(tr,
            iteration,
            state.new_minimum,
            T(NaN),
            dt,
            options.store_trace,
            options.show_trace,
            options.show_every,
            options.callback)
end


function solve()
    n = parse(Int, readline())
    as = [parse(Int, x) for x in split(readline())]

    ret = 0
    m = 0
    for x=2:1000
        s = sum([a%x==0 for a in as])
        if s > m
            m = s
            ret = x
        end
    end
    ret
end

println(solve())


"""
    scaled_squared_distance(log_ell::AbstractArray{Float64, 1},
        x::AbstractArray{Float64, 2}, z::AbstractArray{Float64, 2})

Compute the scaled squared distance between `x` and `z`:
```math
r_{ij} = \\sum_{k=1}^d\\frac{(x_{ik}-z_{jk})^2}{l_k^2}
```
The gradient of this function with respect to
length scale hyperparameter(s) is returned by [`scaled_squared_distance_grad`](@ref).

# Arguments
- `x, z`: Input data, reshaped into 2-d arrays.
  `x` must have dimensions ``n \\times d``; `z` must have dimensions ``m \\times d``.
- `log_ell`: logarithm of length scale(s). Can either be an array of size one (isotropic),
  or an array of size `d` (ARD)

# Return
An ``n \\times m`` matrix of scaled squared distances
"""
function scaled_squared_distance(log_ell::AbstractArray{Float64, 1},
        x1::AbstractArray{Float64, 2}, x2::AbstractArray{Float64, 2})
    ell = exp.(-log_ell .* 2)
    if x1 === x2
        mu = mean(x1, dims=1)
    else
        n1 = size(x1, 1)
        n2 = size(x2, 1)
        mu = n1 / (n1 + n2) * mean(x1, dims=1) + n2 / (n1 + n2) * mean(x2, dims=1)
    end
    x1 = x1 .- mu               # subtract mean for stability, in case x1 or x2 are huge
    ax1 = x1 .* ell'            # a * x1
    sax1 = sum(x1 .* ax1, dims=2)    # a * x1 ^ 2
    if x1 === x2                # a shortcut if we need K(x, x) for log likelihood
        ax2 = ax1
        sax2 = sax1
    else                        # repeat computations for x2
        x2 = x2 .- mu
        ax2 = x2 .* ell'
        sax2 = sum(x2 .* ax2, dims=2)
    end
    D2 = sax1 .- 2 * x1 * ax2' .+ sax2'   # a((b-c)^2) = ab^2 - 2abc + ac^2
end

"""
    scaled_squared_distance_grad(log_ell::AbstractArray{Float64, 1},
        x::AbstractArray{Float64, 2}, z::AbstractArray{Float64, 2}, R::AbstractArray{Float64, 2})

Return the gradient of the [`scaled_squared_distance`](@ref) function with respect to
logarigthms of length scales, based on the provided direction matrix.

# Arguments
- `x, z`: Input data, reshaped into 2-d arrays.
  `x` must have dimensions ``n \\times d``; `z` must have dimensions ``m \\times d``.
- `log_ell`: logarithm of length scale(s). Can either be an array of size one (isotropic),
  or an array of size `d` (ARD)
- `R` the direction matrix, ``n \\times m``. This can be used to compute the gradient
  of a function that depends on [`scaled_squared_distance`](@ref) via the chain rule.

# Return
A vector of size `length(log_ell)`, whose ``k``'th element is equal to
```math
\\text{tr}(R \\frac{\\partial K}{\\partial l_k})
```
"""
function scaled_squared_distance_grad(log_ell::AbstractArray{Float64, 1},
        x1::AbstractArray{Float64, 2}, x2::AbstractArray{Float64, 2}, R::AbstractArray{Float64, 2})
    n1 = size(x1, 1)
    n2 = size(x2, 2)
    if x1 === x2                                        # a shortcut if x1 and x2 are the same - handy for log likelihood optimisation
        x1 = x1 .- mean(x1, dims=1)                          # subtract mean for stability if x1 is huge
        x12_2 = (sum(R, dims=2)' + sum(R, dims=1)) * (x1 .^ 2)    # R * x1 ^ 2 + R * x2 ^ 2
        x1x2 = sum((R * x1) .* x1, dims=1)                   # R * x1 * x2
        rk = x12_2' - 2 * x1x2'                         # R * (x1 - x2) ^ 2 = R * x1 ^ 2 - 2 * R * x1 * x2 + R * x2 ^ 2
    else
        mu = n1 / (n1 + n2) * mean(x1, dims=1) + n2 / (n1 + n2) * mean(x2, dims=1)
        x1 = x1 .- mu                                   # subtract mean for stability if x1 or x2 is huge
        x2 = x2 .- mu
        x1_2 = sum(R, dims=2)' * (x1 .^ 2)                   # R * x1 ^ 2
        x2_2 = sum(R, dims=1) * (x2 .^ 2)                    # R * x2 ^ 2
        x1x2 = sum((R * x2) .* x1, dims=1)                   # R * x1 * x2
        rk = x1_2' - 2 * x1x2' + x2_2'                  # R * (x1 - x2) ^ 2 = R * x1 ^ 2 - 2 * R * x1 * x2 + R * x2 ^ 2
    end
    ell = exp.(-2 * log_ell)
    return reshape(-2 * ell .* rk, length(ell))
end


using RCall, JLD

R"
library(gridExtra)
library(cowplot)
library(tidyverse)
library(latex2exp)
library(viridis)
"

DF = load("/Users/dariabystrova/Documents/GitHub/approximatingBNPpriors/saves_for_figures/figure_6_1.jld")
EV__NGG = DF["NGG"]

DF_NGG = DataFrame(Exp = first.(EV__NGG),Var =last.(EV__NGG), beta = first.(grid_vec), sigma = last.(grid_vec))
DF_NGG.std = sqrt.(DF_NGG.Var)
DF_NGG.beta_log = log.(DF_NGG.beta)
DF_NGG

R"
range01 <- function(x){(x-min(x))/(max(x)-min(x))}
df_r= $DF_NGG
df_r$beta_scaled= range01(df_r$beta_log)
df_r$color = rgb(df_r$beta_scaled, df_r$sigma,0.5,1)
df_r$color_sigma = rgb(0, df_r$sigma,0.5,1)
df_r$color_beta = rgb(df_r$beta_scaled, 0,0.5,1)
p <- df_r %>% ggplot(aes(x=Exp, y = std, group=sigma))  + geom_line( alpha = 0.8,linetype = 'longdash') + geom_point(aes(colour=color), size=4)
p_ngg <- p  + geom_line(data =df_r, aes(x=Exp, y = std, group=beta), alpha = 0.8, linetype='dotted') +
  xlim(0, 100)+ ylim(0,15)+ labs(y='Std', x='Expectation')+scale_colour_identity()+
  theme_classic()+ theme(plot.title = element_text(hjust = 0.5,size = 15), axis.text.x = element_text(size=15), axis.text.y = element_text(size=15), plot.margin = unit(c(0,0, 0, 0),'pt'))
"

R"p_ngg
pdf('/Users/dariabystrova/Documents/GitHub/approximatingBNPpriors/figures/figure_6_1.pdf', width= 4, height = 4)
plot(p_ngg)
dev.off()"


using LaTeXStrings
include("workspace.jl")
pyplot()

N = 27
a = 4

#Calculate the plotting data
xvals = 0:N-1
yvals = [mod(a^x,N) for x ∈ xvals]

plot(xvals, yvals, linewidth=1, ls = :dot, marker = :x,
xlabel = "\$x\$", ylabel = "\$$a^x\$ mod \$$N\$",
title = "Periodicity of modular exponentiation", legend = false, grid = false, mc = :red, lc = :red)
xticks!(0:5:xvals[end]+1)
gui()

r = 0
i = 2
while i < length(xvals)
    if yvals[1] == yvals[i]
        return r = i-1
        break
    end
    if i == length(xvals)
        println("No period was found")
        break
    end
    i += 1
end
plot!([0:1:r,r:-1:0],ones(length(0:1:r),2),
arrow = :arrow, lc = :black)
annotate!((r/2,1.5, Plots.text("\$r = $r\$", 8,:black,:center)))

##
function controlbits(a::Int, n::Int)
    # a is the starting random integer, n is the number of qubits in the circuit
    digits(a,base=2,pad=n)
end

function phaseadd!(C::Circuit, controlbits::AbstractArray{Int})
    n = length(controlbits)
    for m in 1:n
        for i in m:n
            if controlbits[n-i+1] == 1
                k = i-m+1
                λ = 2π/(2^(k))
                U!(C,λ,n-m+1)
            end
        end
    end

   # for qubit in 1:div(n,2)-1
   #     Swap!(C, qubit, n-qubit)
   # end
end

function phaseadd_dagger!(C::Circuit, controlbits::AbstractArray{Int})
    n = length(controlbits)
    for m in n:-1:1
        for i in n:-1:m
            if controlbits[n-i+1] == 1
                λ = 2π/2^(i-m+1)
                U!(C,-λ,n-m+1)
            end
        end
    end
end

#for m in 1:n
#    for i in m:n
#        println(n-i+1," ",i-m+1)
#    end
#end
#
#for m in 5:-1:1
#    for i in 5:-1:m
#        println(5-i+1," ",i-m+1)
#    end
#end

function qft_dagger!(C::Circuit, n::Int)
    #Don't forget the swaps!
    #for qubit in 1:div(n,2)-1
    #    Swap!(C, qubit, n-qubit)
    #end
    for j in 1:n-1
        for m in 1:j-1
            CU!(C,-π/(2^(j-m)), m, j)
        end
        H!(C,j)
    end
end

function qft!(C::Circuit, n::Int)
    #Don't forget the swaps!
    for j in n-1:-1:1
        H!(C,j)
        for m in j-1:-1:1
            CU!(C,π/(2^(j-m)), m, j)
        end
    end
    #for qubit in 1:div(n,2)-1
    #    Swap!(C, qubit, n-qubit)
    #end
end

n = 3
c1 = Circuit(n)
a = 3

#X!(c1,1)
#X!(c1,2)
controls = controlbits(a,n-1)
qft!(c1,n)
phaseadd!(c1,controls)
#phaseadd_dagger!(c1,controls)
qft_dagger!(c1,n)
#qft_dagger!(c1,n)
#qft!(c1,n)
dat = shots(c1, 10_000)
graph!(dat)


using MinimaxCalibratedEBayes
using ECOS

const MCEB = MinimaxCalibratedEBayes

@testset "Check Proposition 1" begin

prior_dbn =  MixtureModel(Normal, [(-2,0.4), (2, 0.4), (0, 1.0)])
target = PriorDensityTarget(1.0)
prior_classes = [GaussianMixturePriorClass(0.4, -3:0.1:3, ECOS.Optimizer),
                 HermitePriorClass(qmax=40, sobolev_order=2.0, sobolev_bound= 1.0, solver=ECOS.Optimizer)]

for prior_class in prior_classes
	Zs_discr = DiscretizedStandardNormalSamples(-4:0.1:4)
	C∞ = 0.02
	Zs_discr_nbhood = set_neighborhood(Zs_discr, prior_dbn; C∞ = C∞)
	delta_tuner= MCEB.FixedDelta(0.001)

	steinminimax_fit = SteinMinimaxEstimator(Zs_discr, prior_class,
								         target, delta_tuner)

    steinminimax_fit_nbhood = SteinMinimaxEstimator(Zs_discr_nbhood, prior_class,
								         target, delta_tuner)

    # Check if Delta is correct
    @test steinminimax_fit_nbhood.δ == delta_tuner.δ

    @test MCEB.get_δ(steinminimax_fit_nbhood.model, recalculate_δ=false) ≈ delta_tuner.δ atol=1e-6
    @test MCEB.get_δ(steinminimax_fit_nbhood.model; recalculate_δ=true ) ≈ delta_tuner.δ atol=1e-6

    # Check statements about bias
   bias_at_g2 = worst_case_bias(steinminimax_fit_nbhood.Q, Zs_discr_nbhood,
                             steinminimax_fit_nbhood.g2, target)[:bias]

   bias_at_g1 = worst_case_bias(steinminimax_fit_nbhood.Q, Zs_discr_nbhood,
							steinminimax_fit_nbhood.g1, target)[:bias]

   worst_case_bias_stein = worst_case_bias(steinminimax_fit_nbhood.Q, Zs_discr_nbhood,
                                        prior_class, target)

   @test 	worst_case_bias_stein.max_bias ≈ -worst_case_bias_stein.min_bias atol=0.001
   @test 	worst_case_bias_stein.max_bias ≈ worst_case_bias_stein.max_abs_bias atol=0.001
   @test 	abs2(worst_case_bias_stein.max_bias) ≈ worst_case_bias_stein.max_squared_bias atol=0.001

   @test 	worst_case_bias_stein.max_bias ≈ bias_at_g1 atol=0.001
   @test 	worst_case_bias_stein.min_bias ≈ bias_at_g2 atol=0.001
   @test   worst_case_bias_stein.max_bias ≈ steinminimax_fit_nbhood.max_bias atol=0.001
   @test   2*worst_case_bias_stein.max_bias ≈ steinminimax_fit_nbhood.ω_δ - steinminimax_fit_nbhood.δ*steinminimax_fit_nbhood.ω_δ_prime atol=0.001
end
end


## Root finding for Chebyshev expansions
#
#  Contains code that is based in part on Chebfun v5's chebfun/@chebteck/roots.m,
# which is distributed with the following license:

# Copyright (c) 2015, The Chancellor, Masters and Scholars of the University
# of Oxford, and the Chebfun Developers. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#     * Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#     * Neither the name of the University of Oxford nor the names of its
#       contributors may be used to endorse or promote products derived from
#       this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

function complexroots{C<:Chebyshev}(f::Fun{C})
    if ncoefficients(f)==0 || (ncoefficients(f)==1 && isapprox(f.coefficients[1],0))
        warn("Tried to take roots of a zero function.  Returning [].")
        Complex128[]
    elseif ncoefficients(f)==1
        Complex128[]
    elseif ncoefficients(f)==2
        Complex128[-f.coefficients[1]/f.coefficients[2]]
    else
        fromcanonical(f,colleague_eigvals(f.coefficients))
    end
end

#function roots(f::Fun)
#    irts=map(real,filter!(x->abs(x)<=1.+10eps(),filter!#(isreal,complexroots(f.coefficients))))
#
#    map!(x->x>1.?1.:x,irts)
#    map!(x->x<-1.?-1.:x,irts)
#
#    if length(irts)==0
#        Float64[]
#    else
#        fromcanonical(f,irts)
#    end
#end

function roots(f::Fun)
    f2=Fun(f,domain(f)) # default is to convert to Chebyshev/Fourier
    if space(f2)==space(f)
        error("roots not implemented for "*string(typeof(f)))
    else
        roots(f2)
    end
end

function roots{C<:Chebyshev}( f::Fun{C} )
# FIND THE ROOTS OF AN IFUN.

    d = domain(f)
    c = f.coefficients
    vscale = maxabs(values(f))
    if vscale == 0
        warn("Tried to take roots of a zero function.  Returning [].")
        return eltype(domain(f))[]
    end

    hscale = maximum( [abs(first(d)), abs(last(d))] )
    htol = eps(2000.)*max(hscale, 1)  # TODO: choose tolerance better

    if eltype(f) == BigFloat
        r = rootsunit_coeffs(convert(Vector{Float64},c./vscale), Float64(htol))
        # Map roots from [-1,1] to domain of f:
        rts = fromcanonical(d,r)
        fp = differentiate(f)
        while norm(f(rts)) > 1000eps(eltype(f))
            rts .-=f(rts)./fp(rts)
        end
    elseif eltype(f) == Complex{BigFloat}
        r = rootsunit_coeffs(convert(Vector{Complex{Float64}},c./vscale), Float64(htol))
        # Map roots from [-1,1] to domain of f:
        rts = fromcanonical(d,r)
        fp = differentiate(f)
        while norm(f(rts)) > 1000eps(eltype(f))
            rts .-=f(rts)./fp(rts)
        end
    else
        cvscale=c./vscale
        r = rootsunit_coeffs(cvscale, Float64(htol))

        # Check endpoints, as these are prone to inaccuracy
        # which can be deadly.
        if (isempty(r) || !isapprox(last(r),1.)) && abs(sum(cvscale)) < htol
            push!(r,1.)
        end
        if (isempty(r) || !isapprox(first(r),-1.)) && abs(alternatingsum(cvscale)) < htol
            insert!(r,1,-1.)
        end

        # Map roots from [-1,1] to domain of f:
        rts = fromcanonical(d,r)
    end
    return rts
end

function colleague_matrix{T<:Number}( c::Vector{T} )
#TODO: This is command isn't typed correctly
# COMPUTE THE ROOTS OF A LOW DEGREE POLYNOMIAL BY USING THE COLLEAGUE MATRIX:
    n = length(c) - 1
    A=zeros(T,n,n)

    for k=1:n-1
        A[k+1,k]=0.5
        A[k,k+1]=0.5
    end
    for k=1:n
        A[1,end-k+1]-=0.5*c[k]/c[end]
    end
    A[n,n-1]=one(T)

    # TODO: can we speed things up because A is lower-Hessenberg
    # Standard colleague matrix (See [Good, 1961]):
    A
end

function colleague_balance!(M)
    n=size(M,1)
    if n<2
        return M
    end

    d=1+abs(M[1,2])
    M[1,2]/=d;M[2,1]*=d

    for k=3:n
        M[k-1,k]*=d;M[k,k-1]/=d
        d+=abs(M[1,k])
        M[1,k]/=d;M[k-1,k]/=d;M[k,k-1]*=d
    end
    M
end

# colleague_eigvals( c::Vector )=hesseneigvals(colleague_balance!(colleague_matrix(c)))

colleague_eigvals( c::Vector )=eigvals(colleague_matrix(c))

function PruneOptions( r, htol::Float64 )
# ONLY KEEP ROOTS IN THE INTERVAL

    # Remove dangling imaginary parts:
    r = real( r[ abs(imag(r)) .< htol ] )
    # Keep roots inside [-1 1]:
    r = sort( r[ abs(r) .<= 1+htol ] )
    # Put roots near ends onto the domain:
    r = min( max( r, -1 ), 1)

    # Return the pruned roots:
    return r
end

rootsunit_coeffs{T<:Number}(c::Vector{T}, htol::Float64)=rootsunit_coeffs(c,htol,ClenshawPlan(T,Chebyshev(),length(c),length(c)))
function rootsunit_coeffs{S,T<:Number}(c::Vector{T}, htol::Float64,clplan::ClenshawPlan{S,T})
# Computes the roots of the polynomial given by the coefficients c on the unit interval.

    # If recursive subdivision is used, then subdivide [-1,1] into [-1,splitPoint] and [splitPoint,1].
    const splitPoint = -0.004849834917525

    # Simplify the coefficients by chopping off the tail:
    nrmc=norm(c, 1)
    @assert nrmc > 0  # There's an error if we make it this far
    c = chop(c,eps()*nrmc)
    n = length(c)

    if n == 0

        # EMPTY FUNCTION
        r = Float64[]

    elseif n == 1

        # CONSTANT FUNCTION
        r = ( c[1] == 0.0 ) ? Float64[0.0] : Float64[]

    elseif n == 2

        # LINEAR POLYNOMIAL
        r1 = -c[1]/c[2];
        r = ( (abs(imag(r1))>htol) | (abs(real(r1))>(1+htol)) ) ? Float64[] : Float64[max(min(real(r1),1),-1)]

    elseif n <= 70

        # COLLEAGUE MATRIX
        # The recursion subdividing below will keep going until we have a piecewise polynomial
        # of degree at most 50 and we likely end up here for each piece.

        # Adjust the coefficients for the colleague matrix
        # Prune roots depending on preferences:
        r = PruneOptions( colleague_eigvals(c), htol )::Vector{Float64}
    else

        #  RECURSIVE SUBDIVISION:
        # If n > 50, then split the interval [-1,1] into [-1,splitPoint], [splitPoint,1].
        # Find the roots of the polynomial on each piece and then concatenate. Recurse if necessary.

        # Evaluate the polynomial at Chebyshev grids on both intervals:
        #(clenshaw! overwrites points, which only makes sence if c is real)

        v1 = isa(c,Vector{Float64})?clenshaw!( c, points([-1,splitPoint],n),clplan):clenshaw( c, points([-1,splitPoint],n),clplan)
        v2 = isa(c,Vector{Float64})?clenshaw!( c, points([splitPoint,1] ,n),clplan):clenshaw( c, points([splitPoint,1] ,n),clplan)

        # Recurse (and map roots back to original interval):
        p = plan_chebyshevtransform( v1 )
        r = [ (splitPoint - 1)/2 + (splitPoint + 1)/2*rootsunit_coeffs( chebyshevtransform(v1,p), 2*htol,clplan) ;
                 (splitPoint + 1)/2 + (1 - splitPoint)/2*rootsunit_coeffs( chebyshevtransform(v2,p), 2*htol,clplan) ]

    end

    # Return the roots:
    r
end

function extremal_args{S<:PiecewiseSpace}(f::Fun{S})
    return cat(1,[extremal_args(fp) for fp in vec(f)]...)
end

function extremal_args(f::Fun)
    d=domain(f)
    if isa(d,PeriodicInterval)
        roots(differentiate(f))
    elseif isa(d,PeriodicDomain)  # avoid complex domains
        S=typeof(space(f))
        fromcanonical(f,extremal_args(setcanonicaldomain(f)))
    else
        dab=∂(domain(f))
        if ncoefficients(f) <=2 #TODO this is only relevant for Polynomial bases
            dab
        else
            [dab;roots(differentiate(f))]
        end
    end
end

for op in (:(Base.maximum),:(Base.minimum),:(Base.extrema),:(Base.maxabs),:(Base.minabs))
    @eval function $op{S<:RealSpace,T<:Real}(f::Fun{S,T})
        pts = extremal_args(f)

        $op(f(pts))
    end
end

for op in (:(Base.maxabs),:(Base.minabs))
    @eval function $op(f::Fun)
        # complex spaces/types can have different extrema
        pts = extremal_args(abs(f))

        $op(f(pts))
    end
end

for op in (:(Base.maximum),:(Base.minimum),:(Base.maxabs),:(Base.minabs))
    @eval begin
        $op{SV,DD<:UnionDomain,d,T<:Real}(f::Fun{PiecewiseSpace{SV,RealBasis,DD,d},T}) =
            $op(map($op,vec(f)))
    end
end

Base.extrema{SV,DD<:UnionDomain,d,T<:Real}(f::Fun{PiecewiseSpace{SV,RealBasis,DD,d},T}) =
    mapreduce(extrema,(x,y)->extrema([x...;y...]),vec(f))

for op in (:(Base.indmax),:(Base.indmin))
    @eval begin
        function $op{S<:RealSpace,T<:Real}(f::Fun{S,T})
            # the following avoids warning when differentiate(f)==0
            pts = extremal_args(f)
            # the extra real avoids issues with complex round-off
            pts[$op(real(f(pts)))]
        end

        function $op{S,T}(f::Fun{S,T})
            # the following avoids warning when differentiate(f)==0
            pts = extremal_args(f)
            fp=f(pts)
            @assert norm(imag(fp))<100eps()
            pts[$op(real(fp))]
        end
    end
end

for op in (:(Base.findmax),:(Base.findmin))
    @eval begin
        function $op{S<:RealSpace,T<:Real}(f::Fun{S,T})
            # the following avoids warning when differentiate(f)==0
            pts = extremal_args(f)
            ext,ind = $op(f(pts))
	        ext,pts[ind]
        end
    end
end

## Root finding for Laurent expansion

function companion_matrix{T}(c::Vector{T})
    n=length(c)-1

    if n==0
        zeros(T,0,0)
    else
        A=zeros(T,n,n)
        for k=1:n
            A[k,end]=-c[k]/c[end]
        end
        for k=2:n
            A[k,k-1]=one(T)
        end
        A
    end
end

if isdir(Pkg.dir("AMVW"))
    using AMVW
    function complexroots(coefficients::Vector)
        c=chop(coefficients,10eps())

        # Only use special routine for large roots
        if length(c)≥70
            Main.AMVW.rootsAMVW(c)
        else
            hesseneigvals(companion_matrix(c))
        end
    end
else
    complexroots(coefficients::Vector)=hesseneigvals(companion_matrix(chop(coefficients,10eps())))
end

complexroots(neg::Vector,pos::Vector)=complexroots([flipdim(chop(neg,10eps()),1);pos])
complexroots{DD}(f::Fun{Laurent{DD}})=mappoint(Circle(),domain(f),complexroots(f.coefficients[2:2:end],f.coefficients[1:2:end]))
complexroots{DD}(f::Fun{Taylor{DD}})=mappoint(Circle(),domain(f),complexroots(f.coefficients))

function roots{DD}(f::Fun{Laurent{DD}})
    irts=filter!(z->in(z,Circle()),complexroots(Fun(f.coefficients,Laurent(Circle()))))
    if length(irts)==0
        Complex{Float64}[]
    else
        rts=fromcanonical(f,tocanonical(Circle(),irts))
        if isa(domain(f),PeriodicInterval)
            sort!(real(rts))  # Make type safe?
        else
            rts
        end
    end
end

roots{D}(f::Fun{Fourier{D}})=roots(Fun(f,Laurent))

function roots{P<:PiecewiseSpace}(f::Fun{P})
    rts=mapreduce(roots,vcat,vec(f))
    k=1
    while k < length(rts)
        if isapprox(rts[k],rts[k+1])
            rts=rts[[1:k,k+2:end]]
        else
            k+=1
        end
    end

    rts
end

## Root finding for JacobiWeight expansion

# Add endpoints for JacobiWeight
# TODO: what about cancellation?
function roots{S<:JacobiWeight,T}(f::Fun{S,T})
    sp=space(f)
    d=domain(sp)
    rts=roots(Fun(f.coefficients,sp.space))
    if sp.α > 0
        rts=[first(d);rts]
    end
    if sp.β > 0
        rts=[rts;last(d)]
    end
    rts
end


# constants in cgs units
const c = 2.997824582e10     # cm s^-1
const σ = 5.6704e-5          # erg cm^-2 s^-1 K^-4
const kB = 1.3806485e-16     # erg K^-1
const mH = 1.6737236e-24     # gm

# constants in other units
const c_ms = 2.99782458e8    # m s^-1
const c_kms = 2.99782458e5   # km s^-1


abstract Position

type GridPosition
  x :: Float64
  y :: Float64
end
uGrid = GridPosition(0.0,0.5)
vGrid = GridPosition(0.5,0.0)
mGrid = GridPosition(0.0,0.0)
sGrid = GridPosition(0.5,0.5)

type GridMoving
  xshift :: Float64
  yshift :: Float64
  xband :: Float64
  yband :: Float64
end

function GridMoving( o :: GridPosition, t :: GridPosition, xband :: Float64, yband :: Float64)
  xshift = t.x - o.x
  yshift = t.y - o.y
  return GridMoving(xshift, yshift, xband, yband)
end

function intercept( ss :: Array{Float64,2}, g :: GridMoving; keep_when_notmoving :: Bool = false)
  signal = ss
  xband = g.xband
  yband = g.yband
  xshift = g.xshift
  yshift = g.yshift
  kx = collect(1:(size(signal,1)/xband)) / size(signal,1)
  ky = collect(1:(size(signal,2)/yband)) / size(signal,2)

  fresx = zeros(signal)
  if ((xshift == 0.0) & keep_when_notmoving)
    fresx = signal
  else
    for ix = kx
      psix = ix*pi*collect(1:size(signal,1))
      psix1 = ix*pi*(xshift + collect(1:size(signal,1)))
      fsinx  = mean(sin(psix) .* signal,1).* sin(psix1)/(0.5)
      fcosx  = mean(cos(psix) .* signal,1).* cos(psix1)/(0.5)
      fresx += fsinx + fcosx
    end
  end

  signal = fresx
  fresy = zeros(size(signal))
  if ((yshift == 0.0) & keep_when_notmoving)
    fresy = signal
  else
    for iy = ky
      psiy = iy*pi*collect(1:size(signal,2))
      psiy1 = iy*pi*(yshift + collect(1:size(signal,2)))
      psiy = psiy'
      psiy1 = psiy1'
      fsiny  = mean(sin(psiy) .* signal,2).* sin(psiy1)/(0.5)
      fcosy  = mean(cos(psiy) .* signal,2).* cos(psiy1)/(0.5)
      fresy += fsiny + fcosy
    end
  end

  return fresy
end

function intercept(ss:: Array{Float64,3}, g:: GridMoving; kws...)
  res = zeros(ss)
  for iz = 1:size(ss,3)
    res[:,:,iz] = intercept(ss[:,:,iz],g;kws...)
  end
  return res
end
export intercept
export GridMoving


# Profiles transit_poly.jl

r = 0.1  # radius ratio is 0.1
u = [0.2,0.3]; n = length(u)
N_c = n+2

const c_n = zeros(typeof(r),n+3)
const dfdrbc = zeros(typeof(r),n+3)
const a_n = zeros(typeof(r),n+1)
const dadu = zeros(typeof(r),n+1,n)
const dcdu = zeros(typeof(r),n+3,n)
const sn = zeros(typeof(r),N_c+1)
const dsndr = zeros(typeof(r),N_c+1)
const dsndb = zeros(typeof(r),N_c+1)
if iseven(N_c)
  v_max = round(Int64,N_c/2)+2
else
  v_max = round(Int64,(N_c-1)/2)+2
end
const Iv = zeros(typeof(r),v_max+1)
const s2_grad = zeros(typeof(r),2)
const Jv = zeros(typeof(r),v_max+1)
const dIvdk = zeros(typeof(r),v_max+1)
const dJvdk = zeros(typeof(r),v_max+1)

#include("../../../src/transit_poly_prealloc.jl")
include("../../../src/transit_poly_struct.jl")
#include("../../../src/dJv_seriesdk.jl")
using PyPlot

# Now, run the same with transit_poly:

function profile_transit_poly()
nb = 1000001; b=zeros(nb)
flux = zeros(nb); fgrad = zeros(4)
flux_big = zeros(nb); fgrad_big = zeros(BigFloat,4)
flux_grad = zeros(nb,4)
flux_grad_big = zeros(nb,4)
tic()
for i=0:1000000
  b[i+1] = sqrt(((i-500000.0)/500000.0*(1.0+2.0*trans.r))^2)
  flux[i+1]= transit_poly!(r,b[i+1],u,fgrad)
  flux_grad[i+1,:] = fgrad
end
toc()
plot(b,flux)
println("b: ",minimum(b)," ",maximum(b))
println("f: ",minimum(flux)," ",maximum(flux))
return
end


using ApproxFun, SingularIntegralEquations, CUDA

# select a CUDA device
dev = CuDevice(0)

# create a context (like a process in CPU) on the selected device
ctx = create_context(dev)

# create the PTX module
run(`make CUDAlogkernel`)

# load the PTX module (each module can contain multiple kernel functions)
md = CuModule("CUDAlogkernel.ptx")

# retrieve the kernel function from the module
CUDAlogkernel = CuFunction(md, "CUDAlogkernel")

# generate data and load them to GPU
T = Float64
a,b = -1.,1.
u = [1.0,1/2,1/3,1/4,1/5]
n = 2^4
x = linspace(-3,3,n)
y = 0.0x+0.5
ret = Array(T,n)

@time begin
# load data arrays onto GPU
    @time gu = CuArray(u)
    @time gx = CuArray(x)
    @time gy = CuArray(y)

# create an array on GPU to store results
    @time gret = CuArray(T,n)

# run the kernel
# syntax: launch(kernel, grid_size, block_size, arguments)
# here, grid_size and block_size can be an integer or a tuple of integers
    @time launch(CUDAlogkernel, 1, n, (a, b, length(u), gu, gx, gy, gret))
    println("Second launch time: ")
    @time launch(CUDAlogkernel, 1, n, (a, b, length(u), gu, gx, gy, gret))

# download the results from GPU
    @time ret = to_host(gret)   # gret is a Julia array on CPU (host)
    println("Second retrieval time: ")
    @time to_host(gret)   # gret is a Julia array on CPU (host)

# release GPU memory
    free(gu)
    free(gx)
    free(gy)
    free(gret)

# finalize: unload module and destroy context
    unload(md)
    destroy(ctx)
end
# print the results

println("The logkernel results are: \n",ret)

sp = JacobiWeight(-.5,-.5,ChebyshevDirichlet{1,1}(Interval(a,b)))
f = Fun(u,sp)

println("The error is: ",norm(@time logkernel(f,complex(x,y))-ret))


# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

#Pkg.add("ReverseDiffSource")
#Pkg.update()
#Pkg.checkout("ReverseDiffSource","devl")
#Pkg.status()
using ReverseDiffSource

typealias Mat Matrix{Float64}

function function2expression(f::Function, sig::Tuple; order::Int=1, evalmod=Main, debug=false, allorders=true)
    fs = methods(f, map( typeof, sig ))
    length(fs) == 0 && error("ardiff: no function '$f' found for signature $sig")
    length(fs) > 1  && error("ardiff: several functions $f found for signature $sig")  # is that possible ?

    fdef  = fs[1].func.code
    fcode = Base.uncompressed_ast(fdef)
    fargs = fcode.args[1]  # function parameters

    cargs = [ (fargs[i], sig[i]) for i in 1:length(sig) ]
    expr = ReverseDiffSource.streamline(fcode.args[3])
    expr,cargs
end

function logsumexp(x::Mat)
  mx = maximum(x)
  log(sum(exp(x .- mx))) + mx
end

gmm_objective_code =
  :(begin
      n = size(x,2)
      d = size(x,1)
      CONSTANT = -n*d*0.5*log(2 * pi)

      sum_qs = similar(alphas)
      for ik in 1:k
        sum_qs[ik] = sum(icf[1:d,ik])
      end

      main_term = similar(alphas)
      slse = 0.
      for ix in 1:n
        for ik in 1:k
          Qxcentered = get_Q(d,icf[:,ik])*(x[:,ix] - means[:,ik])
          main_term[ik] = -0.5*sum(Qxcentered.^2)
        end
        slse += logsumexp(alphas + sum_qs + main_term)
      end
      return CONSTANT + slse - n*logsumexp(alphas)
  end)

k=3
d=2
n=3
icf_sz=div(d*(d + 1),2)
alphas = randn(1,k)
means = randn(d,k)
icf = randn(icf_sz,k)
x = randn(d,n)

logsumexp_code,cargs = function2expression(logsumexp, (alphas,))
logsumexp_d_code = rdiff(logsumexp_code; x=alphas, allorders=false)
@eval logsumexp_d(x) = $logsumexp_d_code
@deriv_rule logsumexp(x) x logsumexp_d(x)
@deriv_rule similar(x) x zeros(size(x))

function make_row(r::Int, diag_elem, L)
      elems=zeros(eltype(L),1,r-1)
      for i=1:r-1
        elems[i]=L[i]
      end
      return [elems diag_elem zeros(d-r)']
end

@deriv_rule make_row(r,diag_elem,L) r 0.
@deriv_rule make_row(r,diag_elem,L) diag_elem make_row(r,0.,ones(L)).*ds
@deriv_rule make_row(r,diag_elem,L) L make_row(r,1.,zeros(L)).*ds
make_row_code_d_code = rdiff(make_row,(1,2,icf[d+1:end,:]))

function ltri_unpack(D,LT)
  d=length(D)
  sum(D)+sum(LT)
  function make_row(r::Int, L)
    elems=zeros(eltype(L),1,r-1)
    for i=1:r-1
      elems[i]=L[i]
    end
    hcat(elems, D[r], zeros(1,d-r))
  end
  #row_start(r::Int) = div((r-1)*(r-2),2)
  #inds(r) = row_start(r)+(1:r-1)
  #vcat([ make_row(r, LT[inds(r)]) for r=1:d ]...)
end
ltri_unpack_code,cargs = function2expression(ltri_unpack, (icf[1:d,1],icf[d+1:end,1]))
ltri_unpack_d_code = rdiff(ltri_unpack_code; D=icf[1:d,1], LT=icf[d+1:end,1])
@eval logsumexp_d(x) = $ltri_unpack_d_code
@deriv_rule ltri_unpack(x) x ltri_unpack_d(x)

#gmm_objective_AD_code,cargs = function2expression(gmm_objective_AD, (alphas,x))
gmm_objective_AD_d_code = rdiff(gmm_objective_code;alphas=alphas,means=means,icf=icf)
@eval g(alphas,means,icf) = $gmm_objective_AD_d_code
g(alphas,means,icf)

include("common.jl")

# input should be 1 dimensional
function logsumexp(x)
  mx = maximum(x)
  log(sum(exp(x - mx))) + mx
end

@inbounds function gmm_objective(alphas,means,icf,x,wishart::Wishart)
  d = size(x,1)
  n = size(x,2)
  CONSTANT = -n*d*0.5*log(2 * pi)

  sum_qs = sum(icf[1:d,:],1)
  slse = sum(sum_qs)
  Qs = [get_Q(d,icf[:,ik]) for ik in 1:k]
  main_term = zeros(eltype(alphas),1,k)

  slse = 0.
  for ix=1:n
    for ik=1:k
      main_term[ik] = -0.5*sumabs2(Qs[ik] * (x[:,ix] - means[:,ik]))
    end
    slse += logsumexp(alphas + sum_qs + main_term)
  end

  CONSTANT + slse - n*logsumexp(alphas) + log_wishart_prior(wishart, sum_qs, Qs, icf)
end

# Read instance
dir_in = ARGS[1]
dir_out = ARGS[2]
fn = ARGS[3]
nruns_f = parse(Int,ARGS[4])
nruns_J = parse(Int,ARGS[5])
replicate_point = size(ARGS,1) >= 6 && ARGS[6] == "-rep"

fn_in = string(dir_in, fn)
fn_out = string(dir_out, fn)

alphas,means,icf,x,wishart = read_gmm_instance(string(fn_in,".txt"),replicate_point)
d = size(means,1)
k = size(means,2)
n = size(x,2)

# Objective
precompile(gmm_objective,(typeof(alphas),typeof(means),typeof(icf),typeof(x),Wishart))
tic()
for i in 1:nruns_f
  err = gmm_objective(alphas,means,icf,x,wishart)
end
tf = toq()/nruns_f;
@printf "tf: %g\n" tf
#@printf "err: %f\n" err

# Gradient helper
function wrapper_gmm_objective(packed)
  alphas,means,icf = unpack(d,k,packed)
  gmm_objective(alphas,means,icf,x,wishart)
end

# Gradient
g = ForwardDiff.gradient(wrapper_gmm_objective)
precompile(pack,(typeof(alphas),typeof(means),typeof(icf)))
precompile(g,(typeof(alphas),typeof(means),typeof(icf)))
J = zeros(1,length(alphas)+length(means)+length(icf))
tic()
for i in 1:nruns_J
  J = g(pack(alphas,means,icf))
end
tJ = toq()/nruns_J;
@printf "tJ: %g\n" tJ
#println("J:")
#println(J)

name = "Julia_F"

write_J(string(fn_out,"_J_",name,".txt"),J)
write_times(string(fn_out,"_times_",name,".txt"),tf,tJ)


# # Aerosol Activation Example

# ## Overview

# This example is based on AerosolActivation module which is a part of the
#   [CloudMicrophysics.jl](https://github.com/CliMA/CloudMicrophysics.jl) package.
# The AerosolActivation module computes the total number and mass
#   of aerosol particles that get activated and become cloud droplets,
#   given the atmospheric conditions and
#   the initial aerosol size distribution and properties.
# See the AerosolActivation module
#   [docs](https://clima.github.io/CloudMicrophysics.jl/stable/Aerosol_activation/)
#   for derivation and description of all input parameters.
#
# In this example we use the ensemble Kalman methods to learn
#   two parameters that describe the chemical composition of aerosol particles
#   based on the observed total number and mass of activated particles.
# The AerosolActivation model is used here in a "perfect model" setting,
#   meaning that the observations are generated by the same module
#   we are calibrating.

# ## Prerequisites

# The example depends on some standard Julia libraries,
#   as well as the CliMA packages:
#   [EnsembleKalmanProcess.jl](https://github.com/CliMA/EnsembleKalmanProcesses.jl),
#   [CLIMAParameters.jl](https://github.com/CliMA/CLIMAParameters.jl) and
#   [CloudMicrophysics.jl](https://github.com/CliMA/CloudMicrophysics.jl).
# To ensure that all the dependencies are met
#   start Julia using the Project.toml file provided in the example
#   and run the Julia package manager to download all the dependecies.

# ## Example

# We begin by importing some standard Julia modules,
# the Ensemble Kalman Process modules,
# CLIMA Parameter modules and Aerosol Activation modules.
using Plots
using Distributions
using LinearAlgebra
using Random
rng_seed = 44

import EnsembleKalmanProcesses
const EKP = EnsembleKalmanProcesses

import CLIMAParameters
const CP = CLIMAParameters
struct EarthParameterSet <: CP.AbstractEarthParameterSet end
const param_set = EarthParameterSet()

import CloudMicrophysics
const AM = CloudMicrophysics.AerosolModel
const AA = CloudMicrophysics.AerosolActivation

import Thermodynamics
const TD = Thermodynamics
nothing # hide

# Next, we provide the information about the priors of the parameters we want to learn.
# We are calibrating two parameters decribing the aerosol properties -
# namely the aerosol molar mass and the osmotic coefficient.
parameter_names = ["molar_mass", "osmotic_coeff"]
nothing # hide

# In this test we do know the parameter values.
# We use them to test the convergence of EKP for aerosol activation.
molar_mass_true = 0.058443
osmotic_coeff_true = 0.9
default_params = [molar_mass_true, osmotic_coeff_true]
nothing # hide

# Both parameters have to be positive definite, therefore we
# define the constraints to be bounded below by zero.
c1 = EKP.ParameterDistributions.bounded_below(0.0)
c2 = EKP.ParameterDistributions.bounded_below(0.0)
nothing # hide

# We don't have any other prior knowledge about the calibrated parameters.
# Defining the prior distribution in the constrained space to be normal
# with the mean equal to zero and the standard deviation equal to one,
# ensures that the default parameter values are well within the
# assumed prior pdf.
d1 = EKP.ParameterDistributions.Parameterized(Distributions.Normal(0, 1))
d2 = EKP.ParameterDistributions.Parameterized(Distributions.Normal(0, 1))
nothing # hide

# This concludes the setup of priors.
prior1 = Dict("distribution" => d1, "constraint" => c1, "name" => parameter_names[1])
prior2 = Dict("distribution" => d2, "constraint" => c2, "name" => parameter_names[2])
priors = EKP.ParameterDistributions.ParameterDistribution([prior1, prior2])
nothing # hide

# Next we define the atmospheric conditions for which the calibration will take place,
# (air temperature in K, air pressure in Pa vertical velocity in m/s
# and vapor specific humidity assuming its saturated in kg/kg)
# This can be changed later to include more than one ``(T, p, w)`` combination in the calibration process
T = 283.15
p = 1e5
w = 5.0
p_vs = TD.saturation_vapor_pressure(param_set, T, TD.Liquid())
q_vs = 1 / (1 - CP.Planet.molmass_ratio(param_set) * (p_vs - p) / p_vs)
q = TD.PhasePartition(q_vs, 0.0, 0.0)
nothing # hide

# We also define the aerosol size distribution (lognormal, 1 mode)
# with (mean radius in m, geometric stdev, number concentration 1/m³).
# These can also be changed later to include different initial size distributions.
r_dry = 0.243e-6
stdev = 1.4
N = 100.0 * 1e6 # since 1/cm³ = 1e6 1/m³
nothing # hide

# Finally, we define additional parameters that describe the aerosol properties.
# The chosen aerosol is sea salt.
dissoc_seasalt = 2.0
soluble_mass_frac_seasalt = 1.0
rho_seasalt = 2170.0;
nothing # hide

# We define a wrapper function that runs the aerosol activation module
# with two input parameters that will be calibrated by EKP.
# The output observations are the number and mass of activated aerosol.
function run_activation_model(molar_mass_calibrated, osmotic_coeff_calibrated)

    accum_mode_seasalt = AM.Mode_B(
        r_dry,
        stdev,
        N,
        (1.0,),
        (soluble_mass_frac_seasalt,),
        (osmotic_coeff_calibrated,),
        (molar_mass_calibrated,),
        (dissoc_seasalt,),
        (rho_seasalt,),
        1,
    )

    aerosol_distr = AM.AerosolDistribution((accum_mode_seasalt,))
    N_act = AA.total_N_activated(param_set, aerosol_distr, T, p, w, q)
    M_act = AA.total_M_activated(param_set, aerosol_distr, T, p, w, q)
    return [N_act, M_act]
end
nothing # hide

# This example is run in a "perfect model setting", meaning the model that we
# calibrate is also used to generate observations.
# We use the total number and mass of activated aerosol particles as
# our observational data.
observation_data_names = ["N_act", "M_act"];
nothing # hide

# We generate artificial truth samples based on the default values
# of parameters we are calibrating.
n_samples = 10
G_t = run_activation_model(molar_mass_true, osmotic_coeff_true)
y_t = zeros(length(G_t), n_samples)

Γy = convert(Array, LinearAlgebra.Diagonal([0.01 * G_t[1], 0.01 * G_t[2]]))
μ = zeros(length(G_t));
nothing # hide

# And add noise to the generated truth sample.
for i in 1:n_samples
    y_t[:, i] = G_t .+ rand(Distributions.MvNormal(μ, Γy))
end

truth_array = EKP.Observations.Observation(y_t, Γy, observation_data_names)
nothing # hide

# One could try for the truth to be a mean of the generated array.
# Or do the calibration for all individual truth samples and then
# compute the mean of calibrated parameters.
# For now we are just taking one truth array member.
truth_sample = truth_array.samples[1];
nothing # hide

# We use 50 ensemble members and do 10 iterations.
N_ens = 50
N_iter = 10

initial_par = EKP.construct_initial_ensemble(priors, N_ens; rng_seed)
ekiobj = EKP.EnsembleKalmanProcess(initial_par, truth_sample, truth_array.obs_noise_cov, EKP.Inversion(), Δt = 1)
nothing # hide

# Finally, we can run the Ensemble Kalman Process calibration.
ϕ_n_values = []
for n in 1:N_iter
    θ_n = EKP.get_u_final(ekiobj)

    ϕ_n = mapslices(x -> EKP.ParameterDistributions.transform_unconstrained_to_constrained(priors, x), θ_n; dims = 1)

    G_n = [run_activation_model(ϕ_n[:, i]...) for i in 1:N_ens]

    G_ens = hcat(G_n...)
    EKP.update_ensemble!(ekiobj, G_ens)

    global ϕ_n_values = vcat(ϕ_n_values, [ϕ_n])
end
nothing # hide

# We define some simple functions for plotting the data.
function plot_ensemble_scatter(id)

    ensemble_member = 1:N_ens

    if id == 1
        ylabel = "Molar mass [kg/mol]"
        filename = "molar_mass_scatter.svg"
    end
    if id == 2
        ylabel = "Osmotic coefficient [-]"
        filename = "osmotic_coeff_scatter.svg"
    end

    plot(
        ensemble_member,
        ϕ_n_values[1][id, 1:N_ens],
        seriestype = :scatter,
        xlabel = "Ensemble Number",
        ylabel = ylabel,
        legend = false,
    )

    for it in 2:N_iter
        plot!(ensemble_member .+ ((it - 1) * 50), ϕ_n_values[it][id, 1:N_ens], seriestype = :scatter, legend = false)
    end

    current()
    savefig(filename)
end

function plot_ensemble_means(id)

    number_of_iters = 1:N_iter
    means = zeros(N_iter)

    for it in 1:N_iter
        means[it] = mean(ϕ_n_values[it][id, 1:N_ens])
    end

    if id == 1
        ylabel = "Molar mass [kg/mol]"
        filename = "molar_mass_average.svg"
    end
    if id == 2
        ylabel = "Osmotic coefficient [-]"
        filename = "osmotic_coeff_average.svg"
    end

    plot(
        number_of_iters,
        means,
        markershape = :star5,
        xticks = number_of_iters,
        xlabel = "Iteration Number",
        ylabel = ylabel,
        label = "Ensemble Mean",
    )
    hline!([default_params[id]], label = "true value")

    savefig(filename)
end
nothing # hide

# We plot the ensemble members and the ensemble mean for the
# molar mass and osmotic coefficient.
plot_ensemble_scatter(1)
plot_ensemble_means(1)
plot_ensemble_scatter(2)
plot_ensemble_means(2)
# ![](molar_mass_scatter.svg)
# ![](osmotic_coeff_scatter.svg)
# ![](molar_mass_average.svg)
# ![](osmotic_coeff_average.svg)

# Finally, we test that the parameter values obtained via EnsembleKalmanProcesses.jl are close
# to the known true parameter values.
molar_mass_ekp = round(mean(ϕ_n_values[N_iter][1, 1:N_ens]), digits = 6)
osmotic_coeff_ekp = round(mean(ϕ_n_values[N_iter][2, 1:N_ens]), digits = 6)

println("Molar mass [kg/mol]: ", molar_mass_ekp, " vs ", molar_mass_true)
println("Osmotic coefficient [-]: ", osmotic_coeff_ekp, " vs ", osmotic_coeff_true)


### A Pluto.jl notebook ###
# v0.12.18

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 2069c2c0-01b4-11eb-3cc9-a7c60467a92a
# Load packages
begin
	using Pkg
	Pkg.activate(mktempdir())
	Pkg.add("Plots")
	Pkg.develop(path="/home/sari/.julia/dev/FractalTools/")
	using Plots
	plotly()
	using FractalTools
end

# ╔═╡ 24b8a340-01b5-11eb-0551-318932fe4a10
@bind d html"<input type=range min=0 max=1 step=1e-3>"

# ╔═╡ 2caf47f8-047f-11eb-1f15-db8e661dc9ce
@bind niter html"<input type=range min=1 max=30 step=1>"

# ╔═╡ 69dbf05e-047f-11eb-3d6e-892350bec24f
niter, d

# ╔═╡ 402f7b62-01b5-11eb-2789-7bdb04c0d694
begin
	# Consrucct data
	f(x) = sin(2π*x)
	g(x) = cos(2π*x)
	xi, dx, xf = 0, 0.1, 1
	x = collect(xi : dx : xf)
	y = f.(x)
	z = g.(x)
	dvar = d * ones(length(x) - 1)
	h = d * ones(length(x) - 1)
	l = d * ones(length(x) - 1)
	m = d * ones(length(x) - 1)

	# Interplate
	interp = hiddenfif(x, y, z, d=dvar, h=h, l=l, m=m, niter=niter)

	# Plots
	plt = plot()
	xd = collect(xi : 0.1dx : xf)
	yd = interp.(xd)
	plot!(xd, getindex.(yd, 1))
	scatter!(x, y, marker=(:circle, 3))
end

# ╔═╡ Cell order:
# ╠═2069c2c0-01b4-11eb-3cc9-a7c60467a92a
# ╠═24b8a340-01b5-11eb-0551-318932fe4a10
# ╠═2caf47f8-047f-11eb-1f15-db8e661dc9ce
# ╠═69dbf05e-047f-11eb-3d6e-892350bec24f
# ╠═402f7b62-01b5-11eb-2789-7bdb04c0d694


abstract type MultiClassLikelihood{T<:Real} <: Likelihood{T} end

""" Return the labels in a vector of vectors for multiple outputs"""
function treat_labels!(y::AbstractArray{T,N},likelihood::L) where {T,N,L<:MultiClassLikelihood}
    @assert N <= 1 "Target should be a vector of labels"
    likelihood = init_multiclass_likelihood(likelihood,y)
    return likelihood.Y,length(likelihood.Y),likelihood
end

function init_multiclass_likelihood(likelihood::L,y::AbstractVector) where {L<:MultiClassLikelihood}
    L(one_of_K_mapping(y)...)
end

""" Given the labels, return one hot encoding, and the mapping of each class """
function one_of_K_mapping(y)
    y_values = unique(y)
    Y = [falses(length(y)) for i in 1:length(y_values)]
    y_class = zeros(Int64,length(y))
    for i in 1:length(y)
        for j in 1:length(y_values)
            if y[i]==y_values[j]
                Y[j][i] = true;
                y_class[i] = j;
                break;
            end
        end
    end
    ind_values = Dict(value => key for (key,value) in enumerate(y_values))
    return Y,y_values,ind_values,y_class
end

function compute_proba(l::MultiClassLikelihood{T},μ::AbstractVector{<:AbstractVector{T}},σ²::AbstractVector{<:AbstractVector{T}},nSamples::Integer=200) where {T<:Real}
    K = length(μ)
    n = length(μ[1])
    μ = hcat(μ...)
    μ = [μ[i,:] for i in 1:n]
    σ² = hcat(σ²...)
    σ² = [σ²[i,:] for i in 1:n]
    pred = zeros(T,n,K)
    for i in 1:n
            p = MvNormal(μ[i],sqrt.(abs.(σ²[i])))
            # p = MvNormal(μ[i],sqrt.(max.(eps(T),σ²[i])))
            for _ in 1:nSamples
                pred[i,:] += pdf(l,rand(p))/nSamples
            end
    end
    return DataFrame(pred,Symbol.(l.class_mapping))
end

function expecLogLikelihood(model::AbstractGP{<:MultiClassLikelihood,<:NumericalVI})
    compute_log_expectations(model)
end

include("softmax.jl")
include("logisticsoftmax.jl")


"""
    MOSUM(data, G[, var_est_method="mosum", alpha=0.1, criterion="eta", eta=0.4, epsilon=0.2])

Runs the MOSUM procedure for the univariate array `data` with bandwidth `G`, and returns the number and position of found changepoints.

Optionally, `var_est_method` specifies the variance estimator to normalise by; this can be the average `mosum` (default) or minimum `mosum.min` across windows.
`alpha` determines the signicance level (default 0.1).
`criterion` determines whether to use the `eta` (default) or `epsilon` location procedure (see references).
`eta` and `epsilon` are tuning parameters for the mentioned procedures (default 0.4 and 0.2).

See also: [`@BS`](@ref), [`MOSUM_multi_scale`](@ref), [`mosum_plot`](@ref)

# Returns
* `out::Dict{String,Union{Bool, Float64, Int64, Array}}`:
  * `out["number"]` : Integer number of detected changes
  * `out["Reject"]` : Boolean hypothesis test outcome
  * `out["threshold"]` : Float of testing threshold
  * `out["detector"]` : vector of mosum detector values
  * `out["var.estimation"]` : vector of estimated variance values

# Example
```julia-repl
# Sample Normal time series with changing mean
n = 1000
λ = 300
μ, σ = Normal(0.0, 10.0), 1.0
data, cps = @changepoint_sampler n λ Normal(μ, σ)
G = 150
# Run MOSUM procedure
MOSUM_out = MOSUM(data, G)
# Plot MOSUM detector
mosum_plot(MOSUM_out)
```

# References
Eichinger, Birte, and Claudia Kirch. "A MOSUM procedure for the estimation of multiple random change points." Bernoulli 24.1 (2018): 526-564.
Meier, Aledataander, Claudia Kirch, and Haeran Cho. "mosum: A package for moving sums in change point analysis." (2018).
"""
function MOSUM( data::Array{Float64}, G::Int64; var_est_method="mosum", alpha=0.1, criterion="eta", eta=0.4, epsilon=0.2)
    var_est_method in ("mosum", "mosum.min") || error("Keyword argument var_est_method must be either \"mosum\" or \"mosum.min\"")

    CP = Array{Int64}(undef,0)
    q = 0
    Reject = false
    n = length(data)

    csum = cumsum(data)
    sqsum = cumsum(data.^2)
    variance = zeros(Float64, n)
    lvar = zeros(Float64, n)
    rvar = zeros(Float64, n)
    lsqsums = zeros(Float64, n)
    rsqsums =  zeros(Float64, n)
    detector = zeros(Float64, n)
    #jump_size = zeros(Float64, n)
    lsums = zeros(Float64, n)
    rsums = zeros(Float64, n)

    #populate
    for t in 1:(n-G)
        rsums[t] = csum[t+G] - csum[t]
        lsums[t+G] = rsums[t]

        rsqsums[t]  = sqsum[t+G] - sqsum[t]
        lsqsums[t+G] = rsqsums[t]

        rlmean = rsums[t]/G
        rvar[t] = rsqsums[t]/G -  (1/G - 2)*(rlmean^2)
        lvar[t+G] = rvar[t]
    end

    # normalise
    for t in (G+1):(n-G)
        if var_est_method == "mosum"
            variance[t] = (rvar[t] + lvar[t])/ 2
        elseif var_est_method == "mosum.min"
            variance[t] = min( rvar[t], lvar[t])
        end
        #jump_size[t] =
        detector[t] = abs(rsums[t] - lsums[t])/ sqrt(2*G*variance[t])
    end

    #threshold
    a = sqrt(2*n/G); b = 2log(n/G) + log(log(n/G))/2 + log(3/2) - log(pi)/2
    c = - log(log( (1- alpha)^(-.5) ))
    D = (b+c)/a

    mosum_stat, k = findmax(detector)

    if mosum_stat > D
        got_cps = get_cps_mosum(detector, D, G, criterion, eta, epsilon)
        CP = got_cps[1]
        q = got_cps[2]
        Reject = got_cps[3]
    end

#    cost = 0.0
#    CP = sort(CP)
#    for j in 1:(length(CP)-1)
#        cost = cost + segment_cost(CP[j]+1,CP[j+1]) + pen
#    end
#    cost = cost + segment_cost(CP[end]+1,n)
    # organise output into a dictionary
    out = Dict{String, Union{Array, Bool, Int64, Float64} }()
    out["number"] = q
    out["detector"] = detector
    out["Reject"] = Reject
    out["changepoints"] = CP
    out["threshold"] = D
    out["var.estimation"] = variance
    #multi-scale
    #if(return_jump_size) out["jump_size"] = jump_size
    return out
end

function get_cps_mosum(stat::Array{Float64} , D_n::Float64, G::Int64,  criterion::String="epsilon",
                       eta::Float64 = 0.4, epsilon::Float64 = 0.2 )
    cps = Array{Int64}(undef,0) ##assign empty cps
    Reject = false
    if criterion == "epsilon"
        n = length(stat)
        rshift = stat[2:end]
        push!(rshift, 0.0)
        lshift = stat[1:(end-1)]
        pushfirst!(lshift, 0.0)
        over =  stat .> D_n #findall(data -> data > D_n , stat ) #indices greater than D_n
        vv = over .& (lshift .< D_n )#& lshift >0 ) #lowers
        v = findall(vv .== 1)
        ww = over .& (rshift .< D_n)
        w = findall(ww .== 1)
        nu_remove = findall(w-v .>= epsilon * G) #nu(epsilon) test for distance between
        v_nu = v[nu_remove]; w_nu = w[nu_remove] #c(w[nu_remove],n)
        sub_pairs = [v_nu w_nu]

        q = size(sub_pairs)[1]
        if q>0
            for ii in 1:q
                interval = sub_pairs[ii,1]:sub_pairs[ii,2]
                kk = findmadata(stat[interval])[2] #internal cp location
                push!(cps, kk + 1 + sub_pairs[ii,1]) #- G-p
            end
        end
    end

    if criterion == "eta"
        n = length(stat)
        window = Int(floor(eta*G))
        for t in (G+1):(n-G)
            if (stat[t] > D_n) & (stat[t] == maximum(stat[(t .- window):(t .+ window)]) )
                push!(cps, t) ##add to cps
            end
        end
    end

    q = length(cps)
    if q>0
        Reject = true
    end

    return cps, q, Reject
end

"""
    MOSUM_multi_scale(data, Gset[, var_est_method="mosum", alpha=0.1, criterion="eta", eta=0.4, epsilon=0.2])

Runs the Multiple Filtre MOSUM procedure for the univariate array `data` with multiple bandwidths `Gset`, and returns the position of found changepoints.
For optional arguments, see `?MOSUM`.

See also: [`MOSUM`](@ref)

# Returns
* `cps` : Integer array of estimated change points

# Edataample
```julia-repl
# Sample Normal time series with changing mean
n = 1000
λ = 300
μ, σ = Normal(0.0, 10.0), 1.0
data, cps = @changepoint_sampler n λ Normal(μ, σ)
Gset = [25, 50, 100, 150]
# Run MOSUM procedure
multi_scale_out = MOSUM_multi_scale(data, Gset; alpha = 0.05)
# Plot change points
changepoint_plot(data, multi_scale_out)
```

# References
Messer M, Kirchner M, Schiemann J, Roeper J, Neininger R, Schneider G (2014). “A Multiple Filter Test for the Detection of Rate Changes in Renewal Processes with Varying Variance.” The Annals of Applied Statistics, 8(4), 2027–2067.
"""
function MOSUM_multi_scale(data::Array{Float64}, Gset::Array{Int64}; kwargs...)

    sort!(Gset)
    cps = Array{Int64}[]
    for G in Gset
        mosum_G = MOSUM(data, G; kwargs...)
        cps_G = mosum_G["changepoints"]
        if G == Gset[1]
            cps = cps_G
        else
            for k in cps_G
                min_dist = minimum(abs.(cps .- k))
                if min_dist >= G
                    push!(cps, k)
                end
            end # for
        end
    end

#    #step 1
#    PairSet = Array{Union{Float64, Int64}}[]
#    for G in Gset
#        mosum_G = MOSUM(data, G; default_args)
#        cps = mosum_G["changepoints"]
#        for k in cps
#            h = mosum_G["detector"][k]
#            push!(PairSet, [h, k, G])
#        end # for
#    end
#    sort!(Pairset)
#    #step 2
#    D = Array{Union{Float64, Int64}}[] #conflicting
#    Cset = PairSet
#    while length(PairSet) > 0
#        candidate = pop!(Pairset) #current max jump
#        for i in 1:length(PairSet)
#            if  abs(PairSet[i,2] -. candidate[2]) < min(candidate[3], PairSet[i,3])
#                take = popat(PairSet, i)
#                push!(D, take)
#            end
#        end
#    end
    return(cps)
    #2
end


module Models

# Import all QuantEcon names
using QuantEcon

# 3rd party dependencies
using Distributions
using Optim: optimize
using Grid: CoordInterpGrid, BCnan, BCnearest, InterpLinear

export
# types
    ArellanoEconomy,
    AssetPrices,
    CareerWorkerProblem,
    ConsumerProblem,
    JvWorker,
    LucasTree,
    SearchProblem,
    GrowthModel,
    UncertaintyTrapEcon,

# functions
    vfi!,                                                # arellano_vfi
    tree_price, consol_price, call_option,               # asset_pricing
    get_greedy, get_greedy!,                             # career, odu, optgrowth
    coleman_operator, coleman_operator!, init_values,    # ifp
    compute_lt_price, lucas_operator,                    # lucastree
    res_wage_operator, res_wage_operator!,               # odu
    bellman_operator, bellman_operator!,                 # many
    psi, update_beliefs!, update_theta!, gen_aggregates  # uncertainty_traps

____bellman_main_docstring = """
Apply the Bellman operator for a given model and initial value
"""

____greedy_main_docstring = """
Extract the greedy policy (policy function) of the model
"""

____see_methods_docstring = """
See the specific methods of the mutating function for more details on arguments
"""

____mutate_last_positional_docstring = """
The last positional argument passed to this function will be over-written
"""

____kwarg_note = """
There is also a version of this function that accepts keyword arguments for
each parameter
"""

include("arellano_vfi.jl")
include("asset_pricing.jl")
include("career.jl")
include("ifp.jl")
include("jv.jl")
include("lucastree.jl")
include("odu.jl")
include("optgrowth.jl")
include("uncertainty_traps.jl")

"""
$(____bellman_main_docstring). $(____see_methods_docstring)
"""
bellman_operator

"""
$(____bellman_main_docstring). $(____see_methods_docstring)

$(____mutate_last_positional_docstring)
"""
bellman_operator!

"""
$(____greedy_main_docstring). $(____see_methods_docstring)
"""
get_greedy

"""
$(____greedy_main_docstring). $(____see_methods_docstring)

$(____mutate_last_positional_docstring)
"""
get_greedy!

end  # module


using QuantileRegressions
using DataFrames, CSV, LinearAlgebra
using Test

include("engel_test.jl")
include("npqreg.jl")