text:
export BatchedSampler

mutable struct BatchedValSampler{BN, P, S, Sc, Sv, Pv, Gv, Gvv, Gva, LC, Lv, Pa, Pc, Pd} <: AbstractIterativeSampler
    bnet::BN
    Ĉ::P # the operator defining the observable to minimise

    sampler::S
    sampler_cache::Sc
    samples::Sv

    ψvals::Pv
    ∇vals::Gv
    ∇vals_vec::Gvv
    ∇vec_avg::Gva

    accum::LC

    local_vals::Lv

    precond_alg::Pa
    precond_cache::Pc

    observables_sampler

    parallel_cache::Pd
end

function BatchedValSampler(net,
                        sampl,
                        prob,
                        algo=Gradient();
                        batch_sz=2^4,
                        local_batch_sz=batch_sz,
                        par_type=automatic_parallel_type())
    if net isa CachedNet
        throw("Only takes standard network.")
    end

    par_cache      = parallel_execution_cache(par_type)

    bnet           = cached(net, batch_sz)
    v              = state(prob, bnet)
    sampler_cache  = cache(sampl, basis(prob), bnet, par_cache)

    ch_len         = chain_length(sampl, sampler_cache)

    samples        = state(prob, bnet, ch_len)
    ψvals          = similar(trainable_first(bnet), out_type(bnet), 1, batch_sz, ch_len)
    ∇vals, ∇vecs   = grad_cache(bnet, batch_sz, ch_len)
    ∇vec_avg       = tuple([similar(∇vec, size(∇vec, 1)) for ∇vec=∇vecs ]...)

    local_acc      = AccumulatorObsScalar(net, basis(prob), v, local_batch_sz)
    Llocal_vals    = collect(similar(ψvals, size(ψvals)[2:end]...)) #ensure it's on cpu

    precond        = algorithm_cache(algo, prob, net, par_cache)

    obs            = BatchedObsKetSampler(samples, ψvals, local_acc)

    nq = BatchedValSampler(bnet, prob,
            sampl, sampler_cache, samples,
            ψvals, ∇vals, ∇vecs, ∇vec_avg,
            local_acc, Llocal_vals,
            algo, precond,
            obs,
            par_cache)

    return nq
end

function _sample_compute_obs!(is::BatchedValSampler)
    ch_len         = chain_length(is.sampler, is.sampler_cache)
    batch_sz       = size(is.local_vals, 1)

    # If we have gpu, this thing is indexed by value so better
    # to copy it to the cpu first in one go...
    if typeof(is.ψvals) <: GPUArray
        ψvals   = collect(is.ψvals)
        samples = collect(is.samples)
    else
        ψvals   = is.ψvals
        samples = is.samples
    end

    # Compute terms C^{loc} = ⟨σ|Ĉ|ψ⟩/⟨σ|ψ⟩
    for (i,σv)=enumerate(states(is.samples))
        init!(is.accum, σv, ψvals[i])
        accumulate_connections!(is.accum, is.Ĉ, σv)
        is.local_vals[i] = NeuralQuantum.finalize!(is.accum)
    end

    # Perform some simple analysis to estimate error and autocorrelations
    L_stat = stat_analysis(is.local_vals, is.parallel_cache)

    return L_stat
end

function _compute_gradient!(is::BatchedValSampler)
    ch_len         = chain_length(is.sampler, is.sampler_cache)
    batch_sz       = size(is.local_vals, 1)

    # Flatten the batches and the iterations
    ∇vr = reshape(is.∇vals_vec[1], :, ch_len*batch_sz)
    Ĉr  = reshape(is.local_vals, 1, :)

    if typeof(is.ψvals) <: GPUArray
        Ĉr = adapt(typeof(is.ψvals), Ĉr)
    end

    # Compute the gradient
    ∇C  = Ĉr*∇vr'
    ∇C ./= (ch_len*batch_sz) # MPI
    workers_mean!(∇C, is.parallel_cache)

    return ∇C, ∇vr
end

function sample!(is::BatchedValSampler; sample = true)
    # Sample the new configurations
    sample && _sample_state!(is)

    # Compute logψ and ∇logψ
    logψ_and_∇logψ!(is.∇vals, is.ψvals, is.bnet, is.samples)

    # Compute the cost function
    L_stat = _sample_compute_obs!(is)

    # Center the gradient
    _center_gradient!(is)

    # Compute the cost-gradient for optimization
    ∇C, ∇vr = _compute_gradient!(is)

    # Setup the algorithm.
    # IF we are doing gradient descent, this does nothing, otherwise initializes
    # the SR/Natural gradient structures.
    setup_algorithm!(is.precond_cache, reshape(∇C,:), ∇vr, is.parallel_cache)

    return L_stat, is.precond_cache
end

function compute_observables(is::BatchedValSampler)
    return compute_observables(is.observables_sampler)
end

function add_observable!(is::BatchedValSampler, name::String, obs)
    return add_observable!(is.observables_sampler, name, obs)
end

Base.show(io::IO, is::BatchedValSampler) = print(io,
    "BatchedValSampler for :"*
    "\n\tnet\t\t: $(is.bnet)"*
    "\n\tproblem\t: $(is.Ĉ)"*
    "\n\tsampler\t: $(is.sampler)")


text:
"""
    ib_cavity(set, ps, vs, gas, Um = 0.15, Vm = 0.0, Tm = 1.0)

Initialize lid-driven cavity
"""
function ib_cavity(
    set::AbstractSetup,
    ps::AbstractPhysicalSpace,
    vs::Union{AbstractVelocitySpace,Nothing},
    gas::AbstractProperty,
    Um = 0.15,
    Vm = 0.0,
    Tm = 1.0,
)

    if set.nSpecies == 1

        prim = [1.0, 0.0, 0.0, 1.0]
        w = prim_conserve(prim, gas.γ)
        h = maxwellian(vs.u, vs.v, prim)
        b = h .* gas.K / 2.0 / prim[end]

        primU = [1.0, Um, Vm, Tm]
        primD = deepcopy(prim)
        primL = deepcopy(prim)
        primR = deepcopy(prim)

        fw = function (args...)
            return w
        end

        bc = function (x, y)
            if y == ps.y1
                return primU
            else
                return prim
            end
        end

        if set.space[1:4] == "2d0f"
            return fw, bc
        elseif set.space == "2d1f2v"
            ff = function (args...)
                return h
            end

            return fw, ff, bc
        elseif set.space == "2d2f2v"
            ff = function (args...)
                return h, b
            end

            return fw, ff, bc
        end

    end

    return nothing

end

ext: jl


text:
struct Meth
  func
  args::Tuple
  graph
end

jscall(args...) = vcall(args...)

jscall(::typeof(identity), x) = x

cvalue(v) = DataFlow.isconstant(v) ? v.value.value : v

function inline_methods(v)
  DataFlow.prewalkλ(v) do v
    cvalue(v) isa Meth ? vcall(identity, cvalue(v).graph) : v
  end
end

function lower(v)
  v = inline_methods(v)
  v = DataFlow.prewalk(v -> DataFlow.islambda(v) ? DataFlow.λopen(v) : v, v)
  v = DataFlow.postwalk(v) do v
    v.value isa DataFlow.Line && return v[1]
    DataFlow.iscall(v) || return v
    jscall(cvalue.(v[:])...)
  end
  v = DataFlow.postwalk(v -> v.value isa DataFlow.OLambda ? DataFlow.λclose(v) : v, v)
  DataFlow.prewalkλ(v) do v
    DataFlow.iscall(v, control) ? v[2] : v
  end
end

function liftweights(v, weights=[])
  v = DataFlow.prewalkλ(v) do x
    cvalue(x) isa AbstractArray || return x
    push!(weights, Float32.(data(cvalue(x))))
    DataFlow.constant(:(model.weights[$(length(weights)-1)]))
  end
  v, weights
end

liftweights(c::Nothing, w=[]) = (c, w)

# Julia Code Stuff

using MacroTools: alias_gensyms, flatten, striplines

function inline_blocks(ex)
  ex = MacroTools.postwalk(ex) do ex
    @capture(ex, x_ = (body__; y_)) || return ex
    unblock(:($(body...); $x = $y))
  end
  ex = MacroTools.postwalk(ex) do ex
    @capture(ex, name_Symbol = (args__,) -> body_) || return ex
    :(function $name($(args...)) $body end)
  end
end

function insert_return(ex)
  isexpr(ex, :block) ? :($(ex.args[1:end-1]...);$(insert_return(ex.args[end]))) :
  isexpr(ex, :if) ? Expr(:if, ex.args[1], map(x -> insert_return(x), ex.args[2:end])...) :
  isexpr(ex, :return) ? ex :
  :(return $ex)
end

function insert_returns(ex)
  MacroTools.prewalk(ex) do x
    isexpr(x, :->, :function) ?
      Expr(x.head, x.args[1], insert_return(x.args[2])) :
      x
  end
end

function prepare(ex, name, states = nothing)
  decl_state = states == nothing ? :(;) :
    :(init = []; states = [])
  state_setup = states == nothing ? :(;) :
    :(global init = $states; global states = init.slice())
  reset_method = states == nothing ? :(;) :
    :(model.reset = () -> (global states = init.slice(); return);)
  get_states = states == nothing ? :(;) : :(model.getStates = () -> (return states);)
  set_weights = quote
    model.setWeights = (ws) -> begin
      model.weights = ws;
      $(state_setup.args...)
      return;
    end
  end
  quote
    model = (() -> begin
      math = tf;
      model.weights = []
      $(decl_state.args...)
      model = $(ex)
      $(reset_method.args...)
      $(get_states.args...)
      $(set_weights.args...)
      model
    end)()
    flux.fetchWeights($"$name.bson").then(model.setWeights)
  end |> insert_returns |> inline_blocks |> alias_gensyms |> flatten |> striplines
end

function compile(v::IVertex, name, states = [])
  statesv = states == [] ? nothing :
    unwrap((states...,)) |> lower
  v, weights = liftweights(lower(v))
  statesv, weights = liftweights(statesv, weights)
  statesv == nothing || (statesv = DataFlow.syntax(statesv))
  jsexpr(prepare(DataFlow.syntax(v), name, statesv)), weights
end

function compile(f, args...; name = "model")
  ctx = Trace()
  v = traceλ(f, args..., meta = ctx)
  compile(v, name, ctx.states)
end

function compile(name::AbstractString, f, args...)
  code, weights = compile(f, args..., name = name)
  open(io -> write(io, code), "$name.js", "w")
  BSON.@save "$name.bson" weights
  return
end

macro code_js(ex)
  @capture(ex, f_(args__)) || error("@code_js f(args...)")
  quote
    Text(compile($(esc(f)), $(esc.(args)...))[1])
  end
end


text:
#Mean absolute error (MAE)
function mae(model::Persa.Model, dataset::Persa.AbstractDataset)
    total = 0.0
    elements = 0

    for (u, v, r) in dataset
        value  = model[u, v]

        if !isnan(value)
            total += abs(r - value)
            elements += 1
        end
    end

    return total / elements
end

#Root mean squared error (RMSE)
function rmse(model::Persa.Model, dataset::Persa.AbstractDataset)
    total = 0.0
    elements = 0

    for (u, v, r) in dataset
        value  = model[u, v]

        if !isnan(value)
            total += (r - value)^2
            elements += 1
        end
    end

    return sqrt(total / elements)
end

#Coverage
function coverage(model::Persa.Model, dataset::Persa.AbstractDataset)
    elements = 0

    for (u, v, r) in dataset
        !isnan(model[u, v]) ? elements += 1 : nothing
    end

    return elements / length(dataset)
end


text:
using Plots
pgfplotsx()
col = [RGB(i) for i in [0x00ab9f,0x7409b7,0xd2cf00,0xff6f00] ] # mah precious colors
using DelimitedFiles
path = "/home/freshstart/DiplomaThesisData/images/"

#---- cpu times ----
cpu_ti = readdlm("cpu_times.dat")
median = (data) -> (s=0;for i in data;s+=i;end;return s/length(data))

cpu = plot(cpu_ti, label = "CPU times",c = col[1],
                title = "CPU time per aifoil variant simulation",
                ylabel = "Time [min]",yguidefontrotation = -90,
                xlabel = "Variant number",xguidefontrotation = 0,)

plot!(cpu,[0,length(cpu_ti)],[median(cpu_ti),median(cpu_ti)],
        label = "Median CPU time value",c = col[4])
savefig(cpu,path*"cpu_times.pdf")

println("Median time per airfoil = ",median(median(cpu_ti)))


text:
export Pass

abstract type Pass end
reftype(::Type{T}) where {T<:Pass} = API.LLVMPassRef

#
# Module passes
#

export ModulePass

@checked struct ModulePass <: Pass
    ref::reftype(Pass)
    root::Any

    function ModulePass(name::String, runner::Core.Function)
        function callback(ptr::Ptr{Cvoid})::Core.Bool
            mod = Module(convert(reftype(Module), ptr))
            return runner(mod)::Core.Bool
        end

        return new(API.LLVMCreateModulePass(name, callback), callback)
    end
end

#
# Function passes
#

export FunctionPass

@checked struct FunctionPass <: Pass
    ref::reftype(Pass)
    root::Any

    function FunctionPass(name::String, runner::Core.Function)
        function callback(ptr::Ptr{Cvoid})::Core.Bool
            fn = Function(convert(reftype(Function), ptr))
            return runner(fn)::Core.Bool
        end

        return new(API.LLVMCreateFunctionPass(name, callback), callback)
    end
end

#
# Basic block passes
#

export BasicBlockPass

@checked struct BasicBlockPass <: Pass
    ref::reftype(Pass)
    root::Any

    function BasicBlockPass(name::String, runner::Core.Function)
        VERSION >= v"1.4.0-DEV.589" && error("BasicBlockPass functionality has been removed from Julia and LLVM")

        function callback(ptr::Ptr{Cvoid})::Core.Bool
            bb = BasicBlock(convert(reftype(BasicBlock), ptr))
            return runner(bb)::Core.Bool
        end

        return new(API.LLVMCreateBasicBlockPass(name, callback), callback)
    end
end