text:
using EntanglementDist
using SQLite
using Convex

# Since the solver crashes from time to time, one can
# run `./keepGenerating.sh` to restart this script
# automatically after crashing.

### Settings ###
# Database where output is saved
db = SQLite.DB("test.sqlite")
ps = collect(0.4:0.1:1)
δs = collect(0.1:0.05:1)
ϵ = 1e-4 # relative accuracy

# This is the name under which the results are saved
# N.B.
# The program type should be included in the state name as otherwise
# there is no way to tell them apart. `./plotData.jl` assumes
# that the state name is of the following form:
# `{statename}` for PPT program
# `1ext {statename}` for 1 extension
stateName = "pptRelax: Double Werner"
n_copies = 2
#Local dimensions of individual copies:
nA = 2
nB = 2
# local dimension of the desired maximally entangled state
k=2

"""
Generate one copy of the state. Change this to perform numerics on another state!
"""
function generateState(parameter)
  return wernerState(parameter)
  #return rStateCorrPhase(parameter)
end

# No need to change this
entries = length(ps)*length(δs)
times = []
firstRun = true
maxRetries = 20
retries = 0
foundNaN = false

# """
# Generate n copies of the state
# """
# function generateState(parameter, n_copies::Int)
#     return copies(generateState(parameter),  n_copies)
# end

"""
Save found (F, p_succ) to the database
"""
function saveData(F, p_succ, δ, stateName, p, eprF, ϵ)
  q = string("INSERT INTO `pptRelax` (`fidelity`, `p_succ`, `delta`, `state`, `p`, `eprFidelity`, `eps`)
    VALUES (" , F , ", ", p_succ ,", ", δ ,", '", stateName, "', ", p ,", ", eprF, ", ", ϵ ,")")
  SQLite.query(db, q)
end

"""
Check in the database if the combination (state, δ, p_succ) already occurs
"""
function hasBeenComputed(δ, stateName, p, ϵ)
  q = string("SELECT COUNT(*) FROM `pptRelax` WHERE `delta` = ", δ ," AND p = ", p ," AND `state` = '", stateName, "' AND eps <= '", ϵ,"'")
  result = SQLite.query(db, q)
  count = result[1][1].value

  return count > 0
end

"""
Convert seconds to hours, minuts and seconds
"""
function hms(timeleft)
  timeleft = round(Int, timeleft)
  hours = floor(Int, timeleft/3600)
  timeleft = timeleft - hours * 3600

  minutes = floor(Int, timeleft/60)
  timeleft = timeleft - minutes * 60

  return string(hours, "h ", minutes, "m ", timeleft, "s")
end

### Run the programs ###
# sometimes the solved gives an error for unknown reasons
# therefore, retry if an error occurs
while firstRun == true || (foundNaN == true && retries < maxRetries)
  entriesDone = 0
  firstRun = false
  foundNaN = false
  retries = retries + 1

  for p in ps
    for δ in δs
      # start timer
      starttime = time()

      # check if it is already computed
      computed = hasBeenComputed(δ, stateName, p, ϵ)

      # only compute data if no entry is available
      if !computed
        state = generateState(p)
        println("computing (δ, p, state) = (", δ, ", ", p, ", ", stateName, ")")

        # Alter the line below to run another program (e.g. PPT or k-ext)
        # The program type should be included in the state name as otherwise
        # there is no way to tell them apart.
        (problem, F, p_succ) = pptRelaxCopies(state, n_copies, nA, nB, k, δ, verbose = false, eps = ϵ, max_iters = 1e7)
        # (problem, F, p_succ) = PPTprogrammeNoTwirling1ExtPermSymOnlySym(state, 2^n_copies, 2^n_copies, n_copies, 2, δ, verbose = true, eps = 1e-4)
        eprF = eprFidelity(generateState(p))

        println("found (F, p_succ) = " , (F, p_succ))

        if isnan(F)
          foundNaN = true
        else
          saveData(F, p_succ, δ, stateName, p, real(eprF), ϵ)

          stoptime = time()
          times = [times; stoptime - starttime]
          timeleft = (entries - entriesDone) * mean(times)
          entriesDone = entriesDone + 1
          println("Time to go: ", hms(timeleft), ". computation time: ", hms(stoptime-starttime), " | ", entriesDone, "/" , entries , "\n")
        end
      else
        println("already computed (δ, p, state, ϵ) = (", δ, ", ", p, ", ", stateName, ", ", ϵ ,"). \n")
        entriesDone = entriesDone + 1
      end
    end
  end
end


text:
module WelcomeStuff

export WelcomeController

using Bukdu
import Bukdu.Actions: index
import Documenter.Utilities.DOM: @tags

struct WelcomeController <: ApplicationController
    conn::Conn
end

function index(c::WelcomeController)
    @tags h3 div a

    render(HTML, div(
        a[:href => "/"]("Back"),
        h3("Welcome"),
        div(
            "Hello World"
        )
    ))
end

end # module WelcomeStuff

module CalculatorStuff

export Calculator

using Bukdu
import Bukdu.Actions: index
import Documenter.Utilities.DOM: @tags

struct Calculator <: ApplicationController
    conn::Conn
end

function index(c::Calculator)
    @tags h3 div a

    x = get(c.params, :x, 0)
    y = get(c.params, :y, 0)

    rand_x = rand(1:10)
    rand_y = rand(1:10)

    render(HTML, div(
        a[:href => "/"]("Back"),
        h3("Calculator"),
        div(
            string(x, +, y, '=', x+y),
        ),
        div(
            a[:href => "/calculator?x=$rand_x&y=$rand_y"]("/calculator?x=$rand_x&y=$rand_y")
        )
    ))
end

end # module CalculatorStuff

module MainStuff

export MainController

using Bukdu
import Bukdu.Actions: index
import Documenter.Utilities.DOM: @tags

struct MainController <: ApplicationController
    conn::Conn
end

function index(::MainController)
    @tags h3 div a
    render(HTML, div(
        h3("Main"),
        div(
            a[:href => "/welcome"]("/welcome")
        ),
        div(
            a[:href => "/calculator"]("/calculator")
        ),
    ))
end

end # module MainStuff

if PROGRAM_FILE == basename(@__FILE__)

using Bukdu
import Bukdu.Actions: index

using .WelcomeStuff
using .CalculatorStuff
using .MainStuff

routes() do
    get("/", MainController, index)
    get("/welcome", WelcomeController, index)
    get("/calculator", Calculator, index)
end

Bukdu.start(8080)

Router.call(get, "/") #
# CLI.routes()

Base.JLOptions().isinteractive==0 && wait()

# Bukdu.stop()

end # if

ext: jl

max_stars_repo_name: wookay/TestBukdu

max_issues_repo_name: wookay/TestBukdu

max_forks_repo_name: wookay/TestBukdu


text:
# ## Getting started

using MLJ, RDatasets, Random
MLJ.color_off() # hide

data = dataset("datasets", "USArrests")
names(data)

# Let's have a look at the mean and standard deviation of each feature:

describe(data, :mean, :std)

# Let's extract the numerical component and coerce

X = select(data, Not(:State))
X = coerce(X, :UrbanPop=>Continuous);

# ## PCA pipeline
#
# PCA is usually best done after standardization but we won't do it here:

@load PCA pkg=MultivariateStats

pca_mdl = PCA(pratio=1)
pca = machine(pca_mdl, X)
fit!(pca)

W = transform(pca, X);

# W is the PCA'd data; here we've used default settings for PCA and it has recovered 2 components:

schema(W).names

# Let's inspect the fit:

r = report(pca)
cumsum(r.principalvars ./ r.tvar)

# In the second line we look at the explained variance with 1 then 2 PCA features and it seems that with 2 we almost completely recover all of the variance.

# ## More interesting data...

# Instead of just playing with toy data, let's load the orange juice data and extract only the columns corresponding to price data:

data = dataset("ISLR", "OJ")

X = select(data, [:PriceCH, :PriceMM, :DiscCH, :DiscMM, :SalePriceMM,
                  :SalePriceCH, :PriceDiff, :PctDiscMM, :PctDiscCH]);

# ### PCA pipeline

Random.seed!(1515)

@pipeline SPCA(std = Standardizer(),
               pca = PCA(pratio=1-1e-4))
spca_mdl = SPCA()
spca = machine(spca_mdl, X)
fit!(spca)
W = transform(spca, X)
names(W)

# What kind of variance can we explain?

r  = report(spca).reports[1]
cs = cumsum(r.principalvars ./ r.tvar)

# Let's visualise this

using PyPlot

figure(figsize=(8,6))

bar(1:length(cs), cs)
plot(1:length(cs), cs, color="red", marker="o")

xlabel("Number of PCA features", fontsize=14)
ylabel("Ratio of explained variance", fontsize=14)

savefig(joinpath(@OUTPUT, "ISL-lab-10-g1.svg")) # hide

# \figalt{PCA explained variance}{ISL-lab-10-g1.svg}

# So 4 PCA features are enough to recover most of the variance.

# ### Clustering

Random.seed!(1515)

@load KMeans pkg=Clustering
@pipeline SPCA2(std = Standardizer(),
                pca = PCA(),
                km = KMeans(k=3))

spca2_mdl = SPCA2()
spca2 = machine(spca2_mdl, X)
fit!(spca2)

assignments = report(spca2).reports[1].assignments
mask1 = assignments .== 1
mask2 = assignments .== 2
mask3 = assignments .== 3;

# Now we can  try visualising this

using PyPlot

figure(figsize=(8, 6))
for (m, c) in zip((mask1, mask2, mask3), ("red", "green", "blue"))
    plot(W[m, 1], W[m, 2], ls="none", marker=".", markersize=10, color=c)
end

xlabel("PCA-1", fontsize=13)
ylabel("PCA-2", fontsize=13)
legend(["Group 1", "Group 2", "Group 3"], fontsize=13)

savefig(joinpath(@OUTPUT, "ISL-lab-10-cluster.svg")) # hide

# \fig{ISL-lab-10-cluster.svg}
PyPlot.close_figs() # hide

max_stars_repo_name: ven-k/MLJTutorials

max_issues_repo_name: ven-k/MLJTutorials

max_forks_repo_name: ven-k/MLJTutorials


text:
function MOI.get(model::Optimizer, ::MOI.TerminationStatus)
	return model.status
end

function MOI.get(model::Optimizer, ::MOI.ObjectiveValue)
    if model.status == MOI.OPTIMIZE_NOT_CALLED
        @error "optimize! not called"
    end
    return model.inner.best_sol
end

function MOI.get(model::Optimizer, ::MOI.ObjectiveBound)
    if model.status == MOI.OPTIMIZE_NOT_CALLED
        @error "optimize! not called"
    end
    return model.inner.best_bound
end

function MOI.get(model::Optimizer, ::MOI.VariablePrimal, vi::MOI.VariableIndex)
    if model.status == MOI.OPTIMIZE_NOT_CALLED
        @error "optimize! not called"
    end
    check_inbounds(model, vi)
    return CS.value(model.inner.search_space[vi.value])
end

function set_status!(model::Optimizer, status::Symbol)
    if status == :Solved
        model.status = MOI.OPTIMAL
    elseif status == :Infeasible
        model.status = MOI.INFEASIBLE
    else
        model.status = MOI.OTHER_LIMIT
    end
end


text:
using HTTP

struct Fetcher
    cache::Dict{String,Any}
    Fetcher(cache::Nothing) = new(Dict{String,Any}())
    Fetcher(cache) = new(cache)
end

function geturl(testnet::Bool=false)
    if testnet
        return "http://localhost:18332"
    else
        return "http://localhost:8332"
    end
end

"""
    txfetch(tx_id::String) -> Tx

Returns the bitcoin transaction given its ID as an hexadecimal string.
"""
function txfetch(tx_id::String, testnet::Bool=false, fresh::Bool=false, fetcher::Fetcher=Fetcher(nothing))
    if fresh || !haskey(fetcher.cache, tx_id)
        url = string(geturl(testnet), "/rest/tx/", tx_id, ".bin")
        response = HTTP.request("GET", url)
        try
            response.status == 200
        catch
            error("Unexpected status: ", response.status)
        end
        raw = response.body
        if raw[5:6] == [0x00, 0x01]
            deleteat!(raw, 5:6)
            # flag = true # If present, always 0001, and indicates the presence of witness data
            tx = txparse(IOBuffer(raw), testnet)
            tx.locktime = bytes2int(raw[end-3:end], true)
        else
            tx = txparse(IOBuffer(raw), testnet)
        end
        if txid(tx) != tx_id
            error("not the same id : ", txid(tx),
                "\n             vs : ", tx_id)
        end
        fetcher.cache[tx_id] = tx
    end
    fetcher.cache[tx_id].testnet = testnet
    return fetcher.cache[tx_id]
end

abstract type TxComponent end

mutable struct TxIn <: TxComponent
    prev_tx::Array{UInt8,1}
    prev_index::Integer
    script_sig::Script
    sequence::Integer
    TxIn(prev_tx, prev_index) = new(prev_tx, prev_index, Script(nothing), 0xffffffff)
    TxIn(prev_tx, prev_index, script_sig, sequence=0xffffffff) = new(prev_tx, prev_index, script_sig, sequence)
end

function show(io::IO, z::TxIn)
    print(io, "\n", bytes2hex(z.prev_tx), ":", z.prev_index, "\n", z.script_sig)
end

"""
Takes a byte stream and parses the tx_input at the start
return a TxIn object
"""
function txinparse(s::Base.GenericIOBuffer)
    prev_tx = UInt8[]
    readbytes!(s, prev_tx, 32)
    reverse!(prev_tx)
    bytes = UInt8[]
    readbytes!(s, bytes, 4)
    prev_index = bytes2int(bytes, true)
    script_sig = scriptparse(s)
    readbytes!(s, bytes, 4)
    sequence = bytes2int(bytes, true)
    return TxIn(prev_tx, prev_index, script_sig, sequence)
end

"""
Returns the byte serialization of the transaction input
"""
function txinserialize(tx::TxIn)
    result = copy(tx.prev_tx)
    reverse!(result)
    append!(result, int2bytes(tx.prev_index, 4, true))
    append!(result, scriptserialize(tx.script_sig))
    append!(result, int2bytes(tx.sequence, 4, true))
    return result
end

function txin_fetchtx(tx::TxIn, testnet::Bool=false)
    return txfetch(bytes2hex(tx.prev_tx), testnet)
end

"""
Get the outpoint value by looking up the tx hash
Returns the amount in satoshi
"""
function txinvalue(txin::TxIn, testnet::Bool=false)
    tx = txin_fetchtx(txin, testnet)
    return tx.tx_outs[txin.prev_index + 1].amount
end

"""
Get the scriptPubKey by looking up the tx hash
Returns a Script object
"""
function txin_scriptpubkey(txin::TxIn, testnet::Bool=false)
    tx = txin_fetchtx(txin, testnet)
    return tx.tx_outs[txin.prev_index + 1].script_pubkey
end

struct TxOut <: TxComponent
    amount::Integer
    script_pubkey::Script
    TxOut(amount, script_pubkey) = new(amount, script_pubkey)
end

function show(io::IO, z::TxOut)
    print(io, "\n", z.script_pubkey, "\namout (BTC) : ", z.amount / 100000000)
end

"""
Takes a byte stream and parses the tx_output at the start
return a TxOut object
"""
function txoutparse(s::Base.GenericIOBuffer)
    bytes = UInt8[]
    readbytes!(s, bytes, 8)
    amount = bytes2int(bytes, true)
    script_pubkey = scriptparse(s)
    return TxOut(amount, script_pubkey)
end

"""
Returns the byte serialization of the transaction output
"""
function txoutserialize(tx::TxOut)
    result = int2bytes(tx.amount, 8, true)
    append!(result, scriptserialize(tx.script_pubkey))
    return result
end

mutable struct Tx <: TxComponent
    version::Integer
    tx_ins::Array{TxIn, 1}
    tx_outs::Array{TxOut, 1}
    locktime::Integer
    testnet::Bool
    Tx(version, tx_ins, tx_outs, locktime, testnet=false) = new(version, tx_ins, tx_outs, locktime, testnet)
end

function show(io::IO, z::Tx)
    print(io, "Transaction\n--------\nTestnet : ", z.testnet,
            "\nVersion : ", z.version,
            "\nLocktime : ", z.locktime,
            "\n--------\n",
            "\n", z.tx_ins,
            "\n--------\n",
            "\n", z.tx_outs)
end

"""
    txparse(s::Base.GenericIOBuffer, testnet::Bool=false) -> Tx

Returns a Tx object given a byte stream
"""
function txparse(s::Base.GenericIOBuffer, testnet::Bool=false)
    bytes = UInt8[]
    readbytes!(s, bytes, 4)
    version = bytes2int(bytes, true)
    num_inputs = read_varint(s)
    inputs = []
    for i in 1:num_inputs
        input = txinparse(s)
        push!(inputs, input)
    end
    num_outputs = read_varint(s)
    outputs = []
    for i in 1:num_outputs
        output = txoutparse(s)
        push!(outputs, output)
    end
    readbytes!(s, bytes, 4)
    locktime = bytes2int(bytes, true)
    return Tx(version, inputs, outputs, locktime, testnet)
end

"""
    txserialize(tx::Tx) -> Array{UInt8,1}

Returns the byte serialization of the transaction
"""
function txserialize(tx::Tx)
    result = int2bytes(tx.version, 4, true)
    append!(result, encode_varint(length(tx.tx_ins)))
    for tx_in in tx.tx_ins
        append!(result, txinserialize(tx_in))
    end
    append!(result, encode_varint(length(tx.tx_outs)))
    for tx_out in tx.tx_outs
        append!(result, txoutserialize(tx_out))
    end
    append!(result, int2bytes(tx.locktime, 4, true))
    return result
end

"""
Binary hash of the legacy serialization
"""
function txhash(tx::Tx)
    return reverse(hash256(txserialize(tx)))
end

"""
    txid(tx::Tx) -> String

Returns an hexadecimal string of the transaction hash
"""
function txid(tx::Tx)
    return bytes2hex(txhash(tx))
end

"""
    txfee(tx::Tx) -> Integer

Returns the fee of this transaction in satoshi
"""
function txfee(tx::Tx)
    input_sum, output_sum = 0, 0
    for tx_in in tx.tx_ins
        input_sum += txinvalue(tx_in, tx.testnet)
    end
    for tx_out in tx.tx_outs
        output_sum += tx_out.amount
    end
    return input_sum - output_sum
end

"""
    txsighash256(tx::Tx, input_index::Integer) -> Array{UInt8,1}

Returns the hash that needs to get signed for index input_index
"""
function txsighash256(tx::Tx, input_index::Integer)
    alt_tx_ins = TxIn[]
    for tx_in in tx.tx_ins
        alt_tx_in = TxIn(tx_in.prev_tx, tx_in.prev_index, Script(nothing), tx_in.sequence)
        push!(alt_tx_ins, alt_tx_in)
    end
    signing_input = alt_tx_ins[input_index + 1]
    script_pubkey = txin_scriptpubkey(signing_input, tx.testnet)
    signing_input.script_sig = script_pubkey
    alt_tx = Tx(
        tx.version,
        alt_tx_ins,
        tx.tx_outs,
        tx.locktime)
    result = UInt8[]
    append!(result, txserialize(alt_tx))
    append!(result, int2bytes(SIGHASH_ALL, 4, true))
    return hash256(result)
end

"""
    txsighash(tx::Tx, input_index::Integer) -> Integer

Returns the integer representation of the hash that needs to get
signed for index input_index
"""
function txsighash(tx::Tx, input_index::Integer)
    return bytes2int(txsighash256(tx, input_index))
end

"""
Returns whether the input has a valid signature
"""
function txinputverify(tx::Tx, input_index)
    tx_in = tx.tx_ins[input_index + 1]
    z = txsighash(tx, input_index)
    combined_script = Script(copy(tx_in.script_sig.instructions))
    append!(combined_script.instructions,
            txin_scriptpubkey(tx_in, tx.testnet).instructions)
    return scriptevaluate(combined_script, z)
end

"""
Verify this transaction
"""
function txverify(tx::Tx)
    if txfee(tx) < 0
        return false
    end
    for i in 1:length(tx.tx_ins)
        if !txinputverify(tx, i - 1)
            return false
        end
    end
    return true
end

"""
Signs the input using the private key
"""
function txsigninput(tx::Tx, input_index::Integer, private_key::PrivateKey)
    z = txsighash(tx, input_index)
    sig = pksign(private_key, z)
    txpushsignature(tx, input_index, z, sig, private_key.𝑃)
end

function txpushsignature(tx::Tx, input_index::Integer, z::Integer, sig::Signature, pubkey::S256Point)
    der = sig2der(sig)
    append!(der, int2bytes(SIGHASH_ALL))
    sec = point2sec(pubkey)
    script_sig = Script([der, sec])
    tx.tx_ins[input_index + 1].script_sig = script_sig
    return txinputverify(tx, input_index)
end


text:
function fastphase()
    return nothing # todo
end

"""
    fastphase_estim_param(xdata; ...)

Runs fastPHASE to estimate `r`, `θ`, `α` parameters. Will run different initial
EM starts in parallel if Julia is started with multiple threads. `θ` (emission)
probabilities) values will automatically be flipped so that `θ[i, j]` is the
probability of observing allele A2 (usually major) in the PLINK bam file at
SNP `i` of haplotype `j`.

# Inputs
`xdata`: A `String` for binary PLINK file (without `.bed/bim/fam` extensions)
    or a `SnpData` (see SnpArrays.jl).

# Optional inputs
`n`: Number of samples used to fit HMM in fastPHASE. Defaults to sample size in `xdata`.
`T`: Number of different initial conditions for EM. Different initial conditions
    will be run in parallel in `Threads.nthreads()` number of threads.
`K`: Number of haplotype clusters. Defaults to 12
`C`: Number of EM iterations before convergence. Defaults to 10.
`outfile`: Extension of output alpha, theta, and r file names. Defaults to `fastphase_out`
`outdir`: Output directory. By default all output will be stored in new folder in
    `knockoffs` in the current directory.
`fastphase_infile`: Filename of fastPHASE's input, which is the decompressed PLINK
    genotypes readable by fastPHASE. Defaults to `fastphase.inp`.
"""
function fastphase_estim_param(
    xdata::SnpData;
    n::Int = size(xdata.snparray, 1),
    T::Int = Threads.nthreads(),
    K::Int = 12,
    C::Int = 10,
    outfile::AbstractString = "fastphase_out",
    outdir::AbstractString = "knockoffs",
    fastphase_infile::AbstractString = joinpath(outdir, "fastphase.inp"),
    )
    isdir(outdir) || mkdir(outdir)
    x = xdata.snparray
    n ≤ size(x, 1) || error("n must be smaller than the number of samples!")
    sampleid = xdata.person_info[!, :iid]
    # create input format for fastPHASE software
    p = size(x, 2)
    open(fastphase_infile, "w") do io
        println(io, n)
        println(io, p)
        for i in 1:n
            println(io, "ID ", sampleid[i])
            # print genotypes for each sample on 2 lines. The "1" for heterozygous
            # genotypes will always go on the 1st line.
            for j in 1:p
                if x[i, j] == 0x00
                    print(io, 0)
                elseif x[i, j] == 0x02 || x[i, j] == 0x03
                    print(io, 1)
                else
                    print(io, '?')
                end
            end
            print(io, "\n")
            for j in 1:p
                if x[i, j] == 0x00 || x[i, j] == 0x02
                    print(io, 0)
                elseif x[i, j] == 0x03
                    print(io, 1)
                else
                    print(io, '?')
                end
            end
            print(io, "\n")
        end
    end
    # run fastPHASE, each thread runs a different start
    Threads.@threads for i in 1:T
        tmp_out = joinpath(outdir, "tmp$(i)")
        seed = i
        s = 0 # no simulation
        U = 0 # no simulation
        H = 0
        run(`$fastPHASE_EXE -T1 -S$seed -K$K -C$C -o$tmp_out -Pp -H$H -s$s -U$U $fastphase_infile`)
    end
    # aggregate results into final output
    r, θ, α = _combine_fastphase_output(outdir, T, outfile=outfile)
    # clean up
    for i in 1:T
        # rm(joinpath(outdir, "tmp$(i)_rhat.txt"), force=true)
        # rm(joinpath(outdir, "tmp$(i)_thetahat.txt"), force=true)
        # rm(joinpath(outdir, "tmp$(i)_alphahat.txt"), force=true)
        # rm(joinpath(outdir, "tmp$(i)_origchars"), force=true)
        rm(joinpath(outdir, "tmp$(i)_hapsfrommodel.out"), force=true)
        rm(joinpath(outdir, "tmp$(i)_hapguess_switch.out"), force=true)
        rm(joinpath(outdir, "tmp$(i)_sampledHgivG.txt"), force=true)
    end
    return r, θ, α
end

fastphase_estim_param(xdata::AbstractString; args...) =
    fastphase_estim_param(SnpData(plinkfile); args...)

"""
    process_fastphase_output(filename; [T])

Reads r, θ, α into memory, averaging over `T` simulations. θ (emission probabilities)
must be flipped sometimes depending on which allele was defined as "allele 1".
fastPHASE simply uses whichever allele was observed first in sample 1 haplotype 1.
Thus, in sample 1, genotypes that start with "10" or "1?" must be flipped, and
this info is provided in the "_origchars" file.

# Inputs
`filename`: Path to the `rhat.txt`, `thetahat.txt`, `alphahat.txt` and
    `_origchars` files are stored. E.g. Use `out`
    if your output files are `out_rhat.txt`, `out_thetahat.txt`, `out_alphahat.txt`
`T`: Number of different runs excuted by fastPHASE. This is the number of different
    initial conditions used for EM algorithm. All files `rhat.txt`, `thetahat.txt`,
    `alphahat.txt` would therefore have `T × p` rows
"""
function process_fastphase_output(
    filename::AbstractString;
    T::Int=1,
    )
    # read full data
    rfile = "$(filename)_rhat.txt" # T*p × 1
    θfile = "$(filename)_thetahat.txt" # T*p × K
    αfile = "$(filename)_alphahat.txt" # T*p × K
    charfile = "$(filename)_origchars" # p × 2 file
    isfile(rfile) && isfile(θfile) && isfile(αfile) && isfile(charfile) ||
        error("Files not found!")
    r_full = readdlm(rfile, comments=true, comment_char = '>', header=false)
    θ_full = readdlm(θfile, comments=true, comment_char = '>', header=false)
    α_full = readdlm(αfile, comments=true, comment_char = '>', header=false)
    flip_idx = flip_θ_index(charfile)
    # compute averages across T simulations as suggested by Scheet et al 2006
    p = Int(size(r_full, 1) / T)
    K = size(θ_full, 2)
    r, θ, α = zeros(p), zeros(p, K), zeros(p, K)
    for i in 1:T
        rows = (i - 1) * p + 1:p*i
        r .+= @view(r_full[rows])
        α .+= @view(α_full[rows, :])
        # make sure to flip θ values if the allele 1 and 2 was flipped by fastPHASE
        θtmp = @view(θ_full[rows, :])
        for i in eachindex(flip_idx)
            if flip_idx[i]
                for j in 1:K
                    θtmp[i, j] = 1 - θtmp[i, j]
                end
            end
        end
        θ .+= θtmp
    end
    r ./= T
    θ ./= T
    α ./= T
    α ./= sum(α, dims = 2) # normalize rows to sum to 1
    return r, θ, α
end

function flip_θ_index(charfile::AbstractString)
    cfile = CSV.read(charfile, DataFrame, delim='\t', header=false, skipto=2)
    flip_idx = falses(size(cfile, 1))
    for i in 1:length(flip_idx)
        if cfile[i, 2] == "10" || cfile[i, 2] == "1?" || cfile[i, 2] == 10
            flip_idx[i] = true
        end
    end
    return flip_idx
end

function _combine_fastphase_output(
    outdir::AbstractString,
    T::Int;
    outfile::AbstractString="fastphase_out"
    )
    r, θ, α = process_fastphase_output(joinpath(outdir, "tmp1"); T=1)
    for i in 2:T
        rtmp, θtmp, αtmp = process_fastphase_output(joinpath(outdir, "tmp$i"); T=1)
        r .+= rtmp
        θ .+= θtmp
        α .+= αtmp
    end
    r ./= T
    θ ./= T
    α ./= T
    α ./= sum(α, dims = 2) # normalize rows to sum to 1
    # save averaged results
    writedlm(joinpath(outdir, outfile * "_rhat.txt"), r, ' ')
    writedlm(joinpath(outdir, outfile * "_thetahat.txt"), θ, ' ')
    writedlm(joinpath(outdir, outfile * "_alphahat.txt"), α, ' ')
    open(joinpath(outdir, outfile * "_origchars"), "w") do io
        println(io, length(r))
        for i in 1:length(r)
            println(io, "2\t01") # since all theta are flipped, all origchar will be 01
        end
    end
    return r, θ, α
end

# function merge_knockoffs_with_original(
#     xdata::SnpData,
#     x̃data::SnpData,
#     des::AbstractString
#     )
#     d = Dict{AbstractString, SnpData}("original"=>xdata, "knockoff"=>x̃data)
#     write_plink(des, merge_plink(d))
# end

function merge_knockoffs_with_original(
    xdata::SnpData,
    x̃data::SnpData;
    des::AbstractString = "knockoff"
    )
    n, p = size(xdata)
    x, x̃ = xdata.snparray, x̃data.snparray
    xfull = SnpArray(des * ".bed", n, 2p)
    original, knockoff = sizehint!(Int[], p), sizehint!(Int[], p)
    for i in 1:p
        # decide which of original or knockoff SNP comes first
        orig, knoc = rand() < 0.5 ? (2i - 1, 2i) : (2i, 2i - 1)
        copyto!(@view(xfull[:, knoc]), @view(x̃[:, i]))
        copyto!(@view(xfull[:, orig]), @view(x[:, i]))
        push!(original, orig)
        push!(knockoff, knoc)
    end
    # copy fam files
    cp(xdata.srcfam, des * ".fam", force=true)
    # copy bim file, knockoff SNPs end in ".k"
    new_bim = copy(xdata.snp_info)
    empty!(new_bim)
    for i in 1:p
        if original[i] < knockoff[i]
            push!(new_bim, xdata.snp_info[i, :])
            push!(new_bim, x̃data.snp_info[i, :])
        else
            push!(new_bim, x̃data.snp_info[i, :])
            push!(new_bim, xdata.snp_info[i, :])
        end
    end
    CSV.write(des * ".bim", new_bim, delim='\t', header=false)
end


text:
using PermPlain

# Test performance of permpower vs permpower2.

n = 100
p = randperm(n)
npow = 100

m = 100000
ts = 0.0
for k in 1:m
    t1 = @elapsed PermPlain.permpower2(p,np)
    t2 = @elapsed PermPlain.permpower(p,np)
    ts = ts + t1/t2
end
println(ts/m)


text:
abstract type TTM_Potential end
include("constants.jl")

@inline function compute_M_site_crd!(ttm::TTM_Potential, coords::AbstractMatrix{Float64})
    """
    Takes the 3xN coordinates you get coming in and stores the coordinates with M site in the TTM_Potential struct.
    """
    num_M_sites::Int = 0
    for i in 1:Int(size(coords, 2) / 3) * 4
        if i % 4 == 0
            @inbounds ttm.M_site_coords[i] = convert(SVector{3, Float64}, ttm.constants.γ_1 * coords[:,i-3 - num_M_sites] + ttm.constants.γ_2 * (coords[:,i-2 - num_M_sites] + coords[:, i-1 - num_M_sites]))
            num_M_sites += 1
        else
            @inbounds ttm.M_site_coords[i] = convert(SVector{3, Float64}, coords[:, (i - num_M_sites)])
        end
    end
end

@inline function triu_linear_index(i::Int, j::Int, n::Int)
    return Int(n*(n-1)/2 - (n-i+1)*(n-i)/2 + j - i)
end

function distance_vectors(coords::Vector{SVector{3, Float64}})
    # make this take the storage as most time will be spent allocating
    n::Int = length(coords)
    distances = [@SVector zeros(3) for _ in 1:(n*(n-1)/2)]
    for i in 1:n-1
        for j in i+1:n
            distances[triu_linear_index(i,j,n)] = coords[i]- coords[j]
        end
    end
    return distances
end

function get_dipoles(ttm::TTM_Potential)
    dipoles = zeros(3, ttm.num_waters)
    for i in 1:length(ttm.M_site_coords)
        if i % 4 != 1 # static dipoles from Partridge-Schwenke
            dipoles[:, (i-1)÷4+1] += ttm.elec_data.q[i] * ttm.M_site_coords[i]
        else # add in the induced dipoles
            dipoles[:, i÷4+1] += ttm.elec_data.dipoles[i]   # O
            dipoles[:, i÷4+1] += ttm.elec_data.dipoles[i+1] # H
            dipoles[:, i÷4+1] += ttm.elec_data.dipoles[i+2] # H
            dipoles[:, i÷4+1] += ttm.elec_data.dipoles[i+3] # M
        end
    end
    return dipoles / chargecon() * debye()
end

function get_dipole_magnitudes(ttm::TTM_Potential)
    return norm.(eachcol(get_dipoles(ttm)))
end

function PS_energies_and_gradients!(ttm::TTM_Potential, coords::AbstractMatrix{Float64}, grads_q::Vector{MArray{Tuple{3, 3, 3}, Float64, 3, 27}}, grads::Union{Matrix{Float64}, Nothing}=nothing, ttm3f::Bool=false)
    E_int::Float64 = 0.0
    q3 = @MVector zeros(3)
    dq3 = @MArray zeros(3,3,3)
    for n in 1:ttm.num_waters

        if (grads !== nothing)
            @inbounds E_int += pot_nasa(convert(SMatrix{3, 3, Float64}, @view(coords[:, (3*n-2):3*n])), @view(grads[:, (3*n-2):3*n]))
            @inbounds dms_nasa!(convert(SMatrix{3, 3, Float64}, @view(coords[:, (3*n-2):3*n])), q3, dq3, ttm3f)
        else
            @inbounds E_int += pot_nasa(convert(SMatrix{3, 3, Float64}, @view(coords[:, (3*n-2):3*n])), nothing)
            @inbounds dms_nasa!(convert(SMatrix{3, 3, Float64}, @view(coords[:, (3*n-2):3*n])), q3, nothing, ttm3f)
        end

        # TTM_Potential2.1-F assignment
        tmp::Float64 = 0.5 * ttm.constants.γ_M / (1.0 - ttm.constants.γ_M)
        @inbounds ttm.elec_data.q[4*(n-1) + 1] = 0.0                                  # O
        @inbounds ttm.elec_data.q[4*(n-1) + 2] = q3[2] + tmp * (q3[2] + q3[3])        # H1
        @inbounds ttm.elec_data.q[4*(n-1) + 3] = q3[3] + tmp * (q3[2] + q3[3])        # H2
        @inbounds ttm.elec_data.q[4*(n-1) + 4] = q3[1] / (1.0 - ttm.constants.γ_M) # M

        if (grads === nothing)
            continue
        end
        # charge gradients
        @inbounds grads_q[n][:, 1, :] = dq3[:, 1, :] +   tmp*(dq3[:, 1, :] + dq3[:, 2, :])
        @inbounds grads_q[n][:, 2, :] = dq3[:, 2, :] +   tmp*(dq3[:, 2, :] + dq3[:, 1, :])
        @inbounds grads_q[n][:, 3, :] = dq3[:, 3, :] - 2*tmp*(dq3[:, 1, :] + dq3[:, 2, :])
    end

    return E_int
end

function add_electrostatic_and_dms_derivatives!(ttm::TTM_Potential, grads::AbstractMatrix{Float64}, grads_E::Vector{MVector{3, Float64}}, grads_q::Vector{MArray{Tuple{3, 3, 3}, Float64, 3, 27}})
    # add electrostatic derivatives
    for n in 1:ttm.num_waters
        # add O H H gradients from PS and electrostatics (which has M site)
        for k in 1:3
            @inbounds @views grads[:, 3 * (n-1) + k] += grads_E[4 * (n-1) + k]
        end

        # redistribute the M-site derivatives
        @inbounds @views grads[:, 3 * (n-1) + 1] += (1.0 - ttm.constants.γ_M) * grads_E[4*(n-1) + 4] # O
        @inbounds @views grads[:, 3 * (n-1) + 2] +=  0.5 * ttm.constants.γ_M  * grads_E[4*(n-1) + 4] # H
        @inbounds @views grads[:, 3 * (n-1) + 3] +=  0.5 * ttm.constants.γ_M  * grads_E[4*(n-1) + 4] # H
    end

    # derivatives from the adjustable charges of the NASA's PES
    for n in 1:ttm.num_waters
        i_O  = 3 * (n-1) + 1
        i_H1 = 3 * (n-1) + 2
        i_H2 = 3 * (n-1) + 3

        @inbounds grads[:, i_H1] += (grads_q[n][:, 1, 1]*ttm.elec_data.ϕ[4*(n-1) + 2]   # phi(h1)
                         + grads_q[n][:, 2, 1]*ttm.elec_data.ϕ[4*(n-1) + 3]   # phi(h2)
                         + grads_q[n][:, 3, 1]*ttm.elec_data.ϕ[4*(n-1) + 4])  # phi(M)

        @inbounds grads[:, i_H2] += (grads_q[n][:, 1, 2]*ttm.elec_data.ϕ[4*(n-1) + 2]   # phi(h1)
                         + grads_q[n][:, 2, 2]*ttm.elec_data.ϕ[4*(n-1) + 3]   # phi(h2)
                         + grads_q[n][:, 3, 2]*ttm.elec_data.ϕ[4*(n-1) + 4])  # phi(M)

        @inbounds grads[:, i_O]  += (grads_q[n][:, 1, 3]*ttm.elec_data.ϕ[4*(n-1) + 2]   # phi(h1)
                         + grads_q[n][:, 2, 3]*ttm.elec_data.ϕ[4*(n-1) + 3]   # phi(h2)
                         + grads_q[n][:, 3, 3]*ttm.elec_data.ϕ[4*(n-1) + 4])  # phi(M)
    end
end

max_stars_repo_name: heindelj/MolecularPotentials

max_issues_repo_name: heindelj/MolecularPotentials

max_forks_repo_name: heindelj/MolecularPotentials


text:
_has_aarch64_sve() = (Base.libllvm_version ≥ v"11") && (Base.BinaryPlatforms.CPUID.test_cpu_feature(Base.BinaryPlatforms.CPUID.JL_AArch64_sve))

if Int === Int64
    @noinline vscale() = ccall("llvm.vscale.i64", llvmcall, Int64, ())
else
    @noinline vscale() = ccall("llvm.vscale.i32", llvmcall, Int32, ())
end

# TODO: find actually support SVE
# _dynamic_register_size() = _has_aarch64_sve() ? 16vscale() : 16
function _dynamic_register_size()
  _has_aarch64_sve() || return 16
  llvmargs = get(ENV, "JULIA_LLVM_ARGS", nothing)
  llvmargs === nothing && return 16
  m = match(r"--aarch64-sve-vector-bits-min=(\d+)", llvmargs)
  m === nothing && return 16
  length(m.captures) == 0 && return 16
  return max(16, parse(Int, last(m.captures)) ÷ 8)
end

function _set_sve_vector_width!(bytes = _dynamic_register_size())
    @eval begin
        register_size() = StaticInt{$bytes}()
        simd_integer_register_size() = StaticInt{$bytes}()
    end
    nothing
end

if _has_aarch64_sve()# && !(Bool(has_feature(Val(:aarch64_sve))))
    has_feature(::Val{:aarch64_sve_cpuid}) = True()
    _set_sve_vector_width!()
else
    # has_feature(::Val{:aarch64_svejl}) = False()
    register_size() = StaticInt{16}()
    simd_integer_register_size() = StaticInt{16}()
end

function reset_extra_features!()
    drs = _dynamic_register_size()
    register_size() ≠ drs && _set_sve_vector_width!(drs)
    hassve = _has_aarch64_sve()
    if hassve ≠ has_feature(Val(:aarch64_sve_cpuid))
        @eval has_feature(::Val{:aarch64_sve_cpuid}) = $(Expr(:call, hassve ? :True : :False))
    end
end

fma_fast() = True()
register_count() = StaticInt{32}()
has_opmask_registers() = has_feature(Val(:aarch64_sve_cpuid))

fast_int64_to_double() = True()

fast_half() = False()


text:
struct Strang <: QuadratureName end
const strang = Strang()

function Quadrature(p::Polytope,::Strang,degree::Integer)
  msg = """\n
  `strang` quadrature rule only available for simplices.
  Use `tensor_product` for n-cubes.
  """
  @assert is_simplex(p) msg
  if num_dims(p) == 2
    quad = _strang_quad_tri(degree)
  elseif num_dims(p) == 3
    quad = _strang_quad_tet(degree)
  else
    msg = """\n
    `strang` quadrature rule only available for tris and tets.
    Use `duffy` for other simplices.
    """
    @unreachable msg
  end
  quad
end

const _strang_tri_k2n = Dict(1=>1,2=>3,3=>4,4=>6,5=>7,7=>13,9=>19,11=>28)

function _strang_quad_tri(degree)
  if ! haskey(_strang_tri_k2n,degree)
    msg = """\n
      `strang` quadrature rule not implemented for degree = $degree on a triangle.
      Implemented degrees are in $(sort(collect(keys(_strang_tri_k2n)))).
      Use `duffy` instead.
    """
    error(msg)
  end
  n = _strang_tri_k2n[degree]
  x = Vector{VectorValue{2,Float64}}(undef,n)
  w = Vector{Float64}(undef,n)
  if degree == 1
    a = 1.0/3.0
    b = 1.0/2.0
    x[1] = (a,a)
    w[1] = b
  elseif degree == 2
    a = 2.0/3.0
    b = 1.0/6.0
    x[1] = (b,b)
    x[2] = (a,b)
    x[3] = (b,a)
    w[1] = b
    w[2] = b
    w[3] = b
  elseif degree == 3
    a = 1.0/3.0
    b = 1.0/5.0
    c = 3.0/5.0
    d = -27.0/96.0
    e = 25.0/96.0
    x[1] = (a,a)
    x[2] = (b,b)
    x[3] = (c,b)
    x[4] = (b,c)
    w[1] = d
    w[2] = e
    w[3] = e
    w[4] = e
  elseif degree == 4
    ex1 = 0.816847572980459
    et1 = 0.091576213509771
    ez1 = 0.091576213509771
    ex2 = 0.108103018168070
    et2 = 0.445948490915965
    ez2 = 0.445948490915965
    a = 0.054975870996713638
    b = 0.1116907969117165
    x[1] = (et1,ez1)
    x[2] = (ez2,ex2)
    x[3] = (ex1,et1)
    x[4] = (ex2,et2)
    x[5] = (et2,ez2)
    x[6] = (ez1,ex1)
    w[1]  = a
    w[2]  = b
    w[3]  = a
    w[4]  = b
    w[5]  = b
    w[6]  = a
  elseif degree == 5
    a = 1.0 / 3.0
    b = (9.0+2.0*sqrt(15.0))/21.0
    c = (6.0-sqrt(15.0))/21.0
    d = (9.0-2.0*sqrt(15.0))/21.0
    e = (6.0+sqrt(15.0))/21.0
    w1 = 0.1125
    w2 = (155.0-sqrt(15.0))/2400.0
    w3 = (155.0+sqrt(15.0))/2400.0
    x[1] = (a,a)
    x[2] = (b,c)
    x[3] = (c,b)
    x[4] = (c,c)
    x[5] = (d,e)
    x[6] = (e,d)
    x[7] = (e,e)
    w[1] = w1
    w[2] = w2
    w[3] = w2
    w[4] = w2
    w[5] = w3
    w[6] = w3
    w[7] = w3
  elseif degree == 7
    a = 0.333333333333333
    b = 0.479308067841923
    c = 0.869739794195568
    d = 0.638444188569809
    e = 0.260345966079038
    f = 0.065130102902216
    g = 0.312865496004875
    h = 0.048690315425316
    w1=-0.149570044467670/2.0
    w2= 0.175615257433204/2.0
    w3= 0.053347235608839/2.0
    w4= 0.077113760890257/2.0
    x[ 1] = (a,a)
    x[ 2] = (e,e)
    x[ 3] = (b,e)
    x[ 4] = (e,b)
    x[ 5] = (f,f)
    x[ 6] = (c,f)
    x[ 7] = (f,c)
    x[ 8] = (d,g)
    x[ 9] = (d,h)
    x[10] = (g,d)
    x[11] = (g,h)
    x[12] = (h,d)
    x[13] = (h,g)
    w[ 1] = w1
    w[ 2] = w2
    w[ 3] = w2
    w[ 4] = w2
    w[ 5] = w3
    w[ 6] = w3
    w[ 7] = w3
    w[ 8] = w4
    w[ 9] = w4
    w[10] = w4
    w[11] = w4
    w[12] = w4
    w[13] = w4
  elseif degree == 9
    a = 1.0 / 3.0
    b = 0.02063496160252593
    c = 0.4896825191987370
    d = 0.1258208170141290
    e = 0.4370895914929355
    f = 0.6235929287619356
    g = 0.1882035356190322
    r = 0.9105409732110941
    s = 0.04472951339445297
    t = 0.7411985987844980
    u = 0.03683841205473626
    v = 0.22196298916076573
    w1 = 0.09713579628279610/2.0
    w2 = 0.03133470022713983/2.0
    w3 = 0.07782754100477543/2.0
    w4 = 0.07964773892720910/2.0
    w5 = 0.02557767565869810/2.0
    w6 = 0.04328353937728940/2.0
    x[ 1] = (a,a)
    x[ 2] = (b,c)
    x[ 3] = (c,b)
    x[ 4] = (c,c)
    x[ 5] = (d,e)
    x[ 6] = (e,d)
    x[ 7] = (e,e)
    x[ 8] = (f,g)
    x[ 9] = (g,f)
    x[10] = (g,g)
    x[11] = (r,s)
    x[12] = (s,r)
    x[13] = (s,s)
    x[14] = (t,u)
    x[15] = (t,v)
    x[16] = (u,t)
    x[17] = (u,v)
    x[18] = (v,t)
    x[19] = (v,u)
    w[ 1] = w1
    w[ 2] = w2
    w[ 3] = w2
    w[ 4] = w2
    w[ 5] = w3
    w[ 6] = w3
    w[ 7] = w3
    w[ 8] = w4
    w[ 9] = w4
    w[10] = w4
    w[11] = w5
    w[12] = w5
    w[13] = w5
    w[14] = w6
    w[15] = w6
    w[16] = w6
    w[17] = w6
    w[18] = w6
    w[19] = w6
  elseif degree == 11
    a = 1.0 / 3.0
    b = 0.9480217181434233
    c = 0.02598914092828833
    d = 0.8114249947041546
    e = 0.09428750264792270
    f = 0.01072644996557060
    g = 0.4946367750172147
    p = 0.5853132347709715
    q = 0.2073433826145142
    r = 0.1221843885990187
    s = 0.4389078057004907
    t = 0.6779376548825902
    u = 0.04484167758913055
    v = 0.27722066752827925
    h = 0.8588702812826364
    z = 0.0
    y = 0.1411297187173636
    w1 = 0.08797730116222190/2.0
    w2 = 0.008744311553736190/2.0
    w3 = 0.03808157199393533/2.0
    w4 = 0.01885544805613125/2.0
    w5 = 0.07215969754474100/2.0
    w6 = 0.06932913870553720/2.0
    w7 = 0.04105631542928860/2.0
    w8 = 0.007362383783300573/2.0
    x[ 1] = (a,a)
    x[ 2] = (b,c)
    x[ 3] = (c,b)
    x[ 4] = (c,c)
    x[ 5] = (d,e)
    x[ 6] = (e,d)
    x[ 7] = (e,e)
    x[ 8] = (f,g)
    x[ 9] = (g,f)
    x[10] = (g,g)
    x[11] = (p,q)
    x[12] = (q,p)
    x[13] = (q,q)
    x[14] = (r,s)
    x[15] = (s,r)
    x[16] = (s,s)
    x[17] = (t,u)
    x[18] = (t,v)
    x[19] = (u,t)
    x[20] = (u,v)
    x[21] = (v,t)
    x[22] = (v,u)
    x[23] = (h,z)
    x[24] = (h,y)
    x[25] = (z,h)
    x[26] = (z,y)
    x[27] = (y,h)
    x[28] = (y,z)
    w[ 1] = w1
    w[ 2] = w2
    w[ 3] = w2
    w[ 4] = w2
    w[ 5] = w3
    w[ 6] = w3
    w[ 7] = w3
    w[ 8] = w4
    w[ 9] = w4
    w[10] = w4
    w[11] = w5
    w[12] = w5
    w[13] = w5
    w[14] = w6
    w[15] = w6
    w[16] = w6
    w[17] = w7
    w[18] = w7
    w[19] = w7
    w[20] = w7
    w[21] = w7
    w[22] = w7
    w[23] = w8
    w[24] = w8
    w[25] = w8
    w[26] = w8
    w[27] = w8
    w[28] = w8
  else
    @unreachable
  end
  GenericQuadrature(x,w,"Strang quadrature of degree $degree on the reference triangle.")
end

const _strang_tet_k2n = Dict(1=>1,2=>4,3=>5,4=>11,5=>14)

function _strang_quad_tet(degree)
  if ! haskey(_strang_tet_k2n,degree)
    msg = """\n
      `strang` quadrature rule not implemented for degree = $degree on a tet.
      Implemented degrees are in $(sort(collect(keys(_strang_tet_k2n)))).
      Use `duffy` instead.
    """
    error(msg)
  end
  n = _strang_tet_k2n[degree]
  x = Vector{VectorValue{3,Float64}}(undef,n)
  w = Vector{Float64}(undef,n)
  if degree==1
    a = 1.0/4.0
    b = 1.0/6.0
    x[1] = (a,a,a)
    w[1] = b
  elseif degree==2
    a = 0.5854101966249685
    b = 0.1381966011250105
    c = 1.0/24.0
    x[1] = (b,b,b)
    x[2] = (a,b,b)
    x[3] = (b,a,b)
    x[4] = (b,b,a)
    w[1] = c
    w[2] = c
    w[3] = c
    w[4] = c
  elseif degree==3
    a = 1.0/4.0
    b = 1.0/6.0
    c = 1.0/2.0
    d = -2.0/15.0
    e = 1.5/20.0
    x[1] = (a,a,a)
    x[2] = (b,b,b)
    x[3] = (c,b,b)
    x[4] = (b,c,b)
    x[5] = (b,b,c)
    w[1] = d
    w[2] = e
    w[3] = e
    w[4] = e
    w[5] = e
  elseif degree== 4
    a = 0.3994035761667992
    b = 0.1005964238332008
    c = (343.0/7500.0)/6.0
    d = (56.0/375.0)/6.0
    e = 1.0/4.0
    f = 11.0/14.0
    g = 1.0/14.0
    h = (-148.0/1875.0)/6.0
    x[ 1] = (e,e,e)
    x[ 2] = (f,g,g)
    x[ 3] = (g,f,g)
    x[ 4] = (g,g,f)
    x[ 5] = (g,g,g)
    x[ 6] = (a,a,b)
    x[ 7] = (a,b,a)
    x[ 8] = (a,b,b)
    x[ 9] = (b,a,a)
    x[10] = (b,a,b)
    x[11] = (b,b,a)
    w[ 1] = h
    w[ 2] = c
    w[ 3] = c
    w[ 4] = c
    w[ 5] = c
    w[ 6] = d
    w[ 7] = d
    w[ 8] = d
    w[ 9] = d
    w[10] = d
    w[11] = d
  elseif degree == 5
    a = 0.0673422422100983
    b = 0.3108859192633005
    c = 0.7217942490673264
    d = 0.0927352503108912
    e = 0.4544962958743506
    f = 0.0455037041256494
    p = 0.1126879257180162/6.0
    q = 0.0734930431163619/6.0
    r = 0.0425460207770812/6.0
    x[ 1] = (a,b,b)
    x[ 2] = (b,a,b)
    x[ 3] = (b,b,a)
    x[ 4] = (b,b,b)
    x[ 5] = (c,d,d)
    x[ 6] = (d,c,d)
    x[ 7] = (d,d,c)
    x[ 8] = (d,d,d)
    x[ 9] = (e,e,f)
    x[10] = (e,f,e)
    x[11] = (e,f,f)
    x[12] = (f,e,e)
    x[13] = (f,e,f)
    x[14] = (f,f,e)
    w[ 1] = p
    w[ 2] = p
    w[ 3] = p
    w[ 4] = p
    w[ 5] = q
    w[ 6] = q
    w[ 7] = q
    w[ 8] = q
    w[ 9] = r
    w[10] = r
    w[11] = r
    w[12] = r
    w[13] = r
    w[14] = r
  else
    @unreachable
  end
  GenericQuadrature(x,w,"Strang quadrature of degree $degree on the reference tet.")
end


text:
module FEMRefinements

using JuAFEM
# using JuAFEM:AbstractCell, Cell

using LinearAlgebra
using Logging

import JuAFEM:vertices, edges, faces

export addcellset!, addfaceset!, addnodeset!,
compute_vertex_values, default_interpolation, edges, faces, getcells, getcellset, getcellsets,
getcelltype, getcoordinates!, getdim, getfaceset, getfacesets,
getnnodes, getnodeset, getnodesets, n_faces_per_cell, nnodes_per_cell,
vertices, generate_grid, getncells, getnodes, getcoordinates

export generate_grid2D, generate_grid3D, refine!, derefine!, Element, AdaptiveGrid,
getncinfo, getedgeindex, getfaceidx, getmuchmorelogicalncinfo

export facedofs, find_local_coordinate

# a few of the utility methods from JuAFEM are ugly [if not wrong]. Also filter instead of findall anod some more advanced collections operations at some places might be useful
include("elements.jl")
include("grid.jl")
include("grid_generators.jl")
include("refinements.jl")
include("utils.jl")

using ColorSchemes # do we need that?
import Makie, AbstractPlotting # do we need that?

include("visualization.jl")

# the current implementation of JuAFEM iterators does not allow for grids with variable eltype
import Base.eltype
function Base.eltype(v::Vector{Element})
    if length(v) == 0 return Element end
    return typeof(v[1])
end
export eltype

end


text:
# Basic Model Interface (BMI) implementation based on
# https://github.com/Deltares/BasicModelInterface.jl

# BMI grid type based on grid identifier
const gridtype = Dict{Int,String}(
    0 => "unstructured",
    1 => "unstructured",
    2 => "unstructured",
    3 => "unstructured",
    4 => "unstructured",
)

# Mapping of grid identifier to a key, to get the active indices of the model domain.
# See also function active_indices(network, key::Tuple).
const grids =
    Dict{Int,String}(0 => "reservoir", 1 => "lake", 2 => "river", 3 => "drain", 4 => "land")

"""
    BMI.initialize(::Type{<:Wflow.Model}, config_file)

Initialize the model. Reads the input settings and data as defined in the Config object
generated from the configuration file `config_file`. Will return a Model that is ready to
run.
"""
function BMI.initialize(::Type{<:Model}, config_file)
    config = Config(config_file)
    modeltype = config.model.type
    model = if modeltype == "sbm"
        initialize_sbm_model(config)
    elseif modeltype == "sbm_gwf"
        initialize_sbm_gwf_model(config)
    elseif modeltype == "hbv"
        initialize_hbv_model(config)
    elseif modeltype == "sediment"
        initialize_sediment_model(config)
    else
        error("unknown model type")
    end
    load_fixed_forcing(model)
    return model
end

"""
    BMI.update(model::Model; run = nothing)

Update the model for a single timestep.
# Arguments
- `run = nothing`: to update a model partially.
"""
function BMI.update(model::Model; run = nothing)
    @unpack network, config = model
    if isnothing(run)
        update_func = update
    elseif run == "sbm_until_recharge"
        update_func = update_until_recharge
    elseif run == "sbm_after_subsurfaceflow"
        update_func = update_after_subsurfaceflow
    end
    load_dynamic_input!(model)
    return update_func(model)
end

function BMI.update_until(model::Model, time::Float64)
    @unpack network, config = model
    curtime = BMI.get_current_time(model)
    n_iter = Int(max(0, (time - curtime) / model.clock.Δt.value))
    end_time = curtime + n_iter * config.timestepsecs
    @info "Updating model until $end_time."
    for _ = 1:n_iter
        load_dynamic_input!(model)
        update(model)
    end
    return model
end

"Write state output to netCDF and close files."
function BMI.finalize(model::Model)
    @unpack config, writer, clock = model
    rewind!(clock)
    write_netcdf_timestep(model, writer.state_dataset, writer.state_parameters)
    reset_clock!(model.clock, config)
    close_files(model, delete_output = false)
end

function BMI.get_component_name(model::Model)
    @unpack config = model
    return config.model.type
end

function BMI.get_input_item_count(model::Model)
    length(BMI.get_input_var_names(model))
end

function BMI.get_output_item_count(model::Model)
    length(BMI.get_output_var_names(model))
end

"""
    BMI.get_input_var_names(model::Model)

Returns model input variables, based on the `API` section in the model configuration file.
This `API` sections contains a list of `Model` components for which variables can be
exchanged.
"""
function BMI.get_input_var_names(model::Model)
    @unpack config = model
    if haskey(config, "API")
        var_names = Vector{String}()
        for c in config.API.components
            append!(
                var_names,
                collect(string.(c, ".", fieldnames(typeof(param(model, c))))),
            )
        end
        return var_names
    else
        @warn("TOML file does not contain section [API] to extract model var names")
        return []
    end
end

"Returns input variables from `BMI.get_input_var_names(model::Model)`, there is no
distinction between input - and output variables."
function BMI.get_output_var_names(model::Model)
    BMI.get_input_var_names(model)
end

function BMI.get_var_grid(model::Model, name::String)
    if name in BMI.get_input_var_names(model)
        if occursin("reservoir", name)
            0
        elseif occursin("lake", name)
            1
        elseif occursin("river", name)
            2
        elseif occursin("drain", name)
            3
        else
            4
        end
    else
        error("Model variable $name not listed as input of output variable")
    end
end

function BMI.get_var_type(model::Model, name::String)
    string(typeof(param(model, name)))
end

function BMI.get_var_units(model::Model, name::String)
    s = split(name, ".")
    get_units(param(model, join(s[1:end-1], ".")), Symbol(s[end]))
end

function BMI.get_var_itemsize(model::Model, name::String)
    sizeof(param(model, name)[1])
end

function BMI.get_var_nbytes(model::Model, name::String)
    sizeof(param(model, name))
end

function BMI.get_var_location(model::Model, name::String)
    if name in BMI.get_input_var_names(model)
        return "node"
    else
        error("$name not in get_input_var_names")
    end
end

function BMI.get_current_time(model::Model)
    datetime2unix(reinterpret(DateTime, model.clock.time))
end

function BMI.get_start_time(model::Model)
    datetime2unix(DateTime(model.config.starttime))
end

function BMI.get_end_time(model::Model)
    datetime2unix(DateTime(model.config.endtime))
end

function BMI.get_time_units(::Type{<:Model})
    string("seconds since ", unix2datetime(0))
end

function BMI.get_time_step(model::Model)
    Float64(model.config.timestepsecs)
end

function BMI.get_value(model::Model, name::String)
    copy(BMI.get_value_ptr(model, name))
end

function BMI.get_value_ptr(model::Model, name::String)
    if name in BMI.get_input_var_names(model)
        return param(model, name)
    else
        error("$name not defined as an output BMI variable")
    end
end

"""
    BMI.get_value_at_indices(model::Model, name::String, inds::Vector{Int})

Returns values of a model variable `name` at indices `inds`.
"""
function BMI.get_value_at_indices(model::Model, name::String, inds::Vector{Int})
    BMI.get_value_ptr(model, name)[inds]
end

"""
    BMI.set_value(model::Model, name::String, src::Vector{T}) where T<:AbstractFloat

Set a model variable `name` to the values in vector `src`, overwriting the current contents.
The type and size of `src` must match the model’s internal array.
"""
function BMI.set_value(model::Model, name::String, src::Vector{T}) where {T<:AbstractFloat}
    BMI.get_value_ptr(model, name) .= src
end

"""
    BMI.set_value_at_indices(model::Model, name::String, inds::Vector{Int}, src::Vector{T})
    where T<:AbstractFloat

    Set a model variable `name` to the values in vector `src`, at indices `inds`.
"""
function BMI.set_value_at_indices(
    model::Model,
    name::String,
    inds::Vector{Int},
    src::Vector{T},
) where {T<:AbstractFloat}
    BMI.get_value_ptr(model, name)[inds] .= src
end

function BMI.get_grid_type(model::Model, grid::Int)
    gridtype[grid]
end

function BMI.get_grid_rank(model::Model, grid::Int)
    if grid in keys(gridtype)
        return 2
    else
        error("unknown grid type $grid")
    end
end

function BMI.get_grid_shape(model::Model, grid::Int)
    size(active_indices(model.network, symbols(grids[grid])), 1)
end

function BMI.get_grid_size(model::Model, grid::Int)
    BMI.get_grid_shape(model, grid)
end

function BMI.get_grid_x(model::Model, grid::Int)
    @unpack reader, config = model
    @unpack dataset = reader
    sel = active_indices(model.network, symbols(grids[grid]))
    inds = [sel[i][1] for i in eachindex(sel)]
    x_nc = read_x_axis(dataset)
    return x_nc[inds]
end

function BMI.get_grid_y(model::Model, grid::Int)
    @unpack reader, config = model
    @unpack dataset = reader
    sel = active_indices(model.network, symbols(grids[grid]))
    inds = [sel[i][2] for i in eachindex(sel)]
    y_nc = read_y_axis(dataset)
    return y_nc[inds]
end

function BMI.get_grid_node_count(model::Model, grid::Int)
    BMI.get_grid_size(model, grid)
end


text:
using CMake, Test

@test isfile(cmake)
@test occursin("cmake version", read(`$cmake --version`, String))


text:
module MIPVerify

using Base.Cartesian
using JuMP
using MathOptInterface
using Memento
using DocStringExtensions
using ProgressMeter

# TODO: more reliable way to determine location for dependencies
const dependencies_path = joinpath(@__DIR__, "..", "deps")

export find_adversarial_example, frac_correct, interval_arithmetic, lp, mip

@enum TighteningAlgorithm interval_arithmetic = 1 lp = 2 mip = 3
@enum AdversarialExampleObjective closest = 1 worst = 2
const DEFAULT_TIGHTENING_ALGORITHM = mip

# importing vendor/ConditionalJuMP.jl first as the remaining files use functions
# defined in it. we're unsure if this is necessary.
include("vendor/ConditionalJuMP.jl")
include("net_components.jl")
include("models.jl")
include("utils.jl")
include("logging.jl")

function get_max_index(x::Array{<:Real,1})::Integer
    return findmax(x)[2]
end

function get_default_tightening_options(optimizer)::Dict
    optimizer_type_name = string(typeof(optimizer()))
    if optimizer_type_name == "Gurobi.Optimizer"
        return Dict("OutputFlag" => 0, "TimeLimit" => 20)
    elseif optimizer_type_name == "Cbc.Optimizer"
        return Dict("logLevel" => 0, "seconds" => 20)
    else
        return Dict()
    end
end

"""
$(SIGNATURES)

Finds the perturbed image closest to `input` such that the network described by `nn`
classifies the perturbed image in one of the categories identified by the
indexes in `target_selection`.

`optimizer` specifies the optimizer used to solve the MIP problem once it has been built.

The output dictionary has keys `:Model, :PerturbationFamily, :TargetIndexes, :SolveStatus,
:Perturbation, :PerturbedInput, :Output`.
See the [tutorial](https://nbviewer.jupyter.org/github/vtjeng/MIPVerify.jl/blob/master/examples/03_interpreting_the_output_of_find_adversarial_example.ipynb)
on what individual dictionary entries correspond to.

*Formal Definition*: If there are a total of `n` categories, the (perturbed) output vector
`y=d[:Output]=d[:PerturbedInput] |> nn` has length `n`.
We guarantee that `y[j] - y[i] ≥ 0` for some `j ∈ target_selection` and for all
`i ∉ target_selection`.

# Named Arguments:
+ `invert_target_selection::Bool`: Defaults to `false`. If `true`, sets `target_selection` to
    be its complement.
+ `pp::PerturbationFamily`: Defaults to `UnrestrictedPerturbationFamily()`. Determines
    the family of perturbations over which we are searching for adversarial examples.
+ `norm_order::Real`: Defaults to `1`. Determines the distance norm used to determine the
    distance from the perturbed image to the original. Supported options are `1`, `Inf`
    and `2` (if the `optimizer` used can solve MIQPs.)
+ `tightening_algorithm::MIPVerify.TighteningAlgorithm`: Defaults to `mip`. Determines how we
    determine the upper and lower bounds on input to each nonlinear unit.
    Allowed options are `interval_arithmetic`, `lp`, `mip`.
    (1) `interval_arithmetic` looks at the bounds on the output to the previous layer.
    (2) `lp` solves an `lp` corresponding to the `mip` formulation, but with any integer constraints
         relaxed.
    (3) `mip` solves the full `mip` formulation.
+ `tightening_options`: Solver-specific options passed to optimizer when used to determine upper and
    lower bounds for input to nonlinear units. Note that these are only used if the
    `tightening_algorithm` is `lp` or `mip` (no solver is used when `interval_arithmetic` is used
    to compute the bounds). Defaults for Gurobi and Cbc to a time limit of 20s per solve,
    with output suppressed.
+ `solve_if_predicted_in_targeted`: Defaults to `true`. The prediction that `nn` makes for the
    unperturbed `input` can be determined efficiently. If the predicted index is one of the indexes
    in `target_selection`, we can skip the relatively costly process of building the model for the
    MIP problem since we already have an "adversarial example" --- namely, the input itself. We
    continue build the model and solve the (trivial) MIP problem if and only if
    `solve_if_predicted_in_targeted` is `true`.
"""
function find_adversarial_example(
    nn::NeuralNet,
    input::Array{<:Real},
    target_selection::Union{Integer,Array{<:Integer,1}},
    optimizer,
    main_solve_options::Dict;
    invert_target_selection::Bool = false,
    pp::PerturbationFamily = UnrestrictedPerturbationFamily(),
    norm_order::Real = 1,
    adversarial_example_objective::AdversarialExampleObjective = closest,
    tightening_algorithm::TighteningAlgorithm = DEFAULT_TIGHTENING_ALGORITHM,
    tightening_options::Dict = get_default_tightening_options(optimizer),
    solve_if_predicted_in_targeted = true,
)::Dict

    total_time = @elapsed begin
        d = Dict()

        # Calculate predicted index
        predicted_output = input |> nn
        num_possible_indexes = length(predicted_output)
        predicted_index = predicted_output |> get_max_index

        d[:PredictedIndex] = predicted_index

        # Set target indexes
        d[:TargetIndexes] = get_target_indexes(
            target_selection,
            num_possible_indexes,
            invert_target_selection = invert_target_selection,
        )
        notice(
            MIPVerify.LOGGER,
            "Attempting to find adversarial example. Neural net predicted label is $(predicted_index), target labels are $(d[:TargetIndexes])",
        )

        # Only call optimizer if predicted index is not found among target indexes.
        if !(d[:PredictedIndex] in d[:TargetIndexes]) || solve_if_predicted_in_targeted
            merge!(d, get_model(nn, input, pp, optimizer, tightening_options, tightening_algorithm))
            m = d[:Model]

            if adversarial_example_objective == closest
                set_max_indexes(m, d[:Output], d[:TargetIndexes])

                # Set perturbation objective
                # NOTE (vtjeng): It is important to set the objective immediately before we carry
                # out the solve. Functions like `set_max_indexes` can modify the objective.
                @objective(m, Min, get_norm(norm_order, d[:Perturbation]))
            elseif adversarial_example_objective == worst
                (maximum_target_var, nontarget_vars) =
                    get_vars_for_max_index(d[:Output], d[:TargetIndexes])
                maximum_nontarget_var = maximum_ge(nontarget_vars)
                # Introduce an additional variable since Gurobi ignores constant terms in objective,
                # but we explicitly need these if we want to stop early based on the value of the
                # objective (not simply whether or not it is maximized).
                # See discussion in https://github.com/jump-dev/Gurobi.jl/issues/111 for more
                # details.
                v_obj = @variable(m)
                @constraint(m, v_obj == maximum_target_var - maximum_nontarget_var)
                @objective(m, Max, v_obj)
            else
                error("Unknown adversarial_example_objective $adversarial_example_objective")
            end
            set_optimizer(m, optimizer)
            set_optimizer_attributes(m, main_solve_options...)
            optimize!(m)
            d[:SolveStatus] = JuMP.termination_status(m)
            d[:SolveTime] = JuMP.solve_time(m)
        end
    end

    d[:TotalTime] = total_time
    return d
end

function get_label(y::Array{<:Real,1}, test_index::Integer)::Int
    return y[test_index]
end

function get_image(x::Array{T,4}, test_index::Integer)::Array{T,4} where {T<:Real}
    return x[test_index:test_index, :, :, :]
end

"""
$(SIGNATURES)

Returns the fraction of items the neural network correctly classifies of the first
`num_samples` of the provided `dataset`. If there are fewer than
`num_samples` items, we use all of the available samples.

# Named Arguments:
+ `nn::NeuralNet`: The parameters of the neural network.
+ `dataset::LabelledDataset`:
+ `num_samples::Integer`: Number of samples to use.
"""
function frac_correct(nn::NeuralNet, dataset::LabelledDataset, num_samples::Integer)::Real

    num_correct = 0.0
    num_samples = min(num_samples, MIPVerify.num_samples(dataset))
    @showprogress 1 "Computing fraction correct..." for sample_index in 1:num_samples
        input = get_image(dataset.images, sample_index)
        actual_label = get_label(dataset.labels, sample_index)
        predicted_label = (input |> nn |> get_max_index) - 1
        if actual_label == predicted_label
            num_correct += 1
        end
    end
    return num_correct / num_samples
end

function get_norm(norm_order::Real, v::Array{<:Real})
    if norm_order == 1
        return sum(abs.(v))
    elseif norm_order == 2
        return sqrt(sum(v .* v))
    elseif norm_order == Inf
        return Base.maximum(Iterators.flatten(abs.(v)))
    else
        throw(DomainError("Only l1, l2 and l∞ norms supported."))
    end
end

function get_norm(norm_order::Real, v::Array{<:JuMPLinearType,N}) where {N}
    if norm_order == 1
        abs_v = abs_ge.(v)
        return sum(abs_v)
    elseif norm_order == 2
        return sum(v .* v)
    elseif norm_order == Inf
        return MIPVerify.maximum_ge(permute_and_flatten(abs_ge.(v), N:-1:1))
    else
        throw(DomainError("Only l1, l2 and l∞ norms supported."))
    end
end

include("batch_processing_helpers.jl")

end

max_stars_repo_name: vtjeng/MIPVerify

max_issues_repo_name: vtjeng/MIPVerify

max_forks_repo_name: vtjeng/MIPVerify