text:
# compiler driver and main interface

# NOTE: the keyword arguments to compile/codegen control those aspects of compilation that
#       might have to be changed (e.g. set libraries=false when recursing, or set
#       strip=true for reflection). What remains defines the compilation job itself,
#       and those values are contained in the CompilerJob struct.

# (::CompilerJob)
const compile_hook = Ref{Union{Nothing,Function}}(nothing)

"""
    compile(target::Symbol, job::CompilerJob;
            libraries=true, deferred_codegen=true,
            optimize=true, strip=false, ...)

Compile a function `f` invoked with types `tt` for device capability `cap` to one of the
following formats as specified by the `target` argument: `:julia` for Julia IR, `:llvm` for
LLVM IR and `:asm` for machine code.

The following keyword arguments are supported:
- `libraries`: link the GPU runtime and `libdevice` libraries (if required)
- `deferred_codegen`: resolve deferred compiler invocations (if required)
- `optimize`: optimize the code (default: true)
- `strip`: strip non-functional metadata and debug information (default: false)
- `validate`: validate the generated IR before emitting machine code (default: true)
- `only_entry`: only keep the entry function, remove all others (default: false).
  This option is only for internal use, to implement reflection's `dump_module`.

Other keyword arguments can be found in the documentation of [`cufunction`](@ref).
"""
function compile(target::Symbol, @nospecialize(job::CompilerJob);
                 libraries::Bool=true, deferred_codegen::Bool=true,
                 optimize::Bool=true, strip::Bool=false, validate::Bool=true,
                 only_entry::Bool=false)
    if compile_hook[] !== nothing
        compile_hook[](job)
    end

    return codegen(target, job;
                   libraries, deferred_codegen, optimize, strip, validate, only_entry)
end

function codegen(output::Symbol, @nospecialize(job::CompilerJob);
                 libraries::Bool=true, deferred_codegen::Bool=true, optimize::Bool=true,
                 strip::Bool=false, validate::Bool=true, only_entry::Bool=false,
                 parent_job::Union{Nothing, CompilerJob} = nothing)
    ## Julia IR

    mi, mi_meta = emit_julia(job)

    if output == :julia
        return mi, mi_meta
    end

    ## LLVM IR

    ir, ir_meta = emit_llvm(job, mi; libraries, deferred_codegen, optimize, only_entry)

    if output == :llvm
        if strip
            @timeit_debug to "strip debug info" strip_debuginfo!(ir)
        end

        return ir, ir_meta
    end

    ## machine code

    format = if output == :asm
        LLVM.API.LLVMAssemblyFile
    elseif output == :obj
        LLVM.API.LLVMObjectFile
    else
        error("Unknown assembly format $output")
    end
    asm, asm_meta = emit_asm(job, ir; strip, validate, format)

    if output == :asm || output == :obj
        return asm, asm_meta
    end

    error("Unknown compilation output $output")
end

@locked function emit_julia(@nospecialize(job::CompilerJob))
    @timeit_debug to "validation" check_method(job)

    @timeit_debug to "Julia front-end" begin

        # get the method instance
        meth = which(job.source.f, job.source.tt)
        sig = Base.signature_type(job.source.f, job.source.tt)::Type
        (ti, env) = ccall(:jl_type_intersection_with_env, Any,
                          (Any, Any), sig, meth.sig)::Core.SimpleVector
        meth = Base.func_for_method_checked(meth, ti, env)
        method_instance = ccall(:jl_specializations_get_linfo, Ref{Core.MethodInstance},
                      (Any, Any, Any, UInt), meth, ti, env, job.source.world)

        for var in env
            if var isa TypeVar
                throw(KernelError(job, "method captures a typevar (you probably use an unbound type variable)"))
            end
        end
    end

    return method_instance, ()
end

# primitive mechanism for deferred compilation, for implementing CUDA dynamic parallelism.
# this could both be generalized (e.g. supporting actual function calls, instead of
# returning a function pointer), and be integrated with the nonrecursive codegen.
const deferred_codegen_jobs = Dict{Int, Union{FunctionSpec, CompilerJob}}()

# We make this function explicitly callable so that we can drive OrcJIT's
# lazy compilation from, while also enabling recursive compilation.
Base.@ccallable Ptr{Cvoid} function deferred_codegen(ptr::Ptr{Cvoid})
    ptr
end

@generated function deferred_codegen(::Val{f}, ::Val{tt}) where {f,tt}
    id = length(deferred_codegen_jobs) + 1
    deferred_codegen_jobs[id] = FunctionSpec(f,tt)

    pseudo_ptr = reinterpret(Ptr{Cvoid}, id)
    quote
        # TODO: add an edge to this method instance to support method redefinitions
        ccall("extern deferred_codegen", llvmcall, Ptr{Cvoid}, (Ptr{Cvoid},), $pseudo_ptr)
    end
end

const __llvm_initialized = Ref(false)

@locked function emit_llvm(@nospecialize(job::CompilerJob), @nospecialize(method_instance);
                           libraries::Bool=true, deferred_codegen::Bool=true, optimize::Bool=true,
                           only_entry::Bool=false)
    if !__llvm_initialized[]
        InitializeAllTargets()
        InitializeAllTargetInfos()
        InitializeAllAsmPrinters()
        InitializeAllAsmParsers()
        InitializeAllTargetMCs()
        __llvm_initialized[] = true
    end

    @timeit_debug to "IR generation" begin
        ir, compiled = irgen(job, method_instance)
        ctx = context(ir)
        entry_fn = compiled[method_instance].specfunc
        entry = functions(ir)[entry_fn]
    end

    # always preload the runtime, and do so early; it cannot be part of any timing block
    # because it recurses into the compiler
    if libraries
        runtime = load_runtime(job; ctx)
        if haskey(globals(runtime), "llvm.used")
            # the runtime shouldn't link-in stuff that gets preserved in the output. this is
            # a hack to get rid of the device function slots emitted by the PTX back-end,
            # but it also makes sense.
            gv = globals(runtime)["llvm.used"]
            LLVM.unsafe_delete!(runtime, gv)
        end
        runtime_fns = LLVM.name.(defs(runtime))
    end

    @timeit_debug to "Library linking" begin
        if libraries
            # target-specific libraries
            undefined_fns = LLVM.name.(decls(ir))
            @timeit_debug to "target libraries" link_libraries!(job, ir, undefined_fns)

            # GPU run-time library
            if any(fn -> fn in runtime_fns, undefined_fns)
                @timeit_debug to "runtime library" link_library!(ir, runtime)
            end
        end
    end

    # mark everything internal except for the entry and any exported global variables.
    # this makes sure that the optimizer can, e.g., touch function signatures.
    ModulePassManager() do pm
        # NOTE: this needs to happen after linking libraries to remove unused functions,
        #       but before deferred codegen so that all kernels remain available.
        exports = String[entry_fn]
        for gvar in globals(ir)
            if linkage(gvar) == LLVM.API.LLVMExternalLinkage
                push!(exports, LLVM.name(gvar))
            end
        end
        internalize!(pm, exports)
        run!(pm, ir)
    end

    # finalize the current module. this needs to happen before linking deferred modules,
    # since those modules have been finalized themselves, and we don't want to re-finalize.
    entry = finish_module!(job, ir, entry)

    # deferred code generation
    do_deferred_codegen = !only_entry && deferred_codegen &&
                          haskey(functions(ir), "deferred_codegen")
    if do_deferred_codegen
        dyn_marker = functions(ir)["deferred_codegen"]

        cache = Dict{CompilerJob, String}(job => entry_fn)

        # iterative compilation (non-recursive)
        changed = true
        while changed
            changed = false

            # find deferred compiler
            # TODO: recover this information earlier, from the Julia IR
            worklist = Dict{CompilerJob, Vector{LLVM.CallInst}}()
            for use in uses(dyn_marker)
                # decode the call
                call = user(use)::LLVM.CallInst
                id = convert(Int, first(operands(call)))

                global deferred_codegen_jobs
                dyn_job = deferred_codegen_jobs[id]
                if dyn_job isa FunctionSpec
                    dyn_job = similar(job, dyn_job)
                end
                push!(get!(worklist, dyn_job, LLVM.CallInst[]), call)
            end

            # compile and link
            for dyn_job in keys(worklist)
                # cached compilation
                dyn_entry_fn = get!(cache, dyn_job) do
                    dyn_ir, dyn_meta = codegen(:llvm, dyn_job; optimize=false,
                                               deferred_codegen=false, parent_job=job)
                    dyn_entry_fn = LLVM.name(dyn_meta.entry)
                    merge!(compiled, dyn_meta.compiled)
                    @assert context(dyn_ir) == ctx
                    link!(ir, dyn_ir)
                    changed = true
                    dyn_entry_fn
                end
                dyn_entry = functions(ir)[dyn_entry_fn]

                # insert a pointer to the function everywhere the entry is used
                T_ptr = convert(LLVMType, Ptr{Cvoid}; ctx)
                for call in worklist[dyn_job]
                    Builder(ctx) do builder
                        position!(builder, call)
                        fptr = ptrtoint!(builder, dyn_entry, T_ptr)
                        replace_uses!(call, fptr)
                    end
                    unsafe_delete!(LLVM.parent(call), call)
                end
            end
        end

        # all deferred compilations should have been resolved
        @compiler_assert isempty(uses(dyn_marker)) job
        unsafe_delete!(ir, dyn_marker)
    end

    @timeit_debug to "IR post-processing" begin
        if optimize
            @timeit_debug to "optimization" begin
                optimize!(job, ir)

                # deferred codegen has some special optimization requirements,
                # which also need to happen _after_ regular optimization.
                # XXX: make these part of the optimizer pipeline?
                do_deferred_codegen && ModulePassManager() do pm
                    # inline and optimize the call to e deferred code. in particular we want
                    # to remove unnecessary alloca's created by pass-by-ref semantics.
                    instruction_combining!(pm)
                    always_inliner!(pm)
                    scalar_repl_aggregates_ssa!(pm)
                    promote_memory_to_register!(pm)
                    gvn!(pm)

                    # merge duplicate functions, since each compilation invocation emits everything
                    # XXX: ideally we want to avoid emitting these in the first place
                    merge_functions!(pm)

                    run!(pm, ir)
                end
            end

            # optimization may have replaced functions, so look the entry point up again
            entry = functions(ir)[entry_fn]
        end

        @timeit_debug to "clean-up" begin
            # we can only clean-up now, as optimization may lower or introduce calls to
            # functions from the GPU runtime (e.g. julia.gc_alloc_obj -> gpu_gc_pool_alloc)
            ModulePassManager() do pm
                # eliminate all unused internal functions
                global_optimizer!(pm)
                global_dce!(pm)
                strip_dead_prototypes!(pm)

                # merge constants (such as exception messages)
                constant_merge!(pm)

                run!(pm, ir)
            end
        end

        # replace non-entry function definitions with a declaration
        # NOTE: we can't do this before optimization, because the definitions of called
        #       functions may affect optimization.
        if only_entry
            for f in functions(ir)
                f == entry && continue
                isdeclaration(f) && continue
                LLVM.isintrinsic(f) && continue
                empty!(f)
            end
        end

        if ccall(:jl_is_debugbuild, Cint, ()) == 1
            @timeit_debug to "verification" verify(ir)
        end
    end

    return ir, (; entry, compiled)
end

@locked function emit_asm(@nospecialize(job::CompilerJob), ir::LLVM.Module;
                          strip::Bool=false, validate::Bool=true, format::LLVM.API.LLVMCodeGenFileType)
    if validate
        @timeit_debug to "validation" begin
            check_invocation(job)
            check_ir(job, ir)
        end
    end

    # NOTE: strip after validation to get better errors
    if strip
        @timeit_debug to "strip debug info" strip_debuginfo!(ir)
    end

    @timeit_debug to "LLVM back-end" begin
        @timeit_debug to "preparation" prepare_execution!(job, ir)

        code = @timeit_debug to "machine-code generation" mcgen(job, ir, format)
    end

    return code, ()
end


text:
__precompile__(false)

module GraphicsManager

using ModernGL

using CodeGeneration
using RessourceManager
using FileManager
using LogManager
using WindowManager

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

using GLLists
using GLDebug
using GLExtensions
using GLInfo

export @GLCHECK
export glCheckError
export getExtensions

export createcontextinfo
export get_glsl_version_string

export glBindBuffersBase
export glBindBuffersRange

export glGetValue
export glGenOne
export glDelOne
export glGenBuffer
export glDeleteBuffer
export glGenFramebuffer
export glDeleteFramebuffer
export glGenRenderbuffer
export glDeleteRenderbuffer
export glGenVertexArray
export glDeleteVertexArray
export glGenTexture
export glDeleteTexture
export glGenSampler
export glDeleteSampler
export glGetIntegerval
export glArrayFunc

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

CodeGeneration.module_listeners(@__MODULE__)
WindowManager.addListener(@__MODULE__)

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

DRIVER_INITALIZED = false

init() = global DRIVER_INITALIZED = true

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

ListElementType = GLuint
ListType = Dict{Symbol,Array{ListElementType, 1}}
lists = Dict{Symbol,Dict{Symbol,Any}}()
lists[:BUFFER] = Dict(:LIST=>ListType(),:CREATE=>glGenBuffers,:DELETE=>glDeleteBuffers)
lists[:FRAMEBUFFER] = Dict(:LIST=>ListType(),:CREATE=>glGenFramebuffers,:DELETE=>glDeleteFramebuffers)
lists[:RENDERBUFFER] = Dict(:LIST=>ListType(),:CREATE=>glGenRenderbuffers,:DELETE=>glDeleteRenderbuffers)
lists[:VERTEXARRAY] = Dict(:LIST=>ListType(),:CREATE=>glGenVertexArrays,:DELETE=>glDeleteVertexArrays)
lists[:TEXTURE] = Dict(:LIST=>ListType(),:CREATE=>glGenTextures,:DELETE=>glDeleteTextures)
lists[:SAMPLER] = Dict(:LIST=>ListType(),:CREATE=>glGenSamplers,:DELETE=>glDeleteSamplers)
#lists[:PROGRAM] = Dict(:LIST=>GLuint[],:CREATE=>(count,objs)->glCreate(count,objs,glCreateProgram),:DELETE=>(len,ids)->glDelete(ids,glDeleteProgram))
##################################################

"""
TODO
"""
function create(typ::Symbol, id::Symbol, count::Number)
  if count <=0 return nothing end

  l=lists[typ]
  biglist=l[:LIST]

  #objs=zeros(GLuint,count)
  #l[:CREATE](count, objs)
  #for obj in objs push!(l[:LIST], obj) end

  if !haskey(biglist, id) biglist[id] = ListElementType[] end
  list = biglist[id]

  if count > length(list)
    objs=zeros(ListElementType,count)
    glArrayFunc(objs, l[:CREATE]) #(count, objs)
    for obj in objs push!(list, obj) end
  else
    objs = ListElementType[]
    index=0
    checkList = Dict{ListElementType,Nothing}()
    for obj in list
      if haskey(checkList, obj) warn("$typ : $id has duplicated id = $obj")
      elseif obj <= 0  warn("$typ : $id has wrong id <= 0")
      else
        checkList[obj] = nothing
        push!(objs, obj)
        index+=1
      end
      #end
      if index>=count break end
    end
  end

  objs
end

"""
TODO
"""
function delete(typ::Symbol, ids::Array{ListElementType,1})
  error("Not implemented yet!")
  #global lists; l=lists[typ]
  #l[:DELETE](length(ids), ids)
  #l[:LIST] = filter!(e->e∉ids,l[:LIST])
end

"""
TODO
"""
create(typ::Symbol, id::Symbol) = create(typ, id, 1)[1]

"""
TODO
"""
delete(typ::Symbol, id::ListElementType) = delete(typ, ListElementType[id])

"""
TODO
"""
function delete(typ::Symbol;freememory=true)
  l=lists[typ]
  rmlist=ListElementType[]
  for (_, list) in l[:LIST] rmlist = vcat(rmlist, list...) end
  if freememory glArrayFunc(rmlist, l[:DELETE]) end
  l[:LIST] = typeof(l[:LIST])()
end

"""
TODO
"""
cleanLists(hascontext=true) = for typ in keys(lists) delete(typ;freememory=hascontext) end

"""
TODO
"""
freeMemory() = cleanLists()

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

"""
TODO
"""
function cleanUp()
  #println("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
  global DRIVER_INITALIZED
  if DRIVER_INITALIZED == false return end
  callOnListeners(:cleanUp)
  cleanLists(false) #has no context, because window is already closed
  global DRIVER_INITALIZED = false
end

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

BackupSource=Dict{Symbol,String}()

"""
TODO
"""
function createShader(name::Symbol, infos::Dict{Symbol,Any})
  pname = stringColor(name;color=:yellow)
  key = Symbol(infos[:path])
  file = infos[:file]
  typ = infos[:type]
  source = infos[:content]
  err=false
  source = string(strip(replace(source,"\r"=>"")))

  typname="?"
  if typ == :VSH typ = GL_VERTEX_SHADER; typname="VERTEX SHADER"
  elseif typ == :FSH typ = GL_FRAGMENT_SHADER; typname="FRAGMENT SHADER"
  elseif typ == :GSH typ = GL_GEOMETRY_SHADER; typname="GEOMETRY SHADER"
  elseif typ == :CSH typ = GL_COMPUTE_SHADER; typname="COMPUTE SHADER"
  end

  # Create the shader
  shader = @GLCHECK glCreateShader(typ)::GLuint
  if shader == 0 LogManager.error("[$pname] Error creating $typname: ", glErrorMessage()); return -1 end

  shdir=RessourceManager.getPath(:SHADERS)
  tmpdir=joinpath(shdir,"tmp/") #joinpath(@__DIR__,"../shaders/tmp/")
  backupdir=joinpath(shdir,"backup/") #joinpath(@__DIR__,"../shaders/backup/")

  if !isdir(tmpdir) mkdir(tmpdir) end
  if !isdir(backupdir) mkdir(backupdir) end

  tmpfile=abspath(tmpdir*file)
  backupfile=abspath(backupdir*file)

  # Compile the shader
  #try
  # load previous backup
  if !haskey(BackupSource, key) && isfile(backupfile)
    BackupSource[key]=fileGetContents(backupfile)
  end

  open(tmpfile, "w") do f write(f,source) end
  result=compileShader(shader, source)

  #binaryLength=GLint[0]
  #glGetShaderiv(shader, GL_SHADER_SOURCE_LENGTH, binaryLength)
  #num_formats = GLint[0]
  #glGetIntegerv(GL_NUM_SHADER_BINARY_FORMATS,num_formats)
  #formats=zeros(GLint,num_formats[])
  #glGetIntegerv(GL_SHADER_BINARY_FORMATS,formats)
  #binary = zeros(GLbyte,binaryLength)
  #glShaderBinary(1, [shader], GL_SHADER_BINARY_FORMAT_SPIR_V_ARB, binary, binaryLength[])

  if !result
    LogManager.warn("[$pname] $typname ($file) compile error: ", getInfoLog(shader))

    # get backup
    if haskey(BackupSource, key)
      result=compileShader(shader, BackupSource[key])
       if !result LogManager.error("[$pname] Backup of $typname ($file) compile error: ", getInfoLog(shader)); return -1 end #lost backup?
      LogManager.debug("[$pname] Backup of $typname ($file) is compiled.")
    else
      glDeleteShader(shader)
      LogManager.error("[$pname] No valid backup for $typname.")
      return -1
    end
  else
    # save backup
    mv(tmpfile, backupfile; force=true) #move into backups
    BackupSource[key]=source
    LogManager.debug("[$pname] $typname ($file) is compiled.")
  end
  #catch(ex)
  #  Base.showerror(stderr, ex, catch_backtrace())
  #  err=true
  #end

  shader
end

export createShader

"""
TODO
"""
function createShaderProgram(shaders::AbstractArray; transformfeedback=false, name::Union{Nothing,Symbol}=nothing)
  # Create, link then return a shader program for the given shaders.
  # Create the shader program
  prog = @GLCHECK glCreateProgram()

  if name == nothing name = symbol("("*string(prog)*")") end
  pname = stringColor(name;color=:yellow)

  if prog == 0 LogManager.error("Error creating shader program $pname: ", glErrorMessage()) end
  LogManager.debug("Created shader program $pname.")

  # attributes
  #glBindAttribLocation(prog,0,"iVertex") # bind attribute always

  shaderIDs = Int32[]

  for shader in shaders
    shaderID = createShader(name, shader)
    if shaderID > 0
      @GLCHECK glAttachShader(prog, shaderID)
      push!(shaderIDs, shaderID)
    end
  end

  if length(shaderIDs) == 0
    glDeleteProgram(prog)
    prog = -1
    LogManager.error("No valid shaders for shader program $pname found.")
  else

    #binaryLength=GLint[0]
    #binaryFormat=GLint[0]
    #binary=zeros(GLbyte,binaryLength)
    #glGetProgramiv(prog, GL_PROGRAM_BINARY_LENGTH, binaryLength)
    #glGetProgramBinary(prog, binaryLength[], C_NULL, binaryFormat, binary)
    #glProgramBinary( prog, binaryFormat, binary, binaryLength ) # upload
    #success = GLint[0]; glGetProgramiv( prog, GL_LINK_STATUS, success )

    if transformfeedback
      #Ptr{Ptr{GLchar}}
      r = [
        convert(Ptr{GLchar}, pointer("iInstancePos")),#pointer(collect("iInstancePos\x00"))
        convert(Ptr{GLchar}, pointer("iInstanceFlags")) #pointer(collect("iInstanceFlags\x00"))
      ]
      glTransformFeedbackVaryings(prog, 2, r, GL_INTERLEAVED_ATTRIBS)
    end

    # link the program and check for errors
    glLinkProgram(prog)

    status = GLint[0]
    glGetProgramiv(prog, GL_LINK_STATUS, status)

    (id->@GLCHECK glDeleteShader(id)).(shaderIDs) # remove shaders

    if status[] == GL_FALSE
      msg = getInfoLog(prog)
      glDeleteProgram(prog)
      LogManager.error("Shader Program($prog) $pname: Error Linking: $msg")
      prog=-1
    end

    LogManager.debug("Shader Program($prog) $pname: Initalized.")
  end

  prog
end
export createShaderProgram

end #GraphicsManager

ext: jl

max_stars_repo_name: Gilga/JuliaOpenGL

max_issues_repo_name: Gilga/JuliaOpenGL

max_forks_repo_name: Gilga/JuliaOpenGL


text:
@testset "BFloat16sr seeding" begin
    StochasticRounding.seed(12312311)
    O = ones(BFloat16sr,10000)
    one_third = O/3

    # now reseed
    StochasticRounding.seed(12312311)
    @test O/3 == one_third
end

@testset "Float16sr seeding" begin
    StochasticRounding.seed(12312311)
    O = ones(Float16sr,10000)
    one_third = O/3

    # now reseed
    StochasticRounding.seed(12312311)
    @test O/3 == one_third
end

@testset "Float32sr seeding" begin
    StochasticRounding.seed(12312311)
    O = ones(Float32sr,10000)
    one_third = O/3

    # now reseed
    StochasticRounding.seed(12312311)
    @test O/3 == one_third
end


text:
module CMBLensing

using Adapt
using Base.Broadcast: AbstractArrayStyle, ArrayStyle, Broadcasted,
    DefaultArrayStyle, preprocess_args, Style, result_style, Unknown
using Base.Iterators: flatten, product, repeated, cycle, countfrom, peel, partition
using Base.Threads
using Base: @kwdef, @propagate_inbounds, Bottom, OneTo, showarg, show_datatype,
    show_default, show_vector, typed_vcat, typename
using Combinatorics
using DataStructures
using DelimitedFiles
using Distributed: pmap, nworkers, myid, workers, addprocs, @everywhere, remotecall_wait,
    @spawnat, pgenerate, procs, @fetchfrom, default_worker_pool
using FileIO
using FFTW
using InteractiveUtils
using IterTools: flagfirst
using JLD2
using JLD2: jldopen, JLDWriteSession
using KahanSummation
using Loess
using LinearAlgebra
using LinearAlgebra: diagzero, matprod, promote_op
using MacroTools: @capture, combinedef, isdef, isexpr, postwalk, prewalk, rmlines, splitdef
using Match
using Markdown
using Measurements
using Memoization
using NamedTupleTools: select, delete
using OptimKit
using Pkg
using Printf
using ProgressMeter
using QuadGK
using Random
using Random: seed!, AbstractRNG
using Roots
using Requires
using Serialization
using Setfield
using SparseArrays
using StaticArrays: @SMatrix, @SVector, SMatrix, StaticArray, StaticArrayStyle,
    StaticMatrix, StaticVector, SVector, SArray, SizedArray
using Statistics
using StatsBase
using TimerOutputs: @timeit, get_defaulttimer, reset_timer!
using UnPack
using Zygote
using Zygote: unbroadcast, Numeric, @adjoint, @nograd

import Adapt: adapt_structure
import Base: +, -, *, \, /, ^, ~, ≈, <, <=, |, &, ==, !,
    abs, adjoint, all, any, axes, broadcast, broadcastable, BroadcastStyle, conj, copy, convert,
    copy, copyto!, eltype, eps, fill!, getindex, getproperty, hash, hcat, hvcat, inv, isfinite,
    iterate, keys, lastindex, length, literal_pow, mapreduce, materialize!,
    materialize, one, permutedims, print_array, promote, promote_rule,
    promote_rule, promote_type, propertynames, real, setindex!, setproperty!, show,
    show_datatype, show_vector, similar, size, sqrt, string, sum, summary,
    transpose, zero
import Base.Broadcast: materialize, preprocess, broadcasted
import LinearAlgebra: checksquare, diag, dot, isnan, ldiv!, logdet, mul!, norm,
    pinv, StructuredMatrixStyle, structured_broadcast_alloc, tr
import Measurements: ±
import Statistics: std

export
    @⌛, @show⌛, @ismain, @namedtuple, @repeated, @unpack, @cpu!, animate,
    argmaxf_lnP, BandPassOp, BaseDataSet, batch, batch_index, batch_length, beamCℓs, cache,
    CachedLenseFlow, camb, cov_to_Cℓ, cpu, Cℓ_2D, Cℓ_to_Cov, DataSet, DerivBasis,
    diag, Diagonal, DiagOp, dot, EBFourier, EBMap, expnorm, Field, FieldArray, fieldinfo,
    FieldMatrix, FieldOrOpArray, FieldOrOpMatrix, FieldOrOpRowVector,
    FieldOrOpVector, FieldRowVector, FieldTuple, FieldVector, FieldVector,
    firsthalf, fixed_white_noise, FlatEB, FlatEBFourier, FlatEBMap, FlatField,
    FlatFourier, FlatIEBCov, FlatIEBFourier, FlatIEBMap,
    FlatIQUFourier, FlatIQUMap, FlatMap, FlatQU, FlatQUFourier, FlatQUMap, FlatS0,
    FlatS02, FlatS2, FlatS2Fourier, FlatS2Map, Fourier, FuncOp, get_Cℓ,
    get_Cℓ, get_Dℓ, get_αℓⁿCℓ, get_ρℓ, get_ℓ⁴Cℓ, gradhess, gradient, HighPass,
    IEBFourier, IEBMap, InterpolatedCℓs, IQUFourier, IQUMap, kde,
    lasthalf, LazyBinaryOp, LenseBasis, LenseFlow, FieldOp, lnP, load_camb_Cℓs,
    load_chains, load_sim, LowPass, make_mask, Map, MAP_joint, MAP_marg,
    mean_std_and_errors, MidPass, mix, nan2zero, noiseCℓs,
    ParamDependentOp, pixwin, PowerLens, ProjLambert, QUFourier, QUMap, resimulate!,
    resimulate, RK4Solver, sample_joint, shiftℓ,
    simulate, SymmetricFuncOp, symplectic_integrate, Taylens, toCℓ, toDℓ,
    ud_grade, unbatch, unmix, white_noise, Ð, Ł,
    ℓ², ℓ⁴, ∇, ∇², ∇ᵢ, ∇ⁱ

# bunch of sampling-related exports
export gibbs_initialize_f!, gibbs_initialize_ϕ!, gibbs_initialize_θ!,
    gibbs_sample_f!, gibbs_sample_ϕ!, gibbs_sample_slice_θ!,
    gibbs_mix!, gibbs_unmix!, gibbs_postprocess!,
    once_every, start_after_burnin, mass_matrix_ϕ, hmc_step

# generic stuff
include("util.jl")
include("util_fft.jl")
include("util_parallel.jl")
include("numerical_algorithms.jl")
include("generic.jl")
include("cls.jl")
include("field_tuples.jl")
include("field_vectors.jl")
include("specialops.jl")
include("flowops.jl")
include("batched_reals.jl")
include("base_fields.jl")

# # lensing
include("lenseflow.jl")
include("powerlens.jl")

# # flat-sky maps
include("flat_proj.jl")
include("flat_fields.jl")
include("masking.jl")
include("taylens.jl")
include("bilinearlens.jl")

# plotting
function animate end
@init @require PyPlot="d330b81b-6aea-500a-939a-2ce795aea3ee" include("plotting.jl")

# # sampling and maximizing the posteriors
include("dataset.jl")
include("posterior.jl")
include("maximization.jl")
include("sampling.jl")
include("chains.jl")

# other estimates
include("quadratic_estimate.jl")

# curved-sky
include("healpix.jl")

include("autodiff.jl")

# gpu
is_gpu_backed(x) = false
@init @require CUDA="052768ef-5323-5732-b1bb-66c8b64840ba" include("gpu.jl")

# misc init
# see https://github.com/timholy/ProgressMeter.jl/issues/71 and links therein
@init if ProgressMeter.@isdefined ijulia_behavior
    ProgressMeter.ijulia_behavior(:clear)
end

end


text:
function check_registry()
    basedir = joinpath(dirname(pathof(MLJModels)), "registry")
    Pkg.activate(basedir)

    # Read Metadata.toml
    dict = TOML.parsefile(joinpath(basedir, "Metadata.toml"))

    # There will be warnings for ambiguous things, ignore them
    for (package, model_dict) in dict
        for (model, meta) in model_dict
            # check if new entry or changed entry, otherwise don't test
            key = "$package.$model"
            try
                load(model; pkg=package, allow_ambiguous=true)
                # add/refresh entry
                print(rpad("Entry for $key was loaded properly ✓", 79)*"\r")
            catch
                println("⚠ there was an issue trying to load $key")
            end
        end
    end
end


text:
using ClimateEasy
using Test

@testset "ClimateEasy.jl" begin
    # Write your own tests here.
end


text:
using LinearAlgebra

import Base: iszero

pcTol(::Type{T}) where {T <: Integer}       = zero(T)
pcTol(::Type{T}) where {T <: Rational}      = zero(T)
pcTol(::Type{T}) where {T <: Float32}       = T(1f-3)
pcTol(::Type{T}) where {T <: Float64}       = T(1e-6)

iszero(n::T, tol::T=pcTol(T)) where {T <: Number} = -tol <= n <= tol

const notvoid = Base.notnothing
const _nv = notvoid

"""
```
    parallelogram_area(m::Matrix) -> Number
```
Area of the parallelogram. The matrix is a 2x3 matrix.
"""
function parallelogram_area(p1::Tuple{T, T},
                            p2::Tuple{T, T},
                            p3::Tuple{T, T}) where T <: Number
    v1 = (p2[1]-p3[1], p2[2]-p3[2])
    v2 = (p3[1]-p1[1], p3[2]-p1[2])
    v3 = (p1[1]-p2[1], p1[2]-p2[2])
    v = max(dot(v1, v1), dot(v2, v2), dot(v3, v3))

    d = v3[1]*v1[2] - v1[1]*v3[2]

    v*pcTol(T)*pcTol(T) >= d*d && return zero(T)
    return d
end


text:
# FIXME: type stability broken. The following should NOT error
A = @inferred(AxisArray(reshape(1:24, 2,3,4), .1:.1:.2, .1:.1:.3, .1:.1:.4))
@test_throws ArgumentError AxisArray(reshape(1:24, 2,3,4), .1:.1:.1, .1:.1:.3, .1:.1:.4)
@test_throws ArgumentError AxisArray(reshape(1:24, 2,3,4), .1:.1:.1, .1:.1:.3)
@test_throws ArgumentError AxisArray(reshape(1:24, 2,3,4), .1:.1:.2, .1:.1:.3, .1:.1:.4, 1:1)
@test parent(A) === reshape(1:24, 2,3,4)
# Test iteration
for (a,b) in zip(A, A.data)
    @test a == b
end
# Cartesian indexing
for idx in eachindex(A)
    @test A[idx] == A.data[idx]
end
# Conversion and similar
@test Array(A) == A.data
@test reshape(A, length(A)) == A.data[:]
B = similar(A, Float64)
for i in eachindex(A)
    B[i] = A[i]
end
@test A == B
for i=1:length(A)
    @test float(A[i]) === B[i]
end
C = similar(A, 0)
@test isa(C, Array{Int,1})
@test C == []
D = similar(A)
@test size(A) == size(D)
@test eltype(A) == eltype(D)

# permutedims and transpose
@test axisnames(permutedims(A, (2,1,3))) == (:col, :row, :page)
@test axisnames(permutedims(A, (2,3,1))) == (:col, :page, :row)
@test axisnames(permutedims(A, (3,2,1))) == (:page, :col, :row)
@test axisnames(permutedims(A, (3,1,2))) == (:page, :row, :col)
for perm in ((:col, :row, :page), (:col, :page, :row),
             (:page, :col, :row), (:page, :row, :col),
             (:row, :page, :col), (:row, :col, :page))
    @test axisnames(permutedims(A, perm)) == perm
end
@test axisnames(permutedims(A, (:col,)))  == (:col, :row, :page)
@test axisnames(permutedims(A, (:page,))) == (:page, :row, :col)
A2 = @inferred(AxisArray(reshape(1:15, 3, 5)))
A1 = AxisArray(1:5, :t)
for f in (transpose, adjoint)
    @test f(A2).data == f(A2.data)
    @test axisnames(f(A2)) == (:col, :row)
    @test f(A1).data == f(A1.data)
    @test axisnames(f(A1)) == (:transpose, :t)
end

# Test modifying a particular axis
E = similar(A, Float64, Axis{:col}(1:2))
@test size(E) == (2,2,4)
@test eltype(E) == Float64
F = similar(A, Axis{:row}())
@test size(F) == size(A)[2:end]
@test eltype(F) == eltype(A)
@test axisvalues(F) == axisvalues(A)[2:end]
@test axisnames(F) == axisnames(A)[2:end]
G = similar(A, Float64)
@test size(G) == size(A)
@test eltype(G) == Float64
@test axisvalues(A) == axisvalues(G)
@test axisnames(A) == axisnames(G)
H = similar(A, 1,1,1)
@test size(H) == (1,1,1)
@test eltype(H) == eltype(A)
@test typeof(H) <: Array
H = similar(A, Float64, 1,1,1)
@test size(H) == (1,1,1)
@test eltype(H) == Float64
@test typeof(H) <: Array

# Size
@test size(A, 1) == size(A, Axis{1}) == size(A, Axis{:row}) == size(A, Axis{:row}())

## Test constructors
# No axis or time args
A = AxisArray(1:3)
@test A.data == 1:3
@test axisnames(A) == (:row,)
@inferred(axisnames(A))
@test axisvalues(A) == (1:3,)
A = AxisArray(reshape(1:16, 2,2,2,2))
@test A.data == reshape(1:16, 2,2,2,2)
@test axisnames(A) == (:row,:col,:page,:dim_4)
@inferred(axisnames(A))
@test axisvalues(A) == (1:2, 1:2, 1:2, 1:2)
# Just axis names
A = AxisArray(1:3, :a)
@test A.data == 1:3
@test axisnames(A) == (:a,)
@inferred(axisnames(A))
@test axisvalues(A) == (1:3,)
A = AxisArray([1 3; 2 4], :a)
@test A.data == [1 3; 2 4]
@test axisnames(A) == (:a, :col)
@inferred(axisnames(A))
@test axisvalues(A) == (1:2, 1:2)
# Just axis values
A = @inferred(AxisArray(1:3, .1:.1:.3))
@test A.data == 1:3
@test axisnames(A) == (:row,)
@inferred(axisnames(A))
@test axisvalues(A) == (.1:.1:.3,)
# FIXME: reintroduce inferred
A = @inferred(AxisArray(reshape(1:16, 2,2,2,2), .5:.5:1))
@test A.data == reshape(1:16, 2,2,2,2)
@test axisnames(A) == (:row,:col,:page,:dim_4)
@inferred(axisnames(A))
@test axisvalues(A) == (.5:.5:1, 1:2, 1:2, 1:2)
A = AxisArray([0]', :x, :y)
@test axisnames(dropdims(A, dims=1)) == (:y,)
@test axisnames(dropdims(A, dims=2)) == (:x,)
@test axisnames(dropdims(A, dims=(1,2))) == axisnames(dropdims(A, dims=(2,1))) == ()
@test axisnames((dropdims(A, dims=Axis{:x}))) == (:y,)
@test axisnames((dropdims(A, dims=Axis{:x,UnitRange{Int}}))) == (:y,)
@test axisnames((dropdims(A, dims=Axis{:y}))) == (:x,)
@test axisnames((dropdims(dropdims(A, dims=Axis{:x}), dims=Axis{:y}))) == ()
@test_broken @inferred(dropdims(A, dims=Axis{:x}))
@test_broken @inferred(dropdims(A, dims=Axis{:x,UnitRange{Int}}))
@test_broken @inferred(dropdims(A, dims=Axis{:y}))
@test_broken @inferred(dropdims(dropdims(A, dims=Axis{:x}), dims=Axis{:y}))
# Names, steps, and offsets
B = AxisArray([1 4; 2 5; 3 6], (:x, :y), (0.2, 100))
@test axisnames(B) == (:x, :y)
@test axisvalues(B) == (0:0.2:0.4, 0:100:100)
B = AxisArray([1 4; 2 5; 3 6], (:x, :y), (0.2, 100), (-3,14))
@test axisnames(B) == (:x, :y)
@test axisvalues(B) == (-3:0.2:-2.6, 14:100:114)

@test AxisArrays.HasAxes(A)   == AxisArrays.HasAxes{true}()
@test AxisArrays.HasAxes([1]) == AxisArrays.HasAxes{false}()

@test_throws ArgumentError AxisArray(reshape(1:24, 2,3,4),
                                     Axis{1}(.1:.1:.2),
                                     Axis{2}(1//10:1//10:3//10),
                                     Axis{3}(["a", "b", "c", "d"])) # Axis need to be symbols
@test_throws ArgumentError AxisArray(reshape(1:24, 2,3,4),
                                     Axis{:x}(.1:.1:.2),
                                     Axis{:y}(1//10:1//10:3//10),
                                     Axis{:z}(["a", "b", "c", "d"]),
                                     Axis{:_}(1:1)) # Too many Axes

A = @inferred(AxisArray(reshape(1:24, 2,3,4),
              Axis{:x}(.1:.1:.2),
              Axis{:y}(1//10:1//10:3//10),
              Axis{:z}(["a", "b", "c", "d"])))

# recursive constructor
@test A === @inferred AxisArray(A)
@test axisnames(AxisArray(A, Axis{:yoyo}(1:length(A[Axis{:x}])))) == (:yoyo, :y, :z)
@test AxisArray(A, Axis{:yoyo}(1:length(A[Axis{:x}]))).data === A.data
@test AxisArray(A, (Axis{:yoyo}(1:length(A[Axis{:x}])),)).data === A.data
@test axisnames(AxisArray(A, :something, :in, :the)) == (:something, :in, :the)
@test AxisArray(A, :way, :you, :move).data === A.data
@test axisnames(AxisArray(A, (:c, :a, :b), (2, 3, 4))) == (:c, :a, :b)
@test AxisArray(A, (:c, :a, :b), (2, 3, 4)).data === A.data
@inferred AxisArray(A, Axis{:yoyo}(1:length(A[Axis{:x}])))
@inferred AxisArray(A, (Axis{:yoyo}(1:length(A[Axis{:x}])),))

# Test axisdim
@test axisdim(A, Axis{:x}) == axisdim(A, Axis{:x}()) == 1
@test axisdim(A, Axis{:y}) == axisdim(A, Axis{:y}()) == 2
@test axisdim(A, Axis{:z}) == axisdim(A, Axis{:z}()) == 3
# Test axes
@test @inferred(AxisArrays.axes(A)) == (Axis{:x}(.1:.1:.2), Axis{:y}(1//10:1//10:3//10), Axis{:z}(["a", "b", "c", "d"]))
@test @inferred(AxisArrays.axes(A, Axis{:x})) == @inferred(AxisArrays.axes(A, Axis{:x}())) == Axis{:x}(.1:.1:.2)
@test @inferred(AxisArrays.axes(A, Axis{:y})) == @inferred(AxisArrays.axes(A, Axis{:y}())) == Axis{:y}(1//10:1//10:3//10)
@test @inferred(AxisArrays.axes(A, Axis{:z})) == @inferred(AxisArrays.axes(A, Axis{:z}())) == Axis{:z}(["a", "b", "c", "d"])
@test AxisArrays.axes(A, 2) == Axis{:y}(1//10:1//10:3//10)
Aplain = rand(2,3)
@test @inferred(AxisArrays.axes(Aplain)) === AxisArrays.axes(AxisArray(Aplain))
@test AxisArrays.axes(Aplain, 1) === AxisArrays.axes(AxisArray(Aplain))[1]
@test AxisArrays.axes(Aplain, 2) === AxisArrays.axes(AxisArray(Aplain))[2]

@test Axis{:col}(1) == Axis{:col}(1)
@test Axis{:col}(1) != Axis{:com}(1)
@test Axis{:x}(1:3) == Axis{:x}(Base.OneTo(3))
@test hash(Axis{:col}(1)) == hash(Axis{:col}(1.0))
@test hash(Axis{:row}()) != hash(Axis{:col}())
@test hash(Axis{:x}(1:3)) == hash(Axis{:x}(Base.OneTo(3)))
@test AxisArrays.axistype(Axis{1}(1:2)) == typeof(1:2)
@test AxisArrays.axistype(Axis{1,UInt32}) == UInt32
@test axisnames(Axis{1}, Axis{2}, Axis{3}) == (1,2,3)
@test Axis{:row}(2:7)[4] == 5
@test eltype(Axis{:row}(1.0:1.0:3.0)) == Float64
@test size(Axis{:row}(2:7)) === (6,)
T = A[AxisArrays.Axis{:x}]
@test T[end] == 0.2
@test Base.axes(Axis{:row}(2:7)) === (Base.OneTo(6),)
@test Base.axes(Axis{:row}(-1:1), 1) === Base.OneTo(3)
@test length(Axis{:col}(-1:2)) === 4
@test AxisArrays.axisname(Axis{:foo}(1:2)) == :foo
@test AxisArrays.axisname(Axis{:foo})      == :foo

# Test Timetype axis construction
dt, vals = DateTime(2010, 1, 2, 3, 40), randn(5,2)
A = @inferred(AxisArray(vals, Axis{:Timestamp}(dt-Dates.Hour(2):Dates.Hour(1):dt+Dates.Hour(2)), Axis{:Cols}([:A, :B])))
@test A[:, :A].data == vals[:, 1]
@test A[dt, :].data == vals[3, :]

@test AxisArrays.axistrait(A.axes[1]) == AxisArrays.Dimensional
@test AxisArrays.axistrait(typeof(A.axes[1])) == AxisArrays.Dimensional
@test AxisArrays.axistrait(A.axes[1].val) == AxisArrays.Dimensional
@test AxisArrays.axistrait(typeof(A.axes[1].val)) == AxisArrays.Dimensional
@test AxisArrays.axistrait(A.axes[2]) == AxisArrays.Categorical
@test AxisArrays.axistrait(typeof(A.axes[2])) == AxisArrays.Categorical
@test AxisArrays.axistrait(A.axes[2].val) == AxisArrays.Categorical
@test AxisArrays.axistrait(typeof(A.axes[2].val)) == AxisArrays.Categorical

@test_throws ArgumentError AxisArrays.checkaxis(Axis{:x}(10:-1:1))
@test_throws ArgumentError AxisArrays.checkaxis(10:-1:1)

# Simply run the display method to ensure no stupid errors
show(IOBuffer(),MIME("text/plain"),A)

# With unconventional indices
import OffsetArrays  # import rather than using because OffsetArrays has a deprecation for ..
A = AxisArray(OffsetArrays.OffsetArray([5,3,4], -1:1), :x)
@test AxisArrays.axes(A) == (Axis{:x}(-1:1),)
@test A[-1] == 5
A[0] = 12
@test A.data[0] == 12
@test Base.axes(A) == Base.axes(A.data)
@test LinearIndices(A) == LinearIndices(A.data)
A = AxisArray(OffsetArrays.OffsetArray(rand(4,5), -1:2, 5:9), :x, :y)
@test Base.axes(A) == Base.axes(A.data)
@test LinearIndices(A) == LinearIndices(A.data)

@test AxisArrays.matchingdims((A, A))

f1(x) = x < 0
A2 = map(f1, A)
@test isa(A2, AxisArray)
@test A2.axes == A.axes
@test A2.data == map(f1, A.data)

map!(~, A2, A2)
@test isa(A2, AxisArray)
@test A2.axes == A.axes
@test A2.data == map(~, map(f1, A).data)

A2 = map(+, A, A)
@test isa(A2, AxisArray)
@test A2.axes == A.axes
@test A2.data == A.data .+ A.data

map!(*, A2, A, A)
@test isa(A2, AxisArray)
@test A2.axes == A.axes
@test A2.data == A.data .* A.data

# Reductions (issue #55)
A = AxisArray(collect(reshape(1:15,3,5)), :y, :x)
B = @inferred(AxisArray(collect(reshape(1:15,3,5)), Axis{:y}(0.1:0.1:0.3), Axis{:x}(10:10:50)))
arrays = (A, B)
functions = (sum, minimum)
for C in arrays
    local C
    for op in functions  # together, cover both reduced_indices and reduced_indices0
        axv = axisvalues(C)
        @test_broken @inferred(op(C; dims=1))
        C1 = op(C; dims=1)
        @test typeof(C1) == typeof(C)
        @test axisnames(C1) == (:y,:x)
        @test axisvalues(C1) === (oftype(axv[1], Base.OneTo(1)), axv[2])
        @test_broken @inferred(op(C, dims=2))
        C2 = op(C, dims=2)
        @test typeof(C2) == typeof(C)
        @test axisnames(C2) == (:y,:x)
        @test axisvalues(C2) === (axv[1], oftype(axv[2], Base.OneTo(1)))
        @test_broken @inferred(op(C, dims=(1,2)))
        C12 = op(C, dims=(1,2))
        @test typeof(C12) == typeof(C)
        @test axisnames(C12) == (:y,:x)
        @test axisvalues(C12) === (oftype(axv[1], Base.OneTo(1)), oftype(axv[2], Base.OneTo(1)))
        if op == sum
            @test C1 == [6 15 24 33 42]
            @test C2 == reshape([35,40,45], 3, 1)
            @test C12 == reshape([120], 1, 1)
        else
            @test C1 == [1 4 7 10 13]
            @test C2 == reshape([1,2,3], 3, 1)
            @test C12 == reshape([1], 1, 1)
        end
        # TODO: add @inferred
        @test (op(C, dims=Axis{:y})) == C1
        @test (op(C, dims=Axis{:x})) == C2
        @test (op(C, dims=(Axis{:y},Axis{:x}))) == C12
        @test (op(C, dims=Axis{:y}())) == C1
        @test (op(C, dims=Axis{:x}())) == C2
        @test (op(C, dims=(Axis{:y}(),Axis{:x}()))) == C12
    end
end

function typeof_noaxis(::AxisArray{T,N,D}) where {T,N,D}
    AxisArray{T,N,D}
end

# uninferrable
C = AxisArray(collect(reshape(1:15,3,5)), Axis{:y}([:a,:b,:c]), Axis{:x}(["a","b","c","d","e"]))
for op in functions  # together, cover both reduced_indices and reduced_indices0
    axv = axisvalues(C)
    C1 = op(C, dims=1)
    @test typeof_noaxis(C1) == typeof_noaxis(C)
    @test axisnames(C1) == (:y,:x)
    @test axisvalues(C1) === (Base.OneTo(1), axv[2])
    C2 = op(C, dims=2)
    @test typeof_noaxis(C2) == typeof_noaxis(C)
    @test axisnames(C2) == (:y,:x)
    @test axisvalues(C2) === (axv[1], Base.OneTo(1))
    C12 = op(C, dims=(1,2))
    @test typeof_noaxis(C12) == typeof_noaxis(C)
    @test axisnames(C12) == (:y,:x)
    @test axisvalues(C12) === (Base.OneTo(1), Base.OneTo(1))
    if op == sum
        @test C1 == [6 15 24 33 42]
        @test C2 == reshape([35,40,45], 3, 1)
        @test C12 == reshape([120], 1, 1)
    else
        @test C1 == [1 4 7 10 13]
        @test C2 == reshape([1,2,3], 3, 1)
        @test C12 == reshape([1], 1, 1)
    end
    # TODO: These should be @inferred, but are currently broken
    @test (op(C, dims=Axis{:y})) == C1
    @test (op(C, dims=Axis{:x})) == C2
    # Unfortunately the type of (Axis{:y},Axis{:x}) is Tuple{UnionAll,UnionAll} so methods will not specialize
    @test (op(C, dims=(Axis{:y},Axis{:x}))) == C12
    @test (op(C, dims=Axis{:y}())) == C1
    @test (op(C, dims=Axis{:x}())) == C2
    @test (op(C, dims=(Axis{:y}(),Axis{:x}()))) == C12
end

C = AxisArray(collect(reshape(1:15,3,5)), Axis{:y}([:a,:b,:c]), Axis{:x}(["a","b","c","d","e"]))
@test occursin(r"axes:\n\s+:y,", summary(C))


text:
ENV["GKSwstype"]="100" #src

# ## 1D1V Vlasov–Ampere system
#md #
#md # ```math
#md # \\frac{\\partial f}{\\partial t} + \\upsilon \\frac{\\partial f}{\\partial x} - E(t,x) \\frac{\\partial f}{\\partial \\upsilon} = 0
#md # ```
#md #
#md # ```math
#md # \\frac{\\partial E}{\\partial t} = - J = \\displaystyle \\int  f\\upsilon \\; d\\upsilon
#md # ```

#md # ---

using ProgressMeter, FFTW, Plots, LinearAlgebra
using BenchmarkTools, Statistics

"""
    UniformMesh(start, stop, length)

1D uniform mesh data for periodic domain (end point is removed)
"""
struct UniformMesh

   start    :: Float64
   stop     :: Float64
   length   :: Int64
   step     :: Float64
   points   :: Vector{Float64}

   function UniformMesh(start, stop, length)

       points = range(start, stop=stop, length=length+1)[1:end-1]
       step = (stop - start) / length

       new( start, stop, length, step, points)

   end

end

#md # ---
#
# ## Compute charge density ρ(x)
#

"""
    compute_rho( mesh, f)

Compute charge density ρ(x,t) = ∫ f(x,v,t) dv

returns ρ - ρ̄

"""
function compute_rho(meshv::UniformMesh, f)

   dv  = meshv.step
   ρ = dv .* vec(sum(real(f), dims=2))
   ρ .- mean(ρ)

end

#md # ---
#
# ## Compute electric field from ρ(x)
#

"""
    compute_e(mesh, ρ)

compute electric field from ρ
"""
function compute_e(mesh::UniformMesh, ρ)

   n = mesh.length
   k =  2π / (mesh.stop - mesh.start)
   modes = [1.0; k .* vcat(1:n÷2-1,-n÷2:-1)...]
   ρ̂ = fft(ρ)./modes
   vec(real(ifft(-1im .* ρ̂)))

end

#
#md # ---
#
# ## Callable struct `Advection`
#
"""
    advection! = AmpereAdvection( mesh, kx)

For every created struct, two methods are available
- Advection method along v
- Advection method along x and e computation

"""
struct AmpereAdvection

    mesh :: UniformMesh
    kx   :: Vector{Float64}

    function AmpereAdvection( mesh )

        nx  = mesh.length
        dx  = mesh.step
        Lx  = mesh.stop - mesh.start
        kx  = zeros(Float64, nx)
        kx .= 2π/Lx .* [0:nx÷2-1;-nx÷2:-1]
        new( mesh, kx)

    end

end

#md # ---

function (adv :: AmpereAdvection)( fᵗ  :: Array{ComplexF64,2},
                                   e   :: Vector{ComplexF64},
                                   dt  :: Float64 )
    fft!(fᵗ, 1)
    fᵗ .= fᵗ .* exp.(-1im * dt * adv.kx * transpose(e))
    ifft!(fᵗ, 1)

end

#md # --

function (adv :: AmpereAdvection)( f   :: Array{ComplexF64,2},
                                   e   :: Vector{ComplexF64},
                                   v   :: Vector{Float64},
                                   dt  :: Float64 )

    ev = exp.(-1im*dt * adv.kx * transpose(v))

    fft!(f,1)
    f .= f .* ev
    dv = v[2]-v[1]
    ρ = dv * vec(sum(f,dims=2))
    for i in 2:length(e)
        e[i] = -1im * ρ[i] / adv.kx[i]
    end
    e[1] = 0.0
    ifft!(f, 1)
    ifft!(e)
    e .= real(e)
end

#
#
#md # ----
#
# ## Initial distribution function
#md #
#md # ```math
#md # f(x,v) = \\frac{1}{\\sqrt{2\\pi}}(1+ ϵ \\cdot cos(k_x x)) e^{-v^2/2}
#md # ```
#md #
"""
    landau( ϵ, kx, x, v )

Landau damping initialisation function

[Wikipedia](https://en.wikipedia.org/wiki/Landau_damping)

"""
function landau( ϵ, kx, x, v )

    (1.0 .+ ϵ*cos.(kx*x))/sqrt(2π) .* transpose(exp.(-0.5*v.*v))

end

#md # ---

nx, nv = 256, 256
xmin, xmax =  0., 4*π
vmin, vmax = -6., 6.
tf = 60
nt = 600
meshx = UniformMesh(xmin, xmax, nx)
meshv = UniformMesh(vmin, vmax, nv);

#md # --

# Initialize distribution function
x = meshx.points
v = meshv.points
ϵ, kx = 0.001, 0.5

# Allocate arrays for distribution function and its transposed
f = zeros(Complex{Float64},(nx,nv))
fᵀ= zeros(Complex{Float64},(nv,nx))

f .= landau( ϵ, kx, x, v)
#md nothing # hide

#md # ---

#nb # Plot the distribution

#nb surface( x, v, real(f))

#nb #-

transpose!(fᵀ,f)

ρ  = compute_rho(meshv, f)
e  = zeros(ComplexF64, nx)
e .= compute_e(meshx, ρ)

mod_e = Float64[]

dt = tf / nt

advection_x! = AmpereAdvection( meshx )
advection_v! = AmpereAdvection( meshv );

#md # ---

@showprogress 1 for i in 1:nt

    advection_v!(fᵀ, e, 0.5dt)

    transpose!(f, fᵀ)

    advection_x!( f, e, v, dt)

    push!(mod_e, log(sqrt((sum(e.^2))*meshx.step))) # diagnostic

    transpose!(fᵀ, f)

    advection_v!(fᵀ, e, 0.5dt)

end

#md # ---

t = range(0, stop=tf, length=nt)
plot(t, -0.1533*t.-5.48)
plot!(t, mod_e , label=:ampere )
#md savefig("mod_e.svg"); nothing # hide

@test true #src

#md # ![](mod_e.svg)

#md # Next: [Rotation on GPU](/Numkin2019/03/build/index.html)


text:
module Solver

using LinearAlgebra
using SparseArrays
using Logging

using HCubature
using Arpack
using DSP

const USE_PARDISO=false
if USE_PARDISO
    import Pardiso
end

using ..FEStructure
using ..LoadCase
using ..FEAssembly

include("./static_solver.jl")
include("./dynamic_solver.jl")
include("../util/mmp.jl")

export solve

function DFSsolve(lc_node::FEAssembly.LCNode,structure::Structure,restrainedDOFs::Vector{Int},path)
    if lc_node.case!="root"
        solve_case(structure,lc_node.case,restrainedDOFs,path)
    end
    if isempty(lc_node.children)
        return
    else
        for child in lc_node.children
            DFSsolve(child,structure,restrainedDOFs,path)
        end
    end
end

function solve_case(structure,loadcase,restrainedDOFs::Vector{Int},path)
    working_dir=joinpath(path,".analysis")
    if is_static(loadcase)
        if loadcase.nl_type=="1st"
            u,F=solve_linear_static(structure,loadcase,restrainedDOFs)
            write_vector(working_dir,loadcase.id*"_u.v",loadcase.hid,u)
            write_vector(working_dir,loadcase.id*"_F.v",loadcase.hid,F)
        elseif loadcase.nl_type=="2nd"
            u,F=solve_2nd_static(structure,loadcase,restrainedDOFs,path=path)
            write_vector(working_dir,loadcase.id*"_u.v",loadcase.hid,u)
            write_vector(working_dir,loadcase.id*"_F.v",loadcase.hid,F)
        elseif loadcase.nl_type=="3rd"

        end
    elseif is_buckling(loadcase)
        println("Based on "*loadcase.plc)
        ω²,ϕ=solve_linear_buckling(structure,loadcase,restrainedDOFs,path=path)
        write_vector(working_dir,loadcase.id*"_o.v",loadcase.hid,ω²)
        write_matrix(working_dir,loadcase.id*"_p.m",loadcase.hid,ϕ)
    elseif is_modal(loadcase)
        if loadcase.modal_type=="eigen"
            ω²,ϕ=solve_modal_eigen(structure,loadcase,restrainedDOFs,path=path)
        elseif loadcase.modal_type=="ritz"
            ω²,ϕ=solve_modal_Ritz(structure,loadcase,restrainedDOFs,path=path)
        else
            throw("Modal method can only be eigen or ritz!")
        end
        write_vector(working_dir,loadcase.id*"_o.v",loadcase.hid,ω²)
        write_matrix(working_dir,loadcase.id*"_p.m",loadcase.hid,ϕ)
    elseif is_time_history(loadcase)
        if loadcase.algorithm=="central_diff"
            u,v,a=solve_central_diff(structure,loadcase,restrainedDOFs,path=path)
            write_matrix(working_dir,loadcase.id*"_u.m",loadcase.hid,u)
            write_matrix(working_dir,loadcase.id*"_v.m",loadcase.hid,v)
            write_matrix(working_dir,loadcase.id*"_a.m",loadcase.hid,a)
        elseif loadcase.algorithm=="newmark"
            u,v,a=solve_newmark_beta(structure,loadcase,restrainedDOFs,path=path)
            write_matrix(working_dir,loadcase.id*"_u.m",loadcase.hid,u)
            write_matrix(working_dir,loadcase.id*"_v.m",loadcase.hid,v)
            write_matrix(working_dir,loadcase.id*"_a.m",loadcase.hid,a)
        elseif loadcase.algorithm=="wilson"
            u,v,a=solve_wilson_theta(structure,loadcase,restrainedDOFs,path=path)
            write_matrix(working_dir,loadcase.id*"_u.m",loadcase.hid,u)
            write_matrix(working_dir,loadcase.id*"_v.m",loadcase.hid,v)
            write_matrix(working_dir,loadcase.id*"_a.m",loadcase.hid,a)
        elseif loadcase.algorithm=="modal_decomp"
            u,v,a=solve_modal_decomposition(structure,loadcase,restrainedDOFs,path=path)
            write_matrix(working_dir,loadcase.id*"_u.m",loadcase.hid,u)
            write_matrix(working_dir,loadcase.id*"_v.m",loadcase.hid,v)
            write_matrix(working_dir,loadcase.id*"_a.m",loadcase.hid,a)
        elseif loadcase.algorithm=="HHT"
            u,v,a=solve_HHT_alpha(structure,loadcase,restrainedDOFs,path=path)
            write_matrix(working_dir,loadcase.id*"_u.m",loadcase.hid,u)
            write_matrix(working_dir,loadcase.id*"_v.m",loadcase.hid,v)
            write_matrix(working_dir,loadcase.id*"_a.m",loadcase.hid,a)
        end
    elseif is_response_spectrum(loadcase)
    end
end

"""
    solve(assembly,run="all",path=pwd())
求解assembly实例
# 参数
- `assembly::Assembly`: Assembly类型实例
- `run`: 运行的工况集合，默认为全部all
- `path`: 工作路径
# 返回
- `assembly`: Assembly对象
"""
function solve(assembly,run="all")
    @info "STRUCTURE INFO" Total_DOF=assembly.nDOF Free_DOF=assembly.nfreeDOF N_Nodes=assembly.node_count N_Beams=assembly.beam_count N_Quads=assembly.quad_count N_Trias=assembly.tria_count
    path=assembly.working_path
    mkpath(joinpath(path,".analysis"))
    structure=assembly.structure
    restrainedDOFs=assembly.restrainedDOFs
    DFSsolve(assembly.lc_tree,structure,restrainedDOFs,path)
end

end


text:
using Test
using LatticeTools

@testset "linpath" begin
    @test_throws DimensionMismatch linpath([0.0, 1.0], [0.0, 1.0, 2.0])
    @test_throws ArgumentError linpath([0.0, 1.0])

    r = linpath([0.0, 0.0], [1.0, 0.0], [1.0, 1.0])
    @test size(r) == (2, 100*2+1)
    @test r[1,:] ≈ vcat(0:0.01:0.99, [1.0 for i in 1:100], [1.0])
    @test r[2,:] ≈ vcat([0.0 for i in 1:100], 0:0.01:0.99, [1.0])

    r = linpath([0.0, 0.0], [1.0, 0.0], [1.0, 1.0]; nseg=10)
    @test size(r) == (2, 10*2+1)
    @test r[1,:] ≈ vcat(0:0.1:0.9, [1.0 for i in 1:10], [1.0])
    @test r[2,:] ≈ vcat([0.0 for i in 1:10], 0:0.1:0.9, [1.0])
end


text:
mutable struct InfPlcNfAbsNS
  field::NfAbsNS
  index::Vector{Int}
  absolute_index::Int
  isreal::Bool
  roots::Vector{acb}
end

absolute_index(P::InfPlcNfAbsNS) = P.absolute_index

number_field(P::InfPlcNfAbsNS) = P.field

place_type(::Type{NfAbsNS}) = InfPlcNfAbsNS

place_type(::NfAbsNS) = InfPlcNfAbsNS

isreal(P::InfPlcNfAbsNS) = P.isreal

function real_places(K::NfAbsNS)
  r, s = signature(K)
  return infinite_places(K)[1:r]
end

function infinite_places(K::NfAbsNS)
  c = conjugate_data_arb_roots(K, 32, copy = false)

  res = InfPlcNfAbsNS[]

  l = ngens(K)

  j = 1

  for v in c[2]
    push!(res, InfPlcNfAbsNS(K, v, j, true, acb[c[1][i].roots[v[i]] for i in 1:l]))
    j += 1
  end

  for v in c[3]
    push!(res, InfPlcNfAbsNS(K, v, j, false, acb[c[1][i].roots[v[i]] for i in 1:l]))
    j += 1
  end

  return res
end

function Base.show(io::IO, P::InfPlcNfAbsNS)
  print(io, "Infinite place of\n")
  print(IOContext(io, :compact => true), P.field)
  print(io, "Corresponding to roots\n")
  print(io, P.roots)
end

function conjugates_data_roots(K::NfAbsNS)
  cache = get_special(K, :conjugates_data_roots)
  if cache !== nothing
    return cache
  end
  pols = fmpq_poly[isunivariate(x)[2] for x in K.pol]
  ctxs = acb_root_ctx[acb_root_ctx(x) for x in pols]
  set_special(K, :conjugates_data_roots => ctxs)
  return ctxs
end

function conjugate_data_arb_roots(K::NfAbsNS, p::Int; copy = true)

  cache = get_special(K, :conjugates_data)
  if cache !== nothing
    if haskey(cache, p)
      return cache[p]
    end
  end
  ctxs = conjugates_data_roots(K)
  acb_roots_vec = Vector{acb_roots}(undef, length(ctxs))
  for i = 1:length(ctxs)
    c = ctxs[i]
    while c.prec < p
      refine(c)
    end
    real_roots = deepcopy(c.real_roots)
    complex_roots = deepcopy(c.complex_roots)
    for z in real_roots
      expand!(z, -p)
    end
    for z in complex_roots
      expand!(z, -p)
    end
    CC = parent(c.roots[1])
    all_roots = Vector{acb}(undef, length(c.roots))
    for i = 1:length(real_roots)
      all_roots[i] = CC(real_roots[i])
    end
    for i = 1:length(complex_roots)
      all_roots[i+length(real_roots)] = complex_roots[i]
      all_roots[i+length(real_roots)+length(complex_roots)] = conj(complex_roots[i])
    end
    acb_roots_vec[i] = acb_roots(p, all_roots, real_roots, complex_roots)
  end
  ind_real, ind_complex = enumerate_conj_prim(acb_roots_vec)
  set_special(K, :conjugates_data => Dict(p => (acb_roots_vec, ind_real, ind_complex)))
  return acb_roots_vec, ind_real, ind_complex

end

function enumerate_conj_prim(v::Vector{acb_roots})
  indices = collect(cartesian_product_iterator([1:length(v[i].roots) for i in 1:length(v)], inplace = false))
  #I have the indices, now I need to order them.
  complex_indices = Int[]
  for i = 1:length(v)
    indc = length(v[i].real_roots)+1
    push!(complex_indices, indc)
  end
  real_combinations = Int[]
  for i = 1:length(indices)
    isreal = true
    for j = 1:length(indices[i])
      if indices[i][j] >= complex_indices[j]
        isreal = false
        break
      end
    end
    if isreal
      push!(real_combinations, i)
    end
  end
  res_real = indices[real_combinations]
  res_complex = typeof(indices)()
  for i = 1:length(indices)
    if i in real_combinations
      continue
    end
    s = indices[i]
    ind_complex = Int[]
    for t = 1:length(s)
      if s[t] >= complex_indices[t]
        push!(ind_complex, t)
      end
    end
    found = false
    for t = 1:length(res_complex)
      found = _is_complex_conj(res_complex[t], s, ind_complex, v)
      if found
        break
      end
    end
    if found
      continue
    end

    push!(res_complex, indices[i])
  end
  return res_real, res_complex
end

function _is_complex_conj(v::Vector, w::Vector, pos::Vector, roots::Vector)
  i = 1
  for x in v
    if i in pos
      if v[i] <= length(roots[i].real_roots)
        return false
      end
      lc = length(roots[i].complex_roots)
      if v[i] != w[i] + lc && v[i] != w[i] - lc
        return false
      end
    elseif v[i] != w[i]
      return false
    end
    i += 1
  end
  return true
end

function conjugates_arb(a::NfAbsNSElem, p::Int, work_tol::Int = p)

  K = parent(a)
  conjs, ind_real, ind_complex = conjugate_data_arb_roots(K, work_tol)
  res = Vector{acb}(undef, degree(K))
  pol_a = data(a)
  r, s = signature(K)
  for i = 1:r
    res[i] = _evaluate(pol_a, acb[conjs[j].roots[ind_real[i][j]] for j = 1:ngens(K)])
    if !isfinite(res[i]) || !radiuslttwopower(res[i], -p)
      return conjugates_arb(a, p, 2*work_tol)
    end
  end
  ind = r+1
  for i = 1:length(ind_complex)
    ev = acb[conjs[j].roots[ind_complex[i][j]] for j = 1:ngens(K)]
    res[ind] = _evaluate(pol_a, ev)
    if !isfinite(res[ind]) || !radiuslttwopower(res[ind], -p)
      return conjugates_arb(a, p, 2*work_tol)
    end
    res[ind+s] = conj(res[ind])
    ind += 1
  end
  return res
end

function _evaluate(f::fmpq_mpoly, vals::Vector{acb})
  S = parent(vals[1])
  powers = [Dict{Int, acb}() for i in 1:length(vals)]
  r = acb[zero(S)]
  i = UInt(1)
  cvzip = zip(coeffs(f), exponent_vectors(f))
  for (c, v) in cvzip
    t = one(S)
    for j = 1:length(vals)
      exp = v[j]
      if iszero(exp)
        continue
      end
      if !haskey(powers[j], exp)
        powers[j][exp] = vals[j]^exp
      end
      mul!(t, t, powers[j][exp])
      #t = t*powers[j][exp]
    end
    push!(r, c*t)
    j = i = i + 1
    while iseven(j) && length(r) > 1
      top = pop!(r)
      r[end] = addeq!(r[end], top)
      j >>>= 1
    end
  end
  while length(r) > 1
    top = pop!(r)
    r[end] = addeq!(r[end], top)
  end
  return r[1]
end

function evaluate(a::NfAbsNSElem, P::InfPlcNfAbsNS, prec::Int)
  return conjugates(a, prec)[absolute_index(P)]
end

function signature(K::NfAbsNS)
  if K.signature[1] != -1
    return K.signature
  end
  signatures = Tuple{Int, Int}[signature(isunivariate(f)[2]) for f in K.pol]
  r = prod(x[1] for x in signatures)
  s = div(degree(K) - r, 2)
  K.signature = (r, s)
  return (r, s)
end


text:
using .JSort

using SharedArrays

using JLD
using DelimitedFiles
using Interpolations
using Polynomials
import Base.show
import YAML

function constructzrange(path)
    data = readdlm(path, skipstart=2)
    itp = LinearInterpolation(convert(Array{Float64, 1}, 1000data[:, 1]),
                              convert(Array{Float64, 1}, data[:, 2]),  # Convert to keV
                              extrapolation_bc=Line())
    return itp
end

function readgainshift(path::Nothing)
    gainshift = [Float32.([0.0, 1.0]) for f in 1:8, b in 1:8]
end

function readgainshift(path)
    println(path)
    gainshift, header = readdlm(path, ',', header=true)
    # Check the file
    b = f = 1
    gainshiftarray = [Float32[] for f in 1:8, b in 1:8]
    for (i, row) in enumerate(eachrow(gainshift))
        # if Int(row[1]) ≠ f || Int(row[2]) ≠ b
        #     error("Gainshift file is corrupt at row $i = $row")
        # end
        f = Int(row[1])
        b = Int(row[2])
        gainshiftarray[f, b] = row[3:end]
    #     f += 1
    #     if f > 8
    #         f = 1
    #         b += 1
    #     end
    end
    gainshiftarray
end

readgainshift_gamma(paths::Vector{Nothing}) = []
readgainshift_gamma(paths::Nothing) = []
function readgainshift_gamma(paths::Vector{T}) where T
    [readgainshift_gamma(path) for path in paths]
end

function readgainshift_gamma(path::AbstractString)
    gainshift, header = readdlm(path, ',', header=true)
    gainshiftarray = zeros((33, 3))
    channels = Set(collect(1:33))
    for (i, row) in enumerate(eachrow(gainshift))
        try
            channel = round(Int, row[1])
            gainshiftarray[channel, :] = row[2:end]
            pop!(channels, channel)
        catch e
            println(e)
            error("Gamma gainshift file corrupt at row $i = $row")
        end
    end
    if length(channels) > 0
        @warn "Channels $channels are unspecified $path"
    end
    gainshiftarray
end

struct Parameters
    parameterpath::String
    raw  # The YAML file
    calibrator::Dict{Symbol, AbstractArray{Poly}}
    isotope
    datafiles
    exfromeΔe::Array{Float64, 2}
    excorr::Array{Float64, 2}
    thickrange::Array{Float64, 1}
    tnaicorresi::Array{T, 1} where T<:Number
    eΔerect
    percent
    range::Interpolations.Extrapolation{Float64,1,Interpolations.GriddedInterpolation{Float64,1,Float64,Interpolations.Gridded{Interpolations.Linear},Tuple{Array{Float64,1}}},Interpolations.Gridded{Interpolations.Linear},Interpolations.Line{Nothing}}
    cores
    savepath::String
    gates::Dict{Symbol,Gate}

    function Parameters(parameterpath::T) where T <: AbstractString
        join(path) = joinpath(splitdir(parameterpath)[1], path)
        join(path::Nothing) = nothing
        parameters    = YAML.load_file(parameterpath)
        isotope       = getparameter(parameters, "isotope")
        datafiles     = getparameter(parameters, "datafiles", T=Vector{String}) .|> join
        exfromeΔe     = getparameter(parameters, "ex from ede", repeat([0.0 0.0 0.0], 8))
        excorr        = getparameter(parameters, "excitation energy correction", repeat([0.0 1.0], 8))
        thickrange    = getparameter(parameters, "thick range", T=Array{Float64})
        tnaicorresi   = getparameter(parameters, "tnai corr enai", [0, 0, 1, 0])
        eΔerect       = getparameter(parameters, "ede rectangle")
        percent       = getparameter(parameters, "percent to read", 1.0)
        cores         = getparameter(parameters, "cpu cores", 1)
        savepath      = getparameter(parameters, "save path", "sirius") |> join

        # Calibration
        egainshift  = readgainshift(join(parameters["e gainshift file"]))
        Δegainshift = readgainshift(join(parameters["de gainshift file"]))
        γgainshift  = parameters["gamma gainshift file"] .|> join |> readgainshift_gamma
        tgainshift  = parameters["time gainshift file"]  .|> join |> readgainshift_gamma
        calibrator  = Dict{Symbol, Calibrator}(
            :e  => combinegainshift(egainshift, 1:8, 1:8),
            :Δe => combinegainshift(Δegainshift, 1:8, 1:8, initial=[0.0, 2.5]),
            :γ  => combinegainshift(γgainshift, 1:33),
            :t  => combinegainshift(tgainshift, 1:33, initial=[0.0, 1.0])
        )
        # Gates # Fix case for empty gates
        gates = getgates(parameters["gates"])
        range = getparameter(parameters, "range file") |> join |> constructzrange

        new(parameterpath, parameters, calibrator, isotope, datafiles, exfromeΔe, excorr, thickrange,
            tnaicorresi, eΔerect, percent, range, cores, savepath, gates)
    end
end

function show(io::IO, a::Parameters)
    println("Parameters for $(a.isotope)")
    println("Loaded from $(a.parameterpath).")
    println("e-calibration coefficients:")
    show(io, a.calibrator[:e])
    println("\nΔe-calibration coefficients:")
    show(io, a.calibrator[:Δe])
    println("\ngamma calibration coefficients:")
    show(io, a.calibrator[:γ])
    println("\ntime calibration coefficients:")
    show(io, a.calibrator[:t])
    println("Gates:")
    for gate in a.gates
        println("$(gate.first) : $(gate.second.low) to $(gate.second.high)")
    end
end

function show(io::IO, a::AbstractVector{T} where T<: Poly)
    for i in eachindex(a)
        print("$i: ")
        println(a[i])
    end
end

function show(io::IO, a::AbstractMatrix{T} where T<: Poly)
    for row in 1:size(a, 1)
        for col in 1:size(a, 2)
            print("$row, $col: ")
            println(a[row, col])
        end
    end
end

function getparameter(parameters, parameterkey, defaultvalue; constraint=nothing)
    if parameterkey ∉ keys(parameters)
        parameter = defaultvalue
    else
        parameter = parameters[parameterkey]
        # Reshape arrays
        if typeof(defaultvalue) <: Array{Float64}
            parameter = reshape(parameter, size(defaultvalue))
        else
            parameter = convert(typeof(defaultvalue), parameter)
        end
        if !isnothing(constraint)
            if !constraint(parameter)
                error("Parameter $parameterkey does not satisfy $constraint")
            end
        end
    end
    return parameter
end

function getparameter(parameters, parameterkey; T=nothing, constraint=nothing)
    if parameterkey ∉ keys(parameters)
        error("Parameter $parameterkey is required")
    else
        parameter = parameters[parameterkey]
        if !isnothing(T)
            parameter = convert(T, parameter)
        end
        if !isnothing(constraint)
            if !constraint(parameter)
                error("Parameter $parameterkey does not satisfy $constraint")
            end
        end
    end
    return parameter
end

function getpath(parameters::Parameters, key::Symbol)
    basepath = splitdir(parameters.parameterpath)[1]
    path = joinpath(basepath, getfield(parameters, key))
end

function getpath(parameters::Parameters, key)
    basepath = splitdir(parameters.parameterpath)[1]
    filepath = parameters.raw[key]
    path = joinpath(basepath, filepath)
end

### HANDLE GAINSHIFT ###
function combinegainshift(gainshifts::AbstractVector, channels;
                          initial = [0.0, 5.0],
                          docalibrate = true)::Vector{Poly}
    #@assert length(channels) == length(gainshift)
    # Combine all gainshift operations into a single
    # loop for a given channel. Gives a 3x speedup
    # Try to unroll the loop
    calibrators = Poly[]
    for channel in channels
        p = Poly(initial)
        if docalibrate
            for gainshift in gainshifts
                p = Poly(gainshift[channel, :])(p)
            end
        end
        # c = coeffs(p)
        # n = length(c)
        # f = let n=n, c=c
            # x -> begin
                # s = 0.0
                # @inbounds for i in 1:n
                    # s += c[i]*x^(i-1)
                # end
                # s
            # end
        # end
        push!(calibrators, p)
    end
    calibrators
end

function combinegainshift(gainshifts::AbstractMatrix, channelsf, channelsb;
                          initial = [0.0, 5.0],
                          docalibrate = true)::Matrix{Poly}
    combinegainshift([gainshifts], channelsf, channelsb, initial=initial)
end
function combinegainshift(gainshifts::AbstractVector, channelsf, channelsb;
                          initial = [0.0, 5.0],
                          docalibrate = true)::Matrix{Poly}
    # Combine all gainshift operations into a single
    # loop for a given channel. Gives a 3x speedup
    # Try to unroll the loop
    calibrators  = Matrix{Poly}(undef, length(channelsf), length(channelsb))
    for f in channelsf, b in channelsf
        p = Poly(initial)
        if docalibrate
            for gainshift in gainshifts
                p = Poly(gainshift[f, b])(p)
            end
        end
        # TODO why not just use Poly(x)??
        # c = coeffs(p)
        # n = length(c)
        # func = let n=n, c=c
            # x -> begin
                # s = 0.0
                # @inbounds for i in 1:n
                    # s += c[i]*x^(i-1)
                # end
                # s
            # end
        # end
        calibrators[f, b] = p
    end
    calibrators
end

function savegainshift_gamma(path, coefficients)
    println(coefficients)
end

function savegainshift_gamma(parameters::Parameters, coefficients)
    path = getpath(parameters, "gamma gainshift file")
    savegainshift_gamma(path, coefficients)
end

function savegainshift_gamma(parameters::Parameters, filename::String, coefficients)
    path = getpath(parameters, :savepath)
    path = joinpath(path, filename)
    # pad different lengthed polynomials
    N = maximum(map(length, coefficients))
    coefficients = map(x -> begin
            diff = length(x) - N
            if diff < 0
                x = [x..., zeros(abs(diff))...]
            else
                x
            end
        end, coefficients)

    open(path, "w") do io
        header = ["channel" "shift" "gain" "sqgain"]
        @show coefficients
        writedlm(io, [header, coefficients...], ',', header=true)
    end
end

function loadlabr(parameters::Parameters)::Vector{TrueEvent}
    return loadlabr(parameters.savepath)
end

function loadlabr(path::AbstractString)::Vector{TrueEvent}
    fhandle = open(joinpath(path, "labr.bin"))
    data = Base.read(fhandle) |> x -> reinterpret(UInt32, x)
    i = 1
    labr = TrueEvent[]
    while i ≤ length(data)
        f, b = data[i:i+1]
        i += 2
        de, e = convert.(Float32, data[i:i+1])
        i += 2
        labrcount = data[i]
        if labrcount + i >= length(data)
            println("Reading failed!")
        end
        i += 1
        adctdcs = ADCTDC[]
        for j in 1:labrcount
            channel = data[i]
            i += 1
            adc, tdc = convert.(Float32, data[i:i+1])
            push!(adctdcs, ADCTDC(channel, adc, tdc))
            i += 2
        end
        push!(labr, TrueEvent(f, b, de, e, adctdcs))
    end
    labr
end

function savecoefficients(parameters, path, coefficients, channels)
    all_coefficients = [[0.0, 1.0, 0.0] for i in 1:33]
    all_coefficients[channels] = coefficients
    # Tag each coefficient with corresponding channel
    all_coefficients = [[i, val...] for (i, val) in enumerate(all_coefficients)]
    savegainshift_gamma(parameters, path * ".csv", all_coefficients)
end

function savecoefficients_eΔe(parameters, coefficients_e, coefficients_Δe, path::Tuple{String, String}=("egainshift.csv", "degainshift.csv"))
    # Coefficients formated as  front, back, x0   x1   y0   y1
    savepath = getpath(parameters, :savepath)
    spath = joinpath(savepath, path[1])
    open(spath, "w") do io
        header = ["front" "back" "shift" "gain"]
        @show coefficients_e
        writedlm(io, [header, coefficients_e...], ',', header=true)
    end

    spath = joinpath(savepath, path[2])
    open(spath, "w") do io
        header = ["front" "back" "shift" "gain"]
        @show coefficients_Δe
        writedlm(io, [header, coefficients_Δe...], ',', header=true)
    end
end

function readede(path::String; T=Float32)
    data = open(path, "r") do io
        len = read(io, T) |> Int
        data = Matrix{T}(undef, len, 2)
        read!(io, data)
        data
    end
end

max_stars_repo_name: Caronthir/JSort

max_issues_repo_name: Caronthir/JSort

max_forks_repo_name: Caronthir/JSort


text:
"""
Jeans model with line-of-sight velocity dispersion predictions.

* [`JeansModel`](@ref): collection of mass model, tracer density model, and anisotropy model
* [`sigma_los`](@ref): calculate the line-of-sight velocity dispersion
* [`kappa_los`](@ref): calculate the line-of-sight velocity kurtosis
"""
module Jeans

using Interpolations: LinearInterpolation
using DifferentialEquations

using Slomo.Models: Model, DensityModel, NotImplemented
using Slomo.Models: has_analytic_profile, mass, density, density2d
using Slomo.Anisotropy: AnisotropyModel, K_jeans, g_jeans, beta
using Slomo.Constants: G
using Slomo.Integrate: solve, integrate

import Slomo.Models: update

export JeansModel, sigma_los, sigma_los_parallel, kappa_los

"""
    JeansModel(mass_model, tracer_model, anisotropy_model)

Represents the key ingredients for the Jeans modeling.  The first argument can be an array
of `DensityModel` instances, in which case we will consider the mass model to be the sum of
all its components.  The second argument refers to the density profile of the kinematic
tracer population (e.g., the stars or star clusters).
"""
struct JeansModel <: Model
    mass_model::Array{DensityModel, 1}
    density_model::DensityModel
    anisotropy_model::AnisotropyModel
end

JeansModel(mass_model::DensityModel,
           density_model::DensityModel,
           anisotropy_model::AnisotropyModel) = begin
               JeansModel([mass_model], density_model, anisotropy_model)
           end

update(model::JeansModel; parameters...) = begin
    JeansModel(update(model.mass_model; parameters...),
               update(model.density_model; parameters...),
               update(model.anisotropy_model; parameters...))
end

"""
    sigma_los(model::JeansModel, R; <keyword arguments>, parameters...)

Calculate the line-of-sight velocity dispersion by numerically integrating the
spherical Jeans equation.

# Examples
```julia
model = JeansModel(NFWModel(), SersicModel(), ConstantBetaModel())
R = collect(1:0.5:10)
sigma = sigma_los(model, R)
```

# Arguments
- `n_interp::Int = 10`: number of points on the interpolation grid
- `fudge::Real = 1e-6`: fudge factor for tweaking integration bounds
- `interp::Bool = true`: whether or not to interpolate from a grid of R
- `rmax::Real = 1e2`: maximum radius (in kpc) for integrating the Jeans equation
- `parameters`: keywords used to update JeansModel parameter values

See also: [`sigma_los_parallel`](@ref), [`kappa_los`](@ref)
"""
function sigma_los(model::JeansModel, R;
                   n_interp::Int = 10, fudge::Real = 1e-6, interp::Bool = true, rmax::Real = 1e2,
                   parameters...)

    # update models with new parameters
    if length(keys(parameters)) > 0
        model = update(model; parameters...)
    end
    mass_model = model.mass_model
    density_model = model.density_model
    anisotropy_model = model.anisotropy_model

    # check for analytic functions
    has_kernel = has_analytic_profile(K_jeans, anisotropy_model)
    # has_g = has_analytic_profile(g_jeans, anisotropy_model)
    # has_mass = has_analytic_profile(mass, mass_model)
    # has_sd = has_analytic_profile(density2d, density_model)

    # set up interpolation grid
    Rmin = minimum(R)
    Rmax = maximum(R)
    if (length(R) <= n_interp) || ~interp
        Rgrid = R
        interp = false
    else
        Rgrid = exp10.(collect(range(log10(Rmin), stop=log10(Rmax), length=n_interp)))
    end

    # construct required functions
    M(r) = mass(mass_model, r)
    ρ(r) = density(density_model, r)
    Σ(R) = density2d(density_model, R)
    K(r, R) = K_jeans(anisotropy_model, r, R)
    β(r) = beta(anisotropy_model, r)
    g(r) = g_jeans(anisotropy_model, r)

    # if we have an analytic Jeans kernel, use it!
    if has_kernel
        integral = integral_from_kernel(Rgrid, M, ρ, K; fudge = fudge,
                                        rmax = rmax)
    else
        integral = integral_from_vr2(Rgrid, M, ρ, β, g; fudge = fudge,
                                     rmax = rmax)
    end

    # restrict to non-negative quantities
    integral[integral .< eps()] .= 0
    sgrid = @. √(2 / Σ(Rgrid) * integral)
    if interp
        return LinearInterpolation(Rgrid, sgrid)(R)
    else
        return sgrid
    end
end

"""
    sigma_los_parallel(model::JeansModel, R, parameter_sets; <keyword arguments>)

Calculate the line-of-sight velocity dispersion by numerically integrating the
spherical Jeans equation in parallel for different sets of parameters.

`parameter_sets` should be a list of dictionaries where each dictionary contains
an update to the model parameters.

# Examples
```julia
model = JeansModel(NFWModel(), SersicModel(), ConstantBetaModel())
R = collect(1:0.5:10)
betas = collect(0.1:0.1:0.9)
parameter_sets = [Dict(:beta => b) for b in betas]
sigmas = sigma_los(model, R, parameter_sets)
```

# Arguments
- `n_interp::Int = 10`: number of points on the interpolation grid
- `fudge::Real = 1e-6`: fudge factor for tweaking integration bounds
- `interp::Bool = true`: whether or not to interpolate from a grid of R
- `rmax::Real = 1e2`: maximum radius (in kpc) for integrating the Jeans equation

See also: [`sigma_los`](@ref)
"""
function sigma_los_parallel(model::JeansModel, R, parameter_sets;
                            n_interp::Int = 10, fudge::Real = 1e-6,
                            interp::Bool = true, rmax::Real = 1e2)
    if collect(keys(parameter_sets[1]))[1] isa AbstractString
        parameter_sets = [Dict(Symbol(k) => v for (k, v) in dict)
                          for dict in parameter_sets]
    end
    sigma_profiles = zeros(length(parameter_sets), length(R))
    Threads.@threads for i = 1:length(parameter_sets)
        sigma_profiles[i, :] = sigma_los(model, R;
                                         n_interp = n_interp, fudge = fudge,
                                         interp = interp, rmax = rmax,
                                         parameter_sets[i]...)
    end
    return sigma_profiles
end

"""
Calculate the line-of-sight excess kurtosis by numerically integrating the
spherical Jeans equations.

    model : potential and tracer models
    R : projected radii in kpc
    parameters : keywords to pass on to mass and tracer models
    return_sigma : if true, will return both excess kurtosis and dispersion
"""

"""
    kappa_los(model::JeansModel, R; <keyword arguments>, parameters...)

Calculate the line-of-sight velocity kurtosis by numerically integrating the
spherical Jeans equations.

# Examples
```julia
model = JeansModel(NFWModel(), SersicModel(), ConstantBetaModel())
R = collect(1:0.5:10)
kappa = kappa_los(model, R)
```

# Arguments
- `n_interp::Int = 10`: number of points on the interpolation grid
- `fudge::Real = 1e-6`: fudge factor for tweaking integration bounds
- `interp::Bool = true`: whether or not to interpolate from a grid of R
- `return_sigma::Bool = false`: whether or not to additionally return the velocity dispersion
- `rmax::Real = 1e2`: maximum radius (in kpc) for integrating the Jeans equation
- `parameters`: keywords used to update JeansModel parameter values

!!! warning
    This is only valid for the `ConstantBetaModel` anisotropy model.  Using
    some other anisotropy model will give nonsense results.

See also: [`sigma_los`](@ref)
"""
function kappa_los(model::JeansModel, R;
                   n_interp::Int = 10, fudge::Real = 1e-6, interp::Bool = true,
                   return_sigma::Bool = false, rmax::Real = 1e2,
                   parameters...)
    # update models with new parameters
    if length(keys(parameters)) > 0
        model = update(model; parameters...)
    end
    mass_model = model.mass_model
    density_model = model.density_model
    anisotropy_model = model.anisotropy_model

    # set up interpolation grid
    Rmin = minimum(R)
    Rmax = maximum(R)
    if (length(R) <= n_interp) || ~interp
        Rgrid = R
    else
        Rgrid = exp10.(collect(range(log10(Rmin), stop=log10(Rmax), length=n_interp)))
    end

    # construct required functions
    M(r) = mass(mass_model, r)
    ρ(r) = density(density_model, r)
    Σ(R) = density2d(density_model, R)
    β(r) = beta(anisotropy_model, r)
    g(r) = g_jeans(anisotropy_model, r)

    rmin = minimum(Rgrid)

    # integrand for ρ * g * v_r^2
    ρgvr2_integrand(r) = -G * M(r) * ρ(r) * g(r) / r ^ 2
    prob = ODEProblem((u, p, t) -> ρgvr2_integrand(t), 0.0, (rmax, rmin * 1e-2))
    ρgvr2_sol = solve(prob, Tsit5() , isoutofdomain=(u, p, t) -> any(u .<= 0))

    # capture bad inputs, restrict output to non-negative values
    ρgvr2(r) = begin
        if r > rmax
            return 0.0
        end
        return max(ρgvr2_sol(r), 0.0)
    end

    # calculate the second moment of the losvd
    integral = zeros(size(Rgrid))
    for (i, R) in enumerate(Rgrid)
        integrand(r) = (1.0 - β(r) * (R / r)^2) * ρgvr2(r) / g(r) * r / √(r^2 - R^2)
        sol = solve(integrand, rmax, R * (1.0 + fudge))
        integral[i] = sol[1] - sol[end]
    end
    sgrid = @. √(2 / Σ(Rgrid) * integral)

    # integrand for ρ * g * v_r^4 / G
    ρgvr4_integrand(r) = -3 * ρgvr2(r) * G * M(r) / r ^ 2
    prob = ODEProblem((u, p, t) -> ρgvr4_integrand(t), 0.0, (rmax, rmin * 1e-1))
    ρgvr4_sol = solve(prob, Tsit5() , isoutofdomain=(u, p, t) -> any(u .<= 0))

    # capture bad inputs, restrict output to non-negative values
    ρvr4(r) = begin
        if r > rmax
            return 0.0
        end
        y = ρgvr4_sol(r) / g(r)
        return max(y, 0.0)
    end

    # calculate fourth moment of the losvd
    integral = zeros(size(Rgrid))
    for (i, R) in enumerate(Rgrid)
        integrand(r) = begin
            factor1 = (1 - 2 * β(r) * (R / r)^2 + β(r) * (1 + β(r)) / 2 * (R / r)^4)
            factor2 = ρvr4(r) * r / √(r^2 - R^2)
            factor1 * factor2
        end
        sol = solve(integrand, rmax, R * (1.0 + fudge))
        integral[i] = sol[1] - sol[end]
    end
    v4los = @. 2 / Σ(Rgrid) * integral
    kgrid = @. v4los / sgrid ^ 4 - 3

    if return_sigma
        return kgrid, sgrid
    else
        return kgrid
    end
end

"""
    integral_from_kernel(Rgrid, M, ρ, K; <keyword arguments>)

Compute the outer Jeans integral for a model which has an analytic form for the
Jeans projection kernel.
"""
function integral_from_kernel(Rgrid, M, ρ, K; fudge = 1e-6, rmax = 1e2)
    integral = zeros(size(Rgrid))
    for (i, Ri) in enumerate(Rgrid)
        integrand(r) = K(r, Ri) * G * M(r) * ρ(r) / r
        sol = solve(integrand, rmax, Ri * (1.0 + fudge))
        integral[i] = sol[1] - sol[end]
    end
    return integral
end

"""
    integral_from_vr2(Rgrid, M, ρ, K; <keyword arguments>)

Compute the double Jeans integral for a model which has no analytic form for the
Jeans projection kernel.
"""
function integral_from_vr2(Rgrid, M, ρ, β, g; fudge = 1e-6, rmax = 1e2)

    rmin = minimum(Rgrid)

    # integrand for ρ * g * v_r^2
    ρgvr2_integrand(r) = -G * M(r) * ρ(r) * g(r) / r ^ 2
    prob = ODEProblem((u, p, t) -> ρgvr2_integrand(t), 0.0, (rmax, rmin * 1e-2))
    ρgvr2_sol = solve(prob, Tsit5() , isoutofdomain=(u, p, t) -> any(u .<= 0))

    ρgvr2(r) = begin
        if r <= rmax
            return max(ρgvr2_sol(r), 0.0)
        else
            return 0.0
        end
    end

    integral = zeros(size(Rgrid))
    for (i, R) in enumerate(Rgrid)
        integrand(r) = (1.0 - β(r) * (R / r)^2) * ρgvr2(r) / g(r) * r / √(r^2 - R^2)
        sol = solve(integrand, rmax, R * (1.0 + fudge))
        integral[i] = sol[1] - sol[end]
    end
    return integral
end

end


text:
# # Continuous Bag of Words (CBOW) Model

include("bow.jl")

# Instead of mapping words directly to 5 dimensions for 5 classes, we map them to 64
# dimensions to represent features, then go to 5 dimensions after adding up these features.

pred(words) = b_sm .+ W_sm * sum(W_emb[:,words], dims=2)

# Only the prediction function and parameters need to change, the rest of the code is the
# same.

EMB_SIZE = 64
W_emb = param(EMB_SIZE, nwords)
W_sm = param(ntags, EMB_SIZE)
b_sm = param(ntags)

# Bonus question: can this model represent a function that the BOW model cannot represent?

train()


text:
#=

    This is a port of the SPPL model random-sequence.pynb
    """
    In this notebook, we will implement and query the following probabilistic model of a random sequence in SPPL.

    For i=0, X[0]  is a standard normal.
    For i>0
        if  X[i−1] is negative, then  X[i] is a standard normal
        otherwise, let  W[i] be the outcome of a fair coin:
            if  W[i]=0 , then  X[i]=0.5∗(X[i−1])2+X[i−1]
            othewrise(sic!),  X[i] is a standard normal.

    """

    The SPPL model give the following (exact) probabilities
    (adjusted for Python's 0 base):

    * Pr(X[1]>0)		    = 0.50
    * Pr(X[1]>0 | X[2]>0)	= 0.647
    * Pr(X[1]>0 | X[3]>0)	= 0.544

    This model:

    We want the probability of X[1] > 0 given by:
    * No observation:
           t    mean: 0.5046   std: 0.5000

    * We observe X[2] > 0
           t    mean:0.6420    std:0.4797

    * We observe X[3] > 0
           t    mean: 0.5402   std: 0.4984

=#

using Turing, StatsPlots, Distributions, StatsBase
using CSV
include("jl_utils.jl")

@model function random_sequence()
    X = tzeros(3)
    W = tzeros(3)

    X[1] ~ Normal()
    for i in 2:3
        if X[i-1] > 0
            X[i] ~ Normal(0, 1)
            W[i] ~ Dirac(0.0)
        else
            W[i] ~ Bernoulli(0.5)
            if W[i] == 0
                X[i] ~ Dirac( 0.5*X[i-1]^2 + X[i-1])  # Dirac((X[i-1]^2 / 2 + X[i-1])
            else
                X[i] ~ Normal(0, 1)
            end
        end
    end

    # Observations:
    true ~ Dirac(X[2] > 0)
    # true ~ Dirac(X[3] > 0)

    # What is the probability that X[1] > 0?
    t ~ Dirac(X[1] > 0)
end

model = random_sequence()

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
chns = sample(model, PG(10), 1_000)
# chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(), 10_000)

# chns = sample(model, HMC(0.1,6), 1_000)
# chns = sample(model, NUTS(), 1_000)

display(chns[[:t]])
display(chns)

max_stars_repo_name: tias/hakank

max_issues_repo_name: tias/hakank

max_forks_repo_name: tias/hakank