abstract type UncertainElement{C, F} end

struct δ{C, F<:Real} <: UncertainElement{C, F}
    val::C
    radius::F
    name::Symbol
    # function δr(val = 0, interval = Interval(val-1, val+1))
    #     lower ≤ val ≤ upper || throw(ArgumentError("val must be between lower and upper"))
    #     val, lower, upper = promote(val, lower, upper)
    #     new{typeof(val)}(val, lower, upper)
    # end
    function δ(val, radius, name = gensym("δ"))
        radius ≥ 0 || throw(ArgumentError("radius must be positive"))
        new{typeof(val), typeof(radius)}(val, radius, name)
    end
end

Base.zero(::Type{δ{C, F}}) where {C,F} = δ(zero(C), zero(F))

"""
    δr(val::Real = 0.0, radius::Real = 1.0, name)

Create a real, uncertain parameter. If no name is given, a boring name will be generated automatically.
"""
function δr(val::Real = 0.0, radius::Real = 1.0, args...)
    radius ≥ 0 || throw(ArgumentError("radius must be positive"))
    val, radius = promote(val, radius)
    δ(val, radius, args...)
end

"""
    δc(val::Complex = complex(0.0), radius::Real = 1.0, name)

Create a complex, uncertain parameter. If no name is given, a boring name will be generated automatically.
"""
function δc(val = complex(0.0), radius=1.0, args...)
    radius ≥ 0 || throw(ArgumentError("radius must be positive"))
    val = complex(val)
    δ(val, radius, args...)
end

Base.size(::δ) = (1,)
Base.size(::δ, n) = 1
Base.length(::δ) = 1
Base.eltype(::UncertainElement{C}) where C = C
Base.eltype(::Type{<:UncertainElement{C}}) where C = C

function Base.promote_rule(::Type{N}, ::Type{δ{C, F}}) where {N <: Number, F, C}
    δ{promote_type(N, C), promote_type(N, F)}
end

Base.convert(::Type{δ{C, F}}, n::Number) where {C,F} = δ(C(n), zero(F))
Base.convert(::Type{δ{C, F}}, d::δ) where {C,F} = δ(C(d.val), F(d.radius), d.name)

#  + and - we define here, all other ops generate ExtendedStateSpace. THis is to avoid creating a complicated tagging mechanism
function Base.:(+)(d1::δ, d2::δ)
    if d1.name === d2.name
        δ(+(d1.val, d2.val), +(d1.radius, d2.radius), d1.name)
    else
        uss(d1) + uss(d2)
    end
end

function Base.:(+)(n::Number, d2::δ)
    δ(+(n, d2.val), d2.radius, d2.name)
end

Base.:(+)(d2::δ, n::Number) = +(n, d2)

function Base.:(*)(n::Number, d2::δ)
    δ(*(n, d2.val), abs(n)*d2.radius, d2.name)
end

function Base.:(-)(d2::δ)
    δ(-d2.val, d2.radius, d2.name)
end

function Base.:(-)(d1::δ, d2::δ)
    # δr(d1.val+d2.val, d1.lower+d2.lower, d1.upper+d2.upper)
    if d1.name === d2.name
        δ(-(d1.val, d2.val), -(d1.radius, d2.radius), d1.name)
    else
        # δ(-(d1.val, d2.val), +(d1.radius, d2.radius))
        uss(d1) + (-uss(d2))
    end
end

function Base.:(*)(d1::δ, d2::δ)
    uss(d1)*uss(d2)
end

function Base.:(/)(d1::δ, d2::δ)
    uss(d1)*inv(uss(d2))
end

function Base.:(/)(d1::Number, d2::δ)
    d1*inv(uss(d2))
end

function Base.inv(d2::δ)
    inv(uss(d2))
end

"""
    UncertainSS{TE} <: AbstractStateSpace{TE}

Represents LFT_u(M, Diagonal(Δ))
"""
struct UncertainSS{TE} <: AbstractStateSpace{TE}
    sys::ExtendedStateSpace{TE}
    Δ
end

function Base.getproperty(sys::UncertainSS, s::Symbol)
    s ∈ fieldnames(typeof(sys)) && return getfield(sys, s)
    if s === :nz
        # return foldl((l, d)->l + size(d, 1), sys.Δ, init=0)
        return size(sys.C1, 1)
    elseif s === :nw
        # return foldl((l, d)->l + size(d, 2), sys.Δ, init=0)
        return size(sys.B1, 2)
    elseif s === :zinds
        return 1:sys.nz
    elseif s === :yinds
        return sys.nz .+ (1:size(sys.C2, 1))
    elseif s === :winds
        return 1:sys.nw
    elseif s === :uinds
        return sys.nw .+ (1:size(sys.B2, 2))
    elseif s ===:M
        return sminreal(performance_mapping(sys.sys))
        # return sys.sys
    elseif s ===:delta
        return sys.Δ
    end
    getproperty(getfield(sys, :sys), s)
end

for f in [:system_mapping, :performance_mapping, :ss]
    @eval $f(s::UncertainSS, args...; kwargs...) = $f(s.sys, args...; kwargs...)
end

function Base.promote_rule(::Type{UncertainSS{TE}}, ::Type{δ{C, F}}) where {TE, F, C}
    UncertainSS{TE}
end

function Base.promote_rule(::Type{UncertainSS{TE}}, ::Type{L}) where {TE, L <: LTISystem}
    UncertainSS{TE}
end

function Base.promote_rule(::Type{UncertainSS{TE}}, ::Type{L}) where {TE, L <: AbstractArray}
    UncertainSS{TE}
end

function Base.convert(::Type{U}, d::δ) where {U <: UncertainSS}
    uss(d)
end

function Base.convert(::Type{U}, d::AbstractArray) where {U <: UncertainSS}
    uss(d)
end

function Base.convert(::Type{U}, s::LTISystem) where {U <: UncertainSS}
    s isa UncertainSS && return s
    sys = ss(s)
    esys = partition(sys, 0, 0)
    UncertainSS(esys, [])
end

Base.:+(sys::UncertainSS{TE}, n::δ) where TE <: ControlSystems.TimeEvolution = +(sys, uss(n))
Base.:+(n::δ, sys::UncertainSS{TE}) where TE <: ControlSystems.TimeEvolution = +(uss(n), sys)

Base.:+(sys::UncertainSS{TE}, n::Number) where TE <: ControlSystems.TimeEvolution = +(sys, uss(n))
Base.:+(n::Number, sys::UncertainSS{TE}) where TE <: ControlSystems.TimeEvolution = +(uss(n), sys)

Base.:*(G::LTISystem, d::UncertainElement) = uss(ss(G)) * uss(d)
Base.:*(d::UncertainElement, G::LTISystem) = uss(d) * uss(ss(G))

"""
    uss(D11, D12, D21, D22, Δ, Ts = nothing)

Create an uncertain statespace object with only gin matrices.
"""
function uss(D11, D12, D21, D22, Δ, Ts=nothing)
    D11, D12, D21, D22 = vcat.((D11, D12, D21, D22))
    sys = ExtendedStateSpace(zeros(0,0), zeros(0,size(D21,2)), zeros(0,size(D12,2)), zeros(size(D12,1),0), zeros(size(D21,1),0), D11, D12, D21, D22, Ts)
    #                                       B1                   B2                    C1                        C2
    UncertainSS(sys, Δ)
end

"""
    uss(d::δ{C, F}, Ts = nothing)

Convert a δ object to an UncertainSS
"""
function uss(d::δ{C,F}, Ts = nothing) where {C,F}
    # sys = partition(ss([d.val d.radius; 1 0]), u=2, y=2, w=1, z=1)
    uss(0, 1, d.radius, d.val, [normalize(d)], Ts)
end

"""
    uss(d::AbstractVector{<:δ}, Ts = nothing)

Create a diagonal uncertain statespace object with the uncertain elements `d` on the diagonal.
"""
function uss(d::AbstractVector{<:δ}, Ts = nothing)
    vals = getfield.(d, :val)
    rads = getfield.(d, :radius)
    nonzero = rads .> 0
    # vals = vals
    # rads = rads
    nd = count(nonzero)
    n = length(d)
    uss(0I(nd), I(n)[nonzero, :], diagm(rads)[:, nonzero], diagm(vals), normalize.(d[nonzero]), Ts)
end

uss(n::Number, Ts=nothing) = uss(zeros(0,0), zeros(0,1), zeros(1,0), 1, [], Ts)

"""
    uss(D::AbstractArray, Δ, Ts = nothing)

If only a single `D` matrix is provided, it's treated as `D11` if Δ is given, and as `D22` if no Δ is provided.
"""
function uss(D::AbstractArray, Δ=[], Ts=nothing)
    if length(Δ) == size(D,1)
        uss(D, zeros(size(D,1),0), zeros(0,size(D,2)), zeros(0,0), Δ, Ts)
    elseif isempty(Δ)
        uss(zeros(0,0), zeros(0,size(D,2)), zeros(size(D,1),0), D, Δ, Ts)
    else
        throw(DimensionMismatch("length(Δ) != size(D,1)"))
    end
end

uss(s::UncertainSS) = s
uss(s::AbstractStateSpace) = convert(UncertainSS, s)

# Lower LFT
# function Base.inv(s::UncertainSS)
#     s.nx == 0 || error("Haven't considered this case yet")
#     @unpack D11, D12, D21, D22 = s
#     iD11 = inv(D11)
#     sys = uss(iD11, -iD11*D12, D21*iD11, D22-D21*iD11*D12, s.Δ, s.timeevol)
# end

function Base.inv(s::UncertainSS) # Upper LFT
    s.nx == 0 || error("Haven't considered this case yet")
    @unpack D11, D12, D21, D22 = s
    iD22 = inv(D22)
    sys = uss(D11-D12*iD22*D21, D12*iD22, -iD22*D21, iD22, s.Δ, s.timeevol)
end

function Base.:*(s1::UncertainSS, s2::UncertainSS)
    sys1 = s1.sys
    sys2 = s2.sys
    sys = invert_mappings(invert_mappings(sys1)*invert_mappings(sys2))
    UncertainSS(sys, [s1.Δ; s2.Δ])
end

function Base.:*(s1::UncertainSS, n::Number)  # reverse case is handled by switching
    sys = invert_mappings(n*invert_mappings(s1.sys))
    UncertainSS(sys, s1.Δ)
end

function Base.:+(s1::UncertainSS, s2::UncertainSS)
    sys1 = s1.sys
    sys2 = s2.sys
    sys = invert_mappings(invert_mappings(sys1)+invert_mappings(sys2))
    UncertainSS(sys, [s1.Δ; s2.Δ])
end

function Base.:-(s::UncertainSS)
    UncertainSS(-s.sys, s.Δ)
end

function Base.hcat(systems::UncertainSS...)
    # Perform checks
    timeevol = common_timeevol(systems...)

    ny = systems[1].ny
    if !all(s.ny == ny for s in systems)
        error("All systems must have same ouput dimension")
    end

    A = blockdiag([s.A for s in systems]...)

    B1 = blockdiag([s.B1 for s in systems]...)
    B2 = blockdiag([s.B2 for s in systems]...)

    C2 = reduce(hcat, [s.C2 for s in systems])
    C1 = blockdiag([s.C1 for s in systems]...)

    D21 = reduce(hcat, [s.D21 for s in systems])
    D22 = reduce(hcat, [s.D22 for s in systems])
    D11 = blockdiag([s.D11 for s in systems]...)
    D12 = blockdiag([s.D12 for s in systems]...)

    sysnew = ss(A, B1, B2, C1, C2, D11, D12, D21, D22, timeevol)
    return UncertainSS(sysnew, reduce(vcat, [s.Δ for s in systems]))
end

function Base.vcat(systems::UncertainSS...)
    # Perform checks
    timeevol = common_timeevol(systems...)

    nu = systems[1].nu
    if !all(s.nu == nu for s in systems)
        error("All systems must have same input dimension")
    end

    A = blockdiag([s.A for s in systems]...)

    B2 = reduce(vcat, [s.B2 for s in systems])
    B1 = blockdiag([s.B1 for s in systems]...)

    C1 = blockdiag([s.C1 for s in systems]...)
    C2 = blockdiag([s.C2 for s in systems]...)

    D12 = reduce(vcat, [s.D12 for s in systems])
    D22 = reduce(vcat, [s.D22 for s in systems])
    D11 = blockdiag([s.D11 for s in systems]...)
    D21 = blockdiag([s.D21 for s in systems]...)

    sysnew = ss(A, B1, B2, C1, C2, D11, D12, D21, D22, timeevol)
    return UncertainSS(sysnew, reduce(vcat, [s.Δ for s in systems]))
end

function ControlSystems.lft(G::UncertainSS, Δ::AbstractArray=G.Δ, type=:u)
    if ndims(Δ) == 1
        Δ = append(ss.(Δ)...)
    end

    if type !== :u
        error("Invalid type of lft ($type), specify type=:u")
    else
        lft(ss(G.sys), ss(Δ), :u)
    end
end

function ControlSystems.lft(G::UncertainSS, K::LTISystem, type=:l)
    if type === :u
        sys = lft(ss(G.sys), ss(K), :u)
    else
        nu1, ny1 = G.nw, G.nz
        sys = lft(ss(G.sys), ss(K), :l)
        sys = partition(sys, nu1, ny1)
        UncertainSS(sys, G.Δ)
    end
end

## Uncertain dynamics

struct δDyn{F, TE} <: UncertainElement{Complex{F}, F}
    ny::Int
    nu::Int
    nx::Int
    bound::F
    timeevol::TE
    name::Symbol
end

function Base.getproperty(d::δDyn, s::Symbol)
    s ∈ fieldnames(typeof(d)) && return getfield(d, s)
    if s === :Ts
        return d.timeevol.Ts
    else
        throw(ArgumentError("$(typeof(d)) has no property named $s"))
    end
end

function δDyn(ny::Int, nu::Int; bound=1, nx::Int = 2, Ts=0, name=gensym("δDyn"))
    timeevol = Ts > 0 ? Discrete(Ts) : Continuous()
    δDyn(ny,nu,nx,bound,timeevol,name)
end

Base.size(d::δDyn, n=:) = (d.ny, d.nu)[n]
# Base.length(d::δDyn) = d.ny + d.nu # NOTE: feels dangerous not to have prod(size)

function δss(ny, nu; bound=1, kwargs...)
    Δ = δDyn(ny, nu; bound, kwargs...)
    b = sqrt(bound)
    # D11 = Matrix([zeros(nu, ny) b*I(nu); b*I(ny) zeros(ny, nu)])
    # D22 = zeros(ny,nu)
    # D12 = zeros(ny+nu, nu)
    # D21 = zeros(ny, nu+ny)
    D11 = zeros(nu, ny)
    D22 = zeros(ny, nu)
    D12 = b*I(nu)
    D21 = b*I(ny)
    uss(D11, D12, D21, D22, [normalize(Δ)], Δ.timeevol)
end

## Sampling =======================================================
MonteCarloMeasurements.nominal(d::UncertainElement) = d.val

function Base.rand(d::δ{V}, N::Int) where V <: Real
    T = float(V)
    d.radius == 0 && return Particles{T, N}(fill(T(d.val), N))
    Particles(N, Uniform(d.val - d.radius, d.val + d.radius))
end

function Base.rand(D::AbstractArray{<:δ}, N::Int)
    M = rand.(D, N)
    T = eltype(M[1])
    T.(M)
end

function Base.rand(d::δ{Complex{F}}, N::Int) where F
    T = float(F)
    d.radius == 0 && return Particles(fill(T(d.val), N))
    r = Particles(N, Uniform(0, d.radius))
    θ = Particles(N, Uniform(0, 2pi))
    c = r*cis(θ)
    d.val == 0 ? c : c + d.val
end

function Base.rand(d::δDyn, N::Int)
    G = map(1:N) do i
        while true
            G = ssrand(d.ny, d.nu, d.nx; Ts = d.timeevol isa Continuous ? nothing : d.Ts)
            try
                n = norm(G, Inf)
                return (1/n) * G
            catch
            end
        end
    end
    ss2particles(G)
end

normalize(d::δ) = δ(0*d.val, one(d.radius), d.name)
normalize(d::δDyn) = δDyn(d.ny, d.nu, d.nx, 1, d.timeevol, d.name)

function Base.rand(s::UncertainSS, N::Int)
    Δ = rand.(normalize.(s.Δ), N)
    lft(s, Δ, :u)
end

## MCM

# function LinearAlgebra.det(D::AbstractMatrix{<:Complex{<:AbstractParticles}})
#     D0 = similar(D, ComplexF64)
#     parts = map(1:nparticles(D[1].re)) do i
#         for j in eachindex(D0)
#             D0[j] = Complex(D[j].re.particles[i], D[j].im.particles[i])
#         end
#         det(D0)
#     end
#     Complex(StaticParticles(getfield.(parts, :re)), StaticParticles(getfield.(parts, :im)))
# end

function ControlSystems.tzeros(A::AbstractMatrix{T}, B::AbstractMatrix{T}, C::AbstractMatrix{T}, D::AbstractMatrix{T}) where T <: AbstractParticles
    bymap(tzeros, A, B, C, D)
end

using MonteCarloMeasurements: vecindex
function sys_from_particles(P, i)
    A,B,C,D = ssdata(P)
    ss(vecindex(A, i), vecindex(B, i), vecindex(C, i), vecindex(D, i))
end

# function any0det(D::Matrix{<:Complex{<:AbstractParticles}})
#     D0 = similar(D, ComplexF64)
#     maxre = maxim = -1
#     minre = minim = 1
#     for i = 1:nparticles(D[1].re)
#         for j in eachindex(D0)
#             D0[j] = Complex(D[j].re.particles[i], D[j].im.particles[i])
#         end
#         d = det(D0)
#         maxre = max(maxre, d.re)
#         minre = min(minre, d.re)
#         maxim = max(maxim, d.im)
#         minim = min(minim, d.im)
#         if maxre > 0 && minre < 0 && maxim > 0 && minim < 0
#             return true
#         end
#     end
#     false
# end

function ss2particles(G::Vector{<:AbstractStateSpace})
    pdp(x) = Particles(permutedims(x))
    A = reduce(hcat, vec.(getproperty.(G, :A))) |> pdp
    B = reduce(hcat, vec.(getproperty.(G, :B))) |> pdp
    C = reduce(hcat, vec.(getproperty.(G, :C))) |> pdp
    D = reduce(hcat, vec.(getproperty.(G, :D))) |> pdp
    @unpack nx,ny,nu = G[1]
    A = reshape(A, nx, nx)
    B = reshape(B, nx, nu)
    C = reshape(C, ny, nx)
    D = reshape(D, ny, nu)
    ss(A,B,C,D, G[1].timeevol)
end

function δr(N::Int)
    Particles(N, Uniform(-1, 1))
end

function δc(N::Int)
    r = Particles(N, Uniform(0, 1))
    θ = Particles(N, Uniform(0, 2pi))
    c = r*cis(θ)
end

Δ(n, δ) = Diagonal([δ() for _ in 1:n])

## Intervals
IntervalArithmetic.Interval(d::δ{R}) where R <: Real =
    Interval(d.val-d.radius, d.val+d.radius)

function IntervalArithmetic.Interval(d::δ{C}) where C <: Complex
    re = Interval(d.val.re-d.radius, d.val.re+d.radius)
    im = Interval(d.val.im-d.radius, d.val.im+d.radius)
    Complex(re, im)
end

"""
    block_structure(Δ)

Take a vector of uncertain elements and return a vector of vectors with the block
structure of perturbation blocks as described by μ-tools, i.e.
- `[-N, 0]` denotes a repeated real block of size `N`
- `[N, 0]` denotes a repeated complex block of size `N`
- `[ny, nu]` denotes a full complex block of size `ny × nu`
"""
function block_structure(D)
    perm = sortperm(D, by=d->d.name)
    D = D[perm]
    blocks = Vector{Int}[]
    blockstart = 1
    for i in 2:length(D)
        if D[i].name == D[i-1].name
            if blockstart == 0
                blockstart = i-1
            end
        else
            push!(blocks, makeblock(D[blockstart:i-1]))
            blockstart = i
        end
    end
    push!(blocks, makeblock(D[blockstart:length(D)]))
    blocks, perm
end

function makeblock(d) # TODO: full uncertain block that is not a dynamical system is not supported
    N = length(d)
    if N > 1 || (N==1 && d[] isa δ) # Δ = δI -> μ(M) = ρ(M) 8.80
        if eltype(d[1]) <: Complex
            b = [N, 0] # complex*I
        else
            b = [-N, 0] # real*I
        end
    else
        b = [size(d[], 1), size(d[], 2)]
    end
end

"""
    blocks, M = blocksort(P::UncertainSS)

Returns the block structure of `P.Δ` as well as `P.M` permuted according to the sorted block structure.
`blocks` is a vector of vectors with the block
structure of perturbation blocks as described by μ-tools, i.e.
- `[-N, 0]` denotes a repeated real block of size `N`
- `[N, 0]` denotes a repeated complex block of size `N`
- `[ny, nu]` denotes a full complex block of size `ny × nu`
"""
function blocksort(P::UncertainSS)
    blocks, perm = block_structure(P.Δ)
    Poutinds = []
    Pininds = []
    # Build indices into M that corresponds to the Δ blocks and permute those indices with the same permutation as was applied to Δ
    for (i, d) in enumerate(P.Δ)
        if i == 1
            push!(Poutinds, 1:size(d, 2))
            push!(Pininds, 1:size(d, 1))
        else
            push!(Poutinds, (1:size(d, 2)) .+ Poutinds[i-1][end])
            push!(Pininds, (1:size(d, 1)) .+ Pininds[i-1][end])
        end
    end
    Poutinds = reduce(vcat, Poutinds[perm])
    Pininds = reduce(vcat, Pininds[perm])
    M = P.M[Poutinds, Pininds]
    blocks, M
end

function getinds(blocks::Vector{Vector{Int}}, d::Int)
    lengths = [b[2] <= 1 ? b[1] : b[d] for b in blocks]
    endinds = cumsum(lengths)
    startinds = [1; endinds[1:end-1] .+ 1]
    inds = UnitRange.(0 .+ startinds, lengths .+ startinds .- 1)
end

# function getinds(D::AbstractVector, d::Int)
#     lengths = size.(D, d)
#     endinds = cumsum(lengths)
#     startinds = [1; endinds[1:end-1] .+ 1]
#     inds = UnitRange.(0 .+ startinds, lengths .+ startinds .- 1)
# end

# struct BlockDiag{T}
#     D::Vector{T}
#     rinds::Vector{UnitRange{Int}}
#     cinds::Vector{UnitRange{Int}}
# end
# BlockDiag(D::AbstractVector) = BlockDiag(D, getinds(D, 1), getinds(D, 2))

# Base.size(B::BlockDiag) = (B.rinds[end][end], B.cinds[end][end])

# function Base.:*(A, B::BlockDiag)
#     rinds = B.rinds
#     submats = map(enumerate(rinds)) do (i,inds)
#         A[:, inds] * B.D[i]
#     end
#     reduce(hcat, submats)
# end

# function quadform_minus(Q::BlockDiag, M, b, op)
#     # con = M'Q*M - b^2*Q ⪯ I(n)
#     rinds = Q.rinds
#     cinds = Q.cinds
#     map(eachindex(rinds)) do i
#         r = rinds[i]
#         c = cinds[i]
#         q = Q.D[i]
#         op(M[r,:]'q*M[c,:] .- b^2*q, I(size(M,2)))
#     end
# end

#=
D_from_Delta: I-full, full-I, diag-diag
Version that constructs from initial guess
=#

# See Skogestad 8.8.3
# for complex δ:
# Δ = δI -> μ(M) = ρ(M) 8.80
# Full block complex -> μ(M) = σ̄(M) = opnorm(M)
# In general ρ(M) <= μ(M) <= σ̄(M)

# function make_D(blocks)
#     D = []
#     for b in blocks
#         if b[2] == 0
#             n = abs(b[1])
#             T = b[1] > 0 ? ComplexF64 : Float64
#             push!(D, zeros(T, n, n))
#         else
#             push!(D, 0*I)
#         end
#     end
#     BlockDiagonal(D)
# end


using MCMCBenchmarks,Test,Random
Random.seed!(4504)
tests = [
    "generative_model_tests",
    "Gaussian_Tests",
    "SDT_Tests",
    "Poisson_Test",
    "Regression_Tests",
    "Plot_Tests",
    #"LBA_Tests"
]

res = map(tests) do t
    @eval module $(Symbol("Test_", t))
        include($t * ".jl")
    end
end


using RICoverAlg

# Queueing problems
m = 2;  n = 40 ; lmd = 1.0
g = ones(n); g /=sum(g)
U = zeros(m, n)
for i in (1:n)
   U[2,i] = 28 - i
end
# shiftterm = (((maximum(U)+minimum(U))/2))   ## reduce extremes of U for numerical stability
Z = exp.((U.-0)/lmd)

@test convertInputsToFloats(U,g,1) == (U,g,1.)
@test convertInputsToFloats([2 1],g, 1//10) == ([2. 1.],g,.1)
#@test convertInputsToFloats([2; 1],1, 1//10) == ([2.; 1.],1.,.1) # how to convert a number to Array if necessary
#@test cancelZeroStates([4 5; 6 7], [0 1]) == ([5; 7], [1], [2]) # why is this test not working?
@test checkForTriviality([2. 1.],[.1; 0.9],1) == (ones(1), log(2.)*.1+log(1.)*.9, 0, "Trivial: only one action is possible")
@test checkForTriviality([2. ; 1.],[1],2) == ([1. ; 0.], log(2.0), 0, "Trivial: only one state is possible")

@test risolve(U, g, 10)[1] == [0.0, 1.0]
@test risolve(reverse(U,dims=1), g, 10)[3][1:47] == "Found boundary solution (two alternatives only)"

# Matejka & McKay 2015, RBT
U = [0 1 0 1;
    0 0 1 1;
    1/2 1/2 1/2 1/2]
lmd = 0.4
g = [1+1; 1-1; 1-1; 1+1]./4
@test convertInputsToFloats(U,g,lmd) == (U./1.,g,lmd)
@test cancelZeroStates(U,g ) == (U[:,[1,4]],[.5, .5 ],[1, 4])
@test transform_utilities(U[:,[1,4]],lmd) == ( exp.((U[:,[1,4]].-.0)/lmd)   ,  .0   )


@testset "Expected Loss" begin

    @testset "calculate expected loss" begin
        Random.seed!(1234)
        n = 5
        sampleA = [1, 1, 1, 0, 0]
        sampleB = [1, 0, 0, 1, 0]
        @test expectedloss(sampleA, sampleB, upliftloss, 5) == 0.2
        @test expectedloss(sampleB, sampleA, upliftloss, 5) == 0.4
    end

    @testset "Approximate expected loss from posteriors" begin
        Random.seed!(1234)
        n = 10000
        modelA = ConjugateBernoulli(50, 50)
        modelB = ConjugateBernoulli(60, 40)

        exploss = expectedloss(modelA, modelB, [:θ], lossfunc=upliftloss, numsamples=n)
        @test isapprox(exploss, 0.1, atol=0.01)

        exploss = expectedloss(modelB, modelA, [:θ], lossfunc=upliftloss, numsamples=n)
        @test isapprox(exploss, 0.0, atol=0.01)
    end

    @testset "Approximate expected loss from experiment" begin
        Random.seed!(1234)
        n = 10000
        modelA = ConjugateBernoulli(50, 50)
        modelB = ConjugateBernoulli(60, 40)

        exploss = expectedloss(modelA, modelB, lossfunc=upliftloss, numsamples=n)
        @test isapprox(exploss, 0.1, atol=0.01)

        stoppingrule = ExpectedLossThresh(0.001)
        experiment = ExperimentAB([modelA, modelB], stoppingrule)

        (modelnames, explosses) = expectedlosses(experiment, lossfunc=upliftloss, numsamples=n)

        @test modelnames == ["control", "variant 1"]
        @test isapprox(explosses[1], 0.1, atol=0.01)
        @test isapprox(explosses[2], 0.0, atol=0.01)
    end

end

@testset "ExperimentAB" begin
    @testset "ChainedModel with Bernoulli and LogNormal" begin
        Random.seed!(1234)
        eachrunnum = 100000
        truedistA1 = Bernoulli(0.05)
        truedistA2 = LogNormal(1.0, 1.0)
        truedistB1 = Bernoulli(0.05 * 1.05)
        truedistB2 = LogNormal(1.0, 1.0)

        numsims = 50000

        dataA1 = rand(truedistA1, eachrunnum)
        dataA2 = rand(truedistA2, eachrunnum)
        dataB1 = rand(truedistB1, eachrunnum)
        dataB2 = rand(truedistB2, eachrunnum)

        statsA1 = BetaStatistics(dataA1)
        statsA2 = LogNormalStatistics(dataA2)
        statsB1 = BetaStatistics(dataB1)
        statsB2 = LogNormalStatistics(dataB2)

        modelA = ChainedModel(
            [ConjugateBernoulli(1, 1), ConjugateLogNormal(0.0, 1.0, 0.001, 0.001)],
        )
        modelB = ChainedModel(
            [ConjugateBernoulli(1, 1), ConjugateLogNormal(0.0, 1.0, 0.001, 0.001)],
        )

        stoppingrule = ExpectedLossThresh(0.001)
        experiment = ExperimentAB([modelA, modelB], stoppingrule)
        update!(experiment, [[statsA1, statsA2], [statsB1, statsB2]])

        modelnames, explosses = expectedlosses(experiment)

        @testset "Control Group" begin
            model1_list = experiment.models["control"].models
            @test model1_list[1].dist.α == 4972
            @test model1_list[1].dist.β == 95030
            @test model1_list[2].dist.μ ≈ 1.0
            @test model1_list[2].dist.v ≈ 0
            @test model1_list[2].dist.α ≈ 50000
            @test isapprox(model1_list[2].dist.θ, 49373, atol=1.0)
        end

        @testset "Expected loss" begin
            @test modelnames == ["control", "variant 1"]
            @test explosses[1] ≈ 0.016
            @test explosses[2] ≈ 0
        end

        winner, _ = decide!(experiment)
        @test winner == "variant 1"
    end

end

@testset "Number of models" begin
    # Generate sample data
    n = 1000
    dataA = rand(Bernoulli(0.150), n)
    dataB = rand(Bernoulli(0.145), n)
    dataC = rand(Bernoulli(0.180), n)

    # Define the models
    modelA = ConjugateBernoulli(1, 1)
    modelB = ConjugateBernoulli(1, 1)
    modelC = ConjugateBernoulli(1, 1)

    stoppingrule = ProbabilityBeatAllThresh(0.99)
    @test_throws ArgumentError("Number of models needs to be equal to 2.") ExperimentAB([modelA, modelB, modelC], stoppingrule)
end

@testset "ExperiemntBF{NormalEffectSize}" begin
    @testset "Bayes factor calculation" begin
        # the example is taken from
        # https://statswithr.github.io/book/hypothesis-testing-with-normal-populations.html
        n0 = 32.7^2
        m0 = 0.5
        x̄ = 0.500177
        σ = 0.5
        n = 1.0449e8;
        σ0 = 1/sqrt(n0)

        model = NormalEffectSize(m0, σ0)
        stats = NormalStatistics(n=1.0449e8, meanx=x̄, sdx=σ)
        bf10 = bayesfactor(model, stats)
        @test isapprox(bf10, 2.2303, rtol=0.001)
    end

    @testset "Two samples with equal size and sd" begin
        n0 = 32.7^2
        m0 = 0.5 # null hypothesis: difference is 0.5
        σ = 0.5
        n = 1.0449e8;
        σ0 = 1/sqrt(n0)

        # 2*n so that the effective sample size will be `n`
        stats1 = NormalStatistics(n=2*n, meanx=1.000177, sdx=σ)
        stats2 = NormalStatistics(n=2*n, meanx=0.5, sdx=σ)
        stats = TwoNormalStatistics(stats1, stats2)

        model = NormalEffectSize(m0, σ0)
        bf10 = bayesfactor(model, stats)
        @test isapprox(bf10, 2.2303, rtol=0.001)
    end
end

@testset "ExperiemntBF{StudentTEffectSize}" begin
    @testset "Two-Sided null wins, thresh=5" begin
        normalstats = TwoNormalStatistics(
            NormalStatistics(meanx=28.8, sdx=13.5, n=133),
            NormalStatistics(meanx=30.6, sdx=14.3, n=867)
        )
        model = StudentTEffectSize(r=1.0)
        stats = StudentTStatistics(normalstats)
        bf10_model = bayesfactor(model, stats)

        thresh=5
        stoppingrule = TwoSidedBFThresh(thresh)
        experiment = ExperimentBF(model=model, rule=stoppingrule)
        update!(experiment, normalstats)
        bf10_exp = bayesfactor(experiment)

        @test bf10_model == bf10_exp
        @test isapprox(bf10_exp, 1/5.45, rtol=0.01)

        winner, _ = decide!(experiment)
        @test winner == "null"
    end

    @testset "Two-Sided alternative wins, thresh=5" begin
        normalstats = TwoNormalStatistics(
            NormalStatistics(meanx=28.8, sdx=13.5, n=133),
            NormalStatistics(meanx=33.6, sdx=14.3, n=867)
        )
        model = StudentTEffectSize(r=1.0)
        stats = StudentTStatistics(normalstats)
        bf10_model = bayesfactor(model, stats)

        thresh=5
        stoppingrule = TwoSidedBFThresh(thresh)
        experiment = ExperimentBF(model=model, rule=stoppingrule)
        update!(experiment, normalstats)
        bf10_exp = bayesfactor(experiment)

        @test bf10_model == bf10_exp
        @test isapprox(bf10_exp, 46.6077, rtol=0.01)

        winner, _ = decide!(experiment)
        @test winner == "alternative"
    end

    @testset "Two-Sided no winners, thresh=0.01" begin
        normalstats = TwoNormalStatistics(
            NormalStatistics(meanx=28.8, sdx=13.5, n=133),
            NormalStatistics(meanx=30.6, sdx=14.3, n=867)
        )
        model = StudentTEffectSize(r=1.0)
        stats = StudentTStatistics(normalstats)
        bf10_model = bayesfactor(model, stats)

        thresh=20
        stoppingrule = TwoSidedBFThresh(thresh)
        experiment = ExperimentBF(model=model, rule=stoppingrule)
        update!(experiment, normalstats)
        bf10_exp = bayesfactor(experiment)

        @test bf10_model == bf10_exp
        @test isapprox(bf10_exp, 1/5.45, rtol=0.01)

        winner, _ = decide!(experiment)
        @test winner === nothing
    end

    @testset "One-Sided no winner, thresh=5" begin
        normalstats = TwoNormalStatistics(
            NormalStatistics(meanx=28.8, sdx=13.5, n=133),
            NormalStatistics(meanx=30.6, sdx=14.3, n=867)
        )
        model = StudentTEffectSize(r=1.0)
        stats = StudentTStatistics(normalstats)
        bf10_model = bayesfactor(model, stats)

        thresh=5
        stoppingrule = OneSidedBFThresh(thresh)
        experiment = ExperimentBF(model=model, rule=stoppingrule)
        update!(experiment, normalstats)
        bf10_exp = bayesfactor(experiment)

        @test bf10_model == bf10_exp
        @test isapprox(bf10_exp, 1/5.45, rtol=0.01)

        winner, _ = decide!(experiment)
        @test winner === nothing
    end

    @testset "One-Sided alternative wins, thresh=5" begin
        normalstats = TwoNormalStatistics(
            NormalStatistics(meanx=28.8, sdx=13.5, n=133),
            NormalStatistics(meanx=33.6, sdx=14.3, n=867)
        )
        model = StudentTEffectSize(r=1.0)
        stats = StudentTStatistics(normalstats)
        bf10_model = bayesfactor(model, stats)

        thresh=5
        stoppingrule = OneSidedBFThresh(thresh)
        experiment = ExperimentBF(model=model, rule=stoppingrule)
        update!(experiment, normalstats)
        bf10_exp = bayesfactor(experiment)

        @test bf10_model == bf10_exp
        @test isapprox(bf10_exp, 46.6077, rtol=0.01)

        winner, _ = decide!(experiment)
        @test winner == "alternative"
    end
end


#=
   Can one not have two models in the same file?
   Yes, it don't seems to be any problem...

=#
using Turing

# Define a simple Normal model with unknown mean and variance.
a = []
@model test_model(n) = begin
        d = TArray{Int}(undef, n)
        # d = Vector{Int}(undef, n)
        for i in 1:n
            d[i] ~ DiscreteUniform(0,9)
        end
        s = sum(d)
        push!(a, s)
end

#  Run sampler, collect results
chn = sample(test_model(10), MH(), 1000)
println("mean(a): $(mean(a))")

a = []
@model test_model2(n) = begin
        d = TArray{Int}(undef, n)
        # d = Vector{Int}(undef, n)
        for i in 1:n
            d[i] ~ DiscreteUniform(0,9)
        end
        s = sum(d)
        push!(a, s)
end

#  Run sampler, collect results
chn = sample(test_model2(10), MH(), 1000)
println("mean(a) 2: $(mean(a))")


# function simulation_1frame_traces(coef_stim::Array{Float64}, internal_noise:: Array{Float64},T:: Float64,tau::Float64,x0::Float64)
function simulation_1frame_traces(coef_stim, internal_noise,T,tau,x0)
    a=size(internal_noise)
    #println(a)
    NT=a[2]
    Nstim=a[1]
    deltat=T/NT
    dt_sim=deltat/tau
    dt_sqrt=sqrt(dt_sim)
    d=zeros(Nstim)
    x=zeros(Nstim,NT)
    for itrial in 1:Nstim
         for it in 1:NT-1
             x[itrial,it+1]=x[itrial,it]-( -coef_stim[1]-coef_stim[2]*x[itrial,it]+coef_stim[3]*x[itrial,it]^3)*dt_sim+dt_sqrt*internal_noise[itrial,it]
         end
    end

    return x
 end

 function simulation_1frame(coef_stim, internal_noise,T,tau,x0)
     a=size(internal_noise)
     println(a)
     NT=a[2]
     Nstim=a[1]
     deltat=T/NT
     dt_sim=deltat/tau
     dt_sqrt=sqrt(dt_sim)
     d=zeros(Nstim)
     x=zeros(Nstim,NT)
     for itrial in 1:Nstim
         x=x0
         for it in 1:NT-1
             x=x-( -coef_stim[1]-coef_stim[2]*x+coef_stim[3]*x^3)*dt_sim+dt_sqrt*internal_noise[itrial,it]
         end
         if x<0
             d[itrial]=-1
        else
            d[itrial]=1
        end
     end

     return d
end

function simulation_Nframe_traces(coef_stim,stim,internal_noise,T,tau,x0,urgency=0.)

    # println("hola")
    # println("puta")
    aux2=size(internal_noise)
    #println(a)
    aux=size(stim)
    Nframes=aux[2]
    Nstim=aux[1]
    NT=aux2[2]/Nframes
    deltat=T/NT
    dt_sim=deltat/tau
    dt_sqrt=sqrt(dt_sim)
    d=zeros(Nstim)
    x=zeros(Nstim,Nframes*NT+1)
    println(Nframes)
    for itrial in 1:Nstim
        it=2
        x[itrial,1]=x0

        for iframe in 1:Nframes
            u=urgency*(iframe-1)
            for i in 1:NT
                # println("it: ",it-1, " itrial ", itrial," iframe ", iframe)
                x[itrial,it]=x[itrial,it-1]-( -stim[itrial,iframe]-(coef_stim[2]+u)*x[itrial,it-1]+coef_stim[3]*x[itrial,it-1]^3)*dt_sim+dt_sqrt*internal_noise[itrial,it-1]
                it=it+1
            end
        end
    end

      return x
end

function simulation_Nframe(coef_stim,stim,internal_noise,T,tau,x0;urgency=0.)

    # println("hola")
    # println("puta")
    aux2=size(internal_noise)
    #println(a)
    aux=size(stim)
    Nframes=aux[2]
    Nstim=aux[1]
    NT=aux2[2]/Nframes
    deltat=T/NT
    dt_sim=deltat/tau
    dt_sqrt=sqrt(dt_sim)
    d=zeros(Float64,Nstim)
    # println(Nframes)
    # println("hola soc ",dt_sim)
    for itrial in 1:Nstim
        it=2
        x=x0

        for iframe in 1:Nframes
            u=urgency*(iframe-1)
            for i in 1:NT
                # println("it: ",it-1, " itrial ", itrial," iframe ", iframe)
                x=x-( -stim[itrial,iframe]-(coef_stim[2]+u)*x+coef_stim[3]*x^3)*dt_sim+dt_sqrt*internal_noise[itrial,it-1]
                it=it+1
            end
        end
        d[itrial]=sign(x)
    end

      return d
end

###########URGENCY+STIM_TRANSFORM+SOFT_BOUND################
function read_out_soft_sim(x,ro_p)
    pr=1./(1+exp.(-10*ro_p[1]*(x-ro_p[2])))
    if rand()<pr
        d=1
    else
        d=-1
    end
    return d
end

function read_out_perfect_sim(x)
    return sign(x)
end

function read_out_TW_sim(x,c2,c4,c6)
    delta=c4^2+4*c2*c6
    #println(delta," ", delta>0)

    if  (c2<0) & (c4<0) & (c6>0) & (delta>0)
        b=sqrt( (-c4-sqrt(delta))/(2*c6)   )
        if abs(x)<b
            a=sign( rand() -0.5)
            return sign( rand() -0.5)
        else
            return sign(x)
        end
    else
        return sign(x)
    end
end

function simulation_Nframe_general(param,stim,internal_noise,T,tau,model_f)
    """
    param: Dict with the parameter models
    internal_noise: interncal_noise= sigma*chi(t)
    T is total time T=Tframe*Nframes
    """
    Nstim,Nframes=size(stim)
    aux2=size(internal_noise)
    NT=aux2[2]
    NTframes=NT/Nframes
    deltat=T/NT
    dt_sim=deltat/tau
    dt_sqrt=sqrt(dt_sim)
    d=zeros(Float64,Nstim)
    xfinal=zeros(Float64,Nstim)
    c2=zeros(Nframes)
    c4=zeros(Nframes)
    c6=zeros(Nframes)
    #println("hola ",dt_sim," ", deltat," ",T," ", NT)
    model_f[:c2_f](c2,param["c2"])
    model_f[:c4_f](c4,param["c4"])
    if "c6" in keys(param)
        model_f[:c6_f](c6,param["c6"])
    else
        model_f[:c6_f](c6)
    end
    xfinal=zeros(Nstim)
    hbias_d=param["hbias_d"][1]
    bias_d=param["bias_d"][1]

    s=zeros(Nframes)
    hb=0
    #println(hb,hbias_d)
    for itrial in 1:Nstim
        #println(itrial)
        it=2
        x=param["x0"][1]+hb*param["x0"][1]

        model_f[:st_f](stim[itrial,:],s,param["st"])
        #println(s)
        for iframe in 1:Nframes
            for i in 1:NTframes

                #println("it: ",it-1, " itrial ", itrial," iframe ", iframe)
                x=x-(-(s[iframe]+hb*hbias_d+bias_d)-c2[iframe]*x+c4[iframe]*x^3+c6[iframe]*x^5)*dt_sim+dt_sqrt*internal_noise[itrial,it-1]
                it=it+1
            end
        end
        xfinal[itrial]=x
        if model_f[:ro_f]==read_out_TW_sim
            d[itrial]=model_f[:ro_f](x,c2[end],c4[end],c6[end])
        elseif model_f[:ro_f]==read_out_perfect_sim
            d[itrial]=model_f[:ro_f](x)
        else
            d[itrial]=model_f[:ro_f](x,param["ro"])
        end
        #hb=d[itrial]
    end
    return d,xfinal
end

function simulation_Nframe_general_trace(param,stim,internal_noise,T,tau,model_f)
    """
    T is total time T=Tframe*Nframes
    """
    #println("hola")
    Nstim,Nframes=size(stim)
    aux2=size(internal_noise)
    NT=aux2[2]
    NTframes=NT/Nframes
    deltat=T/NT
    dt_sim=deltat/tau
    dt_sqrt=sqrt(dt_sim)
    d=zeros(Float64,Nstim)
    x=zeros(Float64,Nstim,NT+1)
    c2=zeros(Nframes)
    c4=zeros(Nframes)
    c6=zeros(Nframes)
    #println("hola ",dt_sim," ", deltat," ",T," ", NT)
    model_f[:c2_f](c2,param["c2"])
    model_f[:c4_f](c4,param["c4"])
    if "c6" in keys(param)
        model_f[:c6_f](c6,param["c6"])
    else
        model_f[:c6_f](c6)
    end
    hbias_d=param["hbias_d"][1]
    bias_d=param["bias_d"][1]
    # println(Nframes)
    # println("hola soc ",dt_sim)
    s=zeros(Nframes)
    hb=0 ###########canvia això per posar histroy bias
    #println(hb,hbias_d)
    Bound=100
    for itrial in 1:Nstim
        it=2
        x[itrial,1]=param["x0"][1]+hb*param["x0"][1]

        model_f[:st_f](stim[itrial,:],s,param["st"])
        #println(s)
        for iframe in 1:Nframes
            for i in 1:NTframes

                # println("it: ",it-1, " itrial ", itrial," iframe ", iframe)
                if abs(x[itrial,it-1])>Bound
                    x[itrial,it]=x[itrial,it-1]
                else

                    x[itrial,it]=x[itrial,it-1]-(-(s[iframe]+hb*hbias_d+bias_d)-c2[iframe]*x[itrial,it-1]+c4[iframe]*x[itrial,it-1]^3+c6[iframe]*x[itrial,it-1]^5)*dt_sim+dt_sqrt*internal_noise[itrial,it-1]
                end
                it=it+1

            end
        end
        if model_f[:ro_f]==read_out_TW_sim
            d[itrial]=model_f[:ro_f](x[itrial,end],c2[end],c4[end],c6[end])
        elseif model_f[:ro_f]==read_out_perfect_sim
            d[itrial]=model_f[:ro_f](x[itrial,end])
        else
            d[itrial]=model_f[:ro_f](x[itrial,end],param["ro"])
        end
        #hb=d[itrial]
    end
    return d,x
end

function simulation_Nframe_PI(param,stim,internal_noise,T,tau,model_f)

    Nstim,Nframes=size(stim)
    aux2=size(internal_noise)
    NT=aux2[2]
    NTframes=NT/Nframes
    deltat=T/NT
    dt_sim=deltat/tau
    dt_sqrt=sqrt(dt_sim)
    d=zeros(Float64,Nstim)
    x=zeros(Float64,Nstim,NT+1)

    s=zeros(Nframes)
    hb=0 ###########canvia això per posar histroy bias
    #println(hb,hbias_d)
    for itrial in 1:Nstim
        it=2
        x[itrial,1]=param["x0"][1]#+hb*param["x0"][1]

        model_f[:st_f](stim[itrial,:],s,param["st"])
        #println(s)
        for iframe in 1:Nframes
            for i in 1:NTframes
                x[itrial,it]=x[itrial,it-1]+s[iframe]*dt_sim+dt_sqrt*internal_noise[itrial,it-1]
                it=it+1
            end
        end
        d[itrial]=sign(x[itrial,end])

    end
    return d,x
end

function DW_potential_prima(x,coef)
    return -coef[1]-coef[2]*x+coef[3]*x^3
end

function compute_kernel(stim,d)
    right=find(x->x==1,d)
    left=find(x->x==-1,d)
    aux=size(stim)
    kernel=zeros(aux[2])
    for iframe in 1:aux[2]
        println(iframe)
        r=roc(stim[left,iframe],stim[right,iframe])
        kernel[iframe]=auc(r)
    end
    return kernel
end

###### Simulations DDMA ####
function simulation_Nframe_general_DDMA_trace(param,stim,internal_noise,T,tau,model_f)
    #println("hola")
    Nstim,Nframes=size(stim)
    aux2=size(internal_noise)
    NT=aux2[2]
    NTframes=NT/Nframes
    deltat=T/NT
    dt_sim=deltat/tau
    dt_sqrt=sqrt(dt_sim)
    d=zeros(Float64,Nstim)
    x=zeros(Float64,Nstim,NT+1)
    hbias_d=param["hbias_d"][1]
    bias_d=param["bias_d"][1]
    # println(Nframes)
    # println("hola soc ",dt_sim)
    s=zeros(Nframes)
    hb=0 ###########canvia això per posar histroy bias
    #println(hb,hbias_d)
    Bound=param["B"]
    for itrial in 1:Nstim
        it=2
        x[itrial,1]=param["x0"][1]+hb*param["x0"][1]

        model_f[:st_f](stim[itrial,:],s,param["st"])
        #println(s)
        for iframe in 1:Nframes
            for i in 1:NTframes

                if abs(x[itrial,it-1])>Bound
                    x[itrial,it]=x[itrial,it-1]
                else
                    x[itrial,it]=x[itrial,it-1]-(-(s[iframe]+hb*hbias_d+bias_d))*dt_sim+dt_sqrt*internal_noise[itrial,it-1]
                end
                it=it+1

            end
        end

        if model_f[:ro_f]==read_out_perfect_sim
            d[itrial]=model_f[:ro_f](x[itrial,end])
        else
            d[itrial]=model_f[:ro_f](x[itrial,end],param["ro"])
        end
    end
    return d,x
end

###### Simulations DDMA ####
function simulation_Nframe_general_DDMA(param,stim,internal_noise,T,tau,model_f)
    #println("hola")
    Nstim,Nframes=size(stim)
    aux2=size(internal_noise)
    NT=aux2[2]
    NTframes=NT/Nframes
    deltat=T/NT
    dt_sim=deltat/tau
    dt_sqrt=sqrt(dt_sim)
    d=zeros(Float64,Nstim)
    x=0.0
    xfinal=zeros(Nstim)

    hbias_d=param["hbias_d"][1]
    bias_d=param["bias_d"][1]
    # println(Nframes)
    # println("hola soc ",dt_sim)
    s=zeros(Nframes)
    hb=0 ###########canvia això per posar histroy bias
    #println(hb,hbias_d)
    Bound=param["B"]
    for itrial in 1:Nstim
        it=2
        x=param["x0"][1]+hb*param["x0"][1]

        model_f[:st_f](stim[itrial,:],s,param["st"])
        #println(s)
        for iframe in 1:Nframes
            for i in 1:NTframes

                if abs(x)>Bound
                    x=x
                else
                    x=x-(-(s[iframe]+hb*hbias_d+bias_d))*dt_sim+dt_sqrt*internal_noise[itrial,it-1]
                end
                it=it+1

            end
        end

        if model_f[:ro_f]==read_out_perfect_sim
            d[itrial]=model_f[:ro_f](x)
        else
            d[itrial]=model_f[:ro_f](x,param["ro"])
        end
        xfinal[itrial]=x
    end
    return d,xfinal
end

###### Simulations DDMR ####
function simulation_Nframe_general_DDMR(param,stim,internal_noise,T,tau,model_f)
    #println("hola")
    Nstim,Nframes=size(stim)
    aux2=size(internal_noise)
    NT=aux2[2]
    NTframes=NT/Nframes
    deltat=T/NT
    dt_sim=deltat/tau
    dt_sqrt=sqrt(dt_sim)
    d=zeros(Float64,Nstim)
    x=0.0
    xfinal=zeros(Nstim)

    hbias_d=param["hbias_d"][1]
    bias_d=param["bias_d"][1]
    # println(Nframes)
    # println("hola soc ",dt_sim)
    s=zeros(Nframes)
    hb=0 ###########canvia això per posar histroy bias
    #println(hb,hbias_d)
    Bound=param["B"]
    for itrial in 1:Nstim
        it=2
        x=param["x0"][1]+hb*param["x0"][1]

        model_f[:st_f](stim[itrial,:],s,param["st"])
        #println(s)
        for iframe in 1:Nframes
            for i in 1:NTframes

                if x>Bound
                    x=Bound
                elseif  x<-Bound
                    x=-Bound
                end

                x=x-(-(s[iframe]+hb*hbias_d+bias_d))*dt_sim+dt_sqrt*internal_noise[itrial,it-1]
                it=it+1

            end
        end

        if model_f[:ro_f]==read_out_perfect_sim
            d[itrial]=model_f[:ro_f](x)
        else
            d[itrial]=model_f[:ro_f](x,param["ro"])
        end
        xfinal[itrial]=x
    end
    return d,xfinal
end

function simulation_Nframe_general_DDMR_trace(param,stim,internal_noise,T,tau,model_f)
    #println("hola")
    Nstim,Nframes=size(stim)
    aux2=size(internal_noise)
    NT=aux2[2]
    NTframes=NT/Nframes
    deltat=T/NT
    dt_sim=deltat/tau
    dt_sqrt=sqrt(dt_sim)
    d=zeros(Float64,Nstim)
    x=zeros(Float64,Nstim,NT+1)
    hbias_d=param["hbias_d"][1]
    bias_d=param["bias_d"][1]
    # println(Nframes)
    # println("hola soc ",dt_sim)
    s=zeros(Nframes)
    hb=0 ###########canvia això per posar histroy bias
    #println(hb,hbias_d)
    Bound=param["B"]
    for itrial in 1:Nstim
        it=2
        x[itrial,1]=param["x0"][1]+hb*param["x0"][1]

        model_f[:st_f](stim[itrial,:],s,param["st"])
        #println(s)
        for iframe in 1:Nframes
            for i in 1:NTframes

                if x[itrial,it-1]>Bound
                    x[itrial,it-1]=Bound
                elseif  x[itrial,it-1]<-Bound
                    x[itrial,it-1]=-Bound
                end

                x[itrial,it]=x[itrial,it-1]-(-(s[iframe]+hb*hbias_d+bias_d))*dt_sim+dt_sqrt*internal_noise[itrial,it-1]
                it=it+1

            end
        end

        if model_f[:ro_f]==read_out_perfect_sim
            d[itrial]=model_f[:ro_f](x[itrial,end])
        else
            d[itrial]=model_f[:ro_f](x[itrial,end],param["ro"])
        end
    end
    return d,x
end


using DiffEqSensitivity,OrdinaryDiffEq, RecursiveArrayTools, DiffEqBase,
      ForwardDiff, Calculus, QuadGK, LinearAlgebra
using Test

function fb(du,u,p,t)
  du[1] = dx = p[1]*u[1] - p[2]*u[1]*u[2]
  du[2] = dy = -p[3]*u[2] + p[4]*u[1]*u[2]
end
function jac(J,u,p,t)
  (x, y, a, b, c) = (u[1], u[2], p[1], p[2], p[3])
  J[1,1] = a + y * b * -1
  J[2,1] = y
  J[1,2] = b * x * -1
  J[2,2] = c * -1 + x
end

f = ODEFunction(fb,jac=jac)
p = [1.5,1.0,3.0,1.0]
prob = ODEProblem(f,[1.0;1.0],(0.0,10.0),p)
sol = solve(prob,Tsit5(),abstol=1e-14,reltol=1e-14)
probb = ODEProblem(fb,[1.0;1.0],(0.0,10.0),p)

solb = solve(probb,Tsit5(),abstol=1e-14,reltol=1e-14)
sol_end = solve(probb,Tsit5(),abstol=1e-14,reltol=1e-14,
          save_everystep=false,save_start=false)

# Do a discrete adjoint problem
println("Calculate discrete adjoint sensitivities")
t = 0.0:0.5:10.0
# g(t,u,i) = (1-u)^2/2, L2 away from 1
function dg(out,u,p,t,i)
  (out.=2.0.-u)
end

easy_res = adjoint_sensitivities(sol,Tsit5(),dg,t,abstol=1e-14,
                                 reltol=1e-14,iabstol=1e-14,ireltol=1e-12)
easy_res2 = adjoint_sensitivities(solb,Tsit5(),dg,t,abstol=1e-14,
                                 reltol=1e-14,iabstol=1e-14,ireltol=1e-12,sensealg=SensitivityAlg(quad=true,backsolve=false))
easy_res3 = adjoint_sensitivities(solb,Tsit5(),dg,t,abstol=1e-14,
                                  reltol=1e-14,iabstol=1e-14,ireltol=1e-12,sensealg=SensitivityAlg(quad=false,backsolve=false))
easy_res4 = adjoint_sensitivities(solb,Tsit5(),dg,t,abstol=1e-14,
                                  reltol=1e-14,iabstol=1e-14,ireltol=1e-12,sensealg=SensitivityAlg(backsolve=true))
easy_res5 = adjoint_sensitivities(sol,Kvaerno5(nlsolve=NLAnderson(), smooth_est=false),dg,t,abstol=1e-12,
                                 reltol=1e-10,iabstol=1e-14,ireltol=1e-12,sensealg=SensitivityAlg(backsolve=true))
easy_res6 = adjoint_sensitivities(solb,Tsit5(),dg,t,abstol=1e-14,
                                  reltol=1e-14,iabstol=1e-14,ireltol=1e-12,
                                  sensealg=SensitivityAlg(checkpointing=true,quad=true),
                                  checkpoints=sol.t[1:5:end])
easy_res7 = adjoint_sensitivities(solb,Tsit5(),dg,t,abstol=1e-14,
                                  reltol=1e-14,iabstol=1e-14,ireltol=1e-12,
                                  sensealg=SensitivityAlg(checkpointing=true,quad=false),
                                  checkpoints=sol.t[1:5:end])

adj_prob = ODEAdjointProblem(sol,dg,t)
adj_sol = solve(adj_prob,Tsit5(),abstol=1e-14,reltol=1e-14)
integrand = AdjointSensitivityIntegrand(sol,adj_sol)
res,err = quadgk(integrand,0.0,10.0,atol=1e-14,rtol=1e-12)

@test isapprox(res, easy_res, rtol = 1e-10)
@test isapprox(res, easy_res2, rtol = 1e-10)
@test isapprox(res, easy_res3, rtol = 1e-10)
@test isapprox(res, easy_res4, rtol = 1e-10)
@test isapprox(res, easy_res5, rtol = 1e-7)
@test isapprox(res, easy_res6, rtol = 1e-9)
@test isapprox(res, easy_res7, rtol = 1e-9)

println("Calculate adjoint sensitivities ")

easy_res8 = adjoint_sensitivities(solb,Tsit5(),dg,t,abstol=1e-14,
                                  reltol=1e-14,iabstol=1e-14,ireltol=1e-12,
                                  save_everystep=false,save_start=false,
                                  sensealg=SensitivityAlg(backsolve=true))

@test isapprox(res, easy_res8, rtol = 1e-9)

end_only_res = adjoint_sensitivities(sol_end,Tsit5(),dg,t,abstol=1e-14,
                                  reltol=1e-14,iabstol=1e-14,ireltol=1e-12,
                                  save_everystep=false,save_start=false,
                                  sensealg=SensitivityAlg(backsolve=true))

@test isapprox(res, end_only_res, rtol = 1e-9)

println("Calculate adjoint sensitivities from autodiff & numerical diff")
function G(p)
  tmp_prob = remake(prob,u0=convert.(eltype(p),prob.u0),p=p)
  sol = solve(tmp_prob,Tsit5(),abstol=1e-10,reltol=1e-10,saveat=t)
  A = convert(Array,sol)
  sum(((2 .- A).^2)./2)
end
G([1.5,1.0,3.0,1.0])
res2 = ForwardDiff.gradient(G,[1.5,1.0,3.0,1.0])
res3 = Calculus.gradient(G,[1.5,1.0,3.0,1.0])

using Tracker
res4 = Tracker.gradient(G,[1.5,1.0,3.0,1.0])[1]

using ReverseDiff
res5 = ReverseDiff.gradient(G,[1.5,1.0,3.0,1.0])

@test norm(res' .- res2) < 1e-8
@test norm(res' .- res3) < 1e-5
@test norm(res' .- res4) < 1e-6
@test norm(res' .- res5) < 1e-6

# check other t handling

t2 = [0.5, 1.0]
t3 = [0.0, 0.5, 1.0]
t4 = [0.5, 1.0, 10.0]

easy_res2 = adjoint_sensitivities(sol,Tsit5(),dg,t2,abstol=1e-14,
                                  reltol=1e-14,iabstol=1e-14,ireltol=1e-12)
easy_res3 = adjoint_sensitivities(sol,Tsit5(),dg,t3,abstol=1e-14,
                                  reltol=1e-14,iabstol=1e-14,ireltol=1e-12)
easy_res4 = adjoint_sensitivities(sol,Tsit5(),dg,t4,abstol=1e-14,
                                  reltol=1e-14,iabstol=1e-14,ireltol=1e-12)

function G(p,ts)
  tmp_prob = remake(prob,u0=convert.(eltype(p),prob.u0),p=p)
  sol = solve(tmp_prob,Tsit5(),abstol=1e-10,reltol=1e-10,saveat=ts)
  A = convert(Array,sol)
  sum(((2 .- A).^2)./2)
end
res2 = ForwardDiff.gradient(p->G(p,t2),[1.5,1.0,3.0,1.0])
res3 = ForwardDiff.gradient(p->G(p,t3),[1.5,1.0,3.0,1.0])
res4 = ForwardDiff.gradient(p->G(p,t4),[1.5,1.0,3.0,1.0])

@test easy_res2' ≈ res2
@test easy_res3' ≈ res3
@test easy_res4' ≈ res4

# Do a continuous adjoint problem

# Energy calculation
g(u,p,t) = (sum(u).^2) ./ 2
# Gradient of (u1 + u2)^2 / 2
function dg(out,u,p,t)
  out[1]= u[1] + u[2]
  out[2]= u[1] + u[2]
end

adj_prob = ODEAdjointProblem(sol,g,nothing,dg)
adj_sol = solve(adj_prob,Tsit5(),abstol=1e-14,reltol=1e-10)
integrand = AdjointSensitivityIntegrand(sol,adj_sol)
res,err = quadgk(integrand,0.0,10.0,atol=1e-14,rtol=1e-10)

println("Test the `adjoint_sensitivities` utility function")
easy_res = adjoint_sensitivities(sol,Tsit5(),g,nothing,dg,abstol=1e-14,
                                 reltol=1e-14,iabstol=1e-14,ireltol=1e-12)
println("2")
easy_res2 = adjoint_sensitivities(sol,Tsit5(),g,nothing,dg,abstol=1e-14,
                                  reltol=1e-14,iabstol=1e-14,ireltol=1e-12,
                                  sensealg=SensitivityAlg(quad=false))
println("3")
easy_res3 = adjoint_sensitivities(sol,Tsit5(),g,nothing,abstol=1e-14,
                                 reltol=1e-14,iabstol=1e-14,ireltol=1e-12)
println("4")
easy_res4 = adjoint_sensitivities(sol,Tsit5(),g,nothing,abstol=1e-14,
                                  reltol=1e-14,iabstol=1e-14,ireltol=1e-12,
                                  sensealg=SensitivityAlg(autodiff=false))
println("5")
easy_res5 = adjoint_sensitivities(sol,Tsit5(),g,nothing,abstol=1e-14,
                                  reltol=1e-14,iabstol=1e-14,ireltol=1e-12,
                                  sensealg=SensitivityAlg(checkpointing=true))
@test norm(easy_res .- res) < 1e-8
@test norm(easy_res2 .- res) < 1e-8
@test norm(easy_res3 .- res) < 1e-8
@test norm(easy_res4 .- res) < 1e-8
@test norm(easy_res5 .- res) < 1e-8

println("Calculate adjoint sensitivities from autodiff & numerical diff")
function G(p)
  tmp_prob = remake(prob,u0=eltype(p).(prob.u0),p=p,
                    tspan=eltype(p).(prob.tspan))
  sol = solve(tmp_prob,Tsit5(),abstol=1e-14,reltol=1e-14)
  res,err = quadgk((t)-> (sum(sol(t)).^2)./2,0.0,10.0,atol=1e-14,rtol=1e-10)
  res
end
res2 = ForwardDiff.gradient(G,[1.5,1.0,3.0,1.0])
res3 = Calculus.gradient(G,[1.5,1.0,3.0,1.0])

@test norm(res' .- res2) < 1e-8
@test norm(res' .- res3) < 1e-6

# Buffer length test
f = (du, u, p, t) -> du .= 0
p = zeros(3); u = zeros(50)
prob = ODEProblem(f,u,(0.0,10.0),p)
sol = solve(prob,Tsit5(),abstol=1e-14,reltol=1e-14)
@test_nowarn res = adjoint_sensitivities(sol,Tsit5(),dg,t,abstol=1e-14,
                                 reltol=1e-14,iabstol=1e-14,ireltol=1e-12)

println("Adjoints of u0")

function dg(out,u,p,t,i)
  out .= 1 .- u
end

ū0,adj = adjoint_sensitivities_u0(sol,Tsit5(),dg,t,abstol=1e-14,
                         reltol=1e-14,iabstol=1e-14,ireltol=1e-12)

adj2 = adjoint_sensitivities(sol,Tsit5(),dg,t,abstol=1e-14,
                        reltol=1e-14,iabstol=1e-14,ireltol=1e-12)

ū02,adj3 = adjoint_sensitivities_u0(sol,Tsit5(),dg,t,abstol=1e-14,
                         sensealg=SensitivityAlg(backsolve=true),
                         reltol=1e-14,iabstol=1e-14,ireltol=1e-12)

ū03,adj4 = adjoint_sensitivities_u0(sol,Tsit5(),dg,t,abstol=1e-14,
                        save_everystep=false, save_start=false,
                        sensealg=SensitivityAlg(backsolve=true),
                        reltol=1e-14,iabstol=1e-14,ireltol=1e-12)

ū03,adj4 = adjoint_sensitivities_u0(sol,Tsit5(),dg,t,abstol=1e-14,
                        save_everystep=false, save_start=false,
                        sensealg=SensitivityAlg(backsolve=true),
                        reltol=1e-14,iabstol=1e-14,ireltol=1e-12)

res = ForwardDiff.gradient(prob.u0) do u0
  tmp_prob = remake(prob,u0=u0)
  sol = solve(tmp_prob,Tsit5(),abstol=1e-14,reltol=1e-14,saveat=t)
  A = convert(Array,sol)
  sum(((1 .- A).^2)./2)
end

@test ū0 ≈ res rtol = 1e-10
@test ū02 ≈ res rtol = 1e-10
@test ū03 ≈ res rtol = 1e-10
@test adj ≈ adj2 rtol = 1e-10
@test adj ≈ adj3 rtol = 1e-10
@test adj ≈ adj4 rtol = 1e-10


function to_str(input)
    out = []
    for i in input
        push!(out, parse(Int, i))
    end
    return out
end

function process_input(adapters)
    sort!(adapters)

    prepend!(adapters, 0)
    append!(adapters, adapters[end] + 3)

    println(is_valid(adapters))
    count_jumps(adapters)
    # count_combs(adapters)
end

function jumps(adapters)
    jumps = []
    a = adapters[1]
    for i in adapters[2:length(adapters)]
        push!(jumps, i-a)
        a = i
    end
    return jumps
end

# I was stuck on part 2, figured this out thanks to
# dynamic programming tip from  Jonathan Paulson
# https://youtu.be/cE88K2kFZn0
comb_dict = Dict()
function num_comb(i)
    println(adapters)
    if i == length(adapters)
        return 1
    end

    if haskey(comb_dict, i)
        return comb_dict[i]
    end

    n = 0
    for j = (i+1):length(adapters)
        if adapters[j] - adapters[i] <= 3
            n += num_comb(j)
        end
    end
    comb_dict[i] = n
    return n
end

# adapters = [16, 10, 15, 5, 1, 11, 7, 19, 6, 12, 4]
# adapters = to_str(readlines("./test_input.txt"))
adapters = to_str(readlines("./adapters.txt"))

process_input(adapters)
println("Number of combinations: ", num_comb(1))


@eval mutable struct Counter
    $((:($(Symbol(op[1], typ[2])) ::Int) for op in ops for typ in typs)...)
    Counter() = new($((0 for _ in 1:length(ops)*length(typs))...))
end

# Relatively inefficient, but there should be no need for performance here...

function flop(c::Counter)
    total = 0
    for (typ, suffix) in typs
        for (name, op, cnt) in ops
            fieldname = Symbol(name, suffix)
            total += cnt * getfield(c, fieldname)
        end
    end
    total
end

import Base: ==, *, show

function Base.show(io::IO, c::Counter)
    type_names  = [typ    for (typ, _)    in typs]
    type_suffix = [suffix for (_, suffix) in typs]
    op_names    = [name   for (name, _)   in ops]

    mat = [getfield(c, Symbol(name, suffix)) for
           name   in op_names,
           suffix in type_suffix]

    fl = flop(c)
    print(io, "Flop Counter: $fl flop")
    fl == 0 && return

    print(io, "\n")
    fc(data, i) = any(data[:,i] .> 0)
    fr(data, i) = any(data[i,:] .> 0)
    pretty_table(io, mat,
                 header = type_names,
                 row_names = op_names,
                 filters_col = (fc,),
                 filters_row = (fr,),
                 newline_at_end = false)
end

function ==(c1::Counter, c2::Counter)
    all(getfield(c1, field)==getfield(c2, field) for field in fieldnames(Counter))
end

function *(n::Int, c::Counter)
    ret = Counter()
    for field in fieldnames(Counter)
        setfield!(ret, field, n*getfield(c, field))
    end
    ret
end


##

using IncrementalInference
using Test
using TensorCast
# using Statistics
# using Manifolds # should be done within regular exports

# using Pkg
# Pkg.resolve()

##

@testset "test Mixture sampling" begin

##

fg = initfg()
addVariable!(fg, :x0, ContinuousScalar)

mp = Mixture(Prior, (Normal(), Normal(10,1)),(1/2,1/2) )
addFactor!(fg, [:x0], mp)

##

pts_ = approxConv(fg, :x0f1, :x0)
@cast pts[i,j] := pts_[j][i]

N = size(pts,2)
@test 0.2*N < sum( -5 .< pts .< 5 )
@test 0.2*N < sum( 5 .< pts .< 15 )
@test sum( 15 .< pts  ) < 0.1*N
@test sum( pts .< -5 ) < 0.1*N
@test sum( 3 .< pts .< 7 ) < 0.1*N

# using KernelDensityEstimatePlotting, Gadfly
# Gadfly.set_default_plot_size(25cm,20cm)

# plotKDE(kde!(pts))

##

fg = initfg()
addVariable!(fg, :x0, ContinuousScalar)
addVariable!(fg, :x1, ContinuousScalar)

addFactor!(fg, [:x0], Prior(Normal()), graphinit=false)
initManual!(fg, :x0, [zeros(1) for _ in 1:100])

mlr = Mixture(LinearRelative, (Normal(), Normal(10,1)),(1/2,1/2) )
addFactor!(fg, [:x0;:x1], mlr, graphinit=false)

##

pts_ = approxConv(fg, :x0x1f1, :x1)
@cast pts[i,j] := pts_[j][i]

# plotKDE(kde!(pts))

##

N = size(pts,2)
@test 0.2*N < sum( -5 .< pts .< 5 )
@test 0.2*N < sum( 5 .< pts .< 15 )
@test sum( 15 .< pts  ) < 0.1*N
@test sum( pts .< -5 ) < 0.1*N
@test sum( 3 .< pts .< 7 ) < 0.1*N

##

end

@testset "test packing of Mixture" begin

##

fg = initfg()
addVariable!(fg, :x0, ContinuousScalar)
addVariable!(fg, :x1, ContinuousScalar)

mp = Mixture(Prior, [Normal(); Normal(10,1)], [0.5;0.5])
f0 = addFactor!(fg, [:x0;], mp)

mr = Mixture(LinearRelative, (fancy=manikde!(ContinuousEuclid(1), [randn(1) for _ in 1:75]), naive=Normal(0,10)), [0.4;0.6])
f1 = addFactor!(fg, [:x0;:x1], mr)

##

pf0 = DFG.packFactor(fg, f0)
pf1 = DFG.packFactor(fg, f1)

# now test unpacking
fg_ = initfg();
addVariable!(fg_, :x0, ContinuousScalar)
addVariable!(fg_, :x1, ContinuousScalar)

##

f0_ = DFG.unpackFactor(fg_, pf0)
f1_ = DFG.unpackFactor(fg_, pf1)

##

# ENV["JULIA_DEBUG"] = "DistributedFactorGraphs"
@warn("Skipping pack/unpack compareFactor test for `timezone` and `zone`")
@show typeof(f1)
@show typeof(f1_)

@show  typeof(getSolverData(f1).fnc.params);
@show typeof(getSolverData(f1_).fnc.params);

@test DFG.compareFactor(f1, f1_, skip=[:components;:labels;:timezone;:zone;:vartypes;])

@test IIF._getCCW(f1).usrfnc!.components.naive == IIF._getCCW(f1).usrfnc!.components.naive

# already ManifoldKernelDensity
A = IIF._getCCW(f1).usrfnc!.components.fancy
B = IIF._getCCW(f1_).usrfnc!.components.fancy

# A = ManifoldBelief(Euclid, IIF._getCCW(f1).usrfnc!.components.fancy )
# B = ManifoldBelief(Euclid, IIF._getCCW(f1_).usrfnc!.components.fancy )

@test mmd(A,B) < 1e-6

##

end

@testset "test simple Mixture" begin

##

fg = initfg()

addVariable!(fg, :x0, ContinuousScalar)
addVariable!(fg, :x1, ContinuousScalar)
addFactor!(fg, [:x0], Prior(Normal(0.0,0.1)))

##

# require default ::UnitScaling constructor from all factors
mlr = Mixture(LinearRelative(I), [Normal(-1.0, 0.1), Normal(1.0, 0.1)], Categorical([0.5; 0.5]))

# test serialization while we are here
pmlr = convert(PackedMixture, mlr)
umlr = convert(Mixture, pmlr)

@test mlr.mechanics == umlr.mechanics
@test mlr.components == umlr.components
@test mlr.diversity == umlr.diversity

mlr = Mixture(LinearRelative, [Normal(-1.0, 0.1), Normal(1.0, 0.1)], Categorical([0.5; 0.5]))

addFactor!(fg, [:x0,:x1], mlr)

# To look at your factor graph
# if false
# using GraphPlot
# using DistributedFactorGraphs
# plotDFG(fg)
# end

tree = solveTree!(fg)

##

btd = getBelief(getVariable(fg, :x0))
@test isapprox(mean(getKDEfit(btd,distribution=Normal)), 0.0; atol=0.1)

@test isapprox(std(getKDEfit(btd,distribution=Normal)), 0.1; atol=0.05)

btd = getBelief(getVariable(fg, :x1))
pts_ = getPoints(btd)
@cast pts[i,j] := pts_[j][i]
pts_p = pts[pts .>= 0]
pts_n = pts[pts .< 0]

nfit_p = fit(Normal, pts_p)
@test isapprox(mean(nfit_p), 1.0; atol=0.1)
@test isapprox(std(nfit_p), 0.14; atol=0.05) #TODO confirm the correct value and tolerance

nfit_n = fit(Normal, pts_n)
@test isapprox(mean(nfit_n), -1.0; atol=0.1)
@test isapprox(std(nfit_n), 0.14; atol=0.05) #TODO confirm the correct value and tolerance

# To look at your results
# if false
#   using RoMEPlotting
#   plotKDE(fg, ls(fg))
# end

##

end


using Test
using CubedSphere

B_Rancic_correct = [
    0.00000000000000,
    0.67698819751739,
    0.11847293456554,
    0.05317178134668,
    0.02965810434052,
    0.01912447304028,
    0.01342565621117,
    0.00998873323180,
    0.00774868996406,
    0.00620346979888,
    0.00509010874883,
    0.00425981184328,
    0.00362308956077,
    0.00312341468940,
    0.00272360948942,
    0.00239838086555,
    0.00213001905118,
    0.00190581316131,
    0.00171644156404,
    0.00155493768255,
    0.00141600715207,
    0.00129556597754,
    0.00119042140226,
    0.00109804711790,
    0.00101642216628,
    0.00094391366522,
    0.00087919021224,
    0.00082115710311,
    0.00076890728775,
    0.00072168382969,
    0.00067885087750
]

@testset "CubedSphere" begin
    @testset "Rančić et al. (1996) Taylor coefficients" begin
        for k in 1:length(B_Rancic_correct)
            @test CubedSphere.B_Rancic[k] ≈ B_Rancic_correct[k]
        end
    end
    @testset "Rančić et al. (1996) inverse mapping" begin
        xv=collect(0:0.1:1)
        yv=collect(0:0.1:1)
        errmax=0;
        tol = 1.e-11
        for x in xv, y in yv
          err= (x,y) .- conformal_cubed_sphere_inverse_mapping(conformal_cubed_sphere_mapping(x,y)...)
          errmax=max(errmax,sum(abs.(err)))
        end
        @test errmax < tol
    end

end


"""Filtration of an abstract complex

We call this sequence of complexes the **filtration** of `f` and
think of it as a construction by adding chunks of simplices at a time `t::FI`.

∅ = K₀ ⊆ K₁ ⊆ . . . ⊆ Kₙ = K.
"""
mutable struct Filtration{C<:AbstractComplex, FI<:AbstractFloat}
    # underlying abstract cell complex
    complex::C
    # total order of simplices as array of (dim, simplex id, filtration value)
    total::Vector{Tuple{Int,UInt,FI}}
    divisions::Number
end

order(flt::Filtration) = flt.total
complex(flt::Filtration) = flt.complex
show(io::IO, flt::Filtration{C,FI}) where {C <: AbstractComplex, FI} = print(io, "Filtration($(complex(flt)), $FI)")
eltype(flt::Filtration{C,FI}) where {C <: AbstractComplex, FI} = C
valtype(flt::Filtration{C,FI}) where {C <: AbstractComplex, FI} = FI
length(flt::Filtration) = isinf(flt.divisions) ? length(unique(e->e[3], order(flt))) : flt.divisions

#
# Constructors
#
Filtration(::Type{C}, ::Type{FI}) where {C <: AbstractComplex, FI} =
    Filtration(C(), Vector{Tuple{Int,UInt,FI}}(), Inf)
Filtration(::Type{C}) where {C <: AbstractComplex} = Filtration(C, Float64)
similar(flt::Filtration{C,FI}) where {C <: AbstractComplex, FI} = Filtration(C, FI)

"""Construct filtration from a cell complex using the order of their appearence in the complex"""
function filtration(cplx::AbstractComplex)
    idx = Vector{Tuple{Int,UInt,Float64}}()
    i = 1
    for d in 0:dim(cplx)
        for c in cells(cplx, d)
            push!(idx, (d, hash(c), i))
            i += 1
        end
    end
    return Filtration(cplx, idx, Inf)
end

"""Construct filtration from a cell complex and a complex weight function"""
function filtration(cplx::C, w::Dict{Int,Vector{FI}}; divisions::Number=Inf) where {C<:AbstractComplex, FI}
    idx = Vector{Tuple{Int,UInt,FI}}()
    for d in 0:dim(cplx)
        for c in cells(cplx, d)
            ci = position(cplx, c)
            push!(idx, (d, hash(c), w[d][ci]))
        end
    end
    sort!(idx, by=x->(x[3], x[1])) # sort by dimension & filtration value
    return Filtration(cplx, idx, divisions)
end

function push!(flt::Filtration{C,FI}, cl::AbstractCell, v::FI; recursive=false) where {C<:AbstractComplex, FI}
    cplx = complex(flt)
    @assert isa(cl, eltype(cplx)) "Complex $(cplx) does not accept $(typeof(cl))"
    cls = push!(cplx, cl, recursive=recursive)
    ord = order(flt)
    idx = length(ord) == 0 ? 1 : findlast(e->e[3]<=v, ord)
    for c in sort!(cls, by=s->dim(s))
        if idx == length(ord)
            push!(ord, (dim(c), hash(c), v))
        else
            insert!(ord, idx, (dim(c), hash(c), v))
        end
        idx += 1
    end
    return flt
end

#
# Auxiliary function
#
minimum(flt::Filtration) = order(flt)[1][3]
maximum(flt::Filtration) = order(flt)[end][3]

"""
    complex(flt, val)

Return a complex from the filtration `flt` at the filtration value `val`
"""
function complex(flt::Filtration{C,FI}, val::FI) where {C <: AbstractComplex, FI}
    cplx = complex(flt)
    res = eltype(flt)()
    for (d,si,fv) in order(flt)
        fv > val && break
        push!(res, cplx[si, d])
    end
    return res
end

"""
    boundary(flt::Filtration) -> Vector{<:AbsrtactSet}

Generate a boundary matrix from the filtration `flt` for the persistent homology calculations.
"""
function boundary(flt::Filtration; reduced=false)
    ridx = reduced ? 1 : 0
    # initialize boundary matrix
    cplx = complex(flt)
    sz = sum(size(cplx))+ridx

    # find suitable type for BM storage
    STs = [Int8, Int16, Int32, Int64]
    STi = findfirst(i->i>sz, map(typemax, STs))
    BT = STi == 1 ? BitSet : Set{STs[STi]}

    # create empty BM
    bm = map(i->BT(), 1:sz)

    # fill boundary matrix
    ord = order(flt)
    revidx = Dict((ci, d) => i for (i, (d, ci, fv)) in enumerate(ord))
    for (i, (d, ci, fv)) in enumerate(ord)
        if d > 0
            for face in faces(cplx, ci)
                push!(bm[i+ridx], revidx[(face, d-1)]+ridx)
            end
        elseif reduced
            push!(bm[i+ridx], 1)
        end
    end
    return bm
end

function sparse(∂::Vector{<:AbstractSet})
    m = length(∂)
    ret = spzeros(Int, m, m)
    for i in 1:m
        bm = ∂[i]
        for (l, j) in enumerate(bm)
            ret[j,i] = j # (-1)^((l-1)%2) # coefficients require exact order of faces in provides simplex
        end
    end
    return ret
end

"""
    similarity(flt::Filtration)

Returns a similarity matrix created from 1-subcomplex of simplicial complex and distance weights of the filtration `flt`.
"""
function similarity(flt::Filtration)
    cplx = complex(flt)
    ord = order(flt)
    C0 = map(hash, cells(cplx, 0))
    N = length(C0)
    @assert N > 0 "Complex should not be empty"
    adj = spzeros(valtype(flt),N,N)
    for c in cells(cplx, 1)
        i, j = map(h->findfirst(isequal(h), C0), vertices(c))
        idx = findfirst(v->v[1]==1 && v[2] == hash(c), ord)
        adj[i, j] = adj[j, i] = ord[idx][3]
    end
    return adj
end

#
# I/O
#
function write(io::IO, flt::Filtration)
    cplx = complex(flt)
    for (d, ci, fv) in order(flt)
        for k in indices(cplx[ci,d])
            write(io, "$k,")
        end
        write(io, "$fv\n")
    end
end

function read(io::IO, ::Type{Filtration{C,FI}}, ::Type{CL}) where {C <: AbstractComplex, FI, CL<:AbstractCell}
    flt = Filtration(C,FI)
    while !eof(io)
        l = readline(io)
        vals = split(l, ',')
        sval = parse(CL, join(vals[1:end-1], ' '))
        fval = parse(FI, vals[end])
        push!(flt, sval, fval)
    end
    return flt
end

"Write a combined boundary matrix to a text file"
function writeboundarymatrix(io::IO, bm::Vector, zeroindex = true)
    for smplx in bm
        if length(smplx) == 0
            write(io, "0")
        else
            write(io, "$(length(smplx)-1)")
        end
        for i in smplx
            write(io, " $(zeroindex ? i-1 : i)")
        end
        write(io, 0x0A)
    end
end

function writebin(io::IO, flt::Filtration)
    cplx = complex(flt)
    for (d, ci, fv) in order(flt)
        write(io, UInt8(d))
        for k in indices(cplx[ci,d])
            write(io, UInt32(k))
        end
        write(io, fv)
    end
end

function readbin(io::IO, ::Type{Filtration{C,FI}}, ::Type{CL}, ::Type{ST}; maxoutdim=1) where {C <: AbstractComplex, FI, CL<:AbstractCell, ST}
    flt = Filtration(C,FI)
    while !eof(io)
        d = read(io, UInt8)
        vals = ST[read(io, UInt32) for i in 1:(d+1)]
        fval = read(io, FI)
        if d <= maxoutdim
            push!(flt, CL(vals), fval)
        end
    end
    return flt
end


using Test, DataFrames, PairsMacros
df = DataFrame(x = [1, 2], y = [3, 4], z = [5, 6])

# one arg
@test (@cols(z = sum(x))) == (:x => sum => :z)
@test (@cols(sum(x))) == (:x => sum)
@test (@cols(z = x)) == (:x => :z)
@test (@cols(x)) == :x

@test select(df, @cols(x)).x == df.x
@test (@cols(z = sum(skipmissing(x)))) == (:x => sum ∘ skipmissing => :z)
@test transform(df, @cols(z = exp.(x))).z == exp.(df.x)
@test transform(df, @cols(z = 1)).z == fill(1, size(df, 1))
@test transform(df, @cols(z = sum([1, 2, 3]))).z == fill(sum([1, 2, 3]), size(df, 1))
@test size(filter(@cols(x > 1), df), 1) == 1
@test size(filter(@cols((x > 1) & (y < 3)), df), 1) == 0

# multiple args
@test transform(df, @cols sum(x) first(x)).x_sum == [3, 3]
@test transform(df, @cols xfirst = first(x) xlast = last(x)).xfirst == [1, 1]

@test transform(df, @rows(z = exp.(x))).z == exp.(df.x)
@test isequal(transform(df, @rows(z2 = x >= 2 ? missing : z)).z2, [5, missing])

# rows
@test (@rows(z = rand())) == (Any[] => (DataFrames.ByRow{typeof(rand)}(rand) => :z))
@test transform(df, @rows(z = x - y)).z  == df.x .- df.y
@test transform(df, @rows(z = y - x)).z  == df.y .- df.x
@test transform(df, @rows(z = x - x)).z  == df.x .- df.x
@test transform(df, @rows abs(x) x^2).x_abs == abs.(df.x)
@test transform(df, @rows(z = x^2)).z  == df.x.^2
@test transform(df, @rows abs(x) x^2).x_abs  == abs.(df.x)

df = DataFrame(x = [[-1, -2], [-3, -4]])
@test transform(df, @rows(z = abs.(x))).z == [abs.(x) for x in df.x]

# test $
u = :y
@test (@cols(z = sum($u))) == (:y => sum => :z)
@test (@cols(z = sum($u))) == (:y => sum => :z)
@test (@rows($u = x)) == (:x => :y)
@test (@cols(z = sum(skipmissing($u)))) == (:y => sum ∘ skipmissing => :z)
u = [:y]
@test (@cols(z = sum($(u[1])))) == (:y => sum => :z)

# test obj
u = [3, 4]
@test combine(df, @cols(z = sum(esc(u)))).z == [sum(u)]

# test hygiene
mean = x -> x
@test transform(df, @cols(z = mean(x))).z == df.x

# weird symbols
df = DataFrame(x = [1, 2])
@test transform!(df, @cols($("ok now") = x))."ok now" == df.x
@test transform(df, @cols(z = sum($("ok now")))).z == [3, 3]


include("velocity.jl")

mutable struct PSO <: AbstractParameters
    N::Int
    C1::Float64
    C2::Float64
    ω::Float64
    v::Array{Float64} # velocity
    flock::Array
end

"""
    PSO(;
        N  = 0,
        C1 = 2.0,
        C2 = 2.0,
        ω  = 0.8,
        information = Information(),
        options = Options()
    )

Parameters for Particle Swarm Optimization (PSO) algorithm: learning rates `C1` and `C2`,
`N` is the population size and `ω` controls the inertia weight.
# Example

```jldoctest
julia> f(x) = sum(x.^2)
f (generic function with 1 method)

julia> optimize(f, [-1 -1 -1; 1 1 1.0], PSO())
+=========== RESULT ==========+
  iteration: 1000
    minimum: 1.40522e-49
  minimizer: [3.0325415595139883e-25, 1.9862212295897505e-25, 9.543772256546461e-26]
    f calls: 30000
 total time: 0.1558 s
+============================+

julia> optimize(f, [-1 -1 -1; 1 1 1.0], PSO(N = 100, C1=1.5, C2=1.5, ω = 0.7))
+=========== RESULT ==========+
  iteration: 300
    minimum: 2.46164e-39
  minimizer: [-3.055334698085433e-20, -8.666986835846171e-21, -3.8118413472544027e-20]
    f calls: 30000
 total time: 0.1365 s
+============================+
```

"""
function PSO(;
    N::Int = 0,
    C1 = 2.0,
    C2 = 2.0,
    ω = 0.8,
    v = Float64[],
    flock = xf_indiv[],
    information = Information(),
    options = Options(),
)

parameters = PSO(N, promote(Float64(C1), C2, ω)..., v, flock)

Algorithm(
    parameters,
    information = information,
    options = options,
)
end

function update_state!(
        status,
        parameters::PSO,
        problem::AbstractProblem,
        information::Information,
        options::Options,
        args...;
        kargs...
    )
    xGBest = get_position(status.best_sol)

    X_new = zeros(parameters.N, size(problem.bounds, 2))

    # For each elements in population
    for i in 1:parameters.N
        x = get_position(parameters.flock[i])
        xPBest = get_position(status.population[i])
        parameters.v[i, :] = velocity(x, parameters.v[i, :], xPBest, xGBest, parameters)
        x += parameters.v[i, :]
        reset_to_violated_bounds!(x, problem.bounds)
        X_new[i,:] = x
    end

    for (i, sol) in enumerate(create_solutions(X_new, problem;ε = options.h_tol))
        if is_better(sol, status.population[i])
            status.population[i] = sol
            if is_better(sol, status.best_sol)
                status.best_sol = sol
                xGBest = get_position(status.best_sol)
            end
        end

        parameters.flock[i] = sol

        # stop condition
        stop_criteria!(status, parameters, problem, information, options)
        status.stop && break
    end

end

function initialize!(
    status,
    parameters::PSO,
    problem::AbstractProblem,
    information::Information,
    options::Options,
    args...;
    kargs...
       )
    D = size(problem.bounds, 2)

    if parameters.N == 0
        parameters.N = 10 * D
    end

    if options.f_calls_limit == 0
        options.f_calls_limit = 10000D
        options.debug &&
        @warn( "f_calls_limit increased to $(options.f_calls_limit)")
    end

    if options.iterations == 0
        options.iterations = div(options.f_calls_limit, parameters.N) + 1
    end

    status = gen_initial_state(problem,parameters,information,options,status)

    parameters.v = zeros(parameters.N, D)

    parameters.flock = status.population

    status

end

function final_stage!(
    status,
    parameters::PSO,
    problem::AbstractProblem,
    information::Information,
    options::Options,
    args...;
    kargs...
    )
    status.final_time = time()

end


module GuessworkQuantumSideInfo

using LinearAlgebra, SparseArrays # standard libraries
using Random: Random, randperm!, AbstractRNG # standard library
using Convex # Julia SDP solvers / interfaces
using Combinatorics: multiset_permutations, combinations # used for POVM outcomes
using UnPack # helper

# Basic functions for quantum states and examples
export ket, bra, ⊗, I, dm, randdm, randprobvec, iid_copies, BB84_states
include("quantum_states.jl")

# The basic SDP routine
export guesswork
include("SDP_formulation.jl")

export guesswork_lower_bound
include("lower_bound.jl")

# Re-run the SDP adding constraints one at a time
export guesswork_upper_bound
include("upper_bound.jl")

# reformulate as a mixed-integer SDP
export guesswork_MISDP
include("MISDP_formulation.jl")

# Utilities for analyzing the POVMs resulting from the SDPs
export pmfN
include("analyze_measurements.jl")

end


using MathOptInterface, LinearAlgebra, Statistics, Plots, QuantEcon, Interpolations, NLsolve, Optim, Random, IterTools, JuMP, Ipopt, MathOptInterface;

u(x) = sum(log(x[i]) for i in 1:numSectors)
w(x) = sum(5*log(x[i]) for i in 1:numSectors)

β = 0.96
α = 0.4
f(x) = x^α

numSectors = 2

numPoints1D = 2

grid = Vector{Vector{Float64}}(undef,(numPoints1D)^numSectors)

gridMax = 3
gridMin = 1

iter=1
if numSectors == 1
	grid = LinRange(gridMin, gridMax, numPoints1D)
elseif numSectors == 2
for p in product(LinRange(gridMin,gridMax,numPoints1D),LinRange(gridMin,gridMax,numPoints1D))
    grid[iter] = collect(p)
    global iter += 1
end
else
	print("Currently we only allow one or two sectors")
end
wVal = w.(grid)

function interp(x,y)
    return interpolate((LinRange(gridMin,gridMaxnumPoints1D),LinRange(gridMin,gridMax,numPoints1D)), wVals, Gridded(Linear()))(x,y)
end

function T(wVal, grid, β, f ; compute_policy = false)
	if numSectors == 1
		wVals = zeros(numPoints1D)
		wVals = wVal
		wFunc(y) = interpolate(grid)
	elseif numSectors == 2
		wVals = zeros(numPoints1D,numPoints1D);
		for j in numPoints1D
		    for i in numPoints1D
                        wVals[i,j] = wVal[i+(i-1)*j]
                    end
                end
		wFunc(x,y) = log(x) + log(y)
		#wFunc(x, y) = interpolate((LinRange(gridMin,gridMax,numPoints1D),LinRange(gridMin,gridMax,numPoints1D)),			  wVals, Gridded(Linear()))(x,y)
		print("wFunc created")
		print(wFunc(1.5,1.5))
	else
		print("Can only handle upto 2 sectors")
	end
    print(wFunc(1.5, 1.5))
    global Tw = zeros(length(grid))
    global σ = similar(grid)
    for n in 1:length(grid)
        y = grid[n]
        modTrial = Model(Ipopt.Optimizer);
        @variable(modTrial,  c[1:numSectors] >= 0.0001)
        @variable(modTrial, k[1:numSectors] >= 0.0001, start=2)
        for i in 1:numSectors
            @constraint(modTrial, gridMin <= c[i] <= y[i])
	    @constraint(modTrial, k[i] == y[i] - c[i])
        end
        register(modTrial, :wFunc, numSectors, wFunc; autodiff = true)
        register(modTrial, :f, 1, f; autodiff = true)
	if numSectors == 1
		@NLobjective(modTrial, Max, sum(log(c[i]) for i in 1:numSectors) + β*wFunc(f(k[1])))
	elseif numSectors == 2
        @NLobjective(modTrial, Max, sum(log(c[i]) for i in 1:numSectors) + β*wFunc(f(k[1]),f(k[2])))
	else
	print("check the number of sectors")
	end
        optimize!(modTrial)
        Tw[n] = JuMP.objective_value(modTrial)
        if compute_policy
            σ[n] = value.(c)
        end
    end
    if compute_policy
        return Tw, σ
    end

    return Tw

end

wVal = T(wVal, grid, β, f; compute_policy = false)
#=
u(c) = log(c[1]^0.5*c[2]^0.5);

numSectors == 2;

w(x,y) = x^2+y^2;

valGrid = zeros(numPoints1D,numPoints1D)
for i in 1:numPoints1D
    for j in 1:numPoints1D
        valGrid[i,j] = w(grid[j+(i-1)*numPoints1D][1],grid[j+(i-1)*numPoints1D][2])
    end
end

valGrid = zeros(numPoints1D,numPoints1D)
for i in 1:numPoints1D
    for j in 1:numPoints1D
        valGrid[i,j] = w(grid[j+(i-1)*numPoints1D][1],grid[j+(i-1)*numPoints1D][2])
    end
end

α = 0.4
β = 0.96
μ = 0
s = 0.1
ϵ_D = 1
ϵ_X = 0.1
Γ = [0.7 0.3; 0.3 0.7]

#initial capital flows
x= [1 2; 3 4]
#Investment is a function from R^2J to R^J
X = zeros(numSectors,1)
XInter = zeros(numSectors,numSectors)
for J in numSectors;
    for I in numSectors;
        XInter[I,J] = (Γ[I,J]^(1/ϵ_X)*x[I,J]^((ϵ_X-1)/ϵ_X))
    end
    X[J] = sum(XInter[:,J])^(ϵ_X/(ϵ_X-1))
end

K = zeros(numSectors,1)
Q = zeros(numSectors,1)
C = zeros(numSectors,1)

#Kapital is function from R^J to R^J
K = X;

Q = K.^α

C = Q - sum(x[:,J] for J in 1:numSectors)

#=
c1 = log(1 - α * β) / (1 - β)
c2 = (μ + α * log(α * β)) / (1 - α)
c3 = 1 / (1 - β)
c4 = 1 / (1 - α * β)
=#

# Utility is a function from R^J to R
u(C) = log(sum(C[j]^((ϵ_D-1)/ϵ_D) for j in 1:numSectors)^((ϵ_D)/(ϵ_D*0.999-1)));

# Deterministic part of production function a function from R^J to R^J
f(k) = k.^α

#= True optimal policy
c_star(y) = (1 - α * β) * y

# True value function
v_star(y) = c1 + c2 * (c3 - c4) + c4 * log(y)
=#

Random.seed!(42) # For reproducible results.

#=
grid = Array{Float64}[]
for n in 1:numSectors
    x = 1:numSectors
    push!(grid, x)
end
=#

#=
grid[j] = 1:numSectors
y = 1:numSectors
x' .* ones(5)
ones(3)' .* y
g         # Largest grid point
grid_size = 200      # Number of grid points
shock_size = [250,250]     # Number of shock draws in Monte Carlo integral

grid_y = range.(1e-5,  gridMax, length = grid_size)
shocks = exp.(μ .+ s * randn(shock_size))
w = T(v_star.(grid_y), grid_y, β, (log∘sum), k -> k^α, shocks)

#=
plt = plot(ylim = (-35,-24))
plot!(plt, grid_y, w, linewidth = 2, alpha = 0.6, label = "T(v_star)")
plot!(plt, v_star, grid_y, linewidth = 2, alpha=0.6, label = "v_star")
plot!(plt, legend = :bottomright)
=#
w = 5 * log.(grid_y)  # An initial condition -- fairly arbitrary
n = 35
#=
plot(xlim = (extrema(grid_y)), ylim = (-50, 10))
lb = "initial condition"
plt = plot(grid_y, w, color = :black, linewidth = 2, alpha = 0.8, label = lb)
for i in 1:n
    global w = T(w, grid_y, β, log, k -> k^α, shocks)
    plot!(grid_y, w, color = RGBA(i/n, 0, 1 - i/n, 0.8), linewidth = 2, alpha = 0.6,
          label = "")
end

lb = "true value function"
plot!(plt, v_star, grid_y, color = :black, linewidth = 2, alpha = 0.8, label = lb)
plot!(plt, legend = :bottomright)
=#

function solve_optgrowth(initial_w; tol = 1e-6, max_iter = 500)
    fixedpoint(w -> T(w, grid_y, β, u, f, shocks), initial_w).zero # gets returned
end

#=
initial_w = 5 * log.(grid_y)
v_star_approx = solve_optgrowth(initial_w)

plt = plot(ylim = (-35, -24))
plot!(plt, grid_y, v_star_approx, linewidth = 2, alpha = 0.6,
      label = "approximate value function")
plot!(plt, v_star, grid_y, linewidth = 2, alpha = 0.6, label = "true value function")
plot!(plt, legend = :bottomright)

Tw, σ = T(v_star_approx, grid_y, β, log, k -> k^α, shocks;
                         compute_policy = true)
cstar = (1 - α * β) * grid_y

plt = plot(grid_y, σ, lw=2, alpha=0.6, label = "approximate policy function")
plot!(plt, grid_y, cstar, lw = 2, alpha = 0.6, label = "true policy function")
plot!(plt, legend = :bottomright)
=#
=#
=#


using ValueOrientedRiskManagementInsurance

using Distributions
using DataFrames

include("S2NonLife_Input.jl")

#################################################################
println("Start S2NonLife.jl")

lobs = Array(NLLob, 0)

for lob in select_lob_names
  push!(lobs, NLLob(df_lobs[df_lobs[:name] .== lob, :],
                    convert(Array, β_raw[lob]),
                    rf,
                    corr_prem_res, df_lobs_prem_risk))
end

## premium reserve risk #########################################
scr_prem_res = 3 * premrestotalsd(lobs, corr_lob)[1]
scr_lapse = 0.0

## catastrophe risk #############################################
## proportional reinsurance seems to be ignored for cat scr

## natural catastrophe risk -------------------------------------
scr_cat_nat = 0.0   ## no natural catastrophe risk in portfolio

## man-made catastrophe risk ------------------------------------
## only fire risk and liability risk in "man-made" portfolio

## fire
scr_cat_fire = sum(cat_fire_is) ## no risk management mechanisms

## liability
prem_liab_vec = Array(Float64, nrow(df_cat_liability))
scr_cat_liab_vec = Array(Float64, nrow(df_cat_liability))
cat_liab_grp = Array(Int, 0)
for lob in lobs
  if lob.name == :liability
    prem_liab_vec = lob.prem_gross_cy * cat_liability_mix
    scr_cat_liab_vec =
      convert(Array,
              prem_liab_vec .* df_cat_liability[:,:f_liab])
    for i = 1:length(prem_liab_vec)
      if prem_liab_vec[i] > 0
        push!(cat_liab_grp, i)
      end
    end
  end
end
scr_cat_liab =
  sqrt(scr_cat_liab_vec' * corr_liab * scr_cat_liab_vec)[1,1]

## total man-made catastrophe risk ------------------------------
scr_cat_man_made_vec = [scr_cat_fire, scr_cat_liab]
scr_cat_man_made =
  sqrt(scr_cat_man_made_vec ⋅ scr_cat_man_made_vec)

## other catastrophe risk ---------------------------------------
scr_cat_other = 0.0

## catastrophe risk for non-prop reinsurance --------------------
scr_cat_nprop = 0.0  ## XYZ does not act as reinsurer

## total catastrophe risk ---------------------------------------
scr_cat = sqrt((scr_cat_nat + scr_cat_nprop)^2 +
                 scr_cat_man_made^2 +
                 scr_cat_other^2)

## Total SCR ####################################################
scr_nl_vec = [scr_prem_res, scr_lapse, scr_cat]
scr_nl = sqrt(scr_nl_vec ⋅ (corr_scr * scr_nl_vec))

## main results #################################################
println("scr_prem_risk        :  $(round(scr_prem_res, 2))")
println("scr_lapse            :  $(round(scr_lapse, 2))")
println("scr_cat_man_made_fire:  $(round(scr_cat_fire, 2))")
println("scr_cat_man_made_liab:  $(round(scr_cat_liab, 2))")
println("scr_cat_man_made     :  $(round(scr_cat_man_made, 2))")
println("scr_cat_other        :  $(round(scr_cat_other, 2))")
println("scr_cat_nprob        :  $(round(scr_cat_nprop, 2))")
println("scr_cat              :  $(round(scr_cat, 2))")
println("scr_total            :  $(round(scr_nl, 2))")

println("End S2NonLife.jl")


p1 = PetriCospanOb(1)
p2 = PetriCospanOb(2)
pc1 = id(p1)
pc2 = id(p2)
I = munit(PetriCospanOb)

@test dom(pc1 ⊗ pc2) == dom(pc1) ⊗ dom(pc2)
@test codom(pc1 ⊗ pc2) == codom(pc1) ⊗ codom(pc2)
@test dom(σ(p1,p2)) == p1 ⊗ p2
@test codom(σ(p1,p2)) == p2 ⊗ p1
compare_petricospan(σ(p1, p2), braid(p1,p2))

@test dom(Δ(p1)) == p1
@test codom(Δ(p1)) == p1 ⊗ p1
@test dom(◊(p1)) == p1
@test codom(◊(p1)) == I

compare_petricospan(pair(pc1,pc1), Δ(p1) ⋅ (pc1 ⊗ pc1))
compare_petricospan(proj1(p1,p2), pc1 ⊗ ◊(p2))
compare_petricospan(proj2(p1,p2), ◊(p1) ⊗ pc2)


using Measurements

d(a_in, a_out) = (1 - a_in*a_in) / (1 - a_in*a_out) * a_out - a_in
d̃(a_in, a_out) = a_out - a_in

# Au(111)
# a_in, a_out...

P_in = [
    # Au(111)
     22    0  -22
    # Au(poly.)
     -4    2    6
     -4    2    6
     -4    2    6
    # Au(poly.) (not reported)
     -4    2    6
    # Al(poly.)
      0    0    0
    # Au(poly.) (not reported)
     -4    2    6
     -4    2    6
     -4    2    6
    # Cu(332)
      0    0    0
      0    0    0
    # Ag(110)
      3    0   -3
      3    0   -3
    # Au(110)
     27    0  -24
     27    0  -24
    # Au(110)
    missing  0  missing
    missing  0  missing
] ./ 100

σ_in = [
    5  5  5
    5  5  5
    5  5  5
    5  5  5
    5  5  5
    5  5  5
    5  5  5
    5  5  5
    5  5  5
    5  5  5
    5  5  5
    5  5  5
    5  5  5
    5  5  5
    5  5  5
    missing  5  missing
    missing  5  missing
] ./ 100

P_out = [
    # Au(111)
    -29 -31 -35
    # Au(poly.)
    -40 -38 -35
    -39 -36 -32
    -61 -57 -55
    # Au(poly.)
      8  14  17
    # Al(poly.)
     14.9 15 15.1
    # Au(poly.)
    -17 -11 -6
    -20 -14 -10
    -22 -18 -14
    # Cu(332)
    -6.3 -6.7 -6
     11  12  12
    # Ag(110)
     -7  -9 -12
    6.2 7.1 6.3
    # Au(111)
     22  -8 -35
     34   8 -16
    # Au(111)
    missing  -12  missing
    missing    4  missing
] ./ 100

σ_out = [
    5  5  5
    5  5  5
    5  5  5
    5  5  5
    5  5  5
    5  5  5
    5  5  5
    5  5  5
    5  5  5
    5  5  5
    5  5  5
    5  5  5
    5  5  5
    5  5  5
    5  5  5
    missing  5  missing
    missing  5  missing
] ./ 100

# Polycrystalline
# a_in, a_out...

D(a_in, a_out, idx=2) = (a_out[idx] - a_in[idx]) / (1 - a_in[idx]*a_out[idx])
A(idx, a_in, a_out) = (1 - a_in[idx]*a_out[idx]) * D(a_in, a_out) - (a_out[idx] - a_in[idx])
ΔT(a_in, a_out) = (A(3, a_in, a_out) + A(1, a_in, a_out)) / (a_out[3] - a_out[1])
ΔD(a_in, a_out) = ΔT(a_in, a_out) * (a_out[3] + a_out[1])/2 - (A(3, a_in, a_out) - A(1, a_in, a_out)) / 2

D2(idx, a_in, a_out) = (a_out[idx] - a_in[idx] + (idx-2)*(ΔT(a_in, a_out)*a_out[idx] - ΔD(a_in, a_out))) / (1 - a_in[idx]*a_out[idx])

Dcons(a_in, a_out) = (D(a_in, a_out, 1) + D(a_in, a_out, 3)) / 2 - D(a_in, a_out, 2)


using JLD2
using FileIO

# tensor for classical 2-d model
const isingβc = log(1+sqrt(2))/2

"""
    tensorfromclassical(h::Matrix)

given a classical 2-body hamiltonian `h`, return the corresponding tensor
for use in e.g. `trg` for a two-dimensional square-lattice.

# Example
```julia
julia> model_tensor(Ising(),β) ≈ tensorfromclassical([β -β; -β β])

true
```
"""
function tensorfromclassical(ham::Matrix)
    wboltzmann = exp.(ham)
    q = sqrt(wboltzmann)
    ein"ij,ik,il,im -> jklm"(q,q,q,q)
end

"""
    model_tensor(::Ising,β)

return the isingtensor at inverse temperature `β` for a two-dimensional
square lattice tensor-network.
"""
function model_tensor(::Ising, β::Real)
    a = reshape(Float64[1 0 0 0; 0 0 0 0; 0 0 0 0; 0 0 0 1], 2,2,2,2)
    cβ, sβ = sqrt(cosh(β)), sqrt(sinh(β))
    q = 1/sqrt(2) * [cβ+sβ cβ-sβ; cβ-sβ cβ+sβ]
    ein"abcd,ai,bj,ck,dl -> ijkl"(a,q,q,q,q)
end

"""
    mag_tensor(::Ising,β)

return the operator for the magnetisation at inverse temperature `β`
at a site in the two-dimensional ising model on a square lattice in tensor-network form.
"""
function mag_tensor(::Ising, β)
    a = reshape(Float64[1 0 0 0; 0 0 0 0; 0 0 0 0; 0 0 0 -1] , 2,2,2,2)
    cβ, sβ = sqrt(cosh(β)), sqrt(sinh(β))
    q = 1/sqrt(2) * [cβ+sβ cβ-sβ; cβ-sβ cβ+sβ]
    ein"abcd,ai,bj,ck,dl -> ijkl"(a,q,q,q,q)
end

"""
    energy_tensor(::Ising,β)

return the operator for the energy at inverse temperature `β`
at a site in the two-dimensional ising model on a square lattice in tensor-network form.
"""
function energy_tensor(::Ising, β)
    ham = [-1 1;1 -1]
    wboltzmann = exp.(-β .* ham)
    wenergy = ham .* wboltzmann
    qb = sqrt(wboltzmann)
    ein"as,si,bi,ci,di -> abcd"(qb^(-1),wenergy,qb,qb,qb)
end

"""
    vumps_env(model<:HamiltonianModel, β, D)

return the vumps environment of the `model` as a function of the inverse
temperature `β` and the environment bonddimension `D` as calculated with
vumps. Save `env` in file `./data/model_β_D.jld2`. Requires that `model_tensor` are defined for `model`.
"""
function vumps_env(model::MT, β, D; tol=1e-10, maxiter=20, verbose = false) where {MT <: HamiltonianModel}
    M = model_tensor(model, β)
    mkpath("./data/")
    chkp_file = "./data/$(model)_β$(β)_D$(D).jld2"
    if isfile(chkp_file)
        rt = SquareVUMPSRuntime(M, chkp_file, D; verbose = verbose)
    else
        rt = SquareVUMPSRuntime(M, Val(:random), D; verbose = verbose)
    end
    env = vumps(rt; tol=tol, maxiter=maxiter, verbose = verbose)
    # save(chkp_file, "env", env) # if forward steps is too small, backward go to this way will make mistake!!!
    return env
end

"""
    Z(env::SquareVUMPSRuntime)

return the partition function of the `env`.
"""
function Z(env::SquareVUMPSRuntime)
    M,AL,C,FL,FR = env.M,env.AL,env.C,env.FL,env.FR
    AC = ein"asc,cb -> asb"(AL,C)
    z = ein"αcβ,βsη,cpds,ηdγ,αpγ ->"(FL,AC,M,FR,conj(AC))[]
    λ = ein"αcβ,βη,ηcγ,αγ ->"(FL,C,FR,conj(C))[]
    return z/λ
end

"""
    magnetisation(env::SquareVUMPSRuntime, model::MT, β)

return the magnetisation of the `model`. Requires that `mag_tensor` are defined for `model`.
"""
function magnetisation(env::SquareVUMPSRuntime, model::MT, β) where {MT <: HamiltonianModel}
    M,AL,C,FL,FR = env.M,env.AL,env.C,env.FL,env.FR
    Mag = mag_tensor(model, β)
    AC = ein"asc,cb -> asb"(AL,C)
    mag = ein"αcβ,βsη,cpds,ηdγ,αpγ ->"(FL,AC,Mag,FR,conj(AC))[]
    λ = ein"αcβ,βsη,cpds,ηdγ,αpγ ->"(FL,AC,M,FR,conj(AC))[]

    return abs(mag/λ)
end

"""
    energy(env::SquareVUMPSRuntime, model::MT, β)

return the energy of the `model` as a function of the inverse
temperature `β` and the environment bonddimension `D` as calculated with
vumps. Requires that `model_tensor` are defined for `model`.
"""
function energy(env::SquareVUMPSRuntime, model::MT, β) where {MT <: HamiltonianModel}
    M,AL,C,FL,FR = env.M,env.AL,env.C,env.FL,env.FR
    Ene = energy_tensor(model, β)
    AC = ein"asc,cb -> asb"(AL,C)
    energy = ein"αcβ,βsη,cpds,ηdγ,αpγ ->"(FL,AC,Ene,FR,conj(AC))[]
    λ = ein"αcβ,βsη,cpds,ηdγ,αpγ ->"(FL,AC,M,FR,conj(AC))[]

    return energy/λ*2 # factor 2 for counting horizontal and vertical links
end

"""
    magofβ(::Ising,β)

return the analytical result for the magnetisation at inverse temperature
`β` for the 2d classical ising model.
"""
magofβ(::Ising, β) = β > isingβc ? (1-sinh(2*β)^-4)^(1/8) : 0.

"""
    magofdβ(::Ising,β)

return the analytical result for the derivative of magnetisation at inverse temperature
`β` for the 2d classical ising model.
"""
magofdβ(::Ising, β) = β > isingβc ? (coth(2*β)*csch(2*β)^4)/(1-csch(2*β)^4)^(7/8) : 0.


const pi = 3.14159265358979323846264338327   #Pi
const grav = 9.8                               #Gravitational acceleration (m / s^2)
const cp = 1004.0                             #Specific heat of dry air at constant pressure
const cv = 717.0                              #Specific heat of dry air at constant volume
const rd = 287.0                              #Dry air constant for equation of state (P=rho*rd*T)
const p0 = 1.e5                              #Standard pressure at the surface in Pascals
const C0 = 27.5629410929725921310572974482   #Constant to translate potential temperature into pressure (P=C0*(rho*theta)**gamma)
const gamma = 1.40027894002789400278940027894   #gamma=cp/Rd
#Define domain and stability-related constants
const xlen = 2.e4    #Length of the domain in the x-direction (meters)
const zlen = 1.e4    #Length of the domain in the z-direction (meters)
const hv_beta = 0.25     #How strong to diffuse the solution: hv_beta \in [0:1]
const cfl = 1.50    #"Courant, Friedrichs, Lewy" number (for numerical stability)
const max_speed = 450        #Assumed maximum wave speed during the simulation (speed of sound + speed of wind) (meter / sec)
const hs = 2          #"Halo" size: number of cells beyond the MPI tasks's domain needed for a full "stencil" of information for reconstruction
const sten_size = 4          #Size of the stencil used for interpolation

#Parameters for indexing and flags
const NUM_VARS = 4           #Number of fluid state variables
const ID_DENS = 1           #index for density ("rho")
const ID_UMOM = 2           #index for momentum in the x-direction ("rho * u")
const ID_WMOM = 3           #index for momentum in the z-direction ("rho * w")
const ID_RHOT = 4           #index for density * potential temperature ("rho * theta")
const DIR_X = 1              #Integer constant to express that this operation is in the x-direction
const DIR_Z = 2              #Integer constant to express that this operation is in the z-direction