text:
# implementation of the GPUCompiler interfaces for generating Metal code

const Metal_LLVM_Tools_jll = LazyModule("Metal_LLVM_Tools_jll", UUID("0418c028-ff8c-56b8-a53e-0f9676ed36fc"))

## target

export MetalCompilerTarget

Base.@kwdef struct MetalCompilerTarget <: AbstractCompilerTarget
    macos::VersionNumber
end

function Base.hash(target::MetalCompilerTarget, h::UInt)
    hash(target.macos, h)
end

source_code(target::MetalCompilerTarget) = "text"

# Metal is not supported by our LLVM builds, so we can't get a target machine
llvm_machine(::MetalCompilerTarget) = nothing

llvm_triple(target::MetalCompilerTarget) = "air64-apple-macosx$(target.macos)"

llvm_datalayout(target::MetalCompilerTarget) =
    "e-p:64:64:64"*
    "-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64"*
    "-f32:32:32-f64:64:64"*
    "-v16:16:16-v24:32:32-v32:32:32-v48:64:64-v64:64:64-v96:128:128-v128:128:128-v192:256:256-v256:256:256-v512:512:512-v1024:1024:1024"*
    "-n8:16:32"

needs_byval(job::CompilerJob{MetalCompilerTarget}) = false

## job

# TODO: encode debug build or not in the compiler job
#       https://github.com/JuliaGPU/CUDAnative.jl/issues/368
runtime_slug(job::CompilerJob{MetalCompilerTarget}) = "metal-macos$(job.target.macos)"

isintrinsic(@nospecialize(job::CompilerJob{MetalCompilerTarget}), fn::String) =
    return startswith(fn, "air.")

const LLVMMETALFUNCCallConv   = LLVM.API.LLVMCallConv(102)
const LLVMMETALKERNELCallConv = LLVM.API.LLVMCallConv(103)

const metal_struct_names = [:MtlDeviceArray, :MtlDeviceMatrix, :MtlDeviceVector]

# Initial mapping of types - There has to be a better way
const jl_type_to_c = Dict(
                    Float32 => "float",
                    Float16 => "half",
                    Int64   => "long",
                    UInt64  => "ulong",
                    Int32   => "int",
                    UInt32  => "uint",
                    Int16   => "short",
                    UInt16  => "ushort",
                    Int8    => "char",
                    UInt8   => "uchar",
                    Bool    => "char" # xxx: ?
                )

function process_module!(job::CompilerJob{MetalCompilerTarget}, mod::LLVM.Module)
    # calling convention
    for f in functions(mod)
        #callconv!(f, LLVMMETALFUNCCallConv)
        # XXX: this makes InstCombine erase kernel->func calls.
        #      do we even need this? if we do, do so in metallib-instead.
    end
end

function process_entry!(job::CompilerJob{MetalCompilerTarget}, mod::LLVM.Module, entry::LLVM.Function)
    entry = invoke(process_entry!, Tuple{CompilerJob, LLVM.Module, LLVM.Function}, job, mod, entry)

    if job.source.kernel
        # calling convention
        callconv!(entry, LLVMMETALKERNELCallConv)
    end

    return entry
end

# TODO: why is this done in finish_module? maybe just in process_entry?
function finish_module!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod::LLVM.Module, entry::LLVM.Function)
    entry = invoke(finish_module!, Tuple{CompilerJob, LLVM.Module, LLVM.Function}, job, mod, entry)

    ctx = context(mod)
    entry_fn = LLVM.name(entry)

    if job.source.kernel
        entry = pass_by_reference!(job, mod, entry)

        arguments = add_input_arguments!(job, mod, entry)
        entry = LLVM.functions(mod)[entry_fn]

        add_argument_metadata!(job, mod, entry, arguments)

        add_module_metadata!(job, mod)
    end

    return functions(mod)[entry_fn]
end

function finish_ir!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod::LLVM.Module,
                                  entry::LLVM.Function)
    entry = invoke(finish_ir!, Tuple{CompilerJob, LLVM.Module, LLVM.Function}, job, mod, entry)

    ctx = context(mod)
    entry_fn = LLVM.name(entry)

    if job.source.kernel
        add_address_spaces!(job, mod, entry)
    end

    return functions(mod)[entry_fn]
end

@unlocked function mcgen(job::CompilerJob{MetalCompilerTarget}, mod::LLVM.Module,
                         format=LLVM.API.LLVMObjectFile)
    # translate to metallib
    input = tempname(cleanup=false) * ".bc"
    translated = tempname(cleanup=false) * ".metallib"
    write(input, mod)
    Metal_LLVM_Tools_jll.metallib_as() do assembler
        proc = run(ignorestatus(`$assembler -o $translated $input`))
        if !success(proc)
            error("""Failed to translate LLVM code to MetalLib.
                     If you think this is a bug, please file an issue and attach $(input).""")
        end
    end

    output = if format == LLVM.API.LLVMObjectFile
        read(translated)
    else
        # disassemble
        Metal_LLVM_Tools_jll.metallib_dis() do disassembler
            read(`$disassembler -o - $translated`, String)
        end
    end

    rm(input)
    rm(translated)

    return output
end

# generic pointer removal
#
# every pointer argument (e.g. byref objs) to a kernel needs an address space attached.
function add_address_spaces!(@nospecialize(job::CompilerJob), mod::LLVM.Module, f::LLVM.Function)
    ctx = context(mod)
    ft = eltype(llvmtype(f))
    @compiler_assert return_type(ft) == LLVM.VoidType(ctx) job

    function remapType(src)
        # TODO: recurse in structs
        dst = if src isa LLVM.PointerType && addrspace(src) == 0
            LLVM.PointerType(remapType(eltype(src)), #=device=# 1)
        else
            src
        end
        # TODO: cache
        return dst
    end

    # generate the new function type & definition
    new_types = LLVMType[remapType(typ) for typ in parameters(ft)]
    new_ft = LLVM.FunctionType(return_type(ft), new_types)
    new_f = LLVM.Function(mod, "", new_ft)
    linkage!(new_f, linkage(f))
    for (arg, new_arg) in zip(parameters(f), parameters(new_f))
        LLVM.name!(new_arg, LLVM.name(arg))
    end

    # map the parameters
    value_map = Dict{LLVM.Value, LLVM.Value}(
        param => new_param for (param, new_param) in zip(parameters(f), parameters(new_f))
    )
    value_map[f] = new_f

    # before D96531 (part of LLVM 13), clone_into! wants to duplicate debug metadata
    # when the functions are part of the same module. that is invalid, because it
    # results in desynchronized debug intrinsics (GPUCompiler#284), so remove those.
    if LLVM.version() < v"13"
        removals = LLVM.Instruction[]
        for bb in blocks(f), inst in instructions(bb)
            if inst isa LLVM.CallInst && LLVM.name(called_value(inst)) == "llvm.dbg.declare"
                push!(removals, inst)
            end
        end
        for inst in removals
            @assert isempty(uses(inst))
            unsafe_delete!(LLVM.parent(inst), inst)
        end
        changes = LLVM.API.LLVMCloneFunctionChangeTypeGlobalChanges
    else
        changes = LLVM.API.LLVMCloneFunctionChangeTypeLocalChangesOnly
    end

    function type_mapper(typ)
        remapType(typ)
    end

    clone_into!(new_f, f; value_map, changes, type_mapper)
    replace_uses!(f, new_f) # to update metadata
    @assert isempty(uses(f))

    # remove the old function
    fn = LLVM.name(f)
    unsafe_delete!(mod, f)
    LLVM.name!(new_f, fn)

    return new_f
end

# value-to-reference conversion
#
# Metal doesn't support passing valuse, so we need to convert those to references instead
function pass_by_reference!(@nospecialize(job::CompilerJob), mod::LLVM.Module, f::LLVM.Function)
    ctx = context(mod)
    ft = eltype(llvmtype(f))
    @compiler_assert return_type(ft) == LLVM.VoidType(ctx) job

    # generate the new function type & definition
    args = classify_arguments(job, ft)
    new_types = LLVM.LLVMType[]
    for arg in args
        if arg.cc == BITS_VALUE && !(arg.typ <: Ptr || arg.typ <: Core.LLVMPtr)
            # pass the value as a reference instead
            push!(new_types, LLVM.PointerType(arg.codegen.typ, #=Constant=# 1))
            # TODO: other attributes (nocapturem nonnull, readonly, align, dereferenceable)?
        elseif arg.cc != GHOST
            push!(new_types, arg.codegen.typ)
        end
    end
    new_ft = LLVM.FunctionType(return_type(ft), new_types)
    new_f = LLVM.Function(mod, "", new_ft)
    linkage!(new_f, linkage(f))
    for (arg, new_arg) in zip(parameters(f), parameters(new_f))
        LLVM.name!(new_arg, LLVM.name(arg))
    end

    # emit IR performing the "conversions"
    new_args = LLVM.Value[]
    Builder(ctx) do builder
        entry = BasicBlock(new_f, "entry"; ctx)
        position!(builder, entry)

        # perform argument conversions
        for arg in args
            if arg.cc == BITS_VALUE && !(arg.typ <: Ptr || arg.typ <: Core.LLVMPtr)
                # load the reference to get a value back
                val = load!(builder, parameters(new_f)[arg.codegen.i])
                push!(new_args, val)
            elseif arg.cc != GHOST
                push!(new_args, parameters(new_f)[arg.codegen.i])
                for attr in collect(parameter_attributes(f, arg.codegen.i))
                    push!(parameter_attributes(new_f, arg.codegen.i), attr)
                end
            end
        end

        # map the arguments
        value_map = Dict{LLVM.Value, LLVM.Value}(
            param => new_args[i] for (i,param) in enumerate(parameters(f))
        )

        value_map[f] = new_f
        clone_into!(new_f, f; value_map,
                    changes=LLVM.API.LLVMCloneFunctionChangeTypeLocalChangesOnly)

        # fall through
        br!(builder, blocks(new_f)[2])
    end

    # remove the old function
    # NOTE: if we ever have legitimate uses of the old function, create a shim instead
    fn = LLVM.name(f)
    @assert isempty(uses(f))
    unsafe_delete!(mod, f)
    LLVM.name!(new_f, fn)

    return new_f
end

# kernel input arguments
#
# hardware index counters (thread id, group id, etc) aren't accessed via intrinsics,
# but using special arguments to the kernel function.

const kernel_intrinsics = Dict()
for intr in [
        "dispatch_quadgroups_per_threadgroup", "dispatch_simdgroups_per_threadgroup",
        "quadgroup_index_in_threadgroup", "quadgroups_per_threadgroup",
        "simdgroup_index_in_threadgroup", "simdgroups_per_threadgroup",
        "thread_index_in_quadgroup", "thread_index_in_simdgroup", "thread_index_in_threadgroup",
        "thread_execution_width", "threads_per_simdgroup"],
    (intr_typ, air_typ, julia_typ) in [
        ("i32",   "uint",  UInt32),
        ("i16",   "ushort",  UInt16),
    ]
    push!(kernel_intrinsics,
          "julia.air.$intr.$intr_typ" =>
          (air_intr="$intr.$air_typ", air_typ, air_name=intr, julia_typ))
end
for intr in [
        "dispatch_threads_per_threadgroup",
        "grid_origin", "grid_size",
        "thread_position_in_grid", "thread_position_in_threadgroup",
        "threadgroup_position_in_grid", "threadgroups_per_grid",
        "threads_per_grid", "threads_per_threadgroup"],
    (intr_typ, air_typ, julia_typ) in [
        ("i32",   "uint",  UInt32),
        ("v2i32", "uint2", NTuple{2, VecElement{UInt32}}),
        ("v3i32", "uint3", NTuple{3, VecElement{UInt32}}),
        ("i16",   "ushort",  UInt16),
        ("v2i16", "ushort2", NTuple{2, VecElement{UInt16}}),
        ("v3i16", "ushort3", NTuple{3, VecElement{UInt16}}),
    ]
    push!(kernel_intrinsics,
          "julia.air.$intr.$intr_typ" =>
          (air_intr="$intr.$air_typ", air_typ, air_name=intr, julia_typ))
end

function add_input_arguments!(@nospecialize(job::CompilerJob), mod::LLVM.Module,
                              entry::LLVM.Function)
    ctx = context(mod)
    entry_fn = LLVM.name(entry)

    # figure out which intrinsics are used and need to be added as arguments
    used_intrinsics = filter(keys(kernel_intrinsics)) do intr_fn
        haskey(functions(mod), intr_fn)
    end |> collect
    # TODO: Figure out how to not be inefficient with these changes
    nargs = length(used_intrinsics)

    # determine which functions need these arguments
    worklist = Set{LLVM.Function}([entry])
    for intr_fn in used_intrinsics
        push!(worklist, functions(mod)[intr_fn])
    end
    worklist_length = 0
    while worklist_length != length(worklist)
        # iteratively discover functions that use an intrinsic or any function calling it
        worklist_length = length(worklist)
        additions = LLVM.Function[]
        for f in worklist, use in uses(f)
            inst = user(use)::Instruction
            bb = LLVM.parent(inst)
            new_f = LLVM.parent(bb)
            in(new_f, worklist) || push!(additions, new_f)
        end
        for f in additions
            push!(worklist, f)
        end
    end
    for intr_fn in used_intrinsics
        delete!(worklist, functions(mod)[intr_fn])
    end

    # add the arguments
    # NOTE: we could be more fine-grained, only adding the specific intrinsics used by this function.
    #       not sure if that's worth it though.
    workmap = Dict{LLVM.Function, LLVM.Function}()
    for f in worklist
        fn = LLVM.name(f)
        ft = eltype(llvmtype(f))
        LLVM.name!(f, fn * ".orig")
        # create a new function
        new_param_types = LLVMType[parameters(ft)...]

        for intr_fn in used_intrinsics
            llvm_typ = convert(LLVMType, kernel_intrinsics[intr_fn].julia_typ; ctx)
            push!(new_param_types, llvm_typ)
        end
        new_ft = LLVM.FunctionType(return_type(ft), new_param_types)
        new_f = LLVM.Function(mod, fn, new_ft)
        linkage!(new_f, linkage(f))
        for (arg, new_arg) in zip(parameters(f), parameters(new_f))
            LLVM.name!(new_arg, LLVM.name(arg))
        end
        for (intr_fn, new_arg) in zip(used_intrinsics, parameters(new_f)[end-nargs+1:end])
            LLVM.name!(new_arg, kernel_intrinsics[intr_fn].air_name)
        end

        workmap[f] = new_f
    end

    # clone and rewrite the function bodies.
    # we don't need to rewrite much as the arguments are added last.
    for (f, new_f) in workmap
        # map the arguments
        value_map = Dict{LLVM.Value, LLVM.Value}()
        for (param, new_param) in zip(parameters(f), parameters(new_f))
            LLVM.name!(new_param, LLVM.name(param))
            value_map[param] = new_param
        end

        value_map[f] = new_f
        clone_into!(new_f, f; value_map,
                    changes=LLVM.API.LLVMCloneFunctionChangeTypeLocalChangesOnly)

        # we can't remove this function yet, as we might still need to rewrite any called,
        # but remove the IR already
        empty!(f)
    end

    # drop unused constants that may be referring to the old functions
    # XXX: can we do this differently?
    for f in worklist
        for use in uses(f)
            val = user(use)
            if val isa LLVM.ConstantExpr && isempty(uses(val))
                LLVM.unsafe_destroy!(val)
            end
        end
    end

    # update other uses of the old function, modifying call sites to pass the arguments
    function rewrite_uses!(f, new_f)
        # update uses
        Builder(ctx) do builder
            for use in uses(f)
                val = user(use)
                if val isa LLVM.CallInst || val isa LLVM.InvokeInst || val isa LLVM.CallBrInst
                    callee_f = LLVM.parent(LLVM.parent(val))
                    # forward the arguments
                    position!(builder, val)
                    new_val = if val isa LLVM.CallInst
                        call!(builder, new_f, [arguments(val)..., parameters(callee_f)[end-nargs+1:end]...], operand_bundles(val))
                    else
                        # TODO: invoke and callbr
                        error("Rewrite of $(typeof(val))-based calls is not implemented: $val")
                    end
                    callconv!(new_val, callconv(val))

                    replace_uses!(val, new_val)
                    @assert isempty(uses(val))
                    unsafe_delete!(LLVM.parent(val), val)
                elseif val isa LLVM.ConstantExpr && opcode(val) == LLVM.API.LLVMBitCast
                    # XXX: why isn't this caught by the value materializer above?
                    target = operands(val)[1]
                    @assert target == f
                    new_val = LLVM.const_bitcast(new_f, llvmtype(val))
                    rewrite_uses!(val, new_val)
                    # we can't simply replace this constant expression, as it may be used
                    # as a call, taking arguments (so we need to rewrite it to pass the input arguments)

                    # drop the old constant if it is unused
                    # XXX: can we do this differently?
                    if isempty(uses(val))
                        LLVM.unsafe_destroy!(val)
                    end
                else
                    error("Cannot rewrite unknown use of function: $val")
                end
            end
        end
    end
    for (f, new_f) in workmap
        rewrite_uses!(f, new_f)
        @assert isempty(uses(f))
        unsafe_delete!(mod, f)
    end

    # replace uses of the intrinsics with references to the input arguments
    for (i, intr_fn) in enumerate(used_intrinsics)
        intr = functions(mod)[intr_fn]
        for use in uses(intr)
            val = user(use)
            callee_f = LLVM.parent(LLVM.parent(val))
            if val isa LLVM.CallInst || val isa LLVM.InvokeInst || val isa LLVM.CallBrInst
                replace_uses!(val, parameters(callee_f)[end-nargs+i])
            else
                error("Cannot rewrite unknown use of function: $val")
            end

            @assert isempty(uses(val))
            unsafe_delete!(LLVM.parent(val), val)
        end
        @assert isempty(uses(intr))
        unsafe_delete!(mod, intr)
    end

    return used_intrinsics
end

# argument metadata generation
#
# module metadata is used to identify buffers that are passed as kernel arguments.

# Recursively generate metadata for normal kernel arguments
function add_md(arg; ctx, field_info=nothing, level=1)
    structinfo(T,i) = (fieldoffset(T,i), fieldname(T,i), fieldtype(T,i))

    if arg.codegen.typ isa LLVM.PointerType
        # Process pointer to structs as argument buffers
        if eltype(arg.codegen.typ) isa LLVM.StructType
            argbuf_info = Metadata[]

            # Get information about struct elements first
            new_codegen_typ = eltype(arg.codegen.typ)
            struct_info = add_md((typ=arg.typ, codegen=(typ=new_codegen_typ, i=1)); ctx, level=level+1)

            # Create argument buffer's type metadata
            struct_type_info = Metadata[]

            # Struct element metadata format:
                # If element is itself a struct directly (not a reference to a struct)
                    # air.struct_type_info keyword
                    # Metadata node of the struct (element struct - NOT a self-reference)
                # Offset in bytes from start of struct
                # Size of element in bytes (8 for buffers (pointer size))
                # Length of element (0 for buffers...always?)
                # Field type
                # Field name
                # Field argument type? (mainly air.indirect_argument)
                    # With structs with the threadgroup addresspace, this metadata node is not present
                    # Because each threadgroup gets the struct data directly without redirection??
                # If the element is itself a struct directly (not a reference to a struct)
                    # Location index of the element struct in the higher-level struct
                # Metadata node to more details about element
                # If argument is unused
                    # air.arg_unused

            # Return the struct element type name and field name from metadata
            function parse_struct_names(md)
                for (i, item) in enumerate(md)
                    if item isa MDString && string(item) == "air.arg_type_name"
                        return (md[i+1], md[i+3])
                    end
                end
                error("Struct element metadata keyword 'air.arg_type_name' not found in $(md)")
            end

            for (i,struct_field_info) in enumerate(struct_info)

                type_name, field_name = parse_struct_names(struct_field_info)
                field_is_struct = arg.typ isa LLVM.StructType

                if field_is_struct
                    push!(struct_type_info, MDString("air.struct_type_info"; ctx))
                    push!(struct_type_info, MDNode(struct_field_info; ctx))
                end

                push!(struct_type_info, Metadata(ConstantInt(Int32(fieldoffset(arg.typ,i)); ctx)))

                # If field is (arg)buffer, set size to 8 and length to 0
                if string(struct_field_info[2]) in ["air.buffer", "air.indirect_buffer"]
                    push!(struct_type_info, Metadata(ConstantInt(Int32(8); ctx))) # Field element size
                    push!(struct_type_info, Metadata(ConstantInt(Int32(0); ctx))) # Length of field
                else
                    field_type = fieldtype(arg.typ, i)
                    push!(struct_type_info, Metadata(ConstantInt(Int32(sizeof(eltype(field_type))); ctx))) # Field element size
                    push!(struct_type_info, Metadata(ConstantInt(Int32(length(field_type.parameters)); ctx))) # Length of field
                end
                push!(struct_type_info, type_name) # Field type
                push!(struct_type_info, field_name) # Field name
                push!(struct_type_info, MDString("air.indirect_argument"; ctx))

                if field_is_struct
                    push!(struct_type_info, Metadata(ConstantInt(Int32(i); ctx)))
                else
                    struct_field_info = MDNode(struct_field_info; ctx)
                    push!(struct_type_info, struct_field_info)
                end
            end

            struct_type_info = MDNode(struct_type_info; ctx)

            # Add argument buffer details
            # Create the argument buffer main metadata
            push!(argbuf_info, Metadata(ConstantInt(Int32(arg.codegen.i-1); ctx))) # Argument index
            push!(argbuf_info, MDString("air.indirect_buffer"; ctx))
            push!(argbuf_info, MDString("air.buffer_size"; ctx))
            push!(argbuf_info, Metadata(ConstantInt(Int32(sizeof(arg.typ)); ctx)))
            push!(argbuf_info, MDString("air.location_index"; ctx))
            push!(argbuf_info, Metadata(ConstantInt(Int32(arg.codegen.i-1); ctx)))
            push!(argbuf_info, Metadata(ConstantInt(Int32(1); ctx)))
            # TODO: Check for const array and put to air.read
            push!(argbuf_info, MDString("air.read_write"; ctx))
            push!(argbuf_info, MDString("air.struct_type_info"; ctx))
            push!(argbuf_info, struct_type_info) # Argument buffer type info
            push!(argbuf_info, MDString("air.arg_type_size"; ctx))
            push!(argbuf_info, Metadata(ConstantInt(Int32(sizeof(arg.typ)); ctx))) # Arg type size
            push!(argbuf_info, MDString("air.arg_type_align_size"; ctx))
            push!(argbuf_info, Metadata(ConstantInt(Int32(Base.datatype_alignment(arg.typ)); ctx)))
            push!(argbuf_info, MDString("air.arg_type_name"; ctx))
            push!(argbuf_info, MDString(string(arg.typ); ctx))
            push!(argbuf_info, MDString("air.arg_name"; ctx))
            push!(argbuf_info, MDString("arg_$(arg.codegen.i-1)"; ctx)) # TODO: How to get this? Does the compiler job have it somewhere?
            # Ignore unused flag for now

            # Make argument buffer metadata node
            argbuf_info = MDNode(argbuf_info; ctx)
            return argbuf_info

        # Process other references (e.g. to scalars) or plain pointers as simple buffers
        else
            ptr_datatype = if arg.typ <: Core.LLVMPtr || arg.typ <: Ptr
                arg.typ.parameters[1]
            else
                arg.typ
            end

            arg_info_ptr = Metadata[]
            push!(arg_info_ptr, Metadata(ConstantInt(Int32(arg.codegen.i-1); ctx))) # Argument index
            push!(arg_info_ptr, MDString("air.buffer"; ctx))
            push!(arg_info_ptr, MDString("air.location_index"; ctx))
            push!(arg_info_ptr, Metadata(ConstantInt(Int32(arg.codegen.i-1); ctx)))
            push!(arg_info_ptr, Metadata(ConstantInt(Int32(1); ctx)))
            push!(arg_info_ptr, MDString("air.read_write"; ctx)) # TODO: Check for const array
            push!(arg_info_ptr, MDString("air.arg_type_size"; ctx))
            push!(arg_info_ptr, Metadata(ConstantInt(Int32(sizeof(ptr_datatype)); ctx))) # TODO: Get properly
            push!(arg_info_ptr, MDString("air.arg_type_align_size"; ctx))
            push!(arg_info_ptr, Metadata(ConstantInt(Int32(Base.datatype_alignment(ptr_datatype)); ctx)))
            push!(arg_info_ptr, MDString("air.arg_type_name"; ctx))
            # Handle naming for pointer to ArrayType
            if eltype(arg.codegen.typ) isa LLVM.ArrayType
                arg_type_name = jl_type_to_c[eltype(eltype(arg.typ))] * string(length(eltype(arg.codegen.typ)))
            else
                arg_type_name = string(eltype(arg.codegen.typ))
            end
            push!(arg_info_ptr, MDString(arg_type_name; ctx)) # TODO: Get properly
            push!(arg_info_ptr, MDString("air.arg_name"; ctx))
            # TODO: Properly get top-level argument names
            arg_name = field_info != nothing ? string(field_info[2]) : "arg_$(arg.codegen.i)"
            push!(arg_info_ptr, MDString(arg_name; ctx))
            return arg_info_ptr
        end

    elseif arg.codegen.typ isa LLVM.StructType
        arg_info_struct = []
        for (i, elem) in enumerate(collect(elements(arg.codegen.typ)))
            field_info = structinfo(arg.typ, i)
            push!(arg_info_struct, add_md((typ=field_info[3], codegen=(typ=elem, i=i)); ctx, field_info, level=level+1))
        end
        return arg_info_struct

    # Process as basic leaf argument
    else
        arg_info = Metadata[]

        # If it's a top-level arg, encode as a buffer
        if level == 1

            ## Simple Buffer Argument Metadata layout
                # Kernel argument index
                # air.buffer keyword
                # air.location_index keyword
                # Kernel argument location index (NOT the same as argument index)
                    # Values are determined as explained by pages 46 and 79 of the Metal docs
                        # https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf
                    # Note that these indices are unique to each resource group type [buffer, threadgroup, sampler, texture]
                # Unknown value - Has always been 1
                    # Vertex/stag_in info? Something with rasters?
                # Resource usage status
                # air.arg_type_size keyword
                # Buffer element size
                # air.arg_type_align_size keyword
                # Buffer element alignment
                # air.arg_type_name keyword
                # Buffer element type name
                # air.arg_name keyword
                # Kernel argument name

            datatype = arg.typ

            arg_info_ptr = Metadata[]
            push!(arg_info_ptr, Metadata(ConstantInt(Int32(arg.codegen.i-1); ctx)))
            push!(arg_info_ptr, MDString("air.buffer"; ctx))
            push!(arg_info_ptr, MDString("air.location_index"; ctx))
            push!(arg_info_ptr, Metadata(ConstantInt(Int32(arg.codegen.i-1); ctx)))
            push!(arg_info_ptr, Metadata(ConstantInt(Int32(1); ctx)))
            push!(arg_info_ptr, MDString("air.read_write"; ctx)) # TODO: Check for const array
            push!(arg_info_ptr, MDString("air.arg_type_size"; ctx))
            push!(arg_info_ptr, Metadata(ConstantInt(Int32(sizeof(datatype)); ctx)))
            push!(arg_info_ptr, MDString("air.arg_type_align_size"; ctx))
            push!(arg_info_ptr, Metadata(ConstantInt(Int32(Base.datatype_alignment(datatype)); ctx)))
            push!(arg_info_ptr, MDString("air.arg_type_name"; ctx))
            push!(arg_info_ptr, MDString(string(eltype(arg.codegen.typ)); ctx))
            push!(arg_info_ptr, MDString("air.arg_name"; ctx))
            # TODO: Properly get top-level argument names
            arg_name = field_info != nothing ? string(field_info[2]) : "arg_$(arg.codegen.i)"
            push!(arg_info_ptr, MDString(arg_name; ctx))

        # Else process as indirect_constant (vector type)
        # TODO: Need to check upstream that we're only passing valid vector types
        elseif arg.codegen.typ isa LLVM.ArrayType
            arg_length = length(arg.codegen.typ)

            push!(arg_info, Metadata(ConstantInt(Int32(arg.codegen.i-1); ctx)))
            push!(arg_info, MDString("air.indirect_constant"; ctx))
            push!(arg_info, MDString("air.location_index"; ctx))
            push!(arg_info, Metadata(ConstantInt(Int32(arg.codegen.i-1); ctx)))
            push!(arg_info, Metadata(ConstantInt(Int32(1); ctx)))
            push!(arg_info, MDString("air.arg_type_name"; ctx))
            arg_type_name = jl_type_to_c[eltype(arg.typ)] * string(arg_length)
            push!(arg_info, MDString(arg_type_name; ctx))
            push!(arg_info, MDString("air.arg_name"; ctx))
            arg_name = field_info != nothing ? string(field_info[2]) : "arg_$(arg.codegen.i)"
            push!(arg_info, MDString(arg_name; ctx))
        else
            error("Unknown Metal kernel argument type $(arg.typ)")
        end

        return arg_info
    end
end

function add_argument_metadata!(@nospecialize(job::CompilerJob), mod::LLVM.Module,
                                entry::LLVM.Function, used_intrinsics::Vector)
    ctx = context(mod)

    ## argument info
    arg_infos = Metadata[]

    # Iterate through arguments and create metadata for them
    for arg in classify_arguments(job, eltype(llvmtype(entry)))
        # Ignore ghost type
        if arg.cc != GHOST
            arg_info = add_md(arg; ctx)
            # Ensure returned type is a metadata node
            if !isa(arg_info, MDTuple)
                arg_info = MDNode(arg_info; ctx)
            end
            push!(arg_infos, arg_info)
        end
    end

    # Create metadata for argument intrinsics last
    for (i, intr_fn) in enumerate(used_intrinsics)
        arg_info = Metadata[]
        push!(arg_info, Metadata(ConstantInt(Int32(length(parameters(entry))-i); ctx)))
        push!(arg_info, MDString("air." * kernel_intrinsics[intr_fn].air_name; ctx))
        push!(arg_info, MDString("air.arg_type_name"; ctx))
        push!(arg_info, MDString(kernel_intrinsics[intr_fn].air_typ; ctx))
        push!(arg_info, MDString("air.arg_name"; ctx))
        push!(arg_info, MDString(kernel_intrinsics[intr_fn].air_name; ctx))
        arg_info = MDNode(arg_info; ctx)
        push!(arg_infos, arg_info)
    end
    arg_infos = MDNode(arg_infos; ctx)
    ## stage info
    stage_infos = Metadata[]
    stage_infos = MDNode(stage_infos; ctx)

    kernel_md = MDNode([entry, stage_infos, arg_infos]; ctx)
    push!(metadata(mod)["air.kernel"], kernel_md)

    return
end

# module-level metadata

# TODO: determine limits being set dynamically
function add_module_metadata!(@nospecialize(job::CompilerJob), mod::LLVM.Module)
    ctx = context(mod)

    # register max device buffer count
    max_buff = Metadata[]
    push!(max_buff, Metadata(ConstantInt(Int32(7); ctx)))
    push!(max_buff, MDString("air.max_device_buffers"; ctx))
    push!(max_buff, Metadata(ConstantInt(Int32(31); ctx)))
    max_buff = MDNode(max_buff; ctx)
    push!(metadata(mod)["llvm.module.flags"], max_buff)

    # register max constant buffer count
    max_const_buff_md = Metadata[]
    push!(max_const_buff_md, Metadata(ConstantInt(Int32(7); ctx)))
    push!(max_const_buff_md, MDString("air.max_constant_buffers"; ctx))
    push!(max_const_buff_md, Metadata(ConstantInt(Int32(31); ctx)))
    max_const_buff_md = MDNode(max_const_buff_md; ctx)
    push!(metadata(mod)["llvm.module.flags"], max_const_buff_md)

    # register max threadgroup buffer count
    max_threadgroup_buff_md = Metadata[]
    push!(max_threadgroup_buff_md, Metadata(ConstantInt(Int32(7); ctx)))
    push!(max_threadgroup_buff_md, MDString("air.max_threadgroup_buffers"; ctx))
    push!(max_threadgroup_buff_md, Metadata(ConstantInt(Int32(31); ctx)))
    max_threadgroup_buff_md = MDNode(max_threadgroup_buff_md; ctx)
    push!(metadata(mod)["llvm.module.flags"], max_threadgroup_buff_md)

    # register max texture buffer count
    max_textures_md = Metadata[]
    push!(max_textures_md, Metadata(ConstantInt(Int32(7); ctx)))
    push!(max_textures_md, MDString("air.max_textures"; ctx))
    push!(max_textures_md, Metadata(ConstantInt(Int32(128); ctx)))
    max_textures_md = MDNode(max_textures_md; ctx)
    push!(metadata(mod)["llvm.module.flags"], max_textures_md)

    # register max write texture buffer count
    max_rw_textures_md = Metadata[]
    push!(max_rw_textures_md, Metadata(ConstantInt(Int32(7); ctx)))
    push!(max_rw_textures_md, MDString("air.max_read_write_textures"; ctx))
    push!(max_rw_textures_md, Metadata(ConstantInt(Int32(8); ctx)))
    max_rw_textures_md = MDNode(max_rw_textures_md; ctx)
    push!(metadata(mod)["llvm.module.flags"], max_rw_textures_md)

    # register max sampler count
    max_samplers_md = Metadata[]
    push!(max_samplers_md, Metadata(ConstantInt(Int32(7); ctx)))
    push!(max_samplers_md, MDString("air.max_samplers"; ctx))
    push!(max_samplers_md, Metadata(ConstantInt(Int32(16); ctx)))
    max_samplers_md = MDNode(max_samplers_md; ctx)
    push!(metadata(mod)["llvm.module.flags"], max_samplers_md)

    # add compiler identification
    llvm_ident_md = Metadata[]
    push!(llvm_ident_md, MDString("Julia $(VERSION) with Metal.jl"; ctx))
    llvm_ident_md = MDNode(llvm_ident_md; ctx)
    push!(metadata(mod)["llvm.ident"], llvm_ident_md)

    # add AIR version
    air_md = Metadata[]
    push!(air_md, Metadata(ConstantInt(Int32(2); ctx)))
    push!(air_md, Metadata(ConstantInt(Int32(4); ctx)))
    push!(air_md, Metadata(ConstantInt(Int32(0); ctx)))
    air_md = MDNode(air_md; ctx)
    push!(metadata(mod)["air.version"], air_md)

    # add language version
    air_lang_md = Metadata[]
    push!(air_lang_md, MDString("Metal"; ctx))
    push!(air_lang_md, Metadata(ConstantInt(Int32(2); ctx)))
    push!(air_lang_md, Metadata(ConstantInt(Int32(4); ctx)))
    push!(air_lang_md, Metadata(ConstantInt(Int32(0); ctx)))
    air_lang_md = MDNode(air_lang_md; ctx)
    push!(metadata(mod)["air.language_version"], air_lang_md)

    return
end


text:
"""
    BalanceLaw

An abstract type representing a PDE balance law of the form

elements for balance laws of the form

```math
q_{,t} + Σ_{i=1,...d} F_{i,i} = s
```

Subtypes `L` should define the following methods:
- `vars_aux(::L)`: a tuple of symbols containing the auxiliary variables
- `vars_state(::L)`: a tuple of symbols containing the state variables
- `vars_gradient(::L)`: a tuple of symbols containing the transformed variables of which gradients are computed
- `vars_gradient_laplacian(::L)`: a tuple of symbols containing the transformed variables of which gradients of laplacian are computed, they have to be a subset of `vars_gradient`
- `vars_diffusive(::L)`: a tuple of symbols containing the diffusive variables
- `vars_hyperdiffusive(::L)`: a tuple of symbols containing the hyperdiffusive variables
- `flux_nondiffusive!(::L, flux::Grad, state::State, auxstate::State, t::Real)`
- `flux_diffusive!(::L, flux::Grad, state::State, diffstate::State, hyperdiffstate::State, auxstate::State, t::Real)`
- `gradvariables!(::L, transformstate::State, state::State, auxstate::State, t::Real)`: transformation of state variables to variables of which gradients are computed
- `diffusive!(::L, diffstate::State, ∇transformstate::Grad, auxstate::State, t::Real)`: transformation of gradients to the diffusive variables
- `hyperdiffusive!(::L, hyperdiffstate::State, ∇Δtransformstate::Grad, auxstate::State, t::Real)`: transformation of laplacian gradients to the hyperdiffusive variables
- `source!(::L, source::State, state::State, diffusive::Vars, auxstate::State, t::Real, direction::Direction)
- `wavespeed(::L, n⁻, state::State, aux::State, t::Real)`
- `boundary_state!(::NumericalFluxGradient, ::L, state⁺::State, aux⁺::State, normal⁻, state⁻::State, aux⁻::State, bctype, t)`
- `boundary_state!(::NumericalFluxNonDiffusive, ::L, state⁺::State, aux⁺::State, normal⁻, state⁻::State, aux⁻::State, bctype, t)`
- `boundary_state!(::NumericalFluxDiffusive, ::L, state⁺::State, diff⁺::State, aux⁺::State, normal⁻, state⁻::State, diff⁻::State, aux⁻::State, bctype, t)`
- `init_aux!(::L, aux::State, coords, args...)`
- `init_state!(::L, state::State, aux::State, coords, args...)`

"""
abstract type BalanceLaw end # PDE part

# function stubs
function vars_state end
function vars_aux end
function vars_gradient end
vars_gradient_laplacian(::BalanceLaw, FT) = @vars()
function vars_diffusive end
vars_hyperdiffusive(::BalanceLaw, FT) = @vars()
vars_integrals(::BalanceLaw, FT) = @vars()
vars_reverse_integrals(::BalanceLaw, FT) = @vars()

num_aux(m::BalanceLaw, FT) = varsize(vars_aux(m, FT))
num_state(m::BalanceLaw, FT) = varsize(vars_state(m, FT)) # nstate
num_gradient(m::BalanceLaw, FT) = varsize(vars_gradient(m, FT))  # number_gradient_states
num_gradient_laplacian(m::BalanceLaw, FT) =
    varsize(vars_gradient_laplacian(m, FT))
num_diffusive(m::BalanceLaw, FT) = varsize(vars_diffusive(m, FT)) # number_viscous_states
num_hyperdiffusive(m::BalanceLaw, FT) = varsize(vars_hyperdiffusive(m, FT))
num_integrals(m::BalanceLaw, FT) = varsize(vars_integrals(m, FT))
num_reverse_integrals(m::BalanceLaw, FT) =
    varsize(vars_reverse_integrals(m, FT))

function update_aux! end
function integral_load_aux! end
function integral_set_aux! end
function reverse_integral_load_aux! end
function reverse_integral_set_aux! end
function flux_nondiffusive! end
function flux_diffusive! end
function gradvariables! end
function diffusive! end
function hyperdiffusive! end
function source! end
function wavespeed end
function boundary_state! end
function init_aux! end
function init_state! end

using ..Courant
"""
    calculate_dt(dg, model, Q, Courant_number, direction, t)

For a given model, compute a time step satisying the nondiffusive Courant number
`Courant_number`
"""
function calculate_dt(dg, model, Q, Courant_number, t, direction)
    Δt = one(eltype(Q))
    CFL = courant(nondiffusive_courant, dg, model, Q, Δt, t, direction)
    return Courant_number / CFL
end

function create_state(bl::BalanceLaw, grid)
    topology = grid.topology
    # FIXME: Remove after updating CUDA
    h_vgeo = Array(grid.vgeo)
    FT = eltype(h_vgeo)
    Np = dofs_per_element(grid)
    DA = arraytype(grid)

    weights = view(h_vgeo, :, grid.Mid, :)
    weights = reshape(weights, size(weights, 1), 1, size(weights, 2))

    V = vars_state(bl, FT)
    state = MPIStateArray{FT, V}(
        topology.mpicomm,
        DA,
        Np,
        num_state(bl, FT),
        length(topology.elems),
        realelems = topology.realelems,
        ghostelems = topology.ghostelems,
        vmaprecv = grid.vmaprecv,
        vmapsend = grid.vmapsend,
        nabrtorank = topology.nabrtorank,
        nabrtovmaprecv = grid.nabrtovmaprecv,
        nabrtovmapsend = grid.nabrtovmapsend,
        weights = weights,
    )
    return state
end

function create_auxstate(bl, grid)
    topology = grid.topology
    Np = dofs_per_element(grid)

    h_vgeo = Array(grid.vgeo)
    FT = eltype(h_vgeo)
    DA = arraytype(grid)

    weights = view(h_vgeo, :, grid.Mid, :)
    weights = reshape(weights, size(weights, 1), 1, size(weights, 2))

    V = vars_aux(bl, FT)
    auxstate = MPIStateArray{FT, V}(
        topology.mpicomm,
        DA,
        Np,
        num_aux(bl, FT),
        length(topology.elems),
        realelems = topology.realelems,
        ghostelems = topology.ghostelems,
        vmaprecv = grid.vmaprecv,
        vmapsend = grid.vmapsend,
        nabrtorank = topology.nabrtorank,
        nabrtovmaprecv = grid.nabrtovmaprecv,
        nabrtovmapsend = grid.nabrtovmapsend,
        weights = weights,
    )

    dim = dimensionality(grid)
    polyorder = polynomialorder(grid)
    vgeo = grid.vgeo
    device = typeof(auxstate.data) <: Array ? CPU() : CUDA()
    nrealelem = length(topology.realelems)
    event = Event(device)
    event = initauxstate!(device, Np, Np * nrealelem)(
        bl,
        Val(dim),
        Val(polyorder),
        auxstate.data,
        vgeo,
        topology.realelems,
        dependencies = (event,),
    )
    event = MPIStateArrays.begin_ghost_exchange!(auxstate; dependencies = event)
    event = MPIStateArrays.end_ghost_exchange!(auxstate; dependencies = event)
    wait(device, event)

    return auxstate
end

function create_diffstate(bl, grid)
    topology = grid.topology
    Np = dofs_per_element(grid)

    h_vgeo = Array(grid.vgeo)
    FT = eltype(h_vgeo)
    DA = arraytype(grid)

    weights = view(h_vgeo, :, grid.Mid, :)
    weights = reshape(weights, size(weights, 1), 1, size(weights, 2))

    # TODO: Clean up this MPIStateArray interface...
    V = vars_diffusive(bl, FT)
    diffstate = MPIStateArray{FT, V}(
        topology.mpicomm,
        DA,
        Np,
        num_diffusive(bl, FT),
        length(topology.elems),
        realelems = topology.realelems,
        ghostelems = topology.ghostelems,
        vmaprecv = grid.vmaprecv,
        vmapsend = grid.vmapsend,
        nabrtorank = topology.nabrtorank,
        nabrtovmaprecv = grid.nabrtovmaprecv,
        nabrtovmapsend = grid.nabrtovmapsend,
        weights = weights,
    )

    return diffstate
end

function create_hyperdiffstate(bl, grid)
    topology = grid.topology
    Np = dofs_per_element(grid)

    h_vgeo = Array(grid.vgeo)
    FT = eltype(h_vgeo)
    DA = arraytype(grid)

    weights = view(h_vgeo, :, grid.Mid, :)
    weights = reshape(weights, size(weights, 1), 1, size(weights, 2))

    ngradlapstate = num_gradient_laplacian(bl, FT)
    # TODO: Clean up this MPIStateArray interface...
    V = vars_gradient_laplacian(bl, FT)
    Qhypervisc_grad = MPIStateArray{FT, V}(
        topology.mpicomm,
        DA,
        Np,
        3ngradlapstate,
        length(topology.elems),
        realelems = topology.realelems,
        ghostelems = topology.ghostelems,
        vmaprecv = grid.vmaprecv,
        vmapsend = grid.vmapsend,
        nabrtorank = topology.nabrtorank,
        nabrtovmaprecv = grid.nabrtovmaprecv,
        nabrtovmapsend = grid.nabrtovmapsend,
        weights = weights,
    )

    V = vars_hyperdiffusive(bl, FT)
    Qhypervisc_div = MPIStateArray{FT, V}(
        topology.mpicomm,
        DA,
        Np,
        ngradlapstate,
        length(topology.elems),
        realelems = topology.realelems,
        ghostelems = topology.ghostelems,
        vmaprecv = grid.vmaprecv,
        vmapsend = grid.vmapsend,
        nabrtorank = topology.nabrtorank,
        nabrtovmaprecv = grid.nabrtovmaprecv,
        nabrtovmapsend = grid.nabrtovmapsend,
        weights = weights,
    )
    return Qhypervisc_grad, Qhypervisc_div
end

ext: jl

max_stars_repo_name: yairchn/CLIMA

max_issues_repo_name: yairchn/CLIMA

max_forks_repo_name: yairchn/CLIMA


text:
export CloneStatePreprocessor, ComposedPreprocessor

(p::AbstractPreprocessor)(obs) = StateOverriddenObs(obs = obs, state = p(get_state(obs)))

"""
    ComposedPreprocessor(p::AbstractPreprocessor...)

Compose multiple preprocessors.
"""
struct ComposedPreprocessor{T} <: AbstractPreprocessor
    preprocessors::T
end

ComposedPreprocessor(p::AbstractPreprocessor...) = ComposedPreprocessor(p)
(p::ComposedPreprocessor)(obs) = reduce((x, f) -> f(x), p.preprocessors, init = obs)

#####
# CloneStatePreprocessor
#####

"""
    CloneStatePreprocessor()

Do `deepcopy` for the state in an observation.
"""
struct CloneStatePreprocessor <: AbstractPreprocessor end

(p::CloneStatePreprocessor)(obs) = StateOverriddenObs(obs, deepcopy(get_state(obs)))


text:
function get_xyz_vector_from_rφz_field_vector_at_rφz(field,r,φ,z,ir,iφ,iz)::Vector
    startpoint_vector = get_xyz_vector_from_rφz_vector([r,φ,z])
    endpoint_vector = get_xyz_vector_from_rφz_vector([r,φ,z]+field[ir,iφ,iz])
    xyz_vector = endpoint_vector-startpoint_vector
    for ic in 1:size(xyz_vector,1)
        isapprox(xyz_vector[ic],0.0) ? xyz_vector[ic] = 0.0 : nothing
    end
    xyz_vector
end

function get_xy_magnitude(xyz_vector::AbstractArray)
    sqrt(xyz_vector[1]^2+xyz_vector[2]^2)
end

@recipe function f(electrical_field::Array{SVector{3,T},3}, grid::CylindricalGrid{T}; view=:Components, plane=:rz, i_fixed=3, spacing = 8, vectorscale = 0.0018, SI_factor=1/1000.) where{T <: SSDFloat}

    vectorfield = electrical_field.*SI_factor
    units = Dict(1e-3=>"mm",1e-2=>"cm",1e-1=>"dm",1=>"m",1.0=>"m")
    if view == :Components
        if plane == :rz
            layout := (2,2)
            size := (800,900)
            vectorfield_r = get_component_field(vectorfield);
            vectorfield_φ = get_component_field(vectorfield,:phi);
            vectorfield_z = get_component_field(vectorfield,:z);
            vectorfield_magn = get_magnitude_of_rφz_vector.(vectorfield);

            st := :heatmap
            # colorbar_title := "Field Strength / V m\$\^\{-1\}\$"
            @series begin
                subplot := 1
                title := "r_component"
                ylabel --> "z ["*units[SI_factor]*"]"
                grid.r ./ SI_factor, grid.z ./ SI_factor, vectorfield_r[:,i_fixed,:]'
            end
            @series begin
                subplot := 2
                colorbar_title --> "Field Strength [V / "*units[SI_factor]*"]"
                title := "φ_component"

                grid.r ./SI_factor, grid.z ./SI_factor, vectorfield_φ[:,i_fixed,:]'
            end
            @series begin
                subplot := 3
                title := "z_component"
                ylabel --> "z ["*units[SI_factor]*"]"
                xlabel --> "r ["*units[SI_factor]*"]"
                grid.r ./SI_factor, grid.z ./SI_factor, vectorfield_z[:,i_fixed,:]'
            end
            @series begin
                subplot := 4
                colorbar_title --> "Field Strength [V / "*units[SI_factor]*"]"
                xlabel --> "r ["*units[SI_factor]*"]"
                title:= "magnitude"
                grid.r ./SI_factor, grid.z ./SI_factor, vectorfield_magn[:,i_fixed,:]'
            end
        end
    elseif view == :ef
        if plane == :rφ
            vectorfield_xyz = Array{Vector{Float32}}(undef,size(vectorfield,1),size(vectorfield,2),size(vectorfield,3));
            for (iz,z) in enumerate(grid.z)
                for (iφ,φ) in enumerate(grid.φ)
                    for (ir,r) in enumerate(grid.r)
                        vectorfield_xyz[ir,iφ,iz]=get_xyz_vector_from_rφz_field_vector_at_rφz(vectorfield,r,φ,z,ir,iφ,iz)
                    end
                end
            end
        elseif plane == :xy
            vectorfield_xyz = electrical_field.field
        end
        vectorfield_xy_magn = map(x->get_xy_magnitude(x),vectorfield_xyz[:,:,i_fixed])
        max_magn = maximum(vectorfield_xy_magn)
        diff_magn = max_magn-minimum(vectorfield_xy_magn)

        size := (800,800)
        line := (:arrow,:blue,2)
        label := ""
        ylabel := "y "
        xlabel := "x "
        title := "z = $(round(grid.z[i_fixed]/SI_factor,digits=2)) / mm"
        xlims := (-1.2/SI_factor*maximum(grid.r),1.2/SI_factor*maximum(grid.r))
        ylims := (-1.2/SI_factor*maximum(grid.r),1.2/SI_factor*maximum(grid.r))
        for (ir,r) in enumerate(grid.r[1:spacing:end])
        for (iφ,φ) in enumerate(grid.φ)
                x= r*cos(φ)
            y= r*sin(φ)
            ir_actual=findfirst(x->x==r,grid.r)
            iφ_actual=findfirst(x->x==φ,grid.φ)
                xy_magn = vectorfield_xy_magn[ir_actual,iφ_actual]
                vector=vectorfield_xyz[ir_actual,iφ_actual,i_fixed]/xy_magn
                vector*=((vectorfield_xy_magn[ir_actual,iφ_actual]-0.8*minimum(vectorfield_xy_magn))/diff_magn)
                vector*=vectorscale
                @series begin
                    [x-0.5*vector[1],x+0.5*vector[1]]/SI_factor, [y-0.5*vector[2],y+0.5*vector[2]]/SI_factor
                end
            end
        end
    end
end

using Base.Math
@recipe function f( electric_field::Array{ <:StaticVector{3, T}, 3}, det::SolidStateDetector; φ=missing, z = missing, scaling_factor=1.0, spacings = [3,20]) where T
    ismissing(φ) && !ismissing(z) ? iz = searchsortednearest(det.zs,T(z)) : nothing
    !ismissing(φ) && ismissing(z) ? iφ = searchsortednearest(det.φs,T(φ)) : nothing
    guidefontsize --> 15
    tickfontsize -->15
    # labelfontsize --> 13
    if ismissing(φ)
        xy = T[0, 0]
        values = T[0, 0]
        magnitudes_xy = []
        for (iφ ,φ)  in enumerate(det.φs)
            if (iφ+spacings[1]-1)%spacings[1]==0
                for (ir, r) in enumerate(det.rs)
                    if (ir+spacings[2]-1)%spacings[2] == 0
                        xy = hcat(xy,[ r * cos(φ), r * sin(φ)])
                        values = hcat(values, electric_field[ir, iφ, iz][1:2])
                        push!(magnitudes_xy,sqrt(electric_field[ir, iφ, iz][1]^2 + electric_field[ir, iφ, iz][2]^2))
                    end
                end
            end
        end

        num = minimum([mean(diff(det.φs)*spacings[1]),mean(diff(det.rs))*spacings[2]])
        nom = maximum(magnitudes_xy)
        scaling = num/nom * scaling_factor
        for i in 1:size(xy,2)
            @series begin
                title --> "z = $(round(det.zs[iz] .* 1000,digits=2)) mm"
                xlabel --> "x / mm"
                ylabel --> "y / mm"
                aspect_ratio --> 1
                label --> ""
                color --> :red
                arrow --> true
                [xy[1,i] - 0.5*scaling*values[1,i], xy[1,i] + 0.5*scaling*values[1,i]] .* 1000, [xy[2,i] - 0.5*scaling*values[2,i], xy[2,i] + 0.5*scaling*values[2,i]] .* 1000
            end
        end

    elseif ismissing(z)
        rz = T[0, 0]
        values = T[0, 0]
        magnitudes_rz = []
        for (iz ,z)  in enumerate(det.zs)
            if (iz+spacings[1])%spacings[1]==0
                for (ir, r) in enumerate(det.rs)
                    if (ir+spacings[2]-1)%spacings[2] == 0
                        rz = hcat(rz,[ r, z])
                        values = hcat(values, electric_field[ir, iφ, iz][1:2:3])
                        push!(magnitudes_rz,sqrt(electric_field[ir, iφ, iz][1]^2 + electric_field[ir, iφ, iz][3]^2))
                    end
                end
            end
        end
        num = minimum([mean(diff(det.zs)*spacings[1]),mean(diff(det.rs))*spacings[2]])
        nom = maximum(magnitudes_rz)
        scaling = num/nom * scaling_factor
        for i in 1:size(rz,2)
            @series begin
                title --> "φ = $(round(rad2deg(det.φs[iφ]) , digits=2)) deg"
                xlabel --> "r [mm]"
                ylabel --> "z [mm]"
                aspect_ratio --> 1
                label --> ""
                color --> :red
                arrow --> true
                [rz[1,i] - 0.5*scaling*values[1,i], rz[1,i] + 0.5*scaling*values[1,i]] .* 1000, [rz[2,i] - 0.5*scaling*values[2,i], rz[2,i] + 0.5*scaling*values[2,i]] .* 1000
            end
        end
    end
end

@userplot MyQuiver
@recipe function f(gdd::MyQuiver; scaling=1.0)
    xy::Matrix = gdd.args[1]
    values = gdd.args[2]
    for i in 1:size(xy,2)
        @series begin
            arrow --> true
            [xy[1,i] - 0.5*scaling*values[1,i], xy[1,i] + 0.5*scaling*values[1,i]], [xy[2,i] - 0.5*scaling*values[2,i], xy[2,i] + 0.5*scaling*values[2,i]]
        end
    end
end

@userplot Plot_electric_field
@recipe function f(gdd::Plot_electric_field; φ = missing, r = missing, x = missing, y = missing, z = missing,
                    spacing = 4, grid_spacing=[0.0005, deg2rad(1.0), 0.0005], n_steps=3000,
                    potential=true, contours_equal_potential=true, offset = (5e-5))
    sim = gdd.args[1]
    S = get_coordinate_system(sim.electric_field.grid)
    T = Float32

    dim_array = [φ, r, x, y, z]
    dim_symbols_array = [:φ, :r, :x, :y, :z]
    if isempty(skipmissing(dim_array))
        if S == :cylindrical
            v::T = 0
            dim_number = 2
            dim_symbol = :φ
        elseif S == :cartesian
            dim_number = 1
            dim_symbol = :x
            v = sim.electric_field.grid[dim_number][ div(length(sim.electric_field.grid[dim_number]), 2) ]
        end
    elseif sum(ismissing.(dim_array)) == 4
        dim_number = findfirst(x -> !ismissing(x), dim_array)
        dim_symbol = dim_symbols_array[dim_number]
        v = dim_array[dim_number]
        if dim_symbol == :φ v = deg2rad(v) end
    else
        throw(ArgumentError("Only one keyword for a certain dimension is allowed. One of 'φ', 'r', 'x', 'y', 'z'"))
    end

    S::Symbol = get_coordinate_system(sim.detector)
    aspect_ratio --> 1
    title --> "Electric Field Lines @$(dim_symbol)=$(round(dim_symbol == :φ ? rad2deg(v) : v, digits=2))"
    xlabel --> (S == :cylindrical ? (dim_symbol == :r ? L"$\varphi$ / °" : L"$r$ / m") : (dim_symbol == :x ? L"$y$ / m" : L"$x$ / m"))
    ylabel --> L"$z$ / m"

    if potential == true
        @series begin
            # contours_equal_potential --> contours_equal_potential
            if dim_symbol == :φ
                φ --> v
            elseif dim_symbol == :z
                if S == :cylindrical error("Electric field plot not yet implemented for z in cylindrical coordinates") end
                z --> v
                if S == :cylindrical
                    proj --> :polar
                end
            elseif dim_symbol == :r
                error("Electric field plot not yet implemented for r")
                r --> v
            elseif dim_symbol == :x
                x --> v
            elseif dim_symbol == :y
                y --> v
            end
            sim.electric_potential
        end
    end
    PT = S == :cylindrical ? CylindricalPoint{T} : CartesianPoint{T}

    contacts_to_spawn_charges_for = filter!(x -> x.id !=1, SSD.Contact{T}[c for c in sim.detector.contacts])
    spawn_positions = PT[]
    grid = sim.electric_field.grid# Grid(sim.detector, init_grid_spacing = grid_spacing, full_2π = true)
    pt_offset = T[offset,0.0,offset]

    # for c in contacts_to_spawn_charges_for
    #     ongrid_positions= map(x-> CylindricalPoint{T}(grid[x...]), SSD.paint_object(sim.detector, c, grid, Val(dim_symbol), v))
    #     for position in ongrid_positions
    #         push!(spawn_positions, CylindricalPoint{T}((position + pt_offset)...))
    #         push!(spawn_positions, CylindricalPoint{T}((position - pt_offset)...))
    #     end
    # end

    ax1 = unique(range(sim.electric_field.grid[1][1], stop = sim.electric_field.grid[1][end], step = spacing * 0.001))
    ax2 = unique(range(sim.electric_field.grid[2][1], stop = sim.electric_field.grid[2][end], step = spacing * 0.001))
    if ismissing(φ) && S == :cylindrical
        ax2 = T[0]
    elseif !ismissing(φ)
        ax2 = T[φ]
    end
    ax3 = unique(range(sim.electric_field.grid[3][1], stop = sim.electric_field.grid[3][end], step = spacing * 0.001))
    for x1 in ax1
        for x2 in ax2
            for x3 in ax3
                pt::PT = PT(x1, x2, x3)
                push!(spawn_positions, pt)
            end
        end
    end
    filter!(x -> x in sim.detector && !in(x, sim.detector.contacts), spawn_positions)

    @info "$(length(spawn_positions)) drifts are now being simulated..."

    el_field_itp     = SSD.get_interpolated_drift_field(sim.electric_field.data,       sim.electric_field.grid)
    el_field_itp_inv = SSD.get_interpolated_drift_field(sim.electric_field.data .* -1, sim.electric_field.grid)
    @showprogress for (ipos, pos) in enumerate(spawn_positions)
        if ((spacing-1)+ipos)%spacing == 0

            path = CartesianPoint{T}[CartesianPoint{T}(0.0,0.0,0.0) for i in 1:n_steps]
            SSD.drift_charge!(path, sim.detector, sim.point_types, sim.electric_potential.grid, CartesianPoint(pos), T(2e-9), el_field_itp, verbose = false )
            @series begin
                c --> :white
                if dim_symbol == :z && S == :cylindrical proj --> :polar end
                label --> ""
                x, y = if dim_symbol == :φ
                    map(x -> sqrt(x[1]^2+x[2]^2), path),
                    map(x -> x[3], path)
                elseif dim_symbol == :x
                    map(x -> x[2], path),
                    map(x -> x[3], path)
                elseif dim_symbol == :y
                    map(x -> x[1], path),
                    map(x -> x[3], path)
                elseif dim_symbol == :z
                    map(x -> x[1], path),
                    map(x -> x[2], path)
                end
                x, y
            end

            path = CartesianPoint{T}[CartesianPoint{T}(0.0,0.0,0.0) for i in 1:n_steps]
            SSD.drift_charge!(path, sim.detector, sim.point_types, sim.electric_potential.grid, CartesianPoint(pos), T(2e-9), el_field_itp_inv, verbose = false )
            @series begin
                c --> :white
                if dim_symbol == :z && S == :cylindrical proj --> :polar end
                label --> ""
                x, y = if dim_symbol == :φ
                    map(x -> sqrt(x[1]^2+x[2]^2), path),
                    map(x -> x[3], path)
                elseif dim_symbol == :x
                    map(x -> x[2], path),
                    map(x -> x[3], path)
                elseif dim_symbol == :y
                    map(x -> x[1], path),
                    map(x -> x[3], path)
                elseif dim_symbol == :z
                    map(x -> x[1], path),
                    map(x -> x[2], path)
                end
                x, y
            end

        end
    end
end


text:
function load(filename)

    open(filename) do f
        return [(x[1],parse(Int,join(x[2:end]))) for x in collect.(readlines(f))]
    end

end

mutable struct Ship

    d :: Int
    x :: Int
    y :: Int

    Ship() = new( 0, 0, 0)

end

function move!( ship, cmd )

    d, l = cmd

    if d == 'N'
       ship.y += l
    elseif d == 'S'
       ship.y -= l
    elseif d == 'E'
       ship.x += l
    elseif d == 'W'
       ship.x -= l
    elseif d == 'L'
       ship.d += l
    elseif d == 'R'
       ship.d -= l
    elseif d == 'F'
       ship.x += round(Int,cos(ship.d * π / 180) * l)
       ship.y += round(Int,sin(ship.d * π / 180) * l)
    end

end

manhattan_distance( ship ) = abs(ship.x) + abs(ship.y)

rotate(x, y, θ) = round.(Int, (x * cos(θ*π/180) - y * sin(θ*π/180),
                               x * sin(θ*π/180) + y * cos(θ*π/180)))

function move!( ship, cmd, waipoint )

    d, l = cmd
    x, y = waypoint

    if d == 'N'
       y += l
    elseif d == 'S'
       y -= l
    elseif d == 'E'
       x += l
    elseif d == 'W'
       x -= l
    elseif d == 'L'
       x, y = rotate( x, y, l)
    elseif d == 'R'
       x, y = rotate(x, y, -l)
    elseif d == 'F'
       ship.x += l * x
       ship.y += l * y
    end

    return (x, y)

end

ship = Ship()
for cmd in load("input12.txt")
   move!(ship, cmd)
end
println(manhattan_distance( ship ))

ship = Ship()
waypoint = (10, 1)
for cmd in load("input12.txt")
   global waypoint
   waypoint = move!(ship, cmd, waypoint)
end
println(manhattan_distance( ship ))

max_stars_repo_name: pnavaro/AdventOfCode

max_issues_repo_name: pnavaro/AdventOfCode

max_forks_repo_name: pnavaro/AdventOfCode


text:
module Prinz

using ModelingToolkit
using ..NeuronBuilder

import ..get_parameters, ..get_states, ..default_params, ..default_states

include("channels.jl")
include("calc_dynamics.jl")

export get_parameters, get_states, default_params, default_states

end


text:
using FEDVR
using LinearAlgebra
using Test

function vecdist(a::AbstractVector, b::AbstractVector,
                 ϵ = eps(eltype(a)))
    δ = √(sum(abs2, a-b))
    δ, δ/√(sum(abs2, a .+ ϵ))
end

@testset "grid" begin
    N = 11
    n = 5
    breaks = range(0, stop=1, length=N)
    grid = FEDVR.Grid(breaks, n)

    @test minimum(grid) == 0
    @test maximum(grid) == 1

    @testset "dirichlet0" begin
        # We want the endpoints of the finite elements to match up exactly
        @test grid.X[1:end-1,end] == grid.X[2:end,1]

        @test grid.N[1] == 1/√(grid.W[1])
        @test grid.N[2,1] == 1/√(grid.W[1,end]+grid.W[2,1])

        @test elcount(grid) == N-1
        @test elems(grid) == 1:N-1
        @test order(grid) == n
        @test basecount(grid) == (N-1)*n - (N-2) - 2 # Dirichlet0 boundary conditions, hence -2
        Xp = locs(grid)
        @test length(Xp) == basecount(grid)
        @test Xp[1] > breaks[1]
        @test Xp[end] < breaks[end]

        gW = weights(grid)
        gN = FEDVR.boundary_sel(grid, [grid.N[:,1:end-1]'[:]..., grid.N[end]])
        @test vecdist(sqrt.(gW[2:end-1]),
                      1 ./ gN[2:end-1])[1] < eps(Float64)
    end

    @testset "intervals" begin
        for nn = 300:301
            x = range(breaks[1], stop=breaks[end], length=nn)
            sel = 1:1
            sel = FEDVR.find_interval(grid.X, x, 1, sel)
            @test x[sel[1]] == breaks[1]
            for i = elems(grid)[1:end-1]
                sel = FEDVR.find_interval(grid.X, x, i, sel)
                @test x[sel[1]] >= breaks[i]
                @test x[sel[end]] < breaks[i+1]
            end
            i = elcount(grid)
            sel = FEDVR.find_interval(grid.X, x, i, sel)
            @test x[sel[1]] >= breaks[i]
            @test x[sel[end]] == breaks[end]
        end
    end

    @testset "misc" begin
        @test !isempty(string(grid))

        @test begin
            using RecipesBase
            RecipesBase.apply_recipe(Dict{Symbol,Any}(), grid)
            true
        end
    end

    @testset "dirichlet1" begin
        N = 11
        n = 5
        breaks = range(0, stop=1, length=N)
        grid = FEDVR.Grid(breaks, n, :dirichlet1, :dirichlet1)

        # We want the endpoints of the finite elements to match up exactly
        @test grid.X[1:end-1,end] == grid.X[2:end,1]

        @test grid.N[1] == 1/√(grid.W[1])
        @test grid.N[2,1] == 1/√(grid.W[1,end]+grid.W[2,1])

        @test elcount(grid) == N-1
        @test elems(grid) == 1:N-1
        @test order(grid) == n
        @test basecount(grid) == (N-1)*n - (N-2)
        Xp = locs(grid)
        @test length(Xp) == basecount(grid)
        @test Xp[1] == breaks[1]
        @test Xp[end] == breaks[end]
    end
end

@testset "lagrange" begin
    N = 11
    n = 5
    breaks = range(0, stop=1, length=N)
    grid = FEDVR.Grid(breaks, n)
    e = m -> vec([zeros(m-1);1;zeros(n-m)])
    for i = elems(grid)
        for m = 1:n
            @test FEDVR.lagrange(grid.X[i,:], m, grid.X[i,m]) == 1
            @test FEDVR.lagrange(grid.X[i,:], m, grid.X[i,:]) == e(m)
        end
    end

    @test FEDVR.δ(-1, -1) == 1
    @test FEDVR.δ(-1, 1) == 0
    @test FEDVR.δ(0, 0) == 1
    @test FEDVR.δ(0, 1) == 0

    f = x -> x^4
    g = x -> 4x^3

    L′ = FEDVR.lagrangeder(grid)

    for i = elems(grid)
        c = f.(grid.X[i,:])
        c′ = L′[i,:,:]'c
        c′a = g.(grid.X[i,:])
        @test vecdist(c′, c′a)[2] < 1e-13
    end
end

@testset "basis" begin
    N = 11
    n = 5
    breaks = range(0, stop=1, length=N)
    for bdrl = [:dirichlet0, :dirichlet1]
        for bdrr = [:dirichlet0, :dirichlet1]
            basis = FEDVR.Basis(breaks, n, bdrl, bdrr)

            x = locs(basis.grid)
            χ = basis(x)
            dχ = diag(χ)
            gN = [basis.grid.N[:,1:end-1]'[:]..., basis.grid.N[end]]
            @test dχ == FEDVR.boundary_sel(basis.grid, gN)

            @testset "misc" begin
                @test !isempty(string(basis))

                @test begin
                    using RecipesBase
                    RecipesBase.apply_recipe(Dict{Symbol,Any}(), basis)
                    true
                end
            end
        end
    end

    @testset "eval on subset" begin
        basis = FEDVR.Basis(breaks, n)
        x₁ = range(0, stop=1, length=35)
        x₂ = x₁[8:19]
        x₃ = range(-1, stop=-0.5, length=11)
        χ₁ = basis(x₁)
        χ₂ = basis(x₂)
        χ₃ = basis(x₃)
        δχ = χ₁[8:19,:] - χ₂
        @test norm(δχ) == 0
        @test norm(χ₃) == 0
    end
end

@testset "projections" begin
    breaks = range(0, stop=1, length=11)
    n = 5
    basis = FEDVR.Basis(breaks, n, :dirichlet1, :dirichlet1)
    x = range(minimum(breaks), stop=maximum(breaks), length=301)
    χ = basis(x)

    f = x -> x^3 - 7x^2 + x^4 + 2
    ϕ = project(f, basis)

    @test vecdist(f.(x), χ*ϕ)[2] < 10eps(Float64)
end

# @testset "derivatives" begin
#     N = 30
#     n = 10
#     L = 5.0
#     xx = range(0, stop=L, length=N+1)
#     basis = FEDVR.Basis(xx, n)
#     T = kinop(basis)
#     x = locs(basis.grid)
#     χ = basis(x)

#     mmax = 30
#     λ,ϕ = eigs(T,which=:SR,nev=mmax)
#     m = eachindex(λ)
#     λa = 0.5*(π*m/L).^2

#     @test vecdist(real.(λ), λa)[2] < 1e-11

#     for m = 1:mmax
#         ϕa = √(2/L) * sin.(m*π*x/L)
#         ϕm = real(ϕ[:,m])
#         ϕm *= sign(ϕm[1])
#         @test vecdist(ϕa, χ*ϕm)[2] < 1e-8
#     end
# end


text:
using Weber
using Base.Test
include("find_timing.jl")

const check_timing = get(ENV,"WEBER_TIMING_TESTS","Yes") != "No"
const moment_eps = 2.5e-3

_,many_times,_ = find_timing() do
  addtrial(repeated(moment(0.5ms,() -> record(:a)),1000))
end

if check_timing
  @testset "Timing" begin
    diffs = abs.(0.0005 - diff(many_times))
    middle99 = diffs[quantile(diffs,0.005) .< diffs .< quantile(diffs,0.995)]
    @test mean(middle99) < moment_eps
  end
end


text:
using ContextExprRules

grammar = @cgrammar begin
	Real = 1 | 2 | 3
	Real = Real + Real
	Real = Real * Real
end

# iter = ConstrExpressionIterator(grammar, 2, :Real)
# elems = collect(iter)

constraint = ComesAfter(3,5)

addconstraint!(grammar, constraint)

iter2 = ContextExpressionIterator(grammar, 2, :Real)
elems2 = collect(iter2)


text:
abstract type InverterComponent <: DynamicComponent end

mutable struct DynamicInverter{
    C <: Converter,
    O <: OuterControl,
    VC <: VSControl,
    DC <: DCSource,
    P <: FrequencyEstimator,
    F <: Filter,
} <: DynamicInjection
    number::Int64
    name::String
    bus::Bus
    ω_ref::Float64
    V_ref::Float64
    P_ref::Float64
    Q_ref::Float64
    MVABase::Float64
    converter::C
    outercontrol::O
    vscontrol::VC
    dc_source::DC
    freq_estimator::P
    filter::F #add MVAbase here
    n_states::Int64
    states::Vector{Symbol}
    ext::Dict{String, Any}
    function DynamicInverter(
        number::Int64,
        name::String,
        bus::Bus,
        ω_ref::Float64,
        V_ref::Float64,
        P_ref::Float64,
        Q_ref::Float64,
        MVABase::Float64,
        converter::C,
        outercontrol::O,
        vscontrol::VC,
        dc_source::DC,
        freq_estimator::P,
        filter::F,
    ) where {
        C <: Converter,
        O <: OuterControl,
        VC <: VSControl,
        DC <: DCSource,
        P <: FrequencyEstimator,
        F <: Filter,
    }

        n_states = (
            converter.n_states +
            outercontrol.n_states +
            vscontrol.n_states +
            dc_source.n_states +
            freq_estimator.n_states +
            filter.n_states
        )

        states = vcat(
            converter.states,
            outercontrol.states,
            vscontrol.states,
            dc_source.states,
            freq_estimator.states,
            filter.states,
        )

        new{C, O, VC, DC, P, F}(
            number,
            name,
            bus,
            ω_ref,
            V_ref,
            P_ref,
            Q_ref,
            MVABase,
            converter,
            outercontrol,
            vscontrol,
            dc_source,
            freq_estimator,
            filter,
            n_states,
            states,
            Dict{String, Any}(),
        )

    end
end

get_inverter_Sbase(device::DynamicInverter) = device.converter.s_rated
get_inverter_Vref(device::DynamicInverter) = device.V_ref
get_number(device::DynamicInverter) = device.number
get_name(device::DynamicInverter) = device.name
get_ω_ref(device::DynamicInverter) = device.ω_ref
get_V_ref(device::DynamicInverter) = device.V_ref
get_P_ref(device::DynamicInverter) = device.P_ref
get_Q_ref(device::DynamicInverter) = device.Q_ref
get_bus(device::DynamicInverter) = device.bus
get_ext(device::DynamicInverter) = device.ext
get_states(device::DynamicInverter) = device.states
get_n_states(device::DynamicInverter) = device.n_states
get_converter(device::DynamicInverter) = device.converter
get_outercontrol(device::DynamicInverter) = device.outercontrol
get_vscontrol(device::DynamicInverter) = device.vscontrol
get_dc_source(device::DynamicInverter) = device.dc_source
get_freq_estimator(device::DynamicInverter) = device.freq_estimator
get_filter(device::DynamicInverter) = device.filter


text:
"""

        composite_and_radialaverage_withmask(azimuthal_average,composite_addition,buf,radiusbins,surface_pressure,variable_of_interest)
    """

function composite_and_radialaverage_in_chunks!(azimuthal_average,composite_addition,buf1,buf2,radiusbins,variable_name,pre_iterator_time,nchunks,file_data,diagnostics_data_dir,exp_name,centers_labels_and_cyclones;maskcyclones = true)
    totalcyclonecount = 0
    variables_diag = ["convec_heating_anomaly","rad_heating_anomaly","buoyancy_anomaly"]
    variables_orig_3d = ["U","V","QV","TABS","QRAD","PP","W"]
    variables_orig_2d = ["PSFC","USFC","VSFC","PW","Prec","LHF","SHF"]
    chunk_size_2d = 2*length(pre_iterator_time)÷nchunks
    chunk_size_3d = length(pre_iterator_time)÷nchunks

    nd = ndims = if (variable_name ∈ variables_orig_3d)
            nd = ndims = 4
        elseif (variable_name ∈ variables_diag)
            nd = ndims = 4
        elseif (variable_name ∈ variables_orig_2d)
            nd = ndims = 3
    end
    for chunk in 1:nchunks
        iterator_time_2d = (1200 + (chunk - 1)*chunk_size_2d + 1 ): 2 : (1200 + chunk*chunk_size_2d)
        iterator_time_data =  ((chunk - 1)*chunk_size_3d + 1 ) : (chunk*chunk_size_3d)
        iterator_time_3d = 600 .+ iterator_time_data
        if (variable_name ∈ variables_orig_3d)
            variable_of_interest = Dataset(file_data) do ds
                variable(ds,variable_name)[:,:,:,iterator_time_3d]     :: Array{Float32,4}
            end
        elseif (variable_name ∈ variables_diag)
            variable_of_interest = gather_diagnostic_3d_in_chunks(diagnostics_data_dir,exp_name,variable_name,chunk,chunk_size_3d,nchunks)
        elseif (variable_name ∈ variables_orig_2d)
            variable_of_interest = Dataset(file_data) do ds
                variable(ds,variable_name)[:,:,iterator_time_2d]     :: Array{Float32,3}
            end
        end

        #centers_labels_and_cyclones = jldopen(file_centers) do file
        #    read(file,"centers_data")
        #end

        #@info nchunks size(variable_of_interest) #size(centers_labels_and_cyclones)

        for timeindex in 1:size(variable_of_interest,nd)
            centers_and_labels,cyclones = centers_labels_and_cyclones[timeindex]
            if !isnothing(centers_and_labels[1])
                count = add_allcyclones!(composite_addition,buf1,buf2,selectdim(variable_of_interest,nd,timeindex),cyclones,centers_and_labels;maskcyclones)
                totalcyclonecount += count
            end
        end
    end

    if !iszero(totalcyclonecount)
        composite_addition ./= totalcyclonecount
        if nd == 4
            for idx in 1:(length(radiusbins) - 1)
                azimuthal_average[idx,:] .= averageallindistance((radiusbins[idx],radiusbins[idx+1]),composite_addition,(256,256), 2000 )
            end
        elseif nd == 3
            for idx in 1:(length(radiusbins) - 1)
                azimuthal_average[idx] = averageallindistance((radiusbins[idx],radiusbins[idx+1]),composite_addition,(256,256), 2000 )
            end
        end

    end
end

function gather_diagnostic_3d_in_chunks(data_dir,exp_name,desired_var,chunk,chunksize,nchunks)
    diagnostic = zeros(Float32,512,512,80,chunksize)
    howmanyfiles = 12÷nchunks
    @info "how many files" howmanyfiles
    for iter in 1:howmanyfiles
        filenumber = ((chunk - 1)*howmanyfiles + iter)

        @info "this is iter" iter
        @info "file number and chunk:" filenumber chunk
        file_name = string(exp_name,"_nosmoothing",filenumber+12,".jld")
        @info "I am opening the file $file_name and reading the variable $desired_var"
        file_path = joinpath(data_dir,file_name)
        file = jldopen(file_path, "r")
        index_initial = (iter - 1)*(50)+1
        index_end = iter*50
        @info "indices to write:" index_initial,index_end
        var_from_file = read(file, desired_var)
        diagnostic[:,:,:,index_initial:index_end] .=  var_from_file[:,:,:,:]
        close(file)
    end
    return diagnostic
end

"""
        getcomposites(file2d,file3d,diagnostics_dir,exp_,initial_timeindex,final_timeindex,partnumber,outputInterval)

"""
function get_composites_in_chunks(file2d,file3d,exp_name)

    @info "Starting composite routine"
    diagnostics_data_dir = "/global/cscratch1/sd/aramreye/for_postprocessing/ApeBudgetOutputs_nosmoothing/"
    composites_data_dir = "/global/cscratch1/sd/aramreye/for_postprocessing/CompositeOutputs_nosmoothing/"
    azimuthal_averages_data_dir = "/global/cscratch1/sd/aramreye/for_postprocessing/AzimuthualAverageOutputs_nosmoothing/"
    centers_data_file = string(composites_data_dir,"centers_detection_",exp_name,".jld")

    #variables_diag = ["convec_heating_anomaly","rad_heating_anomaly","buoyancy_anomaly"]
    #variables_orig_3d = ["U","V","QV","TABS","QRAD","PP","W"]
    #variables_orig_2d = ["PSFC","USFC","VSFC","PW","Prec","LHF","SHF"]

    initial_timeindex = 601
    #    initial_timeindex = 1195
    final_timeindex = 1200
    #final_timeindex = 650

    pre_iterator_time_2d    = initial_timeindex*2:2:final_timeindex*2
    pre_iterator_time = pre_iterator_time_3d   = initial_timeindex:final_timeindex

    radiusbins = 1000:2000:512000;
    ds3d                = Dataset(file3d)
    ds2d                = Dataset(file2d)
    z                   = variable(ds3d,"z")[:]                       :: Array{Float32,1}
    t                   = variable(ds3d,"time")[pre_iterator_time_3d]     :: Array{Float32,1}
    surface_pressure    = variable(ds2d,"PSFC")[:,:,pre_iterator_time_2d] :: Array{Float32,3}
    close(ds2d)
    close(ds3d)
    pressure_anomaly    = surface_pressure .- mean(surface_pressure,dims=(1,2))

    @info size(surface_pressure)

    ##Create buffers to reuse on allocating functions

    size_var_of_interest_2d = (512,512)
    size_var_of_interest_3d = (size_var_of_interest_2d...,length(z))
    ## Create 3D buffers
    composite_addition_3d   = zeros(Float32,size_var_of_interest_3d)
    buf_3d_2                = similar(composite_addition_3d)
    buf_3d_1                = similar(composite_addition_3d)
    azimuthal_average_3d    = zeros(Float32,length(radiusbins) - 1,length(z))
    ## Create 2D buffers
    composite_addition_2d   = zeros(Float32,size_var_of_interest_2d)
    buf_2d_1                = similar(composite_addition_2d)
    buf_2d_2                = similar(composite_addition_2d)
    azimuthal_average_2d    = zeros(Float32,length(radiusbins) - 1)
    @info size_var_of_interest_2d size_var_of_interest_3d

    centers_labels_and_cyclones = @views [
        detect_cyclones!(buf_2d_1, buf_2d_2, pressure_anomaly[:,:,timeindex],-9.0,2000) for
        timeindex in 1:size(surface_pressure,3)]; pressure_anomaly = []; surface_pressure = [];

    composite_filename_withmask = string(composites_data_dir,exp_name        ,"_withmask.jld")
    composite_filename_nomask   = string(composites_data_dir,exp_name        ,"_nomask.jld")
    azimuthal_filename_withmask = string(azimuthal_averages_data_dir,exp_name,"_withmask.jld")
    azimuthal_filename_nomask   = string(azimuthal_averages_data_dir,exp_name,"_nomask.jld")

    ### *** add allcyclones ***
    @info "Creating composite for surface pressure "
    flush(stdout)
    composite_addition_2d .= 0.0
    buf_2d_1              .= 0.0
    buf_2d_2              .= 0.0
    azimuthal_average_2d  .= 0.0

    nchunks=6

    #create files:
    ds =  jldopen(composite_filename_withmask, "w") do file end
    ds =  jldopen(composite_filename_nomask, "w") do file end
    ds =  jldopen(azimuthal_filename_withmask, "w") do file end
    ds =  jldopen(azimuthal_filename_nomask, "w") do file end

    for current_variable in variables_orig_3d

         composite_addition_3d .= 0.0
         buf_3d_1              .= 0.0
         buf_3d_2              .= 0.0
         azimuthal_average_3d  .= 0.0

    #     ## Create products with mask
         @info "Creating composite for: " current_variable
         flush(stdout)

        composite_and_radialaverage_in_chunks!(azimuthal_average_3d,composite_addition_3d,buf_3d_1,buf_3d_2,radiusbins,current_variable,pre_iterator_time_3d,nchunks,file3d,diagnostics_data_dir,exp_name,centers_labels_and_cyclones;maskcyclones = true)

         @info "Writing"
         flush(stdout)
         jldopen(composite_filename_withmask, "r+",mmaparrays=true) do file
             write(file,current_variable,composite_addition_3d)
         end
         jldopen(azimuthal_filename_withmask, "r+",mmaparrays=true) do file
             write(file,current_variable,azimuthal_average_3d)
         end
         GC.gc()
         GC.gc()
    #     ## Create products without mask
         @info "Creating composite without mask for: " current_variable
         flush(stdout)
         composite_addition_3d .= 0.0
         buf_3d_1              .= 0.0
         buf_3d_2              .= 0.0
         azimuthal_average_3d  .= 0.0
        composite_and_radialaverage_in_chunks!(azimuthal_average_3d,composite_addition_3d,buf_3d_1,buf_3d_2,radiusbins,current_variable,pre_iterator_time_3d,nchunks,file3d,diagnostics_data_dir,exp_name,centers_labels_and_cyclones;maskcyclones = false)
         @info "Writing"
         flush(stdout)
         jldopen(composite_filename_nomask, "r+",mmaparrays=true) do file
             write(file,current_variable,composite_addition_3d)
         end
         jldopen(azimuthal_filename_nomask, "r+",mmaparrays=true) do file
             write(file,current_variable,azimuthal_average_3d)
         end
         GC.gc()
    end

    for current_variable in variables_orig_2d
         composite_addition_2d .= 0.0
         buf_2d_1              .= 0.0
         buf_2d_2              .= 0.0
         azimuthal_average_2d  .= 0.0

         ## Create products with mask
         @info "Creating composite for: " current_variable
         flush(stdout)

        composite_and_radialaverage_in_chunks!(azimuthal_average_2d,composite_addition_2d,buf_2d_1,buf_2d_2,radiusbins,current_variable,pre_iterator_time_3d,nchunks,file2d,diagnostics_data_dir,exp_name,centers_labels_and_cyclones;maskcyclones = true)
         @info "Writing"
         flush(stdout)
         jldopen(composite_filename_withmask, "r+",mmaparrays=true) do file
             write(file,current_variable,composite_addition_2d)
         end
         jldopen(azimuthal_filename_withmask, "r+",mmaparrays=true) do file
             write(file,current_variable,azimuthal_average_2d)
         end
         GC.gc()
         GC.gc()
         ## Create products without mask
         @info "Creating composite without mask for: " current_variable
         flush(stdout)
         composite_addition_2d .= 0.0
         buf_2d_1              .= 0.0
         buf_2d_2              .= 0.0
         azimuthal_average_2d  .= 0.0

         composite_and_radialaverage_in_chunks!(azimuthal_average_2d,composite_addition_2d,buf_2d_1,buf_2d_2,radiusbins,current_variable,pre_iterator_time_3d,nchunks,file2d,diagnostics_data_dir,exp_name,centers_labels_and_cyclones;maskcyclones = false)
         @info "Writing"
         flush(stdout)
         jldopen(composite_filename_nomask, "r+",mmaparrays=true) do file
             write(file,current_variable,composite_addition_2d)
         end
         jldopen(azimuthal_filename_nomask, "r+",mmaparrays=true) do file
             write(file,current_variable,azimuthal_average_2d)
         end
         GC.gc()
    end

    for current_variable in variables_diag

         composite_addition_3d .= 0.0
         buf_3d_1              .= 0.0
         buf_3d_2              .= 0.0
         azimuthal_average_3d  .= 0.0
         ## Create products with mask
         @info "Creating composite for: " current_variable
         flush(stdout)

        #try
            composite_and_radialaverage_in_chunks!(azimuthal_average_3d,composite_addition_3d,buf_3d_1,buf_3d_2,radiusbins,current_variable,pre_iterator_time_3d,nchunks,file3d,diagnostics_data_dir,exp_name,centers_labels_and_cyclones;maskcyclones = true)
        #catch
        #    nothing
        #end
         @info "Writing"
         flush(stdout)
         jldopen(composite_filename_withmask, "r+",mmaparrays=true) do file
             write(file,current_variable,composite_addition_3d)
         end
         jldopen(azimuthal_filename_withmask, "r+",mmaparrays=true) do file
             write(file,current_variable,azimuthal_average_3d)
         end
         GC.gc()
         GC.gc()
         ## Create products without mask
         @info "Creating composite without mask for: " current_variable
         flush(stdout)
         composite_addition_3d .= 0.0
         buf_3d_1              .= 0.0
         buf_3d_2              .= 0.0
          azimuthal_average_3d  .= 0.0
        #try
            composite_and_radialaverage_in_chunks!(azimuthal_average_3d,composite_addition_3d,buf_3d_1,buf_3d_2,radiusbins,current_variable,pre_iterator_time_3d,nchunks,file3d,diagnostics_data_dir,exp_name,centers_labels_and_cyclones;maskcyclones = false)
        #catch
        #    nothing
        #end

         @info "Writing"
         flush(stdout)
         jldopen(composite_filename_nomask, "r+",mmaparrays=true) do file
             write(file,current_variable,composite_addition_3d)
         end
         jldopen(azimuthal_filename_nomask, "r+",mmaparrays=true) do file
             write(file,current_variable,azimuthal_average_3d)
         end
         GC.gc()
         GC.gc()
    end

    return nothing

end


text:
abstract type AbstractRegionGraphNode end

function setscope!(n::AbstractRegionGraphNode, dims::Vector{Int})
    fill!(n.scopeVecs, false)
    n.scopeVecs[dims] .= true
end
nscope(d::AbstractRegionGraphNode) = sum(d.scopeVecs)
hasscope(d::AbstractRegionGraphNode) = any(d.scopeVecs)
scope(d::AbstractRegionGraphNode) = findall(d.scopeVecs)

function setscope!(d::AbstractRegionGraphNode, i::Int, s::Bool)
    d.scopeVecs[i] = s
end

length(d::AbstractRegionGraphNode) = size(d.obsVecs,2)
id(d::AbstractRegionGraphNode) = d.id

"""
    topoligicalorder(root::AbstractRegionGraphNode)

Return topological order of the nodes in the SPN.
"""
function topoligicalorder(root::AbstractRegionGraphNode)
  visitedNodes = Vector{AbstractRegionGraphNode}()
  visitNode!(root, visitedNodes)
  return visitedNodes
end

function visitNode!(node::AbstractRegionGraphNode, visitedNodes)
  # check if we have already visited this node
  if !(node in visitedNodes)

    # visit node
    if haschildren(node)
        for child in children(node)
            visitNode!(child, visitedNodes)
        end
    end
    push!(visitedNodes, node)
  end
end

"""
    Region node.

Parameters:

* ids::Symbol                                           Id
* scopeVecs::Vector{Bool}                               Active dimensions (D)
* obsVecs::Matrix{Bool}                                 Active observations (N x K)
* logweights::Matrix{<:Real}                            Log weights of sum nodes (Ch x K)
* prior::Dirichlet                                      Prior for sum nodes
* children::Vector{AbstractRegionGraphNode}             Children of region

"""
struct RegionGraphNode{T<:Real} <: AbstractRegionGraphNode
    id::Symbol
    scopeVecs::Vector{Bool}
    obsVecs::Matrix{Bool}
    logweights::Matrix{T}
    active::Matrix{Bool}
    prior::Dirichlet
    children::Vector{<:AbstractRegionGraphNode}
end

"""
    Partition node.

Parameters:

* ids::Symbol                                   Id
* scopeVecs::Vector{Bool}                       Active dimensions (D)
* obsVecs::Matrix{Bool}                         Active observations (N x K)
* prior::Dirichlet                              Prior on product nodes
* children::Vector{<:AbstractRegionGraphNode}   Child region nodes

"""
struct PartitionGraphNode <: AbstractRegionGraphNode
    id::Symbol
    scopeVecs::Vector{Bool}
    obsVecs::Matrix{Bool}
    prior::Dirichlet
    children::Vector{<:AbstractRegionGraphNode}
end

haschildren(d::RegionGraphNode) = !isempty(d.children)
children(d::RegionGraphNode) = d.children

function updatescope!(d::RegionGraphNode)
    updatescope!.(children(d))
    setscope!(d, scope(first(d.children)))
end

function _logsumexp(y::AbstractMatrix{T}, lw::AbstractMatrix{T}) where {T<:Real}
    r = zeros(T,size(lw,2),size(y,2))
    Threads.@threads for j in 1:size(y,2)
        for k in 1:size(lw,2)
            @inbounds begin
                yi_max = y[1,j] + lw[1,k]
                for i in 2:size(y,1)
                    yi_max = max(y[i,j]+lw[i,k], yi_max)
                end
                s = zero(T)
                for i in 1:size(y,1)
                    s += exp(y[i,j]+lw[i,k] - yi_max)
                end
                r[k,j] = log(s) + yi_max
            end
        end
    end
    return transpose(r)
end

function _getchildlogpdf(child::AbstractRegionGraphNode, out::AxisArray{V}) where {V<:AbstractMatrix}
    return out[id(child)]
end

function _getchildlogpdf(child::PartitionGraphNode, out::AxisArray{V}) where {V<:AbstractMatrix}
    childids = map(id, children(child))
    lp_ = out[childids]
    return reduce(_cross_prod, lp_)
end

"""
    logpdf(d::RegionGraphNode, x::AbstractMatrix{T})

Log pdf of a region node in a region graph.
"""
function logpdf(d::RegionGraphNode, x::AbstractMatrix{T}) where {T<:Real}
    lp_ = logpdf.(d.children, Ref(x))
    return _logsumexp(reduce(hcat, lp_)', d.logweights)
end

"""
    logpdf!(d::RegionGraphNode, x::AbstractMatrix{T}, out::AxisArray{V})

Log pdf of a region node in a region graph. (in-place)
"""
function logpdf!(d::RegionGraphNode, x::AbstractMatrix{T}, out::AxisArray{V}) where {T<:Real,V<:AbstractMatrix}
    lp_ = _getchildlogpdf.(children(d), Ref(out))
    out[id(d)] = _logsumexp(reduce(hcat, lp_)', d.logweights)
    return out
end

"""
    logmllh!(n::RegionGraphNode, x::AbstractMatrix{T}, out::AxisArray{V}, y::AxisArray, d::Int)

Log marginal likelihood of a region node in a region graph for dimension d. (in-place)
"""
function logmllh!(n::RegionGraphNode,
                  x::AbstractMatrix{T},
                  out::AxisArray{V},
                  y::AxisArray,
                  d::Int) where {T<:Real,V<:AbstractVector}
    K = length(n)
    lp_ = out[id.(children(n))]
    lp_ = reduce(vcat, lp_) .* n.active
    out[id(n)] = vec(sum(lp_, dims = 1))
    return out
end

## Partition Graph Node

haschildren(d::PartitionGraphNode) = !isempty(d.children)
children(d::PartitionGraphNode) = d.children

function updatescope!(d::PartitionGraphNode)
    s_ = mapreduce(updatescope!, vcat, children(d))
    setscope!(d, unique(s_))
end

function _cross_prod(x1::AbstractMatrix{T}, x2::AbstractMatrix{T}) where {T<:Real}
    nx, ny = size(x1,2), size(x2,2)
    r = zeros(T, size(x1,1), nx*ny)
    Threads.@threads for j in 1:nx*ny
        @inbounds begin
            j1, j2 = Base._ind2sub((nx,ny), j)
            r[:,j] .= x1[:,j1] .+ x2[:,j2]
        end
    end
    return r
end

"""
    logpdf(d::PartitionGraphNode, x::AbstractMatrix{T})

Log pdf of a partition node in a region graph.
"""
function logpdf(d::PartitionGraphNode, x::AbstractMatrix{T}) where {T<:Real}
    childrn_ = children(d)
    lp_ = logpdf.(childrn_, Ref(x))
    return reduce(_cross_prod, lp_)
end

"""
    logpdf!(d::RegionGraphNode, x::AbstractMatrix{T}, out::AxisArray{V})

Log pdf of a partition node in a region graph. (in-place)
"""
function logpdf!(d::PartitionGraphNode, x::AbstractMatrix{T}, out::AxisArray{V}) where {T<:Real,V<:AbstractMatrix}
    return out
end

"""
    logmllh!(n::PartitionGraphNode, x::AbstractMatrix{T}, out::AxisArray{V}, y::AxisArray, d::Int)

Log marginal likelihood of a partition node in a region graph for dimension d. (in-place)
"""
function logmllh!(n::PartitionGraphNode,
                  x::AbstractMatrix{T},
                  out::AxisArray{V},
                  y::AxisArray,
                  d::Int) where {T<:Real,V<:AbstractVector}
    # TODO: draw y_d ~ p(y_d | y_\d, x_d)
    cs = map(c -> nscope(c), children(n))
    α = postparams(n.prior, DirichletSufficientStats(cs))

    lp_ = map(c -> sum(out[id(c)]), children(n))
    lp_ += log.(α) .- sum(α)

    k = randomindex(map(exp, lp_))
    y[id(n)] = k
    cid = id(children(n)[k])

    out[id(n)] = zeros(length(n))
    dims = tuple(map(c -> length(c), children(n))...)

    for j in 1:length(n)
        @inbounds begin
            idxs = Base._ind2sub(dims, j)
            jc_ = idxs[k]
            out[id(n)][j] = out[cid][jc_]
        end
    end
    return out
end

"""
    Factorized distribution of the scope.

Parameters:

* ids::Symbol                               Id
* scopeVecs::Vector{Bool}                   Active dimensions (D)
* obsVecs::Matrix{Bool}                     Active observations (N x K)
* priors::Vector{Distribution}              Priors for each dimension (D)
* likelihoods::Vector{Distribution}         Likelihood functions (D)
* sstats::Matrix{AbstractSufficientStats}   Sufficient stats (D x K)

"""
struct FactorizedDistributionGraphNode <: AbstractRegionGraphNode
    id::Symbol
    scopeVecs::Vector{Bool}
    obsVecs::Matrix{Bool}
    priors::Vector{<:Distribution}
    likelihoods::Vector{<:Distribution}
    sstats::Matrix{<:AbstractSufficientStats}
end

haschildren(d::FactorizedDistributionGraphNode) = false
updatescope!(d::FactorizedDistributionGraphNode) = scope(d)
obs(d::FactorizedDistributionGraphNode, k::Int) = findall(d.obsVecs[:,k])

function _logpdf!(n::FactorizedDistributionGraphNode,
                  x::AbstractMatrix{T},
                  out::AbstractMatrix{V}) where {T<:Real,V<:Real}
    fill!(out, zero(V))
    ds = scope(n)
    #Threads.@threads
    for d in ds
        sstats_ = n.sstats[d,:]
        θ = postparams.(Ref(n.priors[d]), sstats_)
        lpd_ = _logpostpdf.(Ref(n.likelihoods[d]), Ref(view(x,:,d)), θ)
        out .+= reduce(hcat, lpd_)
    end
    return out
end

"""
    logpdf(d::FactorizedDistributionGraphNode, x::AbstractMatrix{T})

Log pdf of an atomic region node in a region graph.
"""
function logpdf(n::FactorizedDistributionGraphNode, x::AbstractMatrix{T}) where {T<:Real}
    (N, D) = size(x)
    K = length(n)
    lp_ = zeros(Float64, N, K)
    return _logpdf!(n, x, lp_)
end

"""
    logpdf!(d::FactorizedDistributionGraphNode, x::AbstractMatrix{T}, out::AxisArray{V})

Log pdf of an atomic node in a region graph. (in-place)
"""
function logpdf!(n::FactorizedDistributionGraphNode, x::AbstractMatrix{T}, out::AxisArray{V}) where {T<:Real,V<:AbstractMatrix}
    (N, D) = size(x)
    K = length(n)

    i = findfirst(out.axes[1].val .== id(n))
    if !isassigned(out, i)
        out[id(n)] = zeros(Float64, N, K)
    end

    _logpdf!(n, x, out[id(n)])
    return out
end

function _logmllh!(n::FactorizedDistributionGraphNode,
                   x::AbstractMatrix{T},
                   out::AbstractVector{V},
                   d::Int) where {T<:Real,V<:Real}
    fill!(out, zero(V))
    K = length(n)

    sstats_ = n.sstats[d,:]
    θ = postparams.(Ref(n.priors[d]), sstats_)
    out .= map(k -> _logmllh(n.likelihoods[d], n.priors[d], x[obs(n,k),d], θ[k]), 1:K)
    return out
end

"""
    logmllh!(n::FactorizedDistributionGraphNode, x::AbstractMatrix{T}, out::AxisArray{V}, y::AxisArray, d::Int)

Log marginal likelihood of an atomic node in a region graph for dimension d. (in-place)
"""
function logmllh!(n::FactorizedDistributionGraphNode,
                  x::AbstractMatrix{T},
                  out::AxisArray{V},
                  y::AxisArray,
                  d::Int) where {T<:Real,V<:AbstractVector}
    (N,D) = size(x)
    K = length(n)

    i = findfirst(out.axes[1].val .== id(n))
    if !isassigned(out, i)
        out[id(n)] = zeros(Float64, K)
    end

    _logmllh!(n, x, out[id(n)], d)
    return out
end

"""
    Factorized mixture distribution of the scope.

Parameters:

* ids::Symbol                                       Id
* scopeVecs::Vector{Bool}                           Active dimensions (D)
* obsVecs::Vector{Matrix{Bool}}                     LV of active observations (N x K)
* priors::Vector{Vector{Distribution}}              Priors for each dimension D x (M)
* likelihoods::Vector{Vector{Distribution}}         Likelihood functions D x (m)
* sstats::Vector{Matrix{AbstractSufficientStats}}   Sufficient stats D x (M x K)

"""
struct FactorizedMixtureGraphNode <: AbstractRegionGraphNode
    id::Symbol
    scopeVecs::Vector{Bool}
    obsVecs::Matrix{Int}
    weightpriors::Vector{<:Dirichlet}
    logweights::Vector{Matrix{<:AbstractFloat}}
    priors::Vector{Vector{<:Distribution}}
    likelihoods::Vector{Vector{<:Distribution}}
    sstats::Vector{Matrix{<:AbstractSufficientStats}}
end

haschildren(d::FactorizedMixtureGraphNode) = false
updatescope!(d::FactorizedMixtureGraphNode) = scope(d)
length(d::FactorizedMixtureGraphNode) = size(first(d.sstats),2)

function obs(n::FactorizedMixtureGraphNode, k::Int, m::Int)
    if any(n.obsVecs[:,k] .== m)
        return findall(n.obsVecs[:,k] .== m)
    else
        return []
    end
end

"""
    In-place logsumexp for factorized mixture graph nodes.

Parameters:
* `y`: log-predictions of leaves in (N x K x M)
* `lw`: log weights in (M x K)
* `out`: intial out + log-sum-exp result in (N x K)

"""
function _logsumexp!(y::AbstractArray{T,3}, lw::AbstractMatrix{T}, out::AbstractMatrix{T}) where {T<:Real}
    Threads.@threads for j in 1:size(out,1)
        for k in 1:size(out,2)
            @inbounds begin
                yi_max = y[j,k,1] + lw[1,k]
                for i in 2:size(y,3)
                    yi_max = max(y[j,k,i]+lw[i,k], yi_max)
                end
                s = zero(T)
                for i in 1:size(y,3)
                    s += exp(y[j,k,i]+lw[i,k] - yi_max)
                end
                out[j,k] += log(s) + yi_max
            end
        end
    end
    return out
end

function _logpdf!(n::FactorizedMixtureGraphNode,
                  x::AbstractMatrix{T},
                  out::AbstractMatrix{V}) where {T<:Real,V<:Real}
    fill!(out, zero(V))
    ds = scope(n)
    for d in ds
        lp_ = Vector{Matrix}(undef, size(n.sstats[d],1))
        for m in 1:size(n.sstats[d],1)
            @inbounds begin
                sstats_ = n.sstats[d][m,:]
                θ = postparams.(Ref(n.priors[d][m]), sstats_)
                lpd_ = _logpostpdf.(Ref(n.likelihoods[d][m]), Ref(view(x,:,d)), θ)
                lp_[m] = reduce(hcat, lpd_)
            end
        end

        if length(lp_) > 1
            _logsumexp!(reduce((x,y) -> cat(x, y, dims=3), lp_), n.logweights[d], out)
        elseif length(lp_) == 1
            out[:,:] .= first(lp_)
        else
            out[:,:] .= -Inf
            @warn "No leaf distributions found..."
        end
    end
    return out
end

"""
    logpdf(d::FactorizedMixtureGraphNode, x::AbstractMatrix{T})

Log pdf of an atomic region node (mixture over likelihoods) in a region graph.
"""
function logpdf(n::FactorizedMixtureGraphNode, x::AbstractMatrix{T}) where {T<:Real}
    (N, D) = size(x)
    K = length(n)
    lp_ = zeros(Float64, N, K)
    return _logpdf!(n, x, lp_)
end

"""
    logpdf!(d::FactorizedMixtureGraphNode, x::AbstractMatrix{T}, out::AxisArray{V})

Log pdf of an atomic node (mixture over likelihoods) in a region graph. (in-place)
"""
function logpdf!(n::FactorizedMixtureGraphNode, x::AbstractMatrix{T}, out::AxisArray{V}) where {T<:Real,V<:AbstractMatrix}
    (N, D) = size(x)
    K = length(n)

    i = findfirst(out.axes[1].val .== id(n))
    if !isassigned(out, i)
        out[id(n)] = zeros(Float64, N, K)
    end

    _logpdf!(n, x, out[id(n)])
    return out
end

function _logmllh!(n::FactorizedMixtureGraphNode,
                   x::AbstractMatrix{T},
                   out::AbstractVector{V},
                   d::Int) where {T<:Real,V<:Real}
    fill!(out, zero(V))
    K = length(n)

    sstats_ = n.sstats[d]
    for m in 1:size(sstats_,1)
        θ = postparams.(Ref(n.priors[d][m]), sstats_[m,:])
        out .+= map(k -> _logmllh(n.likelihoods[d][m], n.priors[d][m], x[obs(n,k,m),d], θ[k]), 1:K)
    end

    return out
end

"""
    logmllh!(n::FactorizedMixtureGraphNode, x::AbstractMatrix{T}, out::AxisArray{V}, y::AxisArray, d::Int)

Log marginal likelihood of an atomic node (mixture over likelihoods) in a region graph for dimension d. (in-place)
"""
function logmllh!(n::FactorizedMixtureGraphNode,
                  x::AbstractMatrix{T},
                  out::AxisArray{V},
                  y::AxisArray,
                  d::Int) where {T<:Real,V<:AbstractVector}
    (N,D) = size(x)
    K = length(n)

    i = findfirst(out.axes[1].val .== id(n))
    if !isassigned(out, i)
        out[id(n)] = zeros(Float64, K)
    end

    _logmllh!(n, x, out[id(n)], d)
    return out
end

max_stars_repo_name: crinders/BayesianSumProductNetworks

max_issues_repo_name: crinders/BayesianSumProductNetworks

max_forks_repo_name: crinders/BayesianSumProductNetworks