text:
using LinearAlgebra
using ForwardDiff
using Zygote: hessian_dual, hessian_reverse

@testset "hessian: $hess" for hess in [hessian_dual, hessian_reverse]

  if hess == hessian_dual
    @test hess(x -> x[1]*x[2], randn(2)) ≈ [0 1; 1 0]
    @test hess(((x,y),) -> x*y, randn(2)) ≈ [0 1; 1 0]  # original docstring version
  else
    @test_broken hess(x -> x[1]*x[2], randn(2)) ≈ [0 1; 1 0]  # can't differentiate ∇getindex
    @test_broken hess(((x,y),) -> x*y, randn(2)) ≈ [0 1; 1 0]
  end
  @test hess(x -> sum(x.^3), [1 2; 3 4]) ≈ Diagonal([6, 18, 12, 24])
  @test hess(sin, pi/2) ≈ -1

  @test_throws Exception hess(sin, im*pi)
  @test_throws Exception hess(x -> x+im, pi)
  @test_throws Exception hess(identity, randn(2))
end

@testset "diagonal hessian" begin
  @test diaghessian(x -> x[1]*x[2]^2, [1, pi]) == ([0, 2],)

  xs, y = randn(2,3), rand()
  f34(xs, y) = xs[1] * (sum(xs .^ (1:3)') + y^4)  # non-diagonal Hessian, two arguments
  dx, dy = diaghessian(f34, xs, y)
  @test size(dx) == size(xs)
  @test vec(dx) ≈ diag(hessian(x -> f34(x,y), xs))
  @test dy ≈ hessian(y -> f34(xs,y), y)

  zs = randn(7,13)  # test chunk mode
  @test length(zs) > ForwardDiff.DEFAULT_CHUNK_THRESHOLD
  @test length(zs) % ForwardDiff.DEFAULT_CHUNK_THRESHOLD != 0
  f713(zs) = sum(vec(zs)' .* exp.(vec(zs)))
  @test vec(diaghessian(f713, zs)[1]) ≈ diag(hessian(f713, zs))

  @test_throws Exception diaghessian(sin, im*pi)
  @test_throws Exception diaghessian(x -> x+im, pi)
  @test_throws Exception diaghessian(identity, randn(2))
end

@testset "jacobian(f, args...)" begin
  @test jacobian(identity, [1,2])[1] == [1 0; 0 1]
  @test withjacobian(identity, [1,2]) == (val = [1,2], grad = ([1 0; 0 1],))

  j1 = jacobian((a,x) -> a.^2 .* x, [1,2,3], 1)
  @test j1[1] ≈ Diagonal([2,4,6])
  @test j1[2] ≈ [1, 4, 9]
  @test j1[2] isa Vector

  j2 = jacobian((a,t) -> sum(a .* t[1]) + t[2], [1,2,3], (4,5))  # scalar output is OK
  @test j2[1] == [4 4 4]
  @test j2[1] isa Matrix
  @test j2[2] === nothing  # input other than Number, Array is ignored

  j3 = jacobian((a,d) -> prod(a, dims=d), [1 2; 3 4], 1)
  @test j3[1] ≈ [3 1 0 0; 0 0 4 2]
  @test j3[2] ≈ [0, 0]  # pullback is always Nothing, but array already allocated

  j4 = jacobian([1,2,-3,4,-5]) do xs
    map(x -> x>0 ? x^3 : 0, xs)  # pullback gives Nothing for some elements x
  end
  @test j4[1] ≈ Diagonal([3,12,0,48,0])

  j5 = jacobian((x,y) -> hcat(x[1], y), fill(pi), exp(1))  # zero-array
  @test j5[1] isa Matrix
  @test vec(j5[1]) == [1, 0]

  @test_throws ArgumentError jacobian(identity, [1,2,3+im])
  @test_throws ArgumentError jacobian(sum, [1,2,3+im])  # scalar, complex

  f6(x,y) = abs2.(x .* y)
  g6 = gradient(first∘f6, [1+im, 2], 3+4im)
  j6 = jacobian((x,y) -> abs2.(x .* y), [1+im, 2], 3+4im)
  @test j6[1][1,:] ≈ g6[1]
  @test j6[2][1] ≈ g6[2]
end

@testset "jacobian(loss, ::Params)" begin
  xs = [1 2; 3 4]
  ys = [5,7,9];
  Jxy = jacobian(() -> ys[1:2] .+ sum(xs.^2), Params([xs, ys]))
  @test Jxy[ys] ≈ [1 0 0; 0 1 0]
  @test Jxy[xs] ≈ [2 6 4 8; 2 6 4 8]

  z, grad = withjacobian(() -> ys[1:2] .+ sum(xs.^2), Params([xs, ys]))
  @test z == [35, 37]
  @test grad[ys] ≈ [1 0 0; 0 1 0]
end

using ForwardDiff

@testset "adjoints of ForwardDiff functions" begin
  f1(x) = ForwardDiff.gradient(x -> sum(exp.(x.+1)), x)
  x1 = randn(3,7)
  @test Zygote.jacobian(f1, x1)[1] ≈ ForwardDiff.jacobian(f1, x1)

  f2(x) = ForwardDiff.jacobian(x -> log.(x[1:3] .+ x[2:4]), x)
  x2 = rand(5) .+ 1
  @test Zygote.jacobian(f2, x2)[1] ≈ ForwardDiff.jacobian(f2, x2)

  f3(x) = sum(ForwardDiff.hessian(x -> sum(x .^2 .* x'), x)[1:4:end])
  x3 = rand(3)
  @test Zygote.gradient(f3, x3)[1] ≈ ForwardDiff.gradient(f3, x3)

  @test gradient(x -> ForwardDiff.derivative(x -> x^4, x), 7) == (4 * 3 * 7^2,)

  f4(x) = ForwardDiff.derivative(x -> [x,x^2,x^3], x)
  @test Zygote.jacobian(f4, pi)[1] ≈ ForwardDiff.derivative(f4, pi)

  # Tests from https://github.com/FluxML/Zygote.jl/issues/769
  f(x) = [2x[1]^2 + x[1],x[2]^2 * x[1]]
  g1(x) = sum(ForwardDiff.jacobian(f,x))
  out,back = Zygote.pullback(g1,[2.0,3.2])
  stakehouse = back(1.0)[1]
  @test typeof(stakehouse) <: Vector
  @test size(stakehouse) == (2,)
  @test stakehouse ≈ ForwardDiff.gradient(g1,[2.0,3.2])

  g2(x) = prod(ForwardDiff.jacobian(f,x))
  out,back = Zygote.pullback(g2,[2.0,3.2])
  @test_skip back(1.0)[1] == ForwardDiff.gradient(g2,[2.0,3.2])  # contains NaN, @adjoint prod isn't careful

  g3(x) = sum(abs2,ForwardDiff.jacobian(f,x))
  out,back = Zygote.pullback(g3,[2.0,3.2])
  @test back(1.0)[1] == ForwardDiff.gradient(g3,[2.0,3.2])
end


text:
using SPN
using Base.Test

function test_inner_node_construction()
	# Create a product node and check the values of its fields:
	p = ProdNode()
	@test p.logval == -Inf
	@test p.parents == InnerNode[]
	@test p.children == Node[]
	@test p.scope == Int[]

	# Create a sum node and check the values of its fields:
	s = SumNode()
	@test s.logval == -Inf
	@test s.parents == InnerNode[]
	@test s.children == Node[]
	@test s.scope == Int[]
	@test s.weights == Float64[]
end

function test_indicator_node_construction()
	# Create an indicator node and check the values of its fields:
	i = IndicatorNode(1,3.0)
	@test i.logval == -Inf
	@test i.parents == InnerNode[]
	@test length(i.scope) == 1
	@test i.scope == [1]
	@test i.indicates == 3.0

	# Indicator nodes for integer and boolean variables:
	i2 = IndicatorNode(2,1)
	@test i2.indicates == 1.0
	i3 = IndicatorNode(3,true)
	@test i3.indicates == 1.0
	i4 = IndicatorNode(3,false)
	@test i4.indicates == 0.0
end

function test_connect_nodes()
	s = SumNode()
	p1 = ProdNode()
	p2 = ProdNode()
	i1 = IndicatorNode(1,0)
	i2 = IndicatorNode(1,1)
	i3 = IndicatorNode(1,2)

	connect!(s,p1,weight=0.3)
	connect!(s,p2,weight=0.7)
	connect!(p1,i1)
	connect!(p1,i2)
	connect!(p2,i2)
	connect!(p2,i3)

	@test length(s.children) == 2
	@test length(p1.children) == 2
	@test length(p2.children) == 2
	@test p1 in s.children
	@test p2 in s.children
	@test i1 in p1.children
	@test !(i1 in p2.children)
end

function test_setinput_indicator_node()
	i1 = IndicatorNode(1, 1.0)
	i2 = IndicatorNode(2, 2.0)
	i3 = IndicatorNode(3, 3.0)
	x1 = Float64[1.0, 5.0]
	x2 = Float64[1.0, 2.0]
	setInput!(i1, x1)
	setInput!(i2, x1)
	@test i1.logval == 0.0
	@test i2.logval == -Inf

	setInput!(i1, x2)
	@test i1.logval == 0.0
	setInput!(i2, x2)
	@test i2.logval == 0.0
	@test_throws AssertionError setInput!(i3, x1)
end

function test_setinput_indicator_node_partial_evidence()
	i1 = IndicatorNode(1, 2.0)
	i2 = IndicatorNode(1, 3.0)
	i3 = IndicatorNode(2, 2.0)
	i4 = IndicatorNode(2, 3.0)

	x = Float64[2.0, 3.0]
	e = BitVector([true, false])

	setInput!(i1, x, e)
	setInput!(i2, x, e)
	setInput!(i3, x, e)
	setInput!(i4, x, e)

	@test i1.logval == 0.0
	@test i2.logval == -Inf
	@test i3.logval == 0.0
	@test i4.logval == 0.0
end

function test_eval_inner_nodes()
	s, p1, p2, p3, s1, s2, s3, s4, i1, i2, i3, i4 = create_toy_spn()

	x = [true, false]

	# Calculate value of nodes manually (not in logspace):
	s1_val = 0.6*1+0.4*0
	s2_val = 0.9*1+0.1*0
	s3_val = 0.3*0+0.7*1
	s4_val = 0.2*0+0.8*1
	p1_val = s1_val*s3_val
	p2_val = s1_val*s4_val
	p3_val = s2_val*s4_val
	s_val = 0.5*p1_val + 0.2*p2_val + 0.3*p3_val

	# set input of indicator nodes:
	setInput!(i1, x)
	setInput!(i2, x)
	setInput!(i3, x)
	setInput!(i4, x)

	# Compare results:
	@test eval!(i1) ≈ log(1)
	@test eval!(i2) ≈ log(0)
	@test eval!(i3) ≈ log(0)
	@test eval!(i4) ≈ log(1)
	@test eval!(s1) ≈ log(s1_val)
	@test eval!(s2) ≈ log(s2_val)
	@test eval!(s3) ≈ log(s3_val)
	@test eval!(s4) ≈ log(s4_val)
	@test eval!(p1) ≈ log(p1_val)
	@test eval!(p2) ≈ log(p2_val)
	@test eval!(p3) ≈ log(p3_val)
	@test eval!(s) ≈ log(s_val)
end

"""
Test the evaluation with the max keyword argument set to true.
Every sum node becomes a max node whose value is its maximum weighted child value.
"""
function test_eval_max()
	s, p1, p2, p3, s1, s2, s3, s4, i1, i2, i3, i4 = create_toy_spn()

	x = [true, false]
	e = BitVector([true, false])  # we only know the state of the first variable

	# Calculate value of nodes manually (not in logspace):
	s1_val = 0.6
	s1_idx = 1
	s2_val = 0.9
	s2_idx = 1
	s3_val = 0.7
	s3_idx = 2
	s4_val = 0.8
	s4_idx = 2
	p1_val = s1_val*s3_val
	p2_val = s1_val*s4_val
	p3_val = s2_val*s4_val
	s_val = 0.3*p3_val
	s_idx = 3

	setInput!(i1, x, e)
	setInput!(i2, x, e)
	setInput!(i3, x, e)
	setInput!(i4, x, e)

	@test eval!(i1, max=true) ≈ log(1)
	@test eval!(i2, max=true) ≈ log(0)
	@test eval!(i3, max=true) ≈ log(1)
	@test eval!(i4, max=true) ≈ log(1)
	@test eval!(s1, max=true) ≈ log(s1_val)
	@test eval!(s2, max=true) ≈ log(s2_val)
	@test eval!(s3, max=true) ≈ log(s3_val)
	@test eval!(s4, max=true) ≈ log(s4_val)
	@test eval!(p1, max=true) ≈ log(p1_val)
	@test eval!(p2, max=true) ≈ log(p2_val)
	@test eval!(p3, max=true) ≈ log(p3_val)
	@test eval!(s, max=true) ≈ log(s_val)
	@test s1.maxidx == s1_idx
	@test s2.maxidx == s2_idx
	@test s3.maxidx == s3_idx
	@test s4.maxidx == s4_idx
	@test s.maxidx == s_idx
end

function test_normalize()
	s, p1, p2, p3, s1, s2, s3, s4, i1, i2, i3, i4 = create_toy_spn()
	s.weights *= 3
	normalize!(s)
	@test s.weights[1] ≈ 0.5
	@test s.weights[2] ≈ 0.2
	@test s.weights[3] ≈ 0.3

	s1.weights *= 5
	normalize!(s1)
	@test s1.weights[1] ≈ 0.6
	@test s1.weights[2] ≈ 0.4
end

test_inner_node_construction()
test_indicator_node_construction()
test_connect_nodes()
test_setinput_indicator_node()
test_setinput_indicator_node_partial_evidence()
test_eval_inner_nodes()
test_eval_max()
test_normalize()

ext: jl


text:
# United States calendars

# In the United States, if a holiday falls on Saturday, it's observed on the preceding Friday.
# If it falls on Sunday, it's observed on the next Monday.
function adjustweekendholidayUS(dt::Date)

    if dayofweek(dt) == Dates.Saturday
        return dt - Dates.Day(1)
    end

    if dayofweek(dt) == Dates.Sunday
        return dt + Dates.Day(1)
    end

    return dt
end

function isholiday(::USSettlement , dt::Date)

    const yy = Dates.year(dt)
    const mm = Dates.month(dt)
    const dd = Dates.day(dt)

    if (
            # New Year's Day
            adjustweekendholidayUS(Date(yy, 1, 1)) == dt
            ||
            # New Year's Day on the previous year when 1st Jan is Saturday
            (mm == 12 &&  dd == 31 && dayofweek(dt) == Friday)
            ||
            # Birthday of Martin Luther King, Jr.
            (yy >= 1983 && adjustweekendholidayUS(findweekday(Dates.Monday, yy, 1, 3, true)) == dt)
            ||
            # Washington's Birthday
            adjustweekendholidayUS(findweekday(Dates.Monday, yy, 2, 3, true)) == dt
            ||
            # Memorial Day
            adjustweekendholidayUS(findweekday(Dates.Monday, yy, 5, 1, false)) == dt
            ||
            # Independence Day
            adjustweekendholidayUS(Date(yy, 7, 4)) == dt
            ||
            # Labor Day
            adjustweekendholidayUS(findweekday(Dates.Monday, yy, 9, 1, true)) == dt
            ||
            # Columbus Day
            adjustweekendholidayUS(findweekday(Dates.Monday, yy, 10, 2, true)) == dt
            ||
            # Veterans Day
            adjustweekendholidayUS(Date(yy, 11, 11)) == dt
            ||
            # Thanksgiving Day
            adjustweekendholidayUS(findweekday(Dates.Thursday, yy, 11, 4, true)) == dt
            ||
            # Christmas
            adjustweekendholidayUS(Date(yy, 12, 25)) == dt
        )
        return true
    end

    return false
end

function isholiday(::USNYSE , dt::Date)

    const yy = Dates.year(dt)
    const mm = Dates.month(dt)
    const dd = Dates.day(dt)

    const dt_rata::Int = Dates.days(dt)
    const e_rata::Int = easter_rata(Dates.Year(yy))

    if (
            # New Year's Day
            adjustweekendholidayUS(Date(yy, 1, 1)) == dt
            ||
            # Birthday of Martin Luther King, Jr.
            (yy >= 1998 && adjustweekendholidayUS(findweekday(Dates.Monday, yy, 1, 3, true)) == dt)
            ||
            # Washington's Birthday
            adjustweekendholidayUS(findweekday(Dates.Monday, yy, 2, 3, true)) == dt
            ||
            # Good Friday
            dt_rata == ( e_rata - 2 )
            ||
            # Memorial Day
            adjustweekendholidayUS(findweekday(Dates.Monday, yy, 5, 1, false)) == dt
            ||
            # Independence Day
            adjustweekendholidayUS(Date(yy, 7, 4)) == dt
            ||
            # Labor Day
            adjustweekendholidayUS(findweekday(Dates.Monday, yy, 9, 1, true)) == dt
            ||
            # Thanksgiving Day
            adjustweekendholidayUS(findweekday(Dates.Thursday, yy, 11, 4, true)) == dt
            ||
            # Christmas
            adjustweekendholidayUS(Date(yy, 12, 25)) == dt
        )
        return true
    end

    # Presidential election days
    if (yy <= 1968 || (yy <= 1980 && yy % 4 == 0)) && mm == 11 && dd <= 7 && Dates.istuesday(dt)
        return true
    end

    # Special Closings
    if (
        # Hurricane Sandy
        yy == 2012 && mm == 10 && (dd == 29 || dd == 30)
        ||
        # Predient Ford's funeral
        dt == Date(2007,1,2)
        ||
        # President Reagan's funeral
        dt == Date(2004,6,11)
        ||
        # Sep 11th
        yy == 2001 && mm == 9 && ( 11 <= dd && dd <= 14)
        ||
        # President Nixon's funeral
        dt == Date(1994,4,27)
        ||
        # Hurricane Gloria
        dt == Date(1985,9,27)
        ||
        # 1977 Blackout
        dt == Date(1977,7,14)
        ||
        # Funeral of former President Lyndon B. Johnson
        dt == Date(1973,1,25)
        ||
        # Funeral of former President Harry S. Truman
        dt == Date(1972,12,28)
        ||
        # National Day of Participation for the lunar exploration
        dt == Date(1969,7,21)
        ||
        # Eisenhower's funeral
        dt == Date(1969,3,31)
        ||
        # Heavy snow
        dt == Date(1969,2,10)
        ||
        # Day after Independence Day
        dt == Date(1968,7,5)
        ||
        # Paperwork Crisis
        yy == 1968 && dayofyear(dt) >= 163 && Dates.iswednesday(dt)
        ||
        # Mourning for Martin Luther King Jr
        dt == Date(1968,4,9)
        ||
        # President Kennedy's funeral
        dt == Date(1963,11,25)
        ||
        # Day before Decoration Day
        dt == Date(1961,5,29)
        ||
        # Day after Christmas
        dt == Date(1958,12,26)
        ||
        # Christmas Eve
        dt in [Date(1954,12,24), Date(1956,12,24), Date(1965,12,24)]
        )
        return true
    end

    return false
end

function isholiday(::USGovernmentBond , dt::Date)

    const yy = Dates.year(dt)
    const mm = Dates.month(dt)
    const dd = Dates.day(dt)

    const dt_rata::Int = Dates.days(dt)
    const e_rata::Int = easter_rata(Dates.Year(yy))

    if (
            # New Year's Day
            adjustweekendholidayUS(Date(yy, 1, 1)) == dt
            ||
            # Birthday of Martin Luther King, Jr.
            yy >= 1983 && adjustweekendholidayUS(findweekday(Dates.Monday, yy, 1, 3, true)) == dt
            ||
            # Washington's Birthday
            adjustweekendholidayUS(findweekday(Dates.Monday, yy, 2, 3, true)) == dt
            ||
            # Good Friday
            dt_rata == ( e_rata - 2 )
            ||
            # Memorial Day
            adjustweekendholidayUS(findweekday(Dates.Monday, yy, 5, 1, false)) == dt
            ||
            # Independence Day
            adjustweekendholidayUS(Date(yy, 7, 4)) == dt
            ||
            # Labor Day
            adjustweekendholidayUS(findweekday(Dates.Monday, yy, 9, 1, true)) == dt
            ||
            # Columbus Day
            adjustweekendholidayUS(findweekday(Dates.Monday, yy, 10, 2, true)) == dt
            ||
            # Veterans Day
            adjustweekendholidayUS(Date(yy, 11, 11)) == dt
            ||
            # Thanksgiving Day
            adjustweekendholidayUS(findweekday(Dates.Thursday, yy, 11, 4, true)) == dt
            ||
            # Christmas
            adjustweekendholidayUS(Date(yy, 12, 25)) == dt
        )
        return true
    end

    return false
end


text:
using Oceananigans.AbstractOperations: UnaryOperation, BinaryOperation, MultiaryOperation

# Field and FunctionField (both fields with "grids attached")
const MultiRegionUnaryOperation{LX, LY, LZ, O, A, I, G, T} = UnaryOperation{LX, LY, LZ, O, A, I, <:MultiRegionGrid, T} where {LX, LY, LZ, O, A, I, T}
const MultiRegionBinaryOperation{LX, LY, LZ, O, A, B, IA, IB, G, T} = BinaryOperation{LX, LY, LZ, O, A, B, IA, IB, <:MultiRegionGrid, T} where {LX, LY, LZ, O, A, B, IA, IB, T}
const MultiRegionMultiaryOperation{LX, LY, LZ, N, O, A, I, G, T} = MultiaryOperation{LX, LY, LZ, N, O, A, I, <:MultiRegionGrid, T} where {LX, LY, LZ, N, O, A, I, T}

const MultiRegionAbstractOperation = Union{MultiRegionBinaryOperation, MultiRegionUnaryOperation, MultiRegionMultiaryOperation}

# Utils
Base.size(f::MultiRegionAbstractOperation) = size(getregion(f.grid, 1))

@inline isregional(f::MultiRegionAbstractOperation) = true
@inline devices(f::MultiRegionAbstractOperation)    = devices(f.grid)
sync_all_devices!(f::MultiRegionAbstractOperation)  = sync_all_devices!(devices(f.grid))

@inline switch_device!(f::MultiRegionAbstractOperation, d) = switch_device!(f.grid, d)
@inline getdevice(f::MultiRegionAbstractOperation, d)      = getdevice(f.grid, d)

## Functions applied regionally
compute_at!(f::MultiRegionAbstractOperation, time) = apply_regionally!(compute_at!, f, time)
compute!(f::MultiRegionAbstractOperation)          = apply_regionally!(compute!, f, time)

for T in [:BinaryOperation, :UnaryOperation, :MultiaryOperation]
    @eval begin
        @inline getregion(f::$T{LX, LY, LZ}, r) where {LX, LY, LZ} =
                        $T{LX, LY, LZ}(Tuple(_getregion(getproperty(f, n), r) for n in fieldnames($T))...)

        @inline _getregion(f::$T{LX, LY, LZ}, r) where {LX, LY, LZ} =
                        $T{LX, LY, LZ}(Tuple(getregion(getproperty(f, n), r) for n in fieldnames($T))...)
    end
end


text:
import SwitchOnSafety: isinfeasible, isfeasible, isdecided, usestep
export sosbound_γ, sosbound_gamma

function sos_bisection_alg(s::StateDepDiscreteSwitchedLinearSystem, d, optimizer, lo, hi; verbose=1, tol=1e-5, step=0.5)
    @assert hi > lo
    stuck_count = 0
    while hi - lo ≥ tol
        γ_mid = midpoint(hi, lo, step)
        status = soslyap_feasible_prog(γ_mid, s, d, optimizer, verbose=verbose-1)
        if verbose ≥ 1
            println("$γ_mid")
        end
        if !isdecided(status, true)
            if usestep(lo, hi)
                step *= 2
                continue
            end
            stuck_count += 1
            γ_midlo = min(midpoint(lo, γ_mid, step), shift(γ_mid, -tol/2))
            statuslo = soslyap_feasible_prog(γ_midlo, s, d, optimizer, verbose=verbose-1)
            if isdecided(statuslo, true)
                γ_mid = γ_midlo
                status = statuslo
            else
                γ_midhi = max(midpoint(γ_mid, hi, step), shift(γ_mid, tol/2))
                statushi = soslyap_feasible_prog(γ_midhi, s, d, optimizer, verbose=verbose-1)
                if isdecided(statushi, true)
                    γ_mid = γ_midhi
                    status = statushi
                end
            end
        end
        if isinfeasible(status, true)
            lo = γ_mid
        end
        if isfeasible(status, true)
            hi = γ_mid
        end
        if stuck_count >= 5
            @warn("Making slow progress. Reporting upper bound with reduced accuracy.")
            return lo, hi
        end
    end
    return lo, hi
end

"""
    sosbound_γ(s::StateDepDiscreteSwitchedLinearSystem, d)

Compute minimum value of ``γ`` for which the sum-of-squares program is feasible.

# Keywords
- `optimizer=...`: choose an SDP solver. The default solver is CSDP.
- `tol=...`: set tolerance. The default is `tol=1e-5`.
- `verbose=...`: set level of detail reported during calculation process. The default is `verbose=0`.
- `initstep=...`: set initial step size. The default value is `initstep=1.1`.
"""
function sosbound_γ(s::StateDepDiscreteSwitchedLinearSystem, d; optimizer=optimizer_with_attributes(CSDP.Optimizer, MOI.Silent() => true), tol=1e-5, verbose=0, initstep=1.1)
    lo_init, hi_init = 0, initstep; hi_stat = (MOI.INFEASIBLE, MOI.INFEASIBLE_POINT, MOI.INFEASIBILITY_CERTIFICATE, "")
    while isinfeasible(hi_stat, true)
        hi_stat = soslyap_feasible_prog(hi_init, s, d, optimizer, verbose=verbose-1)
        if isinfeasible(hi_stat, true)
            lo_init = copy(hi_init)
            hi_init += initstep
        end
        initstep *= 2
    end
    lo, hi = sos_bisection_alg(s, d, optimizer, lo_init, hi_init, verbose=verbose)
    return hi
end

"""
    sosbound_gamma(s::StateDepDiscreteSwitchedLinearSystem, d; optimizer=nothing, tol=1e-5, verbose=0, initstep=1.1)

Alias for [`sosbound_γ`](@ref).
"""
sosbound_gamma(s::StateDepDiscreteSwitchedLinearSystem, d; optimizer=optimizer_with_attributes(CSDP.Optimizer, MOI.Silent() => true), tol=1e-5, verbose=0, initstep=1.1) = sosbound_γ(s::StateDepDiscreteSwitchedLinearSystem, d; optimizer=optimizer, tol=tol, verbose=verbose, initstep=initstep)


text:
using Resonance
omni, tps = startup(; dfs = [:omni, :tps])

using CairoMakie
using AlgebraOfGraphics

hist(collect(skipmissing(tps.ageMonths)))

##

labels = [
    "White",
    "Black",
    "Asian",
    "Mixed",
    "Other",
    "Unknown"
]

counts = [
    "White"   =>   1142,
    "Black"   =>   290,
    "Asian"   =>   1,
    "Mixed"   =>   343,
    "Other"   =>   30,
    "Unknown" =>   594
]

colors = [
    :navajowhite3,
    :coral,
    :darkorange3,
    :darkseagreen3,
    :cyan4,
    :dimgray
]

fig, ax, p = pie(counts, color = colors, axis=(;aspect=1))
hidedecorations!(ax)
hidespines!(ax)

Legend(fig[2,1], [MarkerElement(color=c, marker=:circle) for c in colors],
                labels, "Race",
        orientation=:horizontal,
        tellheight=true, tellwidth=false, nbanks=3
)

# save("figures/demo_race.png", fig)
fig
##

ed = skipmissing(unique(tps, [:subject]).ed) |> collect

fig, ax, p = pie(map(levels(ed)) do lev
        count(==(lev), ed)
    end;
    color = Makie.resample_cmap(:viridis, 6), axis=(;aspect=1)
)
hidedecorations!(ax)
hidespines!(ax)

Legend(fig[2,1],
    [MarkerElement(color=c, marker=:circle) for c in Makie.resample_cmap(:viridis, 6)],
    levels(ed),
    "Maternal Education";
    orientation=:horizontal,
    tellheight=true, tellwidth=false, nbanks=3
)

save("figures/demo_education.png", fig)
fig

##

rce = skipmissing(unique(tps, [:subject]).rce) |> collect

fig, ax, p = pie(map(levels(rce)) do lev
        count(==(lev), skipmissing(rce))
    end;
    color = [:navajowhite3,:coral,:darkorange3,:darkseagreen3,:cyan4,:dimgray],
    axis  =(; aspect=1)
)

hidedecorations!(ax)
hidespines!(ax)

Legend(fig[2,1],
    [MarkerElement(color=c, marker=:circle) for c in  [:navajowhite3,:coral,:darkorange3,:darkseagreen3,:cyan4,:dimgray]],
    levels(rce),
    "Race";
    orientation=:horizontal,
    tellheight=true, tellwidth=false, nbanks=3
)

save("figures/demo_race.png", fig)
fig

##

unique!(omni, :sample)

dyads = @chain omni begin
    groupby(:subject)
    combine(:Mother_Child => (mc-> length(unique(mc)) == 2) => :hasdyad)
    @rsubset!(:hasdyad)
end

toplot = @chain omni begin
    @rsubset(:subject in dyads.subject)
    groupby(:subject)
    combine(:Mother_Child=> (mc-> count(==("C"), mc))=> :nchild, :Mother_Child=> (mc -> count(==("M"), mc)) => :nmother)
end

fig, ax1, h1 = hist(toplot.nchild, axis=(; xlabel="number of child samples", ylabel="number of dyads"))
ax2 = Axis(fig[1,2], xlabel="number of mother samples")
h2 = hist!(ax2, toplot.nmother)

linkyaxes!(ax1, ax2)

save("figures/dyads.png", fig)
fig

max_stars_repo_name: Klepac-Ceraj-Lab/ResonanceAnalysis

max_issues_repo_name: Klepac-Ceraj-Lab/ResonanceAnalysis

max_forks_repo_name: Klepac-Ceraj-Lab/ResonanceAnalysis


text:
struct ODEQuadratureAdjointSensitivityFunction{C<:AdjointDiffCache,Alg<:QuadratureAdjoint,
                                               uType,SType,fType<:DiffEqBase.AbstractDiffEqFunction} <: SensitivityFunction
  diffcache::C
  sensealg::Alg
  discrete::Bool
  y::uType
  sol::SType
  f::fType
end

function ODEQuadratureAdjointSensitivityFunction(g,sensealg,discrete,sol,dg)
  diffcache, y = adjointdiffcache(g,sensealg,discrete,sol,dg,sol.prob.f;quad=true)
  return ODEQuadratureAdjointSensitivityFunction(diffcache,sensealg,discrete,
                                                 y,sol,sol.prob.f)
end

# u = λ'
function (S::ODEQuadratureAdjointSensitivityFunction)(du,u,p,t)
  @unpack sol, discrete = S
  f = sol.prob.f

  λ,grad,y,dλ,dgrad,dy = split_states(du,u,t,S)

  vecjacobian!(dλ, y, λ, p, t, S)
  dλ .*= -one(eltype(λ))

  discrete || accumulate_cost!(dλ, y, p, t, S)
  return nothing
end

function split_states(du,u,t,S::ODEQuadratureAdjointSensitivityFunction;update=true)
  @unpack y, sol = S

  if update
    if typeof(t) <: ForwardDiff.Dual && eltype(y) <: AbstractFloat
      y = sol(t, continuity=:right)
    else
      sol(y,t, continuity=:right)
    end
  end

  λ  = u
  dλ = du

  λ,nothing,y,dλ,nothing,nothing
end

# g is either g(t,u,p) or discrete g(t,u,i)
@noinline function ODEAdjointProblem(sol,sensealg::QuadratureAdjoint,g,
                                     t=nothing,dg=nothing,
                                     callback=CallbackSet())

  @unpack f, p, u0, tspan = sol.prob
  terminated = false
  if hasfield(typeof(sol),:retcode)
    if sol.retcode == :Terminated
      tspan = (tspan[1], sol.t[end])
      terminated = true
    end
  end
  tspan = reverse(tspan)

  discrete = t !== nothing

  len = length(u0)
  λ = similar(u0, len)
  λ .= false
  sense = ODEQuadratureAdjointSensitivityFunction(g,sensealg,discrete,sol,dg)

  init_cb = t !== nothing && tspan[1] == t[end]
  z0 = vec(zero(λ))
  cb, duplicate_iterator_times = generate_callbacks(sense, g, λ, t, callback, init_cb, terminated)

  jac_prototype = sol.prob.f.jac_prototype
  adjoint_jac_prototype = !sense.discrete || jac_prototype === nothing ? nothing : copy(jac_prototype')

  original_mm = sol.prob.f.mass_matrix
  if original_mm === I || original_mm === (I,I)
    odefun = ODEFunction(sense, jac_prototype=adjoint_jac_prototype)
  else
    odefun = ODEFunction(sense, mass_matrix=sol.prob.f.mass_matrix', jac_prototype=adjoint_jac_prototype)
  end
  return ODEProblem(odefun,z0,tspan,p,callback=cb)
end

struct AdjointSensitivityIntegrand{pType,uType,lType,rateType,S,AS,PF,PJC,PJT,DGP,G}
  sol::S
  adj_sol::AS
  p::pType
  y::uType
  λ::lType
  pf::PF
  f_cache::rateType
  pJ::PJT
  paramjac_config::PJC
  sensealg::QuadratureAdjoint
  dgdp_cache::DGP
  dgdp::G
end

function AdjointSensitivityIntegrand(sol,adj_sol,sensealg,dgdp=nothing)
  prob = sol.prob
  @unpack f, p, tspan, u0 = prob
  numparams = length(p)
  y = zero(sol.prob.u0)
  λ = zero(adj_sol.prob.u0)
  # we need to alias `y`
  f_cache = zero(y)
  f_cache .= false
  isautojacvec = get_jacvec(sensealg)

  dgdp_cache = dgdp === nothing ? nothing : zero(p)

  if DiffEqBase.has_paramjac(f) || sensealg.autojacvec isa ReverseDiffVJP || (sensealg.autojacvec isa Bool && sensealg.autojacvec && DiffEqBase.isinplace(prob))
    tape = if DiffEqBase.isinplace(prob)
      ReverseDiff.GradientTape((y, prob.p, [tspan[2]])) do u,p,t
        du1 = similar(p, size(u))
        du1 .= false
        f(du1,u,p,first(t))
        return vec(du1)
      end
    else
      ReverseDiff.GradientTape((y, prob.p, [tspan[2]])) do u,p,t
        vec(f(u,p,first(t)))
      end
    end
    if compile_tape(sensealg)
      paramjac_config = ReverseDiff.compile(tape)
    elseif sensealg.autojacvec isa Bool && sensealg.autojacvec
      compile = try
          if DiffEqBase.isinplace(prob)
            !hasbranching(prob.f,copy(u0),u0,p,prob.tspan[1])
          else
            !hasbranching(prob.f,u0,p,prob.tspan[1])
          end
      catch
          false
      end
      if compile
          paramjac_config = ReverseDiff.compile(tape)
      else
          paramjac_config = tape
      end
    else
      paramjac_config = tape
    end
    pf = nothing
    pJ = nothing
  elseif sensealg.autojacvec isa EnzymeVJP
      paramjac_config = zero(y),zero(y)
      pf = let f = f.f
          if DiffEqBase.isinplace(prob) && prob isa RODEProblem
              function (out,u,_p,t,W)
                  f(out, u, _p, t, W)
                  nothing
              end
          elseif DiffEqBase.isinplace(prob)
              function (out,u,_p,t)
                  f(out, u, _p, t)
                  nothing
              end
          elseif !DiffEqBase.isinplace(prob) && prob isa RODEProblem
              function (out,u,_p,t,W)
                  out .= f(u, _p, t, W)
                  nothing
              end
          else !DiffEqBase.isinplace(prob)
              function (out,u,_p,t)
                  out .= f(u, _p, t)
                  nothing
              end
          end
      end
      pJ = nothing
  elseif isautojacvec # Zygote
    paramjac_config = nothing
    pf = nothing
    pJ = nothing
  else
    pf = DiffEqBase.ParamJacobianWrapper(f,tspan[1],y)
    pJ = similar(u0,length(u0),numparams)
    paramjac_config = build_param_jac_config(sensealg,pf,y,p)
  end
  AdjointSensitivityIntegrand(sol,adj_sol,p,y,λ,pf,f_cache,pJ,paramjac_config,sensealg,dgdp_cache,dgdp)
end

function (S::AdjointSensitivityIntegrand)(out,t)
  @unpack y, λ, pJ, pf, p, f_cache, dgdp_cache, paramjac_config, sensealg, sol, adj_sol = S
  f = sol.prob.f
  sol(y,t)
  adj_sol(λ,t)
  λ .*= -one(eltype(λ))
  isautojacvec = get_jacvec(sensealg)
  # y is aliased

  if !isautojacvec
    if DiffEqBase.has_paramjac(f)
      f.paramjac(pJ,y,p,t) # Calculate the parameter Jacobian into pJ
    else
      pf.t = t
      jacobian!(pJ, pf, p, f_cache, sensealg, paramjac_config)
    end
    mul!(out',λ',pJ)
  elseif (sensealg.autojacvec isa Bool && DiffEqBase.isinplace(sol.prob)) || sensealg.autojacvec isa ReverseDiffVJP
    tape = paramjac_config
    tu, tp, tt = ReverseDiff.input_hook(tape)
    output = ReverseDiff.output_hook(tape)
    ReverseDiff.unseed!(tu) # clear any "leftover" derivatives from previous calls
    ReverseDiff.unseed!(tp)
    ReverseDiff.unseed!(tt)
    ReverseDiff.value!(tu, y)
    ReverseDiff.value!(tp, p)
    ReverseDiff.value!(tt, [t])
    ReverseDiff.forward_pass!(tape)
    ReverseDiff.increment_deriv!(output, λ)
    ReverseDiff.reverse_pass!(tape)
    copyto!(vec(out), ReverseDiff.deriv(tp))
  elseif (sensealg.autojacvec isa Bool && sensealg.autojacvec) || sensealg.autojacvec isa ZygoteVJP
    _dy, back = Zygote.pullback(p) do p
      vec(f(y, p, t))
    end
    tmp = back(λ)
    out[:] .= vec(tmp[1])
  elseif sensealg.autojacvec isa EnzymeVJP
      tmp3,tmp4 = paramjac_config
      tmp4 .= λ
      out .= 0
      Enzyme.autodiff(pf,Enzyme.Duplicated(tmp3,tmp4),
                      y,Enzyme.Duplicated(p, out),t)
  end

  # TODO: Add tracker?

  if S.dgdp !== nothing
    S.dgdp(dgdp_cache, y, p, t)
    out .+= dgdp_cache
  end
  out'
end

function (S::AdjointSensitivityIntegrand)(t)
  out = similar(S.p)
  S(out,t)
end

function _adjoint_sensitivities(sol,sensealg::QuadratureAdjoint,alg,g,
                                t=nothing,dg=nothing;
                                abstol=sensealg.abstol,reltol=sensealg.reltol,
                                callback = nothing,
                                kwargs...)
  dgdu, dgdp = dg isa Tuple ? dg : (dg, nothing)
  adj_prob = ODEAdjointProblem(sol,sensealg,g,t,dgdu,callback)
  adj_sol = solve(adj_prob,alg;abstol=abstol,reltol=reltol,
                               save_everystep=true,save_start=true,kwargs...)

  p = sol.prob.p
  if p === nothing || p === DiffEqBase.NullParameters()
    return -adj_sol[end],nothing
  else
    integrand = AdjointSensitivityIntegrand(sol,adj_sol,sensealg,dgdp)

    if t === nothing
      res,err = quadgk(integrand,sol.prob.tspan[1],sol.prob.tspan[2],
                       atol=abstol,rtol=reltol)
    else
      res = zero(integrand.p)'

      if callback!==nothing
        cur_time = length(t)
        dλ = similar(integrand.λ)
        dλ .*= false
        dgrad = similar(res)
        dgrad .*= false
      end

      for i in length(t)-1:-1:1
        res .+= quadgk(integrand,t[i],t[i+1],
                       atol=abstol,rtol=reltol)[1]
        if t[i]==t[i+1]
          for cb in callback.discrete_callbacks
            if t[i] ∈ cb.affect!.event_times
              integrand = update_integrand_and_dgrad(res,sensealg,cb,integrand,adj_prob,sol,dgdu,g,dλ,dgrad,t[i],cur_time)
            end
          end
          for cb in callback.continuous_callbacks
            if t[i] ∈ cb.affect!.event_times || t[i] ∈ cb.affect_neg!.event_times
              integrand = update_integrand_and_dgrad(res,sensealg,cb,integrand,adj_prob,sol,dgdu,g,dλ,dgrad,t[i],cur_time)
            end
          end
        end
        callback!==nothing && (cur_time -= one(cur_time))
      end
      if t[1] != sol.prob.tspan[1]
        res .+= quadgk(integrand,sol.prob.tspan[1],t[1],
                       atol=abstol,rtol=reltol)[1]
      end
    end
    return -adj_sol[end], res
  end
end

function update_p_integrand(integrand::AdjointSensitivityIntegrand,p)
  @unpack sol, adj_sol, y, λ, pf, f_cache, pJ, paramjac_config, sensealg, dgdp_cache, dgdp = integrand
  AdjointSensitivityIntegrand(sol,adj_sol,p,y,λ,pf,f_cache,pJ,paramjac_config,sensealg,dgdp_cache,dgdp)
end

function update_integrand_and_dgrad(res,sensealg::QuadratureAdjoint,cb,integrand,adj_prob,sol,dgdu,g,dλ,dgrad,t,cur_time)

  function wp(dp,p,u,t)
    fakeinteg = FakeIntegrator([x for x in u],[x for x in p],t)
    _affect! = get_affect!(cb.affect!.affect!)
    _affect!(fakeinteg)
    dp .= fakeinteg.p
  end

  _p = similar(integrand.p, size(integrand.p))
  wp(_p,integrand.p,integrand.y,t)

  if _p != integrand.p
    fakeSp = CallbackSensitivityFunction(wp,sensealg,adj_prob.f.f.diffcache,sol.prob)
    #vjp with Jacobin given by dw/dp before event and vector given by grad
    vecjacobian!(res, integrand.p, res, integrand.y, t, fakeSp;
                        dgrad=nothing, dy=nothing)
    integrand = update_p_integrand(integrand,_p)
  end

  function w(du,u,p,t)
    fakeinteg = FakeIntegrator([x for x in u],[x for x in p],t)
    _affect! = get_affect!(cb.affect!.affect!)
    _affect!(fakeinteg)
    du .= fakeinteg.u
  end

  # Create a fake sensitivity function to do the vjps needs to be done
  # to account for parameter dependence of affect function
  fakeS = CallbackSensitivityFunction(w,sensealg,adj_prob.f.f.diffcache,sol.prob)
  if dgdu === nothing
    g(dλ,integrand.y,integrand.p,t,cur_time)
  else
    dgdu(dλ,integrand.y,integrand.p,t,cur_time)
  end

  # account for implicit events

  dλ .= dλ-integrand.λ
  vecjacobian!(dλ, integrand.y, dλ, integrand.p, t, fakeS; dgrad=dgrad)
  res .-= dgrad
  return integrand
end


text:
using Revise
using ADCME
using ADCMEKit
using NNFEM
using PyPlot
using ProgressMeter

close("all")
include("load_domain_function.jl")

NT = 100
Δt = 30/NT
domain = load_crack_domain()

# visualize_mesh(domain)
H = domain.elements[1].mat[1].H
EBC_func = nothing
FBC_func = nothing
body_func = nothing
globaldata = GlobalData(missing, missing, missing, missing, domain.neqs, EBC_func, FBC_func, body_func)
# globaldata, domain = LinearStaticSolver(globaldata, domain)
Fext = compute_external_force(globaldata, domain)
d = LinearStaticSolver(globaldata, domain, domain.state, H, Fext)

sess = Session(); init(sess)
domain.state = run(sess, d)

figure()
subplot(211)
visualize_scalar_on_scoped_body(domain.state[1:domain.nnodes], zeros(size(domain.state)...), domain)
subplot(212)
visualize_scalar_on_scoped_body(domain.state[domain.nnodes+1:end], zeros(size(domain.state)...), domain)

strain = getStrain(domain)
stress = getStress(domain)
state = domain.state

# error()

domain = load_crack_domain(option="mixed")
# domain = load_crack_domain()

d0 = state[:]
v0 = zeros(2domain.nnodes)
a0  = zeros(2domain.nnodes)
# σ0 = zeros(getNGauss(domain),3)
# ϵ0 = zeros(getNGauss(domain),3)
σ0 = stress
ϵ0 = strain
domain.state = state[:]
p = 1
for i = 1:domain.neles
  for k = 1:length(domain.elements[i].weights)
    domain.elements[i].mat[k].σ0 = stress[p,:]
    domain.elements[i].mat[k].ε0 = strain[p,:]
    # @info stress[p,:] - H * strain[p,:]
    global p += 1
  end
end

# d, v, a, σ, ϵ = ViscoelasticitySolver(
#   globaldata, domain, d0, v0, a0, σ0, ϵ0, Δt, NT, μ, λ, η, Fext, ubd, abd
# )

Dstate = state[:]
globaldata = GlobalData(d0[domain.dof_to_eq], Dstate[domain.dof_to_eq], v0[domain.dof_to_eq], a0[domain.dof_to_eq], domain.neqs, EBC_func, FBC_func, nothing)
assembleMassMatrix!(globaldata, domain)
@showprogress for i = 1:NT
    global globaldata, domain = GeneralizedAlphaSolverStep(globaldata, domain, Δt, maxiterstep=1)
end

# d_, σ_ = run(sess, [d, σ])
# domain.history["state"] = []
# domain.history["stress"] = []
# for i = 1:NT+1
#   push!(domain.history["state"], d_[i,:])
#   push!(domain.history["stress"], σ_[i,:,:])
# end

d_ = hcat(domain.history["state"]...)'|>Array
σ_ = vcat([reshape(x, 1, size(x,1), size(x,2)) for x in domain.history["stress"]]...)
# figure()
# p = visualize_scalar_on_scoped_body(σ_[:, 1:domain.nnodes, 2], d_, domain, scale_factor=1.0)
# # p = visualize_von_mises_stress_on_scoped_body(d_, domain, scale_factor=10.0)
# # p = visualize_total_deformation_on_scoped_body(d_, domain, scale_factor=10.0)
# saveanim(p, "ad_solver.gif")
# close("all")

###################### check surface ########################
idx = []
k = 0
pts = getGaussPoints(domain)
xcoord = Float64[]
for e in domain.elements
    global k
    if sum(map(x->x≈0.0, e.coords[:,2]))≠2
        k += length(e.weights)
        continue
    end
    l = argmin(pts[k+1:k+length(e.weights), 2])
    push!(idx, k + l)
    push!(xcoord, pts[k+l, 1])
    k += length(e.weights)
end
ii = sortperm(xcoord)
xcoord = xcoord[ii]
idx = idx[ii]

σ_surface = σ_[:, idx, :]
s = zeros(size(σ_surface, 1), size(σ_surface, 2))
for i = 1:size(σ_surface, 1)
    for j = 1:size(σ_surface, 2)
        s[i,j] = NNFEM.postprocess_stress(σ_surface[i,j,:], "vonMises")
    end
end

figure(figsize=(8,3))
semilogy(s[1, :], "-", label="x0")
semilogy(s[NT, :], "--", label="x")
legend()
title("Stress")
savefig("Stress_test.png")

figure(figsize=(8,3))
semilogy(s[:, 13], "-")
title("Stress")

# visualize_total_deformation_on_scoped_body(Array(d_), domain; scale_factor=100)

slip_idx = findall(domain.nodes[:,2] .≈ 0.0)
coords = domain.nodes[slip_idx,1]
ii = sortperm(coords)
coords = coords[ii]
x_ = d_[:, slip_idx[ii]]
y_ = d_[:, slip_idx[ii] .+ domain.nnodes]
# plot(coords, x_)

x0 = d0[slip_idx[ii]]
y0 = d0[slip_idx[ii] .+ domain.nnodes]

figure(figsize=(8,3))
subplot(211)
plot((x_' .- x0)[:,1:10:end])
subplot(212)
plot(x0, label="x")
# plot(x_[1, :] .- x0, "-", label="x0")
# plot(x_[NT+1, :] .- x0, "--", label="x")
legend()

figure(figsize=(8,3))
subplot(211)
S = 1:20:NT
for i = 1:length(S)
  plot(-(y_' .- y0)[:,S[i]],"C$i",  label="$i")
end
legend()

subplot(212)
S = 1:20:NT
for i = 1:length(S)
  plot(-y_'[:,S[i]],"C$i",  label="$i")
end
plot(-y0, label="y")
# plot(y_[1, :] .- y0, "-", label="x0")
# plot(y_[NT+1, :] .- y0, "--", label="x")
# plot(y0, ":", label="static")
legend()
savefig("disp.png")

figure()
plot(y_[:, 10])
# plot(y_[1, :] .- y0, "-", label="x0")
# plot(y_[NT+1, :] .- y0, "--", label="x")
# plot(y0, ":", label="static")
legend()

max_stars_repo_name: ADCMEMarket/ADCMEImages

max_issues_repo_name: ADCMEMarket/ADCMEImages

max_forks_repo_name: ADCMEMarket/ADCMEImages


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

# %% [markdown]
# Subset of https://gist.github.com/genkuroki/8c7e9e484a5877d9747a4819bbab645d

# %%
function conseclist_forloop(x)
    b, e = firstindex(x), lastindex(x)
    s = b - 1
    S = [s]
    while s < e
        t = s + 1
        @inbounds while t < e
            x[t+1] != x[t] + 1 && break
            t += 1
        end
        push!(S, t)
        s = t
    end
    [@view(x[S[i]+1:S[i+1]]) for i in 1:length(S)-1]
end

# %%
function conseclist_bitvector(x)
    b, e = firstindex(x), lastindex(x)
    @views idxs = x[b+1:e] .- x[b:e-1] .!= 1
    S = [b-1; (b:e-1)[idxs]; e]
    [@view(x[S[i]+1:S[i+1]]) for i in 1:length(S)-1]
end

# %%
module LazyConsecLists

struct ConsecList{T<:AbstractVector{<:Integer}} <:
        AbstractVector{SubArray{eltype(T), 1, T, Tuple{UnitRange{Int64}}, true}}
    a::T
    S::Vector{Int}
end
function ConsecList(x)
    b, e = firstindex(x), lastindex(x)
    idx = @views (x[b+1:e] .- x[b:e-1]) .!= 1
    S = [b-1; (b:e-1)[idx]; e]
    ConsecList{typeof(x)}(x, S)
end

Base.length(x::ConsecList) = length(x.S) - 1
Base.size(x::ConsecList) = (length(x),)
function Base.eltype(x::ConsecList{T}) where T<:AbstractVector{<:Integer}
    SubArray{eltype(T), 1, T, Tuple{UnitRange{Int64}}, true}
end

Base.getindex(x::ConsecList, i::Integer) = @view(x.a[x.S[i]+1:x.S[i+1]])
Base.getindex(x::ConsecList, r::AbstractRange) = [x[i] for i in r]
Base.getindex(x::ConsecList, ::Colon) = collect(x)

end

# %%
using BenchmarkTools
A = unique(sort(rand(1:10^4, 10^4)))

# %%
@show conseclist_forloop(A) == conseclist_bitvector(A) == LazyConsecLists.ConsecList(A)
@show length(conseclist_forloop(A)) == length(conseclist_bitvector(A)) == length(LazyConsecLists.ConsecList(A))
@btime conseclist_forloop($A)
@btime conseclist_bitvector($A)
@btime collect(LazyConsecLists.ConsecList($A))

# %%
@btime sum(length, conseclist_forloop($A))
@btime sum(length, conseclist_bitvector($A))
@btime sum(length, LazyConsecLists.ConsecList($A))

# %%
using OffsetArrays
x = OffsetArray([1, 2, 5, 9, 10, 11, 14, 15, 16, 19], -4:5)

# %%
conseclist_forloop(x)

# %%
conseclist_bitvector(x)

# %%
LazyConsecLists.ConsecList(x)

# %%

max_stars_repo_name: genkuroki/public

max_issues_repo_name: genkuroki/public

max_forks_repo_name: genkuroki/public


text:
"""
Copyright Marcus Greiff 2020

Tests of the attitude conversions and distances pertaining to SU(2)
"""

# Include modules and recompile if running test standalone
if !isdefined(Main, :runAllTests)
    using Test, AerialVehicleControl, LinearAlgebra
    recompile();
end

@verbose(1, "Testing the SU(2) maps...")
@testset "Attitude maps on SU(2)" begin
    @testset "Test the conjugate operation" begin
        @testset "Good inputs" begin
            # Fixture
            q, qc, qm = get_matrix_and_copy(rand(4,1))
            p, pc, pm = get_matrix_and_copy(rand(4,1))
            # Test
            status = ccall((:cont_SU2_conjugate, AerialVehicleControl.CONT_LIB_PATH),
                Cint,
                (Ref{matrix_double_t}, Ref{matrix_double_t}),
                qm,
                pm)
            # Assert
            @test isapprox(status,1);
            @test isapprox(q, qc);
            @test !isapprox(p, pc);
            @test isapprox(q[1], p[1]);
            @test isapprox(q[2:4], -p[2:4]);
        end
        @testset "Bad inputs" begin
            # Fixture
            q, qc, qm = get_matrix_and_copy(rand(3,1))
            p, pc, pm = get_matrix_and_copy(rand(4,1))
            # Test
            status = ccall((:cont_SU2_conjugate, AerialVehicleControl.CONT_LIB_PATH),
                Cint,
                (Ref{matrix_double_t}, Ref{matrix_double_t}),
                qm,
                pm)
            # Assert
            @test isapprox(status, 0);
        end
    end
    @testset "Test the product operation" begin
        @testset "Good inputs" begin
            # Fixture
            q, qc, qm = get_matrix_and_copy(rand(4,1))
            p, pc, pm = get_matrix_and_copy(rand(4,1))
            o, oc, om = get_matrix_and_copy(rand(4,1))
            qpprod    = quaternion_left_product(q)*p;

            # Test
            status = ccall((:cont_SU2_product, AerialVehicleControl.CONT_LIB_PATH),
                Cint,
                (Ref{matrix_double_t}, Ref{matrix_double_t}, Ref{matrix_double_t}),
                qm,
                pm,
                om)

            # Assert
            @test  isapprox(status,1);
            @test  isapprox(q, qc);
            @test  isapprox(p, pc);
            @test !isapprox(o, oc);
            @test  isapprox(o, qpprod);
        end
        @testset "Bad inputs" begin
            # Fixture
            q, qc, qm = get_matrix_and_copy(rand(4,1))
            p, pc, pm = get_matrix_and_copy(rand(4,1))
            o, oc, om = get_matrix_and_copy(rand(3,1))
            qpprod    = quaternion_left_product(q)*p;

            # Test
            status = ccall((:cont_SU2_product, AerialVehicleControl.CONT_LIB_PATH),
                Cint,
                (Ref{matrix_double_t}, Ref{matrix_double_t}, Ref{matrix_double_t}),
                qm,
                pm,
                om)
            # Assert
            @test isapprox(status, 0);
        end
    end
    @testset "Test the vee operation" begin
        @testset "Good inputs" begin
            # Fixture
            q, qc, qm = get_matrix_and_copy(rand(4,1))
            o, oc, om = get_matrix_and_copy(rand(3,1))
            X         = quat_2_SU2(q);
            output    = SU2_vee(X);

            # Test
            status = ccall((:cont_SU2_vee, AerialVehicleControl.CONT_LIB_PATH),
                Cint,
                (Ref{matrix_double_t}, Ref{matrix_double_t}),
                qm,
                om)

            # Assert
            @test  isapprox(status,1);
            @test  isapprox(q, qc);
            @test !isapprox(o, oc);
            @test  isapprox(o, output);
        end
        # Todo write tests for bad inputs
    end
    @testset "Test the hat operation" begin
        @testset "Good inputs" begin
            # Fixture
            u, uc, um = get_matrix_and_copy(rand(3,1))
            o, oc, om = get_matrix_and_copy(rand(4,1))
            output    = SU2_hat(u);

            # Test
            status = ccall((:cont_SU2_hat, AerialVehicleControl.CONT_LIB_PATH),
                Cint,
                (Ref{matrix_double_t}, Ref{matrix_double_t}),
                um,
                om)

            # Assert
            @test  isapprox(status,1);
            @test  isapprox(u, uc);
            @test !isapprox(o, oc);
            @test  isapprox(quat_2_SU2(o), output);
        end
        # Todo write tests for bad inputs
    end
    @testset "Test the Exponential map on SU(2)" begin
        @testset "Good inputs" begin
            # Fixture
            u, uc, um = get_matrix_and_copy(rand(3,1))
            q, qc, qm = get_matrix_and_copy(rand(4,1))

            R_J = exp(SO3_hat(u))

            # Test
            status = ccall((:cont_SU2_Exp, AerialVehicleControl.CONT_LIB_PATH),
                Cint,
                (Ref{matrix_double_t}, Ref{matrix_double_t}),
                um,
                qm)

            R_C = quat_2_SO3(q);

            # Assert
            @test  isapprox(status, 1);
            @test  isapprox(u, uc);
            @test !isapprox(q, qc);
            @test  isapprox(R_J, R_C);
        end
        # Todo write tests for bad inputs
    end
    @testset "Test the triple product function on SU(2)" begin
        @testset "Good inputs" begin
            # Fixture
            a = 2*pi*rand()-pi;
            b = 2*pi*rand()-pi;
            c = 2*pi*rand()-pi;

            q, qc, qm = get_matrix_and_copy(rand(4,1))

            R_A = exp(SO3_hat([a;0;0]))
            R_B = exp(SO3_hat([0;b;0]))
            R_C = exp(SO3_hat([0;0;c]))
            R_J = R_A * R_B * R_C;

            # Test
            status = ccall((:cont_SU2_triple, AerialVehicleControl.CONT_LIB_PATH),
                Cint,
                (Cdouble, Cdouble, Cdouble, Ref{matrix_double_t}),
                a,
                b,
                c,
                qm)

            R_C = quat_2_SO3(q);

            # Assert
            @test  isapprox(status, 1);
            @test !isapprox(q, qc);
            @test  isapprox(R_J, R_C);
        end
        # Todo write tests for bad inputs
    end
    @testset "Test the distance on SU(2)" begin
        @testset "Good inputs" begin
            # Fixture
            q = rand(4,1); q = q/norm(q);
            p = rand(4,1); p = p/norm(p);
            Xq = quat_2_SU2(q);
            Xp = quat_2_SU2(p);
            Gamma = real(tr(I - Xq'*Xp) / 2.0)

            q, qc, qm = get_matrix_and_copy(q)
            p, pc, pm = get_matrix_and_copy(p)

            # Test
            distance = ccall((:cont_SU2_distance, AerialVehicleControl.CONT_LIB_PATH),
                Cdouble,
                (Ref{matrix_double_t}, Ref{matrix_double_t}),
                qm,
                pm)

            # Assert
            @test isapprox(q, qc);
            @test isapprox(p, pc);
            @test isapprox(distance, Gamma);
        end
        # Todo write tests for bad inputs
    end
end


text:
const calllike = (:call, :foreigncall)

const compiled_calls = Dict{Any,Any}()

function extract_inner_call!(stmt, idx, once::Bool=false)
    isa(stmt, Expr) || return nothing
    (stmt.head == :toplevel || stmt.head == :thunk) && return nothing
    once |= stmt.head ∈ calllike
    for (i, a) in enumerate(stmt.args)
        isa(a, Expr) || continue
        # Make sure we don't "damage" special syntax that requires literals
        if i == 1 && stmt.head == :foreigncall
            continue
        end
        if i == 2 && stmt.head == :call && stmt.args[1] == :cglobal
            continue
        end
        ret = extract_inner_call!(a, idx, once) # doing this first extracts innermost calls
        ret !== nothing && return ret
        iscalllike = a.head ∈ calllike
        if once && iscalllike
            stmt.args[i] = NewSSAValue(idx)
            return a
        end
    end
    return nothing
end

function replace_ssa!(stmt, ssalookup)
    isa(stmt, Expr) || return nothing
    for (i, a) in enumerate(stmt.args)
        if isa(a, SSAValue)
            stmt.args[i] = SSAValue(ssalookup[a.id])
        elseif isa(a, NewSSAValue)
            stmt.args[i] = SSAValue(a.id)
        else
            replace_ssa!(a, ssalookup)
        end
    end
    return nothing
end

function renumber_ssa!(stmts::Vector{Any}, ssalookup)
    for (i, stmt) in enumerate(stmts)
        if isa(stmt, GotoNode)
            stmts[i] = GotoNode(ssalookup[stmt.label])
        elseif isa(stmt, SSAValue)
            stmts[i] = SSAValue(ssalookup[stmt.id])
        elseif isa(stmt, NewSSAValue)
            stmts[i] = SSAValue(stmt.id)
        elseif isa(stmt, Expr)
            replace_ssa!(stmt, ssalookup)
            if (stmt.head == :gotoifnot || stmt.head == :enter) && isa(stmt.args[end], Int)
                stmt.args[end] = ssalookup[stmt.args[end]]
            end
        end
    end
    return stmts
end

function compute_ssa_mapping_delete_statements!(code::CodeInfo, stmts::Vector{Int})
    stmts = unique!(sort!(stmts))
    ssalookup = collect(1:length(code.codelocs))
    cnt = 1
    for i in 1:length(stmts)
        start = stmts[i] + 1
        stop = i == length(stmts) ? length(code.codelocs) : stmts[i+1]
        ssalookup[start:stop] .-= cnt
        cnt += 1
    end
    return ssalookup
end

# Pre-frame-construction lookup
function lookup_stmt(stmts, arg)
    if isa(arg, SSAValue)
        arg = stmts[arg.id]
    end
    if isa(arg, QuoteNode)
        arg = arg.value
    end
    return arg
end

function smallest_ref(stmts, arg, idmin)
    if isa(arg, SSAValue)
        idmin = min(idmin, arg.id)
        return smallest_ref(stmts, stmts[arg.id], idmin)
    elseif isa(arg, Expr)
        for a in arg.args
            idmin = smallest_ref(stmts, a, idmin)
        end
    end
    return idmin
end

function lookup_global_ref(a::GlobalRef)
    if isdefined(a.mod, a.name) && isconst(a.mod, a.name)
        r = getfield(a.mod, a.name)
        return QuoteNode(r)
    else
        return a
    end
end

function lookup_global_refs!(ex::Expr)
    (ex.head == :isdefined || ex.head == :thunk || ex.head == :toplevel) && return nothing
    for (i, a) in enumerate(ex.args)
        ex.head == :(=) && i == 1 && continue # Don't look up globalrefs on the LHS of an assignment (issue #98)
        if isa(a, GlobalRef)
            ex.args[i] = lookup_global_ref(a)
        elseif isa(a, Expr)
            lookup_global_refs!(a)
        end
    end
    return nothing
end

# See https://github.com/JuliaLang/julia/pull/32800
const foreigncall_version = VERSION < v"1.3.0-alpha.108" ? 0 : 1

"""
    optimize!(code::CodeInfo, mod::Module)

Perform minor optimizations on the lowered AST in `code` to reduce execution time
of the interpreter.
Currently it looks up `GlobalRef`s (for which it needs `mod` to know the scope in
which this will run) and ensures that no statement includes nested `:call` expressions
(splitting them out into multiple SSA-form statements if needed).
"""
function optimize!(code::CodeInfo, scope)
    mod = moduleof(scope)
    evalmod = mod == Core.Compiler ? Core.Compiler : CompiledCalls
    sparams = scope isa Method ? Symbol[sparam_syms(scope)...] : Symbol[]
    code.inferred && error("optimization of inferred code not implemented")
    replace_coretypes!(code)
    # TODO: because of builtins.jl, for CodeInfos like
    #   %1 = Core.apply_type
    #   %2 = (%1)(args...)
    # it would be best to *not* resolve the GlobalRef at %1
    ## Replace GlobalRefs with QuoteNodes
    for (i, stmt) in enumerate(code.code)
        if isa(stmt, GlobalRef)
            code.code[i] = lookup_global_ref(stmt)
        elseif isa(stmt, Expr)
            if stmt.head == :call && stmt.args[1] == :cglobal  # cglobal requires literals
                continue
            else
                lookup_global_refs!(stmt)
            end
        end
    end

    # Replace :llvmcall and :foreigncall with compiled variants. See
    # https://github.com/JuliaDebug/JuliaInterpreter.jl/issues/13#issuecomment-464880123
    foreigncalls_idx = Int[]
    delete_idxs = Int[]
    for (idx, stmt) in enumerate(code.code)
        # Foregincalls can be rhs of assignments
        if isexpr(stmt, :(=))
            stmt = stmt.args[2]
        end
        if isexpr(stmt, :call)
            # Check for :llvmcall
            arg1 = stmt.args[1]
            if (arg1 == :llvmcall || lookup_stmt(code.code, arg1) == Base.llvmcall) && isempty(sparams) && scope isa Method
                nargs = length(stmt.args)-4
                delete_idx = build_compiled_call!(stmt, Base.llvmcall, code, idx, nargs, sparams, evalmod)
                delete_idx === nothing && error("llvmcall must be compiled, but exited early from build_compiled_call!")
                push!(foreigncalls_idx, idx)
                append!(delete_idxs, delete_idx)
            end
        elseif isexpr(stmt, :foreigncall) && scope isa Method
            nargs = foreigncall_version == 0 ? stmt.args[5] : length(stmt.args[3])
            delete_idx = build_compiled_call!(stmt, :ccall, code, idx, nargs, sparams, evalmod)
            if delete_idx !== nothing
                push!(foreigncalls_idx, idx)
                append!(delete_idxs, delete_idx)
            end
        end
    end

    if !isempty(delete_idxs)
        ssalookup = compute_ssa_mapping_delete_statements!(code, delete_idxs)
        foreigncalls_idx = map(x -> ssalookup[x], foreigncalls_idx)
        deleteat!(code.codelocs, delete_idxs)
        deleteat!(code.code, delete_idxs)
        code.ssavaluetypes = length(code.code)
        renumber_ssa!(code.code, ssalookup)
    end

    ## Un-nest :call expressions (so that there will be only one :call per line)
    # This will allow us to re-use args-buffers rather than having to allocate new ones each time.
    old_code, old_codelocs = code.code, code.codelocs
    code.code = new_code = eltype(old_code)[]
    code.codelocs = new_codelocs = Int32[]
    ssainc = fill(1, length(old_code))
    for (i, stmt) in enumerate(old_code)
        loc = old_codelocs[i]
        if !is_type_definition(stmt)  # https://github.com/timholy/Revise.jl/issues/417
            inner = extract_inner_call!(stmt, length(new_code)+1)
            while inner !== nothing
                push!(new_code, inner)
                push!(new_codelocs, loc)
                ssainc[i] += 1
                inner = extract_inner_call!(stmt, length(new_code)+1)
            end
        end
        push!(new_code, stmt)
        push!(new_codelocs, loc)
    end
    # Fix all the SSAValues and GotoNodes
    ssalookup = cumsum(ssainc)
    renumber_ssa!(new_code, ssalookup)
    code.ssavaluetypes = length(new_code)

    # Insert the foreigncall wrappers at the updated idxs
    methodtables = Vector{Union{Compiled,DispatchableMethod}}(undef, length(code.code))
    for idx in foreigncalls_idx
        methodtables[ssalookup[idx]] = Compiled()
    end

    return code, methodtables
end

function parametric_type_to_expr(t::Type)
    t isa Core.TypeofBottom && return t
    t isa UnionAll && (t = t.body)
    if t <: Vararg
        return Expr(:(...), t.parameters[1])
    end
    if t.hasfreetypevars
        params = map(t.parameters) do p
            isa(p, TypeVar) ? p.name :
            isa(p, DataType) && p.hasfreetypevars ? parametric_type_to_expr(p) : p
        end
        return Expr(:curly, scopename(t.name), params...)
    end
    return t
end

# Handle :llvmcall & :foreigncall (issue #28)
function build_compiled_call!(stmt, fcall, code, idx, nargs, sparams, evalmod)
    TVal = evalmod == Core.Compiler ? Core.Compiler.Val : Val
    delete_idx = Int[]
    if fcall == :ccall
        cfunc, RetType, ArgType = lookup_stmt(code.code, stmt.args[1]), stmt.args[2], stmt.args[3]
        # The result of this is useful to have next to you when reading this code:
        # f(x, y) =  ccall(:jl_value_ptr, Ptr{Cvoid}, (Float32,Any), x, y)
        # @code_lowered f(2, 3)
        args = []
        for (atype, arg) in zip(ArgType, stmt.args[6:6+nargs-1])
            if atype === Any
                push!(args, arg)
            else
                @assert arg isa SSAValue
                unsafe_convert_expr = code.code[arg.id]
                push!(delete_idx, arg.id) # delete the unsafe_convert
                cconvert_stmt = unsafe_convert_expr.args[3]
                push!(delete_idx, cconvert_stmt.id) # delete the cconvert
                cconvert_expr = code.code[cconvert_stmt.id]
                push!(args, cconvert_expr.args[3])
            end
        end
    else
        # Run a mini-interpreter to extract the types
        framecode = FrameCode(CompiledCalls, code; optimize=false)
        frame = Frame(framecode, prepare_framedata(framecode, []))
        idxstart = idx
        for i = 2:4
            idxstart = smallest_ref(code.code, stmt.args[i], idxstart)
        end
        frame.pc = idxstart
        if idxstart < idx
            while true
                pc = step_expr!(Compiled(), frame)
                pc == idx && break
                pc === nothing && error("this should never happen")
            end
        end
        cfunc, RetType, ArgType = @lookup(frame, stmt.args[2]), @lookup(frame, stmt.args[3]), @lookup(frame, stmt.args[4])
        args = stmt.args[5:end]
    end
    dynamic_ccall = false
    if isa(cfunc, Expr)   # specification by tuple, e.g., (:clock, "libc")
        cfunc = something(static_eval(cfunc), cfunc)
    end
    if isa(cfunc, Symbol)
        cfunc = QuoteNode(cfunc)
    elseif isa(cfunc, String) || isa(cfunc, Ptr) || isa(cfunc, Tuple)
        # do nothing
    else
        dynamic_ccall = true
        oldcfunc = cfunc
        cfunc = gensym("ptr")
    end
    if isa(RetType, SimpleVector)
        @assert length(RetType) == 1
        RetType = RetType[1]
    end
    # When the ccall is dynamic we pass the pointer as an argument so can reuse the function
    cc_key = (dynamic_ccall ? :ptr : cfunc, RetType, ArgType, evalmod, length(sparams))  # compiled call key
    f = get(compiled_calls, cc_key, nothing)
    argnames = Any[Symbol(:arg, i) for i = 1:nargs]
    if f === nothing
        if fcall == :ccall
            ArgType = Expr(:tuple, [parametric_type_to_expr(t) for t in ArgType]...)
        end
        RetType = parametric_type_to_expr(RetType)
        # #285: test whether we can evaluate an type constraints on parametric expressions
        # this essentially comes down to having the names be available in CompiledCalls,
        # if they are not then executing the method will fail
        try
            isa(RetType, Expr) && Core.eval(CompiledCalls, wrap_params(RetType, sparams))
            isa(ArgType, Expr) && Core.eval(CompiledCalls, wrap_params(ArgType, sparams))
        catch
            return nothing
        end
        wrapargs = copy(argnames)
        if dynamic_ccall
            pushfirst!(wrapargs, cfunc)
        end
        for sparam in sparams
            push!(wrapargs, :(::$TVal{$sparam}))
        end
        methname = gensym("compiledcall")
        calling_convention = stmt.args[foreigncall_version == 0 ? 4 : 5]
        if calling_convention == :(:llvmcall)
            def = :(
                function $methname($(wrapargs...)) where {$(sparams...)}
                    return $fcall($cfunc, llvmcall, $RetType, $ArgType, $(argnames...))
                end)
        elseif calling_convention == :(:stdcall)
            def = :(
                function $methname($(wrapargs...)) where {$(sparams...)}
                    return $fcall($cfunc, stdcall, $RetType, $ArgType, $(argnames...))
                end)
        else
            def = :(
                function $methname($(wrapargs...)) where {$(sparams...)}
                    return $fcall($cfunc, $RetType, $ArgType, $(argnames...))
                end)
        end
        f = Core.eval(evalmod, def)
        compiled_calls[cc_key] = f
    end
    stmt.args[1] = QuoteNode(f)
    stmt.head = :call
    deleteat!(stmt.args, 2:length(stmt.args))
    if dynamic_ccall
        push!(stmt.args, oldcfunc)
    end
    append!(stmt.args, args)
    for i in 1:length(sparams)
        push!(stmt.args, :($TVal($(Expr(:static_parameter, i)))))
    end
    return delete_idx
end

function replace_coretypes!(src; rev::Bool=false)
    if isa(src, CodeInfo)
        replace_coretypes_list!(src.code; rev=rev)
    elseif isa(src, Expr)
        replace_coretypes_list!(src.args; rev=rev)
    end
    return src
end

function replace_coretypes_list!(list; rev::Bool)
    for (i, stmt) in enumerate(list)
        if isa(stmt, rev ? SSAValue : Core.SSAValue)
            list[i] = rev ? Core.SSAValue(stmt.id) : SSAValue(stmt.id)
        elseif isa(stmt, rev ? SlotNumber : Core.SlotNumber)
            list[i] = rev ? Core.SlotNumber(stmt.id) : SlotNumber(stmt.id)
        elseif isa(stmt, Expr)
            replace_coretypes!(stmt; rev=rev)
        end
    end
    return nothing
end


text:
"""
Handles command line arguments for the genie.jl script.
"""
module Commands

import Sockets
import ArgParse
import Genie
using Logging

"""
    execute(config::Settings) :: Nothing

Runs the requested Genie app command, based on the `args` passed to the script.
"""
function execute(config::Genie.Configuration.Settings; server::Union{Sockets.TCPServer,Nothing} = nothing) :: Nothing
  parsed_args = parse_commandline_args(config)::Dict{String,Any}

  # overwrite env settings with command line arguments
  Genie.config.app_env = ENV["GENIE_ENV"]
  Genie.config.server_port = haskey(ENV, "PORT") ? parse(Int, ENV["PORT"]) : parse(Int, parsed_args["p"])
  Genie.config.server_host = parsed_args["l"]

  if called_command(parsed_args, "s")
    Genie.config.run_as_server = true
    Genie.up(Genie.config.server_port, Genie.config.server_host, server = server)

  elseif called_command(parsed_args, "si")
    Genie.up(Genie.config.server_port, Genie.config.server_host, server = server)

  elseif called_command(parsed_args, "r")
    endswith(parsed_args["r"], "Task") || (parsed_args["r"] *= "Task")
    Genie.Toolbox.loadtasks(Main.UserApp)
    taskname = parsed_args["r"]

    @info "Running task $taskname with args $(parsed_args["a"])"

    try
      task = getfield(Main.UserApp, Symbol(taskname))
      if parsed_args["a"] !== nothing
        @info Base.invokelatest(task.runtask, parsed_args["a"])
      else
        @info Base.invokelatest(task.runtask)
      end
    catch ex
      @error ex
    end
  end

  nothing
end

"""
    parse_commandline_args() :: Dict{String,Any}

Extracts the command line args passed into the app and returns them as a `Dict`, possibly setting up defaults.
Also, it is used by the ArgParse module to populate the command line help for the app `-h`.
"""
function parse_commandline_args(config::Genie.Configuration.Settings) :: Dict{String,Any}
  settings = ArgParse.ArgParseSettings()

  settings.description = "Genie web framework CLI"
  settings.epilog = "Visit https://genieframework.com for more info"
  settings.version = string(Genie.Configuration.GENIE_VERSION)
  settings.add_version = true

  ArgParse.@add_arg_table settings begin
    "s"
    help = "starts HTTP server"

    "-p"
    help = "HTTP server port"
    default = "$(config.server_port)"

    "-l"
    help = "Host IP to listen on"
    default = "$(config.server_host)"

    "si"
    help = "starts HTTP server and enters REPL"

    "-r"
    help = "runs Genie.Toolbox task"

    "-a"
    help = "additional arguments passed into the Genie.Toolbox `runtask` function"
    default = nothing
  end

  ArgParse.parse_args(settings)
end

"""
    called_command(args::Dict, key::String) :: Bool

Checks whether or not a certain command was invoked by looking at the command line args.
"""
function called_command(args::Dict{String,Any}, key::String) :: Bool
  haskey(args, key) && (args[key] == "true" || args["s"] == key || args[key] !== nothing)
end

end


text:
using RecyclePermanent
using RNGPool
using JLD2
using Random
import Statistics: mean, var
using Test

@testset "Bernoulli test" begin
  @time include("bernoulli_test.jl")
end

@testset "BP test" begin
  @time include("bp_test.jl")
end

@testset "Kou--McCullagh test" begin
  @time include("km_test.jl")
end


text:
# This file is a part of Julia. License is MIT: https://julialang.org/license

export threadid, nthreads, @threads

"""
    Threads.threadid()

Get the ID number of the current thread of execution. The master thread has ID `1`.
"""
threadid() = Int(ccall(:jl_threadid, Int16, ())+1)

# Inclusive upper bound on threadid()
"""
    Threads.nthreads()

Get the number of threads available to the Julia process. This is the inclusive upper bound
on [`threadid()`](@ref).
"""
nthreads() = Int(unsafe_load(cglobal(:jl_n_threads, Cint)))

function threading_run(func)
    ccall(:jl_enter_threaded_region, Cvoid, ())
    n = nthreads()
    tasks = Vector{Task}(undef, n)
    for i = 1:n
        t = Task(func)
        t.sticky = true
        ccall(:jl_set_task_tid, Cvoid, (Any, Cint), t, i-1)
        tasks[i] = t
        schedule(t)
    end
    try
        for i = 1:n
            wait(tasks[i])
        end
    finally
        ccall(:jl_exit_threaded_region, Cvoid, ())
    end
end

function _threadsfor(iter, lbody, schedule)
    lidx = iter.args[1]         # index
    range = iter.args[2]
    quote
        local threadsfor_fun
        let range = $(esc(range))
        function threadsfor_fun(onethread=false)
            r = range # Load into local variable
            lenr = length(r)
            # divide loop iterations among threads
            if onethread
                tid = 1
                len, rem = lenr, 0
            else
                tid = threadid()
                len, rem = divrem(lenr, nthreads())
            end
            # not enough iterations for all the threads?
            if len == 0
                if tid > rem
                    return
                end
                len, rem = 1, 0
            end
            # compute this thread's iterations
            f = firstindex(r) + ((tid-1) * len)
            l = f + len - 1
            # distribute remaining iterations evenly
            if rem > 0
                if tid <= rem
                    f = f + (tid-1)
                    l = l + tid
                else
                    f = f + rem
                    l = l + rem
                end
            end
            # run this thread's iterations
            for i = f:l
                local $(esc(lidx)) = @inbounds r[i]
                $(esc(lbody))
            end
        end
        end
        if threadid() != 1 || ccall(:jl_in_threaded_region, Cint, ()) != 0
            $(if schedule === :static
              :(error("`@threads :static` can only be used from thread 1 and not nested"))
              else
              # only use threads when called from thread 1, outside @threads
              :(Base.invokelatest(threadsfor_fun, true))
              end)
        else
            threading_run(threadsfor_fun)
        end
        nothing
    end
end

"""
    Threads.@threads [schedule] for ... end

A macro to parallelize a `for` loop to run with multiple threads. Splits the iteration
space among multiple tasks and runs those tasks on threads according to a scheduling
policy.
A barrier is placed at the end of the loop which waits for all tasks to finish
execution.

The `schedule` argument can be used to request a particular scheduling policy.
The only currently supported value is `:static`, which creates one task per thread
and divides the iterations equally among them. Specifying `:static` is an error
if used from inside another `@threads` loop or from a thread other than 1.

The default schedule (used when no `schedule` argument is present) is subject to change.

!!! compat "Julia 1.5"
    The `schedule` argument is available as of Julia 1.5.
"""
macro threads(args...)
    na = length(args)
    if na == 2
        sched, ex = args
        if sched isa QuoteNode
            sched = sched.value
        elseif sched isa Symbol
            # for now only allow quoted symbols
            sched = nothing
        end
        if sched !== :static
            throw(ArgumentError("unsupported schedule argument in @threads"))
        end
    elseif na == 1
        sched = :default
        ex = args[1]
    else
        throw(ArgumentError("wrong number of arguments in @threads"))
    end
    if !(isa(ex, Expr) && ex.head === :for)
        throw(ArgumentError("@threads requires a `for` loop expression"))
    end
    if !(ex.args[1] isa Expr && ex.args[1].head === :(=))
        throw(ArgumentError("nested outer loops are not currently supported by @threads"))
    end
    return _threadsfor(ex.args[1], ex.args[2], sched)
end

"""
    Threads.@spawn expr

Create and run a [`Task`](@ref) on any available thread. To wait for the task to
finish, call [`wait`](@ref) on the result of this macro, or call [`fetch`](@ref)
to wait and then obtain its return value.

Values can be interpolated into `@spawn` via `\$`, which copies the value directly into the
constructed underlying closure. This allows you to insert the _value_ of a variable,
isolating the aysnchronous code from changes to the variable's value in the current task.

!!! note
    See the manual chapter on threading for important caveats.

!!! compat "Julia 1.3"
    This macro is available as of Julia 1.3.

!!! compat "Julia 1.4"
    Interpolating values via `\$` is available as of Julia 1.4.
"""
macro spawn(expr)
    letargs = Base._lift_one_interp!(expr)

    thunk = esc(:(()->($expr)))
    var = esc(Base.sync_varname)
    quote
        let $(letargs...)
            local task = Task($thunk)
            task.sticky = false
            if $(Expr(:islocal, var))
                put!($var, task)
            end
            schedule(task)
            task
        end
    end
end

max_stars_repo_name: sloede/julia

max_issues_repo_name: sloede/julia

max_forks_repo_name: sloede/julia


text:
using CFMG
using Test

include("test1.jl")

# write a default test -
function default_cfmg_test()
    return true
end

# setup the test set for the code generator -
@testset "cfmg_test_set" begin
    @test default_cfmg_test() == true
    # @test CFMG_test1() == true
end

# @testset "example_model" begin
#     @test CFMG_test1() == true
# end

@time begin CFMG_test1() end

max_stars_repo_name: varnerlab/CFMG

max_issues_repo_name: varnerlab/CFMG

max_forks_repo_name: varnerlab/CFMG


text:
# Get traces

import Base.Profile: LineInfo

typealias IP UInt
typealias RawData Vector{IP}
typealias Trace Vector{LineInfo}

const lidict = Dict{IP,LineInfo}()
lookup(ip::IP) = haskey(lidict, ip) ? lidict[ip] : (lidict[ip] = Profile.lookup(ip))
lookup(ips::RawData) = map(lookup, ips)

pruneC(trace::Trace) = filter(line->!line.fromC, trace)

traces(data::Vector{UInt}) =
  @>> split(data, 0, keep=false) map(lookup) map!(pruneC) map!(reverse) filter!(t->!isempty(t))

# Tree Implementation

immutable Node{T}
  data::T
  children::Vector{Node{T}}
end

Node{T}(x::T) = Node(x, Node{T}[])
Node{T}(x::T, children::Node{T}...) = Node(x, [children...])

Base.push!(parent::Node, child::Node) = push!(parent.children, child)
isleaf(node::Node) = isempty(node.children)

# Profile Trees

type ProfileNode
  line::LineInfo
  count::Int
end

ProfileNode(line::LineInfo) = ProfileNode(line, 1)

typealias ProfileTree Node{ProfileNode}

tree(trace::Trace) =
  length(trace) ≤ 1 ?
    Node(ProfileNode(trace[1])) :
    Node(ProfileNode(trace[1]), tree(trace[2:end]))

# Conceptually, a trace is a tree with no branches
# We merge trees by (a) increasing the count of the common nodes
# and (b) adding any new nodes as children.
function Base.merge!(node::ProfileTree, trace::Trace)
  @assert !isempty(trace) && node.data.line == trace[1]
  node.data.count += 1
  length(trace) == 1 && return node
  for child in node.children
    if child.data.line == trace[2]
      merge!(child, trace[2:end])
      return node
    end
  end
  push!(node, tree(trace[2:end]))
  return node
end

function tree(traces::Vector{Trace})
  root = Node(ProfileNode(Profile.UNKNOWN))
  traces = map(trace -> [Profile.UNKNOWN, trace...], traces)
  for trace in traces
    merge!(root, trace)
  end
  return root
end

depth(node::Node) =
  isleaf(node) ? 1 : 1 + maximum(map(depth, node.children))

# Remove redundant lines

childwidths(node::ProfileTree) =
  map(child -> child.data.count/node.data.count, node.children)

function trimroot(tree::ProfileTree)
  validchildren = tree.children[childwidths(tree) .> 0.1]
  length(validchildren) == 1 ? trimroot(validchildren[1]) : tree
end

function sortchildren!(tree::ProfileTree)
  sort!(map!(sortchildren!,tree.children), by = node->node.data.line.line)
  tree
end

# Flatten the tree

function addmerge!(a::Associative, b::Associative)
  for (k, v) in b
    a[k] = haskey(a, k) ? a[k]+b[k] : b[k]
  end
  return a
end

flatlines(tree::ProfileTree; total = tree.data.count) =
  reduce(addmerge!,
         [tree.data.line=>tree.data.count/total],
         map(t->flatlines(t, total = total), tree.children))

function fetch()
  data = Profile.fetch()
  isempty(data) && error("You need to do some profiling first.")
  @> data traces tree trimroot sortchildren!
end