text:
@testset "Save and Load" begin
    @testset "Basic ops" begin
        args = (rand(3, 4), rand(3, 4))
        ort_test(ONNX.add, args...)
        ort_test(ONNX.mul, args...)
    end

    @testset "Gemm" begin
        A, B, C = (rand(3, 4), rand(3, 4), rand(3, 3))
        ort_test(ONNX.onnx_gemm, A, B')
        ort_test(ONNX.onnx_gemm, A', B)
        ort_test(ONNX.onnx_gemm, A', B, C)
        ort_test(ONNX.onnx_gemm, A, B, C; tA=1)
        ort_test(ONNX.onnx_gemm, A, B; tB=1)
        ort_test(ONNX.onnx_gemm, A', B; α=2.0)
        ort_test(ONNX.onnx_gemm, A', B, C; α=2.0, β=0.5)
        # make sure Gemm with just 2 matrices and no keyword arguments
        # is recorded as just *
        before, after = ort_test(*, A', B)
        @test before[V(3)].fn == after[V(3)].fn
        @test before[V(3)].fn == *
    end

    @testset "Conv" begin
        # 2D, keywords
        args = (rand(Float32, 32, 32, 3, 1), rand(Float32, 3, 3, 3, 6))
        ort_test(ONNX.conv, args...)
        ort_test(ONNX.conv, args...; pad=1, stride=(1, 1), dilation=(1, 1), groups=1)
        ort_test(ONNX.conv, args...; pad=1, stride=(1, 2), dilation=(2, 1), groups=1)
        ort_test(ONNX.conv, args...; stride=1, dilation=1)
        ort_test(ONNX.conv, args...; pad=(1, 2))
        ort_test(ONNX.conv, args...; pad=(1, 2, 3, 4))

        # 2D, with bias
        ort_test(ONNX.conv, args..., rand(Float32, 6))
        ort_test(ONNX.conv, args..., rand(Float32, 6); pad=(1, 1))

        # 2D, non-square kernel
        args = (rand(Float32, 32, 32, 3, 1), rand(Float32, 5, 3, 3, 6))
        ort_test(ONNX.conv, args...)

        # 1D
        args = (rand(Float32, 32, 3, 1), rand(Float32, 3, 3, 6))
        ort_test(ONNX.conv, args...)
        ort_test(ONNX.conv, args...; pad=(1, 2))

        # 3D
        args = (rand(Float32, 32, 32, 32, 3, 1), rand(Float32, 3, 3, 3, 3, 6))
        ort_test(ONNX.conv, args...)
        ort_test(ONNX.conv, args...; pad=(1, 2, 3))
    end

    @testset "Pooling" begin
        x = rand(Float32, 32, 32, 3, 1)
        k = (2, 2)
        ort_test(ONNX.maxpool, x; kernel=k)
        ort_test(ONNX.maxpool, x; kernel=k, stride=(3, 3))
        ort_test(ONNX.maxpool, x; kernel=k, stride=(3, 3), pad=1)
    end

    @testset "Activations" begin
        x = rand(3, 4)
        ort_test(ONNX.relu, x)
    end

end


text:
module Lerche

using DataStructures  #For Stack when parsing
using Logging         #Because Lark does

export Tree, Token, Interpreter,Visitor,Transformer, visit_children, visit,transform
export Transformer_InPlace, Transformer_InPlaceRecursive,Visitor_Recursive
export GrammarError, ParseError, UnexpectedToken, UnexpectedCharacters
export UnexpectedInput
export Lark
export visit_tokens   #So users can override default trait for Transformers
export @inline_rule, @rule
#==

Lerche ("Lark" in German) is a port of the Python Lark package to
native Julia. No I don't speak German.

A number of simplifications have been made:
(1) There is no choice of Tree class. Lark has SlottedTree and Tree,
with the former preferred for speed. This is not a consideration for
Julia.
(2) The individual parsers in the parser modules are given types instead
of using the separate module approach of Python.

==#

include("utils.jl")
include("tree.jl")
include("visitors.jl")
include("exceptions.jl")
include("grammar.jl")
include("lexer.jl")
include("parse_tree_builder.jl")
include("parsers/grammar_analysis.jl")
include("parsers/lalr_analysis.jl")
include("parsers/lalr_parser.jl")
include("parser_frontends.jl")
include("load_grammar.jl")
include("lark.jl")

# Compatibility for Julia < 1.4

try
    TaskFailedException
catch ex
    struct TaskFailedException <: Exception
    end
end

include("precompile.jl")

@warnpcfail _precompile_()

const _lark_grammar_f = GrammarLoader(false)
const _lark_grammar_t = GrammarLoader(true)

#load_grammar(text::String;options...) = load_grammar(_lark_grammar,text;options...)

end # module

ext: jl


text:
export graph_bbcs_cheb_exp

"""
    (graph,cref)=graph_bbcs_cheb_exp(k;T=Complex{BigFloat})

Computes a polynomial evaluation approximating the exponential for
skew-Hermitian matrices using `k` matrix multiplications following the procedure
in the reference.

For `k > 5` it resorts to scaling-and-squaring of the `k = 5` graph.
The graph is in the degopt format, see [`graph_degopt`](@ref).

**Reference**

1. P. Bader, S. Blanes, F. Casas, M. Seydaoglu. "An efficient algorithm to
   compute the exponential of skew-Hermitian matrices for the time integration
   of the Schrödinger equation".
   [arXiv:2103.10132 [math.NA]](https://arxiv.org/abs/2103.10132), 2021.
"""
function graph_bbcs_cheb_exp(k; T = Complex{BigFloat})
    CBF = Complex{BigFloat}
    if (k <= 5) # Basic implementation from paper
        if (k == 1)
            x1a = vcat(zero(CBF), one(CBF))
            x1b = vcat(zero(CBF), one(CBF))
            xv = [(x1a, x1b)]
            y = [
                BigFloat("0.9999999999999999999998"),
                BigFloat("-0.9999999999761950000001")im,
                BigFloat("-0.4999999999920650000000"),
            ]
        elseif (k == 2)
            x1a = vcat(zero(CBF), one(CBF))
            x1b = vcat(zero(CBF), one(CBF))
            x2a = vcat(zeros(CBF, 2), one(CBF))
            x2b = vcat(
                zero(CBF),
                BigFloat("0.16666657785001893215")im,
                BigFloat("0.04166664890333648869"),
            )
            xv = [(x1a, x1b), (x2a, x2b)]
            y = [
                BigFloat("0.99999999999999999997"),
                BigFloat("-0.99999999999981067844")im,
                BigFloat("-0.49999999999994320353"),
                one(CBF),
            ]
        elseif (k == 3) # Equation (22)
            x1a = vcat(zero(CBF), one(CBF))
            x1b = vcat(zero(CBF), one(CBF))
            x2a = vcat(zeros(CBF, 2), one(CBF))
            x2b = vcat(
                zero(CBF),
                BigFloat("431") / BigFloat("4000"),
                BigFloat("-0.02693906873598870733")im,
            )
            x3a = vcat(
                zeros(CBF, 2),
                BigFloat("0.66321004441662438593")im,
                one(CBF),
            )
            x3b = vcat(
                BigFloat("0.54960853911436015786")im,
                BigFloat("0.16200952846773660904"),
                BigFloat("-0.01417981805211804396")im,
                BigFloat("-0.03415953916892111403"),
            )
            xv = [(x1a, x1b), (x2a, x2b), (x3a, x3b)]
            y = [
                BigFloat("0.99999999999999999928"),
                BigFloat("-0.99999999999999233987")im,
                BigFloat("-0.13549409636220703066"),
                zero(CBF),
                one(CBF),
            ]
        elseif (k == 4) # Equation (23)
            x1a = vcat(zero(CBF), one(CBF))
            x1b = vcat(zero(CBF), one(CBF))
            x2a = vcat(zeros(CBF, 2), one(CBF))
            x2b = vcat(zero(CBF), one(CBF), zero(CBF))
            # B_4^2
            x3a = vcat(
                zero(CBF),
                BigFloat("0.13340427306445612526")im,
                BigFloat("0.02022602029818310774"),
                BigFloat("-0.00674638241111650999")im,
            )
            x3b = x3a
            # (B_2 + A_6)*A_6 = (B_2 + B_3 + B_4^2)*(B_3 + B_4^2)
            b2 = vcat(
                zero(CBF),
                BigFloat("1.41183797496250375498")im,
                zero(CBF),
                BigFloat("-0.00866935318616372016")im,
            )
            b3 = vcat(
                BigFloat("2.69584306915332564689"),
                BigFloat("-1.35910926168869260391")im,
                BigFloat("-0.09896214548845831754"),
                BigFloat("0.01596479463299466666")im,
            )
            x4a = vcat((b2 + b3), one(CBF))
            x4b = vcat(b3, one(CBF))
            xv = [(x1a, x1b), (x2a, x2b), (x3a, x3b), (x4a, x4b)]
            # B_1 + (B_2 + A_6)*A_6 = B_1 + previous
            b1 = vcat(
                BigFloat("-6.26756985350202252845"),
                BigFloat("2.52179694712098096140")im,
                BigFloat("0.05786296656487001838"),
                BigFloat("-0.07766686408071870344")im,
            )
            y = vcat(b1, zero(CBF), one(CBF))
        elseif (k == 5) # Equation (25)
            x1a = vcat(zero(CBF), one(CBF))
            x1b = vcat(zero(CBF), one(CBF))
            x2a = vcat(zeros(CBF, 2), one(CBF))
            x2b = vcat(zero(CBF), one(CBF), zero(CBF))
            x3a = vcat(zeros(CBF, 3), one(CBF))
            x3b = vcat(zeros(CBF, 3), one(CBF))
            # B_1 * B_5
            x4a = vcat(
                zero(CBF),
                BigFloat("3") / BigFloat("25"),
                BigFloat("-0.00877476096879703859")im,
                BigFloat("-0.00097848453523780954"),
                zero(CBF),
            )
            x4b = vcat(
                zeros(CBF, 2),
                BigFloat("-0.123953695858283131480")im,
                BigFloat("-0.011202694841085592373"),
                BigFloat("-0.000012367240538259896")im,
            )

            # (B_3 + A_9)*A_9 = (B_3 + B_4 + previous)*(B_4 + previos)
            b3 = vcat(
                BigFloat("-2.58175430371188142440"),
                BigFloat("-1.73033278310812419209")im,
                BigFloat("-0.07673476833423340755"),
                BigFloat("-0.00261502969893897079")im,
                BigFloat("-0.00003400011993049304"),
            )
            b4 = vcat(
                BigFloat("2.92377758396553673559"),
                BigFloat("1.44513300347488268510")im,
                BigFloat("0.12408183566550450221"),
                BigFloat("-0.01957157093642723948")im,
                BigFloat("0.00002425253007433925"),
            )
            x5a = vcat((b3 + b4), one(CBF))
            x5b = vcat(b4, one(CBF))
            xv = [(x1a, x1b), (x2a, x2b), (x3a, x3b), (x4a, x4b), (x5a, x5b)]
            # B_2 + (B_3 + A_9)*A_9 = B_2 + previous
            b2 = vcat(
                zero(CBF),
                BigFloat("-0.66040840760771318751")im,
                BigFloat("-1.09302278471564897987"),
                BigFloat("0.25377155817710873323")im,
                BigFloat("0.00054374267434731225"),
            )
            y = vcat(b2, zero(CBF), one(CBF))
        end
        # Force convert to type
        xv = map(
            i -> (convert.(T, xv[i][1]), convert.(T, xv[i][2])),
            1:size(xv, 1),
        )
        y = convert.(T, y)
        return graph_degopt(xv, y)
    else  # Scalign-and-squaring phase
        s = k - 5
        degopt = Degopt(graph_bbcs_cheb_exp(5; T = T)[1])
        scale!(degopt, BigFloat("2.0")^(-big(s)))
        for i = 1:s
            square!(degopt)
        end
        return graph_degopt(degopt)
    end
end


text:
immutable PersistentQueue{T}
    in::AbstractList{T}
    out::AbstractList{T}
    length::Int

    PersistentQueue(in::AbstractList{T}, out::AbstractList{T}, length::Int) =
        new(in, out, length)

    PersistentQueue() = new(EmptyList{T}(), EmptyList{T}(), 0)
end

PersistentQueue{T}(v::AbstractVector{T}) =
    PersistentQueue{T}(EmptyList{T}(), reverse(PersistentList(v)), length(v))

queue = PersistentQueue

Base.length(q::PersistentQueue) = q.length
Base.isempty(q::PersistentQueue) = is(q.length, 0)

peek(q::PersistentQueue) = isempty(q.out) ? head(reverse(q.in)) : head(q.out)

pop{T}(q::PersistentQueue{T}) =
    if isempty(q.out)
        PersistentQueue{T}(EmptyList{T}(), tail(reverse(q.in)), length(q) - 1)
    else
        PersistentQueue{T}(q.in, tail(q.out), length(q) - 1)
    end

enq{T}(q::PersistentQueue{T}, val) =
    if isempty(q.in) && isempty(q.out)
        PersistentQueue{T}(q.in, val..EmptyList{T}(), 1)
    else
        PersistentQueue{T}(val..q.in, q.out, length(q) + 1)
    end

Base.start(q::PersistentQueue) = (q.in, q.out)
Base.done{T}(::PersistentQueue{T}, state::(@compat Tuple{EmptyList{T}, EmptyList{T}})) = true
Base.done(::PersistentQueue, state) = false

function Base.next{T}(::PersistentQueue{T}, state::(@compat Tuple{AbstractList{T},
                                                                  PersistentList{T}}))
    in, out = state
    (head(out), (in, tail(out)))
end
function Base.next{T}(q::PersistentQueue{T}, state::(@compat Tuple{PersistentList{T},
                                                                   EmptyList{T}}))
    in, out = state
    next(q, (EmptyList{T}(), reverse(in)))
end


text:
@cenum Changed::UInt32 begin
    CHANGED_NO = 0
    CHANGED_YES = 1
end

@cenum Status::Int32 begin
    STATUS_SUCCESS = 0
    STATUS_ERROR = -1
    STATUS_FATAL = -2
    STATUS_ERROR_NOT_SUPPORTED = -3
    STATUS_ERROR_INVALID_ARGUMENT = -4
    STATUS_ERROR_INVALID_ARGUMENT_SIZE = -5
    STATUS_ERROR_ALREADY_INITIALIZED = -6
    STATUS_ERROR_NOT_INITIALIZED = -7
    STATUS_ERROR_VERSION_MISMATCH = -8
    STATUS_ERROR_ALREADY_ATTACHED = -9
    STATUS_ERROR_INVALID_ARCHITECTURE_ID = -10
    STATUS_ERROR_ILLEGAL_INSTRUCTION = -11
    STATUS_ERROR_INVALID_CODE_OBJECT_ID = -12
    STATUS_ERROR_INVALID_ELF_AMDGPU_MACHINE = -13
    STATUS_ERROR_INVALID_PROCESS_ID = -14
    STATUS_ERROR_INVALID_AGENT_ID = -15
    STATUS_ERROR_INVALID_QUEUE_ID = -16
    STATUS_ERROR_INVALID_DISPATCH_ID = -17
    STATUS_ERROR_INVALID_WAVE_ID = -18
    STATUS_ERROR_WAVE_NOT_STOPPED = -19
    STATUS_ERROR_WAVE_STOPPED = -20
    STATUS_ERROR_WAVE_OUTSTANDING_STOP = -21
    STATUS_ERROR_WAVE_NOT_RESUMABLE = -22
    STATUS_ERROR_INVALID_DISPLACED_STEPPING_ID = -23
    STATUS_ERROR_DISPLACED_STEPPING_BUFFER_UNAVAILABLE = -24
    STATUS_ERROR_INVALID_WATCHPOINT_ID = -25
    STATUS_ERROR_NO_WATCHPOINT_AVAILABLE = -26
    STATUS_ERROR_INVALID_REGISTER_CLASS_ID = -27
    STATUS_ERROR_INVALID_REGISTER_ID = -28
    STATUS_ERROR_INVALID_LANE_ID = -29
    STATUS_ERROR_INVALID_ADDRESS_CLASS_ID = -30
    STATUS_ERROR_INVALID_ADDRESS_SPACE_ID = -31
    STATUS_ERROR_MEMORY_ACCESS = -32
    STATUS_ERROR_INVALID_ADDRESS_SPACE_CONVERSION = -33
    STATUS_ERROR_INVALID_EVENT_ID = -34
    STATUS_ERROR_INVALID_SHARED_LIBRARY_ID = -35
    STATUS_ERROR_INVALID_BREAKPOINT_ID = -36
    STATUS_ERROR_CLIENT_CALLBACK = -37
    STATUS_ERROR_INVALID_CLIENT_PROCESS_ID = -38
    STATUS_ERROR_PROCESS_EXITED = -39
    STATUS_ERROR_LIBRARY_NOT_LOADED = -40
    STATUS_ERROR_SYMBOL_NOT_FOUND = -41
    STATUS_ERROR_INVALID_ADDRESS = -42
end

@cenum ArchitectureInfo::UInt32 begin
    ARCHITECTURE_INFO_NAME = 1
    ARCHITECTURE_INFO_ELF_AMDGPU_MACHINE = 2
    ARCHITECTURE_INFO_LARGEST_INSTRUCTION_SIZE = 3
    ARCHITECTURE_INFO_MINIMUM_INSTRUCTION_ALIGNMENT = 4
    ARCHITECTURE_INFO_BREAKPOINT_INSTRUCTION_SIZE = 5
    ARCHITECTURE_INFO_BREAKPOINT_INSTRUCTION = 6
    ARCHITECTURE_INFO_BREAKPOINT_INSTRUCTION_PC_ADJUST = 7
    ARCHITECTURE_INFO_PC_REGISTER = 8
    ARCHITECTURE_INFO_EXECUTION_MASK_REGISTER = 9
    ARCHITECTURE_INFO_WATCHPOINT_COUNT = 10
    ARCHITECTURE_INFO_WATCHPOINT_SHARE = 11
    ARCHITECTURE_INFO_DEFAULT_GLOBAL_ADDRESS_SPACE = 12
    ARCHITECTURE_INFO_PRECISE_MEMORY_SUPPORTED = 13
end

@cenum ProcessInfo::UInt32 begin
    PROCESS_INFO_NOTIFIER = 1
end

@cenum Progress::UInt32 begin
    PROGRESS_NORMAL = 0
    PROGRESS_NO_FORWARD = 1
end

@cenum CodeObjectInfo::UInt32 begin
    CODE_OBJECT_INFO_URI_NAME = 1
    CODE_OBJECT_INFO_LOAD_ADDRESS = 2
end

@cenum AgentInfo::UInt32 begin
    AGENT_INFO_NAME = 1
    AGENT_INFO_ARCHITECTURE = 2
    AGENT_INFO_PCIE_SLOT = 3
    AGENT_INFO_PCIE_VENDOR_ID = 4
    AGENT_INFO_PCIE_DEVICE_ID = 5
    AGENT_INFO_SHADER_ENGINE_COUNT = 6
    AGENT_INFO_COMPUTE_UNIT_COUNT = 7
    AGENT_INFO_NUM_SIMD_PER_COMPUTE_UNIT = 8
    AGENT_INFO_MAX_WAVES_PER_SIMD = 9
end

@cenum QueueInfo::UInt32 begin
    QUEUE_INFO_AGENT = 1
    QUEUE_INFO_ARCHITECTURE = 2
    QUEUE_TYPE = 3
    QUEUE_INFO_STATE = 4
    QUEUE_INFO_ERROR_REASON = 5
end

@cenum QueueType::UInt32 begin
    QUEUE_TYPE_UNKNOWN = 0
    QUEUE_TYPE_HSA_KERNEL_DISPATCH_MULTIPLE_PRODUCER = 1
    QUEUE_TYPE_HSA_KERNEL_DISPATCH_SINGLE_PRODUCER = 2
    QUEUE_TYPE_AMD_PM4 = 3
end

@cenum QueueState::UInt32 begin
    QUEUE_STATE_VALID = 1
    QUEUE_STATE_ERROR = 2
end

@cenum QueueErrorReason::UInt32 begin
    QUEUE_ERROR_REASON_INVALID_PACKET = 1
    QUEUE_ERROR_REASON_MEMORY_VIOLATION = 2
    QUEUE_ERROR_REASON_ASSERT_TRAP = 4
    QUEUE_ERROR_REASON_WAVE_ERROR = 8
end

@cenum DispatchInfo::UInt32 begin
    DISPATCH_INFO_QUEUE = 1
    DISPATCH_INFO_AGENT = 2
    DISPATCH_INFO_ARCHITECTURE = 3
    DISPATCH_INFO_PACKET_ID = 4
    DISPATCH_INFO_BARRIER = 5
    DISPATCH_INFO_ACQUIRE_FENCE = 6
    DISPATCH_INFO_RELEASE_FENCE = 7
    DISPATCH_INFO_GRID_DIMENSIONS = 8
    DISPATCH_INFO_WORK_GROUP_SIZES = 9
    DISPATCH_INFO_GRID_SIZES = 10
    DISPATCH_INFO_PRIVATE_SEGMENT_SIZE = 11
    DISPATCH_INFO_GROUP_SEGMENT_SIZE = 12
    DISPATCH_INFO_KERNEL_ARGUMENT_SEGMENT_ADDRESS = 13
    DISPATCH_INFO_KERNEL_ENTRY_ADDRESS = 14
end

@cenum DispatchBarrier::UInt32 begin
    DISPATCH_BARRIER_NONE = 0
    DISPATCH_BARRIER_PRESENT = 1
end

@cenum DispatchFenceScope::UInt32 begin
    DISPATCH_FENCE_SCOPE_NONE = 0
    DISPATCH_FENCE_SCOPE_AGENT = 1
    DISPATCH_FENCE_SCOPE_SYSTEM = 2
end

@cenum WaveInfo::UInt32 begin
    WAVE_INFO_STATE = 1
    WAVE_INFO_STOP_REASON = 2
    WAVE_INFO_WATCHPOINTS = 3
    WAVE_INFO_DISPATCH = 4
    WAVE_INFO_QUEUE = 5
    WAVE_INFO_AGENT = 6
    WAVE_INFO_ARCHITECTURE = 7
    WAVE_INFO_PC = 8
    WAVE_INFO_EXEC_MASK = 9
    WAVE_INFO_WORK_GROUP_COORD = 10
    WAVE_INFO_WAVE_NUMBER_IN_WORK_GROUP = 11
    WAVE_INFO_LANE_COUNT = 12
end

@cenum WaveState::UInt32 begin
    WAVE_STATE_RUN = 1
    WAVE_STATE_SINGLE_STEP = 2
    WAVE_STATE_STOP = 3
end

@cenum WaveStopReason::UInt32 begin
    WAVE_STOP_REASON_NONE = 0
    WAVE_STOP_REASON_BREAKPOINT = 1
    WAVE_STOP_REASON_WATCHPOINT = 2
    WAVE_STOP_REASON_SINGLE_STEP = 4
    WAVE_STOP_REASON_QUEUE_ERROR = 8
    WAVE_STOP_REASON_FP_INPUT_DENORMAL = 16
    WAVE_STOP_REASON_FP_DIVIDE_BY_0 = 32
    WAVE_STOP_REASON_FP_OVERFLOW = 64
    WAVE_STOP_REASON_FP_UNDERFLOW = 128
    WAVE_STOP_REASON_FP_INEXACT = 256
    WAVE_STOP_REASON_FP_INVALID_OPERATION = 512
    WAVE_STOP_REASON_INT_DIVIDE_BY_0 = 1024
    WAVE_STOP_REASON_DEBUG_TRAP = 2048
    WAVE_STOP_REASON_ASSERT_TRAP = 4096
    WAVE_STOP_REASON_TRAP = 8192
    WAVE_STOP_REASON_MEMORY_VIOLATION = 16384
    WAVE_STOP_REASON_ILLEGAL_INSTRUCTION = 32768
    WAVE_STOP_REASON_ECC_ERROR = 65536
    WAVE_STOP_REASON_FATAL_HALT = 131072
    WAVE_STOP_REASON_XNACK_ERROR = 262144
end

@cenum ResumeMode::UInt32 begin
    RESUME_MODE_NORMAL = 0
    RESUME_MODE_SINGLE_STEP = 1
end

@cenum WatchpointShareKind::UInt32 begin
    WATCHPOINT_SHARE_KIND_UNSUPPORTED = 0
    WATCHPOINT_SHARE_KIND_UNSHARED = 1
    WATCHPOINT_SHARE_KIND_SHARED = 2
end

@cenum WatchpointKind::UInt32 begin
    WATCHPOINT_KIND_LOAD = 1
    WATCHPOINT_KIND_STORE_AND_RMW = 2
    WATCHPOINT_KIND_RMW = 3
    WATCHPOINT_KIND_ALL = 4
end

@cenum RegisterClassInfo::UInt32 begin
    REGISTER_CLASS_INFO_NAME = 1
end

@cenum RegisterInfo::UInt32 begin
    REGISTER_INFO_NAME = 1
    REGISTER_INFO_SIZE = 2
    REGISTER_INFO_TYPE = 3
end

@cenum RegisterClassState::UInt32 begin
    REGISTER_CLASS_STATE_NOT_MEMBER = 0
    REGISTER_CLASS_STATE_MEMBER = 1
end

@cenum AddressClassInfo::UInt32 begin
    ADDRESS_CLASS_INFO_NAME = 1
    ADDRESS_CLASS_INFO_ADDRESS_SPACE = 2
end

@cenum AddressSpaceAccess::UInt32 begin
    ADDRESS_SPACE_ACCESS_ALL = 1
    ADDRESS_SPACE_ACCESS_PROGRAM_CONSTANT = 2
    ADDRESS_SPACE_ACCESS_DISPATCH_CONSTANT = 3
end

@cenum AddressSpaceInfo::UInt32 begin
    ADDRESS_SPACE_INFO_NAME = 1
    ADDRESS_SPACE_INFO_ADDRESS_SIZE = 2
    ADDRESS_SPACE_INFO_NULL_ADDRESS = 3
    ADDRESS_SPACE_INFO_ACCESS = 4
end

@cenum AddressSpaceAlias::UInt32 begin
    ADDRESS_SPACE_ALIAS_NONE = 0
    ADDRESS_SPACE_ALIAS_MAY = 1
end

@cenum AddressClassState::UInt32 begin
    ADDRESS_CLASS_STATE_NOT_MEMBER = 0
    ADDRESS_CLASS_STATE_MEMBER = 1
end

@cenum MemoryPrecision::UInt32 begin
    MEMORY_PRECISION_NONE = 0
    MEMORY_PRECISION_PRECISE = 1
end

@cenum EventKind::UInt32 begin
    EVENT_KIND_NONE = 0
    EVENT_KIND_WAVE_STOP = 1
    EVENT_KIND_WAVE_COMMAND_TERMINATED = 2
    EVENT_KIND_CODE_OBJECT_LIST_UPDATED = 3
    EVENT_KIND_BREAKPOINT_RESUME = 4
    EVENT_KIND_RUNTIME = 5
    EVENT_KIND_QUEUE_ERROR = 6
end

@cenum RuntimeState::UInt32 begin
    RUNTIME_STATE_LOADED_SUPPORTED = 1
    RUNTIME_STATE_LOADED_UNSUPPORTED = 2
    RUNTIME_STATE_UNLOADED = 3
end

@cenum EventInfo::UInt32 begin
    EVENT_INFO_KIND = 1
    EVENT_INFO_WAVE = 2
    EVENT_INFO_BREAKPOINT = 3
    EVENT_INFO_CLIENT_THREAD = 4
    EVENT_INFO_RUNTIME_STATE = 5
    EVENT_INFO_RUNTIME_VERSION = 6
end

@cenum LogLevel::UInt32 begin
    LOG_LEVEL_NONE = 0
    LOG_LEVEL_FATAL_ERROR = 1
    LOG_LEVEL_WARNING = 2
    LOG_LEVEL_INFO = 3
    LOG_LEVEL_VERBOSE = 4
end

@cenum SharedLibraryState::UInt32 begin
    SHARED_LIBRARY_STATE_LOADED = 1
    SHARED_LIBRARY_STATE_UNLOADED = 2
end

@cenum BreakpointAction::UInt32 begin
    BREAKPOINT_ACTION_RESUME = 1
    BREAKPOINT_ACTION_HALT = 2
end

@cenum BreakpointState::UInt32 begin
    BREAKPOINT_STATE_DISABLE = 1
    BREAKPOINT_STATE_ENABLE = 2
end


text:
using Revise
using ADCME
using ADSeismic
using PyPlot
using DelimitedFiles

##=======================
reset_default_graph()

scale = 100
param = AcousticPropagatorParams(NX=scale, NY=scale, NSTEP=3, DELTAT=1e-4,  DELTAX=1.0, DELTAY=1.0)

rc = Gauss(param, 1/pi, 100, 1e6)
srci = [div(param.NX,2)]
srcj = [div(param.NY,2)]
srcv = reshape(rc, :, 1)
src = AcousticSource(srci, srcj, srcv)

C = 3300*ones(param.NX+2, param.NY+2)
model = AcousticPropagatorSolver(param, src, C)

sess = Session(); init(sess)

run_profile(sess, model.u)
save_profile("acoustic-profile-$(has_gpu() ? "gpu" : "cpu").json")

##=======================
reset_default_graph()

scale = 100
param = ElasticPropagatorParams(NX=scale, NY=scale, NSTEP=2, DELTAT=1e-4,  DELTAX=1.0, DELTAY=1.0)

source = Gauss(param, 1/pi, 100, 1e6)
srci = [div(param.NX,2)]
srcj = [div(param.NY,2)]
srctype = [0]
srcv = reshape(source, :, 1)
src = ElasticSource(srci, srcj, srctype, srcv)

vp = 3300.
vs = 3300. / 1.732
rho = 2800.
λ, μ, ρ = compute_lame_parameters(param.NX, param.NY, vp, vs, rho)
model = ElasticPropagatorSolver(param, src, ρ, λ, μ)

sess = Session(); init(sess)

run_profile(sess, model.vx)
save_profile("elastic-profile-$(has_gpu() ? "gpu" : "cpu").json")


text:
using Test
using ModelingToolkit
using AnalyticDiffEq
# Laguerre
# xy'' + (α + 1 - x)* y' + ny = 0

# Hermite
# y'' - 2xy' + λy = 0

# Legendre
# (1 - x^2) * y'' -2xy' + l(l+1)y = 0

L0(x) = 1.0
    L1(x) = -x + 1.0
    L2(x) = .5 * (x^2 - 4x + 2)
    L3(x) = (1/6)*(-x^3 +9x^2 -18x + 6)
    L4(x) = (1/24)*(x^4 - 16x^3 + 72x^2 - 96x + 24)
    L5(x) = (1/120)*(-x^5 + 25x^4 -200x^3 + 600x^2 - 600x + 120)
    L6(x) = (1/720)*(x^6 - 36x^5 + 450x^4 - 2400x^3 + 5400x^2 - 4320x + 720)
    L7(x) = (1/5040)*(-x^7 + 49x^6 - 882x^5 + 7350x^4 - 29_400x^3 + 52_920x^2 - 35_280x + 5040)
    L8(x) = (1/40320)*(x^8 - 64x^7 + 1568x^6 - 18_816x^5 + 117_600x^4 - 376_320x^3 + 564_480x^2 - 322_560x + 40320)
    L9(x) = (1/362880)*(-x^9 + 81x^8 - 2592x^7 + 42_336x^6 - 381_024x^5 + 1_905_120x^4 - 5_080_320x^3 + 6_531_840x^2 - 3_265_920x + 362_880)
    L10(x) = (1/3628800)*(x^10 - 100x^9 + 4_050x^8 - 86_400x^7 + 1_058_400x^6 - 7_620_480x^5 + 31_752_000x^4 - 72_576_000x^3  + 81_648_000x^2 - 36_288_000x + 3_628_800)
    L11(x) = (1/39916800)*(-x^11 + 121x^10 - 6_050x^9 + 163_350x^8 - 2_613_600x^7 + 25_613_280x^6 - 153_679_680x^5 + 548_856_000x^4 - 1_097_712_000x^3 + 1_097_712_000x^2 - 439_084_800x + 39_916_800)
    L12(x) = (1/479001600)*(x^12 - 144x^11 + 8_712x^10 - 290_400x^9 + 5_880_600x^8 - 75_271_680x^7 + 614_718_720x^6 - 3_161_410_560x^5 + 9_879_408_000x^4 - 17_563_392_000x^3 + 15_807_052_800x^2 - 5_748_019_200x + 479_001_600)

    Ls = [L0, L1, L2, L3, L4, L5, L6, L7, L8, L9, L10, L11, L12]

    # Credit to @yingboma
    function laguerre(n, x)

        n == 0 ? (return 1) : nothing

        l0, l1 = 1, 1-x
        for k in 1:n-1
            l0, l1 = l1, ( (2k+1-x)*l1 - k*l0 )/(k+1)
        end
        l1
    end

H0(x) = 1.0
    H1(x) = 2x
    H2(x) = -2 + 4x^2
    H3(x) = -12x + 8x^3
    H4(x) = 12 - 48x^2 + 16x^4
    H5(x) = 120x -160x^3 + 32x^5
    H6(x) = -120 + 720x^2 - 480x^4 + 64x^6
    H7(x) = -1680x + 3360x^3 - 1344x^5 + 128x^7
    H8(x) =1680 - 13440x^2 + 13440x^4 - 3584x^6 + 256x^8
    H9(x) =30240x - 80640x^3 + 48384x^5 - 9216x^7 + 512x^9
    H10(x) =-30240 + 302400x^2 - 403200x^4 + 161280x^6 - 23040x^8 + 1024x^10
    H11(x) =-665280x + 2217600x^3 - 1774080x^5 + 506880x^7 - 56320x^9 + 2048x^11
    H12(x) =665280 - 7983360x^2 + 13305600x^4 - 7096320x^6 + 1520640x^8 - 135168x^10 + 4096x^12

    Hs = [H0, H1, H2, H3, H4, H5, H6, H7, H8, H9, H10, H11, H12]

    function hermite(x,n)
        if n == 0
            return 1
        elseif n == 1
            return 2x
        else
            return 2x*hermite(x,n-1) - 2*(n-1)*hermite(x,n-2)
        end
    end

P0(x) = 1
    P1(x) =  x
    P2(x) = .5*(3x^2 -1)
    P3(x) = .5*(5x^3 -3x)
    P4(x) = (1/8)*(35x^4 - 30x^2 + 3)
    P5(x) = (1/8)*(63x^5 - 70x^3 + 15x)
    P6(x) = (1/16)*(231x^6 - 315x^4 + 105x^2 - 5)
    P7(x) = (1/16)*(429x^7 - 693x^5 + 315x^3 - 35x)
    P8(x) = (1/128)*(6435x^8 - 12012x^6 + 6930x^4 - 1260x^2 + 35)
    P9(x) = (1/128)*(12155x^9 - 25740x^7 + 18018x^5 - 4620x^3 + 315x)
    P10(x) = (1/256)*(46189x^10 - 109395x^8 + 90090x^6- 30030x^4 + 3465x^2 - 63)
    P11(x) = (1/256)*(2^11)*(705432x^11 - 1847560x^9 + 1750320x^7 - 720720x^5 + 120120x^3 - 5544x)
    P12(x) = (1/512)*(2^12)*(2704156x^12 - 7759752x^10 + 8314020x^8 - 4084080x^6 + 900900x^4 - 72072x^2 + 924)

    Ps = [P0, P1, P2, P3, P4, P5, P6, P7, P8, P9, P10, P11, P12]
    testPs = [P2, P3, P9, P10]

    function legendre(x, n)
        if n == 0
            return 1
        elseif n == 1
            return x
        else
            return ((2n-1)*x*legendre(x,n-1) - (n - 1)*legendre(x, n-2))/n
        end
    end

@testset "AnalyticDiffEq.jl" begin

    @testset "Analytic Detection" begin

        # Laguerre
        # xy'' + (α + 1 - x)* y' + ny = 0
        # eq = x*D2(y) + (α + 1 - x) * D(y) + n*y ~ 0
        @testset "Laguerre" begin

            @parameters t
            @variables y(t) x(t)
            @derivatives D'~t D2''~t

            eqtest1 = x*D2(y) + (1 + 1 - x) * D(y) + 1*y ~ 0
            eqtest2 = x*D2(y) + (6 - x) * D(y) + 3*y ~ 0
            eqtest3 = x*D2(y) + (1 - x) * D(y) + 7*y ~ 0
            # Weird corner cases
            # eqtest4 = x*D2(y) + (7 - x) * D(y) ~ 0
            # eqtest0 = x*D2(y) + (1 - x) * D(y) ~ 0
            # eqtest5 = x*D2(y) + (1 - x) * D(y) + y ~ 0

            LAE = LaguerreAbstractEquation

            @test isit(eqtest1, LAE)
            @test isit(eqtest2, LAE)
            @test isit(eqtest3, LAE)
            # @test_broken isit(eqtest4, LaguerreAbstractEquation)
            # @test isit(eqtest5, LaguerreAbstractEquation)
            # @test_broken isit(eqtest0, LaguerreAbstractEquation)

       end

       # Hermite
       # y'' - 2xy' + λy = 0
       # eqtest1 = D2(y) - 2*x*D(y) + λ*y ~ 0
        @testset "Hermite" begin

            @parameters t
            @derivatives D'~t D2''~t
            @variables y(t) x(t)

            eqtest1 = D2(y) - 2*x*D(y) + 1*y ~ 0
            eqtest2 = D2(y) - 2*x*D(y) + 7*y ~ 0
            eqtest3 = D2(y) - 2*x*D(y) + 12*y ~ 0

            HAE = HermiteAbstractEquation

            @test isit(eqtest1, HAE)
            @test isit(eqtest2, HAE)
            @test isit(eqtest3, HAE)

        end

        # Legendre
        # (1 - x^2) * y'' -2xy' + l(l+1)y = 0
        # eqtest1 = (1-x^2) * D2(y) - 2*x*D(y) + l*(l+1)*y ~ 0
        @testset "Legendre" begin

            @parameters t
            @derivatives D'~t D2''~t
            @variables y(t) x(t)

            eqtest1 = (1-x^2) * D2(y) - 2*x*D(y) + 2*y ~ 0
            eqtest2 = (1-x^2) * D2(y) - 2*x*D(y) + 6*y ~ 0
            eqtest3 = (1-x^2) * D2(y) - 2*x*D(y) + 12*y ~ 0

            @test AnalyticDiffEq.llplus1checker(6) == true
            @test AnalyticDiffEq.llplus1checker(2) == true
            @test AnalyticDiffEq.llplus1checker(12) == true
            @test AnalyticDiffEq.llplus1checker(20) == true

            LegAE = LegendreAbstractEquation
            @test isit(eqtest1, LegAE)
            @test isit(eqtest2, LegAE)
            @test isit(eqtest3, LegAE)

        end
    end

    @testset "Analytic Extraction" begin

        @parameters t
        @derivatives D'~t D2''~t
        @variables y(t) x(t)

        eqlag = x*D2(y) + (0 + 1 - x) * D(y) + 6*y ~ 0
        eqher = D2(y) - 2*x*D(y) + 5*y ~ 0
        eqleg = (1-x^2) * D2(y) - 2*x*D(y) + 4*(4+1)*y ~ 0
        @test extract(eqlag, LaguerreAbstractEquation) == LaguerreEquation{0,6}()
        @test extract(eqher, HermiteAbstractEquation) == HermiteEquation{5}()
        @test extract(eqleg, LegendreAbstractEquation) == LegendreEquation{4}()

    end

    @testset "Analytic Construction" begin

        @parameters t
        @derivatives D'~t D2''~t
        @variables y(t) x(t)

        eqlag = x*D2(y) + (0 + 1 - x) * D(y) + 6*y ~ 0
        eqher = D2(y) - 2*x*D(y) + 5*y ~ 0
        eqleg = (1-x^2) * D2(y) - 2*x*D(y) + 4*(4+1)*y ~ 0
        eqerr = 0 ~ x^2
        LagEQ = extract(eqlag, LaguerreAbstractEquation)
        HerEQ = extract(eqher, HermiteAbstractEquation)
        LegEQ = extract(eqleg, LegendreAbstractEquation)

        @test build(LagEQ)(.3) ≈ Ls[7](.3)
        @test build(HerEQ)(.3) ≈ Hs[6](.3)
        @test build(LegEQ)(.3) ≈ Ps[5](.3)

        @test analyze(eqlag)(.3) ≈ Ls[7](.3)
        @test analyze(eqher)(.3) ≈ Hs[6](.3)
        @test analyze(eqleg)(.3) ≈ Ps[5](.3)
        @test_broken analyze(eqerr) == 3.14

    end

    @testset "Solve Interface" begin
        @test_broken 1 == 2
    end
end


text:
# 2D FWI on Overthrust model using minConf library
# Author: Philipp Witte, pwitte@eoas.ubc.ca
# Date: December 2017
#

using Statistics, Random, LinearAlgebra
using JUDI.TimeModeling, JUDI.SLIM_optim, HDF5, SegyIO, PyPlot

# Load starting model
n,d,o,m0 = read(h5open("../../data/overthrust_model.h5","r"), "n", "d", "o", "m0")
model0 = Model((n[1],n[2]), (d[1],d[2]), (o[1],o[2]), m0)

# Bound constraints
v0 = sqrt.(1f0 ./ model0.m)
vmin = ones(Float32,model0.n) .* 1.3f0
vmax = ones(Float32,model0.n) .* 6.5f0
vmin[:,1:21] .= v0[:,1:21]   # keep water column fixed
vmax[:,1:21] .= v0[:,1:21]

# Slowness squared [s^2/km^2]
mmin = vec((1f0 ./ vmax).^2)
mmax = vec((1f0 ./ vmin).^2)

# Load data
block = segy_read("../../data/overthrust_shot_records.segy")
d_obs = judiVector(block)

# Set up wavelet
src_geometry = Geometry(block; key="source")
wavelet = ricker_wavelet(src_geometry.t[1],src_geometry.dt[1],0.008f0)  # 8 Hz wavelet
q = judiVector(src_geometry,wavelet)

############################### FWI ###########################################

# Optimization parameters
niterations = 10
batchsize = 16
fhistory_SGD = zeros(Float32,niterations)

# Projection operator for bound constraints
proj(x) = reshape(median([vec(mmin) vec(x) vec(mmax)]; dims=2),model0.n)

# Main loop
for j=1:niterations

    # get fwi objective function value and gradient
    i = randperm(d_obs.nsrc)[1:batchsize]
    fval, gradient = fwi_objective(model0,q[i],d_obs[i])
    println("FWI iteration no: ",j,"; function value: ",fval)
    fhistory_SGD[j] = fval

	# linesearch
	step = backtracking_linesearch(model0, q[i], d_obs[i], fval, gradient, proj; alpha=1f0)

	# Update model and bound projection
	model0.m = proj(model0.m + reshape(step,model0.n))
end

figure(); imshow(sqrt.(1f0./adjoint(model0.m))); title("FWI with SGD")


text:
"""
This module contains function definitions
for using custom types as Terms.
"""
module TermInterface

using ..Util
using Base.Meta

# TODO document this interface for extending for custom types
# Interface for Expr. Implement these methods for your own type to
# Use it instead of Expr in egraphs!
# gethead(e::Expr) = isexpr(e, :call) ? e.args[1] : e.head
gethead(e::Expr) = e.head
# getargs(e::Expr) = @view e.args[(isexpr(e, :call) ? 2 : 1):end]
getargs(e::Expr) = e.args
istree(e::Expr) = true
# getmetadata(e::Expr) = (iscall=isexpr(e, :call),)
getmetadata(e::Expr) = (;)
arity(e::Expr) = length(getargs(e)) # optional
preprocess(e::Expr) = cleanast(e)

# Fallback implementation for all other types
gethead(e) = e
getargs(e) = []
istree(a) = false
getmetadata(e) = (;) # empty NamedTuple
arity(e) = length(getargs(e)) # optional
preprocess(e) = e

export gethead
export getargs
export istree
export getmetadata
export preprocess
export arity
end


text:
# Box constrained problems

using Test, LinearAlgebra, SparseArrays, Random
using COSMO

# This test is precision agnostic and will be run with T precision
if !@isdefined(UnitTestFloats)
    UnitTestFloats = [Float64] #if not run in full test setup, just do it for one float type
end
@testset "Box Problems" begin

for T in UnitTestFloats

    @testset "Box Problems (T = $(T))" begin
        @testset "Box Problems - feasible" begin
            A = SparseMatrixCSC{T, Int64}([1. 0; 0. 1.])
            b = T[0.;0.]

            P = Matrix(Diagonal(ones(T, 2)))
            q = T[1.;-1.]
            l = T[0.;0]
            u = T[1.;1.]
            constraint = COSMO.Constraint(A,b, COSMO.Box(l,u))
            constraints = [constraint]

            model = COSMO.Model{T}()
            assemble!(model, P, q, constraints)

            res = COSMO.optimize!(model)
            @test res.status == :Solved
            @test isapprox(res.obj_val, -0.5, atol=1e-5)
        end
        nothing

        @testset "Box Problems - primal infeasible" begin
            A = SparseMatrixCSC{T, Int64}([1. 0; 1. 0.])
            b = T[2.;0.]

            P = Matrix(Diagonal(ones(T, 2)))
            q = T[1.;-1.]
            l = T[0.;0]
            u = T[1.;1.]
            constraint = COSMO.Constraint(A,b, COSMO.Box(l,u))
            constraints = [constraint]

            model = COSMO.Model{T}()
            assemble!(model, P, q, constraints)

            res = COSMO.optimize!(model)
            @test res.status == :Primal_infeasible
        end

        @testset "Box Problems - primal infeasible" begin
            A = SparseMatrixCSC{T, Int64}([1. 0; 1. 0.])
            b = T[0.;0.]

            P = Matrix(Diagonal(ones(T, 2)))
            q = T[1.;-1.]
            l = T[0.;2.]
            u = T[1.;3.]
            constraint = COSMO.Constraint(A,b, COSMO.Box(l,u))
            constraints = [constraint]

            model = COSMO.Model{T}()
            assemble!(model, P, q, constraints)

            res = COSMO.optimize!(model)
            @test res.status == :Primal_infeasible
        end

        @testset "Box Problems - dual infeasible (unscaled)" begin
            A = SparseMatrixCSC{T, Int64}([1. 0; 0. 1.])
            b = T[1.;1.]

            P = Matrix(Diagonal(zeros(T, 2)))
            q = T[1.;1.]
            l = T[0.;-Inf]
            u = T[1.;3.]
            constraint = COSMO.Constraint(A,b, COSMO.Box(l,u))
            constraints = [constraint]

            model = COSMO.Model{T}()
            settings = COSMO.Settings{T}(check_infeasibility=20,scaling=0)
            assemble!(model, P, q, constraints,settings  = settings)

            res = COSMO.optimize!(model)
            @test res.status == :Dual_infeasible
        end

        @testset "Box Problems - dual infeasible (scaled)" begin
            A = SparseMatrixCSC{T, Int64}([1. 0; 0. 1.])
            b = T[1.;1.]

            P = Matrix(Diagonal(zeros(T, 2)))
            q = T[1.;1.]
            l = T[0.;-Inf]
            u = T[1.;3.]
            constraint = COSMO.Constraint(A,b, COSMO.Box(l,u))
            constraints = [constraint]
            model = COSMO.Model{T}()
            settings = COSMO.Settings{T}(check_infeasibility=40,scaling=10)
            assemble!(model, P, q, constraints,settings = settings)

            res = COSMO.optimize!(model)
            @test res.status == :Dual_infeasible
        end
    end
end
end
nothing


text:
using Iris
using Test

@testset "Simulations" begin
    @testset "1D" begin
        @testset "Cartesian" begin
            cavity = Interval(-.5,.5)
            pump = Interval(-.5,0)
            bnd = Boundary(Interval(-1,1),PML)
            dom = NondispersiveDomain(cavity,2)
            fdom = DispersiveDomain(pump,TwoLevelSystem(10,.1,2))
            lat = Cartesian(.05)
            latdom = LatticeDomain(bnd,lat)
            sim = @test_nowarn Simulation(10,dom,latdom,fdom)
            @test_nowarn smooth!(sim)
            dom.n=3
            fdom.F=2
            @test_nowarn update!(sim)
            @test_nowarn smooth!(sim)
            @test_nowarn update_dielectric!(sim,rand(ComplexF64,40))
            @test_nowarn update_pump!(sim,rand(Float64,40))
            @test_nowarn smooth!(sim)
        end
    end
    @testset "2D" begin
        @testset "Cartesian" begin
            cavity = Circle(.5,0,0)
            pump = Square(.3,0,0)
            bnd = Boundary(Rectangle(1,.7,0,0),PML)
            dom = NondispersiveDomain(cavity,2)
            fdom = DispersiveDomain(pump,TwoLevelSystem(10,.1,2))
            lat = Cartesian(.025,.035)
            latdom = LatticeDomain(bnd,lat)
            sim = @test_nowarn Simulation(10,dom,latdom,fdom)
            # @test_nowarn smooth!(sim)
            dom.n=3
            fdom.F=2
            @test_nowarn update!(sim)
            # @test_nowarn smooth!(sim)
            @test_nowarn update_dielectric!(sim,rand(ComplexF64,40))
            @test_nowarn update_pump!(sim,rand(Float64,40))
            # @test_nowarn smooth!(sim)
        end
    end
end


text:
#############################################################################
# ECOS.jl
# Wrapper around the ECOS solver https://github.com/ifa-ethz/ecos
# See http://github.com/JuliaOpt/ECOS.jl
#############################################################################
# ECOS.jl
# Contains the wrapper itself
#############################################################################

module ECOS

using SparseArrays
using LinearAlgebra

# Try to load the binary dependency
if isfile(joinpath(dirname(@__FILE__),"..","deps","deps.jl"))
    include("../deps/deps.jl")
else
    error("ECOS not properly installed. Please run Pkg.build(\"ECOS\")")
end

# ver  [not exported]
# Returns the version of ECOS in use
function ver()
    ver_ptr = ccall((:ECOS_ver, ECOS.ecos), Ptr{UInt8}, ())
    return unsafe_string(ver_ptr)
end

macro ecos_ccall(func, args...)
    args = map(esc,args)
    f = "ECOS_$(func)"
    quote
        @unix_only ret = ccall(($f,ECOS.ecos), $(args...))
        @windows_only ret = ccall(($f,ECOS.ecos), stdcall, $(args...))
        ret
    end
end

function __init__()
    vnum = VersionNumber(ver())
    if !(vnum.major == 2 && vnum.minor == 0)
        depsdir = realpath(joinpath(dirname(@__FILE__),"..","deps"))
        error("Current ECOS version installed is $(ver()), but we require version 2.0.*. On Linux and Windows, delete the contents of the `$depsdir` directory except for the files `build.jl` and `.gitignore`, then rerun Pkg.build(\"ECOS\"). On OS X, run `using Homebrew; Homebrew.update()` in Julia.")
    end
end

include("types.jl")  # All the types and constants defined in ecos.h

# A julia container for matrices passed to ECOS.
# This makes it easy for the Julia GC to track these arrays.
# The format matches ECOS's spmat: 0-based compressed sparse column (CSC)
struct ECOSMatrix
    pr::Vector{Cdouble} # nzval
    jc::Vector{Clong}   # colptr
    ir::Vector{Clong}   # rowval
end

function ECOSMatrix(mat::AbstractMatrix)
    sparsemat = sparse(mat)
    return ECOSMatrix(sparsemat.nzval, sparsemat.colptr .- 1, sparsemat.rowval .- 1)
end

# setup  (direct interface)
# Provide ECOS with a problem in the form
# min c'x
# st  A*x = b
#     G*x in_K h,
#       or equivalently  h - G*x in K
#       or equivalently  G*x + s = h, s in K
# K is the product of R+ and multiple second-order cones, with the
# positive orthant first.
# Inputs:
#   n       Number of variables == length(x)
#   m       Number of inequality constraints == length(h)
#   p       Number of equality constraints == length(b)  (can be 0)
#   l       Dimension of the positive orthant,
#               i.e. the first l elements of s are >= 0
#   ncones  Number of second-order cones present in K
#   q       ncones-long vector of the length of the index sets of the SOCs
#               e.g. cone 1 is indices 4:6, cone 2 is indices 7:10
#                    ->  q = [3, 4]
#   e       Number of exponential cones present in problem
#   G       The G matrix in ECOSMatrix format.
#   A       The A matrix in ECOSMatrix format.
#           Can be all nothing if no equalities are present.
#   c       Objective coefficients, length(c) == n
#   h       RHS for inequality constraints, length(h) == m
#   b       RHS for equality constraints, length(b) == b (can be nothing)
# Returns a pointer to the ECOS pwork structure (Cpwork in ECOS.jl). See
# types.jl for more information.
# **NOTE**: ECOS retains references to the problem data passed in here.
# You *must* ensure that G, A, c, h, and b are not freed until after cleanup(), otherwise
# memory corruption will occur.
function setup(n::Int, m::Int, p::Int, l::Int, ncones::Int, q::Union{Vector{Int},Nothing}, e::Int,
               G::ECOSMatrix, A::Union{ECOSMatrix, Nothing},
               c::Vector{Float64}, h::Vector{Float64}, b::Union{Vector{Float64},Nothing}; kwargs...)
    # Convert to canonical forms
    q = (q == nothing) ? convert(Ptr{Clong}, C_NULL) : convert(Vector{Clong},q)
    Apr = (A == nothing) ? convert(Ptr{Cdouble}, C_NULL) : A.pr
    Ajc = (A == nothing) ? convert(Ptr{Cdouble}, C_NULL) : A.jc
    Air = (A == nothing) ? convert(Ptr{Cdouble}, C_NULL) : A.ir
    b = (b == nothing) ? convert(Ptr{Cdouble}, C_NULL) : b
    problem_ptr = ccall((:ECOS_setup, ECOS.ecos), Ptr{Cpwork},
        (Clong, Clong, Clong, Clong, Clong, Ptr{Clong}, Clong,
         Ptr{Cdouble}, Ptr{Clong}, Ptr{Clong},
         Ptr{Cdouble}, Ptr{Clong}, Ptr{Clong},
         Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}),
        n, m, p, l, ncones, q, e,
        G.pr, G.jc, G.ir,
        Apr, Ajc, Air,
        c, h, b)

    problem_ptr != C_NULL || error("ECOS failed to construct problem.")

    if !isempty(kwargs)
        problem = unsafe_load(problem_ptr)
        problem.stgs != C_NULL || error("ECOS returned a malformed settings struct.")
        settings = unsafe_load(problem.stgs)

        options = Dict{Symbol, Any}()
        for (k,v) in kwargs
            options[k] = v
        end

        new_settings = Csettings([
            setting in keys(options) ?
            convert(fieldtype(typeof(settings), setting), options[setting]) :
            getfield(settings, setting)
            for setting in fieldnames(typeof(settings))]...)

        unsafe_store!(problem.stgs, new_settings)
    end

    problem_ptr
end

# solve
# Solves the provided problem. Results are stored inside the structure,
# but currently there is no convenient interface-provided way to access
# this - use MathProgBase interface.
function solve(problem::Ptr{Cpwork})
    exitflag = ccall((:ECOS_solve, ECOS.ecos), Clong, (Ptr{Cpwork},), problem)
end

# cleanup
# Frees memory allocated by ECOS for the problem.
# The optional keepvars argument is number of variables to NOT free.
function cleanup(problem::Ptr{Cpwork}, keepvars::Int = 0)
    ccall((:ECOS_cleanup, ECOS.ecos), Cvoid, (Ptr{Cpwork}, Clong), problem, keepvars)
end

include("MPB_wrapper.jl")
include("MOI_wrapper.jl")

end # module


text:
__precompile__(true)

"""
*Extensions for the Julia docsystem.*

# Introduction

This package provides a collection of useful extensions for Julia's built-in docsystem.
These are features that are still regarded as "experimental" and not yet mature enough to be
considered for inclusion in `Base`, or that have sufficiently niche use cases that including
them with the default Julia installation is not seen as valuable enough at this time.

Currently `DocStringExtensions.jl` exports a collection of so-called "abbreviations", which
can be used to add useful automatically generated information to docstrings. These include
information such as:

  * simplified method signatures;
  * documentation for the individual fields of composite types;
  * import and export lists for modules;
  * and source-linked lists of methods applicable to a particular docstring.

Users are most welcome to suggest additional abbreviation ideas, or implement and submit
them themselves. Julia's strong support for program introspection makes this a reasonably
straight forward process.

Details of the currently available abbreviations can be viewed in their individual
docstrings listed below in the "Exports" section.

# Examples

In simple terms an abbreviation can be used by simply interpolating it into a suitable
docstring. For example:

```julia
using DocStringExtensions

\"""
A short summary of `func`...

\$(SIGNATURES)

where `x` and `y` should both be positive.

# Details

Some details about `func`...
\"""
func(x, y) = x + y
```

`\$(SIGNATURES)` will be replaced in the above docstring with

````markdown
# Signatures

```julia
func(x, y)
```
````

The resulting generated content can be viewed via Julia's `?` mode or, if `Documenter.jl` is
set up, the generated external documentation.

The advantage of using [`SIGNATURES`](@ref) (and other abbreviations) is that docstrings are
less likely to become out-of-sync with the surrounding code. Note though that references to
the argument names `x` and `y` that have been manually embedded within the docstring are, of
course, not updated automatically.

# Exports
$(EXPORTS)

# Imports
$(IMPORTS)

!!! note

    This package is installable on Julia 0.4, but does not provide any features.
    Abbreviations can still be interpolated into docstrings but will expand to an empty
    string rather than the expected auto-generated content that it would on Julia 0.5 and
    above.

    This allows the package to be added as a dependancy for packages that still support
    Julia 0.4 to allow them to begin using abbreviations without needing to conditionally
    import this pacakage.

"""
module DocStringExtensions

# Imports.

using Compat

# Exports.

export FIELDS, EXPORTS, METHODLIST, IMPORTS, SIGNATURES, TYPEDEF

# Note:
#
# This package is installable on Julia 0.4, but will not provide any features. All the
# exported "abbreviations" are just defined as `""` so that interpolating them into
# docstrings will not be visible.
#
# Only on Julia 0.5 and above is the real package defined. "Abbreviations" will expand to
# their correct definitions and all utility functions will be defined as well.
#
# This is done so that the package can be used with other packages that still support Julia
# 0.4 without requiring this package as a "conditional" dependancy. The effects of this
# package will be invisible to users on 0.4.
#
# Once `DocStringExtensions` drops support for Julia 0.4 this condition can be removed.
#
if VERSION < v"0.5.0-dev"

const FIELDS, EXPORTS, METHODLIST, IMPORTS, SIGNATURES, TYPEDEF = "", "", "", "", "", ""

else # VERSION < v"0.5.0-dev"

# Includes.

include("utilities.jl")
include("abbreviations.jl")

#
# Bootstrap abbreviations.
#
# Within the package itself we would like to be able to use the abbreviations that have been
# implemented. To do this we need to delay evaluation of the interpolated abbreviations
# until they have all been defined. We use `Symbol`s in place of the actual constants, such
# as `METHODLIST` which is written as `:METHODLIST` instead.
#
# The docstring for the module itself, defined at the start of the file, does not need to
# use `Symbol`s since with the way `@doc` works the module docstring gets inserted at the
# end of the module definition and so has all the definitions already defined.
#
let λ = s -> isa(s, Symbol) ? getfield(DocStringExtensions, s) : s
    for (binding, multidoc) in Docs.meta(DocStringExtensions)
        for (typesig, docstr) in multidoc.docs
            docstr.text = Core.svec(map(λ, docstr.text)...)
        end
    end
end

end # VERSION < v"0.5.0-dev"

end # module


text:
module WorldOceanAtlasTools

using DataDeps
using Downloads
using NCDatasets
using DataFrames
using Match
using Statistics
using StatsBase
using Unitful
using OceanGrids
using NearestNeighbors

include("names.jl")
include("citations.jl")
#include("interpolations.jl")
include("convert_to_Unitful.jl")
include("functions.jl")

#=====================================
    TODO list?
=====================================#

# - Make a loop and register everything
#    (Only the data to be read or info'd will be downloaded.)
# - Update to WOA 2018

end # module


text:
# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENCE in the project root.
# ------------------------------------------------------------------

"""
    Remainder([total])

The transform that takes a table with columns `x₁, x₂, …, xₙ`
and returns a new table with an additional column containing the
remainder value `xₙ₊₁ = total .- (x₁ + x₂ + ⋯ + xₙ)` If the `total`
value is not specified, then default to the maximum sum across rows.
"""
struct Remainder <: Transform
  total::Union{Float64,Nothing}
end

Remainder() = Remainder(nothing)

isrevertible(::Type{Remainder}) = true

assertions(::Type{Remainder}) = [TT.assert_continuous]

function _cache(transform::Remainder, table)
  # design matrix
  X = Tables.matrix(table)

  # find total across rows
  if !isnothing(transform.total)
    transform.total
  else
    maximum(sum(X, dims=2))
  end
end

function _apply(transform::Remainder, table, cache)
  # basic checks
  for assertion in assertions(transform)
    assertion(table)
  end

  # design matrix
  X = Tables.matrix(table)

  # retrieve the total
  total = cache

  # make sure that the total is valid
  @assert all(x -> x ≤ total, sum(X, dims=2)) "the sum for each row must be less than total"

  # original column names
  names = Tables.columnnames(table)

  # create a column with the remainder
  S = sum(X, dims=2)
  Z = [X (total .- S)]

  # create new column name
  rname = :remainder
  while rname ∈ names
    rname = Symbol(rname,:_)
  end
  names = (names..., rname)

  # table with new column
  𝒯 = (; zip(names, eachcol(Z))...)
  newtable = 𝒯 |> Tables.materializer(table)

  newtable, total
end

apply(transform::Remainder, table) =
  _apply(transform, table, _cache(transform, table))

function revert(::Remainder, newtable, cache)
  names = Tables.columnnames(newtable)
  TT.Reject(last(names))(newtable)
end

reapply(transform::Remainder, table, cache) =
  _apply(transform, table, cache) |> first


text:
__precompile__(false)

module TinyCassette

using Logging

include("overdub.jl")

end # module


text:
function split_id(sid::AbstractString; c::String=".")
  id = String.(split(sid, c, keepempty=true))
  L = length(id)
  if L < 4
    id2 = Array{String, 1}(undef, 4-L)
    fill!(id2, "")
    append!(id, id2)
  elseif L > 4
    deleteat!(id, 5:L)
  end
  return id
end


text:
#### DC Optimal Power Shutoff ####

# This file provides a basic implentation of the DC Optimal Power Shutoff
# problem for experimentation on algorithm and problem variations.

# This file can be run by calling `include("dc-ops.jl")` from the Julia REPL

# Developed by Noah Rhodes(@noahrhodes) and Line Roald (@lroald)

###############################################################################
# 0. Initialization
###############################################################################
using Pkg; Pkg.activate(".")
# Load Julia Packages
#--------------------
using PowerModels
using PowerModelsWildfire
using GLPK
using JuMP

# Load System Data
# ----------------
powermodels_path = joinpath(dirname(pathof(PowerModelsWildfire)), "..")

# Datasets with risk data, available in the PowerModelsWildfire.jl package
file_name = "$(powermodels_path)/test/networks/RTS_GMLC_risk.m"
# file_name = "$(powermodels_path)/test/networks/case14_risk.m"
# file_name = "$(powermodels_path)/test/networks/case5_risk_sys1.m"
# file_name = "$(powermodels_path)/test/networks/case5_risk_sys2.m "

# load the data file
data = PowerModels.parse_file(file_name)

# Modify risk weighting
# data["risk_weight"] = 0.2

# Add zeros to turn linear objective functions into quadratic ones
# so that additional parameter checks are not required
PowerModels.standardize_cost_terms!(data, order=2)

# Adds reasonable rate_a values to branches without them
PowerModels.calc_thermal_limits!(data)

# use build_ref to filter out inactive components
ref = PowerModels.build_ref(data)
ref_add_on_off_va_bounds!(ref, data)
ref = ref[:it][:pm][:nw][0]
# Note: ref contains all the relevant system parameters needed to build the OPS model
# When we introduce constraints and variable bounds below, we use the parameters in ref.

###############################################################################
# 1. Building the Optimal Power Shutoff Model
###############################################################################

# Initialize a JuMP Optimization Model
#-------------------------------------
model = Model(GLPK.Optimizer)

# Add Optimization and State Variables
# ------------------------------------

# Add voltage angles va for each bus
@variable(model, va[i in keys(ref[:bus])])
# note: [i in keys(ref[:bus])] adds one `va` variable for each bus in the network

# Add active power generation variable pg for each generator
@variable(model, pg[i in keys(ref[:gen])])

# Add power flow variables p to represent the active power flow for each branch
@variable(model,  p[(l,i,j) in ref[:arcs_from]])

# Build JuMP expressions for the value of p[(l,i,j)] and p[(l,j,i)] on the branches
p_expr = Dict([((l,i,j), 1.0*p[(l,i,j)]) for (l,i,j) in ref[:arcs_from]])
p_expr = merge(p_expr, Dict([((l,j,i), -1.0*p[(l,i,j)]) for (l,i,j) in ref[:arcs_from]]))
# note: this is used to make the definition of nodal power balance simpler

# Add binary variables to represent that state of all components
@variable(model, z_bus[i in keys(ref[:bus])], Bin)
@variable(model, z_gen[i in keys(ref[:gen])], Bin)
@variable(model, z_branch[(l,i,j) in ref[:arcs_from]], Bin)

# Add continous load shed variable
@variable(model, 0.0 <= x_load[l in keys(ref[:load])] <= 1.0)
@variable(model, 0.0 <= x_shunt[s in keys(ref[:shunt])] <= 1.0)

# Add Objective Function
# ----------------------
# Maximize power delivery while minimizing wildfire risk

if haskey(ref, :risk_weight)
    alpha = ref[:risk_weight]
else
    @warn "network data should specify risk_weight, using 0.5 as a default"
    alpha = 0.5
end

for comp_type in [:gen, :load, :bus, :branch]
    for (id,comp) in  ref[comp_type]
        if ~haskey(comp, "power_risk")
            @warn "$(comp_type) $(id) does not have a power_risk value, using 0.1 as a default"
            comp["power_risk"] = 0.1
        end
    end
end

JuMP.@objective(model, Max,
    (1-alpha)*(
            sum(x_load[i]*load["pd"] for (i,load) in ref[:load])
    )
    - alpha*(
        sum(z_gen[i]*gen["power_risk"] for (i,gen) in ref[:gen])
        + sum(z_bus[i]*bus["power_risk"] for (i,bus) in ref[:bus])
        + sum(z_branch[(l,i,j)]*ref[:branch][l]["power_risk"] for (l,i,j) in ref[:arcs_from])
        + sum(x_load[i]*load["power_risk"] for (i,load) in ref[:load])
    )
)

# Add Constraints
# ---------------

# Fix the voltage angle to zero at the reference bus
for (i,bus) in ref[:ref_buses]
    @constraint(model, va[i] == 0)
end

# Nodal power balance constraints
for (i,bus) in ref[:bus]
    # Build a list of the loads and shunt elements connected to the bus i
    bus_loads = [(l,ref[:load][l]) for l in ref[:bus_loads][i]]
    bus_shunts = [(s,ref[:shunt][s]) for s in ref[:bus_shunts][i]]

    # Active power balance at node i
    @constraint(model,
        sum(p_expr[a] for a in ref[:bus_arcs][i]) ==                    # sum of active power flow on lines from bus i +
        sum(pg[g] for g in ref[:bus_gens][i]) -                         # sum of active power generation at bus i -
        sum(x_load[l]*load["pd"] for (l,load) in bus_loads) -           # sum of active load * load shed  at bus i -
        sum(x_shunt[s]*shunt["gs"] for (s,shunt) in bus_shunts)*1.0^2   # sum of active shunt element injections * shed at bus i
    )
end

# Branch power flow physics and limit constraints
for (i,branch) in ref[:branch]
    # Build the from variable id of the i-th branch, which is a tuple given by (branch id, from bus, to bus)
    f_idx = (i, branch["f_bus"], branch["t_bus"])
    p_fr = p[f_idx]                     # p_fr is a reference to the optimization variable p[f_idx]
    z_br = z_branch[f_idx]
    va_fr = va[branch["f_bus"]]         # va_fr is a reference to the optimization variable va on the from side of the branch
    va_to = va[branch["t_bus"]]         # va_fr is a reference to the optimization variable va on the to side of the branch
    # Compute the branch parameters and transformer ratios from the data
    g, b = PowerModels.calc_branch_y(branch)

    # Voltage angle difference limit
    JuMP.@constraint(model, va_fr - va_to <= branch["angmax"]*z_br + ref[:off_angmax]*(1-z_br))
    JuMP.@constraint(model, va_fr - va_to >= branch["angmin"]*z_br + ref[:off_angmin]*(1-z_br))

    # DC Power Flow Constraint
    if b <= 0
        JuMP.@constraint(model, p_fr <= -b*(va_fr - va_to + ref[:off_angmax]*(1-z_br)) )
        JuMP.@constraint(model, p_fr >= -b*(va_fr - va_to + ref[:off_angmin]*(1-z_br)) )
    else # account for bound reversal when b is positive
        JuMP.@constraint(model, p_fr >= -b*(va_fr - va_to + ref[:off_angmax]*(1-z_br)) )
        JuMP.@constraint(model, p_fr <= -b*(va_fr - va_to + ref[:off_angmin]*(1-z_br)) )
    end

    # Thermal limit
    JuMP.@constraint(model, p_fr <=  branch["rate_a"]*z_br)
    JuMP.@constraint(model, p_fr >= -branch["rate_a"]*z_br)

    # Connectivity constraint
    JuMP.@constraint(model, z_br <= z_bus[branch["f_bus"]])
    JuMP.@constraint(model, z_br <= z_bus[branch["t_bus"]])
end

# Generator constraints
for (i,gen) in ref[:gen]
    # Power limit
    JuMP.@constraint(model, pg[i] <= gen["pmax"]*z_gen[i])
    JuMP.@constraint(model, pg[i] >= gen["pmin"]*z_gen[i])

    # Connectivity constraint
    JuMP.@constraint(model, z_gen[i] <= z_bus[gen["gen_bus"]])
end

# Load constraints
for (i, load) in ref[:load]
    # Connectivity constraints
    JuMP.@constraint(model, x_load[i] <= z_bus[load["load_bus"]])
end

###############################################################################
# 3. Solve the Optimal Power Flow Model and Review the Results
###############################################################################

# Solve the optimization problem
optimize!(model)

# Check that the solver terminated without an error
println("The solver termination status is $(termination_status(model))")

# Check the value of the objective function
println("The objective value is $(objective_value(model))")
total_load = sum(value(x_load[i])*load["pd"] for (i,load) in ref[:load])
println("The load served is $total_load ")
total_risk = sum(value(z_branch[(l,i,j)])*ref[:branch][l]["power_risk"] for (l,i,j) in ref[:arcs_from]) +
             sum(value(z_gen[g])*gen["power_risk"] for (g,gen) in ref[:gen]) +
             sum(value(z_bus[i])*bus["power_risk"] for (i,bus) in ref[:bus]) +
             sum(value(x_load[i])*load["power_risk"] for (i,load) in ref[:load])
println("The system risk is $total_risk")

# Get inactive devices
for l in sort(collect(keys(ref[:branch])))
    f_idx = (l, ref[:branch][l]["f_bus"], ref[:branch][l]["t_bus"])
    value(z_branch[f_idx])==0 ? println("Branch $l is disabled.") : false
end
for g in sort(collect(keys(ref[:gen])))
    value(z_gen[g])==0 ? println("Generator $g is disabled.") : false
end
for i in sort(collect(keys(ref[:bus])))
    value(z_bus[i])==0 ? println("Bus $i is disabled.") : false
end
for i in sort(collect(keys(ref[:load])))
    value(x_load[i])!=1 ? println("Load $i is serving $(value(x_load[i])*100)% of demand.") : false
end

max_stars_repo_name: WISPO-POP/PowerModelsWildfireAnnex

max_issues_repo_name: WISPO-POP/PowerModelsWildfireAnnex

max_forks_repo_name: WISPO-POP/PowerModelsWildfireAnnex


text:
## Field access through these functions to reserve dot-getting for keys
@inline getaxes(x::VarAxes) = getaxes(typeof(x))
@inline getaxes(Ax::Type{<:Axes}) where {Axes<:VarAxes} = map(x->x(), (Ax.types...,))

getaxes(x) = ()

# Since we aren't really using the standard approach to indexing, this will forward things to
# the correct methods
Base.to_indices(x::ComponentArray, i::Tuple) = i
Base.to_indices(x::ComponentArray, i::Tuple{Vararg{Union{Integer, CartesianIndex}, N}}) where N = i
Base.to_indices(x::ComponentArray, i::Tuple{Vararg{Int64}}) where N = i
Base.to_index(x::ComponentArray, i) = i

# Get AbstractAxis index
@inline Base.getindex(::AbstractAxis, idx::FlatIdx) = ComponentIndex(idx)
@inline Base.getindex(ax::AbstractAxis, ::Colon) = ComponentIndex(:, ax)
@inline Base.getindex(::AbstractAxis{IdxMap}, s::Symbol) where IdxMap =
    ComponentIndex(getproperty(IdxMap, s))

# Get ComponentArray index
Base.@propagate_inbounds Base.getindex(x::ComponentArray, idx::CartesianIndex) = getdata(x)[idx]
Base.@propagate_inbounds Base.getindex(x::ComponentArray, idx::FlatIdx...) = getdata(x)[idx...]
Base.@propagate_inbounds function Base.getindex(x::ComponentArray, idx::FlatOrColonIdx...)
    axs = map((ax, i) -> getindex(ax, i).ax, getaxes(x), idx)
    axs = remove_nulls(axs...)
    return ComponentArray(getdata(x)[idx...], axs...)
end
Base.@propagate_inbounds Base.getindex(x::ComponentArray, ::Colon) = getdata(x)[:]
@inline Base.getindex(x::ComponentArray, ::Colon...) = x
Base.@propagate_inbounds Base.getindex(x::ComponentArray, idx...) = getindex(x, toval.(idx)...)
@inline Base.getindex(x::ComponentArray, idx::Val...) = _getindex(x, idx...)

# Set ComponentArray index
Base.@propagate_inbounds Base.setindex!(x::ComponentArray, v, idx::FlatIdx...) = setindex!(getdata(x), v, idx...)
Base.@propagate_inbounds Base.setindex!(x::ComponentArray, v, ::Colon) = setindex!(getdata(x), v, :)
Base.@propagate_inbounds Base.setindex!(x::ComponentArray, v, idx...) = setindex!(x, v, toval.(idx)...)
@inline Base.setindex!(x::ComponentArray, v, idx::Val...) = _setindex!(x, v, idx...)

# Property access for CVectors goes through _get/_setindex
@inline Base.getproperty(x::ComponentVector, s::Symbol) = _getindex(x, Val(s))
@inline Base.getproperty(x::ComponentVector, s::Val) = _getindex(x, s)

@inline Base.setproperty!(x::ComponentVector, s::Symbol, v) = _setindex!(x, v, Val(s))
@inline Base.setproperty!(x::ComponentVector, s::Val, v) = _setindex!(x, v, s)

# Generated get and set index methods to do all of the heavy lifting in the type domain
@generated function _getindex(x::ComponentArray, idx...)
    ci = getindex.(getaxes(x), getval.(idx))
    inds = map(i -> i.idx, ci)
    axs = map(i -> i.ax, ci)
    axs = remove_nulls(axs...)
    # the index must be valid after computing `ci`
    :(Base.@_inline_meta; @inbounds ComponentArray(Base.maybeview(getdata(x), $inds...), $axs...))
end

@generated function _setindex!(x::ComponentArray, v, idx...)
    ci = getindex.(getaxes(x), getval.(idx))
    inds = map(i -> i.idx, ci)
    # the index must be valid after computing `ci`
    return :(Base.@_inline_meta; @inbounds setindex!(getdata(x), v, $inds...))
end

Base.@propagate_inbounds Base.view(x::ComponentArray, idx::ComponentArrays.FlatIdx...) = view(getdata(x), idx...)
Base.@propagate_inbounds Base.view(x::ComponentArray, idx...) = _getindex(x, toval.(idx)...)


text:
Base.@kwdef struct SparseSheeps{L} <: Sheeps
    gridsize::Int
    capacity::Int
    available_sheeps::Int = 0
end


text:
@testset "FFTGS" begin
  𝒟 = RegularGrid(100,100)
  problem = SimulationProblem(𝒟, :z=>Float64, 3)

  Random.seed!(2019)
  solver = FFTGS(:z => (variogram=GaussianVariogram(range=10.),))
  solution = solve(problem, solver)

  if visualtests
    @plottest plot(solution,size=(900,300)) joinpath(datadir,"FFT-iso.png") !isCI
  end

  Random.seed!(2019)
  solver = FFTGS(:z => (variogram=GaussianVariogram(distance=aniso2distance([20.,5.],[0.])),))
  solution = solve(problem, solver)

  if visualtests
    @plottest plot(solution,size=(900,300)) joinpath(datadir,"FFT-aniso.png") !isCI
  end
end


text:
using Printf

analytical(n::Integer) = sum(factorial(n) / big(n) ^ i / factorial(n - i) for i = 1:n)

function test(n::Integer, times::Integer = 1000000)
    c = 0
    for i = range(0, times)
        x, bits = 1, 0
        while (bits & x) == 0
            c += 1
            bits |= x
            x = 1 << rand(0:(n - 1))
        end
    end
    return c / times
end

function main(n::Integer)
    println(" n\tavg\texp.\tdiff\n-------------------------------")
    for n in 1:n
        avg = test(n)
        theory = analytical(n)
        diff = (avg / theory - 1) * 100
        @printf(STDOUT, "%2d %8.4f %8.4f %6.3f%%\n", n, avg, theory, diff)
    end
end

main(20)

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
#  Copyright 2017, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
#  This Source Code Form is subject to the terms of the Mozilla Public
#  License, v. 2.0. If a copy of the MPL was not distributed with this
#  file, You can obtain one at http://mozilla.org/MPL/2.0/.
#############################################################################
# JuMP
# An algebraic modeling language for Julia
# See http://github.com/JuliaOpt/JuMP.jl
#############################################################################
# test/utilities.jl
# Utility functions used across JuMP's tests.
#############################################################################

macro test_expression(expr)
    esc(quote
            @test JuMP.isequal_canonical(@expression(model, $expr), $expr)
    end)
end

macro test_expression_with_string(expr, str)
    esc(quote
            @test string(@inferred $expr) == $str
            @test_expression $expr
    end)
end

# Test that the macro call `m` throws an error exception during pre-compilation
macro test_macro_throws(errortype, m)
    # See https://discourse.julialang.org/t/test-throws-with-macros-after-pr-23533/5878
    :(@test_throws $errortype try @eval $m catch err; throw(err.error) end)
end


text:
using Documenter, MuJoCo

const PAGES = [
    "Home" => "index.md",
    "API" => [
        "api/mujoco.md",
        "api/mjcore.md",
    ],
]

makedocs(;
    modules = [MuJoCo],
    format = Documenter.HTML(prettyurls = get(ENV, "GITHUB_ACTIONS", nothing) == "true"),
    pages = PAGES,
    sitename = "MuJoCo.jl",
    authors = "Colin Summers",
    clean = true,
    doctest = true,
    checkdocs = :exports,
    linkcheck = :true,
    linkcheck_ignore = [r"^https://github.com/Lyceum/.*/actions"],
)

deploydocs(repo = "github.com/Lyceum/MuJoCo.jl.git")


text:
@testset "Total degree" begin
    @polyvar x y z

    H, solutions = totaldegree(StraightLineHomotopy{Complex128}, [x^2+y+1, x^3*y-2])
    @test length(solutions) == 8
    @test eltype(solutions) == Vector{Complex128}
    for sol in solutions
        @test norm(H(sol, 1.0)) ≈ 0 atol=1e-14
    end
    H, solutions = totaldegree(StraightLineHomotopy{Complex128}, [x^2+y+1, x^3*y-2], unitroots=true)
    for sol in solutions
        @test norm(sol) ≈ norm(ones(2)) atol=1e-14
    end

    f = x^2+y^2+y*z
    g = x*z + y * z

    H, s = totaldegree(GeodesicOnTheSphere, [f, g])
    @test nvariables(H) == 3

end

@testset "randomsystem" begin
    f = randomsystem(Complex128, 1, 4, density=1.0)[]
    d = FP.degree(f)
    @test length(f.coefficients) == binomial(d + 4, 4)

    f = randomsystem(Complex128, 1, 4, density=0.5)[]
    d = FP.degree(f)
    @test length(f.coefficients) == ceil(Int, binomial(d + 4, 4) * 0.5)

    @test_throws AssertionError randomsystem(Complex128, 1, 4, density=0.0)
    @test_throws AssertionError randomsystem(Complex128, 1, 4, density=-12.0)
    @test_throws AssertionError randomsystem(Complex128, 1, 4, density=213.0)
end

@testset "randomhomotopy" begin
    H, solutions = randomhomotopy(StraightLineHomotopy{Complex64}, 3)
    @test H isa StraightLineHomotopy{Complex64}

    H, solutions = randomhomotopy(StraightLineHomotopy, 3)
    @test H isa StraightLineHomotopy{Complex128}
    @test length(H) == 3
    @test nvariables(H) == 3
    @test length(solutions) == prod(FP.degree.(H.start))

    for sol in solutions
        @test norm(H(sol, 1.0)) ≈ 0 atol=1e-14
    end
end

@testset "gammatrick" begin
    @polyvar x y

    H = StraightLineHomotopy{Complex128}([x^3*y-2, x^2+y+1], [x^2+y+1, x^3*y-2])
    G = GeodesicOnTheSphere{Complex128}([x^3*y-2, x^2+y+1], [x^2+y+1, x^3*y-2])
    H2 = deepcopy(H)

    x = rand(2)
    t = 0.535

    @test H(x, t) == H2(x, t)

    gammatrick!(H)
    @test H(x, t) != H2(x, t)

    H = deepcopy(H2)
    H1 = deepcopy(H2)
    gammatrick!(H, 21312)
    gammatrick!(H1, 21312)

    @test H(x, t) == H1(x, t)

    H = deepcopy(H2)
    gammatrick!(H, 0.5)
    @test H(x, 1.0) ≈ 0.5 * H2(x, 1.0)

    G2 = deepcopy(G)
    x = rand(3)
    gammatrick!(G)
    @test G(x, t) != G2(x, t)
end


text:
export AbstractSys, AbstractSystem, AbstractDiffSystem, AbstractVelSystem
export System, DiffSystem, StabilizedDiffSystem, NominalDiffSystem, VelSystem, RecurrentSystem, RecurrentDiffSystem
export simulate, predict, jacobians

abstract type AbstractSys <: LTVModelsBase.AbstractModel end
abstract type AbstractSystem <: AbstractSys end
abstract type AbstractDiffSystem <: AbstractSys end
abstract type AbstractVelSystem <: AbstractSys end

import LTVModelsBase: simulate, predict
import Flux.Tracker.data

@with_kw struct System{T} <: AbstractSystem
    m::T
    nx::Int
    nu::Int
    h::Float64 = 1.0
end
function System(nx::Int,nu::Int, num_params::Int, activation::Function, h=1)
    ny = nx
    np = num_params
    m  = Chain(Dense(nx+nu,np, activation), Dense(np,np, activation), Dense(np, ny))
    System(m, nx, nu, h)
end
(m::System)(x) = m.m(x)

@with_kw struct DiffSystem{T} <: AbstractDiffSystem
    m::T
    nx::Int
    nu::Int
    h::Float64 = 1.0
end
function DiffSystem(nx::Int,nu::Int, num_params::Int, activation::Function, h=1)
    ny = nx
    np = num_params
    m  = Chain(Dense(nx+nu,np, activation), Dense(np,np, activation), Dense(np, ny))
    DiffSystem(m, nx, nu, h)
end
(m::DiffSystem)(x) = m.m(x)+x[1:m.nx,:]

@with_kw struct StabilizedDiffSystem{T} <: AbstractDiffSystem
    τ::Float64
    m::T
    nx::Int
    nu::Int
    h::Float64 = 1.0
end
function StabilizedDiffSystem(τ::Float64,nx::Int,nu::Int, num_params::Int, activation::Function, h=1)
    ny = nx
    np = num_params
    m  = Chain(Dense(nx+nu,np, activation), Dense(np, ny))
    StabilizedDiffSystem(τ, m, nx, nu, h)
end
(m::StabilizedDiffSystem)(x) = m.m(x)+ m.τ*x[1:m.nx,:]

@with_kw struct NominalDiffSystem{T} <: AbstractDiffSystem
    A::Matrix{Float64}
    B::Matrix{Float64}
    m::T
    nx::Int
    nu::Int
    h::Float64 = 1.0
end
function NominalDiffSystem(A::AbstractMatrix,B::AbstractMatrix,nx::Int,nu::Int, num_params::Int, activation::Function, h=1)
    ny = nx
    np = num_params
    m  = Chain(Dense(nx+nu,np, activation), Dense(np, ny))
    NominalDiffSystem(A, B, m, nx, nu, h)
end
(m::NominalDiffSystem)(x) = m.m(x) .+ m.A*x[1:m.nx,:] .+ m.B*x[1+m.nx:end,:]

@with_kw struct VelSystem{T} <: AbstractVelSystem
    m::T
    nx::Int
    nu::Int
    h::Float64 = 1.0
end
function VelSystem(nx::Int,nu::Int, num_params::Int, activation::Function, h=1)
    ny = nx
    np = num_params
    m = Chain(Dense(nx+nu,np, activation),  Dense(np, ny))
    VelSystem(m, nx, nu, h)
end
(m::VelSystem)(x) = m.m(x)

# Non-flux systems =========================================================================
tovec(w::Vector{<:Matrix}) = vcat([vec(w) for w in w]...)

function pred(w,x,nx)
    state = x
    for i=1:2:length(w)-2
        state = tanh.(w[i]*state .+ w[i+1])
    end
    return w[end-1]*state .+ w[end]
end
pred(m,x) = pred(m.w,x,m.nx)
predd(m,x) = predd(m.w,x,m.nx)
predd(w,x,nx) = pred(w,x,nx) .+ x[1:nx,:]

function pred_jac(w,x,nx)
    l = x
    J = Matrix{eltype(w[1])}(I,length(x),length(x))
    for i = 1:2:length(w)-2
        W,b = w[i], w[i+1]
        l   = W*l .+ b
        ∇σ  = ∇tanh(l)
        ∇a  = W
        l  .= tanh.(l) # Reused l to save allocations
        J   = ∇σ * ∇a * J
    end
    J = w[end-1] * J # Linear output layer
    # J += [Matrix{Float64}(I, nx, nx) zeros(nx)] # Not needed since we'll only use the difference
    # @assert isapprox(J, hcat(fdjac(x->predd(w,x,nx),x)...), atol=1e-4)
    return w[end-1]*l .+ w[end] , J#.+ x[1:nx,:], J
end
∇tanh(x) = Matrix(Diagonal((sech.(x).^2)[:]))
∇relu(x) = Matrix(Diagonal(vec(ifelse.(x .<= 0, 0., 1.))))
function ∇swish!(x)
    σx = σ.(x)
    x .= x.*σx
    σx .= x .+ σx .* (1 .- x)
    Matrix(Diagonal(vec(σx)))
end

@with_kw struct ADSystem <: AbstractSystem
    w::NTuple{6, Matrix{Float64}}
    sizes = ((num_params,nx+nu), (num_params,1), (nx,num_params), (nx,1))
    nx::Int
    nu::Int
    h::Float64 = 1.0
end

@with_kw struct ADDiffSystem <: AbstractDiffSystem
    w::NTuple{6, Matrix{Float64}}
    sizes = ((num_params,nx+nu), (num_params,1), (nx,num_params), (nx,1))
    nx::Int
    nu::Int
    h::Float64 = 1.0
end
function ADSystem(nx::Int,nu::Int, num_params::Int, activation::Function, h=1)
    sizes = ((num_params,nx+nu), (num_params,1), (num_params,num_params), (num_params,1), (nx,num_params), (nx,1))
    w = ( Flux.glorot_uniform(sizes[1]...), zeros(Float64,sizes[2]),Flux.glorot_uniform(sizes[3]...),  zeros(Float64,sizes[4]), Flux.glorot_uniform(sizes[5]...),  zeros(Float64,sizes[6]) )
    model(x)    = pred(w,x,nx)
    ADSystem(w,sizes,nx,nu,h)
end
function ADDiffSystem(nx::Int,nu::Int, num_params::Int, activation::Function, h=1)
    sizes = ((num_params,nx+nu), (num_params,1), (num_params,num_params), (num_params,1), (nx,num_params), (nx,1))
    w = ( Flux.glorot_uniform(sizes[1]...), zeros(Float64,sizes[2]),Flux.glorot_uniform(sizes[3]...),  zeros(Float64,sizes[4]), Flux.glorot_uniform(sizes[5]...),  zeros(Float64,sizes[6]) )
    model(x)    = predd(w,x,nx)
    ADDiffSystem(w,sizes,nx,nu,h)
end
(m::ADSystem)(x) = pred(m, x)
(m::ADSystem)(w,x) = pred(w, x, m.nx)
(m::ADDiffSystem)(x) = predd(m, x)
(m::ADDiffSystem)(w,x) = predd(w, x, m.nx)

Flux.testmode!(m::ADSystem, b=false) = nothing
Flux.testmode!(m::ADDiffSystem, b=false) = nothing

# Recurrent systems ========================================================================
@with_kw struct RecurrentSystem{T} <: AbstractSystem
    m::T
    nx::Int
    nu::Int
    h::Float64 = 1.0
end
function RecurrentSystem(nx::Int,nu::Int, num_params::Int, activation::Function, h=1)
    ny = nx
    np = num_params
    m  = Chain(Dense(nx+nu,np, activation), RNN(np,np), Dense(np,np,activation), Dense(np, ny))
    RecurrentSystem(m, nx, nu, h)
end
(m::RecurrentSystem)(x) = m.m(x)

@with_kw struct RecurrentDiffSystem{T} <: AbstractDiffSystem
    m::T
    nx::Int
    nu::Int
    h::Float64 = 1.0
end
function RecurrentDiffSystem(nx::Int,nu::Int, num_params::Int, activation::Function, h=1)
    ny = nx
    np = num_params
    m  = Chain(Dense(nx+nu,np, activation), RNN(np,np), Dense(np,np,activation), Dense(np, ny))
    RecurrentDiffSystem(m, nx, nu, h)
end
(m::RecurrentDiffSystem)(x) = m.m(x)+x[1:m.nx,:]

# ==========================================================================================

loss(m::AbstractSys) = (x,y) -> sum((m(x).-y).^2)/size(x,2)

# cost(w,x,y)  = sum(abs2, pred(w,x) .- y)/size(y,2)
function cost(f,x,y)
    yh = f(x)
    sum(abs2, yh .- y)/size(y,2) #+ 0.1sum(norm.(w).^2)
end
# cost(w,data) = sum(cost(w, d...) for d in data)/length(data)
cost(f,data) = cost(f, data...)

# WARNING: don't assign to any vector with .= in the inner loss function closure
function loss(w,x,y,mt::ADModelTrainer)
    model = mt.model
    nx, nu = model.nx, model.nu
    predfun = model isa ADSystem ? pred : predd
    function lf(w...)
        # println("Entering loss function, typeof(w): ", typeof(w), " length(w): ", length(w))
        @unpack λ,normalizer = mt
        f(x)   = predfun(w,x,nx)
        l      = cost(f,x,y)
        jac(x) = fdjac(f,x)
        J1 = jac(x[:,1])
        # sd = fill(typeof(w[1])(0.),size(J1))
        for t = 2:size(x,2)
            J2 = jac(x[:,t])
            for i in eachindex(J1)
                l += λ*sum(abs2.(J1[i].-J2[i]))
            end
            # copy!(J1,J2)
            J1 = J2
        end
        l
    end
end

predjac(mt::ADModelTrainer{<:ADSystem,<:Any}) = pred_jac
predjac(mt::ADModelTrainer{<:ADDiffSystem,<:Any}) = function (w,x,nx)
    y,J = pred_jac(w,x,nx)
    y + x[1:nx,:], J
end

function loss2(w,x,y,mt::ADModelTrainer)
    model = mt.model
    nx, nu = model.nx, model.nu
    pj = predjac(mt)
    function lf(w...)
        @unpack λ,normalizer = mt
        yh, J1 = pj(w,x[:,1],nx)
        l = sum(abs2, yh .- y[:,1])
        for t = 2:size(x,2)
            yh, J2 = pj(w,x[:,t],nx)
            l += sum(abs2, yh .- y[:,t])/size(y,2) + λ*sum(abs2.(J1.-J2))
            # copy!(J1,J2)
            J1 = J2
        end
        # l += λ/2*norm(w[1])^2 + λ/2*norm(w[2])^2
        # l += λ*maximum(abs.(w[3]*w[1]))#λ*norm(w[1])^2 + λ*norm(w[2])^2
        l
    end
end

function fdjac(f::Function,x::AbstractVector,epsilon=sqrt(eps()))
    n = length(x)
    f_x = f(x)
    shifted_x = copy(x)
    map(eachindex(x)) do i
        shifted_x[i] += epsilon
        J = @fastmath (f(shifted_x) .- f_x) ./ epsilon
        shifted_x[i] = x[i]
        J
    end
end

function find_normalizer(w,data,mt::ADModelTrainer)
    x,y = data
    model = mt.model
    sizes, nx, nu = model.sizes, model.nx, model.nu
    f(x)          = predd(w,x,nx)
    jcfg          = Diff.JacobianConfig(f, x[:,1])
    jacobian(x) = Diff.jacobian(f, x, jcfg)
    mJ = zeros(nx,nx+nu)
    for t = 1:size(x,2)
        J = jacobian(x[:,t])
        mJ .= max.(abs.(J), mJ)
    end
    mJ
end

function i2m(w,i,sizes)
    s = [1; cumsum([prod.(sizes)...]) .+ 1]
    (w,i)->reshape(view(w,s[i]:(s[i+1]-1)), sizes[i])
end

const AbstractEnsembleSystem = Vector{T} where T <: AbstractSys
const EnsembleSystem = Vector{T} where T <: AbstractSystem
const EnsembleDiffSystem = Vector{T} where T <: AbstractDiffSystem
const EnsembleVelSystem = Vector{T} where T <: AbstractVelSystem

xy(::Type{<:AbstractSys}, x, N)        = (x[:,1:N], x[:,2:N+1])
xy(::Type{<:AbstractVelSystem}, x, N)  = (x[:,1:N], x[3:4,2:N+1]) # TODO magic numbers

function get_minimum_loss(results, key)
    array = getindex.(get.(results[key]),2)
    return isempty(array) ? Inf : minimum(array)
end

function simulate(ms::AbstractEnsembleSystem,xu::AbstractMatrix, testmode=true)
    Flux.testmode!.(ms, testmode)
    xsim = copy(xu)
    nx = _nx(ms)
    for t = 2:size(xu,2)
        @ensemble xsimt = ms(xsim[:,t-1])
        xsim[1:nx,t] = data(mean(xsimt))
    end
    xsim[1:nx,:]
end

function simulate(ms::AbstractSys,xu::AbstractMatrix, testmode=true)
    ms isa Union{ADSystem,ADDiffSystem} || Flux.testmode!.(ms, testmode)
    xsim = copy(xu)
    nx = _nx(ms)
    for t = 2:size(xu,2)
        xsimt = ms(xsim[:,t-1])
        xsim[1:nx,t] = xsimt
    end
    xsim[1:nx,:]
end

function simulate(ms::EnsembleVelSystem,xu::AbstractMatrix, testmode=true)
    Flux.testmode!.(ms, testmode)
    xsim = copy(xu)
    for t = 2:size(xu,2)
        @ensemble xsimt = ms(xsim[:,t-1])
        h = ms[1].h
        xsim[1:2,t] = xsim[1:2,t-1] + h*xsim[3:4,t-1] # TODO magic numbers
        xsim[3:4,t] = data(mean(xsimt)) # TODO magic numbers
    end
    xsim[1:4,:]
end

simulate(ms::Union{AbstractSys,AbstractEnsembleSystem}, t::Trajectory; kwargs...) = simulate(ms, t.xu; kwargs...)

function predict(ms::Vector, x::AbstractMatrix, testmode=true)
    Flux.testmode!.(ms, testmode)
    @ensemble y = ms(x)
    data(mean(y)), extrema(extrema.(y))
end

predict(ms::Union{AbstractSys,Vector}, x::AbstractMatrix, testmode=true) = ms(x), nothing
function predict(ms::EnsembleVelSystem, x::AbstractArray, testmode=true)
    Flux.testmode!.(ms, testmode)
    @ensemble y = ms(x)
    yh = data(mean(y))
    h = ms[1].h
    yh = [x[1:2,1] .+ h*cumsum(yh,2); yh] # TODO magic numbers
    bounds = data.(extrema(y))
    yh, bounds
end

predict(ms::Union{AbstractSys,AbstractEnsembleSystem}, t::Trajectory; kwargs...) = predict(ms, t.xu; kwargs...)

function Flux.jacobian(ms::AbstractEnsembleSystem, x::AbstractArray, testmode=true)
    Flux.testmode!.(ms, testmode)
    # @show Flux.Tracker.grad(params.(ms[1]))
    # Flux.Tracker.zero_grad!.(params.(ms))
    jacs = [Flux.jacobian(m,x) for m in ms]
    # Flux.Tracker.zero_grad!.(params.(ms))
    jacmat = smartcat3(jacs)
    dropdims(mean(jacmat, dims=3), dims=3), dropdims(std(jacmat, dims=3), dims=3)
end

function Flux.jacobian(ms::EnsembleVelSystem, x::AbstractArray, testmode=true)
    Flux.testmode!.(ms, testmode)
    h = ms[1].h
    jacs = [[[I h*I h^2/2*eye(2)];Flux.jacobian(m,x)] for m in ms] # The h²/2*I in ∇ᵤ is an approximation since there are (very small) cross terms.
    jacmat = smartcat3(jacs)
    squeeze(mean(jacmat, 3), 3), squeeze(std(jacmat, 3), 3)
end

function Flux.jacobian(m::ADSystem, x)
    f(x) = pred(m.w, x, m.nx)
    j(x) = hcat(fdjac(f,x)...)
    j(x), nothing
end
function Flux.jacobian(m::ADDiffSystem, x)
    f(x) = predd(m.w, x, m.nx)
    j(x) = hcat(fdjac(f,x)...)
    j(x), nothing
end

Flux.jacobian(ms, t::Trajectory; kwargs...) = Flux.jacobian(ms, t.xu; kwargs...)

models(mt::ModelTrainer) = mt.models
models(mt::ADModelTrainer) = mt.model
models(systems::Vector{<:AbstractSys}) = getfield.(systems, :m)
models(results::AbstractVector) = [r[:m] for r in results]
models(results::AbstractDict) = [r[:m]]
_nx(s::AbstractSys) = s.nx
_nx(v::AbstractVector) = _nx(v[1])
_nx(trainer::ModelTrainer) = _nx(trainer.models)
_nx(trainer::ADModelTrainer) = _nx(trainer.model)

function LTVModels.KalmanModel(ms::AbstractEnsembleSystem, t::Trajectory)
    xu = t.xu
    N  = size(xu,2)
    J  = map(1:length(t)) do i
        Jm, Js = jacobian(ms, xu[:,i])
        Jm, Js
    end
    At = cat([J[i][1][:,1:t.nx] for i = 1:length(t)]..., dims=3)
    Bt = cat([J[i][1][:,t.nx+1:end] for i = 1:length(t)]..., dims=3)
    Pt = cat([diagm(J[i][2][:]).^2 for i = 1:length(t)]..., dims=3)
    LTVModels.KalmanModel(At,Bt,Pt,true)
end

function LTVModels.KalmanModel(ms::AbstractSys, t::Trajectory)
    xu = t.xu
    N  = size(xu,2)
    J  = map(1:length(t)) do i
        jacobian(ms, xu[:,i])[1]
    end
    eye(n) = Matrix{Float64}(I,n,n)
    At = cat([J[i][:,1:t.nx] for i = 1:length(t)]..., dims=3)
    Bt = cat([J[i][:,t.nx+1:end] for i = 1:length(t)]..., dims=3)
    Pt = cat([eye(length(J[1])) for i = 1:length(t)]..., dims=3)
    LTVModels.KalmanModel(At,Bt,Pt,true)
end

function jacobians(ms::AbstractVector, t, ds=1)
    msc = deepcopy(ms) # Obs, this is to not fuck up the gradients of the model parameters ??
    xu = t.xu
    N = size(xu,2)
    J = map(1:ds:N) do evalpoint
        Jm, Js = jacobian(msc, xu[:,evalpoint])
        Jm[:], Js[:]
    end
    Jm = smartcat2(getindex.(J,1))
    Js = smartcat2(getindex.(J,2))
    Jm, Js
end

function jacobians(ms, t, ds=1)
    xu = t.xu
    N = size(xu,2)
    J = map(1:ds:N) do evalpoint
        Jm, Js = jacobian(ms, xu[:,evalpoint])
        Jm[:]
    end
    Jm = smartcat2(J)
    Jm, nothing
end

function eval_pred(trainer::AbstractModelTrainer, vt)
    @unpack x,u,y = vt
    ms = models(trainer)
    nx = _nx(trainer)
    yh, bounds = predict(ms, vt.xu)
    pred = LTVModelsBase.rms(x[:,2:end].-yh[:,1:end-1]) # TODO: check time alignment
    pred
end

function eval_sim(trainer::AbstractModelTrainer, vt)
    @unpack x,u,y,xu = vt
    ms = models(trainer)
    nx = _nx(trainer)
    xh = simulate(ms, xu)
    sim = LTVModelsBase.rms(x[:,:].-xh) # TODO: check time alignment
    sim
end

function eval_jac(trainer::AbstractModelTrainer, vt, truejacfun, ds=1)
    m = models(trainer)
    mean(i-> mean(abs2.(jacobian(m,vt.xu[:,i])[1] .- truejacfun(vt.x[:,i],vt.u[:,i]))), 1:ds:length(vt)) |> √
end

function eval_jac2(trainer::AbstractModelTrainer, vt, truejacfun, ds=1)
    m = models(trainer)
    nx = vt.nx
    mean(1:ds:length(vt)) do i
        J1, J2 = jacobian(m,vt.xu[:,i])[1], truejacfun(vt.x[:,i],vt.u[:,i])
        err = 0mean(abs2.(J1[:,nx+1:end] .- J2[:,nx+1:end]))
        e1,e2 = eigvals(J1[1:nx,1:nx]), eigvals(J2[1:nx,1:nx])
        dm = abs2.(e1 .- e2')
        dists = minimum(dm, dims=1)[:] + minimum(dm, dims=2)[:]
        err  += sum(dists)
        err
    end |> √
end

function plotresults(trainer::AbstractModelTrainer, vt)
    @unpack x,u,y = vt
    ms = models(trainer)
    nx = ms[1].nx
    fig = plot(x[1:nx,2:end]', lab="True", layout=nx)
    testmode!.(ms, true)
    plot_prediction(fig, results, eval)
    plot_simulation(fig, results, eval)
    testmode!.(ms, false)
    fig
end

get_res(res,n) = getindex.(res,n)

try
    foreach(T->treelike(JacProp, T), [System, DiffSystem, StabilizedDiffSystem, NominalDiffSystem, VelSystem, RecurrentSystem, RecurrentDiffSystem])
catch err
    @error(err)
end

function smartcat2(vv)
    dim2 = length(vv)
    dim1 = length(vv[1])
    A = zeros(dim1,dim2)
    for i in eachindex(vv)
        A[:,i] = vv[i]
    end
    A
end

function smartcat3(vv)
    dim3 = length(vv)
    dim1,dim2 = size(vv[1])
    A = zeros(dim1,dim2,dim3)
    for i in eachindex(vv)
        A[:,:,i] = vv[i]
    end
    A
end


text:
function _one_iterate(cutoff_multiplier::Real, target::Real, draws::Array{T,2}, covar_matrix, tuning_parameter::Real) where T<:Real
    cutoffs = cutoff_multiplier .* sqrt.(diag(covar_matrix))
    number_of_draws = size(draws)[1]
    in_confidence_area = 0
    for i in 1:number_of_draws
        in_confidence_area += all(abs.(draws[i,:]) .< cutoffs)
    end
    mass_in_area = in_confidence_area/number_of_draws
    confidence_gap = target - mass_in_area
    return cutoff_multiplier + confidence_gap * tuning_parameter
end

function _randoms(chol, num::Integer, Seed::Integer)
    twist = MersenneTwister(Seed)
    dims = size(chol)[1]
    float_type = typeof(chol[1,1])
    array = Array{float_type,2}(undef, num, dims)
    for i in 1:num
        sobs = rand(twist, dims)
        normal_draw = quantile.(Ref(Normal()), sobs)
        scaled_draw = chol * normal_draw
        array[i,:] = scaled_draw
    end
    return array
end

"""
    get_confidence_hypercube(covar::ForwardCovariance, confidence_level::Real, data::Array{T,2}; tuning_parameter::Real = 1.0)
This returns the endpoints of a hypercube that contains confidence_level (%) of the dataset.
"""
function get_confidence_hypercube(covar::ForwardCovariance, confidence_level::Real, data::Array{T,2}; tuning_parameter::Real = 1.0, ConvergenceMetricThreshold::Real = 1e-10) where T<:Real
    # Using a univariate guess as we can get these pretty cheaply.
    guess = quantile(Normal(), 0.5*(1+confidence_level))
    # This runs once so that any error is explictly thrown and is traceable rather than being obscured by FixedPoint's try-catch
    #_ = _one_iterate.(guess, Ref(confidence_level), Ref(data), Ref(covar.covariance_), Ref(tuning_parameter))
    FP = fixed_point(x -> _one_iterate.(x, Ref(confidence_level), Ref(data), Ref(covar.covariance_), Ref(tuning_parameter)), [guess]; ConvergenceMetricThreshold = ConvergenceMetricThreshold, MaxIter = 10000)
    cutoff_multiplier = FP.FixedPoint_[1]
    cutoffs = vcat(zip(-cutoff_multiplier .* sqrt.(diag(covar.covariance_)) , cutoff_multiplier .* sqrt.(diag(covar.covariance_)))...)
    return Dict{Symbol,Tuple{T,T}}(covar.covariance_labels_ .=> cutoffs)
end

function get_confidence_hypercube(covar::ForwardCovariance, confidence_level::Real, num::Integer; tuning_parameter::Real = 1.0,  ConvergenceMetricThreshold::Real = 1e-10)
    dims = length(covar.covariance_labels_)
    data = _randoms(covar.chol_, num, 1)
    return get_confidence_hypercube(covar, confidence_level, data; tuning_parameter = tuning_parameter, ConvergenceMetricThreshold = ConvergenceMetricThreshold)
end


text:
# Tests a L2-projection of integration point values (to nodal values),
# determined from the function y = 1 + x[1]^2 + (2x[2])^2
function test_projection(order, refshape)
    element = refshape == RefCube ? Quadrilateral : Triangle
    if order == 1
        grid = generate_grid(element, (1, 1), Vec((0.,0.)), Vec((1.,1.)))
    elseif order == 2
        # grid = generate_grid(QuadraticQuadrilateral, (1, 1), Vec((0.,0.)), Vec((1.,1.)))
        grid = generate_grid(element, (1, 1), Vec((0.,0.)), Vec((1.,1.)))
    end

    dim = 2
    ip = Lagrange{dim, refshape, order}()
    ip_geom = Lagrange{dim, refshape, 1}()
    qr = Ferrite._mass_qr(ip)
    cv = CellScalarValues(qr, ip, ip_geom)

    # Create node values for the cell
    f(x) = 1 + x[1]^2 + (2x[2])^2
    # Nodal approximations for this simple grid when using linear interpolation
    f_approx(i) = refshape == RefCube ?
        [0.1666666666666664, 1.166666666666667, 4.166666666666666, 5.166666666666667][i] :
        [0.444444444444465, 1.0277777777778005, 4.027777777777753, 5.444444444444435][i]

    # analytical values
    function analytical(f)
        qp_values = []
        for cell in CellIterator(grid)
            reinit!(cv, cell)
            r = [f(spatial_coordinate(cv, qp, getcoordinates(cell))) for qp in 1:getnquadpoints(cv)]
            push!(qp_values, r)
        end
        return identity.(qp_values) # Tighten the type
    end

    qp_values = analytical(f)

    # Now recover the nodal values using a L2 projection.
    proj = L2Projector(ip, grid; geom_ip=ip_geom)
    point_vars = project(proj, qp_values, qr)
    qp_values_matrix = reduce(hcat, qp_values)
    point_vars_2 = project(proj, qp_values_matrix, qr)
    ## Old API with fe values as first arg
    proj2 = @test_deprecated L2Projector(cv, ip, grid)
    point_vars_3 = @test_deprecated project(qp_values, proj2)
    ## Old API with qr as first arg
    proj3 = @test_deprecated L2Projector(qr, ip, grid)
    point_vars_4 = @test_deprecated project(qp_values, proj3)

    @test point_vars ≈ point_vars_2 ≈ point_vars_3 ≈ point_vars_4

    if order == 1
        # A linear approximation can not recover a quadratic solution,
        # so projected values will be different from the analytical ones
        ae = [f_approx(i) for i in 1:4]
    elseif order == 2
        # For a quadratic approximation the analytical solution is recovered
        ae = compute_vertex_values(grid, f)
    end
    @test point_vars[1:4] ≈ ae

    # Vec
    f_vector(x) = Vec{1,Float64}((f(x),))
    qp_values = analytical(f_vector)
    point_vars = project(proj, qp_values, qr)
    if order == 1
        ae = [Vec{1,Float64}((f_approx(j),)) for j in 1:4]
    elseif order == 2
        ae = compute_vertex_values(grid, f_vector)
    end
    @test point_vars[1:4] ≈ ae

    # Tensor
    f_tensor(x) = Tensor{2,2,Float64}((f(x),2*f(x),3*f(x),4*f(x)))
    qp_values = analytical(f_tensor)
    qp_values_matrix = reduce(hcat, qp_values)::Matrix
    point_vars = project(proj, qp_values, qr)
    point_vars_2 = project(proj, qp_values_matrix, qr)

    @test point_vars ≈ point_vars_2

    if order == 1
        ae = [Tensor{2,2,Float64}((f_approx(i),2*f_approx(i),3*f_approx(i),4*f_approx(i))) for i in 1:4]
    elseif order == 2
        ae = compute_vertex_values(grid, f_tensor)
    end
    @test point_vars[1:4] ≈ ae

    # SymmetricTensor
    f_stensor(x) = SymmetricTensor{2,2,Float64}((f(x),2*f(x),3*f(x)))
    qp_values = analytical(f_stensor)
    qp_values_matrix = reduce(hcat, qp_values)
    point_vars = project(proj, qp_values, qr)
    point_vars_2 = project(proj, qp_values_matrix, qr)

    @test point_vars ≈ point_vars_2

    if order == 1
        ae = [SymmetricTensor{2,2,Float64}((f_approx(i),2*f_approx(i),3*f_approx(i))) for i in 1:4]
    elseif order == 2
        ae = compute_vertex_values(grid, f_stensor)
    end
    @test point_vars[1:4] ≈ ae

    # Test error-path with bad qr
    if refshape == RefTetrahedron && order == 2
        bad_order = 2
    else
        bad_order = 1
    end
    @test_throws LinearAlgebra.PosDefException L2Projector(ip, grid; qr_lhs=QuadratureRule{dim,refshape}(bad_order), geom_ip=ip_geom)
end

# Test a mixed grid, where only a subset of the cells contains a field
function test_projection_mixedgrid()
    # generate a mesh with 1 quadrilateral and 2 triangular elements
    dim = 2
    nodes = Node{dim, Float64}[]
    push!(nodes, Node((0.0, 0.0)))
    push!(nodes, Node((1.0, 0.0)))
    push!(nodes, Node((2.0, 0.0)))
    push!(nodes, Node((0.0, 1.0)))
    push!(nodes, Node((1.0, 1.0)))
    push!(nodes, Node((2.0, 1.0)))

    cells = Ferrite.AbstractCell[]
    push!(cells, Quadrilateral((1,2,5,4)))
    push!(cells, Triangle((2,3,6)))
    push!(cells, Triangle((2,6,5)))

    mesh = Grid(cells, nodes)

    order = 2
    ip = Lagrange{dim, RefCube, order}()
    ip_geom = Lagrange{dim, RefCube, 1}()
    qr = QuadratureRule{dim, RefCube}(order+1)
    cv = CellScalarValues(qr, ip, ip_geom)

    # Create node values for the 1st cell
    # use a SymmetricTensor here for testing the symmetric version of project
    f(x) = SymmetricTensor{2,2,Float64}((1 + x[1]^2, 2x[2]^2, x[1]*x[2]))
    xe = getcoordinates(mesh, 1)
    ae = compute_vertex_values(mesh, f)
    # analytical values
    qp_values = [[f(spatial_coordinate(cv, qp, xe)) for qp in 1:getnquadpoints(cv)]]
    qp_values_matrix = reduce(hcat, qp_values)

    # Now recover the nodal values using a L2 projection.
    # Assume f would only exist on the first cell, we project it to the nodes of the
    # 1st cell while ignoring the rest of the domain. NaNs should be stored in all
    # nodes that do not belong to the 1st cell
    proj = L2Projector(ip, mesh; geom_ip=ip_geom, set=1:1)
    point_vars = project(proj, qp_values, qr)
    point_vars_2 = project(proj, qp_values_matrix, qr)
    ## Old API with fe values as first arg
    proj = @test_deprecated L2Projector(cv, ip, mesh, 1:1)
    point_vars_3 = @test_deprecated project(qp_values, proj)
    ## Old API with qr as first arg
    proj = @test_deprecated L2Projector(qr, ip, mesh, 1:1)
    point_vars_4 = @test_deprecated project(qp_values, proj)

    # In the nodes of the 1st cell we should recover the field
    for node in mesh.cells[1].nodes
        @test ae[node] ≈ point_vars[node] ≈ point_vars_2[node] ≈ point_vars_3[node] ≈
                         point_vars_4[node]
    end

    # in all other nodes we should have NaNs
    for node in setdiff(1:getnnodes(mesh), mesh.cells[1].nodes)
        for d1 = 1:dim, d2 = 1:dim
             @test isnan(point_vars[node][d1, d2])
             @test isnan(point_vars_2[node][d1, d2])
             @test isnan(point_vars_3[node][d1, d2])
             @test isnan(point_vars_4[node][d1, d2])
         end
    end
end

function test_node_reordering()
    grid = generate_grid(Quadrilateral, (1, 1), Vec((0.,0.)), Vec((2.,2.)))
    dim = 2
    ip = Lagrange{dim, RefCube, 2}()
    ip_geo = Lagrange{dim, RefCube,1}()
    qr = QuadratureRule{dim, RefCube}(3)
    cv = CellScalarValues(qr, ip, ip_geo)

    f(x) = x[1]+x[2]

    qp_values = [[f(spatial_coordinate(cv, qp, getcoordinates(cell))) for qp in 1:getnquadpoints(cv)] for cell in CellIterator(grid)]

    projector = L2Projector(ip, grid)
    projected_vals_nodes = project(projector, qp_values, qr)
    projected_vals_dofs = project(projector, qp_values, qr; project_to_nodes=false)

    tol = 1e-12
    @test all(projected_vals_nodes - [0.0, 2.0, 2.0, 4.0] .< tol)
    @test all(projected_vals_dofs - [0., 2., 4., 2., 1., 3., 3., 1., 2.] .< tol)
end

@testset "Test L2-Projection" begin
    test_projection(1, RefCube)
    test_projection(1, RefTetrahedron)
    test_projection(2, RefCube)
    test_projection(2, RefTetrahedron)
    test_projection_mixedgrid()
    test_node_reordering()
end