# integration with LinearAlgebra.jl

CuMatOrAdj{T} = Union{CuMatrix, LinearAlgebra.Adjoint{T, <:CuMatrix{T}}, LinearAlgebra.Transpose{T, <:CuMatrix{T}}}
CuOrAdj{T} = Union{CuVecOrMat, LinearAlgebra.Adjoint{T, <:CuVecOrMat{T}}, LinearAlgebra.Transpose{T, <:CuVecOrMat{T}}}

# matrix division

function Base.:\(_A::CuMatOrAdj, _B::CuOrAdj)
    A, B = copy(_A), copy(_B)
    A, ipiv = CUSOLVER.getrf!(A)
    return CUSOLVER.getrs!('N', A, ipiv, B)
end

# patch JuliaLang/julia#40899 to create a CuArray
# (see https://github.com/JuliaLang/julia/pull/41331#issuecomment-868374522)
if VERSION >= v"1.7-"
_zeros(::Type{T}, b::AbstractVector, n::Integer) where {T} = CUDA.zeros(T, max(length(b), n))
_zeros(::Type{T}, B::AbstractMatrix, n::Integer) where {T} = CUDA.zeros(T, max(size(B, 1), n), size(B, 2))
function Base.:\(F::Union{LinearAlgebra.LAPACKFactorizations{<:Any,<:CuArray},
                          Adjoint{<:Any,<:LinearAlgebra.LAPACKFactorizations{<:Any,<:CuArray}}},
                 B::AbstractVecOrMat)
    m, n = size(F)
    if m != size(B, 1)
        throw(DimensionMismatch("arguments must have the same number of rows"))
    end

    TFB = typeof(oneunit(eltype(B)) / oneunit(eltype(F)))
    FF = Factorization{TFB}(F)

    # For wide problem we (often) compute a minimum norm solution. The solution
    # is larger than the right hand side so we use size(F, 2).
    BB = _zeros(TFB, B, n)

    if n > size(B, 1)
        # Underdetermined
        copyto!(view(BB, 1:m, :), B)
    else
        copyto!(BB, B)
    end

    ldiv!(FF, BB)

    # For tall problems, we compute a least squares solution so only part
    # of the rhs should be returned from \ while ldiv! uses (and returns)
    # the complete rhs
    return LinearAlgebra._cut_B(BB, 1:n)
end
end

# qr

using LinearAlgebra: AbstractQ

# AbstractQ's `size` is the size of the full matrix,
# while `Matrix(Q)` only gives the compact Q.
# See JuliaLang/julia#26591 and JuliaGPU/CUDA.jl#969.
CuMatrix{T}(Q::AbstractQ{S}) where {T,S} = convert(CuArray, Matrix{T}(Q))
CuMatrix(Q::AbstractQ{T}) where {T} = CuMatrix{T}(Q)
CuArray{T}(Q::AbstractQ) where {T} = CuMatrix{T}(Q)
CuArray(Q::AbstractQ) = CuMatrix(Q)

# dot

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

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

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

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

        return
    end

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

        res = zeros(T, 1)

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

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

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

        @allowscalar res[]
    end
end


using StochasticDiffEq, Test, Random
using DiffEqProblemLibrary.SDEProblemLibrary: importsdeproblems; importsdeproblems()
using DiffEqProblemLibrary.SDEProblemLibrary: prob_sde_stiffquadito

Random.seed!(100)
prob = prob_sde_stiffquadito
prob = remake(prob;p=(1e5,2.))
alg = AutoSOSRA2(SKenCarp(), maxstiffstep=5, maxnonstiffstep=2, stiffalgfirst=false)
@test StochasticDiffEq.isadaptive(prob,alg)
@test StochasticDiffEq.isadaptive(alg)
@time sol = solve(prob, alg)
@test typeof(alg.algs[sol.alg_choice[end]]) <: SKenCarp
@test length(unique(sol.alg_choice)) == 2

Random.seed!(100)
prob = prob_sde_stiffquadito
prob = remake(prob;p=(1e5,2.))
@time sol = solve(prob, AutoSOSRI2(ImplicitRKMil(), maxstiffstep=1, maxnonstiffstep=10, stiffalgfirst=false))
@test length(unique(sol.alg_choice)) == 2


# take unaligned seed -> make a msa

function build_model(fileseed::String, filefull::String, ctype::Symbol, L::Int64;
		    filename_ins::String="LambdaOpen_LambdaExt.dat",
		    filename_par::String="Parameters_PlmDCA.dat",
		    filename_gap::String="Gap_Ext_Int.dat",
		    Mtest::Int64=0,
		    verbose::Bool=true)

	if ctype != :amino && ctype != :nbase
		error("Wrong second argument: choose between :amino and :nbase")
		return
	end
	println("### Reading seed ###")
	if ctype == :nbase
		print_pos = true
	else
		print_pos = false
	end
	seed = DCAlign.readfull(fileseed, ctype=ctype, pos = true)
	println("### Inferring insertions penalties ###")
	l_o, l_e = infer_ins_pen(seed, L)

	println("### Inferring a Potts model using PlmDCA ###")
	aligntmp = filter_insertions(seed)
	println("Temporary FASTA alignment in ", aligntmp)
	PlmData = PlmDCA.plmdca(aligntmp, theta=0.20)

	println("### Finding gap penalties ###")
	println("WARNING: Reasonable values are obtained when using many (> 500) sequences")

	full = DCAlign.readfull(filefull, ctype=ctype, pos = print_pos)
	if Mtest == 0
		Mtest = length(seed)
		println("Using all seed sequences to get the gap penalties")
	else
		println("Using ", Mtest, " out of ", length(seed), " to get the gap penalties")
	end
	fulltmp = extract_full_seq(full, seed, Mtest, ctype=ctype)
	println("Temporary full length sequences in ", fulltmp)

	if ctype == :amino
		q = 21
	elseif ctype == :nbase
		q = 5
	end

	print_results(filename_ins, l_o, l_e, filename_par, PlmData, ctype, L, filename_gap, 0.0, 0.0) # tmp
	mu = 0.00:0.50:4.00
	muint = 0.00:0.50:4.00
	d = zeros(length(mu),length(muint))
	aseed = DCAlign.readfull(aligntmp, ctype=ctype, pos = true)
	for a in 1:length(mu)
		for b in 1:length(muint)
			println("#### Aligning ", Mtest, " sequences using (μext, μint) = (", mu[a], ", ", muint[b],")")
			filename_out = tempname()
			filename_flag = tempname()
			DCAlign.align_all(q, L, filename_par, fulltmp, filename_ins, mu[a], muint[b]; typel=:plm, filename_flag=filename_flag, filename_align=aligntmp, filename_ins=fileseed, filename_out=filename_out, verbose = verbose, maxiter = 300)
			tmpseed = DCAlign.readfull(filename_out, ctype=ctype, pos = true)
			d[a,b] = compute_average_dist(aseed, tmpseed)
			rm(filename_out)
			rm(filename_flag)
		end
	end
	#println(d)
	aux = argmin(d)
	muext_best = mu[aux[1]]
	muint_best = muint[aux[2]]
	print_results(filename_ins, l_o, l_e, filename_par, PlmData, ctype, L, filename_gap, muext_best, muint_best)
	rm(fulltmp)
	rm(aligntmp)
	println("Done!")

end


"""
    MixedModelPermutation{T<:AbstractFloat}

Object returned by the various permutation methods with fields
- `fits`: the parameter estimates from the permutation replicates as a vector of named tuples.
- `λ`: `Vector{LowerTriangular{T,Matrix{T}}}` containing copies of the λ field from `ReMat` model terms
- `inds`: `Vector{Vector{Int}}` containing copies of the `inds` field from `ReMat` model terms
- `lowerbd`: `Vector{T}` containing the vector of lower bounds (corresponds to the identically named field of [`OptSummary`](@ref))
- `fcnames`: NamedTuple whose keys are the grouping factor names and whose values are the column names

The schema of `fits` is, by default,
```
Tables.Schema:
 :objective  T
 :σ          T
 :β          NamedTuple{β_names}{NTuple{p,T}}
 :se         StaticArrays.SArray{Tuple{p},T,1,p}
 :θ          StaticArrays.SArray{Tuple{k},T,1,k}
```
where the sizes, `p` and `k`, of the `β` and `θ` elements are determined by the model.

Characteristics of the permutation replicates can be extracted as properties.  The `σs` and
`σρs` properties unravel the `σ` and `θ` estimates into estimates of the standard deviations
and correlations of the random-effects terms.
"""
struct MixedModelPermutation{T<:AbstractFloat}
    fits::Vector
    λ::Vector{LowerTriangular{T,Matrix{T}}}
    inds::Vector{Vector{Int}}
    lowerbd::Vector{T}
    fcnames::NamedTuple
end

#####
##### Convert and Delegate Methods to MixedModelBootstrap
#####

# XXX This can be removed when MixedModels.jl introduces the abstract fit
#     collection type and broadens their method signatures

# non copying converstion to a bootstrap object.
_perm2boot(p::MixedModelPermutation) = MixedModelBootstrap(p.fits, p.λ, p.inds, p.lowerbd, p.fcnames)

# delegate methods
for f in (:allpars, :coefpvalues, :issingular, :setθ!, :tidyβ, :tidyσs)

    @eval begin
        MixedModels.$f(p::MixedModelPermutation) = MixedModels.$f(_perm2boot(p))
    end
end

Base.propertynames(p::MixedModelPermutation) = propertynames(_perm2boot(p))

function Base.getproperty(p::MixedModelPermutation, s::Symbol)
    if s ∈ [:objective, :σ, :θ, :se]
        getproperty.(getfield(p, :fits), s)
    elseif s == :β
        tidyβ(p)
    elseif s == :coefpvalues
        coefpvalues(p)
    elseif s == :σs
        tidyσs(p)
    elseif s == :allpars
        allpars(p)
    else
        getfield(p, s)
    end
end


module ProtStructs

using StaticArrays
import LinearAlgebra: norm, ⋅, ×
import Chemfiles    ## don't bring out names for now

# Write your package code here.
include("constants.jl")

include("Model.jl")
include("Parser.jl")
include("Assignments.jl")

include("GeomUtils.jl")
include("FrameGeom.jl")
# something on contact maps/solvent accessibility

end


"""
Please add the latest Kinetic.jl package as it fixes the issue with the regime classifier.
Backend requirement:
- KitBase: v0.7.6
- KitML: v0.4.4
"""

using Kinetic, NPZ, Plots
using KitBase.ProgressMeter: @showprogress
cd(@__DIR__)

set = Setup(case = "sod", space = "1d1f1v", maxTime = 0.15)
ps = PSpace1D(0.0, 1.0, 100, 1)
vs = VSpace1D(-5.0, 5.0, 64)
gas = Gas(Kn = 1e-3, K = 0, γ = 3)
ib = IB1F(ib_sod(set, ps, vs, gas)...)
ks = SolverSet(set, ps, vs, gas, ib)
ctr, face = init_fvm(ks)

t = 0.0
dt = timestep(ks, ctr, t)
nt = Int(ks.set.maxTime ÷ dt) + 1
res = zero(ctr[1].w)

X, Y, Z, I = begin
    prim = [1.0, 0.0, 1.0]
    w = prim_conserve(prim, ks.gas.γ)
    sw = zeros(3)
    M = maxwellian(ks.vs.u, prim)
    [[1e-4, 0.0, 1e-4]; zeros(3); 1.0], 0.0, M, 0
end

@showprogress for iter = 1:nt
    evolve!(ks, ctr, face, dt)
    update!(ks, ctr, face, dt, res)

    if iter % 3 == 0
        for i = 1:ks.ps.nx
            sw = (ctr[i+1].w .- ctr[i-1].w) / ks.ps.dx[i] / 2.0
            x, y = regime_data(ks, ctr[i].w, sw, ctr[i].f)
            global X = hcat(X, x)
            global Y = hcat(Y, y)
            global Z = hcat(Z, ctr[i].f)
            global I = hcat(I, iter)
        end
    end
end

#plot(ks, ctr)

npzwrite("data.npz", Dict("I" => I, "X" => X, "Y" => Y, "Z" => Z))


abstract type AbstractLoadFormulation <: AbstractDeviceFormulation end

abstract type AbstractControllablePowerLoadFormulation <: AbstractLoadFormulation end

struct StaticPowerLoad <: AbstractLoadFormulation end

struct InterruptiblePowerLoad <: AbstractControllablePowerLoadFormulation end

struct DispatchablePowerLoad <: AbstractControllablePowerLoadFormulation end

########################### dispatchable load variables ####################################
function activepower_variables!(canonical::Canonical,
                               devices::IS.FlattenIteratorWrapper{L}) where L<:PSY.ElectricLoad
    add_variable(canonical,
                 devices,
                 Symbol("P_$(L)"),
                 false,
                 :nodal_balance_active, -1.0;
                 ub_value = x -> PSY.get_maxactivepower(x),
                 lb_value = x -> 0.0)

    return

end

function reactivepower_variables!(canonical::Canonical,
                                 devices::IS.FlattenIteratorWrapper{L}) where L<:PSY.ElectricLoad
    add_variable(canonical,
                 devices,
                 Symbol("Q_$(L)"),
                 false,
                 :nodal_balance_reactive, -1.0;
                 ub_value = x -> PSY.get_maxreactivepower(x),
                 lb_value = x -> 0.0)

    return

end

function commitment_variables!(canonical::Canonical,
                              devices::IS.FlattenIteratorWrapper{L}) where L<:PSY.ElectricLoad

    add_variable(canonical,
                 devices,
                 Symbol("ON_$(L)"),
                 true)

    return

end

####################################### Reactive Power Constraints #########################
"""
Reactive Power Constraints on Loads Assume Constant PowerFactor
"""
function reactivepower_constraints!(canonical::Canonical,
                                   devices::IS.FlattenIteratorWrapper{L},
                                   device_formulation::Type{<:AbstractControllablePowerLoadFormulation},
                                   system_formulation::Type{<:PM.AbstractPowerModel}) where L<:PSY.ElectricLoad

    time_steps = model_time_steps(canonical)
    key = Symbol("reactive_$(L)")
    canonical.constraints[key] = JuMPConstraintArray(undef, (PSY.get_name(d) for d in devices), time_steps)

    for t in time_steps, d in devices
        name = PSY.get_name(d)
        pf = sin(atan((PSY.get_maxreactivepower(d)/PSY.get_maxactivepower(d))))
        canonical.constraints[key][PSY.get_name(d), t] = JuMP.@constraint(canonical.JuMPmodel,
                        canonical.variables[Symbol("Q_$(L)")][name, t] == canonical.variables[Symbol("P_$(L)")][name, t]*pf)
    end

    return

end

######################## output constraints without Time Series ############################
function _get_time_series(canonical::Canonical,
                          devices::IS.FlattenIteratorWrapper{<:PSY.ElectricLoad})

    initial_time = model_initial_time(canonical)
    use_forecast_data = model_uses_forecasts(canonical)
    time_steps = model_time_steps(canonical)
    device_total = length(devices)
    ts_data_active = Vector{Tuple{String, Int64, Float64, Vector{Float64}}}(undef, device_total)
    ts_data_reactive = Vector{Tuple{String, Int64, Float64, Vector{Float64}}}(undef, device_total)

    for (ix, device) in enumerate(devices)
        bus_number = PSY.get_number(PSY.get_bus(device))
        name = PSY.get_name(device)
        active_power = use_forecast_data ? PSY.get_maxactivepower(device) : PSY.get_activepower(device)
        reactive_power = use_forecast_data ? PSY.get_maxreactivepower(device) : PSY.get_reactivepower(device)
        if use_forecast_data
            forecast = PSY.get_forecast(PSY.Deterministic,
                                        device,
                                        initial_time,
                                        "get_maxactivepower",
                                        length(time_steps))
            ts_vector = TS.values(PSY.get_data(forecast))
        else
            ts_vector = ones(time_steps[end])
        end
        ts_data_active[ix] = (name, bus_number, active_power, ts_vector)
        ts_data_reactive[ix] = (name, bus_number, reactive_power, ts_vector)
    end

    return ts_data_active, ts_data_reactive

end

function activepower_constraints!(canonical::Canonical,
                                 devices::IS.FlattenIteratorWrapper{L},
                                 device_formulation::Type{DispatchablePowerLoad},
                                 system_formulation::Type{<:PM.AbstractPowerModel}) where L<:PSY.ElectricLoad

    parameters = model_has_parameters(canonical)
    use_forecast_data = model_uses_forecasts(canonical)

    if !parameters && !use_forecast_data
        range_data = [(PSY.get_name(d), (min = 0.0, max = PSY.get_activepower(d))) for d in devices]
        device_range(canonical,
                    range_data,
                    Symbol("activerange_$(L)"),
                    Symbol("P_$(L)"))
        return
    end

    ts_data_active, _ = _get_time_series(canonical, devices)
    if parameters
        device_timeseries_param_ub(canonical,
                                   ts_data_active,
                                   Symbol("active_$(L)"),
                                   UpdateRef{L}("get_maxactivepower"),
                                   Symbol("P_$(L)"))
    else
        device_timeseries_ub(canonical,
                            ts_data_active,
                            Symbol("active_$(L)"),
                            Symbol("P_$(L)"))
    end

    return

end

function activepower_constraints!(canonical::Canonical,
                                 devices::IS.FlattenIteratorWrapper{L},
                                 device_formulation::Type{InterruptiblePowerLoad},
                                 system_formulation::Type{<:PM.AbstractPowerModel}) where L<:PSY.ElectricLoad

    parameters = model_has_parameters(canonical)
    use_forecast_data = model_uses_forecasts(canonical)

    if !parameters && !use_forecast_data
        range_data = [(PSY.get_name(d), (min = 0.0, max = PSY.get_activepower(d))) for d in devices]
        device_range(canonical,
                    range_data,
                    Symbol("activerange_$(L)"),
                    Symbol("P_$(L)"))
        return
    end

    ts_data_active, _ = _get_time_series(canonical, devices)
    if parameters
        device_timeseries_ub_bigM(canonical,
                                 ts_data_active,
                                 Symbol("active_$(L)"),
                                 Symbol("P_$(L)"),
                                 UpdateRef{L}("get_maxactivepower"),
                                 Symbol("ON_$(L)"))
    else
        device_timeseries_ub_bin(canonical,
                                ts_data_active,
                                Symbol("active_$(L)"),
                                Symbol("P_$(L)"),
                                Symbol("ON_$(L)"))
    end

    return

end

########################## Addition of to the nodal balances ###############################
function nodal_expression!(canonical::Canonical,
                           devices::IS.FlattenIteratorWrapper{L},
                           system_formulation::Type{<:PM.AbstractPowerModel}) where L<:PSY.ElectricLoad

    ts_data_active, ts_data_reactive = _get_time_series(canonical, devices)
    parameters = model_has_parameters(canonical)

    if parameters
        include_parameters(canonical,
                        ts_data_active,
                        UpdateRef{L}("get_maxactivepower"),
                        :nodal_balance_active,
                        -1.0)
        include_parameters(canonical,
                        ts_data_reactive,
                        UpdateRef{L}("get_maxactivepower"),
                        :nodal_balance_reactive,
                        -1.0)
        return
    end

    for t in model_time_steps(canonical)
        for device_value in ts_data_active
            _add_to_expression!(canonical.expressions[:nodal_balance_active],
                            device_value[2],
                            t,
                            -device_value[3]*device_value[4][t])
        end
        for device_value in ts_data_reactive
            _add_to_expression!(canonical.expressions[:nodal_balance_reactive],
                            device_value[2],
                            t,
                            -device_value[3]*device_value[4][t])
        end
    end

    return

end

function nodal_expression!(canonical::Canonical,
                           devices::IS.FlattenIteratorWrapper{L},
                           system_formulation::Type{<:PM.AbstractActivePowerModel}) where L<:PSY.ElectricLoad

    parameters = model_has_parameters(canonical)
    ts_data_active, _ = _get_time_series(canonical, devices)

    if parameters
        include_parameters(canonical,
                        ts_data_active,
                        UpdateRef{L}("get_maxactivepower"),
                        :nodal_balance_active,
                        -1.0)
        return
    end

    for t in model_time_steps(canonical), device_value in ts_data_active
        _add_to_expression!(canonical.expressions[:nodal_balance_active],
                            device_value[2],
                            t,
                            -device_value[3]*device_value[4][t])
    end

    return
end

############################## FormulationControllable Load Cost ###########################
function cost_function(canonical::Canonical,
                       devices::IS.FlattenIteratorWrapper{L},
                       device_formulation::Type{DispatchablePowerLoad},
                       system_formulation::Type{<:PM.AbstractPowerModel}) where L<:PSY.ControllableLoad

    add_to_cost(canonical,
                devices,
                Symbol("P_$(L)"),
                :variable,
                -1.0)

    return

end

function cost_function(canonical::Canonical,
                       devices::IS.FlattenIteratorWrapper{L},
                       device_formulation::Type{InterruptiblePowerLoad},
                       system_formulation::Type{<:PM.AbstractPowerModel}) where L<:PSY.ControllableLoad

    add_to_cost(canonical,
                devices,
                Symbol("ON_$(L)"),
                :fixed,
                -1.0)

    return

end


B_data = [ -2.5   0.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   3.0        -2.0           1.0          0.0   0.0   0.0   0.0     0.0   0.0    0.0          0.0           0.0            0.0            0.0            0.0           0.0     0.0    0.0     0.0   0.0
1.0  -3.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   3.0        -2.0           1.0          0.0   0.0   0.0   0.0     0.0   0.0    0.0          0.0           0.0            0.0            0.0            0.0           0.0     0.0    0.0     0.0   0.0
0.0   0.0  -2.5   0.0  -0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0         0.0           0.0          0.0   0.0   0.0   0.0     0.0   0.0    0.5          0.0           0.0            0.0            0.0            0.0           0.0     0.0    0.0     0.0   0.0
0.0   0.0  -0.0  -2.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0         2.0           0.0          0.0   0.0   0.0   0.0     0.0   0.0    0.0          0.5           0.0            0.0            0.0            0.0           0.0     0.0    0.0     0.0   0.0
0.0   0.0   0.0  -0.0  -2.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0         0.0           2.0          0.0   0.0   0.0   0.0     0.0   0.0    0.0          0.0           0.5            0.0            0.0            0.0           0.0     0.0    0.0     0.0   0.0
0.0   0.0   0.0   0.0   0.0  -2.5   0.0  -0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0         0.0           0.0          2.0   0.0   0.0   0.0     0.0   0.0    0.0          0.0           0.0            0.5            0.0            0.0           0.0     0.0    0.0     0.0   0.0
0.0   0.0   0.0   0.0   0.0  -0.0  -2.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0         0.0           0.0          0.0   2.0   0.0   0.0     0.0   0.0    0.0          0.0           0.0            0.0            0.5            0.0           0.0     0.0    0.0     0.0   0.0
0.0   0.0   0.0   0.0   0.0   0.0  -0.0  -2.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0         0.0           0.0          0.0   0.0   2.0   0.0     0.0   0.0    0.0          0.0           0.0            0.0            0.0            0.5           0.0     0.0    0.0     0.0   0.0
0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -2.5   0.0  -0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0         0.0           0.0          0.0   0.0   0.0   2.0     0.0   0.0    0.0          0.0           0.0            0.0            0.0            0.0           0.5     0.0    0.0     0.0   0.0
0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -0.0  -2.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0         0.0           0.0          0.0   0.0   0.0   0.0     2.0   0.0    0.0          0.0           0.0            0.0            0.0            0.0           0.0     0.5    0.0     0.0   0.0
0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -0.0  -2.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0         0.0           0.0          0.0   0.0   0.0   0.0     0.0   2.0    0.0          0.0           0.0            0.0            0.0            0.0           0.0     0.0    0.5     0.0   0.0
0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -2.5   0.0  -0.0   0.0   0.0   0.0   0.0   0.0   0.0   2.0         0.0           0.0          0.0   0.0   0.0   0.0     0.0   0.0    0.5          0.0           0.0            0.0            0.0            0.0           0.0     0.0    0.0     0.0   0.0
0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -0.0  -2.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0         2.0           0.0          0.0   0.0   0.0   0.0     0.0   0.0    0.0          0.5           0.0            0.0            0.0            0.0           0.0     0.0    0.0     0.0   0.0
0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -0.0  -2.5   0.0   0.0   0.0   0.0   0.0   0.0   0.0         0.0           2.0          0.0   0.0   0.0   0.0     0.0   0.0    0.0          0.0           0.5            0.0            0.0            0.0           0.0     0.0    0.0     0.0   0.0
0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -2.5   0.0  -0.0   0.0   0.0   0.0   0.0         0.0           0.0          2.0   0.0   0.0   0.0     0.0   0.0    0.0          0.0           0.0            0.5            0.0            0.0           0.0     0.0    0.0     0.0   0.0
0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -0.0  -2.5   0.0   0.0   0.0   0.0   0.0         0.0           0.0          0.0   2.0   0.0   0.0     0.0   0.0    0.0          0.0           0.0            0.0            0.5            0.0           0.0     0.0    0.0     0.0   0.0
0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -0.0  -2.5   0.0   0.0   0.0   0.0         0.0           0.0          0.0   0.0   2.0   0.0     0.0   0.0    0.0          0.0           0.0            0.0            0.0            0.5           0.0     0.0    0.0     0.0   0.0
0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -2.5   0.0  -0.0   0.0         0.0           0.0          0.0   0.0   0.0   2.0     0.0   0.0    0.0          0.0           0.0            0.0            0.0            0.0           0.5     0.0    0.0     0.0   0.0
0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -0.0  -2.5   0.0   0.0         0.0           0.0          0.0   0.0   0.0   0.0     2.0   0.0    0.0          0.0           0.0            0.0            0.0            0.0           0.0     0.5    0.0     0.0   0.0
0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -0.0  -2.5   0.0         0.0           0.0          0.0   0.0   0.0   0.0     0.0   2.0    0.0          0.0           0.0            0.0            0.0            0.0           0.0     0.0    0.5     0.0   0.0
0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -6.0         6.66667      -2.66667      0.0   0.0   0.0   0.0     0.0   0.0    1.0          0.0           0.0            0.0            0.0            0.0           0.0     0.0    0.0     0.0   0.0
0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  -0.666667   -2.0           0.666667     0.0   0.0   0.0   0.0     0.0   0.0    0.0          1.0           0.0            0.0            0.0            0.0           0.0     0.0    0.0     0.0   0.0
0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.666667   -2.66667      -4.0          6.0  -4.0   2.0   0.0     0.0   0.0    0.0          0.0           1.0            0.0            0.0            0.0           0.0     0.0    0.0     0.0   0.0
0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0         0.0           0.0        -14.0  20.0  -8.0   0.0     0.0   0.0    0.0          0.0           0.0            1.0            0.0            0.0           0.0     0.0    0.0     0.0   0.0
0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0         0.0           0.0         -2.0  -2.0   2.0   0.0     0.0   0.0    0.0          0.0           0.0            0.0            1.0            0.0           0.0     0.0    0.0     0.0   0.0
0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0         0.0           0.0          2.0  -8.0  -8.0  18.0   -12.0   6.0    0.0          0.0           0.0            0.0            0.0            1.0           0.0     0.0    0.0     0.0   0.0
0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0         0.0           0.0          0.0   0.0   0.0  -3.5     2.5  -1.0    0.0          0.0           0.0            0.0            0.0            0.0           1.0     0.0    0.0     0.0   0.0
0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0         0.0           0.0          0.0   0.0   0.0  -0.25   -2.0   0.25   0.0          0.0           0.0            0.0            0.0            0.0           0.0     1.0    0.0     0.0   0.0
0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0         0.0           0.0          0.0   0.0   0.0   0.25   -1.0  -2.75   0.0          0.0           0.0            0.0            0.0            0.0           0.75   -1.5    3.25    0.0   0.0
0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0  11.0        -6.0           3.0          0.0   0.0   0.0   0.0     0.0   0.0   -6.0         -4.0           1.0            0.0            0.0            0.0           0.0     0.0    0.0     0.0   0.0
0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   3.747e-16   2.0           1.249e-16    0.0   0.0   0.0   0.0     0.0   0.0    1.0         -3.0          -1.0            0.0            0.0            0.0           0.0     0.0    0.0     0.0   0.0
-0.0  -0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0  -2.498e-15   1.66533e-15   2.0          0.0   0.0   0.0   0.0     0.0   0.0   -4.0         10.0          -9.0            0.0            0.0            0.0           0.0     0.0    0.0     0.0   0.0
0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0         0.0           0.0          2.0   0.0   0.0   0.0     0.0   0.0    9.0        -18.0          27.0          -12.0          -12.0            3.0           0.0     0.0    0.0     0.0   0.0
0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0   0.0         0.0           0.0          0.0   2.0   0.0   0.0     0.0   0.0    3.747e-16   -7.49401e-16   1.1241e-15     3.0           -3.0           -3.0           0.0     0.0    0.0     0.0   0.0
0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0   0.0         0.0           0.0          0.0   0.0   2.0   0.0     0.0   0.0   -2.498e-15    4.996e-15    -7.49401e-15  -12.0           30.0          -21.0           0.0     0.0    0.0     0.0   0.0
0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0   0.0         0.0           0.0          0.0   0.0   0.0   2.0     0.0   0.0    0.0          0.0           0.0            1.125         -2.25           3.375        -4.125  -1.5    0.375   0.0   0.0
0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0   0.0         0.0           0.0          0.0   0.0   0.0   0.0     2.0   0.0    0.0          0.0           0.0            4.68375e-17   -9.36751e-17    1.40513e-16   0.375  -3.0   -0.375   0.0   0.0
0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   1.0   0.0         0.0           0.0          0.0   0.0   0.0   0.0     0.0   2.0    0.0          0.0           0.0           -3.1225e-16     6.245e-16     -9.36751e-16  -1.5     3.75  -5.25    0.0   0.0
0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0         0.0           0.0          0.0   0.0   0.0   0.0     0.0   0.0    0.0          0.0           0.0            0.0            0.0            0.0           0.25   -0.5    0.75   -2.5   2.0
0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0         0.0           0.0          0.0   0.0   0.0   0.0     0.0   0.0    0.0          0.0           0.0            0.0            0.0            0.0           0.5    -1.0    1.5     1.0  -2.0]


test_bike_week_1_env = {
    "shift_time_min" => 3,
    "shift_time_max" => 8,
    "coverage"  => Array[
        # Monday
        [3, 4, 5 * ones(Int, 6), 4 * ones(Int, 2), 2 * ones(Int, 2)],
        # Tuesday
        [3, 4, 5 * ones(Int, 6), 4 * ones(Int, 2), 2 * ones(Int, 2)],
        # Wednesday
        [3, 4, 5 * ones(Int, 6), 4 * ones(Int, 2), 2 * ones(Int, 2)],
        # Thursday
        [3, 4, 5 * ones(Int, 6), 4 * ones(Int, 2), 2 * ones(Int, 2)],
        # Friday
        [3, 4, 5 * ones(Int, 6), 4 * ones(Int, 2), 2 * ones(Int, 2)],
        # Saturday
        [1, 2 * ones(Int, 11)],
        # Sunday
        [1, 2 * ones(Int, 11)],
    ],
    "time_between_coverage" => 12,
    "intershift" => 13, # Time periods between shifts   # only serve 1 shift/day
}

# Coverage requires 291 man hours

test_bike_week_1_employees = {
    "Cody" => {
        "hours_min" => 18, # hours worked over simulation
        "hours_max" => 32, # hours worked over simulation
        "shift_time_min" => 4,
        "shift_count_min" => 3,
        "shift_count_max" => 5,
        "availability" => Array[
            [ones(Int, 12)],
            [ones(Int, 12)],
            [ones(Int, 12)],
            [ones(Int, 12)],
            [ones(Int, 12)],
            [ones(Int, 12)],
            [ones(Int, 12)],
        ],
    },
    "Damian" => {
        "hours_min" => 12, # hours worked over simulation
        "hours_max" => 28, # hours worked over simulation
        "intershift" => 4,
        "shift_count_min" => 2,
        "shift_count_max" => 4,
        "availability" => Array[
            [ones(Int, 5), zeros(Int, 7)],
            [ones(Int, 5), zeros(Int, 7)],
            [ones(Int, 5), zeros(Int, 7)],
            [ones(Int, 5), zeros(Int, 7)],
            [ones(Int, 5), zeros(Int, 7)],
            [zeros(Int, 12)],
            [zeros(Int, 12)],
        ],
    },
    "David" => {
        "hours_min" => 12, # hours worked over simulation
        "hours_max" => 18, # hours worked over simulation
        "shift_count_min" => 2,
        "shift_count_max" => 3,
        "availability" => Array[
            [zeros(Int, 12)],
            [zeros(Int, 2), ones(Int, 10)],
            [zeros(Int, 2), ones(Int, 10)],
            [zeros(Int, 2), ones(Int, 10)],
            [zeros(Int, 12)],
            [zeros(Int, 12)],
            [zeros(Int, 12)],
        ],
    },
    "Douglas" => {
        "hours_min" => 12, # hours worked over simulation
        "hours_max" => 28, # hours worked over simulation
        "shift_count_min" => 2,
        "shift_count_max" => 4,
        "availability" => Array[
            [ones(Int, 12)],
            [ones(Int, 12)],
            [ones(Int, 12)],
            [ones(Int, 12)],
            [ones(Int, 12)],
            [ones(Int, 12)],
            [ones(Int, 12)],
        ],
    },
    "Dylan" => {
        # deliberately no shift count
        "hours_min" => 28, # hours worked over simulation
        "hours_max" => 32, # hours worked over simulation
        "shift_time_min" => 6,
        "shift_count_min" => 4,
        "shift_count_max" => 5,
        "availability" => Array[
            [ones(Int, 7), zeros(Int, 5)],
            [zeros(Int, 12)],
            [zeros(Int, 12)],
            [ones(Int, 7), zeros(Int, 5)],
            [ones(Int, 7), zeros(Int, 5)],
            [ones(Int, 7), zeros(Int, 5)],
            [ones(Int, 7), zeros(Int, 5)],
        ],
    },
    "Jaeger" => {
        "hours_min" => 18, # hours worked over simulation
        "hours_max" => 32, # hours worked over simulation
        "shift_count_min" => 3,
        "shift_count_max" => 5,
        "availability" => Array[
            [ones(Int, 12)],
            [ones(Int, 12)],
            [ones(Int, 12)],
            [ones(Int, 12)],
            [ones(Int, 12)],
            [ones(Int, 12)],
            [ones(Int, 12)],
        ],
    },
    "Julian" => {
        "hours_min" => 28, # hours worked over simulation
        "hours_max" => 32, # hours worked over simulation
        "shift_count_min" => 4,
        "shift_count_max" => 5,
        "availability" => Array[
            [ones(Int, 12)],
            [ones(Int, 12)],
            [ones(Int, 12)],
            [ones(Int, 12)],
            [ones(Int, 12)],
            [ones(Int, 12)],
            [zeros(Int, 4), ones(Int, 8)],
        ],
    },
    "Leonardo" => {
        "hours_min" => 18, # hours worked over simulation
        "hours_max" => 32, # hours worked over simulation
        "shift_count_min" => 3,
        "shift_count_max" => 5,
        "availability" => Array[
            [ones(Int, 7), zeros(Int, 5)],
            [ones(Int, 7), zeros(Int, 5)],
            [ones(Int, 7), zeros(Int, 5)],
            [ones(Int, 7), zeros(Int, 5)],
            [ones(Int, 7), zeros(Int, 5)],
            [ones(Int, 7), zeros(Int, 5)],
            [ones(Int, 7), zeros(Int, 5)],
        ],
    },
    "Patrick" => {
        "hours_min" => 12, # hours worked over simulation
        "hours_max" => 32, # hours worked over simulation
        "shift_count_min" => 2,
        "shift_count_max" => 5,
        "availability" => Array[
            [ones(Int, 12)],
            [ones(Int, 12)],
            [ones(Int, 12)],
            [ones(Int, 12)],
            [ones(Int, 12)],
            [ones(Int, 12)],
            [ones(Int, 12)],
        ],
    },
    "R" => {
        "hours_min" => 12, # hours worked over simulation
        "hours_max" => 32, # hours worked over simulation
        "shift_count_min" => 2,
        "shift_count_max" => 5,
        "availability" => Array[
            [ones(Int, 8), zeros(Int, 4)],
            [ones(Int, 8), zeros(Int, 4)],
            [ones(Int, 8), zeros(Int, 4)],
            [ones(Int, 8), zeros(Int, 4)],
            [ones(Int, 8), zeros(Int, 4)],
            [ones(Int, 12)],
            [ones(Int, 12)],
        ],
    },
    "Shaun" => {
        "hours_min" => 18, # hours worked over simulation
        "hours_max" => 32, # hours worked over simulation
        "shift_count_min" => 3,
        "shift_count_max" => 5,
        "availability" => Array[
            [ones(Int, 12)],
            [ones(Int, 12)],
            [ones(Int, 12)],
            [ones(Int, 12)],
            [ones(Int, 12)],
            [zeros(Int, 2), ones(Int, 10)],
            [zeros(Int, 2), ones(Int, 10)],
        ],
    },
}


using Neuroimaging, Test

@testset "Sensors" begin

    fname = joinpath(dirname(@__FILE__), "..", "..", "data", "test.sfp")
    s = read_sfp(fname)

    @testset "Read file" begin

        s = read_sfp(fname)
        @test length(s) == 3
    end

    @testset "Show" begin

        show(s)
        show(s[1])
    end

    @testset "Info" begin

        @test label(s) == ["Fp1", "Fpz", "Fp2"]
        @test label(s[1]) == "Fp1"
        @test labels(s) == ["Fp1", "Fpz", "Fp2"]
        @test labels(s[1]) == "Fp1"

        @test Neuroimaging.x(s) == [-27.747648u"m", -0.085967u"m", 27.676888u"m"]
        @test Neuroimaging.x(s[1]) == -27.747648u"m"

        @test Neuroimaging.y(s) == [98.803864u"m", 103.555275u"m", 99.133354u"m"]
        @test Neuroimaging.y(s[1]) == 98.803864u"m"

        @test Neuroimaging.z(s) == [34.338360u"m", 34.357265u"m", 34.457005u"m"]
        @test Neuroimaging.z(s[1]) == 34.338360u"m"
    end

    @testset "Matching" begin

        a, b = match_sensors(s, ["Fp1", "Fp2"])

    end

    @testset "Leadfield matching" begin

        L = ones(4, 3, 5)
        L[:, :, 2] = L[:, :, 2] * 2
        L[:, :, 3] = L[:, :, 3] * 3
        L[:, :, 4] = L[:, :, 4] * 4
        L[:, :, 5] = L[:, :, 5] * 5

        L, idx = match_sensors(L, ["meh", "Cz", "Fp1", "Fp2", "eh"], ["Fp2", "Cz"])

        @test size(L) == (4, 3, 2)
        @test L[:, :, 1] == ones(4, 3) * 4
        @test L[:, :, 2] == ones(4, 3) * 2
        @test idx == [4, 2]

    end

    @testset "Standard sets" begin

        @test length(EEG_64_10_20) == 64
        @test length(EEG_Vanvooren_2014) == 18
        @test length(EEG_Vanvooren_2014_Left) == 9
        @test length(EEG_Vanvooren_2014_Right) == 9

    end

    @testset "Optodes" begin

        label = "source1"
        coord = Talairach(68.3, -26.9, 8.3)
        s = Source(label, coord, Dict())
        @test isa(s, Neuroimaging.Sensor)
        @test isa(s, Neuroimaging.Optode)
        @test isa(s, Neuroimaging.Source)
        @test !isa(s, Neuroimaging.Detector)
        @test !isa(s, Neuroimaging.Electrode)
        d = Detector(label, coord, Dict())
        @test isa(d, Neuroimaging.Sensor)
        @test isa(d, Neuroimaging.Optode)
        @test isa(d, Neuroimaging.Detector)
    end

end


"""
`struct ProductEvaluator` : specifies a ProductEvaluator, which is basically defined
through a PIBasis and its coefficients. The n-correlations are evaluated directly
via a naive product of the atomic base.
"""
mutable struct ProductEvaluator{T, TPI <: PIBasis, REAL}
   pibasis::TPI        # AA basis from ACE papers
   coeffs::Vector{T}   # c̃ coefficients from ACE papers
   real::REAL          # the real operation stored in the SymmetricBasis
end

==(V1::ProductEvaluator, V2::ProductEvaluator) =
      (V1.pibasis == V2.pibasis) && (V1.coeffs == V2.coeffs)

# ----------- FIO

write_dict(ev::ProductEvaluator) = Dict( "__id__" => "ACE_ProductEvaluator" )

read_dict(::Val{:ACE_ProductEvaluator}, D::Dict, basis, c) = ProductEvaluator(basis, c)

# ------------------------------------------------------------
#   Initialisation and Parameter manipulation code
# ------------------------------------------------------------

ProductEvaluator(basis::SymmetricBasis, c::AbstractVector) =
      ProductEvaluator(basis.pibasis, _get_eff_coeffs(basis, c), basis.real)

# basic setter interface without any checks
function set_params!(ev::ProductEvaluator, c̃::AbstractVector)
   ev.coeffs[:] .= c̃
   return ev
end

# trivial setter when the parameters already come in the c̃ format (AA-basis)
function set_params!(ev::ProductEvaluator, basis::PIBasis, c̃::AbstractVector)
   @assert ev.pibasis === basis
   set_params!(ev, c̃)
end

# if parameters come in c (B-basis) format then they first need to be converted
# to c̃ format (AA basis)
function set_params!(ev::ProductEvaluator, basis::SymmetricBasis, c::AbstractVector)
   len_AA = length(ev.pibasis)
   @assert len_AA == size(basis.A2Bmap, 2)
   c̃ = _acquire_ctilde(basis, len_AA, c)
   _get_eff_coeffs!(c̃, basis, c)
   set_params!(ev, basis.pibasis, c̃)
   release!(basis.B_pool, c̃)
   return ev
end

_get_eff_coeffs!(c̃, basis::SymmetricBasis, c::AbstractVector) =
   genmul!(c̃, transpose(basis.A2Bmap), c, *)

function _get_eff_coeffs(basis::SymmetricBasis, c::AbstractVector)
   # c̃ = acquire_B!(basis, size(basis.A2Bmap, 2))
   c̃ = _alloc_ctilde(basis,c)
   return _get_eff_coeffs!(c̃, basis, c)
end

# TODO: we may need a second pool to allocate ctilde vectors...

struct _One <: Number
end

import Base: *
*(x, ::_One) = x
*(::_One, x) = x
*(x::AbstractProperty, ::_One) = x
*(::_One, x::AbstractProperty) = x
*(x::StaticArray, ::_One) = x
*(::_One, x::StaticArray) = x
*(::ACE._One, x::ACE.XState) = x
*(x::ACE.XState, ::ACE._One) = x

_acquire_ctilde(basis::SymmetricBasis, len_AA, c::AbstractVector{<: Number}) =
      zeros(promote_type(eltype(basis.A2Bmap), eltype(c)), len_AA)

_acquire_ctilde(basis::SymmetricBasis, len_AA, c::AbstractVector{<: SVector}) =
      zeros(SVector{length(c[1]),
                    promote_type(eltype(basis.A2Bmap), eltype(c[1]))
                   } , len_AA )

_alloc_ctilde(basis::SymmetricBasis, c::AbstractVector{<: SVector}) =
      zeros(SVector{length(c[1]),eltype(basis.A2Bmap)}, size(basis.A2Bmap, 2))

_alloc_ctilde(basis::SymmetricBasis, c::AbstractVector{<: Number}) =
      zeros(eltype(basis.A2Bmap), size(basis.A2Bmap, 2))

_alloc_dAco(dAAdA::AbstractVector, A::AbstractVector, c̃::AbstractArray, args...) =
            _alloc_dAco(dAAdA, A, c̃[1], args...)

function _alloc_dAco(dAAdA::AbstractVector, A::AbstractVector,
                     c̃::Union{TP, SVector{N, TP}}, dp = _One()
                     ) where {N, TP <: AbstractProperty}
   c̃_dp = contract(c̃, dp)
   zeros( promote_type(eltype(dAAdA), typeof(c̃_dp)), length(A) )
end

# ------------------------------------------------------------
#   Standard Evaluation code
# ------------------------------------------------------------

evaluate(::LinearACEModel, V::ProductEvaluator, cfg::AbstractConfiguration) =
      evaluate(V::ProductEvaluator, cfg)

# compute one "site energy"
function evaluate(V::ProductEvaluator, cfg::AbstractConfiguration)
   A = acquire_B!(V.pibasis.basis1p, cfg)
   evaluate!(A, V.pibasis.basis1p, cfg)
   spec = V.pibasis.spec
   pireal = V.pibasis.real
   symreal = V.real
   # initialize output with a sensible type
   val = symreal(zero(eltype(V.coeffs)) * pireal(zero(eltype(A))))

   # constant (0-order)
   if spec.orders[1] == 0
      val += V.coeffs[1]
      iAAinit = 2
   else
      iAAinit = 1
   end

   @inbounds for iAA = iAAinit:length(spec)
      aa = A[spec.iAA2iA[iAA, 1]]
      for t = 2:spec.orders[iAA]
         aa *= A[spec.iAA2iA[iAA, t]]
      end
      val += symreal(pireal(aa) * V.coeffs[iAA])
   end

   release_B!(V.pibasis.basis1p, A)
   return val
end

# compute one site energy gradient
grad_config!(g, m::LinearACEModel, V::ProductEvaluator,
                     cfg::AbstractConfiguration) =
      _rrule_evaluate!(g, _One(), m, V, cfg)

function _rrule_evaluate(dp, model::LinearACEModel, cfg::AbstractConfiguration)
   g = acquire_grad_config!(model, cfg, dp)
   return _rrule_evaluate!(g, dp, model, model.evaluator, cfg)
end

# NB - testing shows that pre-allocating everything gains about 10% for small
#      configs and ca 20% for larger, more realistic configs.
#      worth doing at some point, but not really an immediate priority!
function _rrule_evaluate!(g, dp, m::LinearACEModel, V::ProductEvaluator,
                          cfg::AbstractConfiguration)
   basis1p = V.pibasis.basis1p
   pireal = V.pibasis.real
   symreal = V.real
   A = acquire_B!(V.pibasis.basis1p, cfg)
   dA = acquire_dB!(V.pibasis.basis1p, cfg)    # MAJOR ALLOCATION!!
   dAAdA = _acquire_dAAdA!(V.pibasis)

   # stage 1: precompute all the A values
   evaluate_ed!(A, dA, basis1p, cfg)

   # stage 2: compute the coefficients for the ∇A_{nlm} = ∇ϕ_{nlm}
   # dAco[nlm] = coefficient of ∇A_{nlm} (via adjoints)
   c̃ = V.coeffs
   dAco =  _alloc_dAco(dAAdA, A, c̃, dp)        # TODO: ALLOCATION
   spec = V.pibasis.spec

   if spec.orders[1] == 0; iAAinit = 2; else iAAinit = 1; end

   fill!(dAco, zero(eltype(dAco)))
   @inbounds for iAA = iAAinit:length(spec)
      _AA_local_adjoints!(dAAdA, A, spec.iAA2iA, iAA, spec.orders[iAA], pireal)
      @fastmath for t = 1:spec.orders[iAA]
         dAco[spec.iAA2iA[iAA, t]] += dAAdA[t] * contract(dp, c̃[iAA])
      end
   end

   # stage 3: get the gradients
   fill!(g, zero(eltype(g)))
   for iX = 1:length(cfg), iA = 1:length(basis1p)
      g[iX] += symreal( coco_o_daa(dAco[iA], dA[iA, iX]) )
   end

   release_B!(V.pibasis.basis1p, A)
   release_dB!(V.pibasis.basis1p, dA)

   return g
end

function adjoint_EVAL_D(m::LinearACEModel, V::ProductEvaluator, cfg, w)
   basis1p = V.pibasis.basis1p
   TDX = gradtype(m.basis, cfg)
   _real = V.real
   A = acquire_B!(V.pibasis.basis1p, cfg)
   dA = acquire_dB!(V.pibasis.basis1p, cfg)
   dAAdA = _acquire_dAAdA!(V.pibasis)

   # dAw = acquire_B!(V.pibasis.basis1p, cfg)
   # dAAw = acquire_B!(V.pibasis, cfg)
   _dAw = contract(w[1], dA[1])
   dAw = zeros(typeof(_dAw), length(V.pibasis.basis1p))

   _dAAw = _real(_dAw * dAAdA[1])
   dAAw = zeros(typeof(_dAAw), length(V.pibasis))

   # [1] dA_t = ∑_j ∂ϕ_t / ∂X_j
   evaluate_ed!(A, dA, basis1p, cfg)
   # fill!(dAw, 0)
   for k = 1:length(basis1p), j = 1:length(w)
      dAw[k] += contract(w[j], dA[k, j])
   end

   # [2] dAA_k
   spec = V.pibasis.spec
   # fill!(dAAw, 0)
   if spec.orders[1] == 0; iAAinit=2; else; iAAinit=1; end
   @inbounds for iAA = iAAinit:length(spec)
      ord = spec.orders[iAA]
      _AA_local_adjoints!(dAAdA, A, spec.iAA2iA, iAA, ord, _real)
      @fastmath for t = 1:ord
         vt = spec.iAA2iA[iAA, t]
         dAAw[iAA] += _real(dAw[vt] * dAAdA[t])
      end
   end

   # dB = similar(m.c)
   # @show m.basis.A2Bmap[1]
   # @show dAAw[1]
   # @show m.basis.A2Bmap[1].val * dAAw[1]
   # genmul!(dB, m.basis.A2Bmap, dAAw, (a, x) -> a * x)
   dB = m.basis.A2Bmap * dAAw

   release_B!(V.pibasis.basis1p, A)
   release_dB!(V.pibasis.basis1p, dA)
   # release_B!(V.pibasis.basis1p, dAw)
   # release_B!(V.pibasis, dAAw)

   # [3] dB_k
   return dB
end

# #for multiple properties. dispatch on the pullback input being a matrix.
# #Basically the same code, except for some parts where we loop over all properties.
# #We generate a list of size "nprop" and keep the same objects as for a single property
# #inside the list.
# function adjoint_EVAL_D(m::LinearACEModel, V::ProductEvaluator, cfg, wt::Matrix)
#    _contract = ACE.contract

#    basis1p = V.pibasis.basis1p
#    dAAdA = zero(MVector{10, ComplexF64})   # TODO: VERY RISKY -> FIX THIS
#    A = zeros(ComplexF64, length(basis1p))
#    TDX = gradtype(m.basis, cfg)
#    dA = zeros(complex(TDX) , length(A), length(cfg))
#    _real = V.real
#    dAAw = [acquire_B!(V.pibasis, cfg) for _ in 1:length(m.c[1])]
#    dAw = [similar(A) for _ in 1:length(m.c[1])]
#    dB = similar(m.c)

#    # [1] dA_t = ∑_j ∂ϕ_t / ∂X_j
#    evaluate_ed!(A, dA, basis1p, cfg)
#    for i in 1:length(m.c[1])
#       fill!(dAw[i], 0)
#    end
#    for prop in 1:length(m.c[1])
#       w = wt[:,prop]
#       for k = 1:length(basis1p), j = 1:length(w)
#          dAw[prop][k] += _contract(w[j], dA[k, j])
#       end
#    end

#    # [2] dAA_k
#    spec = V.pibasis.spec
#    for i in 1:length(m.c[1])
#       fill!(dAAw[i], 0)
#    end
#    for prop in 1:length(m.c[1])
#       if spec.orders[1] == 0; iAAinit=2; else; iAAinit=1; end
#       @inbounds for iAA = iAAinit:length(spec)
#          _AA_local_adjoints!(dAAdA, A, spec.iAA2iA, iAA, spec.orders[iAA], _real)
#          @fastmath for t = 1:spec.orders[iAA]
#             vt = spec.iAA2iA[iAA, t]
#             dAAw[prop][iAA] += _real(dAw[prop][vt] * dAAdA[t])
#          end
#       end
#    end

#    adjointgenmul!(dB, m.basis.A2Bmap, dAAw, (a, x) -> a.val * x)

#    for i in 1:length(m.c[1])
#       release_B!(V.pibasis, dAAw[i]) #TODO check that this indeed releases
#    end

#    # [3] dB_k
#    return dB
# end


#__precompile__(true)

module UNSflow

import Dierckx: Spline1D, derivative, evaluate

import ForwardDiff

import DelimitedFiles

import Serialization

import NLsolve: nlsolve, not_in_place

import Statistics: mean

import PyPlot: plot, scatter, figure, xlabel, ylabel, xlim, ylim,
    xticks, yticks, subplot, subplot2grid, legend, axis, savefig,
    close, tight_layout

import Plots: @layout

import LaTeXStrings: @L_str

#For use in development and debugging
import Revise

export
    # kinematics types and funtions
    MotionDef,
    KinemPar,
    KinemDef,
    EldUpDef,
    EldUptstartDef,
    ConstDef,
    EldRampReturnDef,
    EldUpIntDef,
    EldUpInttstartDef,
    SinDef,
    CosDef,

    # 2D low-order solver types
    TwoDSurf,
    TwoDOFPar,
    KinemPar2DOF,
    TwoDVort,
    TwoDFlowField,
    KelvinCondition,
    KelvinKutta,

    # vortex count control utility
    delVortDef,
    delNone,
    delSpalart,

    # 3D low-order solver types
    ThreeDSurfSimple,
    KinemDef3D,
    ThreeDFieldSimple,

    # utility functions
    simpleTrapz,
    camber_calc,
    find_tstep,
    simpleInterp,
    cleanWrite,

    # 2D low-order solver function
    update_boundpos,
    update_kinem,
    update_indbound,
    update_downwash,
    update_a0anda1,
    place_tev,
    place_lev,
    update_a2toan,
    mutual_ind,
    update_a2a3adot,
    update_bv,
    ind_vel,
    wakeroll,
    update_adot,
    update_externalvel,
    controlVortCount,
    update_kinem,
    update_kinem2DOF,

    #2D low-order solver methods
    lautat,
    #lautatRoll,
    ldvm,
    ldvmLin,

    #3D low-order solver methods
    QSLLT_lautat,
    QSLLT_ldvm,

    # Postprocessing functions
    calc_forces,
    writeStamp,

    # 2D plotting functions
    viewVort2D,
    viewVortConnect2D,

    # 2D plot output functions
    makeForcePlots2D,
    makeVortPlots2D,

    # 3D plot output functions
    makeForcePlots3D,
    makeVortPlots3D,
    makeInfoPlots3D

### source files

# kinematic types
include("kinem.jl")

# utility functions
include("utils.jl")

# vortex count control utility
include("delVort.jl")

# low-order 2D solvers
include("lowOrder2D/typedefs.jl")            # type definitions
include("lowOrder2D/calcs.jl")               # calculation functions
include("lowOrder2D/solvers.jl")             # solver methods
include("lowOrder2D/postprocess.jl")         # postprocessing functions

# low-order 3D solvers
include("lowOrder3D/typedefs.jl")            # type definitions
include("lowOrder3D/calcs.jl")               # calculation functions
include("lowOrder3D/solvers.jl")             # solver methods
include("lowOrder3D/postprocess.jl")         # postprocessing functions

# 2D plotting functions
include("plots/plots2D.jl")
include("plots/plots3D.jl")

end


module GlobalComShadow

using Printf
import SearchBall, GeometryTypes
const sb = SearchBall
const gt = GeometryTypes

function init_strategies(config; kwargs...)
    sb.init_text!(config, "shadow", "")
    sb.add_check_button!(config, "with hfov", "with hfov") do widget
        run_strategies(config)
    end
end

function calc_myPlayers(config::sb.Config, dt::Number, set_type::Symbol, moved_object::Union{sb.Object{<:Number}, Nothing}=nothing)
    ws = config.ws

    sb.remove_graphics_filter!(config) do s
        startswith(s, "coreShadow")
    end

    all_shadows = [sb.get_shadow_polygons(player, config, sb.get_manip_value(config, "with hfov")) for player in ws.myPlayers if sb.is_inside(player, config.field_corners)]

    if !isempty(all_shadows)
        core_shadows = all_shadows[1]
        for off_shadow in @view(all_shadows[2:end])
            new_core_shadows = Vector{gt.Point2{Float64}}[]
            for s1 in core_shadows, s2 in off_shadow
                core_shads = sb.polygon_intersection(s1, s2)
                for core in core_shads
                    if length(core) > 2
                        push!(new_core_shadows, core)
                    end
                end
            end
            core_shadows = new_core_shadows
        end
    else
        core_shadows = []
    end

    area = 0.0
    for (idx, shadow) in enumerate(core_shadows)
        isempty(shadow) && continue
        push!(shadow, shadow[1])
        sb.show_graphic!(config, string("coreShadow", idx), shadow, sb.GraphicLine(fill=true, fillcolor=[0.0, 0.0, 0.0, 0.2]))
        area += sb.polygon_area(shadow)
    end
    sb.show_text!(config, "shadow", @sprintf("%.2fm²", area))
    return
end

function calc_opponend(config::sb.Config, dt::Number, set_type::Symbol, moved_object::Union{sb.Object{<:Number}, Nothing}=nothing)
    ws = config.ws

    # do stuff

    return
end

function run_strategies(config::sb.Config, dt::Number=0, set_type::Symbol=:default, moved_object::Union{sb.Object{<:Number}, Nothing}=nothing)
    calc_myPlayers(config, dt, set_type, moved_object)
    calc_opponend(config, dt, set_type, moved_object)
    return
end

end


using Celeste.SDSSIO
import Celeste.SDSSIO: SkyIntensity

@testset "sky interpolations" begin
    small_sky = [1.  2.  3.  4.;
                 5.  6.  7.  8.;
                 9. 10. 11. 12]
    sky_x = [0.1, 2.5]
    sky_y = [0.5, 2.5, 4.]
    sky = SkyIntensity(small_sky, sky_x, sky_y, ones(2))

    @test sky[1, 1] ≈ 1.0
    @test sky[2, 1] ≈ 7.0
    @test sky[1, 2] ≈ 2.5
    @test sky[2, 2] ≈ 8.5
    @test sky[1, 3] ≈ 4.0
    @test sky[2, 3] ≈ 10.0
end

@testset "test coordinates out of bounds" begin
    small_sky = [1.  2.  3.  4.;
            5.  6.  7.  8.;
            9. 10. 11. 12]
    sky_x = [-5.0, 4.0]
    sky_y = [-4.0, 5.0]
    sky = SkyIntensity(small_sky, sky_x, sky_y, ones(2))

    @test sky[1,1] == 1.0
    @test sky[1,2] == 4.0
    @test sky[2,1] == 9.0
    @test sky[2,2] == 12.0
end

@testset "read_photoobj handles missing table extensions" begin
    # get an example of a photoobj file with a missing table
    strategy = PlainFITSStrategy(joinpath(Pkg.dir("Celeste"), "test", "data"))
    rcf = RunCamcolField(6597, 4, 25)
    fname = SDSSIO.compute_fname(strategy, SDSSIO.PhotoObj(rcf))[1]
    if !isfile(fname)
        run(`curl --create-dirs -o $fname https://data.sdss.org/sas/dr12/boss/photoObj/301/6597/4/photoObj-006597-4-0025.fits`)
    end

    catalog = SDSSIO.read_photoobj(strategy, rcf)
end


using GeometryTypes, Visualize, Colors, GLAbstraction
using Visualize: Resolution, Open, Visible
using Visualize.GLRasterization: glwindow, renderloop
using Visualize: Model, Color, Position, LineVertex, LineSegments, Camera, BasicCamera

resolution = (1024, 1024)

window = glwindow(
    Resolution => resolution,
    Camera => BasicCamera(),
    Visualize.Debugging => true
);

scale = 10f0
res = Vec2f0(resolution)
radius = (min(resolution...) / 2f0) - scale

N = 32
ls = LineSegments((
    Position => [rand(Vec2f0) .* res for x in 1:N],
    Color => Vec4f0(1, 0,0,1)
))
ls
x = Visualize.visualize(window, ls)

renderloop(window)
window[Camera]
map(x-> window[Camera][ProjectionView] * Vec4f0(x..., 0, 1.0),  [rand(Vec2f0) .* res for x in 1:N])
# Visualize.destroy!(window)

using Visualize, GeometryTypes, Colors, FileIO, NIfTI

window = Visualize.GLRasterization.glwindow(
    Resolution => (1024, 1024),
    Visualize.Color => Vec4f0(0.0, 0.0, 0.0, 0.0),
)

niivol = Array(niread("/media/s/94BC3648BC362560/Users/sdani/juliastuff/demo/t2.nii"))
vol = niivol ./ maximum(niivol)

colorfun2(color, uniforms, intensities) = color

function accum2(accum, pos, stepdir, intensities, uniforms)
    alpha = intensities[pos][1] * 0.1f0
    sample = Vec4f0(0.9f0, 0f0, 0.4f0, alpha)
    alpha = 1f0 - ((1f0 - alpha) ^ 1.2f0)
    sample = Vec4f0(
        sample[1] * alpha,
        sample[2] * alpha,
        sample[3] * alpha,
        alpha
    )
    sample = (1f0 - accum[4]) * sample
    tmp = accum + sample
    tmp, tmp[4] > 0.95f0
end

volume = Visualize.Volume((
    Data => vol,
    AccumulationFunction => accum2,
    StartValue => Vec4f0(0),
    ColorFunction => colorfun2,
))

wf = wireframe(
    AABB(Vec3f0(-0.01), Vec3f0(1.02)), ()
)

visualize(window, wf)
visualize(window, volume)
@async renderloop(window)


"""
NQCModels define the potentials and derivatives that
govern the dynamics of the particles.
These can exist as analytic models or as interfaces to other codes.
"""
module NQCModels

using Reexport: @reexport

export potential
export derivative
export nstates
export ndofs

"""
Top-level type for models.

# Implementation
When adding new models, this should not be directly subtyped. Instead, depending on
the intended functionality of the model, one of the child abstract types should be
subtyped.
If an appropriate type is not already available, a new abstract subtype should be created.
"""
abstract type Model end

Base.broadcastable(model::Model) = Ref(model)

"""
    potential(model::Model, R::AbstractMatrix)

Evaluate the potential at position `R` for the given `model`.
"""
function potential(model::Model, R::AbstractMatrix)
    if ndofs(model) == 1
        if size(R, 2) == 1
            return potential(model, R[1])
        else
            return potential(model, @view R[1,:])
        end
    elseif size(R, 2) == 1
        return potential(model, @view R[:,1])
    else
        throw(MethodError(potential, (model, R)))
    end
end

"""
    potential!(model::Model, V, R::AbstractMatrix)

In-place version of `potential`, used only when mutable arrays are preferred.

Currently used only for `LargeDiabaticModels`, see `diabatic/DiabaticModels.jl`.
"""
function potential!(model::Model, V, R::AbstractMatrix)
    if ndofs(model) == 1
        if size(R, 2) == 1
            return potential!(model, V, R[1])
        else
            return potential!(model, V, view(R, 1, :))
        end
    elseif size(R, 2) == 1
        return potential!(model, V, view(R, :, 1))
    else
        throw(MethodError(potential!, (model, V, R)))
    end

end

"""
    derivative!(model::Model, D, R::AbstractMatrix)

Fill `D` with the derivative of the electronic potential as a function of the positions `R`.

This must be implemented for all models.
"""
function derivative!(model::Model, D, R::AbstractMatrix)
    if ndofs(model) == 1
        if size(R, 2) == 1
            D[1] = derivative(model, R[1])
            return D
        else
            derivative!(model, view(D, 1, :), view(R, 1, :))
            return D
        end
    elseif size(R, 2) == 1
        derivative!(model, view(D, :, 1), view(R, :, 1))
        return D
    else
        throw(MethodError(derivative!, (model, D, R)))
    end
end

"""
    derivative(model::Model, R)

Allocating version of `derivative!`, this definition should be suitable for all models.

Implement `zero_derivative` to allocate an appropriate array then implement `derivative!`
to fill the array.
"""
function derivative(model::Model, R)
    D = zero_derivative(model, R)
    derivative!(model, D, R)
    return D
end

"""
    zero_derivative(model::Model, R)

Create an zeroed array of the right size to match the derivative.
"""
function zero_derivative end

"""
    nstates(::Model)

Get the number of electronic states in the model.
"""
nstates(::Model) = error("This should return the total number of electronic states.")

"""
    ndofs(::Model)

Get the number of degrees of freedom for every atom in the model. Usually 1 or 3.
"""
ndofs(::Model) = error("This should return the number of degrees of freedom for each atom.")

dofs(model::Model) = Base.OneTo(ndofs(model))

state_independent_potential(model, r) = 0.0
state_independent_derivative!(model, derivative, r) = fill!(derivative, zero(eltype(r)))
nelectrons(::Model) = error("This should return the total number of electrons.")
fermilevel(model) = 0.0

eachelectron(model::Model) = Base.OneTo(nelectrons(model))
eachstate(model::Model) = Base.OneTo(nstates(model))

mobileatoms(::Model, n::Int) = Base.OneTo(n)

include("adiabatic/AdiabaticModels.jl")
@reexport using .AdiabaticModels

include("friction/FrictionModels.jl")
@reexport using .FrictionModels

include("diabatic/DiabaticModels.jl")
@reexport using .DiabaticModels

include("plot.jl")

end # module


using NQCDynamics
using Test
using RecursiveArrayTools
using DiffEqBase: CallbackSet

atoms = Atoms([:H])
model = NQCModels.Harmonic()
sim = Simulation(atoms, model)

positions = [randn(1, length(atoms)) for i=1:10]
velocities = [randn(1, length(atoms)) for i=1:10]
distribution = DynamicalDistribution(positions, velocities, (1, 1))

u = rand(distribution)
prob = DynamicsMethods.create_problem(u, (0.0, 1.0), sim)
@testset "OrderedSelection" begin
    selector = Ensembles.OrderedSelection(distribution, 1:10)
    new_prob = selector(prob, 3, false)
    u = getindex(distribution, 3)
    @test new_prob.u0 == ArrayPartition(u.v, u.r)
    selector = Ensembles.OrderedSelection(distribution, 6:10)
    new_prob = selector(prob, 3, false)
    u = getindex(distribution, 8)
    @test new_prob.u0 == ArrayPartition(u.v, u.r)
end

@testset "RandomSelection" begin
    selector = Ensembles.RandomSelection(distribution)
    new_prob = selector(prob, 3, false)
end

@testset "SelectWithCallbacks" begin
    selector = Ensembles.RandomSelection(distribution)
    Ensembles.SelectWithCallbacks(selector, CallbackSet(), (:position,), 10)
end


@sumProductRule(:node_type     => Nonlinear{Sampling},
                :outbound_type => Message{SampleList},
                :inbound_types => (Nothing, Message),
                :name          => SPNonlinearSOutNM)

@sumProductRule(:node_type     => Nonlinear{Sampling},
                :outbound_type => Message{Function},
                :inbound_types => (Message, Nothing),
                :name          => SPNonlinearSIn1MN)

mutable struct SPNonlinearSOutNGX <: SumProductRule{Nonlinear{Sampling}} end
outboundType(::Type{SPNonlinearSOutNGX}) = Message{SampleList}
function isApplicable(::Type{SPNonlinearSOutNGX}, input_types::Vector{<:Type})
    total_inputs = length(input_types)
    (total_inputs > 2) || return false
    (input_types[1] == Nothing) || return false

    for input_type in input_types[2:end]
        matches(input_type, Message{Gaussian}) || return false
    end

    return true
end

mutable struct SPNonlinearSInGX <: SumProductRule{Nonlinear{Sampling}} end
outboundType(::Type{SPNonlinearSInGX}) = Message{GaussianWeightedMeanPrecision}
function isApplicable(::Type{SPNonlinearSInGX}, input_types::Vector{<:Type})
    total_inputs = length(input_types)
    (total_inputs > 2) || return false
    (input_types[1] != Nothing) || return false

    nothing_inputs = 0
    factorfunction_input = false
    gaussian_inputs = 0
    if matches(input_types[1], Message{FactorFunction})
        factorfunction_input = true
    end
    for input_type in input_types[2:end]
        if input_type == Nothing
            nothing_inputs += 1
        elseif matches(input_type, Message{Gaussian})
            gaussian_inputs += 1
        end
    end

    return (nothing_inputs == 1) && (gaussian_inputs == total_inputs-2) && factorfunction_input
end

mutable struct SPNonlinearSOutNFactorX <: SumProductRule{Nonlinear{Sampling}} end
outboundType(::Type{SPNonlinearSOutNFactorX}) = Message{SampleList}
function isApplicable(::Type{SPNonlinearSOutNFactorX}, input_types::Vector{<:Type})
    total_inputs = length(input_types)
    (total_inputs > 2) || return false
    (input_types[1] == Nothing) || return false

    factorNode_input = false
    gaussian_inputs = 0
    for input_type in input_types[2:end]
        if matches(input_type, Message{SampleList})
            return false
        elseif matches(input_type, Message{FactorNode})
            factorNode_input += 1
            if matches(input_type, Message{Gaussian})
                gaussian_inputs += 1
            end
        else
            return false
        end
    end

    return (gaussian_inputs < total_inputs-1) && (factorNode_input == total_inputs-1)
end

mutable struct SPNonlinearSInFactorX <: SumProductRule{Nonlinear{Sampling}} end
outboundType(::Type{SPNonlinearSInFactorX}) = Message{Function}
function isApplicable(::Type{SPNonlinearSInFactorX}, input_types::Vector{<:Type})
    total_inputs = length(input_types)
    (total_inputs > 2) || return false
    (input_types[1] != Nothing) || return false

    nothing_inputs = 0
    factorfunction_input = false
    factorNode_input = 0
    gaussian_inputs = 0
    if matches(input_types[1], Message{FactorFunction})
        factorfunction_input = true
    end
    for input_type in input_types[2:end]
        if input_type == Nothing
            nothing_inputs += 1
        elseif matches(input_type, Message{PointMass})
            return false
        elseif matches(input_type, Message{SampleList})
            return false
        elseif matches(input_type, Message{FactorNode})
            factorNode_input += 1
            if matches(input_type, Message{Gaussian})
                gaussian_inputs += 1
            end
        else
            return false
        end
    end

    return (nothing_inputs == 1) && (gaussian_inputs < total_inputs-2) && (factorNode_input == total_inputs-2) && factorfunction_input
end

mutable struct MNonlinearSInGX <: MarginalRule{Nonlinear{Sampling}} end
function isApplicable(::Type{MNonlinearSInGX}, input_types::Vector{<:Type})
    total_inputs = length(input_types)
    (total_inputs > 2) || return false
    (input_types[1] == Nothing) || return false # Indicates marginalization over outbound variable

    for input_type in input_types[2:end]
        matches(input_type, Message{Gaussian}) || return false
    end

    return true
end


function check_why_there_is_no_sell(closed_order, id_str, open_orders, client, config, bot_path)
    t = Pinguu.examine_order(closed_order, id_str, client, config, bot_path)
    fees = config["current_fees"][config["symbol"]]
    new_id    = "SELL_" * id_str
    new_quant = floor(t["use_vol"], digits=config["precision"])

    new_price = Pinguu.price2string(round((abs(t["current_profit"]) / new_quant *
        (1 + fees["maker"] +
            parse(Float64, config["take_profit"])/100)
            * t["price_correction"]), digits=config["sell_precision"]), config["sell_precision"])
    if abs(t["current_profit"]) > config["min_qty"]
        order = try client.order_limit_sell(
            symbol=config["symbol"],
            quantity=new_quant,
            price=new_price);
            catch
            nothing
        end
        open(joinpath(bot_path, "trigger.json"), "w") do f
            JSON.print(f, Dict(), 4)
        end
        @info(order)
        sleep(2)
        if order !== nothing
            order["time"] = string(now())
            open_orders["SELL"][new_id] = order
            open(joinpath(bot_path, "open_orders.json"), "w") do f
                JSON.print(f, open_orders, 4)
            end
        else
            open_orders["SELL"][new_id] = Dict()
            open_orders["SELL"][new_id]["origQty"] = new_quant
            open_orders["SELL"][new_id]["price"]   = new_price
            open_orders["SELL"][new_id]["status"]  = "place order"
            open_orders["SELL"][new_id]["time"]    = string(now())
            open_orders["SELL"][new_id]["fail"]    = 0
            open_orders["SELL"][new_id]["if_fail"] = new_quant - 1/10^config["precision"]

            open(joinpath(bot_path, "open_orders.json"), "w") do f
                JSON.print(f, open_orders, 4)
            end
        end
    end
end
export check_why_there_is_no_sell


module NNlibCUDA

using NNlib
using CUDA
using Random, Statistics

const IntOrIntTuple = Union{Integer, NTuple{N,<:Integer} where N}

include("upsample.jl")
include("activations.jl")
include("batchedmul.jl")
include("scatter.jl")
include("gather.jl")
include("cudnn/cudnn.jl")
include("cudnn/conv.jl")
include("cudnn/pooling.jl")
include("cudnn/softmax.jl")
include("cudnn/activations.jl")
include("cudnn/batchnorm.jl")

end # module


using IntcodeVM;
using DelimitedFiles;
using Combinatorics;

using Sockets;

listeners = [listen(60000 + i) for i in 1:5]

readers = Array{TCPSocket,1}(undef, 5)
for i in 1:5
    @async begin
        readers[i] = accept(listeners[i])
    end
end

senders = [connect(60000 + i) for i in 1:5]

const examples = ["3,15,3,16,1002,16,10,16,1,16,15,15,4,15,99,0,0",
"3,23,3,24,1002,24,10,24,1002,23,-1,23,101,5,23,23,1,24,23,23,4,23,99,0,0",
"3,31,3,32,1002,32,10,32,1001,31,-2,31,1007,31,0,33,1002,33,7,33,1,33,31,31,1,32,31,31,4,31,99,0,0,0"]

function combine_amplifiers(code, p)
    for i in 1:4
        @async begin
            inputs = (i == 1) ? [p[i] - 1, 0] : [p[i] - 1];
            IntcodeVM.run(code, in=readers[p[i]],
                out=senders[p[i + 1]], inputs=inputs);
        end
    end
    # last one is synchronous
    return last(last(IntcodeVM.run(code, in=readers[p[5]], out=devnull, inputs=[p[5] - 1])))
end

for example in examples
    code = parse.(Int,split(example,","))
    println(first(findmax(map(p -> combine_amplifiers(code, p), permutations(1:5)))))
end

print("Day 7 Part 1:    ");
code = vec(readdlm("input.txt", ',', Int, '\n'))
println(first(findmax(map(p -> combine_amplifiers(code, p), permutations(1:5)))))

foreach(close, senders)
foreach(close, readers)
foreach(close, listeners)


# this script produces Figure 3
# GMT was installed using miniconda3
# If GMT is working,

using SeisNoise, Libdl, Glob, GMT
push!(Libdl.DL_LOAD_PATH, expanduser("~/miniconda3/lib/"))

CORRDIR = expanduser("~/SeisNoise_paper/SINGLECORR")
files = glob("*",CORRDIR)

# load first file and allocate
lat = zeros(Float64,188)
lon = zeros(Float64,188)

C = load_corr(files[1],"ZZ")
lon[1] = C.loc.lon
lat[1] = C.loc.lat
# loop through and plot
for ii = 2:188
    C = load_corr(files[ii],"ZZ")
    loc = get_loc(C.loc,C.azi,C.dist)
    lon[ii] = loc.lon
    lat[ii] = loc.lat
end

# create figure
coast(region=:d, proj=:Winkel, frame=:g, res=:low, area=10000, land=:seashell4,
                 water=:skyblue, figsize=18, shore=:black)
scatter!(lon,lat,fill=:gold,marker=:triangle,markeredgecolor=:black,fmt=:png,
         savefig=expanduser("~/SeisNoise_paper/stationmap.png"))


## tree
struct Tree{newicktype, Teltype, Pareltype, idseltype, nameseltype}
    newick::newicktype   #   (:id, :data, :parent, :children)
    T::Vector{Teltype}
    Par::Vector{Pareltype}
    ids::Vector{idseltype}
    names::Vector{nameseltype}
    lids::Vector{idseltype}
    lnames::Vector{nameseltype}
    lastone::Vector{Bool} # true if it is the last child when running forwards
    n::Int64
end

function Tree(S::String)
    nwtree = readnw(S)

    ids = map(x->x.id, prewalk(nwtree))
    names = map(x->x.data.name, prewalk(nwtree))
    n = length(ids)
    # leaves
    l = getleaves(nwtree)
    lids = map(x -> Int(x.id), l)
    lnames = map(x->x.data.name, l)
    #@assert issubset(tipslist[:,1], lnames)

    # extract parents and time spans (on segments)
    Par = zeros(Int, n)  # parent ids
    T = zeros(n)  # time spans
    for node in prewalk(nwtree)
        isroot(node) && continue
        p = node.parent.id
        Par[node.id] = p
        T[node.id] = T[p] + node.data.distance
    end
    @assert all(Par .< 1:n)

    lastone = zeros(Bool,n)
    parentdone = zeros(Bool,n)
    for i ∈ n:-1:2
        if !parentdone[Par[i]]
            lastone[i] = true
            parentdone[Par[i]] = true
        end
    end
    Tree{typeof(nwtree), eltype(T), eltype(Par), eltype(ids), eltype(names)}(nwtree,  T, Par, ids, names, lids, lnames,lastone, n)
end

## read tree (very simple example)
S = "(A:0.6,B:0.2,(C:0.3,D:0.4)E:0.5)F;"
S20 = "(((t20:0.1377639426,t8:0.3216105914):0.461882598,(((t14:0.09650261351,t12:0.5344695284):0.1058950284,(t1:0.7008816295,t13:0.8587384557):0.7271672823):0.8026451592,(t6:0.5240857233,(t16:0.451590274,t10:0.4654036982):0.8569358122):0.140687139):0.7207935331):0.06727131596,((((t17:0.03115223325,t18:0.8930883685):0.6862687438,t9:0.4896655723):0.3056753092,t3:0.2549535078):0.07862622454,(((t4:0.03973490326,t5:0.4096179632):0.3020580804,t15:0.06550096255):0.2660011558,(t7:0.3895205685,((t19:0.165174433,t2:0.6469203965):0.2730371242,t11:0.2700119715):0.8393620371):0.02117527928):0.206376995):0.6012785432);"
S50 = "((((((t23:0.5766009684,t33:0.9917489749):0.3572315329,(t40:0.4535627225,t25:0.1296473157):0.9540394924):0.4740875189,(t6:0.946340692,t47:0.386682119):0.1535291784):0.07034701738,((((t49:0.6612032412,t27:0.1606025242):0.1527153626,t15:0.1854858669):0.5782065149,t50:0.3087045723):0.0759275395,((t17:0.258283512,t35:0.7651973176):0.2003453802,t2:0.4033703087):0.2418845498):0.6834833182):0.4268306631,(((t26:0.1744227773,(t16:0.6275201498,t10:0.09013780276):0.1857208994):0.8804157388,(t28:0.626011228,t46:0.8026940082):0.7066663445):0.6334459039,(t45:0.3075582262,t42:0.2759609232):0.05279792193):0.4003514573):0.5333229457,((((t30:0.05955170072,t41:0.2500256761):0.04172409023,((t19:0.3332289669,t31:0.6650944015):0.2450174955,t9:0.3372890623):0.2041792406):0.3337947708,t1:0.9396123393):0.0896910606,(((((t12:0.1379979164,t43:0.9007850592):0.703310129,(t11:0.3347885907,(t22:0.7190679635,t14:0.8988074451):0.1574828231):0.9644375315):0.137505424,((t37:0.7556577872,t8:0.3481238757):0.7712922962,(t7:0.7946650374,t48:0.8638419164):0.338934626):0.7507952554):0.5193078597,(t4:0.2408366969,(((t29:0.7862726359,t38:0.4319011718):0.8655222375,((t20:0.09156441619,t3:0.5975058537):0.3240908289,(t34:0.7324223334,t32:0.1348467385):0.4636785307):0.7708284601):0.7991992962,t39:0.2996040138):0.4087550903):0.1688272723):0.7225795928,(((t44:0.02794037666,t5:0.09033886855):0.5778779227,(t24:0.8744501146,((t21:0.2940180032,t18:0.03399693617):0.8231791619,t13:0.5192477896):0.5669933448):0.9547001948):0.8561708115,t36:0.9225172002):0.02867545444):0.1595393196):0.1254606678);"

# S20coal = "(((((t5:0.01408886814,t6:0.01408886814):0.03926824924,t3:0.05335711738):0.3769529348,t17:0.4303100522):0.6904821712,((t2:0.05400246327,(t19:0.04515631227,t7:0.04515631227):0.008846151006):0.02843109882,t15:0.08243356209):1.038358661):0.6759374458,(t12:0.5094458799,((t11:0.09067183363,(t4:0.08830667144,t1:0.08830667144):0.002365162193):0.2392304005,((t13:0.1499179245,((t10:0.01410801788,t20:0.01410801788):0.01360477778,t8:0.02771279566):0.1222051288):0.1061152139,(t16:0.1612987658,((t18:0.03933683292,t9:0.03933683292):0.008828557137,t14:0.04816539006):0.1131333758):0.09473437257):0.07386909574):0.1795436458):1.287283789);"

S20coal =  "(((t19:1.201139146,(t16:0.2402278292,t11:0.2402278292):0.9609113167):2.642506121,(t12:3.603417438,((t14:1.921822633,t17:1.921822633):0.9609113167,(t2:2.642506121,t7:2.642506121):0.2402278292):0.7206834875):0.2402278292):0.7206834875,((t20:3.122961779,((t6:0.4804556584,t4:0.4804556584):1.201139146,((t10:0.7206834875,t3:0.7206834875):0.2402278292,t9:0.9609113167):0.7206834875):1.441366975):1.201139146,(((t18:2.162050463,(t5:1.441366975,t8:1.441366975):0.7206834875):1.201139146,t13:3.363189608):0.7206834875,(t15:2.402278292,t1:2.402278292):1.681594804):0.2402278292):0.2402278292);"


function intersects_geometry(m::AbstractAccretionGeometry{T}, line_element) where {T}
    if in_nearby_region(m, line_element)
        return has_intersect(m, line_element)
    end
    false
end

in_nearby_region(d::AbstractAccretionGeometry{T}, line_element) where {T} =
    error("Not implemented for $(typeof(d))")
has_intersect(d::AbstractAccretionGeometry{T}, line_element) where {T} =
    error("Not implemented for $(typeof(d))")

# Jiménez, Segura, Feito. Computation Geometry 43 (2010) 474-492
function jsr_algorithm(V₁::T, V₂::T, V₃::T, Q₁::V, Q₂::V; ϵ = 1e-8) where {T,V}
    A = Q₁ .- V₃
    B = V₁ .- V₃
    C = V₂ .- V₃
    W₁ = B × C
    w = A ⋅ W₁
    if w > ϵ
        D = Q₂ .- V₃
        s = D ⋅ W₁
        s > ϵ && return false
        W₂ = A × D
        t = W₂ ⋅ C
        t < -ϵ && return false
        u = -W₂ ⋅ B
        u < -ϵ && return false
        w < s + t + u && return false
    elseif w < -ϵ
        return false
    else # w == 0
        D = Q₂ .- V₃
        s = D ⋅ W₁
        if s > ϵ
            return false
        elseif s < -ϵ
            W₂ = D × A
            t = W₂ ⋅ C
            t > ϵ && return false
            u = -W₂ ⋅ B
            u > ϵ && return false
            -s > t + u && return false
        else
            return false
        end
    end
    true
end


using Printf

# @todo does Node need to be mutable?
mutable struct Node
    visit_count::Int32
    prior::Float32
    value_sum::Float32
    children::Dict{Int32,Node}
    hidden_state::Union{Nothing,AbstractArray{Float32, 3}}
    reward::Float32
    info::String
    state::AbstractEnvState # @todo - remove state altogether
    prior_policy::Dict
    action_scores::Dict
    function Node(prior, state; visit_count=0, value_sum=0,
        children=Dict(), hidden_state=nothing, reward=0)
        new(visit_count, prior, value_sum, children, hidden_state, reward, "", state, Dict(), Dict())
    end
end

gethiddenstate(n::Node) = n.hidden_state

gettoplay(n::Node) = n.state.to_play

function json(x::Node, visit_pr=0.0)
    d = Dict()
    noisy_prior_policy = Dict()
    for (action, node) in x.children
        noisy_prior_policy[action] = node.prior
    end
    d["visit_pr"] = visit_pr
    d["visit_count"] = x.visit_count
    d["to_play"] = get_player_label(x)
    d["prior"] = x.prior
    d["value_sum"] = x.value_sum
    d["prior_policy"] = x.prior_policy
    d["noisy_prior_policy"] = noisy_prior_policy
    d["reward"] = x.reward
    d["children"] = Dict()
    d["action_scores"] = x.action_scores
    total_child_visit_count = sum([child.visit_count for child in values(x.children)])
    for (action, child) in x.children
        visit_pr = child.visit_count / total_child_visit_count
        d["children"][action] = json(child, visit_pr)
    end
    d
end

get_player_label(x::Node) = x.state.to_play.idx == PLAYER_ID ? "A" : "B"

function Base.show(io::IO, x::Node)
    show_indented(io, x)
end

function show_indented(io::IO, x::Node, level=0, action=nothing)
    #show(io, x.channels)
    print(" "^(level*5))
    action_label = isnothing(action) ? "" : "A($action)"
    hex_state = "hex_str" #hex_str(x.state)
    @printf "%s %s pr: %.2f {c: %i, r: %i, vs: %.3f}|%s %s" action_label hex_state x.prior x.visit_count x.reward x.value_sum get_player_label(x) x.info
    println(io)
    actions = [action for action in keys(x.children)]
    for action in sort(actions)
        child = x.children[action]
        show_indented(io, child, level+1, action)
    end
end

expanded(x::Node) = length(x.children) > 0

function calc_value(x::Node)
    Float32(x.visit_count == 0 ? 0.0 : x.value_sum / x.visit_count)
end

function priors(x::Node)
    sort([(action, child.prior) for (action, child) in x.children])
end

#=
We expand a node using the value, reward and policy prediction obtained from
the neural network.
=#

function expand_node!(node::Node, actions::Vector, output::SingleNetworkOutput, reward::Number)
    p_logits = policy_logits(output)
    hidden = hidden_state(output)
    node.hidden_state = hidden
    node.reward = Float32(reward)

    for (action, pr) in enumerate(softmax(p_logits))
        node.prior_policy[action] = pr
    end

    if !isempty(actions)
        policy = Dict([a=>exp(p_logits[a.idx]) for a in actions])
        policy_sum = sum(values(policy))
        for (action, p) in policy
            node.children[action.idx] = Node(p / policy_sum, getnextstate(node.state, action)) #reward=get_reward(node.state, action)
        end
    end

end

function expand_node!(
    nodes::Vector{Node},
    actionSets::Vector,
    output::NetworkOutput,
    reward::Vector{Float32}
    )
    setfield!.(
        nodes,
        :hidden_state,
        eachslice(hidden_state(output), dims=4)
    )
    setfield!.(
        nodes,
        :reward,
        Float32.(reward)
    )
    #### @todo - efficiency could be much improved
    p_logits_tensor = policy_logits(output)
    for (idx, actions) in enumerate(actionSets)
        node = nodes[idx]
        p_logits = view(p_logits_tensor, :, idx)
        for (action, pr) in enumerate(softmax(p_logits))
            node.prior_policy[action] = pr
        end
        if !isempty(actions)
            policy = Dict([a=>exp(p_logits[a.idx]) for a in actions])
            policy_sum = sum(values(policy))
            for (action, p) in policy
                node.children[action.idx] = Node(p / policy_sum, getnextstate(node.state, action))
            end
        end
    end
    ####
end


Base.@kwdef mutable struct EvaluateDQNAgent{M, AP<:AbstractValuePolicy, Φ, IR<:HistImageReplay} <: AbstractAgent
    model::M
    ap::AP
    replay::IR
    prev_s::Φ
    hist::Int
    action::Int = 0
    INFO::Dict{Symbol, Any} = Dict{Symbol, Any}()
end

EvaluateDQNAgent(model, ϵ, hist, action_set, img_manip, img_norm, img_size) =
    EvaluateDQNAgent(
        model = model,
        ap = ϵGreedy(0.1, action_set),
        replay = HistImageReplay(1, img_size, img_manip, img_norm, hist, 1),
        prev_s = zeros(Int, hist),
        hist = hist
    )

function RLCore.start!(agent::EvaluateDQNAgent,
                       env_s_tp1,
                       rng::AbstractRNG;
                       kwargs...)

    agent.prev_s .= add!(agent.replay, env_s_tp1)
    state = agent.replay.img_norm(
        reshape(getindex(agent.replay.image_buffer, agent.prev_s),
                agent.replay.image_buffer.img_size..., agent.replay.hist,
                1))
    agent.action = sample(agent.ap,
                          agent.model(state),
                          rng)#::eltype(agent.ap)
    return agent.action
end

function RLCore.step!(agent::EvaluateDQNAgent,
                      env_s_tp1,
                      r,
                      terminal,
                      rng::AbstractRNG;
                      kwargs...)

    add_ret = add!(agent.replay,
                   (agent.prev_s,
                    findfirst((a)->a==agent.action,
                              agent.ap.action_set)::Int,
                    env_s_tp1,
                    r,
                    terminal))
    agent.prev_s .= add_ret
    prev_s =
        agent.replay.img_norm(
            reshape(getindex(agent.replay.image_buffer, agent.prev_s),
                    agent.replay.image_buffer.img_size..., agent.replay.hist,
                    1))

    agent.action = sample(agent.ap,
                          agent.model(prev_s),
                          rng)
    return agent.action
end


module CustomConfigHexTest

using Hashids
using Test

@testset "Custom Config Hex" begin

config = Hashids.configure(
    salt="this is my salt",
    min_length=30,
    alphabet="xzal86grmb4jhysfoqp3we7291kuct5iv0nd"
)

@testset "encodehex and decodehex back" begin

@test decodehex(config, encodehex(config, "deadbeef")) == "deadbeef"
@test decodehex(config, encodehex(config, "abcdef123456")) == "abcdef123456"
@test decodehex(config, encodehex(config, "ABCDDD6666DDEEEEEEEEE")) == lowercase("ABCDDD6666DDEEEEEEEEE")
@test decodehex(config, encodehex(config, "507f1f77bcf86cd799439011")) == "507f1f77bcf86cd799439011"
@test decodehex(config, encodehex(config, "f00000fddddddeeeee4444444ababab")) == "f00000fddddddeeeee4444444ababab"
@test decodehex(config, encodehex(config, "abcdef123456abcdef123456abcdef123456")) == "abcdef123456abcdef123456abcdef123456"
@test decodehex(config, encodehex(config, "f000000000000000000000000000000000000000000000000000f")) == "f000000000000000000000000000000000000000000000000000f"
@test decodehex(config, encodehex(config, "fffffffffffffffffffffffffffffffffffffffffffffffffffff")) == "fffffffffffffffffffffffffffffffffffffffffffffffffffff"

end  # of "encodehex and decodehex back"

@testset "encodehex" begin

@test encodehex(config, "deadbeef") == "0dbq3jwa8p4b3gk6gb8bv21goerm96"
@test encodehex(config, "abcdef123456") == "190obdnk4j02pajjdande7aqj628mr"
@test encodehex(config, "ABCDDD6666DDEEEEEEEEE") == "a1nvl5d9m3yo8pj1fqag8p9pqw4dyl"
@test encodehex(config, "507f1f77bcf86cd799439011") == "1nvlml93k3066oas3l9lr1wn1k67dy"
@test encodehex(config, "f00000fddddddeeeee4444444ababab") == "mgyband33ye3c6jj16yq1jayh6krqjbo"
@test encodehex(config, "abcdef123456abcdef123456abcdef123456") == "9mnwgllqg1q2tdo63yya35a9ukgl6bbn6qn8"
@test encodehex(config, "f000000000000000000000000000000000000000000000000000f") == "edjrkn9m6o69s0ewnq5lqanqsmk6loayorlohwd963r53e63xmml29"
@test encodehex(config, "fffffffffffffffffffffffffffffffffffffffffffffffffffff") == "grekpy53r2pjxwyjkl9aw0k3t5la1b8d5r1ex9bgeqmy93eata0eq0"

end  # of "encodehex"

end  # of "Custom Config Hex"

end  # module


abstract type Gate end

struct QulacsGate <: Gate
    pyobj::PyObject
end

function RX(target::Int, angle::Float64)
    QulacsGate(qulacs.gate.RX(target, angle))
end

function RY(target::Int, angle::Float64)
    QulacsGate(qulacs.gate.RY(target, angle))
end

function RZ(target::Int, angle::Float64)
    QulacsGate(qulacs.gate.RZ(target, angle))
end

function get_matrix(gate::QulacsGate)
    gate.pyobj.get_matrix()
end


# functions for PETSc TS (time stepping) algorithms
export TS, set_ic, set_times, solve!, set_rhs_function, set_rhs_jac
export set_lhs_function, set_lhs_jac
export ComputeRHSFunctionLinear, ComputeRHSJacobianConstant

mutable struct TS{T}
  p::C.TS{T}
  data::Array{Any, 1}  # hold the various ctx tuples
  function TS{T}(p::C.TS{T}, data=nothing; first_instance=false) where {T}
    data_arr = Any[]
    ts = new{T}(p, data_arr)
    push!(ts.data, data)

    if first_instance
      if (!deactivate_finalizers)
        finalizer(PetscDestroy, ts)
      end
    end

    return ts
  end
end

function PetscDestroy(ts::TS{T}) where {T}
  if !PetscFinalized(T) && !isfinalized(ts)
    ts_ref = Ref(ts.p)
    chk(C.TSDestroy(ts_ref))
    ts.p = C.TS{T}(C_NULL)
  end
end

function isfinalized(ts::TS)
  return isfinalized(ts.p)
end

function isfinalized(ts::C.TS)
  return ts.pobj == C_NULL
end

#get the internal KSP object for this TS
function KSP(ts::TS{T}) where {T}
  ksp_c = Ref{C.KSP{T}}()
  chk(C.TSGetKSP(ts.p, ksp_c))
  return KSP{T}(ksp_c[])
end

"""
  Most preferred constructor: take ProblemType, method from options
  database
"""
function TS(::Type{T} ;comm=MPI.COMM_WORLD) where {T<:Scalar}
  ts = Ref{C.TS{T}}()
  chk(C.TSCreate(comm, ts))
  return TS{T}(ts[])
end

"""
  Preferred constructor: set problem type explicitly, get method from
  options database
"""
function TS(::Type{T}, tsptype::C.TSProblemType; comm=MPI.COMM_WORLD) where {T<:Scalar}

  ts = TS(T, comm=comm)
  chk(C.TSSetProblemType(ts.p, tsptype))

  return ts
end

"""
  More explicit constructor: set problem type, method directly

`tsptype` sets the problem type and can be one of the following:
* `TS_LINEAR` - a linear set of ODEs
* `TS_NONLINEAR` - a nonlinear set of ODEs or DEAs

`tstype` sets the method used to solve the problem. More information
about the possible methods is available at the official PETSc [docs](http://www.mcs.anl.gov/petsc/petsc-current/docs/manualpages/TS/TSType.html).
"""
function TS(::Type{T}, tsptype::C.TSProblemType, tstype::C.TSType;
                       comm=MPI.COMM_WORLD) where {T<:Scalar}

  ts = TS(T, tsptype, comm=comm)
  chk(C.TSSetType(ts.p, tstype))

  return ts
end

function set_ic(ts::TS{T}, u::Vec{T}) where {T<:Scalar}
  chk(C.TSSetSolution(ts.p, u.p))
end

"""
  Set the times related quantities:
    t0 : initial time value
    dt0: initial time step
    nsteps: maximum number of steps
    tmax: maximum time value
"""
function set_times(ts::TS{T}, t0, dt0,  nsteps::Integer, tmax) where {T<:Scalar}
  TR = real(T)  # PetscReal

  chk(C.TSSetInitialTimeStep(ts.p, TR(t0), TR(dt0)))
  chk(C.TSSetDuration(ts.p, nsteps, TR(tmax)))
end

function petscview(ts::TS{T}) where {T<:Scalar}
  viewer = C.PetscViewer{T}(C_NULL)
  chk(C.TSView(ts.p, viewer))
end

"""
  Solve the system using the initial condition provided in vec
"""
function solve!(ts::TS{T}, vec::Vec{T}) where {T}

  chk(C.TSSolve(ts.p, vec.p))
end

"""
  Solve the system using the intitial condition proived by set_ic
"""
function solve!(ts::TS{T}) where {T}

  vecp = C.Vec{T}(C_NULL)
  chk(C.TSSolve(ts.p, vecp))
end

###############################################################################
# right hand side function

"""
  Wrapper for the right hand side function.  This function is always passed
  to PETSc as the right hand side function, and calls the user supplied
  function internally.  The user supplied function must be the first
  component of the ctx tuple

"""
function rhs_wrapper(ts::C.TS{T}, t, u::C.Vec{T}, F::C.Vec{T}, ctx_ptr::Ptr{Cvoid}) where {T}

  Treal = real(T)
  # transform into high level objects
  bigts = TS{T}(ts, first_instance=false)
  bigu = Vec{T}(u, first_instance=false)
  bigF = Vec{T}(F, first_instance=false)

  ctx = unsafe_pointer_to_objref(ctx_ptr)
  func = ctx[1]
  ctx_inner = ctx[2]  # the user provided ctx

  ret_status = func(bigts, Treal(t), bigu, bigF, ctx_inner)
  return PetscErrorCode(ret_status)
end

"""
  Sets the function that evalutes u_t = g(u, t) for an ODE.
  The function must have the signature:

  f(TS, t, U, F, ctx)

  where TS is a TS object,
  t is the current time
  u is the current state vector
  F is the vector to be populated with u_t
  ctx is the user supplied context tuple (empty tuple if not provided)
"""
function set_rhs_function(ts::TS{T}, r::Vec{T}, f::Function, ctx=()) where {T}

  ctx_outer = (f, ctx)
  push!(ts.data, ctx_outer)
  #ctx_ptr = pointer_from_objref(ctx_outer)
  Treal = real(T)
  # this this pre-compilable?
  #fptr = @cfunction(rhs_wrapper, PetscErrorCode, (C.TS{T}, Treal, C.Vec{T}, C.Vec{T}, Ptr{Cvoid}))
  fptr = @cfunction(rhs_wrapper, PetscErrorCode, (C.TS{T}, T, C.Vec{T}, C.Vec{T}, Ptr{Cvoid}))

  chk(C.TSSetRHSFunction(ts.p, r.p, fptr, ctx_ptr))
end

# a PETSc provided rhs function for the linear, time invarient coefficient
# matrix case

function ComputeRHSFunctionLinear(ts::TS, t, u::Vec, F::Vec, ctx::Tuple)

  # this is amusing, wrapping things just to unwrap them again
  C.TSComputeRHSFunctionLinear(ts.p, t, u.p, F.p, C_NULL)
end

###############################################################################
# right hand side jacobian
function rhs_jac_wrapper(ts::C.TS{T}, t, u::C.Vec{T}, A::C.Mat{T}, B::C.Mat{T}, ctx_ptr::Ptr{Cvoid}) where {T}

  Treal = real(T)
  bigts = TS{T}(ts, first_instance=false)
  bigu = Vec{T}(u, first_instance=false)
  bigA = Mat{T}(A, first_instance=false)
  bigB = Mat{T}(B, first_instance=false)

  ctx = unsafe_pointer_to_objref(ctx_ptr)
  func = ctx[1]
  ctx_inner = ctx[2]

  ret_status = func(bigts, Treal(t), bigu, bigA, bigB, ctx_inner)

  return PetscErrorCode(ret_status)
end

function set_rhs_jac(ts::TS{T}, A::Mat{T}, B::Mat{T}, f::Function, ctx=()) where {T}

  ctx_outer = (f, ctx)
  push!(ts.data, ctx_outer)
  ctx_ptr = pointer_from_objref(ctx_outer)
  Treal = real(T)
  # this this pre-compilable?
  #fptr = @cfunction(rhs_jac_wrapper, PetscErrorCode, (C.TS{T}, Treal, C.Vec{T}, C.Mat{T}, C.Mat{T}, Ptr{Cvoid}))
  fptr = @cfunction(rhs_jac_wrapper, PetscErrorCode, (C.TS{T}, T, C.Vec{T}, C.Mat{T}, C.Mat{T}, Ptr{Cvoid}))

  chk(C.TSSetRHSJacobian(ts.p, A.p, B.p, fptr, ctx_ptr))
end

# a PETSc provided rhs jacobian function for time invarient jacobians
function ComputeRHSJacobianConstant(ts::TS, t, u::Vec, A::Mat, B::Mat, ctx::Tuple)

  chk(C.TSComputeRHSJacobianConstant(ts.p, t, u.p, A.p, B.p, C_NULL))
end

###############################################################################
# left hand side function

function set_lhs_function(ts::TS{T}, res::Vec{T}, f::Function, ctx::Tuple=()) where {T}

  ctx_outer = (f, ctx)
  push!(ts.data, ctx_outer)
  ctx_ptr = pointer_from_objref(ctx_outer)
  Treal = real(T)

  #fptr = @cfunction(lhs_wrapper, PetscErrorCode, (C.TS{T}, Treal, C.Vec{T}, C.Vec{T}, C.Vec{T}, Ptr{Cvoid}))
  fptr = @cfunction(rhs_jac_wrapper, PetscErrorCode, (C.TS{T}, T, C.Vec{T}, C.Mat{T}, C.Mat{T}, Ptr{Cvoid}))

  chk(C.TSSetIFunction(ts.p, res.p, fptr, ctx_ptr))
  return nothing
end

function lhs_wrapper(ts::C.TS{T}, t, u::C.Vec{T}, ut::C.Vec{T}, F::C.Vec{T}, ctx_ptr::Ptr{Cvoid}) where {T}

  Treal = real(T)
  # transform into high level objects
  bigts = TS{T}(ts, first_instance=false)
  bigu = Vec{T}(u, first_instance=false)
  bigut = Vec{T}(ut, first_instance=false)
  bigF = Vec{T}(F, first_instance=false)

  ctx = unsafe_pointer_to_objref(ctx_ptr)
  func = ctx[1]
  ctx_inner = ctx[2]  # the user provided ctx

  ret_status = func(bigts, Treal(t), bigu, bigut, bigF, ctx_inner)
  return PetscErrorCode(ret_status)
end

###############################################################################
# left hand side jacobian

function lhs_jac_wrapper(ts::C.TS{T}, t, u::C.Vec{T}, ut::C.Vec{T}, a, A::C.Mat{T}, B::C.Mat{T}, ctx_ptr::Ptr{Cvoid}) where {T}

  Treal = real(T)
  bigts = TS{T}(ts, first_instance=false)
  bigu = Vec{T}(u, first_instance=false)
  bigut = Vec{T}(ut, first_instance=false)
  bigA = Mat{T}(A, first_instance=false)
  bigB = Mat{T}(B, first_instance=false)

  ctx = unsafe_pointer_to_objref(ctx_ptr)
  func = ctx[1]
  ctx_inner = ctx[2]  # the user provided ctx

  ret_status = func(bigts, bigu, bigut, Treal(a), bigA, bigB, ctx_inner)
  return PetscErrorCode(ret_status)
end

function set_lhs_jac(ts::TS{T}, A::Mat{T}, B::Mat{T}, f::Function, ctx::Tuple=()) where {T}

  ctx_outer = (f, ctx)
  push!(ts.data, ctx_outer)
  ctx_ptr = pointer_from_objref(ctx_outer)
  Treal = real(T)

  #fptr = @cfunction(lhs_jac_wrapper, PetscErrorCode, (C.TS{T}, Treal, C.Vec{T}, C.Vec{T}, Treal, C.Mat{T}, C.Mat{T}, Ptr{Cvoid}))
  fptr = @cfunction(lhs_jac_wrapper, PetscErrorCode, (C.TS{T}, T, C.Vec{T}, C.Vec{T}, T, C.Mat{T}, C.Mat{T}, Ptr{Cvoid}))

  chk(C.TSSetIJacobian(ts.p, A.p, B.p, fptr, ctx_ptr))
end


using ShaderToy
shadertoy("rayprimitive.frag")