text:
# Use baremodule to shave off a few KB from the serialized `.ji` file
baremodule Perl_jll
using Base
using Base: UUID
import JLLWrappers

JLLWrappers.@generate_main_file_header("Perl")
JLLWrappers.@generate_main_file("Perl", UUID("83958c19-0796-5285-893e-a1267f8ec499"))
end  # module Perl_jll


text:
using Functors, Test

struct Foo
  x
  y
end

@functor Foo

struct Bar
  x
end

@functor Bar

model = Bar(Foo(1, [1, 2, 3]))

model′ = fmap(float, model)

@test model.x.y == model′.x.y
@test model′.x.y isa Vector{Float64}

ext: jl


text:
using Test
using NQCDynamics
using LinearAlgebra

@test_nowarn RingPolymers.RingPolymerParameters{Float64}(10, 1, 1)
@test_nowarn RingPolymers.RingPolymerParameters{Float64}(1, 1, 1)
@test_nowarn RingPolymers.RingPolymerParameters{Float64}(2, 1, 1)
rp = RingPolymers.RingPolymerParameters{Float64}(10, 1.0, vcat(fill(:C, 3), fill(:H, 5)), [:H])
@test rp.quantum_atoms == [4, 5, 6, 7, 8]
rp = RingPolymers.RingPolymerParameters{Float64}(10, 1.0, 10)
@test rp.quantum_atoms == collect(1:10)

@test sort(eigvals(rp.springs)) ≈ sort(rp.normal_mode_springs)

atoms = Atoms{Float64}(vcat(fill(:H, 10), :O))
model = NQCModels.Free(3)
sim = RingPolymerSimulation{Classical}(atoms, model, 10)
@test sim.beads.quantum_atoms == collect(1:11)
sim = RingPolymerSimulation{Classical}(atoms, model, 10; quantum_nuclei=[:H])
@test sim.beads.quantum_atoms == collect(1:10)

U = sim.beads.transformation.U
S = sim.beads.springs
λ = sim.beads.normal_mode_springs
@test U'U ≈ I
@test abs(det(U)) ≈ 1
@test U'S*U ≈ diagm(λ)

half = RingPolymers.cayley_propagator(sim.beads, 0.1; half=true)
full = RingPolymers.cayley_propagator(sim.beads, 0.1; half=false)
@test full ≈ half .* half


text:
using NonlinearEigenproblems
using LinearAlgebra
using SuiteSparse

export LinSolverCache, solve

struct LinSolverCache{T<:Number}
    solver::Dict{T,LinSolver}
    nep::NEP
    linsolvercreator::LinSolverCreator
end

LinSolverCache(::Type{T}, nep, linsolvercreator) where T<:Number =
    LinSolverCache(Dict{T,LinSolver}(), nep, linsolvercreator)

function solve(cache, σ, y, add_to_cache)
    @debug "Lin solve for $σ, cache=$add_to_cache, size=$(length(cache.solver))"
    if add_to_cache
        solver = get!(() -> (
            @debug "Cache miss, creating lin solver for $σ";
            create_linsolver(cache.linsolvercreator, cache.nep, σ)), cache.solver, σ)
    else
        solver = create_linsolver(cache.linsolvercreator, cache.nep, σ)
    end
    lin_solve(solver, y)
end


text:
"""
    findidx(A,val)

Helper function for getting the indices for tiles. A is a gcmfaces struct, val is a
    numeric value to get the indices of. Returns a Dict of the indices for each face.
    Not currently exported. Maybe move to MeshArrays?
"""
function findidx(A,val)
    tileidx = Dict()
    if isa(A,MeshArrays.gcmfaces)
        nfaces = A.nFaces
    elseif isa(A,MeshArray)
        nfaces = A.grid.nFaces
    end
    for i = 1:nfaces
        idx = findall(x->x==val,A.f[i])
        tileidx[i] = [x for x in idx]
    end
    return tileidx
end

"""
    gettile(fldvals,tileinfo,tilesize,tilenum::Int)

Helper function for retrieving a tile from a gcmfaces struct as a numeric Array. Not
    currently exported.
"""
function gettile(fldvals,tileinfo,tilesize,tilenum::Int)
    tilidx = findidx(tileinfo["tileNo"],tilenum)

    if isa(fldvals,MeshArrays.gcmfaces)
        is3D = length(size(fldvals)) == 3
        if is3D; n3 = size(fldvals)[3]; end
    else
        is3D = length(size(fldvals)) == 2 && size(fldvals)[2] > 1
        if is3D; n3 = size(fldvals)[2]; end
    end
    prec = eltype(fldvals)
    if is3D # has depth
        tilfld = Array{prec,3}(undef,0,tilesize[2],n3)
    else
        tilfld = Array{prec,2}(undef,0,tilesize[2])
    end
    if isa(fldvals,MeshArrays.gcmfaces)
        nfaces = fldvals.nFaces
    elseif isa(fldvals,MeshArray)
        nfaces = fldvals.grid.nFaces
    end
    for iF = 1:nfaces
        if ~isempty(tilidx[iF])
            imin = minimum(tilidx[iF])[1]; imax = maximum(tilidx[iF])[1]
            jmin = minimum(tilidx[iF])[2]; jmax = maximum(tilidx[iF])[2]
            if is3D
                if isa(fldvals,MeshArray)
                    tilfld = [tilfld; cat([fldvals.f[iF,d][imin:imax,jmin:jmax] for d in 1:n3]...,dims=3)]
                else
                    tilfld = [tilfld; fldvals.f[iF][imin:imax,jmin:jmax,:]]
                end
            else
                tilfld = [tilfld; fldvals.f[iF][imin:imax,jmin:jmax]]
            end
        end
    end

    if is3D
        tilfld = reshape(tilfld,tilesize[1],tilesize[2],:)
    else
        tilfld = reshape(tilfld,tilesize[1],tilesize[2])
    end

    return tilfld
end

"""
    gettiles(fldvals,tilenum::Int)

Helper function for retrieving a tile from a gcmfaces struct as a numeric Array along
    with associated latitude and longitude. Not currently exported.
"""
function gettiles(tilfld,tilenum::Int)
    tilesize = tilfld.tilesize

    tilfld = gettile(tilfld.vals,tilfld.tileinfo,tilesize,tilenum)
    tillat = gettile(tilfld.tileinfo["XC"],tilfld.tileinfo,tilfld.tilesize,tilenum)
    tillon = gettile(tilfld.tileinfo["YC"],tilfld.tileinfo,tilfld.tilesize,tilenum)

    return tilfld,tillat,tillon
end


text:
using IntervalOptimisation
using Test

const SortedVector = IntervalOptimisation.SortedVector

@testset "SortedVector" begin

    @testset "With standard ordering" begin

        v = SortedVector([5, 2])
        @test v.data == [2, 5]

        push!(v, 3)
        @test v.data == [2, 3, 5]

        push!(v, 4, 6)
        @test v.data == [2, 3, 4, 5, 6]

        push!(v, 1)
        @test v.data == [1, 2, 3, 4, 5, 6]
    end

    @testset "With ordering function" begin

        v = SortedVector([(5, "a"), (2, "b")], x->x[1])
        @test v.data == [(2, "b"), (5, "a")]

        push!(v, (3, "z"))
        @test v.data == [(2, "b"), (3, "z"), (5, "a")]

        push!(v, (6, "z"), (4, "y"))
        @test v.data == [(2, "b"), (3, "z"), (4, "y"), (5, "a"), (6, "z")]

        push!(v, (1, "x"))
        @test v.data == [(1, "x"), (2, "b"), (3, "z"), (4, "y"), (5, "a"), (6, "z")]

    end
   @testset "With ordering fucntion which compares values of index other than 1" begin

        v = SortedVector([(2,"z"), (1,"a")], x->x[2])
        @test v.data ==[(1,"a"), (2,"z")]

        push!(v, (3, "x"))
        @test v.data == [(1, "a"), (3, "x"), (2, "z")]

        push!(v, (4, "c"), (5, "y"))
        @test v.data == [(1, "a"), (4, "c"), (3, "x"), (5, "y"), (2, "z")]
    end

end


text:
mutable struct Song{Wave, MakeEnvelope}
    wave::Wave
    make_envelope::MakeEnvelope
    beat_duration::TIME
    initial_key::FREQUENCY
    ramp::TIME
    sample_rate::FREQUENCY
    volume::Float64
    chords::Vector{Chord}
end

function Song(;
    wave = SawTooth(7),
    make_envelope = pedal,
    beat_duration = 0.3s,
    initial_key = 220.0Hz,
    ramp = 0.1s,
    sample_rate = 44100.0Hz,
    volume = 0.1,
    # need to allocate a new vector
    chords = Chord[],
)
    Song(wave, make_envelope, beat_duration, initial_key, ramp, sample_rate, volume, chords)
end

function print(io::IO, song::Song)
    print(io, song.initial_key)
    print(io, "; ")
    print(io, 60s / song.beat_duration)
    print(io, " bpm")
    println(io)
    for chord in song.chords
        print(io, chord)
    end
end

const INTRO_REGEX = r"(?<initial_key>.*) Hz; (?<tempo>.*) bpm"
const WORDS_REGEX = r"# (?<words>.*)"

function add_initial_key_and_tempo!(song, line_number, line)
    a_match = match(INTRO_REGEX, line)
    if a_match === nothing
        throw_parse_error(line, "initial key and tempo", line_number)
    else
        initial_key_text = a_match["initial_key"]
        initial_key = tryparse(Float64, initial_key_text)
        if initial_key === nothing
            throw_parse_error(initial_key_text, "initial key", line_number)
        else
            song.initial_key = (initial_key)Hz
        end
        tempo_text = a_match["tempo"]
        tempo = tryparse(Float64, tempo_text)
        if tempo === nothing
            throw_parse_error(tempo_text, "tempo", line_number)
        else
            song.beat_duration = 60s / tempo
        end
    end
    nothing
end

function read_justly!(io, song)
    chords = song.chords
    first_one = true
    words = ""
    for (line_number, line) in enumerate(eachline(io))
        if !isempty(line)
            if first_one
                add_initial_key_and_tempo!(song, line_number, line)
                first_one = false
            else
                words_match = match(WORDS_REGEX, line)
                if words_match === nothing
                    push!(
                        chords,
                        parse(Chord, line; line_number = line_number, words = words),
                    )
                    words = ""
                else
                    words = words_match["words"]
                end
            end
        end
    end
end

"""
    function read_justly(song_file;
        wave = SawTooth(7),
        make_envelope = pedal,
        ramp = 0.1s,
        sample_rate = 44100.0Hz,
        volume = 0.1
    )

Create a `Song` from a song file.

- `make_envelope` is a function to make an envelope, like [`pedal`](@ref).
- `ramp` is the onset/offset time, in time units (like `s`).
- `sample_rate` is the sample rate, in frequency units (like `Hz`).
- `volume`, ranging from 0-1, is the volume that a single voice is played at.
- `wave` is a function which takes an angle in radians and returns an amplitude between -1 and 1.

```jldoctest read_justly
julia> using Justly

julia> song = read_justly(joinpath(pkgdir(Justly), "test", "song.justly"));

julia> print(song)
220.0 Hz; 800.0 bpm
# first chord
1 for 5: 1 for 1, 1 for 3, 5/4 for 5, 3/2 for 8
1 for 14: 5/4o1 for 14
1 for -1: 1 for 1
```

You can create an `AudioSchedule` from a song.

```jldoctest read_justly
julia> using AudioSchedules: AudioSchedule

julia> AudioSchedule(song)
1.55 s 44100.0 Hz AudioSchedule
```
"""
function read_justly(file; keyword_arguments...)
    song = Song(; keyword_arguments...)
    open(let song = song
        io -> read_justly!(io, song)
    end, file)
    song
end

export read_justly

function update_beats_per_minute!(song::Song, beats_per_minute)
    song.beat_duration = (60 / beats_per_minute)s
    nothing
end

function update_initial_midi_code!(song::Song, midi_code)
    song.initial_key = get_frequency(midi_code)
    nothing
end

function get_beats_per_minute(song::Song)
    60s / song.beat_duration
end

function get_initial_midi_code(song::Song)
    get_nearest_midi_code(song.initial_key)
end

max_stars_repo_name: bramtayl/Justly

max_issues_repo_name: bramtayl/Justly

max_forks_repo_name: bramtayl/Justly


text:
module Parser
  export to_push

  # Produces a readily executable Push program from a given definition.
  to_push(s::String) = parse(lex(s))

  # Lex source code into a sequence of tokens.
  lex(s::String) = String[ss for ss in matchall(r"\(|\)|[^\(\)\s]+", s)]

  # Parse tokens into a Push program.
  function parse(tokens::Vector{String})
    prog = Any[]
    prog_stack = Vector{Any}[]

    while !isempty(tokens)
      t = shift!(tokens)
      if t == "("
        p2 = Any[]
        push!(prog, p2)
        push!(prog_stack, prog)
        prog = p2
      elseif t == ")"
        prog = pop!(prog_stack)
      elseif ismatch(r"^\-?\d+\.\d+$", t)
        push!(prog, convert(Float32, float(t)))
      elseif ismatch(r"^\-?\d+$", t)
        push!(prog, convert(Int32, int(t)))
      elseif t == "FALSE"
        push!(prog, false)
      elseif t == "TRUE"
        push!(prog, true)
      elseif ismatch(r"^\w\S*$", t)
        push!(prog, convert(Symbol, uppercase(t)))
      else
        error("PARSING ERROR - Unrecognised token: $(t).")
      end
    end

    if !isempty(prog_stack)
      error("PARSING ERROR - Unclosed parenthesis.")
    end
    if length(prog) != 1
      error("PARSING ERROR - Missing top-level parenthesis.")
    end

    return prog[1]
  end

end


text:
## We use Monte Carlo methods, because: 1. The size of the state space
## >> parameter space.  So just explore parameter space.  2. We cannot
## run from the last period reliably without knowing the state space.

## Two algorithms available:
## biological: https://github.com/robertfeldt/BlackBoxOptim.jl
## sampled: Deterministic optimization under different Monte Carlo
## draws.

##### Biological Genetics #####

using BlackBoxOptim

##########

mutable struct UncertainOptimizationProblem
    model::Model
    components::Vector{Symbol}
    names::Vector{Symbol}
    montecarlo::Function
    objective::Function
    lowers::Vector{Float64}
    uppers::Vector{Float64}
end

mutable struct UncertainOptimizationSolution
    problem::UncertainOptimizationProblem
    fmean::Float64
    fserr::Float64
    xmean::Vector{Float64}
    xserr::Vector{Float64}
    extra
end

"""Setup an optimization over Monte Carlo uncertainty."""
function uncertainproblem(model::Model, components::Vector{Symbol}, names::Vector{Symbol}, lowers::Vector{Float64}, uppers::Vector{Float64}, objective::Function, montecarlo::Function=(model) -> nothing)
    my_lowers, my_uppers, totalvars = expandlimits(model, components, names, lowers, uppers)
    my_objective = unaryobjective(model, components, names, objective)

    UncertainOptimizationProblem(model, components, names, montecarlo, my_objective, my_lowers, my_uppers)
end

"""Solve an optimization over Monte Carlo uncertainty."""
function solution(optprob::UncertainOptimizationProblem, generator::Function, algorithm::Symbol, mcperlife::Int64, samples::Int64)
    function sample_objective(parameters::Vector{Float64})
        total = 0
        for iter in 1:mcperlife
            optprob.montecarlo(optprob.model)
            total += optprob.objective(parameters)
        end

        total / mcperlife
    end

    if algorithm == :biological
        res = bboptimize(p -> -sample_objective(p); SearchRange=collect(zip(optprob.lowers, optprob.uppers)), MaxFuncEvals=samples, Method=:separable_nes)

        fmean = NaN
        fserr = NaN
        xmean = best_candidate(res)
        xserr = repmat([NaN], length(optprob.lowers))
        extra = res
    elseif algorithm == :sampled
        opt = Opt(:LN_COBYLA, length(optprob.lowers))
        lower_bounds!(opt, optprob.lowers)
        upper_bounds!(opt, optprob.uppers)
        xtol_rel!(opt, minimum(1e-6 * (optprob.uppers - optprob.lowers)))

        sampleseed = 0
        function hold_objective(parameters::Vector{Float64}, grad::Vector{Float64})
            srand(sampleseed)
            sample_objective(parameters)
        end

        max_objective!(opt, hold_objective)

        sampleprob = OptimizationProblem(optprob.model, optprob.components, optprob.names, opt, Function[])

        allparams = Vector{Float64}[]

        for sample in 1:samples
            sampleseed = sample

            minf, minx = solution(sampleprob, generator, maxiter=10000)
            push!(allparams, Float64[minf; minx])
        end

        allparams = transpose(hcat(allparams...))

        fmean = mean(allparams[:, 1])
        fserr = std(allparams[:, 1]) / sqrt(samples)

        xmean = vec(mean(allparams[:, 2:end], 1))
        xserr = vec(std(allparams[:, 2:end], 1) / sqrt(samples))

        extra = nothing
    else
        error("Unknown algorithm")
    end

    UncertainOptimizationSolution(optprob, fmean, fserr, xmean, xserr, extra)
end

function evaluate(optprob::UncertainOptimizationProblem, param::Vector{Float64}, mcperlife::Int64, sampleseed=1)
    srand(sampleseed)

    total = 0
    for iter in 1:mcperlife
        optprob.montecarlo(optprob.model)
        total += optprob.objective(param)
    end

    total / mcperlife
end


text:
function designs_pprb(ypred, D, p_prop, β_prop)

	M = size(ypred)[2]
	L = length(D)
	N = size(p_prop)[1]

	# Scores
	vars = fill(0.0, N, L, M)
	means = fill(0.0, N, L, M)

	@progress for j in 1:M
		for l in 1:L

			ynew = ypred[D[l], j]

			updated = probit_pprb(ynew, D[l], p_prop, β_prop, 100000)

			vars[:, l, j] = var(updated.y, dims = 2)
			means[:, l, j] = mean(updated.y, dims = 2)

		end
	end

	return (vars = vars, means = means, score = score)

end

max_stars_repo_name: clint-leach/optimal-design

max_issues_repo_name: clint-leach/optimal-design

max_forks_repo_name: clint-leach/optimal-design


text:
const JacType{TF<:AbstractFloat} = Union{AbstractMatrix{TF}, Shift{TF}}

struct TotalJacobian{TF<:AbstractFloat}
    parent::SequenceSpaceModel
    blks::Vector{AbstractBlock}
    vars::Set{Symbol}
    varvals::Dict{Symbol,ValType{TF}}
    nT::Int
    srcs::Set{Symbol}
    tars::Vector{Symbol}
    excluded::Union{Set{BlockOrVar},Nothing}
    parts::IdDict{Symbol,Vector{Pair{Symbol,JacType{TF}}}}
    totals::IdDict{Symbol,IdDict{Symbol,LinearMap}}
end

_tomap(S::Shift, T::Int) = ShiftMap(S, T)
_tomap(M::AbstractMatrix, ::Int) = LinearMap(M)

function TotalJacobian(m::SequenceSpaceModel, sources, targets,
        varvals::Dict{Symbol,ValType{TF}}, nT::Int, excluded=nothing) where TF
    sources isa Symbol && (sources = (sources,))
    targets isa Symbol && (targets = (targets,))
    excluded isa BlockOrVar && (excluded = (excluded,))
    isempty(sources) && throw(ArgumentError("sources cannot be empty"))
    isempty(targets) && throw(ArgumentError("targets cannot be empty"))
    sources = Set{Symbol}(sources)
    # The order of the targets need to be maintained
    targets = collect(Symbol, targets)
    excluded === nothing || (excluded = Set{BlockOrVar}(excluded))
    vars = copy(sources)
    blks = AbstractBlock[]
    parts = IdDict{Symbol,Vector{Pair{Symbol,JacType{TF}}}}()
    Dmap = IdDict{Symbol,LinearMap}
    totals = IdDict{Symbol,Dmap}(u=>Dmap() for u in vars)
    for v in m.order
        isblock(m, v) || continue
        blk = m.pool[v]
        blk isa Symbol && continue
        excluded !== nothing && blk in excluded && continue
        pushed = false
        for (i, vi) in enumerate(inputs(blk))
            if vi in vars
                J = jacobian(blk, i, varvals)
                for (r, vo) in enumerate(outputs(blk))
                    excluded !== nothing && vo in excluded && continue
                    sh = shift(invars(blk)[i])
                    js = get!(valtype(parts), parts, vo)
                    # ! to do: consider vector inputs/outputs and block types
                    j = Shift(sh, J[r,1])
                    push!(js, vi=>j)
                    iszero(j) && continue
                    push!(vars, vo)
                    j = _tomap(j, nT)
                    # Handle the case when vi is a source of the DAG
                    unknown = get(totals, vi, nothing)
                    if unknown === nothing
                        for d in values(totals)
                            maplast = get(d, vi, nothing)
                            if maplast !== nothing
                                jcomp = j * maplast
                                # vo may exist when temporal terms are involved
                                map = get(d, vo, nothing)
                                d[vo] = map === nothing ? jcomp : map+jcomp
                            end
                        end
                    else
                        map = get(unknown, vo, nothing)
                        unknown[vo] = map === nothing ? j : map+j
                    end
                end
                if !pushed
                    push!(blks, blk)
                    pushed = true
                end
            end
        end
    end
    return TotalJacobian(m, blks, vars, varvals, nT, sources, targets, excluded, parts, totals)
end

struct GEJacobian{TF<:AbstractFloat}
    jacs::TotalJacobian{TF}
    exovars::Vector{Symbol}
    unknowns::Vector{Symbol}
    H_U::Union{LU{TF, Matrix{TF}},Nothing}
    Gs::IdDict{Symbol,IdDict{Symbol,Union{Matrix{TF},LinearMap{TF}}}}
end

function GEJacobian(jacs::TotalJacobian{TF}, exovars; keepH_U::Bool=false) where TF
    exovars isa Symbol && (exovars = (exovars,))
    isempty(exovars) && throw(ArgumentError("exovars cannot be empty"))
    for var in exovars
        var in jacs.srcs || throw(ArgumentError("$var is not a source variable"))
    end
    exovars = collect(Symbol, exovars)
    unknowns = collect(setdiff(jacs.srcs, exovars))
    nT = jacs.nT
    nZ = length(exovars)
    nU = length(unknowns)
    ntar = length(jacs.tars)
    H_U = Matrix(hvcat(((nU for _ in 1:ntar)...,),
        (jacs.totals[v][t] for t in jacs.tars for v in unknowns)...))
    H_Z = Matrix(hvcat(((nZ for _ in 1:ntar)...,),
        (jacs.totals[v][t] for t in jacs.tars for v in exovars)...))
    H_U = lu!(H_U)
    ldiv!(H_U, H_Z)
    G_U = H_Z
    G_U .*= -one(eltype(G_U))
    Gs = IdDict{Symbol,IdDict{Symbol,Union{Matrix{TF},LinearMap{TF}}}}()
    for (j, z) in enumerate(exovars)
        Gs[z] = IdDict{Symbol,Union{Matrix{TF},LinearMap{TF}}}()
        for (i, u) in enumerate(unknowns)
            Gs[z][u] = G_U[1+(i-1)*nT:i*nT, 1+(j-1)*nT:j*nT]
        end
    end
    return GEJacobian(jacs, exovars, unknowns, keepH_U ? H_U : nothing, Gs)
end

function getG!(gejac::GEJacobian{TF}, exovar::Symbol, endovar::Symbol) where TF
    z = get(gejac.Gs, exovar, nothing)
    z === nothing && throw(ArgumentError("$exovar is not an exogenous variable"))
    G = get(z, endovar, nothing)
    if G === nothing
        # endovar does not have to be a source
        endovar in gejac.jacs.vars ||
            throw(ArgumentError("$endovar is not an endogenous variable"))
        # M_U combines all indirect effects while M_u is for a specific channel
        M_U = LinearMap(UniformScaling(zero(TF)), gejac.jacs.nT)
        M_Z = M_U
        for (u, ms) in gejac.jacs.totals
            if u === exovar
                M = get(ms, endovar, nothing)
                M === nothing || (M_Z = M)
            else
                M_u = get(ms, endovar, nothing)
                M_u === nothing && continue
                M_U += M_u * gejac.Gs[exovar][u]
            end
        end
        G = Matrix(M_U + M_Z)
        gejac.Gs[exovar][endovar] = G
    end
    return G
end


text:
##  Copyright 2022 Myrl Marmarelis
##
##  Licensed under the Apache License, Version 2.0 (the "License");
##  you may not use this file except in compliance with the License.
##  You may obtain a copy of the License at
##
##    http://www.apache.org/licenses/LICENSE-2.0
##
##  Unless required by applicable law or agreed to in writing, software
##  distributed under the License is distributed on an "AS IS" BASIS,
##  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
##  See the License for the specific language governing permissions and
##  limitations under the License.

using Flux, Zygote
using Random
using ProgressMeter
using StatsFuns

const Optimizer = Flux.Optimise.AbstractOptimiser

function ensemble_log_likelihood(predictors::Vector{P},
    parametrization::PredictionParam{T}, instance::Instance{In,Out,T},
    )::T where {In,Out,POut,T, P <: Predictor{In,POut,T}}
  ensemble_size = length(predictors)
  # `logsumexp` needs an intermediary array... (is reduce(logaddexp, ...) as stable?
  #   yes because logsumexp is implemented that way basically)
  log_likelihoods = Iterators.map(predictors) do predictor
    log_likelihood(predictor, parametrization, instance)
  end
  (reduce(logaddexp, log_likelihoods) - log(ensemble_size)) / Out
end

# INTERESTING THOUGHT, BUT UNFITTING
#  ensemble_size = length(predictors)
#  log_likelihoods = Iterators.map(zip(predictor_weights, predictors)) do (weights, predictor)
#    mean(zip(weights, instances)) do (weight, instance)
#      weight * log_likelihood(predictor, parametrization, instance)
#    end
#  end
#  (logsumexp(log_likelihoods) - log(ensemble_size)) / Out

using Base.Threads

# is "parameterization" more correct?
function train_directly!(predictors::Vector{P},
    instances::Vector{Instance{In,Out,T}}, parametrization::PredictionParam{T};
    weights::WeightScheme{T}, n_iterations::Int, optimizer::Optimizer, batch_size::Int,
    constraints::NTuple{C,Constraint{T}}=(), constraint_weight::T=T(0), n_constraint_draws::Int=0,
    validation_instances::Vector{Instance{In,Out,T}}=Instance{In,Out,T}[], verbose::Bool=false) where
      { In, Out, POut, C, T, P <: Predictor{In,POut,T} }
  @assert POut == n_predictor_outputs(parametrization, Out)
  @assert length(predictors) > 0
  sample_size = length(instances)
  validation_size = length(validation_instances)
  n_predictors = length(predictors)
  sample_indices = collect(1:sample_size) # filter for nonzero weights HERE
  predictor_weights = [
    draw_weights(weights, sample_size) for p in 1:n_predictors ]
  n_batches = div(sample_size, batch_size)
  n_constraints = length(constraints) # C
  CIs = [ constraint_instance_type(constraint, predictors[1], parametrization)
    for constraint in constraints ] # broadcasting doesn't seem to work here
  constraint_instances = [ Matrix{CI}(undef, n_constraint_draws, n_batches) for CI in CIs ]
  scores = zeros(T, n_iterations)
  validations = zeros(T, n_iterations)
  progress = Progress(n_iterations, enabled=verbose)
  for iteration in 1:n_iterations
    shuffle!(sample_indices)
    batch_indices = [
      @view sample_indices[(b-1)*batch_size+1 : b*batch_size]
      for b in 1:n_batches ]
    for (constraint_sample, constraint) in zip(constraint_instances, constraints)
      for draw in 1:n_constraint_draws, batch in 1:n_batches
        constraint_sample[draw, batch] = draw_instance(
          constraint, predictors[1], parametrization, instances)
      end
    end
    @threads for (weights, predictor) in # `collect` is needed for `@threads`
        zip(predictor_weights, predictors) |> collect
      Flux.train!(get_params(predictor),
          enumerate(batch_indices), optimizer) do batch, indices
        batch_weights = @view weights[indices]
        batch_instances = @view instances[indices]
        fit = mean(zip(batch_weights, batch_instances)) do (weight, instance)
          weight * log_likelihood(predictor, parametrization, instance)
        end
        penalty = n_constraints > 0 ? (
          mean(zip(constraint_instances, constraints)) do (constraint_sample, constraint)
            mean(@view constraint_sample[:, batch]) do constraint_instance
              apply_operator(
                constraint, predictor, parametrization, constraint_instance)
            end
          end) : T(0)
        -fit + constraint_weight*penalty
      end
    end
    score = mean(instances) do instance
      ensemble_log_likelihood(
        predictors, parametrization, instance)
    end
    validation = ( validation_size == 0 ? T(NaN) :
      mean(validation_instances) do instance
        ensemble_log_likelihood(
          predictors, parametrization, instance)
      end )
    scores[iteration] = score
    validations[iteration] = validation
    report = () -> [
      (:score, Float16(score)),
      (:validation, Float16(validation)) ]
    next!(progress, showvalues=report)
  end
  (; scores, validations, predictor_weights,
     predictors, instances, validation_instances ) # these ones are just returned back as passed in
end

max_stars_repo_name: marmarelis/TreatmentCurves

max_issues_repo_name: marmarelis/TreatmentCurves

max_forks_repo_name: marmarelis/TreatmentCurves


text:
# ---
# title: JuliaRL\_BasicDQN\_CartPole
# cover: assets/JuliaRL_BasicDQN_CartPole.png
# description: The simplest example to demonstrate how to use BasicDQN
# date: 2022-06-04
# author: "[Jun Tian](https://github.com/findmyway)"
# ---

#+ tangle=true
using ReinforcementLearning
using Flux
using Flux: glorot_uniform

using StableRNGs: StableRNG
using Flux.Losses: huber_loss

function RL.Experiment(
    ::Val{:JuliaRL},
    ::Val{:BasicDQN},
    ::Val{:CartPole};
    seed=123
)
    rng = StableRNG(seed)
    env = CartPoleEnv(; T=Float32, rng=rng)
    ns, na = length(state(env)), length(action_space(env))

    agent = Agent(
        policy=QBasedPolicy(
            learner=BasicDQNLearner(
                approximator=Approximator(
                    model=Chain(
                        Dense(ns, 128, relu; init=glorot_uniform(rng)),
                        Dense(128, 128, relu; init=glorot_uniform(rng)),
                        Dense(128, na; init=glorot_uniform(rng)),
                    ) |> gpu,
                    optimiser=ADAM(),
                ),
                loss_func=huber_loss,
            ),
            explorer=EpsilonGreedyExplorer(
                kind=:exp,
                ϵ_stable=0.01,
                decay_steps=500,
                rng=rng,
            ),
        ),
        trajectory=Trajectory(
            container=CircularArraySARTTraces(
                capacity=1000,
                state=Float32 => (ns,),
            ),
            sampler=BatchSampler{SS′ART}(
                batch_size=32,
                rng=rng
            ),
            controller=InsertSampleRatioController(
                threshold=100,
                n_inserted=-1
            )
        )
    )
    stop_condition = StopAfterStep(10_000, is_show_progress=!haskey(ENV, "CI"))
    hook = TotalRewardPerEpisode()
    Experiment(agent, env, stop_condition, hook)
end

#+ tangle=false
using Plots
pyplot() #hide
ex = E`JuliaRL_BasicDQN_CartPole`
run(ex)
plot(ex.hook.rewards)
savefig("assets/JuliaRL_BasicDQN_CartPole.png") #hide

#=
## Watch a demo episode with the trained agent

```julia
demo = Experiment(ex.policy,
                  CartPoleEnv(),
                  StopWhenDone(),
                  RolloutHook(plot, closeall),
                  "DQN <-> Demo")
run(demo)
```
=#

# ![](assets/JuliaRL_BasicDQN_CartPole.png)


text:
using CSV, DataFrames
using Gadfly
import Cairo

julia_bmk_file = joinpath("benchmark", "scatter_julia.tsv")
python_bmk_file = joinpath("benchmark", "scatter_pytorch.tsv")

bmk_jl = CSV.read(julia_bmk_file; delim='\t')
bmk_py = CSV.read(python_bmk_file; delim='\t')

bmk_jl[!, :framework] .= "ScatterNNlib.jl"
bmk_py[!, :framework] .= "Pytorch-scatter"

bmk = vcat(bmk_jl, bmk_py)

bmk[!, :min_time] .= bmk[!, :min_time]/1000
bmk[!, :mean_time] .= bmk[!, :mean_time]/1000
bmk[!, :max_time] .= bmk[!, :max_time]/1000

function plot_benchmark(device)
        DEVICE = uppercase(device)
        p = plot(bmk[bmk[!,:device] .== device, :], x="sample", y="mean_time", color="framework",
                 Geom.point, Geom.line, Scale.x_log2, Scale.y_log10,
                 Guide.title("Scatter add performance on $(DEVICE)"),
                 Guide.xlabel("Matrix Size"), Guide.ylabel("Time (μs)"),
                 Coord.cartesian(xmin=4, xmax=21, ymin=0, ymax=6))

        draw(SVG(joinpath("benchmark", "pics", "$(device)_scatter.svg"), 9inch, 6inch), p)
        draw(PNG(joinpath("benchmark", "pics", "$(device)_scatter.png"), 9inch, 6inch), p)
end

plot_benchmark("cpu")
plot_benchmark("gpu")


text:
title = "Examples"
weight = 10


text:
mutable struct Connection <: DBInterface.Connection
    dbc::API.Handle
    dsn::String
    quoteidentifierchar::Char
    types::Dict{Type, String}
    alltypes::Any
    cursorstmt::Any # keep track of most recent stmt cursor to free appropriately before next execute
    stmts::WeakKeyDict{Any, Any}
end

Connection(dbc::API.Handle, dsn) = Connection(dbc, dsn, '\0', Dict{Type, String}(), nothing, nothing, WeakKeyDict())

Base.show(io::IO, conn::Connection) = print(io, "ODBC.Connection($(conn.dsn))")

# make sure any previous open cursors are closed on the connection
function clear!(conn::Connection)
    if conn.cursorstmt !== nothing
        API.freestmt(conn.cursorstmt)
        conn.cursorstmt = nothing
    end
    return
end

# format usr, pwd, and extraauth into UID=user;PWD=pass;extraauth
function getextraauth(usr::Union{AbstractString, Nothing}, pwd::Union{AbstractString, Nothing}, extraauth::Union{AbstractString, Nothing}) :: String
    parts = [
        usr === nothing ? nothing : "UID={$usr}",
        pwd === nothing ? nothing : "PWD={$pwd}",
        extraauth === nothing ? nothing : extraauth
    ]
    strings = filter(s -> s !== nothing, parts)
    return join(strings, ";")
end

"""
    ODBC.Connection(dsn_or_connectionstring; user, password, extraauth)

Construct a `Connection` type by connecting to a valid ODBC Connection or by specifying a datasource name or valid connection string.
1st argument `dsn` can be either the name of a pre-defined ODBC Connection or a valid connection string.
A great resource for building valid connection strings is [http://www.connectionstrings.com/](http://www.connectionstrings.com/).
Takes optional keyword arguments `username`, `password`, and `extraauth`, which are used to specify auth parameters. `extraauth` is
to allow you to pass a sensitive string to be appended verbatim to the end of the connection string, e.g. DB-specific auth token
parameters.

Note that connecting will use the currently "set" ODBC driver manager, which by default is iODBC on OSX, unixODBC on Linux, and
the system driver manager on Windows. If you experience cryptic connection errors, it's probably worth checking with your ODBC
driver documentation to see if it requires a specific driver manager. For example, Microsoft-provided ODBC driver for SQL Server
requires unixODBC on OSX, but by default, ODBC.jl sets the driver manager to iODBC, so before connecting, you would need to do:
```julia
ODBC.setunixODBC()
conn = ODBC.Connection(...)
```
"""
function Connection(dsn::AbstractString; user=nothing, password=nothing, extraauth=nothing)
    connstr = occursin('=', dsn) ? dsn : "DSN=$dsn"
    extraauth = getextraauth(user, password, extraauth)
    internalconnection = API.connect(connstr, extraauth)
    return Connection(internalconnection, dsn)
end

# back compat.
# TODO: docstring for this?
Connection(dsn::AbstractString, usr, pwd) = Connection(dsn; user=usr, password=pwd)

"""
    DBInterface.connect(ODBC.Connection, dsn_or_connectionstring; user, password, extraauth, connectionstring::Bool=false)

Construct a `Connection` type by connecting to a valid ODBC Connection or by specifying a datasource name or valid connection string.
1st argument `dsn` can be either the name of a pre-defined ODBC Connection or a valid connection string.
A great resource for building valid connection strings is [http://www.connectionstrings.com/](http://www.connectionstrings.com/).
Takes optional keyword arguments `username`, `password`, and `extraauth`, which are used to specify auth parameters. `extraauth` is
to allow you to pass a sensitive string to be appended verbatim to the end of the connection string, e.g. DB-specific auth token
parameters.

Note that connecting will use the currently "set" ODBC driver manager, which by default is iODBC on OSX, unixODBC on Linux, and
the system driver manager on Windows. If you experience cryptic connection errors, it's probably worth checking with your ODBC
driver documentation to see if it requires a specific driver manager. For example, Microsoft-provided ODBC driver for SQL Server
requires unixODBC on OSX, but by default, ODBC.jl sets the driver manager to iODBC, so before connecting, you would need to do:
```julia
ODBC.setunixODBC()
conn = ODBC.Connection(...)
```
"""
DBInterface.connect(::Type{Connection}, args...; kw...) = Connection(args...; kw...)

Base.isopen(c::Connection) = c.dbc.ptr != C_NULL && API.getconnectattr(c.dbc, API.SQL_ATTR_CONNECTION_DEAD) == API.SQL_CD_FALSE

"disconnect a connected `Connection`"
function disconnect!(conn::Connection)
    # for #292, it seems we need to finalize any stmts on the connection
    # before disconnect/freehandle, even though disconnect should do that for us
    for stmt in keys(conn.stmts)
        finalize(stmt)
    end
    API.disconnect(conn.dbc)
    return nothing
end

"""
    DBInterface.close!(conn)

Close an open connection. In general, statements and open cursors will not be valid
once a connection has been closed.
"""
DBInterface.close!(dsn::Connection) = disconnect!(dsn)

"Struct for working with prepared statements and parameter binding; see `DBInterface.prepare` and `DBInterface.execute` for more details"
mutable struct Statement <: DBInterface.Statement
    dsn::Connection
    stmt::API.Handle
    sql::String
    bindings::Union{Nothing, Vector{Binding}}
    nparams::Int
end

"""
    DBInterface.close!(stmt)

Close a prepared statement. Further parameter binding or execution will not be valid.
"""
DBInterface.close!(stmt::Statement) = finalize(stmt.stmt)

"""
    DBInterface.prepare(conn, sql) -> ODBC.Statement

Prepare a query string, optionally including parameters to bind upon execution (with `?` markers).
Please refer to individual dbms documentation for the exact level of parameter binding support.

The returned prepared statement can then be passed to `DBInterface.execute(stmt, params)` with
`params` that will be bound before execution. This allows preparing the statement once,
and re-using it many times with different parameters (or the same) efficiently.
"""
function DBInterface.prepare(conn::Connection, sql::AbstractString)
    clear!(conn)
    stmt = API.prepare(conn.dbc, sql)
    conn.stmts[stmt] = 0
    return Statement(conn, stmt, sql, nothing, API.numparams(stmt))
end

@noinline paramcheck(stmt, params) = (stmt.nparams == 0 || length(params) == stmt.nparams) || error("stmt requires $(stmt.nparams) params, only $(length(params)) provided")

"""
    DBInterface.execute(stmt, params=(); iterate_rows::Bool=false, ignore_driver_row_count::Bool=false, normalizenames::Bool=false, debug::Bool=false) -> ODBC.Cursor

Execute a prepare statement, binding any parameters beforehand. Returns a `Cursor`
object, even if the statement is not resultset-producing (cursor will have zero rows and/or columns).
The `Cursor` object satisfies the [Tables.jl](https://juliadata.github.io/Tables.jl/dev/)
interface as a source, so any valid sink can be used for inspecting results (a list of integrations
is maintained [here](https://github.com/JuliaData/Tables.jl/blob/master/INTEGRATIONS.md)).

Supported keyword arguments include:
  * `iterate_rows::Bool`: for forcing row iteration of the resultset
  * `ignore_driver_row_count::Bool`: for ignoring the row count returned from the database driver; in some cases (Netezza), the driver may return an incorrect or "prefetched" number for the row count instead of the actual row count; this allows ignoring those numbers and fetching the resultset until truly exhausted
  * `normalizenames::Bool`: normalize column names to valid Julia identifiers; this can be convenient when working with the results in, for example, a `DataFrame` where you can access columns like `df.col1`
  * `debug::Bool`: for printing additional debug information during the query/result process.
"""
function DBInterface.execute(stmt::Statement, params=(); debug::Bool=false, kw...)
    API.freestmt(stmt.stmt)
    clear!(stmt.dsn)
    paramcheck(stmt, params)
    stmt.dsn.cursorstmt = stmt.stmt
    stmt.bindings = bindparams(stmt.stmt, params, stmt.bindings)
    debug && println("executing prepared statement: $(stmt.sql)")
    API.execute(stmt.stmt)
    c = Cursor(stmt.stmt; debug=debug, kw...)
    return c
end

"""
    DBInterface.execute(conn, sql, params=(); iterate_rows::Bool=false, ignore_driver_row_count::Bool=false, normalizenames::Bool=false, debug::Bool=false) -> ODBC.Cursor

Send a query directly to connection for execution. Returns a `Cursor`
object, even if the statement is not resultset-producing (cursor will have zero rows and/or columns).
The `Cursor` object satisfies the [Tables.jl](https://juliadata.github.io/Tables.jl/dev/)
interface as a source, so any valid sink can be used for inspecting results (a list of integrations
is maintained [here](https://github.com/JuliaData/Tables.jl/blob/master/INTEGRATIONS.md)).

Supported keyword arguments include:
  * `iterate_rows::Bool`: for forcing row iteration of the resultset
  * `ignore_driver_row_count::Bool`: for ignoring the row count returned from the database driver; in some cases (Netezza), the driver may return an incorrect or "prefetched" number for the row count instead of the actual row count; this allows ignoring those numbers and fetching the resultset until truly exhausted
  * `normalizenames::Bool`: normalize column names to valid Julia identifiers; this can be convenient when working with the results in, for example, a `DataFrame` where you can access columns like `df.col1`
  * `debug::Bool`: for printing additional debug information during the query/result process.

This is an alternative execution path to `DBInterface.execute` with a prepared statement.
This method is faster/less overhead for one-time executions, but prepared statements will
have more benefit for repeated executions (even with different parameters).
"""
function DBInterface.execute(conn::Connection, sql::AbstractString, params=(); debug::Bool=false, kw...)
    clear!(conn)
    stmt = API.Handle(API.SQL_HANDLE_STMT, API.getptr(conn.dbc))
    conn.stmts[stmt] = 0
    conn.cursorstmt = stmt
    API.enableasync(stmt)
    bindings = bindparams(stmt, params, nothing)
    debug && println("executing statement: $sql")
    GC.@preserve bindings (API.execdirect(stmt, sql))
    return Cursor(stmt; debug=debug, kw...)
end

mutable struct Cursor{columnar, knownlength}
    stmt::API.Handle
    rows::Int
    cols::Int
    names::Vector{Symbol}
    types::Vector{Type}
    lookup::Dict{Symbol, Int}
    current_rownumber::Int
    current_resultsetnumber::Int
    bindings::Vector{Binding}
    columns::Vector{AbstractVector}
    metadata::Any
end

# takes a recently executed statement handle and handles any produced resultsets
function Cursor(stmt; iterate_rows::Bool=false, ignore_driver_row_count::Bool=false, normalizenames::Bool=false, debug::Bool=false)
    rows = API.numrows(stmt)
    cols = API.numcols(stmt)
    debug && println("rows = $rows, cols = $cols")
    # Allocate arrays to hold each column's metadata
    names = Vector{Symbol}(undef, cols)
    types = Vector{Type}(undef, cols)
    namelengths = Vector{API.SQLSMALLINT}(undef, cols)
    sqltypes = Vector{API.SQLSMALLINT}(undef, cols)
    ctypes = Vector{API.SQLSMALLINT}(undef, cols)
    columnsizes = Vector{API.SQLULEN}(undef, cols)
    decimaldigits = Vector{API.SQLSMALLINT}(undef, cols)
    nullables = Vector{API.SQLSMALLINT}(undef, cols)
    longtexts = Vector{Bool}(undef, cols)
    cname = Vector{API.sqlwcharsize()}(undef, 1024)
    for i = 1:cols
        API.SQLDescribeCol(API.getptr(stmt), i, cname, namelengths, sqltypes, columnsizes, decimaldigits, nullables)
        nm = API.str(cname, namelengths[i])
        names[i] = normalizenames ? normalizename(nm) : Symbol(nm)
        sqltype = sqltypes[i]
        ctype, jltype = fetchtypes(sqltype, columnsizes[i])
        ctypes[i] = ctype
        types[i] = nullables[i] == API.SQL_NO_NULLS ? jltype : Union{Missing, jltype}
        # Some drivers return 0 size for variable length or large fields
        longtexts[i] = (sqltype == API.SQL_LONGVARCHAR || sqltype == API.SQL_LONGVARBINARY || sqltype == API.SQL_WLONGVARCHAR) || (columnsizes[i] == 0 && sqltype in (API.SQL_VARCHAR, API.SQL_WVARCHAR, API.SQL_VARBINARY))
        if ctype == API.SQL_C_CHAR || ctype == API.SQL_C_WCHAR || ctype == API.SQL_C_BINARY
            if longtexts[i] || columnsizes[i] == 0 || columnsizes[i] > 2^22
                longtexts[i] = true
                columnsizes[i] = 255
            end
            columnsizes[i] += 1
        end
    end
    metadata = [["column name", names...] ["column type", types...] ["sql type", map(x->API.SQL_TYPES[x], sqltypes)...] ["c type", map(x->API.C_TYPES[x], ctypes)...] ["sizes", map(Int, columnsizes)...] ["nullable", map(x->x != API.SQL_NO_NULLS, nullables)...] ["long data", longtexts...]]
    columnar = knownlength = true
    if any(longtexts) || rows <= 0 || iterate_rows || ignore_driver_row_count
        rowset = 1
        columnar = false
        knownlength = rows > 0 && !ignore_driver_row_count
    else
        rowset = rows
    end
    debug && println("columnar = $columnar, rowset = $rowset")
    API.setrowset(stmt, rowset)
    # we need bindings regardless of row vs. column fetching
    bindings = getbindings(stmt, columnar, ctypes, sqltypes, columnsizes, nullables, longtexts, rowset)
    if columnar && cols > 0
        # will be populated by call to SQLFetchScroll
        rowsfetchedref = API.setrowsfetched(stmt)
        API.fetch(stmt)
        rowsfetched = rowsfetchedref[]
        debug && @show rowsfetched
        columns = Vector{AbstractVector}(undef, cols)
        for (i, binding) in enumerate(bindings)
            ctype = binding.valuetype
            if ctype == API.SQL_C_CHAR || ctype == API.SQL_C_WCHAR || ctype == API.SQL_C_BINARY
                T = types[i]
                A = Vector{T}(undef, rowsfetched)
                data = binding.value.buffer::Vector{UInt8}
                inds = binding.strlen_or_indptr
                cur = 1
                elsize = columnsizes[i]
                for j = 1:rowsfetched
                    @inbounds ind = inds[j]
                    A[j] = ind == API.SQL_NULL_DATA ? missing : jlcast(Base.nonmissingtype(T), unsafe_wrap(Array, pointer(data, cur), ind))
                    cur += elsize
                end
                columns[i] = A
            elseif ctype == API.SQL_C_TYPE_DATE || ctype == API.SQL_C_TYPE_TIME || ctype == API.SQL_C_TYPE_TIMESTAMP
                specialize(binding.value.buffer) do data
                    T = types[i]
                    A = Vector{T}(undef, rowsfetched)
                    inds = binding.strlen_or_indptr
                    @simd for j = 1:rowsfetched
                        @inbounds A[j] = ifelse(inds[j] == API.SQL_NULL_DATA, missing, Base.nonmissingtype(T)(data[j]))
                    end
                    columns[i] = A
                end
            else
                if nullables[i] == API.SQL_NO_NULLS
                    columns[i] = binding.value.buffer
                else
                    specialize(binding.value.buffer) do A
                        inds = binding.strlen_or_indptr
                        @simd for j = 1:rowsfetched
                            @inbounds A[j] = ifelse(inds[j] == API.SQL_NULL_DATA, missing, A[j])
                        end
                        columns[i] = A
                    end
                end
            end
        end
    else
        columns = AbstractVector[]
    end
    lookup = Dict(nm => i for (i, nm) in enumerate(names))
    return Cursor{columnar, knownlength}(stmt, rows, cols, names, types, lookup, 0, 1, bindings, columns, metadata)
end

# Tables.jl interface
Tables.istable(::Type{<:Cursor}) = true
Tables.schema(x::Cursor) = Tables.Schema(x.names, x.types)

# columnar source
Tables.columnaccess(::Type{Cursor{true, T}}) where {T} = true
Tables.columns(x::Cursor{true}) = x

Tables.columnnames(x::Cursor{true}) = x.names
Tables.getcolumn(x::Cursor{true}, nm::Symbol) = Tables.getcolumn(x, x.lookup[nm])
Tables.getcolumn(x::Cursor{true}, i::Int) = x.columns[i]

# row source
Tables.rowaccess(::Type{Cursor{false, T}}) where {T} = true
Tables.rows(x::Cursor{false}) = x

struct Row <: Tables.AbstractRow
    cursor::Cursor
    rownumber::Int
    resultsetnumber::Int
end

getcursor(x::Row) = getfield(x, :cursor)
getrownumber(x::Row) = getfield(x, :rownumber)
getresultsetnumber(x::Row) = getfield(x, :resultsetnumber)

Tables.columnnames(x::Row) = getfield(x, :cursor).names

@noinline wrongrow(i) = throw(ArgumentError("row $i is no longer valid; odbc results are forward-only iterators where each row is only valid when iterated"))

Tables.getcolumn(x::Row, nm::Symbol) = Tables.getcolumn(x, getcursor(x).lookup[nm])
Tables.getcolumn(x::Row, i::Int) = Tables.getcolumn(x, getcursor(x).types[i], i, getcursor(x).names[i])

function Tables.getcolumn(x::Row, ::Type{T}, i::Int, nm::Symbol) where {T}
    c = getcursor(x)
    (getrownumber(x) == c.current_rownumber && getresultsetnumber(x) == c.current_resultsetnumber) || wrongrow(getrownumber(x))
    b = c.bindings[i]
    if T >: Missing && b.strlen_or_indptr[1] == API.SQL_NULL_DATA
        return missing
    elseif b.valuetype == API.SQL_C_CHAR || b.valuetype == API.SQL_C_WCHAR || b.valuetype == API.SQL_C_BINARY
        data = b.value.buffer::Vector{UInt8}
        bytes = data[1:b.totallen]
        return jlcast(Base.nonmissingtype(T), bytes)
    elseif b.valuetype == API.SQL_C_TYPE_DATE || b.valuetype == API.SQL_C_TYPE_TIME || b.valuetype == API.SQL_C_TYPE_TIMESTAMP
        return specialize(x -> Base.nonmissingtype(T)(x[1]), b.value.buffer)
    else
        return specialize(x -> x[1], b.value.buffer)
    end
end

Base.IteratorSize(::Type{Cursor{false, knownlength}}) where {knownlength} = knownlength ? Base.HasLength() : Base.SizeUnknown()
Base.length(x::Cursor{false, true}) = x.rows
Base.IteratorEltype(::Type{Cursor{false, knownlength}}) where {knownlength} = Base.HasEltype()
Base.eltype(x::Cursor{false, knownlength}) where {knownlength} = Row

function Base.iterate(x::Cursor{false}, st=1)
    status = API.SQLFetchScroll(API.getptr(x.stmt), API.SQL_FETCH_NEXT, 0)
    while status == API.SQL_STILL_EXECUTING
        status = API.SQLFetchScroll(API.getptr(x.stmt), API.SQL_FETCH_NEXT, 0)
    end
    status == API.SQL_SUCCESS || status == API.SQL_SUCCESS_WITH_INFO || return nothing
    x.current_rownumber = st
    for (i, binding) in enumerate(x.bindings)
        getdata(x.stmt, i, binding)
    end
    return Row(x, st, x.current_resultsetnumber), st + 1
end

# support returning multiple results from single execute calls
mutable struct Cursors
    cursor::Cursor
    kw
end

Base.eltype(c::Cursors) = Cursor
Base.IteratorSize(::Type{<:Cursors}) = Base.SizeUnknown()

function Base.iterate(cursors::Cursors, first=true)
    cursors.cursor.stmt.ptr == C_NULL && return nothing
    if !first
        if API.moreresults(cursors.cursor.stmt) == API.SQL_SUCCESS
            cursors.cursor = Cursor(cursors.cursor.stmt; cursors.kw...)
        else
            return nothing
        end
    end
    return cursors.cursor, false
end

DBInterface.executemultiple(conn::Connection, sql::AbstractString, params=(); kw...) =
    Cursors(DBInterface.execute(conn, sql, params; kw...), kw)


text:
using CmdStan, Turing, ContinuousBenchmarks

#using Mamba: describe

include(joinpath(ContinuousBenchmarks.STAN_MODELS_DIR, "binormal-stan.model.jl"))

global stanmodel, rc, sim1, sim, stan_time
# stanmodel = Stanmodel(name="binormal", model=binorm, Sample(save_warmup=true));
stanmodel = Stanmodel(Sample(algorithm=CmdStan.Hmc(CmdStan.Static(0.5*5),CmdStan.diag_e(),0.5,0.0),
  save_warmup=true,adapt=CmdStan.Adapt(engaged=false)),
  num_samples=2000, num_warmup=0, thin=1,
  name="binormal", model=binorm, nchains=1);

rc, sim1 = stan(stanmodel, CmdStanDir=ContinuousBenchmarks.CMDSTAN_HOME, summary=false)

if rc == 0
  ## Subset Sampler Output
  sim = sim1[1:size(sim1, 1), ["lp__", "y.1", "y.2"], 1:size(sim1, 3)]

  # describe(sim)
end # cd


text:
module RessourceManager

paths = Dict{Symbol,String}(:ROOT => "")
use_backslash = false

#string(dirname(Base.source_path()),"/"))
#currentDay() = replace(replace(string(Dates.now()),":","-"),"T","_")

""" TODO """
useBackslash(use::Bool) = global use_backslash = use

""" TODO """
printPaths() = for e in paths println(string(e)) end

""" TODO """
frontslash(path::String) = replace(path,"\\"=>"/")

""" TODO """
backslash(path::String) = replace(path,"/"=>"\\")

""" TODO """
slash(path::String) = !use_backslash ? frontslash(path) : backslash(path)

""" TODO """
setRoot() = paths[:ROOT] = slash(abspath(joinpath(@__DIR__,"../")))
setRoot() # set root

""" TODO """
ressourcepath(path::String) = slash(isabspath(path) ? path : string(getRoot(), path, last(path) != '/' ? "/" : ""))

""" TODO """
setPath(id::Symbol, path::String) = paths[id] = ressourcepath(path)
export setPath

""" TODO """
getPath(key::Symbol,file::String="") = paths[haskey(paths, key) ? key : :ROOT]*file
export getPath

""" TODO """
getRoot() = getPath(:ROOT)
export getRoot

end # RessourceManager

max_stars_repo_name: Gilga/JuliaOpenGL

max_issues_repo_name: Gilga/JuliaOpenGL

max_forks_repo_name: Gilga/JuliaOpenGL


text:
function LinearAlgebra.mul!(
    C::GBVecOrMat,
    A::GBArray,
    B::GBArray,
    op = (+, *);
    mask = nothing,
    accum = nothing,
    desc = nothing
)
    mask, accum = _handlenothings(mask, accum)
    desc = _handledescriptor(desc; in1=A, in2=B)
    size(A, 2) == size(B, 1) || throw(DimensionMismatch("size(A, 2) != size(B, 1)"))
    size(A, 1) == size(C, 1) || throw(DimensionMismatch("size(A, 1) != size(C, 1)"))
    size(B, 2) == size(C, 2) || throw(DimensionMismatch("size(B, 2) != size(C, 2)"))
    op = getoperator(op, optype(A, B))
    accum = getaccum(accum, eltype(C))
    op isa TypedSemiring || throw(ArgumentError("$op is not a valid TypedSemiring"))
    libgb.GrB_mxm(C, mask, accum, op, parent(A), parent(B), desc)
    return C
end

"""
    mul(A::GBArray, B::GBArray, op=(+,*); kwargs...)::GBArray

Multiply two `GBArray`s `A` and `B` using a semiring, which defaults to the arithmetic semiring `+.*`.

Either operand may be transposed using `'` or `transpose(A)` provided the dimensions match.

The mutating form, `mul!(C, A, B, op; kwargs...)` is identical except it stores the result in `C::GBVecOrMat`.

The operator syntax `A * B` can be used when the default semiring is desired, and `*(max, +)(A, B)` can be used otherwise.

# Arguments
- `A, B::GBArray`: A GBVector or GBMatrix, possibly transposed.
- `op::Union{Tuple{Function, Function}, AbstractSemiring}`: the semiring used for matrix multiplication. May be passed as a tuple of functions, or an `AbstractSemiring` found in the `Semirings` submodule.
# Keywords
- `mask::Union{Nothing, GBArray} = nothing`: optional mask which determines the output pattern.
- `accum::Union{Nothing, Function, AbstractBinaryOp} = nothing`: optional binary accumulator
    operation such that `C[i,j] = accum(C[i,j], T[i,j])` where T is the result of this function before accum is applied.
- `desc::Union{Nothing, Descriptor}`

# Returns
- `GBArray`: The output matrix whose `eltype` is determined by `A` and `B` or the semiring
    if a type specific semiring is provided.
"""
function mul(
    A::GBArray,
    B::GBArray,
    op = (+, *);
    mask = nothing,
    accum = nothing,
    desc = nothing
)
    t = inferoutputtype(A, B, op)
    if A isa GBMatOrTranspose && B isa GBVector
        C = GBVector{t}(size(A, 1))
    elseif A isa GBVector && B isa GBMatOrTranspose
        C = GBVector{t}(size(B, 2))
    elseif A isa Transpose{<:Any, <:GBVector} && B isa GBVector
        C = GBVector{t}(1)
    else
        C = GBMatrix{t}(size(A, 1), size(B, 2))
    end
    mul!(C, A, B, op; mask, accum, desc)
    return C
end

function Base.:*(
    A::GBArray,
    B::GBArray;
    mask = nothing,
    accum = nothing,
    desc = nothing
)
    return mul(A, B, Semirings.PLUS_TIMES; mask, accum, desc)
end

function Base.:*((⊕)::Function, (⊗)::Function)
    return function(A::GBArray, B::GBArray; mask=nothing, accum=nothing, desc=nothing)
        mul(A, B, (⊕, ⊗); mask, accum, desc)
    end
end

function Base.:*(rig::AbstractSemiring)
    return function(A::GBArray, B::GBArray; mask=nothing, accum=nothing, desc=nothing)
        mul(A, B, rig; mask, accum, desc)
    end
end


text:
# geocoding.jl
module geocoding
include("HTTP")
include("JSON")

function getCoordinates(addressList)
    """
    getCoordinates(addressList::Array{String,1})

    returns a Vector of coordinates in (lat,lon) format as well as a data Array that contains all the data returned by nominatim
    """
    #nominatim requires as user agent and will block you otherwise
    HTTP.setuseragent!("Mozilla/5.0 (Windows NT 10.0; rv:68.0) Gecko/20100101 Firefox/68.0")
    coordinatesList = Array{Tuple{Float64,Float64},1}(undef, length(addressList)) # will contain (lat,lon) of addresses in addressList
    data = Array{Any,1}(undef, length(addressList)) # will contain all data nominatim gives for each request
    for (i, address) in enumerate(addressList)
        addressFormatted = join(split(address), "+")
        rawdata = HTTP.get(string(
            "https://nominatim.openstreetmap.org/search?q=",
            addressFormatted,
            "&format=json&limit=1",
        ))
        data[i] = JSON.parse(String(rawdata.body))[1]
        coordinatesList[i] = parse(Float64, data[i]["lat"]), parse(Float64, data[i]["lon"])
        sleep(1) #nominatim will block you if you have more than 1 request per second
    end
    return coordinatesList, data
end


text:
using Pkg
Pkg.activate(".")
using Test, LinearAlgebra, Plots, Revise, RecursiveClusterAlgorithm
gr(grid=false,legend=false)

## get_dipoles
# some vortices
xi = [-0.549 1.160 0.335 -1.773]
yi = [0.123 -0.603 -0.079 0.220]

scatter(xi[1:2],yi[1:2],m=:circle,c=:red)
scatter!(xi[3:4],yi[3:4],m=:circle,c=:blue)

n = 4
nplus = 2
Lx = 1.2
Ly = 0.3

κi = [ones(nplus); -ones(n-nplus)]

xd,yd,is_dipole,num_found_in_dipoles = get_dipoles(xi,yi,n,nplus,Lx,Ly)

@test num_found_in_dipoles == 2

## seed_clusters
xi = [-0.549 -0.548 0.160 0.335 -1.772 -1.78]
yi = [0.123 0.16 -0.103 -0.079 0.20 0.220]
xi,yi=xi[:],yi[:]
scatter(xi[1:3],yi[1:3],m=:circle,c=:red)
scatter!(xi[4:6],yi[4:6],m=:circle,c=:blue)
n = 6
nplus = 3
Lx = 1.2
Ly = 0.3
xs,ys,is_seed,seeds_index,num_found_in_seeds = seed_clusters(xi,yi,n,nplus,Lx,Ly)

@test num_found_in_seeds == 4

## grow_plus_clusters

## RCA
# some vortices
xp,yp = xi[1:3],yi[1:3]
xn,yn = xi[4:6],yi[4:6]

# RCA body code
nplus = length(xp)
nminus = length(xn)
n = nplus + nminus
kappai = [ones(nplus); -ones(nminus)]
vortex_index = 1:n

#
#   @info "All vortices are in dipoles: nothing left to do..."
#   PCLUSTERS = struct()
#   NCLUSTERS = struct()
#   DIPOLES.x = x_dipoles
#   DIPOLES.y = y_dipoles
#   DIPOLES.n = size(x_dipoles,1)
#
#   if PLOT
#  PlotDipoles(x_dipoles,y_dipoles,Lx,Ly)
#   end
#   return

# elseif num_found_in_dipoles == n-2 && sum(kappanew == 1) == 1 && sum(kappanew == -1) ==1
#   # NOTE!! This part is a HACK, to fix a BUG currently lurking in MATLAB
#   # 2012b (16/07/2014). RCA fails when there are only two vortices left, each
#   # of different signs (which necessarily must form a dipole). For some reason
#   # MATLAB thinks that the arrays are not the same size when using ==, even
#   # when they are. The problem only occurs when the routine is a function,
#   # so keyboard doesn"t help to debug either. It works fine in the command
#   # line...
#    x_dipoles = [x_dipoles xnew(1) xnew(2)]
#    y_dipoles = [y_dipoles ynew(1) ynew(2)]
#    PCLUSTERS = struct()
#    NCLUSTERS = struct()
#    DIPOLES.x = x_dipoles
#    DIPOLES.y = y_dipoles
#    DIPOLES.n = size(x_dipoles,1)
#
#   # if PLOT
#   #   PlotDipoles(x_dipoles,y_dipoles,Lx,Ly)
#   # end
#
#   @info "All vortices are in dipoles: nothing left to do..."
#   return
# end

## TODO Recursively look for dipoles, until no more can be found
# initial dipole search
x_dipoles,y_dipoles,is_dipole,num_found_in_dipoles = get_dipoles(xi,yi,n,nplus,Lx,Ly)

xnew = @. xi[~Bool(is_dipole)]
ynew = @. yi[~Bool(is_dipole)]
kappanew = @. kappai[~Bool(is_dipole)]
remaining_vortices_index =  @. vortex_index[~Bool(is_dipole)]

np = sum(@. kappanew==1.0)
ntot = length(kappanew)

if num_found_in_dipoles == n
  @info "All vortices are in dipoles. Finished."
end

# while (num_found_in_dipoles >0) && ~isempty(xnew)
  xd,yd,is_dipole2,num_found_in_dipoles = get_dipoles(xnew,ynew,ntot,np,Lx,Ly)

  add_to_dipoles = @. remaining_vortices_index[Bool(is_dipole2)]

  is_dipole = Bool.(is_dipole)
  @. is_dipole[add_to_dipoles] = true

  xnew = @. xi[~is_dipole]
  ynew = @. yi[~is_dipole]
  kappanew = @. kappai[~is_dipole]
  remaining_vortices_index = @. vortex_index[~is_dipole]
  np = sum(@. kappanew==1.0)
  ntot = length(kappanew)
  x_dipoles = [x_dipoles; xd] # ok
  y_dipoles = [y_dipoles; yd] # ok. A little sloppy here ...
# end


text:
"""
    Assignable

`Assignable` represents any data type that can be associated with traits.
It essentially covers all data types including parametric types e.g. `AbstractArray`
"""
const Assignable = Union{UnionAll, DataType}

"""
    Contract{T <: DataType, F <: Function, N}

A contract refers to a function defintion `func` that is required to satisfy
a trait. The function `func` must accepts `args` and returns `ret`.

# Fields
- `trait`: a trait type e.g. `Can{Fly}` or `Cannot{Fly}`
- `func`: function that must be implemented to satisfy this trait
- `args`: argument types of the function `func`
- `kwargs`: keyword argument names of the function `func`
- `ret`: return type of the function `func`
"""
struct Contract{T <: DataType, F <: Function}
    trait::T
    func::F
    args::Tuple
    kwargs::Tuple
    ret::Type
end

function Base.show(io::IO, c::Contract)
    # Represent the trait type with a special symbol
    args = [x == c.trait ? "🔹" : "::$x" for x in c.args]
    args = string("(", join(args, ", "))
    if length(c.kwargs) > 0
        args = string(args, "; ", join(c.kwargs, ", "))
    end
    args = string(args, ")")
    trait = supertype(c.trait)
    print(io, "$(trait): $(c.trait) ⇢ $(c.func)$(args)")
    c.ret !== Any && print(io, "::$(c.ret)")
end

"""
    InterfaceReview

An InterfaceReview object contains the validation results of an interface.

# Fields
- `data_type`: the type being checked
- `result`: true if the type fully implements all required contracts
- `implemented`: an array of implemented contracts
- `misses`: an array of unimplemented contracts
"""
@Base.kwdef struct InterfaceReview
    data_type::Assignable
    result::Bool
    implemented::Vector{Contract}
    misses::Vector{Contract}
    miss_reasons::Vector{String}
end

function Base.show(io::IO, ir::InterfaceReview)
    irtype = ir.data_type
    if length(ir.implemented) == length(ir.misses) == 0
        print(io, "✅ $(irtype) has no interface contract requirements.")
        return
    end
    if length(ir.implemented) > 0
        println(io, "✅ $(irtype) has implemented:")
        for (i, c) in enumerate(ir.implemented)
            println(io, "$(i). $c")
        end
    end
    if length(ir.misses) > 0
        println(io, "❌ $(irtype) is missing these implementations:")
        for (i, (c,reason)) in enumerate(zip(ir.misses, ir.miss_reasons))
            println(io, "$(i). $c ($reason)")
        end
    end
end

"""
    SyntaxError

Syntax error for macros.
"""
struct SyntaxError <: Exception
    msg
end

const MyDict = VERSION < v"1.2.0" ? Dict : IdDict

"""
    TraitsMap

Map a data type to the its assigned traits.
For example, `Dog => Set([Can{Swim}, Can{Run}, Cannot{Fly}])`.
"""
const TraitsMap = MyDict{Assignable,Set{DataType}}

"""
    InterfaceMap

Map a trait to a set of interface contracts.  See [`Contract`](@ref).
"""
const InterfaceMap = MyDict{DataType,Set{Contract}}

"""
    CompositeTraitMap

Maps a composite positive trait type to a set of its underlying traits,
which could be positive or negative.
e.g. `Is{Fish} => Set([Can{Swim}, Cannot{Fly}])`.
"""
const CompositeTraitMap = MyDict{DataType,Set{DataType}}

"""
    TraitsStorage

Keeps all traits-related dynamic data.

# Fields
- `traits_map`: see [`TraitsMap`](@ref)
- `interface_map`: see [`InterfaceMap`](@ref)
- `composite_map`: see [`CompositeTraitMap`](@ref)
"""
struct TraitsStorage
    traits_map::TraitsMap
    interface_map::InterfaceMap
    composite_map::CompositeTraitMap
    TraitsStorage() = new(TraitsMap(), InterfaceMap(), CompositeTraitMap())
end


text:
# -------- Main method:
"""
    generate_counterfactual(
        x::Union{AbstractArray,Int}, target::Union{AbstractFloat,Int}, data::CounterfactualData, M::Models.AbstractFittedModel, generator::AbstractGenerator;
        γ::AbstractFloat=0.75, T=1000
    )

The core function that is used to run counterfactual search for a given factual `x`, target, counterfactual data, model and generator. Keywords can be used to specify the desired threshold for the predicted target class probability and the maximum number of iterations.

# Examples

## Generic generator

```julia-repl
using CounterfactualExplanations

# Data:
using CounterfactualExplanations.Data
using Random
Random.seed!(1234)
xs, ys = Data.toy_data_linear()
X = hcat(xs...)
counterfactual_data = CounterfactualData(X,ys')

# Model
using CounterfactualExplanations.Models: LogisticModel, probs
# Logit model:
w = [1.0 1.0] # true coefficients
b = 0
M = LogisticModel(w, [b])

# Randomly selected factual:
x = select_factual(counterfactual_data,rand(1:size(X)[2]))
y = round(probs(M, x)[1])
target = round(probs(M, x)[1])==0 ? 1 : 0

# Counterfactual search:
generator = GenericGenerator()
counterfactual = generate_counterfactual(x, target, counterfactual_data, M, generator)
```
"""
function generate_counterfactual(
    x::Union{AbstractArray,Int}, target::Union{AbstractFloat,Int}, data::CounterfactualData, M::Models.AbstractFittedModel, generator::AbstractGenerator;
    γ::AbstractFloat=0.75, T=1000
)
    # Initialize:
    counterfactual = CounterfactualExplanation(x, target, data, M, generator, γ, T)
    initialize!(counterfactual)

    # Search:
    while !counterfactual.search[:terminated]
        update!(counterfactual)
    end

    return counterfactual

end


text:
module TightBinding

    include("./TBfromk2r.jl")
    #using Plots
    #using .Plotfuncs
    using .TBfromk2r
    using LinearAlgebra
    using Requires
    export set_Lattice,add_atoms!,add_hoppings!,add_diagonals!,hamiltonian_k,
    dispersion,get_position,calc_band,get_position_kspace,hamiltonian_k_1d,
    set_Klines,show_neighbors,add_Kpoints!,set_onsite!,set_μ!
    #,calc_band_plot,plotfuncs,
    #plot_lattice_2d,calc_band_plot_finite,plot_DOS

    export surfaceHamiltonian,σx,σy,σz,σ0
    export get_DOS

    struct Hopping
        amplitude
        ijatoms::Array{Int64,1}
        ijpositions::Array{Float64,1}
    end

    struct Kpoints
        kmin::Array{<:Number,1}
        kmax::Array{<:Number,1}
        nk::Int64
        name_start
        name_end
    end

    mutable struct Klines
        numlines::Int64
        kpoints::Array{Kpoints,1}
    end

    function __init__()
        @require Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" begin
            include("./Plotfuncs.jl")
            export Plotfuncs,plot_lattice_2d,calc_band_plot,plot_DOS,calc_band_plot_finite #From Plotfuncs.jl
            export plot_fermisurface_2D
        end
    end

    function set_Klines()
        kpoint = Array{Kpoints,1}[]
        klines = Klines(0,kpoint)
        return klines
    end

    function add_Kpoints!(klines,kmin,kmax,name_start,name_end;nk=20)
        kpoint = Kpoints(kmin,kmax,nk,name_start,name_end)
        klines.numlines += 1
        push!(klines.kpoints,kpoint)
    end

    mutable struct Lattice
        dim::Int8
        vectors::Array{Array{Float64,1},1}
        numatoms::Int64
        positions::Array{Array{Float64,1},1}
        numhopps::Int64
        hoppings::Array{Hopping,1}
        diagonals::Array{Float64,1}
        rvectors::Array{Array{Float64,1},1} #reciplocal lattice vectors
        μ::Float64
    end

    """
        set_Lattice(dim::Integer,vectors::Array{Array{Float64,1},1})
    Initialize lattice.
    We have to call this before making Hamiltonian.
    dim: Dimension of the system.
    vector:: Primitive vectors.

    Example:

    1D system

    ```julia
    la1 = set_Lattice(1,[[1.0]])
    ```

    2D system

    ```julia
    a1 = [sqrt(3)/2,1/2]
    a2 = [0,1]
    la2 = set_Lattice(2,[a1,a2])
    ```

    3D system

    ```julia
    a1 = [1,0,0]
    a2 = [0,1,0]
    a3 = [0,0,1]
    la2 = set_Lattice(3,[a1,a2,a3])
    ```
    """
    function set_Lattice(dim,vectors)
        positions = Array{Float64,1}[]
        hoppings = Array{Hopping,1}[]
        numatoms = 0
        numhopps=0
        diagonals = Float64[]
        μ = 0.0

        #making primitive vectors
        pvector_1 = zeros(Float64,3)
        pvector_2 = zeros(Float64,3)
        pvector_3 = zeros(Float64,3)
        if dim ==1
            pvector_1[1] = vectors[1][1]
            pvector_2[2] = 1.0 # 0 1 0
            pvector_3[3] = 1.0 # 0 0 1
        elseif dim == 2
            pvector_1[1:2] = vectors[1][1:2]
            pvector_2[1:2] = vectors[2][1:2]
            pvector_3[3] = 1.0 # 0 0 1
        elseif dim == 3
            pvector_1[1:3] = vectors[1][1:3]
            pvector_2[1:3] = vectors[2][1:3]
            pvector_3[1:3] = vectors[3][1:3]
        end
        #making reciplocal lattice vectors
        area = dot(pvector_1,cross(pvector_2,pvector_3))
        rvector_1 = 2π*cross(pvector_2,pvector_3)/area
        rvector_2 = 2π*cross(pvector_3,pvector_1)/area
        rvector_3 = 2π*cross(pvector_1,pvector_2)/area

        if dim ==1
            rvectors = [[rvector_1[1]]]
        elseif dim == 2
            rvectors = [rvector_1[1:2],rvector_2[1:2]]
        elseif dim == 3
            rvectors = [rvector_1[1:3],rvector_2[1:3],rvector_3[1:3]]
        end

        lattice = Lattice(dim,vectors,numatoms,positions,numhopps,hoppings,diagonals,rvectors,μ)
        return lattice
    end

    function add_atoms!(lattice,position::Array{<:Number,1})
        lattice.numatoms += 1
        push!(lattice.positions,position)
        add_diagonals!(lattice,0.0)
    end

    function add_atoms!(lattice,positions::Array{<:Array})
        for i=1:length(positions)
            lattice.numatoms += 1
            push!(lattice.positions,positions[i])
            add_diagonals!(lattice,0.0)
        end
    end

    function set_μ!(lattice,μ)
        lattice.μ = μ
    end

    function set_onsite!(lattice,onsite::Array{<:Number,1})
        lattice.diagonals = onsite[:]
    end

    function add_hoppings!(lattice,amplitude,iatom,jatom,hopping)
        hop = Hopping(amplitude,[iatom,jatom],hopping)
        lattice.numhopps += 1
        push!(lattice.hoppings,hop)
    end

    function add_diagonals!(lattice,diagonal::Number)
        push!(lattice.diagonals,diagonal)
    end

    function add_diagonals!(lattice,diagonals::Array{<:Number,1})
        for i=1:length(diagonals)
            push!(lattice.diagonals,diagonals[i])
        end
    end

    function get_position(lattice,vec)
        position = zeros(Float64,lattice.dim)
        for i=1:lattice.dim
            position += vec[i]*lattice.vectors[i]
        end
        return position
    end

    function get_position_kspace(lattice,vec)
        position = zeros(Float64,lattice.dim)
        for i=1:lattice.dim
            position += vec[i]*lattice.rvectors[i]
        end
        return position
    end

    function get_hopindex(lattice)
        hopindices = []
        for ihop=1:lattice.numhopps
            hopping = lattice.hoppings[ihop]
            i = hopping.ijatoms[1]
            ri = lattice.positions[i]
            rhop = ri+hopping.ijpositions
            iabhop = round.(Int,rhop).-1
            #j = hopping.ijatoms[2]
            push!(hopindices,iabhop)
        end
        return hopindices
    end

    function hamiltonian_k_1d(lattice,direction;nsites = 20,periodic=true)
        numatoms = lattice.numatoms
        diagonals = zeros(Float64,numatoms)
        vectors = lattice.vectors
        hopindices = get_hopindex(lattice)
        numhopps = lattice.numhopps
        μ = lattice.μ

        for i=1:length(lattice.diagonals)
            diagonals[i] = lattice.diagonals[i]-μ
        end

        function calc_ham(k)
            N = numatoms*nsites
            realpart_ham = zeros(Float64,N,N)
            imagpart_ham = zeros(Float64,N,N)
            for isite=1:nsites
                for i=1:numatoms
                    ii = (isite-1)*numatoms+i
                    realpart_ham[ii,ii] = diagonals[i]
                end
                for ihop=1:numhopps
                    δ = hopindices[ihop][direction]
                    hopping = lattice.hoppings[ihop]
                    i = hopping.ijatoms[1]
                    ii = (isite-1)*numatoms+i
                    jsite = isite + δ
                    if periodic
                        while jsite > nsites || jsite < 1
                            jsite += ifelse(jsite < 1,nsites,0)
                            jsite += ifelse(jsite > nsites,-nsites,0)
                        end
                    end
                    if 1 <= jsite <= nsites
                        j= hopping.ijatoms[2]
                        jj = (jsite-1)*numatoms+j

                        ak = 0.0
                        ik = 0
                        hop = hopping.ijpositions
                        for idim = 1:lattice.dim
                            if idim != direction
                                ik += 1
                                ak += k[ik]*hop[idim]
                            end
                        end

                        ampcos = hopping.amplitude*cos(ak)
                        realpart_ham[ii,jj] +=  ampcos
                        realpart_ham[jj,ii] +=  ampcos
                        ampsin = hopping.amplitude*sin(ak)
                        imagpart_ham[ii,jj] +=  ampsin
                        imagpart_ham[jj,ii] +=  -ampsin

                    end
                end
            end
            if sum(abs.(imagpart_ham)) == 0.0
                ham = realpart_ham[:,:]
            else
                ham = realpart_ham[:,:]+im*imagpart_ham[:,:]
            end
            ham
        end
        k -> calc_ham(k)
    end

    function hamiltonian_k(lattice)
        numatoms = lattice.numatoms
        diagonals = zeros(Float64,numatoms)
        vectors = lattice.vectors
        μ = lattice.μ
        for i=1:length(lattice.diagonals)
            diagonals[i] = lattice.diagonals[i]-μ
        end

        function calc_ham(k)
            realpart_ham = zeros(Float64,numatoms,numatoms)
            imagpart_ham = zeros(Float64,numatoms,numatoms)
            for i=1:numatoms
                realpart_ham[i,i] = diagonals[i]
            end
            if lattice.numhopps != 0
                for δ=1:lattice.numhopps
                    hopping = lattice.hoppings[δ] #Type:
                    i=hopping.ijatoms[1]
                    j=hopping.ijatoms[2]
                    hop = hopping.ijpositions
                    vec_a = zeros(Float64,lattice.dim)
                    for i=1:length(lattice.vectors)
                        vec_a += hop[i].*lattice.vectors[i]
                    end
                    ak = 0.0
                    for idim = 1:lattice.dim
                        ak += k[idim]*vec_a[idim]
                    end

                    ampcos = hopping.amplitude*cos(ak)
                    realpart_ham[i,j] +=  ampcos
                    realpart_ham[j,i] +=  ampcos
                    ampsin = hopping.amplitude*sin(ak)
                    imagpart_ham[i,j] +=  ampsin
                    imagpart_ham[j,i] +=  -ampsin

                end
            end
            if sum(abs.(imagpart_ham)) == 0.0
                ham = realpart_ham[:,:]
            else
                ham = realpart_ham[:,:]+im*imagpart_ham[:,:]
            end

            ham

        end
        k -> calc_ham(k)
    end

    function dispersion(dim,n,ham,k)
        if n==1
            return [ham(k)[1]]
        else
            energy = eigen(ham(k)).values
            return energy
        end
    end

    function calc_band(kmin_real,kmax_real,nk,lattice,ham)
        n = lattice.numatoms
        dim = lattice.dim
        energies = zeros(Float64,n,nk)
        vec_k = zeros(Float64,dim,nk)
        for idim=1:dim
            k = range(kmin_real[idim],length=nk,stop=kmax_real[idim])
            vec_k[idim,:] = k[:]
        end
        for ik = 1:nk
            energies[:,ik] = dispersion(dim,n,ham,vec_k[:,ik])
        end

        return vec_k,energies
    end

    function show_neighbors(lattice;maxn = 1,cutoff=2.0)
        dim = lattice.dim
        println("Possible hoppings")
        if dim==1
            for i1 = -maxn:maxn
                for i=1:lattice.numatoms
                    r0 = [i1]
                    for i=1:lattice.numatoms
                        ri = lattice.positions[i]
                        for j=i:lattice.numatoms
                            rj = lattice.positions[j]+r0
                            r = rationalize.(rj-ri)
                            r_pos = get_position(lattice,r)
                            if norm(r_pos) < cutoff
                                println("($i,$j), x:$(r[1])")
                            end
                        end
                    end
                end
            end

        elseif dim == 2
            for i1 = -maxn:maxn
                for i2 = -maxn:maxn
                    r0 = [i1,i2]
                    for i=1:lattice.numatoms
                        ri = lattice.positions[i]
                        for j=i:lattice.numatoms
                            rj = lattice.positions[j]+r0
                            r = rationalize.(rj-ri)
                            r_pos = get_position(lattice,r)
                            if norm(r_pos) < cutoff
                                println("($i,$j), x:$(r[1]), y:$(r[2])")
                            end
                        end
                    end
                end
            end
        elseif dim == 3
            for i1 = -maxn:maxn
                for i2 = -maxn:maxn
                    for i3 = -maxn:maxn
                        r0 = [i1,i2,i3]
                        for i=1:lattice.numatoms
                            ri = lattice.positions[i]
                            for j=i:lattice.numatoms
                                rj = lattice.positions[j]+r0
                                r = rationalize.(rj-ri)
                                r_pos = get_position(lattice,r)
                                if norm(r_pos) < cutoff
                                    println("($i,$j), x:$(r[1]), y:$(r[2]),z:$(r[3])")
                                end
                            end
                        end

                    end
                end
            end
        end
        return

    end

    function get_DOS(lattice,nk;nbins=100)
        dim = lattice.dim
        ham = hamiltonian_k(lattice)
        n = lattice.numatoms

        if dim == 1
            energies = zeros(Float64,n,nk)
            dk = 1/(nk-1)
            for ik = 1:nk
                k = (ik-1)*dk
            end
        elseif dim == 2
            energies = zeros(Float64,n,nk,nk)
            dk1 = 1/(nk-1)
            dk2 = 1/(nk-1)
            for ik1 = 1:nk
                k1 = (ik1-1)*dk1
                for ik2 = 1:nk
                    k2 = (ik2-1)*dk2
                    kx,ky = get_position_kspace(lattice,[k1,k2])
                    energy = dispersion(dim,n,ham,[kx,ky])
                    energies[:,ik1,ik2] = energy[:]

                end
            end
        elseif dim == 3
            energies = zeros(Float64,nk,nk,nk)
            dk1 = 1/(nk-1)
            dk2 = 1/(nk-1)
            dk3 = 1/(nk-1)

            for ik1 = 1:nk
                k1 = (ik1-1)*dk1
                for ik2 = 1:nk
                    k2 = (ik2-1)*dk2
                    for ik3=1:nk
                        k3= (ik3-1)*dk3
                        kx,ky,kz = get_position_kspace(lattice,[k1,k2,k3])
                        energy = dispersion(dim,n,ham,[kx,ky,kz])
                        energies[:,ik1,ik2,ik3] = energy[:]
                    end

                end
            end
        end
        energies = vec(energies)
#            println(energies)
        #pls = histogram(energies,bins=nbins,normalize=true)
        hist = fit(Histogram,energies,nbins=nbins)

        return hist
    end

    include("./Testfuncs.jl")

end # module