import Base:*, broadcast, reshape, exp, log, tanh, sum,
    adjoint, inv, argmax, argmin, ^, max, maximum, min, minimum,
    vec, \, cos, sin, sign, map, prod, reverse
import LinearAlgebra: tr, diag, det, norm, diagm, dot, I, svd, tril, triu
import Statistics: mean, std
import FFTW: fft, ifft
export
*,
^,
einsum,
sigmoid,
tanh,
mean,
log,
exp,
softplus,
softmax,
softsign,
sum,
relu,
relu6,
squeeze,
adjoint,
diag,
diagm,
det,
inv,
triangular_solve,
argmin,
argmax,
max,
min,
group_assign,
assign,
maximum,
minimum,
cast,
group,
clip,
scatter_add,
scatter_sub,
scatter_update,
stack,
concat,
unstack,
norm,
cvec,
rvec,
vec,
sqrt,
mean,
pad,
leaky_relu,
fft,
ifft,
I,
svd,
vector,
pmap,
std,
lgamma,
topk,
argsort,
batch_matmul,
dot,
set_shape,
selu,
elu,
tr,
tril,
triu,
solve_batch,
swish, hard_sigmoid, hard_swish, concat_elu, concat_hard_swish, concat_relu, fourier

@doc raw"""
    batch_matmul(o1::PyObject, o2::PyObject)

Computes `o1[i,:,:] * o2[i, :]` or `o1[i,:,:] * o2[i, :, :]` for each index `i`.
"""
function batch_matmul(o1::PyObject, o2::PyObject)
    flag = false
    if length(size(o2))==2
        flag = true
        o2 = tf.expand_dims(o2, 2)
    end
    if length(size(o1))!=3 || length(size(o2))!=3
        error("The size of o1 or o2 is not valid.")
    end
    out = tf.matmul(o1, o2)
    if flag
        squeeze(out)
    else
        out
    end
end

batch_matmul(o1::Array{<:Real}, o2::PyObject) = batch_matmul(constant(o1), o2)
batch_matmul(o1::PyObject, o2::Array{<:Real}) = batch_matmul(o1, constant(o2))
batch_matmul(o1::Array{<:Real}, o2::Array{<:Real}) = batch_matmul(constant(o1), constant(o2))

function PyCall.:*(o1::PyObject, o2::PyObject)
    s1 = size(o1)
    s2 = size(o2)
    if s1==nothing || s2==nothing
        error("o1 and o2 should be tensors of rank 0, 1, 2")
    end
    if length(s1) == 0 || length(s2)==0
        return tf.multiply(o1, o2)
    end
    if length(s1)==2 && length(s2)==2
        return tf.matmul(o1, o2)
    elseif length(s1)==2 && length(s2)==1
        return tf.einsum("nm,m->n", o1, o2)
    elseif length(s1)==2 && length(s2)==0
        return tf.multiply(o1, o2)
    elseif length(s1)==1 && length(s2)==2
        error("[rand 1] x [rank 2] not defined")
    elseif length(s1)==1 && length(s2)==1
        return tf.multiply(o1, o2)
    else
        @warn("Unusual usage of multiplication. Check carefully")
        tf.matmul(o1,o2)
    end
end

Base.:*(o1::PyObject, o2::AbstractArray{<:Real}) = *(o1, constant(Array(o2), dtype=get_dtype(o1)))
Base.:*(o1::AbstractArray{<:Real}, o2::PyObject) = *(constant(Array(o1), dtype=get_dtype(o2)), o2)
Base.:*(o1::Number, o2::PyObject) = *(constant(o1, dtype=get_dtype(o2)), o2)
Base.:*(o1::PyObject, o2::Number) = *(o1, constant(o2, dtype=get_dtype(o1)))

Base.Broadcast.broadcasted(::typeof(*), o1::PyObject, o2::AbstractArray{<:Real}) = tf.multiply(o1, Array(o2))
Base.Broadcast.broadcasted(::typeof(*), o1::AbstractArray{<:Real}, o2::PyObject) = tf.multiply(Array(o1), o2)
Base.Broadcast.broadcasted(::typeof(*), o1::PyObject, o2::PyObject) = tf.multiply(o1, o2)
Base.Broadcast.broadcasted(::typeof(*), o1::PyObject, o2::Number) = tf.multiply(o1, o2)
Base.Broadcast.broadcasted(::typeof(*), o1::Number, o2::PyObject) = tf.multiply(o1, o2)

Base.Broadcast.broadcasted(::typeof(/), o1::PyObject, o2::AbstractArray{<:Real}) = tf.divide(o1, Array(o2))
Base.Broadcast.broadcasted(::typeof(/), o1::AbstractArray{<:Real}, o2::PyObject) = tf.divide(Array(o1), o2)
Base.Broadcast.broadcasted(::typeof(/), o1::PyObject, o2::PyObject) = tf.divide(o1, o2)
Base.Broadcast.broadcasted(::typeof(/), o1::PyObject, o2::Number) = tf.divide(o1, o2)
Base.Broadcast.broadcasted(::typeof(/), o1::Number, o2::PyObject) = tf.divide(o1, o2)

Base.Broadcast.broadcasted(::typeof(+), o1::PyObject, o2::AbstractArray{<:Real}) = o1 + Array(o2)
Base.Broadcast.broadcasted(::typeof(+), o1::AbstractArray{<:Real}, o2::PyObject) = Array(o1) + o2
Base.Broadcast.broadcasted(::typeof(+), o1::PyObject, o2::PyObject) = o1 + o2
Base.Broadcast.broadcasted(::typeof(+), o1::PyObject, o2::Number) = o1 + o2
Base.Broadcast.broadcasted(::typeof(+), o1::Number, o2::PyObject) = o1 + o2

Base.Broadcast.broadcasted(::typeof(-), o1::PyObject, o2::AbstractArray{<:Real}) = o1 - Array(o2)
Base.Broadcast.broadcasted(::typeof(-), o1::AbstractArray{<:Real}, o2::PyObject) = Array(o1) - o2
Base.Broadcast.broadcasted(::typeof(-), o1::PyObject, o2::PyObject) = o1 - o2
Base.Broadcast.broadcasted(::typeof(-), o1::PyObject, o2::Number) = o1 - o2
Base.Broadcast.broadcasted(::typeof(-), o1::Number, o2::PyObject) = o1 - o2

warn_broadcast_pow() = error(".^ is disabled due to eager evaluation. Use ^ instead.")
Base.Broadcast.broadcasted(::typeof(^), o1::PyObject, o2::Union{AbstractArray{<:Real},Number}) = warn_broadcast_pow()
Base.Broadcast.broadcasted(::typeof(^), o1::PyObject, o2::PyObject) = warn_broadcast_pow()
Base.Broadcast.broadcasted(::typeof(^), o1::Union{AbstractArray{<:Real},Number}, o2::PyObject) = warn_broadcast_pow()

function einsum(equation, args...; kwargs...)
    tf.einsum(equation, args...; kwargs...)
end

"""
    reshape(o::PyObject, s::Union{Array{<:Integer}, Tuple{Vararg{<:Integer, N}}}) where N
    reshape(o::PyObject, s::Integer; kwargs...)
    reshape(o::PyObject, m::Integer, n::Integer; kwargs...)
    reshape(o::PyObject, ::Colon, n::Integer)
    reshape(o::PyObject, n::Integer, ::Colon)

Reshapes the tensor according to row major if the "TensorFlow style" syntax is used; otherwise
reshaping according to column major is assumed.

# Example
```julia
reshape(a, [10,5]) # row major
reshape(a, 10, 5) # column major
```
"""
function reshape(o::PyObject, s::Union{Array{<:Integer}, Tuple{Vararg{<:Integer, N}}}) where N
    tf.reshape(o, s)
end

function reshape(o::PyObject, s::Integer; kwargs...)
    if length(size(o))>=2
        return tf.reshape(o', [s]; kwargs...)
    end
    tf.reshape(o, [s]; kwargs...)
end

function reshape(o::PyObject, m::Integer, n::Integer; kwargs...)
    if length(size(o))==1
        return tf.reshape(o, [n,m]; kwargs...)'
    elseif length(size(o))==2
        return tf.reshape(o', [n,m]; kwargs...)'
    end
    tf.reshape(o, [m, n]; kwargs...)
end

reshape(o::PyObject, ::Colon, n::Integer) = reshape(o, -1, n)
reshape(o::PyObject, n::Integer, ::Colon) = reshape(o, n, -1)

"""
    rvec(o::PyObject; kwargs...)

Vectorizes the tensor `o` to a row vector, assuming column major.
"""
function rvec(o::PyObject; kwargs...)
    s = size(o)
    if length(s)==0
        return reshape(o, 1, 1, kwargs...)
    elseif length(s)==1
        return reshape(o, 1, s[1], kwargs...)
    elseif length(s)==2
        return reshape(o, 1, s[1]*s[2], kwargs...)
    else
        error("Invalid argument")
    end
end

"""
    rvec(o::PyObject; kwargs...)

Vectorizes the tensor `o` to a column vector, assuming column major.
"""
function cvec(o::PyObject;kwargs...)
    s = size(o)
    if length(s)==0
        return reshape(o, 1, 1,kwargs...)
    elseif length(s)==1
        return reshape(o, s[1], 1,kwargs...)
    elseif length(s)==2
        return reshape(o, s[1]*s[2], 1,kwargs...)
    else
        error("Invalid argument")
    end
end

"""
    vec(o::PyObject;kwargs...)

Vectorizes the tensor `o` assuming column major.
"""
function vec(o::PyObject;kwargs...)
    s = size(o)
    if length(s)==0
        return reshape(o, 1, kwargs...)
    elseif length(s)==1
        return o
    elseif length(s)>=2
        return reshape(o, s[1]*s[2])
    end
end

"""
    set_shape(o::PyObject, s::Union{Array{<:Integer}, Tuple{Vararg{<:Integer, N}}}) where N
    set_shape(o::PyObject, s::Integer...)

Sets the shape of `o` to `s`. `s` must be the actual shape of `o`. This function is used to convert a
tensor with unknown dimensions to a tensor with concrete dimensions.

# Example
```julia
a = placeholder(Float64, shape=[nothing, 10])
b = set_shape(a, 3, 10)
run(sess, b, a=>rand(3,10)) # OK
run(sess, b, a=>rand(5,10)) # Error
run(sess, b, a=>rand(10,3)) # Error
```
"""
function set_shape(o::PyObject, s::Union{Array{<:Integer}, Tuple{Vararg{<:Integer, N}}}) where N
    o.set_shape(s)
    return o
end

set_shape(o::PyObject, s::Integer...) = set_shape(o, s)

function sigmoid(x::Real)
    return 1/(1+exp(-x))
end

function sigmoid(o::PyObject; kwargs...)
    tf.math.sigmoid(o; kwargs...)
end

relu(x::Real) = max(zero(x), x)
function relu(o::PyObject; kwargs...)
    tf.nn.relu(o; kwargs...)
end

relu6(x::Real) = min(relu(x), one(x)*oftype(x, 6))
relu6(o::PyObject; kwargs...) = tf.nn.relu6(o; kwargs...)

function tan(o::PyObject; kwargs...)
    tf.math.tan(o; kwargs...)
end

function leaky_relu(x::Real, a = oftype(x / 1, 0.2))
    max(a * x, x / one(x))
end

function leaky_relu(o::PyObject; kwargs...)
    tf.nn.leaky_relu(o; kwargs...)
end

function tanh(o::PyObject; kwargs...)
    tf.tanh(o; kwargs...)
end

function selu(x::Real)
    λ = oftype(x / 1, 1.0507009873554804934193349852946)
    α = oftype(x / 1, 1.6732632423543772848170429916717)
    λ * ifelse(x > 0, x / one(x), α * (exp(x) - one(x)))
end

selu(o::PyObject; kwargs...) = tf.nn.selu(o; kwargs...)

elu(x, α = one(x)) = ifelse(x ≥ 0, x / one(x), α * (exp(x) - one(x)))
elu(o::PyObject; kwargs...) = tf.nn.elu(o; kwargs...)

softsign(x::Real) = x / (one(x) + abs(x))
softsign(o::PyObject; kwargs...) = tf.nn.softsign(o; kwargs...)

function argmax(o::PyObject; kwargs...)
    kwargs = jlargs(kwargs)
    tf.argmax(o; kwargs...) + 1
end

function Base.:sqrt(o::PyObject; kwargs...)
    kwargs = jlargs(kwargs)
    tf.sqrt(o)
end

function argmin(o::PyObject; kwargs...)
    kwargs = jlargs(kwargs)
    tf.argmin(o; kwargs...) + 1
end

function max(o1::PyObject, o2::PyObject; kwargs...)
    tf.maximum(o1, o2; kwargs...)
end

function min(o1::PyObject, o2::PyObject; kwargs...)
    tf.minimum(o1, o2; kwargs...)
end

function maximum(o::PyObject; kwargs...)
    kwargs = jlargs(kwargs)
    tf.reduce_max(o; kwargs...)
end

function minimum(o::PyObject; kwargs...)
    kwargs = jlargs(kwargs)
    tf.reduce_min(o; kwargs...)
end

function std(o::PyObject; kwargs...)
    kwargs = jlargs(kwargs)
    tf.math.reduce_std(o; kwargs...)
end

function cast(x::PyObject, dtype::Type;kwargs...)
    dtype = DTYPE[dtype]
    tf.cast(x, dtype; kwargs...)
end

function cast(dtype::Type, x::PyObject;kwargs...)
    dtype = DTYPE[dtype]
    tf.cast(x, dtype; kwargs...)
end

softplus(x::Real) = ifelse(x > 0, x + log1p(exp(-x)), log1p(exp(x)))
function softplus(x;kwargs...)
    tf.math.softplus(x; kwargs...)
end

function log(o::PyObject; kwargs...)
    tf.math.log(o; kwargs...)
end

function exp(o::PyObject; kwargs...)
    tf.exp(o; kwargs...)
end

function cos(o::PyObject; kwargs...)
    tf.cos(o; kwargs...)
end

function sin(o::PyObject; kwargs...)
    tf.sin(o; kwargs...)
end

function sign(o::PyObject; kwargs...)
    tf.sign(o; kwargs...)
end

function softmax(o::PyObject; kwargs...)
    kwargs = jlargs(kwargs)
    tf.math.softmax(o; kwargs...)
end

function sum(o::PyObject; kwargs...)
    kwargs = jlargs(kwargs)
    tf.reduce_sum(o; kwargs...)
end

function mean(o::PyObject; kwargs...)
    kwargs = jlargs(kwargs)
    tf.reduce_mean(o; kwargs...)
end

function prod(o::PyObject; kwargs...)
    kwargs = jlargs(kwargs)
    tf.reduce_prod(o; kwargs...)
end
function squeeze(o::PyObject; kwargs...)
    kwargs = jlargs(kwargs)
    tf.squeeze(o;kwargs...)
end

"""
    pad(o::PyObject, paddings::Array{Int64, 2}, args...; kwargs...)

Pads `o` with values on the boundary.
# Example
```julia
o = rand(3,3)
o = pad(o, [1 4      # first dimension
             2 3])   # second dimension
run(sess, o)
```
Expected:
```
8×8 Array{Float64,2}:
 0.0  0.0  0.0       0.0       0.0       0.0  0.0  0.0
 0.0  0.0  0.250457  0.666905  0.823611  0.0  0.0  0.0
 0.0  0.0  0.23456   0.625145  0.646713  0.0  0.0  0.0
 0.0  0.0  0.552415  0.226417  0.67802   0.0  0.0  0.0
 0.0  0.0  0.0       0.0       0.0       0.0  0.0  0.0
 0.0  0.0  0.0       0.0       0.0       0.0  0.0  0.0
 0.0  0.0  0.0       0.0       0.0       0.0  0.0  0.0
 0.0  0.0  0.0       0.0       0.0       0.0  0.0  0.0
```
"""
function pad(o::Union{Array{<:Real}, PyObject}, paddings::Array{Int64, 2}, args...; kwargs...)
    o = constant(o)
    tf.pad(o, paddings, args...; kwargs...)
end

function assign(o::PyObject,value, args...; kwargs...)
    tf.compat.v1.assign(o, value, args...;kwargs...)
end

function group(args...; kwargs...)
    tf.group(args...; kwargs...)
end

assign(o::Array{PyObject}, value::Array, args...;kwargs...) = group_assign(o, value, args...; kwargs...)

@deprecate group_assign assign
function group_assign(os::Array{PyObject}, values, args...; kwargs...)
    ops = Array{PyObject}(undef, length(os))
    for i = 1:length(os)
        ops[i] = tf.compat.v1.assign(os[i], values[i], args...; kwargs...)
    end
    ops
end

"""
    adjoint(o::PyObject; kwargs...)

Returns the conjugate adjoint of `o`.
When the dimension of `o` is greater than 2, only the last two dimensions are permuted, i.e., `permutedims(o, [1,2,...,n,n-1])`
"""
function adjoint(o::PyObject; kwargs...)
    if length(size(o))==0
        return o
    elseif length(size(o))==1
        return rvec(o)
    else
        return tf.linalg.adjoint(o; kwargs...)
    end
end
diagm(o::PyObject; kwargs...) = tf.linalg.diag(o; kwargs...)
diag(o::PyObject; kwargs...) = tf.linalg.diag_part(o; kwargs...)
det(o::PyObject; kwargs...) = tf.linalg.det(o; kwargs...)
inv(o::PyObject; kwargs...) = tf.linalg.inv(o; kwargs...)

@doc raw"""
    solve_batch(A::Union{PyObject, Array{<:Real, 2}}, rhs::Union{PyObject, Array{<:Real,2}})

Solves $$Ax = b$$ for a batch of right hand sides.

- `A`: a $m\times n$ matrix, where $m\geq n$
- `rhs`: a $n_b\times m$ matrix. Each row is a new right hand side to solve.

The returned value is a $n_b\times n$ matrix.

# Example
```julia
a = rand(10,5)
b = rand(100, 10)
sol = solve_batch(a, b)
@assert run(sess, sol) ≈ (a\b')'
```

!!! note
    Internally, the matrix $A$ is factorized first and then the factorization is used to solve multiple right hand side.
"""
function solve_batch(A::Union{PyObject, Array{<:Real, 2}}, rhs::Union{PyObject, Array{<:Real,2}})
    solve_batched_rhs_ = load_system_op("solve_batched_rhs"; multiple=false)
    a,rhs = convert_to_tensor([A,rhs], [Float64,Float64])
    sol = solve_batched_rhs_(a,rhs)
    if size(a, 2)!=nothing && size(rhs,1) != nothing
        sol = set_shape(sol, (size(rhs,1), size(a,2)))
    end
    return sol
end

function solve(matrix, rhs; kwargs...)
    rhs = constant(rhs)
    matrix = constant(matrix)
    flag = false
    if length(size(rhs))==1
        flag = true
        rhs = reshape(rhs, size(rhs, 1), 1)
    end
    if size(matrix,1)==size(matrix,2)
        ret = tf.linalg.solve(matrix, rhs; kwargs...)
    else
        # @show matrix, rhs
        ret = tf.linalg.lstsq(matrix, rhs;kwargs...)
    end
    if flag
        ret = squeeze(ret, dims=2)
    end
    return ret
end

Base.:\(o1::PyObject, o2::PyObject) = solve(o1, o2)
Base.:\(o1::PyObject, o2::Array) = solve(o1, o2)
Base.:\(o1::Array, o2::PyObject) = solve(o1, o2)

function triangular_solve(matrix, rhs; kwargs...)
    flag = false
    if length(size(rhs))==1
        flag = true
        rhs = reshape(rhs, size(rhs, 1), 1)
    end
    ret = tf.linalg.triangular_solve(matrix, rhs; kwargs...)
    if flag
        ret = squeeze(ret, dims=2)
    end
    return ret
end

# reference: https://blog.csdn.net/LoseInVain/article/details/79638183
function concat(o::Union{PyObject,Array{PyObject}}, args...;kwargs...)
    if isa(o, PyObject)
        @warn "Only one input is consumed by concat" maxlog=1
        return o
    end
    kwargs = jlargs(kwargs)
    if length(size(o[1]))==0
        return tf.stack(o)
    end
    tf.concat(o, args...; kwargs...)
end

function stack(o::Array{PyObject}, args...;kwargs...)
    kwargs = jlargs(kwargs)
    tf.stack(o, args...; kwargs...)
end

Base.:vcat(args::PyObject...) = concat([args...],0)
function Base.:hcat(args::PyObject...)
    if length(size(args[1]))>=2
        concat([args...],1)
    elseif length(size(args[1]))==1
        stack([args...],dims=2)
    else
        vcat(args...)'
    end
end

"""
    stack(o::PyObject)

Convert a `TensorArray` `o` to a normal tensor. The leading dimension is the size of the tensor array.
"""
function stack(o::PyObject)
    o.stack()
end

function unstack(o::PyObject, args...;kwargs...)
    kwargs = jlargs(kwargs)
    tf.unstack(o, args...; kwargs...)
end

function _jlindex2indices(len::Int64,
    indices::Union{PyObject, Colon, Int64, Array{Int64}, BitArray{1}, Array{Bool,1}, UnitRange{Int64}, StepRange{Int64, Int64}})
    if isa(indices, BitArray{1}) || isa(indices, Array{Bool,1})
        indices = findall(indices)
    elseif isa(indices, UnitRange{Int64}) || isa(indices, StepRange{Int64, Int64})
        indices = collect(indices)
    elseif isa(indices, Colon)
        indices = Array(1:len)
    elseif isa(indices, Int64)
        indices = [indices]
    end
    if isa(indices, PyObject)
        return indices - 1
    else
        return constant(indices .- 1)
    end
end

for (op1, op2) = [(:_scatter_add, :tensor_scatter_nd_add), (:_scatter_sub, :tensor_scatter_nd_sub),
                    (:_scatter_update,:tensor_scatter_nd_update)]
    @eval begin
        function $op1(ref::PyObject,
            indices::Union{Colon, Int64, Array{Int64}, BitArray{1}, Array{Bool,1}, UnitRange{Int64}, StepRange{Int64, Int64}, PyObject},
            updates::Union{Array{<:Real}, Real, PyObject})
            updates = convert_to_tensor(updates, dtype=get_dtype(ref))
            @assert length(size(updates)) <= 1
            @assert length(size(ref))==1
            if length(size(updates))==0
                updates = reshape(updates, (-1,))
            end
            indices = _jlindex2indices(length(ref),indices)
            indices = reshape(indices, (-1,1))
            # @info ref, indices, updates
            tf.$op2(ref, indices, updates)
        end
    end
end

"""
    scatter_update(a::PyObject,
        indices::Union{Colon, Int64, Array{Int64}, BitArray{1}, Array{Bool,1}, UnitRange{Int64}, StepRange{Int64, Int64}, PyObject},
        updates::Union{Array{<:Real}, Real, PyObject})

Updates array `a`
```
a[indices] = updates
```

# Example
Julia:
```julia
A[[1;2;3]] = rand(3)
A[2] = 1.0
```

ADCME:
```
A = scatter_update(A, [1;2;3], rand(3))
A = scatter_update(A, 2, 1.0)
```
"""
scatter_update(a::PyObject,
    indices::Union{Colon, Int64, Array{Int64}, BitArray{1}, Array{Bool,1}, UnitRange{Int64}, StepRange{Int64, Int64}, PyObject},
    updates::Union{Array{<:Real}, Real, PyObject}) = _scatter_update(a, indices, updates)

"""
    scatter_sub(a::PyObject,
        indices::Union{Colon, Int64, Array{Int64}, BitArray{1}, Array{Bool,1}, UnitRange{Int64}, StepRange{Int64, Int64}, PyObject},
        updates::Union{Array{<:Real}, Real, PyObject})

Updates array `a`
```
a[indices] -= updates
```

# Example
Julia:
```julia
A[[1;2;3]] -= rand(3)
A[2] -= 1.0
```

ADCME:
```
A = scatter_sub(A, [1;2;3], rand(3))
A = scatter_sub(A, 2, 1.0)
```
"""
scatter_sub(a::PyObject,
    indices::Union{Colon, Int64, Array{Int64}, BitArray{1}, Array{Bool,1}, UnitRange{Int64}, StepRange{Int64, Int64}, PyObject},
    updates::Union{Array{<:Real}, Real, PyObject}) = _scatter_sub(a, indices, updates)

"""
    scatter_add(a::PyObject,
        indices::Union{Colon, Int64, Array{Int64}, BitArray{1}, Array{Bool,1}, UnitRange{Int64}, StepRange{Int64, Int64}, PyObject},
        updates::Union{Array{<:Real}, Real, PyObject})

Updates array `add`
```
a[indices] += updates
```

# Example
Julia:
```julia
A[[1;2;3]] += rand(3)
A[2] += 1.0
```

ADCME:
```
A = scatter_add(A, [1;2;3], rand(3))
A = scatter_add(A, 2, 1.0)
```
"""
scatter_add(a::PyObject,
    indices::Union{Colon, Int64, Array{Int64}, BitArray{1}, Array{Bool,1}, UnitRange{Int64}, StepRange{Int64, Int64}, PyObject},
    updates::Union{Array{<:Real}, Real, PyObject}) = _scatter_add(a, indices, updates)

for (op1, op2) = [(:scatter_update2, :scatter_update), (:scatter_add2, :scatter_add),
    (:scatter_sub2,:scatter_sub)]
    @eval begin
        function $op1(A::PyObject,
            xind::Union{Colon, Int64, Array{Int64}, BitArray{1}, Array{Bool,1}, UnitRange{Int64}, StepRange{Int64, Int64}, PyObject},
            yind::Union{Colon, Int64, Array{Int64}, BitArray{1}, Array{Bool,1}, UnitRange{Int64}, StepRange{Int64, Int64}, PyObject},
            updates::Union{Array{<:Real}, Real, PyObject})
            m, n = size(A)
            updates = convert_to_tensor(updates, dtype=get_dtype(A))
            @assert length(size(A))==2
            xind = _jlindex2indices(m,xind)
            yind = _jlindex2indices(n,yind)
            if length(size(updates))==0
                updates = reshape(updates, (1,1))
            end
            indices = reshape(repeat(xind, 1, length(yind)), (-1,)) * n + repeat(yind, length(xind))
            out = $op2(reshape(A, (-1,)), indices + 1, reshape(updates, (-1,)))
            reshape(out, (m, n))
        end
    end
end

"""
    scatter_update(A::PyObject,
        xind::Union{Colon, Int64, Array{Int64}, BitArray{1}, Array{Bool,1}, UnitRange{Int64}, StepRange{Int64, Int64}, PyObject},
        yind::Union{Colon, Int64, Array{Int64}, BitArray{1}, Array{Bool,1}, UnitRange{Int64}, StepRange{Int64, Int64}, PyObject},
        updates::Union{Array{<:Real}, Real, PyObject})

```julia
A[xind, yind] = updates
```
"""
scatter_update(A::PyObject,
    xind::Union{Colon, Int64, Array{Int64}, BitArray{1}, Array{Bool,1}, UnitRange{Int64}, StepRange{Int64, Int64}, PyObject},
    yind::Union{Colon, Int64, Array{Int64}, BitArray{1}, Array{Bool,1}, UnitRange{Int64}, StepRange{Int64, Int64}, PyObject},
    updates::Union{Array{<:Real}, Real, PyObject}) = scatter_update2(A, xind, yind, updates)

"""
    scatter_add(A::PyObject,
        xind::Union{Colon, Int64, Array{Int64}, BitArray{1}, Array{Bool,1}, UnitRange{Int64}, StepRange{Int64, Int64}, PyObject},
        yind::Union{Colon, Int64, Array{Int64}, BitArray{1}, Array{Bool,1}, UnitRange{Int64}, StepRange{Int64, Int64}, PyObject},
        updates::Union{Array{<:Real}, Real, PyObject})

```julia
A[xind, yind] += updates
```
"""
scatter_add(A::PyObject,
    xind::Union{Colon, Int64, Array{Int64}, BitArray{1}, Array{Bool,1}, UnitRange{Int64}, StepRange{Int64, Int64}, PyObject},
    yind::Union{Colon, Int64, Array{Int64}, BitArray{1}, Array{Bool,1}, UnitRange{Int64}, StepRange{Int64, Int64}, PyObject},
    updates::Union{Array{<:Real}, Real, PyObject}) = scatter_add2(A, xind, yind, updates)

"""
    scatter_add(A::PyObject,
        xind::Union{Colon, Int64, Array{Int64}, BitArray{1}, Array{Bool,1}, UnitRange{Int64}, StepRange{Int64, Int64}, PyObject},
        yind::Union{Colon, Int64, Array{Int64}, BitArray{1}, Array{Bool,1}, UnitRange{Int64}, StepRange{Int64, Int64}, PyObject},
        updates::Union{Array{<:Real}, Real, PyObject})

```julia
A[xind, yind] -= updates
```
"""
scatter_sub(A::PyObject,
    xind::Union{Colon, Int64, Array{Int64}, BitArray{1}, Array{Bool,1}, UnitRange{Int64}, StepRange{Int64, Int64}, PyObject},
    yind::Union{Colon, Int64, Array{Int64}, BitArray{1}, Array{Bool,1}, UnitRange{Int64}, StepRange{Int64, Int64}, PyObject},
    updates::Union{Array{<:Real}, Real, PyObject}) = scatter_sub2(A, xind, yind, updates)

function norm(o::PyObject, args...;kwargs...)
    kwargs = jlargs(kwargs)
    tf.norm(o, args...;kwargs...)
end

function Base.:diff(o::PyObject; dims::Union{Int64,Nothing}=nothing)
    if dims==nothing
        if length(size(o))!=1
            error("expect rank=1")
        end
        return o[2:end]-o[1:end-1]
    elseif length(size(o))==2
        if dims==1
            return o[2:end,:]-o[1:end-1,:]
        elseif dims==2
            return o[:,2:end]-o[:,1:end-1]
        end
    else
        error("Arguments not understood")
    end
end

clip(o::PyObject, vmin, vmax, args...;kwargs...) = tf.clip_by_value(o, vmin, vmax, args...;kwargs...)

@doc raw"""
    clip(o::Union{Array{Any}, Array{PyObject}}, vmin, vmax, args...;kwargs...)

Clips the values of `o` to the range [`vmin`, `vmax`]

# Example
```julia
a = constant(3.0)
a = clip(a, 1.0, 2.0)
b = constant(rand(3))
b = clip(b, 0.5, 1.0)
```
"""
function clip(o::Union{Array{Any}, Array{PyObject}}, vmin, vmax, args...;kwargs...)
    out = Array{PyObject}(undef, length(o))
    for i = 1:length(o)
        out[i] = clip(o[i], vmin, vmax, args...;kwargs...)
    end
    out
end

function fft(o::PyObject, args...; kwargs...)
    if length(size(o))==1
        tf.fft(o, args...; kwargs...)
    elseif length(size(o))==2
        tf.fft2d(o, args...; kwargs...)
    elseif length(size(o))==3
        tf.fft3d(o, args...; kwargs...)
    else
        error("FFT for d>=4 not supported")
    end
end

# mimic the Julia SVD
struct TFSVD
    S::PyObject
    U::PyObject
    V::PyObject
    Vt::PyObject
end

"""
    svd(o::PyObject, args...; kwargs...)

Returns a `TFSVD` structure which holds the following data structures
```julia
S::PyObject
U::PyObject
V::PyObject
Vt::PyObject
```
We have the equality
``o = USV'``

# Example
```julia
A = rand(10,20)
r = svd(constant(A))
A2 = r.U*diagm(r.S)*r.Vt # The value of `A2` should be equal to `A`
```
"""
function svd(o::PyObject, args...; kwargs...)
    s,u,v = tf.linalg.svd(o)
    TFSVD(s, u, v, v')
end

function ifft(o::PyObject, args...; kwargs...)
    if length(size(o))==1
        tf.ifft(o, args...; kwargs...)
    elseif length(size(o))==2
        tf.ifft2d(o, args...; kwargs...)
    elseif length(size(o))==3
        tf.ifft3d(o, args...; kwargs...)
    else
        error("IFFT for d>=4 not supported")
    end
end

function Base.:real(o::PyObject, args...; kwargs...)
    tf.real(o, args...; kwargs...)
end

function Base.:imag(o::PyObject, args...; kwargs...)
    tf.imag(o, args...; kwargs...)
end

@doc raw"""
    map(fn::Function, o::Union{Array{PyObject},PyObject};
    kwargs...)

Applies `fn` to each element of `o`.
- `o`∈`Array{PyObject}` : returns `[fn(x) for x in o]`
- `o`∈PyObject : splits `o` according to the first dimension and then applies `fn`.

# Example
```julia
a = constant(rand(10,5))
b = map(x->sum(x), a) # equivalent to `sum(a, dims=2)`
```

!!! note
    If `fn` is a multivariate function, we need to specify the output type using `dtype` keyword. For example,
    ```julia
    a = constant(ones(10))
    b = constant(ones(10))
    fn = x->x[1]+x[2]
    c = map(fn, [a, b], dtype=Float64)
    ```
"""
function map(fn::Function, o::Union{Array{PyObject},PyObject};
         kwargs...)
    # if `o` is not a tensorflow tensor, roll back to normal `map`
    if (isa(o, Array) && !hasproperty(o[1], :graph)) ||
        (isa(o, PyObject) && !hasproperty(o, :graph))
        return fn.(o)
    end
    kwargs = jlargs(kwargs)
    tf.map_fn(fn, o;kwargs...)
end

"""
    pmap(fn::Function, o::Union{Array{PyObject}, PyObject})

Parallel for loop. There should be no data dependency between different iterations.

# Example
```julia
x = constant(ones(10))
y1 = pmap(x->2.0*x, x)
y2 = pmap(x->x[1]+x[2], [x,x])
y3 = pmap(1:10, x) do z
    i = z[1]
    xi = z[2]
    xi + cast(Float64, i)
end
run(sess, y1)
run(sess, y2)
run(sess, y3)
```
"""
function pmap(fn::Function, o::Union{Array{PyObject}, PyObject})
    tf.compat.v1.vectorized_map(fn, o)
end

function pmap(fn::Function, range_::Union{Array{Int64},UnitRange{Int64},StepRange{Int64}, PyObject},
        o::Union{Array{PyObject}, PyObject})
    ipt = convert_to_tensor(collect(range_))
    if isa(o, PyObject)
        return tf.compat.v1.vectorized_map(fn, [ipt,o])
    end
    if length(o)==0
        return tf.compat.v1.vectorized_map(fn, ipt)
    end
    if length(o)>0
        tf.compat.v1.vectorized_map(fn, [ipt, o...])
    end
end

dot(x::PyObject, y::PyObject) = sum(x.*y)
dot(x::PyObject, y::AbstractArray{<:Real}{<:Real}) = sum(x.*constant(y))
dot(x::AbstractArray{<:Real}{<:Real}, y::PyObject) = sum(constant(x).*y)

import PyCall: +, -
function +(o::PyObject, I::UniformScaling{Bool})
    @assert size(o,1)==size(o,2)
    o + diagm(0=>ones(size(o,1)))
end

function -(o::PyObject, I::UniformScaling{Bool})
    @assert size(o,1)==size(o,2)
    o - diagm(0=>ones(size(o,1)))
end

Base.:+(I::UniformScaling{Bool}, o::PyObject) = o+I
Base.:-(I::UniformScaling{Bool}, o::PyObject) = -(o-I)

function Base.:findall(o::PyObject)
    if !(length(size(o)) in [1,2])
        error("ADCME: input tensor must have rank 1 or 2")
    end
    if !(eltype(o) <: Bool)
        error("ADCME: input tensor must have boolean types")
    end
    if length(size(o))==2
        tf.compat.v2.where(o) + 1
    else
        o = reshape(o, :, 1)
        res = findall(o)
        res'[1,:]
    end
end

"""
    vector(i::Union{Array{T}, PyObject, UnitRange, StepRange}, v::Union{Array{Float64},PyObject},s::Union{Int64,PyObject})

Returns a vector `V` with length `s` such that
```
V[i] = v
```
"""
function vector(i::Union{Array{T}, PyObject, UnitRange, StepRange}, v::Union{Array{Float64},PyObject},
                 s::Union{Int64,PyObject}) where T<:Integer
    if isa(i, UnitRange) || isa(i, StepRange)
        i = collect(i)
    end
    i = convert_to_tensor(i)
    s = convert_to_tensor(s)
    i = reshape(i - 1,:,1)
    s = reshape(s, 1)
    tf.scatter_nd(i, v, s)
end

function Base.:repeat(o::PyObject, i::Int64, j::Int64)
    if length(size(o))==0
        return o*ones(eltype(o), i, j)
    end
    if length(size(o))==1
        o = reshape(o, :, 1)
    end
    if length(size(o))!=2
        error("ADCME: size of `o` must be 0, 1, 2")
    end
    tf.tile(o, (i,j))
end

function Base.:repeat(o::PyObject, i::Int64)
    if length(size(o))==0
        o * ones(eltype(o),i)
    elseif length(size(o))==1
        squeeze(repeat(o, i, 1))
    elseif length(size(o))==2
        repeat(o, i, 1)
    else
        error("ADCME: rank of `o` must be 0, 1, 2")
    end
end

function lgamma(o::Union{T, PyObject}, args...;kwargs...) where T<:Real
    tf.math.lgamma(convert_to_tensor(o), args...;kwargs...)
end

@doc raw"""
    Base.:sort(o::PyObject;
    rev::Bool=false, dims::Integer=-1, name::Union{Nothing,String}=nothing)

Sort a multidimensional array `o` along the given dimension.
- `rev`: `true` for DESCENDING and `false` (default) for ASCENDING
- `dims`: `-1` for last dimension.
"""
function Base.:sort(o::PyObject;
    rev::Bool=false, dims::Integer=-1, name::Union{Nothing,String}=nothing)
    direction = rev == false ? "ASCENDING" : "DESCENDING"
    axis = -1
    if dims!=-1
        axis = dims - 1
    end
    tf.compat.v1.sort(o, axis=axis, direction=direction, name = name)
end

"""
    topk(o::PyObject, k::Union{PyObject,Integer}=1;
        sorted::Bool=true, name::Union{Nothing,String}=nothing)

Finds values and indices of the `k` largest entries for the last dimension.
If `sorted=true` the resulting k elements will be sorted by the values in descending order.
"""
function topk(o::PyObject, k::Union{PyObject,Integer}=1;
        sorted::Bool=true, name::Union{Nothing,String}=nothing)
    k = convert_to_tensor(k, dtype=Int32)
    tf.compat.v1.math.top_k(o, k, sorted=sorted, name = name)
end

"""
    argsort(o::PyObject;
    stable::Bool = false, rev::Bool=false, dims::Integer=-1, name::Union{Nothing,String}=nothing)

Returns the indices of a tensor that give its sorted order along an axis.
"""
function argsort(o::PyObject;
    stable::Bool = false, rev::Bool=false, dims::Integer=-1, name::Union{Nothing,String}=nothing)
    direction = rev == false ? "ASCENDING" : "DESCENDING"
    axis = -1
    if dims!=-1
        axis = dims + 1
    end
    tf.compat.v1.argsort(
        o,
        axis=axis,
        direction=direction,
        stable=stable,
        name=name
    ) + 1
end

function tr(o::PyObject)
    tf.linalg.trace(o)
end

function tril(o::PyObject, num::Int64 = 0)
    flag = false
    if length(size(o))==2
        flag = true
        o = tf.expand_dims(o, 0)
    end
    tri_lu_ = load_system_op("tri_lu"; multiple=false)
    u,num,lu = convert_to_tensor([o,num,1], [Float64,Int64,Int64])
    out = tri_lu_(u,num,lu)
    if flag
        return out[1]
    else
        return out
    end
end

function triu(o::PyObject, num::Int64 = 0)
    flag = false
    if length(size(o))==2
        flag = true
        o = tf.expand_dims(o, 0)
    end
    tri_lu_ = load_system_op("tri_lu"; multiple=false)
    u,num,lu = convert_to_tensor([o,num,0], [Float64,Int64,Int64])
    out = tri_lu_(u,num,lu)
    if flag
        return out[1]
    else
        return out
    end
end

"""
    split(o::PyObject,
        num_or_size_splits::Union{Integer, Array{<:Integer}, PyObject}; kwargs...)

Splits `o` according to `num_or_size_splits`

# Example 1
```julia
a = constant(rand(10,8,6))
split(a, 5)
```
Expected output:
```
5-element Array{PyCall.PyObject,1}:
 PyObject <tf.Tensor 'split_5:0' shape=(2, 8, 6) dtype=float64>
 PyObject <tf.Tensor 'split_5:1' shape=(2, 8, 6) dtype=float64>
 PyObject <tf.Tensor 'split_5:2' shape=(2, 8, 6) dtype=float64>
 PyObject <tf.Tensor 'split_5:3' shape=(2, 8, 6) dtype=float64>
 PyObject <tf.Tensor 'split_5:4' shape=(2, 8, 6) dtype=float64>
```

# Example 2
```julia
a = constant(rand(10,8,6))
split(a, [4,3,1], dims=2)
```
Expected output:
```
3-element Array{PyCall.PyObject,1}:
 PyObject <tf.Tensor 'split_6:0' shape=(10, 4, 6) dtype=float64>
 PyObject <tf.Tensor 'split_6:1' shape=(10, 3, 6) dtype=float64>
 PyObject <tf.Tensor 'split_6:2' shape=(10, 1, 6) dtype=float64>
```

# Example 3
```julia
a = constant(rand(10,8,6))
split(a, 3, dims=3)
```
Expected output:
```
3-element Array{PyCall.PyObject,1}:
 PyObject <tf.Tensor 'split_7:0' shape=(10, 8, 2) dtype=float64>
 PyObject <tf.Tensor 'split_7:1' shape=(10, 8, 2) dtype=float64>
 PyObject <tf.Tensor 'split_7:2' shape=(10, 8, 2) dtype=float64>
```
"""
function Base.:split(o::PyObject,
    num_or_size_splits::Union{Integer, Array{<:Integer}, PyObject}; kwargs...)
    kwargs = jlargs(kwargs)
    tf.split(o, num_or_size_splits; kwargs...)
end

"""
    reverse(o::PyObject, kwargs...)

Given a tensor `o`, and an index `dims` representing the set of dimensions of tensor to reverse.

# Example
```julia
a = rand(10,2)
A = constant(a)
@assert run(sess, reverse(A, dims=1)) == reverse(a, dims=1)
@assert run(sess, reverse(A, dims=2)) == reverse(a, dims=2)
@assert run(sess, reverse(A, dims=-1)) == reverse(a, dims=2)
```
"""
function reverse(o::PyObject; dims::Integer = -1)
    if dims==-1
        dims = length(size(o))
    end
    dims = convert_to_tensor([dims-1], dtype=Int32)
    tf.reverse(o, dims)
end

function swish(x)
    return x*sigmoid(x)
end

function hard_swish(x)
    return x*sigmoid(100*x)
end

function hard_sigmoid(x)
    return min(1.0, maximum(0, x+0.5))
end

function concat_relu(x)
    relu([x -x])
end

function concat_elu(x)
    elu([x -x])
end

function concat_hard_swish(x)
    hard_swish([x -x])
end

function fourier(x, terms::Int64=10)
    list = []
    for i = 1:terms
        push!(list, sin(x))
        push!(list, cos(x))
    end
    hcat(list...)
end


using OrdinaryDiffEq, DiffEqBase, Test

## Start on Number
f = (u,p,t) -> (2u)
#(::typeof(f))(::Type{Val{:analytic}},u0,p,t) = u0*exp(t)
prob = ODEProblem(f,1/2+(1/4)im,(0.0,1.0))

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

u0 = rand(ComplexF32,5,5,5)
f = (du,u,p,t) -> (du.=2u)
#(::typeof(f))(::Type{Val{:analytic}},u0,p,t) = u0*exp(t)
prob = ODEProblem(f,u0,(0.0,1.0))

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

@test typeof(sol[1]) == Array{ComplexF32,3}


# olsvar.m
# Luca Brugnolini 2014
# Based on Kilian code

function olsvar(y,p)

# Input : data y, lag order p
# Output: coefficient estimate B, estimate of Cov(U), residual U

(T,K) = size(y);
Y = (y[p+1:T,:])';                     # LHS variable
X = ones(1,T-p);                       # constant term

for j=1:p                              # RHS variable
    Z = [y[p+1-j:T-j,:]]';
    X = [X;Z];

end

B = (Y*X')/(X*X');                     # OLS estimator
U = Y-B*X;                             # estimated residuals
SIGMA = U*U'/(T-p-p*K-1);              # estimated variance of disturbances

return B, SIGMA, U

end


import KernelDensity.UnivariateKDE
import Distributions.Distribution
import StatsBase: AbstractWeights, Weights
import Distributions

"""
    UncertainValue(x::T) where T <: Real

Create a `CertainValue` instance from a scalar with no uncertainty.
"""
UncertainValue(x::T) where T <: Real = CertainValue(x)

# Identity constructor
UncertainValue(uval::AbstractUncertainValue) = uval

# From Measurements.jl
UncertainValue(m::Measurement{T}) where T = UncertainValue(Normal, m.val, m.err)

"""
    UncertainValue(values::Vector{<:Number}, probs::Vector{<:Number})

From a numeric vector, construct an `UncertainPopulation` whose
members are scalar values.
"""
function UncertainValue(values::Vector{<:Number}, probs::Vector{<:Number})
    UncertainScalarPopulation(float.(values), probs)
end

"""
    UncertainValue(values::Vector{<:Number}, probs::Vector{<:Number})

From a numeric vector, construct an `UncertainPopulation` whose
members are scalar values.
"""
function UncertainValue(values::Vector{<:Number}, probs::W) where {W <: AbstractWeights}
    UncertainScalarPopulation(float.(values), probs)
end

"""
    UncertainValue(values::Vector, probs::Union{Vector, AbstractWeights})

Construct a population whose members are given by `values` and whose sampling
probabilities are given by `probs`. The elements of `values` can be either
numeric or uncertain values of any type.
"""
function UncertainValue(values::VT, probs) where VT <: Vector{ELTYPE} where {ELTYPE<:POTENTIAL_UVAL_TYPES}
    UncertainScalarPopulation(UncertainValue.(values), probs)
end

function UncertainValue(values::VT, probs::Vector{Number}) where VT <: Vector{ELTYPE} where {ELTYPE<:POTENTIAL_UVAL_TYPES}
    UncertainScalarPopulation(UncertainValue.(values), probs)
end

"""
    UncertainValue(data::Vector{T};
        kernel::Type{D} = Normal,
        npoints::Int=2048) where {D <: Distributions.Distribution, T}

Construct an uncertain value by a kernel density estimate to `data`.

Fast Fourier transforms are used in the kernel density estimation, so the
number of points should be a power of 2 (default = 2048).
"""
function UncertainValue(data::Vector{T};
        kernel::Type{D} = Normal,
        bandwidth = KernelDensity.default_bandwidth(data),
        npoints::Int = 2048) where {D <: Distributions.Distribution, T}

    # Kernel density estimation
    KDE = kde(data, npoints = npoints, kernel = kernel, bandwidth = bandwidth)

    # Get the x value for which the density is estimated.
    xrange = KDE.x

    # Normalise estimated density
    density = KDE.density ./ sum(KDE.density)

    # Create an uncertain value
    UncertainScalarKDE(KDE, data, xrange, Weights(density))
end

"""
    UncertainValue(kerneldensity::Type{K}, data::Vector{T};
        kernel::Type{D} = Normal,
        npoints::Int=2048) where {K <: UnivariateKDE, D <: Distribution, T}

Construct an uncertain value by a kernel density estimate to `data`.

Fast Fourier transforms are used in the kernel density estimation, so the
number of points should be a power of 2 (default = 2048).
"""
function UncertainValue(kerneldensity::Type{K}, data::Vector{T};
        kernel::Type{D} = Normal,
        bandwidth = KernelDensity.default_bandwidth(data)/4,
        npoints::Int = 2048) where {K <: UnivariateKDE, D <: Distribution, T}

    # Kernel density estimation
    KDE = kde(data, npoints = npoints, kernel = kernel, bandwidth = bandwidth)

    # Get the x value for which the density is estimated.
    xrange = KDE.x

    # Normalise estimated density
    density = KDE.density ./ sum(KDE.density)

    # Create an uncertain value
    UncertainScalarKDE(KDE, data, xrange, Weights(density))
end

# For vectors of zero-dimensional arrays.
UncertainValue(x::Vector{Array{<:Real, 0}}) = UncertainValue([el[] for el in x])

"""
    UncertainValue(empiricaldata::AbstractVector{T},
        d::Type{D}) where {D <: Distribution}

# Constructor for empirical distributions.

Fit a distribution of type `d` to the data and use that as the
representation of the empirical distribution. Calls `Distributions.fit` behind
the scenes.

## Arguments
- **`empiricaldata`**: The data for which to fit the `distribution`.
- **`distribution`**: A valid univariate distribution from `Distributions.jl`.

"""
function UncertainValue(d::Type{D},
        empiricaldata::Vector{T}) where {D<:Distribution, T}

    distribution = FittedDistribution(Distributions.fit(d, empiricaldata))
    UncertainScalarTheoreticalFit(distribution, empiricaldata)
end

"""

    UncertainValue(distribution::Type{D}, a::T1, b::T2;
        kwargs...) where {T1<:Number, T2 <: Number, D<:Distribution}

# Constructor for two-parameter distributions

`UncertainValue`s are currently implemented for the following two-parameter
distributions: `Uniform`, `Normal`, `Binomial`, `Beta`, `BetaPrime`, `Gamma`,
and `Frechet`.

### Arguments

- **`a`, `b`**: Generic parameters whose meaning varies depending
    on what `distribution` is provided. See the list below.
- **`distribution`**: A valid univariate distribution from `Distributions.jl`.

Precisely what  `a` and `b` are depends on which distribution is provided.

- `UncertainValue(Normal, μ, σ)` returns an `UncertainScalarNormallyDistributed` instance.
- `UncertainValue(Uniform, lower, upper)` returns an `UncertainScalarUniformlyDistributed` instance.
- `UncertainValue(Beta, α, β)` returns an `UncertainScalarBetaDistributed` instance.
- `UncertainValue(BetaPrime, α, β)` returns an `UncertainScalarBetaPrimeDistributed` instance.
- `UncertainValue(Gamma, α, θ)` returns an `UncertainScalarGammaDistributed` instance.
- `UncertainValue(Frechet, α, θ)` returns an `UncertainScalarFrechetDistributed` instance.
- `UncertainValue(Binomial, n, p)` returns an `UncertainScalarBinomialDistributed` instance.

### Keyword arguments

- **`nσ`**: If `distribution <: Distributions.Normal`, then how many standard
    deviations away from `μ` does `lower` and `upper` (i.e. both, because
    they are the same distance away from `μ`) represent?
- **`tolerance`**: A threshold determining how symmetric the uncertainties
    must be in order to allow the construction of  Normal distribution
    (`upper - lower > threshold` is required).
- **`trunc_lower`**: Lower truncation bound for distributions with infinite
    support. Defaults to `-Inf`.
- **`trunc_upper`**: Upper truncation bound for distributions with infinite
    support. Defaults to `Inf`.

## Examples

### Normal distribution

Normal distributions are formed by using the constructor
`UncertainValue(μ, σ, Normal; kwargs...)`. This gives a normal distribution with
mean μ and standard deviation σ/nσ (nσ must be given as a keyword argument).

```julia
# A normal distribution with mean = 2.3 and standard deviation 0.3.
UncertainValue(2.3, 0.3, Normal)

# A normal distribution with mean 2.3 and standard deviation 0.3/2.
UncertainValue(2.3, 0.3, Normal, nσ = 2)

# A normal distribution with mean 2.3 and standard deviation = 0.3,
truncated to the interval `[1, 3]`.
UncertainValue(2.3, 0.3, Normal, trunc_lower = 1.0, trunc_upper = 3.0)
```

### Uniform distribution

Uniform distributions are formed using the
`UncertainValue(lower, upper, Uniform)` constructor.

```julia
#  A uniform distribution on `[2, 3]`
UncertainValue(-2, 3, Uniform)
```

"""
function UncertainValue(distribution::Type{D}, a::T1, b::T2;
        kwargs...) where {T1<:Number, T2 <: Number, D<:Distribution}

    if distribution == Uniform
        dist = assigndist_uniform(a, b)
        UncertainScalarUniformlyDistributed(dist, a, b)
    elseif distribution == Binomial
        dist = assigndist_binomial(a, b)
        UncertainScalarBinomialDistributed(dist, a, b)
    elseif distribution == Normal
        dist = assigndist_normal(a, b; kwargs...)
        UncertainScalarNormallyDistributed(dist, a, b)
    elseif distribution == Beta
        dist = assigndist_beta(a, b; kwargs...)
        UncertainScalarBetaDistributed(dist, a, b)
    elseif distribution == BetaPrime
        dist = assigndist_betaprime(a, b; kwargs...)
        UncertainScalarBetaPrimeDistributed(dist, a, b)
    elseif distribution == Gamma
        if !((a > 0) & (b > 0))
            error("α and θ must both be > 0")
        end
        dist = assigndist_gamma(a, b; kwargs...)
        UncertainScalarGammaDistributed(dist, a, b)
    elseif distribution == Frechet
        dist = assigndist_frechet(a, b; kwargs...)
        UncertainScalarFrechetDistributed(dist, a, b)
    else
        throw(DomainError("Two-parameter $dist is not implemented."))
    end
end

"""
    UncertainValue(distribution::Type{D}, a::T1, b::T2, c::T3;
        kwargs...) where {T1<:Number, T2<:Number, T3<:Number, D<:Distribution}

## Constructor for three-parameter distributions

Currently implemented distributions are `BetaBinomial`.

### Arguments
- **`a`, `b`, `c`**: Generic parameters whose meaning varies depending
    on what `distribution` is provided. See the list below.
- **`distribution`**: A valid univariate distribution from `Distributions.jl`.

Precisely what `a`, `b` and `c` are depends on which distribution is provided.

- `UncertainValue(BetaBinomial, n, α, β)` returns an `UncertainScalarBetaBinomialDistributed` instance.

### Keyword arguments
- **`nσ`**: If `distribution <: Distributions.Normal`, then how many standard
    deviations away from `μ` does `lower` and `upper` (i.e. both, because
    they are the same distance away from `μ`) represent?
- **`tolerance`**: A threshold determining how symmetric the uncertainties
    must be in order to allow the construction of  Normal distribution
    (`upper - lower > threshold` is required).
- **`trunc_lower`**: Lower truncation bound for distributions with infinite
    support. Defaults to `-Inf`.
- **`trunc_upper`**: Upper truncation bound for distributions with infinite
    support. Defaults to `Inf`.

## Examples
### BetaBinomial distribution

Normal distributions are formed by using the constructor
`UncertainValue(μ, σ, Normal; kwargs...)`. This gives a normal distribution with
mean μ and standard deviation σ/nσ (nσ must be given as a keyword argument).

```julia
# A beta binomial distribution with n = 100 trials and parameters α = 2.3 and
# β = 5
UncertainValue(100, 2.3, 5, BetaBinomial)
```
"""
function UncertainValue(distribution::Type{D}, a::T1, b::T2, c::T3;
        kwargs...) where {T1<:Number, T2<:Number, T3<:Number, D<:Distribution}

    if distribution == BetaBinomial
        dist = assigndist_betabinomial(a, b, c; kwargs...)
        UncertainScalarBetaBinomialDistributed(dist, a, b, c)
    else
        throw(DomainError("Three-parameter $dist is not implemented."))
    end
end

"""
    nested_disttype(t::Distributions.Truncated)

Get the type of the untruncated distribution for a potentially nested
truncated distribution.
"""
function untruncated_dist(t::Distributions.Truncated)
    t_untrunc = t

    while typeof(t_untrunc) <: Distributions.Truncated
        t_untrunc = t_untrunc.untruncated
    end

    return t_untrunc
end

"""
    untruncated_disttype(t::Distributions.Truncated)

Get the type of the untruncated distribution for a potentially nested
truncated distribution.
"""
function untruncated_disttype(t::Distributions.Truncated)
    t_untrunc = t

    while typeof(t_untrunc) <: Distributions.Truncated
        t_untrunc = t_untrunc.untruncated
    end

    return typeof(t_untrunc)
end

"""
    UncertainValue(t::Distributions.Truncated)

Construct an uncertain value from an instance of a distribution. If a specific
uncertain value type has not been implemented, the number of parameters is
determined from the distribution and an instance of one of the following types
is returned:

- `ConstrainedUncertainScalarValueOneParameter`
- `ConstrainedUncertainScalarValueTwoParameter`
- `ConstrainedUncertainScalarValueThreeParameter`

## Examples

```julia
# Normal distribution truncated to the interval [0.5, 0.7]
t = truncated(Normal(0, 1), 0.5, 0.7)
UncertainValue(t)

# Gamma distribution truncated to the interval [0.5, 3.5]
t = Truncate(Gamma(4, 5.1), 0.5, 3.5)
UncertainValue(t)

# Binomial distribution truncated to the interval [2, 7]
t = Truncate(Binomial(10, 0.4), 2, 7)
UncertainValue(t)
```
"""
function UncertainValue(t::Distributions.Truncated)
    dist_type = untruncated_disttype(t)
    original_dist = untruncated_dist(t)
    params = fieldnames(dist_type)
    param_values = [getfield(original_dist, p) for p in params]

    n_params = length(params)
    if n_params == 1
        return ConstrainedUncertainScalarValueOneParameter(t, param_values...)
    elseif n_params == 2
        return ConstrainedUncertainScalarValueTwoParameter(t, param_values...)
    elseif n_params == 3
        return ConstrainedUncertainScalarValueThreeParameter(t, param_values...)
    end
end

"""
    UncertainValue(d::Distributions.Distribution)

Construct an uncertain value from an instance of a distribution. If a specific
uncertain value type has not been implemented, the number of parameters is
determined from the distribution and an instance of one of the following types
is returned:

- `UncertainScalarTheoreticalOneParameter`
- `UncertainScalarTheoreticalTwoParameter`
- `UncertainScalarTheoreticalThreeParameter`

## Examples

```julia
UncertainValue(Normal(0, 1))
UncertainValue(Gamma(4, 5.1))
UncertainValue(Binomial, 8, 0.2)
```
"""
function UncertainValue(d::Distributions.Distribution)
    params = fieldnames(typeof(d))
    n_params = length(params)
    param_values = [getfield(d, p) for p in params]

    if d isa Uniform
        UncertainScalarUniformlyDistributed(d, param_values...)
    elseif d isa Binomial
        UncertainScalarBinomialDistributed(d, param_values...)
    elseif d isa Normal
        UncertainScalarNormallyDistributed(d, param_values...)
    elseif d isa Beta
        UncertainScalarBetaDistributed(d, param_values...)
    elseif d isa BetaPrime
        UncertainScalarBetaPrimeDistributed(d, param_values...)
    elseif d isa Gamma
        UncertainScalarGammaDistributed(d, param_values...)
    elseif d isa Frechet
        UncertainScalarFrechetDistributed(d, param_values...)
    # if no specific type is implemented for this distribution, just create
    # a generic one
    else
        if n_params == 1
            return UncertainScalarTheoreticalOneParameter(d, param_values...)
        elseif n_params == 2
            return UncertainScalarTheoreticalTwoParameter(d, param_values...)
        elseif n_params == 3
            return UncertainScalarTheoreticalThreeParameter(d, param_values...)
        else
            msg = "uncertain value type for $n_params-parameter $d not implemented."
            throw(DomainError(msg))
        end
    end
end

##############################
# Macro constructors
##############################
# macro uncertainvalue(μ, lower, upper, d,
#         nσ = 2, trunc_lower = -Inf, trunc_upper = Inf, tolerance = 1e-3)
#     return :(UncertainValue($μ, $lower, $upper, $d,
#                 nσ = $nσ,
#                 trunc_lower = $trunc_lower,
#                 trunc_upper = $trunc_upper,
#                 tolerance = $tolerance))
# end
#
#macro uncertainvalue(empiricaldata, dist)
#    return :(UncertainValue($empiricaldata, $dist))
#end
#
# """
# A macro for the construction of uncertain values. Calls
#     [`UncertainValue`](@ref) with the provided arguments.
#
# - **`@uncertainvalue(μ, lower, upper, dist, kwargs...)`**:
#     Fit a distribution of type `dist` to `μ` and `lower`/`upper` uncertainty
#     bounds.
# - **`@uncertainvalue(empiricaldata, dist)`**:
#     Fit a distribution of type `dist` to a vector of empirical data.
# """
#:(@uncertainvalue)
#
#
# """
#     @evalue(values, d)
#
# Construct an uncertain value from an empirical distribution.
# """
#macro uncertain(values::AbstractVector, d)
#    :(UncertainScalarTheoreticalFit($values, $d))
#end
#uncertain, @uncertain

export
UncertainValue
#uncertainvalue,
#@uncertainvalue


"""
    catmullrom_by_arclength( points )
    catmullrom_by_arclength( points, (min_between_points, max_between_points))

Calculates a Catmull-Rom fit through given `points`, where the number of interpoint knots
varys inversely with the approximate arclength of that arc.

`arcpoints_min` and `arcpoints_max` refer any one arc between two points.
""" catmullrom_by_arclength

function catmullrom_by_arclength(points::P, arcpoints::Tuple{I,I}=(ArcpointsMin, ArcpointsMax); extend::Bool=true) where {P, I<:Integer}
    catmullrom_requirement(npoints(points))

    crpoints = deepcopy(points)
    if extend
        extend_seq(crpoints)
    end

    return catmullrombyarclength(crpoints, minmax(arcpoints...))
end

function catmullrombyarclength(points::P, arcpoints::Tuple{I,I}=(ArcpointsMin, ArcpointsMax)) where {P, I<:Integer}
    n_points = npoints(points)
    pointsperarc = arclength_interpolants(points, arcpoints)
    total_points = sum(pointsperarc) + n_points
    n_coords = ncoords(points)
    T = coordtype(points)
    result = []
    for i=1:n_coords
        push!(result, T[])
    end
    point_seqs = [points[i:i+3] for i=1:(length(points)-3)]

    n_points_3 = n_points - 3
    for idx=1:(n_points-3)
        fourpoints = point_seqs[idx]
        n_between  = pointsperarc[idx]
        fitted = catmullrom_splines(fourpoints, n_between)
        for j=1:n_coords
            append!(result[j], fitted[j][1:(end-1+(idx==n_points_3))])
        end
    end

    return result
end

"""
    arclength_interpolants(points; arcpoints_min=2, arcpoints_max=64)
Convert a sequence of points and a count of additional points to interpolate
to a sequence of arc-length specific point counts where each arc has at least
`arcpoints_min` and at most `arcpoints_max` points.
"""
function arclength_interpolants(points::P, arcpoints::Tuple{I,I}=(ArcpointsMin, ArcpointsMax)) where {P, I<:Integer}
    normalized_arclengths = normalized_catmullrom_arclengths(points)

    arcpoints_min, arcpoints_max = arcpoints
    arclength_min = minimum(normalized_arclengths)
    multiplier = nextfloat(arcpoints_min / arclength_min)
    arcpoints = floor.(Int, normalized_arclengths .* multiplier)
    arcpoints = arcpoints .+ isodd.(arcpoints)

    most_arcpoints = maximum(arcpoints)
    if most_arcpoints > arcpoints_max
        multiplier = prevfloat(arcpoints_max / (most_arcpoints + arcpoints_min))
        arcpoints  = floor.(Int, muladd.(arcpoints .- arcpoints_min, multiplier, arcpoints_min))
        arcpoints = arcpoints .+ isodd.(arcpoints)
    end

    return arcpoints
end

"""
    normalized_catmullrom_arclengths(points)
Convert a sequence of CatmullRom points, those given
as data and those interpolated through the data,
to a corresponding sequence of normalized arclengths.
The sum of the normalized arclengths is 1.0.  Each
arc's approximated length is a portion of unity.
The most probable arclength, and the mean of these
normalized arclengths are roughly equal; given by:
`1//number_of_arcs`.
The way to use normalized arclengths is to multiply
the fractional lengths by the number of "plottable"
points (the given data + the interpoint interpolants).
```
smallest_arc = minimum(normalized_arclengths)
rational_arc = rationalize(Float16(smallest_arc))
scalefactor = points_to_realize / denominator(rational_arc)
# scaled_arclengths_are_point_counts counting points on the arc
points_per_arc = round.(Int, normalized_arclengths .* scalefactor)
```
"""
function normalized_catmullrom_arclengths(points::P) where P
    arclengths = approx_catmullrom_arclengths(points)
    sum_of_arcs = sum(arclengths)
    arclengths[:] = arclengths ./ sum_of_arcs
    return arclengths
end

"""
    approx_catmullrom_arclengths(points)
Convert a sequence of CatmullRom points, those given
as data and those interpolated through the data,
to a corresponding sequence of arclength approximations.
"""
function approx_catmullrom_arclengths(points::P) where P
    L = float(coordtype(P))
    n_points = npoints(points)
    # the first two points and the last two points are boundary + anchor
    # boundary anchor first_internal_point ... final_internal_point anchor boundary
    #       noarc   arc                    arcs                   arc   noarc
    #    pt1     pt2     pt3                            ptN-2       ptN-1     ptN
    #                                   pt4 .. ptN-3
    #               1 arc               ((N-3)-4 arcs)          1 arc
    #  N-3-4+1+1 = N-1-4 = Npts-5 arcs
    n_arcs = n_points - 3
    result = Array{L, 1}(undef, n_arcs)

    # points[1:4], points[2:5], .. points[idx:idx+3] .., points[N-3:N]
    for idx = 1:n_points-3
        arc = approx_catmullrom_arclength(points[idx:idx+3]...,)
        result[idx] = arc
    end
    return result
end

"""
    approx_catmullrom_arclength(pt0, pt1, pt2, pt3)
Given 4 ND points along a centripetal Catmull-Rom span,
roughly approximate the arclength of the curvilinear segment
that would be determined by the two bounding points
and the tangents they determine [the arc between p1 and p2].
This well-behaved approximation was developed by Jens Gravesen
```
    (2*corddist + (n-1)*bezdist)/(n+1), n is degree of the curve
    deg=2 --> (2*corddist + bezdist)/(3)
    deg=3 --> (2*corddist + 2*bezdist)/(4) --> (corddist + bezdist)/2
    deg=4 --> (2*corddsit + 3*bezdist)/(5)
```
"""
function approx_catmullrom_arclength(p0::T, p1::T, p2::T, p3::T) where T
    b0, b1, b2, b3 = catmullrom_as_bezier(p0, p1, p2, p3)
    cordal_dist = norm(b3 .- b0)
    bezier_dist = norm(b1 .- b0) + norm(b2 .- b1) + norm(b3 .- b2)

    return (cordal_dist + bezier_dist) * 0.5
end

function catmullrom_as_bezier(p0::T, p1::T, p2::T, p3::T) where T
    b0 = p1
    b3 = p2

    d1 = norm(p1 .- p0); d1a = sqrt(d1)
    d2 = norm(p2 .- p1); d2a = sqrt(d2)
    d3 = norm(p3 .- p2); d3a = sqrt(d3)

    b1n = @. (d1 * p2) - (d2 * p0) + ((2*d1 + 3*d1a*d2a+d2) * p1)
    b1d = 3*d1a*(d1a+d2a)
    b1 = b1n ./ b1d
    b2n = @. (d3 * p1) - (d2 * p3) + ((2*d3 + 3*d3a*d2a+d2) * p2)
    b2d = 3*d3a*(d3a+d2a)
    b2 = b2n ./ b2d

    return b0, b1, b2, b3
end


# Invertible Hyperbolic network from Keegan et al (2018)
# Author: Gabrio Rizzuti, rizzuti.gabrio@gatech.edu
# Date: August 2020

export NetworkHyperbolic1D, clear_grad!, get_params

"""
    H = NetworkHyperbolic1D(nx, ny, n_in, batchsize, n_hidden, depth; k1=3, k2=3, p1=1, p2=1, s1=1, s2=1)
 Create an hyperbolic network for data-driven generative modeling based
 on the change of variables formula.
 *Input*:
 - `nx`, `ny`, `n_in`, `batchsize`: spatial dimensions, number of channels and batchsize of input tensors `X`
 - `n_hidden`: number of hidden units in residual blocks
 - `depth`: number network layers
 - `k1`, `k2`: kernel size for first and third residual layer (`k1`) and second layer (`k2`)
 - `p1`, `p2`: respective padding sizes for residual block layers
 - `s1`, `s2`: respective strides for residual block layers
 *Output*:
 - `H`: HINT network
 *Usage:*
 - Forward mode: `Z, logdet = H.forward(X)`
 - Inverse mode: `X = H.inverse(Z)`
 - Backward mode: `ΔX, X = H.backward(ΔZ, Z)`
 *Trainable parameters:*
 - None in `H` itself
 - Trainable parameters in activation normalizations `H.AN[i]`, and in coupling layers `H.CL[i]`, where `i` ranges from `1` to `depth`.
 See also: [`ActNorm`](@ref), [`CouplingLayerHINT`](@ref), [`get_params`](@ref), [`clear_grad!`](@ref)
"""
mutable struct NetworkHyperbolic1D <: InvertibleNetwork
    AN::AbstractArray{ActNorm, 1}
    CL::AbstractArray{HyperbolicLayer, 2}
    logdet::Bool
    is_reversed::Bool
end

@Flux.functor NetworkHyperbolic1D

# Constructor
function NetworkHyperbolic1D(n_in::Int64, batchsize::Int64, n_hidden::Int64, depth::Int64; logdet=true, depthHL=1)

    AN = Array{ActNorm}(undef, depth)
    CL = Array{HyperbolicLayer, 2}(undef, depth, depthHL)

    # Create layers
    for j = 1:depth
        AN[j] = ActNorm(n_in; logdet=logdet)
        for k = 1:depthHL
            CL[j, k] = HyperbolicLayer(1, 1, Int(n_in/2), batchsize, 1, 1, 0; action="same", α=1f0, hidden_factor=Int(n_hidden/n_in))
        end
    end

    return NetworkHyperbolic1D(AN, CL, logdet, false)
end

# Forward pass and compute logdet
function forward(X, H::NetworkHyperbolic1D; logdet=nothing)
    isnothing(logdet) ? logdet = (H.logdet && ~H.is_reversed) : logdet = logdet

    depth, depthHL = size(H.CL)
    logdet_ = 0f0
    for j=1:depth
        logdet ? (X, logdet1) = H.AN[j].forward(X) : X = H.AN[j].forward(X)
        X_prev, X_curr = tensor_split(X)
        for k = 1:depthHL
            X_prev, X_curr = H.CL[j, k].forward(X_prev, X_curr)
        end
        X = tensor_cat(X_prev, X_curr)
        logdet && (logdet_ += logdet1)
    end
    logdet ? (return X, logdet_) : (return X)
end

# Inverse pass and compute gradients
function inverse(Z, H::NetworkHyperbolic1D; logdet=nothing)
    isnothing(logdet) ? logdet = (H.logdet && H.is_reversed) : logdet = logdet

    depth, depthHL = size(H.CL)
    logdet_ = 0f0
    for j = depth:-1:1
        Z_curr, Z_new = tensor_split(Z)
        for k = depthHL:-1:1
            Z_curr, Z_new = H.CL[j, k].inverse(Z_curr, Z_new)
        end
        Z = tensor_cat(Z_curr, Z_new)
        logdet ? (Z, logdet2) = H.AN[j].inverse(Z; logdet=true) : Z = H.AN[j].inverse(Z; logdet=false)
        logdet && (logdet_ += (logdet1 + logdet2))
    end
    logdet ? (return Z, logdet_) : (return Z)
end

# Backward pass and compute gradients
function backward(ΔZ, Z, H::NetworkHyperbolic1D)
    depth, depthHL = size(H.CL)
    for j = depth:-1:1
        ΔZ_curr, ΔZ_new = tensor_split(ΔZ)
        Z_curr, Z_new = tensor_split(Z)
        for k = depthHL:-1:1
            ΔZ_curr, ΔZ_new, Z_curr, Z_new = H.CL[j, k].backward(ΔZ_curr, ΔZ_new, Z_curr, Z_new)
        end
        ΔZ = tensor_cat(ΔZ_curr, ΔZ_new)
        Z = tensor_cat(Z_curr, Z_new)
        ΔZ, Z = H.AN[j].backward(ΔZ, Z)
    end
    return ΔZ, Z
end

# Clear gradients
function clear_grad!(H::NetworkHyperbolic1D)
    depth, depthHL = size(H.CL)
    for j=1:depth
        clear_grad!(H.AN[j])
        for k = 1:depthHL
            clear_grad!(H.CL[j, k])
        end
    end
end

# Get parameters
function get_params(H::NetworkHyperbolic1D)
    depth, depthHL = size(H.CL)
    p = []
    for j=1:depth
        p = cat(p, get_params(H.AN[j]); dims=1)
        for k = 1:depthHL
            p = cat(p, get_params(H.CL[j, k]); dims=1)
        end
    end
    return p
end

# Set is_reversed flag in full network tree
function tag_as_reversed!(H::NetworkHyperbolic1D, tag::Bool)
    depth, depthHL = size(H.CL)
    H.is_reversed = tag
    for j=1:depth
        tag_as_reversed!(H.AN[j], tag)
        for k = 1:depthHL
            tag_as_reversed!(H.CL[j, k], tag)
        end
    end
    return H
end


module ComplementaryDistributions

using Distributions
using Random
using SpecialFunctions
using StatsFuns: twoπ, halfπ, inv2π, fourinvπ

export GeneralizedInverseGaussian, PolyaGamma, LaplaceTransformDistribution
include("generalizedinversegaussian.jl")
include("polyagamma.jl")
include("lap_transf_dist.jl")

end


#=
Here I show the exact correspondence
my spiking implementation and code by
Litwin-Kumar / Auguste Schulz
For one I neuron connected to one E neuron
=#

push!(LOAD_PATH, abspath(@__DIR__,".."))

using LinearAlgebra,Statistics,StatsBase,Distributions
using Plots,NamedColors ; theme(:dark)
using SparseArrays
using SpikingRNNs; const global S = SpikingRNNs

function onesparsemat(w::Real)
  return sparse(fill(w,(1,1)))
end

##

# this is a bare-bone version of the Schulz code
function runsimulation_static_simplified2(Ne::Integer,Ni::Integer,
  weights::Array{Float64,2},spiketimes::Array{Int32,2}; dt = 0.1, T = 2000)

		# Naming convention
		# e corresponds to excitatory
		# i corresponds to inhibitory
		# x corresponds to external

		#membrane dynamics
		taue = 30 #e membrane time constant ms
		taui = 20 #i membrane time constant ms
		vreste = -60 #e resting potential mV
		vresti = -70 #i resting potential mV
		vpeak = 20 #cutoff for voltage.  when crossed, record a spike and reset mV
		eifslope = 2 #eif slope parameter mV
		C = 300 #capacitance pF
		erev = 0 #e synapse reversal potential mV
		irev = -75 #i synapse reversal potntial mV
		vth0 = -52 #initial spike voltage threshold mV
		thrchange = false # can be switched off to have vth constant at vth0
		ath = 10 #increase in threshold post spike mV
		tauth = 30 #threshold decay timescale ms
		vreset = -70 #reset potential mV
		taurefrac = 50 #absolute refractory period ms

		# total number of neurons
		Ncells = Ne+Ni

    # constant input
    in_const = fill(0.0,Ncells)
    in_const[2] = -51.0 - vresti

		# synaptic kernel
		tauerise = 1 #e synapse rise time
		tauedecay = 6 #e synapse decay time
		tauirise = .5 #i synapse rise time
		tauidecay = 2 #i synapse decay time

		# Initialisation of spike arrays ------------------------------
		# spike count and spike times
		# Scope problems - changing arrays are therefore defined here directly

		totalspikes = zeros(Int,Ncells)
		totsp::Int64 = 0; # total numberof spikes
		spmax::Int64 = size(spiketimes,1); # maximum recorded spikes

		# further arrays for storing inputs and synaptic integration
		forwardInputsE = zeros(Float64,Ncells) #sum of all incoming weights from excitatpry inputs both external and EE and IE
		forwardInputsI = zeros(Float64,Ncells) #sum of all incoming weights from inhibitory inputs both II and EI
		forwardInputsEPrev = zeros(Float64,Ncells) #as above, for previous timestep
		forwardInputsIPrev = zeros(Float64,Ncells)

		xerise = zeros(Float64,Ncells) #auxiliary variables for E/I currents (difference of exponentials)
		xedecay = zeros(Float64,Ncells)
		xirise = zeros(Float64,Ncells)
		xidecay = zeros(Float64,Ncells)

		v = zeros(Float64,Ncells) #membrane voltage

		# initialisation of membrane potentials and poisson inputs
		for cc = 1:Ncells
			v[cc] = vreset # + (vth0-vreset)*rand()
		end

		vth = vth0*ones(Float64,Ncells) #adaptive threshold
		lastSpike = -100*ones(Ncells) #last time the neuron spiked

		# ---------------------------------- set up storing avg weights -----------------

		# total number of simulation steps, steps when to normalise, and save
		Nsteps = round(Int,T/dt)

		# true time
		t::Float64 = 0.0

		# bool counter if a neuron has just had a spike
		spiked = zeros(Bool,Ncells)
  v_out = Matrix{Float64}(undef,Ncells,Nsteps)
	@time	for tt = 1:Nsteps

				if mod(tt,Nsteps/100) == 1  #print percent complete
					print("\r",round(Int,100*tt/Nsteps))
				end

				forwardInputsE[:] .= 0.
				forwardInputsI[:] .= 0.
				t = dt*tt

				fill!(spiked,zero(Bool)) # reset spike bool without new memory allocation

				for cc = 1:Ncells
					xerise[cc] += -dt*xerise[cc]/tauerise + forwardInputsEPrev[cc]
					xedecay[cc] += -dt*xedecay[cc]/tauedecay + forwardInputsEPrev[cc]
					xirise[cc] += -dt*xirise[cc]/tauirise + forwardInputsIPrev[cc]
					xidecay[cc] += -dt*xidecay[cc]/tauidecay + forwardInputsIPrev[cc]

					if cc <= Ne # excitatory
						if thrchange
						vth[cc] += dt*(vth0 - vth[cc])/tauth;
						end
					end

					if t > (lastSpike[cc] + taurefrac) #not in refractory period
						# update membrane voltage

						ge = (xedecay[cc] - xerise[cc])/(tauedecay - tauerise);
						gi = (xidecay[cc] - xirise[cc])/(tauidecay - tauirise);

						if cc <= Ne #excitatory neuron (eif), has adaptation
							dv = (vreste - v[cc] + eifslope*exp((v[cc]-vth[cc])/eifslope))/taue +
                  ge*(erev-v[cc])/C + gi*(irev-v[cc])/C +
                  in_const[cc]/taue
							v[cc] += dt*dv
							if v[cc] > vpeak
								spiked[cc] = true
							end
						else
							dv = (vresti - v[cc])/taui +
                 ge*(erev-v[cc])/C + gi*(irev-v[cc])/C +
                  in_const[cc]/taui
							v[cc] += dt*dv
							if v[cc] > vth0
								spiked[cc] = true
							end
						end

						if spiked[cc] #spike occurred
							spiked[cc] = true;
							v[cc] = vreset;
							lastSpike[cc] = t;
							totalspikes[cc] += 1;
							totsp += 1;
							if totsp < spmax
								spiketimes[totsp,1] = tt; # time index as a sparse way to save spiketimes
								spiketimes[totsp,2] = cc; # cell id
							elseif totsp == spmax
								spiketimes[totsp,1] = tt; # time index
								spiketimes[totsp,2] = cc; # cell id

								totsp = 0 # reset counter total number of spikes
							end

							if cc <= Ne && thrchange # only change for excitatory cells and when thrchange == true
								vth[cc] = vth0 + ath;
							end

							#loop over synaptic projections
							for dd = 1:Ncells # postsynaptic cells dd  - cc presynaptic cells
								if cc <= Ne #excitatory synapse
									forwardInputsE[dd] += weights[cc,dd];
								else #inhibitory synapse
									forwardInputsI[dd] += weights[cc,dd];
								end
							end

						end #end if(spiked)
					end #end if(not refractory)
        v_out[cc,tt] = v[cc]
			end #end loop over cells

			forwardInputsEPrev = copy(forwardInputsE)
			forwardInputsIPrev = copy(forwardInputsI)

		end # tt loop over time
	print("\r")

	println("simulation finished")

	return totalspikes,v_out

end

##
vtest = let dt=0.1,T=500,
  _w = 123.0,
  Ne=1,Ni=1,weights=[0.0 0.0 ; _w 0.0];
  spiketimes=Matrix{Int32}(undef,1000,1000)
  runsimulation_static_simplified2(Ne,Ni,weights,spiketimes;
   dt=dt, T=T)[2]
end

_ = let plt = plot()
  ts = (1:size(vtest,2)).*0.1E-3
  plot!(plt,ts,vtest[1,:];linewidth=2)
  plot!(plt,ts,vtest[2,:];linewidth=2,linestyle=:dash)
end

##
#=

		#membrane dynamics
		taue = 30 #e membrane time constant ms
		taui = 20 #i membrane time constant ms
		vreste = -60 #e resting potential mV
		vresti = -70 #i resting potential mV
		vpeak = 20 #cutoff for voltage.  when crossed, record a spike and reset mV
		eifslope = 2 #eif slope parameter mV
		C = 300 #capacitance pF
		erev = 0 #e synapse reversal potential mV
		irev = -75 #i synapse reversal potntial mV
		vth0 = -52 #initial spike voltage threshold mV
		thrchange = false # can be switched off to have vth constant at vth0
		ath = 10 #increase in threshold post spike mV
		tauth = 30 #threshold decay timescale ms
		vreset = -70 #reset potential mV
		taurefrac = 50 #absolute refractory period ms
		aw_adapt = 4 #adaptation parameter a nS conductance
		bwfactor = 100
		bw_adapt = bwfactor*0.805 #adaptation parameter b pA current

    # constant input
    in_const = fill(0.0,Ncells)
    in_const[2] = -51.0 - vreste

		# synaptic kernel
		tauerise = 1 #e synapse rise time
		tauedecay = 6 #e synapse decay time
		tauirise = .5 #i synapse rise time
		tauidecay = 2 #i synapse decay time

=#

##
# time in seconds, voltage in mV
dt = 0.1E-3
Ttot = 0.5
myτe = 30E-3 # seconds
myτi = 20E-3 # seconds
τrefr= 50E-3 # refractoriness
vth = 20.   # mV
vthexp = -52.0 # actual threshold for spike-generation
vth_i = vthexp
eifslope = 2.0
v_reset = -70.0
v_rev_e = 0.0
v_rev_i = -75.0
vreste = -60.0 #e resting potential mV
v_leak_e = vreste
vresti = -70.0 #i resting potential mV
in_const_e = 0.0
Cap = 300.0 #capacitance mF
in_const_i = (-51.0 - vresti)*Cap/myτi

# synaptic kernel
tauerise = 1E-3 #e synapse rise time
tauedecay = 6E-3 #e synapse decay time
#taueplus,taueminus  = tauerise, tauedecay
taueplus,taueminus  = tauedecay, tauerise
tauirise = 0.5E-3 #i synapse rise time
tauidecay = 2E-3 #i synapse decay time
tauiplus,tauiminus = tauidecay,tauirise

##
nt_e = let sker = S.SKExpDiff(taueplus,taueminus)
  sgen = S.SpikeGenEIF(vthexp,eifslope)
  S.NTLIFConductance(sker,sgen,myτe,Cap,
   vth,v_reset,vreste,τrefr,v_rev_e)
end
nt_i = let sker = S.SKExpDiff(tauiplus,tauiminus)
  sgen = S.SpikeGenNone()
  S.NTLIFConductance(sker,sgen,myτi,Cap,
   vth_i,v_reset,vresti,τrefr,v_rev_i)
end
ps_e = S.PSLIFConductance(nt_e,1)
ps_i = S.PSLIFConductance(nt_i,1)

# one static input
ps_input1 = S.PSSimpleInput(S.InputSimpleOffset(in_const_i))

# I to E connection
conn_ei = let  _w = 123.0
  wmat = sparse(fill(_w,(1,1)))
  S.ConnGeneralIF2(wmat)
end

pop_i = S.Population(ps_i,(S.FakeConnection(),ps_input1))
pop_e = S.Population(ps_e,(conn_ei,ps_i))

##
# that's it, let's make the network
myntw = S.RecurrentNetwork(dt,pop_e,pop_i)

# record spiketimes and internal potential
krec = 1
rec_state_e = S.RecStateNow(ps_e,krec,dt,Ttot)
rec_state_i = S.RecStateNow(ps_i,krec,dt,Ttot)
rec_spikes_e = S.RecSpikes(ps_e,100.0,Ttot)
rec_spikes_i = S.RecSpikes(ps_i,100.0,Ttot)

## Run

times = (0:myntw.dt:Ttot)
nt = length(times)
# clean up
S.reset!.([rec_state_e,rec_spikes_e])
S.reset!.([rec_state_i,rec_spikes_i])
S.reset!.([ps_e,ps_i])
S.reset!(conn_ei)
# initial conditions
ps_i.state_now .= v_reset
ps_e.state_now .= v_reset

for (k,t) in enumerate(times)
  rec_state_e(t,k,myntw)
  rec_state_i(t,k,myntw)
  rec_spikes_e(t,k,myntw)
  rec_spikes_i(t,k,myntw)
  S.dynamics_step!(t,myntw)
end

S.add_fake_spikes!(1.0vth,rec_state_e,rec_spikes_e)
S.add_fake_spikes!(1.0vth_i,rec_state_i,rec_spikes_i)
##
_ = let plt=plot()
  plot!(plt,rec_state_e.times,rec_state_e.state_now[1,:];linewidth=2,leg=false)
  ts = (1:size(vtest,2)).*0.1E-3
  plot!(plt,ts,vtest[1,:];linewidth=2,linestyle=:dash)
end

##

_ = let plt = plot()
  plot!(plt,rec_state_i.times,rec_state_i.state_now[1,:];linewidth=2,leg=false)
  ts = (1:size(vtest,2)).*0.1E-3
  plot!(plt,ts,vtest[2,:];linewidth=2,linestyle=:dash)
end


"""
Adapted from 2D Types:
Geometric predicates for 3D Vectors, cartesian points (px,py,pz), homogeneuos points (pw,px,py,pz), and planes
	Julia can inline these.
"""

Vector3 = Tuple{Real,Real,Real}
Point3C = Tuple{Real,Real,Real}
Plane3 = Tuple{Real,Real,Real,Real}
Point3H = Tuple{Real,Real,Real,Real}
Point3 = Union{Point3C,Point3H}
Point3L = Tuple{Real,Real,Real,Real,Real} # last cord is lift
Sphere3 = Tuple{Real,Real,Real,Real,Real} # x,y,z,q

toPoint3H((x,y,z)::Point3C)::Point3H = (one(x),x,y,z)
lift(x::Float32,y::Float32,z::Float32)::Point3L = (x,y,z,muladd(widen(z),widen(z), muladd(widen(y),widen(y), widemul(x,x))) ) # promote to double
lift((x,y,z)::Point3C)::Point3L = (x,y,z, muladd(z,z, muladd(y,y, x*x)) ) # promote to double
lift(a::Array{<:Real,1}) = (a[1], a[2], a[3], sum(a.*widen.(a)))
dot(u,v) = sum(u.*v)
# cartesian and homogeneous orientation determinants

# cartesian plane equation
function plane(p::Point3C, q::Point3C, r::Point3C)::Plane3
	pq = q.-p
	pr = r.-p
	return (det3(p,pq,pr,1,2,3), -det2(pq,pr,2,3), det2(pq,pr,1,3), -det2(pq,pr,1,2))
end

# computing homogeneous plane equations
plane(p::Point3H,q::Point3H,r::Point3H)::Plane3 = (det3(p,q,r,2,3,4),
													-det3(p,q,r,1,3,4),
													det3(p,q,r,1,2,4),
													-det3(p,q,r,1,2,3))

orientptpl((x,y,z)::Point3C,(pw,px,py,pz)::Plane3) = orientptpl(toPoint3H(x,y,z),(pw,px,py,pz))
orientptpl((w,x,y,z)::Point3H,(pw,px,py,pz)::Plane3) = muladd(px,x,muladd(py,y,muladd(pz,z,pw)))
abovePl(q::Point3,p::Plane3)::Bool = orientptpl(q,p)>zero(p[1])
belowPl(q::Point3,p::Plane3)::Bool = orientptpl(q,p)<zero(p[1])


@doc raw"""
    Ising{T<:Int} <: AbstractGraphPointProcess{T}

The [Ising model](https://en.wikipedia.org/wiki/Ising_model) characterizes a point process defined on the vertices of a graph ``(V, E)`` with joint density proportional to

```math
    \mathbb{P}\!\left[ \mathcal{X}=X \right]
    \propto
    \prod_{i \in V}
        \exp(h_i x_i)
    \prod_{\{i, j\} \in E}
        \exp(J x_i x_j)
```

where ``(h_i)_{i\in V}`` are called magnetization paremeters and ``J`` the interaction coefficient (``J \gtreqless 0`` characterizes ferro/antiferro magnetic interactions).

# Example

A realization from a ``5\times 5`` grid graph, with ``h = 0.2, J = 0.1``.

![assets/ising_grid_prs.png](assets/ising_grid_prs.png)
"""
struct Ising{T<:Int} <: AbstractGraphPointProcess{T}
    graph::LG.SimpleGraph{T}
    "Interaction"
    J::Float64
    "Magnetization"
    h::Union{Float64,Vector{Float64}}
end

function Base.show(io::IO, pp::Ising{T}) where {T}
    print(io, "Ising{$T}\n- graph = $(pp.graph)\n- J = $(pp.J) (interaction)\n- h = $(pp.h) (magnetization)")
end

"""
    Ising(
        dims::Vector{T},
        J::Real,
        h::Union{Real,Vector{Real}}=0;
        periodic::Bool=true
    ) where {T<:Int}

Construct [`PRS.Ising`](@ref) on a grid graph with dimension `dims`, interaction parameter `J` and magnetization `h`.

Periodic boundary conditions on the grid graph are set according to `periodic`.

```jldoctest; output = true
using PartialRejectionSampling

dims = [5, 5]
J, h = 0.01, 0
PRS.Ising(dims, J, h; periodic=true)

# output

Ising{Int64}
- graph = {25, 50} undirected simple Int64 graph
- J = 0.01 (interaction)
- h = 0.0 (magnetization)
```
"""
function Ising(
    dims::Vector{T},
    J::Real,
    h::Union{Real,Vector{Real}}=0;
    periodic::Bool=true
) where {T<:Int}
    g = LG.grid(dims; periodic=periodic)
    return Ising(g, J, float.(h))
end

"""
    Ising(
        graph::LG.SimpleGraph{T},
        J::Real,
        h::Union{Real,Vector{Real}}=0
    ) where {T<:Int}

Construct [`PRS.Ising`](@ref) on `graph` with interaction parameter `J` and magnetization `h`.

```jldoctest; output = true
using PartialRejectionSampling
using LightGraphs; const LG = LightGraphs

graph = LG.grid([5, 5])
J, h = 0.01, 0
PRS.Ising(graph, J, h)

# output

Ising{Int64}
- graph = {25, 40} undirected simple Int64 graph
- J = 0.01 (interaction)
- h = 0.0 (magnetization)
```
"""
function Ising(
    graph::LG.SimpleGraph{T},
    J::Real,
    h::Union{Real,Vector{Real}}=0
) where {T<:Int}
    h isa Vector && @assert length(h) == LG.nv(graph)
    return Ising{T}(graph, J, float.(h))
end

## Sampling

@doc raw"""
    generate_sample_prs(
        [rng::Random.AbstractRNG,]
        pp::Ising
    )

Generate an exact sample form `pp` using Partial Rejection Sampling (PRS), see Section 4.2 of [GuJeLi19](@cite).

Default sampler is [`PRS.generate_sample_grid_prs`](@ref).

**See also**

- [FeViYi19](@cite),
- [FeGuYi19](@cite) [`PRS.generate_sample_gibbs_perfect`](@ref).
"""
generate_sample_prs(rng::Random.AbstractRNG, pp::Ising) = generate_sample_grid_prs(rng, pp)
# generate_sample_prs(pp::Ising) = generate_sample_prs(Random.default_rng(), pp)

@doc raw"""
    generate_sample(
        [rng::Random.AbstractRNG,]
        pp::Ising{T},
        idx::T
    )::T where {T<:Int}

Generate an exact sample from the marginal distribution of state ``i=`` `idx` of `pp`.

More specifically,

```math
    x_i
        \sim
        \operatorname{Bernoulli}_{-1, 1}
            (\sigma(h_i)),
```

where ``\sigma`` denotes the [sigmoid function](https://en.wikipedia.org/wiki/Sigmoid_function).
"""
function generate_sample(
    rng::Random.AbstractRNG,
    pp::Ising{T},
    idx::T
)::T where {T<:Int}
    hᵢ = pp.h isa Real ? pp.h : pp.h[idx]
    return rand(rng) < sigmoid(hᵢ) ? one(T) : -one(T)
end

# function generate_sample(
#     pp::Ising{T},
#     idx::T
# )::T where {T<:Int}
#     return generate_sample(Random.default_rng(), pp, idx)
# end

@doc raw"""
    generate_sample!(
        rng::Random.AbstractRNG,
        state::AbstractVector{T},
        indices,
        ising::Ising{T}
    ) where {T<:Int}

Generate an exact sample from the marginal distribution of each states of the [`PRS.Ising`](@ref) model `ising` at the prescribed `indices`.
"""
function generate_sample!(
    rng::Random.AbstractRNG,
    state::AbstractVector{T},
    indices,
    ising::Ising{T}
) where {T<:Int}
    for i in indices
        state[i] = generate_sample(rng, ising, i)
    end
end

@doc raw"""
    generate_sample_conditional!(
        rng::Random.AbstractRNG,
        state::AbstractVector{T},
        i::T,
        ising::Ising{T}
    ) where {T<:Int}

Generate an exact sample from the conditional distribution of the state ``x_i`` given its neighboring states in `ising.graph`.
More specifically,
```math
    x_i \mid x_{N(i)}
        \sim
        \operatorname{Bernoulli}_{-1, 1}
            (\sigma(h_i + J \sum_{j \in N(i)} x_j)),

where ``\sigma`` denotes the [sigmoid function](https://en.wikipedia.org/wiki/Sigmoid_function).
```
"""
function generate_sample_conditional!(
    rng::Random.AbstractRNG,
    state::AbstractVector{T},
    i::T,
    ising::Ising{T}
) where {T<:Int}
    hᵢ = ising.h isa Number ? ising.h : ising.h[i]
    proba = sigmoid(hᵢ + ising.J * sum(state[j] for j in LG.neighbors(ising.graph, i)))
    state[i] = rand(rng) < proba ? 1 : -1
end

include("ising_sampling_gibbs_perfect.jl")
include("ising_sampling_grid_prs.jl")


using ForwardDiff

struct ControlLimits{T}
    ω_max::T
    ω_min::T
    a_max::T
    a_min::T

    #TODO: These are new - move somewhere else later?
    vs_max::T
    vs_min::T
end

struct DynamicsModel{T}
    n::Int64
    m::Int64
    u_limits::ControlLimits{T}
end

function linearize_dynamics(dynamics, x0, u0, Δt)
    ForwardDiff.jacobian(x -> dynamics(x, u0, Δt), x0), ForwardDiff.jacobian(u -> dynamics(x0, u, Δt), u0)
end

function discrete_dynamics(x, u, dt=0.02)
    ω = u[1]
    a = u[2]
    if abs(ω) > 1E-3
        return [x[1] + x[4]/ω * (sin(x[3] + ω*dt) - sin(x[3])) + a/ω * dt * sin(x[3] + ω*dt) + a/ω^2 * (cos(x[3] + ω*dt) - cos(x[3])),
                x[2] - x[4]/ω * (cos(x[3] + ω*dt) - cos(x[3])) - a/ω * dt * cos(x[3] + ω*dt) + a/ω^2 * (sin(x[3] + ω*dt) - sin(x[3])),
                x[3] + ω * dt,
                x[4] + a * dt]
    else
        return [x[1] + x[4] * dt * cos(x[3]) + 0.5 * a * dt^2 * cos(x[3]),
                x[2] + x[4] * dt * sin(x[3]) + 0.5 * a * dt^2 * sin(x[3]),
                x[3],
                x[4] + a * dt]
    end
end

function simulate_path_dynamics(s, u, dt=0.02)
    u * dt + s
end

function simulate(discrete_dynamics, x0, us, Δt)
    xs = [x0]
    for t in 1:size(us)[2]
        x_curr = xs[end]
        x_next = discrete_dynamics(x_curr, us[:,t], Δt)
        push!(xs, x_next)
    end
    xs
end


@testset "weak scatterers" begin

    # angular frequencies
    ωs = [0.001,20.]
    # ωs = [20.] # use only one frequency to speed up test

    # tolerance
    tol = 1e-5;
    basis_order = 2;

    # background medium
    medium = Acoustic(2; ρ=1.0, c=1.0)
    ms = MultipleScattering

    p1 = Particle(Acoustic(2; ρ=10.0, c=12.),ms.Circle(0.01))
    p2 = Particle(Acoustic(2; ρ=3.0, c=2.0),ms.Circle(0.2))

    species = [
        Specie(p1; volume_fraction=0.03),
        Specie(p2; volume_fraction=0.04)
    ]

    # wavenumbers
    eff_medium = effective_medium(medium, species)

    k_eff_lows = ωs./eff_medium.c
    k_eff_φs = wavenumber_low_volumefraction(ωs, medium, species; basis_order = basis_order)

    k_effs_arr = [
        wavenumbers(ω, medium, species;
            num_wavenumbers=1, tol = tol, basis_order = basis_order)
    for ω in ωs]

    inds = [argmin(abs.(k_effs_arr[i] .- k_eff_φs[i])) for i in eachindex(ωs)]
    k_effs = [k_effs_arr[i][inds[i]] for i in eachindex(inds)]

    # k_effs =
    # 0.0010101407549128578 + 7.824113157942236e-13im
    # 20.207827596241156 + 0.11344062283733775im

    @test norm( (k_effs - k_eff_φs) ./ norm.(k_eff_φs) ) < 0.001
    @test norm( (k_effs - k_eff_lows) ./ norm.(k_eff_lows) ) < 0.01
    @test norm(k_effs[1] - k_eff_lows[1])/norm(k_effs[1]) < 1e-7

    normal = [-1.0,0.0] # an outward normal to the surface
    material = Material(Halfspace(normal),species)

    # define a plane wave source travelling at a 45 degree angle in relation to the material
    source = PlaneSource(medium, -normal)

    wave_effs_φs = [WaveMode(ωs[i], k_eff_φs[i], source, material; tol = tol) for i in eachindex(ωs)]

    # reflection coefficient
    Rs2 = reflection_coefficients(ωs, source, material; tol = tol, num_wavenumbers = 1, basis_order = basis_order)

    # if the wavenumber k_effs are already calculated, the below is faster
    # pairs each ω in ωs with each wave in wave_effs2 to calculate the refleciton coefficients
    Rs = map(eachindex(ωs)) do i
        w_effs = [
            WaveMode(ωs[i], k_eff, source, material; tol = tol, basis_order=basis_order)
        for k_eff in k_effs_arr[i]]
        reflection_coefficient(ωs[i], w_effs, source, material; tol = tol)
    end

    @test Rs ≈ Rs2 # same keywords should lead to the same reflection coefficient, except for numerically rounding

    Rs1 = map(eachindex(ωs)) do i
        w_eff = WaveMode(ωs[i], k_effs[i][1], source, material; tol = 1e-9)
        reflection_coefficient(ωs[i], w_eff, source, material)
    end

    # Direct incidence
    R_low = reflection_transmission_coefficients(ωs[1], source, eff_medium, material.shape)[1]
    Rs_φs = reflection_coefficients(ωs, wave_effs_φs, source, material; tol=1e-9)
    # the below takes a low-volfrac expansion for both the wavenumber and reflection coefficient
    Rs_φs2 = reflection_coefficient_low_volumefraction(ωs, source, material; basis_order = basis_order)

    @test maximum(abs.(Rs_φs - Rs1)) < 1e-4 # already relative to incident wave amplitude = 1
    @test maximum(abs.(Rs_φs2 - Rs1)) < 1e-3
    @test abs(R_low - Rs1[1]) < 1e-7

    # Vary angle of incidence θin. Beware of extreme angles θ ~= -pi/2 or pi/2.
    θs = (-π/2.1):0.3:(π/2.1)
    sources = [PlaneSource(medium, [cos(θ),sin(θ)]) for θ in θs]

    R_low = map(sources) do s
        reflection_transmission_coefficients(ωs[1], s, eff_medium, material.shape)[1]
    end

    Rs = map(sources) do s
        wave_effs_2 = [
            WaveMode(ωs[i], k_effs[i], s, material; tol = tol, basis_order = basis_order)
        for i in eachindex(ωs)]
        reflection_coefficients(ωs, wave_effs_2, s, material)
    end
    Rs_φs = [
        reflection_coefficient_low_volumefraction(ωs, s, material; basis_order = basis_order)
    for s in sources]

    len = length(ωs)

    @test maximum(abs(R_low[i] - Rs[i][1]) for i in eachindex(R_low)) < 1e-7
    @test maximum(norm(R)/len for R in (Rs_φs - Rs)[2:end-2]) < 5e-3
end


# This file is part of the IntervalArithmetic.jl package; MIT licensed

const unary_operations = (:-, :inv, :inv_inner, :abs, :abs2, :sign)
for f in unary_operations
    @eval ($f)(x::EmptyInterval) = x
end

# these binary operations with empty intervals yield empty intervals
for f in (:+, :-, :*, :min, :max)
    @eval ($f)(::ValidInterval{T}, ::ValidInterval{T}) where T =
        emptyinterval(T)
end

## Addition and subtraction

+(a::ValidInterval) = a
+(a::Interval{T}, b::Interval{T}) where T = @round(a.lo + b.lo, a.hi + b.hi)

-(a::Interval) = unsafe_interval(-a.hi, -a.lo)
-(a::Interval{T}, b::Interval{T}) where T = @round(a.lo - b.hi, a.hi - b.lo)

min_ignore_nans(args...) = minimum(Iterators.filter(!isnan, args))
max_ignore_nans(args...) = maximum(Iterators.filter(!isnan, args))

## Multiplication
function *(a::Interval{T}, b::Interval{T}) where T
    (a == zero(a) || b == zero(b)) && return zero(a)
    lo = min_ignore_nans(
        @round_down(a.lo * b.lo),
        @round_down(a.lo * b.hi),
        @round_down(a.hi * b.lo),
        @round_down(a.hi * b.hi)
    )
    hi = max_ignore_nans(
        @round_up(a.lo * b.lo),
        @round_up(a.lo * b.hi),
        @round_up(a.hi * b.lo),
        @round_up(a.hi * b.hi)
    )
    return unsafe_interval(lo, hi)
end

## Division
inv_inner(a::Interval{T}) where T =
    a == zero(a) ? emptyinterval(T) : @round(inv(a.hi), inv(a.lo))

inv(a::Interval{T}) where T = hull(
        inv_inner(a ∩ nonnegatives(T)),
        inv_inner(a ∩ -nonnegatives(T))
    )

/(::ValidInterval{T}, ::ValidInterval{T}) where T = emptyinterval(T)

function /(a::Interval{T}, b::Interval{T}) where T
    if zero_in(b)
        return div_by_zero(a, b)
    else
        return b ≺ zero(b) ? -div_by_pos(a, -b) : div_by_pos(a, b)
    end
end

function div_by_zero(a::Interval{T}, b::Interval{T}) where T
    b == zero(b) && return emptyinterval(T) # excluded from domain
    a == zero(a) && return a
    contains_zero(b) && return entireinterval(T)
    contains_zero(a) && return entireinterval(T)
    return hull(
        div_by_pos(a, b ∩ nonnegatives(T)),
        -div_by_pos(a, -b ∩ nonnegatives(T))
    )
end

function div_by_pos(a::Interval{T}, b::Interval{T}) where T
    b == zero(b) && return emptyinterval(T)
    if contains_zero(a)
        return hull(
            div_pos_by_pos(a ∩ nonnegatives(T), b),
            -div_pos_by_pos(-a ∩ nonnegatives(T), b)
        )
    elseif a ≼ zero(a)
        return -div_pos_by_pos(-a, b)
    else
        return div_pos_by_pos(a, b)
    end
end

div_pos_by_pos(a::Interval, b::Interval) = @round(a.lo / b.hi, a.hi / b.lo)

mul_rev_to_pair(a::ValidInterval, b::ValidInterval) =
    extended_div(b, a)

extended_div(::ValidInterval{S}, ::ValidInterval{S}) where S =
    (emptyinterval(S), emptyinterval(S))

function extended_div(a::Interval{T}, b::Interval{T}) where T
    if contains_zero(b)
        return extended_div_by_zero(a, b)
    elseif zero_in(b) && zero_in(a)
        return (entireinterval(T), emptyinterval(T))
    else
        return (a / b, emptyinterval(T))
    end
end

ordered(x, y) = precedes(x, y) ? (x, y) : (y, x)

function extended_div_by_zero(a::Interval{T}, b::Interval{T}) where T
    zero_in(a) && return (entireinterval(T), emptyinterval(T))
    return ordered(
        a * unsafe_interval(T(-Inf), @round_up(inv(b.lo))),
        a * unsafe_interval(@round_down(inv(b.hi)), T(Inf))
    )
end

## fma: fused multiply-add
fma(a::Interval{T}, b::Interval{T}, c::Interval{T}) where T = a * b + c

function fma(a::Interval{T}, b::Interval{T}, c::Interval{T}) where T <: BigFloat
    if a == zero(a) || b == zero(b)
        return c
    end
    if isentire(a) || isentire(b)
        return entireinterval(T)
    end

    lo = setrounding(T, RoundDown) do
        lo1 = fma(a.lo, b.lo, c.lo)
        lo2 = fma(a.lo, b.hi, c.lo)
        lo3 = fma(a.hi, b.lo, c.lo)
        lo4 = fma(a.hi, b.hi, c.lo)
        min_ignore_nans(lo1, lo2, lo3, lo4)
    end

    hi = setrounding(T, RoundUp) do
        hi1 = fma(a.lo, b.lo, c.hi)
        hi2 = fma(a.lo, b.hi, c.hi)
        hi3 = fma(a.hi, b.lo, c.hi)
        hi4 = fma(a.hi, b.hi, c.hi)
        max_ignore_nans(hi1, hi2, hi3, hi4)
    end

    unsafe_interval(lo, hi)
end
fma(::ValidInterval{T}, ::ValidInterval{T}, ::ValidInterval{T}) where T =
    emptyinterval(T)

sign(a::Interval)= unsafe_interval(sign(a.lo), sign(a.hi))

abs(a::Interval) = unsafe_interval(mig(a), mag(a))

abs2(a::Interval) = sqr(a)

min(a::Interval{T}, b::Interval{T}) where T =
    unsafe_interval(min(a.lo, b.lo), min(a.hi, b.hi))

max(a::Interval{T}, b::Interval{T}) where T =
    unsafe_interval(max(a.lo, b.lo), max(a.hi, b.hi))

# cancelplus and cancelminus
"""
    cancelminus(a, b)

Return the unique interval `c` such that `b+c=a`.

See Section 12.12.5 of the IEEE-1788 Standard for
Interval Arithmetic.
"""
function cancelminus(a::Interval{T}, b::Interval{T}) where T <: Real
    diam_smaller(a, b) = a.hi + b.lo < b.hi + a.lo
    hasnovalue(a, b) =
        isunbounded(a) || isunbounded(b) || diam_smaller(a, b)

    hasnovalue(a, b) && return entireinterval(T)

    return @round(a.lo - b.lo, a.hi - b.hi)
end
cancelminus(::EmptyInterval{T}, b::Interval{T}) where T =
    isbounded(b) ? emptyinterval(T) : entireinterval(T)
cancelminus(::Interval{T}, ::EmptyInterval{T}) where T = entireinterval(T)
cancelminus(::EmptyInterval{T}, ::EmptyInterval{T}) where T = emptyinterval(T)
cancelminus(a::ValidInterval, b::ValidInterval) =
    cancelminus(promote(a, b)...)

"""
    cancelplus(a, b)

Returns the unique interval `c` such that `b-c=a`;
it is equivalent to `cancelminus(a, −b)`.
"""
cancelplus(a::ValidInterval, b::ValidInterval) = cancelminus(a, -b)


using StaticArrays

raw_data = readlines("../inputs/day04_example.txt")

bingo_number_str = raw_data[1]
boards_str = raw_data[3:end]
num_boards = Int((length(boards_str) + 1) / 6)

function vecvec_to_matrix(vecvec)
    dim1 = length(vecvec)
    dim2 = length(vecvec[1])
    my_array = zeros(Int64, dim1, dim2)
    for i = 1:dim1
        for j = 1:dim2
            my_array[i, j] = vecvec[i][j]
        end
    end
    return my_array
end

function board_str2num(lines)
    board_vecvec = [[parse(Int, val) for val in split(line, " ", keepempty = false)] for line in lines]
    return [vecvec_to_matrix(board) for board in board_vecvec]
end

function check_numbers(numbers, board)
    return [board .== num for num in numbers]
end

function check_bingo(numbers, board)
    board_check = check_numbers(numbers, board)
    col_sum = sum(board_check, dims = 1)
    row_sum = sum(board_check, dims = 2)

    return row_sum
end

boards = [board_str2num(boards_str[6*(ii-1)+1:6*(ii-1)+5]) for ii in range(1, num_boards)]

display(boards[1])
println("")
display(boards[2])
println("")
display(boards[3])
println("")

display(check_numbers(boards[2], [10, 16, 15, 9, 19]))


"""
    Acrobot
"""

struct Acrobot{I, T} <: Model{I, T}
    n::Int
    m::Int
    d::Int

    m1   # mass link 1
    J1    # inertia link 1
    l1    # length link 1
    lc1   # length to COM link 1

    m2    # mass link 2
    J2    # inertia link 2
    l2    # length link 2
    lc2   # length to COM link 2

    g     # gravity

    b1    # joint friction
    b2
end

function M(model::Acrobot, x, w)
    a = (model.J1 + model.J2 + model.m2 * model.l1 * model.l1
         + 2.0 * model.m2 * model.l1 * model.lc2 * cos(x[2]))

    b = model.J2 + model.m2 * model.l1 * model.lc2 * cos(x[2])

    c = model.J2

    @SMatrix [a b;
              b c]
end

function τ(model::Acrobot, x, w)
    a = (-1.0 * model.m1 * model.g * model.lc1 * sin(x[1])
         - model.m2 * model.g * (model.l1 * sin(x[1])
         + model.lc2 * sin(x[1] + x[2])))

    b = -1.0 * model.m2 * model.g * model.lc2 * sin(x[1] + x[2])

    @SVector [a,
              b]
end

function C(model::Acrobot, x, w)
    a = -2.0 * model.m2 * model.l1 * model.lc2 * sin(x[2]) * x[4]
    b = -1.0 * model.m2 * model.l1 * model.lc2 * sin(x[2]) * x[4]
    c = model.m2 * model.l1 * model.lc2 * sin(x[2]) * x[3]
    d = 0.0

    @SMatrix [a b;
              c d]
end

function B(model::Acrobot, x, w)
    @SMatrix [0.0;
              1.0]
end

function f(model::Acrobot, x, u, w)
    q = view(x, 1:2)
    v = view(x, 3:4)
    qdd = M(model, q, w) \ (-1.0 * C(model, x, w) * v
            + τ(model, q, w) + B(model, q, w) * u[1:1] - [model.b1; model.b2] .* v)
    @SVector [x[3],
              x[4],
              qdd[1],
              qdd[2]]
end

function k_mid(model::Acrobot, x)
    @SVector [model.l1 * sin(x[1]),
              -1.0 * model.l1 * cos(x[1])]
end

function k_ee(model::Acrobot, x)
    @SVector [model.l1 * sin(x[1]) + model.l2 * sin(x[1] + x[2]),
              -1.0 * model.l1 * cos(x[1]) - model.l2 * cos(x[1] + x[2])]
end

n, m, d = 4, 1, 4
model = Acrobot{Midpoint, FixedTime}(4, 1, 4, 1.0, 0.33, 1.0, 0.5, 1.0, 0.33, 1.0, 0.5, 9.81, 0.1, 0.1)

# visualization
function visualize!(vis, model::Acrobot, x;
        color=RGBA(0.0, 0.0, 0.0, 1.0),
        r = 0.1, Δt = 0.1)

    default_background!(vis)

    i = 1
    l1 = Cylinder(Point3f0(0.0, 0.0, 0.0), Point3f0(0.0, 0.0, model.l1),
        convert(Float32, 0.025))
    setobject!(vis["l1$i"], l1, MeshPhongMaterial(color = color))
    l2 = Cylinder(Point3f0(0.0,0.0,0.0), Point3f0(0.0, 0.0, model.l2),
        convert(Float32, 0.025))
    setobject!(vis["l2$i"], l2, MeshPhongMaterial(color = color))

    setobject!(vis["elbow$i"], Sphere(Point3f0(0.0),
        convert(Float32, 0.05)),
        MeshPhongMaterial(color = RGBA(0.0, 1.0, 0.0, 1.0)))
    setobject!(vis["ee$i"], Sphere(Point3f0(0.0),
        convert(Float32, 0.05)),
        MeshPhongMaterial(color = RGBA(0.0, 1.0, 0.0, 1.0)))

    anim = MeshCat.Animation(convert(Int, floor(1.0 / Δt)))

    T = length(x)
    for t = 1:T

        MeshCat.atframe(anim,t) do
            p_mid = [k_mid(model, x[t])[1], 0.0, k_mid(model, x[t])[2]]
            p_ee = [k_ee(model, x[t])[1], 0.0, k_ee(model, x[t])[2]]

            settransform!(vis["l1$i"], cable_transform(zeros(3), p_mid))
            settransform!(vis["l2$i"], cable_transform(p_mid, p_ee))

            settransform!(vis["elbow$i"], Translation(p_mid))
            settransform!(vis["ee$i"], Translation(p_ee))
        end
    end

    settransform!(vis["/Cameras/default"],
       compose(Translation(0.0 , 1.0 , -1.0), LinearMap(RotZ(pi / 2.0))))

    MeshCat.setanimation!(vis, anim)
end


using ITensors
using ITensors.ITensorNetworkMaps
using KrylovKit
using LinearAlgebra

include("utils.jl")

N = 4
s = siteinds("S=1/2", N; conserve_qns=true)
χ = 2
ψ = randomMPS(s, n -> isodd(n) ? "↑" : "↓"; linkdims=χ)
ℋ = OpSum()
for n in 1:(N - 1)
  ℋ .+= 0.5, "S+", n, "S-", n + 1
  ℋ .+= 0.5, "S-", n, "S+", n + 1
  ℋ .+= "Sz", n, "Sz", n + 1
end
H = MPO(ℋ, s)

T = [transfer_matrix(ψ, n) for n in 1:N]
Tᴴ = [transfer_matrix((H, ψ), n) for n in 1:N]
@show (T[1] * T[2] * T[3] * T[4] * ITensor(1))[] ≈ inner(ψ, ψ)
@show (ITensor(1) * T[1] * T[2] * T[3] * T[4])[] ≈ inner(ψ, ψ)
@show (ITensor(1) * (T[1] * T[2] * T[3] * T[4]))[] ≈ inner(ψ, ψ)
@show contract(T[1] * T[2] * T[3] * T[4])[] ≈ inner(ψ, ψ)

@show contract(prod(T))[] ≈ inner(ψ, ψ)
@show contract(prod(Tᴴ))[] ≈ inner(ψ, H, ψ)

Ttot = transfer_matrix(ψ, 1:N)
Tᴴtot = transfer_matrix((H, ψ), 1:N)
@show contract(Ttot)[] ≈ inner(ψ, ψ)

println("Time to contract transfer matrix")
@time contract(Tᴴtot)[]
println("Time to contract using `inner(ψ, H, ψ)`")
@time inner(ψ, H, ψ)
@show contract(Tᴴtot)[] ≈ inner(ψ, H, ψ)

v = cache(T)
vᴴ = cache(Tᴴ)

@show (v[2 => 3] * v[3 => 2])[] ≈ inner(ψ, ψ)
@show (v[1 => 2] * v[2 => 1])[] ≈ inner(ψ, ψ)

@show (vᴴ[2 => 3] * vᴴ[3 => 2])[] ≈ inner(ψ, H, ψ)
@show (vᴴ[1 => 2] * vᴴ[2 => 1])[] ≈ inner(ψ, H, ψ)

# One-site effective Hamiltonian
H² = hessian(H, vᴴ, 2)
ψ² = ψ[2]
@show (dag(ψ²) * H²(ψ²))[] ≈ inner(ψ, H, ψ)

# Two-site effective Hamiltonian
H²⁻³ = hessian(H, vᴴ, 2:3)
ψ²⁻³ = prod(ψ[2:3])
@show (dag(ψ²⁻³) * H²⁻³(ψ²⁻³))[] ≈ inner(ψ, H, ψ)

# Three-site effective Hamiltonian
H²⁻⁴ = hessian(H, vᴴ, 2:4)
ψ²⁻⁴ = prod(ψ[2:4])
@show (dag(ψ²⁻⁴) * H²⁻⁴(ψ²⁻⁴))[] ≈ inner(ψ, H, ψ)


using JSON, StanSample, Random, Statistics

stan_chris = "
data {
    int n_rows;
    int<lower=1> n_cols;
    matrix<lower=0>[n_rows,n_cols] x;
}

parameters {
   real mu;
}

model {
    mu ~ normal(0, 1);
}";

seed = 65445;
Random.seed!(seed)

n_rows = 10
n_cols = 2
x = fill(0, n_rows, n_cols)
data = Dict(:x => x, :n_rows => n_rows, :n_cols => n_cols)

tmpdir = joinpath(@__DIR__, "tmp")
sm = SampleModel("chris", stan_chris, tmpdir);
rc = stan_sample(sm; data, seed);

if success(rc)
	df = read_samples(sm , :dataframe)
    #df |> display
    [mean(df.mu), std(df.mu)] |> display
end