text:
using SpecialFunctions
using Test
using LinearAlgebra
using Base.MathConstants
using SymPy
using SymPy.Introspection
import PyCall

@testset "Core" begin
    ## Symbol creation
    x = Sym("x")
    #x = sym"x" # deprecated
    x = Sym(:x)
     x,y = Sym(:x, :y)
    x,y = symbols("x,y")

    @vars u1 u2 u3
    @vars u positive=true
    @test length(solve(u+1)) == 0
    # make sure @vars defines in a local scope
    let
        @vars w
    end
    @test_throws UndefVarError isdefined(w)
    @vars a b c

    # Renaming with @vars
    @vars a=>"α₁"
    @test a.name == "α₁"

    ## extract symbols
    @vars z
    ex = x*y*z
    @test isa(free_symbols(ex), Vector{Sym})
    @test free_symbols(ex) == [x, y, z]

    ## number conversions
    @test Sym(2) == 2
    @test Sym(2.0) == 2.0
    @test Sym(2//1) == 2
    @test Sym(im) == 1im
    @test Sym(2im) == 2im
    @test Sym(1 + 2im) == 1 + 2im

    pi, e, catalan = Base.MathConstants.pi, Base.MathConstants.e, Base.MathConstants.catalan
    @test N(Sym(pi)) == pi
    @test N(Sym(ℯ)) ==  ℯ
    @test N(Sym(catalan)) == catalan

    ## function conversion
    f1 = convert(Function, x^2)
    @test f1(2) == 4

    ### Subs
    ## subs, |> (x == number)
    f(x) = x^2 - 2
    y = f(x)
    @test float(y.subs( x, 1)) == f(1)
    ### @test float( y |> subs(x,1) ) == f(1) no subs method

    ## interfaces
    ex = (x-1)*(y-2)
    @test ex.subs(x, 1) == 0
    @test ex.subs(((x,1),)) == 0
    @test ex.subs(((x,2),(y,2))) == 0

    @test subs(ex, x=>1) == 0       # removed
    @test subs(ex, x=>2, y=>2) == 0
    @test subs(ex, Dict(x=>1)) == 0
    @test ex(x=>1) == 0
    @test ex(x=>2, y=>2) == 0
    @test ex.subs(Dict(x=>1)) == 0

    ### doit
    @syms x f() g()
    D = Differential(x)
    df = D(f(x))
    dfx = subs(df, f(x), x^2)
    @test dfx.doit() == 2*x
    @test doit(dfx) == 2*x
    @test dfx |> doit == 2*x
    # use deep=true to force nested evaluations
    dgfx = g(dfx)
    @test dgfx.doit(deep=true) == g(2*x)
    @test doit(dgfx, deep=true) == g(2*x)
    @test dgfx |> doit(deep=true) == g(2*x)

    ## match, replace, xreplace, rewrite
    x,y,z = symbols("x, y, z")
    a,b,c = map(Wild, (:a,:b,:c))
    ## match: we have pattern, expression to follow `match`
    d = match(a^a, (x+y)^(x+y))
    @test d[a] == x+y
    d = match(a^b, (x+y)^(x+y))
    @test d[b] == x + y
    ex = (2x)^2
    pat = a*b^c
    d = match(pat, ex)
    @test d[a] == 4 && d[b] == x && d[c] == 2
    @test pat.xreplace(d) == 4x^2

    ## replace
    a = Wild("a")
    ex = log(sin(x)) + tan(sin(x^2))
    ##XXX    @test replace(ex, func(sin(x)), u ->  sin(2u)) == log(sin(2x)) + tan(sin(2x^2))
    @test replace(ex, func(sin(x)), func(tan(x))) == log(tan(x)) + tan(tan(x^2))
    @test replace(ex, sin(a), tan(a)) ==  log(tan(x)) + tan(tan(x^2))
    @test replace(ex, sin(a), a) == log(x) + tan(x^2)
    @test replace(x*y, a*x, a) == y

    ## xreplace
    @test (1 + x*y).xreplace(Dict(x => PI)) == 1 + PI*y
    @test (x*y + z).xreplace(Dict(x*y => PI)) == z + PI
    @test (x*y * z).xreplace(Dict(x*y => PI)) == x* y * z
    @test (x + 2 + exp(x + 2)).xreplace(Dict(x+2=>y)) == x + exp(y) + 2

    # Test subs on simple numbers
    @vars x y
    @test Sym(2)(x=>300, y=>1.2) == 2

    #Test subs for pars and dicts
    ex = 1
    dict1 = Dict{String,Any}()
    dict2 = Dict{Any,Any}()
    #test subs
    for i=1:4
        x = Sym("x$i")
        ex=ex*x
        dict2[x] = i
        dict1[string(x)] = i
    end
    for d in [dict1, dict2]
        @test ex |> subs(d) == factorial(4)
        @test subs(ex, d) == factorial(4)
        @test subs(ex, d...) == factorial(4)
        @test ex |> subs(d...) == factorial(4)
        @test ex(d) == factorial(4)
        @test ex(d...) == factorial(4)
    end

    a = Sym("a")
    b = Sym("b")
    line = x -> a + b * x
    sol = solve([line(0)-1, line(1)-2],[a,b])
    ex = line(10)
    @test ex(sol) == ex(sol...) == 11

    ## Simplify (issue 343)
    @vars x
    @test simplify(sin(x)^2 + cos(x)^2) == 1
    @test simplify(sympy.gamma(x)/sympy.gamma(x-2)) == (x-1)*(x-2)
    _ones = (1, 1.0, big"1", "1", one)
    @test simplify.(_ones) == _ones

    ## Conversion
    x = Sym("x")
    p = x.subs(x,pi)
    q = x.subs(x,Sym(1)/2)
    r = x.subs(x,1.2)
    z = x.subs(x,1)
    @test isa(N(p), Float64)
##XXX    @test isa(N(p, 60), BigFloat)
    @test isa(p.evalf(), Sym)
    @test isa(N(x), Sym)
    @test isa(N(q), Rational)
    @test isa(N(r), Float64)
    @test isa(N(z), Integer)

    ## method calls via getproperty
    p = (x-1)*(x-2)
    @test sympy.roots(p) == Dict{Any,Any}(Sym(1) => 1, Sym(2)=> 1) # sympy.roots
    p = sympy.Poly(p, x)
    @test p.coeffs() == Any[1,-3,2] # p.coeffs

    ## algebra
    @test expand((x + 1)*(x + 2)) == x^2 + 3x + 2  # v0.7 deprecates expand, in v1.0 this is fine w/o qualifacation
    x1 = (x + 1)*(x + 2)
    @test expand(x1) == x^2 + 3x + 2
    @test sympy.expand_trig(sin(2x)) == 2sin(x)*cos(x)

    ## math functions
    u = abs(x^2 - 2)
    @test u(x=>0) == 2
    u = min(x, x^2, x^3, x^4)
    @test u(x=>2) == 2
    @test u(x=>1//2) == 1//2^4

    ## solve
    x,y,a = symbols("x,y,a", real=true)
    solve(x^2 - 2x)
    solve(x^2 - 2a, x)
    solve(x^2 - 2a, a)
    solve(Lt(x-2, 0))
    solve( x-2 ≪ 0)
    exs = [x-y-1, x+y-2]
    di = solve(exs)
    @test di[x] == 3//2
    @test map(ex -> subs(ex, di), exs) == [0,0]
    solve([x-y-a, x+y], [x,y])

    ## linsolve
    M=Sym[1 2 3; 2 3 4]
    as = linsolve(M, x, y)
    @test length(elements(as)) == 1
    @vars a b; eqs = (a*x+2y-3, 2b*x + 3y - 4)
    as = linsolve(eqs, x, y)
    @test length(elements(as)) == 1

    ## nsolve -- not method for arrays, issue 268
    @vars z1 z2z1 positive=true
    @test_throws MethodError nsolve([z1^2-1, z1+z2z1-2], [z1,z2z1], [1,1])  # non symbolic first argument
    @test all(N.(sympy.nsolve([z1^2-1, z1+z2z1-2], [z1,z2z1], (1,1))) .≈ [1.0, 1.0])

    ## issue 355: direct definition of LinearAlgebra.:\
    @test_throws SingularException Sym[1 1; 1 1] \ [1, 2]
    @vars  a b c d e f
    A, b= [a b; c d], [e, f]
    x = A \ b
    @test  simplify.(A*x-b) == [0,0]
    out =  Sym[1 1; 1 1] \ [1,1]
    u =   free_symbols(out)[1]
    @test out ==  [1-u,u]

    # Just a made-up example to test  if manageable
    @vars a1 a2
    n = 7
    A = diagm(0 => ones(Int, n), 1 => fill(a1, n-1), 2 => fill(1, n-2), -1 => fill(a2, n-1))
    b = Vector{Sym}(1:n)
    x = A \ b
    @test length(free_symbols(x)) ==  2

    ## limits
    @vars x
    @test limit(x -> sin(x)/x, 0) == 1
    @test limit(sin(x)/x, x, 0) |> float == 1
    @test limit(sin(x)/x, x => 0) == 1
    @vars x h
    out = limit((sin(x+h) - sin(x))/h, h, 0)
    @test (out.replace(x, pi) |> float) == -1.0

    ## diff
    diff(sin(x), x)
    out = diff(sin(x), x, 2)
    @test abs((out.replace(x, pi/4) |> float) - - sin(pi/4)) < sqrt(eps())

    # partial derivatives
    @vars x y
    @test diff(x^2 + x*y^2, x, 1) == 2x + y^2

    t = symbols("t", real=true)     # vector-valued functions
    r1(t) = [sin(t), cos(t), t]
    u = r1(t)
    kappa = norm(diff.(u) × diff.(u,t,2)) / norm(diff.(u))^3 |> simplify
    @test convert(Rational,kappa) == 1//2

    u = SymFunction("u")
    eqn = Eq(x^2 + u(x)^2, x^3 - u(x))
    # diff(eqn) doesn't evaluate over Eq:
    @test func(eqn)(diff.(args(eqn))...) == Eq(2x + 2u(x) * diff(u(x),x), 3x^2 - diff(u(x),x))

    ## integrate
    @test integrate(sin(x)) == -cos(x)
    @test integrate(sin(x), (x, 0, pi)) == 2.0
    a, b, t = symbols("a, b, t")
    @test integrate(sin(x), (x, a, b)) == cos(a) - cos(b)
    @test integrate(sin(x), (x, a, b)).replace(a, 0).replace(b, pi) == 2.0
    @test integrate(sin(x) * DiracDelta(x)) == sin(Sym(0))
    @test integrate(Heaviside(x), (x, -1, 1)) == 1
    curv = sympy.Curve([exp(t)-1, exp(t)+1], (t, 0, log(Sym(2))))
    @test line_integrate(x + y, curv, [x,y]) == 3 * sqrt(Sym(2))

    ## summation
    summation(1/x^2, (x, 1, 10))
    out = summation(1/x^2, (x, 1, 10))
    out1 = sum([1//x^2 for  x in 1:10])
    @test round(Integer, out.p) == out1.num
    @test round(Integer, out.q) == out1.den

    ## Ops
    s = 3
    x = Sym("x")
    v = [x, 1]
    rv = [x 1]
    a = [x 1; 1 x]
    b = [x 1 2; 1 2 x]

    DIMERROR =  DimensionMismatch
    DimensionOrMethodError =  Union{MethodError, DimensionMismatch}
    ## scalar, [vector, matrix]
    @test s .+ v == [x+3, 4]
    @test v .+ s == [x+3, 4]
    @test s .+ rv == [x+3 4]
    @test rv .+ s == [x+3 4]
    @test s .+ a == [x+3 4; 4 x+3]
    @test a .+ s == [x+3 4; 4 x+3]

    @test s .- v == [3-x, 2]
    @test v .- s == [x-3, -2]
    @test s .- rv == [3-x 2]
    @test rv .- s == [x-3 -2]
    @test s .- a == [3-x 2; 2 3-x]
    @test a .- s == [x-3 -2; -2 x-3]

    2v
    2rv
    2a
    s .* v
    v .* s
    s .* rv
    rv .* s
    s .* a
    a .* s

    ## s / v  ## broadcasts s Depreacated
    @test s ./ v == [3/x, 3]
    @test v / s == [x/3, Sym(1)/3]
    @test v .\ s == s ./ v
    @test s \ v == v / s
    ## s / rv ## broadcasts s Deprecated
    @test s ./ rv == [3/x 3]
    @test rv / s == [x/3 Sym(1)/3]
    @test rv .\ s == s ./ rv
    @test s \ rv == rv / s
    ## s / a ## broadcasts s Deprecated
    @test s ./ a == [3/x 3; 3 3/x]
    @test a / s == [x/3 Sym(1)/3; Sym(1)/3 x/3]
    @test a .\ s == s ./ a
    @test s \ a == a / s

    @test_throws MethodError  s ^ v ## error
    @test s .^ v == [3^x, 3]
    @test_throws DimensionOrMethodError v ^ s ## error
    v .^ s
    @test_throws MethodError  s ^ rv ## error
    s .^ rv
    @test_throws DimensionMismatch  rv ^ s ## error
    rv .^ s
    @test_throws MethodError  s ^ a ## error
    s .^ a
    a ^ s
    a .^ s

    ## vector vector
    @test v .+ v == 2v
    ##@test_throws MethodError  v .+ rv ##  no longer an error, broadcase
    @test v .- v == [0, 0]
    @test_throws DIMERROR  v - rv
    @test_throws DimensionOrMethodError   v * v ## error
    @test v .* v == [x^2,1]
    @test dot(v, v) == 1 + conj(x)*x
    v * rv
    rv * v ## 1x2 2 x 1 == 1x1
    v .* rv ## XXX ?? should be what? -- not 2 x 2
    rv .* v ## XXX ditto
    ## @test_throws MethodError  v / v ## error
    v ./ v ## ones()
    v .\ v
    ## @test_throws MethodError  v / rv ## error
    v ./ rv  ## ??
    rv .\ v
    @test_throws MethodError  v ^ v ## error
    v .^ v
    @test_throws MethodError  v ^ rv ## error
    v .^ rv ## ??

    ## vector matrix
    @test_throws DIMERROR v + a ## error (Broadcast?)
    @test_throws DIMERROR a + v ## error
    v .+ a ## broadcasts
    a .+ v
    @test_throws DIMERROR  v - a ## error
    v .- a
    @test_throws DimensionMismatch  v * a ## error
    v .* a
    #@test_throws MethodError  v / a ## error
    v ./ a
    a .\ v
    @test_throws MethodError  v ^ a ## error
    v .^ a
    v
    ## matrix matrix
    a + a
    @test_throws DIMERROR  a + b ## error
    a + 2a
    a - a
    @test_throws DIMERROR  a - b ## error
    a * a
    a .* a
    a * b ## 2x2 * 2*3 -- 2x3
    @test_throws DIMERROR  a .* b ## error -- wrong size
    #@test_throws MethodError  a / a
    a ./ a ## ones
    a .\ a
    ##@test_throws MethodError  a / b ## error
    @test_throws DIMERROR  a ./ b ## error
    @test_throws DIMERROR  b .\ a ## error
    @test_throws MethodError  a ^ a ## error
    a .^ a
    @test_throws MethodError  a ^ b ## error
    @test_throws DIMERROR  a .^ b ## error

    ## Number theory
    #@test isprime(100) == isprime(Sym(100))
    #@test factorint(Sym(100)) == factor(100)
    @test prime(Sym(100)) == 541
    @test multiplicity(Sym(10), 100) == 2

    ## polynomials
    @vars x y
    f1 = 5x^2  + 10x + 3
    g1 = 2x + 2
    q,r = sympy.div(f1,g1, domain="QQ") # use sympy.div to dispatch; o/w we can't disambiguate div(Sym(7), 5)) to do integer division
    @test r == Sym(-2)
    @test simplify(q*g1 + r - f1) == Sym(0)
    ## sympy.interpolate as first arg is not symbolic
    @test sympy.interpolate([1,2,4], x) == sympy.interpolate(collect(zip([1,2,3], [1,2,4])), x)
    @test sympy.interpolate(collect(zip([-1,0,1], [0,1,0])), x) == 1 - x^2

    ## piecewise
    x = Sym("x")
    # sympy.Piecewise is a FunctionClass,  we qualify, as args not Symbolic
    p = sympy.Piecewise((x, Ge(x,0)), (0, Lt(x,0)), (1, Eq(x,0)))
    ## using infix \ll<tab>, \gt<tab>, \Equal<tab>
    p = sympy.Piecewise((x, (x ≫ 0)), (0, x ≪ 0), (1, x ⩵ 0))
    @test p.subs(x,2) == 2
    @test p.subs(x,-1) == 0
    @test p.subs(x,0) == 1

    ## if VERSION < v"0.7.0-" # ifelse changed
    ##     u = ifelse(Lt(x, 0), "neg", ifelse(Gt(x, 0), "pos", "zero"))
    ##     @test u.subs(u,x,-1) == Sym("neg")
    ##     @test subs(u,x, 0) == Sym("zero")
    ##     @test subs(u,x, 1) == Sym("pos")
    ## end
    p = sympy.Piecewise((-x, x ≪ 0), (x, x ≧ 0))

    ## relations
    x,y=symbols("x, y")
    ex  = Eq(x^2, x)
    @test ex.lhs() == x^2
    @test ex.rhs() == x
    @test args(ex) == (x^2, x)

    # alternative operators
    for (ex1,ex2) ∈ ((Eq(x^2, x), x^2 ⩵ x),
                   (Eq(x^2, x), x^2 ~ x),
                   (Lt(x^2, x), x^2 ≪ x),
                   (Le(x^2, x), x^2 ≦ x),
                   (Ge(x^2, x), x^2 ≧ x),
                   (Gt(x^2, x), x^2 ≫ x))
        @test lhs(ex1) == lhs(ex2)
        @test rhs(ex1) == rhs(ex2)
    end

    ## mpmath functions
#    if @isdefined mpmath
    if isdefined(SymPy, :mpmath)
        x = Sym("x")
        Sym(big(2))
        Sym(big(2.0))                   # may need mpmath (e.g., conda install mpmath)

        sympy.__version__ != "1.9" && (@test limit(besselj(Sym(1),1/x), x, 0) == Sym(0))
        complex(N(SymPy.mpmath.hankel2(2, pi)))
        SymPy.mpmath.bei(2, 3.5)
        SymPy.mpmath.bei(1+im, 2+3im)
    end

    ## Assumptions
    @test ask(𝑄.even(Sym(2))) == true
    @test ask(𝑄.even(Sym(3))) == false
    @test ask(𝑄.nonzero(Sym(3))) == true
    @vars x_real real=true
    @vars x_real_positive real=true positive=true
    @test ask(𝑄.positive(x_real)) == nothing
    @test ask(𝑄.positive(x_real_positive)) == true
    @test ask(𝑄.nonnegative(x_real^2)) == true
    @test ask(𝑄.upper_triangular([x_real 1; 0 x_real])) == true
    @test ask(𝑄.positive_definite([x_real 1; 1 x_real])) == nothing

              ## sets
    s = sympy.FiniteSet("H","T")
    s1 = s.powerset()
    VERSION >= v"0.4.0" && @test length(collect(convert(Set, s1))) == length(collect(s1.__pyobject__))
    a, b = sympy.Interval(0,1), sympy.Interval(2,3)
    @test a.is_disjoint(b) == true
    @test a.union(b).measure() == 2

    ## test cse output
    @test cse(x) == (Any[], Sym[x])
    @test sympy.cse([x]) == (Any[], [reshape([x],1,1)])
    @test sympy.cse([x, x]) == (Any[],  [reshape([x,x], 2, 1)] )
    @test sympy.cse([x x; x x]) == (Any[], [[x x; x x]])

    ## sympy"..."(...)
    ## removed
    #@vars x
    #@test sympy"sin"(1) == sin(Sym(1))
end

@testset "Syms macro" begin
    @syms u
    @test isa(u, Sym)

    ret = @syms a, b, c
    @test isa(ret, Tuple{Sym, Sym, Sym})

    @syms x::(real,positive)=>"x₀", y, z::complex, n::integer
    @test isa(x, Sym)
    @test ask(And(𝑄.real(x), 𝑄.positive(x)))
    @test string(x) == "x₀"

    @test isa(y, Sym)

    @test isa(z, Sym)
    @test ask(𝑄.complex(z))

    @test isa(n, Sym)
    @test ask(𝑄.integer(n))

    @syms f()::(real, positive), g(), h()::complex=>"h̄"
    @test isa(f, SymFunction)
    @test ask(And(𝑄.real(f(x)), 𝑄.positive(f(x))))

    @test isa(g, SymFunction)

    @test isa(h, SymFunction)
    @test ask(𝑄.complex(h(x)))
    @test string(h) == "h̄"

    @syms X[1:20]
    @test isa(X, Vector{Sym})
    @test size(X) == (20,)
    @test string(X[11]) == "X₁₁"

    @syms bigy[1:5]=>"Y"
    @test string(bigy[3]) == "Y₃"

    @syms Z[1:5, 1:6]
    @test isa(Z, Matrix{Sym})
    @test size(Z) == (5, 6)
    @test string(Z[2,4]) == "Z₂_₄"

    @syms F[1:2](), G()[1:2]
    @test isa(F, Vector{SymFunction})
    @test isa(G, Vector{SymFunction})

    @syms WOW[1:3, 1:2:4]()::(real, positive)=>"f"
    @test isa(WOW, Matrix{SymFunction})
    @test size(WOW) == (3, 2)
    @test ask(And(𝑄.real(WOW[1,2](x)), 𝑄.positive(WOW[1,2](x))))
    @test string(WOW[1,2]) == "f₁_₃"
end

@testset "SymFunctions" begin
    @syms x::real y::real
    @symfuns f g real=true

    @test isreal(f(x))
    @test isreal(g(y))

    @symfuns h real=true positive=true
    @test h(x)>0
end

@testset "Fix past issues" begin
    @vars x y z
    ## Issue # 56
    @test Sym(1+2im) == 1+2IM
    @test convert(Sym, 1 + 2im) == 1 + 2IM

    ## Issue #59
    cse(sin(x)+sin(x)*cos(x))
    sympy.cse([sin(x), sin(x)*cos(x)])
    sympy.cse( [sin(x), sin(x)*cos(x), cos(x), sin(x)*cos(x)])

    ## Issue #60, lambidfy
    x, y = symbols("x, y")
    lambdify(sin(x)*cos(2x) * exp(x^2/2))
    fn = lambdify(sin(x)*asin(x)*sinh(x)); fn(0.25)
    lambdify(real(x)*imag(x))
    @test lambdify(Min(x,y))(3,2) == 2

    ex = 2 * x^2/(3-x)*exp(x)*sin(x)*sind(x)
    fn = lambdify(ex); map(fn, rand(10))
    ex = x - y
    @test lambdify(ex, (x,y))(3,2) == 1

    Indicator(x, a, b) = sympy.Piecewise((1, Lt(x, b) & Gt(x,a)), (0, Le(x,a)), (0, Ge(x,b)))
    i = Indicator(x, 0, 1)
    u = lambdify(i)
    @test u(.5) == 1
    @test u(1.5) == 0

#    i2 = SymPy.lambdify_expr(x^2,name=:square)
#    @test i2.head == :function
#    @test i2.args[1].args[1] == :square
    ## @test i2.args[2] == :(x.^2) # too fussy

    ## issue #67
    @test N(Sym(4)/3) == 4//3
    @test N(convert(Sym, 4//3)) == 4//3

    ## issue #71
    @test log(Sym(3), Sym(4)) == log(Sym(4)) / log(Sym(3))

    ## issue #103 # this does not work for `x` (which has `classname(x) == "Symbol"`), but should work for other expressions
    for ex in (sin(x), x*y^2*x, sqrt(x^2 - 2y))
        @test SymPy.Introspection.func(ex)(SymPy.Introspection.args(ex)...) == ex
    end

    ## properties (Issue #119)
    @test (sympify(3).is_odd) == true
    @test sympy.Poly(x^2 -2, x).is_monic == true

    ## test round (Issue #153)
    y = Sym(eps())
    @test round(N(y), digits=5) == 0
    @test round(N(y), digits=16) != 0

    ## lambdify over a matrix #218
    @vars x y
    s = [1 0; 0 1]
    @test lambdify(x*s)(2) == 2 * s
    U = [x-y x+y; x+y x-y]
    @test lambdify(U, [x,y])(2,4) == [-2 6;6 -2]
    @test lambdify(U, [y,x])(2,4) == [ 2 6;6  2]

    @test eltype(lambdify([x 0; 1 x])(0)) <: Integer
    @test eltype(lambdify([x 0; 1 x], T=Float64)(0)) == Float64

    # issue 222 type of eigvals
    A = [Sym("a") 1; 0 1]
    @test typeof(eigvals(A)) <: Vector{Sym}

    # issue 231 Q.complex
    @vars x_maybecomplex
    @vars x_imag imaginary=true
    @test ask(𝑄.complex(x_maybecomplex)) == nothing
    @test ask(𝑄.complex(x_imag)) == true

    # issue 242 lambdify and conjugate
    @vars x
    expr = conjugate(x)
    fn = lambdify(expr)
    @test fn(1.0im) == 0.0 - 1.0im
    fn = lambdify(expr, use_julia_code=true)
    @test fn(1.0im) == 0.0 - 1.0im

    # issue 245 missing sincos
    @test applicable(sincos, x)
    @test sincos(x)[1] == sin(x)

    # issue 256 det
    @vars rho phi theta real=true
    xs = [rho*cos(theta)*sin(phi), rho*sin(theta)*sin(phi), rho*cos(phi)]
    J = [diff(x, u) for x in xs, u in (rho, phi, theta)]
    J.det()

    # issue #273 x[i]
    x = sympy.IndexedBase("x")
    i,j = sympy.symbols("i j", integer=True)
    @test x[i] == PyCall.py"$x[$i]"

    # issue 298 lambdify for results of dsolve
    @vars t
    F = SymFunction("F")
    diffeq = diff(F(t),t) - 3*F(t)
    res = dsolve(diffeq, F(t), ics=Dict(F(0) => 2))  # 2exp(3t)
    @test lambdify(res)(1) ≈ 2*exp(3*1)

    # issue 304 wrong values for sind, ...
    a = Sym(45)
    @test sind(a) == sin(PI/4)

    # issue #319  with   use   of  Dummy, but
    #  really  a  lambdify issue
    dummy = sympy.Dummy
    # Symbolic differentiation of functions
    function D(f)
        x = dummy("x")
        lambdify(diff.(f(x), x), (x,))
    end
    @test D(t -> t^2)(1) == 2

    # issue   #320  with integrate(f) when
    # f  is consant
    @test integrate(x -> 1, 0, 1)  == 1
    @test limit(x->1,  x, 0) == 1
    @test diff(x->1)  ==   0

    ## Issue 324 with inference of matrix operations
    A = fill(Sym("a"), 2, 2)
    @test eltype(A*A) == Sym
    @test eltype(A*ones(2,2)) == Sym
    @test eltype(A*Diagonal([1,1])) == Sym
    VERSION >= v"1.2.0"  && @test eltype(A * I(2)) == Sym

    ## Issue 328 with E -> e
    @vars x
    ex = 3 * sympy.E * x
    fn = lambdify(ex)
    @test fn(1) ≈ 3*exp(1) * 1

    ## Issue 332 with `abs2`
    @vars x real=true
    @test abs2(x) == x*x
    @vars x
    @test abs2(x) == x*conj(x)

    ## Issue 376 promote to Sym Before pycall
    f(x) = x^2 + 1 +log(abs( 11*x-15 ))/99
    @test limit(f, 15//11) == limit(f(x), x, 15//11) == limit(f(x), x=>15//11) == -oo

    ## Issue #390 on div (__div__ was depracated, use __truediv__)
    @test Sym(2):-Sym(2):-Sym(2) |> collect == [2, 0, -2]

    ## Lambda function to create a lambda
    @vars x
    ex = x^2 - 2
    fn1 = Lambda(x, ex)
    fn2 = lambdify(ex)
    @test fn1(3) == fn2(3)

    ## issue 402 with lamdify and Order
    @vars x
    t = series(exp(x), x, 0, 2)
    @test lambdify(t)(1/2) == 1 + 1/2

    ## issue #411 with Heaviside
    @vars t
    u = Heaviside(t)
    λ = lambdify(u)
    @test all((iszero(λ(-1)), isone(λ(1))))
    if parse(Float64, sympy.__version__) >= 1.9
        @test λ(0) == 1//2
    else
        @test isnan(λ(0))
    end
    u = Heaviside(t, 1)
    λ = lambdify(u)
    @test all((iszero(λ(-1)), isone(λ(0)), isone(λ(1))))

    ## issue #295 with piecewise function
    @syms x
    p = sympy.Piecewise((x,Gt(x,0)), (x^2, Le(x,0)))
    @test lambdify(p)(2) == 2
    @test lambdify(p)(-2) == (-2)^2

    ## Issue catch all for N; floats only
    @syms x
    ex = integrate(sqrt(1 + (1/x)^2), (x, 1/sympy.E, sympy.E))
    @test N(ex) ≈ 3.196198513599507
end

@testset "generic programming, issue 223" begin
    # arose in issue 223
    @vars xreal real=true
    @vars xcomplex
    zreal = sympify(1)
    zcomplex = sympify(1) + sympify(2)*IM

    @test isreal(xreal)     # is_real(xreal) is also true, but xreal is Sym, not a Julia object
    @test !isreal(xcomplex) # is_real(xcomplex) is nothing
    @test isreal(zreal)
    @test !isreal(zcomplex)

    # conversions
    @test complex(xreal) == xreal
    @test complex(xreal, xreal) == xreal + IM*xreal
    @test complex(xcomplex) != xcomplex
    @test complex(zreal) == zreal
    @test complex(zreal) !== zreal      # Complex{Int} !== Sym
    @test complex(zcomplex) == zcomplex
    @test complex(zcomplex) !== zcomplex

    ## issue 284 N(PI,50)
    @test N(PI, 50) ≈ pi
    @test length(string(N(PI,50))) == 2 + 50
    ## issue 284 lambdify of Pi
    mpi = SymPy.PyCall.pyimport("sympy.parsing.mathematica")."mathematica"("Pi")
    @test SymPy.walk_expression(mpi) == :pi
    @test lambdify(PI^4*xreal)(256) == 256 * pi^4

    ## Issue 351 booleans  and arithmetic operations
    @test Sym(1) + true == Sym(2) == true +  Sym(1)
    @test Sym(1) - true == Sym(0) == true -  Sym(1)
    @test Sym(1) * true == Sym(1) == true * Sym(1)
    @test Sym(1) / true == Sym(1) == true / Sym(1)
    @test true^Sym(1)   == Sym(1) == Sym(1)^true

    ## issue with `pycall_hasproperty` and nothing values.
    @test !SymPy.is_rational(Sym(2.5))

    ## Issue #405 with ambigous methods
    @vars α
    M = SymMatrix([1 2; 3 4])
    @test α * M == M * α
    @test 2 * M == M * 2
    @test isa(M/α, SymMatrix)
    @test isa(α * inv(M), SymMatrix)

    ## issue #408 with inv
    @vars n integer=true positive=true
    A = sympy.MatrixSymbol("A", n, n)
    @test inv(A) == A.I

    # ceil broken
    @syms x
    @test limit(ceil(x), x=>0, dir="+") != limit(ceil(x), x=>0, dir="-")
    @test limit(floor(x), x=>0, dir="+") != limit(floor(x), x=>0, dir="-")

    ## Issue #433 add sympy docstrings, clean up docstring
    sprint(io -> show(io, SymPy.Doc(:sin)))

end

@test SymPy.convert_expr(sympy.Indexed(sympy.IndexedBase(:x), 1, -2)) == :(x[1, -2])


text:
module SimplePovertyCounts

using DataFrames

export GroupPoverty, line_count, calc_child_poverty

   struct GroupPoverty{T<:Real}
        total :: T
        affected :: T
        prop :: T
    end

    function line_count(
        line    :: T,
        data  :: DataFrame;
        counter :: Symbol,
        measure :: Symbol,
        weight  :: Symbol ) :: GroupPoverty{T} where T <: Real
        tot :: Real = 0.0
        aff :: Real = 0.0
        for r in eachrow( data )
            inc = r[measure]
            w = r[weight]
            m = r[counter]
            wm = w*m
            tot += wm
            if inc <= line
                aff += wm
            end
        end
        return GroupPoverty{T}(tot, aff, aff/tot )
    end

    function calc_child_poverty(
        line    :: T,
        hhdata  :: DataFrame;
        measure :: Symbol ) :: GroupPoverty{T} where T <: Real
        return line_count( line, hhdata;
            counter = :num_children,
            measure = measure,
            weight = :weight )
    end

end

ext: jl

max_stars_repo_name: grahamstark/ScottishTaxBenefitModel

max_issues_repo_name: grahamstark/ScottishTaxBenefitModel

max_forks_repo_name: grahamstark/ScottishTaxBenefitModel


text:
q1 = nominal_configuration(hopper)
v1 = zeros(hopper.nq)

T = 100
h = 0.01

s = Simulator(hopper, T, h=h)
simulate!(s, q1, v1)

vis = Visualizer()
open(vis)
visualize!(vis, s)


text:
using DataFrames, Query

df1 = DataFrame(a=[1,2,3], b=[1.,2.,3.])
df2 = DataFrame(c=[2.,4.,2.], d=["John", "Jim","Sally"])

x = @from i in df1 begin
    @join j in df2 on i.a equals convert(Int,j.c) into k
    @where i.a>1
    @select {t1=i,t2=k}
    @collect DataFrame
end

println(x)


text:
function FunctionalDetector()
    return FunctionalDetector(())
end

function g(obj::FunctionalDetector, arg0::SpacecraftState)
    return jcall(obj, "g", jdouble, (SpacecraftState,), arg0)
end

function get_function(obj::FunctionalDetector)
    return jcall(obj, "getFunction", ToDoubleFunction, ())
end

function with_function(obj::FunctionalDetector, arg0::ToDoubleFunction)
    return jcall(obj, "withFunction", FunctionalDetector, (ToDoubleFunction,), arg0)
end


text:
@testset "Simulator: Quadruped" begin
    # Reference trajectory
	s = get_simulation("quadruped", "flat_2D_lc", "flat")
	model = s.model
	env = s.env
    model.μ_world = 0.5

    ref_traj = deepcopy(ContactImplicitMPC.get_trajectory(s.model, s.env,
		joinpath(module_dir(), "src/dynamics/quadruped/gaits/gait2.jld2"),
		load_type = :split_traj_alt))
    ContactImplicitMPC.update_friction_coefficient!(ref_traj, model, env)

    T = ref_traj.H
    h = ref_traj.h

    for t = 1:T
    	r = ContactImplicitMPC.residual(model, env, ref_traj.z[t], ref_traj.θ[t], 0.0)
    	@test norm(r) < 1.0e-4
    end

    # initial conditions
    q0 = SVector{model.dim.q}(ref_traj.q[1])
    q1 = SVector{model.dim.q}(ref_traj.q[2])

    # simulator
    sim = ContactImplicitMPC.simulator(s, q0, q1, h, T,
        p = ContactImplicitMPC.open_loop_policy([SVector{model.dim.u}(ut) for ut in ref_traj.u]),
        ip_opts = ContactImplicitMPC.InteriorPointOptions(
    		r_tol = 1.0e-8, κ_tol = 1.0e-8, solver = :lu_solver),
        sim_opts = ContactImplicitMPC.SimulatorOptions(warmstart = true))

    # simulate
    @test status = ContactImplicitMPC.simulate!(sim, verbose = false)
    # @show sim.traj.q[end][1:3]
    @test norm(ref_traj.q[end][1:3] - sim.traj.q[end][1:3], Inf) < 0.025
end


text:
##############
# GitLabType #
##############
# A `GitLabType` is a Julia type representation of a JSON object defined by the
# GitLab API. Generally:
#
# - The fields of these types should correspond to keys in the JSON object. In
#   the event the JSON object has a "type" key, the corresponding field name
#   used should be `typ` (since `type` is a reserved word in Julia).
#
# - The method `name` should be defined on every GitLabType. This method
#   returns the type's identity in the form used for URI construction. For
#   example, `name` called on an `Owner` will return the owner's login, while
#   `name` called on a `Commit` will return the commit's sha.
#
# - A GitLabType's field types should be Nullables of either concrete types, a
#   Vectors of concrete types, or Dicts.

abstract GitLabType

typealias GitLabString Compat.UTF8String

function @compat(Base.:(==))(a::GitLabType, b::GitLabType)
    if typeof(a) != typeof(b)
        return false
    end

    for field in fieldnames(a)
        aval, bval = getfield(a, field), getfield(b, field)
        if isnull(aval) == isnull(bval)
            if !(isnull(aval)) && get(aval) != get(bval)
                return false
            end
        else
            return false
        end
    end

    return true
end

# `namefield` is overloaded by various GitLabTypes to allow for more generic
# input to AP functions that require a name to construct URI paths via `name`
name(val) = val
name(g::GitLabType) = get(namefield(g))

########################################
# Converting JSON Dicts to GitLabTypes #
########################################

function extract_nullable{T}(data::Dict, key, ::Type{T})
    if haskey(data, key)
        val = data[key]
        if !(isa(val, Void))
            if T <: Vector
                V = eltype(T)
                return Nullable{T}(V[prune_gitlab_value(v, V) for v in val])
            else
                return Nullable{T}(prune_gitlab_value(val, T))
            end
        end
    end
    return Nullable{T}()
end

prune_gitlab_value{T}(val, ::Type{T}) = T(val)
prune_gitlab_value(val, ::Type{Dates.DateTime}) = Dates.DateTime(chopz(val))

# ISO 8601 allows for a trailing 'Z' to indicate that the given time is UTC.
# Julia's Dates.DateTime constructor doesn't support this, but GitLab's time
# strings can contain it. This method ensures that a string's trailing 'Z',
# if present, has been removed.
function chopz(str::AbstractString)
    if !(isempty(str)) && last(str) == 'Z'
        return chop(str)
    end
    return str
end

# Calling `json2gitlab(::Type{G<:GitLabType}, data::Dict)` will parse the given
# dictionary into the type `G` with the expectation that the fieldnames of
# `G` are keys of `data`, and the corresponding values can be converted to the
# given field types.
@generated function json2gitlab{G<:GitLabType}(::Type{G}, data::Dict)
    types = G.types
    fields = fieldnames(G)
    args = Vector{Expr}(length(fields))
    for i in eachindex(fields)
        field, T = fields[i], first(types[i].parameters)
        key = field == :typ ? "type" : string(field)
        args[i] = :(extract_nullable(data, $key, $T))
    end
    return :(G($(args...))::G)
end

#############################################
# Converting GitLabType Dicts to JSON Dicts #
#############################################

gitlab2json(val) = val
gitlab2json(uri::HttpCommon.URI) = string(uri)
gitlab2json(dt::Dates.DateTime) = string(dt) * "Z"
gitlab2json(v::Vector) = [gitlab2json(i) for i in v]

function gitlab2json(g::GitLabType)
    results = Dict()
    for field in fieldnames(g)
        val = getfield(g, field)
        if !(isnull(val))
            key = field == :typ ? "type" : string(field)
            results[key] = gitlab2json(get(val))
        end
    end
    return results
end

function gitlab2json{K}(data::Dict{K})
    results = Dict{K,Any}()
    for (key, val) in data
        results[key] = gitlab2json(val)
    end
    return results
end

###################
# Pretty Printing #
###################

function Base.show(io::IO, g::GitLabType)
    print(io, "$(typeof(g)) (all fields are Nullable):")
    for field in fieldnames(g)
        val = getfield(g, field)
        if !(isnull(val))
            gotval = get(val)
            println(io)
            print(io, "  $field: ")
            if isa(gotval, Vector)
                print(io, typeof(gotval))
            else
                showcompact(io, gotval)
            end
        end
    end
end

function Base.showcompact(io::IO, g::GitLabType)
    uri_id = namefield(g)
    if isnull(uri_id)
        print(io, typeof(g), "(…)")
    else
        print(io, typeof(g), "($(repr(get(uri_id))))")
    end
end


text:
# TEST SINGLE VORTEX RING

include("vm2d.jl")

# create circle geometry
center = [0.0,0.0]
radius = 0.1
thetastart = 0.0
thetaend = 2*pi
refinement = 9
points = circlexy(center,radius,refinement,thetastart,thetaend)

# set vectorial vortex strengths
gamma = [0,0,1.0]
gammas = Array{Float64,1}[]
for point in points
    push!(gammas,gamma)
end

# set initial particle velocities
vs = Array{Float64,1}[]
for point in points
    v = Uinf(point)
    push!(vs,v)
end

particles = place(points,gammas,vs)

# set simulation variables
name = "vortexrings"
global currenttime
currenttime = 0.0
timestep = 0.01
numtimesteps = 3

# plot
## OPTION 1: toggle 2D x-y plot
plt.plot()
plt.hot()
plt.axis("equal")
for particle in particles
    plt.scatter(particle.x[1],particle.x[2],c=particle.gamma[3])
end
plt.title("time: "*string(round(currenttime; digits=3)))
## save frame
for i in range(1,length=numtimesteps)
    ## advance timestep
    advance(particles,timestep,Uinf)
    currenttime += timestep
    ## update plot
    plt.clf()
    for particle in particles
        plt.scatter(particle.x[1],particle.x[2],c=particle.gamma[3])
    end
    plt.title("time: "*string(round(currenttime; digits=3)))
    plt.axis("equal")
    ## save frame
    # println(particles[1])
    plt.pause(0.0001)
end


text:
############################################################################################
#Import
include("phase.jl")
include("proposal.jl")
include("stepsize.jl")
include("diagnostics.jl")
include("tune.jl")

############################################################################################
# Export


text:
function mdl_outer_ode!(device_states,
                        output_ode,
                        f0,
                        device::DynInverter{C,VirtualInertiaQdroop{VirtualInertia,ReactivePowerDroop},VC,DC,P,F}) where {C <: Converter,
                                                                VC<: VSControl,
                                                                DC<: DCSource,
                                                                P <: FrequencyEstimator,
                                                                F <: Filter}

      #Obtain external states inputs for component
      external_ix = device.input_port_mapping[device.outercontrol]
      vpll_d = device_states[external_ix[1]]
      vpll_q = device_states[external_ix[2]]
      ϵ_pll = device_states[external_ix[3]]
      vod = device_states[external_ix[4]]
      voq = device_states[external_ix[5]]
      iod = device_states[external_ix[6]]
      ioq = device_states[external_ix[7]]

      #Obtain inner variables for component
      ω_pll =  device.inner_vars[ω_freq_estimator_var]

      #Get Active Power Controller parameters
      Ta = device.outercontrol.active_power.Ta #VSM Inertia constant
      kd = device.outercontrol.active_power.kd #VSM damping constant
      kω = device.outercontrol.active_power.kω #Frequency droop gain
      ωb = device.outercontrol.active_power.ωb #Rated angular frequency

      #Get Reactive Power Controller parameters
      kq = device.outercontrol.reactive_power.kq #Reactive power droop gain
      ωf = device.outercontrol.reactive_power.ωf #Reactive power filter cutoff frequency

      #Obtain external parameters
      kp_pll = device.freq_estimator.kp_pll
      ki_pll = device.freq_estimator.ki_pll
      p_ref = device.P_ref
      ω_ref = device.ω_ref
      V_ref = device.V_ref
      q_ref = device.Q_ref
      ωg = 1.0

      #Obtain indices for component w/r to device
      local_ix = device.local_state_ix[device.outercontrol]

      #Define internal states for frequency estimator
      internal_states = @view device_states[local_ix]
      δω_vsm = internal_states[1]
      δθ_vsm = internal_states[2]
      qm = internal_states[3]

      #Compute 3 states ODEs
      output_ode[local_ix[1]] = (- iod*vod/Ta
                                    - ioq*voq/Ta
                                    + kd*kp_pll*atan(vpll_q/vpll_d)/Ta
                                    + kd*ki_pll*ϵ_pll/Ta
                                    - (kd+kω)*δω_vsm/Ta
                                    + p_ref/Ta
                                    + kω*ω_ref/Ta
                                    - kω*ωg/Ta)
      output_ode[local_ix[2]] = ωb*δω_vsm
      output_ode[local_ix[3]] = (- ωf*ioq*vod
                                 + ωf*iod*voq
                                 - ωf*qm)

    #Update inner vars
    device.inner_vars[δdqRI_var] = δθ_vsm
    device.inner_vars[ω_control_var] = δω_vsm + 1.0
    device.inner_vars[v_control_var] = V_ref + kq*(q_ref - qm)
end


text:
using CSV
using DataFrames
using DelimitedFiles
using Glob
using Plots; pgfplotsx()

# Create plots
function plot_spmv_nthreads()

	# Dictionary to convert command line arguments to label parts
	# for all solvergroups
	labels = Dict(
		# kernels
		"ser" => "SparseArrays", "par" => "MtSpMV.jl", "mkl" => "MKLSparse.jl",
		# formats
		"csr" => ", CSR", "csc" => ", CSC",
	)

	# Linestyles and markershapes
	styles = [:solid, :dash, :dot, :xcross, :cross]

	# Iterate over systems
	for sys in ["sys1", "sys2"]

		# Iterate over matrix sizes
		for matsize in ["1e6", "1e7"]

			# Create all plots
			plt1 = plot(ylabel="Runtime in seconds", xlabel="Number of threads",
						legend=:topright, dpi=300, framestyle=:box)
			plt2 = plot(ylabel="Parallel speedup", xlabel="Number of threads",
						legend=:none, dpi=300, framestyle=:box)
			plts = [plt1, plt2]

			# Plot for all benchmark types
			for (j, metric) in enumerate(["runtimes", "speedups"])

				# Get folder with benchmarks
				folder = joinpath(sys, matsize)
				folder = joinpath(folder, readdir(folder)[1])

				# Get all csv files for metric in sorted order
				files_ser = glob(string(metric, "*ser*"), joinpath(folder, "csv"))
				files_par = glob(string(metric, "*par*"), joinpath(folder, "csv"))
				files_mkl = glob(string(metric, "*mkl*"), joinpath(folder, "csv"))
				files = vcat(reverse(files_par), reverse(files_mkl), reverse(files_ser))

				# Add results from all files to plots
				for (i, file) in enumerate(files)

					# Get benchmark results
				    df = DataFrame!(CSV.File(file, header=true))
				    x = convert(Vector, df[!, 1])
				    y = convert(Vector, df[!, 5])

				    # # Convert from seconds to milliseconds
				    # if metric == "runtimes"
				    # 	y = y .* 1000
				    # end

				    # Create label from filename
				    infos = file[findlast('/', file) + 1 : findlast('.', file) - 1]
				    infos = infos[findfirst('_', infos) + 1 : length(infos)]
				    infos = split(infos, "_")
				    label = ""
				    for (j, info) in enumerate(infos)
						label = string(label, labels[info])
					end

					# Add to plot
					if length(x) == 1
						scatter!(plts[j], x, y, label=label, markershape=styles[i])
					else
						plot!(plts[j], x, y, label=label, linestyle=styles[i])
					end

				end # for

				# Plot side by side and save
				width = 530
				plot(plts[1], plts[2], layout=(1,2), size=(width, width/2))
				filename = string("bench_spmv_nthreads_", sys, "_", matsize)
				savefig(string(filename, ".png"))
				savefig(string(filename, ".tex"))

			end # for

		end # for

	end # for

end # function

max_stars_repo_name: jkrch/JuliaParGMRES-Plots

max_issues_repo_name: jkrch/JuliaParGMRES-Plots

max_forks_repo_name: jkrch/JuliaParGMRES-Plots


text:
export GainObject
export assign_pumpstr!

# GainObject has
# - Shape
# - scale factor to be multiplied to the true/false array.
#
# Later, to put an object with a spatially changing pump strength, I could define a new
# GainObject.  Different types of GainObjects will need to define different assignment
# function.  Eventually, I will need to write something like MaxwellFDFD's assign function
# such that it internally uses the individual shape's assign function.  But let's not think
# about that right now.
#
# This may require defining GainObject as an abstract type later, and defining its concrete
# types.

# See Julia issue #26321 for using S<:Shape{K}.
mutable struct GainObject{K,Kₑ,S<:Shape{K},F}  # Kₑ is number of Cartesian components of E-field
    shape::S
    D₀fun::F  # population inversion in absence of lasing; function of pump parameter d
    GainObject{K,Kₑ,S,F}(shape, D₀fun) where {K,Kₑ,S<:Shape{K},F} = new(shape, D₀fun)
end

GainObject{K,Kₑ}(shape::S, D₀fun::F) where {K,Kₑ,S<:Shape{K},F} = GainObject{K,Kₑ,S,F}(shape, D₀fun)
GainObject{K,Kₑ}(shape::Shape{K}) where {K,Kₑ} = GainObject{K,Kₑ}(shape, d::Real->d)

add_gainobj!(gobj_vec::AbsVec{<:GainObject}, shapes::Shape...; D₀fun::Any=identity) = add_gainobj!(gobj_vec, [shapes...]; D₀fun)
function add_gainobj!(gobj_vec::AbsVec{<:GainObject{K,Kₑ}}, shapes::AbsVec{<:Shape{K}}; D₀fun::Any=identity) where {K,Kₑ} # identity: d::Real->d
    for shape = shapes  # shapes is tuple
        gobj = GainObject{K,Kₑ}(shape, D₀fun)
        push!(gobj_vec, gobj)
    end

    return nothing
end

# Assign the pump strength specified by gain objects to an array.
# The implementation depends on MaxwellFDFD, so this function is moved from SALTBase to
# MaxwellSALT.
function assign_pumpstr!(D₀::AbsVecReal,  # output vector to write on
                         gobj_vec::AbsVec{<:GainObject{K,Kₑ}},  # vector of gain objects; later objects overwrite earlier objects
                         d::Real,  # pump parameter at which D₀ is evaluated
                         gt₀::SVec{K,GridType},  # grid type of E-field voxel corners; generated by ft2gt.(EE, boundft)
                         l::Tuple2{NTuple{K,AbsVecReal}},  # l[PRIM][k], l[DUAL][k]: primal, dual vertex locations in k-direction
                         ise˔shp::Bool,  # true if E-field dimension is orthogonal to shape dimension; false otherwise
                         order_cmpfirst::Bool  #  if true, D₀ is indexed D₀[w,x,y,z] for K = 3; otherwise indexed D₀[x,y,z,w]
                         ) where {K,Kₑ}
    N = SVec(length.(l[nPR]))  # SInt{K}
    D₀ .= 0  # initialize (pump strength is nonzero only inside gain media)
    D₀array = order_cmpfirst ? reshape(D₀, Kₑ, N.data...) : reshape(D₀, N.data..., Kₑ)  # related: MaxwellBase/field.jl/field_arr2vec().
    for nw = 1:Kₑ
        gt_cmp = ise˔shp ? gt₀ : gt_w(nw, gt₀)
        lcmp = t_ind(l, gt_cmp)
        D₀cmp = order_cmpfirst ? selectdim(D₀array, 1, nw) :  # D₀array[nw,:,:,:] for K = 3 (note 1 ≤ nw ≤ Kₑ)
                                 selectdim(D₀array, K+1, nw)  # D₀array[:,:,:,nw] for K = 3 (note 1 ≤ nw ≤ Kₑ)
        for gobj = gobj_vec
            shape = gobj.shape
            D₀fun = gobj.D₀fun
            D₀val = D₀fun(d)
            assign_val_shape!(D₀cmp, D₀val, shape, lcmp)  # exported from MaxwellBase
        end
    end

    return nothing
end


text:
@doc raw"""
    SquaredExponentialKernel([α=1])

The squared exponential kernel, or alternatively the Gaussian kernel, is identical to the
exponential kernel except that the Euclidean distance is squared:

```math
\kappa(\mathbf{x},\mathbf{y}) = \exp\left(-\alpha ||\mathbf{x} - \mathbf{y}||^2\right)
\qquad \alpha > 0
```

where ``\alpha`` is a scaling parameter of the squared Euclidean distance.
Just like the exponential kernel, the squared exponential kernel is an
isotropic Mercer kernel. The squared exponential kernel is more commonly known
as the radial basis kernel within machine learning communities.
"""
struct SquaredExponentialKernel{T<:AbstractFloat} <: MercerKernel{T}
    alpha::HyperParameter{T}
    SquaredExponentialKernel{T}(α::Real) where {T<:AbstractFloat} = new{T}(
        HyperParameter(convert(T,α), interval(OpenBound(zero(T)), nothing))
    )
end
SquaredExponentialKernel(α::T=1.0) where {T<:Real} = SquaredExponentialKernel{floattype(T)}(α)
GaussianKernel = SquaredExponentialKernel
RadialBasisKernel = SquaredExponentialKernel

@inline squaredexponentialkernel(z::T, α::T) where {T<:AbstractFloat} = exp(-α*z)

@inline basefunction(::SquaredExponentialKernel) = SquaredEuclidean()
@inline function kappa(κ::SquaredExponentialKernel{T}, z::T) where {T}
    squaredexponentialkernel(z, getvalue(κ.alpha))
end


text:
# This file assumes that the directories are organized as follows:
#
# ./SomeDirectory/RRTMGP.jl/
# ./SomeDirectory/rte-rrtmgp/
#

if Sys.iswindows()
    # ENV["NCHOME"] = joinpath("C:","Program Files","netCDF 4.7.1","bin")
    # ENV["NCHOME"] = joinpath("C:","Program Files","netCDF 4.7.1","lib")
    ENV["NCHOME"] = joinpath("C:", "Program Files", "netCDF 4.7.1")
    # ENV["NFHOME"] = joinpath("C:","MinGW","bin")
    ENV["NFHOME"] = joinpath("C:", "MinGW")
else
    # ENV["NCHOME"] = "/usr/local/Cellar/netcdf/4.6.1_4"
    ENV["NCHOME"] = "/usr/local"
    ENV["NFHOME"] = "/usr/local"
    # ENV["NCHOME"] = joinpath("/usr","local","Cellar","netcdf","4.6.1_4")
    # ENV["NFHOME"] = joinpath("/usr","local","Cellar","netcdf","4.6.1_4")
end

root_dir = dirname(pwd())
rte_rrtmgp_dir = joinpath(root_dir, "rte-rrtmgp")
rte_rrtmgp_rfmip_clear_sky = joinpath(root_dir, "rte-rrtmgp", "examples", "rfmip-clear-sky")
rte_rrtmgp_rfmip_all_sky = joinpath(root_dir, "rte-rrtmgp", "examples", "all-sky")
rte_rrtmgp_test_val_dir = joinpath(root_dir, "rte-rrtmgp-test-val")

# rte-rrtmgp ENV variables
ENV["RTE_KERNELS"] = joinpath(rte_rrtmgp_dir, "rte", "kernels")

# rte-rrtmgp and rte-rrtmgp-test-val ENV variables
ENV["RTE_DIR"] = joinpath(rte_rrtmgp_dir, "rte")
ENV["RTE_KERNEL_DIR"] = joinpath(rte_rrtmgp_dir, "rte", "kernels")
ENV["RRTMGP_KERNEL_DIR"] = joinpath(rte_rrtmgp_dir, "rrtmgp", "kernels")
ENV["RRTMGP_DIR"] = joinpath(rte_rrtmgp_dir, "build")
ENV["FCFLAGS"] = "-fimplicit-none -ffree-line-length-none -Wuninitialized -cpp -std=gnu"
ENV["FC"] = "gfortran"

# rte-rrtmgp-test-val ENV variables
ENV["TEST_ROOT"] = rte_rrtmgp_test_val_dir
ENV["RRTMGP_ROOT"] = rte_rrtmgp_dir
ENV["RRTMGP_BUILD"] = joinpath(rte_rrtmgp_dir, "build")

max_stars_repo_name: climate-machine/JRRTMGP

max_issues_repo_name: climate-machine/JRRTMGP

max_forks_repo_name: climate-machine/JRRTMGP


text:
export armijo_ls

"""
An armijo backtracking algorithm.
"""
function armijo_ls(h :: LineModel, stop_ls :: LS_Stopping;
                   f_meta = LS_Function_Meta(),
                   φ_dφ        :: Function = (x, y) -> phi_dphi(x, y),
                   verboseLS   :: Bool = false,  kwargs...)
    state = stop_ls.current_state
    update!(state, x = 1.0)

    stop_ls.optimality_check = (x, y) -> armijo(x,y)

    # First try to increase t
    h1 = obj(h, state.x)
    slope1 = grad(h, state.x)

    OK = update_and_start!(stop_ls, ht = h1, gt = slope1, tmps = time())

    while !OK
        t = state.x
        update!(state, x = 0.5 * t)
        ht = obj(h, state.x)
        slope = grad(h, state.x)
        OK = update_and_stop!(stop_ls, ht = ht, gt = slope)
        verboseLS && @printf(" iter  %4d  t  %7.1e slope %7.1e\n", stop_ls.meta.nb_of_stop, state.x, slope);
    end

    return (state, stop_ls.meta.optimal)
end # function


text:
## Implement (E)CBS functions
function MultiAgentPathFinding.add_constraint!(cbase::MAPFTransitConstraints,
                                               cadd::MAPFTransitConstraints)
    union!(cbase.avoid_vertex_map[towards], cadd.avoid_vertex_map[towards])
    union!(cbase.avoid_vertex_map[from], cadd.avoid_vertex_map[from])
end

function MultiAgentPathFinding.overlap_between_constraints(cbase::MAPFTransitConstraints,
                                                           cadd::MAPFTransitConstraints)
    return ~(isempty(intersect(cbase.avoid_vertex_idxs, cadd.avoid_vertex_idxs)))
end

function MultiAgentPathFinding.get_mapf_state_from_idx(env::MAPFTransitEnv, idx::Int64)
    return env.state_graph.vertices[idx]
end

function get_mapf_transit_action(env::MAPFTransitEnv, vtx_u::MAPFTransitVertexState, vtx_v::MAPFTransitVertexState)

    usplit = split(vtx_u.vertex_str, "-")
    vsplit = split(vtx_v.vertex_str, "-")

    if usplit[1] == "r" && vsplit[1] == "r" &&
        usplit[2] == vsplit[2]
        return MAPFTransitAction(Stay::ActionType)
    elseif vsplit[1] != "r"
        return MAPFTransitAction(Fly::ActionType)
    else
        return MAPFTransitAction(Board::ActionType)
    end

end

function MultiAgentPathFinding.get_mapf_action(env::MAPFTransitEnv, u::Int64, v::Int64)

    vtx_u = env.state_graph.vertices[u]
    vtx_v = env.state_graph.vertices[v]

    return get_mapf_transit_action(env, vtx_u, vtx_v)

end

function MultiAgentPathFinding.set_low_level_context!(::MAPFTransitEnv, ::Int64, ::MAPFTransitConstraints)
    # Not sure what to do here
end

MultiAgentPathFinding.get_empty_constraint(::Type{MAPFTransitConstraints}) = MAPFTransitConstraints()

function MultiAgentPathFinding.get_first_conflict(env::MAPFTransitEnv,
                                                  solution::Vector{PR}) where {PR <: PlanResult}

    # First look for boarding/alighting constraints
    # But simultaneously note transit capacity usage
    for (i, sol_i) in enumerate(solution[1:end-1])
        for (j, sol_j) in enumerate(solution[i+1:end])

            # Now we have the two solutions
            for (si, ((state_i, _), (act_i, _))) in enumerate(zip(sol_i.states[2:end], sol_i.actions))
                for (sj, ((state_j, _), (act_j, _))) in enumerate(zip(sol_j.states[2:end], sol_j.actions))

                    # Conflict - boarding vertex at same time
                    if state_i.state == state_j.state && act_i == Board::ActionType && act_j == Board::ActionType

                        conflict = MAPFTransitConflict(type = Transfer::ConflictType,
                                                   overlap_vertices = Set{String}(state_i.idx),
                                                   agent_to_state_idx = Dict(i => si+1, j => sj+1)) # +1 because enumerating from 2
                        return conflict
                    end
                end
            end
        end
    end

    # Now look for capacity constraints ONLY
    # Dict of
    transit_vtx_agents = Dict()
    for (i, sol_i) in enumerate(solution)
        for (si, (state_i, _)) in enumerate(sol_i.states)

            isplit = split(state_i.vertex_str, "-")

            if isplit[1] == "r"
                rid = parse(Int64, isplit[2])
                wpid = parse(Int64, isplit[3])

                if ~(haskey(transit_vtx_agents, rid))
                    transit_vtx_agents[rid] = Dict{Int64,Dict{Int64,Int64}}()
                end

                if ~(haskey(transit_vtx_agents[rid], wpid))
                    transit_vtx_agents[rid][wpid] = Dict{Int64,Int64}()
                end

                # Create a mini dict for waypoint
                transit_vtx_agents[rid][wpid][i] = si
            end
        end
    end

    # Now loop over vtx_agents and report first over-capacity link
    # TODO: Find a way to use all conflicts??? For now, returning first
    for (rid, wpid_agent_dicts) in transit_vtx_agents

        route_cap = env.transit_graph.transit_capacity[rid]
        wpid_list = sort(collect(keys(wpid_agent_dicts)))
        trip_to_vtx_range = env.trip_to_vtx_range[rid]

        for i = 1:length(wpid_list)-1

            wpid = wpid_list[i]
            agent_dict = wpid_agent_dicts[wpid]

            # Check if this marks the start of a link with excess capacity
            # If it does, get the particulars and return conflict
            if length(agent_dict) > route_cap

                overlap_idx = trip_to_vtx_range[1] + wpid - 1
                overlap_verts = Set{Int64}(overlap_idx)
                agents = Set(keys(agent_dict))

                for j = i+1:length(wpid_list)

                    next_wpid = wpid_list[j]
                    next_agent_dict = wpid_agent_dicts[next_wpid]
                    next_overlap = trip_to_vtx_range[1] + next_wpid - 1
                    push!(overlap_verts, next_overlap)

                    if Set(keys(next_agent_dict)) != agents
                        break
                    end
                end

                # Return conflict
                # @show overlap_verts
                # @show agent_dict
                # readline()
                conflict = MAPFTransitConflict(type = Capacity::ConflictType,
                                           overlap_vertices = overlap_verts,
                                           agent_to_state_idx = agent_dict,
                                           cap = route_cap)
                # @show conflict.type
                return conflict
            end
        end
    end

    # No conflict thus far
    return nothing

end

function MultiAgentPathFinding.create_constraints_from_conflict(env::MAPFTransitEnv, conflict::MAPFTransitConflict)

    # Increment the number of global conflicts and throw error if above threshold
    env.num_global_conflicts += 1
    if env.num_global_conflicts > env.threshold_global_conflicts
        throw(DomainError("Too many conflicts!"))
    end

    # Do for each type of conflict
    if conflict.type ==  Transfer::ConflictType

        res_constraints = Dict{Int64,MAPFTransitConstraints}()

        # Obtain the kind of subpath
        for (agt, state_idx) in conflict.agent_to_state_idx
            if state_idx < env.curr_site_points[agt] # Must be greater or less
                res_constraints[agt] = MAPFTransitConstraints(Dict(towards => conflict.overlap_vertices,
                                                                   from => Set{Int64}()))
            else
                res_constraints[agt] = MAPFTransitConstraints(Dict(from => conflict.overlap_vertices,
                                                                   towards => Set{Int64}()))
            end
        end

        # @show [res_constraints]
        return [res_constraints]

    else

        route_cap = conflict.cap
        overlap_agents = Set(keys(conflict.agent_to_state_idx))
        n_excess = length(overlap_agents) - route_cap
        res_constraint_set = Vector{Dict}(undef, 0)

        # In one pass, get agent to constraint mapping
        agt_to_constraints_dict = Dict()

        for agt in overlap_agents
            state_idx = conflict.agent_to_state_idx[agt]

            if state_idx < env.curr_site_points[agt] # Must be greater or less
                agt_to_constraints_dict[agt] = MAPFTransitConstraints(Dict(towards => conflict.overlap_vertices,
                                                                   from => Set{Int64}()))
            else
                agt_to_constraints_dict[agt] = MAPFTransitConstraints(Dict(from => conflict.overlap_vertices,
                                                                   towards => Set{Int64}()))
            end
        end

        # Generate all subsets of excess agents and generate Constraints
        for agt_subset in subsets(collect(overlap_agents), n_excess)

            res_constraints = Dict{Int64,MAPFTransitConstraints}()

            # For each agent in subset, find if its halfway point crossed
            for agt in agt_subset
                res_constraints[agt] = agt_to_constraints_dict[agt]
            end

            push!(res_constraint_set, res_constraints)
        end

        # @show res_constraint_set[1]
        # @show res_constraint_set[2]
        # readline()

        # @show res_constraint_set
        return res_constraint_set
    end
end

"""
    elapsed_time(env::MAPFTransitEnv,
                 s1::MAPFTransitVertexState, s2::MAPFTransitVertexState)

Compute the time difference between time-stamped vertices.
If any not time-stamped, computes the average traversal time.
"""
function elapsed_time(env::MAPFTransitEnv,
                      s1::MAPFTransitVertexState, s2::MAPFTransitVertexState)

    if startswith(s2.vertex_str, "r") == true
        return (s2.state.time - s1.state.time)
    else
        # Compute distance between s1 and s2 and divide by avg speed
        dist = env.dist_fn(s1.state.location, s2.state.location)
        return dist/env.drone_params.avg_speed
    end
end

"""
    elapsed_time_heuristic(env::MAPFTransitEnv, s::MAPFTransitVertexState)

Computes the trivial heuristic for traversal time from a vertex to the goal.
"""
function elapsed_time_heuristic(env::MAPFTransitEnv, s::MAPFTransitVertexState)

    # Just direct flight time to goal
    goal_vtx = env.state_graph.vertices[env.curr_goal_idx]
    return (env.dist_fn(goal_vtx.state.location, s.state.location)/env.drone_params.avg_speed)

end

"""
    distance_traversed(env::MAPFTransitEnv, s1::MAPFTransitVertexState, s2::MAPFTransitVertexState)

Computes the flight distance between two operation graph vertices (0 for transit edges).
"""
function distance_traversed(env::MAPFTransitEnv, s1::MAPFTransitVertexState, s2::MAPFTransitVertexState)

    split1 = split(s1.vertex_str, "-")
    split2 = split(s2.vertex_str, "-")

    if split1[1] == "r" && split2[1] == "r" && split1[2] == split2[2]
        return 0.0
    else
        return env.dist_fn(s1.state.location, s2.state.location)
    end
end

"""
    distance_heuristic(env::MAPFTransitEnv, s::MAPFTransitVertexState)

Computes the admissible heuristic on flight distance between vertices. See Appendix II-B
"""
function distance_heuristic(env::MAPFTransitEnv, s::MAPFTransitVertexState)

    ssplit = split(s.vertex_str, "-")
    goal_vtx = env.state_graph.vertices[env.curr_goal_idx]

    n_depots = length(env.off_transit_graph.depots)
    n_sites = length(env.off_transit_graph.sites)

    # NOTE : Should be correct since state graph is first populated with depots and then sites
    depot_site_idx = goal_vtx.idx

    if ssplit[1] == "r"
        trip_id = parse(Int64, ssplit[2])
        stop_id = env.transit_graph.transit_trips[trip_id][parse(Int64, ssplit[3])].stop_id

        mindist = Inf
        for tid in env.stop_idx_to_trips[stop_id]
            trip_to_ds = env.aug_trips_fws_dists[n_depots + n_sites + tid, depot_site_idx]
            mindist = (trip_to_ds < mindist) ? trip_to_ds : mindist
        end

        return mindist
    else
        return env.aug_trips_fws_dists[s.idx, depot_site_idx]
    end
end

## N.B The focal heuristic functions that follow are the lowest hanging fruit for
## improving empirical performance of ECBS for search.
function focal_transition_heuristic_transit(env::MAPFTransitEnv, solution::Vector{PR},
                                            agent_idx::Int64, u::MAPFTransitVertexState, v::MAPFTransitVertexState) where {PR <: PlanResult}

    num_conflicts = 0

    this_act = get_mapf_transit_action(env, u, v)

    for (i, agt_soln) in enumerate(solution)
        if i != agent_idx && ~(isempty(agt_soln))

            # Count conflicts
            for (vtx_state, act) in zip(agt_soln.states[2:end], agt_soln.actions)
                if vtx_state[1].state == v && act[1] == MAPFTransitAction(Board) && this_act == MAPFTransitAction(Board)
                    num_conflicts += 1
                end
            end
        end
    end

    return num_conflicts
end

# Just count the boarding conflicts, if any?
function MultiAgentPathFinding.focal_heuristic(env::MAPFTransitEnv, solution::Vector{PR}) where {PR <: PlanResult}

    num_potential_conflicts = 0

    # First look for boarding/alighting constraints
    # But simultaneously note transit capacity usage
    for (i, sol_i) in enumerate(solution[1:end-1])
        for (j, sol_j) in enumerate(solution[i+1:end])

            # Now we have the two solutions
            for (si, ((state_i, _), (act_i, _))) in enumerate(zip(sol_i.states[2:end], sol_i.actions))
                for (sj, ((state_j, _), (act_j, _))) in enumerate(zip(sol_j.states[2:end], sol_j.actions))

                    # Conflict - boarding vertex at same time
                    if state_i.state == state_j.state
                        num_potential_conflicts += 1
                    end
                end
            end
        end
    end

    return num_potential_conflicts
end

function reachable_by_agent(env::MAPFTransitEnv, s1::MAPFTransitState, s2::MAPFTransitState)
    # Check that dist is less than avg_speed x time
    return env.drone_params.avg_speed * (s2.time - s1.time) > env.dist_fn(s1.location, s2.location)
end

"""
Implements the actual single-agent low-level search for MultiAgentPathFinding.jl's Enhanced CBS.
Computes the dpd' path with Focal-MCSP twice. Does some hacking to concat paths and to avoid
searching for a subpath that has no constraints.
"""
function MultiAgentPathFinding.low_level_search!(solver::ECBSSolver, agent_idx::Int64,
                                                 s::MAPFTransitVertexState, constraints::MAPFTransitConstraints,
                                                 solution::Vector{PR}) where {PR <: PlanResult}

    env = solver.env
    agent_state = env.agent_states[agent_idx]

    # set up heuristics regardless
    # Set up the weight and heuristic functions
    edge_wt_fn(u, v) = elapsed_time(env, u, v)

    edge_constr_fn(u, v) = distance_traversed(env, u, v)
    edge_constr_functions = [edge_constr_fn]

    edge_constraints = [env.drone_params.max_distance - agent_state.dist_flown]

    admissible_heuristic(u) = elapsed_time_heuristic(env, u)

    dist_heur(u) = distance_heuristic(env, u)

    focal_state_heuristic(u) = 0.0
    focal_transition_heuristic(u, v) = focal_transition_heuristic_transit(env, solution, agent_idx, u, v)

    # Don't reset initial time
    orig_idx = s.idx

    states = Tuple{MAPFTransitVertexState,Float64}[]
    actions = Tuple{MAPFTransitAction,Float64}[]

    first_cost = env.plan_ref_times[agent_idx]

    # ONLY run first search if site not crossed
    if agent_state.site_crossed == false

        # Then goal is the site
        goal_str = goal_str = string("s-", agent_state.task.site)
        env.curr_goal_idx = env.depot_sites_to_vtx[goal_str]
        goal_vtx = env.state_graph.vertices[env.curr_goal_idx]

        # Note - don't really need an sl dist heuristic
        # Since you are beginning the search anyway
        if isempty(solution) || ~(isempty(constraints.avoid_vertex_map[towards]))

            vis = MAPFTransitGoalVisitor(env, constraints.avoid_vertex_map[towards])
            astar_eps_states, tgt_entry = a_star_implicit_constrained_epsilon_path!(env.state_graph,
                                                                                edge_wt_fn,
                                                                                orig_idx, vis, solver.weight,
                                                                                admissible_heuristic,
                                                                                focal_state_heuristic,
                                                                                focal_transition_heuristic,
                                                                                edge_constr_functions,
                                                                                [dist_heur],
                                                                                edge_constraints)

            # IMP - Need to sort out solution
            if tgt_entry.v_idx == orig_idx
                @warn "Agent $(agent_idx) Did not find first sub-path!"

                env.any_invalid_path = true

                # For debugging, make a direct path with distance
                weight = distance_traversed(env, s, goal_vtx)

                append!(states, [(s, first_cost), (goal_vtx, first_cost + elapsed_time(env, s, goal_vtx))])
                append!(actions, [(MAPFTransitAction(Fly), weight)])
                env.curr_site_points[agent_idx] = 2

            else
                sp_idxs, costs, wts = shortest_path_cost_weights(astar_eps_states.a_star_states, env.state_graph, orig_idx, tgt_entry)

                append!(states, [(get_mapf_state_from_idx(env, idx), first_cost + c) for (idx, c) in zip(sp_idxs, costs)])
                append!(actions, [(get_mapf_action(env, u, v), wts[i+1][1] - wts[i][1]) for (i, (u, v)) in enumerate(zip(sp_idxs[1:end-1], sp_idxs[2:end]))])

                env.curr_site_points[agent_idx] = length(sp_idxs)
            end

        else
            # Copy over previous values
            @debug "Copying over subpart towards of agent $(agent_idx)"

            append!(states, solution[agent_idx].states[1 : env.curr_site_points[agent_idx]])
            append!(actions, solution[agent_idx].actions[1 : env.curr_site_points[agent_idx] - 1])
        end

        first_cost = states[end][2]

        # Reset stateful things for next search
        orig_str = string("s-", agent_state.task.site)
        orig_idx = env.depot_sites_to_vtx[orig_str]
        edge_constraints = [env.drone_params.max_distance]

    end

    goal_str = string("d-", agent_state.task.dest)
    env.curr_goal_idx = env.depot_sites_to_vtx[goal_str]
    goal_vtx = env.state_graph.vertices[env.curr_goal_idx]

    if isempty(solution) || ~(isempty(constraints.avoid_vertex_map[from]))

        vis = MAPFTransitGoalVisitor(env, constraints.avoid_vertex_map[from])
        astar_eps_states, tgt_entry = a_star_implicit_constrained_epsilon_path!(env.state_graph,
                                                                               edge_wt_fn,
                                                                               orig_idx, vis, solver.weight,
                                                                               admissible_heuristic,
                                                                               focal_state_heuristic,
                                                                               focal_transition_heuristic,
                                                                               edge_constr_functions,
                                                                               [dist_heur],
                                                                               edge_constraints)

        if tgt_entry.v_idx == orig_idx
            @warn "Agent $(agent_idx) Did not find second sub-path!"

            env.any_invalid_path = true

            weight = distance_traversed(env, env.state_graph.vertices[orig_idx], goal_vtx)

            append!(states, [(goal_vtx, first_cost + elapsed_time(env, env.state_graph.vertices[orig_idx], goal_vtx))])
            append!(actions, [(MAPFTransitAction(Fly),weight)])

        else

            sp_idxs, costs, wts = shortest_path_cost_weights(astar_eps_states.a_star_states, env.state_graph, orig_idx, tgt_entry)

            if agent_state.site_crossed == false
                append!(states, [(get_mapf_state_from_idx(env, idx), first_cost + c) for (idx, c) in zip(sp_idxs[2:end], costs[2:end])])
            else
                # Append the full post soln
                append!(states, [(get_mapf_state_from_idx(env, idx), first_cost + c) for (idx, c) in zip(sp_idxs, costs)])
            end

            append!(actions, [(get_mapf_action(env, u, v), wts[i+1][1] - wts[i][1]) for (i, (u, v)) in enumerate(zip(sp_idxs[1:end-1], sp_idxs[2:end]))])
        end

    else

        @debug "Copying over subpart from of agent $(agent_idx)"

        append!(states, solution[agent_idx].states[env.curr_site_points[agent_idx] + 1 : end])
        append!(actions, solution[agent_idx].actions[env.curr_site_points[agent_idx] : end])
    end

    cost = states[end][2]
    fmin = cost

    # Update the next finish time of agent
    env.agent_states[agent_idx].next_finish_time = cost

    return PlanResult{MAPFTransitVertexState, MAPFTransitAction, Float64}(states, actions, cost, fmin)
end

# NOTE: Assumes that solution[i].states has the correct tuple of (state, cost)
"""
    update_agent_states!(env::MAPFTransitEnv, time_val::Float64, agent_idx::Int64,
                         solution::Vector{PR}) where {PR <: PlanResult}

Returns the next state of each agent based on its currently executing solution. Also updates the site_crossed flag if so
Assumes that agent has not finished yet
Assumes that solution[i].actions tracks the weight costs
"""
function update_agent_states!(env::MAPFTransitEnv, time_val::Float64, agent_idx::Int64,
                              solution::Vector{PR}) where {PR <: PlanResult}

    @assert time_val < env.agent_states[agent_idx].next_finish_time "Agent $(agent_idx) has already finished at time $(time)!"

    agt_soln_states = solution[agent_idx].states
    agt_soln_actions = solution[agent_idx].actions

    state_idx = 0
    state_time = Inf

    # Iterate over agent_states until first state after time
    for (si, (state, t)) in enumerate(agt_soln_states)
        if t >= time_val
            state_idx = si
            state_time = t
            break
        end
    end

    # Now deal with state_idx
    @assert state_idx <= length(agt_soln_states) "Agent $(agent_idx) has inconsistent time trajectory!; $(time_val); $(agt_soln_states)"

    # NOTE: hack to avoid pathological case of agent being at end of solution
    if state_idx == length(agt_soln_states)
        state_idx = state_idx - 1
        state_time = agt_soln_states[state_idx][2]
    end

    dist_start_idx = 1
    if state_idx >= env.curr_site_points[agent_idx]
        env.agent_states[agent_idx].site_crossed = true
        dist_start_idx = env.curr_site_points[agent_idx]
    end

    dist_flown = 0.0
    for i = dist_start_idx:state_idx-1
        dist_flown += agt_soln_actions[i][2]
    end

    env.agent_states[agent_idx].dist_flown = dist_flown

    return agt_soln_states[state_idx][1], state_time # That's the state
end

mutable struct MAPFTransitGoalVisitor <: AbstractDijkstraVisitor
    env::MAPFTransitEnv
    avoid_vertex_idxs::Set{Int64}
end

function Graphs.include_vertex!(vis::MAPFTransitGoalVisitor, u::MAPFTransitVertexState, v::MAPFTransitVertexState,
                                d::Float64, nbrs::Vector{Int64})

    env = vis.env
    avoid_vertex_idxs = vis.avoid_vertex_idxs
    goal_vtx = env.state_graph.vertices[env.curr_goal_idx]

    if goal_vtx.vertex_str == v.vertex_str
        # NOTE: Commented out to have a rolling horizon of trips
        # reset_time(env.state_graph.vertices[env.curr_goal_idx], d)
        return false
    end

    # Now add other neighbours according to whether depot/site or route vertex
    vsplit = split(v.vertex_str, "-")

    if vsplit[1] == "r"

        # Get the action leading to it
        act = get_mapf_transit_action(env, u, v)
        trip_id = parse(Int64, vsplit[2])

        vtx_range = env.trip_to_vtx_range[trip_id]

        next_vtx = vtx_range[1] + parse(Int64, vsplit[3])

        if act == MAPFTransitAction(Board) && next_vtx <= vtx_range[2] &&
            ~(next_vtx in avoid_vertex_idxs)
            # If just boarded, just add the next route vertex
             # say 21-30, seq 4; add vtx 25 (5th in seq)
            push!(nbrs, next_vtx)
            return true
        else
            # Has been on trip - can either continue to next (if any left) and/or must get off
            if next_vtx <= vtx_range[2] &&
                ~(next_vtx in avoid_vertex_idxs)
                push!(nbrs, next_vtx)
            end

            # Cycle through other trips and add those that can be reached in time
            for (tid, trip) in enumerate(env.transit_graph.transit_trips)

                # Ignore the current trip id
                if tid != trip_id

                    # Use non-dominated set
                    subseq = get_non_dominated_trip_points(env, avoid_vertex_idxs, v, tid)
                    append!(nbrs, subseq)
                end
            end
        end
    else
        # A depot/site vertex - just add possible trip verts
        for (tid, trip) in enumerate(env.transit_graph.transit_trips)

            subseq = get_non_dominated_trip_points(env, avoid_vertex_idxs, v, tid)
            append!(nbrs, subseq)
        end
    end

    # Now generate neighbours
    # Always consider goal if possible to reach
    if env.dist_fn(v.state.location, goal_vtx.state.location) < env.drone_params.max_distance
        # @info "Goal added by ", v.idx
        push!(nbrs, env.curr_goal_idx)
    end

    # @show v
    # @show nbrs

    return true
end

"""
    get_depot_to_site_travel_time!(env::MAPFTransitEnv, weight::Float64, orig_str::String, goal_str::String)

Runs a single search over a snapshot of the transit network to get the travel time. If not reachable, then return Inf
"""
function get_depot_to_site_travel_time!(env::MAPFTransitEnv, weight::Float64, orig_str::String, goal_str::String)

    # Get origin idx and destination idx
    # No need to reset as must be zero
    orig_idx = env.depot_sites_to_vtx[orig_str]
    env.curr_goal_idx = env.depot_sites_to_vtx[goal_str]

    edge_wt_fn(u, v) = elapsed_time(env, u, v)
    edge_constr_fn(u, v) = distance_traversed(env, u, v)
    edge_constr_functions = [edge_constr_fn]
    edge_constraints = [env.drone_params.max_distance]
    admissible_heuristic(u) = elapsed_time_heuristic(env, u)
    focal_state_heuristic(u) = 0.0
    focal_transition_heuristic(u, v) = 0.0
    dist_heur(u) = 0.0  # Keep this simple
    edge_constr_heuristics = [dist_heur]

    vis = MAPFTransitGoalVisitor(env, Set{Int64}())

    astar_eps_states, tgt_entry = a_star_implicit_constrained_epsilon_path!(env.state_graph,
                                                                           edge_wt_fn,
                                                                           orig_idx, vis, weight,
                                                                           admissible_heuristic,
                                                                           focal_state_heuristic,
                                                                           focal_transition_heuristic,
                                                                           edge_constr_functions,
                                                                           edge_constr_heuristics,
                                                                           edge_constraints)

    # TODO: Not returning Inf to avoid instability? Hack
    if tgt_entry.v_idx == orig_idx
        @warn "In allocation, edge from $(orig_str) to $(goal_str) is empty!"
        return 100000.
    end

    sp_idxs, costs, wts = shortest_path_cost_weights(astar_eps_states.a_star_states, env.state_graph, orig_idx, tgt_entry)

    return costs[end]
end

"""
    allocation_cost_wrapper_truett(env::MAPFTransitEnv, weight::Float64, n_depots::Int64, n_sites::Int64,
                                    orig_ds_idx::Int64, goal_ds_idx::Int64)

Use actual travel time between locations by running a full search. Not Used.
"""
function allocation_cost_wrapper_truett(env::MAPFTransitEnv, weight::Float64, n_depots::Int64, n_sites::Int64,
                                    orig_ds_idx::Int64, goal_ds_idx::Int64)

    # Compute the origin and goal strings
    if orig_ds_idx > n_depots
        orig_str = string("s-", (orig_ds_idx - n_depots))
    else
        orig_str = string("d-", orig_ds_idx)
    end

    if goal_ds_idx > n_depots
        goal_str = string("s-", (goal_ds_idx - n_depots))
    else
        goal_str = string("d-", goal_ds_idx)
    end

    return get_depot_to_site_travel_time!(env, weight, orig_str, goal_str)
end

"""
    allocation_cost_wrapper_estimate(env::MAPFTransitEnv, weight::Float64, n_depots::Int64, n_sites::Int64,
                                     halton_nn_tree::BallTree, estimate_matrix::Matrix{Float64},
                                     orig_ds_idx::Int64, goal_ds_idx::Int64)

Get the surrogate travel time using the estimate matrix.
"""
function allocation_cost_wrapper_estimate(env::MAPFTransitEnv, weight::Float64, n_depots::Int64, n_sites::Int64,
                                          halton_nn_tree::BallTree, estimate_matrix::Matrix{Float64},
                                          orig_ds_idx::Int64, goal_ds_idx::Int64)

    # Compute the origin and goal strings
    if orig_ds_idx > n_depots
        orig_str = string("s-", (orig_ds_idx - n_depots))
    else
        orig_str = string("d-", orig_ds_idx)
    end
    loc1 = env.state_graph.vertices[env.depot_sites_to_vtx[orig_str]].state.location

    if goal_ds_idx > n_depots
        goal_str = string("s-", (goal_ds_idx - n_depots))
    else
        goal_str = string("d-", goal_ds_idx)
    end
    loc2 = env.state_graph.vertices[env.depot_sites_to_vtx[goal_str]].state.location

    tt = get_travel_time_estimate(halton_nn_tree, loc1, loc2, estimate_matrix)

    return get_travel_time_estimate(halton_nn_tree, loc1, loc2, estimate_matrix)
end

"""
    get_first_finish(solution::Vector{PR}) where {PR <: PlanResult}

Get the agent that will finish first. Also returns the time at which it finishes
"""
function get_first_finish(solution::Vector{PR}) where {PR <: PlanResult}

    min_cost = Inf
    min_idx = 0

    for (i, pr) in enumerate(solution)
        if pr.cost < min_cost
            min_idx = i
            min_cost = pr.cost
        end
    end

    return (min_idx, min_cost)
end

# NOTE: How to handle the time issue here? (Just take route vertex indices)
"""
    get_replan_constraints(env::MAPFTransitEnv, time_val::Float64, agent_idx::Int64,
                          solution::Vector{PR}) where {PR <: PlanResult}

For single agent replanning, we need the set of all vertex constraints of other agents to ignore as constraints to the replanning agent
"""
function get_replan_constraints(env::MAPFTransitEnv, time_val::Float64, agent_idx::Int64,
                              solution::Vector{PR}) where {PR <: PlanResult}

    avoid_vertex_idxs = Set{Int64}()

    # Iterate over all solutions for agents other than agent_idx and find vertices to avoid
    for (other_idx, soln) in enumerate(solution)

        if other_idx != agent_idx
            agt_soln_states = soln.states

            for (si, (state, t)) in enumerate(agt_soln_states)
                if t >= time_val
                    push!(avoid_vertex_idxs, state.idx)
                end
            end
        end
    end

    avoid_vertex_map = Dict(towards=>avoid_vertex_idxs, from=>avoid_vertex_idxs)

    return MAPFTransitConstraints(avoid_vertex_map)
end

"""
    replan_individual!(env::MAPFTransitEnv, solution::Vector{PR},
                       n_depots::Int64, n_sites::Int64,
                       agent_tours::Vector{Vector{Int64}}, weight::Float64) where {PR <: PlanResult}

Get all the info necessary to replan for INDIVIDUAL. Returns false if drone has no further task. Also returns the planning time
"""
function replan_individual!(env::MAPFTransitEnv, solution::Vector{PR},
                            n_depots::Int64, n_sites::Int64,
                            agent_tours::Vector{Vector{Int64}}, weight::Float64) where {PR <: PlanResult}

    # Get first finishing agent and next task for it
    replan_idx, time_val = get_first_finish(solution)
    finished_site = env.agent_states[replan_idx].task.site
    replan_task = get_next_agent_task(agent_tours[replan_idx], n_depots, n_sites, finished_site)

    if replan_task == nothing
        @warn "Agent has no next task!"
        return (false, 0.0)
    end

    env.agent_states[replan_idx] = AgentState(task=replan_task)
    new_init_state = env.state_graph.vertices[env.depot_sites_to_vtx[string("d-", replan_task.origin)]]
    replan_constraints = get_replan_constraints(env, time_val, replan_idx, solution)

    # Just a sham solver to rerun low level search
    solver = ECBSSolver{MAPFTransitVertexState,MAPFTransitAction,Float64,Makespan,MAPFTransitConflict,MAPFTransitConstraints,MAPFTransitEnv}(env = env, weight = weight)
    el_time = @elapsed new_pr = MultiAgentPathFinding.low_level_search!(solver, replan_idx, new_init_state, replan_constraints, solution)
    solution[replan_idx] = new_pr

    return (true, el_time)
end

"""
    replan_collective!(env::MAPFTransitEnv, solution::Vector{PR},
                       n_depots::Int64, n_sites::Int64,
                       agent_tours::Vector{Vector{Int64}}, weight::Float64) where {PR <: PlanResult}

Replan collectively! In this we want to update the agent states for all agents and then re-run the solver.
Returns true if valid, the elapsed time, and the new solution
"""
function replan_collective!(env::MAPFTransitEnv, solution::Vector{PR},
                            n_depots::Int64, n_sites::Int64,
                            agent_tours::Vector{Vector{Int64}}, weight::Float64) where {PR <: PlanResult}

    replan_idx, time_val = get_first_finish(solution)
    finished_site = env.agent_states[replan_idx].task.site
    replan_task = get_next_agent_task(agent_tours[replan_idx], n_depots, n_sites, finished_site)

    if replan_task == nothing
        @warn "Agent has no next task!"
        return (false, 0.0, solution)
    end

    env.agent_states[replan_idx] = AgentState(task=replan_task)

    true_n_agents = length(agent_tours)
    new_initial_states = Vector{MAPFTransitVertexState}(undef, true_n_agents)

    for i = 1:true_n_agents
        if i == replan_idx
            new_initial_states[i] = env.state_graph.vertices[env.depot_sites_to_vtx[string("d-", replan_task.origin)]]
            env.plan_ref_times[i] = 0.0
        else
            new_state, ref_time = update_agent_states!(env, time_val, i, solution)
            new_initial_states[i] = new_state
            env.plan_ref_times[i] = ref_time
        end
    end

    solver = ECBSSolver{MAPFTransitVertexState,MAPFTransitAction,Float64,Makespan,MAPFTransitConflict,MAPFTransitConstraints,MAPFTransitEnv}(env = env, weight = weight)

    el_time = @elapsed new_solution = search!(solver, new_initial_states)

    return (true, el_time, new_solution)
end

"""
    set_solution_diagnostics!(env::MAPFTransitEnv, solution::Vector{PR}) where {PR <: PlanResult}

Set the number of transit options and valid distances
"""
function set_solution_diagnostics!(env::MAPFTransitEnv, solution::Vector{PR}) where {PR <: PlanResult}

    # Iterate over solutions and only update if valid
    for (agent_idx, agent_soln) in enumerate(solution)

        # You know start state and site state - get dist
        agt_task = env.agent_states[agent_idx].task
        start_vtx = env.state_graph.vertices[env.depot_sites_to_vtx[string("d-", agt_task.origin)]]
        site_vtx = env.state_graph.vertices[env.depot_sites_to_vtx[string("s-", agt_task.site)]]
        dist = distance_traversed(env, start_vtx, site_vtx)

        # Now only update IF valid
        if dist < env.drone_params.max_distance || env.curr_site_points[agent_idx] > 2
            # IS VALID
            push!(env.valid_path_dists, dist)

            # Now search for transit options
            transit_opts = 0
            for (act, _) in agent_soln.actions
                if act.action == Board::ActionType
                    transit_opts += 1
                end
            end
            push!(env.valid_transit_options, transit_opts)
        end
    end
end


text:
function latlon(s::Simulation,map_g_point_id::Int64)
	osm_node_ix = s.m.n[map_g_point_id]
	lla = LLA(s.m.nodes[osm_node_ix], s.m.bounds)
    return (lla.lat, lla.lon)
end

function plot_sim(s::Simulation, filename::AbstractString; agentIds=1:min(1000,length(s.agents)), tiles="Stamen Toner" )
	MAP_BOUNDS = [( s.m.bounds.min_y, s.m.bounds.min_x),( s.m.bounds.max_y, s.m.bounds.max_x)]
	map_size = (abs(MAP_BOUNDS[1][1]-MAP_BOUNDS[2][1]), abs(MAP_BOUNDS[1][2]-MAP_BOUNDS[2][2]))
	MAP_BOUNDSx9 = [ (MAP_BOUNDS[1][1]-map_size[1], MAP_BOUNDS[1][2]-map_size[1]),
					 (MAP_BOUNDS[2][1]+map_size[1], MAP_BOUNDS[2][2]+map_size[1])]

    flm = pyimport("folium")
    matplotlib_cm = pyimport("matplotlib.cm")
    matplotlib_colors = pyimport("matplotlib.colors")
    cmap = matplotlib_cm.get_cmap("prism")
    m_plot = flm.Map(tiles=tiles)

	map_g_connected_points_set = Set(s.map_g_connected_points)
    for e in edges(s.m.g)
	    (!(e.src in map_g_connected_points_set) || !(e.dst in map_g_connected_points_set)) && continue;
		flm.PolyLine( 	(latlon(s,e.src), latlon(s,e.dst)),
						color="brown", weight=4, opacity=1).add_to(m_plot)
	end

    for vm in s.vertex_meta
		lla = LLA(vm.enu,s.m.bounds)
        info = 	(vm.randomNode ? "Random node\n
(agent randomly chooses next edge here)\n
 OSM id: $(vm.osmid)" : "Deterministic node
\n(added artificially in the discretization process, agent continous in the same direction)") *
			"\n
Node: $(vm.ix)\n
Lattitude: $(lla.lat)\n
Longitude: $(lla.lon) "
        flm.Circle(
			(lla.lat, lla.lon),
			popup=info,
			tooltip=info,
			radius=10,
			color=(vm.randomNode ? "orange" : "yellow"),
			weight=3,
			fill=true,
			fill_color=(vm.randomNode ? "orange" : "yellow")
      	).add_to(m_plot)
    end

	jitter = 0.5e-4

    for agent in s.agents
		lla = LLA(s.vertex_meta[agent.current_node].enu,s.m.bounds)
        info = "Agent: $(agent.id)\n
Infected: " *
			(agent.infected ? "YES" : "NO")*
			"\n
Current node: $(agent.current_node)"*
			"\n
Previous node: $(agent.last_node)"*
			"\n
Total distance travelled so far $(round(Int,agent.total_route_len))m"
		loc = (lla.lat+jitter*randn(), lla.lon+jitter*randn())
		sizex = agent.infected ? 0.00014 : 0.00016
		sizey = agent.infected ? 0.0002 : 0.00022
        flm.Rectangle(
			[(loc[1]-sizex, loc[2]-sizey), (loc[1]+sizex, loc[2]+sizey)],
			popup=info,
			tooltip=info,
			color=(agent.infected ? "green" : "black"),
			weight=(agent.infected ? 5 : 1.5),
			fill=(agent.infected ? false : true),
			fill_opacity=(agent.infected ? 0.2 : 1.0),
			fill_color=(agent.infected ? "green" : "#FAFAFA"),
		).add_to(m_plot)
    end

    MAP_BOUNDS = [( s.m.bounds.min_y, s.m.bounds.min_x),( s.m.bounds.max_y, s.m.bounds.max_x)]
    flm.Rectangle(MAP_BOUNDS, color="black",weight=4).add_to(m_plot)
    m_plot.fit_bounds(MAP_BOUNDS)
    m_plot.save(filename)
end


text:
using KnetLayers, ArgParse, BytePairEncoding
import KnetLayers: IndexedDict, arrtype, PadSequenceArray
using KnetLayers.Datasets
using KnetLayers.Datasets: StoryCloze

include("gpt.jl")

if isempty(ARGS)
    push!(ARGS,"rocstories")
end

function parse_commandline()
    s = ArgParseSettings()

    @add_arg_table s begin
        "--gpu", "-g"
            help = "use gpu"
            action = :store_true
        "--epoch", "-e"
            help = "epoch"
            arg_type = Int
            default = 3
        "task"
            help = "task name"
            required = true
            range_tester = x -> x ∈ ["rocstories"]
    end

    return parse_args(ARGS, s)
end

const args = parse_commandline()

if args["gpu"]
    KnetLayers.settype!(KnetArray)
end

function transform(s1, s2, s3, s4, c1, c2, y)
    x = [startsym;
         segment(bpe, s1);
         segment(bpe, s2);
         segment(bpe, s3);
         segment(bpe, s4);
         delisym]
    x1 = [x; segment(bpe, c1); clfsym]
    x2 = [x; segment(bpe, c2); clfsym]

    x1, x2, y
end

function preprocess(batch, vocab_in, vocab_out)
    tdb = transform.(batch...)
    b1, b2, y = Datasets.batched(tdb)
    b1_mask = GPT.getmask(b1)
    b2_mask = GPT.getmask(b2)
    c1i = length.(b1)
    c2i = length.(b2)
    #c1i = [(findfirst(isequal(clfsym), x), i) for (i, x) in enumerate(b1)]
    #c2i = [(findfirst(isequal(clfsym), x), i) for (i, x) in enumerate(b2)]
    b1, b2 = map(b->vocab_in[b], b1), map(b->vocab_in[b], b2)
    #y = onehotbatch(y, labels)
    return PadSequenceArray(b1, pad=1)', PadSequenceArray(b2, pad=1)', c1i, c2i, vocab_out[y], b1_mask, b2_mask
end


text:
############################################################
# judiJacobianExQ ###############################################
############################################################

# Authors: Philipp Witte (pwitte@eos.ubc.ca), Henryk Modzelewski (hmodzelewski@eos.ubc.ca)
# Date: January 2017

export judiJacobianExQ, judiJacobianExQException, subsample

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

# Type for linear operator representing  Pr*A(m)^-1*Ps,
# i.e. it includes source and receiver projections
struct judiJacobianExQ{DDT<:Number,RDT<:Number} <: judiAbstractJacobian{DDT,RDT}
    name::String
    m::Integer
    n::Integer
    info::Info
    model::Model
    recGeometry::Geometry
    wavelet
    weights
    options::Options
    fop::Function              # forward
    fop_T::Union{Function, Nothing}  # transpose
end

mutable struct judiJacobianExQException <: Exception
    msg :: String
end

judiJacobian(J::judiJacobianExQ{DDT,RDT}; name=J.name, m=J.m, n=J.n, info=J.info, model=J.model, wavelet=J.wavelet,
     weights=J.weights, geom=J.recGeometry, opt=J.options, fop=J.fop, fop_T=J.fop_T) where {DDT, RDT} =
        judiJacobianExQ{DDT,RDT}(name, m, n, info, model, geom, wavelet, weights, opt, fop, fop_T)

############################################################
## Constructor
"""
    judiJacobianExQ(F,q)
Create a linearized modeling operator from the non-linear modeling operator `F` and \\
the source `q`. `F` is a full modeling operator including source/receiver projections.
Examples
========
1) `F` is a modeling operator without source/receiver projections:
    J = judiJacobianExQ(Pr*F*Ps',q)
2) `F` is the combined operator `Pr*F*Ps'`:
    J = judiJacobianExQ(F,q)
"""
function judiJacobian(F::judiPDEextended, weights::Union{judiWeights, Array}; DDT::DataType=Float32, RDT::DataType=DDT)
# JOLI wrapper for nonlinear forward modeling w/ extended source

    (DDT == Float32 && RDT == Float32) || throw(judiJacobianExQException("Domain and range types not supported"))
    m = n_samples(F.recGeometry, F.info)
    n = F.info.n

    weights = process_input_data(weights, F.model, F.info)  # extract cell array

    return J = judiJacobianExQ{Float32,Float32}("linearized wave equation", m, n, F.info, F.model, F.recGeometry, F.wavelet, weights, F.options,
                                                bornop, adjbornop)
end

############################################################
## Additional overloaded functions

# Subsample Jacobian
function subsample(J::judiJacobianExQ{ADDT,ARDT}, srcnum) where {ADDT,ARDT}
    recGeometry = subsample(J.recGeometry,srcnum)
    nsrc = typeof(srcnum) <: Int ? 1 : length(srcnum)
    info = Info(J.info.n, nsrc, J.info.nt[srcnum])
    m = n_samples(recGeometry, info)

    opt = subsample(J.options, srcnum)
    nsrc == 1 && (srcnum = srcnum:srcnum)
    return judiJacobian(J; m=m, info=info, weights=J.weights[srcnum], wavelet=J.wavelet[srcnum],
                        geom=recGeometry, opt=opt)
end

# *(num,judiJacobian)
*(a::Number,A::judiJacobianExQ{ADDT,ARDT}) where {ADDT,ARDT} =  judiJacobian(A; wavelet=a*A.wavelet)
# -(judiJacobian)
-(A::judiJacobianExQ{DDT,RDT}) where {DDT,RDT} = judiJacobian(A; wavelet=-A.wavelet)

############################################################
## Forward/adjoint function to avoid unecessary extra declaration

function bornop(J::judiJacobianExQ, v)
    srcnum = 1:J.info.nsrc
    return extended_source_modeling(J.model, J.wavelet, J.recGeometry, nothing, J.weights, v, srcnum, 'J', 1, J.options)
end

function adjbornop(J::judiJacobianExQ, w)
    srcnum = 1:J.info.nsrc
    return extended_source_modeling(J.model, J.wavelet, J.recGeometry, process_input_data(w, J.info),
                                    J.weights, nothing, srcnum, 'J', -1, J.options)
end

max_stars_repo_name: nogueirapeterson/JUDI

max_issues_repo_name: nogueirapeterson/JUDI

max_forks_repo_name: nogueirapeterson/JUDI


text:
module SIMDVectors

importall Base

export SIMDVector, load, store!

const VEC_REGISTER_SIZE_BITS = 128

const VECTOR_DATATYPES = Set{DataType}(
	 (Bool,
	  Int8,Int16, Int32, Int64, Int128,
	  UInt8, UInt16, UInt32, UInt64, UInt128,
	  Float16, Float32, Float64))

# List taken from FixedSizeArrays.jl
const UNARY_FUNCS = (:-, :~, :conj, :abs,
                  :sin, :cos, :tan, :sinh, :cosh, :tanh,
                  :asin, :acos, :atan, :asinh, :acosh, :atanh,
                  :sec, :csc, :cot, :asec, :acsc, :acot,
                  :sech, :csch, :coth, :asech, :acsch, :acoth,
                  :sinc, :cosc, :cosd, :cotd, :cscd, :secd,
                  :sind, :tand, :acosd, :acotd, :acscd, :asecd,
                  :asind, :atand, :rad2deg, :deg2rad,
                  :log, :log2, :log10, :log1p, :exponent, :exp,
                  :exp2, :expm1, :cbrt, :sqrt, :erf,
                  :erfc, :erfcx, :erfi, :dawson, :ceil, :floor,
                  :trunc, :round, :significand, :lgamma,
                  :gamma, :lfact, :airy, :airyai,
                  :airyprime, :airyaiprime, :airybi, :airybiprime,
                  :besselj0, :besselj1, :bessely0, :bessely1,
                  :eta, :zeta, :digamma)

const BINARY_FUNCS  = (:.+, :.-,:.*, :./, :.\, :.^,
                   :.==, :.!=, :.<, :.<=, :.>, :.>=, :+, :-,
                   :min, :max,
                   :div, :fld, :rem, :mod, :mod1,
                   :atan2, :besselj, :bessely, :hankelh1, :hankelh2,
                   :besseli, :besselk, :beta, :lbeta)

const REDUCTION_FUNCS = ((:maximum, :max),
                         (:minimum,  :min),
                         (:sum, :.+),
                         (:prod, :.*))

include("tuple_ops.jl")
include("VecRegister.jl")
include("SIMDVector.jl")
include("conversion_promotion.jl")

end # module


text:
using LinearAlgebra

struct ValueOne; end
ValueOne()

# Compute X <- a X + b I.
function matfun_axpby!(X,a,b,Y::UniformScaling)
    m,n=size(X)
    if ~(a isa ValueOne)
        rmul!(X,a)
    end
    @inbounds for i=1:n
        X[i,i]+=(b isa ValueOne) ? 1 : b
    end
end

# Compute X <- a X + b Y.
function matfun_axpby!(X,a,b,Y)
    m,n=size(X)
    if ~(a isa ValueOne)
        rmul!(X,a)
    end
    @inbounds for i=1:m
        @inbounds for j=1:n
            if (b isa ValueOne)
                X[i,j]+=Y[i,j]
            else
                X[i,j]+=b*Y[i,j]
            end
        end
    end
end

@inline function exp_bbc_m5(A)
    T=promote_type(eltype(A),Float64)
    A_copy=similar(A,T); A_copy .= A;
    return exp_bbc_m5!(A_copy)
end

@inline function exp_bbc_m5!(A)
    T=promote_type(eltype(A),Float64) # Make it work for many 'bigger' types (matrices and scalars)
    # max_memslots=7
    n=size(A,1)
    # The first slots are precomputed nodes [:A]
    memslots2 = similar(A,T)
    memslots3 = similar(A,T)
    memslots4 = similar(A,T)
    memslots5 = similar(A,T)
    memslots6 = similar(A,T)
    memslots7 = similar(A,T)
    # Assign precomputed nodes memslots
    memslots1=A # overwrite A
    # Uniform scaling is exploited.
    # No matrix I explicitly allocated.
    value_one=ValueOne()
    # Computation order: B2 B3 Ba5_4 B4 Bb5 B5 Ba6 Bb6 B6 T2k8
    # Computing B2 with operation: mult
    mul!(memslots2,memslots1,memslots1)
    # Computing B3 with operation: mult
    mul!(memslots3,memslots1,memslots2)
    # Computing Ba5_4 = x*A+x*B2+x*B3
    coeff1=-0.10036558103014462
    coeff2=-0.00802924648241157
    coeff3=-0.00089213849804573
    memslots4 .= coeff1.*memslots1 .+ coeff2.*memslots2 .+ coeff3.*memslots3
    # Computing B4 with operation: mult
    mul!(memslots5,memslots3,memslots3)
    # Computing Bb5 = x*B2+x*B3+x*B4
    coeff1=-0.09233646193671186
    coeff2=-0.016936493900208172
    coeff3=-1.400867981820361e-5
    memslots6 .= coeff1.*memslots2 .+ coeff2.*memslots3 .+ coeff3.*memslots5
    # Computing B5 with operation: mult
    mul!(memslots7,memslots4,memslots6)
    # Deallocating Ba5_4 in slot 4
    # Deallocating Bb5 in slot 6
    # Computing Ba6 = x*I+x*A+x*B2+x*B3+x*B4+x*B5
    coeff1=-11.058071288535288
    coeff2=1.6125176868819238
    coeff3=0.12477411482493252
    coeff4=0.02257315581805103
    coeff5=1.9579475957000978e-5
    coeff6=1.0
    memslots4 .= coeff2.*memslots1 .+ coeff3.*memslots2 .+ coeff4.*memslots3 .+ coeff5.*memslots5 .+ coeff6.*memslots7
    mul!(memslots4,true,I*coeff1,true,true)
    # Computing Bb6 = x*I+x*A+x*B2+x*B3+x*B4+x*B5
    coeff1=-0.09043168323908106
    coeff2=-0.06764045190713819
    coeff3=0.06759613017704597
    coeff4=0.029555257042931552
    coeff5=-1.391802575160607e-5
    coeff6=1.0
    # Smart lincomb recycle B5
    memslots7 .= coeff2.*memslots1 .+ coeff3.*memslots2 .+ coeff4.*memslots3 .+ coeff5.*memslots5 .+ coeff6.*memslots7
    mul!(memslots7,true,I*coeff1,true,true)
    # Computing B6 with operation: mult
    mul!(memslots6,memslots4,memslots7)
    # Deallocating Ba6 in slot 4
    # Deallocating Bb6 in slot 7
    # Computing T2k8 = x*A+x*B2+x*B3+x*B4+x*B6
    coeff1=0.3978497494996451
    coeff2=1.3678377846041172
    coeff3=0.49828962252538267
    coeff4=-0.0006378981945947233
    coeff5=1.0
    # Smart lincomb recycle A
    memslots1 .= coeff1.*memslots1 .+ coeff2.*memslots2 .+ coeff3.*memslots3 .+ coeff4.*memslots5 .+ coeff5.*memslots6
    # Deallocating B2 in slot 2
    # Deallocating B3 in slot 3
    # Deallocating B4 in slot 5
    # Deallocating B6 in slot 6
    return memslots1 # Returning T2k8
end

max_stars_repo_name: matrixfunctions/GraphMatFunData

max_issues_repo_name: matrixfunctions/GraphMatFunData

max_forks_repo_name: matrixfunctions/GraphMatFunData


text:
module TestUtils

using Aqua: askwargs, format_diff
using Test

@testset "askwargs" begin
    @test_throws ArgumentError("expect `true`") askwargs(false)
    @test askwargs(true) === NamedTuple()
    @test askwargs(()) === NamedTuple()
    @test askwargs((a = 1,)) === (a = 1,)
end

@testset "format_diff" begin
    @testset "normal" begin
        if Sys.which("diff") === nothing
            @info "Comamnd `diff` not found; skip testing `format_diff`."
        else
            diff = format_diff("LABEL_A" => "TEXT_A", "LABEL_B" => "TEXT_B")
            @test occursin("--- LABEL_A", diff)
            @test occursin("+++ LABEL_B", diff)
        end
    end
    @testset "fallback" begin
        diff = withenv("PATH" => "/") do
            format_diff("LABEL_A" => "TEXT_A", "LABEL_B" => "TEXT_B")
        end
        @test occursin("*** LABEL_A ***", diff)
    end
end

end  # module