text:
# [src/state_to_vector.jl]

"""
Build a state vector by recursively moving through the elements of a more
complex state (matrix, cell array, or struct), saving the numerical
values along the way in a vector. This is used in odehybrid to convert a
complex series of states into a vector to be used for continuous state
updates. The reverse function is vector_to_state.

vector = state_to_vector(state);

Inputs:

    state:   A numeric, array, or struct array type, the contents of
                which consist of numeric, array, or struct array types, etc.
                (Structs MUST be mutable)

Outputs:

    vector:  A vector containing the numerical values from the state (doubles)

Example:

    mutable struct TestStruc
        a
        bcd
    end
    Base.:(==)(a::TestStruc, b::TestStruc) = (Base.:(==)(a.a, b.a) && Base.:(==)(a.bcd, b.bcd)) # Custom

    x = [ [1 3; 4 2], TestStruc([5; 6], [7:9; 10]),  3.14 ]
    v = state_to_vector(x)
    x2 = vector_to_state(v, x)
    x .== x2[1]

See also: vector_to_state.

Online doc: http://www.anuncommonlab.com/doc/odehybrid/state_to_vector.html

Copyright 2014 An Uncommon Lab
"""
function state_to_vector(state)

    # If it shows up here, it should be a struct
    vector = [];
    fields = fieldnames(typeof(state))
    for n = 1:1 # 1:length(state)
        for k = 1:length(fields)
            vector = [vector; state_to_vector(getproperty(state, fields[k]))]
        end
    end

    return vector
end

function state_to_vector(state::Number)
    """
        Converts a number state to a float vector
    """
    # Converts a number to a float and returns it
    vector = float(state)

    return vector
end

function state_to_vector(state::Array{<:Number})
    """
        Converts a number array state to a float array vector
    """
    return float.(state[:])
end

function state_to_vector(state::Char)
    """
        Converts a Char state to a float vector
    """
    vector = float(state)

    return vector
end

function state_to_vector(state::Array{Char})
    """
        Converts a Char arary state to a float array vector
    """
    return float.(state[:])
end

function state_to_vector(state::String)
    """
        Converts a string state to a vector
    """
    # Converts a string to a vector
    vector = []
    for i = 1:length(state)
        vector = [vector; state_to_vector(state[i])]
    end

    return vector
end

function state_to_vector(state::Array) #{Any, 1})
    """
        Converts each element of an Array to the appropriate value and returns in a vector
    """
    vector = []
    for i = 1:length(state)
        vector = [vector; state_to_vector(state[i])]
    end

    return vector
end


text:
import DataFrames
import FITSIO
import WCS
using Celeste.SDSSIO

@testset "the identity WCSTransform works as expected" begin
    rand_coord = rand(2, 10)
    @test WCS.pix_to_world(SampleData.wcs_id, rand_coord) == rand_coord
    @test WCS.world_to_pix(SampleData.wcs_id, rand_coord) == rand_coord
end

@testset "linear_world_to_pix works" begin
    image = SDSSIO.load_field_images(PlainFITSStrategy(datadir),
                SDSSIO.RunCamcolField(3900, 6, 269))[1]
    wcs = image.wcs

    pix_center = Float64[0.5 * image.H, 0.5 * image.W]
    pix_loc = pix_center + [5., 3.]
    world_center = WCS.pix_to_world(wcs, pix_center)
    world_loc = WCS.pix_to_world(wcs, pix_loc)

    function test_jacobian(wcs, pix_center, world_center)
        wcs_jacobian = Model.pixel_world_jacobian(wcs, pix_center);

        pix_loc_test1 = WCS.world_to_pix(wcs, world_loc)
        pix_loc_test2 = Model.linear_world_to_pix(wcs_jacobian, world_center, pix_center,
                                       world_loc)

        # Note that the accuracy of the linear approximation isn't great.
        @test pix_loc_test1 ≈ pix_loc
        @test isapprox(pix_loc_test2, pix_loc, atol=1e-2)
    end

    @test Model.pixel_world_jacobian(SampleData.wcs_id, pix_center) == [1.0 0.0; 0.0 1.0];

    test_jacobian(wcs, pix_center, world_center)
end

ext: jl


text:
struct VarsCart
    t::Num
    x::Vector{Num}
    y::Vector{Num}
    z::Vector{Num}
end

function VarsCart(N)
    t = (@variables t)[1]
    x = (@variables x[1:N])[1] |> Symbolics.scalarize
    y = (@variables y[1:N])[1] |> Symbolics.scalarize
    z = (@variables z[1:N])[1] |> Symbolics.scalarize
    VarsCart(t, x, y, z)
end
mi(i, vars::VarsCart) = [vars.x[i]; vars.y[i]; vars.z[i]]

# struct RHS
#     x::Vector{Num}
#     y::Vector{Num}
#     z::Vector{Num}
# end

# Base.:+(x::RHS, y::RHS) = RHS(x.x .+ y.x, x.y .+ y.y, x.z .+ y.z)

# we use the conventions of SPIRIT: dtm = - m × Beff and A = -Beff

function exchangeEffectiveField(coupling::Matrix, vars::VarsCart)
    # coupling is a NxN array which describes which sites couple to which
    # returns an array with effective field
    Heff = zeros(Num, size(coupling)[1], 3)

    for i in 1:size(coupling)[1]
        for j in 1:size(coupling)[1]
            Heff[i,:] .+= coupling[i,j] .* [vars.x[j], vars.y[j], vars.z[j]]
        end
    end
    Heff
end

function spinTorqueField(field, vars::VarsCart)
    # spin torque is defined as m × (m × J)
    Js = zeros(Num, length(vars.x)[1], 3)
    for i in 1:size(Js)[1]
        Js[i,:] .+= cross([vars.x[i], vars.y[i], vars.z[i]], field[i,:])
    end
    Js
end

function DMIEffectiveField(coupling::Matrix, Dfield::Vector, vars::VarsCart)
    # coupling is a NxN array which describes which sites couple to which
    # returns an array with effective field
    Heff = zeros(Num, size(coupling)[1], 3)
    for i in 1:size(coupling)[1]
        for j in 1:size(coupling)[1]
            Heff[i,:] .-= coupling[i,j] * cross(Dfield, mi(j, vars))
        end
    end
    Heff
end

function anistropyEnergy(anistropy::Vector, vars::VarsCart)
    E = Num(0)
    for i in 1:length(vars.x)
        E -= sum(anistropy .* mi(i, vars).^2)
    end
    return E
end

function exchangeEnergy(coupling::Matrix, vars::VarsCart)
    # we count every pair twice, so multiply by 0.5
    E = Num(0)
    for i in 1:length(vars.x)
        for j in 1:length(vars.x)
            E += -.5coupling[i,j] * mi(i, vars) ⋅ mi(j, vars)
        end
    end
    return E
end

function magneticEnergy(B::Vector, vars::VarsCart)
    E = Num(0)
    for i in 1:length(vars.x)
        E += - B ⋅ mi(i, vars)
    end
    return E
end

function DMIEnergy(coupling::Matrix, Dfield::Vector, vars::VarsCart)
    # we count every pair twice, so multiply by 0.5
    E = Num(0)
    for i in 1:length(vars.x)
        for j in 1:length(vars.x)
            E += -.5coupling[i,j] * Dfield ⋅ cross(mi(i, vars), mi(j, vars))
        end
    end
    return E
end

function DipolarEnergy(R, vars::VarsCart)
    E = Num(0)
    for i in 1:length(vars.x)
        Si = mi(i, vars)
        for j in 1:length(vars.x)
            if i == j
                continue
            end
            Sj = mi(j, vars)
            Rij = R[i] - R[j]
            Rijnorm = Rij.^2 |> sum |> sqrt
            Rijhat = Rij ./ Rijnorm
            E -= 0.5 * Rijnorm^(-3) * (3 * (Si ⋅ Rijhat) * (Sj ⋅ Rijhat) - Si ⋅ Sj)
        end
    end
    return E
end

function effectiveFieldFromEnergy(E, vars::VarsCart)
    Heff = zeros(Num, length(vars.x), 3)
    Heff[:,1] = -Symbolics.gradient(E, vars.x)
    Heff[:,2] = -Symbolics.gradient(E, vars.y)
    Heff[:,3] = -Symbolics.gradient(E, vars.z)

    return Heff

end

# custom functions to make sure the substitute works:

istree(c::Complex{Num}) = true
operation(::Complex{Num}) = Complex{Num}
arguments(c::Complex{Num}) = [real(c), imag(c)]

function HamiltonianFromHeff(Heff, vars)
    N = length(vars.x)
    δm = (@variables δm[1:N])[1] |> Symbolics.scalarize
    δmconj = (@variables δmconj[1:N])[1] |> Symbolics.scalarize
    repl_rules = Dict([])
    for i in 1:N
        repl_rules[vars.x[i]] = (δm[i] + δmconj[i]) / 2
        repl_rules[vars.y[i]] = (δm[i] - δmconj[i]) * -1im / 2
        repl_rules[vars.z[i]] = 1
    end
    # Heffδ = similar(Heff)
    # for i in 1:length(Heff)
    #     Heffδ[i] = substitute(Heff[i], repl_rules)
    # end
    eom_δm = zeros(Complex{Num}, N)
    eom_δm_imag = zeros(Num, N)

    eom_δmconj = zeros(Complex{Num}, N)
    eom_δmconj_imag = zeros(Num, N)

    for i in 1:N
        eom_m = cross(LLG.mi(i, vars), Heff[i,:])
        eom_δm[i] = substitute((eom_m[1] + 1im * eom_m[2]), repl_rules)
        # eom_δm_imag[i] = substitute((eom_m[1] + 1im * eom_m[2]) |> imag, repl_rules)
        eom_δmconj[i] = substitute((eom_m[1] - 1im * eom_m[2]), repl_rules)
        # eom_δmconj_imag[i] = substitute((eom_m[1] - 1im * eom_m[2]) |> imag, repl_rules)
    end
    # A = zeros(Complex{Num}, N, N)
    # B = zeros(Complex{Num}, N, N)
    # for i in 1:N
    #     for j in 1:N
    #         A[i,j] = Symbolics.linear_expansion(real(eom[i,1]), δm[j])[1] + im * Symbolics.linear_expansion(imag(eom[i,1]), δm[j])[1]
    #         B[i,j] = Symbolics.linear_expansion(real(eom[i, 2]), δmconj[j])[1] + im * Symbolics.linear_expansion(imag(eom[i, 2]), δmconj[j])[1]
    #     end
    # end
    # A, B
    i, j = (2, 3)
    eom_m = cross(LLG.mi(i, vars), Heff[i,:])[1]
    # @info eom_m
    # @time Symbolics.linear_expansion(eom_m, vars.z[j])[1]
    @info Symbolics.value(eom_δm[i])
    # @time Symbolics.linear_expansion(eom_δm[i], δm[j])[1]
end

function magneticEffectiveField(Hfield::Vector, vars::VarsCart)
    Heff = zeros(Num, length(vars.x), 3)
    for i in 1:length(vars.x)
        Heff[i,:] = Hfield
    end
    return Heff
end

function generateEnergy(E, vars::VarsCart)
    eval(build_function(E, [vars.x; vars.y; vars.z]))
end

function generateEnergyFromHeff(Heff, vars::VarsCart)
    E = Num(0)
    for i in 1:length(vars.x)
        E -= Heff[:,i] ⋅ mi(i, vars)
    end
    return eval(build_function(E, [vars.x; vars.y; vars.z]))
end

function generateEOM(Heff, α, vars::VarsCart)
    res = zeros(Num, length(vars.x), 3)
    for i in 1:length(vars.x)
        S = [vars.x[i], vars.y[i], vars.z[i]]
        res[i, :] = - cross(S, Heff[:,i]) - α * cross(S, cross(S, Heff[:,i]))
    end
    D = Differential(vars.t)
    return [D.(vars.x) .~ res[:, 1]; D.(vars.y) .~ res[:, 2]; D.(vars.z) .~ res[:, 3]]
end

function generateA(Heff, vars::VarsCart, p = [])
    eval(build_function(-Heff, [vars.x; vars.y; vars.z], p, vars.t)[1])
end
function generateA!(Heff, vars::VarsCart, p = []; parallel = false)
    if parallel
        eval(build_function(-Heff, [vars.x; vars.y; vars.z], p, vars.t, parallel = Symbolics.MultithreadedForm())[2])
    else
        eval(build_function(-Heff, [vars.x; vars.y; vars.z], p, vars.t)[2])
    end
end

function generateA!(Heff, vars::VarsCart, α::Float64, p = []; parallel = false)
    totalH = -Heff
    for i in 1:length(vars.x)
        totalH[i,:] .-= α * cross([vars.x[i], vars.y[i], vars.z[i]], Heff[i,:])
    end
    if parallel
        eval(build_function(totalH, [vars.x; vars.y; vars.z], p,  vars.t, parallel = Symbolics.MultithreadedForm())[2])
    else
        eval(build_function(totalH, [vars.x; vars.y; vars.z], p,  vars.t)[2])
    end
end

function generateA!(Heff, auxfield, vars::VarsCart, α::Float64, p = []; parallel = false)
    totalH = -Heff
    for i in 1:length(vars.x)
        totalH[i,:] .-= α * cross([vars.x[i], vars.y[i], vars.z[i]], Heff[i,:])
    end
    if parallel
        eval(build_function(totalH + auxfield, [vars.x; vars.y; vars.z], p,  vars.t, parallel = Symbolics.MultithreadedForm())[2])
    else
        eval(build_function(totalH + auxfield, [vars.x; vars.y; vars.z], p,  vars.t; expression = Val{false})[2])
    end
end

function generateA!(Heff, auxfield, vars::VarsCart, α::Vector{Float64}, p = []; parallel = false)
    totalH = -Heff
    for i in 1:length(vars.x)
        totalH[i,:] .-= α[i] * cross([vars.x[i], vars.y[i], vars.z[i]], Heff[i,:])
    end
    if parallel
        eval(build_function(totalH + auxfield, [vars.x; vars.y; vars.z], p,  vars.t, parallel = Symbolics.MultithreadedForm())[2])
    else
        eval(build_function(totalH + auxfield, [vars.x; vars.y; vars.z], p, vars.t)[2])
    end
end

function noiseLLG(vars::VarsCart)
    # eval(build_function(0, [vars.x; vars.y; vars.z], vars.t)[1])
    (u, t) -> 0
end

function generateEOM(Heff, Js::Array, α, vars::VarsCart)
    res = zeros(Num, length(vars.x), 3)
    for i in 1:length(vars.x)
        S = [vars.x[i], vars.y[i], vars.z[i]]
        res[i, :] = - cross(S, Heff[i,:]) - α * cross(S, cross(S, Heff[i,:])) + cross(S, cross(S, Js[i,:]))
    end
    D = Differential(vars.t)
    return [D.(vars.x) .~ res[:, 1]; D.(vars.y) .~ res[:, 2]; D.(vars.z) .~ res[:, 3]]
end

function initRandomSpinToolkit(vars::VarsCart, λ)
    # boltzmann is proportional to θ (for now)
    # the angle ϕ is chosen randomly between 0 and 2pi
    N = length(vars.x)
    ϕ = rand(N) * 2π
    θ = rand(N) * π
    mx = sin.(θ) .* cos.(ϕ)
    my = sin.(θ) .* sin.(ϕ)
    mz = cos.(θ)
    return [
        [vars.x[i] => mx[i] for i in 1:N];
        [vars.y[i] => my[i] for i in 1:N];
        [vars.z[i] => mz[i] for i in 1:N]
        ]
end

function initRandomSpin(N::Int, λ)
    ϕ = rand(N) * 2π
    θ = rand(N) * π
    mx = sin.(θ) .* cos.(ϕ)
    my = sin.(θ) .* sin.(ϕ)
    mz = cos.(θ)
    return [
        mx my mz
        ]
end

function initRandomSpin(vars::VarsCart, λ)
    # boltzmann is proportional to θ (for now)
    # the angle ϕ is chosen randomly between 0 and 2pi
    N = length(vars.x)
    initRandomSpin(N, λ)
end

function initBoltzmannSpin(vars::VarsCart, λ)
    # boltzmann is proportional to θ (for now)
    # the angle ϕ is chosen randomly between 0 and 2pi
    N = length(vars.x)
    ϕ = rand(N) * 2π
    dist = Exponential(λ / 2)
    p = rand(dist, N)
    θ = acos.(1 .- 2p)
    mx = sin.(θ) .* cos.(ϕ)
    my = sin.(θ) .* sin.(ϕ)
    mz = cos.(θ)
    m = zeros(Float64, N, 3)
    m[:,1] .= mx
    m[:,2] .= my
    m[:,3] .= mz
    return m
end

function initBoltzmannSpinToolkit(vars::VarsCart, λ)
    # boltzmann is proportional to θ (for now)
    # the angle ϕ is chosen randomly between 0 and 2pi
    N = length(vars.x)
    ϕ = rand(N) * 2π
    dist = Exponential(λ / 2)
    p = rand(dist, N)
    θ = acos.(1 .- 2p)
    mx = sin.(θ) .* cos.(ϕ)
    my = sin.(θ) .* sin.(ϕ)
    mz = cos.(θ)
    return [
        [vars.x[i] => mx[i] for i in 1:N];
        [vars.y[i] => my[i] for i in 1:N];
        [vars.z[i] => mz[i] for i in 1:N]
        ]
end

max_stars_repo_name: pgunnink/LLGTools

max_issues_repo_name: pgunnink/LLGTools

max_forks_repo_name: pgunnink/LLGTools


text:
export ExactSampler

mutable struct ExactSampler{S} <: MonteCarloSampler
    samples_length::Int
    seed::S
end

ExactSampler(n_samples; seed=rand(UInt)) =
    ExactSampler(n_samples, seed)

## Cache
mutable struct ExactSamplerCache{A<:AbstractRNG,T} <: SamplerCache{ExactSampler}
    rng::A
    steps_done::Int
    pdf::Vector{T}
end

_sampler_cache(s::ExactSampler, v, net, part) =
    ExactSamplerCache(MersenneTwister(s.seed),
                      0,
                      zeros(real(eltype(v)), spacedimension(v)))

##
function init_sampler!(sampler::ExactSampler, net::Union{MatrixNet,KetNet}, σ, c::ExactSamplerCache)
    # Compute the distribution
    for i=1:spacedimension(σ)
        set!(σ, i-1)
        c.pdf[i] = exp(log_prob_ψ(net, σ))
    end
    c.pdf ./= sum(c.pdf)
    cumsum!(c.pdf, c.pdf)
    c.steps_done = 0

    samplenext!(σ, sampler, net, c)
    return c
end

done(s::ExactSampler, σ, c) = c.steps_done >= s.samples_length

function samplenext!(σ, s::ExactSampler, net::Union{MatrixNet,KetNet}, c)
    # Check termination condition, and return if verified
    done(s, σ, c) && return false

    # sample and set
    r = rand(c.rng)
    i = searchsortedfirst(c.pdf, r)
    set!(σ, i-1)
    c.steps_done += 1

    return true
end

# return findfirst(probs .> number)
function sample_from_distr(number::Real, probs::Vector)
    @inbounds for (i,p)=enumerate(probs)
        if p > number
            return i
        end
    end
end

######## Multithreaded
function _mt_recompute_sampler_params!(samplers, s::ExactSampler)
    nt = Threads.nthreads()
    _n_samples = Int(ceil(s.samples_length / nt))
    rng = MersenneTwister(s.seed)
    for i=1:Threads.nthreads()
        samplers[i] = ExactSampler(_n_samples, seed=rand(rng, UInt))
    end
end

function _mt_sampler_cache(s::MTSampler{T}, v, net, ::ParallelThreaded) where T<:ExactSampler
    pdf = zeros(real(eltype(v)), spacedimension(v))
    rng = MersenneTwister(s.seed)
    seeds = rand(UInt, Threads.nthreads())
    scs = Vector{ExactSamplerCache{typeof(rng), eltype(pdf)}}()
    for i=1:Threads.nthreads()
        push!(scs, ExactSamplerCache(MersenneTwister(seeds[i]), 0, pdf))
    end
    return MTSamplerCache(s, scs, net, v)
end

function mt_init_sampler(s::MTSampler{T}, net::Union{MatrixNet,KetNet}, σ, c::MTSamplerCache) where T<:ExactSampler
    np     = Threads.nthreads()

    Threads.@threads for i=1:Threads.nthreads()
        np     = Threads.nthreads()
        rank   = i - 1
        _σ = c.σs[i]
        _c = c.caches[i]
        _n = c.nets[i]
        n_min, extra = divrem(spacedimension(_σ), np)
        iter_length  = n_min + (rank < extra ? 1 : 0)
        iter_start   = n_min * rank + min(rank, extra) + 1
        iter_end     = iter_start + iter_length - 1

        for j=iter_start:iter_end
            set_index!(_σ, j)
            _c.pdf[j] = exp(log_prob_ψ(_n, _σ))
        end
    end

    c.caches[1].pdf ./= sum(c.caches[1].pdf)
    cumsum!(c.caches[1].pdf, c.caches[1].pdf)

    Threads.@threads for i=1:np
        _σ   = c.σs[i]
        _c   = c.caches[i]
        _net = c.nets[i]
        _c.steps_done = 0

        samplenext!(_σ, sampler_list(s)[i], _net, _c)
    end
    return c
end


text:
using LinearAlgebra
"""
    (V,H,hist)=iar_lr(prob,startvec,N)

### Low-rank Infinite Arnoldi
Low-rank version of Infinite Arnoldi method as described in Algorithm 2 (Taylor coefficients) in "A rank-exploiting infinite Arnoldi algorithm for nonlinear eigenvalue problems", R. Van Beeumen, E. Jarlebring, and W. Michiels

Runs `N` steps of the infinite Arnoldi method started with `startvec` with target `σ=0`, specifically for the low-rank problem.
The argument `prob` represents a problem `M(λ)` and has associated functions
`compute_Mder`, `compute_Mlincomb`:

* `M=compute_Mder(prob,s)` computes the matrix function value at `s=λ`.
* `z=compute_Mlincomb(prob,s,X)` computes a linear combination of derivatives. The linear combination coefficients are given in the matrix `X`:  `z=M(s)*X[:,1]+M'(s)*X[:,2]+...+M^{(k)}(s)*X[:,k+1]`.

"""
function iar_lr(prob,startvec,N,precomputed_factorization=nothing)

    T=ComplexF64; # Problem type

    n=length(startvec);

    sigma=0;  # Only target implemented

    s=prob.s
    p=size(prob.V,2)
    @show s,p

    x0=zeros(T,n);
    x0[:]=normalize(startvec);
    H=zeros(T,N+1,N);
    print("IAR L-R Iteration:");
    # Assume N>s+1
    V=zeros(T,n*s+p*(N-s+1),N+1);
    V[1:n,1]=x0;

    start_time=time();
    hist=Dict(:time_count => NaN*zeros(N))

    if (isnothing(precomputed_factorization))
        Msolve=factorize(compute_Mder(prob,sigma));
    else
        Msolve=precomputed_factorization;
    end

    for k=1:N
        print("$k ")

        # Sizes of the blocks in V
        # V:
        # n n n n n n  ^
        #   n n n n n  |
        #     n n n n  s
        #       ..     |
        #       n   n  _
        #        p  p
        #
        #

        # X1 in C^nx(s-1)  (independent of k)
        vecX1=V[1:(s-1)*n,k]
        X1=reshape(vecX1,n,s-1);

        # X2 in C^n  (independent of k)
        X2=V[(s-1)*n .+ (1:n),k]

        # X3 in C^p x (k-s)
        vecX3=V[(s*n+1):(s*n+p*(k-s)),k]
        X3=reshape(vecX3,p,max(k-s,0))

        # x[1]=X1[:,1]   in C^n
        # x[2]=X1[:,2]   in C^n
        # ..
        # x[s-s]=X1[:,s-1]   in C^n  (X2)
        # x[s]=X2  in C^n  (X2)
        # x[s+1]=X3[:,1] in C^p
        # x[s+2]=X3[:,2] in C^p
        # ...
        x=Vector{Vector{T}}(undef,max(k,s))
        for i=1:s-1
            # compute C^n vector
            x[i]=X1[:,i];
        end
        x[s]=X2;
        for i=(s+1):k
            # compute C^p vector
            x[i]=X3[:,i-s];
        end

        y0=zeros(T,n);
        for i=1:s
            y0 += (prob.Aders[i+1]*x[i])/i
        end
        for i=(s+1):k
            y0 += (prob.Uv[i+1]*x[i])/i
        end

        xtilde=-(Msolve\y0)

        # In manuscript: D_{s-1,s-1}
        Dsm1=Diagonal( 1 ./ (1:(s-1)));
        # In manuscript: Dtile_{s+1,m}
        Dtilde=Diagonal( 1 ./ ((s+1):k))

        vecY=[xtilde;vec(X1*Dsm1); (1/s)*prob.V'*X2 ; vec(X3*Dtilde)]
        #vecY=[xtilde;vec(Y1)];

        # Orthogonalization

        if (k<(s+1))
            NN=n*k
        else
            NN=n*s+p*(k-s)
        end
        NN1=size(vecY,1);
        #@show NN,NN1
        V0=V[1:NN,1:k];
        h=V0'*vecY[1:NN];
        vecY=vecY-V[1:NN1,1:k]*h;

        g=V0'*vecY[1:NN];   # twice to be sure
        vecY=vecY-V[1:NN1,1:k]*g;
        h=h+g;

        beta=norm(vecY);  # normalize
        vecY=vecY/beta;
        H[1:(k+1),k]=[h;beta];
        V[1:NN1,k+1]=vecY;

        hist[:time_count][k]=time()-start_time;

    end

    println("done");

    return (V,H,hist);

end

max_stars_repo_name: siobhanie/InfGMRES

max_issues_repo_name: siobhanie/InfGMRES

max_forks_repo_name: siobhanie/InfGMRES


text:
#=
    Designer based on a metaheuristic
=#

mutable struct MetaheuristicOptions
    method::Metaheuristics.AbstractMetaheuristic
    iterations::Int64
    multithreads::Bool
    controller::AbstractController
    isnpv::Bool
    reducer::AbstractScenariosReducer
    objective_risk::AbstractRiskMeasure
    share_risk::AbstractRiskMeasure
    lpsp_risk::AbstractRiskMeasure
    lpsp_tol::Float64
    reopt::Bool
    read_reduction::Union{String, Nothing}
    write_reduction::Union{String, Nothing}

    MetaheuristicOptions(; method = Metaheuristics.Clearing(),
                           iterations = 50,
                           multithreads = false,
                           controller = RBC(),
                           isnpv = false,
                           reducer = FeatureBasedReducer(),
                           objective_risk = Expectation(),
                           share_risk = Expectation(),
                           lpsp_risk = WorstCase(),
                           lpsp_tol = 1e-3,
                           reopt = false,
                           read_reduction = nothing,
                           write_reduction = nothing) =
                           new(method, iterations, multithreads, controller, isnpv, reducer, objective_risk, share_risk, lpsp_risk, lpsp_tol, reopt, read_reduction, write_reduction)

end

mutable struct Metaheuristic <: AbstractDesigner
    options::MetaheuristicOptions
    decisions::NamedTuple
    results::Metaheuristics.MetaheuristicResults
    history::AbstractScenarios

    Metaheuristic(; options = MetaheuristicOptions()) = new(options)
end

# Objective functions
function fobj(decisions::Array{Float64,1}, mg::Microgrid, designer::Metaheuristic, ω::Scenarios, probabilities::Array{Float64})
    # Paramters
    nh, ny, ns = size(ω.demands[1].power)
    λ1 = λ2 = λ3 = 1e6

    # Initialize mg
    mg_m = deepcopy(mg)

    # Initialize controller
    controller_m = initialize_controller!(mg_m, designer.options.controller, ω)

    # Initialize with the manual designer
    designer_m = initialize_designer!(mg_m, Manual(generations = [decisions[1:length(mg_m.generations)]...], storages = [decisions[length(mg_m.generations)+1:length(mg_m.generations)+length(mg_m.storages)]...], converters = [decisions[end-length(mg_m.converters)+1:end]...]), ω)

    # Simulate
    simulate!(mg_m, controller_m, designer_m, ω)

    # Metrics
    metrics = Metrics(mg_m, designer_m)

    # Share constraint
    share = max(0., mg.parameters.renewable_share - conditional_value_at_risk([reshape(metrics.renewable_share[2:ny, 1:ns], :, 1)...],  probabilities,  designer.options.share_risk))

    # LPSP constraint for the heat
    metrics.lpsp.heat isa Nothing ? lpsp = 0. : lpsp = max(0., conditional_value_at_risk([reshape(metrics.lpsp.heat[2:ny, 1:ns], :, 1)...], probabilities,  designer.options.lpsp_risk) - designer.options.lpsp_tol)

    # SoC constraint for the seasonal storage
    soc_seasonal = 0.
    for a in mg_m.storages
        if a isa H2Tank
            soc_seasonal += sum(max(0., a.soc[1,y,s] - a.soc[end,y,s]) for y in 2:ny, s in 1:ns)
        end
    end

    # Objective - Algortihm find the maximum
    if designer.options.isnpv
        # NPV
        npv = conditional_value_at_risk([metrics.npv.total...], probabilities, designer.options.objective_risk)
        return npv - λ1 * share - λ2 * lpsp - λ3 * soc_seasonal
    else
        # Equivalent annual cost
        eac = conditional_value_at_risk([metrics.eac.total...], probabilities, designer.options.objective_risk)
        return - eac - λ1 * share - λ2 * lpsp - λ3 * soc_seasonal
    end
end

### Offline
function initialize_designer!(mg::Microgrid, designer::Metaheuristic, ω::Scenarios)
    # Preallocate and assigned values
    preallocate!(mg, designer)

    # Scenario reduction from the optimization scenario pool
    if designer.options.isnpv
        println("Starting scenario reduction...")
        ω_reduced, probabilities = reduce(designer.options.reducer, ω)
    else
        if isa(designer.options.read_reduction, Nothing)
            println("Starting scenario reduction...")
            ω_reduced, probabilities = reduce(designer.options.reducer, ω)
            # Saving
            if !isa(designer.options.write_reduction, Nothing)
                save(designer.options.write_reduction, "scenarios", ω_reduced, "probabilities", probabilities)
            end
        else
            println("Reading scenario reduction from file...")
            ω_reduced = load(designer.options.read_reduction, "scenarios")
            probabilities = load(designer.options.read_reduction, "probabilities")
        end
        # Repeat to simulate 2 years
        ω_reduced = repeat(ω_reduced, 1, 2, 1)
    end

    # Bounds
    lb, ub = set_bounds(mg)

    # Optimize
    designer.results = Metaheuristics.optimize(lb, ub,
                                               designer.options.method,
                                               options = Metaheuristics.Options(iterations = designer.options.iterations, multithreads = designer.options.multithreads)
    ) do decisions
        fobj(decisions, mg, designer, ω_reduced, probabilities)
      end

    # Assign values
    for k in 1:length(mg.generations)
        designer.decisions.generations[k][1,:] .= designer.results.minimizer[k]
    end
    for k in 1:length(mg.storages)
        designer.decisions.storages[k][1,:] .= designer.results.minimizer[length(mg.generations)+k]
    end
    for k in 1:length(mg.converters)
        designer.decisions.converters[k][1,:] .= designer.results.minimizer[end-length(mg.converters)+k]
    end

    # Save history for online optimization
    designer.history = ω_reduced

    return designer
end

### Online
function compute_investment_decisions!(y::Int64, s::Int64, mg::Microgrid, designer::Metaheuristic)
    #TODO
end

### Utils
function set_bounds(mg::Microgrid)
    # Initialization
    lb, ub = [], []
    # Generations
    for a in mg.generations
        push!(lb, a.bounds.lb)
        push!(ub, a.bounds.ub)
    end
    # Storages
    for a in mg.storages
        push!(lb, a.bounds.lb)
        push!(ub, a.bounds.ub)
    end
    # Converters
    for a in mg.converters
        push!(lb, a.bounds.lb)
        push!(ub, a.bounds.ub)
    end
    return lb, ub
end


text:
## MATLAB interfcae to find_zero
## Main functions are
## * fzero(f, ...) to find _a_ zero of f
## * fzeros(f, ...) to attempt to find all zeros of f

"""
    fzero(f, x0; order=0; kwargs...)

Find zero of a function using an iterative algorithm

* `f`: a scalar function or callable object
* `x0`: an initial guess, finite valued.
* `order`

This is a polyalgorithm redirecting to different algorithms based on the value of `order`. Dispatches to `find_zero(f, x0, OrderN(); kwargs...)`.

"""
function fzero(f, x0::Number; kwargs...)
    x = float(x0)
    isinf(x) && throw(ConvergenceFailed("An initial value must be finite"))
    derivative_free(f, x; kwargs...)
end

"""
    fzero(f, a, b; kwargs...)

Find zero of a function within a bracket, [a,b].

Dispatches to `find_zero(f, (a,b), Bisection())`.

"""
fzero(f, a::Number, b::Number; kwargs...) = find_zero(f, (a,b), Bisection(); kwargs...)
fzero(f, bracket::Vector{T}; kwargs...)  where {T <: Number} = find_zero(f, bracket, Bisection(); kwargs...)
fzero(f, bracket::Tuple{T,S}; kwargs...)  where {T <: Number, S<:Number} = find_zero(f, bracket, Bisection();kwargs...)

"""
    fzero(f, fp, x0; kwargs)

Find zero using Newton's method.

Dispatches to `find_zero((f,fp), x0, Roots.Newton(); kwargs...)`.

"""
fzero(f::Function, fp::Function, x0::Real; kwargs...) = newton(f, fp, float(x0); kwargs...)

# match fzero up with find_zero
@noinline function derivative_free(f, x0; order::Int=0, kwargs...)

    if order == 0
        method = Order0()
    elseif order == 1
        method = Order1()
    elseif order == 2
        method = Order2()
    elseif order == 5
        method = Order5()
    elseif order == 8
        method = Order8()
    elseif order == 16
        method = Order16()
    else
        warn("Invalid order. Valid orders are 0, 1, 2, 5, 8, and 16")
        throw(ArgumentError())
    end

    d = Dict(kwargs)
    for (o, n) in ((:ftol, :atol), (:ftolrel, :rtol),
         (:xtol, :xatol), (:xtolrel, :xrtol))
        if haskey(d, o)
            d[n] = d[o]
        end
    end

    find_zero(f, x0, method; d...)
end

# ##
# """
#     fzero(f, x0, bracket; kwargs...)
#
# Find a zero within a bracket with an initial guess to *possibly* speed things along.
#
# Dispatches to the `A42` method.
#
#"""
@deprecate fzero(f, x0::Real, bracket::Vector; kwargs...)  fzero(f, bracket)

## fzeros
"""

`fzeros(f, a, b; kwargs...)`

Searches for all zeros of `f` within an interval `(a,b)`. Assume neither `a` or `b` is a zero.

Dispatches to `find_zeros(f, a, b; kwargs...)`.
"""
function fzeros(f, a::Number, b::Number; kwargs...)
    find_zeros(f, float(a), float(b); kwargs...)
end
fzeros(f, bracket::Vector{T}; kwargs...) where {T <: Number} = fzeros(f, a, b; kwargs...)
fzeros(f, bracket::Tuple{T,S}; kwargs...) where {T <: Number, S<:Number} = fzeros(f, a, b; kwargs...)


text:
"""
$(TYPEDEF)

XYZ Euclidean manifold variable node softtype.

Example
-------
```julia
p3 = Point3()
```
"""
struct Point3 <: IncrementalInference.InferenceVariable end
getDimension(::Point3) = 3
getManifolds(::Point3) = (:Euclid,:Euclid,:Euclid)


text:
using Documenter
using ArrowMacros

PAGES = [
    "Home" => "index.md",
    "Tests" => "tests.md"
]

makedocs(;
    sitename = "ArrowMacros.jl",
    format = Documenter.HTML(),
    modules = [ArrowMacros],
    pages = PAGES,
    authors = "Giancarlo A. Antonucci <giancarlo.antonucci@icloud.com>"
)

deploydocs(;
    repo = "https://github.com/giancarloantonucci/ArrowMacros.jl"
)


text:
# stacks

type Stack{S}   # S is the type of the internal dequeue instance
    store::S
end

Stack{T}(ty::Type{T}) = Stack(Deque{T}())
Stack{T}(ty::Type{T}, blksize::Integer) = Stack(Deque{T}(blksize))

isempty(s::Stack) = isempty(s.store)
length(s::Stack) = length(s.store)

top(s::Stack) = back(s.store)

function push!(s::Stack, x)
    push!(s.store, x)
    s
end

pop!(s::Stack) = pop!(s.store)

max_stars_repo_name: ambujone/DIDebug


text:
#some shortcuts for common grids in polar remote sensing
"""
	NSIDCN(n)

generate the _National Snow and Ice Data Center_ North polar stereographic grid
for a given resolution `n` in m
"""
NSIDCN(n)=Grid(Proj4.Projection(Proj4.epsg[3413]),(-3850,3750).*1e3,(-5350,5850).*1e3,n)

"""
	NSDICS(n)

generate the _National Snow and Ice Data Center_ South polar stereographic grid
for a given resolution `n` in m
"""
NSIDCS(n)=Grid(Proj4.Projection(Proj4.epsg[3976]),(-3950,3950).*1e3,(-3950,4350).*1e3,n)

"""
	EASE2N(n)

generate the _National Snow and Ice Data Center_ North EASE2 grid
for a given resolution `n` in m
"""
EASE2N(n)=Grid(Proj4.Projection(Proj4.epsg[6931]),(-9e6, 9e6),(-9e6, 9e6),n)

"""
	EASE2N(n)

generate the _National Snow and Ice Data Center_ South EASE2 grid
for a given resolution `n` in m
"""
EASE2S(n)=Grid(Proj4.Projection(Proj4.epsg[6932]),(-9e6, 9e6),(-9e6, 9e6),n)


text:
## ------------  Types

# """
# 	$(SIGNATURES)

# School groups in the data.
# """
# abstract type AbstractSchoolGroups end

# # HSG, CD, CG
# struct SchoolGroups3 <: AbstractSchoolGroups end

"""
	$(SIGNATURES)

Settings for experience profiles.

Wage regressions for intercepts regress log fixed effect on:
- schooling dummies
- gpa dummies
- parental dummies
"""
Base.@kwdef mutable struct WageRegressions
	# Max experience in the data
	maxExper :: UInt8 = 12
	# School groups in each experience profile
	experGroupV :: Vector{Vector{Symbol}} = [[:HSG, :SC, :CG]]
	# Highest exponenent in experience profile
	maxExperExponent :: UInt8 = 2
	useParentalDummies :: Bool = true
end

"""
	DataSettings

Settings for data files. These are baked into the data construction (by Oksana), but should still be changeable.
"""
Base.@kwdef mutable struct DataSettings
	name :: Symbol = :default
	# Base directory for all data files
	baseDir :: String
	# Data dir relative to base directory for all data files
	dataSubDir :: String
	# Uses afqt or gpa
	afqtGpa :: String = "afqt"
	# School groups
	sGroups :: AbstractSchoolGroups = Schooling3{SchoolInt}()
	# Wage regression settings
	wageRegressions :: WageRegressions = default_wage_regressions()
	"Number of colleges (qualities)"
    nColleges :: CollInt = 4
    "Number of two year colleges"
    n2Year :: CollInt = 1
	"HS GPA groups"
	hsGpaUbV :: Vector{Float64} = collect(0.25 : 0.25 : 1.0)
	"Parental income groups"
	parentalUbV :: Vector{Float64} = collect(0.25 : 0.25 : 1.0)
	"Number of years for which stats by year are constructed"
	Tmax :: TimeInt = 4
	"Number of data credits per data course"
	creditsPerCourse :: ncInt = 3
	"Year for transfer regression moments"
	transferYear :: TimeInt = 1
	"Year for cross-sectional study time moments"
	studyTimeYear :: TimeInt = 1
	"Year for cross-sectional hours worked moments"
	workTimeYear :: TimeInt = 1
	"Year for cross-sectional tuition moments"
	tuitionYear :: TimeInt = 1
end

"""
	RawDataFile

Raw data file info.

Raw data files sit in a nested directory structure, such as
`SelfReport/dat_files/Means/HS_Char/`
The object contains info for constructing the path.
"""
mutable struct RawDataFile
	# :self or :transcript
	selfOrTranscript  ::  Symbol
	# E.g. :finance
	group  ::  Symbol
	# E.g. :mean or :regression; :std for std deviations
	momentType :: Symbol
	# E.g. "regression.dat"
	rawFile  ::  String
	ds :: DataSettings
end

# DataFrame holding data by (x,y) loaded from a raw data file
struct DataFrameXY
    df :: DataFrame
end

"""
	$(SIGNATURES)

DataFrame that is loaded from the delimited file.
"""
struct DataFrameX
    df :: DataFrame
end

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


text:
using RegressionAndOtherStories
using StanSample

hibbs = CSV.read(ros_datadir("ElectionsEconomy", "hibbs.csv"), DataFrame)

stan1_1 = "
functions {
}
data {
    int<lower=1> N;      // total number of observations
    vector[N] growth;    // Independent variable: growth
    vector[N] vote;      // Dependent variable: votes
}
parameters {
    real b;              // Coefficient independent variable
    real a;              // Intercept
    real<lower=0> sigma; // dispersion parameter
}
model {
    vector[N] mu;
    mu = a + b * growth;

    // priors including constants
    a ~ normal(50, 20);
    b ~ normal(2, 10);
    sigma ~ exponential(1);

    // likelihood including constants
    vote ~ normal(mu, sigma);
}";

data1_1s = (N=16, vote=hibbs.vote, growth=hibbs.growth);

m1_1s = SampleModel("hibbs", stan1_1)
rc1_1s = stan_sample(m1_1s; data=data1_1s)

if success(rc1_1s)
    sdf = read_summary(m1_1s)
    post1_1s_df = read_samples(m1_1s, :dataframe)
    post1_1s_df[!, :chain] = repeat(collect(1:m1_1s.num_chains);
        inner=m1_1s.num_samples)
    post1_1s_df[!, :chain] = categorical(post1_1s_df.chain)
    ā, b̄, σ̄ = mean(Array(post1_1s_df), dims=1)

    @testset "test_hibbs.jl" begin
        @test ā ≈ 3.0 atol = 0.1
        @test b̄ ≈ 46.3 atol = 1.0
        @test σ̄ ≈ 3.6 atol = 0.2
    end

    @testset "model_summary.jl" begin
        mod_sum = model_summary(post1_1s_df, [:a, :b, :sigma])
        @test mod_sum[:a, :median] ≈ 46.5 atol=0.4
    end
end


text:
@testset "Interval" begin
    test_values = [
        (-10, 1000, 1),
        (0.0, 1, 0.01),  # Use different types to test promotion
        ('a', 'z', 1),
        (Date(2013, 2, 13), Date(2013, 3, 13), Day(1)),
        (DateTime(2016, 8, 11, 0, 30), DateTime(2016, 8, 11, 1), Millisecond(1))
    ]

    @testset "constructor" begin
        for T in [Int32, Int64, Float64]
            @test Interval{T}() == Interval{T}(zero(T), zero(T), Inclusivity(false, false))
        end

        @test Interval{Date}() ==
            Interval{Date}(Date(0), Date(0), Inclusivity(false, false))
        @test Interval{DateTime}() ==
            Interval{DateTime}(DateTime(0), DateTime(0), Inclusivity(false, false))
        @test Interval{ZonedDateTime}() == Interval{ZonedDateTime}(
            ZonedDateTime(0, tz"UTC"), ZonedDateTime(0, tz"UTC"), Inclusivity(false, false)
        )

        for (a, b, _) in test_values
            @test a..b == Interval(a, b)
            @test Interval(a, b) == Interval{typeof(a)}(a, b, Inclusivity(true, true))
            @test Interval(a, b, true, false) ==
                Interval{typeof(a)}(a, b, Inclusivity(true, false))
            @test Interval{typeof(a)}(a, b, true, false) ==
                Interval{typeof(a)}(a, b, Inclusivity(true, false))
            @test Interval(a, b, Inclusivity(true, false)) ==
                Interval{typeof(a)}(a, b, Inclusivity(true, false))
            @test Interval(b, a, Inclusivity(true, false)) ==
                Interval{typeof(a)}(a, b, Inclusivity(false, true))
        end

        # The three-argument Interval constructor can generate a StackOverflow if we aren't
        # careful
        @test_throws MethodError Interval(1, 2, 3)
    end

    @testset "non-ordered" begin
        @test_throws ArgumentError Interval(NaN, Inf)
        @test_throws ArgumentError Interval(NaN, NaN)
    end

    @testset "hash" begin
        # Need a complicated enough element type for this test to ever fail
        zdt = now(tz"Europe/London")
        a = Interval(zdt, zdt)
        b = deepcopy(a)
        @test hash(a) == hash(b)
    end

    @testset "conversion" begin
        @test_throws DomainError convert(Int, Interval(10, 10, Inclusivity(false, false)))
        @test_throws DomainError convert(Int, Interval(10, 10, Inclusivity(false, true)))
        @test_throws DomainError convert(Int, Interval(10, 10, Inclusivity(true, false)))
        @test convert(Int, Interval(10, 10, Inclusivity(true, true))) == 10
        @test_throws DomainError convert(Int, Interval(10, 11, Inclusivity(true, true)))

        for T in (Date, DateTime)
            dt = T(2013, 2, 13)
            @test_throws DomainError convert(T, Interval(dt, dt, Inclusivity(false, false)))
            @test_throws DomainError convert(T, Interval(dt, dt, Inclusivity(false, true)))
            @test_throws DomainError convert(T, Interval(dt, dt, Inclusivity(true, false)))
            @test convert(T, Interval(dt, dt, Inclusivity(true, true))) == dt
            @test_throws DomainError convert(T, Interval(dt, dt + Day(1), Inclusivity(true, true)))
        end
    end

    @testset "eltype" begin
        @test eltype(Interval(1,2)) == Int
        @test eltype(Interval{Float64}(1,2)) == Float64
    end

    @testset "accessors" begin
        for (a, b, _) in test_values
            for i in 0:3
                inc = Inclusivity(i)
                interval = Interval(a, b, inc)

                @test first(interval) == a
                @test last(interval) == b
                @test span(interval) == b - a
                @test inclusivity(interval) == inc
                @test isclosed(interval) == (first(inc) && last(inc))
                @test isopen(interval) == !(first(inc) || last(inc))
            end
        end

        # DST transition
        firstpoint = ZonedDateTime(2018, 3, 11, 1, tz"America/Winnipeg")
        endpoint = ZonedDateTime(2018, 3, 11, 3, tz"America/Winnipeg")
        interval = Interval(firstpoint, endpoint)
        @test span(interval) == Hour(1)

        firstpoint = ZonedDateTime(2018, 11, 4, 0, tz"America/Winnipeg")
        endpoint = ZonedDateTime(2018, 11, 4, 2, tz"America/Winnipeg")
        interval = Interval(firstpoint, endpoint)
        @test span(interval) == Hour(3)
    end

    @testset "display" begin
        interval = Interval(1, 2, Inclusivity(false, false))
        @test string(interval) == "(1 .. 2)"
        @test sprint(show, interval, context=:compact=>true) == string(interval)
        @test sprint(show, interval) == "Interval{$Int}(1, 2, Inclusivity(false, false))"

        interval = Interval('a', 'b', Inclusivity(false, true))
        @test string(interval) == "(a .. b]"
        @test sprint(show, interval, context=:compact=>true) == string(interval)
        @test sprint(show, interval) == "Interval{Char}('a', 'b', Inclusivity(false, true))"

        interval = Interval(Date(2012), Date(2013), Inclusivity(true, false))
        shown = if VERSION < v"1.2.0-DEV.29"
             string(
                "Interval{Date}(",
                repr(Date(2012, 1, 1)),
                ", ",
                repr(Date(2013, 1, 1)),
                ", Inclusivity(true, false))",
            )
        else
            "Interval{Date}(Date(2012, 1, 1), Date(2013, 1, 1), Inclusivity(true, false))"
        end

        @test string(interval) == "[2012-01-01 .. 2013-01-01)"
        @test sprint(show, interval, context=:compact=>true) == string(interval)
        @test sprint(show, interval) == shown

        interval = Interval("a", "b", Inclusivity(true, true))
        @test string(interval) == "[a .. b]"
        @test sprint(show, interval, context=:compact=>true) == string(interval)
        @test sprint(show, interval) ==
            "Interval{String}(\"a\", \"b\", Inclusivity(true, true))"
    end

    @testset "equality" begin
        for (a, b, unit) in test_values
            for i in 0:3
                interval = Interval(a, b, Inclusivity(i))
                cp = copy(interval)
                lesser_val = Interval(a - unit, b + unit, Inclusivity(i))
                greater_val = Interval(a + unit, b + unit, Inclusivity(i))
                diff_inc = Interval(a, b, Inclusivity(mod(i + 1, 4)))

                @test interval == cp
                @test interval != lesser_val
                @test interval != diff_inc

                @test isequal(interval, cp)
                @test !isequal(interval, lesser_val)
                @test !isequal(interval, diff_inc)

                @test hash(interval) == hash(cp)
                @test hash(interval) != hash(lesser_val)
                @test hash(interval) != hash(diff_inc)

                @test !isless(interval, cp)
                @test !(interval < cp)
                @test !(interval ≪ cp)
                @test !(interval > cp)
                @test !(interval ≫ cp)
                @test isless(interval, greater_val)
                @test interval < greater_val
                @test !(interval ≪ greater_val)     # Still overlap, so not disjoint
                @test !(interval > greater_val)
                @test !(interval ≫ greater_val)
                @test !isless(greater_val, interval)
                @test !(greater_val < interval)
                @test !(greater_val ≪ interval)     # Still overlap, so not disjoint
                @test greater_val > interval
                @test !(greater_val ≫ interval)     # Still overlap, so not disjoint
                @test isless(lesser_val, interval)
                @test lesser_val < interval
                @test !(lesser_val ≪ interval)
                @test !(lesser_val > interval)
                @test !(lesser_val ≫ interval)
                @test !isless(interval, lesser_val)
                @test !(interval < lesser_val)
                @test !(interval ≪ lesser_val)
                @test interval > lesser_val
                @test !(interval ≫ lesser_val)      # Still overlap, so not disjoint
            end
        end

        @test Interval(2010, 2011) ≪ Interval(2012, 2013)
        @test !(Interval(2012, 2013) ≪ Interval(2010, 2011))

        @test Interval(Date(2010), Date(2011)) ≪ Interval(Date(2012), Date(2013))
        @test !(Interval(Date(2012), Date(2013)) ≪ Interval(Date(2010), Date(2011)))

        @test !(Interval(0, 10) ≪ Interval(10, 20))
        @test !(Interval(10, 20) ≪ Interval(0, 10))

        @test Interval(0, 10, Inclusivity(true, false)) ≪
            Interval(10, 20, Inclusivity(false, true))
        @test !(
            Interval(10, 20, Inclusivity(true, false)) ≪
            Interval(0, 10, Inclusivity(false, true))
        )

        # Comparisons between Interval{T} and T
        @test isless(5, Interval(10, 20))
        @test 5 < Interval(10, 20)
        @test 5 ≪ Interval(10, 20)
        @test !(5 > Interval(10, 20))
        @test !(5 ≫ Interval(10, 20))

        @test isless(10, Interval(10, 20, Inclusivity(false, false)))
        @test 10 < Interval(10, 20, Inclusivity(false, false))
        @test 10 ≪ Interval(10, 20, Inclusivity(false, false))
        @test !(10 > Interval(10, 20, Inclusivity(false, false)))
        @test !(10 ≫ Interval(10, 20, Inclusivity(false, false)))

        @test !isless(10, Interval(10, 20))
        @test !(10 < Interval(10, 20))
        @test !(10 ≪ Interval(10, 20))
        @test !(10 > Interval(10, 20))
        @test !(10 ≫ Interval(10, 20))

        @test !isless(15, Interval(10, 20))
        @test !(15 < Interval(10, 20))
        @test !(15 ≪ Interval(10, 20))
        @test 15 > Interval(10, 20)
        @test !(15 ≫ Interval(10, 20))

        @test !isless(20, Interval(10, 20))
        @test !(20 < Interval(10, 20))
        @test !(20 ≪ Interval(10, 20))
        @test 20 > Interval(10, 20)
        @test !(20 ≫ Interval(10, 20))

        @test !isless(20, Interval(10, 20, Inclusivity(false, false)))
        @test !(20 < Interval(10, 20, Inclusivity(false, false)))
        @test !(20 ≪ Interval(10, 20, Inclusivity(false, false)))
        @test 20 > Interval(10, 20, Inclusivity(false, false))
        @test 20 ≫ Interval(10, 20, Inclusivity(false, false))

        @test !isless(25, Interval(10, 20))
        @test !(25 < Interval(10, 20))
        @test !(25 ≪ Interval(10, 20))
        @test 25 > Interval(10, 20)
        @test 25 ≫ Interval(10, 20)

        @test !isless(Interval(10, 20), 5)
        @test !(Interval(10, 20) < 5)
        @test !isless(Interval(10, 20, Inclusivity(false, false)), 10)
        @test !(Interval(10, 20, Inclusivity(false, false)) < 10)
        @test !isless(Interval(10, 20), 10)
        @test !(Interval(10, 20) < 10)
        @test !(Interval(10, 20) ≪ 10)
        @test isless(Interval(10, 20), 15)
        @test Interval(10, 20) < 15
        @test !(Interval(10, 20) ≪ 15)
        @test isless(Interval(10, 20), 20)
        @test Interval(10, 20) < 20
        @test !(Interval(10, 20) ≪ 20)
        @test isless(Interval(10, 20, Inclusivity(false, false)), 20)
        @test Interval(10, 20, Inclusivity(false, false)) < 20
        @test isless(Interval(10, 20), 25)
        @test Interval(10, 20) < 25

        for lt in (isless, <)
            @test lt(Date(2013), Interval(Date(2014), Date(2016)))
            @test lt(Date(2014), Interval(Date(2014), Date(2016), false, false))
            @test !lt(Date(2014), Interval(Date(2014), Date(2016)))
            @test !lt(Date(2014), Interval(Date(2014), Date(2016)))
            @test !lt(Date(2015), Interval(Date(2014), Date(2016)))
            @test !lt(Date(2016), Interval(Date(2014), Date(2016)))
            @test !lt(Date(2016), Interval(Date(2014), Date(2016), false, false))
            @test !lt(Date(2017), Interval(Date(2014), Date(2016)))
        end

        @test !isless(Interval(Date(2014), Date(2016)), Date(2013))
        @test !(Interval(Date(2014), Date(2016)) < Date(2013))
        @test !isless(Interval(Date(2014), Date(2016), false, false), Date(2014))
        @test !(Interval(Date(2014), Date(2016), false, false) < Date(2014))
        @test !isless(Interval(Date(2014), Date(2016)), Date(2014))
        @test !(Interval(Date(2014), Date(2016)) < Date(2014))
        @test !(Interval(Date(2014), Date(2016)) ≪ Date(2014))
        @test isless(Interval(Date(2014), Date(2016)), Date(2015))
        @test Interval(Date(2014), Date(2016)) < Date(2015)
        @test !(Interval(Date(2014), Date(2016)) ≪ Date(2015))
        @test isless(Interval(Date(2014), Date(2016)), Date(2016))
        @test Interval(Date(2014), Date(2016)) < Date(2016)
        @test !(Interval(Date(2014), Date(2016)) ≪ Date(2016))
        @test isless(Interval(Date(2014), Date(2016), false, false), Date(2016))
        @test Interval(Date(2014), Date(2016), false, false) < Date(2016)
        @test Interval(Date(2014), Date(2016), false, false) ≪ Date(2016)
        @test isless(Interval(Date(2014), Date(2016)), Date(2017))
        @test Interval(Date(2014), Date(2016)) < Date(2017)
    end

    @testset "broadcasting" begin
        # Validate that an Interval is treated as a scalar during broadcasting
        interval = Interval(DateTime(2016, 8, 11, 17), DateTime(2016, 8, 11, 18))
        @test size(interval .== interval) == ()
    end

    @testset "sort" begin
        i1 = 1 .. 10
        i2 = Interval(1, 10, false, false)
        i3 = 2 .. 11

        @test sort([i1, i2, i3]) == [i1, i2, i3]
        @test sort([i1, i2, i3]; rev=true) == [i3, i2, i1]
    end

    @testset "arithmetic" begin
        for (a, b, unit) in test_values
            for i in 0:3
                inc = Inclusivity(i)
                interval = Interval(a, b, inc)
                @test interval + unit == Interval(a + unit, b + unit, inc)
                @test unit + interval == Interval(a + unit, b + unit, inc)
                @test interval - unit == Interval(a - unit, b - unit, inc)

                if a isa Number
                    @test -interval == Interval(-b, -a, last(inc), first(inc))
                    @test unit - interval == Interval(unit-b, unit-a, last(inc), first(inc))
                else
                    @test_throws MethodError -interval
                    @test_throws MethodError unit - interval
                end
            end

            @test_throws MethodError Interval(a, b) + Interval(a, b)
            @test_throws MethodError Interval(a, b) - Interval(a, b)
        end

        # DST transition
        firstpoint = ZonedDateTime(2018, 3, 11, 1, tz"America/Winnipeg")
        endpoint = ZonedDateTime(2018, 3, 11, 3, tz"America/Winnipeg")
        interval = Interval(firstpoint, endpoint) + Hour(1)
        @test first(interval) == ZonedDateTime(2018, 3, 11, 3, tz"America/Winnipeg")
        @test last(interval) == ZonedDateTime(2018, 3, 11, 4, tz"America/Winnipeg")

        firstpoint = ZonedDateTime(2018, 11, 4, 0, tz"America/Winnipeg")
        endpoint = ZonedDateTime(2018, 11, 4, 2, tz"America/Winnipeg")
        interval = Interval(firstpoint, endpoint) + Hour(1)
        @test first(interval) == ZonedDateTime(2018, 11, 4, 1, tz"America/Winnipeg", 1)
        @test last(interval) == ZonedDateTime(2018, 11, 4, 3, tz"America/Winnipeg")
    end

    @testset "isempty" begin
        for T in [Int32, Int64, Float64, Date, DateTime, ZonedDateTime]
            @test isempty(Interval{T}())
        end

        @test !isempty(Interval(0, 1, Inclusivity(false, false)))
        @test !isempty(Interval(0, 1, Inclusivity(false, true)))
        @test !isempty(Interval(0, 1, Inclusivity(true, false)))
        @test !isempty(Interval(0, 1, Inclusivity(true, true)))

        @test isempty(Interval(0, 0, Inclusivity(false, false)))
        @test isempty(Interval(0, 0, Inclusivity(false, true)))
        @test isempty(Interval(0, 0, Inclusivity(true, false)))
        @test !isempty(Interval(0, 0, Inclusivity(true, true)))

        # DST transition
        @test !isempty(
            Interval(
                ZonedDateTime(2018, 11, 4, 1, tz"America/Winnipeg", 1),
                ZonedDateTime(2018, 11, 4, 1, tz"America/Winnipeg", 2),
                Inclusivity(false, false),
            )
        )
    end

    @testset "in" begin
        for (a, b, unit) in test_values
            interval = Interval(a, b)
            @test in(a, interval)
            @test in(a + unit, interval)
            @test !in(a - unit, interval)
            @test in(b, interval)
            @test in(b - unit, interval)
            @test !in(b + unit, interval)

            interval = Interval(a, b, Inclusivity(true, false))
            @test in(a, interval)
            @test in(a + unit, interval)
            @test !in(a - unit, interval)
            @test !in(b, interval)
            @test in(b - unit, interval)
            @test !in(b + unit, interval)

            interval = Interval(a, b, Inclusivity(false, true))
            @test !in(a, interval)
            @test in(a + unit, interval)
            @test !in(a - unit, interval)
            @test in(b, interval)
            @test in(b - unit, interval)
            @test !in(b + unit, interval)

            interval = Interval(a, b, Inclusivity(false, false))
            @test !in(a, interval)
            @test in(a + unit, interval)
            @test !in(a - unit, interval)
            @test !in(b, interval)
            @test in(b - unit, interval)
            @test !in(b + unit, interval)

            @test_throws ArgumentError (in(Interval(a, b), Interval(a, b)))
        end
    end

    @testset "issubset" begin
        @test 0..10 ⊆ 0..10
        @test 0..10 ⊇ 0..10
        @test Interval(0, 10, false, false) ⊆ 0..10
        @test Interval(0, 10, false, false) ⊉ 0..10
        @test 0..10 ⊈ Interval(0, 10, false, false)
        @test 0..10 ⊇ Interval(0, 10, false, false)
        @test 1..9 ⊆ 0..10
        @test 1..9 ⊉ 0..10
        @test 0..10 ⊈ 1..9
        @test 0..10 ⊇ 1..9
        @test 1..11 ⊈ 0..10
        @test 1..11 ⊉ 0..10
        @test -1..9 ⊈ 0..10
        @test -1..9 ⊉ 0..10
        @test 20..30 ⊈ 0..10
        @test 20..30 ⊉ 0..10
    end

    @testset "intersect" begin
        @testset "overlapping" begin
            a = Interval(-10, 5, Inclusivity(true, true))
            b = Interval(-2, 10, Inclusivity(true, true))
            @test intersect(a, b) == Interval(-2, 5, Inclusivity(true, true))
            @test intersect(b, a) == intersect(a, b)

            a = Interval(-10, 5, Inclusivity(true, false))
            b = Interval(-2, 10, Inclusivity(true, true))
            @test intersect(a, b) == Interval(-2, 5, Inclusivity(true, false))
            @test intersect(b, a) == intersect(a, b)

            a = Interval(-10, 5, Inclusivity(true, true))
            b = Interval(-2, 10, Inclusivity(false, true))
            @test intersect(a, b) == Interval(-2, 5, Inclusivity(false, true))
            @test intersect(b, a) == intersect(a, b)

            a = Interval(-10, 5, Inclusivity(true, false))
            b = Interval(-2, 10, Inclusivity(false, true))
            @test intersect(a, b) == Interval(-2, 5, Inclusivity(false, false))
            @test intersect(b, a) == intersect(a, b)
        end

        @testset "adjacent" begin
            a = Interval(-10, 0, Inclusivity(true, true))
            b = Interval(0, 10, Inclusivity(true, true))
            @test intersect(a, b) == Interval(0, 0, Inclusivity(true, true))
            @test intersect(b, a) == intersect(a, b)

            a = Interval(-10, 0, Inclusivity(true, false))
            b = Interval(0, 10, Inclusivity(true, true))
            @test isempty(intersect(a, b))
            @test isempty(intersect(b, a))

            a = Interval(-10, 0, Inclusivity(true, true))
            b = Interval(0, 10, Inclusivity(false, true))
            @test isempty(intersect(a, b))
            @test isempty(intersect(b, a))

            a = Interval(-10, 0, Inclusivity(true, false))
            b = Interval(0, 10, Inclusivity(false, true))
            @test isempty(intersect(a, b))
            @test isempty(intersect(b, a))
        end

        @testset "identical" begin
            for inclusivity in Inclusivity.(0:3)
                x = Interval(1, 10, inclusivity)
                @test intersect(x, x) == x
            end

            x = Interval(0, 0, Inclusivity(false, false))
            @test intersect(x, x) == x
            @test isempty(intersect(x, x))

            # But what if their inclusivities are different?
            expected = Interval(1, 10, Inclusivity(false, false))
            @test intersect(
                Interval(1, 10, Inclusivity(true, true)),
                Interval(1, 10, Inclusivity(false, false)),
            ) == expected
            @test intersect(
                Interval(1, 10, Inclusivity(true, false)),
                Interval(1, 10, Inclusivity(false, true)),
            ) == expected
        end

        @testset "disjoint" begin
            for inclusivity in Inclusivity.(0:3)
                a = Interval(-100, -1, inclusivity)
                b = Interval(1, 100, inclusivity)
                @test isempty(intersect(a, b))
                @test isempty(intersect(b, a))
            end
        end
    end

    @testset "astimezone" begin
        zdt1 = ZonedDateTime(2013, 2, 13, 0, 30, tz"America/Winnipeg")
        zdt2 = ZonedDateTime(2016, 8, 11, 21, tz"America/Winnipeg")

        for inclusivity in Inclusivity.(0:3)
            for tz in (tz"America/Winnipeg", tz"America/Regina", tz"UTC")
                @test isequal(
                    astimezone(Interval(zdt1, zdt2, inclusivity), tz),
                    Interval(astimezone(zdt1, tz), astimezone(zdt2, tz), inclusivity),
                )
            end
        end
    end

    @testset "merge" begin
        a = Interval(-100, -1)
        b = Interval(-3, 10)

        @test merge(a, b) == Interval(-100, 10)
        @test merge(b, a) == Interval(-100, 10)

        b = Interval(1, 10)
        @test_throws ArgumentError merge(a, b)

        a = Interval(-100, -1, Inclusivity(false, false))
        b = Interval(-2, 10, Inclusivity(true, true))
        @test merge(a, b) == Interval(-100, 10, Inclusivity(false, true))

        a = Interval(-100, -1, Inclusivity(true, false))
        b = Interval(-2, 10, Inclusivity(false, false))
        @test merge(a, b) == Interval(-100, 10, Inclusivity(true, false))
    end

    @testset "union" begin
        intervals = [
            Interval(-100, -1, Inclusivity(false, false)),
            Interval(-10, -1, Inclusivity(false, false)),
            Interval(10, 15, Inclusivity(false, false)),
            Interval(13, 20, Inclusivity(false, false))
        ]
        expected = [
            Interval(-100, -1, Inclusivity(false, false)),
            Interval(10, 20, Inclusivity(false, false))
        ]
        @test union(intervals) == expected

        # Ordering
        intervals = [
            Interval(-100, -1, Inclusivity(false, false)),
            Interval(10, 15, Inclusivity(false, false)),
            Interval(-10, -1, Inclusivity(false, false)),
            Interval(13, 20, Inclusivity(false, false))
        ]
        @test union(intervals) == expected
        @test intervals == [
            Interval(-100, -1, Inclusivity(false, false)),
            Interval(10, 15, Inclusivity(false, false)),
            Interval(-10, -1, Inclusivity(false, false)),
            Interval(13, 20, Inclusivity(false, false))
        ]

        @test union!(intervals) == expected
        @test intervals == expected

        # Inclusivity
        intervals = [
            Interval(-100, -1, Inclusivity(false, false)),
            Interval(-10, -1, Inclusivity(true, true))
        ]
        @test union(intervals) == [Interval(-100, -1, Inclusivity(false, true))]
    end
end


text:
# Solves the steady state Williamson2 test case for the shallow water equations on a sphere
# of physical radius 6371220m. Involves a modified coriolis term that exactly balances
# the potential gradient term to achieve a steady state
# reference:
# D. L. Williamson, J. B. Drake, J. J.HackRüdiger Jakob, P. N.Swarztrauber, (1992)
# J Comp. Phys. 102 211-224

# Constants of the Williamson2 test case
const α  = π/4.0              # deviation of the coriolis term from zonal forcing
const U₀ = 38.61068276698372  # velocity scale
const H₀ = 2998.1154702758267 # mean fluid depth

# Modified coriolis term
function f₀(xyz)
   θϕr   = xyz2θϕr(xyz)
   θ,ϕ,r = θϕr
   2.0*Ωₑ*( -cos(θ)*cos(ϕ)*sin(α) + sin(ϕ)*cos(α) )
end

# Initial velocity
function u₀(xyz)
  θϕr   = xyz2θϕr(xyz)
  θ,ϕ,r = θϕr
  u     = U₀*(cos(ϕ)*cos(α) + cos(θ)*sin(ϕ)*sin(α))
  v     = -U₀*sin(θ)*sin(α)
  spherical_to_cartesian_matrix(θϕr)⋅VectorValue(u,v,0)
end

# Initial fluid depth
function h₀(xyz)
  θϕr   = xyz2θϕr(xyz)
  θ,ϕ,r = θϕr
  h  = -cos(θ)*cos(ϕ)*sin(α) + sin(ϕ)*cos(α)
  H₀ - (rₑ*Ωₑ*U₀ + 0.5*U₀*U₀)*h*h/g
end

# Topography
function topography(xyz)
  0.0
end

max_stars_repo_name: santiagobadia/GridapGeosciences

max_issues_repo_name: santiagobadia/GridapGeosciences

max_forks_repo_name: santiagobadia/GridapGeosciences


text:
function main()
    println("Hello, World!!!")
end

main()


text:
# This file is part of EndemicModel project, which is licensed under BDS-3-Clause.
# See file LICENSE.md for more information.

module EndemicModel

import Base.Iterators
import CSV
import JSON
import Measures: mm
import Plots
import ZipFile

using ArgCheck
using Base
using Base.Iterators
using Dates
using DataFrames
using DataFramesMeta
using DataStructures
using DifferentialEquations
using FilePaths
using GZip
using Interpolations
using Optim
using StatsBase
#using Unitful

export DataDict,
    DataFrameDict,
    AbstractDataDict,
    AbstractDataSource,
    DataSourceDesignator,
    AbstractDataPath,
    CsseSource,
    CsvDirectory,
    CsvPath,
    OdsPath,
    option,
    option!,
    get_parameters,
    to_json,
    DATA_SOURCES,
    import_data,
    export_data,
    Database,
    @defcolumn,
    @deftable,
    @defgroup,
    insert_column!,
    insert_table!,
    group_table!,
    column_creator,
    table_creator,
    table_grouper,
    SEIR_PARAMS,
    SEIR_DERIV,
    SEIR_VARS,
    SEIRVariables,
    SEIRDerivatives,
    SEIRParameters,
    SEIRModel,
    SEIRModel!,
    pack_vars,
    pack_params,
    unpack_vars,
    unpack_params,
    SEIR_ODE_fun,
    SEIR_ODEProblem,
    plot,
    model_problem,
    model_solution,
    to_dataframe,
    to_dataframe!,
    model_validate,
    model_step,
    datadict,
    datadict!,
    modeldict,
    modeldict!,
    paramdict,
    paramdict!,
    variables,
    variables!,
    derivatives,
    derivatives!,
    parameters,
    parameters!,
    modeldata,
    modeldata!,
    realdata,
    realdata!,
    phuber_loss,
    diff_phuber_loss,
    l2_loss,
    diff_l2_loss,
    model_loss,
    optimize_parameters,
    optimize_parameters!,
    optimize_variables,
    optimize_variables!,
    estimate_μ,
    estimate_α,
    estimate_γ,
    estimate_β,
    estimate_exposed!,
    covid19_database,
    covid19

const DATA_SOURCES = Dict{Symbol,Type}()

include("utils.jl")
include("options.jl")
include("source.jl")
include("sources/general.jl")
include("sources/covid19.jl")
include("database.jl")
include("subdatabase.jl")
include("database_files.jl")
include("covid19.jl")
include("model.jl")
include("plot.jl")
include("estimate.jl")

end # module


text:
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