text:
using UnPack, OrderedCollections, ForwardDiff

mutable struct WachterDisasterRisk{T <: Real}
    μₐ::T
    σₐ::T
    ν::T
    δ::T
    ρₚ::T
    pp::T
    ϕₚ::T
    ρ::T
    γ::T
    β::T
end

function WachterDisasterRisk(; μₐ::T = .0252 / 4., σₐ::T = .02 / sqrt(4.), ν::T = .3, δ::T = 0., ρₚ::T = .08^(1. / 4.), pp::T = .0355 / 4.,
                             ϕₚ::T = .0114 / 4. / (.02 / sqrt(4.)) / sqrt(.0355 / 4.), ρ::T = 2.0, γ::T = 3.0,
                             β::T = exp(-.012 / 4.)) where {T <: Real}
    return WachterDisasterRisk{T}(μₐ, σₐ, ν, δ, ρₚ, pp, ϕₚ, ρ, γ, β)
end

function inplace_wachter_disaster_risk(m::WachterDisasterRisk{T}) where {T <: Real}
    @unpack μₐ, σₐ, ν, δ, ρₚ, pp, ϕₚ, ρ, γ, β = m

    @assert ρ != 1. # Forcing ρ to be non-unit for this example

    S  = OrderedDict{Symbol, Int}(:p => 1,  :εc => 2, :εξ => 3) # State variables
    J  = OrderedDict{Symbol, Int}(:vc => 1, :xc => 2, :rf => 3) # Jump variables
    SH = OrderedDict{Symbol, Int}(:εₚ => 1, :εc => 2, :εξ => 3) # Exogenous shocks
    Nz = length(S)
    Ny = length(J)
    Nε = length(SH)

    function μ(F, z, y)
        F_type    = eltype(F)
        F[S[:p]]  = (1 - ρₚ) * pp + ρₚ * z[S[:p]]
        F[S[:εc]] = zero(F_type)
        F[S[:εξ]] = zero(F_type)
    end

    function ξ(F, z, y)
        F[J[:vc]] = log(β) - γ * μₐ + γ * ν * z[S[:p]] - (ρ - γ) * y[J[:xc]] + y[J[:rf]]
        F[J[:xc]] = log(1. - β + β * exp((1. - ρ) * y[J[:xc]])) - (1. - ρ) * y[J[:vc]]
        F[J[:rf]] = (1. - γ) * (μₐ - ν * z[S[:p]] - y[J[:xc]])
    end

    Λ = zeros(T, Nz, Ny)

    # The cache is initialized as zeros so we only need to fill non-zero elements
    function Σ(F, z)
        F_type = eltype(F)
        F[SH[:εₚ], SH[:εₚ]] = ((z[S[:p]] < 0) ? 0. : sqrt(z[S[:p]])) * ϕₚ * σₐ
        F[SH[:εc], SH[:εc]] = one(F_type)
        F[SH[:εξ], SH[:εξ]] = one(F_type)
    end

    function ccgf(F, α, z)
        # F .= .5 .* α[:, 1].^2 + .5 * α[:, 2].^2 + (exp.(α[:, 3] + α[:, 3].^2 .* δ^2 ./ 2.) .- 1. - α[:, 3]) * z[S[:p]]
        # Following lines does the above calculation but in a faster way
        sum!(F, view(α, :, 1:2).^2)
        F .*= .5
        F .+= (exp.(α[:, 3] + view(α, :, 3).^2 .* δ^2 ./ 2.) .- 1. - view(α, :, 3)) * z[S[:p]]
    end

    Γ₅ = zeros(T, Ny, Nz)
    Γ₅[J[:vc], S[:εc]] = (-γ * σₐ)
    Γ₅[J[:vc], S[:εξ]] = (γ * ν)
    Γ₅[J[:rf], S[:εc]] = (1. - γ) * σₐ
    Γ₅[J[:rf], S[:εξ]] = -(1. - γ) * ν

    Γ₆ = zeros(T, Ny, Ny)
    Γ₆[J[:vc], J[:vc]] = (ρ - γ)
    Γ₆[J[:rf], J[:vc]] = (1. - γ)

    z = [pp, 0., 0.]
    xc_sss = log((1. - β) / (exp((1. - ρ) * (ν * pp - μₐ)) - β)) / (1. - ρ)
    vc_sss = xc_sss + ν * pp - μₐ
    y = [vc_sss, xc_sss, -log(β) + γ * (μₐ - ν * pp) - (ρ - γ) * (vc_sss - xc_sss)]
    Ψ = zeros(T, Ny, Nz)
    return RiskAdjustedLinearization(μ, Λ, Σ, ξ, Γ₅, Γ₆, ccgf, z, y, Ψ, Nε)
end

function outofplace_wachter_disaster_risk(m::WachterDisasterRisk{T}) where {T}
    @unpack μₐ, σₐ, ν, δ, ρₚ, pp, ϕₚ, ρ, γ, β = m

    @assert ρ != 1. # Forcing ρ to be non-unit for this example

    S  = OrderedDict{Symbol, Int}(:p => 1,  :εc => 2, :εξ => 3) # State variables
    J  = OrderedDict{Symbol, Int}(:vc => 1, :xc => 2, :rf => 3) # Jump variables
    SH = OrderedDict{Symbol, Int}(:εₚ => 1, :εc => 2, :εξ => 3) # Exogenous shocks
    Nz = length(S)
    Ny = length(J)
    Nε = length(SH)

    function μ(z, y)
        return [(1 - ρₚ) * pp + ρₚ * z[S[:p]], 0., 0.]
    end

    function ξ(z, y)
        F = RiskAdjustedLinearizations.dualvector(y, z)

        F[J[:vc]] = log(β) - γ * μₐ + γ * ν * z[S[:p]] - (ρ - γ) * y[J[:xc]] + y[J[:rf]]
        F[J[:xc]] = log(1. - β + β * exp((1. - ρ) * y[J[:xc]])) - (1. - ρ) * y[J[:vc]]
        F[J[:rf]] = (1. - γ) * (μₐ - ν * z[S[:p]] - y[J[:xc]])
        return F
    end

    Λ = zeros(T, Nz, Nz)

    function Σ(z)
        F = zeros(eltype(z), Nz, Nz)
        F[SH[:εₚ], SH[:εₚ]] = sqrt(z[S[:p]]) * ϕₚ * σₐ
        F[SH[:εc], SH[:εc]] = 1.
        F[SH[:εξ], SH[:εξ]] = 1.
        return F
    end

    ccgf(α, z) = .5 * α[:, 1].^2 + .5 * α[:, 2].^2 + (exp.(α[:, 3] + α[:, 3].^2 * δ.^2 ./ 2.) .- 1. - α[:, 3]) * z[S[:p]]

    Γ₅ = zeros(T, Ny, Nz)
    Γ₅[J[:vc], S[:εc]] = (-γ * σₐ)
    Γ₅[J[:vc], S[:εξ]] = (γ * ν)
    Γ₅[J[:rf], S[:εc]] = (1. - γ) * σₐ
    Γ₅[J[:rf], S[:εξ]] = -(1. - γ) * ν

    Γ₆ = zeros(T, Ny, Ny)
    Γ₆[J[:vc], J[:vc]] = (ρ - γ)
    Γ₆[J[:rf], J[:vc]] = (1. - γ)

    z = [pp, 0., 0.]
    xc_sss = log((1. - β) / (exp((1. - ρ) * (ν * pp - μₐ)) - β)) / (1. - ρ)
    vc_sss = xc_sss + ν * pp - μₐ
    y = [vc_sss, xc_sss, -log(β) + γ * (μₐ - ν * pp) - (ρ - γ) * (vc_sss - xc_sss)]
    Ψ = zeros(T, Ny, Nz)
    return RiskAdjustedLinearization(μ, Λ, Σ, ξ, Γ₅, Γ₆, ccgf, z, y, Ψ, Nε)
end


text:
using LCIO_Cxx
using Test

# test iteration
LCIO_Cxx.open("test.slcio") do reader
for event in reader
    mcparts = getCollection(event, "MCParticle")
    for p in mcparts
        e = getEnergy(p)
        println(e)
        @test e > 0
        momentum = getMomentum(p)
        if getPDG(p) != 22
            # tolerance, since we're dealing with floats
            @test getEnergy(p)^2 - sum(momentum.^2) > -1e-12
        else
            @test abs(getEnergy(p)^2 - sum(momentum.^2)) < 1e-12
        end
        parentList = getParents(p)
        for parent in parentList
            @test getEnergy(parent) > 0
        end
        #direct access
        nParents = length(parentList)
        for i = 1:nParents
            @test getEnergy(parentList[i-1]) > 0
        end
    end
end
end
println("First iteration successful")

# test that everything is closed and opened properly
LCIO_Cxx.open("test.slcio") do reader
iEvent = 0
for event in reader
    iEvent += 1
    println(getCollectionNames(event))
    println(length(getCollectionNames(event)))
    @test length(getCollectionNames(event)) == 23
    @test getDetectorName(event) == "sidloi3_scint1x1"
    println("printing names")
    for name in getCollectionNames(event)
        println(name)
    end
    HcalBarrelHits = getCollection(event, "HcalBarrelHits")
    @test getTypeName(HcalBarrelHits) == "SimCalorimeterHit"
    decode = CellIDDecoder(HcalBarrelHits)
    println(decode)
    l = length(HcalBarrelHits)
    iHit = 0
    for h in HcalBarrelHits
        println(decode(h))
        @test 0 <= decode(h)["layer"] <= 39
        iHit += 1
    end
    # test iteration -- julia counting vs. C counting
    @test l == iHit
end
@test iEvent == length(reader)
end
println("Second iteration successful")

# test the stdhep reader
LCIO.openStdhep("test.stdhep") do reader
    iEvent = 0
    for event in reader
        iEvent += 1
        iParticle = 0
        for particle in event
            iParticle += 1
        end
        @test iParticle == length(event)
    end
    @test iEvent == length(reader)
end
println("Stdhep iteration successful")

# test creating a new file and writing out a particle
wrt = LCIO.createLCWriter()

LCIO.open(wrt, "writeTest.slcio", LCIO.WRITE_NEW)
run = LCIO.LCRunHeaderImpl()
LCIO.setRunNumber(run, 0)
LCIO.setValue(LCIO.parameters(run),"Purpose","runTest")
LCIO.setValue(LCIO.parameters(run), "intval", 1 )
LCIO.setValue(LCIO.parameters(run), "floatval", 3.14f0)
LCIO.writeRunHeader(wrt, run)
p = 5.0
pdg = -13
charge = +1.f0
mass =  0.105658f0
theta = 85./180.f0 * pi
for i in 1:1
    evt = LCIO.LCEventImpl()
    col = LCIO.LCCollectionVec(LCIO.MCPARTICLE)
    LCIO.setTransient(col, false)
    LCIO.setEventNumber(evt, i)
    phi = rand() * 2pi
    energy = sqrt(mass^2 + p^2)
    px = p * cos(phi) * sin(theta)
    py = p * sin(phi) * sin(theta)
    pz = p * cos(theta)
    momentum  = [ px, py, pz ]
#--------------- create MCParticle -------------------
    mcp = LCIO.MCParticleImpl()
    LCIO.setGeneratorStatus(mcp, 1)
    LCIO.setMass(mcp, mass)
    LCIO.setPDG(mcp, pdg)
    LCIO.setMomentum(mcp, momentum)
    LCIO.setCharge(mcp, charge)
    LCIO.addElement(col, mcp)
    LCIO.addCollection(evt, col, "genParticles")
    LCIO.writeEvent(wrt, evt)
end
LCIO.close(wrt)
println("WriteEvent successful")

ext: jl


text:
#!/usr/bin/env julia

@show ARGS

max_stars_repo_name: mkanenobu/trashbox

max_issues_repo_name: mkanenobu/trashbox

max_forks_repo_name: mkanenobu/trashbox


text:
using NNlib: DenseConvDims

@testset "convolution" begin
    a, b, c = rand(Float64, 10, 10, 3, 1), rand(Float64, 2, 2, 3, 4), rand(Float64, 9, 9, 4, 1)
    da, db, dc = CuArray(a), CuArray(b), CuArray(c)
    cdims = DenseConvDims(a, b)
    @test NNlib.conv(a, b, cdims) ≈ collect(NNlib.conv(da, db, cdims))
    @test ∇conv_data(c, b, cdims) ≈ collect(∇conv_data(dc, db, cdims))
    @test ∇conv_filter(a, c, cdims) ≈ collect(∇conv_filter(da, dc, cdims))

    # Test for agreement between CPU NNlib and CuDNN versions, across a variety of kwargs
    options = Dict{Any, Any}.((
        (), (:dilation => 2), (:flipkernel => true), (:stride => 2),
        (:padding => 1),
        (:padding => (1,0)),
        (:padding => (0,1)),
        (:padding => (2,3)),
    ))
    C_in_ = 3
    C_out = 4
    batch_size = 1

    for groups in (1, 2, 4), num_spatial_dims in (1, 2, 3)
        # Make `C_in = C_out` when using grouped convolution.
        C_in = groups == 1 ? C_in_ : C_out
        # Initialize data we'll run our tests over
        x = rand(Float64, fill(8, num_spatial_dims)..., C_in, batch_size)
        w = rand(Float64, fill(2, num_spatial_dims)..., C_in ÷ groups, C_out)

        for opts in options
            opts[:groups] = groups

            if :padding in keys(opts)
                padding = opts[:padding]
                if 1 < length(padding) && length(padding) != 2num_spatial_dims
                    opts[:padding] = ntuple(i -> padding[mod1(i,2)] .+ 2div(i-1,2), 2num_spatial_dims)
                end
            end

            cdims = DenseConvDims(x, w; opts...)
            y = NNlib.conv(x, w, cdims)

            # Test that basic convolution is equivalent across GPU/CPU
            gputest((x, w) -> NNlib.conv(x, w, cdims), x, w)
            gputest((y, w) -> NNlib.∇conv_data(y, w, cdims), y, w)
            gputest((x, y) -> NNlib.∇conv_filter(x, y, cdims), x, y, checkgrad=false) # TODO fix grad

            # Scaling factors
            gputest((x, w) -> NNlib.conv(x, w, cdims; alpha=2.0), x, w, checkgrad=false) # TODO
            gputest((y, w) -> NNlib.∇conv_data(y, w, cdims; alpha=2.0), y, w, checkgrad=false) # TODO
            gputest((x, y) -> NNlib.∇conv_filter(x, y, cdims; alpha=2.0), x, y, checkgrad=false) # TODO

            gputest((y, x, w) -> NNlib.conv!(copy(y), x, w, cdims; beta=2.0), y, x, w, checkgrad=false) # TODO
            # @test_broken gputest((x, y, w) -> NNlib.∇conv_data!(copy(x), y, w, cdims; beta=2.0), x, y, w, checkgrad=false) #TODO
            gputest((w, x, y) -> NNlib.∇conv_filter!(copy(w), x, y, cdims; beta=2.0), w, x, y, checkgrad=false) # TODO
        end
    end
end


text:
VERSION >= v"0.4-" && __precompile__()
module UAParser

export parsedevice, parseuseragent, parseos, DeviceResult, OSResult, UAResult, DataFrame

##############################################################################
##
## Dependencies
##
##############################################################################

using YAML, DataFrames
import DataFrames.DataFrame, DataFrames.names!

##############################################################################
##
## Load YAML file
##
##############################################################################

const REGEXES = YAML.load(open(joinpath(dirname(@__FILE__), "..", "regexes.yaml")));

##############################################################################
##
## Create custom types to hold parsed YAML output
##
##############################################################################

immutable UserAgentParser
  user_agent_re::Regex
  family_replacement::Union{AbstractString, Void}
  v1_replacement::Union{AbstractString, Void}
  v2_replacement::Union{AbstractString, Void}
end

immutable OSParser
  user_agent_re::Regex
  os_replacement::Union{AbstractString, Void}
  os_v1_replacement::Union{AbstractString, Void}
  os_v2_replacement::Union{AbstractString, Void}
end

immutable DeviceParser
  user_agent_re::Regex
  device_replacement::Union{AbstractString, Void}
end

##############################################################################
##
## Create custom types to hold function results
##
##############################################################################

immutable DeviceResult
  family::UTF8String
end

immutable UAResult
  family::AbstractString
  major::Union{AbstractString, Void}
  minor::Union{AbstractString, Void}
  patch::Union{AbstractString, Void}
end

immutable OSResult
  family::AbstractString
  major::Union{AbstractString, Void}
  minor::Union{AbstractString, Void}
  patch::Union{AbstractString, Void}
  patch_minor::Union{AbstractString, Void}
end

##############################################################################
##
## Create USER_AGENT_PARSERS, OS_PARSERS, and DEVICE_PARSERS arrays
##
##############################################################################

function loadua()
  #Create empty array to hold user-agent information
  temp = []

  #Loop over entire set of user_agent_parsers, add to USER_AGENT_PARSERS
  for _ua_parser in REGEXES["user_agent_parsers"]
      _user_agent_re = Regex(_ua_parser["regex"])
      _family_replacement = get(_ua_parser, "family_replacement", nothing)
      _v1_replacement = get(_ua_parser, "v1_replacement", nothing)
      _v2_replacement = get(_ua_parser, "v2_replacement", nothing)

    #Add values to array
      push!(temp, UserAgentParser(_user_agent_re,
                                  _family_replacement,
                                  _v1_replacement,
                                  _v2_replacement
                                  ))

  end
  return temp
end #End loadua

const USER_AGENT_PARSERS = loadua()

function loados()
  #Create empty array to hold os information
  temp = []

  #Loop over entire set of os_parsers, add to OS_PARSERS
  for _os_parser in REGEXES["os_parsers"]
    _user_agent_re = Regex(_os_parser["regex"])
    _os_replacement = get(_os_parser, "os_replacement", nothing)
    _os_v1_replacement = get(_os_parser, "os_v1_replacement", nothing)
    _os_v2_replacement = get(_os_parser, "os_v2_replacement", nothing)

    #Add values to array
    push!(temp, OSParser(_user_agent_re,
                         _os_replacement,
                         _os_v1_replacement,
                         _os_v2_replacement
                         ))

  end
  return temp
end #End loados

const OS_PARSERS = loados()

function loaddevice()
  #Create empty array to hold device information
  temp = []

  #Loop over entire set of device_parsers, add to DEVICE_PARSERS
  for _device_parser in REGEXES["device_parsers"]
      _user_agent_re = Regex(_device_parser["regex"])
      _device_replacement = get(_device_parser, "device_replacement", nothing)

    #Add values to array
      push!(temp, DeviceParser(_user_agent_re, _device_replacement))
  end
  return temp
end #End loaddevice

const DEVICE_PARSERS = loaddevice()

##############################################################################
##
## Functions for parsing user agent strings
##
##############################################################################

function parsedevice(user_agent_string::AbstractString)
  for value in DEVICE_PARSERS
    if ismatch(value.user_agent_re, user_agent_string)
      if value.device_replacement != nothing
        if ismatch(r"\$1", value.device_replacement)
          device = replace(value.device_replacement, "\$1", match(value.user_agent_re, user_agent_string).captures[1])
        else
          device = value.device_replacement
        end
      else
        device = match(value.user_agent_re, user_agent_string).captures[1]
      end

        return DeviceResult(device)
    end
  end

return DeviceResult("Other")  #Fail-safe for no match
end #End parsedevice

#Vectorize parsedevice for any array of user-agent strings
Base.@vectorize_1arg AbstractString parsedevice

function parseuseragent(user_agent_string::AbstractString)
  for value in USER_AGENT_PARSERS
    if ismatch(value.user_agent_re, user_agent_string)

      match_vals = match(value.user_agent_re, user_agent_string).captures

      #family
      if value.family_replacement != nothing
        if ismatch(r"\$1", value.family_replacement)
          family = replace(value.family_replacement, "\$1", match_vals[1])
        else
          family = value.family_replacement
        end
      else
        family = match_vals[1]
      end

      #major
      if value.v1_replacement != nothing
        v1 = value.v1_replacement
      elseif length(match_vals) > 1
        v1 = match_vals[2]
      else
        v1 = nothing
      end

      #minor
      if value.v2_replacement != nothing
        v2 = value.v2_replacement
      elseif length(match_vals) > 2
        v2 = match_vals[3]
      else
        v2 = nothing
      end

      #patch
      if length(match_vals) > 3
        v3 = match_vals[4]
      else
        v3 = nothing
      end

      return UAResult(family, v1, v2, v3)

    end
  end

return UAResult("Other", nothing, nothing, nothing) #Fail-safe for no match
end #End parseuseragent

#Vectorize parseuseragent for any array of user-agent strings
Base.@vectorize_1arg AbstractString parseuseragent

function parseos(user_agent_string::AbstractString)
    for value in OS_PARSERS
        if ismatch(value.user_agent_re, user_agent_string)
            match_vals = match(value.user_agent_re, user_agent_string).captures

            #os
            if value.os_replacement != nothing
                os = value.os_replacement
            else
                os = match_vals[1]
            end

            #os_v1
            if value.os_v1_replacement != nothing
                os_v1 = value.os_v1_replacement
            elseif length(match_vals) > 1
                os_v1 = match_vals[2]
            else
                os_v1 = nothing
            end

            #os_v2
            if value.os_v2_replacement != nothing
                os_v2 = value.os_v2_replacement
            elseif length(match_vals) > 2
                os_v2 = match_vals[3]
            else
                os_v2 = nothing
            end

            #os_v3
            if length(match_vals) > 3
                os_v3 = match_vals[4]
            else
                os_v3 = nothing
            end

            #os_v4
            if length(match_vals) > 4
                os_v4 = match_vals[5]
            else
                os_v4 = nothing
            end

            return OSResult(os, os_v1, os_v2, os_v3, os_v4)

        end
    end

return OSResult("Other", nothing, nothing, nothing, nothing) #Fail-safe if no match
end #End parseos

#Vectorize parseos for any array of user-agent strings
Base.@vectorize_1arg AbstractString parseos

##############################################################################
##
## Extend DataFrames to include UAParser methods
##
##############################################################################

#Convenience function to take nothing type to UTF8String of length 0
#This is a hack for sure
function nothing_to_utf8empty(x::Union{AbstractString, Void})
  if x == nothing
    x = ""
  else
    x = convert(UTF8String, x)
  end
end

#DeviceResult to DataFrame method
function DataFrame(x::Array{DeviceResult, 1})
  temp = DataFrame([element.family for element in x])
  names!(temp, ["device"])
  return temp
end

#OSResult to DataFrame method
function DataFrame(x::Array{OSResult, 1})
  #Pre-allocate size of DataFrame based on array passed in
  temp = DataFrame(UTF8String, size(x, 1), 5)
  names!(temp, ["os_family", "os_major", "os_minor", "os_patch", "os_patch_minor"])

  #Family - Can use comprehension since family always UTF8String
  temp["os_family"] = UTF8String[element.family for element in x]

  #Major
  temp["os_major"] = UTF8String[nothing_to_utf8empty(element.major) for element in x]

  #Minor
  temp["os_minor"] = UTF8String[nothing_to_utf8empty(element.minor) for element in x]

  #Patch
  temp["os_patch"] = UTF8String[nothing_to_utf8empty(element.patch) for element in x]

  #Patch_Minor
  temp["os_patch_minor"] = UTF8String[nothing_to_utf8empty(element.patch_minor) for element in x]

  return temp
end

#UAResult to DataFrame method
function DataFrame(x::Array{UAResult, 1})
  #Pre-allocate size of DataFrame based on array passed in
  temp = DataFrame(UTF8String, size(x, 1), 4)
  names!(temp, ["browser_family", "browser_major", "browser_minor", "browser_patch"])

  #Family - Can use comprehension since family always UTF8String
  temp["browser_family"] = UTF8String[element.family for element in x]

  #Major
  temp["browser_major"] = UTF8String[nothing_to_utf8empty(element.major) for element in x]

  #Minor
  temp["browser_minor"] = UTF8String[nothing_to_utf8empty(element.minor) for element in x]

  #Patch
  temp["browser_patch"] = UTF8String[nothing_to_utf8empty(element.patch) for element in x]

  return temp
end

end # module


text:
is_jwt, jwt_decoded = SMARTBackendServices.try_decode_jwt("")
@test !is_jwt
@test jwt_decoded == Dict{String, Any}()
@test isempty(jwt_decoded)


text:
# adapted from Petri.jl
# https://github.com/mehalter/Petri.jl

graph_attrs = Attributes(:rankdir=>"LR")
node_attrs  = Attributes(:shape=>"plain", :style=>"filled", :color=>"white")
edge_attrs  = Attributes(:splines=>"splines")

function edgify(δ, i, reverse::Bool)
    attr = Attributes()
    return map(δ) do p
        val = String(p[1].op.name)
      weight = "$(p[2])"
      attr = Attributes(:label=>weight, :labelfontsize=>"6")
      return Edge(reverse ? ["rx_$i", "$val"] :
                            ["$val", "rx_$i"], attr)
    end
end

# make distinguished edge based on rate constant
function edgifyrates(rxs, specs)
    es = Edge[]
    for (i,rx) in enumerate(rxs)
        deps = rx.rate isa Operation ? get_variables(rx.rate, specs) : Operation[]
        for dep in deps
            val = String(dep.op.name)
            attr = Attributes(:color => "#d91111", :style => "dashed")
            e = Edge(["$val", "rx_$i"], attr)
            push!(es, e)
        end
    end
    es
end

"""
    Graph(rn::ReactionSystem)

Converts a [`ReactionSystem`](@ref) into a
[Catlab.jl](https://github.com/AlgebraicJulia/Catlab.jl/) Graphviz graph.
Reactions correspond to small green circles, and species to blue circles.

Notes:
- Black arrows from species to reactions indicate reactants, and are labelled
  with their input stoichiometry.
- Black arrows from reactions to species indicate products, and are labelled
  with their output stoichiometry.
- Red arrows from species to reactions indicate that species is used within the
  rate expression. For example, in the reaction `k*A, B --> C`, there would be a
  red arrow from `A` to the reaction node. In `k*A, A+B --> C`, there would be
  red and black arrows from `A` to the reaction node.
"""
function Graph(rn::ReactionSystem)
    rxs   = reactions(rn)
    specs = species(rn)
    statenodes = [Node(string(s.name), Attributes(:shape=>"circle", :color=>"#6C9AC3")) for s in specs]
    transnodes = [Node(string("rx_$i"), Attributes(:shape=>"point", :color=>"#E28F41", :width=>".1")) for (i,r) in enumerate(rxs)]

    stmts = vcat(statenodes, transnodes)
    edges = map(enumerate(rxs)) do (i,r)
      vcat(edgify(zip(r.substrates,r.substoich), i, false),
           edgify(zip(r.products,r.prodstoich), i, true))
    end
    es = edgifyrates(rxs, specs)
    (!isempty(es)) && push!(edges, es)

    stmts = vcat(stmts, collect(flatten(edges)))
    g = Digraph("G", stmts; graph_attrs=graph_attrs, node_attrs=node_attrs, edge_attrs=edge_attrs)
    return g
end

"""
    savegraph(g::Graph, fname, fmt="png")

Given a [Catlab.jl](https://github.com/AlgebraicJulia/Catlab.jl/) `Graph`
generated by [`Graph`](@ref), save the graph to the file with name `fname` and
extension `fmt`.

Notes:
- `fmt="png"` is the default output format.
"""
function savegraph(g::Graph, fname, fmt="png")
    open(fname, "w") do io
        run_graphviz(io, g, format=fmt)
    end
    nothing
end


text:
@testset "Midpoints" begin
    @test_throws Exception Midpoints{Int}(2)
    @test eltype(@inferred(Midpoints(2))) == Float64

    for T in Floats
        @test eltype(@inferred(Midpoints{T}(3))) == T
        @test typeof(@inferred(Midpoints{T}(3)[1])) == T

        @test Midpoints{T}(0) ≈ T[]
        @test Midpoints{T}(1) ≈ T[0]  atol = eps(real(T))
        @test Midpoints{T}(2) ≈ T[-0.5,0.5]
        @test Midpoints{T}(3) ≈ T[-T(2)/3,0,T(2)/3]
        @test Midpoints{T}(4) ≈ T[-0.75,-0.25,0.25,0.75]
    end
end


text:
if !in("../src",LOAD_PATH)
    push!(LOAD_PATH,"../src")
end

using GeneMisc
using Base.Test

@test 1 == 1

include("synonym.jl")
#include("location.jl")
#include("exon.jl")
#=
let gene = "KRAS", genesyn = ["KRAS","C-K-RAS","CFC2","K-RAS2A","K-RAS2B",
                              "K-RAS4A","K-RAS4B","KI-RAS","KRAS1","KRAS2",
                              "NS","NS3","RALD","RASK2"]

#    @test query_gene("KRAS") == Set(genesyn)

end

let chr="chr20",pos="45314173"
#    @test query_gene(chr,pos) == "TP53RK"
end
=#


text:
function differentiate_traces(d;dt=0.001,pow=-2,rot=0)

	nt = size(d,1)
	D = fft(d,1)
	dw = 2.*pi/nt/dt
	nw = Int(nt/2) + 1
	eps = pow < 0 ? 10.0 : 0.0
	for iw=1:nw
		D[iw,:] *= exp(rot*1im)*((iw*dw) + eps)^pow
	end
	# symmetries
	for iw=nw+1:nt
		D[iw,:] = conj(D[nt-iw+2,:])
	end
	d = real(ifft(D,1))
	return d

end

function SeisDiff(in,h::Array{Header,1};pow=-2,rot=0)

	out = differentiate_traces(in;dt=h[1].d1,pow=pow,rot=rot)
	return out,h

end

function SeisDiff(in::String,out::String;pow=-2,rot=0)

	@compat parameters = Dict(:pow=>pow,:rot=>rot)
	SeisProcess(in,out,[SeisDiff],[parameters])

end


text:
module CameraTargetTrigger

using ..Ahorn, Maple

const placements = Ahorn.PlacementDict(
    "Camera Target" => Ahorn.EntityPlacement(
        Maple.CameraTargetTrigger,
        "rectangle",
        Dict{String, Any}(),
        function(trigger)
            trigger.data["nodes"] = [(Int(trigger.data["x"]) + Int(trigger.data["width"]) + 8, Int(trigger.data["y"]))]
        end
    )
)

function Ahorn.editingOptions(trigger::Maple.CameraTargetTrigger)
    return Dict{String, Any}(
        "positionMode" => Maple.trigger_position_modes
    )
end

function Ahorn.nodeLimits(trigger::Maple.CameraTargetTrigger)
    return 1, 1
end

end


text:
using Distributions, Turing, Random, Parameters
ProjDir = @__DIR__
cd(ProjDir)
include("LNR.jl")

@model model(data, Nr) = begin
    μ = Array{Real,1}(undef,Nr)
    μ ~ [Normal(0,3)]
    s ~ Uniform(0, pi/2)
    σ = tan(s)
    #σ ~ Truncated(Cauchy(0, 1), 0, Inf)
    for i in 1:length(data)
        data[i] ~ LNR(μ, σ, 0.0)
    end
end

Random.seed!(343)
Nreps = 50
error_count = fill(0.0, Nreps)
Nr = 3
Nobs = 10

for i in 1:Nreps
    μ = rand(Normal(0, 3), Nr)
    σ = rand(Uniform(.2, 1))
    dist = LNR(μ=μ, σ=σ, ϕ=0.0)
    data = rand(dist, Nobs)
    chain = psample(model(data, Nr), NUTS(1000, .8), 2000, discard_adapt = false, progress = false)
    error_count[i] = sum(get(chain,:numerical_error)[1])
end

describe(error_count)

max_stars_repo_name: itsdfish/Miscellaneous

max_issues_repo_name: itsdfish/Miscellaneous

max_forks_repo_name: itsdfish/Miscellaneous


text:
# Create STO Problem
function stoproblem(
  x0::Array{Float64,1},                 # Initial State
  A::Array{Float64,3};                  # Linear Dynamics
  ngrid::Int64=2,                       # Number of Linearization points in the fixed grid (2 for linear case. We do not need them by default)
  t0::Float64=0.0,                      # Initial Time
  tf::Float64=1.0,                      # Final Time
  Q::Array{Float64, 2}=emptyfmat,       # Cost Matrix
  E::Array{Float64, 2}=emptyfmat,      # Final Cost Matrix
  lb::Array{Float64, 1}=emptyfvec,      # Lower Bound on Intervals
  ub::Array{Float64, 1}=emptyfvec,      # Upper Bound on Intervals
  tau0ws::Array{Float64,1}=emptyfvec,   # Warm Start tau vector
  solver::MathOptInterface.AbstractOptimizer=Ipopt.Optimizer())

  # Get Dimensions
  nx = size(A, 1)      # State Dimension
  N = size(A, 3) - 1   # Get number of switches

  # Adjust variables which have not been initalized
  if isempty(Q)
    Q = 1.0 * Matrix(I, nx, nx)
  end

  if isempty(E)
    E = zeros(nx, nx)
  end

  if isempty(lb)
    lb = zeros(N+1)
  end

  if isempty(ub)
    ub = Inf*ones(N+1)
  end

  if isempty(tau0ws)
    tau0ws = collect(range(t0, stop=tf, length=N+2))
    tau0ws = tau0ws[2:end-1]
  end

  # Define warm starting delta0
  delta0ws = tau2delta(tau0ws, t0, tf)

  # Create Discretization Grid
  tgrid = collect(range(t0, stop=tf, length=ngrid))

  # Initialize time vectors
  tvec = Array{Float64}(undef, N + ngrid)    # Complete grid
  tauIdx = Array{Int}(undef, N + 2)      # Indeces of switching times in the complete grid
  tgridIdx = Array{Int}(undef, ngrid)      # Indeces of switching times in the complete grid
  deltacomplete = Array{Float64}(undef, N + ngrid - 1)   # Intervals over the whole grid

  ### Initialize NLP Evaluator
  # Preallocate arrays
  deltafun_prev = Array{Float64}(undef, N+1)
  deltagrad_prev = Array{Float64}(undef, N+1)
  deltahess_prev = Array{Float64}(undef, N+1)
  xpts = Array{Float64}(undef, nx, N+ngrid); xpts[:, 1] = x0   # Set Initial State
  expMat = Array{Float64}(undef, nx, nx, N+ngrid-1)
  Phi = Array{Float64}(undef, nx, nx, N+2, N+2)
  M = Array{Float64}(undef, nx, nx, N+ngrid-1)
  S = Array{Float64}(undef, nx, nx, N+ngrid)
  C = Array{Float64}(undef, nx, nx, N+1)

  # Decompose Dynamics Matrices
  V = Array{Complex{Float64}}(undef, 2*nx, 2*nx, N+1)
  invV = Array{Complex{Float64}}(undef, 2*nx, 2*nx, N+1)
  D =  Array{Complex{Float64}}(undef, 2*nx, N+1)
  isDiag = Array{Bool}(undef, N+1)

  for i = 1:N+1
    D[:, i], V[:, :, i] = eigen([-A[:, :, i]'  Q;
                             zeros(nx, nx) A[:, :, i]])
    if cond(V[:, :, i]) == Inf  # Non diagonalizable matrix
      isDiag[i] = false
    else
      invV[:, :, i] = inv(V[:, :, i])
      isDiag[i] = true
    end
  end

  # Construct Matrix of Indeces for lower triangular Matrix (Hessian)
  IndTril = (LinearIndices(tril(ones(N+1, N+1))))[findall(!iszero,tril(ones(N+1, N+1)))]
  Itril, Jtril, _ = begin
  I_temp = findall(!iszero, tril(ones(N+1, N+1)))
  (getindex.(I_temp, 1), getindex.(I_temp, 2), tril(ones(N+1, N+1))[I_temp])
  end

  # Construct Constraints Matrix (Vector)
  Ag = ones(1, N+1)
  Ig, Jg, Vg = begin
    I_temp = findall(!iszero, Ag)
    (getindex.(I_temp, 1), getindex.(I_temp, 2), Ag[I_temp])
  end
  bg = [tf]   # Only one constraint for the sum of the switching intervals

  # Initialize objective evaluator
  obj = Array{Float64}(undef, 0)
  deltaval = Array{Float64}(undef, N+1, 0)
  nobjeval = 0                           # Number of objective function evaluations
  ngradeval = 0                           # Number of gradient evaluations
  nhesseval = 0                           # Number of hessian evaluations

  # Construct NLP evaluator
  STOev = linSTOev(x0, nx, A, N, t0, tf, tf, Q, E, ngrid, tgrid, tvec,
          tauIdx, tgridIdx, deltacomplete, V, invV, D, isDiag, IndTril,
          Itril, Jtril, Ag, Ig, Jg, Vg, bg, deltafun_prev, deltagrad_prev,
          deltahess_prev, xpts, expMat, Phi, M, S, C,
          obj, deltaval, nobjeval, ngradeval, nhesseval)

model = solver
MathOptInterface.empty!(model)
delta = MathOptInterface.add_variables(model, N+1)
for w_i in delta
  MathOptInterface.add_constraint(model, MathOptInterface.SingleVariable(w_i), MathOptInterface.GreaterThan(0.0))
  MathOptInterface.add_constraint(model, MathOptInterface.SingleVariable(w_i), MathOptInterface.LessThan(tf))
end

lbub = [MathOptInterface.NLPBoundsPair(tf, tf)]
MathOptInterface.set(model, MathOptInterface.VariablePrimalStart(), delta, delta0ws)
block_data = MathOptInterface.NLPBlockData(
  lbub,
  STOev,
  true
)
MathOptInterface.set(model, MathOptInterface.NLPBlock(), block_data)
MathOptInterface.set(model, MathOptInterface.ObjectiveSense(), MathOptInterface.MIN_SENSE)

# Propagate dynamic for new switching times
propagateDynamics!(STOev, delta0ws)

# Create STO
STOproblem = linSTO(model, STOev, delta0ws, delta)
return STOproblem  # Return STO

end

# Create STO Problem for Nonlinear STO
function stoproblem(
  x0::Array{Float64,1},             # Initial State
  nonlin_dyn::Function,             # Nonlinear Dynamics
  nonlin_dyn_deriv::Function,       # Nonlinear Dynamics Derivative
  uvec::Array{Float64, 2};          # Vector of integer Inputs per switching combination
  ngrid::Int64=10,                  # Number of Linearization points in the fixed grid
  t0::Float64=0.0,                  # Initial Time
  tf::Float64=1.0,                  # Final Time
  Q::Array{Float64, 2}=emptyfmat,   # Cost Matrix
  E::Array{Float64, 2}=emptyfmat,  # Final Cost Matrix
  lb::Array{Float64, 1}=emptyfvec,  # Lower Bound on Intervals
  ub::Array{Float64, 1}=emptyfvec,  # Upper Bound on Intervals
  tau0ws::Array{Float64,1}=emptyfvec, # Warm Start tau vector
  solver::MathOptInterface.AbstractOptimizer=Ipopt.Optimizer())

  # Get Dimensions
  nx = length(x0)         # State Dimension

  # Generate artificial switches repeating the input sequence
  N = size(uvec, 2) - 1   # Get total number of switches

  # Adjust variables which have not been initalized
  if isempty(Q)
    Q = 1.0 * Matrix(I, nx, nx)
  end

  if isempty(E)
    E = zeros(nx, nx)
  end

  if isempty(lb)
    lb = zeros(N+1)
  end

  if isempty(ub)
    ub = Inf*ones(N+1)
  end

  if isempty(tau0ws)
    tau0ws = collect(range(t0, stop=tf, length=N+2))  # Currently counting tau_0 and tau_{N+1}. They are removed below after xpts initialization.
    tau0ws = tau0ws[2:end-1]  # Include only switching instants

  end

  # Create Discretization grid
  tgrid = collect(range(t0, stop=tf, length=ngrid))

  # Initialize time vectors
  tvec = Array{Float64}(undef, N + ngrid)    # Complete grid
  tauIdx = Array{Int}(undef, N + 2)      # Indeces of switching times in the complete grid
  tgridIdx = Array{Int}(undef, ngrid)      # Indeces of switching times in the complete grid
  deltacomplete = Array{Float64}(undef, N + ngrid - 1)   # Intervals over the whole grid

  ### Initialize NLP Evaluator
  # Extend System State To get Affine to Linear Dynamics
  nx = nx + 1

  # Extend Initial State and Cost Matrix
  x0 = [x0; 1]
  spz = Array{Float64}(undef, 1,1); spz[1,1] = 0.0; spz = sparse(spz)  # Sparse scalar
  Q = Matrix(blockdiag(sparse(Q), spz))
  E = Matrix(blockdiag(sparse(E), spz))

  # Define Required Matrices and Switching Instants
  A = Array{Float64}(undef, nx, nx, N+ngrid-1)

  # Initialize Switching Instants
  xpts = Array{Float64}(undef, nx, N+ngrid)
  xpts[:, 1] = x0

  # Define warm starting delta0ws from tau0ws
  delta0ws = tau2delta(tau0ws, t0, tf)

  deltafun_prev = Array{Float64}(undef, N+1)
  deltagrad_prev = Array{Float64}(undef, N+1)
  deltahess_prev = Array{Float64}(undef, N+1)
  expMat = Array{Float64}(undef, nx, nx, N+ngrid-1)
  Phi = Array{Float64}(undef, nx, nx, N+2, N+2)
  M = Array{Float64}(undef, nx, nx, N+ngrid-1)
  S = Array{Float64}(undef, nx, nx, N+ngrid)
  C = Array{Float64}(undef, nx, nx, N+1)  # Only at sw times (including initial time)

  # Construct Matrix of Indeces for lower triangular Matrix (Hessian)
  IndTril = (LinearIndices(tril(ones(N+1, N+1))))[findall(!iszero,tril(ones(N+1, N+1)))]
  Itril, Jtril, _ = begin
  I_temp = findall(!iszero, tril(ones(N+1, N+1)))
  (getindex.(I_temp, 1), getindex.(I_temp, 2), tril(ones(N+1, N+1))[I_temp])
  end

  # Construct Constraints Matrix (Vector)
  Ag = ones(1, N+1)
  Ig, Jg, Vg = begin
  I_temp = findall(!iszero, Ag)
  (getindex.(I_temp, 1), getindex.(I_temp, 2), Ag[I_temp])
  end
  bg = [tf]   # Only one constraints for the sum of the switching intervals

  # Initialize objective evaluator
  obj = Array{Float64}(undef, 0)
  deltaval = Array{Float64}(undef, N+1, 0)
  nobjeval = 0                           # Number of objective function evaluations
  ngradeval = 0                           # Number of gradient evaluations
  nhesseval = 0                           # Number of hessian evaluations

  # Construct NLPEvaluator
  STOev = nlinSTOev(x0, nx, A, N, t0, tf, tf, Q, E, uvec, ngrid, tgrid, tvec, tauIdx, tgridIdx, deltacomplete, nonlin_dyn,
                 nonlin_dyn_deriv, IndTril, Itril, Jtril, Ag, Ig, Jg, Vg, bg, deltafun_prev, deltagrad_prev, deltahess_prev,
                 xpts, expMat, Phi, M, S, C, obj, deltaval, nobjeval, ngradeval, nhesseval)

  # Propagate Dynamics to compute matrix exponentials and states at the switching times
  propagateDynamics!(STOev, delta0ws)

  # Generate Model
  model = solver
  MathOptInterface.empty!(model)
  delta = MathOptInterface.add_variables(model, N+1)
  for i=1:N+1
    MathOptInterface.add_constraint(model, MathOptInterface.SingleVariable(delta[i]), MathOptInterface.GreaterThan(lb[i]))
    MathOptInterface.add_constraint(model, MathOptInterface.SingleVariable(delta[i]), MathOptInterface.LessThan(ub[i]))
  end

  lbub = [MathOptInterface.NLPBoundsPair(tf, tf)]
  MathOptInterface.set(model, MathOptInterface.VariablePrimalStart(), delta, delta0ws)
  block_data = MathOptInterface.NLPBlockData(
    lbub,
    STOev,
    true
  )
  MathOptInterface.set(model, MathOptInterface.NLPBlock(), block_data)
  MathOptInterface.set(model, MathOptInterface.ObjectiveSense(), MathOptInterface.MIN_SENSE)

  # Create STO
  STOproblem = nlinSTO(model, STOev, delta0ws, delta)

  return STOproblem  # Return STO

end

# Solve Optimization for Linear System
function solve!(m::STO)

  # Perform STO
  m.soltime = @elapsed MathOptInterface.optimize!(m.model)
  m.stat = MathOptInterface.get(m.model, MathOptInterface.TerminationStatus())
  m.delta = MathOptInterface.get(m.model, MathOptInterface.VariablePrimal(), m.delta_MOI)
  m.objval = MathOptInterface.get(m.model, MathOptInterface.ObjectiveValue())
  m.tau, _ = delta2tau(m.delta, m.STOev.t0)

end

"Set warm starting point"
function setwarmstart!(m::STO, tau0ws::Array{Float64,1})

  # Define warm starting delta0
  delta0ws = tau2delta(tau0ws, m.STOev.t0, m.STOev.tf)

  # Set Warm Starting Point for Nonlinear Solver
  MathOptInterface.set(m.model, VariablePrimalStart, delta_MOI, delta0ws)

  # Reset Cost Function Iterates
  m.STOev.obj = Array{Float64}(undef, 0)
  m.STOev.deltaval = Array{Float64}(undef, m.STOev.N+1, 0)

end

"Set initial state x0"
function setx0!(m::STO, x0::Array{Float64,1})

  # Define initial point
  m.STOev.x0 = x0

end

# Return Variables from STO
gettau(m::STO) = m.tau
getdelta(m::STO) = m.delta
getobjval(m::STO) = m.objval
getstat(m::STO) = m.stat
getsoltime(m::STO) = m.soltime
getnobjeval(m::STO) = m.STOev.nobjeval
getngradeval(m::STO) = m.STOev.ngradeval
getnhesseval(m::STO) = m.STOev.nhesseval

# Convert from Switching Times to Intervals
function tau2delta(tau::Array{Float64, 1}, t0::Float64, tf::Float64)

  # Extend vector of switching times with initial and final time
  tauexp = [t0; tau; tf]  # Extend tau vector to simplify numbering

  # Return Vector of differences
  return  diff(tauexp)

end

# Convert from Intervals to Switching Times
function delta2tau(delta::Array{Float64, 1}, t0::Float64)

  # Define tau vector
  tau = Array{Float64}(undef, length(delta)-1)

  # Initialize first tau element
  tau[1] = t0 + delta[1]

  # Iteratively add all deltas to create switching times vector
  for i = 2:length(tau)
    tau[i] = tau[i-1] + delta[i]
  end

  tfdelta = tau[end] + delta[end]

  return tau, tfdelta

end

function plotSolution(d::nlinSTO, savename::String, title::String, show=true, save=false)

  tauopt = gettau(d)
  xsim, ~, ~, t = simulate(d, tauopt)
  usim, ~ = simulateinput(d, t)

  n_omega = size(usim)[1]
  nx = d.STOev.nx -1
  # Plot solution

  l = @layout [a;b]
  p1 = plot(title=title, titlefont=font(10))
  for i=1:nx
    plot!(t, xsim[i,:], label="x"*string(i))
  end
  plot!(grid=true, ylim=[(1-0.2*sign(minimum(xsim)))*minimum(xsim), (1+0.2*sign(maximum(xsim)))*maximum(xsim)])

  p2 = plot()
  for i=1:n_omega
    plot!(t,usim[i,:], label="u"*string(i), linetype=:steppre)
  end
  plot!(ylim=[(1-0.2*sign(minimum(xsim)))*minimum(usim), (1+0.2*sign(maximum(xsim)))*maximum(usim)], grid=true, xlabel="time")

  if show && save
    display(plot(p1, p2, layout=l))
    savefig(savename*".pdf")
  elseif save && !show
    plot(p1, p2, layout=l)
    savefig(savename*".pdf")
  elseif !save && show
    display(plot(p1, p2, layout=l))
  end
  return
end

function plotSolution(d::linSTO, u::Array{Float64,2}, savename::String, title::String, show=true, save=false)

  tauopt = gettau(d)
  xsim, ~, ~, t = simulate(d, tauopt)
  println(length(t))
  n_omega = size(u)[1]
  nx = d.STOev.nx
  # Plot solution
  l = @layout [a;b]
  p1 = plot(title=title, titlefont=font(10))
  for i=1:nx
    plot!(t, xsim[i,:], label="x"*string(i))
  end
  plot!(grid=true, ylim=[(1-0.2*sign(minimum(xsim)))*minimum(xsim), (1+0.2*sign(maximum(xsim)))*maximum(xsim)])

  p2 = plot()
  tSwitch = [d.STOev.t0; tauopt; d.STOev.tf]
  println(length(tSwitch))
  println(length([u[1,1]; u[1,:]]))
  for i=1:n_omega
    plot!(tSwitch, [u[i,1]; u[i,:]], label="A"*string(i), linetype=:steppre)
  end
  plot!(ylim=[(1-0.2*sign(minimum(xsim)))*minimum(u), (1+0.2*sign(maximum(xsim)))*maximum(u)], grid=true, xlabel="time")

  if show && save
    display(plot(p1, p2, layout=l))
    savefig(savename*".pdf")
  elseif save && !show
    plot(p1, p2, layout=l)
    savefig(savename*".pdf")
  elseif !save && show
    display(plot(p1, p2, layout=l))
  end
  return
end


text:
# Function to generate constraint and/or external forces

# include("RigidBody.jl")
# include("OrientationConversion.jl")
# include("Joint.jl")
using GenericLinearAlgebra
using LinearAlgebra
using ForwardDiff
using Revise
using StaticArrays
using SparseArrays
using FiniteDifferences

function Constraint(x::AbstractArray{T}, u::AbstractArray{S}, j::Vector{Joint},  GravityInInertial::Vector{Float64})::Tuple{AbstractArray{Union{T,S}}, AbstractArray{Union{T,S}}} where {T<:Real, S<:Real}

    Fc, Ffinal = FcFn(x, u, j, GravityInInertial);

    # Partition generated ForceConstr for each rigid body.
    ForceConstr = Matrix{Union{T,S}}(undef,(7,length(j)+1))
    ForceConstr[:,1] = zeros(7);
    for i=1:length(j)
        k = j[i].RB2.bodyID
        ForceConstr[:,k] = Fc[7*(k-2)+1:7*(k-1)]
    end

    # Partition generated Ffinal (unconstrainedF) for each rigid body.
    unconstrF = Matrix{Union{T,S}}(undef, (7,length(j)+1));
    unconstrF[:,1] = zeros(7);
    for i=1:length(j)
        k = j[i].RB2.bodyID
        unconstrF[:,k] = Ffinal[7*(k-2)+1:7*(k-1)]
    end
    # Return unconstrF, ForceConstr.
    return (unconstrF,ForceConstr)

end

function FcFn(x::AbstractArray{T}, u::AbstractArray{S}, j::Vector{Joint}, GravityInInertial::Vector{Float64}) where {T<:Real, S<:Real}

    Afinal, bfinal = Ab_VecOfMat(x,j);
    Afinal = Array(Afinal); bfinal = Array(bfinal);

    Ffinal = assembleF(x,u,j,GravityInInertial)

    Mfinal, Ms, Ms_inv = M_mat(x,j);

    accUnconstr = inv(Mfinal)*Ffinal # Unconstrained Acceleration

    if !isreal(Ms) # insignificant imaginary values showing up
        M1 = real(Ms)
        M2 = real(Ms_inv)

        tempMat = Afinal*M2
        tempMat_pinv = tempMat'*((tempMat*tempMat')\I(size(tempMat)[1]));
        Fc = Ms*tempMat_pinv*(bfinal - Afinal*accUnconstr);
    else
        tempMat = Afinal*Ms_inv;
        tempMat_pinv = tempMat'*((tempMat*tempMat')\I(size(tempMat)[1]));
        Fc = Ms*tempMat_pinv*(bfinal - Afinal*accUnconstr); # using Udwadia's formulation
    end
    # Fc and Ffinal are in the inertial frame
return Fc, Ffinal
end

function Ab_VecOfMat(x::AbstractArray{T}, j::Vector{Joint})::Tuple{AbstractArray{T}, AbstractArray{T}} where T<:Real
    nCols = 7*(j[end].RB2.bodyID-1);
    nConstrEqns_j = zeros(Integer,length(j));
    ## Number of constraint equations for each joint in robot
    for i=1:length(j)
        # Use jointDict to identify number of nConstrEqns in each joint
        if j[i].RB1.bodyID > 1 # First body not the inertial frame
            nConstrEqns_j[i] = get(JointDict, j[i].type, 0)
        else
            nConstrEqns_j[i]= get(JointDictIn, j[i].type, 0)
        end
    end
    nConstrEqns = sum(nConstrEqns_j) # Total number of constraint equations in

    # Afinal = zeros(T,(nConstrEqns,nCols))
    # bfinal = zeros(T,(nConstrEqns))

    Afinal = spzeros(T,nConstrEqns,nCols)
    bfinal = spzeros(T,(nConstrEqns))

    for i = 1:length(j)
        if i > 1
            b1 = j[i].RB1.bodyID-1; b2 = j[i].RB2.bodyID-1;
            nR = sum(nConstrEqns_j[1:i-1])
            Afinal[nR+1:nR + nConstrEqns_j[i],7*(b1-1)+1:7*b1] = ForceCon(x,j[i])[1][:,1:7] # not required for first joint
            Afinal[nR+1:nR + nConstrEqns_j[i],7*(b2-1)+1:7*b2] = ForceCon(x,j[i])[1][:,8:14]
            bfinal[nR+1:nR + nConstrEqns_j[i]] = ForceCon(x,j[i])[2]

            ## old
            # Afinal[nConstrEqns_j[i-1]+1:nConstrEqns_j[i-1] + nConstrEqns_j[i],7*(b1-1)+1:7*b1] = ForceCon(x,j[i])[1][:,1:7] # not required for first joint
            # Afinal[nConstrEqns_j[i-1]+1:nConstrEqns_j[i-1] + nConstrEqns_j[i],7*(b2-1)+1:7*b2] = ForceCon(x,j[i])[1][:,8:14]
            # bfinal[nConstrEqns_j[i-1]+1:nConstrEqns_j[i-1] + nConstrEqns_j[i]] = ForceCon(x,j[i])[2]
        else
            b2 = j[i].RB2.bodyID-1;
            Afinal[1:nConstrEqns_j[i],7*(b2-1)+1:7*b2] = ForceCon(x,j[i])[1][:,8:14]
            bfinal[1:nConstrEqns_j[i]] = ForceCon(x,j[i])[2]
        end

    end

    # Quaternion constraint gets repeated for bodies common to multiple joints
    Afinal, ids = unique_ids(Afinal);
    bfinal = bfinal[ids];
    return Afinal, bfinal
end

function M_mat(x::AbstractArray{T}, j::Vector{Joint})::Tuple{AbstractArray{T}, AbstractArray{T}, AbstractArray{T}} where T<:Real
    Mfinal = Array(assembleM(x,j));
    # if typeof(x[1]) == Float64 # eigendecomposition not required for sqrt
    Ms = real(sqrt(Mfinal)) # insignificant imaginary values showing up
    Ms_inv = Ms\I(size(Ms)[1])

    # else # Dual type variable (for ForwardDiff), applicable when taking jacobian
    #     # Ms = (sqrt.(diag(Mfinal))).*Matrix{T}(I,size(Mfinal))
    #
    #     # diag_Mfinal = diag(Mfinal);
    #     # vals_sq = zeros(T,size(Mfinal,1))
    #     # for i=1:size(Mfinal,1)
    #     #     vals_sq[i] = sqrt(diag_Mfinal[i])
    #     # end
    #     # Ms = diagm(vals_sq);
    #
    #     # Ms = (Mfinal^(0.5))
    #
    #     # vals, vecs = eigen(Mfinal);
    #     # vals_sq = zeros(T,size(vals))
    #     # for i=1:length(vals)
    #     #     vals_sq[i] = sqrt(vals[i])
    #     # end
    #     # Ms = vecs*diagm(vals_sq)*(vecs\I(size(vecs,1)))
    #
    #     F = svd(Mfinal)
    #     Ms = permutedims(F.Vt)*diagm(sqrt.(F.S))*F.Vt
    #     Ms_inv = Ms\I(size(Ms)[1]); #M^(-0.5)
    # end
    return Mfinal, Ms, Ms_inv
end

function assembleM(x::AbstractArray{T}, j::Vector{Joint})::AbstractArray{T} where T<:Real
    maxBodyID = j[end].RB2.bodyID
    v = collect(3:maxBodyID)

    # Mfinal = zeros(T,(7*(maxBodyID-1),7*(maxBodyID-1)))
    Mfinal = spzeros(T,7*(maxBodyID-1),7*(maxBodyID-1))

    # First body always connected to inertial frame
    Mfinal[1:7,1:7] = genMatM(x,j[1].RB2)

    # Iterate over joints
    for i = 2:length(j)
        b1 = j[i].RB1; b2 = j[i].RB2;
        if b2.bodyID in v
            Mfinal[7*(i-1)+1:7*i,7*(i-1)+1:7*i] = genMatM(x,b2)
            # Remove corresponding element from v
            filter!(x->x≠b2.bodyID,v)
        end
    end
    return Mfinal
end

function assembleF(x, u, j::Vector{Joint}, GravityInInertial:: Vector{Float64})
    maxBodyID = j[end].RB2.bodyID
    F = Vector(undef,7*(maxBodyID-1))

    for i=1:length(j)
        k = j[i].RB2.bodyID - 1
        F[7*(k-1)+1:7*k] =
        genExtF(x,j[i].RB2,u[:,k+1],GravityInInertial)

        if j[i].type == "Spring"
            F_b1, F_b2 = ForceConSpr(x,j[i]); # Forces acting on the bodies connected by the spring. Note that a spring does not constrain 2 bodies any more than they previously were, i.e., degrees of freedom remain the same.

            b1_id = j[i].RB1.bodyID -1; b2_id = j[i].RB2.bodyID - 1;
            if j[i].RB1.m == 0.0 # First body is the inertial frame
                F[7*(b2_id-1)+1:7*b2_id] += F_b2;
            else # First body is not the inertial frame
                F[7*(b1_id-1)+1:7*b1_id] += F_b1;
                F[7*(b2_id-1)+1:7*b2_id] += F_b2;
            end
        end
    end
    return F
end

function ForceCon(x::AbstractArray{T}, j::Joint)::Tuple{AbstractArray{T}, AbstractArray{T}} where T<:Real
    if(j.RB1.m == 0.0)
        # Joint to world. Body is connected to the Inertial frame.
        if j.type == "Revolute"
            A,b = ForceConRevIn(x,j)
            # function tempFn(x::AbstractArray{T}, j::Joint) where T<:Real
            #     Ar = ForceConRevIn(x,j)[1]
            #     Arp = pinv(Ar)
            #     return Arp
            # end
            # jb = ForwardDiff.jacobian(z -> tempFn([x[1:14];z],j),x[15:28]);
            # println("jb(Ar) = ", jb)#[:,15:28])
            # sleep(1000);
        elseif j.type == "Revolute2"
            A,b = ForceConRev2In(x,j)
        elseif j.type == "Spherical"
            A,b = ForceConSphIn(x,j)
            # function tempFn2(x::AbstractArray{T}, j::Joint) where T<:Real
            #     Ar = ForceConSphIn(x,j)[1]
            #     Arp = pinv(Ar)
            #     return Arp
            # end
            # jb = ForwardDiff.jacobian(z -> tempFn2([x[1:14];z],j),x[15:28]);
            # println("jb(Ar) = ", jb)#[:,15:28])
            # sleep(1000);
        elseif j.type == "Weld"
            A,b = ForceConWeldAllIn(x,j)
        elseif j.type == "Free" || j.type == "Spring"
            A,b = ForceFreeIn(x,j)
        else
            error("Joint type not prescribed.")
        end
    else
        if j.type == "Revolute"
            A,b = ForceConRev(x,j)
        elseif j.type == "Revolute2"
            A,b = ForceConRev2(x,j)
        elseif j.type == "Spherical"
            A,b = ForceConSph(x,j)
        elseif j.type == "Weld"
            A,b = ForceConWeldAll(x,j)
        elseif j.type == "Free" || j.type == "Spring"
            A,b = ForceFree(x,j)
        else
            error("Joint type not prescribed.")
        end
    end
    return A,b
end

# For inertial frame and body
function ForceFreeIn(x::AbstractArray{T}, j::Joint)::Tuple{AbstractArray{T}, AbstractArray{T}} where T<:Real
    # Constraint Force generated due to quaternion constraint
    b2 = j.RB2

    # Revolute Joint has 7 constraints
    # A = zeros(T,(1,14));b = T[0.0];
    A = spzeros(T,1,14);b = T[0.0];

    QcA, Qcb = QuatNormConstraint(x,j)
    A[1,:] = QcA[2,:]; b[1] = Qcb[2]
    return (A,b)

end

function ForceConRevIn(x::AbstractArray{T},j::Joint)::Tuple{AbstractArray{T}, AbstractArray{T}} where T<:Real#, extF2::extForces, GravityInInertial::Vector{Float64})::Vector{Float64}
    # Constraint Force generated by revolute joint
    b2 = j.RB2

    # Revolute Joint has 6 constraints (wrt inertial frame)
    # A = zeros(T,(6,14)); b = zeros(T,6);
    A = spzeros(T,6,14); b = spzeros(T,6);

    A[1:3,:],b[1:3] = TranslationConstraint(x,j)

    QcA, Qcb = QuatNormConstraint(x,j)
    A[4,:] = QcA[2,:]; b[4] = Qcb[2]

    A[5:end,:] = RevJointConstraint(x,j)[1]; b[5:end] = RevJointConstraint(x,j)[2]

    return (A,b)
end

function ForceConRev2In(x::AbstractArray{T}, j::Joint)::Tuple{AbstractArray{T}, AbstractArray{T}} where T<:Real
    # Constraint Force generated by 2revolute joint

    # Revolute Joint has 5 constraints (wrt inertial frame)
    # A = zeros(T,(5,14)); b = zeros(T,5);
    A = spzeros(T,5,14); b = spzeros(T,5);

    A[1:3,:],b[1:3] = TranslationConstraint(j);

    QcA, Qcb = QuatNormConstraint(x,j)
    A[4,:] = QcA[2,:]; b[4] = Qcb[2]

    A[5,:],b[5]  = RevJointConstraint(j);

    return (A,b)
end

function ForceConWeldAllIn(x::AbstractArray{T}, j::Joint)::Tuple{AbstractArray{T}, AbstractArray{T}} where T<:Real
    # Constraint Force generated by weld joint

    # Weld Joint has 7 constraints
    # A = zeros(T,(8,14)); b = zeros(T,8);
    A = spzeros(T,8,14); b = spzeros(T,8);

    A[1:3,:],b[1:3] = TranslationConstraint(j);

    QcA, Qcb = QuatNormConstraint(x,j);
    A[4,:] = QcA[2,:]; b[4] = Qcb[2];

    A[5:end,:], b[5:end] = WeldJointAllConstraint(j);

    return (A,b)
end

function ForceConSphIn(x::AbstractArray{T}, j::Joint)::Tuple{AbstractArray{T}, AbstractArray{T}} where T<:Real
    # Constraint Force generated by spherical joint

    # Spherical Joint leaves 3 degrees of freedom (remove 4 constraints)
    # A = zeros(T,(4,14)); b = zeros(T,4)
    A = spzeros(T,4,14); b = spzeros(T,4)

    A[1:3,:],b[1:3] = TranslationConstraint(x,j);

    QcA, Qcb = QuatNormConstraint(x,j)
    A[4,:] = QcA[2,:]; b[4] = Qcb[2]

    return (A,b)
end

# For 2 rigid bodies, no inertial frame
function ForceFree(x::AbstractArray{T}, j::Joint)::Tuple{AbstractArray{T}, AbstractArray{T}} where T<:Real
    # Constraint Force generated due to quaternion constraint
    A,b = QuatNormConstraint(j);
    return (A,b)
end

function ForceConRev(x::AbstractArray{T}, j::Joint)::Tuple{AbstractArray{T}, AbstractArray{T}} where T<:Real
    # Constraint Force generated by revolute joint
    # Revolute Joint imposes 7 constraints between two rigid bodies - 3 translational, 2 rotational, 2 quat norm
    # Resultant dof of the two bodies: 7
    # A = zeros(T, (7,14)); b = zeros(T,7);
    A = spzeros(T, 7,14); b = spzeros(T,7);

    A[1:3,:],b[1:3] = TranslationConstraint(x,j)
    A[4:5,:], b[4:5] = QuatNormConstraint(x,j)
    A[6:end,:],b[6:end] = RevJointConstraint(x,j)

    return (A,b)
end

function ForceConRev2(x::AbstractArray{T}, j::Joint)::Tuple{AbstractArray{T}, AbstractArray{T}} where T<:Real
    # Constraint Force generated by revolute2 joint
    # Revolute Joint2 has 6 constraints - 3 translational, 2 quat norm, 1 rotational
    # A = zeros(T,(6,14)); b = zeros(T,6);
    A = spzeros(T,6,14); b = spzeros(T,6);

    A[1:3,:], b[1:3] = TranslationConstraint(j);
    A[4:5,:], b[4:5] = QuatNormConstraint(j);
    A[6,:], b[6] = RevJoint2Constraint(j);

    return (A,b)
end

function ForceConWeldAll(x::AbstractArray{T}, j::Joint)::Tuple{AbstractArray{T}, AbstractArray{T}} where T<:Real
    # Constraint Force generated by weld joint
    # Weld Joint has 8 constraints
    # A = zeros(T,(8,14)); b = zeros(T,8);
    A = spzeros(T,8,14); b = spzeros(T,8);

    A[1:3,:],b[1:3] = TranslationConstraint(x,j)
    A[4:5,:], b[4:5] = QuatNormConstraint(x,j)
    A[6:end,:], b[6:end] = WeldJointAllConstraint(j)

    return (A,b)
end

function ForceConSpr(x::AbstractArray{T}, j::Joint) where T<:Real
    b1_id = j.RB1.bodyID; b2_id = j.RB2.bodyID;
    x1 = x[(b1_id-1)*14+1:b1_id*14]
    x2 = x[(b2_id-1)*14+1:b2_id*14]
    β1 = x1[4:7]; β2 = x2[4:7]
    b1_dcm = quat2dcm(β1); b2_dcm = quat2dcm(β2);

    # To compute Spring force and torque
    posSpr1 = x1[1:3] + transpose(b1_dcm)*j.pos1 # Inertial position of spring connection on first body
    posSpr2 = x2[1:3] + transpose(b2_dcm)*j.pos2 # Inertial position of spring connection on second body
    unitVec = (posSpr2 - posSpr1)/norm(posSpr2-posSpr1,2) # unit Vector in direction of spring (first body to second body)

    # First body
    E1 = genE(β1)
    F1 = j.k*(norm(posSpr1-posSpr2)-j.restLen)*unitVec # Force exerted by spring on the first body (inertial frame)
    F1 += j.c*((x2[8:10] - x1[8:10])); # Damper force (damper coefficient*(velocity difference))
    τ1 = cross(transpose(b1_dcm)*j.pos1,F1) # Torque exerted by spring-damper on the first body (inertial frame)
    Γb1 = [0.0;b1_dcm*τ1] #(body frame)
    Γu1 = 2*transpose(E1)*Γb1 # generalized torque vector

    # Second body
    E2 = genE(β2)
    F2 = -F1 # Force exerted by spring-damper on second body
    τ2 = cross(transpose(b2_dcm)*j.pos2,F2) # Torque exerted by spring on second body
    Γb2 = [0.0;b2_dcm*τ2]
    Γu2 = 2*transpose(E2)*Γb2

    F_b1 = [F1; Γu1]; F_b2 = [F2; Γu2];
    return (F_b1, F_b2)
end

function ForceConSph(x::AbstractArray{T}, j::Joint)::Tuple{AbstractArray{T}, AbstractArray{T}} where T<:Real
    # Constraint Force generated by spherical joint
    # Spherical Joint has 5 constraints (3 translational, 2 quat norm)
    # A = zeros(T,(5,14)); b = zeros(T,5)
    A = spzeros(T,5,14); b = spzeros(T,5)

    A[1:3,:], b[1:3] = TranslationConstraint(j);
    A[4:5,:],b[4:5] = QuatNormConstraint(j)

    return (A,b)
end

function TranslationConstraint(x::AbstractArray{T},j::Joint)::Tuple{AbstractArray{T}, AbstractArray{T}} where T<:Real
    b1_id = j.RB1.bodyID; b2_id = j.RB2.bodyID;
    x1 = x[(b1_id-1)*14+1:b1_id*14]
    x2 = x[(b2_id-1)*14+1:b2_id*14]
    b1 = j.RB1
    b2 = j.RB2
    rj1 = j.pos1
    rj2 = j.pos2

    # A = zeros(T,(3,14)); b = zeros(T,3)
    # A[:,1:3] = Matrix{Float64}(I,3,3)
    # A[:,8:10] = -Matrix{Float64}(I,3,3)

    A = spzeros(T,3,14); b = spzeros(T,3)
    A[:,1:3] = sparse(1.0I,3,3)
    A[:,8:10] = -sparse(1.0I,3,3)

    b1Tcs = TranslationConstraintSupplement(x1, rj1)
    b2Tcs = TranslationConstraintSupplement(x2, rj2)
    A[:,4:7] = b1Tcs[1]
    A[:,11:14] = -b2Tcs[1]
    b = -b1Tcs[2] + b2Tcs[2]

    return (A,b)
end

function TranslationConstraintSupplement(x::AbstractArray{T},pos::Vector{Float64}) where T <: Real
    # To constrain joint wrt body's cm
    β = x[4:7]
    βdot = x[11:14]

    r(y::AbstractArray{T}) where T<:Real = transpose(quat2dcm(y))*pos
    rJac = z->ForwardDiff.jacobian(r,z)

    rdot = y->ForwardDiff.jacobian(r,y)*βdot
    rβdot = y->ForwardDiff.jacobian(r,β)*y
    rddotRHS = ForwardDiff.jacobian(rdot,β)*βdot
    rddotLHS = rJac(β)
    return (rddotLHS,rddotRHS)
end

function QuatNormConstraint(x::AbstractArray{T},j::Joint)::Tuple{AbstractArray{T}, AbstractArray{T}} where T<:Real
    # A = zeros(T,(2,14))
    # b = zeros(T,2)

    A = spzeros(T,2,14)
    b = spzeros(T,2)

    b1 = j.RB1
    b2 = j.RB2

    b1_id = b1.bodyID;
    b2_id = b2.bodyID;
    β1 = x[14*(b1_id-1)+4:14*(b1_id-1)+7]
    β2 = x[14*(b2_id-1)+4:14*(b2_id-1)+7]
    β1dot = x[14*(b1_id-1)+11:14*(b1_id-1)+14]
    β2dot = x[14*(b2_id-1)+11:14*(b2_id-1)+14]

    A[1,4:7] = β1
    A[2,11:14] = β2
    b[1] = -transpose(β1dot)*β1dot
    b[2] = -transpose(β2dot)*β2dot
    return (A,b)
end

function RevJointConstraint(x::AbstractArray{T},j::Joint)::Tuple{AbstractArray{T}, AbstractArray{T}} where T<:Real
    # Attempting relative angular ω between bodies in body 1 frame only has component about the joint axis
    b1 = j.RB1
    b2 = j.RB2
    b1_id = b1.bodyID;
    b2_id = b2.bodyID;
    β1 = x[14*(b1_id-1)+4:14*(b1_id-1)+7]
    β2 = x[14*(b2_id-1)+4:14*(b2_id-1)+7]
    β1dot = x[14*(b1_id-1)+11:14*(b1_id-1)+14]
    β2dot = x[14*(b2_id-1)+11:14*(b2_id-1)+14]
    axis = j.axis
    βaug = [β1;β1dot;β2;β2dot]
    function f(y::AbstractArray{T}) where T <: Real
        y1  = y[1:4]; y1dot = y[5:8]#β1; β1dot
        y2 = y[9:12]; y2dot = y[13:16] #β2;β2dot
        ωb2 = angVel(y2,y2dot)
        ωb1 = angVel(y1,y1dot)
        ωb2inb1 = quat2dcm(y1)*transpose(quat2dcm(y2))*ωb2
        relAngVel = ωb2inb1 - ωb1
        return relAngVel
    end
    fJac = z->ForwardDiff.jacobian(f,z)

    # Ax = zeros(T,(3,14)); bx = zeros(T,3)
    Ax = spzeros(T,3,14); bx = spzeros(T,3)
    Ax[:,4:7] = fJac(βaug)[:,5:8]
    Ax[:,11:14] = fJac(βaug)[:,13:16]
    bx = -fJac(βaug)[:,1:4]*β1dot - fJac(βaug)[:,9:12]*β2dot

    zAxis = findall(iszero,axis)
    # A = zeros(T,(2,14)); b = zeros(T,2)
    A = spzeros(T,2,14); b = spzeros(T,2)
    A[1,:] = Ax[zAxis[1],:]
    A[2,:] = Ax[zAxis[2],:]
    b[1] = bx[zAxis[1]]
    b[2] = bx[zAxis[2]]

    return (A,b)

end

function RevJoint2Constraint(x::AbstractArray{T}, j::Joint)::Tuple{AbstractArray{T}, AbstractArray{T}} where T<:Real
    b1 = j.RB1
    b2 = j.RB2
    b1_id = b1.bodyID;
    b2_id = b2.bodyID;
    x1 = x[14*(b1_id-1)+1:14*(b1_id-1)+14];
    x2 = x[14*(b2_id-1)+1:14*(b2_id-1)+14];

    # A = zeros(T,(1,14));
    A = spzeros(T,1,14);
    b = T(0.0)
    y = RevJoint2ConstraintSupplement(x1,x2)
    A[:,4:7] = y[1][1]
    A[:,11:14] = y[1][2]
    b = y[2][1]
    return (A,b)
end

function RevJoint2ConstraintSupplement(x1::AbstractArray{T},x2::AbstractArray{T}) where T <: Real
    function f(y::AbstractArray{T}) where T <: Real
        # Computes the Hamilton product of β2 and β1^-1 (i.e. β2*β1^-1))
        # Let Q be the product.
        # (Q(1) + Q(2))^2 + (Q(0) + Q(3))^2 = 1.
        q1 = [y[1];-y[2:4]]
        q2 = y[5:8]
        Q = quaternionProduct(q2,q1)
        # Q = [q2[1]*q1[1] - transpose(q2[2:4])*q1[2:4];
        #      q2[1]*q1[2:4] + q1[1]*q2[2:4] + cross(q2[2:4],q1[2:4])]
        x = (Q[2] + Q[3])^2 + (Q[1] + Q[4])^2
        return x
    end
    fGrad = z->ForwardDiff.gradient(f,z)
    fdotJac = z->ForwardDiff.jacobian(fGrad,z)

    β1 = x1[4:7]
    β2 = x2[4:7]
    β1dot = x1[11:14]
    β2dot = x2[11:14]
    β = [β1;β2]

    y1 = f(β)
    y2 = fGrad(β)
    y3 = fdotJac(β)
    y3r = reshape(y3,(1,8,8))
    x1 = y3[1:4,1:4]*β1dot; x1 = reshape(x1,(1,4))
    x2 = y3[1:4,5:8]*β2dot; x2 = reshape(x2,(1,4))
    x3 = y3[5:end,1:4]*β1dot; x3 = reshape(x3,(1,4))
    x4 = y3[5:end,5:8]*β2dot; x4 = reshape(x4,(1,4))
    # β = rand(8); β1dot = rand(4); β2dot = rand(4);
    LHS = (reshape(-fGrad(β)[1:4],(1,4)),reshape(-fGrad(β)[5:8],(1,4)))
    RHS = (x1+x2)*β1dot + (x3+x4)*β2dot
    return (LHS,RHS)

end

function WeldJointAllConstraint(x::AbstractArray{T},j::Joint)::Tuple{AbstractArray{T}, AbstractArray{T}} where T<:Real
    # Attempting relative angular ω between bodies in body 1 frame is zero
    b1 = j.RB1
    b2 = j.RB2
    b1_id = b1.bodyID;
    b2_id = b2.bodyID;
    x1 = x[14*(b1_id-1)+1:14*(b1_id-1)+14];
    x2 = x[14*(b2_id-1)+1:14*(b2_id-1)+14];
    β1 = x1[4:7]
    β2 = x2[4:7]
    β1dot = x1[11:14]
    β2dot = x2[11:14]
    axis = j.axis
    βaug = [β1;β1dot;β2;β2dot]
    function f(y::AbstractArray{T}) where T <: Real
        y1  = y[1:4]; y1dot = y[5:8]#β1
        y2 = y[9:12]; y2dot = y[13:16] #β2;β2dot
        ωb2 = angVel(y2,y2dot)
        ωb1 = angVel(y1,y1dot)
        ωb2inb1 = quat2dcm(y1)*transpose(quat2dcm(y2))*ωb2
        relAngVel = ωb2inb1 - ωb1
        return relAngVel
    end
    fJac = z->ForwardDiff.jacobian(f,z)

    # Ax = zeros(T,(3,14)); bx = zeros(T,3)
    Ax = spzeros(T,3,14); bx = spzeros(T,3)
    Ax[:,4:7] = fJac(βaug)[:,5:8]
    Ax[:,11:14] = fJac(βaug)[:,13:16]
    bx = -fJac(βaug)[:,1:4]*β1dot - fJac(βaug)[:,9:12]*β2dot

    return (Ax,bx)
end

function genMatM(X::AbstractArray{T},b::RigidBody)::AbstractArray{T} where T<:Real
    # Function to generate mass matrix for each rigid body
    # T <: Real
    b_id = b.bodyID;
    β = X[14*(b_id-1)+4:14*(b_id-1)+7]
    E = genE(β)
    J = b.J
    M = Matrix{T}(undef,(7,7))
    M = [b.m*Matrix{T}(I,3,3)          zeros(T,(3,4))
                   zeros(T,(4,3)) 4*permutedims(E)*J*E]
    return sparse(M)
end

function genExtF(X::AbstractArray{T},b::RigidBody,U::AbstractArray{S},GravityInInertial::Vector{Float64}) where {T<:Real, S<:Real}#::AbstractArray{Union{T,S}}
    # Function to generate augmented external Force vector for unconstrained system
    # External Forces are always in the body frame
    b_id = b.bodyID;
    β = X[14*(b_id-1)+4:14*(b_id-1)+7]
    βdot = X[14*(b_id-1)+11:14*(b_id-1)+14]
    dcm = quat2dcm(β)
    # b.ω = angVel(β,βdot)
    E = genE(β)
    Edot = genE(βdot)
    TotalMoment = U[4:6]

    Γb = [0.0;TotalMoment] # In the body frame
    Γu = 2*transpose(E)*Γb

    F = Vector{Union{T,S}}(undef,7)
    F = [transpose(dcm)*U[1:3] + b.m*GravityInInertial
          Γu - 8*transpose(Edot)*b.J*E*βdot - 4*b.J[1,1]*(transpose(βdot)*βdot)*β]
    return F
    # F is a 7x1 vector in the inertial frame
end

function unique_ids(A::AbstractArray{T,2}) where T<:Real
    # Function to remove the redundant rows in constraint matrix A that repeat the quaternion norm constraint for a body connected by multiple joints

    A_unq = unique(A,dims=1); #after removing duplicate rows from A
    inds = Int64[]; # indices of non unique rows
    for i=1:size(A_unq,1)
        for j=1:size(A,1)
            if A_unq[i,:] == A[j,:]
                push!(inds,j)
                break;
            end
        end
    end
    return A_unq, inds
end

max_stars_repo_name: isrlab/Multibody-Dynamics

max_issues_repo_name: isrlab/Multibody-Dynamics

max_forks_repo_name: isrlab/Multibody-Dynamics


text:
using OBC, MPM
using HDF5, JLD
using PyPlot
using Distributions
using Stats

loc = "/home/bana/largeresearch/rnaseq.jld"
if !isfile(loc)
    loc = "/home/bana/largeresearch/seq-data/rnaseq.jld"
end
@load loc luad_norm luad_norm_sampids luad_norm_geneids

d = 10.0
kmax = 1
iters = 15000
burn = 2000
thin = 50
D = 2
rounds = 1
numsamps = 100

Ts = [x->quantile(x,[0.05])[1], x->quantile(x,[0.5])[1], x->quantile(x,[0.95])[1], iqr, var]
Tdesc = ["5th quantile", "Median", "95th quantile", "IQR", "Variance"]

pvals = zeros(D*rounds, length(Ts))
true_ts = similar(pvals)
tvals = zeros(D*rounds, length(Ts), int((iters-burn)/thin + 1 ))

samps = sample(1:size(luad_norm,2), numsamps)
genes = sample(1:size(luad_norm,1), D*rounds)

######################################################################
# Fit Model
######################################################################
cov = eye(D,D)
#cov = [0.5 -0.2; -0.2 1.0]
mu = ones(D,kmax)
#mu = [1.6 -1.0; 1.6 -1.0]
#mu = [0.2 0.2]'
prior = MPM.MPMPrior(D=2, kmax=kmax, kappa=10., S=eye(2)*10)

pmoves = MPM.MPMPropMoves()
pmoves.lammove = 0.005
pmoves.mumove = 0.08
pmoves.priorkappa = 100.0
global data, mymh_a

for i=1:rounds

    #pick two random genes
    data = luad_norm[genes[(i-1)*D+1:i*D],samps]'

    start = MPM.MPMParams(mu, #mu :: Matrix{Float64}
        cov, #sigma :: Matrix{Float64}
        ones(kmax)/kmax, #w :: Vector{Float64}
        clamp(log(data'/d),-10.0,Inf), #lam :: Matrix{Float64}
        1) #k :: Int

    obj_a = MPM.MPMSampler(prior, data, deepcopy(start), pmoves, d)
    obj_a.usepriors = true

    mymh_a = MPM.MHRecord(obj_a,burn=burn,thin=thin)
    @time sample(mymh_a,iters)
    @time temp1,temp2,temp3 = MPM.calc_pvals(Ts, data', mymh_a.db, d)
    pvals[(i-1)*D+1:i*D, :] = temp1
    true_ts[(i-1)*D+1:i*D, :] = temp2
    tvals[(i-1)*D+1:i*D, :, :] = temp3
end

#assert(false)
######################################################################
# Plotting
######################################################################

include("src/plot_utils.jl")
n1, n2, gext, grid = get_grid(data)
#n1, n2, gext, grid = gen_grid([0,10.0,0,10])

close("all")

#@time ga,lams = MPM.calc_g(grid, mymh_a.db, 20, dmean=d)
#imshow(reshape(ga,n1,n2)', extent=gext, origin="lower")
#imshow(reshape(ga,n1,n2)', extent=gext, aspect="equal", origin="lower")
#colorbar()
#contour(reshape(ga,n1,n2)', extent=gext, aspect="equal", origin="lower")
fig = figure()
ax1 = fig[:add_subplot](211)
plot(data[:,1], data[:,2], "g.", alpha=0.5)
plt.grid(true)

title("TCGA Data")
newpts = MPM.gen_posterior_points(size(data,1), d, mymh_a.db)
newpts = newpts[:,sample(1:size(newpts,2), size(data,1))]
fig[:add_subplot](212, sharex=ax1, sharey=ax1)
title("Replicated data from predictive posterior")
plt.grid(true)
plot(newpts[1,:], newpts[2,:], "r.", alpha=0.5)

xlabel(string(luad_norm_geneids[genes[1]], " counts"))
#ylabel(string(luad_norm_geneids[genes[2]], " counts"))
fig[:text](0.06, 0.5, string(luad_norm_geneids[genes[2]], " counts"), ha="center", va="center", rotation="vertical")

function plot_hist(geneid, tid)
    figure()
    plt.hist(vec(tvals[geneid,tid,:]), bins=30)
    axvline(true_ts[geneid,tid], color="k", linewidth=4)
end

#function plot_pval(q)
    #p,trut,tvals = MPM.calc_pval(y->quantile(vec(y[1,:]),[q])[1], data', mymh_a.db, d)
    #figure()
    #plt.hist(tvals, bins=30)
    #axvline(trut, color="k", linewidth=2)
    #@show q,p
#end

#plot_pval(0.05)
#plot_pval(0.5)
#plot_pval(0.95)

#g() = plot(tst_data[:,1]+rand(size(tst_data,1)), tst_data[:,2]+rand(size(tst_data,1)), "g.", alpha=0.3)
#fa() = plot_traces(mymh_a.db, [:mu,:sigma,:lam,:energy])

#for i in 1:size(all,1)
       #for j in 1:size(all,2)
       #if j==1
       #print(all[i,j], " & ")
       #else
       #@printf "%.2f %s" all[i,j] " & "
       #end
       #end
       #println(" \\\\")
       #end


text:
# ---
# title: 540. Single Element in a Sorted Array
# id: problem540
# author: Indigo
# date: 2021-06-23
# difficulty: Medium
# categories: Binary Search
# link: <https://leetcode.com/problems/single-element-in-a-sorted-array/description/>
# hidden: true
# ---
#
# You are given a sorted array consisting of only integers where every element
# appears exactly twice, except for one element which appears exactly once. Find
# this single element that appears only once.
#
# **Follow up:** Your solution should run in O(log n) time and O(1) space.
#
#
#
# **Example 1:**
#
#
#
#     Input: nums = [1,1,2,3,3,4,4,8,8]
#     Output: 2
#
#
# **Example 2:**
#
#
#
#     Input: nums = [3,3,7,7,10,11,11]
#     Output: 10
#
#
#
#
# **Constraints:**
#
#   * `1 <= nums.length <= 10^5`
#   * `0 <= nums[i] <= 10^5`
#
#
## @lc code=start
using LeetCode

function single_non_duplicate(nums::Vector{Int})
    i, j = 1, length(nums)
    while i < j
        mid = (i + j) >> 1
        if isodd(mid)
            nums[mid] == nums[mid + 1] ? (i = mid + 2) : (j = mid)
        else
            nums[mid] == nums[mid + 1] ? (j = mid - 1) : (i = mid + 1)
        end
    end
    return nums[i]
end
## @lc code=end


text:
# Important note: even if one or more tests fail, we will still exit with status code 0.

# The reason for this is that we always want to upload code coverage, even if some of the
# tests fail. Therefore, even if the `coverage_linux64` builder passes, you should not
# assume that all of the tests passed. If you want to know if all of the tests are passing,
# please look at the status of the `tester_*` builders (e.g. `tester_linux64`).

# When running this file, make sure to set all of the following command-line flags:
# 1. `--code-coverage=all`
# 2. `--sysimage-native-code=no`

empty!(Base.DEPOT_PATH)
push!(Base.DEPOT_PATH, mktempdir(; cleanup = true))

const tests = "all"
const ncores = Sys.CPU_THREADS

@info "" Sys.CPU_THREADS
@info "" tests ncores

try
    Base.runtests(tests; ncores)
catch ex
    @error "" exception=(ex, catch_backtrace())
end

max_stars_repo_name: RelationalAI-oss/julia


text:
function group{T}(::Type{T}, atoms::Atoms, prop::Symbol)
    props = Dict{T, Atoms}()

    for i in atoms.indices
        k = (atoms.props[prop]::Array{T})[i]

        if !haskey(props, k)
            props[k] = Atoms(Int[], atoms.props)
        end

        push!(props[k].indices, i)
    end

    return props
end

function group(atoms::Atoms, prop::Symbol)
    T = eltype(atoms.props[prop])
    group(T, atoms, prop)
end

box_dims(f::Frame) = ntuple(i -> f[:box_max][i] - f[:box_min][i], 3)

@inline function unscale{T}(bl::T, bu::T, bw::T, coord::T)
    coord = bw * coord + bl

    if coord >= bu
        coord -= bw
    end

    if coord < bl
        coord = clamp(coord, bl, bu)
    end

    return coord
end

function unscale!{T, U <: AbstractMatrix{T}}(bl::Vector{T}, bu::Vector{T}, bw::Vector{T}, coords::U)
    ncols = size(coords, 2)

    for c in 1:ncols, r in 1:3
        @inbounds coords[r, c] = unscale(bl[r], bu[r], bw[r], coords[r, c])
    end

    return coords
end

unscale!(frame::Frame) = (unscale!(box_lower(frame), box_upper(frame), box_widths(frame), frame[:atoms][:coords]) ; frame)

function unwrap!{T}(atoms::Atoms, box_width::T, max_sep::T = 3.0)
    sep_sq = max_sep ^ 2
    bw = box_width
    hbw = bw / 2

    coords = atoms[:coords]

    function distance_sq(i::Int, j::Int)
        x = coords[1, i] - coords[1, j]
        y = coords[2, i] - coords[2, j]
        z = coords[3, i] - coords[3, j]

        x^2 + y^2 + z^2
    end

    num_atoms = length(atoms)

    b = fld(num_atoms, 2)
    path = Int[b]
    inpath = fill(false, num_atoms)
    inpath[b] = true

    prev_path_length = 1
    while true
        for i in 1:num_atoms
            inpath[i] && continue

            for j in path
                if distance_sq(i, j) < sep_sq
                    push!(path, i)
                    inpath[i] = true
                    break
                end
            end
        end

        length(path) == prev_path_length && break
        prev_path_length = length(path)
    end

    while true
        for i in 1:num_atoms
            inpath[i] && continue

            for j in path
                for k in 1:3
                    d = coords[k, i] - coords[k, j]
                    if d > hbw
                        coords[k, i] -= bw
                    elseif d < -hbw
                        coords[k, i] += bw
                    end
                end

                if distance_sq(i, j) < sep_sq
                    push!(path, i)
                    inpath[i] = true
                    break
                end
            end
        end

        length(path) == prev_path_length && break
        prev_path_length = length(path)
    end

    all(inpath) || error("Unable to correctly unwrap all atoms in fragment")
end

max_stars_repo_name: NickMcNutt/Alpine

max_issues_repo_name: NickMcNutt/Alpine

max_forks_repo_name: NickMcNutt/Alpine


text:
using Test
using IncrementalInference

@testset "Skip downsolve" begin
N=6
fg = generateCanonicalFG_lineStep(N;
                                  graphinit=false,
                                  poseEvery=1,
                                  landmarkEvery=N+1,
                                  posePriorsAt=[0],
                                  landmarkPriorsAt=[],
                                  sightDistance=N+1)

deleteFactor!.(fg, [Symbol("x$(i)lm0f1") for i=1:(N-1)])

getSolverParams(fg).useMsgLikelihoods = true
getSolverParams(fg).downsolve = false

smtasks = Task[]
tree, smt, hists = solveTree!(fg; smtasks=smtasks, recordcliqs=[:x4]);

# See if downsolve was called
@test !(IIF.solveDown_StateMachine in getindex.(hists[2], 3))

#test if values are still correct
for var in sortDFG(ls(fg))
    sppe = getVariable(fg,var) |> getPPE |> IIF.getPPESuggested
    @test isapprox(sppe[1], parse(Int,string(var)[end]), atol=0.2)
end

#now downsolve only
getSolverParams(fg).upsolve = false
getSolverParams(fg).downsolve = true

tree, smt, hists = solveTree!(fg; smtasks=smtasks, recordcliqs=[:x4]);
# See if upsolved was called
@test !(IIF.solveUp_StateMachine in getindex.(hists[2], 3))

#test if values are still correct
for var in sortDFG(ls(fg))
    sppe = getVariable(fg,var) |> getPPE |> IIF.getPPESuggested
    @test isapprox(sppe[1], parse(Int,string(var)[end]), atol=0.2)
end

# recycled downsolve only
tree, smt, hists = solveTree!(fg, tree; smtasks=smtasks, recordcliqs=[:x4]);
@test !(IIF.solveUp_StateMachine in getindex.(hists[2], 3))

end


text:
struct UniformFluxDiffDisc1D{BalanceLaw, T, Basis, Fvol, Fnumint, Fnumext, LeftBC, RightBC,
                                Fluxes, Parallel} <: AbstractSemidiscretisation
    balance_law::BalanceLaw
    meshx::UniformMesh1D{T}
    basis::Basis
    fvol::Fvol
    fnumint::Fnumint
    fnumext::Fnumext
    left_bc::LeftBC
    right_bc::RightBC

    fluxes::Fluxes
    parallel::Parallel

    function UniformFluxDiffDisc1D(balance_law::BalanceLaw, meshx::UniformMesh1D{T},
                                    basis::Basis, fvol::Fvol, fnumint::Fnumint,
                                    fnumext::Fnumext, left_bc::LeftBC, right_bc::RightBC,
                                    fluxes::Fluxes, parallel::Parallel) where
                                        {BalanceLaw, T, Basis, Fvol, Fnumint, Fnumext,
                                        LeftBC, RightBC, Fluxes, Parallel}
        @assert typeof(fluxes) <: AbstractArray{variables(balance_law), 1}
        @assert size(fluxes, 1) == numedges(meshx)
        new{BalanceLaw, T, Basis, Fvol, Fnumint, Fnumext, LeftBC, RightBC, Fluxes, Parallel}(
            balance_law, meshx, basis, fvol, fnumint, fnumext, left_bc, right_bc,
            fluxes, parallel)
    end
end

function UniformFluxDiffDisc1D(balance_law, meshx, basis, fvol, fnumint, fnumext,
                                left_bc, right_bc, parallel=Val{:serial}())
    fluxes = Array{variables(balance_law)}(undef, numedges(meshx))
    UniformFluxDiffDisc1D(balance_law, meshx, basis, fvol, fnumint, fnumext,
                            left_bc, right_bc, fluxes, parallel)
end

@noinline function (semidisc::UniformFluxDiffDisc1D)(du, u, p, t)
  @boundscheck begin
    if size(u) != size(du)
      error("size(u) = $(size(u)) != $(size(du)) = size(du)")
    end
    @assert size(u,1) == length(semidisc.basis.nodes)
    @assert size(u,2) == numcells(semidisc.meshx)
    if eltype(u) != variables(semidisc.balance_law)
      error("eltype(u) == $(eltype(u)) != $(variables(semidisc.balance_law)) == variables(semidisc.balance_law)")
    end
  end

  fill!(du, zero(eltype(du)))
  add_flux_differences!(du, u, semidisc)
  add_numerical_fluxes!(du, u, semidisc, t)

  nothing
end

struct UniformPeriodicFluxDiffDisc1D{BalanceLaw, T, Basis, Fvol, Fnumint,
                                    Fluxes, Parallel} <: AbstractSemidiscretisation
    balance_law::BalanceLaw
    meshx::UniformPeriodicMesh1D{T}
    basis::Basis
    fvol::Fvol
    fnumint::Fnumint

    fluxes::Fluxes
    parallel::Parallel

    function UniformPeriodicFluxDiffDisc1D(balance_law::BalanceLaw, meshx::UniformPeriodicMesh1D{T},
                                            basis::Basis, fvol::Fvol, fnumint::Fnumint,
                                            fluxes::Fluxes, parallel::Parallel) where
                                            {BalanceLaw, T, Basis, Fvol, Fnumint,
                                            Fluxes, Parallel}
        @assert typeof(fluxes) <: AbstractArray{variables(balance_law), 1}
        @assert size(fluxes, 1) == numedges(meshx)
        new{BalanceLaw, T, Basis, Fvol, Fnumint, Fluxes, Parallel}(
            balance_law, meshx, basis, fvol, fnumint, fluxes, parallel)
    end
end

function UniformPeriodicFluxDiffDisc1D(balance_law, meshx, basis, fvol, fnumint,
                                        parallel=Val{:serial}())
    fluxes = Array{variables(balance_law)}(undef, numedges(meshx))
    UniformPeriodicFluxDiffDisc1D(balance_law, meshx, basis, fvol, fnumint,
                                    fluxes, parallel)
end

@noinline function (semidisc::UniformPeriodicFluxDiffDisc1D)(du, u, p, t)
  @boundscheck begin
    if size(u) != size(du)
      error("size(u) = $(size(u)) != $(size(du)) = size(du)")
    end
    @assert size(u,1) == length(semidisc.basis.nodes)
    @assert size(u,2) == numcells(semidisc.meshx)
    if eltype(u) != variables(semidisc.balance_law)
      error("eltype(u) == $(eltype(u)) != $(variables(semidisc.balance_law)) == variables(semidisc.balance_law)")
    end
  end

  fill!(du, zero(eltype(du)))
  add_flux_differences!(du, u, semidisc)
  add_numerical_fluxes!(du, u, semidisc, t)

  nothing
end

function evaluate_coefficients(u, semidisc::Union{UniformFluxDiffDisc1D,UniformPeriodicFluxDiffDisc1D})
    evaluate_coefficients(u, semidisc.meshx, semidisc.basis)
end

function add_flux_differences!(du, u, semidisc::Union{UniformFluxDiffDisc1D,UniformPeriodicFluxDiffDisc1D})
    Nx  = size(u, 2)
    Pp1 = size(u, 1)

    @unpack balance_law, fvol, basis, parallel = semidisc
    @unpack D = basis
    @boundscheck @assert Pp1 == size(D,1) == size(D,2)
    jacx = 2 / semidisc.meshx.Δx

    if Pp1 <= 1
        return nothing
    end

    add_flux_differences_inner_loop!(du, u, balance_law, fvol, Nx, Pp1, D, jacx, parallel)
end

@inline function add_flux_differences_inner_loop!(du, u, balance_law, fvol, Nx, Pp1,
                                                    D, jacx, parallel)
    dims = (Pp1, Nx)

    @inbounds for ix in Base.OneTo(Nx)
        for nx in Base.OneTo(Pp1)
            idx = LinearIndices(dims)[nx, ix]
            u_idx = u[idx]
            # compute x derivative
            # at first for different indices
            for k in (nx+1):Pp1
                idxk = LinearIndices(dims)[k, ix]
                f = fvol(u_idx, u[idxk], balance_law)
                du[idx]  -= 2*jacx*D[nx,k] * f
                du[idxk] -= 2*jacx*D[k,nx] * f
            end
            # then for the diagonal element, using the consistency of the flux
            du[idx] -= 2*jacx*D[nx,nx] * flux(u_idx, balance_law)
        end
    end

    nothing
end
@inline function add_flux_differences_inner_loop!(du, u, balance_law, fvol, Nx, Pp1,
                                                    D, jacx, parallel::Val{:threads})
    dims = (Pp1, Nx)

     @inbounds Threads.@threads for ix in Base.OneTo(Nx)
        for nx in Base.OneTo(Pp1)
            idx = LinearIndices(dims)[nx, ix]
            u_idx = u[idx]
            # compute x derivative
            # at first for different indices
            for k in (nx+1):Pp1
                idxk = LinearIndices(dims)[k, ix]
                f = fvol(u_idx, u[idxk], balance_law)
                du[idx]  -= 2*jacx*D[nx,k] * f
                du[idxk] -= 2*jacx*D[k,nx] * f
            end
            # then for the diagonal element, using the consistency of the flux
            du[idx] -= 2*jacx*D[nx,nx] * flux(u_idx, balance_law)
        end
    end

    nothing
end

function add_numerical_fluxes!(du, u, semidisc::UniformFluxDiffDisc1D, t)
    Nx  = size(u, 2)
    Pp1 = size(u, 1)

    @unpack balance_law, basis, fnumint, fnumext, fluxes, parallel, left_bc, right_bc = semidisc
    @boundscheck @assert Pp1 == length(basis.weights)

    jacx = 2 / semidisc.meshx.Δx

    add_numerical_fluxes_inner_loop1!(du, fluxes, u, t, balance_law, fnumint,
                                        fnumext, left_bc, right_bc, Nx, basis,
                                        includes_boundaries(basis), jacx, parallel)
    add_numerical_fluxes_inner_loop2!(du, fluxes, u, balance_law, Nx, basis,
                                        includes_boundaries(basis), jacx, parallel)
end

@inline function add_numerical_fluxes_inner_loop1!(du, fluxes, u, t, balance_law,
                                                    fnumint, fnumext, left_bc, right_bc,
                                                    Nx, basis, boundaries_included::Val{true},
                                                    jacx, parallel)
    # calculate external numerical fluxes
    @inbounds fluxes[1] = fnumext(left_bc(t), u[1,1], balance_law)
    @inbounds fluxes[end] = fnumext(u[end,end], right_bc(t), balance_law)

    # calculate internal numerical fluxes
    @inbounds for ix in 2:Nx
        # flux x - left
        fluxes[ix] = fnumint(u[end,ix-1], u[1,ix], balance_law)
    end

    nothing
end
@inline function add_numerical_fluxes_inner_loop1!(du, fluxes, u, t, balance_law,
                                                    fnumint, fnumext, left_bc, right_bc,
                                                    Nx, basis, boundaries_included::Val{true},
                                                    jacx, parallel::Val{:threads})
    # calculate external numerical fluxes
    @inbounds fluxes[1] = fnumext(left_bc(t), u[1,1], balance_law)
    @inbounds fluxes[end] = fnumext(u[end,end], right_bc(t), balance_law)

    # calculate internal numerical fluxes
    @inbounds Threads.@threads for ix in 2:Nx
        # flux x - left
        fluxes[ix] = fnumint(u[end,ix-1], u[1,ix], balance_law)
    end

    nothing
end
@inline function add_numerical_fluxes_inner_loop1!(du, fluxes, u, t, balance_law,
                                                    fnumint, fnumext, left_bc, right_bc,
                                                    Nx, basis, boundaries_included::Val{false},
                                                    jacx, parallel)
    # calculate external numerical fluxes
    @views @inbounds fluxes[1] = fnumext(left_bc(t), interpolate(-1, u[:,1], basis), balance_law)
    @views @inbounds fluxes[end] = fnumext(interpolate(1, u[:,end], basis), right_bc(t), balance_law)

    # calculate internal numerical fluxes
    @views @inbounds for ix in 2:Nx
        # flux x - left
        fluxes[ix] = fnumint(interpolate(+1, u[:,ix-1], basis),
                             interpolate(-1, u[:,ix], basis),
                             balance_law)
    end

    nothing
end

@inline function add_numerical_fluxes_inner_loop2!(du, fluxes, u, balance_law,
                                                    Nx, basis, boundaries_included::Val{true},
                                                    jacx, parallel)
    @inbounds i_ω1 = 1 / basis.weights[1]
    @inbounds i_ωend = 1 / basis.weights[end]
    Pp1 = length(basis.nodes)
    dims = (Pp1, Nx)

    # add numerical fluxes
    @inbounds for ix in Base.OneTo(Nx)
        # flux x - left
        idx = LinearIndices(dims)[1, ix]
        du[idx] += ( fluxes[ix] - flux(u[idx], balance_law) ) * jacx * i_ω1

        # flux x - right
        idx = LinearIndices(dims)[Pp1, ix]
        du[idx] -= ( fluxes[ix+1] - flux(u[idx], balance_law) ) * jacx * i_ωend
    end

    nothing
end
@inline function add_numerical_fluxes_inner_loop2!(du, fluxes, u, balance_law,
                                                    Nx, basis, boundaries_included::Val{true},
                                                    jacx, parallel::Val{:threads})
    @inbounds i_ω1 = 1 / basis.weights[1]
    @inbounds i_ωend = 1 / basis.weights[end]
    Pp1 = length(basis.nodes)
    dims = (Pp1, Nx)

    # add numerical fluxes
    @inbounds Threads.@threads for ix in Base.OneTo(Nx)
        # flux x - left
        idx = LinearIndices(dims)[1, ix]
        du[idx] += ( fluxes[ix] - flux(u[idx], balance_law) ) * jacx * i_ω1

        # flux x - right
        idx = LinearIndices(dims)[Pp1, ix]
        du[idx] -= ( fluxes[ix+1] - flux(u[idx], balance_law) ) * jacx * i_ωend
    end

    nothing
end
# NOTE: add_numerical_fluxes_inner_loop2! has to be specialised for each balance_law
# if no boundary nodes are included, e.g. for GaussLegendre bases.

function add_numerical_fluxes!(du, u, semidisc::UniformPeriodicFluxDiffDisc1D, t)
    Nx  = size(u, 2)
    Pp1 = size(u, 1)

    @unpack balance_law, basis, fnumint, fluxes, parallel = semidisc
    @boundscheck @assert Pp1 == length(basis.weights)

    jacx = 2 / semidisc.meshx.Δx

    add_numerical_fluxes_inner_loop1!(du, fluxes, u, t, balance_law, fnumint,
                                        Nx, basis, includes_boundaries(basis),
                                        jacx, parallel)
    add_numerical_fluxes_inner_loop2!(du, fluxes, u, balance_law, Nx, basis,
                                        includes_boundaries(basis),jacx, parallel)
end

@inline function add_numerical_fluxes_inner_loop1!(du, fluxes, u, t, balance_law,
                                                    fnumint, Nx, basis, boundaries_included::Val{true},
                                                    jacx, parallel)
    # calculate boundary numerical fluxes
    @inbounds fluxes[1] = fluxes[end] = fnumint(u[end,end], u[1,1], balance_law)

    # calculate internal numerical fluxes
    @inbounds for ix in 2:Nx
        # flux x - left
        fluxes[ix] = fnumint(u[end,ix-1], u[1,ix], balance_law)
    end

    nothing
end
@inline function add_numerical_fluxes_inner_loop1!(du, fluxes, u, t, balance_law,
                                                    fnumint, Nx, basis, boundaries_included::Val{true},
                                                    jacx, parallel::Val{:threads})
    # calculate boundary numerical fluxes
    @inbounds fluxes[1] = fluxes[end] = fnumint(u[end,end], u[1,1], balance_law)

    # calculate internal numerical fluxes
    @inbounds Threads.@threads for ix in 2:Nx
        # flux x - left
        fluxes[ix] = fnumint(u[end,ix-1], u[1,ix], balance_law)
    end

    nothing
end
@inline function add_numerical_fluxes_inner_loop1!(du, fluxes, u, t, balance_law,
                                                    fnumint, Nx, basis, boundaries_included::Val{false},
                                                    jacx, parallel)
    # calculate external numerical fluxes
    @views @inbounds fluxes[1] = fluxes[end] = fnumint(interpolate(+1, u[:,end], basis),
                                                       interpolate(-1, u[:,1], basis),
                                                       balance_law)

    # calculate internal numerical fluxes
    @views @inbounds for ix in 2:Nx
        # flux x - left
        fluxes[ix] = fnumint(interpolate(+1, u[:,ix-1], basis),
                             interpolate(-1, u[:,ix], basis),
                             balance_law)
    end

    nothing
end

function semidiscretise(semidisc::Union{UniformFluxDiffDisc1D, UniformPeriodicFluxDiffDisc1D},
                        u₀func, tspan)
    @unpack meshx, basis, parallel = semidisc
    u₀ = compute_coefficients(u₀func, meshx, basis, parallel)

    ODEProblem(semidisc, u₀, tspan)
end

function integrate(func, u, semidisc::Union{UniformFluxDiffDisc1D, UniformPeriodicFluxDiffDisc1D})
    integrate(func, u, semidisc.meshx, semidisc.basis)
end

"""
    max_dt(t, u, semidisc::UniformFluxDiffDisc1D, cfl)

Compute the maximal time step `dt` satisfying the CFL condition `dt <= cfl * dx / ((2p+1)*λ)`,
where `dx` is the length of a cell, `λ` the greatest absolute value of the
speed in this cell, and `p` the polynomial degree.
"""
function max_dt(t, u, semidisc::Union{UniformFluxDiffDisc1D, UniformPeriodicFluxDiffDisc1D}, cfl=0.5)
    @unpack balance_law, meshx, basis = semidisc
    @unpack Δx = meshx
    p = length(basis.nodes)-1
    factor = Δx / (2p+1)

    dt = mapreduce(u->factor/max_abs_speed(u, balance_law), min, u, init=typemax(t))
    if dt == Inf
        dt = factor
    end
    cfl * dt
end