struct SymTen{T<:Number} <: AbstractTen{T}
    xx::T
    xy::T
    xz::T
    yy::T
    yz::T
    zz::T
end

Base.IndexStyle(a::Type{<:SymTen}) = Base.IndexLinear()

@inline function Base.getindex(a::SymTen,I::Integer)
    I <= 3 && return getfield(a,I)
    I == 4 && return a.xy
    I == 5 && return a.yy
    I == 6 && return a.yz
    I == 7 && return a.xz
    I == 8 && return a.yz
    I == 9 && return a.zz
end

@inline Base.size(a::SymTen) =
    (3,3)

@inline Base.length(a::SymTen) =
    9

@inline Base.zero(a::Type{SymTen{T}}) where {T} =
    SymTen{T}(zero(T),zero(T),zero(T),zero(T),zero(T),zero(T))

@inline SymTen(xx::T,xy::T,xz::T,yy::T,yz::T,zz::T) where {T} =
    SymTen{T}(xx,xy,xz,yy,yz,zz)

@inline SymTen(q,w,e,r,t,y) =
    SymTen(promote(q,w,e,r,t,y)...)

@inline SymTen(s::SymTen) = SymTen(s.xx,s.xy,s.xz,s.yy,s.yz,s.zz)

@inline Base.:+(a::SymTen{T},b::SymTen{T2}) where {T,T2} =
    SymTen{promote_type(T,T2)}(a.xx+b.xx, a.xy+b.xy, a.xz+b.xz, a.yy+b.yy, a.yz + b.yz, a.zz+b.zz)

@inline Base.:-(A::SymTen{T}) where T = SymTen{T}(-A.xx,-A.xy,-A.xz,-A.yy,-A.yz,-A.zz)

@inline Base.:-(a::SymTen{T},b::SymTen{T2}) where {T,T2} =
    SymTen{promote_type(T,T2)}(a.xx-b.xx, a.xy-b.xy, a.xz-b.xz, a.yy-b.yy, a.yz - b.yz, a.zz - b.zz)

@inline Base.:*(a::T,b::SymTen{T2}) where {T<:Number,T2<:Number} =
    SymTen{promote_type(T,T2)}(a*b.xx, a*b.xy, a*b.xz, a*b.yy, a*b.yz, a*b.zz)

@inline Base.:*(b::SymTen{T},a::T2) where {T<:Number,T2<:Number} =
    a*b

@inline Base.:/(b::SymTen{T},a::T2) where {T<:Number,T2<:Number} =
    inv(a)*b

@inline Base.:(:)(a::SymTen,b::SymTen) =
    muladd(a.xx, b.xx, muladd(a.yy, b.yy, a.zz*b.zz)) + 2muladd(a.xy, b.xy, muladd(a.xz, b.xz, a.yz*b.yz))

@inline Base.:(+)(L::LinearAlgebra.UniformScaling,A::SymTen) =
    SymTen(A.xx + L.λ, A.xy,A.xz,A.yy + L.λ, A.yz, A.zz + L.λ)

@inline Base.:(+)(A::SymTen, L::LinearAlgebra.UniformScaling) = L+A

@inline Base.:(-)(A::SymTen, L::LinearAlgebra.UniformScaling) =
    SymTen(A.xx - L.λ, A.xy,A.xz,A.yy - L.λ, A.yz, A.zz - L.λ)

@inline Base.:(-)(L::LinearAlgebra.UniformScaling,A::SymTen) = -(A-L)

@inline LinearAlgebra.tr(a::SymTen) =
    a.xx + a.yy + a.zz

@inline LinearAlgebra.norm(a::SymTen) =
    @fastmath sqrt(2(a:a))

@inline LinearAlgebra.det(t::SymTen) =
    #-txx*tyz^2 - txz^2*tyy - tzz*txy^2 + 2*txz*txy*tyz + tzz*txx*tyy
    t.xx*t.yy*t.zz + 2*t.xz*t.xy*t.yz - (t.xx*t.yz*t.yz + t.yy*t.xz*t.xz + t.zz*t.xy*t.xy)

@inline LinearAlgebra.dot(a::SymTen{T},b::Vec{T2}) where {T<:Number,T2<:Number} =
    Vec{promote_type(T,T2)}(muladd(a.xx, b.x, muladd(a.xy, b.y, a.xz*b.z)),
        muladd(a.xy, b.x, muladd(a.yy, b.y, a.yz*b.z)),
        muladd(a.xz, b.x, muladd(a.yz, b.y, a.zz*b.z)))

@inline LinearAlgebra.dot(b::Vec,a::SymTen) =
    LinearAlgebra.dot(a,b)

@inline LinearAlgebra.dot(A::AbstractMatrix,v::Vec) = A*v

@inline symouter(a::Vec,b::Vec) = #symmetric part of the outer product of two vectors
    SymTen(a.x*b.x,
           muladd(a.x,b.y,a.y*b.x)/2,
           muladd(a.x,b.z,a.z*b.x)/2,
           a.y*b.y,
           muladd(a.y,b.z,a.z*b.y)/2,
           a.z*b.z)

@inline antisymouter(a::Vec,b::Vec) =
    AntiSymTen((a.x*b.y - a.y*b.x)/2,
               (a.x*b.z - a.z*b.x)/2,
               (a.y*b.z - a.z*b.y)/2)

@inline traceless(S::SymTen) =
    S - inv(3)*tr(S)*I

@inline function Base.:^(t::SymTen{T},n::Integer) where {T}
    λ,e = eigvec(t)
    r = (λ[1]^n)*(symouter(e[1],e[1])) + (λ[2]^n)*(symouter(e[2],e[2])) + (λ[3]^n)*(symouter(e[3],e[3]))
end

@inline square(t::SymTen) =
    SymTen(muladd(t.xx,t.xx,muladd(t.xy,t.xy,t.xz^2)),
              muladd(t.xx,t.xy,muladd(t.yy,t.xy,t.xz*t.yz)),
              muladd(t.xx,t.xz,muladd(t.xy,t.yz,t.zz*t.xz)),
              muladd(t.xy,t.xy,muladd(t.yy,t.yy,t.yz^2)),
              muladd(t.xy,t.xz,muladd(t.yy,t.yz,t.zz*t.yz)),
              muladd(t.xz,t.xz,muladd(t.yz,t.yz,t.zz^2)))

@inline function Base.literal_pow(::typeof(Base.:^),t::SymTen{T},::Val{2}) where {T<:AbstractFloat}
    return square(t)
end

@inline Base.adjoint(a::SymTen{T}) where T<:Complex = SymTen{T}(conj(a.xx),conj(a.xy),conj(a.xz),conj(a.yy),conj(a.yz),conj(a.zz))
@inline Base.adjoint(a::SymTen) = a

####### Full Tensors

struct Ten{T<:Number} <: AbstractTen{T}
    xx::T
    yx::T
    zx::T
    xy::T
    yy::T
    zy::T
    xz::T
    yz::T
    zz::T
end

Base.IndexStyle(a::Type{<:Ten}) = Base.IndexLinear()

@inline Base.@propagate_inbounds function Base.getindex(a::Ten,I::Integer)
    @boundscheck checkbounds(a,I)
    @inbounds getfield(a,I)
end

@inline Base.size(a::Ten) =
    (3,3)

@inline Base.length(a::Ten) =
    9

@inline Base.zero(a::Type{Ten{T}}) where {T} =
    Ten{T}(zero(T),zero(T),zero(T),zero(T),zero(T),zero(T),zero(T),zero(T),zero(T))

@inline Ten(xx::T,xy::T,xz::T,yy::T,yz::T,zz::T) where {T} =
    Ten{T}(xx,xy,xz,yy,yz,zz)

@inline Ten(q,w,e,r,t,y,a,d,f) =
    Ten(promote(q,w,e,r,t,y,a,d,f)...)

@inline Base.:+(a::Ten{T},b::Ten{T2}) where {T,T2} =
    Ten{promote_type(T,T2)}(a.xx+b.xx, a.yx+b.yx, a.zx+b.zx, a.xy+b.xy, a.yy+b.yy, a.zy+b.zy, a.xz+b.xz, a.yz + b.yz, a.zz+b.zz)

@inline Base.:-(A::Ten{T}) where T = Ten{T}(-A.xx,-A.yx,-A.zx,-A.xy,-A.yy,-A.zy,-A.xz,-A.yz,-A.zz)

@inline Base.:-(a::Ten{T},b::Ten{T2}) where {T,T2} =
    Ten{promote_type(T,T2)}(a.xx-b.xx, a.yx-b.yx, a.zx-b.zx, a.xy-b.xy, a.yy - b.yy, a.zy - b.zy, a.xz - b.xz, a.yz - b.yz, a.zz - b.zz)

@inline Base.:*(a::T,b::Ten{T2}) where {T<:Number,T2<:Number} =
    Ten{promote_type(T,T2)}(a*b.xx, a*b.yx, a*b.zx, a*b.xy, a*b.yy, a*b.zy, a*b.xz, a*b.yz, a*b.zz)

@inline Base.:*(b::Ten{T},a::T2) where {T<:Number,T2<:Number} =
    a*b

@inline Base.:/(b::Ten{T},a::T2) where {T<:Number,T2<:Number} =
    inv(a)*b

@inline Base.:(:)(a::Ten,b::Ten) =
    muladd(a.xx, b.xx, muladd(a.yx, b.yx, muladd(a.zx, b.zx, muladd(a.xy, b.xy, muladd(a.yy, b.yy, muladd(a.zy, b.zy, muladd(a.xz, b.xz, muladd(a.yz, b.yz, a.zz*b.zz))))))))

@inline Base.:(:)(a::SymTen,b::Ten) =
    muladd(a.xx, b.xx, muladd(a.xy, b.yx, muladd(a.xz, b.zx, muladd(a.xy, b.xy, muladd(a.yy, b.yy, muladd(a.yz, b.zy, muladd(a.xz, b.xz, muladd(a.yz, b.yz, a.zz*b.zz))))))))

@inline Base.:(:)(b::Ten,a::SymTen) = a:b

@inline Base.:(+)(L::LinearAlgebra.UniformScaling,A::Ten) =
    Ten(A.xx + L.λ, A.yx, A.zx, A.xy, A.yy + L.λ, A.zy, A.xz, A.yz, A.zz + L.λ)

@inline Base.:(+)(A::Ten, L::LinearAlgebra.UniformScaling) = L+A

@inline Base.:(-)(A::Ten, L::LinearAlgebra.UniformScaling) =
    Ten(A.xx - L.λ, A.yx, A.zx, A.xy, A.yy - L.λ, A.zy, A.xz, A.yz, A.zz - L.λ)

@inline Base.:(-)(L::LinearAlgebra.UniformScaling,A::Ten) = -(A-L)

@inline LinearAlgebra.tr(a::Ten) =
    a.xx + a.yy + a.zz

@inline LinearAlgebra.norm(a::Ten) =
    @fastmath sqrt(2(a:a))

@inline LinearAlgebra.det(t::Ten) =
    (t.xy*t.yz - t.xz*t.yy)*t.zx + (t.zy*t.xz - t.zz*t.xy)*t.yx + (t.zz*t.yy - t.zy*t.yz )*t.xx

@inline LinearAlgebra.dot(a::Ten{T},b::Ten{T2}) where {T<:Number,T2<:Number} =
    Ten{promote_type(T,T2)}(muladd(a.xx, b.xx, muladd(a.xy, b.yx, a.xz*b.zx)), muladd(a.yx, b.xx, muladd(a.yy, b.yx, a.yz*b.zx)), muladd(a.zx, b.xx, muladd(a.zy, b.yx, a.zz*b.zx)),
                                      muladd(a.xx, b.xy, muladd(a.xy, b.yy, a.xz*b.zy)), muladd(a.yx, b.xy, muladd(a.yy, b.yy, a.yz*b.zy)), muladd(a.zx, b.xy, muladd(a.zy, b.yy, a.zz*b.zy)),
                                      muladd(a.xx, b.xz, muladd(a.xy, b.yz, a.xz*b.zz)), muladd(a.yx, b.xz, muladd(a.yy, b.yz, a.yz*b.zz)), muladd(a.zx, b.xz, muladd(a.zy, b.yz, a.zz*b.zz)))

@inline LinearAlgebra.dot(a::Ten{T},b::Vec{T2}) where {T<:Number,T2<:Number} =
    Vec{promote_type(T,T2)}(muladd(a.xx, b.x, muladd(a.xy, b.y, a.xz*b.z)),
                                      muladd(a.yx, b.x, muladd(a.yy, b.y, a.yz*b.z)),
                                      muladd(a.zx, b.x, muladd(a.zy, b.y, a.zz*b.z)))

@inline LinearAlgebra.dot(b::Vec{T},a::Ten{T2}) where {T,T2} =
    Vec{promote_type(T,T2)}(muladd(a.xx, b.x, muladd(a.yx, b.y, a.zx*b.z)),
                                      muladd(a.xy, b.x, muladd(a.yy, b.y, a.zy*b.z)),
                                      muladd(a.xz, b.x, muladd(a.yz, b.y, a.zz*b.z)))

@inline LinearAlgebra.dot(a::SymTen{T},b::SymTen{T2}) where {T<:Number,T2<:Number} =
    Ten{promote_type(T,T2)}(muladd(a.xx, b.xx, muladd(a.xy, b.xy, a.xz*b.xz)), muladd(a.xy, b.xx, muladd(a.yy, b.xy, a.yz*b.xz)), muladd(a.xz, b.xx, muladd(a.yz, b.xy, a.zz*b.xz)),
    muladd(a.xx, b.xy, muladd(a.xy, b.yy, a.xz*b.yz)), muladd(a.xy, b.xy, muladd(a.yy, b.yy, a.yz*b.yz)), muladd(a.xz, b.xy, muladd(a.yz, b.yy, a.zz*b.yz)),
    muladd(a.xx, b.xz, muladd(a.xy, b.yz, a.xz*b.zz)), muladd(a.xy, b.xz, muladd(a.yy, b.yz, a.yz*b.zz)), muladd(a.xz, b.xz, muladd(a.yz, b.yz, a.zz*b.zz)))

@inline Base.adjoint(a::Ten{T}) where T = Ten{T}(conj(a.xx),conj(a.xy),conj(a.xz),conj(a.yx),conj(a.yy),conj(a.yz),conj(a.zx),conj(a.zy),conj(a.zz))

@inline outer(a::Vec,b::Vec) = #symmetric part of the outer product of two vectors
    Ten(a.x*b.x, a.y*b.x, a.z*b.x,
                  a.x*b.y, a.y*b.y, a.z*b.y,
                  a.x*b.z, a.y*b.z, a.z*b.z)

const ⊗ = outer

@inline traceless(S::Ten) =
    S - inv(3)*LinearAlgebra.tr(S)*LinearAlgebra.I

function Base.:^(t::Ten{T},n::Integer) where {T}
    n == 1 && return t
    n == 2 && return t⋅t
    return t⋅(t^(n-1))
end

@inline square(t::Ten) = t⋅t

@inline function Base.literal_pow(::typeof(Base.:^),t::Ten{T},::Val{2}) where {T<:AbstractFloat}
    return t⋅t
end

@inline function SymTen(t::Ten)
    if (t.yx == t.xy && t.zx == t.xz && t.zy == t.yz)
       return SymTen(t.xx, t.yx, t.zx, t.yy, t.zy, t.zz)
    else
        throw(InexactError(:SymTen,SymTen,t))
    end
end
@inline Ten(a::SymTen) = Ten(a.xx,a.xy,a.xz,a.xy,a.yy,a.yz,a.xz,a.yz,a.zz)

####### AntiSym Tensors

struct AntiSymTen{T<:Number} <: AbstractTen{T}
    xy::T
    xz::T
    yz::T
end

@inline Ten(a::AntiSymTen) = Ten(0,-a.xy,-a.xz,a.xy,0,-a.yz,a.xz,a.yz,0)

Base.IndexStyle(a::Type{<:AntiSymTen}) = Base.IndexLinear()

@inline Base.@propagate_inbounds function Base.getindex(a::AntiSymTen{T},I::Integer) where T
    @boundscheck checkbounds(a,I)
    (I == 1 || I == 5 || I == 9) && return zero(T)
    I == 2 && return -a.xy
    I == 3 && return -a.xz
    I == 4 && return a.xy
    I == 6 && return -a.yz
    I == 7 && return a.xz
    I == 8 && return a.yz
end

@inline Base.size(a::AntiSymTen) =
    (3,3)

@inline Base.length(a::AntiSymTen) =
    9

@inline Base.zero(a::Type{AntiSymTen{T}}) where {T} =
    AntiSymTen{T}(zero(T),zero(T),zero(T))

@inline AntiSymTen(xx::T,xy::T,xz::T) where {T} =
    AntiSymTen{T}(xx,xy,xz)

@inline AntiSymTen(q,w,e) =
    AntiSymTen(promote(q,w,e)...)

@inline Base.:+(a::AntiSymTen{T},b::AntiSymTen{T2}) where {T,T2} =
    AntiSymTen{promote_type(T,T2)}(a.xy+b.xy, a.xz+b.xz, a.yz + b.yz)

@inline Base.:-(A::AntiSymTen{T}) where T = AntiSymTen{T}(-A.xy,-A.xz,-A.yz)

@inline Base.:-(a::AntiSymTen{T},b::AntiSymTen{T2}) where {T,T2} =
    AntiSymTen{promote_type(T,T2)}(a.xy-b.xy, a.xz - b.xz, a.yz - b.yz)

@inline Base.:*(a::T,b::AntiSymTen{T2}) where {T<:Number,T2<:Number} =
    AntiSymTen{promote_type(T,T2)}(a*b.xy, a*b.xz, a*b.yz)

@inline Base.:*(b::AntiSymTen{T},a::T2) where {T<:Number,T2<:Number} =
    a*b

@inline Base.:/(b::AntiSymTen{T},a::T2) where {T<:Number,T2<:Number} =
    inv(a)*b

@inline Base.:(:)(a::AntiSymTen,b::AntiSymTen) =
    2muladd(a.xy, b.xy, muladd(a.xz, b.xz, a.yz*b.yz))

@inline Base.:(:)(a::Ten,b::AntiSymTen) =
    muladd(a.xy,b.xy,muladd(a.xz,b.xz,a.yz*b.yz)) - muladd(a.yx,b.xy,muladd(a.zx,b.xz, a.zy*b.yz))

@inline Base.:(:)(b::AntiSymTen,a::Ten) = a:b

@inline Base.:(:)(a::AntiSymTen,b::SymTen) = false

@inline Base.:(:)(a::SymTen,b::AntiSymTen) = false

@inline Base.:(+)(L::LinearAlgebra.UniformScaling,A::AntiSymTen) =
    Ten(L.λ, -A.xy, -A.xz, A.xy, L.λ, -A.yz, A.xz, A.yz, L.λ)

@inline Base.:(+)(A::AntiSymTen, L::LinearAlgebra.UniformScaling) = L+A

@inline Base.:(-)(A::AntiSymTen, L::LinearAlgebra.UniformScaling) =
    Ten(-L.λ, -A.xy, -A.xz, A.xy, -L.λ, -A.yz, A.xz, A.yz, -L.λ)

@inline Base.:(-)(L::LinearAlgebra.UniformScaling,A::AntiSymTen) = -(A-L)

@inline LinearAlgebra.tr(a::AntiSymTen{T}) where T = zero(T)

@inline LinearAlgebra.norm(a::AntiSymTen) =
    @fastmath sqrt(2(a:a))

@inline LinearAlgebra.det(a::AntiSymTen) = false

@inline LinearAlgebra.dot(a::AntiSymTen{T},b::AntiSymTen{T2}) where {T<:Number,T2<:Number} =
    Ten{promote_type(T,T2)}(-muladd(a.xy, b.xy, a.xz*b.xz), -a.yz*b.xz, a.yz*b.xy,
                                      -a.xz*b.yz, -muladd(a.xy,b.xy,a.yz*b.yz), -a.xz*b.xy,
                                      a.xy*b.yz, -a.xy*b.xz, -muladd(a.xz,b.xz,a.yz*b.yz))

@inline LinearAlgebra.dot(a::AntiSymTen{T},b::Vec{T2}) where {T<:Number,T2<:Number} =
    Vec{promote_type(T,T2)}(muladd(a.xy, b.y, a.xz*b.z),
                                      muladd(a.yz, b.z, -a.xy*b.x),
                                      -muladd(a.xz, b.x, a.yz*b.y))

@inline LinearAlgebra.dot(b::Vec,a::AntiSymTen) = -a⋅b

@inline Base.adjoint(a::AntiSymTen{T}) where T = AntiSymTen{T}(conj(-a.xy),conj(-a.xz),conj(-a.yz))

@inline traceless(S::AntiSymTen) = S

Base.:^(t::AntiSymTen{T},n::Integer) where {T} = Ten(t)^n

@inline square(a::AntiSymTen) = SymTen(-muladd(a.xy,a.xy,a.xz^2), -a.xz*a.yz, a.xy*a.yz, -muladd(a.xy,a.xy,a.yz^2), -a.xy*a.xz, -muladd(a.xz,a.xz,a.yz^2))

@inline function Base.literal_pow(::typeof(Base.:^),a::AntiSymTen{T},::Val{2}) where {T<:AbstractFloat}
    return square(a)
end

@inline Ten(a::AntiSymTen{T}) where {T} = Ten{T}(zero(T),-a.xy,-a.xz,a.xy,zero(T),-a.yz,a.xz,a.yz,zero(T))
@inline Vec(a::AntiSymTen) = Vec(-a.yz,a.xz,-a.xy)
@inline AntiSymTen(a::Vec) = AntiSymTen(-a.z,a.y,-a.x)

@inline Lie(s::SymTen,w::AntiSymTen) =
    SymTen(-2*(w.xy*s.xy + w.xz*s.xz),
           w.xy*s.xx - w.yz*s.xz - w.xy*s.yy - w.xz*s.yz,
           w.xz*s.xx + w.yz*s.xy - w.xy*s.yz - w.xz*s.zz,
           2*(w.xy*s.xy - w.yz*s.yz),
           w.xy*s.xz + w.xz*s.xy + w.yz*s.yy - w.yz*s.zz,
           2*(w.xz*s.xz + w.yz*s.yz))

@inline Lie(w::AntiSymTen,s::SymTen) = Lie(s,-w)
@inline Lie(t::AbstractTen,l::AbstractTen) = t⋅l - l⋅t

############################## Interoperation

@inline Base.:+(a::AntiSymTen{T},b::SymTen{T2}) where {T,T2} =
    Ten{promote_type(T,T2)}(b.xx, b.xy-a.xy, b.xz-a.xz, a.xy+b.xy, b.yy, b.yz - a.yz, a.xz+b.xz, a.yz+b.yz, b.zz)

@inline Base.:+(a::AntiSymTen{T},b::Ten{T2}) where {T,T2} =
    Ten{promote_type(T,T2)}(b.xx, b.yx-a.xy, b.zx-a.xz, a.xy+b.xy, b.yy, b.zy - a.yz, a.xz+b.xz, a.yz+b.yz, b.zz)

@inline Base.:+(a::SymTen{T},b::Ten{T2}) where {T,T2} =
    Ten{promote_type(T,T2)}(a.xx+b.xx, a.xy+b.yx, a.xz+b.zx, a.xy+b.xy, a.yy+b.yy, a.yz+b.zy, a.xz+b.xz, a.yz + b.yz, a.zz+b.zz)

@inline Base.:+(b::SymTen,a::AntiSymTen) = a+b

@inline Base.:+(b::Ten,a::AntiSymTen) = a+b

@inline Base.:+(b::Ten,a::SymTen) = a+b

@inline Base.:-(a::AntiSymTen{T},b::SymTen{T2}) where {T,T2} =
    Ten{promote_type(T,T2)}(-b.xx, -a.xy-b.xy, -a.xz-b.xz, a.xy-b.xy, -b.yy, -a.yz - b.yz, a.xz - b.xz, a.yz - b.yz, -b.zz)

@inline Base.:-(a::SymTen{T},b::AntiSymTen{T2}) where {T,T2} =
    Ten{promote_type(T,T2)}(a.xx, a.xy+b.xy, a.xz+b.xz, a.xy-b.xy, a.yy, a.yz + b.yz, a.xz - b.xz, a.yz - b.yz, a.zz)

@inline Base.:-(a::Ten{T},b::AntiSymTen{T2}) where {T,T2} =
    Ten{promote_type(T,T2)}(a.xx, a.yx+b.xy, a.zx+b.xz, a.xy-b.xy, a.yy, a.zy + b.yz, a.xz - b.xz, a.yz - b.yz, a.zz)

@inline Base.:-(a::AntiSymTen{T},b::Ten{T2}) where {T,T2} = -(b-a)

@inline Base.:-(a::Ten{T},b::SymTen{T2}) where {T,T2} =
    Ten{promote_type(T,T2)}(a.xx-b.xx, a.yx-b.xy, a.zx-b.xz, a.xy-b.xy, a.yy - b.yy, a.zy - b.yz, a.xz - b.xz, a.yz - b.yz, a.zz - b.zz)

@inline Base.:-(a::SymTen{T},b::Ten{T2}) where {T,T2} = -(b-a)

@inline LinearAlgebra.dot(a::Ten{T},b::SymTen{T2}) where {T<:Number,T2<:Number} =
    Ten{promote_type(T,T2)}(muladd(a.xx, b.xx, muladd(a.xy, b.xy, a.xz*b.xz)), muladd(a.yx, b.xx, muladd(a.yy, b.xy, a.yz*b.xz)), muladd(a.zx, b.xx, muladd(a.zy, b.xy, a.zz*b.xz)),
                                      muladd(a.xx, b.xy, muladd(a.xy, b.yy, a.xz*b.yz)), muladd(a.yx, b.xy, muladd(a.yy, b.yy, a.yz*b.yz)), muladd(a.zx, b.xy, muladd(a.zy, b.yy, a.zz*b.yz)),
                                      muladd(a.xx, b.xz, muladd(a.xy, b.yz, a.xz*b.zz)), muladd(a.yx, b.xz, muladd(a.yy, b.yz, a.yz*b.zz)), muladd(a.zx, b.xz, muladd(a.zy, b.yz, a.zz*b.zz)))

@inline LinearAlgebra.dot(a::SymTen{T},b::Ten{T2}) where {T<:Number,T2<:Number} =
    Ten{promote_type(T,T2)}(muladd(a.xx, b.xx, muladd(a.xy, b.yx, a.xz*b.zx)), muladd(a.xy, b.xx, muladd(a.yy, b.yx, a.yz*b.zx)), muladd(a.xz, b.xx, muladd(a.yz, b.yx, a.zz*b.zx)),
    muladd(a.xx, b.xy, muladd(a.xy, b.yy, a.xz*b.zy)), muladd(a.xy, b.xy, muladd(a.yy, b.yy, a.yz*b.zy)), muladd(a.xz, b.xy, muladd(a.yz, b.yy, a.zz*b.zy)),
    muladd(a.xx, b.xz, muladd(a.xy, b.yz, a.xz*b.zz)), muladd(a.xy, b.xz, muladd(a.yy, b.yz, a.yz*b.zz)), muladd(a.xz, b.xz, muladd(a.yz, b.yz, a.zz*b.zz)))

@inline LinearAlgebra.dot(a::Ten{T},b::AntiSymTen{T2}) where {T<:Number,T2<:Number} =
    Ten{promote_type(T,T2)}(-muladd(a.xy,b.xy,a.xz*b.xz), -muladd(a.yy,b.xy, a.yz*b.xz), -muladd(a.zy,b.xy, a.zz*b.xz),
                                      a.xx*b.xy - a.xz*b.yz, a.yx*b.xy - a.yz*b.yz, a.zx*b.xy - a.zz*b.yz,
                                      muladd(a.xx,b.xz,a.xy*b.yz), muladd(a.yx,b.xz,a.yy*b.yz), muladd(a.zx,b.xz, a.zy*b.yz))

@inline LinearAlgebra.dot(b::AntiSymTen{T2},a::Ten{T}) where {T,T2} =
    Ten{promote_type(T,T2)}(muladd(b.xy, a.yx, b.xz*a.zx), -a.xx*b.xy+a.zx*b.yz, -muladd(a.xx,b.xz,a.yx*b.yz),
                                      muladd(a.yy,b.xy,a.zy*b.xz), -a.xy*b.xy+a.zy*b.yz, -muladd(a.xy,b.xz,a.yy*b.yz),
                                      muladd(a.yz,b.xy,a.zz*b.xz), -a.xz*b.xy+a.zz*b.yz, -muladd(a.xz,b.xz,a.yz*b.yz))

@inline LinearAlgebra.dot(a::SymTen{T},b::AntiSymTen{T2}) where {T<:Number,T2<:Number} =
    Ten{promote_type(T,T2)}(-muladd(a.xy,b.xy,a.xz*b.xz), -muladd(a.yy,b.xy, a.yz*b.xz), -muladd(a.yz,b.xy, a.zz*b.xz),
                                      a.xx*b.xy - a.xz*b.yz, a.xy*b.xy - a.yz*b.yz, a.xz*b.xy - a.zz*b.yz,
                                      muladd(a.xx,b.xz,a.xy*b.yz), muladd(a.xy,b.xz,a.yy*b.yz), muladd(a.xz,b.xz, a.yz*b.yz))

@inline LinearAlgebra.dot(b::AntiSymTen{T2},a::SymTen{T}) where {T,T2} =
    Ten{promote_type(T,T2)}(muladd(a.xy,b.xy,a.xz*b.xz), -a.xx*b.xy+a.xz*b.yz, -muladd(a.xx,b.xz,a.xy*b.yz),
                                      muladd(a.yy,b.xy,a.yz*b.xz), -a.xy*b.xy+a.yz*b.yz, -muladd(a.xy,b.xz,a.yy*b.yz),
                                      muladd(a.yz,b.xy,a.zz*b.xz), -a.xz*b.xy+a.zz*b.yz, -muladd(a.xz,b.xz,a.yz*b.yz))

@inline symmetric(t::SymTen) = t
@inline symmetric(t::AntiSymTen{T}) where {T} = zero(SymTen{T})
@inline symmetric(t::Ten) = SymTen(t.xx, (t.yx+t.xy)/2, (t.xz+t.zx)/2, t.yy, (t.zy+t.yz)/2, t.zz)

@inline antisymmetric(t::SymTen{T}) where {T} = zero(AntiSymTen{T})
@inline antisymmetric(t::AntiSymTen) = t
@inline antisymmetric(t::Ten) = AntiSymTen((t.xy-t.yx)/2, (t.xz-t.zx)/2, (t.yz-t.zy)/2)

include("tensorarrays.jl")


module ConvexBench

using Convex
using Convex.ProblemDepot: foreach_problem, PROBLEMS

function dummy_problem(opt; T = Float64)
    x = Variable(3)
    p = minimize(2*x[2] + x[1] - x[3], [x >= 1]; numeric_type = T)
    solve!(p, opt)
end

handle_problem_function(opt) = p -> solve!(p, opt)

end


################################################################################
# Iterative Best Response Example
################################################################################
# using Algames
using StaticArrays
using LinearAlgebra

T = Float64

# Define the dynamics of the system
p = 3 # Number of players
d = 2
model = DoubleIntegratorGame(p=p, d=d) # game with 3 players with unicycle dynamics
n = model.n
m = model.m

# Define the horizon of the problem
N = 20 # N time steps
dt = 0.1 # each step lasts 0.1 second
probsize = ProblemSize(N,model) # Structure holding the relevant sizes of the problem

# Define the objective of each player
# We use a LQR cost
Q = [Diagonal(50*ones(SVector{model.ni[i],T})) for i=1:p] # Quadratic state cost
R = [Diagonal(0.01*ones(SVector{model.mi[i],T})) for i=1:p] # Quadratic control cost
# Desrired state
xf = [SVector{model.ni[1],T}([ 0, 0*0.0,0,0]),
      SVector{model.ni[1],T}([ 0, 0*0.2,0,0]),
	  SVector{model.ni[1],T}([ 0, 0*0.4,0,0]),
	  SVector{model.ni[1],T}([ 0, 0*0.6,0,0]),
	  SVector{model.ni[1],T}([ 0, 0*0.8,0,0]),
	  SVector{model.ni[1],T}([ 0, 0*1.0,0,0]),
	  SVector{model.ni[1],T}([ 0, 0*1.2,0,0]),
	  SVector{model.ni[1],T}([ 0, 0*1.4,0,0]),
	  SVector{model.ni[1],T}([ 0, 0*1.6,0,0]),
	  SVector{model.ni[1],T}([ 0, 0*1.8,0,0]),
      ]
xf = xf[1:p]
# Desired control
uf = [zeros(SVector{model.mi[i],T}) for i=1:p]
# Objectives of the game
game_obj = GameObjective(Q,R,xf,uf,N,model)
radius = 3.0*ones(p)
μ = 2.0*ones(p)
add_collision_cost!(game_obj, radius, μ)

# Define the constraints that each player must respect
game_con = GameConstraintValues(probsize)
# # Add collision avoidance
radius = 0.25
add_collision_avoidance!(game_con, radius)

# Define the initial state of the system
x0 = [
	[ 0.00,  0.0,  0.0,  0.0],
	[ 0.05,  0.2,  0.0,  0.0],
	[ 0.10,  0.4,  0.0,  0.0],
	[ 0.15,  0.6,  0.0,  0.0],
	[ 0.20,  0.8,  0.0,  0.0],
	[ 0.25,  1.0,  0.0,  0.0],
	[ 0.30,  1.2,  0.0,  0.0],
	[ 0.35,  1.4,  0.0,  0.0],
	[ 0.40,  1.6,  0.0,  0.0],
	[ 0.45,  1.8,  0.0,  0.0],
	]
θ = range(0.0, length=2, stop=2π)
rad = 0.5
x0 = [[(rad+θ/10)*cos(θ), (rad+θ/10)*sin(θ), 0, 0] for θ in range(0.0, length=p, stop=2π*(1-1/p))]
x0 = SVector{model.n,T}(reshape(vcat([x0[i]' for i=1:p]...), model.n))

# Define the Options of the solver
opts_alg = Options(Δ_min=1e-9, inner_print=false)
opts_ibr = Options(dual_reset=false, outer_iter=10, inner_iter=25, Δ_min=1e-9, inner_print=false)
# opts_ibr = Options(dual_reset=false, outer_iter=7, inner_iter=20, Δ_min=1e-5, inner_print=false)
# ibr_opts = IBROptions(Δ_min=1e-9, ibr_iter=100, ordering=shuffle!([1,2,3,4,5,6,7,8,9,10]))
ibr_opts = IBROptions(Δ_min=1e-9, ibr_iter=15, ordering=[1,2,3,4,5,6,7,8,9,10], live_plotting=true)
ibr_opts = IBROptions(Δ_min=1e-9, ibr_iter=100, ordering=[2,1,3,4,5,6,7,8,9,10], live_plotting=false)

# Define the game problem
prob0_alg = GameProblem(N,dt,x0,model,opts_alg,game_obj,game_con)
prob0_ibr = GameProblem(N,dt,x0,model,opts_ibr,game_obj,game_con)
prob_alg = deepcopy(prob0_alg)
prob_ibr = deepcopy(prob0_ibr)

# cost.(prob_ibr.game_obj.obj)
# game_obj.obj

# Solve the problem
@elapsed newton_solve!(prob_alg)
@elapsed ibr_newton_solve!(prob_ibr, ibr_opts=ibr_opts)

residual!(prob_ibr, prob_ibr.pdtraj)
norm(prob_ibr.core.res, 1)/length(prob_ibr.core.res)
residual!(prob_alg, prob_alg.pdtraj)
norm(prob_alg.core.res, 1)/length(prob_alg.core.res)
prob_ibr.stats

plt = plot()
plot!(plt, cumsum(prob_alg.stats.t_elap), log.(10, prob_alg.stats.res))
plot!(plt, cumsum(prob_alg.stats.t_elap), log.(10, [v.max for v in prob_alg.stats.dyn_vio]))
plot!(plt, cumsum(prob_alg.stats.t_elap), log.(10, [v.max for v in prob_alg.stats.sta_vio]))
plot!(plt, cumsum(prob_ibr.stats.t_elap), log.(10, prob_ibr.stats.res))
plot!(plt, cumsum(prob_ibr.stats.t_elap), log.(10, [v.max for v in prob_ibr.stats.dyn_vio]))
plot!(plt, cumsum(prob_ibr.stats.t_elap), log.(10, [v.max for v in prob_ibr.stats.sta_vio]))

# Visualize the Results
plot(prob_alg.model, prob_alg.pdtraj.pr)
plot!(prob_ibr.model, prob_ibr.pdtraj.pr)
plot(prob_alg.stats)
plot!(prob_ibr.stats)
const Algames = Main
using Plots

# IBR
[prob_ibr.game_con.state_conval[1][1].λ[k][1] for k=1:N-1][end]
[prob_ibr.game_con.state_conval[1][2].λ[k][1] for k=1:N-1][end]

[prob_ibr.game_con.state_conval[2][1].λ[k][1] for k=1:N-1][end]
[prob_ibr.game_con.state_conval[2][2].λ[k][1] for k=1:N-1][end]

[prob_ibr.game_con.state_conval[3][1].λ[k][1] for k=1:N-1][end]
[prob_ibr.game_con.state_conval[3][2].λ[k][1] for k=1:N-1][end]

# Algames
[prob_alg.game_con.state_conval[1][1].λ[k][1] for k=1:N-1][end]
[prob_alg.game_con.state_conval[1][2].λ[k][1] for k=1:N-1][end]

[prob_alg.game_con.state_conval[2][1].λ[k][1] for k=1:N-1][end]
[prob_alg.game_con.state_conval[2][2].λ[k][1] for k=1:N-1][end]

[prob_alg.game_con.state_conval[3][1].λ[k][1] for k=1:N-1][end]
[prob_alg.game_con.state_conval[3][2].λ[k][1] for k=1:N-1][end]

#Conclusions
	# - if there is no coupling at all (dynamics, cost, state constraints)
	# IBR marginally better, but should be the same since we exploit sparsity using a sparse linear solver.

	# - if there is cost coupling only
	# IBR converges to a different local NE
	# the solution depends on the players' ordering
	# the solution does not reflect the symmetry of information between players
	# some players having cost large advantages over

	# - if there is state constraints coupling (when state constraint are active)
	# IBR converges to a different local NNE
	# the solution depends on the players' ordering
	# the solution does not reflect the symmetry of information between players
	# some players having large cost advantages over
	# converges to a solution which is also a Stackelberg equilibrium for 2 players,
	# and chain-Stackelberg for multi player scenarios.
	# really hard time converging the solutions are oscillating in practice


using Test, InnovationPrivacy, StatsBase

@testset "Check Transition Probabilities sum up to one" begin
    for N in 5:1:15, n in 0:1:10
       for i in 1:N
           @test sum(Transition_Probability(n, ParsGrowth(N = N), normalize = :no)[i, :]) ≈ 1.0
       end
    end
end

@testset "Check Bayesian Identities" begin
    for σ²_θ in 0.1:0.2:1.5, σ²_ϵ in 0.1:0.2:1.5, n in 0:1:10
        bpar = BPars(n, ParsGrowth(σ²_θ = σ²_θ, σ²_ϵ = σ²_ϵ))
        @test bpar.w_an + bpar.w_pn ≈ 1.0
    end
end

@testset "Check Expected Profit Identity" begin
    for σ²_θ in 0.1:0.5:1.6, σ²_ϵ in 0.1:0.5:1.6, n in 0:2:10, a_bar in -2:0.2:2
        par  = ParsGrowth(σ²_θ = σ²_θ, σ²_ϵ = σ²_ϵ)
        bpar = BPars(n, par)

        e_profits = Eπ(a_bar; par, bpar)
        e_revenue = Er(a_bar; par, bpar)
        wage_bill = labor_costs(a_bar; par, bpar)

        @test e_profits - (e_revenue - wage_bill - par.w*par.f) ≈ 0.
    end
end

@testset "Check Transition Matrix" begin
    for σ²_θ in 0.1:0.5:1.1, σ²_ϵ in 0.1:0.5:1.1, n in 0:5:10

        par      = ParsGrowth(σ²_θ = σ²_θ, σ²_ϵ = σ²_ϵ)
        bpar     = BPars(n, par)

        Π_test   = Transition_Matrix2(par, bpar)
        Π_direct = Transition_Probability(par, bpar)

        @test sum((Π_direct - Π_test).^2) < 1e-4
    end
end


@time include("0_setup.jl")
using IterableTables, FastGroupBy

@everywhere import JuliaDB.extractarray

if !isdefined(:data)
    @time data = load("lc_juliadb")
end

# fields = fieldnames(eltype(data))
fields = names(data)
ftypes = eltype(data).parameters

xx = JuliaDB.extractarray(data, x -> x.loan_status)

numeric_cols = [fields[i] for i = 1:length(fields) if (ftypes[i] <: Number ||
            ftypes[i] <: Nullable{<:Number}) && !(fields[i] in [:id, :member_id, :dti_joint])];

const bad_statuses = ("Late (16-30 days)","Late (31-120 days)","Default","Charged Off")
good_loans = filter(data) do row
    !(row.loan_status in bad_statuses)
end
bad_loans = filter(x->x.loan_status in bad_statuses, data)

using Gadfly
import NullableArrays: dropnull

# Density plot for bad and good loans
plots = Gadfly.Plot[]

good_numbers = good_loans[:loan_amnt]
@which columns(good_loans,:loan_amnt)

for (name, g, b) in zip(numeric_cols, columns(good_loans), columns(bad_loans))
    g = NullableArrays.dropnull(g)
    b = NullableArrays.dropnull(b)
    p = plot(layer(x=g, Geom.density, Theme(default_color=colorant"green")),
             layer(x=b, Geom.density, Theme(default_color=colorant"red")),
              Guide.title(string(name)), Guide.ylabel("density"))
    push!(plots, p)
end

perm = randperm(length(data))
train_till = round(Int, length(data) * 3/4)

training_subidx = sort!(perm[1:train_till])
testing_subidx = sort!(perm[train_till+1:end]);

training_subset = data[training_subidx]
testing_subset = data[testing_subidx]

features_train = [revol_util_train int_rate_train mths_since_last_record_train annual_inc_joint_train total_rec_prncp_train all_util_train]

labels = collect(map(x->x in bad_statuses, data))

import NullableArrays: dropnull
findnonnulls(xs::Columns) = find(x->!any(map(isnull, x)), xs)

function input_matrix(table, fields)
    float_features = collect(values(table, fields))
    tmp = findnonnulls(float_features) # indices of the rows where all fields are non-null
    nzidxs = collect(keys(table,1)[tmp]) # corresponding indices in the table
    reduce(hcat, map(dropnull, columns(float_features[tmp]))), values(labels[nzidxs]) # matrix, label vector
end

training_matrix, train_labels = input_matrix(training_subset, numeric_features)
model = build_forest(train_labels, training_matrix, 3, 10, 0.8, 6)

#     0.512169 seconds (7.67 k allocations: 440.766 KiB)
#   Out[98]:
#   Ensemble of Decision Trees
#   Trees:      10
#   Avg Leaves: 41.8
#   Avg Depth:  8.0

f = open("  loanmodel.jls", "w")
serialize(f, model)
close(f)

features_test, test_labels = input_matrix(testing_subset, numeric_features)

@time predictions = mapslices(features_test, 2) do fs
    DecisionTree.apply_forest(model, fs)
end;

# Receiver Operating Characteristics curve
using ROC
curve = roc(convert(Vector{Float64}, predictions[:]),convert(BitArray{1}, test_labels))

# An ROC plot in Gadfly with data calculuated using ROC.jl

plot(layer(x=curve.FPR, y=curve.TPR, Geom.line),
     layer(x = linspace(0.0,1.0,101), y = linspace(0.0,1.0,101),
       Geom.line, Theme(default_color=colorant"red")), Guide.title("ROC"),
       Guide.xlabel("False Positive Rate"),Guide.ylabel("True Positive Rate"))

# Area Under Curve
AUC(curve)

# => 0.7291976225854385


# Optimizes the code for computing the derivatives
# of I_v and J_v with respect to k.
include("../src/IJv_derivative.jl")

function optimize_IJv_derivative(k2::Float64)
v_max = 20
Iv_lower = zeros(Float64,v_max+1); Jv_lower = zeros(Float64,v_max+1)
Iv_raise = zeros(Float64,v_max+1); Jv_raise = zeros(Float64,v_max+1)
dIvdk_lower = zeros(Float64,v_max+1); dJvdk_lower = zeros(Float64,v_max+1)
dIvdk_raise = zeros(Float64,v_max+1); dJvdk_raise = zeros(Float64,v_max+1)
Iv_big_lower = zeros(BigFloat,v_max+1); Jv_big_lower = zeros(BigFloat,v_max+1)
Iv_big_raise = zeros(BigFloat,v_max+1); Jv_big_raise = zeros(BigFloat,v_max+1)
dIvdk_big_lower = zeros(BigFloat,v_max+1); dJvdk_big_lower = zeros(BigFloat,v_max+1)
dIvdk_big_raise = zeros(BigFloat,v_max+1); dJvdk_big_raise = zeros(BigFloat,v_max+1)
# This computes I_v for the largest v, and then works down to smaller values:
if k2 <= 1
  kc = convert(Float64,sqrt(1.-big(k2)))
  k = convert(Float64,sqrt(big(k2)))
  kck = convert(Float64,sqrt((1.-big(k2))*big(k2)))
  kap = convert(Float64,2*asin(sqrt(big(k2))))
  k2_big = big(k2)
  k_big = sqrt(k2_big)
  kc_big = sqrt(1-big(k2))
  kck_big = k_big*kc_big
  kap_big = 2*asin(k_big)
else
  kc = convert(Float64,sqrt(1.-inv(big(k2))))
  k = convert(Float64,sqrt(big(k2)))
  kap = convert(typeof(k2),pi); kck=zero(k2)
  k2_big = big(k2)
  k_big = sqrt(k2_big)
  kc_big = sqrt(big(1.0)-inv(big(k2)))
  kck_big = big(0.0)
  kap_big = big(pi)
end
if k2 > 0
  dIJv_lower_dk!(v_max,k2,kck,kc,kap,Iv_lower,Jv_lower,dIvdk_lower,dJvdk_lower)
  dIJv_lower_dk!(v_max,k2_big,kck_big,kc_big,kap_big,Iv_big_lower,Jv_big_lower,dIvdk_big_lower,dJvdk_big_lower)
  dIJv_raise_dk!(v_max,k2,kck,kc,kap,Iv_raise,Jv_raise,dIvdk_raise,dJvdk_raise)
  dIJv_raise_dk!(v_max,k2_big,kck_big,kc_big,kap_big,Iv_big_raise,Jv_big_raise,dIvdk_big_raise,dJvdk_big_raise)
#  for v=0:v_max
#    println("v: ",v," k2: ",k2," kc: ",kc," Iv_lower: ",Iv_lower[v+1]," Iv_big_lower: ",convert(Float64,Iv_big_lower[v+1]),
#    " Iv_raise: ",Iv_raise[v+1]," Iv_big_raise: ",convert(Float64,Iv_big_raise[v+1]),
#    " Jv_lower: ",Jv_lower[v+1]," Jv_big_lower: ",convert(Float64,Jv_big_lower[v+1]),
#    " Jv_raise: ",Jv_raise[v+1]," Jv_big_raise: ",convert(Float64,Jv_big_raise[v+1]))
##      " dIvdk_lower: ",dIvdk[v+1]," dIvdk_num: ",dIvdk2_num[v+1]*2*k," diff: ",dIvdk[v+1]-dIvdk2_num[v+1]*2*k," dJvdk: ",dJvdk[v+1]," dJvdk_num: ",dJvdk2_num[v+1]*2*k," diff: ",dJvdk[v+1]-dJvdk2_num[v+1]*2*k)
#  end
end
return Iv_lower,Jv_lower,Iv_raise,Jv_raise,dIvdk_lower,dJvdk_lower,dIvdk_raise,dJvdk_raise,convert(Array{Float64,1},Iv_big_lower),convert(Array{Float64,1},Jv_big_lower),convert(Array{Float64,1},Iv_big_raise),convert(Array{Float64,1},Jv_big_raise),convert(Array{Float64,1},dIvdk_big_lower),convert(Array{Float64,1},dJvdk_big_lower),convert(Array{Float64,1},dIvdk_big_raise),convert(Array{Float64,1},dJvdk_big_raise)
end

nk2 = 100
v_max = 20
Iv_lower = zeros(Float64,nk2,v_max+1); Jv_lower = zeros(Float64,nk2,v_max+1)
Iv_raise = zeros(Float64,nk2,v_max+1); Jv_raise = zeros(Float64,nk2,v_max+1)
dIvdk_lower = zeros(Float64,nk2,v_max+1); dJvdk_lower = zeros(Float64,nk2,v_max+1)
dIvdk_raise = zeros(Float64,nk2,v_max+1); dJvdk_raise = zeros(Float64,nk2,v_max+1)
Iv_big_lower = zeros(Float64,nk2,v_max+1); Jv_big_lower = zeros(Float64,nk2,v_max+1)
Iv_big_raise = zeros(Float64,nk2,v_max+1); Jv_big_raise = zeros(Float64,nk2,v_max+1)
dIvdk_big_lower = zeros(Float64,nk2,v_max+1); dJvdk_big_lower = zeros(Float64,nk2,v_max+1)
dIvdk_big_raise = zeros(Float64,nk2,v_max+1); dJvdk_big_raise = zeros(Float64,nk2,v_max+1)
k2 = logspace(0,2,nk2)
for i=1:nk2
  Iv_lower[i,:],Jv_lower[i,:],Iv_raise[i,:],Jv_raise[i,:],dIvdk_lower[i,:],dJvdk_lower,dIvdk_raise[i,:],dJvdk_raise[i,:],Iv_big_lower[i,:],Jv_big_lower[i,:],Iv_big_raise[i,:],Jv_big_raise[i,:],dIvdk_big_lower[i,:],dJvdk_big_lower[i,:],dIvdk_big_raise[i,:],dJvdk_big_raise = optimize_IJv_derivative(k2[i])
end


#   Copyright (c) 2020-2022 Velexi Corporation
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

"""
Blade.jl defines the Blade type and core methods
"""

# --- Exports

# Types
export Blade

# Functions
export blade_atol

# --- Type definitions

import LinearAlgebra
using LinearAlgebra: det, diag, dot, qr

"""
    struct Blade{T<:AbstractFloat} <: AbstractBlade{T}

Blade represented with the floating-point precision of type `T`. The norm and
orientation of a Blade are encoded by its `volume`. The norm of a Blade is
equal to `abs(volume)` and the orientation of a Blade relative to its `basis`
is equal to `sign(volume)`.

Notes
=====
* The grade of a Blade type is greater than 0 and less than the dimension of
  the space that the blade is embedded in
"""
struct Blade{T<:AbstractFloat} <: AbstractBlade{T}
    #=
     Fields
      ------
      * `dim`: the dimension of the space that the blade is embedded in

      * `grade`: the dimension of the space spanned by the blade

      * `basis`: an orthonormal basis for the space spanned by the blade. Note
        that the order of the columns in `basis` defines the orientation for
        the unit blade represented by `basis`.

      * `volume`: the signed-norm (hypervolume) of the blade. The sign
         of `volume` indicates the orientation of the blade relative to the
         unit blade represented by `basis`. It is positive when the blade has
         the same orientation as `basis` and negative when the blade has the
         opposite orientation.
    =#
    dim::Int
    grade::Int
    basis::Matrix{T}
    volume::T

    """
    Construct a Blade for a geometric algebra in `dim` dimensions having the
    specified data field values. If the absolute value of `volume` is less than
    `atol`, a Scalar representing zero is returned. If the `dim` and `grade`
    are equal, a Pseudoscalar is returned.

    When `enforce_constraints` is true, constraints are enforced. When
    `copy_basis` is true, the basis of the new Blade is a copy of `basis`;
    otherwise, the basis of the new Blade is a reference to `basis`.

    Note: this inner constructor is intended primarily for use by outer
    constructors to enforce constraints.
    """
    function Blade{T}(dim::Int,
                      grade::Int,
                      basis::Matrix{T},
                      volume::Real;
                      atol::Real=blade_atol(T),
                      enforce_constraints::Bool=true,
                      copy_basis::Bool=true) where {T<:AbstractFloat}

        # --- Enforce constraints

        if enforce_constraints
            if dim < 0
                throw(ArgumentError("`dim` must be nonnegative"))
            end

            if grade < 0
                throw(ArgumentError("`grade` must be nonnegative"))
            end

            basis_size = size(basis)

            if dim != basis_size[1]
                message = string("`dim` not equal to the number of ",
                                 "rows in `basis`")
                throw(DimensionMismatch(message))
            end

            if grade != basis_size[2]
                message = string("`grade` not equal to the number of ",
                                 "columns in `basis`")
                throw(DimensionMismatch(message))
            end

            if dim > 0 && grade > 0
                for i in 1:grade
                    if LinearAlgebra.norm(basis[:, i]) ≉ 1
                        message = "columns of `basis` not normalized"
                        throw(ArgumentError(message))
                    end
                end
            end
        end

        # --- Handle edge cases

        # `volume` is effectively zero
        if abs(volume) < atol
            return zero(Blade{T})
        end

        # Return zero if dim or grade is equal to 0
        if dim == 0 || grade == 0
            return zero(Blade{T})
        end

        # Return a Pseudoscalar if the grade of the blade is equal to the
        # dimension of the embedding space.
        if grade == dim
            if det(basis) == 0
                return zero(Blade{T})
            elseif det(basis) > 0
                return Pseudoscalar{T}(dim, volume)
            else
                return Pseudoscalar{T}(dim, -volume)
            end
        end

        # Return new Blade
        copy_basis ?
            new(dim, grade, copy(basis), volume) :
            new(dim, grade, basis, volume)
    end
end

"""
    Blade(vectors::Array{T};
          volume::Union{Real, Nothing}=nothing,
          atol::Real=blade_atol(T)) where {T<:AbstractFloat}

    Blade(vectors::Array{<:Integer};
          volume::Union{Real, Nothing}=nothing,
          atol::Real=blade_atol(Float64))

Construct a Blade from a collection of vectors stored as (1) the columns of a
matrix or (2) a single vector. If the norm of the blade is less than `atol`, a
Scalar representing zero is returned. If the grade of the blade is equal to
the dimension of the space that the blade is embedded in, a Pseudoscalar is
returned.

By default, `vectors` determines the volume of the blade. However, if `volume`
is specified, `vectors` is only used to define the subspace (including
orientation) represented by the blade.

Notes
=====

Orientation
-----------

* _`vectors` contains more than one vector_. When `volume` is positive, the
  orientation of the blade is the same as the orientation of the outer product
  of the columns of `vectors` (taken in order). When `volume` is negative, the
  orientation of the blade is the opposite of the orientation implied by the
  `vectors`.

* _`vectors` contains a single vector `v`_. When `volume` is positive, the
  orientation of the blade is the same as the direction of `v`. When `volume`
  is negative, the orientation of the blade is the opposite of the direction
  of `v`.

Precision
---------

When the precision is not specified, the following rules are applied to set
the precision of the Blade.

* If `vectors` is an Array of floating-point values, the precision of the
  Blade is inferred from the precision of the elements of `vector`.

* If `vectors` is an Array of integers, the precision of the Blade
  defaults to `Float64`.
"""
function Blade{T}(vectors::Matrix{<:Real};
                  volume::Union{Real, Nothing}=nothing,
                  atol::Real=blade_atol(T)) where {T<:AbstractFloat}

    # --- Handle edge cases

    # `volume` is effectively zero
    if volume != nothing && abs(volume) < atol
        return zero(Blade{T})
    end

    # number of vectors == 0
    dims = size(vectors)
    if dims[1] == 0 || dims[2] == 0
        return zero(Blade{T})
    end

    # number of vectors > dimension of column space
    if dims[2] > dims[1]
        if dims[1] == 1
            # `vectors` is a single row vector, so convert it to a column
            # vector and call constructor for single column vector.
            return Blade{T}(reshape(vectors, dims[2]))
        else
            return zero(Blade{T})
        end
    end

    # --- Construct Blade

    # Convert `vectors` to have type `T`
    vectors = convert(Matrix{T}, vectors)

    # Compute QR factorization and signed norm
    F = qr(vectors)
    signed_norm = prod(diag(F.R))

    # Return zero if norm is below tolerance
    if abs(signed_norm) < atol
        return zero(Blade{T})
    end

    # Compute orthonormal basis for subspace
    basis = Matrix(F.Q)

    # Compute volume
    volume = isnothing(volume) ? signed_norm : volume * sign(signed_norm)

    # Return new Blade
    Blade{T}(dims[1], dims[2], basis, volume,
             atol=atol, enforce_constraints=false, copy_basis=false)
end

function Blade{T}(v::Vector{<:Real};
                  volume::Union{Real, Nothing}=nothing,
                  atol::Real=blade_atol(T)) where {T<:AbstractFloat}

    # --- Handle edge cases

    # `volume` is effectively zero
    if volume != nothing && abs(volume) < atol
        return zero(Blade{T})
    end

    # length(v) == 0
    if length(v) == 0
        return zero(Blade{T})
    end

    # --- Construct Blade

    # Convert `v` to have type `T` and compute norm
    v = convert(Vector{T}, v)
    norm_v = LinearAlgebra.norm(v)

    # Return zero if norm is below tolerance
    if abs(norm_v) < atol
        return zero(Blade{T})
    end

    # Compute basis as a column vector
    basis = reshape(v, length(v), 1) / norm_v

    # Compute volume
    volume = isnothing(volume) ? norm_v : volume

    # Return new Blade
    Blade{T}(length(v), 1, basis, volume, atol=atol,
             enforce_constraints=false, copy_basis=false)
end

Blade(vectors::Array{T};
      volume::Union{Real, Nothing}=nothing,
      atol::Real=blade_atol(T)) where {T<:AbstractFloat} =
    Blade{T}(vectors, volume=volume, atol=atol)

Blade(vectors::Array{<:Integer};
      volume::Union{Real, Nothing}=nothing,
      atol::Real=blade_atol(Float64)) =
    Blade(convert(Array{Float64}, vectors), volume=volume, atol=atol)

"""
    Blade(B::Blade;
          volume::Real=volume(B),
          atol::Real=blade_atol(typeof(volume(B))),
          copy_basis=false)

    Blade{T}(B::Blade;
          volume::Real=volume(B),
          atol::Real=blade_atol(T),
          copy_basis=false) where {T<: AbstractFloat}

Copy constructors. Construct a Blade representing the same space as `B` having
a specified oriented volume relative to `B`. A Scalar representing zero is
returned if the absolute value of `volume` is less than `atol`.

When `copy_basis` is true, the `basis` of the new Blade is a copy of the
`basis` of the original Blade; otherwise, the `basis` of the new Blade is
reference to the `basis` of the original Blade.
"""
Blade{T}(B::Blade;
         volume::Real=volume(B),
         atol::Real=blade_atol(T),
         copy_basis::Bool=false) where {T<:AbstractFloat} =

    Blade{T}(dim(B), grade(B), convert(Matrix{T}, basis(B)), volume,
             atol=atol, enforce_constraints=false, copy_basis=copy_basis)

Blade(B::Blade;
      volume::Real=B.volume,
      atol::Real=blade_atol(typeof(B.volume)),
      copy_basis=false) =
    Blade{typeof(B.volume)}(B, volume=volume, atol=atol, copy_basis=copy_basis)

"""
    Blade(x::Real)
    Blade(x::Scalar)

Convenience constructors that return a Scalar with value `x`.
"""
Blade(x::Real) = Scalar(x)
Blade(x::AbstractScalar) = Scalar(value(x))

# --- Method definitions for AbstractMultivector interface functions

dim(B::Blade) = B.dim

-(B::Blade) = Blade(B, volume=-volume(B), copy_basis=false)

Base.reverse(B::Blade) =
    mod(grade(B), 4) < 2 ?  B : Blade(B, volume=-volume(B), copy_basis=false)

function dual(B::Blade)
    # --- Extend basis(B) to an orthonormal basis for entire space.

    F = qr(basis(B))

    # --- Compute volume of dual

    # Account for orientation of Q relative to orientation of I formed from
    # standard basis
    dual_volume = volume(B) * sign(det(F.Q))

    # Account for orientation of first grade(B) columns of Q relative to
    # orientation of basis(B)
    if prod(diag(F.R)) < 0
        dual_volume = -dual_volume
    end

    # Account for sign of I^{-1} relative to I
    if mod(dim(B), 4) >= 2
        dual_volume = -dual_volume
    end

    # Account for reversals required to eliminate B
    if mod(grade(B), 4) >= 2
        dual_volume = -dual_volume
    end

    Blade{typeof(volume(B))}(dim(B), dim(B) - grade(B),
                             F.Q[:, grade(B) + 1:end], dual_volume,
                             enforce_constraints=false,
                             copy_basis=false)
end

# --- Method definitions for AbstractBlade interface functions

grade(B::Blade) = B.grade

basis(B::Blade) = B.basis

volume(B::Blade) = B.volume

inv(B::Blade) =
    mod(grade(B), 4) < 2 ?
        Blade(B, volume=1 / volume(B), copy_basis=false) :
        Blade(B, volume=-1 / volume(B), copy_basis=false)

# --- Comparison methods

==(B::AbstractBlade, C::AbstractBlade) =
    dim(B) == dim(C) && grade(B) == grade(C) &&
    volume(B) == volume(C) && basis(B) == basis(C)

function isapprox(B::Blade{T1}, C::Blade{T2};
  atol::Real=0,
  rtol::Real=atol>0 ? 0 : max(√eps(T1), √eps(T2))) where {T1<:AbstractFloat,
                                                          T2<:AbstractFloat}
    # Check dim, grade, and norm are equal
    if dim(B) != dim(C) || grade(B) != grade(C) ||
        !isapprox(norm(B), norm(C), atol=atol, rtol=rtol)

        return false
    end

    # Check that B and C represent the same space
    projection = det(transpose(basis(B)) * basis(C))
    if ≉(abs(projection), 1, atol=atol, rtol=rtol)
        return false
    end

    # Check that B and C have the same orientation
    return sign(B) * sign(C) == sign(projection)
end

function isapprox(B::Blade{T1}, C::Vector{T2};
  atol::Real=0,
  rtol::Real=atol>0 ? 0 : max(√eps(T1), √eps(T2))) where {T1<:AbstractFloat,
                                                          T2<:AbstractFloat}
    C_blade = Blade(C)
    return isapprox(B, C_blade, atol=atol, rtol=rtol)
end

function isapprox(B::Vector{T1}, C::Blade{T2};
  atol::Real=0,
  rtol::Real=atol>0 ? 0 : max(√eps(T1), √eps(T2))) where {T1<:AbstractFloat,
                                                          T2<:AbstractFloat}
    return isapprox(C, B, atol=atol, rtol=rtol)
end

# --- Utility methods

convert(::Type{T}, B::Blade) where {T<:AbstractFloat} =
    T == typeof(volume(B)) ? B : Blade{T}(B)

#=
 TODO: review numerical error in factorizations to see if a different
       tolerance would be better.
=#
"""
    blade_atol(::Type{T}) where {T<:AbstractFloat}

Return the minimum value of a nonzero blade's norm.
"""
blade_atol(::Type{T}) where {T<:AbstractFloat} = 100*eps(T)


using Test
using Quon
using Yao

function check_inv(a, b, c)
    theta = [a,b,c] * im

    phi = yang_baxter_param_inv(yang_baxter_param(theta[1], theta[2], theta[3])...)
    phi = imag.([phi...])

    circ1 = chain(1, Rz(a), Rx(b), Rz(c))
    circ2 = chain(1, Rz(phi[1]), Rx(phi[2]), Rz(phi[3]))
    return operator_fidelity(circ1, circ2) ≈ 1
end
check_inv() = check_inv(rand(3)...)
check_inv()

@test check_inv()


"""
    VariableLinearAdvectionNonperiodicSemidiscretisation

A semidiscretisation of the linear advection equation
    \$\\partial_t u(t,x) + \\partial_x ( a(x) u(t,x) ) = 0\$
with boundary conditions `left_bc(t)`, `right_bc(t)`.
"""
struct VariableLinearAdvectionNonperiodicSemidiscretisation{T,Derivative<:AbstractDerivativeOperator{T},
                                                            Dissipation,
                                                            SplitForm<:Union{Val{false}, Val{true}},
                                                            LeftBC, RightBC} <: AbstractSemidiscretisation
    derivative::Derivative
    dissipation::Dissipation
    a::Vector{T}
    tmp1::Vector{T}
    tmp2::Vector{T}
    split_form::SplitForm
    left_bc::LeftBC
    right_bc::RightBC

    function VariableLinearAdvectionNonperiodicSemidiscretisation(derivative::Derivative, dissipation::Dissipation, afunc, split_form::SplitForm, left_bc::LeftBC, right_bc::RightBC) where {T, Derivative<:AbstractDerivativeOperator{T}, Dissipation, SplitForm<:Union{Val{false}, Val{true}}, LeftBC, RightBC}
        if dissipation != nothing
            @argcheck size(derivative) == size(dissipation) DimensionMismatch
            @argcheck grid(derivative) == grid(dissipation) ArgumentError
        end
        N = size(derivative, 2)
        a = compute_coefficients(afunc, derivative)
        tmp1 = Array{T}(undef, N)
        tmp2 = Array{T}(undef, N)
        new{T,Derivative,Dissipation,SplitForm,LeftBC,RightBC}(derivative, dissipation, a, tmp1, tmp2, split_form, left_bc, right_bc)
    end
end

function Base.show(io::IO, semidisc::VariableLinearAdvectionNonperiodicSemidiscretisation)
    print(io, "Semidiscretisation of the linear advection equation\n")
    print(io, "  \$ \\partial_t u(t,x) + \\partial_x ( a(x) u(t,x) ) = 0 \$ \n")
    print(io, "with nonperiodic boundaries using")
    if semidisc.split_form == Val{true}()
        print(io, " a split form and: \n")
    else
        print(io, " no split form and: \n")
    end
    print(io, semidisc.derivative)
    print(io, semidisc.dissipation)
end

function godunov_flux_variablelinearadvection(uₗ::T, uᵣ::T, a::T) where {T<:Real}
    ifelse(a > 0, a*uₗ, a*uᵣ)
end

function (disc::VariableLinearAdvectionNonperiodicSemidiscretisation)(du, u, p, t)
    @unpack a, tmp1, tmp2, derivative, dissipation, split_form, left_bc, right_bc = disc
    @boundscheck begin
        @argcheck length(u) == length(a)
        @argcheck length(u) == length(tmp1)
        @argcheck length(u) == length(tmp2)
        @argcheck length(u) == length(du)
    end

    # volume terms
    if typeof(split_form) <: Val{true}
        m1_2 = -one(eltype(u)) / 2

        ## a * D * u
        mul!(tmp1, derivative, u)
        @. du = m1_2 * a * tmp1
        ## u * D * a
        mul!(tmp1, derivative, a)
        @. du += m1_2 * u * tmp1
        ## D * a*u
        @. tmp2 = a*u
        mul!(tmp1, derivative, tmp2)
        @. du += m1_2 * tmp1
    else
        @. tmp2 = a * u
        mul!(tmp1, derivative, tmp2)
        @. du = -tmp1
    end

    # dissipation
    if dissipation != nothing
        mul!(tmp1, dissipation, u)
        @. du += tmp1
    end

    # boundary conditions via Godunov's flux
    @inbounds fnum_left = godunov_flux_variablelinearadvection(left_bc(t), u[1], a[1])
    @inbounds du[1] += (fnum_left - a[1]*u[1]) / left_boundary_weight(derivative)
    @inbounds fnum_right = godunov_flux_variablelinearadvection(u[end], right_bc(t), a[end])
    @inbounds du[end] -= (fnum_right - a[end]*u[end]) / right_boundary_weight(derivative)

    nothing
end


function conditional_gaussian_closure(sys::MomentEquations,
                                      binary_vars::Array{Int,1}=Int[])

    if isempty(binary_vars)
      error("conditional closure does not work if there are no binary species")
    end
    N = sys.N

    sys = bernoulli_moment_eqs(sys, binary_vars)
    # define symbolic raw moment expressions
    μ = typeof(sys) == CentralMomentEquations ? define_μ(N, sys.q_order) : copy(sys.μ)
    # express cumulants in terms of raw moments
    K = cumulants_to_raw_moments(N, sys.q_order)

    # closure of higher order raw moments without explicit form of truncated moments
    # e.g. μ₁₄ would still be a function of μ₁₃ even though μ₁₃ is also truncated
    closure_μ = OrderedDict()
    # closure of higher order raw moments explicitly expanding closed lower order moments
    # the form we use in the end when solving the ODEs
    closure_μ_exp = OrderedDict()

    # perform conditional gaussian closure on raw moments μ

    # by nonbernoulli we denote moments which cannot be written in the conditional form
    nonbernoulli_iters = filter(x -> sum(x[binary_vars]) == 0, sys.iter_all)
    for order in sys.m_order+1:sys.q_order

        # building the closed moment expressions order by order (due to such hierarchical functional dependency)
        iter_order = filter(x -> sum(x) == order, sys.iter_q)

        for iter in iter_order

            # Two cases: (i) conditional (⟨gp^j⟩=⟨p^j|g=1⟩⟨g⟩) or (ii) basic non bernoulli (⟨p^j⟩)"
            # Here g is a bernoulli variable (e.g gene state), p is any other stochastic variable (e.g no. of proteins)
            if sum(iter[binary_vars]) > 0
                # Case (i): ⟨gp^j⟩=⟨p^j|g=1⟩⟨g⟩
                # boils down to manipulation on indices:
                # 1. ignore the bernoulli variable state and write ⟨gp^j⟩ as μ₀ⱼ
                # 2. perform gaussian closure on μ₀ⱼ (assuming cumulant κ₀ⱼ = 0) and multiply by ⟨g⟩ (μ₁₀)
                # 3. add back bernoulli variable indices, for example:
                #   moment ⟨gp^3⟩ = ⟨p^3|g=1⟩⟨g⟩
                # i)  ⟨gp^3⟩ = μ₁₃ -> μ₀₃
                # ii) μ₀₃ = 3μ₀₂²*μ₀₁-2μ₀₁^3
                # iii) 3μ₀₂²*μ₀₁-2μ₀₁^3 -> 3μ₁₂²*μ₁₁-2μ₁₁^3

                # step 1
                bernoulli_iter = fill(0, N)
                for ind in binary_vars
                    if iter[ind] > 0
                        bernoulli_iter[ind] = 1
                    end
                end
                bernoulli_iter = Tuple(bernoulli_iter)
                r = iter.-bernoulli_iter

                # step 2
                conditional_μ = -μ[bernoulli_iter]*(K[r] - μ[r])
                conditional_μ = simplify(conditional_μ, expand=true)

                # step 3
                iter_conditional = filter(x -> 0 < sum(x) <= sum(r), nonbernoulli_iters)
                conditional_sub = Dict([Pair(μ[iter], μ[iter.+bernoulli_iter]/μ[bernoulli_iter])
                                       for iter in iter_conditional])
                conditional_μ = substitute(conditional_μ, conditional_sub)
                conditional_μ = simplify(conditional_μ, expand=true)

                closure_μ[μ[iter]] = conditional_μ
                closure_μ_exp[μ[iter]] = substitute(conditional_μ, closure_μ_exp)
                closure_μ_exp[μ[iter]] = simplify(closure_μ_exp[μ[iter]], expand=true)
            else
                # Case (ii)

                # Conditional Gaussian closure has no well-defined approximation for terms
                # such as ⟨p^j⟩ which are independent of the Bernoulli variables
                # we assume that the marginal distribution P(p) also follows a Gaussian
                # so that ⟨p^j⟩ is truncated according to normal closure

                moment = simplify(-(K[iter]-μ[iter]), expand=true)

                closure_μ[μ[iter]] = moment
                closure_μ_exp[μ[iter]] = substitute(moment, closure_μ_exp)
                closure_μ_exp[μ[iter]] = simplify(closure_μ_exp[μ[iter]], expand=true)
            end
        end
    end

    if typeof(sys) == CentralMomentEquations

        central_to_raw = central_to_raw_moments(N, sys.q_order)
        μ_central = Dict()
        for iter in vcat(sys.iter_m, sys.iter_q)
            μ_central[μ[iter]] = central_to_raw[iter]
        end

        μ_M_exp = define_μ(N, 1)
        for i in sys.iter_m
            μ_M_exp[i] = μ_central[μ[i]]
        end
        μ_M = copy(μ_M_exp)
        for i in sys.iter_q
            μ_M[i] = closure_μ[μ[i]]
            μ_M[i] = substitute(μ_M[i], μ_central)
            μ_M[i] = simplify(μ_M[i], expand=true)
            μ_M_exp[i] = closure_μ_exp[μ[i]]
            μ_M_exp[i] = substitute(μ_M_exp[i], μ_central)
            μ_M_exp[i] = simplify(μ_M_exp[i], expand=true)
        end
        closure = OrderedDict()
        closure_exp = OrderedDict()
        # construct the corresponding truncated expressions of higher order
        # central moments from the obtained raw moment expressions
        raw_to_central_exp = raw_to_central_moments(N, sys.q_order, μ_M_exp, bernoulli=true)
        raw_to_central = raw_to_central_moments(N, sys.q_order, μ_M, bernoulli=true)
        for i in sys.iter_q
            closure_exp[sys.M[i]] = simplify(raw_to_central_exp[i], expand=true)
            closure[sys.M[i]] = simplify(raw_to_central[i], expand=true)
        end

    else
        closure_exp = closure_μ_exp
        closure = closure_μ
    end

    close_eqs(sys, closure_exp, closure, false)

end


# This code is part of QuantumCircuits.
#
# (C) Copyright Rafał Pracht 2022.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

using Test

using QuantumCircuits
using QuantumCircuits.QML

using QuantumCircuits.QCircuits.Math
using QuantumCircuits.QCircuits.Circuit
using QuantumCircuits.QCircuits.Circuit: toQiskit
using QuantumCircuits.QCircuits.Qiskit
using QuantumCircuits.QCircuits.Qiskit: qiskit
using QuantumCircuits.Execute

statevectorsim = qiskit.Aer.get_backend("statevector_simulator")

function calcQiskitStateVector(qc)
    qcirc = toQiskit(qc)

    statevector = qiskit.execute(qcirc.qc, statevectorsim).result().get_statevector()
    return statevector.probabilities()
end

################################################################################
#  U3                                                                          #
################################################################################
circ_N = 1
circ_qr = QuantumRegister(circ_N)
circ_cr = ClassicalRegister(circ_N)
circ = QCircuit(circ_qr, circ_cr)
circ.u3(0, 1.3537, 2.5652, 5.5142)

@test sum(abs.(execute(circ) - calcQiskitStateVector(circ))) < 1e-6


using  jInv.Mesh
using  jInv.Utils
using  jInv.LinearSolvers
using  jInv.InverseSolve
using  jInvVis
using  EikonalInv
using  MAT
using  FWI
using  ForwardHelmholtz
using  PyPlot
#############################################################################################################

####################################################################################################################################
####################################################################################################################################
#######################   Problem definition and data creation #####################################################################
####################################################################################################################################
####################################################################################################################################

include("ex2DFWI/readModelAndGenerateMeshMref.jl");
include("ex2DFWI/prepareFWIDataFiles.jl");

m = readdlm("model/SEGmodel2Dsalt.dat");
m = m*1e-3;
m = m';

newSize          = [256,128];
pad     	     = 10;
ABLPad 		     = pad + 8;
jumpSrc 	 	 = 5
maxBatchSize     = 256;
omega   	     = [0.5,0.75,1.25,1.75]*2*pi;

offset  = ceil(Int64,(newSize[1]*(10.0/13.5)));
println("Offset is: ",offset)
domain = [0.0,13.5,0.0,4.2];

(m,Minv,mref,boundsHigh,boundsLow) = readModelAndGenerateMeshMref(m,pad,newSize,domain);

	limits = [1.5,4.5];
	figure(1,figsize = (22,10))
	plotModel(m,true,Minv,pad,limits);
	figure(2,figsize = (22,10))# ,figsize = (22,10)
	plotModel(mref,true,Minv,pad,limits);

useFilesForFields = false;

# ###################################################################################################################
dataFilenamePrefix = "ex2DFWI/DATA_SEG2D";

#######################################################################################################################

######################## DIRECT SOLVER #################################################
numCores 	= 4;
BLAS.set_num_threads(numCores);
Ainv = getMUMPSsolver([],0,0,2); # Alternatives: Ainv = getJuliaSolver();

##########################################################################################

println("omega*maximum(h): ",omega*maximum(Minv.h)*sqrt(maximum(1./(boundsLow.^2))));

# This is a list of workers for FWI. Ideally they should be on different machines.
workersFWI = [workers()[1]];
prepareFWIDataFiles(m,Minv,mref,boundsHigh,boundsLow,dataFilenamePrefix,omega,one(Complex128)*ones(size(omega)), pad,ABLPad,jumpSrc,
					offset,workersFWI,maxBatchSize,Ainv,useFilesForFields);

####################################################################################################################################
####################################################################################################################################
####################### Inversion setup for a single machine  ######################################################################
####################################################################################################################################
####################################################################################################################################
dataFilenamePrefix = "ex2DFWI/DATA_SEG2D";
file = matread(string(dataFilenamePrefix,"_PARAM.mat"));
n_cells 		= file["n"];
domain 	        = file["domain"];
gamma 			= file["gamma"];
omega 			= file["omega"];
waveCoef 		= file["waveCoef"];
boundsLow 		= file["boundsLow"];
boundsHigh 		= file["boundsHigh"];
mref 			= file["mref"];
file = 0;

Minv = getRegularMesh(domain,n_cells);
if length(omega)==1 # matlab saves a 1 variable array as scalar.
	omega = [omega];
	waveCoef = [waveCoef];
end

### Read receivers and sources files
RCVfile = string(dataFilenamePrefix,"_rcvMap.dat");
SRCfile = string(dataFilenamePrefix,"_srcMap.dat");
srcNodeMap = readSrcRcvLocationFile(SRCfile,Minv);
rcvNodeMap = readSrcRcvLocationFile(RCVfile,Minv);
Q = generateSrcRcvProjOperators(Minv.n+1,srcNodeMap);
Q = Q.*1/(norm(Minv.h)^2);
println("We have ",size(Q,2)," sources");
P = generateSrcRcvProjOperators(Minv.n+1,rcvNodeMap);

### Read the data files to an array of array pointers
Wd   = Array{Array{Complex128,2}}(length(omega))
dobs = Array{Array{Complex128,2}}(length(omega))
for k = 1:length(omega)
	omRound = string(round((omega[k]/(2*pi))*100.0)/100.0);
	(DobsFWIwk,WdFWIwk) =  readDataFileToDataMat(string(dataFilenamePrefix,"_freq",omRound,".dat"),srcNodeMap,rcvNodeMap);
	Wd[k] 	= WdFWIwk;
	dobs[k] = DobsFWIwk;
end

# setup active cells and background for the inversion.
N = prod(Minv.n+1);
Iact = sparse(I,N,N);
mback   = zeros(Float64,N);

########################################################################################################
##### Set up remote workers ############################################################################
########################################################################################################

batch = size(Q,2);         # We solve all the sources at once.
useFilesForFields = false; # a flag whether to store the fields on the disk.
## Set up a MUMPS direct solver #Ainv = getMUMPSsolver([],0,0,2); #
Ainv = getJuliaSolver();
## Choose the workers for FWI (here, its a single worker)
workersFWI = [1];
## Set up workers and division to tasks per frequencies ######################
(pFor,contDiv,SourcesSubInd) = getFWIparam(omega,waveCoef,vec(gamma),Q,P,Minv,Ainv,workersFWI,batch,useFilesForFields);
misfun = SSDFun;
pMis = getMisfitParam(pFor, Wd, dobs, misfun, Iact,mback);

########################################################################################################
# Set up Inversion #################################################################################
########################################################################################################

## models are usually given in velocity (km/sec). We invert here for the slowness, which is 1/v.
mref 				= velocityToSlow(mref)[1];
t    				= copy(boundsLow);
boundsLow 			= velocityToSlow(boundsHigh)[1];
boundsHigh 			= velocityToSlow(t)[1]; t = 0;
modfun 				= slowToSlowSquared;

maxStep				=0.05*maximum(boundsHigh);

regparams 			= [1.0,1.0,1.0,1e-5];
cgit 				= 8;
alpha 				= 1e-10;
pcgTol 				= 1e-1;
maxit 				= 8;
HesPrec 			= getExactSolveRegularizationPreconditioner();
regfun(m,mref,M) 	= wdiffusionRegNodal(m,mref,M,Iact=Iact,C=regparams);

pInv = getInverseParam(Minv,modfun,regfun,alpha,mref[:],boundsLow,boundsHigh,
                         maxStep=maxStep,pcgMaxIter=cgit,pcgTol=pcgTol,
						 minUpdate=1e-3, maxIter = maxit,HesPrec=HesPrec);

plotting = true;
function plotIntermediateResults(mc,Dc,iter,pInv,PMis,resultsFilename="")
	# Models are usually shown in velocity.
	fullMc = slowSquaredToVelocity(reshape(Iact*pInv.modelfun(mc)[1] + mback,tuple((pInv.MInv.n+1)...)))[1];
	if plotting
		close(888);
		figure(888);
		plotModel(fullMc,false,[],0,[1.5,4.8]);
	end
end

# Run one sweep of a frequency continuation procedure.
mc,Dc,flag,His = freqCont(copy(mref[:]), pInv, pMis,contDiv, 4, "",plotIntermediateResults,"Joint",1,1,"projGN");

subplot(1,3,1)
semilogy(His[end].F,"-o")
xlabel("PGNCG iterations")
title("misfit")

subplot(1,3,2)
semilogy(His[end].Rc,"-o")
xlabel("PGNCG iterations")
title("regularizer");

subplot(1,3,3)
semilogy(His[end].dJ,"-o")
xlabel("PGNCG iterations")
title("norm of proj grad");


using DelimitedFiles
using LOWDER

include("generators.jl")

# -----------------------------------------------
# Path to file
# -----------------------------------------------
source_filename = "CUTEr_selected_problems.dat"

# -----------------------------------------------
# LOWDER parameters
# -----------------------------------------------
δinit = 2.0
Δinit = 2.0

problems = readdlm( source_filename, Int64 )

# -----------------------------------------------
# Problem info
# -----------------------------------------------
i = 28
nprob = problems[i, 1]
n = problems[i, 2]
p = problems[i, 3]
rsp = convert( Bool, problems[i, 4] )
n_points = n + 1

# Generates the problem
( x, l, u, fmin ) = problem_generator_mw( nprob, n, p, rsp; unconstrained = true )

# Solves the problem using 'lowder'
sol = LOWDER.lowder( fmin, x, l, u; δ = δinit, Δ = Δinit, m = n_points, verbose = 3, filename="../data_files/simple_runtest_01.dat")


em = ErgodicManagerR2("double gaussian", K=5, bins=100)

# Trajectory manager
x0 = [0.89,0.01]
N = 400
h = 0.1
ci = ConstantInitializer([-0.01,0.01])
tm = TrajectoryManager(x0, h, N, ci)

xd, ud = smc_trajectory(em, tm, umax=.15)

# now test with two agents...

# Set up different domains with different discretizations
d = Domain([1,1], 100)
num_agents = 2

# Set up distribution and ergodic manager
K = 5
means = [[.3,.7], [.7,.3]]
Sigmas = [.025*eye(2), .025*eye(2)]
phi = gaussian(d, means, Sigmas)
em = ErgodicManagerR2(d, phi, K)

# Set up first trajectory manager
x0 = [0.49,0.01]
N = 50
h = 0.6
ci = ConstantInitializer([0.0, 0.0])
tm1 = TrajectoryManager(x0, h, N, ci)
dynamics!(tm1, SingleIntegrator(2,h))

# second tm is like the first, but different starting point
tm2 = deepcopy(tm1)
tm2.x0 = [.79,.99]

# array of trajectory managers
vtm = [tm1, tm2]

# Generate the trajectories
ddc = 1e-4
xd, ud = smc_trajectory(em, vtm, umax=.15)


using Plots
using NPZ
using Random

include("../HCOPE.jl")
include("../HCOVE.jl")
include("../UnO_utils.jl")
include("../UnO_params.jl")
include("../UnO_CDF.jl")

function compute_trend(props,
                        data::Array{Float64, 3},
                        pi_data::Array{Float64, 1},
                        trials::Int,
                        mini::Float64,
                        maxi::Float64,
                        intervals::Int)

    ######################################################
    # Pre-allocate arrays for storing all the estimates
    n_points = 1000      # Number of support points for plotting CDFs
    α_quantile = 0.5     # For inverse CDF based stuff
    α_cvar = 0.25    # For inverse CDF based stuff
    δ = 0.05             # Failure rates

    CDF_estimates = fill(0.0, trials, intervals, n_points)
    mean_estimates = fill(0.0, trials, intervals)
    variance_estimates = fill(0.0, trials, intervals)
    quantile_estimates = fill(0.0, trials, intervals)
    cvar_estimates = fill(0.0, trials, intervals)

    CDF_bounds = fill(0.0, trials, intervals, 2, n_points)
    mean_bounds = fill(0.0, trials, intervals, 2)
    mean_baseline_bounds = fill(0.0, trials, intervals, 4)
    variance_bounds = fill(0.0, trials, intervals, 2)
    variance_baseline_bounds = fill(0.0, trials, intervals, 4)
    quantile_bounds = fill(0.0, trials, intervals, 2)
    cvar_bounds = fill(0.0, trials, intervals, 2)

    CDF_boot_bounds = fill(0.0, trials, intervals, 2, n_points)
    mean_boot_bounds = fill(0.0, trials, intervals, 2)
    variance_boot_bounds = fill(0.0, trials, intervals, 2)
    quantile_boot_bounds = fill(0.0, trials, intervals, 2)
    cvar_boot_bounds = fill(0.0, trials, intervals, 2)

    # True CDF and parameters
    true_CDF = fill(0.0, n_points)
    CDF_plot_xs = fill(0.0, n_points)
    true_params = fill(0.0, 5)      # Order: Mean, quantile, CVaR, variance, entropy

    ###################################################
    # Compute estimates and bounds for CDF and params for different trials
    ###################################################

    step = (maxi - mini) / (intervals - 1)
    full_size, _, _ = size(data)
    trend_xs = fill(0.0, intervals)

    # Run in CLI before executing code
    # Be careful, this changes arrays asynchronously
    # export JULIA_NUM_THREADS=10
    #
    for index in 1:1:intervals

        n = Int(floor( 10^(mini + (index - 1)*step) ))
        trend_xs[index] = n

        perm_idxs = shuffle(collect(1:full_size))   # Shuffle indices
        Threads.@threads for trial in 1:1:trials
            println(n, ":", trial)

            ## Without replacement sampling
            ## Ensure unique data for each trial
            # start = (trial-1)*n + 1
            # stop = n*trial
            # idxs = perm_idxs[start:stop]

            # With replacement sampling
            perm2 = shuffle(perm_idxs)
            idxs = perm2[1:n]

            fragment = data[idxs, :, :]

            ###########################################
            # Estimators for CDF and other parameters
            ###########################################

            xs, ys = IS_CDF(fragment, props)
            wis_xs, wis_ys = WIS_CDF(fragment, props)     # WIS version

            CDF_estimates[trial, index, :] = get_equispaced_CDF(xs, ys, props, n_points)[2]
            mean_estimates[trial, index] = CDF_mean(xs, ys)
            variance_estimates[trial, index] = CDF_variance(xs, ys)
            quantile_estimates[trial, index] = CDF_quantile(xs, ys, α_quantile)
            cvar_estimates[trial, index] = CDF_CVaR(xs, ys, α_cvar)

            ############################################################
            # Guaranteed coverage bounds for CDF and other parameters
            ###########################################################

            # bound_xs, lb_ys, ub_ys = CDF_band(xs, ys, props, δ)
            bound_xs, lb_ys, ub_ys = CDF_band_optimized(xs, ys, props, δ)
            CDF_bounds[trial, index, 1, :] = get_equispaced_CDF(bound_xs, lb_ys, props, n_points)[2]
            CDF_bounds[trial, index, 2, :] = get_equispaced_CDF(bound_xs, ub_ys, props, n_points)[2]

            lb, ub = CDF_mean_bound(bound_xs, lb_ys, ub_ys)
            mean_bounds[trial, index, 1] = lb
            mean_bounds[trial, index, 2] = ub

            # Divide by 8 because of two-sided bounds for 4 statistics
            mean_baseline_bounds[trial, index, 1] = CI_lower(fragment, props, δ/8)
            mean_baseline_bounds[trial, index, 2] = CI_upper(fragment, props, δ/8)
            # mean_baseline_bounds[trial, index, 3] = CI_lower(fragment, props, δ/2)  # w/o union
            # mean_baseline_bounds[trial, index, 4] = CI_upper(fragment, props, δ/2)  # w/o union

            lb, ub = CDF_variance_bound(bound_xs, lb_ys, ub_ys)
            variance_bounds[trial, index, 1] = lb
            variance_bounds[trial, index, 2] = ub

            # Divide by 8 because of two-sided bounds for 4 statistics
            variance_baseline_bounds[trial, index, 1] = HCOVE_lower(fragment, props, δ/8)
            variance_baseline_bounds[trial, index, 2] = HCOVE_upper(fragment, props, δ/8)
            # variance_baseline_bounds[trial, index, 3] = HCOVE_lower(fragment, props, δ/2)   # w/o union
            # variance_baseline_bounds[trial, index, 4] = HCOVE_upper(fragment, props, δ/2)   # w/o union

            lb, ub = CDF_quantile_bound(bound_xs, lb_ys, ub_ys, α_quantile)
            quantile_bounds[trial, index, 1] = lb
            quantile_bounds[trial, index, 2] = ub

            lb, ub = CDF_CVaR_bound(bound_xs, lb_ys, ub_ys, α_cvar)
            cvar_bounds[trial, index, 1] = lb
            cvar_bounds[trial, index, 2] = ub

            ############################
            # Bootstrap based bounds
            ############################
            lb, ub = CDF_mean_boot_bound(wis_xs, wis_ys, δ/4)
            mean_boot_bounds[trial, index, 1] = lb
            mean_boot_bounds[trial, index, 2] = ub

            lb, ub = CDF_variance_boot_bound(wis_xs, wis_ys, δ/4)
            variance_boot_bounds[trial, index, 1] = lb
            variance_boot_bounds[trial, index, 2] = ub

            lb, ub = CDF_quantile_boot_bound(wis_xs, wis_ys, α_quantile, δ/4)
            quantile_boot_bounds[trial, index, 1] = lb
            quantile_boot_bounds[trial, index, 2] = ub

            lb, ub = CDF_CVaR_boot_bound(wis_xs, wis_ys, α_cvar, δ/4)
            cvar_boot_bounds[trial, index, 1] = lb
            cvar_boot_bounds[trial, index, 2] = ub

        end;
    end;

    #######################
    # True CDF and parameters
    ########################

    xs, ys = on_policy_CDF(pi_data)
    CDF_plot_xs, true_CDF = get_equispaced_CDF(xs, ys, props, n_points)
    true_params[1] = CDF_mean(xs, ys)
    true_params[2] = CDF_quantile(xs, ys, α_quantile)
    true_params[3] = CDF_CVaR(xs, ys, α_cvar)
    true_params[4] = CDF_variance(xs, ys)

    ##############################################
    ## Save all the results from the experiments
    #################################################

    npzwrite(props["name"]*"_CDF.npy", CDF_estimates)
    npzwrite(props["name"]*"_mean.npy", mean_estimates)
    npzwrite(props["name"]*"_variance.npy", variance_estimates)
    npzwrite(props["name"]*"_quantile.npy", quantile_estimates)
    npzwrite(props["name"]*"_cvar.npy", cvar_estimates)

    npzwrite(props["name"]*"_CDF_bounds.npy", CDF_bounds)
    npzwrite(props["name"]*"_mean_bounds.npy", mean_bounds)
    npzwrite(props["name"]*"_mean_baseline_bounds.npy", mean_baseline_bounds)
    npzwrite(props["name"]*"_variance_bounds.npy", variance_bounds)
    npzwrite(props["name"]*"_variance_baseline_bounds.npy", variance_baseline_bounds)
    npzwrite(props["name"]*"_quantile_bounds.npy", quantile_bounds)
    npzwrite(props["name"]*"_cvar_bounds.npy", cvar_bounds)

    npzwrite(props["name"]*"_CDF_boot_bounds.npy", CDF_boot_bounds)
    npzwrite(props["name"]*"_mean_boot_bounds.npy", mean_boot_bounds)
    npzwrite(props["name"]*"_variance_boot_bounds.npy", variance_boot_bounds)
    npzwrite(props["name"]*"_quantile_boot_bounds.npy", quantile_boot_bounds)
    npzwrite(props["name"]*"_cvar_boot_bounds.npy", cvar_boot_bounds)

    npzwrite(props["name"]*"_true_CDF.npy", true_CDF)
    npzwrite(props["name"]*"_CDF_plot_xs.npy", CDF_plot_xs)
    npzwrite(props["name"]*"_true_params.npy", true_params)
    npzwrite(props["name"]*"_trend_xs.npy", trend_xs)

    #######################
    # Plots
    ########################

    function average_plot!(xs, ys, label)
        μs = dropdims(mean(ys, dims=1), dims=1)
        σs = dropdims(std(ys, dims=1), dims=1) ./ sqrt(trials)
        plot!(xs,μs, ribbon=σs, fillalpha=.3, label=label, legend=:left)
    end

    ########## CDF Plots ############

    # selected_int = ceil(Int32, intervals/2 + 1)         # Take the middle number of samples
    selected_int = intervals         # Take the middle number of samples
    plotCDF = plot(CDF_plot_xs, true_CDF, label="Fπ")                   # true CDF
    average_plot!(CDF_plot_xs, CDF_estimates[:, selected_int, :], "̂̂F")              # esimtated CDF
    average_plot!(CDF_plot_xs, CDF_bounds[:, selected_int, 1, :], "̂̂F-")             # lower bound CDF
    average_plot!(CDF_plot_xs, CDF_bounds[:, selected_int, 2, :], "̂̂F+")             # upper bound CDF
    average_plot!(CDF_plot_xs, CDF_boot_bounds[:, selected_int, 1, :], "Boot-")     # boot lower bound CDF
    average_plot!(CDF_plot_xs, CDF_boot_bounds[:, selected_int, 2, :], "Boot+")     # boot upper bound CDF

    ######## Mean Plots ###############

    plotMean = plot([true_params[1]], linetype=:hline, label="μπ", xaxis=:log)   # true mean
    average_plot!(trend_xs, mean_estimates[:, :], "̂μ")                              # esimtated mean
    average_plot!(trend_xs, mean_bounds[:, :, 1], "̂μ-")                             # lower bound mean
    average_plot!(trend_xs, mean_bounds[:, :, 2], "̂μ+")                             # upper bound mean
    average_plot!(trend_xs, mean_baseline_bounds[:, :, 1], "Base-")                 # Baseline lower bound mean
    average_plot!(trend_xs, mean_baseline_bounds[:, :, 2], "Base+")                 # Baseline upper bound mean
    average_plot!(trend_xs, mean_boot_bounds[:, :, 1], "Boot-")                     # boot lower bound mean
    average_plot!(trend_xs, mean_boot_bounds[:, :, 2], "Boot+")                     # boot upper bound mean

    ######## Variance Plots ###############

    plotVariance = plot([true_params[4]], linetype=:hline, label="σπ", xaxis=:log)   # true mean
    average_plot!(trend_xs, variance_estimates[:, :], "̂μ")                              # esimtated mean
    average_plot!(trend_xs, variance_bounds[:, :, 1], "̂μ-")                             # lower bound mean
    average_plot!(trend_xs, variance_bounds[:, :, 2], "̂μ+")                             # upper bound mean
    average_plot!(trend_xs, variance_baseline_bounds[:, :, 1], "Base-")                 # Baseline lower bound mean
    average_plot!(trend_xs, variance_baseline_bounds[:, :, 2], "Base+")                 # Baseline upper bound mean
    average_plot!(trend_xs, variance_boot_bounds[:, :, 1], "Boot-")                     # boot lower bound mean
    average_plot!(trend_xs, variance_boot_bounds[:, :, 2], "Boot+")                     # boot upper bound mean

    ######## Quantile Plots ###############

    plotQuantile = plot([true_params[2]],linetype=:hline, label="Qπ", xaxis=:log)   # true quantile
    average_plot!(trend_xs, quantile_estimates[:, :], "̂Qπ")                             # esimtated quantile
    average_plot!(trend_xs, quantile_bounds[:, :, 1], "̂Q-")                             # lower bound quantile
    average_plot!(trend_xs, quantile_bounds[:, :, 2], "̂Q+")                             # upper bound quantile
    average_plot!(trend_xs, quantile_boot_bounds[:, :, 1], "Boot-")                     # boot lower bound quantile
    average_plot!(trend_xs, quantile_boot_bounds[:, :, 2], "Boot+")                     # boot upper bound quantile

    ######## CVaR Plots ###############

    plotCVaR = plot([true_params[3]],linetype=:hline, label="Cπ", xaxis=:log)   # true cvar
    average_plot!(trend_xs, cvar_estimates[:, :], "̂Cπ")                             # esimtated cvar
    average_plot!(trend_xs, cvar_bounds[:, :, 1], "̂C-")                             # lower bound cvar
    average_plot!(trend_xs, cvar_bounds[:, :, 2], "̂C+")                             # upper bound cvar
    average_plot!(trend_xs, cvar_boot_bounds[:, :, 1], "Boot-")                     # boot lower bound cvar
    average_plot!(trend_xs, cvar_boot_bounds[:, :, 2], "Boot+")                     # boot upper bound cvar

    ####### Failure rates ##############

    plotFailure = plot([δ], linetype=:hline, label="δ", xaxis=:log)

    # Simultaneous guaranteed coverage bounds failure
    f1l = mean_bounds[:, :, 1] .> true_params[1]         # Lower mean bound failed
    f2l = quantile_bounds[:, :, 1] .> true_params[2]     # Lower quantile bound failed
    f3l = cvar_bounds[:, :, 1] .> true_params[3]         # Lower cvar bound failed
    f4l = variance_bounds[:, :, 1] .> true_params[4]     # Lower variance bound failed

    f1u = mean_bounds[:, :, 2] .< true_params[1]         # Upper mean bound failed
    f2u = quantile_bounds[:, :, 2] .< true_params[2]     # Upper quantile bound failed
    f3u = cvar_bounds[:, :, 2] .< true_params[3]         # Upper cvar bound failed
    f4u = variance_bounds[:, :, 2] .< true_params[4]     # Upper variance bound failed

    ci_failure = f1l .| f2l .| f3l .| f4l .| f5l .| f1u .| f2u .| f3u .| f4u .| f5u       # Bitwise OR
    average_plot!(trend_xs, ci_failure, "CI")

    # Simultaneous Bootstrap bounds failure
    f1l = mean_boot_bounds[:, :, 1] .> true_params[1]         # Lower mean boot bound failed
    f2l = quantile_boot_bounds[:, :, 1] .> true_params[2]     # Lower quantile boot bound failed
    f3l = cvar_boot_bounds[:, :, 1] .> true_params[3]         # Lower cvar boot bound failed
    f4l = variance_boot_bounds[:, :, 1] .> true_params[4]     # Lower variance boot bound failed

    f1u = mean_boot_bounds[:, :, 2] .< true_params[1]         # Upper mean boot bound failed
    f2u = quantile_boot_bounds[:, :, 2] .< true_params[2]     # Upper quantile boot bound failed
    f3u = cvar_boot_bounds[:, :, 2] .< true_params[3]         # Upper cvar boot bound failed
    f4u = variance_boot_bounds[:, :, 2] .< true_params[4]     # Upper variance boot bound failed

    boot_failure = f1l .| f2l .| f3l .| f4l .| f1u .| f2u .| f3u .| f4u          # Bitwise OR
    average_plot!(trend_xs, boot_failure, "Boot")

    npzwrite(props["name"]*"_ci_failure.npy", ci_failure)
    npzwrite(props["name"]*"_boot_failure.npy", boot_failure)

    ###### Plot everything in a single plot ##########
    plot(plotCDF, plotMean, plotVariance, plotQuantile, plotCVaR, plotFailure, layout=(7,1))

end

function compute_NS_trend(props,
    data::Array{Float64, 5},  # speed, trial, number of eps, horizon, 2
    pi_data::Array{Float64, 2}, # speed, trials
    trials::Int)

    ######################################################
    # Pre-allocate arrays for storing all the estimates
    n_points = 1000      # Number of support points for plotting CDFs
    δ = 0.05             # Failure rates
    speeds = [0,1,2]     # Speeds of non-stationarity

    CDF_bounds = fill(0.0, length(speeds), trials, 2, n_points)

    # True CDF and parameters
    true_CDF = fill(0.0, length(speeds), n_points)
    CDF_plot_xs = fill(0.0, n_points)

    ###################################################
    # Compute estimates and bounds for CDF and params for different trials
    ###################################################

    for speed in speeds
        Threads.@threads for trial in 1:1:trials
            println(speed, ":", trial)

            fragment = data[speed+1, trial, :, :, :]

            # Guaranteed coverage bounds for CDF and other parameters
            # bound_xs, lb_ys, ub_ys = CDF_NS_band(fragment, props, δ)
            bound_xs, lb_ys, ub_ys = CDF_NS_band_optimized(fragment, props, δ)
            CDF_bounds[speed+1, trial, 1, :] = get_equispaced_CDF(bound_xs, lb_ys, props, n_points)[2]
            CDF_bounds[speed+1, trial, 2, :] = get_equispaced_CDF(bound_xs, ub_ys, props, n_points)[2]

        end;

        #######################
        # True CDF and parameters
        ########################
        xs, ys = on_policy_CDF(pi_data[speed+1, :])
        CDF_plot_xs, true_CDF[speed+1, :] = get_equispaced_CDF(xs, ys, props, n_points)
    end;

    ##############################################
    ## Save all the results from the experiments
    #################################################

    npzwrite(props["name"]*"_CDF_bounds.npy", CDF_bounds)
    npzwrite(props["name"]*"_true_CDF.npy", true_CDF)
    npzwrite(props["name"]*"_CDF_plot_xs.npy", CDF_plot_xs)

    #######################
    # Plots
    ########################

    function average_plot!(xs, ys, label)
        μs = dropdims(mean(ys, dims=1), dims=1)
        σs = dropdims(std(ys, dims=1), dims=1) ./ sqrt(trials)
        plot!(xs,μs, ribbon=σs, fillalpha=.3, label=label, legend=:left)
    end

    ########## CDF Plots ############

    plotCDF1 = plot(CDF_plot_xs, true_CDF[1, :], label="Fπ")                   # true CDF
    average_plot!(CDF_plot_xs, CDF_bounds[1, :, 1, :], "̂̂F-")             # lower bound CDF
    average_plot!(CDF_plot_xs, CDF_bounds[1, :, 2, :], "̂̂F+")             # upper bound CDF

    plotCDF2 = plot(CDF_plot_xs, true_CDF[2, :], label="Fπ")                   # true CDF
    average_plot!(CDF_plot_xs, CDF_bounds[2, :, 1, :], "̂̂F-")             # lower bound CDF
    average_plot!(CDF_plot_xs, CDF_bounds[2, :, 2, :], "̂̂F+")             # upper bound CDF

    plotCDF3 = plot(CDF_plot_xs, true_CDF[3, :], label="Fπ")                   # true CDF
    average_plot!(CDF_plot_xs, CDF_bounds[3, :, 1, :], "̂̂F-")             # lower bound CDF
    average_plot!(CDF_plot_xs, CDF_bounds[3, :, 2, :], "̂̂F+")             # upper bound CDF

    ###### Plot everything in a single plot ##########
    plot(plotCDF1, plotCDF2, plotCDF3, layout=(3, 1))

end

####
# For all of the domains, R \in [rmin, rmax] is always required
# If the range for return [Gmin, Gmax] is also known,
# then the code can use that to make bounds tighter as well.

domain_props = Dict(
    "Reco" => Dict(
                    "beta_path" => "../../Experiments/Reco/ActorCritic/DataDefault/0/Results/beta_trajectories_0.5.npy",
                    "pi_path" => "../../Experiments/Reco/ActorCritic/DataDefault/0/Results/eval_returns_0.5.npy",
                    "rmin" => 0.0,
                    "rmax" => 10.0,
                    "Gmin" => 0.0,
                    "Gmax" => 10.0,
                    "horizon" => 1,
                    "gamma" => 0.99,
                    "name" => "Reco"
    ),
    "NS_Reco" => Dict(
                    "beta_path" => "../../Experiments/NS_Reco/beta_trajectories_1000.npy",
                    "pi_path" => "../../Experiments/NS_Reco/eval_returns_1000.npy",
                    "rmin" => -1.0,
                    "rmax" => +1.0,
                    "Gmin" => -1.0,
                    "Gmax" => +1.0,
                    "horizon" => 1,
                    "gamma" => 1.0,
                    "name" => "NS_Reco"
    ),
    "Maze_POMDP" => Dict(
                    # POMDP + multiple behavior policies
                    "beta_path" => "../../Experiments/Maze/ActorCritic/DataDefault/0/Results/beta_trajectories_2.npy",
                    "pi_path" => "../../Experiments/Maze/ActorCritic/DataDefault/0/Results/eval_returns_0.75.npy",
                    "rmin" => -1.0,
                    "rmax" => +1.0,
                    "Gmin" => -6.0,
                    "Gmax" => -1.0,
                    "horizon" => 6,
                    "gamma" => 0.99,
                    "name" => "Maze_POMDP"
    ),
    "Diabetes" => Dict(
                    "beta_path" => "../../Experiments/SimGlucosediscrete-v0/ActorCritic/DataDefault/0/Results/beta_trajectories_0.5.npy",
                    "pi_path" => "../../Experiments/SimGlucosediscrete-v0/ActorCritic/DataDefault/0/Results/eval_returns_0.5.npy",
                    "rmin" => -15.0,
                    "rmax" => +15.0,
                    "Gmin" => -15.0,
                    "Gmax" => +15.0,
                    "horizon" => 1,
                    "gamma" => 0.99,
                    "name" => "Diabetes"
    )

)

function run_UnO(domain)
    props = domain_props[domain]
    beta_data = npzread(joinpath(@__DIR__, props["beta_path"]))
    pi_data = npzread(joinpath(@__DIR__, props["pi_path"]))
    compute_trend(props, beta_data, pi_data, 30, 2.0, 4.5, 6)
    # compute_trend(props, beta_data, pi_data, 3, 2.0, 4.0, 5)
end;

function run_NS_UnO(domain)
    props = domain_props[domain]
    beta_data = npzread(joinpath(@__DIR__, props["beta_path"]))
    pi_data = npzread(joinpath(@__DIR__, props["pi_path"]))
    compute_NS_trend(props, beta_data, pi_data, 30)
end;

#######################
# IMP: To make use threading
#
# Run in CLI before executing code
# export JULIA_NUM_THREADS=10
###########################

# run_UnO("Maze")
# run_UnO("Reco")
# run_UnO("Diabetes")

# run_UnO("Maze_POMDP")
# run_NS_UnO("NS_Reco")


using StaticArrays
using StaticArrays.ImmutableArrays
using PyPlot
export Polygon, isInside, isIntersect

const Vec2f = SVector{2, Float64}
const Vec4f = SVector{4, Float64}

abstract type Polygonic end

struct Polygon <: Polygonic
    N::Int
    V::Vector{Vec2f}
    function Polygon(V_::AbstractVector) #V nust be counterclockwise
        N = length(V_)
        V = Vec2f[]
        for v in V_
            push!(V, Vec2f(v))
        end
        new(N, V)
    end
end

struct Rectangle <: Polygonic
    N::Int
    V::Vector{Vec2f}
    function Rectangle(center::T, width::Float64, height::Float64) where {T<:AbstractVector}
        V1 = Vec2f(center + [-width, -height]*0.5)
        V2 = Vec2f(center + [+width, -height]*0.5)
        V3 = Vec2f(center + [+width, +height]*0.5)
        V4 = Vec2f(center + [-width, +height]*0.5)
        new(4, [V1, V2, V3, V4])
    end
end

@inline function isInside(this::T, q::Vec2f) where {T<:Polygonic}
    # note: n'x = n'b
    @inbounds q_vec = [q[1]; q[2]]
    for n = 1:this.N
        @inbounds p1 = this.V[n]
        @inbounds n<this.N ? p2 = this.V[n+1] : p2 = this.V[1]
        @inbounds u = [p2[1]-p1[1]; p2[2]-p1[2]] #p2-p1
        @inbounds n = [0 1; -1 0]*u # normla vector
        @inbounds px = [q[1]-p1[1]; q[2]-p1[2]] # x-p1
        px'*n>0 && return false
    end
    return true
end

@inline function isIntersect(this::T, q1_::AbstractVector, q2_::AbstractVector) where {T<:Polygonic}
    @views q1 = Vec2f(q1_[1:2])
    @views q2 = Vec2f(q2_[1:2])
    return isIntersect(this, q1, q2)
end

@inline function isIntersect(this::T, q1::Vec2f, q2::Vec2f)::Bool where {T<:Polygonic}
    for n = 1:this.N
        p1 = this.V[n]
        if n<this.N
            p2=this.V[n+1]
        else
            p2=this.V[1]
        end
        isIntersect(p1, p2, q1, q2) && return true
    end
    return false
end

@inline function isIntersect(this::T, q_seq::Vector{Vec2f})::Bool where {T<:Polygonic}
    for q in q_seq
        isInside(this, q) && return true
    end
    return false
end

function show(this::T, color=:red) where {T<:Polygonic}
    for n = 1:this.N
        if n<this.N
            x = [this.V[n][1], this.V[n+1][1]]
            y = [this.V[n][2], this.V[n+1][2]]
        else
            x = [this.V[n][1], this.V[1][1]]
            y = [this.V[n][2], this.V[1][2]]
        end
        plot(x, y, "r-")
    end
end

@inline function isIntersect(p1::Vec2f, p2::Vec2f, q1::Vec2f, q2::Vec2f)::Bool
    # solve [p2-p1, q2-q1]*[s; t]=[q1-p1] or A*[s; t]=B
    # for computational efficiency, not using inv() func
    # A = [a, b; c, d]
    # B = [e; f]
    a = p2[1]-p1[1]
    b = -(q2[1]-q1[1])
    c = p2[2]-p1[2]
    d = -(q2[2]-q1[2])
    e = q1[1]-p1[1]
    f = q1[2]-p1[2]
    det = a*d-b*c
    abs(det)<1e-5 && return false

    s = 1/det*(d*e-b*f)
    t = 1/det*(-c*e+a*f)
    return (0.0<=s<=1.0) && (0.0<=t<=1.0)
end

function test()
    """
    v1 = [0.0, 0.0]
    v2 = [1.0, 0.0]
    v3 = [1.0, 1.0]
    v4 = [0.0, 1.0]
    V = [v1, v2, v3, v4]
    P = Polygon(V)
    """
    P = Rectangle([0.5 0.5], 1.0, 1.0)
    isInside(P, Vec2f([1.2, 0.8]))
end