export @cast, @cast!, @reduce, @reduce!

"""
    @cast Z[i,j,...] := f(A[j,k,...])  options

Macro for broadcasting, reshaping, and slicing of arrays in index notation.
Understands the following things:
* `A[i,j,k]` is a three-tensor with these indices.
* `B[(i,j),k]` is the same thing, reshaped to a matrix. Its first axis (the bracket) is indexed
  by `n = i + (j-1) * N` where `i ∈ 1:N`. This may also be written `B[i\\j,k]` or `B[i⊗j,k]`.
* `C[k][i,j]` is a vector of matrices.
* `D[j,k]{i}` is an ordinary matrix of `SVector`s, which may be reinterpreted from `A[i,j,k]`.
* `E[i,_,k]` has two nontrivial dimensions, and `size(E,2)==1`. On the right hand side
  (or when writing to an existing array) you may also write `E[i,3,k]` meaning `view(E, :,3,:)`,
  or `E[i,\$c,j]` to use a variable `c`. Fixing inner indices, like `C[k][i,_]`, is not allowed.
* `f(x)[i,j,k]` is allowed; `f(x)` must return a 3-tensor (and will be evaluated only once).
* `g(H[:,k])[i,j]` is a generalised `mapslices`, with `g` mapping columns of `H`
    to matrices, which are glued into a 3-tensor `[i,j,k]`.
  Similarly `h(I[k])[i,j]` expects an `h` which maps scalars to matrices.
* `J[i,-j,k]` means roughly `reverse(J, dims=2)`.
* `K[i,j]'` conjugates each element, equivalent to `K'[j,i]` which is the
  conjugate-transpose of the matrix.
* `M[i,i]` means `diag(M)[i]`, but only for matrices: `N[i,i,k]` is an error.

The left and right hand sides must have all the same indices,
with the sole exception of `M[i,i]`.
See `@reduce` and `@mul` for related macros which can sum over things.

If a function of one or more tensors appears on the right hand side,
then this represents a broadcasting operation,
and the necessary re-orientations of axes are automatically inserted.

The following actions are possible:
* `=` writes into an existing array, `copyto!(Z, ...)`.
* `:=` creates a new object. This need not be named: `Z = @cast [i,j,k] := ...` is allowed.
* `==` insists on a view of the old object (error if impossible), and `|=` insists on a copy.
* `=>` creates an anonymous function, for which the output must be on the right.
   For example `@cast A[i] + B[j]' => Z[i,j]` gives `(A,B) -> A .+ B'`.

Re-ordering of indices `Z[k,j,i]` is done lazily with `PermutedDimsArray(A, ...)`.
Reversing of an axis `F[i,-j,k]` is also done lazily, by `Reverse{2}(F)` which makes a `view`.
Using `|=` (or broadcasting) will produce a simple `Array`.

Options can be specified at the end (if several, separated by `,` i.e. `options::Tuple`)
* `i:3` supplies the range of index `i`. Variables and functions like `j:Nj, k:length(K)`
  are allowed.
* `assert` or `!` will turn on explicit dimension checks of the input.
  (Providing ranges may also turn these on, but imperfectly, confirm with `@pretty @cast ...`.)
* `cat` will glue slices by things like `hcat(A...)` instead of the default `reduce(hcat, A)`,
  and `lazy` will instead make a `VectorOfArrays` container.
* `nolazy` disables `PermutedDimsArray` and `Reverse` in favour of `permutedims` and `reverse`,
  and `Diagonal` in favour of `diagm` for `Z[i,i]` output.

These options only work if you load other packages:
* `strided` will place `@strided` in front of broadcasting operations,
  and use `@strided permutedims(A, ...)` instead of `PermutedDimsArray(A, ...)`.
  For this you need `using Strided`.
* Static slices `D[j,k]{i}` need `using StaticArrays`, and to create them you should
  give all slice dimensions explicitly. You may write `D[k]{i:2,j:2}` to specify
  `Size(2,2)` slices. They are made most cleanly from the first indices of the input,
  i.e. this `D` from `A[i,j,k]`.
"""
macro cast(exs...)
    where = (mod=__module__, src=__source__, str=unparse("@cast", exs...))
    _macro(exs...; reduce=false, where=where)
end

"""
    @cast! Z[i...] := A[j...] opt

Variant of `@cast` which effectively runs `@check!()` on each tensor.
"""
macro cast!(exs...)
    where = (mod=__module__, src=__source__, str=unparse("@cast!", exs...))
    _macro(exs...; reduce=false, where=where, icheck=true)
end

"""
    @reduce A[i] := sum(j,k) B[i,j,k]             # A = vec(sum(B, dims=(2,3)))
    @reduce A[i] := prod(j) B[i] + ε * C[i,j]     # A = vec(prod(B .+ ε .* C, dims=2))
    @reduce A[i] = sum(j) exp( C[i,j] / D[j] )    # sum!(A, exp.(C ./ D') )

Tensor reduction macro:
* The reduction function can be anything which works like `sum(B, dims=(1,3))`,
  for instance `prod` and `maximum` and `Statistics.mean`.
* In-place operations `Z[j] = sum(...` will construct the banged version of the given function's name,
  which must work like `sum!(Z, A)`.
* The tensors can be anything that `@cast` understands, including gluing of slices `B[i,k][j]`
  and reshaping `B[i\\j,k]`. See `? @cast` for the complete list.
* If there are several tensors (or scalars) on the right, then this is a broadcasting operation.
* Index ranges may be given afterwards (as for `@cast`) or inside the reduction `sum(i:3, k:4)`.
* All indices appearing on the right must appear either within `sum(...)` etc, or on the left.

    F = @reduce sum(i,j)  B[i] + γ * D[j]         # sum(B .+ γ .* D')
    @reduce G[] := sum(i,j)  B[i] + γ * D[j]      # F == G[]

Complete reduction to a scalar output `F`, or a zero-dim array `G`.
`G[]` involves `sum(A, dims=(1,2))` rather than `sum(A)`.

    @reduce Z[k] := sum(i,j) A[i] * B[j] * C[k]  lazy, i:N, j:N, k:N

The option `lazy` replaces the broadcast expression with a `BroadcastArray`,
to avoid `materialize`ing the entire array before summing. In the example this is of size `N^3`.

The option `strided` will place `@strided` in front of the broadcasting operation.
You need `using Strided` for this to work.

    @reduce sum(i) A[i] * log(@reduce [i] := sum(j) A[j] * exp(B[i,j]))
    @cast W[i] := A[i] * exp(- @reduce S[i] = sum(j) exp(B[i,j]) lazy)

Recursion like this is allowed, inside either `@cast` or `@reduce`.
The array being created need not have a name, like `[i]`,
unless writing into an existing array, like `S[i]` here.
"""
macro reduce(exs...)
    where = (mod=__module__, src=__source__, str=unparse("@reduce", exs...)) # wtf?
    _macro(exs...; reduce=true, where=where)
end

"""
    @reduce! Z[j] := sum(i,k) A[i,j,k]

Variant of `@reduce` which effectively runs `@check!()` on each tensor.
"""
macro reduce!(exs...)
    where = (mod=__module__, src=__source__, str=unparse("@reduce!", exs...))
    _macro(exs...; reduce=true, where=where, icheck=true)
end

#=

This much is now per RHS term. None need exist outside the function:

nameA, indA, indAsub
sizeA      -- known & used for fixing ranges
getB, numB -- for view(A, get), and number of fixed indices
sizeC      -- after reshaping outer indices, in terms of sz_i
codeD      -- for gluing
indEflat   -- after gluing
negF       -- done before permutedims
shiftG     -- ditto, but not yet written
permH      --
codeI      -- to orient, if lacking some indices.

Of these things from LHS, only really store + canon need to be passed to RHS function inputex()
via walker(). But many need to be passed from readleft() to output, packaged into:
outUZ = (redUind, negV, codeW, sizeX, getY, numY, indZ, sizeZ, nameZ)

flags      -- todo, to-done, and options
store      -- information from parse! with sizes etc.
canon      -- canonical list of indices
canonsize  -- filled with size(A,2) etc, from store, at most one (:)

redUind    -- if reducing, this is done before slicing
negV       -- directions to reverse, done here.
codeW      -- for slicing. sizeWstatic is Size() of the slice.
sizeX      -- after slicing, container size, in sz_i
getY, numY -- for view(Z, get) in-place, and number fixed
sizeZ      -- for final reshape, in sz_i
nameZ, indZ, indZsub, indZred -- given LHS

=#

function _macro(exone, extwo=nothing, exthree=nothing;
    reduce=false, icheck=false, where=nothing, recurse=false)

    flags = Set{Symbol}()

    if reduce
        #===== parse @reduce input =====#

        if @capture(exone, left_ := redfun_(redind__) )     # Z[i] := sum(j) A[i] * B[j]
            right = extwo
            options = exthree
            push!(flags, :reduce)
            V && @info "partial reduce" left right redfun Tuple(redind) options
        elseif @capture(exone, left_ = redfun_(redind__) )
            right = extwo
            options = exthree
            push!(flags, :reduce)
            push!(flags, :inplace)
        elseif @capture(exone, left_ |= redfun_(redind__) )
            right = extwo
            options = exthree
            push!(flags, :reduce)
            push!(flags, :mustcopy)

        elseif @capture(exone, redfun_(redind__) )          # sum(i,j) A[i] * B[j]
            right = extwo
            options = exthree
            push!(flags, :reduce)
            push!(flags, :scalar)
            V && @info "full reduce" right redfun Tuple(redind)

        else
            throw(MacroError("don't know what to do with $exone", where))
        end

    else
        #===== parse @cast input =====#

        if @capture(exone, left_ := right_ )                # Z[i,j] := A[i] * B[j]
        elseif @capture(exone, left_ = right_ )
            push!(flags, :inplace)
        elseif @capture(exone, left_ |= right_ )
            push!(flags, :mustcopy)
        elseif @capture(exone, left_ == right_ )
            push!(flags, :mustview)
                                                            # A[i,j] -> Z[j][i]
        elseif @capture(exone, right_ -> left_ ) ||  @capture(exone, right_ => left_ )
            push!(flags, :anonfunc)
        else
            throw(MacroError("@cast doesn't know what to do with $exone", where))
        end
        V && @info "no reduction" left right
        options = extwo
        redind = []
        redfun = identity
        exthree == nothing || throw(MacroError("@cast doesn't know what to do with $exthree", where))

    end

    #===== parse LHS to get canonical list =====#

    store = SizeDict()
    canon, outUZ, nameZ, checkZ = readleft(left, redind, flags, store, icheck, where)

    # parse options both to look for keywords and sizes
    @capture(options, (optvec__,)) || (optvec = Any[options])
    optind, _,_,_ = parse!(store, nothing, [], optvec; allowranges=true, flags=flags)

    if count(i -> i != nothing, setdiff(optind, canon)) > 0
        str = join(something.(setdiff(optind, canon), "nothing"), ", ")
        m_error("attempting to ignore unrecognised options: $str", where)
    end

    #===== parse and process RHS =====#

    outex = quote end

    if @capture(right, AA_[ii__] ) || @capture(right, AA_{ii__} )
        newright = walker(outex, right, canon, flags, store, icheck, where)
    else
        newright = MacroTools.prewalk(
            x -> walker(outex, x, canon, flags, store, icheck, where), right)
        push!(flags, :broadcast)
    end

    notseen = setdiff(canon, unique(store.seen))
    isempty(notseen) || throw(MacroError("did not see index $(join(notseen, ", ")) on the right", where))

    #===== almost done =====#

    packagecheck(flags, where) # disabled for issue #2

    canonsize = sizeinfer(store, canon, where, true)

    V && @info "before in/out choice" store flags Tuple(canonsize)

    if :inplace in flags

        #===== in-place output  =====#

        if checkZ != nothing
            push!(outex.args, checkZ)
        end

        inout = outputinplace(newright, outUZ, redfun, canonsize, canon, flags, store, nameZ, where)
        push!(outex.args, inout)

        if :finalres in flags
            push!(outex.args, nameZ)
        end

    else
        #===== out-of-place output  =====#

        if :broadcast in flags
            if (:reduce in flags) && (:lazy in flags)
                newright = :( Base.@__dot__ TensorCast.lazy($newright) ) # really LazyArrays.lazy
            elseif :strided in flags
                dotted = Broadcast.__dot__(newright) # @strided does not work on @.
                newright = :( Strided.@strided $dotted )
            else
                newright = :( Base.@__dot__ $newright ) # prettier in @pretty
            end
        end

        finalright = outputnew(newright, outUZ, redfun, canonsize, canon, flags, store, where)
        push!(outex.args, :( $nameZ =  $finalright ) )

        if checkZ != nothing
            push!(outex.args, checkZ)
            push!(outex.args, nameZ)
        end
    end

    #===== finalise =====#

    if :needsize in flags
        szcanon = Any[ Symbol(:sz_,i) for i in canon ]
        pushfirst!(outex.args, :( local ($(szcanon...),) = ($(canonsize...),) ) )
    end

    if :assert in flags || :(!) in flags || check_options.size
        for ch in store.checks
            pushfirst!(outex.args, ch)
        end
    end

    for tex in store.topex
        pushfirst!(outex.args, tex)
    end

    if recurse==true # only for @reduce inside something
        indZ = outUZ[end-1]
        return outex, indZ
    end

    if :anonfunc in flags
        leftex = anoninput(store.rightnames, where)
        outex = :( $leftex -> $outex )
        outex.args[2] = MacroTools.unblock(outex.args[2])
        return esc(outex)
    else
        outex = MacroTools.unblock(outex) # removes begin end if it can
        return esc(outex)
    end
end

#="""
    canon, outUZ, nameZ, checkZ = readleft(left, redind, flags, store, icheck, where)

outUZ = (redUind, negV, codeW, sizeX, getY, numY, indZ, sizeZ)
are things passed to output construction.
"""=#
function readleft(left, redind, flags, store, icheck, where)

    if @capture(left, Z_[outer__][inner__]) ||  @capture(left, [outer__][inner__])
        push!(flags, :slice)
    elseif @capture(left, Z_[outer__]{inner__}) || @capture(left, [outer__]{inner__})
        push!(flags, :staticslice)
    elseif @capture(left, Z_[outer__]) || @capture(left, [outer__])
        inner = []
    elseif left==nothing
        @assert :scalar in flags
        inner = []
        outer = []
        Z = nothing
    else
        throw(MacroError("readleft doesn't know what to do with $left", where))
    end

    if Z == nothing
        nameZ = gensym(:Z) # no output name
        if :inplace in flags
            throw(MacroError("can't write in-place into nowhere!", where))
        end
    else
        nameZ = Z
    end

    if !(:inplace in flags)
        Z = nothing # tells parse! not to think about size(Z,...)
    end
    flat12, getY, sizeZ, negV = parse!(store, Z, outer, inner; allowranges=true) # true allows A[i]{j:3}
    redUind, _,_,_ = parse!(store, nothing, [], redind; allowranges=true) # true means sum(i:3) allowed

    if length(flat12)-length(inner) == 2 && flat12[end] == flat12[end-1] # Z[i,i] special case
        length(getY) == 2 || throw(MacroError("can't fix output index with Diagonal output", where)) # because I'm lazy, this would not be impossible
        push!(flags, :diagleft)
        pop!(flat12) # remove index from flat list
        pop!(getY)
        pop!(sizeZ)
    end

    canon = vcat(flat12, redUind) # order here is [inner, outer, reduced]

    checkrepeats(flat12, " on left hand side", where)
    checkrepeats(canon, " in reduction function", where)

    codeW = repeat(Any[*], length(canon) - length(redUind))
    codeW[1:length(inner)] .= (:)

    indW = canon[1:length(inner)] # inner without minuses etc

    indX = flat12[length(inner)+1:end] # flattened outer without minuses, also without fixed
    sizeX = Any[ Symbol(:sz_,i) for i in indX ] # only for in-place

    numY = count(!isequal(:), getY) # the number of fixed indices in output

    if length(getY) - numY != length(sizeX)
        push!(flags, :backshape) # whether reshaping of view(Z, getY) is needed, only in-place
    end

    if (length(sizeZ) + numY) != (length(canon) - length(inner) - length(redUind))
        push!(flags, :outshape) # whether final reshaping is needed, for := case
    end

    indZ = [] # like indX, but with fixed indices re-inserted, for :named / :axis
    iz = 1
    for g in getY
        if g == (:)
            push!(indZ, indX[iz])
            iz += 1
        else
            push!(indZ, g)
        end
    end

    checkZ = nothing
    if icheck
        checknow = check_one(:($nameZ[$(outer...)]), where)
        if check_options.size
            checkZ = checknow # check!(...) to be inserted
        end
    end

    outUZ = (redUind, negV, codeW, indW, sizeX, getY, numY, indZ, sizeZ)
    V && @info "readleft" Tuple(canon) outUZ nameZ
    return canon, outUZ, nameZ, checkZ
end

#="""
    walker(outex, x, canon, flags, store, icheck, where)

Called by `MacroTools.prewalk` on RHS, finds tensors & pushes `:( sym = inputex(this) )` into
`outex.args`, and then replaces them with `sym`.
"""=#
function walker(outex, ex, canon, flags, store, icheck, where)

    # allow @reduce inside RHS, by simply calling the entire macro first
    if @capture(ex, @reduce(redex__))
        Rsym = gensym(:R)
        Rval, Rind = _macro(redex...; reduce=true, where=where, recurse=true)

        # sew the result of that into ... not outex but topex, as may need its size:
        push!(store.topex, :(local $Rsym = $Rval ) )

        # construct name for resulting tensor, for outer @cast to further process
        ex =  :( $Rsym[$(Rind...)] )

    elseif @capture(ex, @mul(mulex__))
        Msym = gensym(:M)
        Mval, Mind = _mul(mulex...; where=where, recurse=true)

        push!(store.topex, :(local $Msym = $Mval ) )

        ex =  :( $Msym[$(Mind...)] )
    end

    # find any indexed tensor, process it, and replace it
    if @capture(ex, A_[ij__][kl__] ) || @capture(ex, A_[ij__]{kl__} ) || @capture(ex, A_[ij__] )

        push!(store.rightnames, A) # used for @cast A[i,j] -> B[i\j]

        if isa(A, Symbol) || @capture(A, Astruct_.field_)

        # allow something like mapslices:
        # @pretty @cast Z[i,j] := fun(X[:,i,:],42)[j] + X[j,i,3] + Y[j]^2
        elseif @capture(A, f_(B_[i_colons__],rest__) )

            kl == nothing || throw(MacroError("can't nest this deeply: $ex", where))

            # i_colons = vcat(rvec(ii1), rvec(i1), Any[:(:)], rvec(i2), rvec(ii2)) # NB contains the symbol :(:) not the function (:)

            A, ex = slicemap(f, B, i_colons, rest, ij, store)

        # if we have f(x)[i,j] then we should evaluate f just once, e.g.
        # @pretty @cast A[i]{j:5} |= rand(15)[i\j]
        else
            Atop = gensym(:A)
            push!(store.topex, :(local $Atop = $A ) )
            A = Atop
        end

        V && @info "walker before inputex" A ex canon store

        Aval = inputex(A, ex, canon, flags, store, icheck, where)

        if isa(Aval, Symbol)
            ex = Aval
        else
            Asym = gensym(:A)
            push!(outex.args, :(local $Asym = $Aval) )
            ex =  Asym
        end
        V && @info "walker" ex Aval

    # catch A[i,j]' as element-wise conj
    elseif @capture(ex, arg_')
        ex = :( Base.conj($arg) )

    # try to protect log(2) etc from @. , but not log(A[i])
    elseif @capture(ex, f_(x_Symbol) ) || @capture(ex, f_(x_Int) ) || @capture(ex, f_(x_Float64) )
        fxsym = gensym(:fx)
        push!(outex.args, :( local $fxsym = $ex ))
        ex = fxsym
    end

    return ex
end

#="""
    slicemap(f, :B, [:i, :(:), :j], rest, kl, store)

This processes `f(B[i,:mj],rest...)[k,l]`, such as in these:
```
@pretty @cast Z[i,j] := fun(X[:,i,:],42)[j] + X[j,i,3] + Y[j]^2
@pretty @cast Z[i,j\k] := fun(X[:,i\k,:])[j] + Y[k]^2
```

* Note that `f(X[i,:], Y[i,:])` is not allowed (yet).

* `g(X[i,:,_])` should work, but notices _ after map,
and thus `g(X[i,:,3])` will evaluate g on many slices which aren't required.

These could be fixed by upgrading this function -- TODO maybe?
"""=#
function slicemap(f, B, i_colons, rest, kl, store)

    if isa(B, Symbol) || @capture(B, Bstruct_.field_)
    else # pull out B = B(x) so that it's evaluated only once
        Btop = gensym(:B)
        push!(store.topex, :(local $Btop = $B ) )
        B = Btop
    end

    # assertions, and size information for resising etc.
    Bdims = length(i_colons)
    str = "expected a $Bdims-tensor $B[" * join(i_colons, ", ") * "]"
    push!(store.checks, :(TensorCast.@assert_ ndims($B)==$Bdims $str) )
    for (dout, ind) in enumerate(i_colons)
        if ind != :(:) && !isconstant(ind)
            savesize!(store, ind, :( size($B, $dout) ) )
        end
    end

    # construct map(f,eachslice(B)) expression
    Atop = gensym(:slicemap)
    i_code = map(i -> i==:(:) ? (:) : (*), Tuple(i_colons))
    if count(isequal(:), i_code) > 0
        sliceex = :( TensorCast.sliceview($B, $i_code) )
    else
        sliceex = B
    end
    if length(rest) == 0
        push!(store.topex, :(local $Atop = map($f, $sliceex) ) )
    else
        Xsymb = gensym(:X)
        push!(store.topex, :(local $Atop = map($Xsymb -> $f($Xsymb, $(rest...)), $sliceex) ) )
    end

    # construct ex for inputex later
    ij = filter(!isequal(:(:)), i_colons)
    ex = :( $Atop[$(ij...)][$(kl...)] )

    return Atop, ex
end

#="""
    inputex(:A, :( A[i,j][k] ), target, flags, store, icheck, where)

Figures out all the steps needed to transform the given tensor to a boring one,
aligned with the given target, and returns the necessary expression.
Writes sizes which can be read from `A` into `store`, and necessary reshapes in terms of `sz_i`.

Now `target` need not be `== canon`, since `sz_i` is independent of that.

Now needs explicitly the name of tensor `A`,
so that when `inex = rand(2,3)[i,j]` this is not evaluated twice.
"""=#
function inputex(A, inex, target, flags, store, icheck, where)

    if @capture(inex, Aa_[outer__][inner__])
        glue = :yes
    elseif @capture(inex, Aa_[outer__]{inner__})
        glue = :static
    elseif @capture(inex, Aa_[outer__])
        inner = []
        glue = :no
    else error("inputex should not have been called")
    end

    flatE, getB, _, negF = parse!(store, A, outer, inner)

    append!(store.seen, flatE)

    ex = A
    if icheck
        ex = check_one(:($A[$(outer...)]), where)           # @check!
    end

    numB = count(!isequal(:), getB)
    if numB > 0                                             # A[_,i]
        if numB == length(getB) || :nolazy in flags
            needview!(getB) # replace _ with 1 if any
            ex = :( $ex[$(getB...)] )
        elseif needview!(getB) # then at least one index fixed and not _
            ex = :(view($ex, $(getB...) ))
        else
            ex = :(TensorCast.rview($ex, $(getB...) )) # really a reshape
        end
    end

    sizeC = Any[ Symbol(:sz_, i) for i in flatE[length(inner)+1 : end] ]

    if (length(sizeC) + numB) != length(outer)              # A[i\j]
        sizeCex = :(($(sizeC...) ,))
        ex = :( reshape($ex, $sizeCex) )
        push!(flags, :needsize)
    end

    if length(flatE)-length(inner) == 2 && flatE[end] == flatE[end-1]  # A[i,i] special case
        if :nolazy in flags
            ex = :( TensorCast.diag($ex) ) # LinearAlgebra might not be loaded by caller
            push!(flags, :havecopied)
        else
            ex = :( TensorCast.diagview($ex) )
        end
        pop!(flatE) # remove repeated index from the end
    end

    checkrepeats(flatE, " in term $inex", where) # after dealing with diag story

    dirs = [ findcheck(i, target, where) for i in flatE ]

    if glue == :yes                                         # A[i][k]
        codeD = repeat(Any[*],length(flatE))
        codeD[1:length(inner)] .= (:)

        if codeD == [:,*] && dirs[1] > dirs[2] # then we can avoid a transpose
            codeD = [*,:]
            dirs = [dirs[2], dirs[1]]
        end
        # you could perform more elaborate versions of that, e.g. for this:
        # @pretty @reduce A[i\j,_] = sum(k) B[i,j][k]
        # however only copy_glue and julienne_glue understand arbitrary codes

        if :lazy in flags || :mustview in flags
            ex = :( TensorCast.lazy_glue($ex, $(codeD...,))  )

        elseif :cat in flags
            ex = :( TensorCast.cat_glue($ex, $(codeD...,))  )
            push!(flags, :havecopied)
        elseif :glue in flags
            ex = :( TensorCast.copy_glue($ex, $(codeD...,))  )
            push!(flags, :havecopied)
        elseif :julienne in flags
            ex = :( TensorCast.julienne_glue($ex, $(codeD...,))  )
            push!(flags, :havecopied)
        else
            ex = :( TensorCast.red_glue($ex, $(codeD...,))  )
            push!(flags, :havecopied)
        end

    elseif glue == :static                                  # A[i]{k}
        ex = :( TensorCast.static_glue($ex)  )
        push!(flags, :staticglue) # for packagecheck
    end

    perm = ntuple(identity, length(dirs))
    if dirs != sort(dirs)                                   # A[j,i]
        perm = Tuple(sortperm(dirs))
        if :nolazy in flags
            ex = :( permutedims($ex, $perm) )
            push!(flags, :havecopied)
        elseif :strided in flags
            ex = :( Strided.@strided permutedims($ex, $perm) )
        elseif perm == (2,1)
            ex = :( TensorCast.PermuteDims($ex) ) # avoids transpose, except on numbers, because it's recursive
        else
            ex = :( PermutedDimsArray($ex, $perm) )
        end
        dirs = sort(dirs) # avoids unnecessary orient(PermuteDims(M), (:,:,*))
    end

    for i in negF                                           # A[-i,j]
        d = invperm(perm)[findcheck(i, flatE, where)]
        if :nolazy in flags
            ex = :( reverse($ex, dims=$d) )
            push!(flags, :havecopied)
        else
            ex = :( TensorCast.Reverse{$d}($ex) )
        end
    end

    if length(flatE) != length(target) && dirs != 1:length(dirs) # A[i] + B[j]
        codeH = repeat(Any[*],length(target))
        codeH[dirs] .= (:)
        while codeH[end] == *
            pop!(codeH)
        end
        ex = :( TensorCast.orient($ex, $(codeH...,)) )
    end

    return ex
end

#="""
     outputnew(newright, outUZ, redfun, canonsize, canon, flags, store, where)

For the case of `:=`, this constructs the expression to do reduction if needed,
and slicing/reshaping/reversing for LHS.

outUZ = (redUind, negV, codeW, indW, sizeX, getY, numY, indZ, sizeZ)
"""=#
function outputnew(newright, (redUind, negV, codeW, indW, sizeX, getY, numY, indZ, sizeZ),
        redfun, canonsize, canon, flags, store, where)

    ex = newright

    for ri in negV                                          # Z[-i,j]
        d = findcheck(ri, canon, where)
        ex = :( reverse($ex, dims=$d) )
    end

    if :reduce in flags && :scalar in flags                 # Z = @reduce sum(i) ...
        ex = :( $redfun($ex) )
    elseif :reduce in flags                                 # Z[i] := sum(j) ...
        rdims = Tuple([findcheck(i, canon, where) for i in redUind])
        if length(rdims)==1
            rdims = first(rdims)
        end
        ex = :( $redfun($ex, dims=$rdims) )
        if :outshape in flags && !(:slice in flags || :staticslice in flags)
            # then we need not dropdims, as reshape will handle it
        else
            ex = :( dropdims($ex, dims=$rdims) )
        end
    end

    if :slice in flags                                      # Z[i][k] :=
        if :mustcopy in flags # && !(:havecopied in flags) # make |= slightly stronger
            ex = :( TensorCast.slicecopy($ex, $(codeW...,)) )
            push!(flags, :havecopied)
        elseif :julienne in flags
            ex = :( TensorCast.julienne_slice($ex, $(codeW...,)) )
        else
            ex = :( TensorCast.sliceview($ex, $(codeW...,)) )
        end
    elseif :staticslice in flags                            # Z[i]{k} :=
        # codeW worked out already, but sizeWstatic must be done here
        # TODO sizeinfer!(store, canon, where, false)
        sizeWstatic = :( StaticArrays.Size($([store.dict[i] for i in indW]...)) )
        if :outshape in flags
            ex = :( TensorCast.static_slice($ex, $sizeWstatic, false) )
        else
            ex = :( TensorCast.static_slice($ex, $sizeWstatic) )
        end
    end

    if :outshape in flags                                   # Z[i\j, _, k]
        sizeZex = :(($(sizeZ...) ,))
        ex = :( reshape($ex, $sizeZex) )
        push!(flags, :needsize)
        for n in filter(!isequal(:), getY)
            n == 1 || n == :_ || throw(MacroError("can't fix output index to $n != 1, when creating a new array", where))
        end
    end

    if :mustcopy in flags                                   # Z[i] |= ...
        if !(:havecopied in flags) && !(:broadcast in flags)
            ex = :( collect($ex) )
        elseif (:slice in flags) && (:strided in flags)     # Z[i][j] |= ... strided
            ex = :( collect.($ex) ) # because slicecopy actually makes views here
        elseif (:staticslice in flags) || (:strided in flags) # i.e. then copy even if there was a broadcast etc.
            ex = :( collect($ex) )
        end

    elseif :mustview in flags                               # Z[i] == ...
        if :havecopied in flags
            m_error("can't do what you ask without copying, sorry", where)
        elseif :broadcast in flags
            m_error("can't broadcast without copying, sorry", where)
        end
    end

    if :diagleft in flags                                   # Z[i,i] := ...
        if :nolazy in flags
            ex = :( TensorCast.diagm(0 => $ex) )
        else
            ex = :( TensorCast.Diagonal($ex) )
        end
    end

    if :namedarray in flags                                 # Z[i] := ... named
        ex = :( TensorCast.namedarray($ex, $(indZ...,) ) )
    elseif :nameddims in flags || :named in flags
        ex = :( NamedDims.NamedDimsArray{$(indZ...,)}($ex) )
    # elseif :axis in flags
    #     ex = :( TensorCast.axisarray($ex, $(indZ...,) ) )
    end

    return ex
end

#="""
    outputinplace(newright, outUZ, redfun, canonsize, canon, flags, store, nameZ, where)

For the case of `=` this figures out how to write RHS into LHS, in one of three ways:
* reduction `sum!(Z, newright)`
* broadcasting `@. Z[...] = newright`
* neither, `copyto!(Z, newright)`

No longer attempts to write `permutedims!(Z, A, ...)`, now just `copyto!(Z, PermutedDimsArray(A, ...))`.
Doesn't really need so many arguments...
"""=#
function outputinplace(newright, (redUind, negV, codeW, indW, sizeX, getY, numY, indZ, sizeZ),
        redfun, canonsize, canon, flags, store, nameZ, where)

    if :slice in flags
        throw(MacroError("can't write to sliced arrays in-place, for now", where))
    end
    if length(negV) > 0
        m_error("can't reverse axes of in-place output, try moving -$(negV[1]) to right hand side", where)
    end

    if :reduce in flags                                     # sum!(revleft, newright)

        if :broadcast in flags
            if :lazy in flags
                newright = :( Base.@__dot__ TensorCast.lazy($newright) ) # really LazyArrays.lazy
            elseif :strided in flags
                dotted = Broadcast.__dot__(newright) # @strided does not work on @.
                newright = :( Strided.@strided $dotted )
            else
                newright = :( Base.@__dot__ $newright ) # prettier in @pretty
            end
        end

        # working backwards
        revleft = nameZ

        if numY > 0
            if needview!(getY)
                revleft = :(view($revleft, $(getY...) ))
            else
                revleft = :(TensorCast.rview($revleft, $(getY...) )) # really a reshape
            end
            push!(flags, :finalres)
        end

        if :diagleft in flags
            revleft = :( TensorCast.diagview($revleft) )
            push!(flags, :finalres)
        end

        if :backshape in flags
            sizeXex = :(($(sizeX...) ,))
            revleft = :( reshape($revleft, $sizeXex) )
            push!(flags, :needsize)
            push!(flags, :finalres)
        end

        if !endswith(string(redfun), '!')
            redfun = Symbol(redfun, '!')
        end

        ex = :( $redfun($revleft, $newright) )

    elseif :broadcast in flags                              # @. revleft[...] = newright

        # working backwards
        revleft = nameZ

        if numY > 0 # when getY has only : and 1, and backshape, then you could skip this
            if needview!(getY)
                revleft = :(view($revleft, $(getY...) ))
            else
                revleft = :(TensorCast.rview($revleft, $(getY...) )) # really a reshape
            end
            push!(flags, :finalres)
        end

        if :diagleft in flags
            revleft = :( TensorCast.diagview($revleft) )
            push!(flags, :finalres)
        end

        if :backshape in flags
            sizeXex = :(($(sizeX...) ,))
            revleft = :( reshape($revleft, $sizeXex) )
            push!(flags, :needsize)
            push!(flags, :finalres)
        end

        if :strided in flags
            dotted = Broadcast.__dot__(newright)
            ex = :( Strided.@strided $revleft .= $dotted )
        else
            ex = :( $revleft .= Base.@__dot__ $newright ) # don't apply @. to reshape(diagview(Z))
        end

    else                                                    # copyto!(revleft, newright)

        if :strided in flags
            newright = :( Strided.@strided $newright )
        end

        # working backwards
        revleft = nameZ

        if numY > 0
            if needview!(getY)
                revleft = :(view($revleft, $(getY...) ))
            else
                revleft = :(TensorCast.rview($revleft, $(getY...) )) # really a reshape
            end
            push!(flags, :finalres)
        end

        if :diagleft in flags
            revleft = :( TensorCast.diagview($revleft) )
            push!(flags, :finalres)
        end

        ex = :( $copyto!($revleft, $newright) )

    end

    return ex
end

#="""
    anoninput(store.rightnames)

Given a vector of names captured from `A_[i...]`, returns ex needed for `(A,B) -> ...`
"""=#
function anoninput(rightnames, where)
    for A in rightnames
        isa(A,Symbol) || throw(MacroError("can't use $A as anonymous function input", where))
        # TODO make this understand that $B means interpolate not arg
    end
    if length(rightnames) == 1
        return rightnames[1]
    else
        return :( ($(rightnames...),) )
    end
end

using LinearAlgebra  # for diag()

function packagecheck(flags, where)
    where === nothing && return
    # now check in caller's scope?
    if :staticslice in flags || :staticglue in flags
        isdefined(where.mod, :StaticArrays) || m_error("can't use static arrays without using StaticArrays", where)
    end
    if :strided in flags
        isdefined(where.mod, :Strided) || m_error("can't use option strided without using Strided", where)
    end
    if :julienne in flags
        isdefined(where.mod, :JuliennedArrays) || m_error("can't use option julienne without using JuliennedArrays", where)
    end
    if :namedarray in flags
        isdefined(where.mod, :NamedArrays) || m_error("can't use option namedarray without using NamedArrays", where)
    end
    if :named in flags || :nameddims in flags
        isdefined(where.mod, :NamedDims) || m_error("can't use option nameddims without using NamedDims", where)
    end
    # if :axis in flags
    #     isdefined(where.mod, :AxisArrays) || m_error("can't use option axis without using AxisArrays", where)
    # end
end


# open up a block to limit scope and lifetime of everything to
# ensure no cross-talk between tests
begin

    N = 500

    A = rand(N)
    B = rand(N)
    C = rand(N)

    # addition
    @et add_et = A + B + C
    add_native = A + B + C

    @test add_et == add_native;

    #subtraction
    @et sub_et = A - B - C
    sub_native = A - B - C

    @test sub_et == sub_native;

    # elementwise mult
    @et dottimes_et = A .* B .* C
    dottimes_native = A .* B .* C

    @test dottimes_et == dottimes_native;

    # elementwise div
    @et dotdiv_et = A ./ B ./ C
    dotdiv_native = A ./ B ./ C

    @test dotdiv_et == dotdiv_native;

end

begin

    N = 50

    A = rand(N,N)
    B = rand(N,N)
    C = rand(N,N)

    #test order-2 arrays
    @et res_et = A+B+C
    res_native = A+B+C

    @test res_et == res_native

end

begin

    N = 15

    A = rand(N,N,N)
    B = rand(N,N,N)
    C = rand(N,N,N)

    #test order-3 arrays
    @et res_et = A+B+C
    res_native = A+B+C

    @test res_et == res_native

    # test .* and ./
    @et res_et = A.*B./C
    res_native = A.*B./C

    @test res_et == res_native
end


# This file reproduces figures 2-3 of
# Identification of LTV Dynamical Models with Smooth or Discontinuous Time Evolution by means of Convex Optimization
# Fredrik Bagge Carlson, Anders Robertsson, Rolf Johansson
# https://arxiv.org/abs/1802.09794
# The script requires DynamicMovementPrimitives.jl to be installed

using OrdinaryDiffEq
using LTVModels

using ValueHistories, IterTools, MLDataUtils, OrdinaryDiffEq, Parameters, InteractNext

# Import robot simulator from DynamicMovementPrimitives
include(Pkg.dir("DynamicMovementPrimitives","src","two_link.jl"))
using TwoLink

@with_kw struct TwoLinkSys
    N  = 1000
    n  = 2
    ns = 2n
    h = 0.02
    σ0 = 0
    sind = 1:ns
    uind = ns+1:(ns+n)
    s1ind = (ns+n+1):(ns+n+ns)
end

# Problem parameters and trajectory
n,m,N,h,σ0 = 4,3,200,0.02,0.0
np         = n*(n+m)
sys        = TwoLinkSys(N=N, h=h, σ0 = σ0)#; warn("Added noise")
seed       = 1

ni         = N+2
traj       = connect_points([0 -0.5; 0.9 -0.1],ni)
q          = inverse_kin(traj[1],:up)
p          = forward_kin(q)
qd         = centraldiff(q')'/sys.h
qdd        = centraldiff(qd')'/sys.h
u0         = torque(q,qd,qdd) # Computed torque
u0       .+= 0.1randn(size(u0)) # Random noise on input

# Define callback to simulate the robot hitting a wall
const wall = -0.2
function affect!(i)
    q = i.u[1:2]
    qd = i.u[3:4]
    p1,p2 = forward_kin(q)[2]
    J = TwoLink.jacobian(q)
    v1,v2 = J*qd
    d = [0, wall-p2]
    δq = J\d
    i.u[3:4] = J\[v1, 0]
    i.u[1:2] .+= δq
end
condition(x,t,i) = forward_kin(x)[2][2] > wall
cb = DiscreteCallback(condition,affect!)

# Simulate robot using ODE-solver
t      = 0:h:N*h
q0     = [q[:,1]; qd[:,1]]
prob   = OrdinaryDiffEq.ODEProblem((x,p,t)->time_derivative(x, u0[:,floor(Int,t/h)+1]),q0,(t[[1,end]]...,))
sol    = solve(prob,Tsit5(),reltol=1e-8,abstol=1e-8, callback=cb)
x      = hcat(sol(t)...)
u      = u0[:,1:N]
u      = [u; ones(1,size(u,2))]
x,y    = x[:,1:N], x[:,2:N+1]

plot(forward_kin(x[1:2,:])', layout=2, subplot=1, title="End-effector positions")
plot!([1,N],[wall, wall], l=(:dash, :black), grid=false, subplot=1, lab="Wall")
plot!(u', subplot=2, title="Control torques")

# Normalize input data to mean 0 and unit variance
sx = std(x,dims=2)
su = std(u[1:2,:])
mx, mu = mean(x,dims=2), mean(u[1:2,:],dims=2)
xt = (x.-mx) ./ sx
ut = copy(u)
ut[1:2,:] = (ut[1:2,:].-mu) ./ su

# Fit model using ADMM
model = LTVModels.fit_statespace_admm(xt,ut,10^-0.5, extend=false, tol=1e-8, iters=20000, D=1, zeroinit=false)

signch = findfirst(u[2,:] .< 0) # Detect when control signal changes sign
contact = findfirst(forward_kin(x[1:2,:])[2,:] .> wall) # Detect first contact with wall

fig=plot(xt', label="States", show=false, layout=grid(3,2), subplot=1:4, size=(1800,1000), title="State trajectories", title=["\$q_1\$" "\$q_2\$" "\$\\dot{q}_1\$" "\$\\dot{q}_2\$"], c=:blue, linewidth=2)
vline!([signch].*ones(1,4), show=false, subplot=1:4, lab="Velocity sign change", l=(:black, :dash))
vline!([contact].*ones(1,4), show=false, subplot=1:4, lab="Stiff contact", l=(:black, :dash))
plot!(forward_kin(x[1:2,:])', subplot=5, title="End-effector positions", lab=["x" "y"], c=[:blue :red])
hline!([wall], l=(:dash, :black), grid=false, subplot=5, lab="Constraint")
vline!([contact], subplot=5, lab="Stiff contact", l=(:black, :dash))
plot!(ut[1:2,:]', subplot=6, title="Control torques", lab="", c=[:blue :red])
vline!([signch], subplot=6, lab="Velocity sign change", l=(:black, :dash))
plot!(LTVModels.simulate(model, xt[:,1], ut)', label="Simulation", show=false, subplot=1:4, l=(:red,), legend=[false false :left false])
gui()

# Validation data. Simulate system again with different input trajectory
uv   = torque(q,qd,qdd)
uv .+= 0.1randn(size(uv))
prob   = OrdinaryDiffEq.ODEProblem((x,p,t)->time_derivative(x, uv[:,floor(Int,t/h)+1]),q0,(t[[1,end]]...,))
sol   = solve(prob,Tsit5(),reltol = 1e-8,abstol = 1e-8, callback = cb)
xv    = hcat(sol(t)...)
uv    = uv[:,1:N]
uv    = [uv; ones(1,size(uv,2))] # Extend to estimate friction
xv,yv = xv[:,1:N], xv[:,2:N+1]

# Normalize data with same parameters as during training
xt2        = (xv.-mx) ./ sx
ut2        = copy(uv)
ut2[1:2,:] = (uv[1:2,:].-mu) ./ su

plot(xt2', label="States", show=false, layout=grid(2,2), subplot=1:4, size=(1800,1000), title="State trajectories", title=["\$q_1\$" "\$q_2\$" "\$\\dot{q}_1\$" "\$\\dot{q}_2\$"], c=:blue, linewidth=2, reuse=false, grid=false)
vline!([signch].*ones(1,4), show=false, subplot=1:4, lab="Original Velocity sign change", l=(:black, :dash))
vline!([contact].*ones(1,4), show=false, subplot=1:4, lab="Original Stiff contact", l=(:black, :dash))

plot!(LTVModels.simulate(model, xt2[:,1], ut2)', label="Simulation", show=false, subplot=1:4, l=(:red,), legend=[false false :left false])
gui()


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.6.2
#     language: julia
#     name: julia-1.6
# ---

# %%
?@generated

# %%
function _staticsum(f, a, b)
    expr = :(f($a))
    for i in a+1:b
        expr = :($expr + f($i))
    end
    expr
end

# %%
_staticsum(:f, -3, 4)

# %%
@generated function staticsum(f, ::Val{a}, ::Val{b}) where {a, b}
    @assert b - a < 100
    expr = :(f($a))
    for i in a+1:b
        expr = :($expr + f($i))
    end
    expr
end

# %%
staticsum(identity, Val(1), Val(10))

# %%
2staticsum(x -> sin(x)/x, Val(1), Val(100)) + 1

# %%
F(n) = 2sum(x -> sin(x)/x, 1:n) + 1
F_static(n) = 2staticsum(x -> sin(x)/x, Val(1), Val(n)) + 1

# %%
using BenchmarkTools
@show F(100) == F_static(100)
@btime F(100)
@btime F_static(100)

# %%
@benchmark F(100)

# %%
@benchmark F_static(100)

# %%
@code_warntype staticsum(x -> sin(x)/x, Val(1), Val(5))

# %%
@code_typed staticsum(x -> sin(x)/x, Val(1), Val(5))

# %%
@code_typed sum(x -> sin(x)/x, 1:5)

# %%


using PyPlot
using PyCall
using MAT
using LinearAlgebra
using Statistics
mpl = pyimport("tikzplotlib")

n = 10
m = 2n
h = 0.01
out = 0.5 * ones(n)
out[1:div(n,3)] .= 1.0

out *= 50
rg = 0:20:200
err = zeros(4, length(rg))
for (j,idx) in enumerate([1])
    for (i,k) in enumerate(rg)
        eta = matread("$idx/eta$k.mat")["eta"]
        eta = eta[1:4m:end]
        err[j, i] = mean( ((out-eta) ./ out).^2)
        @info j, i, err[j,i]
    end
end

close("all")
# tp = ["ro-", "g+-", "b*-", "k>-"]
# labels = []
# for k in [1,2,3,5]
# s = length(1:k:m+1)
# push!(labels, "$s Sensors")
# end
# for i = 1:4
#     plot(rg .+ 1, err[i,:], tp[i], label=labels[i])
# end
semilogy(rg .+ 1, err[1,:], tp[1])
legend()
xlabel("Iterations")
ylabel("Error")
mpl.save("space_loss.tex")
savefig("loss.png")

function visualize_inv_eta(X)
    ηmin = 0.5*50.
    ηmax = 50.
    x = LinRange(0.5h,m*h, m)
    y = LinRange(0.5h,n*h, n)
    V = zeros(m, n)
    for i = 1:m
        for j = 1:n
            elem = (j-1)*m + i
            V[i, j] = mean(X[4(elem-1)+1:4elem])
        end
    end
    close("all")
    pcolormesh(x, y, V', vmin=ηmin-(ηmax-ηmin)/4, vmax=ηmax+(ηmax-ηmin)/4, rasterized=true)
    colorbar(shrink=0.5)
    xlabel("x")
    ylabel("y")
    # title("Iteration = $k")
    axis("scaled")
    gca().invert_yaxis()
end

out = repeat(out, 1, 4m)'[:]
visualize_inv_eta(out)
savefig("space_true.pdf")

eta  = matread("1/eta0.mat")["eta"]
visualize_inv_eta(eta)
savefig("space0.pdf")

eta  = matread("1/eta60.mat")["eta"]
visualize_inv_eta(eta)
savefig("space60.pdf")

eta  = matread("1/eta120.mat")["eta"]
visualize_inv_eta(eta)
savefig("space120.pdf")


# Copyright (c) 2021, Rolfe Power

"""
    GRAVITATIONAL_CONSTANT

Universal constant of gravitation. [km^3 / (kg . s^2)]

Source: [NIST](https://physics.nist.gov/cgi-bin/cuu/Value?bg|search_for=universal_in!)
"""
const GRAVITATIONAL_CONSTANT = 6.6743e-20

"""
    ASTRONOMICAL_UNIT

Astronomical unit as defined by the IAU. [km]

Source: [IAU](https://www.iau.org/static/resolutions/IAU2012_English.pdf)
"""
const ASTRONOMICAL_UNIT = 149_597_870.700

"""
    EARTH_STANDARD_GRAVITY

Standard gravitational acceleration at surface of Earth. [km/s^2]

Source: [NIST](https://physics.nist.gov/cgi-bin/cuu/Value?gn|search_for=adopted_in!)
"""
const EARTH_STANDARD_GRAVITY = 9.80665e-03

"""
    PROTON_MASS

Mass of the proton. [kg]

Source: [NIST](https://physics.nist.gov/cgi-bin/cuu/Value?mp|search_for=atomnuc!)
"""
const PROTON_MASS = 1.67262192369e-27

"""
    NEUTRON_MASS

Mass of the neutron. [kg]

Source: [NIST](https://physics.nist.gov/cuu/Constants/Table/allascii.txt)
"""
const NEUTRON_MASS = 1.67492749804e-27

"""
    ELECTRON_MASS

Mass of the electron. [kg]

Source: [NIST](https://physics.nist.gov/cgi-bin/cuu/Value?mp|search_for=atomnuc!)
"""
const ELECTRON_MASS = 9.1093837015e-31

"""
    SPEED_OF_LIGHT

Speed of light in a vacuum. [km/s]

Source: [NIST](https://physics.nist.gov/cgi-bin/cuu/Value?c|search_for=adopted_in!)
"""
const SPEED_OF_LIGHT = 299792458e-3

"""
    BOLTZMANN_CONSTANT

Boltzmann constant. [J/K]
"""
const BOLTZMANN_CONSTANT = 1.380649e-23


@testset "66.plus-one.jl" begin
    rvdigits = i -> reverse!(digits(i))
    @test all(digits_plus_one(rvdigits(i)) == rvdigits(i + 1) for i in 0:1000)
    @test digits_plus_one([9, 9, 9, 9]) == [1, 0, 0, 0, 0]
    @test all(plus_one(rvdigits(i)) == rvdigits(i + 1) for i in 0:1000)
    @test plus_one([9, 9, 9, 9]) == [1, 0, 0, 0, 0]
end


###### mean rates and other measures

function numerical_rates(ps::PopulationState;
    Tstart::Real=0.0,Tend::Real=Inf)
  return numerical_rate.(ps.trains_history;Tstart=Tstart,Tend=Tend)
end
function numerical_rates(pop::AbstractPopulation;
    Tstart::Real=0.0,Tend::Real=Inf)
  return numerical_rates(pop.state;Tstart=Tstart,Tend=Tend)
end

function numerical_rate(train::Vector{Float64};
    Tstart::Real=0.0,Tend::Real=Inf)
  isempty(train) && return 0.0
  Tend = min(Tend,train[end])
  Δt = Tend - Tstart
  return length(train)/Δt
end

##########################
## covariance density

function bin_spikes(Y::Vector{R},dt::R,Tend::R;
    Tstart::R=0.0) where R
  times = range(Tstart,Tend;step=dt)
  ret = fill(0,length(times)-1)
  for y in Y
    if Tstart < y <= last(times)
      k = searchsortedfirst(times,y)-1
      ret[k] += 1
    end
  end
  return ret
end

# and returns a value in Hz
function instantaneous_rates(idxs_neu::AbstractVector{<:Integer},
    dt::Float64,pop::AbstractPopulation;Tend::Float64=-1.0,Tstart::Float64=0.0)
  trains = pop.state.trains_history[idxs_neu]
  return  instantaneous_rates(dt,trains;
   Tend=Tend,Tstart=Tstart)
end
function instantaneous_rates(dt::Float64,trains::Vector{Vector{Float64}};
  Tend::Float64=-1.0,Tstart::Float64=0.0)
  if Tend < 0
    Tend = minimum(last.(trains))
  end
  tmid = midpoints(range(Tstart,Tend;step=dt))
  counts = fill(0.0,length(tmid))
  for train in trains
    counts .+= bin_spikes(train,dt,Tend;Tstart=Tstart)
  end
  return tmid,counts ./ (dt*length(trains))
end

# time starts at 0, ends at T-dt, there are T/dt steps in total
@inline function get_times(dt::Real,T::Real)
  return (0.0:dt:(T-dt))
end

# frequencies for Fourier transform.
# from -1/dt to 1/dt - 1/T in steps of 1/T
function get_frequencies_centerzero(dt::Real,T::Real)
  dω = inv(T)
  ωmax = 0.5/dt
  f = dω:dω:ωmax
  ret = vcat(-reverse(f),0.,f[1:end-1])
  return ret
end

function get_frequencies(dt::Real,T::Real)
  dω = inv(T)
  ωmax = inv(dt)
  ret = 0:dω:ωmax-dω
  return ret
end

function covariance_self_numerical(Y::Vector{R},dτ::R,τmax::R,
     Tmax::Union{R,Nothing}=nothing) where R
  ret = covariance_density_numerical([Y,],dτ,τmax,Tmax;verbose=false)
  return ret[1,1,:]
end

function covariance_density_numerical(Ys::Vector{Vector{R}},dτ::Real,τmax::R,
   Tmax::Union{R,Nothing}=nothing ; verbose::Bool=false) where R
  Tmax = something(Tmax, maximum(x->x[end],Ys)- dτ)
  ndt = round(Integer,τmax/dτ)
  n = length(Ys)
  ret = Array{Float64}(undef,n,n,ndt)
  if verbose
      @info "The full dynamical iteration has $(round(Integer,Tmax/dτ)) bins ! (too many?)"
  end
  for i in 1:n
    binnedi = bin_spikes(Ys[i],dτ,Tmax)
    fmi = length(Ys[i]) / Tmax # mean frequency
    ndt_tot = length(binnedi)
    _ret_alloc = Vector{R}(undef,ndt)
    for j in 1:n
      if verbose
        @info "now computing cov for pair $i,$j"
      end
      binnedj =  i==j ? binnedi : bin_spikes(Ys[j],dτ,Tmax)
      fmj = length(Ys[j]) / Tmax # mean frequency
      binnedj_sh = similar(binnedj)
      @inbounds @simd for k in 0:ndt-1
        circshift!(binnedj_sh,binnedj,k)
        _ret_alloc[k+1] = dot(binnedi,binnedj_sh)
      end
      @. _ret_alloc = _ret_alloc / (ndt_tot*dτ^2) - fmi*fmj
      ret[i,j,:] = _ret_alloc
    end
  end
  return ret
end

"""
     draw_spike_raster(trains::Vector{Vector{Float64}},
      dt::Real,Tend::Real;
      Tstart::Real=0.0,
      spike_size::Integer = 5,
      spike_separator::Integer = 1,
      background_color::Color=RGB(1.,1.,1.),
      spike_colors::Union{C,Vector{C}}=RGB(0.,0.0,0.0),
      max_size::Real=1E4) where C<:Color

Draws a matrix that contains the raster plot of the spike train.

# Arguments
+ `Trains` :  Vector of spike trains. The order of the vector corresponds to
the order of the plot. First element is at the top, second is second row, etc.
+ `dt` : time interval representing one horizontal pixel
+ `Tend` : final time to be considered

# Optional arguments
+ `Tstart` : starting time
+ `max_size` : throws an error if image is larger than this number (in pixels)
+ `spike_size` : heigh of spike (in pixels)
+ `spike_separator` : space between spikes, and vertical padding
+ `background_color` : self-explanatory
+ `spike_colors` : if a single color, color of all spikes, if vector of colors,
color for each neuron (length should be same as number of neurons)

# returns
`raster_matrix::Matrix{Color}` you can save it as a png file

"""
function draw_spike_raster(trains::Vector{Vector{Float64}},
  dt::Real,Tend::Real;
    Tstart::Real=0.0,
    spike_size::Integer = 5,
    spike_separator::Integer = 1,
    background_color::Color=RGB(1.,1.,1.),
    spike_colors::Union{C,Vector{C}}=RGB(0.,0.0,0.0),
    max_size::Real=1E4) where C<:Color
  nneus = length(trains)
  if typeof(spike_colors) <: Color
    spike_colors = repeat([spike_colors,];outer=nneus)
  else
    @assert length(spike_colors) == nneus "error in setting colors"
  end
  binned_binary  = map(trains) do train
    .! iszero.(bin_spikes(train,dt,Tend;Tstart=Tstart))
  end
  ntimes = length(binned_binary[1])
  ret = fill(background_color,
    (nneus*spike_size + # spike sizes
      spike_separator*nneus + # spike separators (incl. top/bottom padding)
      spike_separator),ntimes)
  @assert all(size(ret) .< max_size ) "The image is too big! Please change time limits"
  for (neu,binv,col) in zip(1:nneus,binned_binary,spike_colors)
    spk_idx = findall(binv)
    _idx_pre = (neu-1)*(spike_size+spike_separator)+spike_separator
    y_color = _idx_pre+1:_idx_pre+spike_size
    ret[y_color,spk_idx] .= col
  end
  return ret
end

#### this is the part which involves reader objects

function get_spiketimes_spikeneurons(rec::RecFullTrain,
    (pop_idx::Integer)=1)
  spiketimes,spikeneurons = rec.timesneurons[pop_idx]
  spkt = filter(isfinite,spiketimes)
  spkn = filter(>(0),spikeneurons)
  return spkt,spkn
end
function spiketn_to_trains(spiketimes::Vector{<:Real},
    spikeneurons::Vector{<:Integer},Nneus::Integer)
  trains = map(_-> Float64[],1:Nneus)
  for (spkt,spkn) in zip(spiketimes,spikeneurons)
    push!(trains[spkn],spkt)
  end
  return trains
end
function get_trains(rec::RecFullTrain,Nneus::Integer,
    (pop_idx::Integer)=1)
  spiketn = get_spiketimes_spikeneurons(rec,pop_idx)
  return spiketn_to_trains(spiketn...,Nneus)
end

function numerical_rates(rec::RecFullTrain,Nneus::Integer,Tend::Real,
  (pop_idx::Integer)=1;Tstart=0.0)
  trains = get_trains(rec,Nneus,pop_idx)
  return numerical_rate.(trains;Tstart=Tstart,Tend=Tend)
end


module testExamples
using LinearAlgebra
using Test

srcDir = dirname(dirname(@__FILE__))*"\\src\\Examples\\"
include(srcDir *"Examples_sep.jl")
include(srcDir *"Examples_ent.jl")
include(srcDir *"Examples.jl")

using .Examples_sep
using .Examples_ent
using .Examples

@testset "Test the call examples" begin
    ρ_ex = Examples.get_example()
    @test ρ_ex["RANDℂb5"] == Examples.get_example("RANDℂb5")
    @test ρ_ex["RANDℂb5"] == Examples.get_example(5)
end

@testset "Test the call example meta data" begin
    ρ_ex = Examples.get_example()
    df = Examples.get_example_meta()
    @test df[df.Example .== "Eq14Hor97a",:] ==  Examples.get_example_meta("Eq14Hor97a")
    @test df[df.Example .== "Eq14Hor97a",:].Example[1] ==  Examples.get_example_meta(1).Example
    length(ρ_ex) == size(df)[1]
end

##
df = ex.get_example_meta()
entdf = df[df.isSeparable .== "ent",:]
ρ_dict = ex.get_example()

for n_ex in 1:23
    exa = entdf.Example[n_ex]
    ρ = ρ_dict[exa]
    d = entdf.Size[n_ex]
    println("$(entdf.Example[n_ex]) $(la.eigmin(ρ))")
    @assert us.isPPᵀ(ρ,d) == entdf.hasPPᵀ[n_ex]
end

##
df = ex.get_example_meta()
sepdf = df[df.isSeparable .== "sep",:]
ρ_dict = ex.get_example()

for n_ex in 1:36
    exa = sepdf.Example[n_ex]
    ρ = ρ_dict[exa]
    d = sepdf.Size[n_ex]
    println("$(sepdf.Example[n_ex]) $(la.eigvals(ρ)[1])")
    @assert us.isPPᵀ(ρ,d) == sepdf.hasPPᵀ[n_ex]
end

end


using OrdinaryDiffEq, Test
using DiffEqProblemLibrary.ODEProblemLibrary: importodeproblems; importodeproblems()
import DiffEqProblemLibrary.ODEProblemLibrary: prob_ode_2Dlinear
prob = prob_ode_2Dlinear

integrator = init(prob,Euler();dt=1//2^(4))
integrator = init(prob,Tsit5();dt=1//2^(4))

solve!(integrator)

sol =solve(prob,Euler();dt=1//2^(4),maxiters=Inf)
sol =solve(prob,Tsit5();dt=1//2^(20),progress=true,adaptive=false,maxiters=Inf)

# plot(sol,plot_analytic=true)

sol =solve(prob,ExplicitRK();dt=1//2^(4))

# out-of-place 2D prob to test vec & reshape
rhs_oop(u,p,t) = u
ft0=ones(2, 2)
tspan=(0.0,30.0)
prob=ODEProblem(rhs_oop,ft0,tspan)
@test_nowarn solve(prob,Rosenbrock23())
@test_nowarn solve(prob,Rodas4())
@test_nowarn solve(prob,Kvaerno5())


function basis(sig::Signature, prefix::Symbol, export_symbols::Bool, export_metadata::Bool, export_pseudoscalar::Bool, modname::Symbol)
    prefix isa Symbol || throw(ArgumentError("Only symbols are supported for the second argument (received $prefix)"))

    n = dimension(sig)
    n > 0 || throw(ArgumentError("Invalid zero-dimensional signature $sig"))

    _combinations = [[], combinations(1:n)...]

    inds = 1:2^n
    blade_symbols = map(x -> Symbol(prefix, join(string.(indices_from_linear_index(_combinations, x)))), inds)
    blades = Blade.(1, inds)
    exprs = map((x, y) -> (:(const $x = $y)), blade_symbols, blades)

    table = reshape([produce_indices(i, j, _combinations, sig) for i ∈ _combinations, j ∈ _combinations], 2^n, 2^n)

    exports = Symbol[]
    export_symbols && append!(exports, blade_symbols)
    export_metadata && append!(exports, [:N, :SIGNATURE, :TABLE, :MultivectorArray])
    export_pseudoscalar && push!(exports, :I)

    mod = :(module $modname
        using GeometricAlgebra: GeometricAlgebra, subscript, reverse_sign, scalar, ⋅, ∧
        using GeometricAlgebra.StaticArrays

        const ga = GeometricAlgebra

        import Base: fill, zero, *, getindex, show, reverse
        import GeometricAlgebra: Blade, KVector, Multivector, grade, grade_projection, geom, op_result_type, kvectors, kvector, dual, ∨, linear_index, indices_from_linear_index

        $(exprs...)

        const N = $n
        const SIGNATURE = $sig
        const TABLE = $table
        const I = $(last(blade_symbols))

        include($("$(@__DIR__)/dimension_dependent.jl"))
        include($("$(@__DIR__)/backends/precomputed.jl"))

        export $(exports...)

    end)

    mod, modname
end

function produce_indices(i, j, combinations, sig::Signature)
    concat_inds = vcat(i, j)
    inds = sort(filter(x -> count(==(x), concat_inds) == 1, concat_inds))
    double_inds = filter(x -> count(==(x), concat_inds) == 2, unique(concat_inds))
    metric_factor = prod(map(ei -> metric(sig, Val(ei)), double_inds))
    linear_index(combinations, inds), metric_factor * permsign(i, j)
end

"""
    @basis <signature> [prefix=v, export_symbols=true, export_metadata=true, export_pseudoscalar=true, modname=GeneratedGA]

Create a module `modname`, fill it with all unit blade symbols from a geometric algebra with a given `signature` prefixed with `prefix`, and import it with `using`.
The exported variables depend on the options `export_symbols` and `export_metadata`.
If `export_symbols` is true, then all unit blade symbols are exported.
If `export_metadata` is true, then the following symbols are exported:
- `N`: the dimension of the algebra
- `TABLE`: the table containing precomputed values for blade products
- `SIGNATURE`: the signature used to build the algebra
- `MultivectorArray`: a concrete multivector array representation, with `2^N` coefficients.
If `export_pseudoscalar` is true, then the alias `I = v₁...ₙ` will be exported.

## Examples

To obtain the unit blades of 𝒢(ℝ³) the geometric algebra over the 3-dimensional vector space ℝ³, you just have to specify a positive signature with "+++":

```julia
julia> @basis "+++" # v is the default prefix
```

To bind the blades to variables with different prefix than the default v, just add the prefix after the signature:

```julia
julia> @basis "+++" prefix=g # assigned variables will be g, g1, g12...
```
"""
macro basis(sig::AbstractString, opts...)
    try
        mod, modname = basis(Signature(sig), parse_macro_opts(collect(opts))...)
        :($(Expr(:toplevel, esc(mod))); using .$modname)
    catch e
        e isa ArgumentError || rethrow()
        :(throw($(esc(e))))
    end
end

const basis_macro_opts = [
    :prefix => :v,
    :export_symbols => true,
    :export_metadata => true,
    :export_pseudoscalar => true,
    :modname => :GeneratedGA,
]

function parse_macro_opts(opts::AbstractVector)
    check_opt.(opts)
    map(basis_macro_opts) do opt
        index = findfirst(==(opt.first), opt_name.(opts))
        isnothing(index) ? opt.second : last(opts[index].args)
    end
end

function check_opt(opt)
    opt isa Expr && opt.head == :(=) || throw(ArgumentError("Macro options must be `(name=value)` pairs"))
    opt_name(opt) ∈ first.(basis_macro_opts) || throw(ArgumentError("Unknown option '$(opt_name(opt))'. Available options are: $(join(first.(basis_macro_opts), ", "))"))
end

opt_name(opt::Expr) = first(opt.args)


using InfiniteOpt, Ipopt, Plots
using KNITRO

# Parmeters
# Time span
t0 = 0
tf = 1
# Initial values
x1 = 5

# Model
model = InfiniteModel(KNITRO.Optimizer)

# infinite_parameter
@infinite_parameter(model, t ∈ [t0, tf], num_supports = 101)

# variable
@variable(model, x , Infinite(t))

@variable(model, 0 ≤ u ≤ 2, Infinite(t))

# objective
@objective(model, Max, ∫(2*x-3*u, t))

# constraint
@constraint(model, x(0) == x1)
@constraint(model, x_constr, ∂(x, t) == x + u)

# Optimization

print(model)

optimize!(model)

# Visulization
ts = value(t)
u = value(u)
x = value(x)

@show plot(ts, u)
@show plot(ts, x)


using JuMP
using DisjunctiveProgramming

m = Model()
@variable(m, -1<=x<=10)

@constraint(m, con1, x<=3)
@constraint(m, con2, 0<=x)
@constraint(m, con3, x<=9)
@constraint(m, con4, 5<=x)

@disjunction(m,(con1,con2),con3,con4,reformulation=:CHR,name=:y)


using Base.Test
# spec.jl
import WaveletScattering: tune_motherfrequency, default_motherfrequency
# meta.jl
import WaveletScattering: get_γ
# spec1d.jl
import WaveletScattering: Spec1D

# tune_motherfrequency
nfos = [1, 2, 4, 8, 12, 16, 24]
pitchforks = [392, 415, 422, 430, 435, 440, 442, 444, 466]
for nfo in nfos, pitchfork in pitchforks
    tuningfrequency = pitchfork / 44100.0
    spec = Spec1D()
    ξ = tune_motherfrequency(tuningfrequency, spec.class, nfo)
    γs = map(get_γ, spec.ψmetas)
    ωs = ξ * exp2.(-γs / nfo)
    @test any(abs.(ωs - ξ) .< 1e-4)
    max_ξ = default_motherfrequency(spec.class, nfo)
    @test ξ < max_ξ
    @test ξ * exp2(inv(nfo)) > max_ξ
end


@doc raw"""
    LatentFactorMOKernel(g::AbstractVector{<:Kernel}, e::MOKernel, A::AbstractMatrix)

Kernel associated with the semiparametric latent factor model.

# Definition

For inputs ``x, x'`` and output dimensions ``p_x, p_{x'}'``, the kernel is defined as[^STJ]
```math
k\big((x, p_x), (x, p_{x'})\big) = \sum^{Q}_{q=1} A_{p_xq}g_q(x, x')A_{p_{x'}q}
                                   + e\big((x, p_x), (x', p_{x'})\big),
```
where ``g_1, \ldots, g_Q`` are ``Q`` kernels, one for each latent process, ``e`` is a
multi-output kernel for ``m`` outputs, and ``A`` is a matrix of weights for the kernels of
size ``m \times Q``.

[^STJ]: M. Seeger, Y. Teh, & M. I. Jordan (2005). [Semiparametric Latent Factor Models](https://infoscience.epfl.ch/record/161465/files/slfm-long.pdf).
"""
struct LatentFactorMOKernel{Tg,Te<:MOKernel,TA<:AbstractMatrix} <: MOKernel
    g::Tg
    e::Te
    A::TA
    function LatentFactorMOKernel(g, e::MOKernel, A::AbstractMatrix)
        all(gi isa Kernel for gi in g) || error("`g` should be an collection of kernels")
        length(g) == size(A, 2) ||
            error("Size of `A` not compatible with the given array of kernels `g`")
        return new{typeof(g),typeof(e),typeof(A)}(g, e, A)
    end
end

function (κ::LatentFactorMOKernel)((x, px)::Tuple{Any,Int}, (y, py)::Tuple{Any,Int})
    cov_f = sum(κ.A[px, q] * κ.g[q](x, y) * κ.A[py, q] for q in 1:length(κ.g))
    return cov_f + κ.e((x, px), (y, py))
end

function kernelmatrix(k::LatentFactorMOKernel, x::MOInput, y::MOInput)
    x.out_dim == y.out_dim || error("`x` and `y` should have the same output dimension")
    x.out_dim == size(k.A, 1) ||
        error("Kernel not compatible with the given multi-output inputs")

    # Weights matrix ((out_dim x out_dim) x length(k.g))
    W = [col * col' for col in eachcol(k.A)]

    # Latent kernel matrix ((N x N) x length(k.g))
    H = [gi.(x.x, permutedims(y.x)) for gi in k.g]

    # Weighted latent kernel matrix ((N*out_dim) x (N*out_dim))
    W_H = sum(kron(Wi, Hi) for (Wi, Hi) in zip(W, H))

    return W_H .+ kernelmatrix(k.e, x, y)
end

function Base.show(io::IO, k::LatentFactorMOKernel)
    return print(io, "Semi-parametric Latent Factor Multi-Output Kernel")
end

function Base.show(io::IO, ::MIME"text/plain", k::LatentFactorMOKernel)
    print(io, "Semi-parametric Latent Factor Multi-Output Kernel\n\tgᵢ: ")
    join(io, k.g, "\n\t\t")
    print(io, "\n\teᵢ: ")
    return join(io, k.e, "\n\t\t")
end


export bsplsurf, bsplsurfu, dbsurf

"""
	bsplsurf(B[], ordx, ordy, npts, mpts, p1, p2)

Calculate a Cartesian product B-spline surface using open uniform knot vectors.

---

# Arguments
- `B::Array{Float64}`: array containing the control net vertices
- `ordx::Int64`: order in the x direction
- `ordy::Int64`: order in the y direction
- `npts::Int64`: the number of control net vertices in x direction
- `mpts::Int64`: the number of control net vertices in y direction
- `p1::Int64`: number of parametric lines in the x direction
- `p2::Int64`: number of parametric lines in the y direction

---

# Examples
```jldoctest
julia> b = [
        0.0 0.0 0.0 1.0 1.0 1.0 2.0 2.0 2.0;
        0.0 2.0 4.0 0.0 2.0 4.0 0.0 2.0 4.0;
        2.0 0.0 2.0 0.0 1.0 0.0 2.0 0.0 2.0
       ]
3×9 Array{Float64,2}:
 0.0  0.0  0.0  1.0  1.0  1.0  2.0  2.0  2.0
 0.0  2.0  4.0  0.0  2.0  4.0  0.0  2.0  4.0
 2.0  0.0  2.0  0.0  1.0  0.0  2.0  0.0  2.0

julia> bsplsurf(b, 2, 2, 3, 3, 5, 5)[1]
3×25 Array{Float64,2}:
 0.0  0.0  0.0  0.0  0.0  0.5  0.5   0.5  0.5   0.5  1.0  …  1.0  1.5  1.5   1.5  1.5   1.5  2.0  2.0  2.0  2.0  2.0
 0.0  1.0  2.0  3.0  4.0  0.0  1.0   2.0  3.0   4.0  0.0     4.0  0.0  1.0   2.0  3.0   4.0  0.0  1.0  2.0  3.0  4.0
 2.0  1.0  0.0  1.0  2.0  1.0  0.75  0.5  0.75  1.0  0.0     0.0  1.0  0.75  0.5  0.75  1.0  2.0  1.0  0.0  1.0  2.0

julia> bsplsurf(b, 2, 2, 3, 3, 5, 5)[2]
44-element Array{Array{Int64,1},1}:
 [1, 2]
 [2, 3]
 [3, 4]
 [4, 5]
 [5, 6]
 [6, 7]
 [7, 8]
 [8, 9]
 [9, 10]
 [10, 11]
 ⋮
 [12, 17]
 [13, 18]
 [14, 19]
 [15, 20]
 [16, 21]
 [17, 22]
 [18, 23]
 [19, 24]
 [20, 25]

julia> bsplsurf(b, 2, 2, 3, 3, 5, 5)[3]
16-element Array{Array{Int64,1},1}:
 [1, 2, 6, 7]
 [2, 3, 7, 8]
 [3, 4, 8, 9]
 [4, 5, 9, 10]
 [6, 7, 11, 12]
 [7, 8, 12, 13]
 [8, 9, 13, 14]
 [9, 10, 14, 15]
 [11, 12, 16, 17]
 [12, 13, 17, 18]
 [13, 14, 18, 19]
 [14, 15, 19, 20]
 [16, 17, 21, 22]
 [17, 18, 22, 23]
 [18, 19, 23, 24]
 [19, 20, 24, 25]
```

---

_By Paolo Macciacchera, Elia Onofri_
"""

function bsplsurf(B::Array{Float64,2}, ordx::Int64, ordy::Int64, npts::Int64, mpts::Int64, p1::Int64, p2::Int64)::Tuple{Array{Float64, 2}, Array{Array{Int64, 1}, 1}, Array{Array{Int64, 1}, 1}}

    nplusc = npts + ordx
    mplusc = mpts + ordy
    x = zeros(nplusc)
    y = zeros(mplusc)
    nbasis = zeros(1, npts)
    mbasis = zeros(1, mpts)
    q = zeros(3, p1*p2)

    #generate the open uniform knot vectors
    x = knot(npts, ordx)
    y = knot(mpts, ordy)

    icount = 0
    #calculate the points on the B-spline surface
    stepu = x[nplusc] / (p1 - 1)
    stepw = y[mplusc] / (p2 - 1)
    for u = 0 : stepu : x[nplusc]
        nbasis = basis(npts, ordx, u, x)
        for w = 0 : stepw : y[mplusc]
            mbasis = basis(mpts, ordy, w, y)
            icount = icount + 1
            for i = 1 : npts
                for j = 1 : mpts
                    j1 = mpts * (i - 1) + j
                    for s = 1 : 3
                        q[s, icount] = q[s, icount] + B[s, j1] * nbasis[i] * mbasis[j]
                    end
                end
            end
        end
    end

    EV = Array{Int64,1}[] #1-dimensional cellular complex used for plotting the curve with Plasm package

    for i = 1 : p1 * p2 - 1
        C = Array{Int64}(2)
        C[1] = i
        C[2] = i+1
        push!(EV,C)
    end

    for i = 1 : p2 * (p1 - 1)
        C = Array{Int64}(2)
        C[1] = i
        C[2] = i+p2
        push!(EV,C)
    end

    FV = Array{Int64,1}[] #1-dimensional cellular complex used for plotting the curve with Plasm package

    for j = 0 : p1 - 2
        for i = 1 : p2 - 1
            C = Array{Int64}(4)
            C[1] = j * p2 + i
            C[2] = j * p2 + i + 1
            C[3] = (j + 1) * p2 + i
            C[4] = (j + 1) * p2 + i + 1
            push!(FV,C)
        end
    end
    return(q, EV, FV)
end

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

"""
	bsplsurfu(B[], ordx, ordy, npts, mpts, p1, p2)

Calculate a Cartesian product B-spline surface using periodic uniform knot vectors.

---

# Arguments
- `B::Array{Float64}`: array containing the control net vertices
- `ordx::Int64`: order in the x direction
- `ordy::Int64`: order in the y direction
- `npts::Int64`: the number of control net vertices in x direction
- `mpts::Int64`: the number of control net vertices in y direction
- `p1::Int64`: number of parametric lines in the x direction
- `p2::Int64`: number of parametric lines in the y direction

---

# Examples
```jldoctest
julia> b = [
        0.0 0.0 0.0 1.0 1.0 1.0 2.0 2.0 2.0;
        0.0 2.0 4.0 0.0 2.0 4.0 0.0 2.0 4.0;
        2.0 0.0 2.0 0.0 1.0 0.0 2.0 0.0 2.0
       ]
3×9 Array{Float64,2}:
 0.0  0.0  0.0  1.0  1.0  1.0  2.0  2.0  2.0
 0.0  2.0  4.0  0.0  2.0  4.0  0.0  2.0  4.0
 2.0  0.0  2.0  0.0  1.0  0.0  2.0  0.0  2.0

julia> bsplsurfu(b, 2, 2, 3, 3, 5, 5)[1]
3×25 Array{Float64,2}:
 0.0  0.0  0.0  0.0  0.0  0.5  0.5   …  1.5   1.5  2.0  2.0  2.0  2.0  2.0
 0.0  1.0  2.0  3.0  4.0  0.0  1.0      3.0   4.0  0.0  1.0  2.0  3.0  4.0
 2.0  1.0  0.0  1.0  2.0  1.0  0.75     0.75  1.0  2.0  1.0  0.0  1.0  2.0

julia> bsplsurfu(b, 2, 2, 3, 3, 5, 5)[2]
44-element Array{Array{Int64,1},1}:
 [1, 2]
 [2, 3]
 [3, 4]
 [4, 5]
 [5, 6]
 [6, 7]
 [7, 8]
 [8, 9]
 [9, 10]
 [10, 11]
 ⋮
 [12, 17]
 [13, 18]
 [14, 19]
 [15, 20]
 [16, 21]
 [17, 22]
 [18, 23]
 [19, 24]
 [20, 25]

julia> bsplsurfu(b, 2, 2, 3, 3, 5, 5)[3]
16-element Array{Array{Int64,1},1}:
 [1, 2, 6, 7]
 [2, 3, 7, 8]
 [3, 4, 8, 9]
 [4, 5, 9, 10]
 [6, 7, 11, 12]
 [7, 8, 12, 13]
 [8, 9, 13, 14]
 [9, 10, 14, 15]
 [11, 12, 16, 17]
 [12, 13, 17, 18]
 [13, 14, 18, 19]
 [14, 15, 19, 20]
 [16, 17, 21, 22]
 [17, 18, 22, 23]
 [18, 19, 23, 24]
 [19, 20, 24, 25]
```

---

_By Paolo Macciacchera, Elia Onofri_
"""

function bsplsurfu(B::Array{Float64,2}, ordx::Int64, ordy::Int64, npts::Int64, mpts::Int64, p1::Int64, p2::Int64)::Tuple{Array{Float64, 2}, Array{Array{Int64, 1}, 1}, Array{Array{Int64, 1}, 1}}

    nplusc = npts + ordx
    mplusc = mpts + ordy
    x = zeros(nplusc)
    y = zeros(mplusc)
    nbasis = zeros(1, npts)
    mbasis = zeros(1, mpts)
    q = zeros(3, p1 * p2)

    #generate the open uniform knot vectors
    x = knotu(npts, ordx)
    y = knotu(mpts, ordy)

    icount = 0
    #calculate the points on the B-spline surface
    stepu = (npts - ordx + 1) / (p1 - 1)
    stepw = (mpts - ordy + 1) / (p2 - 1)
    for u = ordx-1 : stepu : npts
        nbasis = basis(npts, ordx, u, x)
        for w = ordy-1 : stepw : mpts
            mbasis = basis(mpts, ordy, w, y)
            icount = icount + 1
            for i = 1 : npts
                for j = 1 : mpts
                    for s = 1 : 3
                        j1 = mpts * (i - 1) + j
                        q[s, icount] = q[s, icount] + B[s, j1] * nbasis[i] * mbasis[j]
                    end
                end
            end
        end
    end

    EV = Array{Int64,1}[] #1-dimensional cellular complex used for plotting the curve with Plasm package

    for i = 1 : p1 * p2 - 1
        C = Array{Int64}(2)
        C[1] = i
        C[2] = i+1
        push!(EV,C)
    end

    for i = 1 : p2 * (p1 - 1)
        C = Array{Int64}(2)
        C[1] = i
        C[2] = i+p2
        push!(EV,C)
    end

    FV = Array{Int64,1}[] #1-dimensional cellular complex used for plotting the curve with Plasm package

    for j = 0 : p1 - 2
        for i = 1 : p2 - 1
            C = Array{Int64}(4)
            C[1] = j * p2 + i
            C[2] = j * p2 + i + 1
            C[3] = (j + 1) * p2 + i
            C[4] = (j + 1) * p2 + i + 1
            push!(FV,C)
        end
    end
    return(q, EV, FV)
end

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

"""
	dbsurf(B[], ordx, ordy, npts, mpts, p1, p2)

Calculate a Cartesian product B-spline surface and its derivatives using open uniform knot vectors.

---

# Arguments
- `B::Array{Float64}`: array containing the control net vertices
- `ordx::Int64`: order in the x direction
- `ordy::Int64`: order in the y direction
- `npts::Int64`: the number of control net vertices in x direction
- `mpts::Int64`: the number of control net vertices in y direction
- `p1::Int64`: number of parametric lines in the x direction
- `p2::Int64`: number of parametric lines in the y direction

---

# Examples
```jldoctest
julia> b = [
        0.0 0.0 0.0 1.0 1.0 1.0 2.0 2.0 2.0;
        0.0 2.0 4.0 0.0 2.0 4.0 0.0 2.0 4.0;
        2.0 0.0 2.0 0.0 1.0 0.0 2.0 0.0 2.0
       ]
3×9 Array{Float64,2}:
 0.0  0.0  0.0  1.0  1.0  1.0  2.0  2.0  2.0
 0.0  2.0  4.0  0.0  2.0  4.0  0.0  2.0  4.0
 2.0  0.0  2.0  0.0  1.0  0.0  2.0  0.0  2.0

julia> bsplsurf(arr, 4, 4, 5, 5, 100, 100)[1] == dbsurf(arr, 4, 4, 5, 5, 100, 100)[1]
true
julia> bsplsurf(arr, 4, 4, 5, 5, 100, 100)[2] == dbsurf(arr, 4, 4, 5, 5, 100, 100)[7]
true
julia> bsplsurf(arr, 4, 4, 5, 5, 100, 100)[3] == dbsurf(arr, 4, 4, 5, 5, 100, 100)[8]
true
```

---

_By Paolo Macciacchera, Elia Onofri_
"""

function dbsurf(B::Array{Float64,2}, ordx::Int64, ordy::Int64, npts::Int64, mpts::Int64, p1::Int64, p2::Int64)::Tuple{Array{Float64,2}, Array{Float64,2}, Array{Float64,2}, Array{Float64,2}, Array{Float64,2}, Array{Float64,2}, Array{Array{Int64, 1}, 1}, Array{Array{Int64, 1}, 1}}
    nplusc = npts + ordx
    mplusc = mpts + ordy
    x = zeros(nplusc)
    y = zeros(mplusc)
    nbasis = zeros(1, npts)
    mbasis = zeros(1, mpts)
    d1nbasis = zeros(1, npts)
    d1mbasis = zeros(1, mpts)
    d2nbasis = zeros(1, npts)
    d2mbasis = zeros(1, mpts)
    q = zeros(3, p1 * p2)
    qu = zeros(3, p1 * p2)
    qw = zeros(3, p1 * p2)
    quu = zeros(3, p1 * p2)
    quw = zeros(3, p1 * p2)
    qww = zeros(3, p1 * p2)

    x = knot(npts, ordx)
    y = knot(mpts, ordy)

    icount = 0
    #calculate the points on the B-spline surface
    stepu = x[nplusc] / (p1 - 1)
    stepw = y[mplusc] / (p2 - 1)
    for u = 0 : stepu : x[nplusc]
        nbasis, d1nbasis, d2nbasis = dbasis(npts, ordx, u, x)
        for w = 0 : stepw : y[mplusc]
            mbasis, d1mbasis, d2mbasis = dbasis(mpts, ordy, w, y)
            icount = icount + 1
            for i = 1 : npts
                for j = 1 : mpts
                    for s = 1 : 3
                        j1 = mpts * (i - 1) + j
                        q[s, icount] = q[s, icount] + B[s, j1] * nbasis[i] * mbasis[j]
                        qu[s, icount] = qu[s, icount] + B[s, j1] * d1nbasis[i] * mbasis[j]
                        qw[s, icount] = qw[s, icount] + B[s, j1] * nbasis[i] * d1mbasis[j]
                        quu[s, icount] = quu[s, icount] + B[s, j1] * d1nbasis[i] * d1mbasis[j]
                        quw[s, icount] = quw[s, icount] + B[s, j1] * d2nbasis[i] * mbasis[j]
                        qww[s, icount] = qww[s, icount] + B[s, j1] * nbasis[i] * d2mbasis[j]
                    end
                end
            end
        end
    end

    EV = Array{Int64,1}[] #1-dimensional cellular complex used for plotting the curve with Plasm package

    for i = 1 : p1 * p2 - 1
        C = Array{Int64}(2)
        C[1] = i
        C[2] = i+1
        push!(EV,C)
    end

    for i = 1 : p2 * (p1 - 1)
        C = Array{Int64}(2)
        C[1] = i
        C[2] = i+p2
        push!(EV,C)
    end

    FV = Array{Int64,1}[] #1-dimensional cellular complex used for plotting the curve with Plasm package

    for j = 0 : p1 - 2
        for i = 1 : p2 - 1
            C = Array{Int64}(4)
            C[1] = j * p2 + i
            C[2] = j * p2 + i + 1
            C[3] = (j + 1) * p2 + i
            C[4] = (j + 1) * p2 + i + 1
            push!(FV,C)
        end
    end

    return(q, qu, qw, quu, quw, qww, EV, FV)

end


struct ConstrainedAquisition{A<:EGOAquisition,N} <: EGOAquisition
    base::A
    cons::NTuple{N,Function}
end

update_parameters!(a::ConstrainedAquisition, args...) = update_parameters!(a.base, args...)

function objective(a::ConstrainedAquisition, model::Kriging, x)
    for g in a.cons
        if g(x) > eps()
            return 0.0
        end
    end
    return objective(a.base, model, x)
end


using JuMP
using FlexibilityAnalysis
using Gurobi
using Pavito, Ipopt
using LinearAlgebra

# Set the dimensions
n_gens = 5
n_lines = 20
n_dems = 11

# Setup the uncertainty set parameters
β = 240.
covar = Matrix(I, n_dems, n_dems) * 1200.
covar[covar .== 0] = β
box_dev = ones(n_dems) * 2 * sqrt(covar[1])

# Specify the network details
line_cap = 100
gen_cap = [332; 140; 100; 100; 100]

# Setup the model
m = FlexibilityModel(solver = PavitoSolver(mip_solver = GurobiSolver(OutputFlag = 0),
                     cont_solver = IpoptSolver(print_level = 0), log_level = 0, mip_solver_drives = false))

# Define variables
@randomvariable(m, d[i = 1:n_dems], mean = 0) # Temperarily set the mean to 0
@recoursevariable(m, a[1:n_lines])
@recoursevariable(m, g[1:n_gens])

# Set the line capacity constraints
@constraint(m, [line = 1:n_lines], -line_cap <= a[line])
@constraint(m, [line = 1:n_lines], a[line] <= line_cap)

# Set the generator capacity constraints
@constraint(m, [gen = 1:n_gens], 0.0 <= g[gen])
@constraint(m, [gen = 1:n_gens], g[gen] <= gen_cap[gen])

# Set the node balance constraints
@constraint(m, g[1] - a[1] - a[6] == 0)
@constraint(m, a[1] + g[2] - sum(a[i] for i = [2; 4; 5]) - d[1] == 0)
@constraint(m, g[3] + a[2] - a[3] - d[2] == 0)
@constraint(m, sum(a[i] for i = [3; 4; 8]) - sum(a[i] for i = [7; 11]) - d[3] == 0)
@constraint(m, sum(a[i] for i = [5; 6; 7; 12]) - d[4] == 0)
@constraint(m, g[4] + sum(a[i] for i = [16; 18]) - sum(a[i] for i = [12; 19]) - d[5] == 0)
@constraint(m, a[9] - sum(a[i] for i = [8; 10]) == 0)
@constraint(m, g[5] - a[9] == 0)
@constraint(m, sum(a[i] for i = [10; 11]) - sum(a[i] for i = [13; 14]) - d[6] == 0)
@constraint(m, sum(a[i] for i = [13; 20]) - d[7] == 0)
@constraint(m, a[19] - a[20] - d[8] == 0)
@constraint(m, a[17] - a[18] - d[9] == 0)
@constraint(m, a[15] - sum(a[i] for i = [16; 17]) - d[10] == 0)
@constraint(m, a[14] - a[15] - d[11] == 0)

# Define the uncertainty set
setcovariance(m, covar)
# setuncertaintyset(m, :Ellipsoid, covar)
# setuncertaintyset(m, :Hyperbox, [[box_dev]; [box_dev]])
# setuncertaintyset(m, :PNorm, 1)
# setuncertaintyset(m, :PNorm, 2)
setuncertaintyset(m, :PNorm, Inf)

# Compute a center to replace the mean if desired
new_mean = findcenteredmean(m, center = :feasible, update_mean = true)
updated_mean = getmean(m)

# Solve
status = solve(m)

if status == :Optimal
    # Retrieve optimized data
    flexibility_index = getflexibilityindex(m)
    lines = getvalue(a)
    generators = getvalue(g)
    demands = getvalue(d)
    actives = getactiveconstraints(m)

    # Parse data
    data = getflexibilitydata(m)

    # Get the confidence level
    if data.uncertainty_set.name == :Ellipsoid
        conf_lvl = getconfidencelevel(m)
    end

    # Rank the inequality constraints
    rank_data = rankinequalities(m, max_ranks = 3)
end

# Estimate the value of SF
SF = findstochasticflexibility(m, use_flexibility_index = true)


"""
    KGL01(S::T)

Whittaker
"""
function KGL01(S::T) where {T<:NamedTuple}
  return (S.a+S.b+S.c)/((S.a + (S.a+S.b+S.c))/2.0)
end

function KGL02(S::T) where {T<:NamedTuple}
  return KGL01(S).-1.0
end

function KGL03(S::T) where {T<:NamedTuple}
  return (S.b + S.c)/2.0
end

function KGL04(S::T) where {T<:NamedTuple}
  return S.b + S.c
end

function KGL05(S::T) where {T<:NamedTuple}
  return ((S.a+S.b+S.c)^2)/((S.a+S.b+S.c)^2-2*S.b*S.c)-1.0
end

function KGL06(S::T) where {T<:NamedTuple}
  one_over_sum = 1/(2*S.a+S.b+S.c)
  s1 = (S.a+S.b)*log(S.a+S.b)
  s2 = (S.a+S.c)*log(S.a+S.c)
  return log(2*S.a+S.b+S.c)-(one_over_sum*2*S.a*log(2))-(one_over_sum*(s1+s2))
end

function KGL07(S::T) where {T<:NamedTuple}
  return exp(KGL06(S))-1.0
end

function KGL08(S::T) where {T<:NamedTuple}
  return (S.b + S.c) / (2*S.a + S.b + S.c)
end

function KGL09(S::T) where {T<:NamedTuple}
  return KGL08(S)
end

function KGL10(S::T) where {T<:NamedTuple}
  return S.a / (S.a + S.b + S.c)
end

function KGL11(S::T) where {T<:NamedTuple}
  return (2*S.a) / (2*S.a + S.b + S.c)
end

function KGL12(S::T) where {T<:NamedTuple}
  return (2*S.a+S.b+S.c) * ( 1 - (S.a/(S.a+S.b+S.c)))
end

function KGL13(S::T) where {T<:NamedTuple}
  return min(S.b, S.c) / (max(S.b, S.c) + S.a)
end

function KGL14(S::T) where {T<:NamedTuple}
  return 1 - (S.a*(2*S.a+S.b+S.c))/(2*(S.a+S.b)*(S.a+S.c))
end

function KGL15(S::T) where {T<:NamedTuple}
  return (S.b + S.c)/(S.a + S.b + S.c)
end

function KGL16(S::T) where {T<:NamedTuple}
  return KGL15(S)
end

function KGL17(S::T) where {T<:NamedTuple}
  return min(S.b, S.c)/(S.a+S.b+S.c)
end

function KGL18(S::T) where {T<:NamedTuple}
  return (S.b+S.c)/2
end

function KGL19(S::T) where {T<:NamedTuple}
  return ((S.b*S.c)+1)/((S.a+S.b+S.c)^2-(S.a+S.b+S.c)/2)
end

function KGL20(S::T) where {T<:NamedTuple}
  return 1 - (2*S.a)/(2*S.a.+S.b+S.c)
end

function KGL21(S::T) where {T<:NamedTuple}
  return S.a/(S.a+S.c)
end

function KGL22(S::T) where {T<:NamedTuple}
  return min(S.b, S.c)/(min(S.b, S.c) + S.a)
end

function KGL23(S::T) where {T<:NamedTuple}
  return 2*abs(S.b-S.c)/(2*S.a+S.b+S.c)
end

function KGL24(S::T) where {T<:NamedTuple}
  in_par = (2*S.a+S.b+S.c)/(S.a+S.b+S.c)
  return 1-log(in_par)/log(2)
end


using HuberLR
using Base.Test

let
    # true parameters
    w = [1., 2, 3., 4.]
    b = 5.

    # generate sample points
    srand(20150218)
    x = randn(length(w), 150)
    y = mapslices(x, 1) do x
        w ⋅ x + b + randn() + (rand() > .9 ? abs(randn()) * 5 : 0.)
    end |> vec
    #writedlm("data.tsv", vcat(x, y')')

    # JIT compile
    lr = LinearRegression(.5)
    fit!(lr, x[:,1:5], y[1:5])

    # robust regression
    lr = LinearRegression(1.)
    @time fit!(lr, x, y)
    println(lr)
    @show norm(predict(lr, x) .- y)
    Δw = norm(lr.w .- [b, w])

    # quadratic loss (not robust)
    lr = LinearRegression(Inf)
    @time fit!(lr, x, y)
    println(lr)
    @show norm(predict(lr, x) .- y)
    Δw′ = norm(lr.w .- [b, w])

    #@show Δw Δw′
    @test Δw < Δw′ < 1.
end


export RigCategory, SymmetricRigCategory, DistributiveMonoidalCategory,
  DistributiveSemiadditiveCategory, DistributiveCategory

# Distributive categories
#########################

""" Theory of a *rig category*, also known as a *bimonoidal category*

[Rig categories](https://ncatlab.org/nlab/show/rig+category) are the most
general in the hierarchy of [distributive monoidal
structures](https://ncatlab.org/nlab/show/distributivity+for+monoidal+structures).

TODO: Do we also want the distributivty and absorption isomorphisms? Usually we
ignore coherence isomorphisms such as associators and unitors.

FIXME: This theory should also inherit `MonoidalCategory`, but multiple
inheritance is not supported.
"""
@signature RigCategory{Ob,Hom} <: SymmetricMonoidalCategoryAdditive{Ob,Hom} begin
  otimes(A::Ob, B::Ob)::Ob
  otimes(f::(A → B), g::(C → D))::((A ⊗ C) → (B ⊗ D)) ⊣
    (A::Ob, B::Ob, C::Ob, D::Ob)
  @op (⊗) := otimes
  munit()::Ob
end

""" Theory of a *symmetric rig category*

FIXME: Should also inherit `SymmetricMonoidalCategory`.
"""
@signature SymmetricRigCategory{Ob,Hom} <: RigCategory{Ob,Hom} begin
  braid(A::Ob, B::Ob)::((A ⊗ B) → (B ⊗ A))
  @op (σ) := braid
end

""" Theory of a *distributive (symmetric) monoidal category*

Reference: Jay, 1992, LFCS tech report LFCS-92-205, "Tail recursion through
universal invariants", Section 3.2

FIXME: Should also inherit `CocartesianCategory`.
"""
@theory DistributiveMonoidalCategory{Ob,Hom} <: SymmetricRigCategory{Ob,Hom} begin
  plus(A::Ob)::((A ⊕ A) → A)
  zero(A::Ob)::(mzero() → A)

  copair(f::(A → C), g::(B → C))::((A ⊕ B) → C) <= (A::Ob, B::Ob, C::Ob)
  coproj1(A::Ob, B::Ob)::(A → (A ⊕ B))
  coproj2(A::Ob, B::Ob)::(B → (A ⊕ B))

  copair(f,g) == (f⊕g)⋅plus(C) ⊣ (A::Ob, B::Ob, C::Ob, f::(A → C), g::(B → C))
  coproj1(A,B) == id(A)⊕zero(B) ⊣ (A::Ob, B::Ob)
  coproj2(A,B) == zero(A)⊕id(B) ⊣ (A::Ob, B::Ob)

  # Naturality axioms.
  plus(A)⋅f == (f⊕f)⋅plus(B) ⊣ (A::Ob, B::Ob, f::(A → B))
  zero(A)⋅f == zero(B) ⊣ (A::Ob, B::Ob, f::(A → B))
end

""" Theory of a *distributive monoidal category with diagonals*

FIXME: Should also inherit `MonoidalCategoryWithDiagonals`.
"""
@theory DistributiveMonoidalCategoryWithDiagonals{Ob,Hom} <:
    DistributiveMonoidalCategory{Ob,Hom} begin
  mcopy(A::Ob)::(A → (A ⊗ A))
  @op (Δ) := mcopy
  delete(A::Ob)::(A → munit())
  @op (◊) := delete
end

""" Theory of a *distributive semiadditive category*

This terminology is not standard but the concept occurs frequently. A
distributive semiadditive category is a semiadditive category (or biproduct)
category, written additively, with a tensor product that distributes over the
biproduct.

FIXME: Should also inherit `SemiadditiveCategory`
"""
@theory DistributiveSemiadditiveCategory{Ob,Hom} <: DistributiveMonoidalCategory{Ob,Hom} begin
  mcopy(A::Ob)::(A → (A ⊕ A))
  @op (Δ) := mcopy
  delete(A::Ob)::(A → mzero())
  @op (◊) := delete

  pair(f::(A → B), g::(A → C))::(A → (B ⊕ C)) ⊣ (A::Ob, B::Ob, C::Ob)
  proj1(A::Ob, B::Ob)::((A ⊕ B) → A)
  proj2(A::Ob, B::Ob)::((A ⊕ B) → B)

  # Naturality axioms.
  f⋅Δ(B) == Δ(A)⋅(f⊕f) ⊣ (A::Ob, B::Ob, f::(A → B))
  f⋅◊(B) == ◊(A) ⊣ (A::Ob, B::Ob, f::(A → B))
end

""" Theory of a *distributive category*

A distributive category is a distributive monoidal category whose tensor product
is the cartesian product, see [`DistributiveMonoidalCategory`](@ref).

FIXME: Should also inherit `CartesianCategory`.
"""
@theory DistributiveCategory{Ob,Hom} <: DistributiveMonoidalCategoryWithDiagonals{Ob,Hom} begin
  pair(f::(A → B), g::(A → C))::(A → (B ⊗ C)) ⊣ (A::Ob, B::Ob, C::Ob)
  proj1(A::Ob, B::Ob)::((A ⊗ B) → A)
  proj2(A::Ob, B::Ob)::((A ⊗ B) → B)

  pair(f,g) == Δ(C)⋅(f⊗g) ⊣ (A::Ob, B::Ob, C::Ob, f::(C → A), g::(C → B))
  proj1(A,B) == id(A)⊗◊(B) ⊣ (A::Ob, B::Ob)
  proj2(A,B) == ◊(A)⊗id(B) ⊣ (A::Ob, B::Ob)

  # Naturality axioms.
  f⋅Δ(B) == Δ(A)⋅(f⊗f) ⊣ (A::Ob, B::Ob, f::(A → B))
  f⋅◊(B) == ◊(A) ⊣ (A::Ob, B::Ob, f::(A → B))
end


@testset "Running WBS tests..." begin

    n = 1000;        # Number of samples
    λ = 100;         # Frequencey of changes

    ########################
    # Normal mean segments #
    ########################
    μ, σ = Normal(0.0, 10.0), 1.0
    sample, cps = @changepoint_sampler n λ Normal(μ, σ)
    seg_cost_CUSUM = Changepoints.CUSUM(sample);
    seg_cost_sSIC = Changepoints.sSIC(sample);

    @test_nowarn WBS(seg_cost_CUSUM, n);
    @test_nowarn get_WBS_changepoints(seg_cost_sSIC, WBS(seg_cost_CUSUM, n))
    @test_nowarn WBS(seg_cost_CUSUM, n, 1.3, 1.0, 5000, true , 1/sqrt(2));

end


"""
    function stepBoris(x,p,efield,bfield,q,m,dt)

Performs one complete time step of the Boris algorithm for relativistic charged particles in an electromagnetic field. This version is from Ripperda2017
"""
function stepBoris(x, p, efield, bfield, q, m, dt)
    qdt2m = q * dt / (2.0 * m)

    un = p / m
    γn = γFromP(p, m)
    xnph = x + un * dt / (2.0 * γn)

    eph = efield(xnph)
    um = un + qdt2m * eph

    γm = γFromP(m * um, m)
    t = (qdt2m / γm) * bfield(xnph)
    s = (2.0 / (1 + t ⋅ t)) * t

    up = um + cross((um + cross(um, t)), s)

    unp1 = up + eph * q * dt / (2.0 * m)
    pnew = m * unp1

    γnp1 = γFromP(pnew, m)
    xnew = xnph + (dt / (2.0 * γnp1)) * unp1
    return ArrayPartition(xnew, pnew)
end
export stepBoris

"""
    fBoris!(u::ArrayPartition, params::Dict, t)

Boris stepper wrapper for the discrete solver.
Use with single particle coordinates.

Input:
    - u{ArrayPartition{Number}} ≡ ArrayPartition([position],[momentum])
    - params{Dict} ≡ Dict[
            :t => {Number} current time value,
            :efield => {Function} returns electric field at [x,y,z]
            :bfield => {Function} returns magnetic flux density at [x,y,z]
            :q      => {Number} charge of particle in SI
            :m      => {Number} mass of particle in SI
        ]
    - t{Number} ≡ time at next time step

Returns:
    Particle coordinates at time t in the form
    {ArrayPartition{Number}}([[x,y,z],[px,py,pz]]).
"""
function fBoris!(u::ArrayPartition, params::Dict, t)
    dt = t - params[:t]
    params[:t] = t
    stepBoris(
        u.x[1],
        u.x[2],
        params[:efield],
        params[:bfield],
        params[:q],
        params[:m],
        dt
    )
end

"""
    fBoris!(u::Snapshot, params::Dict, t)

Boris stepper wrapper for the discrete solver.
Use with Snapshot branch of ParticleTree.

Input:
    - u{Snapshot}  ≡ level 2 ParticleTree branch
    - params{Dict} ≡ Dict[
            :t      => {Number} current time value,
            :efield => {Function} returns electric field at [x,y,z]
            :bfield => {Function} returns magnetic flux density at [x,y,z]
        ]
    - t{Number}    ≡ time at next time step

Returns:
    Snapshot at time t
"""
function fBoris!(u::Snapshot, params::Dict, t)
    dt = t - params[:t]
    params[:t] = t
    snapOut = deepcopy(u)
    for species in snapOut.nodes
        q = usistrip(charge(species.id))
        m = usistrip(mass(species.id))
        map(species.nodes) do coords
            stepCoords = stepBoris(
                    coords.values.x[1],
                    coords.values.x[2],
                    params[:efield],
                    params[:bfield],
                    q,
                    m,
                    dt
                )
            coords.values .= stepCoords
        end
    end
    return snapOut
end

export fBoris!

# dictionary of all relativistic integrators availble
rintegrator = Dict(
                :boris => fBoris!
)

export rintegrator


@testset "gaussian state bhd" begin
    dim = 100
    ρ = SqueezedState(0.8, π, Matrix, dim=dim)
    d = GaussianStateBHD(ρ)

    θ, x = rand(d)
    @test 0 ≤ θ < 2π

    points = rand(d, 4096)
    @test all(0 ≤ θ < 2π for θ in points[1, :])
end

@testset "operator" begin
    T = Float64
    dim = 100

    ρ = Matrix{T}(I, dim, dim)
    c = QSB.Creation(T, dim=dim)
    a = QSB.Annihilation(T, dim=dim)

    @test QSD.mean_of_create(ρ) ≈ tr(c * ρ)
    @test QSD.mean_of_create²(ρ) ≈ tr(c^2 * ρ)
    @test QSD.mean_of_annihilate(ρ) ≈ tr(a * ρ)
    @test QSD.mean_of_annihilate²(ρ) ≈ tr(a^2 * ρ)
    @test QSD.mean_of_create_annihilate(ρ) ≈ tr(c*a * ρ)

    θ = π/8
    @test QSD.mean_of_π̂ₓ²(ρ, θ) ≈ (
        exp(-2im*θ) * tr(a^2 * ρ) +
        exp(2im*θ) * tr(c^2 * ρ) +
        1 + 2tr(c*a * ρ)
    ) / 4
    @test QSD.mean_of_π̂ₓ(ρ, θ) ≈ (
        exp(-im*θ) * tr(a * ρ) +
        exp(im*θ) * tr(c * ρ)
    ) / 2

    θs = LinRange(0, 2π, 4096)
    @test QSD.mean_of_π̂ₓ²(ρ, θs) ≈ (
        exp.(-2im*θs) .* tr(a^2 * ρ) .+
        exp.(2im*θs) .* tr(c^2 * ρ) .+
        1 .+ 2tr(c*a * ρ)
    ) ./ 4
    @test QSD.mean_of_π̂ₓ(ρ, θs) ≈ (
        exp.(-im*θs) .* tr(a * ρ) .+
        exp.(im*θs) .* tr(c * ρ)
    ) ./ 2
end


using AdalmPluto
using FFTW

function normalizePower!(buffer)
    maxI = (buffer |> real .|> abs |> maximum) * 1.05 # Max scale + backoff of I path
    maxQ = (buffer |> imag .|> abs |> maximum) * 1.05 # Max scale + backoff of Q path
    buffer .= real(buffer) ./ maxI + 1im*imag(buffer) / maxQ
end

function tx()

    # ----------------------------------------------------
    # --- Radio
    # ----------------------------------------------------
    carrierFreq		= 700e6
    samplingRate    = 4e6;

    # ----------------------------------------------------
    # --- Opening radio
    # ----------------------------------------------------
    # global pluto = openPluto(Int(carrierFreq), Int(samplingRate), Int(40),bufferSize=UInt(2048));
    global pluto = openPluto(Int(carrierFreq), Int(samplingRate), Int(40));

    # ----------------------------------------------------
    # --- Create buffer
    # ----------------------------------------------------
    f_c     = 1e6
    nbSamples = 2048
    buffer  = 0.9.*[exp.(2im * π * f_c / samplingRate * n)  for n ∈ (0:nbSamples-1)];
    buffer  = convert.(Complex{Cfloat},buffer);

    # --- Emulate a simple OFDM signal
    nS = 1024
    bData = ComplexF32.( (rand([0 1],nS) +1im*(rand([0 1],nS)))/sqrt(2))
    buffer = zeros(ComplexF32,nbSamples)
    buffer[1:nS÷2] = bData[1:nS÷2]
    buffer[end-nS÷2+1:end] = bData[1+nS÷2:end]
    buffer = ifft(buffer)
    normalizePower!(buffer)

    # ----------------------------------------------------
    # --- Send buffer
    # ----------------------------------------------------
    cnt = 0
    try
        while (true)
            n = send(pluto,buffer)
            cnt += n
            yield()
        end
    catch exception
        @info "Get interruption, close radio after $cnt elements "
        close(pluto)
        rethrow(exception)
    end
    close(radio)
    return cnt
end


############################################################
# joMatrix - miscaleneous constructor-only operators #######
############################################################

# identity operators: joDirac
include("joMatrixConstructors/joDirac.jl")

# identity operators: joEye
include("joMatrixConstructors/joEye.jl")

# diagonal operators: joDiag
include("joMatrixConstructors/joDiag.jl")

# matrix of ones
include("joMatrixConstructors/joOnes.jl")

# matrix of zeros
include("joMatrixConstructors/joZeros.jl")

# matrix of constants
include("joMatrixConstructors/joConstants.jl")

# vector conversion operators: joReal joImag joConj
include("joMatrixConstructors/joVecConvert.jl")

# sinc interpolation
include("joMatrixConstructors/joSincInterp.jl")

# 1D linear interpolation
include("joMatrixConstructors/joLinInterp1D.jl")


module TestUtils

export gradcheck

using YAAD
import YAAD: Value

make_jacobian(input::AbstractArray, num_out::Int) = (zeros(eltype(input), length(input), num_out), )
make_jacobian(input::Variable, num_out::Int) = make_jacobian(value(input), num_out)
make_jacobian(inputs::Tuple, num_out::Int) = Tuple(first(make_jacobian(each, num_out)) for each in inputs)

zero_like(x::T) where {T <: Number} = zero(T)
zero_like(x) = fill!(similar(x), 0)
zero_like(x::Broadcast.Broadcasted) = zero_like(Broadcast.materialize(x))
zero_like(x::Value) = zero_like(value(x))

function get_numerical_jacobian(f, inputs...; target=inputs, eps=1e-3)
    output_size = length(value(f(inputs...)))
    jacobian = make_jacobian(target, output_size)

    x_tensors = [value(t) for t in target]
    for (x_tensor, d_tensor) in zip(x_tensors, jacobian)
        for (d_idx, x_idx) in enumerate(eachindex(x_tensor))
            orig = x_tensor[x_idx]
            x_tensor[x_idx] = orig - eps
            outa = copy(value(f(inputs...)))
            x_tensor[x_idx] = orig + eps
            outb = copy(value(f(inputs...)))
            x_tensor[x_idx] = orig

            r = (outb - outa) / (2 * eps)
            _set_numerical_jacobian_elem!(d_tensor, d_idx, r)
        end
    end
    return jacobian
end

_set_numerical_jacobian_elem!(d_tensor, d_idx, r::AbstractArray) = d_tensor[d_idx, :] = reshape(r, :)
_set_numerical_jacobian_elem!(d_tensor, d_idx, r::Number) = d_tensor[d_idx, 1] = r

function get_analytical_jacobian(inputs, output::Value)
    jacobian = make_jacobian(inputs, length(value(output)))
    jacobian_reentrant = make_jacobian(inputs, length(value(output)))
    grad_output = zero_like(output)
    reentrant = true
    correct_grad_sizes = true

    _analytical_jacobian!(jacobian, jacobian_reentrant, grad_output, output)

    for (jacobian_x, jacobian_reentrant_x) in zip(jacobian, jacobian_reentrant)
        length(jacobian_x) != 0 && maximum(abs.(jacobian_x - jacobian_reentrant_x)) != 0 && (reentrant = false)
    end

    jacobian, reentrant, correct_grad_sizes
end

function _analytical_jacobian!(jacobian::Tuple, jacobian_reentrant::Tuple, grad_output::AbstractArray, output::Value)
    for i in 1:length(grad_output)
        grad_output = fill!(grad_output, 0)
        grad_output[i] = 1
        for jacobian_c in (jacobian, jacobian_reentrant)
            grads_input = gradient(output, grad_output)
            for (jacobian_x, d_x) in zip(jacobian_c, grads_input)
                jacobian_x[:, i] = reshape(d_x, :)
                # TODO: check gradient size
            end
        end
    end
    jacobian, jacobian_reentrant
end

function _analytical_jacobian!(jacobian::Tuple, jacobian_reentrant::Tuple, grad_output::Number, output::Value)
    grad_output = 1
    for jacobian_c in (jacobian, jacobian_reentrant)
        grads_input = gradient(output, grad_output)
        for (jacobian_x, d_x) in zip(jacobian_c, grads_input)
            jacobian_x[:, 1] = reshape(d_x, :)
            # TODO: check gradient size
        end
    end

    jacobian, jacobian_reentrant
end

_differentiable_outputs(x) = Tuple(o for o in _as_tuple(x) if o isa Value)
_as_tuple(x) = (x, )
_as_tuple(x::Tuple) = x

function gradcheck(f, inputs...; eps=1e-6, atol=1e-5, rtol=1e-3, raise_exception=true)
    any_input_requiring_grad = false
    for inp in inputs
        if inp isa Variable
            any_input_requiring_grad = true
        end
    end

    any_input_requiring_grad || error("gradcheck expects at least one input to require gradient")

    output = _differentiable_outputs(f(inputs...))

    # replace this with assert in the future
    function fail_test(msgs...)
        raise_exception && error(msgs...)
        false
    end

    for (i, o) in enumerate(output)
        fn(inputs...) = _as_tuple(f(inputs...))[i]

        analytical, reentrant, correct_grad_sizes = get_analytical_jacobian(inputs, o)
        numerical = get_numerical_jacobian(fn, inputs...; eps=eps)

        correct_grad_sizes || fail_test("Analytical gradient has incorrect size")

        for (j, (a, n)) in enumerate(zip(analytical, numerical))
            if length(a) != 0 || length(n) != 0
                isapprox(a, n; rtol=rtol, atol=atol) ||
                    fail_test("Jacobian mismatch for output $(i) with respect to input $j, \n",
                              "numerical: $n \nanalytical: $a\n")
            end
        end

        # TODO: reentrant is not actually checked
        reentrant || fail_test("Backward is not reentrant, i.e., running backward with same ",
                               "input and grad_output multiple times gives different values, ",
                               "although analytical gradient matches numerical gradient")
    end

    true
end

end # module TestUtils


using FastGaussQuadrature, Compat
using Compat.Test, Compat.LinearAlgebra, Compat.Random, SpecialFunctions

if VERSION < v"0.7-"
    const floatmax = realmax
end

println("Chebyshev tests")
include("test_gausschebyshev.jl")
println("Legendre tests")
include("test_gausslegendre.jl")
println("Jacobi tests")
include("test_gaussjacobi.jl")
println("Radau tests")
include("test_gaussradau.jl")
println("Lobatto tests")
include("test_gausslobatto.jl")
println("Bessel roots tests")
include("test_besselroots.jl")
println("Laguerre tests")
include("test_gausslaguerre.jl")
println("Hermite tests")
include("test_gausshermite.jl")


# Create example data
x = LinRange(-2, 2, 40)
y = x.^3 .+ x.^2 .+ x .+ 6
# Plot x and y, with dashed stems ended in a star
stem(x, y, "--p")
# Move the baseline to 5
stem(x, y, baseline = 5)