text:
module ModifiersWrapper

using JavaCall

const AbstractShapiroBaseModifier = @jimport org.orekit.estimation.measurements.modifiers.AbstractShapiroBaseModifier
const AngularIonosphericDelayModifier = @jimport org.orekit.estimation.measurements.modifiers.AngularIonosphericDelayModifier
const AngularRadioRefractionModifier = @jimport org.orekit.estimation.measurements.modifiers.AngularRadioRefractionModifier
const AngularTroposphericDelayModifier = @jimport org.orekit.estimation.measurements.modifiers.AngularTroposphericDelayModifier
const AtmosphericRefractionModel = @jimport org.orekit.models.AtmosphericRefractionModel
const AttitudeProvider = @jimport org.orekit.attitudes.AttitudeProvider
const Bias = @jimport org.orekit.estimation.measurements.modifiers.Bias
const Class = @jimport java.lang.Class
const DiscreteTroposphericModel = @jimport org.orekit.models.earth.troposphere.DiscreteTroposphericModel
const DynamicOutlierFilter = @jimport org.orekit.estimation.measurements.modifiers.DynamicOutlierFilter
const EstimatedMeasurement = @jimport org.orekit.estimation.measurements.EstimatedMeasurement
const FieldSpacecraftState = @jimport org.orekit.propagation.FieldSpacecraftState
const Gradient = @jimport org.hipparchus.analysis.differentiation.Gradient
const GroundStation = @jimport org.orekit.estimation.measurements.GroundStation
const IonosphericGradientConverter = @jimport org.orekit.estimation.measurements.modifiers.IonosphericGradientConverter
const IonosphericModel = @jimport org.orekit.models.earth.ionosphere.IonosphericModel
const JString = @jimport java.lang.String
const List = @jimport java.util.List
const Object = @jimport java.lang.Object
const OnBoardAntennaInterSatellitesRangeModifier = @jimport org.orekit.estimation.measurements.modifiers.OnBoardAntennaInterSatellitesRangeModifier
const OnBoardAntennaPhaseModifier = @jimport org.orekit.estimation.measurements.modifiers.OnBoardAntennaPhaseModifier
const OnBoardAntennaRangeModifier = @jimport org.orekit.estimation.measurements.modifiers.OnBoardAntennaRangeModifier
const OnBoardAntennaTurnAroundRangeModifier = @jimport org.orekit.estimation.measurements.modifiers.OnBoardAntennaTurnAroundRangeModifier
const OutlierFilter = @jimport org.orekit.estimation.measurements.modifiers.OutlierFilter
const PhaseAmbiguityModifier = @jimport org.orekit.estimation.measurements.modifiers.PhaseAmbiguityModifier
const PhaseIonosphericDelayModifier = @jimport org.orekit.estimation.measurements.modifiers.PhaseIonosphericDelayModifier
const PhaseTroposphericDelayModifier = @jimport org.orekit.estimation.measurements.modifiers.PhaseTroposphericDelayModifier
const RangeIonosphericDelayModifier = @jimport org.orekit.estimation.measurements.modifiers.RangeIonosphericDelayModifier
const RangeRateIonosphericDelayModifier = @jimport org.orekit.estimation.measurements.modifiers.RangeRateIonosphericDelayModifier
const RangeRateTroposphericDelayModifier = @jimport org.orekit.estimation.measurements.modifiers.RangeRateTroposphericDelayModifier
const RangeTroposphericDelayModifier = @jimport org.orekit.estimation.measurements.modifiers.RangeTroposphericDelayModifier
const RealFieldElement = @jimport org.hipparchus.RealFieldElement
const ShapiroInterSatelliteRangeModifier = @jimport org.orekit.estimation.measurements.modifiers.ShapiroInterSatelliteRangeModifier
const ShapiroPhaseModifier = @jimport org.orekit.estimation.measurements.modifiers.ShapiroPhaseModifier
const ShapiroRangeModifier = @jimport org.orekit.estimation.measurements.modifiers.ShapiroRangeModifier
const SpacecraftState = @jimport org.orekit.propagation.SpacecraftState
const TroposphericGradientConverter = @jimport org.orekit.estimation.measurements.modifiers.TroposphericGradientConverter
const TurnAroundRangeIonosphericDelayModifier = @jimport org.orekit.estimation.measurements.modifiers.TurnAroundRangeIonosphericDelayModifier
const TurnAroundRangeTroposphericDelayModifier = @jimport org.orekit.estimation.measurements.modifiers.TurnAroundRangeTroposphericDelayModifier
const Vector3D = @jimport org.hipparchus.geometry.euclidean.threed.Vector3D

export AbstractShapiroBaseModifier
export AngularIonosphericDelayModifier
export AngularRadioRefractionModifier
export AngularTroposphericDelayModifier
export Bias
export DynamicOutlierFilter
export IonosphericGradientConverter
export OnBoardAntennaInterSatellitesRangeModifier
export OnBoardAntennaPhaseModifier
export OnBoardAntennaRangeModifier
export OnBoardAntennaTurnAroundRangeModifier
export OutlierFilter
export PhaseAmbiguityModifier
export PhaseIonosphericDelayModifier
export PhaseTroposphericDelayModifier
export RangeIonosphericDelayModifier
export RangeRateIonosphericDelayModifier
export RangeRateTroposphericDelayModifier
export RangeTroposphericDelayModifier
export ShapiroInterSatelliteRangeModifier
export ShapiroPhaseModifier
export ShapiroRangeModifier
export TroposphericGradientConverter
export TurnAroundRangeIonosphericDelayModifier
export TurnAroundRangeTroposphericDelayModifier
export equals
export get_class
export get_free_state_parameters
export get_parameters
export get_parameters_drivers
export get_sigma
export get_state
export hash_code
export modify
export notify
export notify_all
export range_rate_error_tropospheric_model
export set_sigma
export to_string
export wait

include("abstract_shapiro_base_modifier.jl")
include("angular_ionospheric_delay_modifier.jl")
include("angular_radio_refraction_modifier.jl")
include("angular_tropospheric_delay_modifier.jl")
include("bias.jl")
include("dynamic_outlier_filter.jl")
include("ionospheric_gradient_converter.jl")
include("on_board_antenna_inter_satellites_range_modifier.jl")
include("on_board_antenna_phase_modifier.jl")
include("on_board_antenna_range_modifier.jl")
include("on_board_antenna_turn_around_range_modifier.jl")
include("outlier_filter.jl")
include("phase_ambiguity_modifier.jl")
include("phase_ionospheric_delay_modifier.jl")
include("phase_tropospheric_delay_modifier.jl")
include("range_ionospheric_delay_modifier.jl")
include("range_rate_ionospheric_delay_modifier.jl")
include("range_rate_tropospheric_delay_modifier.jl")
include("range_tropospheric_delay_modifier.jl")
include("shapiro_inter_satellite_range_modifier.jl")
include("shapiro_phase_modifier.jl")
include("shapiro_range_modifier.jl")
include("tropospheric_gradient_converter.jl")
include("turn_around_range_ionospheric_delay_modifier.jl")
include("turn_around_range_tropospheric_delay_modifier.jl")

end


text:
# This file is a part of RRRMC.jl. License is MIT: http://github.com/carlobaldassi/RRRMC.jl/LICENCE.md

module PercXEntr

using Random
using ExtractMacro
using ..Interface
using ..Common

using ..DeltaE.ArraySets

using ...RRRMC # this is silly but it's required for correct cross-linking in docstrings, apparently

export GraphPercXEntr

import ..Interface: energy, delta_energy, update_cache!, neighbors

function gen_ξ(N::Integer, P::Integer)
    ξ = BitArray(undef, P, N)
    ξv = Vector{BitVector}(undef, P)
    for a = 1:P
        ξv[a] = bitrand(N)
        ξ[a,:] = ξv[a]
    end
    # the ξv representation is for energy
    # the ξ representation is for delta_energy and update_cache!
    return ξ, ξv
end

mutable struct Stabilities
    Δs::Vector{Int}
    newΔs::Vector{Int}
    ξsi::BitVector
    last_move::Int
    justΔ::Bool
    function Stabilities(P::Integer)
        Δs = zeros(P)
        newΔs = zeros(P)
        ξsi = BitArray(undef, P)
        return new(Δs, newΔs, ξsi, 0, false)
    end
end

function swap_Δs!(stab::Stabilities)
    stab.Δs, stab.newΔs = stab.newΔs, stab.Δs
    return stab
end

Δ2i(Δ, N) = (Δ + N) >>> 1 + 1

mutable struct GraphPercXEntr{λ} <: SimpleGraph{Float64}
    N::Int
    P::Int
    ξ::BitMatrix
    ξv::Vector{BitVector}
    stab::Stabilities
    Hs::Vector{Float64}
    tmps::BitVector
    function GraphPercXEntr{λ}(ξ::BitMatrix, ξv::Vector{BitVector}) where {λ}
        P, N = size(ξ)
        #TODO: check
        λ isa Float64 || throw(ArgumentError("λ must be a Float64, given a $(typeof(λ))"))
        isodd(N) || throw(ArgumentError("N must be odd, given: $N"))
        stab = Stabilities(P)
        Hs = [log(1 + exp(-2 * λ * Δ / √(N))) for Δ = -N:2:N]
        # for Δ = -N:2:N
        #     @assert Hs[Δ2i(Δ, N)] == log((1 + exp(-2 * λ * Δ / √(N))))
        # end
        tmps = BitVector(undef, N)
        return new{λ}(N, P, ξ, ξv, stab, Hs, tmps)
    end
end

"""
    GraphPercXEntr(N::Integer, P::Integer, λ::Float64) <: SimpleGraph{Float64}

A `SimpleGraph` implementing a single-layer binary perceptron with `N` binary (\$±1\$) synapses,
trained on `P` random i.i.d. \$±1\$ patterns.

The energy of the model is computed for each pattern using the cross-entropy loss (with a `λ`
parameter: ``\\log(1 + \\exp(-2λ σ⋅ξ / √N))``.

See also [`GraphPercStep`](@ref RRRMC.GraphPercStep), [`GraphPercLinear`](@ref RRRMC.GraphPercLinear)
"""
GraphPercXEntr(N::Integer, P::Integer, λ::Float64) = GraphPercXEntr{λ}(gen_ξ(N, P)...)

GraphPercXEntr(X::GraphPercXEntr{λold}, λnew::Float64) where {λold} = GraphPercXEntr{λnew}(X.ξ, X.ξv)

function Base.empty!(stab::Stabilities)
    @extract stab : Δs
    fill!(Δs, 0)
    stab.last_move = 0
    stab.justΔ = false
    return stab
end

function energy(X::GraphPercXEntr{λ}, C::Config) where {λ}
    @assert X.N == C.N
    @extract C : s
    @extract X : N P ξv stab Hs tmps
    @extract stab : Δs

    E = 0.0
    empty!(stab)
    tmps = BitArray(undef, N)

    for a = 1:P
        # tmps[:] = ξ[a,:]
        # map!(⊻, tmps, s, tmps)
        # Δ = N - 2 * sum(tmps)
        map!(⊻, tmps, s, ξv[a])
        Δ = N - 2 * sum(tmps)
        Δs[a] = Δ
        # E += H(λ * Δ)
        E += Hs[Δ2i(Δ, N)]
    end

    return E
end

function update_cache!(X::GraphPercXEntr{λ}, C::Config, move::Int) where {λ}
    @assert X.N == C.N

    @extract X : stab
    @extract stab : last_move justΔ
    # NOTE: this only works if called right after delta_energy
    if justΔ && last_move == move
        swap_Δs!(stab)
        return
    end

    @extract C : s
    @extract X : N P ξ
    @extract stab : Δs ξsi
    @assert N == length(s)

    si = ~s[move]
    last_move ≠ move && unsafe_copyto!(ξsi, 1, ξ, (move-1)*P + 1, P)
    stab.last_move = move
    si && flipbits!(ξsi)
    # @assert ξsi == ξ[:,move] .⊻ si
    @inbounds for a = 1:P
        oldΔ = Δs[a]
        newΔ = oldΔ + (4 * ξsi[a] - 2)
        Δs[a] = newΔ
    end
    si && flipbits!(ξsi)
    stab.justΔ = false

    return
end

function delta_energy_naive(X::GraphPercXEntr, C::Config, move::Int)
    @assert C.N == X.N
    @assert 1 ≤ move ≤ N

    spinflip!(C, move)
    E1 = energy(X, C)
    spinflip!(C, move)
    E0 = energy(X, C)
    X.stab.justΔ = false
    return E1 - E0
end

function delta_energy(X::GraphPercXEntr{λ}, C::Config, move::Int) where {λ}
    @assert C.N == X.N
    @extract C : s
    @extract X : N P ξ stab Hs
    @extract stab : Δs newΔs ξsi last_move
    @assert 1 ≤ move ≤ N

    si = s[move]
    last_move ≠ move && unsafe_copyto!(ξsi, 1, ξ, (move-1)*P + 1, P)
    stab.last_move = move
    si && flipbits!(ξsi)
    # @assert ξsi == ξ[:,move] .⊻ si
    ΔE = 0.0
    @inbounds for a = 1:P
        oldΔ = Δs[a]
        newΔ = oldΔ + (4 * ξsi[a] - 2)
        newΔs[a] = newΔ
        # @assert -N ≤ newΔ ≤ N  (N, oldΔ, newΔ, ξsi[a])

        # ΔE += H(λ * newΔ) - H(λ * oldΔ)
        ΔE += Hs[Δ2i(newΔ, N)] - Hs[Δ2i(oldΔ, N)]
    end
    si && flipbits!(ξsi)
    # td = delta_energy_naive(X, s, stab, move)
    # ΔE == td || @show ΔE, td
    # @assert ΔE == td
    stab.justΔ = true
    return ΔE
end

function check_delta(X::GraphPercXEntr, C::Config, move::Int)
    delta = delta_energy(X, C, move)
    e0 = energy(X, C)
    spinflip!(X, C, move)
    e1 = energy(X, C)
    spinflip!(X, C, move)

    @assert isapprox(e1 - e0, delta, atol=1e-9) (e1, e0, delta, e1-e0)
end

function step_energy(X::GraphPercXEntr)
    @extract X : P stab
    @extract stab : Δs
    E = 0
    @inbounds for a = 1:P
        E += Δs[a] < 0
    end
    return E
end

neighbors(X::GraphPercXEntr, i::Int) = AllButOne(X.N, i)

end

ext: jl


text:
#!/usr/bin/julia
# Looks for a bound.
# I'm not sure if it's currently working.

"""
Ternary search (based on Wikipedia). This searches for an integer
    value maximizing f.
  f: function to maximize
  left, right: bounds
  Returns: value maximizing f
"""
function ternarySearch(f, left, right)
  print(string(left) * " " * string(right) * "\n")
  # left and right are the current bounds;
  # the maximum is between them
  if abs(right - left) <= 1
    return trunc((left+right)/2)
  end
  leftThird = ((2*left + right)/3)
  rightThird = ((left + 2*right)/3)
  if f(trunc(leftThird)) < f(trunc(rightThird))
    return ternarySearch(f, leftThird, right)
  else
    return ternarySearch(f, left, rightThird)
  end
end

"""
Counting bound, based on Shannon's argument.
  m: number of edges
  w: number of 'wires' -- that is, log2(number of functions),
    which is the number of bits needed to specify a function
  Returns: average number of gates required to computea
    any of those functions. (This may not be an integer).
"""
function countingBound(m, w)
  b = m - 0.5
  # the "-1" here is because this is the average, not the max.
  sqrt(2*w + b*b) - b - 1
end

"""
  Number of maximal hypercliques of some size.
  Note that k < r < n .
  k: number of vertices per hyperedge
  r: number of vertices in the clique
  n: vertices in the larger graph
  Returns: number of maximal hypercliques
"""
function numMaximalCliques(k, r, n)
  k = BigInt(k)
  r = BigInt(r)
  n = BigInt(n)
  one = BigInt(1)
  two = BigInt(2)

  # expected number of r-cliques
  numRCliques = Rational(binomial(n, r),
    two ^ binomial(r, k))

  # probability that one of those is not covered by a larger clique
  a = two ^ binomial(r, k-one)
  p = Rational(a-one, a) ^ (n-r)
  # println(a)
  # println(p)
  # result is number of cliques, * prob. they're maximal
  numRCliques * p
end

"""
  Bound, using one size of clique.
    This is approximate, because it's approximating the number of
    edges and cliques.
  k: number of vertices per hyperedge
  r: number of vertices for the clique finder
  n: vertices in the larger graph
  Returns: number of NAND gates
"""
function bound2(k, r, n)
  k = BigInt(k)
  r = BigInt(r)
  n = BigInt(n)
  one = BigInt(1)
  two = BigInt(2)

  # the bound on the total number of gates
  minTotalGates = ceil(countingBound(binomial(k, two), binomial(n, k)))

  # expected number of edges "left over" (not covered
  # by an r-clique)
  # p = (1 - 2 ^ (-(choose(r, k)-1))) ^ choose(n-k, r-k)
  a = two ^ (binomial(r, k) - one)
  # p = Rational(a-one, a) ^ binomial(n-k, r-k)
  # print("0\n")
  b = binomial(n-k, r-k)
  p = Rational((a-one)^b, a^b)
  # print("1\n")
  # num.left.over = (p/2) * choose(n, k)
  numLeftOver = p * binomial(n, k) / two
  # print("2\n")

  # expected number of r-cliques
  numRCliques = Rational(binomial(BigInt(n), BigInt(r)),
    BigInt(2) ^ binomial(r, k))
  # print("3\n")
  # the remaining gates are part of some number of
  # clique detectors
  r = (minTotalGates - numLeftOver) / numRCliques
end

"""
  Bound, using one size of clique.
    This is approximate, as it's estimating the number of edges
    and cliques. It also approximates the "(1 - 1/a)^..." bit.
  k: number of vertices per hyperedge
  r: number of vertices for the clique finder
  n: vertices in the larger graph
  Returns: number of NAND gates
"""
function bound3(k, r, n)
  k = BigInt(k)
  r = BigInt(r)
  n = BigInt(n)
  one = BigInt(1)
  two = BigInt(2)

  # the bound on the total number of gates (in the larger graph)
  minTotalGates = ceil(countingBound(binomial(k, two), binomial(n, k)))

  # expected number of edges "left over" (not covered by an r-clique)
  # First, the probability of a hyperedge being "left over"
  # (not covered by an r-clique) is
  #
  # p = (1 - 2 ^ (-(choose(r, k)-1))) ^ choose(n-k, r-k)
  #
  # We use an approximate upper bound (which should be good
  # for large numbers).
  logP = - binomial(n-k, r-k) / (two ^ (binomial(r, k) - one))
  # (log of) number of edges left over
  logNumLeftOver = logP + log(binomial(n, k) / two)

  # expected number of r-cliques
  # (Using right-shift for division, and adding 1 in case
  # of rounding error)
  numRCliques = (binomial(n, r) >> binomial(r, k)) + 1
  # the remaining gates are part of some number of
  # clique detectors
  r = (minTotalGates - exp(logNumLeftOver)) / numRCliques
end

function printBound1(k, r, n)
  b = bound1(k,r,n)
  print(string(k, ",", r, ",", n, ",", b, "\n"))
end

# print("k,r,n,bound\n")
# for n in 100:20:600
#   printBound1(5, 6, n)
# end

max_stars_repo_name: joshtburdick/misc

max_issues_repo_name: joshtburdick/misc

max_forks_repo_name: joshtburdick/misc


text:
module day1

using DelimitedFiles
readInput() = vec(readdlm(joinpath(@__DIR__, "input.txt"), Int))

fuelRequired(x) = x ÷ 3 - 2
part1(input) = sum(fuelRequired, input)

function fuelRequired2(x)
    fuel = fuelRequired(x)
    fuel > 0 ? fuel + fuelRequired2(fuel) : 0
end
part2(input) = sum(fuelRequired2, input)

end


text:
# using Gadfly
using DataFrames

using QuadGK
using GeneralTools

sigma = 0.01;
mu = 0.0;
# define the acceleration as a function of normal distribution:
fn(x) = (sigma*sqrt(2*pi))^-1 * exp(-0.5*((x-mu)/sigma)^2)

T = -1.0:0.01:3.0; # define x domain
A = fn.(T);

function rsign(v) # sign with tolerance
	tol = 10^-5
	if -tol < v < tol
		absv = 0.0;
	else
		absv = sign(v);
	end
	return absv
end

# Define damping
Fc = 2.0;
Fv = 4.0;
# Calculate integration of A
I_E = quadgk.(fn,minimum(T),collect(T)[2:end]); # (I, E)
v0 = 0.0;
vc = copy(v0);
vv = copy(v0);
V = Float64[v0];
Vc = copy(V);
Vv = copy(V);
dts = diff(T);

for (index, value) in enumerate(I_E)
	global vc, vv
	dt = dts[index];
	push!(V, value[1]);

	vc = vc + A[index]*dt - Fc*rsign(vc)*dt;
	vv = vv + A[index]*dt - Fv*vv*dt;
	push!(Vc, vc);
	push!(Vv, vv);
end
(V, Vc, Vv)

df_wide = DataFrame("time"=>T, "acceleration"=>A, "velocity"=>V, "velocity (with damping)"=> Vc);

println("""A "wide" dataframe: """);
preview(df_wide,10);

# a long dataframe where the three variables are grouped by their column name
df_long = stack(df_wide, ["acceleration", "velocity", "velocity (with damping)"]);

println("""The converted "long" dataframe: """);
preview(df_long,10)

plot(df_long, x = "time", y="value", color="variable", Geom.line)

max_stars_repo_name: okatsn/UnpredictableLife

max_issues_repo_name: okatsn/UnpredictableLife

max_forks_repo_name: okatsn/UnpredictableLife


text:
const SUITESPARSE_LDL_SOLVER = 0
const MKL_PARDISO_SOLVER = 1

# Define OSQP infinity constants
const OSQP_INFTY = 1e20

# OSQP return values
# https://github.com/oxfordcontrol/osqp/blob/master/include/constants.h
const status_map = Dict{Int,Symbol}(4 => :Dual_infeasible_inaccurate,
    3 => :Primal_infeasible_inaccurate,
    2 => :Solved_inaccurate,
    1 => :Solved,
    -2 => :Max_iter_reached,
    -3 => :Primal_infeasible,
    -4 => :Dual_infeasible,
    -5 => :Interrupted,
    -6 => :Time_limit_reached,
    -10 => :Unsolved)

const SOLUTION_PRESENT = [:Solved_inaccurate, :Solved, :Max_iter_reached]

# UPDATABLE_DATA
const UPDATABLE_DATA = [:q, :l, :u, :Px, :Px_idx, :Ax, :Ax_idx]

# UPDATABLE_SETTINGS
const UPDATABLE_SETTINGS = [:max_iter, :eps_aps, :eps_rel, :eps_prim_inf, :eps_dual_inf, :time_limit,
              :rho, :alpha, :delta, :polish, :polish_refine_iter, :verbose,
              :check_termination,:warm_start]


text:
# *********************************************************************************
# REopt, Copyright (c) 2019-2020, Alliance for Sustainable Energy, LLC.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification,
# are permitted provided that the following conditions are met:
#
# Redistributions of source code must retain the above copyright notice, this list
# of conditions and the following disclaimer.
#
# Redistributions in binary form must reproduce the above copyright notice, this
# list of conditions and the following disclaimer in the documentation and/or other
# materials provided with the distribution.
#
# Neither the name of the copyright holder nor the names of its contributors may be
# used to endorse or promote products derived from this software without specific
# prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
# IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
# INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
# OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
# OF THE POSSIBILITY OF SUCH DAMAGE.
# *********************************************************************************
function add_fuel_burn_constraints(m,p)
  	@constraint(m, [t in p.techs.gen, ts in p.time_steps],
		m[:dvFuelUsage][t, ts] == (p.s.generator.fuel_slope_gal_per_kwh * KWH_PER_GAL_DIESEL *
		p.production_factor[t, ts] * p.hours_per_time_step * m[:dvRatedProduction][t, ts]) +
		(p.s.generator.fuel_intercept_gal_per_hr * KWH_PER_GAL_DIESEL * p.hours_per_time_step * m[:binGenIsOnInTS][t, ts])
	)
	@constraint(m,
		sum(m[:dvFuelUsage][t, ts] for t in p.techs.gen, ts in p.time_steps) <=
		p.s.generator.fuel_avail_gal * KWH_PER_GAL_DIESEL
	)
end

function add_binGenIsOnInTS_constraints(m,p)
	# Generator must be on for nonnegative output
	@constraint(m, [t in p.techs.gen, ts in p.time_steps],
		m[:dvRatedProduction][t, ts] <= p.s.generator.max_kw * m[:binGenIsOnInTS][t, ts]
	)
	# Note: min_turn_down_pct is only enforced when off_grid_flag is true and in p.time_steps_with_grid, but not for grid outages for on-grid analyses
	if p.s.settings.off_grid_flag
		@constraint(m, [t in p.techs.gen, ts in p.time_steps_without_grid],
			p.s.generator.min_turn_down_pct * m[:dvSize][t] - m[:dvRatedProduction][t, ts] <=
			p.s.generator.max_kw * (1 - m[:binGenIsOnInTS][t, ts])
		)
	else
		@constraint(m, [t in p.techs.gen, ts in p.time_steps_with_grid],
			p.s.generator.min_turn_down_pct * m[:dvSize][t] - m[:dvRatedProduction][t, ts] <=
			p.s.generator.max_kw * (1 - m[:binGenIsOnInTS][t, ts])
		)
	end

end

function add_gen_can_run_constraints(m,p)
	if p.s.generator.only_runs_during_grid_outage
		for ts in p.time_steps_with_grid, t in p.techs.gen
			fix(m[:dvRatedProduction][t, ts], 0.0, force=true)
		end
	end

	if !(p.s.generator.sells_energy_back_to_grid)
		for t in p.techs.gen, u in p.export_bins_by_tech[t], ts in p.time_steps
			fix(m[:dvProductionToGrid][t, u, ts], 0.0, force=true)
		end
	end
end

function add_gen_rated_prod_constraint(m, p)
	@constraint(m, [t in p.techs.gen, ts in p.time_steps],
		m[:dvSize][t] >= m[:dvRatedProduction][t, ts]
	)
end

"""
    add_gen_constraints(m, p)

Add Generator operational constraints and cost expressions.
"""
function add_gen_constraints(m, p)
    add_fuel_burn_constraints(m,p)
    add_binGenIsOnInTS_constraints(m,p)
    add_gen_can_run_constraints(m,p)
    add_gen_rated_prod_constraint(m,p)

    m[:TotalGenPerUnitProdOMCosts] = @expression(m, p.third_party_factor * p.pwf_om *
        sum(p.s.generator.om_cost_per_kwh * p.hours_per_time_step *
        m[:dvRatedProduction][t, ts] for t in p.techs.gen, ts in p.time_steps)
    )
    m[:TotalGenFuelCosts] = @expression(m, p.pwf_fuel["Generator"] *
        sum(m[:dvFuelUsage][t,ts] * p.s.generator.fuel_cost_per_gallon / KWH_PER_GAL_DIESEL for t in p.techs.gen, ts in p.time_steps)
    )
end

max_stars_repo_name: NREL/REopt

max_issues_repo_name: NREL/REopt

max_forks_repo_name: NREL/REopt


text:
struct TestItem
    field::Int
end

function get_field(item::TestItem)
    return item.field
end

@testset "Test LazyDictFromIterator" begin
    first = [TestItem(x) for x in 1:5]
    second = [TestItem(x) for x in 6:10]
    items = [first, second]
    iter = Iterators.flatten(items)

    container = IS.LazyDictFromIterator(Int, TestItem, iter, get_field)

    # Run through twice because the items must persist in the dict.
    for i in range(1, length=2)
        for x in 1:10
            @test get(container, x) isa TestItem
        end
    end

    # Add an item to an underlying array and ensure it is found.
    push!(first, TestItem(22))
    IS.reset_iterator(container)
    @test get(container, 22) isa TestItem

    @test isnothing(get(container, 25))

    third = [TestItem(x) for x in 11:15]
    items = [third]
    @test isnothing(get(container, 12))
    IS.replace_iterator(container, Iterators.flatten(items))
    @test get(container, 12) isa TestItem
    # The old values should still be there.
    @test get(container, 5) isa TestItem
end


text:
function execute_order!(acc::Account, o::OpenOrder, ba::BidAsk)
    # create new position
    pos = Position(
        o.inst,
        sign(o.dir) * abs(o.size),
        o.dir,
        ba,
        ba.dt,
        open_price(o.dir, ba),
        ba,
        ba.dt,
        close_price(o.dir, ba),
        o.stop_loss,
        o.take_profit,
        Unspecified::CloseReason,
        0.0,        # PnL
        o.data      # user-defined data object
    )
    book_position!(acc, pos, ba)
    return
end

function execute_order!(acc::Account, o::CloseOrder, ba::BidAsk)
    close_position!(acc, o.pos, ba, o.close_reason)
    return
end

function execute_order!(acc::Account, o::CloseAllOrder, ba::BidAsk)
    for pos in acc.open_positions
        close_position!(acc, pos, ba, Unspecified::CloseReason)
    end
    return
end

function book_position!(acc::Account, pos::Position, ba::BidAsk)
    # update P&L of position and account equity
    update_pnl!(acc, pos, ba)

    # add position to portfolio
    push!(acc.open_positions, pos)
    return
end

function update_pnl!(acc::Account, pos::Position, ba::BidAsk)
    # temporarily remove from account equity
    acc.equity -= pos.pnl

    # update P&L of position
    update_pnl!(pos, ba)

    # add updated value to account equity
    acc.equity += pos.pnl
    return
end

function close_position!(acc::Account, pos::Position, ba::BidAsk, close_reason::CloseReason)
    # find in vector of open positions
    idx = findfirst(x -> x === pos, acc.open_positions)
    if isnothing(idx)
        println("WARNING: Position already closed!\n", pos)
        return
    end

    pos.close_reason = close_reason

    # update P&L of position and account equity
    update_pnl!(acc, pos, ba)

    # remove from open positions vector
    deleteat!(acc.open_positions, idx)

    # add to closed positions vector
    push!(acc.closed_positions, pos)

    # update account balance
    acc.balance += pos.pnl
    return
end

function update_account!(acc::Account, inst::Instrument, ba::BidAsk)
    # # close orders for given instrument
    # if length(acc.close_orders) > 0
    #     inst_orders = get(acc.close_orders, inst, nothing)
    #     if !isnothing(inst_orders)
    #         for o in inst_orders
    #             execute_order!(acc, o, nbbo)
    #             acc.close_orders_count -= 1
    #         end
    #         empty!(inst_orders)
    #     end
    # end

    # # open orders for given instrument
    # if length(acc.place_orders) > 0
    #     inst_orders = get(acc.place_orders, inst, nothing)
    #     if !isnothing(inst_orders)
    #         for o in inst_orders
    #             execute_order!(acc, o, nbbo)
    #             acc.place_orders_count -= 1
    #         end
    #         empty!(inst_orders)
    #     end
    # end

    # value open positions for given instrument
    for pos in acc.open_positions
        if pos.inst === inst
            # update P&L of position and account equity
            update_pnl!(acc, pos, ba)
        end
    end

    return
end


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

mutable struct ResourceGroupProperties <: SwaggerModel
    provisioningState::Any # spec type: Union{ Nothing, String } # spec name: provisioningState

    function ResourceGroupProperties(;provisioningState=nothing)
        o = new()
        validate_property(ResourceGroupProperties, Symbol("provisioningState"), provisioningState)
        setfield!(o, Symbol("provisioningState"), provisioningState)
        o
    end
end # type ResourceGroupProperties

const _property_map_ResourceGroupProperties = Dict{Symbol,Symbol}(Symbol("provisioningState")=>Symbol("provisioningState"))
const _property_types_ResourceGroupProperties = Dict{Symbol,String}(Symbol("provisioningState")=>"String")
Base.propertynames(::Type{ ResourceGroupProperties }) = collect(keys(_property_map_ResourceGroupProperties))
Swagger.property_type(::Type{ ResourceGroupProperties }, name::Symbol) = Union{Nothing,eval(Meta.parse(_property_types_ResourceGroupProperties[name]))}
Swagger.field_name(::Type{ ResourceGroupProperties }, property_name::Symbol) =  _property_map_ResourceGroupProperties[property_name]

function check_required(o::ResourceGroupProperties)
    true
end

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


text:
# A module that provides convolution, matrix multiplication, DFT, and IDFT functions
module JuliaDSP
include("DSP_Functions.jl")
end


text:
# module Swot

using Sads
using Distributions
using NCDatasets

"Load hydraulic variables from Pepsi-1 NetCDF file."
function load_data(ncfile)
    ds = Dataset(ncfile)
    xs = NCDatasets.group(ds, "XS_Timeseries")
    ri = NCDatasets.group(ds, "River_Info")
    rbnd = ri["rch_bnd"][:]
    qwbm = ri["QWBM"][1]
    x = xs["X"][:, 1]
    ri = [findmin(sqrt.((r .- x).^2))[2] for r in rbnd]
    x = x[end] .- x; x = x[end:-1:1]
    ri = length(x) .- ri .+ 1; ri = ri[end:-1:1]
    W = xs["W"][:]; W = W[end:-1:1, :]
    H = xs["H"][:]; H = H[end:-1:1, :]
    Q = xs["Q"][:]; Q = Q[end:-1:1, :]
    close(ds)
    Q, H, W, x, qwbm, ri
end

"Read data from SWOT NetCDF file."
function read_data(ncfile)
    ds = Dataset(ncfile)
    x = ds["XS_90m"][:] .* 90.0 .- 90
    xs = NCDatasets.group(ds, "XS_Timseries")
    W = xs["W"][:]
    H = xs["H_1km"][:]
    Q = xs["Q"][:]
    S = abs.(NCDatasets.group(ds, "Reach_Timeseries")["S_90m"][:]) ./ length(x)
    close(ds)
    j = findall((sum(isnan.(H), dims=1)[1, :] .== 0) .& (all(W .> 0, dims=1)[1, :]))
    # H = sort(H[:, j], dims=1)
    W = W[:, j]; W = W[end:-1:1, :]
    Hmin = minimum(H[:, j], dims=1)
    H = zeros(size(W))
    H[1, :] = Hmin
    for c in 2:size(H, 1)
        for t in 1:size(H, 2)
            H[c, t] = H[c-1, t] + S[t] .* (x[c] - x[c-1])
        end
    end
    Q = Q[:, j]; Q = Q[end:-1:1, :]
    qwbm = mean(abs.(Q))
    ri = [1; length(x)]
    Q, H, W, x, qwbm, ri
end

"Write output data to NetCDF file."
function write_data(ncfile, H, W, S, Qa, A0, n)
    nr, nt = size(Qa)
    ds = Dataset(ncfile, "c")
    defDim(ds, "nx", nr)
    defDim(ds, "nt", nt)
    Hvar = defVar(ds, "H", Float32, ("nx", "nt"))
    Hvar.attrib["units"] = "m"
    Hvar.attrib["comments"] = "Water Surface Elevation"
    Wvar = defVar(ds, "W", Float32, ("nx", "nt"))
    Wvar.attrib["units"] = "m"
    Wvar.attrib["comments"] = "Width"
    Svar = defVar(ds, "S", Float32, ("nx", "nt"))
    Svar.attrib["units"] = "m/m"
    Svar.attrib["comments"] = "Water Surface Slope"
    Qvar = defVar(ds, "Q", Float32, ("nx", "nt"))
    Qvar.attrib["units"] = "m3/s"
    Qvar.attrib["comments"] = "Discharge"
    nvar = defVar(ds, "n", Float32, ("nx",))
    nvar.attrib["units"] = "unitless"
    nvar.attrib["comments"] = "Manning's roughness"
    A0var = defVar(ds, "Abase", Float32, ("nx",))
    Hvar[:, :] = H
    Wvar[:, :] = W
    Svar[:, :] = S
    Qvar[:, :] = Qa
    A0var[:] = A0
    nvar[:] = n
    close(ds)
end

"Calculate Nash-Sutcliffe efficiency."
function nse(o, m)
    1.0 .- sum((m .- o).^2) ./ sum((o .- mean(o)).^2)
end

"Calculate relative root mean squared error."
function rrmse(o, m)
    sqrt(mean(((o .- m) ./ o).^2))
end

"Calculate normalized root mean squared error."
function nrmse(o, m)
    sqrt(mean(((o .- m) ./ mean(o)).^2))
end

"Compare theoretical with estimated AHG coefficients."
function compare_ahg(Qa, H, W, x, qwbm, rₚ, nₚ=Uniform(0.02, 0.05), nens=100, nsamples=1000)
    est = zeros(length(x), 6)
    the = zeros(length(x), 6)
    wbf, hbf, hmin = maximum(W, dims=2), maximum(H, dims=2), minimum(H[1, :])
    S = diff(H, dims=1) ./ diff(x)
    S = [S[1, :]'; S]
    zm, zs, dm, ds, zb, qb = Sads.priors(qwbm, H, W, x, nₚ, rₚ, nsamples, nens)
    Qₚ = Truncated(LogNormal(log(qwbm/sqrt(1+dm^2)), log(1+dm^2)), qb[1], qb[2])
    zₚ = Normal(zm, 0.1)
    So = Sads.bed_slope(S, H, W, x, hbf, wbf, Qₚ, nₚ, rₚ, zₚ, nens)
    ze = zeros(length(x), size(H,2))
    ze[1,:] .= zm
    for j in 2:length(x) ze[j, :] = ze[j-1, :] .+ So[j] * (x[j] - x[j-1]); end
    ybf = hbf .- ze[:, 1]
    n = mean(rand(nₚ, nsamples))
    r = mean(rand(rₚ, nsamples))
    he = Sads.gvf_ensemble!(mean(H,dims=2), mean(W,dims=2), So, x, maximum(H, dims=2), maximum(W, dims=2), Qa, [n for _ in 1:size(H,2)], [r for _ in 1:size(H,2)], ze)
    i = findall(he[1, :] .> 0)
    h = he .* (r+1)/r .+ ze
    S = diff(h, dims=1) ./ diff(x)
    S = [S[1, :]'; S]
    for xi=1:length(x)
        py = polyfit(log.(Qa[i]), log.(he[xi,i]), 1)
        w = wbf[end] * (he[xi, i] ./ ybf[xi]).^(1/r)
        pw = polyfit(log.(Qa[i]), log.(w), 1)
        sf = [S[xi,t] > 0 ? S[xi,t] : minimum(S[xi,S[xi,:] .> 0]) for t in 1:size(S,2)]
        U = (1/n) .* he[end,i].^(2/3) .* sf[i].^(1/2)
        pu = polyfit(log.(Qa[i]), log.(U), 1)
        a, c, k, b, f, m = Sads.estimate_ahg(r, n, wbf[xi], maximum(he[xi,i]), mean(sf))
        est[xi, :] = [exp(pw[0]), exp(py[0]), exp(pu[0]), pw[1], py[1], pu[1]]
        the[xi, :] = [a, c, k, b, f, m]
    end
    est, the
end

"Main driver routine."
function main()
    ncpath = ARGS[1]
    rₚ = eval(Symbol(ARGS[2]))(parse(Float64, ARGS[3]), parse(Float64, ARGS[4]))
    ncfiles = joinpath.(abspath(ncpath), readdir(ncpath))
    nt = maximum([length(Dataset(ncfile)["Time steps"]) for ncfile in ncfiles])
    Qa = zeros(length(ncfiles), nt)
    A0 = zeros(length(ncfiles))
    n = zeros(length(ncfiles))
    H = zeros(length(ncfiles), nt)
    W = zeros(length(ncfiles), nt)
    S = zeros(length(ncfiles), nt)
    for (i, ncfile) in enumerate(ncfiles)
        Q, h, w, x, qwbm, ri = read_data(ncfile)
        qa, a0, mann = Sads.estimate(qwbm, h, w, x, rₚ, ri)
        H[i, 1:size(h, 2)] = mean(h, dims=1)
        W[i, 1:size(w, 2)] = mean(w, dims=1)
        S[i, 1:size(w, 2)] = mean(diff(h, dims=1) ./ diff(x), dims=1)
        Qa[i, 1:size(h, 2)], A0[i], n[i] = mean(qa, dims=1), mean(a0), mean(mann)
    end
    write_data("$(ncpath)/sads_output.nc", H, W, S, Qa, A0, n)
end

main()

# Base.@ccallable function julia_main(ARGS::Vector{String})::Cint
#     main(ARGS)
#     return 0
# end

# end


text:
###############################################################################
#
# Stomatal model scheme types
#
###############################################################################
#= EmpiricalStomatalModel type tree
AbstractStomatalModel
---> EmpiricalStomatalModel
    ---> ESMBallBerry    # use RH
    ---> ESMGentine      # use k_leaf
    ---> ESMLeuning      # Use VPD
    ---> ESMMedlyn       # use VPD
---> OptimizationStomatalModel
    ---> OSMEller        # Eller 2018 Model
    ---> OSMSperry       # Sperry 2017 model
    ---> OSMWang         # Wang 2020 model
    ---> OSMWAP          # Wolf-Anderegg-Pacala model
    ---> OSMWAP          # Modified Wolf-Anderegg-Pacala model
=#
"""
    type AbstractStomatalModel

Hierarchy of the `AbstractStomatalModel`:
- [`EmpiricalStomatalModel`](@ref)
- [`OptimizationStomatalModel`](@ref)
"""
abstract type AbstractStomatalModel{FT<:AbstractFloat} end

"""
    type EmpiricalStomatalModel

Hierarchy of the `EmpiricalStomatalModel`:
- [`ESMBallBerry`](@ref)
- [`ESMGentine`](@ref)
- [`ESMLeuning`](@ref)
- [`ESMMedlyn`](@ref)
"""
abstract type EmpiricalStomatalModel{FT} <: AbstractStomatalModel{FT} end

"""
    struct ESMBallBerry{FT}

An empirical model parameter set type for Ball-Berry type model.
The equation used for Ball-Berry type model is
```math
gs = g0 + g1 ⋅ RH ⋅ \\dfrac{A}{Cs}
```

# Fields
$(TYPEDFIELDS)
"""
Base.@kwdef mutable struct ESMBallBerry{FT} <: EmpiricalStomatalModel{FT}
    "minimal stomatal conductance g0 `[mol m⁻² s⁻¹]`"
    g0::FT = FT(0.025)
    "slope of conductance-photosynthesis correlation `[unitless]`"
    g1::FT = FT(9.0  )
end

"""
    struct ESMGentine{FT}

An empirical model parameter set type for Gentine type model.
The equation used for Gentine type model is
```math
gs = g0 + g1 ⋅ \\dfrac{k_{leaf}}{k_{max}} ⋅ \\dfrac{A}{Ca}.
```
Note it that the Gentine model does not require for a `β` function to tune the
    soil drought response, but the use of `k_leaf` also does not permit
    post-drought stomatal response unless `k_leaf` can be recovered.

# Fields
$(TYPEDFIELDS)
"""
Base.@kwdef mutable struct ESMGentine{FT} <: EmpiricalStomatalModel{FT}
    "minimal stomatal conductance g0 `[mol m⁻² s⁻¹]`"
    g0::FT = FT(0.025)
    "slope of conductance-photosynthesis correlation `[unitless]`"
    g1::FT = FT(9.0  )
end

"""
    struct ESMLeuning{FT}

An empirical model parameter set type for Leuning type model.
The equation used for Leuning type model is
```math
gs = g0 + g1 ⋅ \\dfrac{A}{Cs - Γ^{*}} ⋅ \\dfrac{1}{1 + \\dfrac{VPD}{d0}}
```

# Fields
$(TYPEDFIELDS)
"""
Base.@kwdef mutable struct ESMLeuning{FT} <: EmpiricalStomatalModel{FT}
    "minimal stomatal conductance g0 `[mol m⁻² s⁻¹]`"
    g0::FT = FT(0.025 )
    "slope of conductance-photosynthesis correlation `[unitless]`"
    g1::FT = FT(8.0   )
    "fitting parameter of d/d0 below the fraction, same unit as vpd `[Pa]`"
    d0::FT = FT(3000.0)
end

"""
    struct ESMMedlyn{FT}

An empirical model parameter set type for Medlyn type model.
The equation used in Medlyn type model is
```math
gs = g0 + 1.6 ⋅ \\left( 1 + \\dfrac{g1}{\\sqrt{VPD}} \\right) ⋅ \\dfrac{A}{Ca}
```

# Fields
$(TYPEDFIELDS)
"""
Base.@kwdef mutable struct ESMMedlyn{FT} <: EmpiricalStomatalModel{FT}
    "minimal stomatal conductance g0 `[mol m⁻² s⁻¹]`"
    g0::FT = FT(0.025)
    "slope of conductance-photosynthesis correlation `[Pa⁽⁵⁾]`"
    g1::FT = FT(125.0)
end

"""
    type OptimizationStomatalModel

Hierarchy of the `OptimizationStomatalModel`:
- [`OSMEller`](@ref)
- [`OSMSperry`](@ref)
- [`OSMWang`](@ref)
- [`OSMWAP`](@ref)
- [`OSMWAPMod`](@ref)
"""
abstract type OptimizationStomatalModel{FT} <: AbstractStomatalModel{FT} end

"""
    struct OSMEller

An optimization model parameter set type for Eller model.
The equation used for Eller model is
```math
\\dfrac{∂Θ}{∂E} = -\\dfrac{∂K}{∂E} ⋅ \\dfrac{A}{K}
```
where K is ``\\dfrac{∂E}{∂P}``.

# Fields
$(TYPEDFIELDS)
"""
struct OSMEller{FT} <: OptimizationStomatalModel{FT} end

"""
    struct OSMSperry

An optimization model parameter set type for Sperry model.
The equation used for Sperry model is
```math
\\dfrac{∂Θ}{∂E} = -\\dfrac{∂K}{∂E} ⋅ \\dfrac{A_{max}}{K_{max}}
```
where K is ``\\dfrac{∂E}{∂P}``.

# Fields
$(TYPEDFIELDS)
"""
struct OSMSperry{FT} <: OptimizationStomatalModel{FT} end

"""
    struct OSMWang

An optimization model parameter set type for Eller type model.
The equation used for Wang model is
```math
\\dfrac{∂Θ}{∂E} = \\dfrac{A}{E_{crit} - E}
```

# Fields
$(TYPEDFIELDS)
"""
struct OSMWang{FT} <: OptimizationStomatalModel{FT} end

"""
    struct OSMWAP{FT}

An optimization model parameter set type for Wolf-Anderegg-Pacala type model.
The equation used for Wolf-Anderegg-Pacala model is
```math
\\dfrac{∂Θ}{∂E} = \\dfrac{2aP + b}{K}
```
where K is ∂P/∂E.

# Fields
$(TYPEDFIELDS)
"""
Base.@kwdef mutable struct OSMWAP{FT} <: OptimizationStomatalModel{FT}
    "Quadratic equation parameter `[μmol m⁻² s⁻¹ MPa⁻²]`"
    a::FT = FT(0.5)
    "Quadratic equation parameter `[μmol m⁻² s⁻¹ MPa⁻¹]`"
    b::FT = FT(2.0)
end

"""
    struct OSMWAPMod{FT}

An optimization model parameter set type for Wolf-Anderegg-Pacala type model,
    modified by adding a photosynthesis component while set b and c = 0.
The equation used for modified Wolf-Anderegg-Pacala model is
```math
\\dfrac{∂Θ}{∂E} = \\dfrac{aAP}{K}
```
where P is absolute value of leaf xylem pressure.

# Fields
$(TYPEDFIELDS)
"""
Base.@kwdef mutable struct OSMWAPMod{FT} <: OptimizationStomatalModel{FT}
    "Quadratic equation parameter `[mol mol⁻¹ MPa⁻¹]`"
    a::FT = FT(0.1)
end


text:
#---------------------------------------------------------
# # [River discharge](@id river-discharge)
#---------------------------------------------------------

#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/notebooks/5_river_discharge.ipynb)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/notebooks/5_river_discharge.ipynb)

# In this tutorial we will simulate a fictitious radioactive tracer that is injected into the ocean by the 200 largest rivers (by estimated discharge).
# The 200 major rivers dataset from [*Dai and Trenberth* (2002)](https://rda.ucar.edu/datasets/ds551.0/) is available from within the AIBECS.
# Once "born", our ficitious tracer decays with a parameter timescale $\tau$ as it flows through ocean basins.
#
# The 3D tracer equation is:
#
# $$\left[\frac{\partial}{\partial t} + \nabla \cdot (\boldsymbol{u} + \mathbf{K}\nabla)\right] x = s_\mathsf{rivers} - x / \tau$$
#
# where $\nabla \cdot \left[ \boldsymbol{u} - \mathbf{K} \nabla \right]$ represents the ocean circulation transport.
# ([Tracer transport operators are described in the documentation](@ref tracer-transport-operators).)
# The riverine source of the tracer is $s_\mathsf{rivers}$, and $x / \tau$ is the decay rate.

# In AIBECS, we must recast this equation in the generic form
#
# $$\left[\frac{\partial}{\partial t} + \mathbf{T}(\boldsymbol{p})\right] \boldsymbol{x} = \boldsymbol{G}(\boldsymbol{x},\boldsymbol{p}).$$

# We start by telling Julia that we want to use the AIBECS and the OCIM2 transport matrix for the ocean circulation.

using AIBECS
grd, T_OCIM2 = OCIM2.load()

# The transport is

T_radiorivers(p) = T_OCIM2

# For the radioactive decay, we simply use

function decay(x, p)
    @unpack τ = p
    return x / τ
end

# To build the river sources, we will load the geographic locations and discharge (in m³ s⁻¹) from the [*Dai and Trenberth* (2017) dataset](https://rda.ucar.edu/datasets/ds551.0/index.html#!description).

RIVERS = Rivers.load()

# This is an array of rivers, for which the type `River{T}` contains the river's name, lat–lon coordinates, and discharge in m³ s⁻¹.
# For example, the first river is the Amazon

r = RIVERS[1]

# We can check the locations with

using Plots
scatter([r.lon for r in RIVERS], [r.lat for r in RIVERS],
        zcolor=ustrip.([r.VFR for r in RIVERS] / u"m^3/s"),
        colorbartitle="log₁₀(discharge / (1 m³ s⁻¹))")

# We can regrid these into the OCIM2 grid and return the corresponding vector with

rivers = regrid(RIVERS, grd)

# (Note this regridding uses [NearestNeighbors.jl](https://github.com/KristofferC/NearestNeighbors.jl) to assign a wet box as the mouth of each river, which sometimes is not exactly the real location of the river mouth.)

# We control the global magnitude of the river discharge, $\sigma$ (in mol s⁻¹), by making it a parameter of our model.
# For that, we separate the river source
#
# $$s_\mathsf{rivers} = \sigma s_0$$
#
# into global magnitude ($\sigma$) and spatial pattern ($s_0$).
#
# Since $\int s_0 \mathrm{d}V = 1$, `s_0` can be computed by normalizing `rivers`.
# In Julia/AIBECS, this can be done by dividing `rivers` by the dot product `v ⋅ rivers` (or `v'rivers` in matrix form).
# (`v ⋅ x` is the discrete equivalent of the volume integral $\int x \mathrm{d}V$.)

v = vector_of_volumes(grd)
s_0 = rivers / (v'rivers)
function s_rivers(p)
    @unpack σ = p
    return σ * ustrip.(s_0) # we must remove the units in AIBECS here :(
end

# We then write the generic $\boldsymbol{G}$ function, which is

G_radiorivers(x,p) = s_rivers(p) - decay(x,p)

# ##### Parameters

# We specify some initial values for the parameters and also include units.

import AIBECS: @units, units
import AIBECS: @initial_value, initial_value
@initial_value @units struct RadioRiversParameters{U} <: AbstractParameters{U}
    τ::U | 5.0 | u"yr"
    σ::U | 1.0 | u"Gmol/yr"
end

# Finally, thanks to the initial values we provided, we can instantiate the parameter vector succinctly as

p = RadioRiversParameters()

# We generate the state function `F` and its Jacobian `∇ₓF`,

F, ∇ₓF = state_function_and_Jacobian(T_radiorivers, G_radiorivers)

# generate the steady-state problem `prob`,

nb = sum(iswet(grd))
x = ones(nb) # initial guess
prob = SteadyStateProblem(F, ∇ₓF, x, p)

# and solve it

s = solve(prob, CTKAlg()).u * u"mol/m^3"

# Let's now run some visualizations using the plot recipes.
# Taking a horizontal slice of the 3D field at 200m gives

cmap = :viridis
plothorizontalslice(s, grd, zunit=u"μmol/m^3", depth=200, color=cmap, clim=(0,0.5))

# and at 500m:

plothorizontalslice(s, grd, zunit=u"μmol/m^3", depth=500, color=cmap, clim=(0,0.05))

# Or we can change the timescale and watch the tracer fill the oceans:

p = RadioRiversParameters(τ = 50.0u"yr")
prob = SteadyStateProblem(F, ∇ₓF, x, p)
s_τ50 = solve(prob, CTKAlg()).u * u"mol/m^3"
plothorizontalslice(s_τ50, grd, zunit=u"μmol/m^3", depth=500, color=cmap, clim=(0,1))


text:
#=
  Copyright (c) 2021
  https://github.com/Ismael-VC/Varvara.jl/blob/main/CONTRIBUTORS.md

  Permission to use, copy, modify, and distribute this software for any
  purpose with or without fee is hereby granted, provided that the above
  copyright notice and this permission notice appear in all copies.

  THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
  WITH REGARD TO THIS SOFTWARE.
=#

module VarvaraOS

using Printf: @sprintf

import Match
import Base: show

using Match: @match

using ..VarvaraUtils: bool
using ..Uxn: CPU, Memory, Stack, Device, PAGE_PROGRAM, peek, poke, uxn_eval

#! format: off

export load, system_talk, console_talk, file_talk, datetime_talk, inspect,
 uxn_boot, run!, FILDES, UXN_EXCEPTIONS, UXN_ERRORS, console_input,
 UxnUnderflowError, UxnOverflowError, UxnZeroDivisionError

 #! format: on

#= Generics =#
const FILDES = Dict(0 => stdin, 1 => stdout, 2 => stderr)

const UXN_EXCEPTIONS = (
  :UxnInputError,
  :UxnLoadError,
  :UxnInitError,
  :UnderflowErro,
  :UxnUnderflowError,
  :UxnOverflowError,
  :UxnZeroDivisionError,
)

for e ∈ UXN_EXCEPTIONS
  @eval struct $e
    msg::String
  end
end

const UXN_ERRORS = [UxnUnderflowError, UxnOverflowError, UxnZeroDivisionError]

#= Core =#
function console_input(c::CPU, ch::Char)::Int
  dev_console = c.dev[1]
  dev_console.dat[0x2] = ch
  return uxn_eval(c, dev_console.vec)
end

function run!(c::CPU)::Nothing
  try
    dev_system_data = c.dev[0].dat
    dev_console_data = c.dev[1].dat
    while !bool(dev_system_data[0xf]) && read(dev_console_data[0x2]) > 0
      vec = peek(dev_console_data, 0)
      !bool(vec) && (vec = c.ram.ptr)  # continue after last BRK
      uxn_eval(c, vec)
    end
  catch e
    if e isa InterruptException
      @info "bye!"
      return
    else
      throw(e)
    end
  end
end

function inspect(s::Stack, name::AbstractString)::Nothing
  @info name
  head = ""
  for y ∈ 0:3, x ∈ 0:7
    p = y * 8 + x
    sp = s.dat[p]
    head *= p == s.ptr ? @sprintf("[%02x] ", sp) : @sprintf("%02x ", sp)
  end
  @info head
end

function load(c::CPU, filepath::AbstractString)::Int
  begin
    rom = read(filepath)
    used = PAGE_PROGRAM + length(rom) - 1
    free = 65280  # 0x10000 - PAGE_PROGRAM
    @assert free - used > 0x10000 - PAGE_ZERO
    c.ram.dat[PAGE_PROGRAM:PAGE_PROGRAM+length(rom)-1] = rom
    @info "Loaded $filepath"

    return 1
  end

  @error("Load: Failed")
  return 0
end

function uxn_boot(c)::Int
  loaded = 0

  #! format: off
  #= system   =# Device(c, 0x0, system_talk)
  #= console  =# Device(c, 0x1, console_talk)
  #= empty    =# Device(c, 0x2)
  #= empty    =# Device(c, 0x3)
  #= empty    =# Device(c, 0x4)
  #= empty    =# Device(c, 0x5)
  #= empty    =# Device(c, 0x6)
  #= empty    =# Device(c, 0x7)
  #= empty    =# Device(c, 0x8)
  #= empty    =# Device(c, 0x9)
  #= file     =# Device(c, 0xa, file_talk)
  #= datetime =# Device(c, 0xb, datetime_talk)
  #= empty    =# Device(c, 0xc)
  #= empty    =# Device(c, 0xd)
  #= empty    =# Device(c, 0xe)
  #= empty    =# Device(c, 0xf)
  #! format: on

  for rom ∈ ARGS
    if !bool(loaded)
      loaded += 1
      !bool(load(c, rom)) && (@error("Load: Failed"); return 0)
      !bool(uxn_eval(c, PAGE_PROGRAM)) && (@error("Init: Failed"); return 0)
    else
      foreach((ch) -> console_input(c, ch), rom)
      console_input(c, '\n')
    end
  end
  bool(loaded) || (@error("Input: Missing"); return 0)

  return 1
end

#= Devices =#
function system_talk(d::Device, b0::UInt8, w::UInt8)::Int
  if bool(w)  #= read =#
    @match b0 begin
      0x2 => (d.dat[2] = d.c.wst.ptr)
      0x3 => (d.dat[3] = d.c.rst.ptr)
    end

  else #= write =#
    @match b0 begin
      0x2 => (d.c.wst.ptr = d.dat[2])
      0x3 => (d.c.rst.ptr = d.dat[3])
      0xe => begin
        inspect(d.c.wst, "Working-stack")
        inspect(d.c.rst, "Return-stack")
      end
      0xf => return 0
    end
  end

  return 1
end

function console_talk(d::Device, b0::UInt8, w::UInt8)::Int
  b0 == 0x1 && (d.vector = peek(d.dat, 0x0))
  (bool(w) && b0 > 0x7) && write(FILDES[b0-0x7], Char(d.dat[b0]))

  return 1
end

function file_talk(d::Device, b0::UInt8, w::UInt8)::Int
  read = b0 == 0xd

  if w && (read || b0 == 0xf)
    name::Char = Char(d.mem[peek(d.dat, 0x8)])
    result::UInt16 = 0
    length::UInt16 = peek(d.dat, 0xa)
    offset::Int32 = (Int32(peek(d.dat, 0x4) << 16)) + peek(d.dat, 0x6)
    addr::UInt16 = peek(d.dat, b0 - 1)

    open(name, read ? "r" : (offset ? "a" : "w")) do f::IOStream
      fseek(f, offset)
      result = read ? read(d.mem[addr], 1, length, f) : write(d.mem[addr], 1, length, f)
    end

    poke16(d.dat, 0x2, result)
  end

  return 1
end

struct Ctm
  second::Cint
  minute::Cint
  hour::Cint
  dayofmoth::Cint
  month::Cint
  year::Cint
  dayofweek::Cint
  dayofyear::Cint
  isdst::Cint
end

function Base.show(io::IO, t::Ctm)
  #! format: off
  print(io,
    t.year + 1900, "-",
    lpad(t.month, 2, "0"), "-",
    lpad(t.dayofmoth, 2, "0"), "T",
    lpad(t.hour, 2, "0"), ":",
    lpad(t.minute, 2, "0"), ":",
    lpad(t.second, 2, "0")
  )
  #! format: on
end

function datetime_talk(d::Device, b0::UInt8, w::UInt8)::Int
  result = Ref{Int64}(0)

  localtime = ccall(@static(Sys.iswindows() ? :localtime : (:localtime, "libc.so.6")), Ptr{Ctm}, (Ptr{Int64},), result)

  t = unsafe_load(localtime)
  t.year += 1900

  poke16(d.dat, 0x0, t.year)
  d.dat[0x2] = t.month
  d.dat[0x3] = t.dayofmoth
  d.dat[0x4] = t.hour
  d.dat[0x5] = t.minute
  d.dat[0x6] = t.second
  d.dat[0x7] = t.dayofweek
  poke16(d.dat, 0x08, t.dayofyear)
  d.dat[0xa] = t.isdst

  return true
end

end  # module

if abspath(PROGRAM_FILE) == @__FILE__
  using .VarvaraOS: uxn_boot

  c = CPU()
  uxn_boot(c)
  run!(c)
end


text:
using AnimalBehavior
using Turing
using Test

println("Macros")
println("======")
include("models_test.jl")

println("Simulation")
println("==========")
include("simul_test.jl")

println("Inference")
println("=========")
include("sample_test.jl")

max_stars_repo_name: sqwayer/AnimalBehavior

max_issues_repo_name: sqwayer/AnimalBehavior


text:
import Requests

DB_DIR = joinpath(dirname(@__FILE__), "db")

# clear out old database directory
if isdir(DB_DIR)
    rm(DB_DIR; recursive=true)
end
mkdir(DB_DIR)

"""Data source for data.gov.au"""
function get_australia(id)
    api_root = "https://data.gov.au/api/3/action"
    all_regions = Set{String}(["AU-$x" for x in ["ACT", "NSW", "NT", "QLD", "SA", "TAS", "VIC", "WA"]])  # ISO 3166-2:AU

    # get the package to find all resources (years)
    response = try
        Requests.get(joinpath(api_root, "package_show"); query=Dict(:id => id))
    catch e
        error("Unable to retrieve package $id: $e")
    end
    response = Requests.json(response; dicttype=Dict{Symbol, Any})
    if !response[:success]
        error("Unable to retrieve package $id: ", response[:error][:message])
    end
    # process all resources in the package, newest resources first
    SOURCE_DIR = joinpath(DB_DIR, "data.gov.au")
    if isdir(SOURCE_DIR)
        rm(SOURCE_DIR; recursive=true)
    end
    mkdir(SOURCE_DIR)
    # pick the newest resource for each hash
    for resource in unique((x) -> x[:hash], sort(response[:result][:resources]; by=(x) -> x[:created], rev=true))
        response = try
            Requests.get(joinpath(api_root, "datastore_search"); query=Dict(:resource_id => resource[:id]))
        catch e
            error("Unable to retrieve resource $(resource[:id]): $e")
        end
        info("Reading: data.gov.au:$id:$(resource[:id])")
        response = Requests.json(response; dicttype=Dict{Symbol, Any})
        if !response[:success]
            error("Unable to retrieve resource $(resource[:id]): ", response[:error][:message])
        end
        records = response[:result][:records]
        i = 0
        open(joinpath(SOURCE_DIR, resource[:id]), "w") do f
            for record in records
                # remove spaces in key names
                record = Dict{Symbol, Any}(Symbol(replace(string(k), " ", "")) => v for (k, v) in record)
                if record[:Date] == "TBC"
                    # ignore "to be confirmed" holidays
                    continue
                end
                if record[:ApplicableTo] == "NAT"
                    regions = copy(all_regions)
                else
                    regions = intersect(Set{String}(["AU-$x" for x in split(record[:ApplicableTo], "|")]), all_regions)
                end
                date = Date(record[:Date], "yyyymmdd")
                write(f, join([record[:HolidayName], record[:Information], join(regions, '|'), date], '\t'), '\n')
                i += 1
            end
        end
        info("Saved: $i/$(length(records)) records")
    end
end

# get data
get_australia("australian-holidays-machine-readable-dataset")


text:
# TODO: remove initial values for all inputs once IO handling in MTK is in place
"""
    Constant(val; name)

Outputs a constant value `val`.
"""
function Constant(val; name)
    @variables y(t)=val [output=true]
    @parameters val=val
    eqs = [
        y ~ val
    ]
    ODESystem(eqs, t, name=name)
end

"""
    Integrator(; k=1, name)

Outputs `y = ∫k*u dt`, corresponding to the transfer function `1/s`.
"""
function Integrator(; k=1, name)
    @variables x(t)=0 u(t)=0 [input=true] y(t)=0 [output=true]
    @parameters k=k
    eqs = [
        Dₜ(x) ~ k*u
        y ~ x
    ]
    ODESystem(eqs, t, name=name)
end

"""
    Derivative(; k=1, T, name)

Outputs an approximate derivative of the input. The transfer function of this block is
```
k       k
─ - ──────────
T    2 ⎛    1⎞
    T ⋅⎜s + ─⎟
       ⎝    T⎠
```
and a state-space realization is given by `ss(-1/T, 1/T, -k/T, k/T)`
where `T` is the time constant of the filter.
A smaller `T` leads to a more ideal approximation of the derivative.
"""
function Derivative(; k=1, T, name)
    @variables x(t)=0 u(t)=0 [input=true] y(t)=0 [output=true]
    @parameters T=T k=k
    eqs = [
        Dₜ(x) ~ (u - x) / T
        y ~ (k/T)*(u - x)
    ]
    ODESystem(eqs, t, name=name)
end

"""
    FirstOrder(; k=1, T, name)

A first-order filter with a single real pole in `s = -T` and gain `k`. The transfer function
is given by `Y(s)/U(s) = `
```
   k
───────
sT + 1
```
"""
function FirstOrder(; k=1, T, name)
    @variables x(t)=0 u(t)=0 [input=true] y(t) [output=true]
    @parameters T=T k=k
    eqs = [
        Dₜ(x) ~ (-x + k*u) / T
        y ~ x
    ]
    ODESystem(eqs, t, name=name)
end

"""
    SecondOrder(; k=1, w, d, name)

A second-order filter with gain `k`, a bandwidth of `w` rad/s and relative damping `d`. The transfer function
is given by `Y(s)/U(s) = `
```
      k*w^2
─────────────────
s² + 2d*w*s + w^2
```
Critical damping corresponds to `d=1`, which yields the fastest step response without overshoot, d < 1` results in an under-damped filter while `d > 1` results in an over-damped filter.
`d = 1/√2` corresponds to a Butterworth filter of order 2 (maximally flat frequency response).
"""
function SecondOrder(; k=1, w, d, name)
    @variables x(t)=0 xd(t)=0 u(t)=0 [input=true] y(t) [output=true]
    @parameters k=k w=w d=d
    eqs = [
        Dₜ(x) ~ xd
        Dₜ(xd) ~ w*(w*(k*u - x) - 2*d*xd)
        y ~ x
    ]
    ODESystem(eqs, t, name=name)
end

"""
    PID(; k, Ti=false, Td=false, wp=1, wd=1, Ni, Nd=12, y_max=Inf, y_min=-y_max, gains = false, name)

Proportional-Integral-Derivative (PID) controller with output saturation, set-point weighting and integrator anti-windup.

This block has inputs
- `u_r` reference/set point
- `u_y` measurement signal
and output `y` corresponding to the control signal.

The equation for the control signal is roughly
```
k(ep + 1/Ti * ∫e + 1/Td * d/dt(ed))
e = u_r - u_y
ep = wp*u_r - u_y
ed = wd*u_r - u_y
```
where the transfer function for the derivative includes additional filtering, see `? Derivative` for more details.

- `k`: Proportional gain
- `Ti`: Integrator time constant. Set to `false` to turn off integral action.
- `Td`: Derivative time constant. Set to `false` to turn off derivative action.
- `wp`: Set-point weighting in the proportional part.
- `wd`: Set-point weighting in the derivative part.
- `Nd`: Derivative limit, limits the derivative gain to Nd/Td. Reasonable values are ∈ [8, 20]. A higher value gives a better approximation of an ideal derivative at the expense of higher noise amplification.
- `Ni`: `Ni*Ti` controls the time constant `Tₜ` of anti-windup tracking. A common (default) choice is `Tₜ = √(Ti*Td)` which is realized by `Ni = √(Td / Ti)`. Anti-windup can be effectively turned off by setting `Ni = Inf`.
`gains`: If `gains = true`, `Ti` and `Td` will be interpreted as gains with a fundamental PID transfer function on parallel form `ki=Ti, kd=Td, k + ki/s + kd*s`
"""
function PID(; k, Ti=false, Td=false, wp=1, wd=1,
    Ni    = Ti == 0 ? Inf : √(max(Td / Ti, 1e-6)),
    Nd    = 12,
    y_max = Inf,
    y_min = y_max > 0 ? -y_max : -Inf,
    gains = false,
    name
)
    if gains
        Ti = k / Ti
        Td = Td / k
    end
    0 ≤ wp ≤ 1 || throw(ArgumentError("wp out of bounds, got $(wp) but expected wp ∈ [0, 1]"))
    0 ≤ wd ≤ 1 || throw(ArgumentError("wd out of bounds, got $(wd) but expected wd ∈ [0, 1]"))
    Ti ≥ 0     || throw(ArgumentError("Ti out of bounds, got $(Ti) but expected Ti ≥ 0"))
    Td ≥ 0     || throw(ArgumentError("Td out of bounds, got $(Td) but expected Td ≥ 0"))
    y_max ≥ y_min || throw(ArgumentError("y_min must be smaller than y_max"))

    @variables x(t)=0 u_r(t)=0 [input=true] u_y(t)=0 [input=true] y(t) [output=true] e(t)=0 ep(t)=0 ed(t)=0 ea(t)=0

    @named D = Derivative(k = Td, T = Td/Nd) # NOTE: consider T = max(Td/Nd, 100eps()), but currently errors since a symbolic variable appears in a boolean expression in `max`.
    if isequal(Ti, false)
        @named I = Gain(false)
    else
        @named I = Integrator(k = 1/Ti)
    end
    @named sat = Saturation(; y_min, y_max)
    derivative_action = Td > 0
    @parameters k=k Td=Td wp=wp wd=wd Ni=Ni Nd=Nd # TODO: move this line above the subsystem definitions when symbolic default values for parameters works. https://github.com/SciML/ModelingToolkit.jl/issues/1013
    # NOTE: Ti is not included as a parameter since we cannot support setting it to false after this constructor is called. Maybe Integrator can be tested with Ti = false setting k to 0 with IfElse?

    eqs = [
        e ~ u_r - u_y # Control error
        ep ~ wp*u_r - u_y  # Control error for proportional part with setpoint weight
        ea ~ sat.y - sat.u # Actuator error due to saturation
        I.u ~ e + 1/(k*Ni)*ea  # Connect integrator block. The integrator integrates the control error and the anti-wind up tracking. Note the apparent tracking time constant 1/(k*Ni), since k appears after the integration and 1/Ti appears in the integrator block, the final tracking gain will be 1/(Ti*Ni)
        sat.u ~ derivative_action ? k*(ep + I.y + D.y) : k*(ep + I.y) # unsaturated output = P + I + D
        y ~ sat.y
    ]
    systems = [I, sat]
    if derivative_action
        push!(eqs, ed ~ wd*u_r - u_y)
        push!(eqs, D.u ~ ed) # Connect derivative block
        push!(systems, D)
    end
    ODESystem(eqs, t, name=name, systems=systems)
end

"""
    StateSpace(A, B, C, D=0; x0=zeros(size(A,1)), name)

A linear, time-invariant state-space system on the form.
```
ẋ = Ax + Bu
y = Cx + Du
```
Transfer functions can also be simulated by converting them to a StateSpace form.
"""
function StateSpace(A, B, C, D=0; x0=zeros(size(A,1)), name)
    nx = size(A,1)
    nu = size(B,2)
    ny = size(C,1)
    if nx == 0
        length(C) == length(B) == 0 || throw(ArgumentError("Dimension mismatch between A,B,C matrices"))
        return Gain(D; name=name)
    end
    if B isa AbstractVector
        B = reshape(B, length(B), 1)
    end
    if D == 0
        D = zeros(ny, nu)
    end
    @variables x[1:nx](t)=x0 u[1:nu](t)=0 [input=true] y[1:ny](t)=C*x0 [output=true]
    x = collect(x) # https://github.com/JuliaSymbolics/Symbolics.jl/issues/379
    u = collect(u)
    y = collect(y)
    # @parameters A=A B=B C=C D=D # This is buggy
    eqs = [
        Dₜ.(x) .~ A*x .+ B*u
        y      .~ C*x .+ D*u
    ]
    ODESystem(eqs, t, name=name)
end


text:
module adapters_postgresql_public_use_cases

# example for https://discourse.julialang.org/t/generating-insert-into-sql-statement/59106

using Test
using Octo.Adapters.PostgreSQL

Repo.debug_sql()

include("options.jl")

Repo.connect(;
    adapter = Octo.Adapters.PostgreSQL,
    Options.for_postgresql...
)

Repo.execute(Raw("""
CREATE EXTENSION IF NOT EXISTS "uuid-ossp";

CREATE TABLE IF NOT EXISTS public.use_cases
(
    id uuid NOT NULL DEFAULT uuid_generate_v4(),
    name character varying(150) COLLATE pg_catalog."default" NOT NULL,
    description text COLLATE pg_catalog."default",
    input text COLLATE pg_catalog."default",
    output text COLLATE pg_catalog."default",
    business_value text COLLATE pg_catalog."default",
    business_kpi text COLLATE pg_catalog."default",
    accuracy_impact integer,
    workspace_id uuid NOT NULL,
    dataset_id uuid,
    is_favorite boolean NOT NULL DEFAULT false,
    default_f_experiment uuid,
    default_ad_experiment uuid,
    created_at timestamp without time zone NOT NULL DEFAULT timezone('utc'::text, now()),
    created_by uuid NOT NULL,
    updated_at timestamp without time zone,
    updated_by uuid,
    business_objective text COLLATE pg_catalog."default",
    CONSTRAINT use_cases_pkey PRIMARY KEY (id)
)
"""))

struct PublicUseCases
end
Schema.model(PublicUseCases, table_name="public.use_cases")

Repo.insert!(PublicUseCases, [
("ac7e086a-a7d5-4474-8ab4-e6b8c94994c7", "Edited name", "Edited description", "Edited input", "Edited output", "test", "test", 6, "3541ee59-d48e-4aaa-8f38-79554136462d", NULL, true, NULL, NULL, "2021-04-08 13:46:11.549539", "b306ec38-ee19-4af5-9e3a-a36d2c22fe85", "2021-04-08 13:46:11.954205", "b306ec38-ee19-4af5-9e3a-a36d2c22fe85", "test"),
("d1197838-17c7-43c9-80f3-91d745ac32a8", "xyz", NULL, NULL, NULL, NULL, NULL, NULL, "71ecb06c-2555-4f2b-9b99-d1dc9c0876ca", "07b18f7a-35b6-423d-b252-daf26d6e2e1c", false, NULL, NULL, "2021-03-17 13:17:18.743356", "f31e650c-23b7-4886-a0a1-894c5eee9443", NULL, NULL, NULL),
("d726f6de-bd40-4c96-95c9-7bc080dab46d", "2 - with default exps", "something else again", "much data", "", "", "", NULL, "4e8fcfb2-4518-4445-afed-180941a0fd68", NULL, false, NULL, NULL, "2021-03-31 12:24:04.757653", "182b080b-8e42-4b85-82f5-e7e87bf8df0f", "2021-04-01 13:19:31.307656", "182b080b-8e42-4b85-82f5-e7e87bf8df0f", ""),
])

result = Repo.query([SELECT * FROM PublicUseCases])
@test length(result) == 3

Repo.delete!(PublicUseCases, (id="ac7e086a-a7d5-4474-8ab4-e6b8c94994c7",))
Repo.delete!(PublicUseCases, (id="d1197838-17c7-43c9-80f3-91d745ac32a8",))
Repo.delete!(PublicUseCases, (id="d726f6de-bd40-4c96-95c9-7bc080dab46d",))

result = Repo.query([SELECT * FROM PublicUseCases])
@test length(result) == 0

end # module adapters_postgresql_public_use_cases


text:
function movetospiral!(r, n, side)
flag=true
    for i = 1:n
        if ismarker(r)==1
            flag = false
            return flag
        end
    move!(r,side)
    end
end

function task_9(r)
    n=1
    side=Ost
    while ismarker(r)==0
        side = HorizonSide((Int(side)+1)%4)
        flag=movetospiral!(r, c, side)
        side=HorizonSide((Int(side)+1)%4)
        flag=movetospiral!(r, c, side)
        if flag == 1
            break
        end
        n=n+1
    end
end

task_9(r)


text:
using Distributions
using DataFrames
using Test
using RMInsurance
using Dierckx
using CSV

import RMInsurance.zb

### TESTS
curr_dir = @__DIR__

println("start HWBS.jl...")

## All Data sourced from Swiss National Bank (www.snb.ch)
# Read the raw data for market prices.  The raw data only
# constains spot rates for different measurement dates
# (columns :xyyyy_mm, where yyyy is the year and mm the month)
# and the corresponding Durations (column :Duration).
df_chf_spots = CSV.File( curr_dir * "/HWBS_Input_CHF_Spot.csv",
                          header = true) |> DataFrame
# contains various short rates ("SARON", "Overnight",
# "MM_3months") for different measurement dates
df_chf_shorts = CSV.File(curr_dir * "/HWBS_Input_CHF_Short.csv",
                          header = true) |> DataFrame
# parameters for the Hull-White Black-Scholes model
hwpar = CSV.File(curr_dir * "/HWBS_Input_Parameters.csv",
                  header = true) |> DataFrame
# Asset portfolio
assets = CSV.File( curr_dir * "/HWBS_Input_Portfolio.csv",
                    header = true) |> DataFrame

# t₀ refers to the time at which we will base the spot curve
# it is given as a column identifyer for df_chf_spots_sparse
t₀ =  Symbol("2016-09")  #:x2016_09
# We add the 3month money market debt rate (Switzerland - CHF -
# Money market debt register claims of the Swiss Confederation,
# 3-month) as a proxi for the spot rate of duration 0.
# This is used for extrapolation to spot rates with
# duration < 1 year
spot₀ =
  df_chf_shorts[findall(x->x=="SARON",df_chf_shorts[!, :Type])[1], t₀]

df_spot_coarse = DataFrame()
df_spot_coarse[!, :t] = vcat(0.0, df_chf_spots[!, :Duration])
df_spot_coarse[!, :spot] = vcat(spot₀, df_chf_spots[!, t₀])

dict_hwpar =
  Dict{Symbol, Real}( Symbol( hwpar[i, :name]) =>
                      hwpar[i, :value] for i = 1:nrow(hwpar))

function spot_coarse2fine(nʸ, df_spot_coarse)
  δt = 1 / nʸ
  nᵀ = maximum(df_spot_coarse[!, :t]) * nʸ
  func_spot =
    Spline1D( convert(Array, df_spot_coarse[!, :t]),
              convert(Array, df_spot_coarse[!, :spot]))
  # Durations and corresponding spot rates
  DataFrame(t = collect(1:nᵀ) * δt,
            spot = func_spot(collect(1:nᵀ) * δt))
end

hw = HWBS(dict_hwpar, spot_coarse2fine(dict_hwpar[:nʸ],
                                       df_spot_coarse))

asset_id = Array{Symbol}(undef, nrow(assets))
for 𝑟𝑜𝑤 ∈ 1:nrow(assets)
  if assets[𝑟𝑜𝑤,:Type] in ["zb"]
    asset_id[𝑟𝑜𝑤] =
      Symbol(assets[𝑟𝑜𝑤,:Type] *"_" * string(assets[𝑟𝑜𝑤,:Maturity]))
  else
    asset_id[𝑟𝑜𝑤] = Symbol(assets[𝑟𝑜𝑤,:Type])
  end
end

n_scen = 100_000
t₁ = 1
V₀ = zeros(nrow(assets))            ## Value of assets at time 0
V₁ = zeros(nrow(assets), n_scen)    ## Value of assets at time 1
## loss in value from 0 to 1, discounted to time 0
loss = DataFrame()
## random variables for capital market model at time 1
cm₁ = rand_rw(hw, t₁, n_scen)
r₁ = cm₁[1,:]  ## random interest rates at time 1
S₁ = cm₁[2,:]  ## random index values at time 1

discount = zb(hw, 1) ## discount rate from 0 to 1

for 𝑟𝑜𝑤 ∈ 1:nrow(assets)
  if assets[𝑟𝑜𝑤, :Type] == "index"
    V₀[𝑟𝑜𝑤] = assets[𝑟𝑜𝑤, :Nominal]
    V₁[𝑟𝑜𝑤, :] = V₀[𝑟𝑜𝑤] * S₁
  else
    V₀[𝑟𝑜𝑤] =
      zb(hw, assets[𝑟𝑜𝑤, :Maturity]) * assets[ 𝑟𝑜𝑤, :Nominal]
    zb_tmp(r) = zb(hw, r, t₁, assets[𝑟𝑜𝑤, :Maturity])
    V₁[𝑟𝑜𝑤, :] = zb_tmp.(r₁) * assets[𝑟𝑜𝑤, :Nominal]
  end
  loss[!, asset_id[𝑟𝑜𝑤]] = V₀[𝑟𝑜𝑤] .- discount * V₁[𝑟𝑜𝑤,:]
end
loss[!, :total] = sum(loss[!, asset_id[𝑟𝑜𝑤]] for  𝑟𝑜𝑤  ∈ 1:nrow(assets))

es_assets = es(loss[!, :total], 0.99)

println("...end HWBS.jl")


text:
import Pkg

Pkg.add("Polynomials")

using Polynomials
using .Polynoms

function diff(p::Polynomialx, x; ord=2)
    Q = eval_poly(x,p)
    Qn = eval_polyN(x,p,ord)
    return Q-Qn
end

function eval_poly(x, A)
    Q = first(A)
    for a in @view A[2:end]
        Q=Q*x+a
    end
    return Q
end


text:
set_global_client()

@testset "Redis Locks" begin
    start = datetime2unix(now())
    lock_time = 3  # Seconds
    lock_name = "example_lock"

    @async redis_lock(lock_name) do
        sleep(lock_time)  # Lock will exist for 3 seconds
    end

    while !isredislocked(lock_name)
        sleep(0.1)  # Ensure async lock is active before proceeding
    end

    redis_lock(lock_name) do
        println("This message will be delayed by $lock_time seconds!")  # Blocked by first lock
    end

    total_time = datetime2unix(now()) - start
    @test total_time >= lock_time
end

flushall()


text:
# sorting functionality

"""
The main quicksort kernel uses dynamic parallelism. Let's call blocksize `M`. The first part
of the kernel bubble sorts `M` elements with maximal stride between `lo` and `hi`. If the
sublist is <= `M` elements, `stride` = 1 and no recursion happens. Otherwise, we pick
element `lo + M ÷ 2 * stride` as a pivot. This is an efficient choice for random lists and
pre-sorted lists.

Partition is done in stages:
1. For batches of M values, cumsum how many > pivot are left of each index. The comparison
   alternates between < and <= with recursion depth. This makes no difference when there are
   many unique values, but when there are many duplicates, this effectively partitions into
   <, =, and >.
2. Consolidate batches. This runs inside the quicksort kernel.

Sublists (ranges of the list being sorted) are denoted by `lo` and one of `L` and `hi`. `lo`
is an exclusive lower bound, `hi` is an inclusive upperboard, `L` is their difference.
`b_sums` is "batch sums", the number of values in a batch which are >= pivot or > pivot
depending on the relevant `parity`

Originally developed by @xaellison (Alex Ellison).
"""
module Quicksort

export quicksort!

using ..CUDA

# Comparison

@inline function flex_lt(a, b, eq, lt, by)
    a′ = by(a)
    b′ = by(b)
    (eq && a′ == b′) || lt(a′, b′)
end

# Batch partitioning
"""
Performs in-place cumsum using shared memory. Intended for use with indexes
"""
function cumsum!(sums)
    shift = 1

    while shift < length(sums)
        to_add = 0
        @inbounds if threadIdx().x - shift > 0
            to_add = sums[threadIdx().x - shift]
        end

        sync_threads()
        @inbounds if threadIdx().x - shift > 0
            sums[threadIdx().x] += to_add
        end

        sync_threads()
        shift *= 2
    end
end

"""
Partition the region of `values` after index `lo` up to (inclusive) `hi` with
respect to `pivot`. Computes each value's comparison to pivot, performs a cumsum
of those comparisons, and performs one movement using shmem. Comparison is
affected by `parity`. See `flex_lt`. `swap` is an array for exchanging values
and `sums` is an array of Ints used during the merge sort.
Uses block y index to decide which values to operate on.
"""
@inline function batch_partition(values, pivot, swap, sums, lo, hi, parity, lt::F1, by::F2) where {F1,F2}
    sync_threads()
    blockIdx_yz = (blockIdx().z - 1) * gridDim().y + blockIdx().y
    idx0 = lo + (blockIdx_yz - 1) * blockDim().x + threadIdx().x
    val = idx0 <= hi ? values[idx0] : one(eltype(values))
    comparison = flex_lt(pivot, val, parity, lt, by)

    @inbounds if idx0 <= hi
         sums[threadIdx().x] = 1 & comparison
    else
         sums[threadIdx().x] = 1
    end
    sync_threads()

    cumsum!(sums)

    dest_idx = @inbounds comparison ? blockDim().x - sums[end] + sums[threadIdx().x] : threadIdx().x - sums[threadIdx().x]
    @inbounds if idx0 <= hi && dest_idx <= length(swap)
        swap[dest_idx] = val
    end
    sync_threads()

    @inbounds if idx0 <= hi
         values[idx0] = swap[threadIdx().x]
    end
    sync_threads()
end

"""
Each block evaluates `batch_partition` on consecutive regions of length blockDim().x
from `lo` to `hi` of `values`.
"""
function partition_batches_kernel(values::AbstractArray{T}, pivot, lo, hi, parity, lt::F1,
                                  by::F2) where {T,F1,F2}
    sums = @cuDynamicSharedMem(Int, blockDim().x)
    swap = @cuDynamicSharedMem(T, blockDim().x, sizeof(sums))
    batch_partition(values, pivot, swap, sums, lo, hi, parity, lt, by)
    return
end

# Batch consolidation

"""
Finds the index in `array` of the last value <= `pivot` if `parity` = true or the
last value < `pivot` if `parity` = false.
Searches after index `lo` up to (inclusive) index `hi`
"""
function find_partition(array, pivot, lo, hi, parity, lt::F1, by::F2) where {F1,F2}
    low = lo + 1
    high = hi
    @inbounds while low <= high
        mid = (low + high) ÷ 2
        if flex_lt(pivot, array[mid], parity, lt, by)
            high = mid - 1
        else
            low = mid + 1
        end
    end
    return low - 1
end

"""
This assumes the region of `vals` of length `L` starting after `lo`
has been batch partitioned with respect to `pivot`. Further, it assumes that
these batches are of size `blockDim().x`.

Using 1 step per batch, consolidate these partitioned batches such that the
region is fully partitioned. Each step moves at most `blockDim().x` values.

`b_sums`: either shared memory or a global array which serves as scratch space
for storing the partition of each batch.

`parity`: see top docstring

Must only run on 1 SM.
"""
@inline function consolidate_batch_partition(vals::AbstractArray{T}, pivot, lo, L, b_sums,
                                             parity, lt::F1, by::F2) where {T,F1,F2}
    sync_threads()
    @inline N_b() = cld(L , blockDim().x)
    @inline batch(k) = threadIdx().x + k * blockDim().x

    my_iter = 0
    a = 0
    b = 0

    @inbounds for batch_i in 1:N_b()
        if batch_i % blockDim().x == 1
            if batch(my_iter) <= N_b()
                seek_lo = lo + (batch(my_iter) - 1) * blockDim().x
                seek_hi = lo + min(L, batch(my_iter) * blockDim().x)
                b_sums[threadIdx().x] =
                    seek_hi - find_partition(vals, pivot, seek_lo, seek_hi, parity, lt, by)
            end
            my_iter += 1
        end

        function n_eff()
            if batch_i != N_b() || L % blockDim().x == 0
                blockDim().x
            else
                L % blockDim().x
            end
        end

        sync_threads()
        d = b_sums[batch_i - (my_iter - 1) * blockDim().x]
        c = n_eff() - d
        to_move = min(b, c)
        sync_threads()
        swap = if threadIdx().x <= to_move
            vals[lo + a + threadIdx().x]
        else
            zero(eltype(vals))  # unused value
        end
        sync_threads()
        if threadIdx().x <= to_move
            vals[lo + a + threadIdx().x] = vals[lo + a + b + c - to_move + threadIdx().x]
        end
        sync_threads()
        if threadIdx().x <= to_move
            vals[lo + a + b + c - to_move + threadIdx().x] = swap
        end
        sync_threads()
        a += c
        b += d
    end

    sync_threads()
    return lo + a
end

# Sorting
"""
Finds the median of `vals` starting after `lo` and going for `blockDim().x`
elements spaced by `stride`. Performs bitonic sort in shmem, returns middle value.
Faster than bubble sort, but not as flexible. Does not modify `vals`
"""
function bitonic_median(vals :: AbstractArray{T}, swap, lo, L, stride, lt::F1, by::F2) where {T,F1,F2}
    sync_threads()
    bitonic_lt(i1, i2) = @inbounds flex_lt(swap[i1 + 1], swap[i2 + 1], false, lt, by)

    @inbounds swap[threadIdx().x] = vals[lo + threadIdx().x * stride]
    sync_threads()
    old_val = zero(eltype(swap))

    log_blockDim = begin
        out = 0
        k = blockDim().x
        while k > 1
            k = k >> 1
            out += 1
        end
        out
    end

    log_k = 1
    while log_k <= log_blockDim
        k = 1 << log_k
        j = k ÷ 2

        while j > 0
            i = threadIdx().x - 1
            l = xor(i, j)
            to_swap = (i & k) == 0 && bitonic_lt(l, i) || (i & k) != 0 && bitonic_lt(i, l)
            to_swap = to_swap == (i < l)

            if to_swap
                @inbounds old_val = swap[l + 1]
            end
            sync_threads()
            if to_swap
                @inbounds swap[i + 1] = old_val
            end
            sync_threads()
            j = j ÷ 2
        end
        log_k += 1
    end
    sync_threads()
    return @inbounds swap[blockDim().x ÷ 2]
end

"""
Performs bubble sort on `vals` starting after `lo` and going for min(`L`, `blockDim().x`)
elements spaced by `stride`. Good for sampling pivot values as well as short sorts.
"""
@inline function bubble_sort(vals, swap, lo, L, stride, lt::F1, by::F2) where {F1,F2}
    sync_threads()
    L = min(blockDim().x, L)
    @inbounds begin
        if threadIdx().x <= L
            swap[threadIdx().x] = vals[lo + threadIdx().x * stride]
        end
        sync_threads()
        for level in 0:L
            # get left/right neighbor depending on even/odd level
            buddy = threadIdx().x - 1 + 2 * (1 & (threadIdx().x % 2 != level % 2))
            buddy_val = if 1 <= buddy <= L && threadIdx().x <= L
                 swap[buddy]
            else
                zero(eltype(swap)) # unused value
            end
            sync_threads()
            if 1 <= buddy <= L && threadIdx().x <= L
                is_left = threadIdx().x < buddy
                # flex_lt needs to handle equivalence in opposite ways for the
                # two threads in each swap pair. Otherwise, if there are two
                # different values with the same by, one will overwrite the other
                if is_left != flex_lt(swap[threadIdx().x], buddy_val, is_left, lt, by)
                    swap[threadIdx().x] = buddy_val
                end
            end
            sync_threads()
        end
        if threadIdx().x <= L
            vals[lo + threadIdx().x * stride] = swap[threadIdx().x]
        end
    end
    sync_threads()
end

"""
Launch batch partition kernel and sync
"""
@inline function call_batch_partition(vals::AbstractArray{T}, pivot, swap, b_sums, lo, hi,
                                      parity, sync::Val{true}, lt::F1, by::F2) where {T, F1, F2}
    L = hi - lo
    if threadIdx().x == 1
        blocks_y = cld(L, blockDim().x)

        # TODO: add wrappers
        device = Ref{Cint}()
        CUDA.cudaGetDevice(device)
        max_blocks_y = Ref{Cint}()
        CUDA.cudaDeviceGetAttribute(max_blocks_y, CUDA.cudaDevAttrMaxGridDimY, device[])

        blocks_z, blocks_y = fldmod1(blocks_y, max_blocks_y[])

        @cuda(blocks=(1,blocks_y,blocks_z), threads=blockDim().x, dynamic=true,
              shmem=blockDim().x*(sizeof(Int)+sizeof(T)),
              partition_batches_kernel(vals, pivot, lo, hi, parity, lt, by))
        device_synchronize()
    end
end

"""
Partition batches in a loop using a single block
"""
@inline function call_batch_partition(vals::AbstractArray{T}, pivot, swap, b_sums, lo, hi,
                                      parity, sync::Val{false}, lt::F1, by::F2) where {T, F1, F2}
    while lo <= hi
        batch_partition(vals, pivot, swap, b_sums, lo, min(hi, lo + blockDim().x), parity, lt, by)
        lo += blockDim().x
    end
end

"""
Quicksort recursion condition
For a full sort, `partial` is nothing so it shouldn't affect whether recursion
happens.
"""
function partial_range_overlap(lo, hi, partial :: Nothing)
    true
end

"""
Quicksort recursion condition
If the domain to sort `lo` to `hi` overlaps with `partial`, then we should
do recursion on it, and this returns true (if not, then false)
"""
function partial_range_overlap(lo, hi, partial_k)
    return !(lo > last(partial_k) || hi < first(partial_k))
end

"""
Perform quicksort on dimension `dims` of `vals` for the region with `lo` as an exclusive
floor and `hi` as an inclusive ceiling. `parity` is a boolean which says whether to
partition by < or <= with respect to the pivot. `sync_depth` is how many (more) levels of
recursion with `qsort_kernel` can be done before reaching `cudaLimitDevRuntimeSyncDepth`.
From the host, this value must not exceed that limit.

`sync` and enclosed type `S` determine how partition occurs: If `sync` is `true`, the kernel
partitions batches in a child kernel, synchronizes, and then consolidates the batches. The
benefit of this kernel is that it distributes the work of partitioning batches across
multiple SMs. If `sync` is `false`, the kernel partitions without launching any child
kernels, then has recursive `qsort_kernel` children for left and right partitions.
`device_synchronize` is never called from this kernel, so there is no practical limit on
recursion.

To detect the scenario of all values in the region being the same, we have two args:
`prev_pivot` and `stuck`. If two consecutive partitions have the same pivot and both failed
to split the region in two, that means all the values are equal. `stuck` is incremented when
the pivot hasn't changed and partition = `lo` or `hi`. If `stuck` reaches 2, recursion ends.
`stuck` is initialized at -1 because `prev_pivot` must be initialized to some value, and
it's possible that the first pivot will be that value, which could lead to an incorrectly
early end to recursion if we started `stuck` at 0.
"""
function qsort_kernel(vals::AbstractArray{T,N}, lo, hi, parity, sync::Val{S}, sync_depth,
                      prev_pivot, lt::F1, by::F2, ::Val{dims}, partial=nothing, stuck=-1) where {T, N, S, F1, F2, dims}
    b_sums = @cuDynamicSharedMem(Int, blockDim().x)
    swap = @cuDynamicSharedMem(T, blockDim().x, sizeof(b_sums))
    shmem = sizeof(b_sums) + sizeof(swap)
    L = hi - lo

    # extract the dimension that we need to sort (selected from the rest by the block x index)
    slice = if N == 1
        vals
    else
        # dimensions that are not part of the sort; index them using the block index
        otherdims = ntuple(i -> i == dims ? 1 : size(vals, i), N)
        other = CartesianIndices(otherdims)[blockIdx().x]

        # create a view that keeps the sorting dimension but indexes across the others
        slicedims = map(Base.Slice, axes(vals))
        idxs = ntuple(i->i==dims ? slicedims[i] : other[i], N)
        view(vals, idxs...)
    end

    # step 1: single block sort. It'll either finish sorting a subproblem or
    # help select a pivot value

    if L <= blockDim().x
        bubble_sort(slice, swap, lo, L, 1, lt, by)
        return
    end

    pivot = bitonic_median(slice, swap, lo, L, L ÷ blockDim().x, lt, by)

    # step 2: use pivot to partition into batches
    call_batch_partition(slice, pivot, swap, b_sums, lo, hi, parity, sync, lt, by)

    # step 3: consolidate the partitioned batches so that the sublist from [lo, hi) is
    #         partitioned, and the partition is stored in `partition`. Dispatching on P
    #         cleaner and faster than an if statement

    partition = consolidate_batch_partition(slice, pivot, lo, L, b_sums, parity, lt, by)

    # step 4: recursion
    if threadIdx().x == 1
        stuck = (pivot == prev_pivot && partition == lo || partition == hi) ? stuck + 1 : 0

        if stuck < 2 && partition > lo && partial_range_overlap(lo, partition, partial)
            s = CuDeviceStream()
            if S && sync_depth > 1
                @cuda(threads=blockDim().x, dynamic=true, stream=s, shmem=shmem,
                      qsort_kernel(slice, lo, partition, !parity, Val(true), sync_depth - 1,
                      pivot, lt, by, Val(1), partial, stuck))
            else
                @cuda(threads=blockDim().x, dynamic=true, stream=s, shmem=shmem,
                      qsort_kernel(slice, lo, partition, !parity, Val(false), sync_depth - 1,
                      pivot, lt, by, Val(1), partial, stuck))
            end
            CUDA.unsafe_destroy!(s)
        end

        if stuck < 2 && partition < hi && partial_range_overlap(partition, hi, partial)
            s = CuDeviceStream()
            if S && sync_depth > 1
                @cuda(threads=blockDim().x, dynamic=true, stream=s, shmem=shmem,
                      qsort_kernel(slice, partition, hi, !parity, Val(true), sync_depth - 1,
                      pivot, lt, by, Val(1), partial, stuck))
            else
                @cuda(threads=blockDim().x, dynamic=true, stream=s, shmem=shmem,
                      qsort_kernel(slice, partition, hi, !parity, Val(false), sync_depth - 1,
                      pivot, lt, by, Val(1), partial, stuck))
            end
            CUDA.unsafe_destroy!(s)
        end
    end

    return
end

function sort_args(args, partial_k :: Nothing)
    return args
end

function sort_args(args, partial_k)
    return (args..., partial_k)
end

#function sort

function quicksort!(c::AbstractArray{T,N}; lt::F1, by::F2, dims::Int, partial_k=nothing) where {T,N,F1,F2}
    max_depth = CUDA.limit(CUDA.LIMIT_DEV_RUNTIME_SYNC_DEPTH)
    len = size(c, dims)

    1 <= dims <= N || throw(ArgumentError("dimension out of range"))
    otherdims = ntuple(i -> i == dims ? 1 : size(c, i), N)

    my_sort_args = sort_args((c, 0, len, true, Val(N==1 && max_depth > 1),
             max_depth, nothing, lt, by, Val(dims)), partial_k)

    kernel = @cuda launch=false qsort_kernel(my_sort_args...)

    get_shmem(threads) = threads * (sizeof(Int) + sizeof(T))
    config = launch_configuration(kernel.fun, shmem=threads->get_shmem(threads))
    threads = prevpow(2, config.threads)

    kernel(my_sort_args...;
           blocks=prod(otherdims), threads=threads, shmem=get_shmem(threads))

    return c
end

end

# Base interface implementation

using .Quicksort

function Base.sort!(c::AnyCuArray; dims::Integer, lt=isless, by=identity, rev=false)
    # for reverse sorting, invert the less-than function
    if rev
        lt = !lt
    end

    quicksort!(c; lt, by, dims)
    return c
end

function Base.sort!(c::AnyCuVector; lt=isless, by=identity, rev=false)
    # for reverse sorting, invert the less-than function
    if rev
        lt = !lt
    end

    quicksort!(c; lt, by, dims=1)
    return c
end

function Base.sort(c::AnyCuArray; kwargs...)
    return sort!(copy(c); kwargs...)
end

function Base.partialsort!(c::AnyCuVector, k::Union{Integer, OrdinalRange}; lt=isless, by=identity, rev=false)
    # for reverse sorting, invert the less-than function
    if rev
        lt = !lt
    end

    function out(k :: OrdinalRange)
        return copy(c[k])
    end

    # work around disallowed scalar index
    function out(k :: Integer)
        return Array(c[k:k])[1]
    end

    quicksort!(c; lt, by, dims=1, partial_k=k)
    return out(k)
end

function Base.partialsort(c::AnyCuArray, k::Union{Integer, OrdinalRange}; kwargs...)
    return partialsort!(copy(c), k; kwargs...)
end


text:
using Statistics: mean, std
using Printf

struct CFLChecker <: Diagnostic end

struct FieldSummary <: Diagnostic
    diagnostic_frequency::Int
    fields::Array{Field,1}
    field_names::Array{AbstractString,1}
end

function run_diagnostic(model::Model, fs::FieldSummary)
    for (field, field_name) in zip(fs.fields, fs.field_names)
        padded_name = lpad(field_name, 4)
        field_min = minimum(field.data)
        field_max = maximum(field.data)
        field_mean = mean(field.data)
        field_abs_mean = mean(abs.(field.data))
        field_std = std(field.data)
        @printf("%s: min=%.6g, max=%.6g, mean=%.6g, absmean=%.6g, std=%.6g\n",
                padded_name, field_min, field_max, field_mean, field_abs_mean, field_std)
    end
end

struct NaNChecker <: Diagnostic
    diagnostic_frequency::Int
    fields::Array{Field,1}
    field_names::Array{AbstractString,1}
end

function run_diagnostic(model::Model, nc::NaNChecker)
    for (field, field_name) in zip(nc.fields, nc.field_names)
        if any(isnan, field.data)  # This is also fast on CuArrays.
            t, i = model.clock.time, model.clock.iteration
            println("time = $t, iteration = $i")
            println("NaN found in $field_name. Aborting simulation.")
            exit(1)
        end
    end
end

struct VelocityDivergenceChecker <: Diagnostic
    diagnostic_frequency::Int
    warn_threshold::Float64
    abort_threshold::Float64
end

function velocity_div!(grid::RegularCartesianGrid, u, v, w, div)
    @loop for k in (1:grid.Nz; blockIdx().z)
        @loop for j in (1:grid.Ny; (blockIdx().y - 1) * blockDim().y + threadIdx().y)
            @loop for i in (1:grid.Nx; (blockIdx().x - 1) * blockDim().x + threadIdx().x)
                @inbounds div[i, j, k] = div_f2c(grid, u, v, w, i, j, k)
            end
        end
    end

    @synchronize
end

function run_diagnostic(model::Model, diag::VelocityDivergenceChecker)
    u, v, w = model.velocities.u.data, model.velocities.v.data, model.velocities.w.data
    div = model.stepper_tmp.fC1

    velocity_div!(model.grid, u, v, w, div)
    min_div, mean_div, max_div = minimum(div), mean(div), maximum(div)

    if max(abs(min_div), abs(max_div)) >= diag.warn_threshold
        t, i = model.clock.time, model.clock.iteration
        println("time = $t, iteration = $i")
        println("WARNING: Velocity divergence is high! min=$min_div, mean=$mean_div, max=$max_div")
    end

    if max(abs(min_div), abs(max_div)) >= diag.abort_threshold
        t, i = model.clock.time, model.clock.iteration
        println("time = $t, iteration = $i")
        println("Velocity divergence is too high! min=$min_div, mean=$mean_div, max=$max_div. Aborting simulation.")
        exit(1)
    end
end


text:
notchheight(q2, q4, N) = 1.58 * (q4 - q2) / sqrt(N)

#=
Taken from https://github.com/JuliaPlots/StatPlots.jl/blob/master/src/boxplot.jl#L7
The StatPlots.jl package is licensed under the MIT "Expat" License:
    Copyright (c) 2016: Thomas Breloff.
=#
"""
    boxplot(x, y; kwargs...)
Draw a Tukey style boxplot.
The boxplot has 3 components:
- a `crossbar` spanning the interquartile (IQR) range with a midline marking the
    median
- an `errorbar` whose whiskers span `range * iqr`
- points marking outliers, that is, data outside the whiskers
# Arguments
- `x`: positions of the categories
- `y`: variables within the boxes
# Keywords
- `orientation=:vertical`: orientation of box (`:vertical` or `:horizontal`)
- `width=0.8`: width of the box
- `show_notch=false`: draw the notch
- `notchwidth=0.5`: multiplier of `width` for narrowest width of notch
- `show_median=true`: show median as midline
- `range`: multiple of IQR controlling whisker length
- `whiskerwidth`: multiplier of `width` for width of T's on whiskers, or
    `:match` to match `width`
- `show_outliers`: show outliers as points
"""
@recipe(BoxPlot, x, y) do scene
    Theme(
        color = theme(scene, :patchcolor),
        colormap = theme(scene, :colormap),
        colorrange = automatic,
        orientation = :vertical,
        # box and dodging
        width = automatic,
        dodge = automatic,
        n_dodge = automatic,
        x_gap = 0.2,
        dodge_gap = 0.03,
        strokecolor = theme(scene, :patchstrokecolor),
        strokewidth = theme(scene, :patchstrokewidth),
        # notch
        show_notch = false,
        notchwidth = 0.5,
        # median line
        show_median = true,
        mediancolor = theme(scene, :linecolor),
        medianlinewidth = theme(scene, :linewidth),
        # whiskers
        range = 1.5,
        whiskerwidth = 0.0,
        whiskercolor = theme(scene, :linecolor),
        whiskerlinewidth = theme(scene, :linewidth),
        # outliers points
        show_outliers = true,
        marker = theme(scene, :marker),
        markersize = theme(scene, :markersize),
        outliercolor = automatic,
        outlierstrokecolor = theme(scene, :markerstrokecolor),
        outlierstrokewidth = theme(scene, :markerstrokewidth),
        cycle = [:color => :patchcolor],
        inspectable = theme(scene, :inspectable)
    )
end

conversion_trait(x::Type{<:BoxPlot}) = SampleBased()

_cycle(v::AbstractVector, idx::Integer) = v[mod1(idx, length(v))]
_cycle(v, idx::Integer) = v

_flip_xy(p::Point2f) = reverse(p)
_flip_xy(r::Rect{2,T}) where {T} = Rect{2,T}(reverse(r.origin), reverse(r.widths))

function Makie.plot!(plot::BoxPlot)
    args = @extract plot (width, range, show_outliers, whiskerwidth, show_notch, orientation, x_gap, dodge, n_dodge, dodge_gap)

    signals = lift(
        plot[1],
        plot[2],
        plot[:color],
        args...,
    ) do x, y, color, width, range, show_outliers, whiskerwidth, show_notch, orientation, x_gap, dodge, n_dodge, dodge_gap
        x̂, boxwidth = xw_from_dodge(x, width, 1.0, x_gap, dodge, n_dodge, dodge_gap)
        if !(whiskerwidth == :match || whiskerwidth >= 0)
            error("whiskerwidth must be :match or a positive number. Found: $whiskerwidth")
        end
        ww = whiskerwidth == :match ? boxwidth : whiskerwidth * boxwidth
        outlier_points = Point2f[]
        centers = Float32[]
        medians = Float32[]
        boxmin = Float32[]
        boxmax = Float32[]
        notchmin = Float32[]
        notchmax = Float32[]
        t_segments = Point2f[]
        outlier_indices = Int[]
        T = color isa AbstractVector ? eltype(color) : typeof(color)
        boxcolor = T[]
        for (i, (center, idxs)) in enumerate(StructArrays.finduniquesorted(x̂))
            values = view(y, idxs)

            # compute quantiles
            q1, q2, q3, q4, q5 = quantile(values, LinRange(0, 1, 5))

            # notches
            if show_notch
                nh = notchheight(q2, q4, length(values))
                nmin, nmax = q3 - nh, q3 + nh
                push!(notchmin, nmin)
                push!(notchmax, nmax)
            end

            # outliers
            if Float64(range) != 0.0  # if the range is 0.0, the whiskers will extend to the data
                limit = range * (q4 - q2)
                inside = Float64[]
                for (value, idx) in zip(values,idxs)
                    if (value < (q2 - limit)) || (value > (q4 + limit))
                        if show_outliers
                            push!(outlier_points, (center, value))
                            # register outlier box indices
                            push!(outlier_indices, idx)
                        end
                    else
                        push!(inside, value)
                    end
                end
                # change q1 and q5 to show outliers
                # using maximum and minimum values inside the limits
                q1, q5 = extrema_nan(inside)
                # register boxcolor
                push!(boxcolor, getuniquevalue(color, idxs))
            end

            # whiskers
            HW = 0.5 * _cycle(ww, i) # Whisker width
            lw, rw = center - HW, center + HW
            push!(t_segments, (center, q2), (center, q1), (lw, q1), (rw, q1)) # lower T
            push!(t_segments, (center, q4), (center, q5), (rw, q5), (lw, q5)) # upper T

            # box
            push!(centers, center)
            push!(boxmin, q2)
            push!(medians, q3)
            push!(boxmax, q4)
        end

        # for horizontal boxplots just flip all components
        if orientation == :horizontal
            outlier_points = _flip_xy.(outlier_points)
            t_segments = _flip_xy.(t_segments)
        elseif orientation != :vertical
            error("Invalid orientation $orientation. Valid options: :horizontal or :vertical.")
        end

        return (
            centers = centers,
            boxmin = boxmin,
            boxmax = boxmax,
            medians = medians,
            notchmin = notchmin,
            notchmax = notchmax,
            outliers = outlier_points,
            t_segments = t_segments,
            boxwidth = boxwidth,
            outlier_indices = outlier_indices,
            boxcolor = boxcolor,
        )
    end
    centers = @lift($signals.centers)
    boxmin = @lift($signals.boxmin)
    boxmax = @lift($signals.boxmax)
    medians = @lift($signals.medians)
    notchmin = @lift($show_notch ? $signals.notchmin : automatic)
    notchmax = @lift($show_notch ? $signals.notchmax : automatic)
    outliers = @lift($signals.outliers)
    t_segments = @lift($signals.t_segments)
    boxwidth = @lift($signals.boxwidth)
    outlier_indices = @lift($signals.outlier_indices)
    boxcolor = @lift($signals.boxcolor)

    outliercolor = lift(plot[:outliercolor], plot[:color], outlier_indices) do outliercolor, color, outlier_indices
        c = outliercolor === automatic ? color : outliercolor
        if c isa AbstractVector
            return c[outlier_indices]
        else
            return c
        end
    end

    scatter!(
        plot,
        color = outliercolor,
        marker = plot[:marker],
        markersize = plot[:markersize],
        strokecolor = plot[:outlierstrokecolor],
        strokewidth = plot[:outlierstrokewidth],
        outliers,
        inspectable = plot[:inspectable]
    )
    linesegments!(
        plot,
        color = plot[:whiskercolor],
        linewidth = plot[:whiskerlinewidth],
        t_segments,
        inspectable = plot[:inspectable]
    )
    crossbar!(
        plot,
        color = boxcolor,
        colorrange = plot[:colorrange],
        colormap = plot[:colormap],
        strokecolor = plot[:strokecolor],
        strokewidth = plot[:strokewidth],
        midlinecolor = plot[:mediancolor],
        midlinewidth = plot[:medianlinewidth],
        show_midline = plot[:show_median],
        orientation = orientation,
        width = boxwidth,
        show_notch = show_notch,
        notchmin = notchmin,
        notchmax = notchmax,
        notchwidth = plot[:notchwidth],
        inspectable = plot[:inspectable],
        centers,
        medians,
        boxmin,
        boxmax,
    )
end


text:
@testset "1024.video-stitching.jl" begin
    @test video_stitching([[0, 4], [2, 8]], 5) == 2
    @test video_stitching(
        [
            [0, 1],
            [6, 8],
            [0, 2],
            [5, 6],
            [0, 4],
            [0, 3],
            [6, 7],
            [1, 3],
            [4, 7],
            [1, 4],
            [2, 5],
            [2, 6],
            [3, 4],
            [4, 5],
            [5, 7],
            [6, 9],
        ],
        9,
    ) == 3
    @test video_stitching([[0, 1], [1, 2]], 5) == -1
    @test video_stitching([[0, 2], [4, 6], [8, 10], [1, 9], [1, 5], [5, 9]], 10) == 3
end