text:
"""
    Find the prices of capital if the firm is not selling everything it produces.
"""
function find_price_info_unc(par)

    lower = [0.0 + eps(), 0.0 + eps()]
    upper = [5, 5]

    function loss(xy)
        sum((f(xy...) .- c) .^ 2)
    end

    opt = optimize(loss, lower, upper, [0.5, 0.5], Fminbox(LBFGS()))
    f(p_I) = begin
        @set! par.p_I = p_I^2
        k_I_b_A, k_I_b_B = solve_firm_problems_unc(par)
        diff = (1 - par.γ) * k_I_b_B + (1 - par.τ) * par.γ * k_I_b_A
        return diff # What one firm sells needs to be equal to what the other firm buys
    end
    find_zero(f, 0.25, tol = 1e-16)^2
end

"""
    Find capital allocation, starting with guesses for the prices of different kinds of capital.
"""
function solve_firm_problems_bundle(par)
    k_I_P_A = k_I_prod_A_bundle(par)
    k_I_P_B = k_I_prod_B_bundle(par)

    K_A_guess = 2 * k_I_P_A
    K_B_guess = 2 * k_I_P_B

    #! Now I can infer the other variables

    f_A(x) = loss_A(x, par)
    k_I_b_A = find_zero(f_A, [-k_I_P_A + 0.00001, 5])

    # B is always unconstrained, because they are buying capital
    k_I_P_B = k_I_prod_B(par)
    f_B(x) = loss_B(x, par)
    k_I_b_B = find_zero(f_B, [-k_I_P_B + 0.0001, 5])

    return k_I_b_A, k_I_b_B
end

"""
   How much capital A firms produce at an interior solution.
"""
k_I_prod_A_bundle(par) = ((par.p_I_A / (1 - par.ϕ)) / (par.ca_A * par.cb_A))^(1 / (par.cb_A - 1))

"""
   How much capital A firms produce at an interior solution.
"""
k_I_prod_A_bundle(par) = ((par.p_I_B / (1 - par.ϕ)) / (par.ca_B * par.cb_B))^(1 / (par.cb_B - 1))

"""
    How much capital is produced by B firms if they are indifferent between
    producing one more unit or selling one unit less.

        AND

    The firm does not sell all of its capital! Otherwise producing one more unit has a different payoff (sell and keep the rest instead of keeping it fully).
"""
k_I_prod_B(par) = ((par.p_I / (1 - par.τ)) / (par.ca_B * par.cb_B))^(1 / (par.cb_B - 1))


text:
using Clang.Generators
using Usrsctp_jll

cd(@__DIR__)

include_dir = joinpath(Usrsctp_jll.artifact_dir, "include") |> normpath

# wrapper generator options
options = load_options(joinpath(@__DIR__, "generator.toml"))

# add compiler flags, e.g. "-DXXXXXXXXX"
args = get_default_args()
push!(args, "-I$include_dir")

header_dir = include_dir
headers = [joinpath(header_dir, header) for header in readdir(header_dir) if endswith(header, ".h")]

# create context
ctx = create_context(headers, args, options)

# run generator
build!(ctx)

ext: jl


text:
@testset "Raw JSON" begin
    fhir_version = FHIRClient.R4()
    base_url = FHIRClient.BaseURL("https://hapi.fhir.org/baseR4")
    auth = FHIRClient.AnonymousAuth()
    client = FHIRClient.Client(fhir_version, base_url, auth)
    search_request = "/Patient?given=Jason&family=Argonaut"
    json_response_search_results_bundle = FHIRClient.request_json(client, "GET", search_request)
    patient_id = json_response_search_results_bundle.entry[1].resource.id
    patient_request = "/Patient/$(patient_id)"
    json_responses_patient = [
        FHIRClient.request_json(client, "GET", patient_request),
        FHIRClient.request_json(client, "GET", patient_request; body = JSON3.read("{}")),
        FHIRClient.request_json(client, "GET", patient_request; query = Dict{String, String}()),
        FHIRClient.request_json(client, "GET", patient_request; body = JSON3.read("{}"), query = Dict{String, String}()),
    ]
    for json_response_patient in json_responses_patient
        @test json_response_patient isa JSON3.Object
    end
    Base.shred!(auth)
    Base.shred!(client)
end


text:
using TypedCodeUtils
using Test

import TypedCodeUtils: reflect, filter, lookthrough,
                       DefaultConsumer, Reflection,
                       identify_invoke, identify_call,
                       process_invoke, process_call

# Test simple reflection
f(x, y) = x + y

@test !isempty(reflect(f, Tuple{Int, Int}))
@test !isempty(reflect(f, Tuple{Int, Number}))
@generated g(x, y) = :(x + y)
@test !isempty(reflect(g, Tuple{Int, Int}))
if VERSION >= v"1.2.0-rc1"
    @test !isempty(reflect(g, Tuple{Int, Number}))
else
    @test isempty(reflect(g, Tuple{Int, Number}))
end

# Cthulhu's inner loop
function cthulhu(ref::Reflection)

    invokes = filter((c)->lookthrough(identify_invoke,      c), ref.CI.code)
    calls   = filter((c)->lookthrough(identify_call,        c), ref.CI.code)

    invokes = map((arg) -> process_invoke(DefaultConsumer(), ref, arg...), invokes)
    calls   = map((arg) -> process_call(  DefaultConsumer(), ref, arg...), calls)

    callsites = append!(invokes, calls)
    sort!(callsites, by=(c)->first(c))
    return callsites
end

params = TypedCodeUtils.current_params()
ref = reflect(f, Tuple{Int, Int}, params=params)
@test length(ref) == 1
calls = cthulhu(first(ref))
nextrefs = collect(first(reflect(c)) for c in calls if TypedCodeUtils.canreflect(c[2]))

function h(x)
    if x >= 2
        return x ^ 2
    else
        return x + 2
    end
end

params = TypedCodeUtils.current_params()
ref = reflect(h, Tuple{Int}, params=params)
@test length(ref) == 1
calls = cthulhu(first(ref))
nextrefs = collect(first(reflect(c)) for c in calls if TypedCodeUtils.canreflect(c[2]))

if VERSION >= v"1.1.0-DEV.215" && Base.JLOptions().check_bounds == 0
Base.@propagate_inbounds function f(x)
    @boundscheck error()
end
g(x) = @inbounds f(x)

params = TypedCodeUtils.current_params()
ref = reflect(g, Tuple{Vector{Float64}}, params=params)
@test length(ref) == 1
ref = first(ref)
@show ref.CI.code
calls = cthulhu(ref)
@test !isempty(calls)

TypedCodeUtils.preprocess!(DefaultConsumer(), ref, true)
calls = cthulhu(ref)
@test isempty(calls)

end


text:
using ClangCompiler

const CC = ClangCompiler

# source file
src = joinpath(@__DIR__, "lookup.cpp")

# compilation flags
args = get_compiler_args()

cc = IncrementalIRGenerator(src, args)

decl_lookup = DeclFinder(cc.instance)

# @assert decl_lookup(cc, "vector", "std::vector")
@assert decl_lookup(cc, "f", "foo::bar::f")

decl = get_decl(decl_lookup)
CC.dump(decl)

# clean up
dispose(decl_lookup)
dispose(cc)


text:
using Pkg
Pkg.activate(joinpath(@__DIR__, ".."))

using Bump10MSun
using Glob
using HDF5
using Logging
using MCMCChainsStorage
using Printf
using StatsBase
using Turing

Nmcmc = 1000
Nchain = Threads.nthreads()
Npost = 128
Nsel = 4096

log_dN_default = make_log_dN(1.0, -1.0, 10.0, 0.7, 0.5)

samps = []
fnames = []
for fn in glob("GW*[0-9].h5", "/Users/wfarr/Research/o3a_posterior_samples/all_posterior_samples")
    h5open(fn, "r") do f
        push!(samps, read(f, "PublicationSamples/posterior_samples"))
        push!(fnames, fn)
    end
end
for fn in glob("*GW*_nocosmo.h5", "/Users/wfarr/Research/o3b_data/PE")
    h5open(fn, "r") do f
        push!(samps, read(f, "C01:Mixed/posterior_samples"))
        push!(fnames, fn)
    end
end

mask = map(samps) do ss
    median([x.mass_2_source for x in ss]) > mlow && median([x.mass_1_source for x in ss]) < mhigh
end
narrow_samps = samps[mask]
narrow_fnames = fnames[mask]

pop_wts_and_wts = map(narrow_samps) do samps
    map(samps) do s
        m1 = s.mass_1_source
        m2 = s.mass_2_source
        z = s.redshift
        pop_wt = exp.(log_dN_default(m1, m2)).*md_sfr_zwt(z)
        (pop_wt, pop_wt ./ li_prior_wt(m1, m2, z))
    end
end
pop_wts = [[x[1] for x in xx] for xx in pop_wts_and_wts]
wts = [[x[2] for x in xx] for xx in pop_wts_and_wts]

neffs = [sum(w)^2 / sum(w.^2) for w in wts]
@info @sprintf("Minimum PE Neff = %.1f", minimum(neffs))
inds = [sample(1:length(w), weights(w), Npost) for w in wts]

m1s = [[xs[i].mass_1_source for i in is] for (xs, is) in zip(narrow_samps, inds)]
m2s = [[xs[i].mass_2_source for i in is] for (xs, is) in zip(narrow_samps, inds)]
log_wts = [[log(w[i]) for i in is] for (w, is) in zip(pop_wts, inds)]

m1s_sel, m2s_sel, zs_sel, pdraw, Ndraw = read_selection("/Users/wfarr/Research/o3b_data/O3-Sensitivity/endo3_bbhpop-LIGO-T2100113-v12.hdf5")
pdraw = pdraw ./ (md_sfr_zwt.(zs_sel) ./ md_sfr(zref))
m1s_sel, m2s_sel, pdraw, Ndraw = resample_selection(log_dN_default, m1s_sel, m2s_sel, pdraw, Ndraw)
@info "Number of resampled selection function samples = $(Ndraw)"

f = Nsel / length(m1s_sel)
Ndraw = round(Int, f*Ndraw)
m1s_sel = m1s_sel[1:Nsel]
m2s_sel = m2s_sel[1:Nsel]
pdraw = pdraw[1:Nsel]

model = broken_pl_model(m1s, m2s, log_wts, m1s_sel, m2s_sel, log.(pdraw), Ndraw)
trace = sample(model, NUTS(), Nmcmc)
trace = with_logger(NullLogger()) do
    append_generated_quantities(trace, generated_quantities(model, trace))
end
@info @sprintf("Minimum Neff_sel = %.1f", minimum(trace[:Neff_sel]))
@info @sprintf("Minimum Neff_samps = %.1f", minimum([minimum(trace[n]) for n in namesingroup(trace, :Neff_samps)]))

h5open(joinpath(@__DIR__, "..", "chains", "chains.h5"), "w") do f
    write(f, trace)
end


text:
# Exported types

type QueueProperties
  queue_id::Integer
  interarrival::Distribution
  service::Distribution
  num_servers::Integer
  max_capacity::Integer
end

QueueProperties(i::Distribution, s::Distribution, servers::Integer, cap::Integer) = QueueProperties(-1, i, s, servers, cap)
QueueProperties(i::Distribution, s::Distribution, servers::Integer) = QueueProperties(-1, i, s, servers, -1)

type QueueEdge
  to::Integer
  weight::Float64
  function QueueEdge(to, weight::Float64)
    if weight < 0 || weight > 1
      throw(ArgumentError(string("weight must be between 0 and 1")))
    end
    new(to, weight)
  end
end

type QueueNode
  queue_id::Integer
  props::QueueProperties
  is_entering::Bool
  edges::AbstractArray{QueueEdge}
end
QueueNode(id::Integer, props::QueueProperties) = QueueNode(id, props, true, Array{QueueEdge}())

type SimulationArgs
  max_time::Float64
  max_customers::Float64
  topology::AbstractArray{QueueNode}
end

type QueueStats
  num_in_system::Float64
  total_num_waiting::Float64
  total_departures::Float64
  total_wait_time::Float64
  total_system_time::Float64
  num_monitors::Float64
  average_wait_time::Float64
  average_system_time::Float64
end

QueueStats() = QueueStats(0, 0, 0, 0, 0, 0, 0, 0)
+(a::QueueStats, b::QueueStats) = QueueStats(a.num_in_system + b.num_in_system, a.total_num_waiting + b.total_num_waiting, a.total_departures + b.total_departures,
                                  a.total_wait_time + b.total_wait_time, a.total_system_time + b.total_system_time, a.num_monitors + b.num_monitors,
                                  a.average_wait_time + b.average_wait_time, a.average_system_time + b.average_system_time)
/(a::QueueStats, b::Integer) = QueueStats(a.num_in_system / b, a.total_num_waiting / b, a.total_departures / b, a.total_wait_time / b, a.total_system_time / b, a.num_monitors / b, a.average_wait_time / b, a.average_system_time / b)

# Implementation types

type Server
  id::Integer
end

type Event
  queue_id::Integer
  event_type::String
  enter_times::AbstractArray{Float64}
  scheduled_time::Float64
  server::Server
end

Event(a::Integer, b::String, c::Float64, d::Float64, e::Server) = Event(a, b, [c], d, e)

type SingleQueueState
  queue_id::Integer
  server_state::Dict{Server, Any}
  props::QueueProperties
  edges::AbstractArray{QueueEdge}
  waiting::Queue
  num_currently_being_served::Integer
end

type SimulationState
  calendar::MutableBinaryHeap{Event}
  current_event::Event
  system_time::Float64
  entering_queue_ids::AbstractArray{Integer}
  queue_states::AbstractArray{SingleQueueState}
end

<=(a::Event, b::Event) = a.scheduled_time <= b.scheduled_time
==(a::Event, b::Event) = isequal(a.scheduled_time, b.scheduled_time)
>=(a::Event, b::Event) = a.scheduled_time >= b.scheduled_time
<(a::Event, b::Event) = a.scheduled_time < b.scheduled_time
>(a::Event, b::Event) =  a.scheduled_time > b.scheduled_time

# Initialization functions

function initialize_system(args::SimulationArgs)
  queue_states = []
  entering_queue_ids = []
  for i in 1:length(args.topology)
    cur = args.topology[i]
    if cur.is_entering
      push!(entering_queue_ids, i)
    end
    cur.props.queue_id = cur.queue_id
    push!(queue_states, SingleQueueState(cur.queue_id, Dict{Server, Any}(), cur.props, cur.edges, Queue(Event), 0))
  end

  state = SimulationState(mutable_binary_minheap(Event), Event(1, "", [0.0], 0, Server(0)), 0, entering_queue_ids, queue_states)
  initialize_servers(state)
  return state
end

function schedule_initial_events(state::SimulationState)
  for queue_id in state.entering_queue_ids
    next_birth(state, queue_id)
  end
  next_monitoring(state)
end

# Execution functions

function should_continue_simulation(state::SimulationState, args::SimulationArgs, qs::QueueStats)
  stop_from_time = args.max_time > 0 && state.system_time < args.max_time
  stop_from_customers = args.max_customers > 0 && qs.num_departures < args.max_customers
  return stop_from_time || stop_from_customers
end

function system_at_capacity(s::SimulationState, id::Integer)
  c = s.queue_states[id]
  return c.props.max_capacity <= 0 || ((length(c.waiting) + c.num_currently_being_served) < c.props.max_capacity)
end

function handle_event(e::Event, s::SimulationState, qs::QueueStats)
  s.system_time = e.scheduled_time
  if e.event_type == "birth"
    if system_at_capacity(s, e.queue_id)
      return birth(e, s, qs)
    else
      return next_birth(s, e.queue_id)
    end
  elseif e.event_type == "death"
    return death(e, s, qs)
  elseif e.event_type == "monitoring"
    return monitoring(s, qs)
  end
end

# Server related functions

function initialize_servers(s::SimulationState)
  for queue in s.queue_states
    for j in 1:queue.props.num_servers
      queue.server_state[Server(j)] = false
    end
  end
end

function find_server(s::SimulationState, queue_id::Integer)
  for (k, v) in s.queue_states[queue_id].server_state
    if v == false
      return k
    end
  end
  return false
end

function free_server(e::Event, s::SimulationState)
  s.queue_states[e.queue_id].server_state[e.server] = false
  s.queue_states[e.queue_id].num_currently_being_served -= 1
end

function allocate_server(e::Event, server::Server, state::SimulationState)
  state.queue_states[e.queue_id].server_state[server] = e
  state.queue_states[e.queue_id].num_currently_being_served += 1
end

# birth related functions

function birth(e::Event, s::SimulationState, qs::QueueStats)
    enqueue!(s.queue_states[e.queue_id].waiting, e)
    queue_id = e.queue_id
    if find_server(s, queue_id) != false
      next_death(s, queue_id, qs)
    end
    next_birth(s, queue_id)
end

function next_birth(s::SimulationState, id::Integer)
  if id ∈ s.entering_queue_ids
    next_time = s.system_time + rand(s.queue_states[id].props.interarrival)
    push!(s.calendar, Event(id, "birth", next_time, next_time, Server(-1)))
  end
end

# death related functions
function possibly_exit_system(e::Event, s::SimulationState, qs::QueueStats)
  current_queue_id = e.queue_id
  current_queue_state = s.queue_states[current_queue_id]
  sum = 0.0
  r = rand()

  for edge in current_queue_state.edges
    sum += edge.weight
    if r <= sum
      push!(e.enter_times, s.system_time)
      push!(s.calendar, Event(edge.to, "birth", e.enter_times, s.system_time, Server(-1)))
      return
    end
  end

  #assume here that if we find no outgoing edge in the system, then we exit the system
  qs.total_system_time = qs.total_system_time + s.system_time - first(e.enter_times)
  qs.total_departures = qs.total_departures + 1
end

function death(e::Event, s::SimulationState, qs::QueueStats)
  free_server(e, s)
  if length(s.queue_states[e.queue_id].waiting) > 0
    next_death(s, e.queue_id, qs)
  end
  possibly_exit_system(e, s, qs)
end

function next_death(s::SimulationState, queue_id::Integer, qs::QueueStats)
  current_queue_state = s.queue_states[queue_id]
  e = dequeue!(current_queue_state.waiting)
  server = find_server(s, queue_id)
  next_time = s.system_time + rand(current_queue_state.props.service)
  qs.total_wait_time = qs.total_wait_time + s.system_time - last(e.enter_times)
  allocate_server(e, server, s)
  e.server = server
  e.scheduled_time = next_time
  e.event_type = "death"
  push!(s.calendar, e)
end

# monitoring

function next_monitoring(s::SimulationState)
  next_time = s.system_time + rand(Exponential(10))
  push!(s.calendar, Event(1, "monitoring", next_time, next_time, Server(-1)))
end

function monitoring(s::SimulationState, qs::QueueStats)
  for queue in s.queue_states
    num_waiting = length(queue.waiting)
    qs.total_num_waiting += num_waiting
    qs.num_in_system += num_waiting + queue.num_currently_being_served
  end
  qs.num_monitors = qs.num_monitors + 1
  next_monitoring(s)
end

function simulate(args::SimulationArgs)
  state = initialize_system(args)
  qs = QueueStats()
  schedule_initial_events(state)
  while (should_continue_simulation(state, args, qs))
    e = pop!(state.calendar)
    handle_event(e, state, qs)
  end
  qs.average_wait_time = qs.total_wait_time / qs.total_departures
  qs.average_system_time = qs.total_system_time / qs.total_departures
  return qs
end

function aggregate_simulate(args::SimulationArgs, n::Integer)
  qs = @parallel (+) for i=1:n
    simulate(args)
  end
  return qs / n
end

# convenience functions

MM1(λ::Integer, μ::Integer, time::Integer) = simulate(SimulationArgs(time, 0, [QueueNode(1, QueueProperties(Exponential(1/λ), Exponential(1/μ), 1, -1), true, [])]))
MMN(λ::Integer, μ::Integer, servers::Integer, time::Integer) = simulate(SimulationArgs(time, 0, [QueueNode(1, QueueProperties(Exponential(1/λ), Exponential(1/μ), servers, -1), true, [])]))


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct ApplicationGatewayAvailableServerVariablesResult <: SwaggerModel

    function ApplicationGatewayAvailableServerVariablesResult(;)
        o = new()
        o
    end
end # type ApplicationGatewayAvailableServerVariablesResult

const _property_map_ApplicationGatewayAvailableServerVariablesResult = Dict{Symbol,Symbol}()
const _property_types_ApplicationGatewayAvailableServerVariablesResult = Dict{Symbol,String}()
Base.propertynames(::Type{ ApplicationGatewayAvailableServerVariablesResult }) = collect(keys(_property_map_ApplicationGatewayAvailableServerVariablesResult))
Swagger.property_type(::Type{ ApplicationGatewayAvailableServerVariablesResult }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_ApplicationGatewayAvailableServerVariablesResult[name]))}
Swagger.field_name(::Type{ ApplicationGatewayAvailableServerVariablesResult }, property_name::Symbol) =  _property_map_ApplicationGatewayAvailableServerVariablesResult[property_name]

function check_required(o::ApplicationGatewayAvailableServerVariablesResult)
    true
end

function validate_property(::Type{ ApplicationGatewayAvailableServerVariablesResult }, name::Symbol, val)
end


text:
# Tests the Arnoldi relation AV = VH when expanding the search subspace

using Test, LinearAlgebra, SparseArrays
using ArnoldiMethod: reinitialize!, Arnoldi, iterate_arnoldi!

@testset "Initialization" begin
    arnoldi = Arnoldi{Float64}(5, 3)
    reinitialize!(arnoldi)
    @test norm(arnoldi.V[:, 1]) ≈ 1
end

@testset "Arnoldi Factorization" begin
    n = 10
    max = 6
    A = sprand(n, n, .1) + I

    arnoldi = Arnoldi{Float64}(n, max)
    reinitialize!(arnoldi)
    V, H = arnoldi.V, arnoldi.H

    # Do a few iterations
    iterate_arnoldi!(A, arnoldi, 1:3)
    @test A * V[:,1:3] ≈ V[:,1:4] * H[1:4,1:3]
    @test norm(V[:,1:4]' * V[:,1:4] - I) < 1e-10

    # Do the rest of the iterations.
    iterate_arnoldi!(A, arnoldi, 4:max)
    @test A * V[:,1:max] ≈ V * H
    @test norm(V' * V - I) < 1e-10
end

@testset "Invariant subspace" begin
    # Generate a block-diagonal matrix A
    A = [rand(4,4)  zeros(4,4);
         zeros(4,4) rand(4,4)]

    # and an initial vector [1; 0; ... 0]
    vh = Arnoldi{Float64}(8, 5)
    V, H = vh.V, vh.H
    V[:,1] .= 0.0
    V[1,1] = 1.0

    # Then {v, Av, A²v, A³v}
    # is an invariant subspace
    iterate_arnoldi!(A, vh, 1:5)

    @test norm(V' * V - I) < 1e-10
    @test iszero(H[5, 4])
end


text:
# Autogenerated wrapper script for Gettext_jll for x86_64-w64-mingw32-cxx11
export libgettext

using CompilerSupportLibraries_jll
using Libiconv_jll
using XML2_jll
JLLWrappers.@generate_wrapper_header("Gettext")
JLLWrappers.@declare_library_product(libgettext, "libgettextlib-0-21.dll")
function __init__()
    JLLWrappers.@generate_init_header(CompilerSupportLibraries_jll, Libiconv_jll, XML2_jll)
    JLLWrappers.@init_library_product(
        libgettext,
        "bin\\libgettextlib-0-21.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
# Example script of how you can train this UNet and do inference
using Images, FileIO, JLD2, Flux, UNet

using Juno: @progress

function check_img_size(img)
    # checks that image size is compatible with the unet
    height, width = size(img)
    if (height - 4) % 16 != 0
        throw(BoundsError())
    end
end

function load_img_label(number, crop=500)
    train_path = "data/membrane/train/"
    img_path = train_path*"image/"*string(number)*".png"
    label_path = train_path*"label/"*string(number)*".png"

    img = load(img_path)[1:crop, 1:crop]
    label = load(label_path)[1:crop, 1:crop]
    check_img_size(img)

    return img, label
end

# ---------------------------------------------
# Train a new model
# ---------------------------------------------

# build up dataset for training
dataset = []
for i = 0:20
    img, label = load_img_label(i, 340)
    X = img2array(unet_tiling(img))
    y = img2array(label)
    push!(dataset, (X, y))
end

loss(x, y) = Flux.crossentropy(
    reshape(unet(x), size(y)[1]^2), reshape(y, size(y)[1]^2),
)
opt = Flux.ADAM(1e-4)

# get a validation datapoint
img, label = load_img_label(25, 340)
X_valid = img2array(unet_tiling(img))
y_valid = img2array(label)

evalcb = () -> @show loss(X_valid, y_valid)

@progress for i = 0:5
    @info "Epoch $i"
    Flux.train!(loss, params(unet), dataset, opt, cb = evalcb)
end

weights = params(unet)
@save "unet_weights.jld2" weights

# ---------------------------------------------
# And now, we can do some inference at last
# ---------------------------------------------

@load "unet_weights.jld2" weights
Flux.loadparams!(unet, weights)

img, label = load_img_label(21)
y_pred = unet(img2array(unet_tiling(img)))

img_pred = array2img(y_pred)


text:
Base.HOME_PROJECT[] = abspath(Base.HOME_PROJECT[]) # JuliaLang/julia/pull/28625

using CLIMA, Documenter

include("generate.jl")

GENERATED_BL_EXAMPLES =
[joinpath("examples", "DGmethods_old", "generated", f) for f in
 (
  "ex_001_periodic_advection.md",
  "ex_002_solid_body_rotation.md",
 )]

makedocs(
  sitename = "CLIMA",
  doctest = false,
  strict = false,
  format = Documenter.HTML(
    prettyurls = get(ENV, "CI", nothing) == "true",
    # prettyurls = !("local" in ARGS),
    # canonical = "https://climate-machine.github.io/CLIMA/stable/",
  ),
  clean = false,
  modules = [Documenter, CLIMA],
  pages = Any[
    "Home" => "index.md",
    "Common" => Any[
      "MoistThermodynamics" => "Common/MoistThermodynamics.md",
    ],
    "Utilites" => Any[
      "RootSolvers" => "Utilities/RootSolvers.md",
    ],
    "Atmos" => Any[
      "Atmos/SurfaceFluxes.md",
      "Atmos/TurbulenceConvection.md",
      "Atmos/EDMFEquations.md",
      "Microphysics" => "Atmos/Microphysics.md",
    ],
    "ODESolvers" => "ODESolvers.md",
    "LinearSolvers" => "LinearSolvers.md",
    "Mesh" => "Mesh.md",
    "Arrays" => "Arrays.md",
    "DGmethods_old" => "DGmethods_old.md",
    "InputOutput.md",
    "Developer docs" => Any[
      "CodingConventions.md",
      "AcceptableUnicode.md",
      "VariableList.md",
    ],
    "Balance Law Examples" => ["BalanceLawOverview.md",
                               GENERATED_BL_EXAMPLES...]
  ],
)

# make sure there are no *.vtu files left around from the build
cd(joinpath(@__DIR__, "build", "examples", "DGmethods_old", "generated")) do
    foreach(file -> endswith(file, ".vtu") && rm(file), readdir())
end

deploydocs(
           repo = "github.com/climate-machine/CLIMA.git",
           target = "build",
          )

max_stars_repo_name: akshaysridhar/CLIMA

max_issues_repo_name: akshaysridhar/CLIMA

max_forks_repo_name: akshaysridhar/CLIMA


text:
using Econometrics, Plots, Statistics, Calculus, LinearAlgebra, DelimitedFiles

include("SVlib.jl") # library of functions for the SV model

# proposals for MCMC
# uniform random walk in one dimension
function proposal1(current, tuning)
    trial = copy(current)
    i = rand(1:size(trial,1))
    trial[i] += tuning[i].*randn()
    return trial
end
# MVN random walk, or occasional draw from prior
function proposal2(current, cholV)
    current + cholV'*randn(size(current))
end

function main()
# setup
θtrue = TrueParameters() # true param values, on param space
lb, ub = PriorSupport()

# generate a new data set, or load the example data
n = 500
burnin = 100
#y = SVmodel(TrueParameters(),n, burnin)
y = readdlm("svdata.txt")

# target statistics at sample data
m0 =  auxstat(y)

# Estimation by MSM
S = 100 # number of simulation reps
shocks_u = randn(n+burnin,S) # used fixed shocks to control chatter
shocks_e = randn(n+burnin,S)
# define the moments
ms = θ -> m0' .- auxstat(θ, shocks_e, shocks_u)  # moment contributions: sample stats minus simulated stats
# do two-step GMM
# first step
W = eye(size(m0,1))
m = θ -> vec(mean(ms(θ),dims=1)) # use average of moment contributions
obj = θ -> InSupport(θ) ? m(θ)'*W*m(θ) : Inf
θinit = (lb + ub)./2.0
θhat, objvalue, converged, details = samin(obj, θinit, lb, ub; ns = 5, verbosity = 3, rt = 0.5)
# second step
W = inv((1.0 + 1.0/S).*cov(ms(θhat))) # get the optimal weight matrix
obj = θ -> InSupport(θ) ? m(θ)'*W*m(θ) : Inf
θhat, objvalue, converged, details = samin(obj, θhat, lb, ub; ns = 5, verbosity = 3, rt = 0.5)
# compute the estimated standard errors and CIs
D = Calculus.jacobian(m, vec(θhat), :central)
V = inv(D'*W*D)
se = sqrt.(diag(V))

# now do MCMC
prior = θ -> Prior(θ)
obj = θ -> InSupport(θ) ? -m(θ)'*W*m(θ) : Inf
verbosity = true
burnin = 100
ChainLength = 1000
# initial proposal moves one at a time
tuning = 0.2/sqrt(12.0)*(ub-lb) # two tenths of a standard. dev. of prior
proposal = θ -> proposal1(θ, tuning)
chain = mcmc(θhat, ChainLength, burnin, prior, obj, proposal, verbosity)
# now use a MVN random walk proposal
Σ = NeweyWest(chain[:,1:3])
tuning = 1.0
for j = 1:5
    P = (cholesky(Σ)).U
    proposal = θ -> proposal2(θ,tuning*P)
    θinit = vec(mean(chain[:,1:3],dims=1))
    if j == 5
        ChainLength = 10000
    end
    chain = mcmc(θinit, ChainLength, 0, prior, obj, proposal, verbosity)
    accept = mean(chain[:,end])
    #println("tuning: ", tuning, "  acceptance rate: ", accept)
    if accept > 0.35
        tuning *= 1.5
    elseif accept < 0.25
        tuning *= 0.5
    end
    θinit = vec(mean(chain[:,1:3],dims=1))
    Σ = 0.5*Σ + 0.5*cov(chain[:,1:3])
end
# plain MCMC fit
posmean = vec(mean(chain[:,1:3],dims=1))
inci = zeros(3)
for i = 1:3
    lower = quantile(chain[:,i],0.05)
    upper = quantile(chain[:,i],0.95)
    inci[i] = θtrue[i] >= lower && θtrue[i] <= upper
end

# Results
p1 = npdensity(chain[:,1]) # example of posterior plot
p2 = npdensity(chain[:,2]) # example of posterior plot
p3 = npdensity(chain[:,3]) # example of posterior plot
plot(p1,p2,p3)
plot!(legendfontsize=4)
#savefig("mcmc.png")
println("extremum estimation results:")
println("true values, estimates, st. error, and limits of 90% CI")
prettyprint([θtrue θhat se θhat-1.644*se θhat+1.644*se],["true value", "estimate", "std. err.", "CI lower", "CI upper"])
return θhat, chain
end
θhat, chain = main()

max_stars_repo_name: Hiroakiyusheng/Econometrics

max_issues_repo_name: Hiroakiyusheng/Econometrics

max_forks_repo_name: Hiroakiyusheng/Econometrics


text:
function bracket_minimum(f, x=0; s=1e-2, k=2.0)
    a, ya = x, f(x)
    b, yb = a + s, f(a + s)
    if yb > ya
        a, b = b, a
        ya, yb = yb, ya
        s = -s
    end
    while true
        c, yc = b + s, f(b + s)
        if yc > yb
            return a < c ? (a, c) : (c, a)
        end
        a, ya, b, yb = b, yb, c, yc
        s *= k
    end
end

f = x->sin(10x)+cos(3x)
x0 = 3.4
#busca do intervalo
a,c = bracket_minimum(f,x0)

using Plots

plot(f,0.,4.,lw=2,draw_arrow="true")
scatter!([a,c],[f(a),f(c)])


text:
using HigherKindedPolymorphisms
using Test

abstract type Vect end
@def_hkt Vect{T} where T = Vector{T}
@testset "HigherKindedPolymorphisms.jl" begin
    hkt_vect = inj([1, 2, 3])
    @test (hkt_vect |> typeof) == App{Vect, Int}
    @test prj(hkt_vect) == [1, 2, 3]
    # Write your own tests here.
end


text:
using GLMakie
npairs, θinit = 40, 10*π/180
z = 1:npairs
θ = z .* θinit
x(θ; r = 5) = r * cos(θ)
y(θ; r = 5) = r * sin(θ)
x1, y1 = x.(θ), y.(θ)
x2, y2 = x.(θ .+ π), y.(θ .+ π)
colors1 = rand(1:4, npairs) # 1, 2, 3, 4 =>  'A', 'T', 'C', 'G'
colors2 = [i == 1 ? 2 : i==2 ? 1 : i == 3 ? 4 : 3 for i in colors1]
cmap = cgrad(:Set1_4, 4, categorical = true)
lseg, colors = Point3f[], Int64[]
for i in 1:npairs
    push!(lseg, [x1[i], z[i], y1[i]])
    push!(lseg, [x2[i], z[i], y2[i]])
    push!(colors, colors1[i])
    push!(colors, colors2[i])
end

with_theme(theme_dark()) do
    fig, ax, = meshscatter(x1, z, y1, color = colors1, colormap = cmap,
        colorrange = (0.5, 4.5), markersize = 0.35,
        axis=(type=Axis3, aspect = :data, perspectiveness = 0.5,
        azimuth = -0.49, elevation = 0.63),
        figure = (;resolution = (1200,800)))
    meshscatter!(ax, x2, z, y2, color = colors2, colormap = cmap,
        colorrange = (0.85, 4.5), markersize = 0.35)
    linesegments!(lseg; color = colors, colormap = cmap)
    meshscatter!(ax, tuple.(8, 1:10:40,-2); color = 1:4, colormap = cmap,
        markersize = 0.5)
    text!(ax, ["A", "T", "C", "G"], position = tuple.(8, 1:10:40,-4),
        align = (:center, :bottom), color = :white)
    hidedecorations!(ax; grid=false)
    save("helix.png", fig)
    display(fig)
end


text:
module Vortex

import ..Elements
import ..Elements: circulation, flux, impulse, angularimpulse, seed_zeros, seed_position,
                    seed_strength, kind, seed_position!, seed_strength!

import ..Utils: dualize, ComplexGradientConfig, seed!, Partials

import ..Points
import ..Blobs
import ..Sheets

#== Wrapper for a point vortex ==#

"""
    Vortex.Point(z::ComplexF64, Γ::Float64)

A point vortex located at `z` with circulation `Γ`.

A new point vortex can be created from an existing one by treating the
existing point vortex as a function and passing in the parameter you
want to change as keyword arguments.
For example,
```jldoctest
julia> p = Vortex.Point(1.0, 1.0)
Vortex.Point(1.0 + 0.0im, 1.0)

julia> p()
Vortex.Point(1.0 + 0.0im, 1.0)

julia> p(Γ = 2.0)
Vortex.Point(1.0 + 0.0im, 2.0)
```
"""
#const Point = Points.Point{Float64,Float64}
const Point = Points.Point{T,R} where {T<: Real, R<:Real}
Point(z::Complex{R},Γ::T) where {T<:Real,R<:Real} = Points.Point{T}(z,Γ)
Point(z::Real,Γ::T) where {T} = Points.Point{T}(complex(z),Γ)

(p::Point)(; z = p.z, Γ = p.S) = Point(z, Γ)

circulation(p::Point) = p.S
flux(::Point) = 0.0
impulse(p::Point) = -im*p.z*p.S
angularimpulse(p::Point) = -0.5*p.z*conj(p.z)*p.S

Base.show(io::IO, s::Point) = print(io, "Vortex.Point($(s.z), $(s.S))")

#== Wrapper for a vortex blob ==#

"""
    Vortex.Blob(z::ComplexF64, Γ::Float64, δ::Float64)

A regularized point vortex located at `z` with circulation `Γ` and blob radius `δ`.

A new vortex blob can be created from an existing one by treating the
existing blob as a function and passing in the parameter you want to
change as keyword arguments.
For example,
```jldoctest
julia> b = Vortex.Blob(1.0, 1.0, 0.1)
Vortex.Blob(1.0 + 0.0im, 1.0, 0.1)

julia> b()
Vortex.Blob(1.0 + 0.0im, 1.0, 0.1)

julia> b(Γ = 2.0, δ = 0.01)
Vortex.Blob(1.0 + 0.0im, 2.0, 0.01)
```
"""
#const Blob = Blobs.Blob{Float64,Float64}
const Blob = Blobs.Blob{T,R} where {T<:Real,R<:Real}
Blob(z::Complex{R},Γ::T,δ::Float64) where {T<:Real,R<:Real} = Blobs.Blob{T}(z,Γ,δ)
Blob(z::Real,Γ::T,δ) where {T<:Real} = Blobs.Blob{T}(complex(z),Γ,δ)

(b::Blob)(; z = b.z, Γ = b.S, δ = b.δ) = Blob(z, Γ, δ)

circulation(b::Blob) = b.S
flux(::Blob) = 0.0
impulse(b::Blob) = -im*b.z*b.S
angularimpulse(b::Blob) = -0.5*b.z*conj(b.z)*b.S
Base.show(io::IO, s::Blob) = print(io, "Vortex.Blob($(s.z), $(s.S), $(s.δ))")

#== Wrapper for a vortex sheet ==#

"""
    Vortex.Sheet <: Elements.Element

A vortex sheet represented by vortex blob control points

## Fields

- `blobs`: the underlying array of vortex blobs
- `Ss`: the cumulated sum of circulation starting from the first control point
- `δ`: the blob radius of all the vortex blobs
- `zs`: a mapped array that accesses the position of each control point

## Constructors:

- `Vortex.Sheet(blobs, Γs, δ)`
- `Vortex.Sheet(zs, Γs, δ)` where `zs` is an array of positions for the control points
"""
const Sheet = Sheets.Sheet{T,R} where {T<:Real,R<:Real}
Sheet(blobs::Vector{Blob}, Ss::AbstractVector{Float64}, δ::Float64) = Sheets.Sheet(blobs, Ss, δ)
Sheet(zs::AbstractVector,  Ss::AbstractVector{Float64}, δ::Float64) = Sheets.Sheet(zs, Ss, δ)

function Base.show(io::IO, s::Sheet)
    L = Sheets.arclength(s)
    print(io, "Vortex Sheet: L ≈ $(round(L, digits=3)), Γ = $(round(s.Ss[end] - s.Ss[1], digits=3)), δ = $(round(s.δ, digits=3))")
end

circulation(s::Sheet) = s.Ss[end] - s.Ss[1]

include("vortex/diff.jl")

end


text:
using Tokenize
using BenchmarkTools
using Printf

function speed_test()
    tot_files = 0
    tot_tokens = 0
    tot_errors = 0
    basedir = dirname(Base.find_source_file("int.jl"))
    for dir in (basedir, Sys.STDLIB)
        for (root, dirs, files) in walkdir(dir)
            for file in files
                if endswith(file, ".jl")
                    tot_files += 1
                    file = joinpath(root, file)
                    str = read(file, String)::String
                    l = tokenize(str)
                    while !Tokenize.Lexers.eof(l)
                        t = Tokenize.Lexers.next_token(l)
                        tot_tokens += 1
                        if Tokens.iserror(t.kind)
                            tot_errors += 1
                        end
                    end
                end
            end
        end
    end
    tot_files, tot_tokens, tot_errors
end

tot_files, tot_tokens, tot_errors = speed_test()
tot_time_token = @belapsed speed_test()
println("Lexed ", tot_files, " files, with a total of ", tot_tokens,
        " tokens with ", tot_errors, " errors")
println("Time Token: ", @sprintf("%3.4f", tot_time_token), " seconds")


text:
# Odometry example from GTSAM written in Julia
# This example is using GTSAM 4.0 Python modules
# @author: Andrej Kitanov
# ANPL, Technion

using PyCall
#@pyimport __future__.print_function
@pyimport gtsam
@pyimport gtsam_utils
@pyimport numpy as np

# keys(o::PyObject) # returns an array of the available attribute symbols.

fg = gtsam.NonlinearFactorGraph()
initial = gtsam.Values()

x0 = gtsam.Pose2(0.0, 0.0, 0.0)
#priorNoise = gtsam.noiseModel.Diagonal.Sigmas(np.array([0.3, 0.3, 0.1]))
priorNoise = gtsam.noiseModel[:Diagonal][:Sigmas](np.array([0.3, 0.3, 0.1])) # python
#priorNoise = gtsam.noiseModel_Diagonal[:Sigmas](np.array([0.3, 0.3, 0.1])) # cython
fg[:add](gtsam.PriorFactorPose2(0, x0, priorNoise))

# Add odometry factors
odometry = gtsam.Pose2(2.0, 0.0, 0.0)
odometryNoise = gtsam.noiseModel[:Diagonal][:Sigmas](np.array([0.2, 0.2, 0.1])) # python
#odometryNoise = gtsam.noiseModel_Diagonal[:Sigmas](np.array([0.2, 0.2, 0.1])) # cython
# Create odometry (Between) factors between consecutive poses
fg[:add](gtsam.BetweenFactorPose2(0, 1, odometry, odometryNoise))
fg[:add](gtsam.BetweenFactorPose2(1, 2, odometry, odometryNoise))
fg[:print]() # python
#fg[:print_]("") # cython

initial[:insert](0, gtsam.Pose2(0.5, 0.0, 0.2))
initial[:insert](1, gtsam.Pose2(2.3, 0.1, -0.2))
initial[:insert](2, gtsam.Pose2(4.1, 0.1, 0.1))
initial[:print]("\nInitial Estimate:\n")
#initial[:print_]("\nInitial Estimate:\n")

params = gtsam.LevenbergMarquardtParams()
lm = gtsam.LevenbergMarquardtOptimizer(fg, initial, params)
result = lm[:optimize]()
result[:print]("\nFinal Result:\n")
#result[:print_]("\nFinal Result:\n")

# show more rows on screen
# show(IOContext(STDOUT, limit=true, displaysize=(100,4)), "text/plain", zeros(100,1))
# show(IOContext(STDOUT, limit=false), "text/plain", keys(fg))
#gtsam.writeG2o(fg, initial, "odometry_graph.g2o")
#retrieved = gtsam.readG2o("odometry_graph.g2o")
#retrieved[1][:print_]("\nRetrieved graph from file\n")
#retrieved[2][:print_]("\nRetrieved values from file\n")

# test serialization
str = gtsam.serializeGraph(fg)
println(str)
fg2 = gtsam.deserializeGraph(str)
fg2[:print]("Deserialized factor graph")


text:
using Test
using SafeTestsets
using NQCDynamics

const GROUP = get(ENV, "GROUP", "All")

if GROUP == "All" || GROUP == "Core"
    @time @safetestset "Calculator Tests" begin include("Core/calculators.jl") end
    @time @safetestset "Simulation Tests" begin include("Core/simulations.jl") end
    @time @safetestset "Ring Polymer Tests" begin include("Core/ring_polymers.jl") end
    @time @safetestset "Estimator tests" begin include("Core/estimators.jl") end
end

if GROUP == "All" || GROUP == "InitialConditions"
    @time @safetestset "Monte Carlo Tests" begin include("InitialConditions/monte_carlo.jl") end
    @time @safetestset "AdvancedMH Sampling Tests" begin include("InitialConditions/advancedmh_sampling.jl") end
    @time @safetestset "QuantisedDiatomic Tests" begin include("InitialConditions/quantised_diatomic.jl") end
end

if GROUP == "All" || GROUP == "Dynamics"
    @time @safetestset "DynamicsUtils Tests" begin include("Dynamics/DynamicsUtils.jl") end
    @time @safetestset "Classical Tests" begin include("Dynamics/classical.jl") end
    @time @safetestset "Langevin Tests" begin include("Dynamics/langevin.jl") end
    @time @safetestset "MDEF BAOAB Tests" begin include("Dynamics/mdef_baoab.jl") end
    @time @safetestset "MDEF Tests" begin include("Dynamics/mdef.jl") end
    @time @safetestset "RPMDEF Tests" begin include("Dynamics/rpmdef.jl") end
    @time @safetestset "BCBwithTsit5 Tests" begin include("Dynamics/bcbwithtsit5.jl") end
    @time @safetestset "FSSH Tests" begin include("Dynamics/fssh.jl") end
    @time @safetestset "NRPMD Tests" begin include("Dynamics/nrpmd.jl") end
    @time @safetestset "CMM Tests" begin include("Dynamics/cmm.jl") end
    @time @safetestset "RPeCMM Tests" begin include("Dynamics/rpecmm.jl") end
    @time @safetestset "SpinMapping Tests" begin include("Dynamics/spin_mapping.jl") end
    @time @safetestset "Cell Boundary Callback Tests" begin include("Dynamics/cell_boundary_callback.jl") end
    @time @safetestset "Ehrenfest Tests" begin include("Dynamics/ehrenfest.jl") end
    @time @safetestset "Ensemble Tests" begin include("Ensembles/ensembles.jl") end
    @time @safetestset "Ensemble Selection Tests" begin include("Ensembles/selections.jl") end
    @time @safetestset "Ensemble Reduction Tests" begin include("Ensembles/reductions.jl") end
    @time @safetestset "Ensemble Output Tests" begin include("Ensembles/outputs.jl") end
end


text:
const PAPI_NULL = Cint(-1)
const PAPI_OK = Cint(0)
const PAPI_ENOEVNT = Cint(-7)
const PAPI_ECNFLCT = Cint(-8)
const PAPI_MAX_STR_LEN=128
const PAPI_MIN_STR_LEN=64
const PAPI_ENUM_EVENTS=0
const PAPI_ENUM_FIRST=1
const PAPI_PRESET_ENUM_AVAIL=2
const PAPI_NATIVE_MASK=Cuint(0x40000000)
const PAPI_PRESET_MASK=Cuint(0x80000000)


text:
module SynthControl

using Dates, DataFrames, Optim, LinearAlgebra, RecipesBase, Parameters, FixedEffectModels, Statistics
import CSV

using Reexport
@reexport using TreatmentPanels

export SimpleSCM, SyntheticDiD, fit!, isfitted
export load_brexit, load_germany, load_basque, load_smoking
export load_brexit_panel, load_germany_panel, load_basque_panel, load_smoking_panel

include("SimpleSCM.jl")
include("plotrecipes.jl")
include("load_data.jl")
include("SyntheticDiD.jl")

end # module

max_stars_repo_name: nilshg/SynthControl

max_issues_repo_name: nilshg/SynthControl

max_forks_repo_name: nilshg/SynthControl


text:
# Plot helpers

"Axis unit type. Currently supported: `SI`, `EARTH`."
@enum AxisUnit SI EARTH
"Color scales type for 2D plots. Currently supported: `Log`, `Linear`, `SymLog`."
@enum ColorScale Log Linear SymLog

"Plotting arguments."
struct PlotArgs
   "Unit of spatial axis"
   axisunit::AxisUnit
   "data array size"
   sizes::Tuple{Int64, Int64}
   "plotting data range"
   plotrange::Tuple{Float64, Float64, Float64, Float64}
   "cell IDs in the cut plane"
   idlist::Vector{Int}
   "mapping from original cell order to cut plane"
   indexlist::Vector{Int}
   "title"
   str_title::String
   "xlabel"
   strx::String
   "ylabel"
   stry::String
   "colorbar title"
   cb_title::String
end

"""
    set_args(meta::MetaVLSV, var, axisunit::AxisUnit; normal::Symbol=:none, origin=0.0)

Set plot-related arguments of `var` in `axisunit`. `normal` and `origin` are used for 2D
slices of 3D data, as specified in [`pcolormeshslice`](@ref).
"""
function set_args(meta::MetaVLSV, var, axisunit::AxisUnit; normal::Symbol=:none, origin=0.0)
   (;ncells, coordmin, coordmax) = meta

   if normal == :x
      seq = (2,3)
      dir = 1
   elseif normal == :y || (ncells[2] == 1 && ncells[3] != 1) # polar
      seq = (1,3)
      dir = 2
   elseif normal == :z || (ncells[3] == 1 && ncells[2] != 1) # ecliptic
      seq = (1,2)
      dir = 3
   else
      throw(ArgumentError("1D data detected. Please use 1D plot functions."))
   end

   plotrange = (coordmin[seq[1]], coordmax[seq[1]], coordmin[seq[2]], coordmax[seq[2]])
   axislabels = ['X', 'Y', 'Z'][[seq...]]
   # Scale the sizes to the highest refinement level
   sizes = ncells[[seq...]] .<< meta.maxamr # data needs to be refined later

   if normal == :none
      idlist, indexlist = Int[], Int[]
   else
      idlist, indexlist = let sliceoffset = origin - coordmin[dir]
         getslicecell(meta, sliceoffset, dir, coordmin[dir], coordmax[dir])
      end
   end

   unitstr = axisunit == EARTH ? L"$R_E$" : L"$m$"
   strx = axislabels[1]*" ["*unitstr*"]"
   stry = axislabels[2]*" ["*unitstr*"]"

   str_title = @sprintf "t= %4.1fs" meta.time

   datainfo = readvariablemeta(meta, var)

   cb_title = !isempty(datainfo.variableLaTeX) ?
      datainfo.variableLaTeX * " ["*datainfo.unitLaTeX*"]" : ""

   PlotArgs(axisunit, sizes, plotrange, idlist, indexlist, str_title, strx, stry, cb_title)
end

"""
    set_lim(vmin, vmax, data, colorscale=Linear)

Set colormap limits `vmin`, `vmax` for `data` under scale `colorscale`.
"""
function set_lim(vmin::T, vmax::T, data, colorscale::ColorScale=Linear) where T
   if colorscale in (Linear, SymLog)
      v1 = isinf(vmin) ? minimum(x->isnan(x) ? typemax(T) : convert(T,x), data) : vmin
      v2 = isinf(vmax) ? maximum(x->isnan(x) ? typemin(T) : convert(T,x), data) : vmax
   else # logarithmic
      datapositive = data[data .> 0.0f0]
      v1 = isinf(vmin) ? minimum(datapositive) : vmin
      v2 = isinf(vmax) ? maximum(x->isnan(x) ? typemin(T) : convert(T,x), data) : vmax
   end

   v1, v2
end

"Return x and y ranges for 2D."
function get_axis(axisunit::AxisUnit, plotrange, sizes)
   if axisunit == EARTH
      x = LinRange(plotrange[1], plotrange[2], sizes[1]) ./ RE
      y = LinRange(plotrange[3], plotrange[4], sizes[2]) ./ RE
   else
      x = LinRange(plotrange[1], plotrange[2], sizes[1])
      y = LinRange(plotrange[3], plotrange[4], sizes[2])
   end
   x, y
end

get_axis(pArgs::PlotArgs) = get_axis(pArgs.axisunit, pArgs.plotrange, pArgs.sizes)

"""
    prep2d(meta, var, comp=0) -> Array

Obtain data from `meta` of `var` for 2D plotting. Use `comp` to select vector components.
"""
function prep2d(meta::MetaVLSV, var, comp=0)
   dataRaw = Vlasiator.getdata2d(meta, var)

   data =
      if ndims(dataRaw) == 3
         if comp in (:x, 1)
            @view dataRaw[1,:,:]
         elseif comp in (:y, 2)
            @view dataRaw[2,:,:]
         elseif comp in (:z, 3)
            @view dataRaw[3,:,:]
         elseif comp in (:mag, 0)
            @views hypot.(dataRaw[1,:,:], dataRaw[2,:,:], dataRaw[3,:,:])
         end
      else
         dataRaw
      end
   data
end

"""
    prep2dslice(meta::MetaVLSV, var, normal, comp, pArgs::PlotArgs)

Return `data` of `var` on a uniform 2D mesh on the finest AMR level. Use `normal` to select
the plane orientation, and `comp` to select the component of a vector, same as in
[`pcolormeshslice`](@ref).
"""
function prep2dslice(meta::MetaVLSV, var, normal, comp, pArgs::PlotArgs)
   (;idlist, indexlist) = pArgs

   data3D = readvariable(meta, var)

   if startswith(var, "fg_") # field quantities, fsgrid
      throw(ArgumentError("FS grid variable $var plotting in cut currently not supported!"))
   else # moments, dccrg grid
      # vlasov grid, AMR
      if ndims(data3D) == 1
         data2D = data3D[indexlist]

         data = refineslice(meta, idlist, data2D, normal)
      elseif ndims(data3D) == 2
         data2D = data3D[:,indexlist]

         if comp in (:x, :y, :z, 1, 2, 3)
            if comp in (:x, 1)
               slice = @view data2D[1,:]
            elseif comp in (:y, 2)
               slice = @view data2D[2,:]
            elseif comp in (:z, 3)
               slice = @view data2D[3,:]
            end
            data = refineslice(meta, idlist, slice, normal)
         elseif comp in (:mag, 0)
            datax = @views refineslice(meta, idlist, data2D[1,:], normal)
            datay = @views refineslice(meta, idlist, data2D[2,:], normal)
            dataz = @views refineslice(meta, idlist, data2D[3,:], normal)
            data = hypot.(datax, datay, dataz)
         end
      end
   end
   data
end

"""
    prep_vdf(meta::MetaVLSV, location; kwargs...)

Return the cell velocities `v1, v2`, bin ranges `r1, r2`, cell VDF values `fweight`,
and strings of labels and titles for VDF plots.

# Optional arguments
- `unit::AxisUnit`: location unit in `SI`, `EARTH`.
- `unitv::String`: velocity unit in ("km/s", "m/s").
- `limits::Vector{Real}`: velocity space range given in [xmin, xmax, ymin, ymax].
- `slicetype`: symbol for choosing the slice type from `:xy`, `:xz`, `:yz`,
`:bperp`, `:bpar1`, `:bpar2`.
- `center`: symbol for setting the reference frame from `:bulk`, `:peak`.
- `vslicethick`: setting the velocity space slice thickness in the normal direction. If set
to 0, the whole distribution along the normal direction is projected onto a plane. Currently
this is only meaningful when `center` is set such that a range near the bulk/peak normal
velocity is selected!
- `weight::Symbol`: choosing distribution weights from phase space density or particle flux
between `:particle` and `:flux`.
- `flimit`: minimum VDF threshold for plotting.
- `verbose`: display the selection process.
"""
function prep_vdf(meta::MetaVLSV, location;
   species="proton", unit::AxisUnit=SI, unitv="km/s", slicetype=:default, vslicethick=0.0,
   center=:nothing, weight=:particle, flimit=-1.0, verbose=false)

   ncells = meta.ncells

   if haskey(meta.meshes, species)
      vmesh = meta.meshes[species]
   else
      throw(ArgumentError("Unable to detect population $species"))
   end

   if slicetype ∉ (:default, :xy, :xz, :yz, :bperp, :bpar1, :bpar2)
      throw(ArgumentError("Unknown type $slicetype"))
   end

   unit == EARTH && (location .*= RE)

   # Set unit conversion factor
   unitvfactor = unitv == "km/s" ? 1f3 : 1.0f0

   # Get closest cell ID from input coordinates
   cidReq = getcell(meta, location)
   cidNearest = getnearestcellwithvdf(meta, cidReq)

   # Set normal direction
   if slicetype == :default
      slicetype =
         if ncells[2] == 1 && ncells[3] == 1 # 1D, select xz
            :xz
         elseif ncells[2] == 1 # polar
            :xz
         elseif ncells[3] == 1 # ecliptic
            :xy
         end
   end

   if slicetype in (:xy, :yz, :xz)
      if slicetype == :xy
         dir1, dir2, dir3 = 1, 2, 3
         ŝ = @SVector [0., 0., 1.]
      elseif slicetype == :xz
         dir1, dir2, dir3 = 1, 3, 2
         ŝ = @SVector [0., 1., 0.]
      elseif slicetype == :yz
         dir1, dir2, dir3 = 2, 3, 1
         ŝ = @SVector [1., 0., 0.]
      end
      v1size = vmesh.vblocks[dir1] * vmesh.vblock_size[dir1]
      v2size = vmesh.vblocks[dir2] * vmesh.vblock_size[dir2]

      v1min, v1max = vmesh.vmin[dir1], vmesh.vmax[dir1]
      v2min, v2max = vmesh.vmin[dir2], vmesh.vmax[dir2]
   elseif slicetype in (:bperp, :bpar1, :bpar2)
      if hasvariable(meta, "vg_b_vol")
         b̂ = readvariable(meta, "vg_b_vol", cidNearest) |> vec |> normalize!
      elseif hasvariable(meta, "B_vol")
         b̂ = readvariable(meta, "B_vol", cidNearest) |> vec |> normalize!
      end
      if hasvariable(meta, "proton/vg_v")
         v̂ = readvariable(meta, "proton/vg_v", cidNearest) |> vec |> normalize!
      end

      eperp1 = b̂ × v̂
      êperp1 = normalize(eperp1)
      # b̂ ∥ v̂
      norm(eperp1) < 1f-1 && @warn "B is almost parallel to V!"

      êperp2 = normalize(b̂ × êperp1)

      if slicetype == :bperp # slice in b_perp1/b_perp2
         ŝ = SVector{3}(b̂)
         es = SMatrix{3,3}(hcat(êperp1, êperp2, b̂))
      elseif slicetype == :bpar2 # slice in b_parallel/b_perp1 plane
         ŝ = SVector{3}(êperp2)
         es = SMatrix{3,3}(hcat(b̂, êperp1, êperp2))
      else # slice in b_parallel/b_perp2 plane
         ŝ = SVector{3}(êperp1)
         es = SMatrix{3,3}(hcat(êperp2, b̂, êperp1))
      end
      Rinv = let e = @SMatrix Float32[1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 1.0]
         getRotationMatrix(e, es)'
      end
      # heuristic guess
      v1size = vmesh.vblocks[1] * vmesh.vblock_size[1]
      v2size = v1size
      # Compute the ranges along the two perpendicular axis in the slice plane
      r = zeros(Float32, 3, 8)
      r[:,1] = Rinv * Float32[vmesh.vmin[1], vmesh.vmin[2], vmesh.vmin[3]]
      r[:,2] = Rinv * Float32[vmesh.vmin[1], vmesh.vmin[2], vmesh.vmax[3]]
      r[:,3] = Rinv * Float32[vmesh.vmin[1], vmesh.vmax[2], vmesh.vmin[3]]
      r[:,4] = Rinv * Float32[vmesh.vmin[1], vmesh.vmax[2], vmesh.vmax[3]]
      r[:,5] = Rinv * Float32[vmesh.vmax[1], vmesh.vmin[2], vmesh.vmin[3]]
      r[:,6] = Rinv * Float32[vmesh.vmax[1], vmesh.vmin[2], vmesh.vmax[3]]
      r[:,7] = Rinv * Float32[vmesh.vmax[1], vmesh.vmax[2], vmesh.vmin[3]]
      r[:,8] = Rinv * Float32[vmesh.vmax[1], vmesh.vmax[2], vmesh.vmax[3]]

      v1min, v1max = @views extrema(r[1,:])
      v2min, v2max = @views extrema(r[2,:])
   end

   if !isapprox((v1max - v1min) / v1size, (v2max - v2min) / v2size)
      @warn "Noncubic vgrid applied!"
   end
   cellsize = (v1max - v1min) / v1size

   r1 = LinRange(v1min, v1max, v1size+1) / unitvfactor
   r2 = LinRange(v2min, v2max, v2size+1) / unitvfactor

   vcellids, vcellf = readvcells(meta, cidNearest; species)

   V = getvcellcoordinates(meta, vcellids; species)

   if center == :bulk # centered with bulk velocity
      Vbulk =
         if hasvariable(meta, "moments") # From a restart file
            readvariable(meta, "restart_V", cidNearest)
         elseif hasvariable(meta, species*"/vg_v") # Vlasiator 5
            readvariable(meta, species*"/vg_v", cidNearest)
         elseif hasvariable(meta, species*"/V")
            readvariable(meta, species*"/V", cidNearest)
         else
            readvariable(meta, "V", cidNearest)
         end

      for i in eachindex(V)
         V[i] = V[i] .- Vbulk
      end
   elseif center == :peak # centered on highest VDF-value
      Vpeak = maximum(vcellf)
      for i in eachindex(V)
         V[i] = V[i] .- Vpeak
      end
   end

   # Set sparsity threshold
   if flimit < 0
      flimit =
         if hasvariable(meta, species*"/vg_effectivesparsitythreshold")
            readvariable(meta, species*"/vg_effectivesparsitythreshold", cidNearest)
         elseif hasvariable(meta, species*"/EffectiveSparsityThreshold")
            readvariable(meta, species*"/EffectiveSparsityThreshold", cidNearest)
         else
            1e-16
         end
   end

   # Drop velocity cells which are below the sparsity threshold
   findex_ = vcellf .≥ flimit
   fselect = vcellf[findex_]
   Vselect = V[findex_]

   if slicetype ∈ (:xy, :yz, :xz)
      v1select = [v[dir1] for v in Vselect]
      v2select = [v[dir2] for v in Vselect]
      vnormal = [v[dir3] for v in Vselect]

      strx, stry = getindex(["vx", "vy", "vz"], [dir1, dir2])
   elseif slicetype ∈ (:bperp, :bpar1, :bpar2)
      v1select = zeros(eltype(Vselect[1]), length(Vselect))
      v2select = similar(v1select)
      vnormal = similar(v1select)

      for iv in eachindex(Vselect)
         v1select[iv], v2select[iv], vnormal[iv] = Rinv * Vselect[iv]
      end

      if slicetype == :bperp
         strx = L"$v_{B \times V}$"
         stry = L"$v_{B \times (B \times V)}$"
      elseif slicetype == :bpar2
         strx = L"$v_{B}$"
         stry = L"$v_{B \times V}$"
      elseif slicetype == :bpar1
         strx = L"$v_{B \times (B \times V)}$"
         stry = L"$v_{B}$"
      end
   end

   let str = " [$unitv]"
      strx *= str
      stry *= str
   end

   if vslicethick < 0 # Set a proper thickness
      vslicethick =
         if any(ŝ .== 1.0) # Assure that the slice cut through at least 1 vcell
            cellsize
         else # Assume cubic vspace grid, add extra space
            cellsize*(√3+0.05)
      end
   end

   # Weights using particle flux or phase-space density
   fweight = weight == :flux ? fselect*norm([v1select, v2select, vnormal]) : fselect

   # Select cells within the slice area
   if vslicethick > 0.0
      v1, v2, fweight = let ind_ = @. (abs(vnormal) ≤ 0.5*vslicethick)
         v1select[ind_], v2select[ind_], fweight[ind_]
      end
   else
      v1, v2 = v1select, v2select
   end

   v1 ./= unitvfactor
   v2 ./= unitvfactor

   str_title = @sprintf "t = %4.1fs" meta.time

   if verbose
      @info "Original coordinates : $location"
      @info "Original cell        : $(getcellcoordinates(meta, cidReq))"
      @info "Nearest cell with VDF: $(getcellcoordinates(meta, cidNearest))"
      @info "CellID: $cidNearest"

      if center == :bulk
         @info "Transforming to plasma frame, travelling at speed $Vbulk"
      elseif center == :peak
         @info "Transforming to peak f-value frame, travelling at speed $Vpeak"
      end

      @info "Using VDF threshold value of $flimit."

      if vslicethick > 0
         @info "Performing slice with a counting thickness of $vslicethick"
      else
         @info "Projecting total VDF to a single plane"
      end
   end

   v1, v2, r1, r2, fweight, strx, stry, str_title
end


text:
global const ev2K = 1.160451812e4
global const K2eV = 1.0/1.160451812e4
global const kB = 1.38066e-16

"""
    part_weight_one(N::Integer)

Equivalent to no weighting. Returns an Array of ones.
"""
function part_weight_one(N::Integer)
    return ones(N)
end

"""
    part_weight_physical(N::Integer, par::mappingParameters)

Physical weighting function in units of [cm/pix].
"""
function part_weight_physical(N::Integer, par::mappingParameters, x_cgs::Real=3.085678e21)
    return ones(N) .* par.pixelSideLength .* x_cgs
end

"""
    part_weight_emission(rho::Array{<:Real}, T_K::Array{<:Real})

Emission weighted mapping. Takes density in internal untis and temperature in K and computes weights.
"""
function part_weight_emission(rho::Array{<:Real}, T_K::Array{<:Real})
    return @. rho^2 * √(T_K)
end

"""
    part_weight_spectroscopic(rho::Array{<:Real}, T_K::Array{<:Real})

Spectroscopic weighted mapping from Mazotta+ 04. Takes density and temperature and computes weights.
"""
function part_weight_spectroscopic(rho::Array{<:Real}, T_K::Array{<:Real})
    return @. rho^2 * T_K^(0.75 - 1.5)
end

"""
    part_weight_XrayBand(T_K::Array{<:Real}, Emin::Real, Emax::Real)

Computes Xray weighted emission of a defined energy band. Emin and Emax are energies in eV.
"""
function part_weight_XrayBand(T_K::Array{<:Real}, Emin::Real=5.0e4, Emax::Real=1.0e10)
    # convert Kelvin to eV
    T_eV = T_K .* cgs2eV

    @. exp( -Emin / ( kB * T_eV )) - exp( -Emax / ( kB * T_eV ))
end


text:
try_counters = MPCCCounters()

@test getfield(try_counters, :neval_obj)     == 0
@test getfield(try_counters, :neval_grad)    == 0
@test getfield(try_counters, :neval_cons)    == 0
@test getfield(try_counters, :neval_consG)   == 0
@test getfield(try_counters, :neval_consH)   == 0
@test getfield(try_counters, :neval_jcon)    == 0
@test getfield(try_counters, :neval_jgrad)   == 0
@test getfield(try_counters, :neval_jac)     == 0
@test getfield(try_counters, :neval_jacG)    == 0
@test getfield(try_counters, :neval_jacH)    == 0
@test getfield(try_counters, :neval_jprod)   == 0
@test getfield(try_counters, :neval_jGprod)  == 0
@test getfield(try_counters, :neval_jHprod)  == 0
@test getfield(try_counters, :neval_jtprod)  == 0
@test getfield(try_counters, :neval_jGtprod) == 0
@test getfield(try_counters, :neval_jHtprod) == 0
@test getfield(try_counters, :neval_hess)    == 0
@test getfield(try_counters, :neval_hessG)   == 0
@test getfield(try_counters, :neval_hessH)   == 0
@test getfield(try_counters, :neval_hprod)   == 0
@test getfield(try_counters, :neval_hGprod)  == 0
@test getfield(try_counters, :neval_hHprod)  == 0
@test getfield(try_counters, :neval_jhprod)  == 0

@test sum_counters(try_counters) == 0

reset!(try_counters)
@test sum_counters(try_counters) == 0


text:
# interfacing with LinearAlgebra standard library

#
# BLAS 1
#

LinearAlgebra.rmul!(x::StridedCuArray{<:CublasFloat}, k::Number) =
  scal!(length(x), k, x)

# Work around ambiguity with GPUArrays wrapper
LinearAlgebra.rmul!(x::DenseCuArray{<:CublasFloat}, k::Real) =
  invoke(rmul!, Tuple{typeof(x), Number}, x, k)

function LinearAlgebra.dot(x::StridedCuArray{T}, y::StridedCuArray{T}) where T<:Union{Float16, CublasReal}
    n = length(x)
    n==length(y) || throw(DimensionMismatch("dot product arguments have lengths $(length(x)) and $(length(y))"))
    dot(n, x, y)
end

function LinearAlgebra.dot(x::StridedCuArray{T}, y::StridedCuArray{T}) where T<:Union{ComplexF16, CublasComplex}
    n = length(x)
    n==length(y) || throw(DimensionMismatch("dot product arguments have lengths $(length(x)) and $(length(y))"))
    dotc(n, x, y)
end

# generic fallback
function LinearAlgebra.dot(x::AnyCuArray{T1}, y::AnyCuArray{T2}) where {T1,T2}
    n = length(x)
    n==length(y) || throw(DimensionMismatch("dot product arguments have lengths $(length(x)) and $(length(y))"))

    # custom kernel using simple linear indexing and atomic additions,
    # resulting in about 10% speed-up compared to a simple mapreduce.
    function kernel(x, y, res::AbstractArray{T}, shuffle) where {T}
        local_val = zero(T)

        # grid-stride loop
        i = threadIdx().x + (blockIdx().x - 1i32)*blockDim().x
        while i <= length(x)
            @inbounds local_val += LinearAlgebra.dot(x[i], y[i])
            i += blockDim().x * gridDim().x
        end

        val = CUDA.reduce_block(+, local_val, zero(T), shuffle)
        if threadIdx().x == 1i32
            # NOTE: introduces nondeterminism
            @inbounds CUDA.@atomic res[] += val
        end

        return
    end

    dev = device()
    let T = promote_type(T1, T2)
        # only use the above kernel if we don't care about determinism
        # and if atomic operations are supported on these inputs
        atomic = if capability(device()) >= v"7.0"
            T <: Union{Int16, Int32, Int64, Float16, Float32, Float64}
        else
            T <: Union{Int32, Int64, Float32, Float64}
        end
        if CUDA.math_mode() == CUDA.PEDANTIC_MATH || !atomic
            return mapreduce((x,y)->LinearAlgebra.dot(x, y), +, x, y; init=zero(T))
        end

        res = CUDA.zeros(T, 1)

        # be conservative about using shuffle instructions
        shuffle = T <: Union{Bool,
                             UInt8, UInt16, UInt32, UInt64, UInt128,
                             Int8, Int16, Int32, Int64, Int128,
                             Float16, Float32, Float64,
                             ComplexF16, ComplexF32, ComplexF64}

        # how many threads do we want?
        # reduce_block(shuffle=true) requires the block to consist of full warps.
        wanted_threads = shuffle ? nextwarp(dev, n) : n
        function compute_threads(max_threads)
            if wanted_threads > max_threads
                shuffle ? prevwarp(dev, max_threads) : max_threads
            else
                wanted_threads
            end
        end

        # how many threads can we launch?
        kernel = @cuda launch=false kernel(x, y, res, Val(shuffle))
        compute_shmem(threads) = shuffle ? 0 : threads*sizeof(T)
        config = launch_configuration(kernel.fun; shmem=compute_shmem∘compute_threads)
        threads = compute_threads(config.threads)
        blocks = min(config.blocks, cld(n, config.blocks))
        shmem = compute_shmem(threads)
        kernel(x, y, res, Val(shuffle); threads, blocks, shmem)

        CUDA.@allowscalar res[]
    end
end

function LinearAlgebra.:(*)(transx::Transpose{<:Any,<:StridedCuVector{T}}, y::StridedCuVector{T}) where T<:Union{ComplexF16, CublasComplex}
    x = transx.parent
    n = length(x)
    n==length(y) || throw(DimensionMismatch("dot product arguments have lengths $(length(x)) and $(length(y))"))
    return dotu(n, x, y)
end

function LinearAlgebra.norm(x::DenseCuArray{<:Union{Float16, ComplexF16, CublasFloat}}, p::Real=2)
    if p == 2
        return nrm2(x)
    else
        return invoke(norm, Tuple{AbstractGPUArray, Real}, x, p)
    end
end

LinearAlgebra.BLAS.asum(x::StridedCuArray{<:CublasFloat}) = asum(length(x), x)

function LinearAlgebra.axpy!(alpha::Number, x::StridedCuArray{T}, y::StridedCuArray{T}) where T<:Union{Float16, ComplexF16, CublasFloat}
    length(x)==length(y) || throw(DimensionMismatch("axpy arguments have lengths $(length(x)) and $(length(y))"))
    axpy!(length(x), alpha, x, y)
end

function LinearAlgebra.axpby!(alpha::Number, x::StridedCuArray{T}, beta::Number, y::StridedCuArray{T}) where T<:Union{Float16, ComplexF16, CublasFloat}
    length(x)==length(y) || throw(DimensionMismatch("axpby arguments have lengths $(length(x)) and $(length(y))"))
    axpby!(length(x), alpha, x, beta, y)
end

function LinearAlgebra.rotate!(x::StridedCuArray{T}, y::StridedCuArray{T}, c, s) where T<:CublasFloat
    nx = length(x)
    ny = length(y)
    nx==ny || throw(DimensionMismatch("rotate arguments have lengths $nx and $ny"))
    rot!(nx, x, y, c, s)
end

function LinearAlgebra.reflect!(x::StridedCuArray{T}, y::StridedCuArray{T}, c, s) where T<:CublasFloat
    nx = length(x)
    ny = length(y)
    nx==ny || throw(DimensionMismatch("reflect arguments have lengths $nx and $ny"))
    rot!(nx, x, y, c, s)
    scal!(ny, -real(one(T)), y)
    x, y
end

#
# BLAS 2
#

# GEMV

function gemv_dispatch!(Y::CuVector, A, B, alpha::Number=true, beta::Number=false)
    mA, nA = size(A)

    if nA != length(B)
        throw(DimensionMismatch("second dimension of A, $nA, does not match length of B, $(length(B))"))
    end

    if mA != length(Y)
        throw(DimensionMismatch("first dimension of A, $mA, does not match length of Y, $(length(Y))"))
    end

    if mA == 0
        return Y
    end

    if nA == 0
        return rmul!(Y, 0)
    end

    tA, dA = if A isa Transpose
        'T', parent(A)
    elseif A isa Adjoint
        'C', parent(A)
    else
        'N', A
    end

    T = eltype(Y)
    if T <: CublasFloat && A isa StridedCuArray{T} && B isa StridedCuArray{T}
        gemv!(tA, alpha, dA, B, beta, Y)
    else
        gemm_dispatch!(Y, A, B, alpha, beta)
    end
end

for NT in (Number, Real)
    # NOTE: alpha/beta also ::Real to avoid ambiguities with certain Base methods
    @eval begin
        LinearAlgebra.mul!(Y::CuVector, A::StridedCuMatrix, B::StridedCuVector, a::$NT, b::$NT) =
            gemv_dispatch!(Y, A, B, a, b)
        LinearAlgebra.mul!(Y::CuVector, A::Transpose{<:Any, <:StridedCuVecOrMat}, B::StridedCuVector, a::$NT, b::$NT) =
            gemv_dispatch!(Y, A, B, a, b)
        LinearAlgebra.mul!(Y::CuVector, A::Adjoint{<:Any, <:StridedCuVecOrMat}, B::StridedCuVector, a::$NT, b::$NT) =
            gemv_dispatch!(Y, A, B, a, b)
    end
end

# triangular

## direct multiplication/division
for (t, uploc, isunitc) in ((:LowerTriangular, 'L', 'N'),
                            (:UnitLowerTriangular, 'L', 'U'),
                            (:UpperTriangular, 'U', 'N'),
                            (:UnitUpperTriangular, 'U', 'U'))
    @eval begin
        # Multiplication
        LinearAlgebra.lmul!(A::$t{T,<:DenseCuMatrix},
                            b::StridedCuVector{T}) where {T<:CublasFloat} =
            CUBLAS.trmv!($uploc, 'N', $isunitc, parent(A), b)

        # Left division
        LinearAlgebra.ldiv!(A::$t{T,<:DenseCuMatrix},
                            B::StridedCuVector{T}) where {T<:CublasFloat} =
            CUBLAS.trsv!($uploc, 'N', $isunitc, parent(A), B)
    end
end

## adjoint/transpose multiplication ('uploc' reversed)
for (t, uploc, isunitc) in ((:LowerTriangular, 'U', 'N'),
                            (:UnitLowerTriangular, 'U', 'U'),
                            (:UpperTriangular, 'L', 'N'),
                            (:UnitUpperTriangular, 'L', 'U'))
    @eval begin
        # Multiplication
        LinearAlgebra.lmul!(A::$t{<:Any,<:Transpose{T,<:DenseCuMatrix}},
                            b::DenseCuVector{T}) where {T<:CublasFloat} =
            CUBLAS.trmv!($uploc, 'T', $isunitc, parent(parent(A)), b)
        LinearAlgebra.lmul!(A::$t{<:Any,<:Adjoint{T,<:DenseCuMatrix}},
                            b::DenseCuVector{T}) where {T<:CublasReal} =
            CUBLAS.trmv!($uploc, 'T', $isunitc, parent(parent(A)), b)
        LinearAlgebra.lmul!(A::$t{<:Any,<:Adjoint{T,<:DenseCuMatrix}},
                            b::DenseCuVector{T}) where {T<:CublasComplex} =
            CUBLAS.trmv!($uploc, 'C', $isunitc, parent(parent(A)), b)

        # Left division
        LinearAlgebra.ldiv!(A::$t{<:Any,<:Transpose{T,<:DenseCuMatrix}},
                            B::StridedCuVector{T}) where {T<:CublasFloat} =
            CUBLAS.trsv!($uploc, 'T', $isunitc, parent(parent(A)), B)
        LinearAlgebra.ldiv!(A::$t{<:Any,<:Adjoint{T,<:DenseCuMatrix}},
                            B::StridedCuVector{T}) where {T<:CublasReal} =
            CUBLAS.trsv!($uploc, 'T', $isunitc, parent(parent(A)), B)
        LinearAlgebra.ldiv!(A::$t{<:Any,<:Adjoint{T,<:DenseCuMatrix}},
                            B::StridedCuVector{T}) where {T<:CublasComplex} =
            CUBLAS.trsv!($uploc, 'C', $isunitc, parent(parent(A)), B)
    end
end

#
# BLAS 3
#

# GEMM

function gemm_dispatch!(C::CuVecOrMat, A, B, alpha::Number=true, beta::Number=false)
    if ndims(A) > 2
        throw(ArgumentError("A has more than 2 dimensions"))
    elseif ndims(B) > 2
        throw(ArgumentError("B has more than 2 dimensions"))
    end
    mA, nA = size(A,1), size(A,2)
    mB, nB = size(B,1), size(B,2)

    if nA != mB
        throw(DimensionMismatch("A has dimensions ($mA,$nA) but B has dimensions ($mB,$nB)"))
    end

    if C === A || B === C
        throw(ArgumentError("output matrix must not be aliased with input matrix"))
    end

    if mA == 0 || nA == 0 || nB == 0
        if size(C) != (mA, nB)
            throw(DimensionMismatch("C has dimensions $(size(C)), should have ($mA,$nB)"))
        end
        return LinearAlgebra.rmul!(C, 0)
    end

    tA, dA = if A isa Transpose
        'T', parent(A)
    elseif A isa Adjoint
        'C', parent(A)
    else
        'N', A
    end

    tB, dB = if B isa Transpose
        'T', parent(B)
    elseif B isa Adjoint
        'C', parent(B)
    else
        'N', B
    end

    T = eltype(C)
    if dA isa DenseCuArray && dB isa DenseCuArray &&
       gemmExComputeType(eltype(A), eltype(B), eltype(C), mA, nA, nB) !== nothing
        gemmEx!(tA, tB, alpha, dA, dB, beta, C)
    elseif T <: CublasFloat && dA isa DenseCuArray{T} && dB isa DenseCuArray{T}
        gemm!(tA, tB, alpha, dA, dB, beta, C)
    else
        GPUArrays.generic_matmatmul!(C, A, B, alpha, beta)
    end
end

for NT in (Number, Real)
    # NOTE: alpha/beta also ::Real to avoid ambiguities with certain Base methods
    @eval begin
        LinearAlgebra.mul!(C::CuMatrix, A::StridedCuVecOrMat, B::StridedCuVecOrMat, a::$NT, b::$NT) =
            gemm_dispatch!(C, A, B, a, b)

        LinearAlgebra.mul!(C::CuMatrix, A::Transpose{<:Any, <:StridedCuVecOrMat}, B::StridedCuMatrix, a::$NT, b::$NT) =
            gemm_dispatch!(C, A, B, a, b)
        LinearAlgebra.mul!(C::CuMatrix, A::StridedCuMatrix, B::Transpose{<:Any, <:StridedCuVecOrMat}, a::$NT, b::$NT) =
            gemm_dispatch!(C, A, B, a, b)
        LinearAlgebra.mul!(C::CuMatrix, A::Transpose{<:Any, <:StridedCuVecOrMat}, B::Transpose{<:Any, <:StridedCuVecOrMat}, a::$NT, b::$NT) =
            gemm_dispatch!(C, A, B, a, b)

        LinearAlgebra.mul!(C::CuMatrix, A::Adjoint{<:Any, <:StridedCuVecOrMat}, B::StridedCuMatrix, a::$NT, b::$NT) =
            gemm_dispatch!(C, A, B, a, b)
        LinearAlgebra.mul!(C::CuMatrix, A::StridedCuMatrix, B::Adjoint{<:Any, <:StridedCuVecOrMat}, a::$NT, b::$NT) =
            gemm_dispatch!(C, A, B, a, b)
        LinearAlgebra.mul!(C::CuMatrix, A::Adjoint{<:Any, <:StridedCuVecOrMat}, B::Adjoint{<:Any, <:StridedCuVecOrMat}, a::$NT, b::$NT) =
            gemm_dispatch!(C, A, B, a, b)

        LinearAlgebra.mul!(C::CuMatrix, A::Transpose{<:Any, <:StridedCuVecOrMat}, B::Adjoint{<:Any, <:StridedCuVecOrMat}, a::$NT, b::$NT) =
            gemm_dispatch!(C, A, B, a, b)
        LinearAlgebra.mul!(C::CuMatrix, A::Adjoint{<:Any, <:StridedCuVecOrMat}, B::Transpose{<:Any, <:StridedCuVecOrMat}, a::$NT, b::$NT) =
            gemm_dispatch!(C, A, B, a, b)
    end
end

# triangular

## direct multiplication/division
for (t, uploc, isunitc) in ((:LowerTriangular, 'L', 'N'),
                            (:UnitLowerTriangular, 'L', 'U'),
                            (:UpperTriangular, 'U', 'N'),
                            (:UnitUpperTriangular, 'U', 'U'))
    @eval begin
        # Multiplication
        LinearAlgebra.lmul!(A::$t{T,<:DenseCuMatrix},
                            B::DenseCuMatrix{T}) where {T<:CublasFloat} =
            CUBLAS.trmm!('L', $uploc, 'N', $isunitc, one(T), parent(A), B, B)
        LinearAlgebra.rmul!(A::DenseCuMatrix{T},
                            B::$t{T,<:DenseCuMatrix}) where {T<:CublasFloat} =
            CUBLAS.trmm!('R', $uploc, 'N', $isunitc, one(T), parent(B), A, A)

        # optimization: Base.mul! uses lmul!/rmul! with a copy (because of BLAS)
        LinearAlgebra.mul!(X::DenseCuMatrix{T}, A::$t{T,<:DenseCuMatrix},
                           B::DenseCuMatrix{T}) where {T<:CublasFloat} =
            CUBLAS.trmm!('L', $uploc, 'N', $isunitc, one(T), parent(A), B, X)
        LinearAlgebra.mul!(X::DenseCuMatrix{T}, A::DenseCuMatrix{T},
                           B::$t{T,<:DenseCuMatrix}) where {T<:CublasFloat} =
            CUBLAS.trmm!('R', $uploc, 'N', $isunitc, one(T), parent(B), A, X)

        # Left division
        LinearAlgebra.ldiv!(A::$t{T,<:DenseCuMatrix},
                            B::DenseCuMatrix{T}) where {T<:CublasFloat} =
            CUBLAS.trsm!('L', $uploc, 'N', $isunitc, one(T), parent(A), B)

        # Right division
        LinearAlgebra.rdiv!(A::DenseCuMatrix{T},
                            B::$t{T,<:DenseCuMatrix}) where {T<:CublasFloat} =
            CUBLAS.trsm!('R', $uploc, 'N', $isunitc, one(T), parent(B), A)

        # Matrix inverse
        function LinearAlgebra.inv(x::$t{T, <:CuMatrix{T}}) where T<:CublasFloat
            out = CuArray{T}(I(size(x,1)))
            $t(LinearAlgebra.ldiv!(x, out))
        end
    end
end

## adjoint/transpose multiplication ('uploc' reversed)
for (t, uploc, isunitc) in ((:LowerTriangular, 'U', 'N'),
                            (:UnitLowerTriangular, 'U', 'U'),
                            (:UpperTriangular, 'L', 'N'),
                            (:UnitUpperTriangular, 'L', 'U'))
    @eval begin
        # Multiplication
        LinearAlgebra.lmul!(A::$t{<:Any,<:Transpose{T,<:DenseCuMatrix}},
                            B::DenseCuMatrix{T}) where {T<:CublasFloat} =
            CUBLAS.trmm!('L', $uploc, 'T', $isunitc, one(T), parent(parent(A)), B, B)
        LinearAlgebra.lmul!(A::$t{<:Any,<:Adjoint{T,<:DenseCuMatrix}},
                            B::DenseCuMatrix{T}) where {T<:CublasComplex} =
            CUBLAS.trmm!('L', $uploc, 'C', $isunitc, one(T), parent(parent(A)), B, B)
        LinearAlgebra.lmul!(A::$t{<:Any,<:Adjoint{T,<:DenseCuMatrix}},
                            B::DenseCuMatrix{T}) where {T<:CublasReal} =
            CUBLAS.trmm!('L', $uploc, 'T', $isunitc, one(T), parent(parent(A)), B, B)

        LinearAlgebra.rmul!(A::DenseCuMatrix{T},
                            B::$t{<:Any,<:Transpose{T,<:DenseCuMatrix}}) where {T<:CublasFloat} =
            CUBLAS.trmm!('R', $uploc, 'T', $isunitc, one(T), parent(parent(B)), A, A)
        LinearAlgebra.rmul!(A::DenseCuMatrix{T},
                            B::$t{<:Any,<:Adjoint{T,<:DenseCuMatrix}}) where {T<:CublasComplex} =
            CUBLAS.trmm!('R', $uploc, 'C', $isunitc, one(T), parent(parent(B)), A, A)
        LinearAlgebra.rmul!(A::DenseCuMatrix{T},
                            B::$t{<:Any,<:Adjoint{T,<:DenseCuMatrix}}) where {T<:CublasReal} =
            CUBLAS.trmm!('R', $uploc, 'T', $isunitc, one(T), parent(parent(B)), A, A)

        # optimization: Base.mul! uses lmul!/rmul! with a copy (because of BLAS)
        LinearAlgebra.mul!(X::DenseCuMatrix{T}, A::$t{<:Any,<:Transpose{T,<:DenseCuMatrix}},
                           B::DenseCuMatrix{T}) where {T<:CublasFloat} =
            CUBLAS.trmm!('L', $uploc, 'T', $isunitc, one(T), parent(parent(A)), B, X)
        LinearAlgebra.mul!(X::DenseCuMatrix{T}, A::$t{<:Any,<:Adjoint{T,<:DenseCuMatrix}},
                           B::DenseCuMatrix{T}) where {T<:CublasComplex} =
            CUBLAS.trmm!('L', $uploc, 'C', $isunitc, one(T), parent(parent(A)), B, X)
        LinearAlgebra.mul!(X::DenseCuMatrix{T}, A::$t{<:Any,<:Adjoint{T,<:DenseCuMatrix}},
                           B::DenseCuMatrix{T}) where {T<:CublasReal} =
            CUBLAS.trmm!('L', $uploc, 'T', $isunitc, one(T), parent(parent(A)), B, X)
        LinearAlgebra.mul!(X::DenseCuMatrix{T}, A::DenseCuMatrix{T},
                           B::$t{<:Any,<:Transpose{T,<:DenseCuMatrix}}) where {T<:CublasFloat} =
            CUBLAS.trmm!('R', $uploc, 'T', $isunitc, one(T), parent(parent(B)), A, X)
        LinearAlgebra.mul!(X::DenseCuMatrix{T}, A::DenseCuMatrix{T},
                           B::$t{<:Any,<:Adjoint{T,<:DenseCuMatrix}}) where {T<:CublasComplex} =
            CUBLAS.trmm!('R', $uploc, 'C', $isunitc, one(T), parent(parent(B)), A, X)
        LinearAlgebra.mul!(X::DenseCuMatrix{T}, A::DenseCuMatrix{T},
                           B::$t{<:Any,<:Adjoint{T,<:DenseCuMatrix}}) where {T<:CublasReal} =
            CUBLAS.trmm!('R', $uploc, 'T', $isunitc, one(T), parent(parent(B)), A, X)

        # Left division
        LinearAlgebra.ldiv!(A::$t{<:Any,<:Transpose{T,<:DenseCuMatrix}},
                            B::DenseCuMatrix{T}) where {T<:CublasFloat} =
            CUBLAS.trsm!('L', $uploc, 'T', $isunitc, one(T), parent(parent(A)), B)
        LinearAlgebra.ldiv!(A::$t{<:Any,<:Adjoint{T,<:DenseCuMatrix}},
                            B::DenseCuMatrix{T}) where {T<:CublasReal} =
            CUBLAS.trsm!('L', $uploc, 'T', $isunitc, one(T), parent(parent(A)), B)
        LinearAlgebra.ldiv!(A::$t{<:Any,<:Adjoint{T,<:DenseCuMatrix}},
                            B::DenseCuMatrix{T}) where {T<:CublasComplex} =
            CUBLAS.trsm!('L', $uploc, 'C', $isunitc, one(T), parent(parent(A)), B)

        # Right division
        LinearAlgebra.rdiv!(A::DenseCuMatrix{T},
                            B::$t{<:Any,<:Transpose{T,<:DenseCuMatrix}}) where {T<:CublasFloat} =
            CUBLAS.trsm!('R', $uploc, 'T', $isunitc, one(T), parent(parent(B)), A)
        LinearAlgebra.rdiv!(A::DenseCuMatrix{T},
                            B::$t{<:Any,<:Adjoint{T,<:DenseCuMatrix}}) where {T<:CublasReal} =
            CUBLAS.trsm!('R', $uploc, 'T', $isunitc, one(T), parent(parent(B)), A)
        LinearAlgebra.rdiv!(A::DenseCuMatrix{T},
                            B::$t{<:Any,<:Adjoint{T,<:DenseCuMatrix}}) where {T<:CublasComplex} =
            CUBLAS.trsm!('R', $uploc, 'C', $isunitc, one(T), parent(parent(B)), A)
    end
end

# symmetric mul!
# level 2
@inline function LinearAlgebra.mul!(y::CuVector{T}, A::Hermitian{T,<:CuMatrix}, x::CuVector{T},
             α::Number, β::Number) where {T<:CublasReal}
    alpha, beta = promote(α, β, zero(T))
    if alpha isa Union{Bool,T} && beta isa Union{Bool,T}
        return CUBLAS.symv!(A.uplo, alpha, A.data, x, beta, y)
    else
        error("only supports BLAS type, got $T")
    end
end

@inline function LinearAlgebra.mul!(y::CuVector{T}, A::Hermitian{T,<:CuMatrix}, x::CuVector{T},
             α::Number, β::Number) where {T<:CublasComplex}
    alpha, beta = promote(α, β, zero(T))
    if alpha isa Union{Bool,T} && beta isa Union{Bool,T}
        return CUBLAS.hemv!(A.uplo, alpha, A.data, x, beta, y)
    else
        error("only supports BLAS type, got $T")
    end
end

# level 3

@inline function LinearAlgebra.mul!(C::CuMatrix{T}, A::Hermitian{T,<:CuMatrix}, B::CuMatrix{T},
             α::Number, β::Number) where {T<:CublasReal}
    alpha, beta = promote(α, β, zero(T))
    if alpha isa Union{Bool,T} && beta isa Union{Bool,T}
        return CUBLAS.symm!('L', A.uplo, alpha, A.data, B, beta, C)
    else
        error("only supports BLAS type, got $T")
    end
end
@inline function LinearAlgebra.mul!(C::CuMatrix{T}, A::CuMatrix{T}, B::Hermitian{T,<:CuMatrix},
             α::Number, β::Number) where {T<:CublasReal}
    alpha, beta = promote(α, β, zero(T))
    if alpha isa Union{Bool,T} && beta isa Union{Bool,T}
        return CUBLAS.symm!('R', B.uplo, alpha, B.data, A, beta, C)
    else
        error("only supports BLAS type, got $T")
    end
end
@inline function LinearAlgebra.mul!(C::CuMatrix{T}, A::Hermitian{T,<:CuMatrix}, B::CuMatrix{T},
             α::Number, β::Number) where {T<:CublasComplex}
    alpha, beta = promote(α, β, zero(T))
    if alpha isa Union{Bool,T} && beta isa Union{Bool,T}
        return CUBLAS.hemm!('L', A.uplo, alpha, A.data, B, beta, C)
    else
        error("only supports BLAS type, got $T")
    end
end
@inline function LinearAlgebra.mul!(C::CuMatrix{T}, A::CuMatrix{T}, B::Hermitian{T,<:CuMatrix},
             α::Number, β::Number) where {T<:CublasComplex}
    alpha, beta = promote(α, β, zero(T))
    if alpha isa Union{Bool,T} && beta isa Union{Bool,T}
        return CUBLAS.hemm!('R', B.uplo, alpha, B.data, A, beta, C)
    else
        error("only supports BLAS type, got $T")
    end
end


text:
###########Importing Libraries and Defining Structs############
using  LinearAlgebra, SparseArrays, StaticArrays, BenchmarkTools, Base, Plots

mutable struct Qubit #type for holding qubit information
    bitpos::Int
    vec::Vector{ComplexF64}
end

mutable struct Circuit
    qreg::Array{Qubit,1}
    gates::Vector{SparseMatrixCSC{ComplexF64,Int64}}
end

function Circuit(nqubits::Int)
    G = Vector{SparseMatrixCSC{ComplexF64,Int64}}()
    return Circuit(makeregister(nqubits),G)
end

function makeregister(nqubits::Int)
    #function for generating a quantum register of qubits
    zerovec = [1.0+0.0*im, 0.0+0.0*im]
    return [Qubit(i, copy(zerovec)) for i ∈ 1:nqubits]
end

########## Running the Circuit ###########
function Run_circ(C::Circuit)

    maxiter = length(C.gates)
    maxbit = length(C.qreg)

    M = copy(C.gates[1])
    ket = copy(C.qreg[1].vec)
    for i ∈ 2:maxbit
    ket = kron(C.qreg[i].vec,ket)
    end

    for i ∈ 2:maxiter
        M .= C.gates[i]*M
    end

    return M*ket
end

function Run_circ(C::Circuit)

    maxiter = length(C.gates)
    maxbit = length(C.qreg)

    M = copy(C.gates[1])
    ket = copy(C.qreg[1].vec)
    @inbounds for i ∈ 2:maxbit
    ket = kron(C.qreg[i].vec,ket)
    end

    @inbounds for i ∈ 2:maxiter
        M .= C.gates[i]*M
    end

    return M*ket
end

function shots(C::Circuit,nshots::Int)

    finalket = Run_circ(C)
    ntrials = rand(nshots)
    ensemble = zeros(Float64,length(finalket))
    trials!(ensemble, finalket, ntrials)
end

#function shots(C::Circuit,nqubits::AbstractVector{Int},nshots::Int)
#
#    finalket = Run_circ(C)
#    ntrials = rand(nshots)
#    ensemble = zeros(Int,length(finalket))
#    trials!(ensemble, finalket, ntrials)
#end

function trials!(ensemble::AbstractVector{Float64}, ket::AbstractVector{ComplexF64}, probs::AbstractVector{Float64})
    nshots = length(probs)

    @inbounds for j ∈ 1:nshots
        local p = probs[j]
        local sumold = 0.0
        local sumnew = 0.0

        @inbounds for i ∈ 1:length(ket)
            sumnew = sumold + real(conj(ket[i])*ket[i])
            if (p<=sumnew)
                ensemble[i] += 1.0
                break
            end
            sumold = sumnew
        end
    end
    return ensemble
end

#function trials!(ensemble::AbstractVector{Float64}, ket::AbstractVector{ComplexF64}, nqubits::AnstractVector{Int}, probs::AbstractVector{Float64})
#    nshots = length(probs)
#
#    @inbounds for j ∈ 1:nshots
#        local p = probs[j]
#        local sumold = 0.0
#        local sumnew = 0.0
#
#        @inbounds for i ∈ 1:length(ket)
#            sumnew = sumold + real(conj(ket[i])*ket[i])
#            if (p<=sumnew)
#                ensemble[i] += 1.0
#                break
#            end
#            sumold = sumnew
#        end
#    end
#    return ensemble
#end

function graph!(data::AbstractVector)
    pyplot()
    data .= data./sum(data)
    max = length(data)-1
    nqubits = convert(Int,log(2,max+1))
    xs = [string(i,base=2,pad=nqubits) for i ∈ 0:max]
    pl = bar(xs, data, leg=false, title="Simulation of Circuit", xlabel="qubit readout",
    ylabel="Probabilities", xtickfontrotation = 70, dpi = 300, xticks = :all,
    yticks = 0.0:0.05:1.0, grid = :y, minorticks = 2, gridtyle = :dash)
    gui()
end


text:
module RipsererBenchmarks

using BenchmarkTools
using CSV
using DataFrames
using Eirene
using Logging
using ProgressMeter
using Ripserer
using SparseArrays
using TerminalLoggers

export Benchmark, ripser_benchmark, ripserer_benchmark, eirene_benchmark,
    nv, load_data, setup_all, grind_ripser, grind_ripserer, run_and_save, save_csv

include("utils.jl")
include("benchmarks.jl")
include("ripserer.jl")
include("ripser.jl")
include("eirene.jl")
include("benchmarksuites.jl")
include("valgrind.jl")
include("report.jl")

function run_and_save(selected=nothing)
    get_ripsers()

    start_time = time_ns()
    if isnothing(selected)
        suite = setup_all()
        file = "all_results.json"
        @info "running all"
    else
        suite = setup_all()[selected]
        file = "$selected.json"
        @info "running $selected"
    end
    result = with_logger(TerminalLogger()) do
        run(suite)
    end
    BenchmarkTools.save(joinpath(@__DIR__, "..", "results", file), result)
    elapsed = (time_ns() - start_time) / 1e9
    @info "Took $(ProgressMeter.durationstring(elapsed))"
    return result
end

end


text:
######################################################################
# ergodicity.jl
# handles stuff needed for ergodicity
######################################################################

export ErgodicManagerSE2

"""
`ErgodicManagerSE2(L::Float64, K::Int, bins::Int)`

`ErgodicManagerSE2(example_name::String; K::Int=5, bins::Int=100)`

Valid `example_name` entries are:

* "single gaussian"
* "double gaussian"
"""
mutable struct ErgodicManagerSE2 <: ErgodicManager
    domain::Domain
    M::Int
    N::Int
    P::Int
    phi::Array{Float64,3}				# spatial distribution
    phik::Array{Complex{Float64},3}		# spatial Fourier coefficients
    Lambda::Array{Float64,3}			# constants

    function ErgodicManagerSE2(d::Domain, phi::Array{Float64,3}, K::Int=5)
        em = new()
        em.domain = deepcopy(d)
        em.M = K
        em.N = K
        em.P = K
        em.phi = deepcopy(phi)
        em.phik = zeros(em.M+1, em.N+1, em.P+1)
        em.Lambda = zeros(em.M+1, em.N+1, em.P+1)

        Lambda!(em)
        decompose!(em)
        return em
    end

    function ErgodicManagerSE2(d::Domain, K::Int=5)
        em = new()
        em.domain = deepcopy(d)
        em.M = K
        em.N = K
        em.P = K
        em.phik = zeros(em.M+1, em.N+1, em.P+1)
        em.Lambda = zeros(em.M+1, em.N+1, em.P+1)

        Lambda!(em)
        return em
    end
end

# fills the matrix Lambda_{m,n,p}
function Lambda!(em::ErgodicManagerSE2)
    for m = 0:em.M, n = 0:em.N, p = 0:em.P
        den_sqrt = (1.0 + m*m + n*n + p*p)
        em.Lambda[m+1, n+1, p+1] = 1.0 / (den_sqrt * den_sqrt)
    end
end

######################################################################
# Setting and computing phi
######################################################################
function phi!(em::ErgodicManagerSE2, dm::VF, ds::MF)
    # first, generate d
    #d = zeros(em.bins, em.bins, em.bins)
    d = zeros(x_cells(em), y_cells(em), z_cells(em))
    d_sum = 0.0
    for xi = 1:x_cells(em)
		x = x_min(em) + (xi-0.5)*x_size(em)
		println("xi = ", xi)
		for yi = 1:y_cells(em)
			y = y_min(em) + (yi-0.5)*y_size(em)
			for zi = 1:z_cells(em)
				z = z_min(em) + (zi-0.5)*z_size(em)
				d[xi,yi,zi] = my_pdf([x,y,z], dm, ds)
				d_sum += d[xi,yi,zi]
			end
		end
	end
    normalize!(d, em.domain.cell_size)
    em.phi = d
end

######################################################################
# Computing Fourier coefficients
######################################################################
# update the Fourier coefficients based on some distribution
# Here, I assume it is discrete, but I should change this...
# TODO: maybe I should do some bounds checking?
function decompose!(em::ErgodicManagerSE2, d::Array{Float64,3})
    for m = 0:em.M
        #println("m = ", m)
        for n = 0:em.N, p = 0:em.P
            em.phik[m+1,n+1,p+1] = phi_mnp(em, m, n, p, d)
        end
    end
    em.phi = d
end

# iterate over the state space
function phi_mnp(em::ErgodicManagerSE2, m::Int, n::Int, p::Int, d::Array{Float64,3})
	val = 0.0im
	i = float(im)
	for xi = 1:x_cells(em)
		x = x_min(em) + (xi-0.5) * x_size(em)
		for yi = 1:y_cells(em)
			y = y_min(em) + (yi-0.5) * y_size(em)
			r = sqrt(x*x + y*y)
			psi = atan2(y,x)	# atan(y/x)
			bj = besselj(m-n, p*r)
			for zi = 1:z_cells(em)
				z = z_min(em) + (zi-0.5) * z_size(em)
				#val += d[xi,yi,zi] * i^(n-m) * exp(i*(m*psi+(n-m)z)) * bj * em.domain.cell_size
				val += d[xi,yi,zi] * exp(i*(m*psi+(n-m)z)) * bj
			end
		end
	end
	# i^(n-m) and em.domain.cell_size don't depend on x, y, z
	return i^(n-m) * val * em.domain.cell_size
end

function F_mnp(m::Int, n::Int, p::Int, x::Float64, y::Float64, z::Float64)
    # compute psi, r
    r = sqrt(x*x + y*y)
    psi = atan2(y,x)	# atan(y/x)
    i = float(im)
    return i^(n-m) * exp(i*(m*psi + (n-m)z)) * besselj(m-n, p*r)
end

function decompose(em::ErgodicManagerSE2, traj::VVF)
    N = length(traj)-1
    ck = zeros(Complex{Float64}, em.M+1, em.N+1, em.P+1)
    for m = 0:em.M, n = 0:em.N, p = 0:em.P
        fk_sum = 0.0im
        # now loop over time
        for i = 0:N-1
            xi = traj[i + 1]
            fk_sum += F_mnp(m,n,p,xi[1],xi[2],xi[3])
        end
        # TODO: check that this is right
        ck[m+1, n+1, p+1] = fk_sum / N
    end
    return ck
end

# reconstructs from Fourier coefficients in ck
# This comes from last equation in chapter 10.3 of
#  Engineering Applications of Noncommutative Harmonic Analysis
# TODO: This is a first cut. Things that I can/should probably change
#  * probably compute M, N before P
#  * p = 0 seems to add nothing, don't bother there
#  * is the call to real() necessary?
function reconstruct(em::ErgodicManagerSE2, ck::Array{Complex{Float64},3})

    vals = zeros(x_cells(em), y_cells(em), z_cells(em))

    for xi = 1:x_cells(em), yi = 1:y_cells(em), zi=1:z_cells(em)
		x = x_min(em) + (xi-0.5)*x_size(em)
        y = y_min(em) + (yi-0.5)*y_size(em)
        z = z_min(em) + (zi-0.5)*z_size(em)
        for m = 0:em.M
            for n = 0:em.N
                for p = 1:em.P
                    c = ck[m+1,n+1,p+1]
                    f = conj(F_mnp(m,n,p,x,y,z))
                    vals[xi,yi,zi] += real(c * f * p)
                end
            end
        end
	end
	return vals
end


text:
function is_it(u0,H,tol)
    prob = ODEProblem(Eq_of_M,u0,(0.,10.))
    sol=solve(prob, RK4(),reltol=1e-5,abstol=1e-7)
    A=sol[:,:]

    Z=A[1,:]+im *A[2,:]
    W=A[3,:]+im *A[4,:]
    A=zeros(0)
    h=1/(2*H)
    a=(2+h-2*sqrt(h+1))/h;
    N=length(sol.t)
    lin_impulse=im*(1+a )*ones(N)

    z_1pos=.5*(Z+lin_impulse+W)
    z_1neg=.5*(Z-lin_impulse-W)
    z_2pos=.5*(-Z+lin_impulse-W)
    z_2neg=.5*(-Z-lin_impulse+W)

    #1pos and 1neg bound
    d_11=abs.(z_1pos-z_1neg)
    #1pos and 2neg bound
    d_12=abs.(z_1pos-z_2neg)

    delta_11=abs(d_11[1]-d_11[end])
    delta_12=abs(d_12[1]-d_12[end])

    Z=zeros(0)
    W=zeros(0)
    lin_impulse=zeros(0)
    z_1pos=zeros(0)
    z_1neg=zeros(0)
    z_2pos=zeros(0)
    z_2neg=zeros(0)
    d_11=zeros(0)
    d_12=zeros(0)

    return delta_11<tol || delta_12<tol
    # return delta_11, delta_12

end