text:
export AbstractTaxType,
    AbstractTechChange,
    AbstractConstrainedType,
    AbstractLearningType,
    AbstractRoyaltyType,
    NoTaxes,
    WithTaxes,
    GathProcess,
    NoTrend,
    TimeTrend,
    TimeFE,
    Unconstrained,
    Constrained,
    DrillingRevenue,
    Learn,
    NoLearn,
    WithRoyalty,
    NoRoyalty,
    ConstrainedProblem,
    UnconstrainedProblem,
    NoLearningProblem,
    STARTING_α_ψ,
    STARTING_log_ogip,
    STARTING_α_t,
    STARTING_α_D,
    TimeTrendDgt0

# ----------------------------------------------------------------
# Drilling revenue variations
# ----------------------------------------------------------------

abstract type AbstractLearningType    <: AbstractModelVariations end
abstract type AbstractConstrainedType <: AbstractModelVariations end
abstract type AbstractTechChange      <: AbstractModelVariations end
abstract type AbstractTaxType         <: AbstractModelVariations end
abstract type AbstractRoyaltyType     <: AbstractModelVariations end

# drilling revenue
struct DrillingRevenue{
    Cn <: AbstractConstrainedType, Tech <: AbstractTechChange,
    Tax <: AbstractTaxType, Lrn <: AbstractLearningType,
     Roy <: AbstractRoyaltyType
} <: AbstractDrillingRevenue
    constr::Cn
    tech::Tech
    tax::Tax
    learn::Lrn
    royalty::Roy
end

DrillingRevenue(Cn, Tech, Tax) = DrillingRevenue(Cn, Tech, Tax, Learn(), WithRoyalty())

@inline constr( x::DrillingRevenue) = x.constr
@inline tech(   x::DrillingRevenue) = x.tech
@inline tax(    x::DrillingRevenue) = x.tax
@inline learn(  x::DrillingRevenue) = x.learn
@inline royalty(x::DrillingRevenue) = x.royalty

# Constrained parameters
# -----------------------

# Whether we've constrained coefs
struct Unconstrained <: AbstractConstrainedType end
struct Constrained   <: AbstractConstrainedType
    log_ogip::Float64
    α_ψ::Float64
    α_t::Float64
    α_D::Float64
end
Constrained(; log_ogip=STARTING_log_ogip, α_ψ = STARTING_α_ψ, α_t = STARTING_α_t, α_D = STARTING_α_D, args...) = Constrained(log_ogip, α_ψ, α_t, α_D)

@inline theta_g(x::Constrained) = x.log_ogip
@inline theta_ψ(x::Constrained) = x.α_ψ
@inline theta_t(x::Constrained) = x.α_t
@inline theta_D(x::Constrained) = x.α_D

const DrillingRevenueConstrained = DrillingRevenue{Constrained}
const DrillingRevenueUnconstrained = DrillingRevenue{Unconstrained}

# Technology
# -----------------

abstract type AbstractTimeTrend <: AbstractTechChange end

const TIME_TREND_BASE = 2008.5

struct NoTrend <: AbstractTechChange  end
struct TimeTrend <: AbstractTimeTrend
    baseyear::Float64
end
struct TimeTrendDgt0 <: AbstractTimeTrend
    baseyear::Float64
end
(::Type{T})() where {T<:AbstractTimeTrend} = T(TIME_TREND_BASE)
struct TimeFE <: AbstractTechChange
    yearrange::UnitRange{Float64}
end

# Special TimeFE stuff
yearrange(x::TimeFE) = x.yearrange
yearrange(x::AbstractTechChange) = UnitRange(2008.5, 2016.5)
start(x::TimeFE) = first(yearrange(x))
stop(x::TimeFE) = last(yearrange(x))
_nparm(x::TimeFE) = length(yearrange(x))-1
TimeFE(a,b) = TimeFE(UnitRange(a,b))
TimeFE() = TimeFE(2008.5, 2016.5)
function coefnames(x::TimeFE)
    yrng = UnitRange(start(x)+1, stop(x))
    cfnms = ["\\alpha_{" * string(yr) *"}" for yr in yrng]
    return cfnms
end

baseyear(x::AbstractTimeTrend) = x.baseyear
baseyear(x::NoTrend) = TIME_TREND_BASE
baseyear(x::TimeFE) = start(x)
baseyear(x::DrillingRevenue) = baseyear(tech(x))
baseyear(x::DrillReward) = baseyear(revenue(x))

const DrillingRevenueTimeTrend         = DrillingRevenue{Cn,TimeTrend} where {Cn}
const DrillingRevenueTimeTrendDgt0     = DrillingRevenue{Cn,TimeTrendDgt0} where {Cn}
const DrillingRevenueAbstractTimeTrend = DrillingRevenue{Cn,T} where {Cn,T<:AbstractTimeTrend}
const DrillingRevenueNoTrend           = DrillingRevenue{Cn,NoTrend} where {Cn}
const DrillingRevenueTimeFE            = DrillingRevenue{Cn,TimeFE} where {Cn}

@inline centered_time(  x::DrillingRevenue, z::Tuple) = year(z) - baseyear(tech(x))
@inline trend_component(x::DrillingRevenueTimeTrend,     θ, obs) = theta_t(x,θ) * centered_time(x, zchars(obs))
@inline trend_component(x::DrillingRevenueNoTrend,       θ, obs) = 0
@inline trend_component(x::DrillingRevenueTimeTrendDgt0, θ, obs) = theta_t(x,θ)*centered_time(x,zchars(obs)) + _Dgt0(obs)*theta_D(x,θ)

@inline function idx_year(x::DrillingRevenueTimeFE, z)
    timefe = tech(x)
    if year(z) < start(timefe)
        return 0
    elseif year(z) > stop(timefe)
        return _nparm(timefe)
    else
        t = centered_time(x,z)
        return Int(floor(t))
    end
end

function trend_component(x::DrillingRevenueTimeFE, θ, obs)
    i = idx_year(x,zchars(obs))
    if i > 0
        return theta_t(x,θ,i)
    else
        return zero(eltype(θ))
    end
end

@inline function gradt!(grad, x::DrillingRevenueNoTrend, rev, θ, obs)
    nothing
end
@inline function gradt!(grad, x::DrillingRevenueTimeTrend, rev, θ, obs)
    grad[idx_t(x)] = rev*centered_time(x, zchars(obs))
end
@inline function gradt!(grad, x::DrillingRevenueTimeTrendDgt0, rev, θ, obs)
    grad[idx_t(x)] = rev*centered_time(x, zchars(obs))
    grad[idx_D(x)] = rev*_Dgt0(obs)
end
@inline function gradt!(grad, x::DrillingRevenueTimeFE, rev, θ, obs)
    grad[idx_t(x)] .= 0
    i = idx_year(x,zchars(obs))
    if i > 0
        grad[idx_t(x,i)] = rev
    end
end

# chebshev polynomials
# See http://www.aip.de/groups/soe/local/numres/bookcpdf/c5-8.pdf
@inline checkinterval(x::Real,min::Real,max::Real) =  min <= x <= max || throw(DomainError("x = $x must be in [$min,$max]"))
@inline checkinterval(x::Real) = checkinterval(x,-1,1)

@inline cheb0(z::Real) = one(z)
@inline cheb1(z::Real) = clamp(z,-1,1)
@inline cheb2(z::Real) = (x = cheb1(z); return 2*x^2 - 1)
@inline cheb3(z::Real) = (x = cheb1(z); return 4*x^3 - 3*x)
@inline cheb4(z::Real) = (x = cheb1(z); return 8*(x^4 - x^2) + 1)

# learning
# -----------------------

struct Learn       <: AbstractLearningType end
struct NoLearn     <: AbstractLearningType end
struct PerfectInfo <: AbstractLearningType end
struct MaxLearning <: AbstractLearningType end

const DrillingRevenueLearn       = DrillingRevenue{Cn,Tech,Tax,Learn}       where {Cn,Tech,Tax}
const DrillingRevenueNoLearn     = DrillingRevenue{Cn,Tech,Tax,NoLearn}     where {Cn,Tech,Tax}
const DrillingRevenuePerfectInfo = DrillingRevenue{Cn,Tech,Tax,PerfectInfo} where {Cn,Tech,Tax}
const DrillingRevenueMaxLearning = DrillingRevenue{Cn,Tech,Tax,MaxLearning} where {Cn,Tech,Tax}

@inline _ρ(x::AbstractLearningType, σ::Number) = _ρ(σ)
@inline _ρ(x::PerfectInfo         , σ::Number) = 1
@inline _ρ(x::MaxLearning         , σ::Number) = 0
@inline _ρ(x::DrillingRevenue     , σ::Number) = _ρ(learn(x), σ)

# @deprecate _ρ(σ::Number, x::AbstractLearningType) _ρ(x,σ)
# @deprecate _ρ(σ::Number, x::PerfectInfo) _ρ(x,σ)
# @deprecate _ρ(σ::Number, x::MaxLearning) _ρ(x,σ)
# @deprecate _ρ(σ::Number, x::DrillingRevenue) _ρ(x,σ)

# parameter access
# -----------------------

@inline _nparm(x::DrillingRevenue{Constrained}) = 2
@inline _nparm(x::DrillingRevenue{Unconstrained, NoTrend}) = 4
@inline _nparm(x::DrillingRevenue{Unconstrained, TimeTrend}) = 5
@inline _nparm(x::DrillingRevenue{Unconstrained, TimeTrendDgt0}) = 6
@inline _nparm(x::DrillingRevenue{Unconstrained, TimeFE}) = 4 + _nparm(tech(x))

@inline idx_0(x::DrillingRevenue) = 1
@inline idx_g(x::DrillingRevenue) = 2
@inline idx_ψ(x::DrillingRevenue) = 3
@inline idx_t(x::DrillingRevenueAbstractTimeTrend) = 4
@inline idx_t(x::DrillingRevenueTimeFE) = 3 .+ OneTo(_nparm(tech(x)))
function idx_t(x::DrillingRevenueTimeFE, i)
    1 <= i <= _nparm(tech(x)) || throw(DomainError(i))
    return 3 + i
end
@inline idx_D(x::DrillingRevenueTimeTrendDgt0) = 5
@inline idx_ρ(x::DrillingRevenue) = _nparm(x)

@inline theta_0(x::DrillingRevenue, θ) = θ[idx_0(x)]
@inline theta_g(x::DrillingRevenue, θ) = θ[idx_g(x)]
@inline theta_ψ(x::DrillingRevenue, θ) = θ[idx_ψ(x)]
@inline theta_t(x::DrillingRevenue{Unconstrained}, θ) = θ[idx_t(x)]
@inline theta_t(x::DrillingRevenue{Unconstrained, TimeFE}, θ) = uview(θ, idx_t(x))
@inline theta_t(x::DrillingRevenue{Unconstrained, TimeFE}, θ, i) = θ[idx_t(x,i)]
@inline theta_D(x::DrillingRevenue{Unconstrained, TimeTrendDgt0},θ) = θ[idx_D(x)]
@inline theta_D(x, θ) = zero(Float64)
@inline theta_ρ(x, θ) = θ[idx_ρ(x)]
function theta_t(x::DrillingRevenue{Unconstrained, NoTrend}, θ)
    # @warn "No t in $x. setting to $STARTING_α_t"
    return STARTING_α_t
end
@inline theta_g(x::DrillingRevenue{Constrained}, θ) = theta_g(constr(x))
@inline theta_ψ(x::DrillingRevenue{Constrained}, θ) = theta_ψ(constr(x))
@inline theta_t(x::DrillingRevenue{Constrained}, θ) = theta_t(constr(x))

function ConstrainedCoefs(x::DrillingRevenueUnconstrained, theta)
    @assert length(theta) == _nparm(x)
    g   = theta_g(x, theta)
    psi = theta_ψ(x, theta)
    t   = theta_t(x, theta)
    D   = theta_D(x, theta)
    out = (log_ogip=g, α_ψ=psi, α_t=t, α_D=D)
    return out
end

ConstrainedCoefs(x::DrillReward, θ) = ConstrainedCoefs(revenue(x), vw_revenue(x, θ))
ConstrainedIdx(x::AbstractDrillingRevenue) = idx_g(x), idx_ψ(x), idx_t(x)
ConstrainedIdx(x::DrillingRevenueNoTrend) = idx_g(x), idx_ψ(x)
ConstrainedIdx(x::DrillingRevenueTimeTrendDgt0) = idx_g(x), idx_ψ(x), idx_t(x), idx_D(x)

function UnconstrainedFmConstrainedIdx(x::DrillingRevenueConstrained)
    xu = UnconstrainedProblem(x)
    idxu = [idx_0(xu), idx_ρ(xu)]
end

# @deprecate α_0(     x::DrillingRevenue, θ) theta_0(x, θ)
# @deprecate _σ(      x::DrillingRevenue, θ) theta_ρ(x, θ)
# @deprecate log_ogip(x::DrillingRevenue, θ) theta_g(x, θ)
# @deprecate α_ψ(     x::DrillingRevenue, θ) theta_ψ(x, θ)
# @deprecate α_t(     x::DrillingRevenue, θ) theta_t(x, θ)

# taxes
# -----------------------

# From Gulen et al (2015) "Production scenarios for the Haynesville..."
const GATH_COMP_TRTMT_PER_MCF   = 0.42 + 0.07
const MARGINAL_TAX_RATE = 0.402

# other calculations
const REAL_DISCOUNT_AND_DECLINE = 0x1.8e06b611ed4d8p-1  #  0.777394952476835= sum( β^((t+5-1)/12) q(t)/Q(240) for t = 1:240)

const STARTING_α_ψ      = 0x1.7587cc6793516p-2 # 0.365
const STARTING_log_ogip = 0x1.401755c339009p-1 # 0.625
const STARTING_α_t      = 0.01948
const STARTING_α_D      = 0.0

struct NoTaxes <: AbstractTaxType end

struct WithTaxes <: AbstractTaxType
    one_minus_mgl_tax_rate::Float64
    cost_per_mcf::Float64
end

struct GathProcess <: AbstractTaxType
    cost_per_mcf::Float64
end

const DrillingRevenueNoTaxes     = DrillingRevenue{Cn,Tech,NoTaxes}     where {Cn,Tech}
const DrillingRevenueWithTaxes   = DrillingRevenue{Cn,Tech,WithTaxes}   where {Cn,Tech}
const DrillingRevenueGathProcess = DrillingRevenue{Cn,Tech,GathProcess} where {Cn,Tech}

WithTaxes(;mgl_tax_rate = 1-MARGINAL_TAX_RATE, cost_per_mcf = GATH_COMP_TRTMT_PER_MCF * REAL_DISCOUNT_AND_DECLINE) = WithTaxes(mgl_tax_rate, cost_per_mcf)
GathProcess(;cost_per_mcf = GATH_COMP_TRTMT_PER_MCF * REAL_DISCOUNT_AND_DECLINE) = GathProcess(cost_per_mcf)

@inline one_minus_mgl_tax_rate(x::AbstractTaxType) = 1
@inline one_minus_mgl_tax_rate(x::WithTaxes) = x.one_minus_mgl_tax_rate
@inline one_minus_mgl_tax_rate(x) = one_minus_mgl_tax_rate(tax(x))

@inline cost_per_mcf(x::NoTaxes) = 0
@inline cost_per_mcf(x::AbstractTaxType) = x.cost_per_mcf
@inline cost_per_mcf(x) = cost_per_mcf(tax(x))

# royalty
# -----------------------

struct WithRoyalty <: AbstractRoyaltyType end
struct NoRoyalty   <: AbstractRoyaltyType end

@inline (::WithRoyalty)(roy) = roy
@inline (::NoRoyalty)(roy) = 0
@inline oneminusroyalty(x::DrillingRevenue, roy) = 1-royalty(x)(roy)
@inline d_tax_royalty(x::DrillingRevenue, d::Integer, roy::Real) = d * oneminusroyalty(x, roy) * one_minus_mgl_tax_rate(x)

# base functions
# -------------------------------------------

@inline function Eexpψ(x::DrillingRevenueLearn, d, obs, θ, sim)
    ψ = _ψ(obs, sim)
    θ4 = theta_ψ(x,θ)

    if _Dgt0(obs)
        return θ4*ψ
    else
        σ = theta_ρ(x,θ)
        ρ = _ρ(learn(x), σ)
        return θ4*(ψ*ρ + θ4*0.5*(1-ρ^2))
    end
end

@inline function dEexpψdα_ψ(x::DrillingRevenueLearn, d, obs, θ, sim)
    ψ = _ψ(obs, sim)
    if _Dgt0(obs)
        return ψ
    else
        σ = theta_ρ(x,θ)
        ρ = _ρ(x, σ)
        return ψ*ρ + theta_ψ(x,θ) * (1-ρ^2)
    end
end

@inline function Eexpψ(x::DrillingRevenueNoLearn, d, obs, θ, sim)
    ψ = _ψ(obs, sim)
    θ4= theta_ψ(x,θ)
    σ = theta_ρ(x,θ)
    ρ = _ρ(σ)
    return θ4*(ψ*ρ + θ4*0.5*(1-ρ^2))
end

@inline function Eexpψ(x::DrillingRevenueMaxLearning, d, obs, θ, sim)
    θ4= theta_ψ(x,θ)
    if _Dgt0(obs)
        return θ4 * _ψ(obs, sim)
    else
        return 0.5 * θ4^2
    end
end

@inline Eexpψ(x::DrillingRevenuePerfectInfo, d, obs, θ, sim) = theta_ψ(x,θ) * _ψ(obs, sim)

@inline function dEexpψdσ(x::DrillingRevenueLearn, d, obs, θ, sim)
    T = eltype(θ)
    ψ = _ψ(obs, sim)
    if !_Dgt0(obs) && d > 0
        θ4= theta_ψ(x,θ)
        σ = theta_ρ(x,θ)
        out = θ4 *(ψ - θ4*_ρ(σ)) * _dρdθρ(σ)
    else
        out = zero(T)
    end
    return out::T
end

# quantity (for post-estimation SIMULATIONS)
# -----------
function eur_kernel(x::DrillingRevenue, d, obs::ObservationDrill, θ, sim)
    @assert length(θ) == _nparm(x)
    z = zchars(obs)
    logQ = theta_g(x,θ)*geology(obs) + theta_ψ(x,θ)*_ψ2(sim) + trend_component(x, θ, obs)
    return exp(logQ)
end

# ----------------------------------------------------------------
# flow revenue
# ----------------------------------------------------------------

function flow(x::DrillingRevenueNoTaxes, d::Integer, obs::ObservationDrill, θ::AbstractVector, sim)
    z = zchars(obs)
    u = d_tax_royalty(x, d, royalty(obs)) *
    exp(
        theta_0(x,θ) + logprice(obs) + theta_g(x,θ)*geology(obs) +
        Eexpψ(x, d, obs, θ, sim) +
        trend_component(x, θ, obs)
    )
    return u
end

function flow(x::DrillingRevenue, d::Integer, obs::ObservationDrill, θ::AbstractVector, sim)
    z = zchars(obs)
    u = d_tax_royalty(x, d, royalty(obs)) *
    exp(
        theta_0(x,θ) + theta_g(x,θ)*geology(obs) +
        Eexpψ(x, d, obs, θ, sim) +
        trend_component(x, θ, obs)
    ) * (
        price(obs) - cost_per_mcf(x)
    )
    return u
end

# Gradient
# ------------------------------

@inline function flowzero!(grad, θ, dograd)
    dograd && fill!(grad, 0)
    return azero(θ)
end

function flow!(grad, x::DrillingRevenue{Unconstrained}, d, obs, θ, sim, dograd::Bool)
    d == 0 && return flowzero!(grad,θ,dograd)
    z = zchars(obs)
    rev = flow(x, d, obs, θ, sim)
    if dograd
        grad[idx_0(x)] = rev
        grad[idx_g(x)] = rev*geology(obs)
        grad[idx_ψ(x)] = rev*dEexpψdα_ψ(x, d, obs, θ, sim)
        gradt!(grad, x, rev, θ, obs)
        grad[idx_ρ(x)] = rev*dEexpψdσ(x, d, obs, θ, sim)
    end
    return rev
end

# Constrained gradient
# ------------------------------

function flow!(grad, x::DrillingRevenue{Constrained}, d, obs, θ, sim, dograd::Bool)
    d == 0 && return flowzero!(grad,θ,dograd)
    rev = flow(x, d, obs, θ, sim)
    if dograd
        grad[idx_0(x)] = rev
        grad[idx_ρ(x)] = rev*dEexpψdσ(x, d, obs, θ, sim)
    end
    return rev
end

coefnames(x::DrillingRevenue{Unconstrained, NoTrend}) =
    ["\\alpha_0", "\\alpha_g", "\\alpha_\\psi", "\\theta_\\rho"]
coefnames(x::DrillingRevenue{Unconstrained, TimeTrend}) =
    ["\\alpha_0", "\\alpha_g", "\\alpha_\\psi", "\\alpha_t", "\\theta_\\rho"]
coefnames(x::DrillingRevenue{Unconstrained, TimeTrendDgt0}) =
    ["\\alpha_0", "\\alpha_g", "\\alpha_\\psi", "\\alpha_t", "\\alpha_D", "\\theta_\\rho"]
coefnames(x::DrillingRevenue{Constrained}) =
    ["\\alpha_0", "\\theta_\\rho"]
coefnames(x::DrillingRevenue{Unconstrained, TimeFE}) = vcat(
    "\\alpha_0", "\\alpha_g", "\\alpha_\\psi", coefnames(tech(x)), "\\theta_\\rho")

# ----------------------------------------------------------------
# dψ is the same across many functions
# ----------------------------------------------------------------

function flowdψ(rev::Real, x::DrillingRevenue, d, obs, θ, sim)
    dψ = rev * theta_ψ(x,θ)
    if _Dgt0(obs)
        return dψ
    else
        return dψ * _ρ(x,theta_ρ(x,θ))
    end
end

function flowdψ(x::DrillingRevenue, d, obs, θ, sim)
    d == 0 && return azero(θ)
    rev = flow(x, d, obs, θ, sim)
    return flowdψ(rev, x, d, obs, θ, sim)
end

# Conversion between model types
# ----------------------------------------------------------------

# Constrained / Unconstrained
ConstrainedProblem(  x::AbstractPayoffComponent; kwargs...) = x
UnconstrainedProblem(x::AbstractPayoffComponent; kwargs...) = x

UnconstrainedProblem(x::DrillingRevenue; kwargs...) = DrillingRevenue(Unconstrained(;kwargs...), tech(x), tax(x), learn(x), royalty(x))
ConstrainedProblem(  x::DrillingRevenue; kwargs...) = DrillingRevenue(Constrained(;kwargs...),   tech(x), tax(x), learn(x), royalty(x))

ConstrainedProblem(  x::DrillReward; kwargs...) = DrillReward(ConstrainedProblem(revenue(x); kwargs...), ConstrainedProblem(drill(x)), ConstrainedProblem(extend(x)))
UnconstrainedProblem(x::DrillReward; kwargs...) = DrillReward(UnconstrainedProblem(revenue(x); kwargs...), UnconstrainedProblem(drill(x)), UnconstrainedProblem(extend(x)))
ConstrainedProblem(x::DrillReward, theta::AbstractVector) = ConstrainedProblem(x; ConstrainedCoefs(x,theta)...)

# Learning models
NoLearningProblem(x::AbstractPayoffComponent, args...) = x
LearningProblem(  x::AbstractPayoffComponent, args...) = x

NoLearningProblem(x::DrillingRevenue, args...) = DrillingRevenue(constr(x), tech(x), tax(x), NoLearn(),     royalty(x))
LearningProblem(  x::DrillingRevenue, args...) = DrillingRevenue(constr(x), tech(x), tax(x), Learn(),       royalty(x))
PerfectInfo(      x::DrillingRevenue, args...) = DrillingRevenue(constr(x), tech(x), tax(x), PerfectInfo(), royalty(x))

NoLearningProblem(x::DrillReward, args...) = DrillReward(NoLearningProblem(revenue(x), args...), drill(x), extend(x))
LearningProblem(  x::DrillReward, args...) = DrillReward(  LearningProblem(revenue(x), args...), drill(x), extend(x))

# Royalty
NoRoyaltyProblem(  x::AbstractPayoffComponent, args...) = x
WithRoyaltyProblem(x::AbstractPayoffComponent, args...) = x

NoRoyaltyProblem(  x::DrillingRevenue, args...)      = DrillingRevenue(constr(x), tech(x), tax(x), learn(x), NoRoyalty())
WithRoyaltyProblem(x::DrillingRevenue, args...)      = DrillingRevenue(constr(x), tech(x), tax(x), learn(x), WithRoyalty())

NoRoyaltyProblem(  x::DrillReward, args...) = DrillReward(  NoRoyaltyProblem(revenue(x), args...), drill(x), extend(x))
WithRoyaltyProblem(x::DrillReward, args...) = DrillReward(WithRoyaltyProblem(revenue(x), args...), drill(x), extend(x))


text:
"""
    _check_signs(pa, pb, values = ())

Check the signs of `pa`, `pb` and `values`. Returns `1` if all of them
definitely have the same sign, `-1` if there are at least two who
definitely have different signs and `0` otherwise.
"""
function _check_signs(pa, pb, values = ())
    has_negative = false
    has_positive = false
    has_zero = false

    for v in Base.Iterators.flatten(((pa, pb), values))
        sgn = Arblib.sgn_nonzero(v)

        has_negative |= sgn == -1
        has_positive |= sgn == 1
        has_zero |= sgn == 0

        has_negative && has_positive && return -1
    end

    if has_zero
        return 0
    else
        return 1
    end
end

"""
    extrema_polynomial(p::ArbPoly, a::Arf, b::Arf; abs_value = false, verbose = false)

Compute an enclosure of both the minimum and maximum of the polynomial
`p` on the interval `[a, b]` and return them as a 2-tuple.

The extrema are computed by finding the zeros of the derivative of `p`
using [`isolate_roots`](@ref) and refining them using
[`refine_root`](@ref), then evaluating `p` on the potential zeros as
well as the endpoints `a`, `b` of the interval.

If `abs_value = true` the compute the extrema of `abs(p(x))` on the
interval `[a, b]`. For the computation of the maximum this is mostly
the same, only difference is that we have to take the absolute value
of the evaluations. For the minimum we have to take into account that
the polynomial might cross zero, in which case the minimum is zero. To
see if `p` crosses zero we compare the sign of `p` at the endpoints
and all zeros of its derivative. If any of these have different signs
we are guaranteed that `p` crosses zero and the minimum is zero, if
all of them have the same sign we are instead guaranteed that the
polynomial doesn't cross zero.

If `verbose = true` then output information about the process.

TODO: Allow passing arguments to `isolate_roots`?
"""
function extrema_polynomial(p::ArbPoly, a::Arf, b::Arf; abs_value = false, verbose = false)
    check_interval(a, b)

    maybe_abs = abs_value ? abs : identity

    if a == b
        pa = maybe_abs(p(a))
        return pa, pa
    end

    # Compute values at endpoints
    pa, pb = p(a), p(b)

    # Flag to determine if the minimum value has already been
    # determined in the case when computing the absolute value.
    min_done = false

    # Handle the signs not being the same when computing the absolute
    # value
    if abs_value
        sgn = _check_signs(pa, pb)
        if sgn == -1
            # The sign differs - the minimum is zero
            verbose && @info "sign of endpoints differ - minimum is zero"

            min_value = zero(pa)
            min_done = true
        elseif sgn == 0
            # At least one of the enclosures contain zero, can't get a
            # better enclosure than the minimum of these two
            min_value = min(maybe_abs(pa), maybe_abs(pb))
            Arblib.nonnegative_part!(min_value, min_value)

            verbose &&
                @info "sign of endpoints undetermined - minimum on endpoints: $min_value"
            min_done = true
        end
    end

    # Compute the extrema on the endpoints
    min_endpoints, max_endpoints = minmax(maybe_abs(pa), maybe_abs(pb))

    verbose && @info "extrema on endpoints: $((min_endpoints, max_endpoints))"

    # Short circuit non-finite values and return NaN
    !isfinite(min_endpoints) ||
        !isfinite(max_endpoints) && return (
            Arb(NaN, prec = precision(min_endpoints)),
            Arb(NaN, prec = precision(min_endpoints)),
        )

    # Short circuit on constant or linear polynomial where extrema is
    # attained at endpoints (or possibly is zero if computing absolute
    # value)
    if Arblib.degree(p) <= 1
        if min_done
            return min_value, max_endpoints
        else
            return min_endpoints, max_endpoints
        end
    end

    # Compute the roots of the derivative
    dp = Arblib.derivative(p)
    roots, flags = isolate_roots(dp, a, b)
    roots = Arb.(roots) # We want the roots as balls rather than tuples

    verbose &&
        @info "found $(length(roots)) intervals with possible roots of which $(sum(flags)) unique"

    # Evaluate the polynomial on all the roots
    values = p.(roots)

    # Handle sign crossings when computing the absolute value
    if abs_value && !min_done && _check_signs(pa, pb, values) == -1
        verbose && @info "polynomial crosses zero - minimum is zero"
        min_value = zero(pa)
        min_done = true
    end

    if !min_done
        min_value = copy(min_endpoints)
    end
    max_value = copy(max_endpoints)
    for v in values
        v = maybe_abs(v)
        min_done || Arblib.min!(min_value, min_value, v)
        Arblib.max!(max_value, max_value, v)
    end

    abs_value && Arblib.nonnegative_part!(min_value, min_value)
    abs_value && Arblib.nonnegative_part!(max_value, max_value)

    verbose && @info "extrema with unrefined roots: $((min_value, max_value))"

    # Refine the unique roots where the extrema could occur
    count = 0
    for i in eachindex(roots)
        !(maybe_abs(values[i]) < max_value)

        if flags[i] && (
            !(min_done || maybe_abs(values[i]) > min_value) ||
            !(maybe_abs(values[i]) < max_value)
        )
            roots[i] = refine_root(dp, roots[i], strict = false)
            values[i] = p(roots[i])
            count += 1
        end
    end

    verbose && @info "refined $count roots"

    if abs_value && !min_done && _check_signs(pa, pb, values) == -1
        verbose && @info "polynomial crosses zero - minimum is zero"
        min_value = zero(pa)
        min_done = true
    end

    # Recompute extrema on the refined roots
    if !min_done
        min_value = copy(min_endpoints)
    end
    max_value = copy(max_endpoints)
    for v in values
        v = maybe_abs(v)
        min_done || Arblib.min!(min_value, min_value, v)
        Arblib.max!(max_value, max_value, v)
    end

    abs_value && Arblib.nonnegative_part!(min_value, min_value)
    abs_value && Arblib.nonnegative_part!(max_value, max_value)

    verbose && @info "extrema with refined roots: $((min_value, max_value))"

    return min_value, max_value
end

"""
    minimum_polynomial(p::ArbPoly, a::Arf, b::Arf; abs_value = false, verbose = false)

Compute an enclosure of the minimum of the polynomial `p` on the
interval `[a, b]`.

Takes the same arguments as [`extrema_polynomial`](@ref). The
algorithm is also the same except that it only looks for the minimum.
"""
function minimum_polynomial(p::ArbPoly, a::Arf, b::Arf; abs_value = false, verbose = false)
    check_interval(a, b)

    maybe_abs = abs_value ? abs : identity

    if a == b
        return maybe_abs(p(a))
    end

    # Compute values at endpoints
    pa, pb = p(a), p(b)

    # Handle the signs not being the same when computing the absolute
    # value
    if abs_value
        sgn = _check_signs(pa, pb)
        if sgn == -1
            # The sign differs return zero
            verbose && @info "sign of endpoints differ - minimum is zero"

            return zero(pa)
        elseif sgn == 0
            # At least one of the enclosures contain zero, return the
            # minimum of the two enclosures
            m = min(maybe_abs(pa), maybe_abs(pb))
            Arblib.nonnegative_part!(m, m)

            verbose && @info "sign of endpoints undetermined - minimum on endpoints: $m"

            return m
        end
    end

    # Compute the minimum on the endpoints
    m_endpoints = min(maybe_abs(pa), maybe_abs(pb))

    abs_value && Arblib.nonnegative_part!(m_endpoints, m_endpoints)

    verbose && @info "minimum on endpoints: $m_endpoints"

    # Short circuit non-finite values and return NaN
    !isfinite(m_endpoints) && return Arb(NaN, prec = precision(m_endpoints))

    # Short circuit on constant or linear polynomial where minimum is
    # attained at endpoints
    Arblib.degree(p) <= 1 && return m_endpoints

    # Compute the roots of the derivative
    dp = Arblib.derivative(p)
    roots, flags = isolate_roots(dp, a, b)
    roots = Arb.(roots) # We want the roots as balls rather than tuples

    verbose &&
        @info "found $(length(roots)) intervals with possible roots of which $(sum(flags)) unique"

    # Evaluate the polynomial on all the roots and find the minimum
    values = p.(roots)

    # Handle sign crossings when computing the absolute value
    if abs_value && _check_signs(pa, pb, values) == -1
        verbose && @info "polynomial crosses zero - minimum is zero"
        return zero(m_endpoints)
    end

    m = copy(m_endpoints)
    for v in values
        Arblib.min!(m, m, maybe_abs(v))
    end

    # If computing with absolute value keep only non-negative part
    abs_value && Arblib.nonnegative_part!(m, m)

    verbose && @info "minimum with unrefined roots: $m"

    # Refine the unique roots where the minimum could occur
    count = 0
    for i in eachindex(roots)
        if flags[i] && !(values[i] > m)
            roots[i] = refine_root(dp, roots[i], strict = false)
            values[i] = p(roots[i])
            count += 1
        end
    end

    verbose && @info "refined $count roots"

    # Handle sign crossings when computing the absolute value
    if abs_value && _check_signs(pa, pb, values) == -1
        verbose && @info "polynomial crosses zero - minimum is zero"
        return zero(m_endpoints)
    end

    # Recompute minimum on the refined roots
    m = copy(m_endpoints)
    for v in values
        Arblib.min!(m, m, maybe_abs(v))
    end

    # If computing with absolute value keep only non-negative part
    abs_value && Arblib.nonnegative_part!(m, m)

    verbose && @info "minimum with refined roots: $m"

    return m
end

"""
    maximum_polynomial(p::ArbPoly, a::Arf, b::Arf; abs_value = false, verbose = false)

Compute an enclosure of the maximum of the polynomial `p` on the
interval `[a, b]`.

Takes the same arguments as [`extrema_polynomial`](@ref). The
algorithm is also the same except that it only looks for the maximum.
"""
function maximum_polynomial(p::ArbPoly, a::Arf, b::Arf; abs_value = false, verbose = false)
    check_interval(a, b)

    maybe_abs = abs_value ? abs : identity

    if a == b
        return maybe_abs(p(a))
    end

    # Compute the maximum on the endpoints
    m_endpoints = max(maybe_abs(p(a)), maybe_abs(p(b)))

    # If computing with absolute value keep only non-negative part
    abs_value && Arblib.nonnegative_part!(m_endpoints, m_endpoints)

    verbose && @info "maximum on endpoints: $m_endpoints"

    # Short circuit non-finite values and return NaN
    !isfinite(m_endpoints) && return Arb(NaN, prec = precision(m_endpoints))

    # Short circuit on constant or linear polynomial where maximum is
    # attained at endpoints
    Arblib.degree(p) <= 1 && return m_endpoints

    # Compute the roots of the derivative
    dp = Arblib.derivative(p)
    roots, flags = isolate_roots(dp, a, b)
    roots = Arb.(roots) # We want the roots as balls rather than tuples

    verbose &&
        @info "found $(length(roots)) intervals with possible roots of which $(sum(flags)) unique"

    # Evaluate the polynomial on all the roots and find the maximum
    values = p.(roots)

    m = copy(m_endpoints)
    for v in values
        Arblib.max!(m, m, maybe_abs(v))
    end

    # If computing with absolute value keep only non-negative part
    abs_value && Arblib.nonnegative_part!(m, m)

    verbose && @info "maximum with unrefined roots: $m"

    # Refine the unique roots where the maximum could occur
    count = 0
    for i in eachindex(roots)
        if flags[i] && !(maybe_abs(values[i]) < m)
            roots[i] = refine_root(dp, roots[i], strict = false)
            values[i] = p(roots[i])
            count += 1
        end
    end

    verbose && @info "refined $count roots"

    # Recompute maximum on the refined roots
    m = copy(m_endpoints)
    for v in values
        Arblib.max!(m, m, maybe_abs(v))
    end

    # If computing with absolute value keep only non-negative part
    abs_value && Arblib.nonnegative_part!(m, m)

    verbose && @info "maximum with refined roots: $m"

    return m
end

ext: jl


text:
# Create Roadway
roadway_opts = MergingRoadwayOptions(merging_point=1.4, lane_width=0.28)
roadway = build_roadway(roadway_opts)

# Create players
T = Float64
p = 3
model = BicycleGame(p=p)
n = model.n
m = model.m

# Define the horizon of the problem
N = 30 # N time steps
dt = 0.10 # each step lasts 0.1 second
probsize = ProblemSize(N,model) # Structure holding the relevant sizes of the problem

# Define the objective of each player
# We use a LQR cost
Q = [Diagonal(3*SVector{model.ni[i],T}([1., 0.8, 1., 1.])) for i=1:p] # Quadratic state cost
R = [Diagonal(0.3*ones(SVector{model.mi[i],T})) for i=1:p] # Quadratic control cost
# Desrired state
xf = [SVector{model.ni[1],T}([2.0, 0.05, 0.0, 0]),
      SVector{model.ni[2],T}([2.0, 0.05, 0.0, 0]),
      SVector{model.ni[2],T}([2.0, 0.05, 0.0, 0]),
      ]
# Desired control
uf = [zeros(SVector{model.mi[i],T}) for i=1:p]
# Objectives of the game
game_obj = Algames.GameObjective(Q,R,xf,uf,N,model)
radius = 0.20*ones(p)
μ = 10.0*ones(p)
add_collision_cost!(game_obj, radius, μ)

# Define the constraints that each player must respect
game_con = Algames.GameConstraintValues(probsize)
radius = 0.08
add_collision_avoidance!(game_con, radius)

# Define the initial state of the system
x0 = SVector{model.n,T}([
    0.2,     0.2,    0.2,
   -0.45,   -0.13,   0.13,
    0.45,    0.60,   0.60,
    0.15,    0.0,    0.0,
    ])

# Define the Options of the solver
opts = Options()
opts.ls_iter = 15
opts.outer_iter = 20
opts.inner_iter = 20
opts.ρ_0 = 1e0
opts.reg_0 = 1e-5
opts.α_dual = 1.0
opts.λ_max = 1.0*1e7
opts.ϵ_dyn = 1e-6
opts.ϵ_sta = 1e-6
opts.ϵ_con = 1e-6
opts.ϵ_opt = 1e-6
opts.regularize = true
# Define the game problem
prob = Algames.GameProblem(N,dt,x0,model,opts,game_obj,game_con)

# Solve the problem
newton_solve!(prob)

plot!(prob.model, prob.pdtraj.pr)
plot!(prob.stats)

players = Vector{Player{T}}(undef, p)
player_opts = HighwayPlayerOptions{model.ni[1],model.mi[1],T}()
players[1] = Player(model, roadway.lane[3])
players[2] = Player(model, roadway.lane[3])
players[3] = Player(model, roadway.lane[3])

# Create Scenario
sce = Scenario(model, roadway, players)

# function get_figure(prob::Algames.GameProblem, sce::Scenario, vis_sol::Visualizer, vis_init::Visualizer)
#     N = prob.probsize.N
#     model = prob.model
#
#     prob.opts.shift = 2^10
#     prob.opts.α_dual = 0.0
#     prob.opts.αx_dual = zeros(p)
#     prob.opts.f_init = x -> (rand(x) .- 0.5)
#     prob.opts.amplitude_init = 1.0
#
#     # Set tolerances
#     prob.opts.ϵ_opt = 1e-6
#     prob.opts.ϵ_dyn = 1e-6
#     prob.opts.ϵ_con = 1e-6
#     prob.opts.ϵ_sta = 1e-6
#
#     init = []
#     sol = []
#     plt_init = plot()
#     plt_sol = plot()
#     for k = 1:20
#         # Randomize the initialization
#         prob.opts.seed = k
#
#         # Get the initial guess (random controls + dynamics rollout)
#         prob.opts.outer_iter = 0
#         newton_solve!(prob)
#         init_traj = deepcopy(prob.pdtraj)
#
#         # Get solution trajectory
#         prob.opts.outer_iter = 20
#         newton_solve!(prob)
#         sol_traj = deepcopy(prob.pdtraj)
#
#         # Only keep the solution that converged
#         opt = mean(abs.(prob.stats.opt_vio[end].vio))
#         sta = mean(abs.(prob.stats.sta_vio[end].vio))
#         dyn = mean(abs.(prob.stats.dyn_vio[end].vio))
#         con = mean(abs.(prob.stats.con_vio[end].vio))
#         opts = prob.opts
#         if opt < opts.ϵ_opt && sta < opts.ϵ_sta && dyn < opts.ϵ_dyn && con < opts.ϵ_con
#             build_waypoint!(vis_init, sce.player, N, key=k)
#             build_waypoint!(vis_sol,  sce.player, N, key=k)
#             set_waypoint_traj!(vis_init, model, sce, init_traj.pr, key=k)
#             set_waypoint_traj!(vis_sol,  model, sce, sol_traj.pr, key=k)
#
#             plot!(model, init_traj.pr, plt=plt_init)
#             plot!(model, sol_traj.pr, plt=plt_sol)
#             push!(sol, sol_traj)
#             push!(init, init_traj)
#         end
#     end
#     set_state!(vis_init, model, sce, Algames.state(init[1].pr[1]))
#     set_state!(vis_sol,  model, sce, Algames.state(sol[1].pr[1]))
#     return init, sol
# end

# Initialize visualizers
# vis = Visualizer()
# open(vis)

# Visualize trajectories
set_scenario!(vis, sce)
set_env!(vis, VehicleState(1.3, 0.0, 0.0, 0.0))
set_camera_birdseye!(vis, height=5.0)

build_waypoint!(vis, sce.player, N, key=0)
set_waypoint_traj!(vis, model, sce, prob.pdtraj.pr, key=0)
set_traj!(vis, model, sce, prob.pdtraj.pr)

function display_ghost(vis::Visualizer, prob::Algames.GameProblem, sce::Scenario)
    N = prob.probsize.N
    for k ∈ [1,5,10,17,N]
        @show k
        vis_opts = PlayerVisualizationOptions{T}(α=0.3+0.2*k/10)
        set_player!(vis, sce.player, vis_opts=vis_opts, key=k)
        set_state!(vis, prob.model, sce, Algames.state(prob.pdtraj.pr[k]), key=k)
    end
    return nothing
end

display_ghost(vis, prob, sce)

# Verify the invertibility of the Jacobian of the KKT conditions.
residual_jacobian!(prob, prob.pdtraj)
null = nullspace(Matrix(prob.core.jac))
dim = size(null)[2]


text:
if isdefined(OffsetArrays, :centered)
    # Requires OffsetArrays v1.9
    # https://github.com/JuliaArrays/OffsetArrays.jl/pull/242
    # Originally from ImageFiltering https://github.com/JuliaImages/ImageFiltering.jl/pull/214
    OffsetArrays.centered(ax::Axis{name}) where name = Axis{name}(OffsetArrays.centered(ax.val))
    OffsetArrays.centered(a::AxisArray) = AxisArray(OffsetArrays.centered(a.data), OffsetArrays.centered.(AxisArrays.axes(a)))
end


text:
using Formulas
using Base.Test

## from Formulas
f = Formula(nothing, 0)
t = Terms(f)
@test t.intercept == false
@test t.response == false
@test t.terms == []
@test t.eterms == []

t = Terms(a ~ b)
@test t.terms == Term[:b]
@test t.eterms == [:a, :b]
@test t.intercept == true
@test t.response == true

## totally empty
t = Terms(Formula(nothing, 0))
@test t.response == false
@test t.intercept == false
@test t.terms == []
@test t.eterms == []

## empty RHS
t = Terms(y ~ 0)
@test t.intercept == false
@test t.terms == []
@test t.eterms == [:y]
t = Terms(y ~ -1)
@test t.intercept == false
@test t.terms == []

## intercept-only
t = Terms(y ~ 1)
@test t.response == true
@test t.intercept == true
@test t.terms == []
@test t.eterms == [:y]

## terms add
t = Terms(y ~ 1 + x1 + x2)
@test t.intercept == true
@test t.terms == Term[:x1, :x2]
@test t.eterms == [:y, :x1, :x2]

## implicit intercept behavior:
t = Terms(y ~ x1 + x2)
@test t.intercept == true
@test t.terms == Term[:x1, :x2]
@test t.eterms == [:y, :x1, :x2]

## no intercept
t = Terms(y ~ 0 + x1 + x2)
@test t.intercept == false
@test t.terms == Term[:x1, :x2]

@test t == Terms(y ~ -1 + x1 + x2) == Terms(y ~ x1 - 1 + x2) == Terms(y ~ x1 + x2 -1)

## can't subtract terms other than 1
@test_throws ErrorException Terms(y ~ x1 - x2)

t = Terms(y ~ x1 & x2)
@test t.terms == Term[:(x1 & x2)]
@test t.eterms == [:y, :x1, :x2]

## `*` expansion
t = Terms(y ~ x1 * x2)
@test t.terms == Term[:x1, :x2, :(x1 & x2)]
@test t.eterms == [:y, :x1, :x2]

## associative rule:
## +
t = Terms(y ~ x1 + x2 + x3)
@test t.terms == Term[:x1, :x2, :x3]

## &
t = Terms(y ~ x1 & x2 & x3)
@test t.terms == Term[:((&)(x1, x2, x3))]
@test t.eterms == [:y, :x1, :x2, :x3]

## distributive property of + and &
t = Terms(y ~ x1 & (x2 + x3))
@test t.terms == Term[:(x1&x2), :(x1&x3)]

t = Terms(y ~ (x1 + x2) & x3)
@test t.terms == Term[:(x1&x3), :(x2&x3)]

t = Terms(y ~ (x1 + x2) & (x3 + x4))
@test t.terms == Term[:(x1&x3), :(x1&x4), :(x2&x3), :(x2&x4)]

## three-way *
t = Terms(y ~ x1 * x2 * x3)
@test t.terms == Term[:x1, :x2, :x3,
                  :(x1&x2), :(x1&x3), :(x2&x3),
                  :((&)(x1, x2, x3))]
@test t.eterms == [:y, :x1, :x2, :x3]


text:
# Autogenerated wrapper script for libblastrampoline_jll for i686-w64-mingw32
export libblastrampoline

JLLWrappers.@generate_wrapper_header("libblastrampoline")
JLLWrappers.@declare_library_product(libblastrampoline, "libblastrampoline.dll")
function __init__()
    JLLWrappers.@generate_init_header()
    JLLWrappers.@init_library_product(
        libblastrampoline,
        "bin\\libblastrampoline.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
    @static if VERSION < v"1.7.0-DEV.641"
        ccall((:lbt_forward, libblastrampoline), Int32, (Cstring, Int32, Int32),
              Libdl.dlpath(Base.libblas_name) , 1, 0)
    end

end  # __init__()


text:
""" Serialize raw and semantic flow graphs.
"""
module Serialization
export parse_raw_graphml, parse_raw_graph_json,
  parse_semantic_graphml, parse_semantic_graph_json,
  read_raw_graphml, read_raw_graph_json,
  read_semantic_graphml, read_semantic_graph_json

using Compat

using Catlab, Catlab.WiringDiagrams
import Catlab.WiringDiagrams: convert_from_graph_data, convert_to_graph_data
using ..Doctrine
using ..RawFlowGraphs

# Raw flow graphs
#################

parse_raw_graphml(xml) = parse_graphml(RawNode, RawPort, Nothing, xml)
parse_raw_graph_json(json) = parse_json_graph(RawNode, RawPort, Nothing, json)
read_raw_graphml(filename) = read_graphml(RawNode, RawPort, Nothing, filename)
read_raw_graph_json(filename) = read_json_graph(RawNode, RawPort, Nothing, filename)

function convert_from_graph_data(::Type{RawNode}, data::AbstractDict)
  annotation = pop!(data, "annotation", nothing)
  annotation_index = pop!(data, "annotation_index", nothing)
  annotation_kind_str = pop!(data, "annotation_kind", nothing)
  annotation_kind = isnothing(annotation_kind_str) ? FunctionAnnotation :
    convert(RawNodeAnnotationKind, annotation_kind_str)
  RawNode(data, annotation, annotation_index, annotation_kind)
end

function convert_from_graph_data(::Type{RawPort}, data::AbstractDict)
  annotation = pop!(data, "annotation", nothing)
  annotation_index = pop!(data, "annotation_index", nothing)
  value = pop!(data, "value", nothing)
  RawPort(data, annotation, annotation_index, value)
end

# Semantic flow graphs
######################

parse_semantic_graphml(xml) = parse_graphml(
  Union{Monocl.Hom,Nothing}, Union{Monocl.Ob,Nothing}, Nothing, xml)
parse_semantic_graph_json(json) = parse_json_graph(
  Union{Monocl.Hom,Nothing}, Union{Monocl.Ob,Nothing}, Nothing, json)
read_semantic_graphml(filename) = read_graphml(
  Union{Monocl.Hom,Nothing}, Union{Monocl.Ob,Nothing}, Nothing, filename)
read_semantic_graph_json(filename) = read_json_graph(
  Union{Monocl.Hom,Nothing}, Union{Monocl.Ob,Nothing}, Nothing, filename)

function convert_from_graph_data(::Type{Monocl.Ob}, data::AbstractDict)
  parse_json_sexpr(Monocl, data["ob"]; symbols=false)
end
function convert_from_graph_data(::Type{Monocl.Hom}, data::AbstractDict)
  parse_json_sexpr(Monocl, data["hom"]; symbols=false)
end
function convert_from_graph_data(::Type{Union{Monocl.Ob,Nothing}}, data::AbstractDict)
  isempty(data) ? nothing : convert_from_graph_data(Monocl.Ob, data)
end
function convert_from_graph_data(::Type{Union{Monocl.Hom,Nothing}}, data::AbstractDict)
  isempty(data) ? nothing : convert_from_graph_data(Monocl.Hom, data)
end

function convert_to_graph_data(expr::Monocl.Ob)
  Dict(
    "ob" => to_json_sexpr(expr),
    "label" => Dict(
      "text" => text_label(expr)
    )
  )
end
function convert_to_graph_data(expr::Monocl.Hom)
  Dict(
    "hom" => to_json_sexpr(expr),
    "label" => Dict(
      "text" => text_label(expr)
    )
  )
end

function convert_from_graph_data(::Type{MonoclElem}, data::AbstractDict)
  ob = haskey(data, "ob") ?
    parse_json_sexpr(Monocl, data["ob"]; symbols=false) : nothing
  value = get(data, "value", nothing)
  MonoclElem(ob, value)
end

function convert_to_graph_data(elem::MonoclElem)
  data = Dict{String,Any}()
  if !isnothing(elem.ob)
    data["ob"] = to_json_sexpr(elem.ob)
  end
  if !isnothing(elem.value)
    data["value"] = elem.value
  end
  data
end

""" Short, human-readable, plain text label for Moncl expression.
"""
text_label(expr::Monocl.Ob) = string(expr)
text_label(expr::Monocl.Hom) = string(expr)
text_label(expr::Monocl.Hom{:coerce}) = "coerce"
text_label(expr::Monocl.Hom{:construct}) = string(codom(expr))

end

max_stars_repo_name: IBM/semanticflowgraph

max_issues_repo_name: epatters/semanticflowgraph

max_forks_repo_name: IBM/semanticflowgraph


text:
module TwoAgents

using CompEcon
using JLD
using NLsolve

include("../ModelAbstraction.jl")
using .ModelAbstraction

immutable SolutionForState
    c1::Float64
    c2::Float64
    J::Float64
    result::NLsolve.SolverResults{Float64}
end

immutable BCFL21
    agent1::EpsteinZinAgent
    agent2::EpsteinZinAgent
    production1::CESProducer
    production2::CESProducer
    ac1::AdjCost
    ac2::AdjCost
    exog::ConstantVolatility{2}
    gp::Matrix{Float64}
    ss::DefaultStateSpace{3,1,Float64,4}
end

function BCFL21(;ρ1::Real=-1.0 , α1::Real=-9.0, β1::Real=0.999,
                 ρ2::Real=-1.0 , α2::Real=-9.0, β2::Real=0.999,
                 η1::Real=1.0/3, ν1::Real=.1, η2::Real=1.0/3, ν2::Real=.1,
                 γ::Real=.9, δ::Real=0.025, σ1::Real=0.001, σ2::Real=0.001,
                 nϵ1::Int=2, nϵ2::Int=2,
                 nk1::Int=10, nk2::Int=10, nU::Int=8, nξ::Int=5)

    # Agents
    agent1 = EpsteinZinAgent(ρ1, α1, β1)
    agent2 = EpsteinZinAgent(ρ2, α2, β2)

    # Producers
    production1 = CESProducer(η1, ν1)
    production2 = CESProducer(η2, ν2)

    # AdjustmentCosts
    ac1 = AdjCost(δ, .75)
    ac2 = AdjCost(δ, .75)

    # exog
    exog = ConstantVolatility(γ, σ1, σ2, nϵ1, nϵ2)

    # state space
    kmin, kmax = 10., 35.
    ξmin, ξmax = .97, 1.03
    Umin = .1
    Umax = produce(production1, kmax, 1.) + produce(production2, kmax, ξmax)
    basis = Basis(SplineParams(collect(linspace(kmin, kmax, nk1)), 0, 2),
                  SplineParams(collect(linspace(kmin, kmax, nk2)), 0, 2),
                  SplineParams(collect(linspace(Umin , Umax, nU)), 0, 2),
                  LinParams(collect(linspace(ξmin, ξmax, nξ)), 0))

    grid, (_k1grid, _k2grid, _Ugrid, _ξgrid) = nodes(basis)

    lξp = (2γ-1)*log(grid[:, 4]) .+ (σ2*exog.ϵ[:, 2]' - σ1*exog.ϵ[:, 1]')
    lξp = clamp(lξp, log(_ξgrid[1]), log(_ξgrid[end]))'
    lgp = (γ-1)*log(grid[:, 4]) .+ σ1*exog.ϵ[:, 1]'

    _en = tuple(map(EndogenousState, (_k1grid, _k2grid, _Ugrid),
                                     (basis[1], basis[2], basis[2]))...)
    _ex = (ExogenousState(_ξgrid, exp(lξp), basis[4]),)

    ss = DefaultStateSpace(_en, _ex, grid, grid', basis)

    # package it up!
    return BCFL21(agent1, agent2, production1, production2, ac1, ac2, exog,
                  exp(lgp)', ss)
end

_unpack_params(m::BCFL21) = vcat(_unpack(m.agent1)..., _unpack(m.agent2)...,
                             _unpack(m.ac1)..., _unpack(m.ac2)...)

function compute_residuals!(bcfl::BCFL21, state::Vector{Float64}, J::Float64,
                            dJU::Float64, dJk1::Float64, dJk2::Float64,
                            coefs::Vector{Float64},
                            gp::Vector{Float64}, ξp::Vector{Float64},
                            guess::Vector, resid::Vector)

    # Unpack parameters.
    ρ1, α1, β1, ρ2, α2, β2, δ1, η1, δ2, η2 = _unpack_params(bcfl)
    Π = bcfl.exog.Π

    # extract guess and state
    I1, I2 = guess[1:2]
    Up = guess[3:end]
    k1, k2, U, ξ = state

    # Force the constraints I want!!! RAWRRRRRR
    if any(Up .< 0)
        # @show Up statep
        resid[:] = 1e6
        return NaN, NaN, NaN
    end
    if -I1 > k1
        resid[:] = 1e6
        return NaN, NaN, NaN
    end
    if -I2 > k2
        resid[:] = 1e6
        return NaN, NaN, NaN
    end

    # Derivative of Adjustment costs
    # NOTE: I am assuming that Γ_i(k_i, I_i) = (1 - δ_i) k_i + I_i
    dΓ1_dI1 = 1.0  # - _dIac(bcfl.ac1, k1, I1)
    dΓ2_dI2 = 1.0  # - _dIac(bcfl.ac2, k2, I2)
    dΓ1_dk1 = (1 - δ1)  # - _dkac(bcfl.ac1, k1, I1)
    dΓ2_dk2 = (1 - δ2)  # - _dkac(bcfl.ac1, k1, I1)

    # MPK for agent 1. Note we set the `z` arg to one b/c we scaled by z1
    df1dk1 = f_k(bcfl.production1, k1, 1.0)
    df2dk2 = f_k(bcfl.production2, k2, ξ)

    # Get tomorrow's state
    k1p = ((1 - δ1)*k1 + I1) ./ gp  # - _ac(bcfl.ac1, k1, I1)) ./ gp
    k2p = ((1 - δ2)*k2 + I2) ./ gp  # - _ac(bcfl.ac2, k2, I2)) ./ gp
    statep = [k1p k2p Up ξp]

    # Evaluate value function and its partials at next period's state. This is
    # done at the same time so we only have to compute 7 basis matrices instead
    # of the full 16 we will be using.
    out = funeval(coefs, bcfl.ss.basis, statep,
                  [0 0 0 0; 1 0 0 0; 0 1 0 0; 0 0 1 0])
    # out will be (4×1×length(gp))
    Jp    = out[:, 1, 1]
    dJpk1 = out[:, 1, 2]
    dJpk2 = out[:, 1, 3]
    dJpU  = out[:, 1, 4]

    # Evaluate all expectations
    if any(Jp .< 0)
        @show Jp statep
    end

    μ1 = dot(Π, (gp .* Jp).^(α1))^(1.0/α1)
    μ2 = dot(Π, (gp .* Up).^(α2))^(1.0/α2)
    EV11 = dot(Π, (gp .* Jp).^(α1 - 1.) .* dJpk1)
    EV12 = dot(Π, (gp .* Jp).^(α1 - 1.) .* dJpk2)

    # TODO: Remove all shows once we are more convinced it is working
    # get consumption
    c1_num   = dJk1 - β1 * J^(1-ρ1) * μ1^(ρ1-α1) * dΓ1_dk1 * EV11
    c1_num = c1_num > 0 ? c1_num : abs(c1_num)
    c1_denom = J.^(1-ρ1) * (1-β1) * df1dk1
    c1       = (c1_num/c1_denom)^(1/(ρ1-1))

    c2_num   = dJk2 - β1 * J^(1-ρ1) * μ1^(ρ1-α1) * dΓ2_dk2 * EV12
    c2_num = c2_num > 0 ? c2_num : abs(c2_num)
    c2_denom = -dJU * U.^(1-ρ1) * (1-β2) * df2dk2
    c2       = (c2_num/c2_denom)^(1/(ρ2-1))

    # c2 = ((-dJU * U^(1-ρ2) * (1-β2)) / (J^(1-ρ1) * (1-β1) * c1^(ρ1-1)))^(1/(ρ2-1))

    # I residual
    lhsI1 = EV11 .* dΓ1_dI1
    rhsI1 = EV12 .* dΓ2_dI2
    resid[1] = lhsI1 - rhsI1

    # use c1, c2, I1 in budget constraint to get residual for I2
    rhsI2 = c1 + c2 + I1 + I2
    lhsI2 = produce(bcfl.production1, k1, 1.) + produce(bcfl.production2, k2, ξ)
    resid[2] = rhsI2 - lhsI2

    # U residual
    nUp = length(Up)
    for i=1:nUp
        lhs = -dJU * U.^(1-ρ2) * β2 * μ2.^(ρ2-α2) .* gp[i]^α2 .* Up[i].^(α2-1)
        # TODO: check rhs to make sure that we have negative signs in the right
        #       place.
        rhs = J^(1-ρ1) * β1 * μ1.^(ρ1-α1) * gp[i]^α1 * Jp[i]^(α1-1) * (-dJpU[i])
        resid[i+2] = lhs - rhs
    end

    J = utility(bcfl.agent1, c1, μ1)

    # we are safe to return whatever we want because NLsolve just cares that
    # we update resid in place and completely ignores the return value of this
    # function
    return c1, c2, J
end

function initial_coefs(bcfl::BCFL21, bs::BasisStructure)
    #=
    Here we just manufacturing an initial guess that should be in the ballpark of
    the actual domain and has partial derivatives with the right sign. Specifically,
    we are imposing that:

    - `dJ/dU < 0`
    - `dJ/dk1 > 0`
    - `dJ/dk2 > 0`

    everywhere
    =#

    k1 = bcfl.ss.grid[:, 1]
    k2 = bcfl.ss.grid[:, 2]
    U = bcfl.ss.grid[:, 3]
    ξ = bcfl.ss.grid[:, 4]

    y = produce(bcfl.production1, k1, 1.) + produce(bcfl.production2, k2, ξ)

    # Agent 1's value starts at 1/2 of agent 2's in the "opposite" state (hence the reverse)
    # then gets 1/2 of production added to it.
    J = (reverse(U) + y) / 2.0
    coefs = CompEcon.get_coefs(bcfl.ss.basis, bs, J)

    return coefs
end

function create_jld_output_file()
    out_dir = joinpath(dirname(@__FILE__), "output")
    !isdir(out_dir) && mkdir(out_dir)

    matches = filter(x->x != nothing,
                     map(x->match(r"(\d+$)", split(x, ".")[1]),
                     filter(x->endswith(x, ".jld"),
                     readdir(out_dir))))
    if isempty(matches)
        num = 1
    else
        num = maximum(Int[parse(Int, m[1]) for m in matches]) + 1
    end

    fn = joinpath(out_dir, "results$(num).jld")
    jldopen(fn, "w") do f
        nothing
    end
    return fn
end

function brutal_solution(bcfl::BCFL21; tol=1e-4, maxiter=500)

    # Unpack parameters.
    dist, iter = 10., 0
    ρ1, α1, β1, ρ2, α2, β2, δ1, η1, δ2, η2 = _unpack_params(bcfl)

    # BasisStructure for interpolation
    bs = BasisStructure(bcfl.ss.basis, Direct(), bcfl.ss.grid,
                        [0 0 0 0; 1 0 0 0; 0 1 0 0; 0 0 1 0])

    coefs = initial_coefs(bcfl, bs)
    Nϵ = size(bcfl.gp, 1)

    # extract grids to ease notation
    k1_grid = bcfl.ss.grid[:, 1]
    k2_grid = bcfl.ss.grid[:, 2]
    U_grid = bcfl.ss.grid[:, 3]
    ξ_grid = bcfl.ss.grid[:, 4]

    i1 = (δ1-1).*k1_grid .+ (k1_grid .+ k2_grid)./(ξ_grid + 1)
    i2 = (k1_grid.*ξ_grid + k2_grid.*(δ2*ξ_grid + δ2 - 1.)) ./ (ξ_grid + 1)
    prev_soln(i::Int) = [i1[i]; i2[i]; fill(U_grid[i], Nϵ)]

    result_fn = create_jld_output_file()

    local all_sfs  # declare out local so we can return it outside while
    while dist > tol && iter < maxiter
        # Increment counter
        iter += 1

        stuff = funeval(coefs, bs, [0 0 0 0; 1 0 0 0; 0 1 0 0; 0 0 1 0])
        J_all    = stuff[:, 1, 1]
        dJk1_all = stuff[:, 1, 2]
        dJk2_all = stuff[:, 1, 3]
        dJU_all  = stuff[:, 1, 4]

        # function to solve state i
        function ssi(i)
            # prep function for nlsolve
            guess = prev_soln(i)
            resid = similar(guess)
            state = bcfl.ss.grid_transpose[:, i]
            J = J_all[i]
            dJk1 = dJk1_all[i]
            dJk2 = dJk2_all[i]
            dJU = dJU_all[i]
            gp = bcfl.gp[:, i]
            ξp = bcfl.ss.exog[1].gridp[:, i]

            # TODO: performance will probably be significantly better if we can outsourse
            #       most of the guts of compute_residuals! elsewhere. Othwerwise we
            #       will be recompiling all the guts for every i on every iteration :yuck:
            function f!(x, fvec)
                 compute_residuals!(bcfl, state, J, dJU, dJk1, dJk2, coefs, gp,
                                    ξp, x, fvec)
            end

            # lb = [-state[1], -state[2], 0., 0., 0., 0.]
            # ub = [Inf, Inf, Inf, Inf, Inf, Inf]
            soln = nlsolve(f!, guess, factor=.025)

            if maxabs(soln.initial_x - guess) > 1e-12
                println("Sebastian changed our guess on state $i")
            end

            # now call the function one more time to get c1, c2, J. Use
            # the already allocated guess vector as the resid buffer that will
            # be changed but ignored by us
            out = SolutionForState(f!(soln.zero, guess)..., soln)

            if soln.f_converged
                print("  $i converged.")
            else
                println("\n$i failed. residual norm: $(soln.residual_norm)")
            end
            return out
        end

        # if iter == 1
        #     all_sfs = jldopen(result_fn, "r+") do f
        #         read(f, "all_sfs_1")
        #     end
        # else

            all_sfs = pmap(ssi, 1:size(bcfl.ss.grid, 1))
            jldopen(result_fn, "r+") do f
                write(f, "all_sfs_$iter", all_sfs)
            end
        # end

        prev_soln(i::Int) = all_sfs[i].result.zero

        new_J = Float64[s.J for s in all_sfs]
        dist = maxabs(J_all - new_J)

        println("Hallelujah! Finished iteration $iter. dist: $dist")

        # Update coefficients
        coefs = CompEcon.get_coefs(bcfl.ss.basis, bs, new_J)
    end

    all_sfs
end

end  # module

# include("common.jl")
# addprocs(3)
# @everywhere include("common.jl")
# bcfl = TwoAgents.BCFL21();
# out = TwoAgents.brutal_solution(bcfl; maxiter=2)

max_stars_repo_name: NYUEcon/GrowthModels

max_issues_repo_name: NYUEcon/GrowthModels

max_forks_repo_name: NYUEcon/GrowthModels


text:
# ------------------------------------------------------------------------------------------
# # IMDB Movie Review Sentiment Analysis Demo
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## Setup
# 1-Adds required packages to Julia.
# 2-Loads the data and a pretrained model.
# ------------------------------------------------------------------------------------------

include("imdb.jl")

# ------------------------------------------------------------------------------------------
# ## Sample Data
# The model was trained using 25000 movie reviews such as the following (shift-ENTER to see
# random example)
# Each review was tokenized, lowercased, truncated to max 150 words and a 30,000 word
# vocabulary.
# ------------------------------------------------------------------------------------------

r = rand(1:length(xtrn))
println(reviewstring(xtrn[r],ytrn[r]))
flush(STDOUT)

# ------------------------------------------------------------------------------------------
# ## Test
# We test the model on 25000 never before seen reviews on the test set (shift-ENTER to see
# random example)
# The test accuracy is around 86%
# ------------------------------------------------------------------------------------------

r = rand(1:length(xtst))
println(reviewstring(xtst[r],ytst[r]))
println("\nModel prediction: "*predictstring(xtst[r]))
flush(STDOUT)

# ------------------------------------------------------------------------------------------
# ## User Input
# In this cell you can enter your own review and let the model guess the sentiment
# ------------------------------------------------------------------------------------------

userinput = readline(STDIN)
words = split(lowercase(userinput))
ex = [get!(imdbdict,wr,UNK) for wr in words]
ex[ex.>MAXFEATURES]=UNK
println("\nModel prediction: "*predictstring(ex))

max_stars_repo_name: grenkoca/JuliaTutorials

max_issues_repo_name: grenkoca/JuliaTutorials

max_forks_repo_name: grenkoca/JuliaTutorials


text:
using SparseArrays

"""
    struct DeterministicPolicy <: AbstractPolicy
        table::Vector{Int}
        nactions::Int
    end

The action to be adopted is stored in `table`.
"""
struct DeterministicPolicy <: AbstractPolicy
    table::Vector{Int}
    nactions::Int
end

(p::DeterministicPolicy)(s::Int) = p.table[s]
(p::DeterministicPolicy)(s::Int, ::Val{:dist}) = SparseVector(p.nactions, [p(s)], [1.0])
(p::DeterministicPolicy)(s::Int, a::Int) = p.table[s] == a ? 1.0 : 0.0
update!(p::DeterministicPolicy, s::Int, a::Int) = p.table[s] = a


text:
using ReachabilityAnalysis
using Plots

#
# w = 2pi/T
#
A = [0.0 1; -4π 0]
X0 = BallInf([0.0, 10.0], 0.01)
X = Universe(2)
B = hcat([1.0, -1.0])
U = ConstantInput(linear_map(B, Singleton([-3.0])))

P = @ivp(x' = A*x + u, x ∈ X, u ∈ U, x(0) ∈ X0)

# flowpipe (tight approximation)
δ = 0.00001
T = 0.05

alg_fw = GLGM06(δ=δ, approx_model=Forward())
sol_fw = solve(P; T=T, alg=alg_fw);

# discretizations
Ω0_chull = ConvexHullArray(map(set, sol_fw))
Ω0_fw = discretize(P, T, Forward()).x0
Ω0_ch = discretize(P, T, CorrectionHull(order=4)).x0

# random trajectories
x0s = vcat(sample(X0, 0, include_vertices=true),
           sample(X0, 50, sampler=LazySets.FaceSampler(1)))

P_orbit = [@ivp(x' = A*x + u, x ∈ X, u ∈ U, x(0) ∈ x0) for x0 in x0s]
sol_orbit = [solve(Pi, T=T, alg=ORBIT(δ=T/20.)) for Pi in P_orbit]

fig = plot(legend=:bottomright, legendfontsize=15, xticks=[], yticks=[])

plot!(fig, Ω0_ch, alpha=0, color=:blue, linecolor=:blue, lw=2, linealpha=1)
plot!(fig, Ω0_chull, 1e-4, alpha=0.3, color=:purple, linecolor=:purple, lw=2, linealpha=1)

plot!(fig, X0, c=:yellow, alpha=1., lab="")
plot!(fig, sol_fw[1501:2000], vars=(1, 2), lw=0.1, c=:orange, alpha=0.8, lab="")
plot!(fig, sol_fw[end], vars=(1, 2), c=:green, alpha=0.8, lab="")
[plot!(fig, s, vars=(1, 2), c=:magenta, seriestype=:path, marker=:none) for s in sol_orbit]

savefig(fig, "intro_approximate.pdf")


text:
module TestDelete
using Test
using Accessors
using StaticArrays

@testset "test delete" begin
    @testset "function" begin
        @test delete( (a=1, b=2, c=3), @optic(_.a) ) == (b=2, c=3)
        @test delete( (a=1, b=2, c=3), @optic(_.xxxx) ) === (a=1, b=2, c=3)
        @test delete( (a=1, b=(c=2, d=3)), @optic(_.b.c) ) == (a=1, b=(d=3,))
        @test delete( (1,2,3), @optic(last(_)) ) == (1, 2)
        @test delete( SVector(1,2,3), @optic(last(_)) ) === SVector(1, 2)
        let A = [1,2,3]
            @test delete(A, @optic(_[2])) == [1, 3]
            @test_throws Exception delete(A, @optic(_[2, 2]))
            @test_throws BoundsError delete(A, @optic(_[10]))
            @test delete(A, @optic(_[end])) == [1, 2]
            @test A == [1, 2, 3]  # not changed
        end
        let A = Dict("a" => 1, "b" => 2, "c" => 3)
            @test delete(A, @optic(_["a"])) == Dict("b" => 2, "c" => 3)
            @test delete(A, @optic(_["xxxx"])) == A  # follows Base.delete! behavior
            @test A == Dict("a" => 1, "b" => 2, "c" => 3)  # not changed
        end
        @test delete( (a=1, b=(2, 3, 4)), @optic(first(_.b)) ) == (a=1, b=(3, 4))
        @test delete( "path/to/file", @optic(basename(_)) ) == "path/to"
        @test delete( "path/to/file", @optic(dirname(_)) ) == "file"
    end

    @testset "macro" begin
        x = (a=1, b=2, c=3)
        @test @delete(x.c) == (a=1, b=2)
        x = (a=1, b=(c=2, d=3))
        @test @delete(x.b.c) == (a=1, b=(d=3,))
        x = [1, 2, 3]
        @test @delete(x[2]) == [1, 3]
        @test_throws BoundsError @delete(x[10])
    end
end

end


text:
# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

#=
(https://en.wikipedia.org/wiki/Line-plane_intersection)
=#
function intersection(f, s::Segment{3,T}, p::Plane{T}) where {T}
  sᵥ = coordinates.(vertices(s))
  pₒ = coordinates(origin(p))
  n  = normal(p)

  # calculate components
  ln = (sᵥ[2] - sᵥ[1]) ⋅ n

  # if ln is zero, the segment is parallel to the plane
  if isapprox(ln, zero(T), atol=atol(T))
    # if the numerator is zero, the segment is coincident
    if isapprox((pₒ - sᵥ[1]) ⋅ n, zero(T), atol=atol(T))
      return @IT OverlappingSegmentPlane s f
    else
      return @IT NoIntersection nothing f
    end
  else
    # calculate the segment parameter
    λ = ((pₒ - sᵥ[1]) ⋅ n) / ln

    # if λ is approximately 0 or 1, set as so to prevent any domain errors
    if isapprox(λ, zero(T), atol=atol(T))
      return @IT TouchingSegmentPlane s(zero(T)) f
    elseif isapprox(λ, one(T), atol=atol(T))
      return @IT TouchingSegmentPlane s(one(T)) f
    end

    # if λ is out of bounds for the segment, then there is no intersection
    if (λ < zero(T)) || (λ > one(T))
      return @IT NoIntersection nothing f
    else
      return @IT CrossingSegmentPlane s(λ) f
    end
  end
end


text:
"""
    Som

Structure to hold all data of a trained SOM.

# Fields:
- `codes::Array{Float64,2}`: 2D-array of codebook vectors. One vector per row
- `xdim::Int`: number of neurons in x-direction
- `ydim::Int`: number of neurons in y-direction
- `numCodes::Int`: total number of neurons
- `grid::Array{Float64,2}`: 2D-array of coordinates of neurons on the map
          (2 columns (x,y)] for rectangular and hexagonal maps
           3 columns (x,y,z) for spherical maps)
"""
mutable struct Som
    codes::Matrix{Float64}
    xdim::Int
    ydim::Int
    numCodes::Int
    grid::Matrix{Float64}

    Som(;
        codes::Matrix{Float64},
        xdim::Int,
        ydim::Int,
        numCodes::Int = xdim * ydim,
        grid::Matrix{Float64},
    ) = new(codes, xdim, ydim, numCodes, grid)
end

"""
    LoadedDataInfo

The basic structure for working with loaded data, distributed amongst workers. In completeness, it represents a dataset as such:

- `val` is the "value name" under which the data are saved in processes. E.g.
  `val=:foo` means that there is a variable `foo` on each process holding a
  part of the matrix.
- `workers` is a list of workers (in correct order!) that hold the data
  (similar to `DArray.pids`)
"""
struct LoadedDataInfo
    val::Symbol
    workers::Array{Int64}
    LoadedDataInfo(val, workers) = new(val, workers)
end


text:
# This file is a part of BATBase.jl, licensed under the MIT License (MIT).

using BATBase
using Test

using Distributions

@testset "abstract_density" begin
    dist = MvNormal(float([1, 2, 3]))

    @test @inferred(MyMvDistributionDensity(dist)) isa AbstractDistributionDensity
    @test @inferred(MyMvDistributionDensity(dist)) isa DistLikeDensity
    @test @inferred(MyMvDistributionDensity(dist)) isa AbstractDensity

    density = MyMvDistributionDensity(dist)
    v = rand(dist)

    @test @inferred(convert(Distribution, density)) == dist
    @test @inferred BATBase.eval_logval_unchecked(density, v) == logpdf(dist, v)
    @test @inferred(varshape(density)) == ArrayShape{Real}(3)
    @test @inferred(totalndof(density)) == 3
    @test @inferred(sampler(density)) == sampler(dist)

    @test @inferred(convert(AbstractDistributionDensity, dist)) == density
    @test @inferred(convert(DistLikeDensity, dist)) == density
    @test @inferred(convert(AbstractDensity, dist)) == density

    @test @inferred(logdensityof(density, v)) == logpdf(dist, v)
    @test @inferred(logdensityof(density)(v)) == logpdf(dist, v)
end


text:
# Include the source file
using Test
using DynamicalNetworks
using Graphs

# Construction
net = ODENetwork()
net = ODENetwork(Lorenz[])
net = ODENetwork([Lorenz() for i = 1:10])
net = ODENetwork(Lorenz)
net = ODENetwork(Rossler)
mat = rand(3, 3)
net = ODENetwork(Lorenz, mat + mat')
net = ODENetwork(Lorenz, star_graph(10))
net = ODENetwork(Lorenz, cycle_graph(10))
net = ODENetwork(Lorenz, barabasi_albert(10, 3, 2))

# Add nodes
net = ODENetwork(Lorenz)
addnode!(net, Lorenz())
addnode!(net)
addnode!(net, Lorenz(σ = 5.0))

# Remove nodes
mat = rand(3, 3)
net = ODENetwork(Lorenz, mat + mat')
remnode!(net)

# Add edges
net = ODENetwork()
foreach(i -> addnode!(net), 1:5)
addedge!(net, 1, 2, 10.0)

# Solve network
net = ODENetwork(Lorenz, path_graph(3))
prob = netproblem(net, (0.0, 100.0), collectstates(net), ones(10))
sol = solvenet!(prob)


text:
#
# Copyright (c) 2021 Tobias Thummerer, Lars Mikelsons, Josef Kircher
# Licensed under the MIT license. See LICENSE file in the project root for details.
#

###############
# Prepare FMU #
###############

myFMU = fmi3Load("Feedthrough", "ModelicaReferenceFMUs", "0.0.14")
inst = fmi3InstantiateCoSimulation!(myFMU; loggingOn=false)
@test inst != 0

@test fmi3EnterInitializationMode(inst) == 0
# TODO check the value references used and adjust them
# TODO add clocks

float32ValueReferences = ["real_discrete_out", "Float_32_continuous_output"]
float64ValueReferences = ["real_discrete_in", "real_tunable_param"]
integerValueReferences = ["int_in", "int_out"]
booleanValueReferences = ["bool_in", "bool_out"]
stringValueReferences = ["string_param", "string_param"]
binaryValueReferences = ["binary_in", "binary_out"]

#########################
# Testing Single Values #
#########################

rndReal = 100 * rand()
rndInteger = round(Integer, 100 * rand())
rndBoolean = rand() > 0.5
rndString = Random.randstring(12)

cacheReal = 0.0
cacheInteger = 0
cacheBoolean = false
cacheString = ""

# TODO not contained in the FMU
# @test fmi3SetFloat32(inst, float32ValueReferences[1], Float32(rndReal)) == 0
# @test fmi3GetFloat32(inst, float32ValueReferences[1]) == rndReal
# @test fmi3SetFloat32(inst, realValueReferences[1], -rndReal) == 0
# @test fmi3GetFloat32(inst, realValueReferences[1]) == -rndReal

@test fmi3SetFloat64(inst, float64ValueReferences[1], rndReal) == 0
@test fmi3GetFloat64(inst, float64ValueReferences[1]) == rndReal
@test fmi3SetFloat64(inst, float64ValueReferences[1], -rndReal) == 0
@test fmi3GetFloat64(inst, float64ValueReferences[1]) == -rndReal

# TODO not contained in the FMU
# @test fmi3SetInt8(inst, integerValueReferences[1], rndInteger) == 0
# @test fmi3GetInt8(inst, integerValueReferences[1]) == rndInteger
# @test fmi3SetInt8(inst, integerValueReferences[1], -rndInteger) == 0
# @test fmi3GetInt8(inst, integerValueReferences[1]) == -rndInteger

# @test fmi3SetInt16(inst, integerValueReferences[2], rndInteger) == 0
# @test fmi3GetInt16(inst, integerValueReferences[2]) == rndInteger
# @test fmi3SetInt16(inst, integerValueReferences[2], -rndInteger) == 0
# @test fmi3GetInt16(inst, integerValueReferences[2]) == -rndInteger

@test fmi3SetInt32(inst, integerValueReferences[1], Int32(rndInteger)) == 0
@test fmi3GetInt32(inst, integerValueReferences[1]) == rndInteger
@test fmi3SetInt32(inst, integerValueReferences[1], Int32(-rndInteger)) == 0
@test fmi3GetInt32(inst, integerValueReferences[1]) == -rndInteger

# TODO not contained in the FMU
# @test fmi3SetInt64(inst, integerValueReferences[1], rndInteger) == 0
# @test fmi3GetInt64(inst, integerValueReferences[1]) == rndInteger
# @test fmi3SetInt64(inst, integerValueReferences[4], -rndInteger) == 0
# @test fmi3GetInt64(inst, integerValueReferences[4]) == -rndInteger

# @test fmi3SetUInt8(inst, integerValueReferences[5], rndInteger) == 0
# @test fmi3GetUInt8(inst, integerValueReferences[5]) == rndInteger
# @test fmi3SetUInt8(inst, integerValueReferences[5], -rndInteger) == 0
# @test fmi3GetUInt8(inst, integerValueReferences[5]) == -rndInteger

# @test fmi3SetUInt16(inst, integerValueReferences[6], rndInteger) == 0
# @test fmi3GetUInt16(inst, integerValueReferences[6]) == rndInteger
# @test fmi3SetUInt16(inst, integerValueReferences[6], -rndInteger) == 0
# @test fmi3GetUInt16(inst, integerValueReferences[6]) == -rndInteger

# @test fmi3SetUInt32(inst, integerValueReferences[7], rndInteger) == 0
# @test fmi3GetUInt32(inst, integerValueReferences[7]) == rndInteger
# @test fmi3SetUInt32(inst, integerValueReferences[7], -rndInteger) == 0
# @test fmi3GetUInt32(inst, integerValueReferences[7]) == -rndInteger

# @test fmi3SetUInt64(inst, integerValueReferences[8], rndInteger) == 0
# @test fmi3GetUInt64(inst, integerValueReferences[8]) == rndInteger
# @test fmi3SetUInt64(inst, integerValueReferences[8], -rndInteger) == 0
# @test fmi3GetUInt64(inst, integerValueReferences[8]) == -rndInteger

@test fmi3SetBoolean(inst, booleanValueReferences[1], rndBoolean) == 0
@test fmi3GetBoolean(inst, booleanValueReferences[1]) == rndBoolean
@test fmi3SetBoolean(inst, booleanValueReferences[1], !rndBoolean) == 0
@test fmi3GetBoolean(inst, booleanValueReferences[1]) == !rndBoolean

@test fmi3SetString(inst, stringValueReferences[1], rndString) == 0
@test fmi3GetString(inst, stringValueReferences[1]) == rndString

@test fmi3SetBinary(inst, binaryValueReferences[1], Csize_t(length(rndString)), pointer(rndString)) == 0
binary = fmi3GetBinary(inst, binaryValueReferences[1])
@test unsafe_string(binary) == rndString

# TODO conflict with same alias for fmi3 datatypes fmi3Boolean, fmi3UInt8
# fmi3Set(inst,
#         [float64ValueReferences[1], integerValueReferences[1], booleanValueReferences[1], stringValueReferences[1], binaryValueReferences[1]],
#         [rndReal,                rndInteger,                rndBoolean,                rndString,                pointer(rndString)])
# @test fmi3Get(inst,
#                 [float64ValueReferences[1], integerValueReferences[1], booleanValueReferences[1], stringValueReferences[1], binaryValueReferences[1]]) ==
#                 [rndReal,                rndInteger,                rndBoolean,                rndString,                unsafe_string(rndString)]

##################
# Testing Arrays #
##################

rndReal = [100 * rand(), 100 * rand()]
rndInteger = [round(Int32, 100 * rand()), round(Int32, 100 * rand())]
rndBoolean = [(rand() > 0.5), (rand() > 0.5)]
tmp = Random.randstring(8)
rndString = [tmp, tmp]

cacheReal = [0.0, 0.0]
cacheInteger =  [fmi3Int32(0), fmi3Int32(0)]
cacheBoolean = [fmi3Boolean(false), fmi3Boolean(false)]
cacheString = [pointer(""), pointer("")]

# TODO not contained in the FMU
# @test fmi3SetFloat32(inst, realValueReferences, rndReal) == 0
# @test fmi3GetFloat32(inst, realValueReferences) == rndReal
# fmi3GetFloat32!(inst, realValueReferences, cacheReal)
# @test cacheReal == rndReal
# @test fmi3SetFloat32(inst, realValueReferences, -rndReal) == 0
# @test fmi3GetFloat32(inst, realValueReferences) == -rndReal
# fmi3GetFloat32!(inst, realValueReferences, cacheReal)
# @test cacheReal == -rndReal

@test fmi3SetFloat64(inst, float64ValueReferences, rndReal) == 0
@test fmi3GetFloat64(inst, float64ValueReferences) == rndReal
fmi3GetFloat64!(inst, float64ValueReferences, cacheReal)
@test cacheReal == rndReal
@test fmi3SetFloat64(inst, float64ValueReferences, -rndReal) == 0
@test fmi3GetFloat64(inst, float64ValueReferences) == -rndReal
fmi3GetFloat64!(inst, float64ValueReferences, cacheReal)
@test cacheReal == -rndReal

# TODO not contained in the FMU
# @test fmi3SetInt8(inst, integerValueReferences, rndInteger) == 0
# @test fmi3GetInt8(inst, integerValueReferences) == rndInteger
# fmi3GetInt8!(inst, integerValueReferences, cacheInteger)
# @test cacheInteger == rndInteger
# @test fmi3SetInt8(inst, integerValueReferences, -rndInteger) == 0
# @test fmi3GetInt8(inst, integerValueReferences) == -rndInteger
# fmi3GetInt8!(inst, integerValueReferences, cacheInteger)
# @test cacheInteger == -rndInteger

# @test fmi3SetInt16(inst, integerValueReferences, rndInteger) == 0
# @test fmi3GetInt16(inst, integerValueReferences) == rndInteger
# fmi3GetInt16!(inst, integerValueReferences, cacheInteger)
# @test cacheInteger == rndInteger
# @test fmi3SetInt16(inst, integerValueReferences, -rndInteger) == 0
# @test fmi3GetInt16(inst, integerValueReferences) == -rndInteger
# fmi3GetInt16!(inst, integerValueReferences, cacheInteger)
# @test cacheInteger == -rndInteger

# TODO only one variable is settable, we want to set at least two
# @test fmi3SetInt32(inst, integerValueReferences, rndInteger) == 0
# @test fmi3GetInt32(inst, integerValueReferences) == rndInteger
# fmi3GetInt32!(inst, integerValueReferences, cacheInteger)
# @test cacheInteger == rndInteger
# @test fmi3SetInt32(inst, integerValueReferences, -rndInteger) == 0
# @test fmi3GetInt32(inst, integerValueReferences) == -rndInteger
# fmi3GetInt32!(inst, integerValueReferences, cacheInteger)
# @test cacheInteger == -rndInteger

# @test fmi3SetInt64(inst, integerValueReferences, rndInteger) == 0
# @test fmi3GetInt64(inst, integerValueReferences) == rndInteger
# fmi3GetInt64!(inst, integerValueReferences, cacheInteger)
# @test cacheInteger == rndInteger
# @test fmi3SetInt64(inst, integerValueReferences, -rndInteger) == 0
# @test fmi3GetInt64(inst, integerValueReferences) == -rndInteger
# fmi3GetInt64!(inst, integerValueReferences, cacheInteger)
# @test cacheInteger == -rndInteger

# @test fmi3SetUInt8(inst, integerValueReferences, rndInteger) == 0
# @test fmi3GetUInt8(inst, integerValueReferences) == rndInteger
# fmi3GetUInt8!(inst, integerValueReferences, cacheInteger)
# @test cacheInteger == rndInteger
# @test fmi3SetUInt8(inst, integerValueReferences, -rndInteger) == 0
# @test fmi3GetUInt8(inst, integerValueReferences) == -rndInteger
# fmi3GetUInt8!(inst, integerValueReferences, cacheInteger)
# @test cacheInteger == -rndInteger

# @test fmi3SetUInt16(inst, integerValueReferences, rndInteger) == 0
# @test fmi3GetUInt16(inst, integerValueReferences) == rndInteger
# fmi3GetUInt16!(inst, integerValueReferences, cacheInteger)
# @test cacheInteger == rndInteger
# @test fmi3SetUInt16(inst, integerValueReferences, -rndInteger) == 0
# @test fmi3GetUInt16(inst, integerValueReferences) == -rndInteger
# fmi3GetUInt16!(inst, integerValueReferences, cacheInteger)
# @test cacheInteger == -rndInteger

# @test fmi3SetUInt32(inst, integerValueReferences, rndInteger) == 0
# @test fmi3GetUInt32(inst, integerValueReferences) == rndInteger
# fmi3GetUInt32!(inst, integerValueReferences, cacheInteger)
# @test cacheInteger == rndInteger
# @test fmi3SetUInt32(inst, integerValueReferences, -rndInteger) == 0
# @test fmi3GetUInt32(inst, integerValueReferences) == -rndInteger
# fmi3GetUInt32!(inst, integerValueReferences, cacheInteger)
# @test cacheInteger == -rndInteger

# @test fmi3SetUInt64(inst, integerValueReferences, rndInteger) == 0
# @test fmi3GetUInt64(inst, integerValueReferences) == rndInteger
# fmi3GetUInt64!(inst, integerValueReferences, cacheInteger)
# @test cacheInteger == rndInteger
# @test fmi3SetUInt64(inst, integerValueReferences, -rndInteger) == 0
# @test fmi3GetUInt64(inst, integerValueReferences) == -rndInteger
# fmi3GetUInt64!(inst, integerValueReferences, cacheInteger)
# @test cacheInteger == -rndInteger

# TODO only one variable is settable, we want to set at least two
# @test fmi3SetBoolean(inst, booleanValueReferences, rndBoolean) == 0
# @test fmi3GetBoolean(inst, booleanValueReferences) == rndBoolean
# fmi3GetBoolean!(inst, booleanValueReferences, cacheBoolean)
# @test cacheBoolean == rndBoolean
# not_rndBoolean = collect(!b for b in rndBoolean)
# @test fmi3SetBoolean(inst, booleanValueReferences, not_rndBoolean) == 0
# @test fmi3GetBoolean(inst, booleanValueReferences) == not_rndBoolean
# fmi3GetBoolean!(inst, booleanValueReferences, cacheBoolean)
# @test cacheBoolean == not_rndBoolean

# @test fmi3SetString(inst, stringValueReferences, rndString) == 0
# @test fmi3GetString(inst, stringValueReferences) == rndString
# fmi3GetString!(inst, stringValueReferences, cacheString)
# @test unsafe_string.(cacheString) == rndString

# TODO only one variable is settable, we want to set at least two
# @test fmi3SetBinary(inst, binaryValueReferences, Csize_t(length(rndString)), pointer(rndString)) == 0
# binary = FMI.fmi3GetBinary(inst, binaryValueReferences)
# @test unsafe_string(binary) == rndString

# Testing input/output derivatives

# TODO not supported
# myFMU = fmi3Load(pathfmu)
# inst = fmi3InstantiateCoSimulation!(myFMU; loggingOn=false)
# @test inst != 0
# @test fmi3EnterInitializationMode(inst) == 0
# dirs = fmi3GetOutputDerivatives(inst, ["h"], ones(Integer, 1))
# @test dirs == -Inf # at this point, derivative is undefined
# # removed @test fmi3SetRealInputDerivatives(inst, ["u_real"], ones(Int, 1), zeros(1)) == 0
# @test fmi3ExitInitializationMode(inst) == 0
# @test fmi3DoStep!(inst, 0.1) == 0

# dirs = fmi3GetOutputDerivatives(inst, ["h"], ones(Int, 1))
# @test dirs == 0.0

############
# Clean up #
############

fmi3Unload(myFMU)


text:
"""
    acentric_factor(model::EoSModel)

calculates the acentric factor using its definition:

    ω = -log10(psatᵣ) -1, at Tᵣ = 0.7
To do so, it calculates the critical temperature (using `crit_pure`) and performs a saturation calculation (with `sat_pure`)

"""
function acentric_factor(model)
    T_c,p_c,_ = crit_pure(model)
    T = 0.7*T_c
    p = first(saturation_pressure(model,T))
    p_r = p/p_c
    return -log10(p_r) - 1.0
end

include("saturation.jl")
include("crit_pure.jl")

export saturation_pressure, saturation_temperature, crit_pure, enthalpy_vap


text:
mutable struct Map <: AbstractMap
    data::Array{T, 2} where T <: Real
    value_unit::Unitful.FreeUnits
    header::FITSHeader
    wcs::WCSTransform
    xaxis::Array{T, 1} where T <: Real
    yaxis::Array{T, 1} where T <: Real
    xy_unit::Unitful.FreeUnits
    xy_resolution::Real
end

# --------------------------------------------------
# map functions

function create_map(data::Array{T, 2} where T <: Real, header::FITSHeader)
    wcs = wcs2d_from_header(header)

    idx = zeros((2, size(data, 2)))
    idx[2, :] = 1:size(data, 2)
    y_axis = pix_to_world(wcs, idx)[2, :];

    idx = zeros((2, size(data, 1)))
    idx[1, :] = 1:size(data, 1)
    x_axis = pix_to_world(wcs, idx)[1, :];

    xy_unit = u"°";

    val_unit = u"K";

    xy_resolution = abs(header["CDELT1"])

    return Map(data, val_unit, header, wcs, x_axis, y_axis, xy_unit, xy_resolution);
end

function submap(map_::Map, xy_range::MapRange)
    xmin_incube = closest_index(map_.xaxis, xy_range.xmin)
    xmax_incube = closest_index(map_.xaxis, xy_range.xmax)
    ymin_incube = closest_index(map_.yaxis, xy_range.ymin)
    ymax_incube = closest_index(map_.yaxis, xy_range.ymax)

    # make sure the indeces are ascending
    xmin_incube, xmax_incube = sort([xmin_incube, xmax_incube])
    ymin_incube, ymax_incube = sort([ymin_incube, ymax_incube])

    data = map_.data[xmin_incube:xmax_incube, ymin_incube:ymax_incube]

    header = deepcopy(map_.header)
    header["CRPIX1"] = header["CRPIX1"] - xmin_incube + 1
    header["CRPIX2"] = header["CRPIX2"] - ymin_incube + 1
    header["NAXIS"] = 2
    header["NAXIS1"] = size(data, 1)
    header["NAXIS2"] = size(data, 2)

    return create_map(data, header)
end

function Base.string(map_::Map)
    # TODO
    return "Map info"
end


text:
## Define type hierchical parser for use in building expressions with mixed types
## This is tested in test/operators.jl
# MeasureRef
function _var_type_parser(V::T, W::T) where {T<:Type{<:MeasureRef}}
    return MeasureRef
end

# ParameterRef
function _var_type_parser(V::T, W::T) where {T<:Type{<:ParameterRef}}
    return ParameterRef
end

# InfiniteVariableRef
function _var_type_parser(V::T, W::T) where {T<:Type{<:InfiniteVariableRef}}
    return InfiniteVariableRef
end

# PointVariableRef
function _var_type_parser(V::T, W::T) where {T<:Type{<:PointVariableRef}}
    return PointVariableRef
end

# HoldVariableRef
function _var_type_parser(V::T, W::T) where {T<:Type{<:HoldVariableRef}}
    return HoldVariableRef
end

# ReducedInfiniteVariableRef
function _var_type_parser(V::T, W::T) where {T<:Type{<:ReducedInfiniteVariableRef}}
    return ReducedInfiniteVariableRef
end

# FiniteVariableRef
function _var_type_parser(V::Type{<:FiniteVariableRef},
                          W::Type{<:FiniteVariableRef})
    return FiniteVariableRef
end

# MeasureFiniteVariableRef
function _var_type_parser(V::Type{<:MeasureFiniteVariableRef},
                          W::Type{<:MeasureFiniteVariableRef})
    return MeasureFiniteVariableRef
end

# GeneralVariableRef
function _var_type_parser(V::Type{<:GeneralVariableRef},
                          W::Type{<:GeneralVariableRef})
    return GeneralVariableRef
end

## Extend add_to_expression! for some more functionality, tested in test/operators.jl
# Mixed variable addition
function JuMP.add_to_expression!(quad::JuMP.GenericQuadExpr{C, Z}, new_coef::C,
                                 new_var1::V, new_var2::W
                                 )::JuMP.GenericQuadExpr where {C,
                                 Z <: GeneralVariableRef, V <: GeneralVariableRef,
                                 W <: GeneralVariableRef}
    type = _var_type_parser(Z, _var_type_parser(V, W))
    key = JuMP.UnorderedPair{type}(new_var1, new_var2)
    new_quad = convert(JuMP.GenericQuadExpr{C, type}, quad)
    JuMP._add_or_set!(new_quad.terms, key, new_coef)
    return new_quad
end

# var1 is a number
function JuMP.add_to_expression!(quad::JuMP.GenericQuadExpr{C, Z},
                                 new_coef::Number, new_var1::Number, new_var2::V
                                 )::JuMP.GenericQuadExpr where {C,
                                 Z <: GeneralVariableRef, V <: GeneralVariableRef}
    type = _var_type_parser(Z, V)
    new_quad = convert(JuMP.GenericQuadExpr{C, type}, quad)
    return JuMP.add_to_expression!(new_quad, new_coef * new_var1, new_var2)
end

# var2 is a number
function JuMP.add_to_expression!(quad::JuMP.GenericQuadExpr{C, Z},
                                 new_coef::Number, new_var1::V, new_var2::Number
                                 )::JuMP.GenericQuadExpr where {C,
                                 Z <: GeneralVariableRef, V <: GeneralVariableRef}
    type = _var_type_parser(Z, V)
    new_quad = convert(JuMP.GenericQuadExpr{C, type}, quad)
    return JuMP.add_to_expression!(new_quad, new_coef * new_var2, new_var1)
end

# var1 and var2 are numbers
function JuMP.add_to_expression!(quad::JuMP.GenericQuadExpr, new_coef::Number,
                                 new_var1::Number, new_var2::Number
                                 )::JuMP.GenericQuadExpr
    JuMP.add_to_expression!(quad.aff, new_coef * new_var2 * new_var1)
    return quad
end

## Extend for better comparisons than default
# GenericAffExpr
function Base.:(==)(aff1::JuMP.GenericAffExpr{C, V},
                    aff2::JuMP.GenericAffExpr{C, W}) where {C, V <: GeneralVariableRef,
                                                            W <: GeneralVariableRef}
    return aff1.constant == aff2.constant && collect(pairs(aff1.terms)) == collect(pairs(aff2.terms))
end

# GenericQuadExpr
function Base.:(==)(quad1::JuMP.GenericQuadExpr{C, V},
                    quad2::JuMP.GenericQuadExpr{C, W}) where {C, V <: GeneralVariableRef,
                                                              W <: GeneralVariableRef}
    pairs1 = collect(pairs(quad1.terms))
    pairs2 = collect(pairs(quad2.terms))
    if length(pairs1) != length(pairs2)
        return false
    end
    for i in eachindex(pairs1)
        if pairs1[i][1].a != pairs2[i][1].a || pairs1[i][1].b != pairs2[i][1].b || pairs1[i][2] != pairs2[i][2]
            return false
        end
    end
    return quad1.aff == quad2.aff
end

## Determine which variables are present in a function
# GeneralVariableRef
function _all_function_variables(f::GeneralVariableRef)::Vector{<:GeneralVariableRef}
    return [f]
end

# GenericAffExpr
function _all_function_variables(f::JuMP.GenericAffExpr)::Vector{<:GeneralVariableRef}
    return GeneralVariableRef[vref for vref in keys(f.terms)]
end

# GenericQuadExpr
function _all_function_variables(f::JuMP.GenericQuadExpr)::Vector{<:GeneralVariableRef}
    aff_vrefs = _all_function_variables(f.aff)
    vref_pairs = [k for k in keys(f.terms)]
    a_vrefs = GeneralVariableRef[pair.a for pair in vref_pairs]
    b_vrefs = GeneralVariableRef[pair.b for pair in vref_pairs]
    return unique([aff_vrefs; a_vrefs; b_vrefs])
end

# Fallback
function _all_function_variables(f)
    error("Can only use InfiniteOpt variables and expressions.")
end

## Return a tuple of the parameter references in an expr
# FiniteVariableRef
_all_parameter_refs(expr::FiniteVariableRef)::Tuple = ()

# InfiniteVariableRef
_all_parameter_refs(expr::InfiniteVariableRef)::Tuple = parameter_refs(expr)

# ParameterRef
_all_parameter_refs(expr::ParameterRef)::Tuple = (expr, )

# ReducedInfiniteVariableRef
_all_parameter_refs(expr::ReducedInfiniteVariableRef)::Tuple = parameter_refs(expr)

# GenericAffExpr
function _all_parameter_refs(expr::JuMP.GenericAffExpr{C,
                              <:GeneralVariableRef})::Tuple where {C}
    pref_list = []
    for var in keys(expr.terms)
        push!(pref_list, _all_parameter_refs(var)...)
    end
    groups = _group.(pref_list)
    unique_groups = unique(groups)
    return Tuple(pref_list[findfirst(isequal(unique_groups[i]), groups)]
                 for i in eachindex(unique_groups))
end

# GenericQuadExpr
function _all_parameter_refs(expr::JuMP.GenericQuadExpr{C,
                             <:GeneralVariableRef})::Tuple where {C}
    pref_list = Any[i for i in _all_parameter_refs(expr.aff)]
    for pair in keys(expr.terms)
        push!(pref_list, _all_parameter_refs(pair.a)...)
        push!(pref_list, _all_parameter_refs(pair.b)...)
    end
    groups = _group.(pref_list)
    unique_groups = unique(groups)
    return Tuple(pref_list[findfirst(isequal(unique_groups[i]), groups)]
                 for i in eachindex(unique_groups))
end

## Delete variables from an expression
# GenericAffExpr
function _remove_variable(f::JuMP.GenericAffExpr, vref::GeneralVariableRef)
    if haskey(f.terms, vref)
        delete!(f.terms, vref)
    end
    return
end

# GenericQuadExpr
function _remove_variable(f::JuMP.GenericQuadExpr, vref::GeneralVariableRef)
    _remove_variable(f.aff, vref)
    vref_pairs = [k for k in keys(f.terms)]
    for i = 1:length(vref_pairs)
        if vref_pairs[i].a == vref
            delete!(f.terms, vref_pairs[i])
        elseif vref_pairs[i].b == vref
            delete!(f.terms, vref_pairs[i])
        end
    end
    return
end

## Modify linear coefficient of variable in expression
# GeneralVariableRef
function _set_variable_coefficient!(expr::GeneralVariableRef,
                                    var::GeneralVariableRef,
                                    coeff::Real)::JuMP.GenericAffExpr
    # Determine if variable is that of the expression and change accordingly
    if expr == var
        return coeff * var
    else
        return expr + coeff * var
    end
end

# GenericAffExpr
function _set_variable_coefficient!(expr::JuMP.GenericAffExpr,
                                    var::GeneralVariableRef,
                                    coeff::Real)::JuMP.GenericAffExpr
    # Determine if variable is in the expression and change accordingly
    if haskey(expr.terms, var)
        expr.terms[var] = coeff
        return expr
    else
        return expr + coeff * var
    end
end

# GenericQuadExpr
function _set_variable_coefficient!(expr::JuMP.GenericQuadExpr,
                                    var::GeneralVariableRef,
                                    coeff::Real)::JuMP.GenericQuadExpr
    # Determine if variable is in the expression and change accordingly
    if haskey(expr.aff.terms, var)
        expr.aff.terms[var] = coeff
        return expr
    else
        return expr + coeff * var
    end
end

# Fallback
function _set_variable_coefficient!(expr, var::GeneralVariableRef, coeff::Real)
    error("Unsupported function type for coefficient modification.")
end

# Check expression for a particular variable type via a recursive search
# This is tested in test/measures.jl
function _has_variable(vrefs::Vector{<:GeneralVariableRef},
                       vref::GeneralVariableRef; prior=[])
    if vrefs[1] == vref
        return true
    elseif isa(vrefs[1], MeasureRef)
        if length(vrefs) > 1
            return _has_variable(_all_function_variables(measure_function(vrefs[1])),
                          vref, prior = GeneralVariableRef[prior; vrefs[2:end]])
        else
            return _has_variable(_all_function_variables(measure_function(vrefs[1])),
                                 vref, prior = prior)
        end
    elseif length(vrefs) > 1
        return _has_variable(vrefs[2:end], vref, prior = prior)
    elseif length(prior) > 0
        return _has_variable(prior, vref)
    else
        return false
    end
end


text: sumdigits(n, base=10) = sum(digits(n, base))

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
# LSTNet.jl
#
# Layers for LSTNet

# Relu-GRU

mutable struct ReluGRUCell{A,V,S}
    Wi::A
    Wh::A
    b::V
    state0::S
end

ReluGRUCell(in, out; init=Flux.glorot_uniform, initb=Flux.zeros32, init_state=Flux.zeros32) =
  ReluGRUCell(init(out * 3, in), init(out * 3, out), initb(out * 3), init_state(out, 1))

function (m::ReluGRUCell{A,V,<:AbstractMatrix{T}})(h, x::Union{AbstractVecOrMat{T},Flux.OneHotArray}) where {A,V,T}
    b, o = m.b, size(h, 1)
    gx, gh = m.Wi * x, m.Wh * h
    r = σ.(Flux.gate(gx, o, 1) .+ Flux.gate(gh, o, 1) .+ Flux.gate(b, o, 1))
    z = σ.(Flux.gate(gx, o, 2) .+ Flux.gate(gh, o, 2) .+ Flux.gate(b, o, 2))
    h̃ = relu.(Flux.gate(gx, o, 3) .+ r .* Flux.gate(gh, o, 3) .+ Flux.gate(b, o, 3))
    h′ = (1 .- z) .* h̃ .+ z .* h
    sz = size(x)
    return h′, reshape(h′, :, sz[2:end]...)
end

Flux.@functor ReluGRUCell

Base.show(io::IO, l::ReluGRUCell) =
  print(io, "ReluGRUCell(", size(l.Wi, 2), ", ", size(l.Wi, 1) ÷ 3, ")")

"""
    ReluGRU(in::Integer, out::Integer)

Gated Recurrent Unit layer with `relu` as activation function.
"""
ReluGRU(a...; ka...) = Recur(ReluGRUCell(a...; ka...))
Recur(m::ReluGRUCell) = Flux.Recur(m, m.state0)

# Skip-GRU

"""
    SkipGRU(in::Integer, out::Integer, p::Integer)
Skip Gated Recurrent Unit layer with skip length `p`. The hidden state is recalled
from `p` steps prior to the current calculation.
"""
SkipGRU(in, out, p ; ka...) = SeqSkip(ReluGRUCell(in, out; ka...), p)

# LSTnet

mutable struct LSTnetCell{A,B,C,D,G}
    ConvLayer::A
    RecurLayer::B
    RecurSkipLayer::C
    RecurDense::D
    AutoregLayer::G
end

"""
    LSTnet(in, convlayersize, recurlayersize, poolsize, skiplength)
    LSTnet(in, convlayersize, recurlayersize, poolsize, skiplength, Flux.relu)

Create a LSTnet layer based on the architecture described in
[Lai et. al.](https://arxiv.org/abs/1703.07015). `in` specifies the number of
input features. `convlayersize` defines the number of convolutional layers, and
`recurlayersize` defines the number of recurrent layers. `poolsize` gives the
length of the window for the pooled input data, and `skiplength` defines the
number of steps the hidden state of the recurrent layer is taken back in time.

Data is expected as array with dimensions `features x poolsize x 1 x data`, i.e.
for 1000 data points containing 31 features that have been windowed over 6
timesteps, `LSTNet` expects an input size of `(31, 6, 1, 1000)`.

Takes the keyword arguments `init` for the initialization of the recurrent
layers; and `initW` and `bias` for the initialization of the dense layer.
"""
function LSTnet(in::Integer, convlayersize::Integer, recurlayersize::Integer, poolsize::Integer, skiplength::Integer, σ=Flux.relu;
	init=Flux.glorot_uniform, initW=Flux.glorot_uniform, bias=true)

	CL = Chain(Conv((in, poolsize), 1 => convlayersize, σ),
			a -> dropdims(a, dims=(1, 2)))
	RL = Seq(ReluGRU(convlayersize, recurlayersize; init=init))
	RSL = SkipGRU(convlayersize, recurlayersize, skiplength; init=init)
	RD = Chain(Dense(2 * recurlayersize, 1, identity))
	AL = Chain(a -> a[:,end,1,:], Dense(in, 1, identity; init=initW, bias=bias))

    LSTnetCell(CL, RL, RSL, RD, AL)
end

function (m::LSTnetCell)(x)
	modelRL1 = m.RecurLayer(m.ConvLayer(x))
	modelRL2 = m.RecurSkipLayer(m.ConvLayer(x))
	modelRL =  m.RecurDense(cat(modelRL1, modelRL2; dims=1))
	return modelRL + m.AutoregLayer(x)
end

function Base.show(io::IO, l::LSTnetCell)
    print(io, "LSTnet(", size(l.ConvLayer[1].weight, 1))
    print(io, ", ", size(l.RecurLayer.chain.cell.Wi, 2), ", ", size(l.RecurLayer.chain.cell.Wi, 1) ÷ 3)
    print(io, ", ", size(l.ConvLayer[1].weight, 2))
    print(io, ", ", l.RecurSkipLayer.p)
    l.ConvLayer[1].σ == Flux.relu || print(io, ", ", l.ConvLayer[1].σ)
    print(io, ")")
end

Flux.@functor LSTnetCell
Flux.params(m::LSTnetCell) = Flux.params(m.ConvLayer, m.RecurLayer, m.RecurSkipLayer, m.RecurDense, m.AutoregLayer)
Flux.reset!(m::LSTnetCell) = foreach(Flux.reset!, (m.ConvLayer, m.RecurLayer, m.RecurSkipLayer, m.RecurDense, m.AutoregLayer))

# Initialize forget gate bias with -1
function initialize_bias!(l::LSTnetCell)
	l.RecurLayer.chain.cell.b .= 1
	l.RecurSkipLayer.cell.b .= 1
	return nothing
end

max_stars_repo_name: sdobber/FluxArchitectures

max_issues_repo_name: sdobber/FluxArchitectures

max_forks_repo_name: sdobber/FluxArchitectures


text:
module Topologies

using EllipsisNotation
import Base: getindex

using ..Transforms
using ..Elements
using ..Functions

export Lagrange
export StructuredTopology
export boundary, mkbasis, refelems

abstract type Topology{N} <: AbstractArray{Element,N} end

struct Lagrange end

struct StructuredTopology{N} <: Topology{N}
    elemids :: Array{Int,N}
    dimcorr :: Tuple{Vararg{AbstractTransform}}
    parent :: Union{StructuredTopology, Nothing}
end

function StructuredTopology(shape::Vararg{Int,N}) where N
    elemids = Array{Int,N}(undef, shape...)
    elemids[:] .= 1:prod(shape)
    StructuredTopology{N}(elemids, (), nothing)
end

Base.size(self::StructuredTopology) = size(self.elemids)
Base.length(self::StructuredTopology) = length(self.elemids)
Base.IndexStyle(::Type{StructuredTopology{N}}) where N = IndexCartesian()

@inline function Base.getindex(self::StructuredTopology{N}, I::Vararg{Int,N}) where N
    @boundscheck checkbounds(self, I...)
    if self.parent == nothing
        trans = (Shift([float(i-1) for i in I]),)
        Element(Simplex{N}(), self.elemids[I...], (self.dimcorr, trans))
    else
        super = self.parent[self.elemids[I...]]
        Element(Simplex{N}(), super.index, ((self.dimcorr..., super.dimcorr...), super.transform))
    end
end

function boundary(self::StructuredTopology{N}, axis::Int, loc::Symbol) where N
    index = Any[(:) for _ in 1:N-1]
    insert!(index, axis, loc == :first ? 1 : size(self.elemids, axis))
    elemids = self.elemids[index...]
    dimcorr = Updim{N-1,axis}(loc == :first ? 0.0 : 1.0)
    StructuredTopology{N-1}(elemids, (dimcorr,), self)
end

refelems(::StructuredTopology{N}) where N = (Simplex{N}(),)

function mkbasis(self::StructuredTopology{N}, ::Type{Lagrange}, degree::Int) where N
    # Generate N single-dimensional Lagrangian bases of the right degree
    poly = Monomials(point(N), degree)
    coeffs = hcat((
        [binomial(degree,nu) * binomial(degree-nu,k-nu) * (isodd(k-nu) ? -1 : 1) for nu in 0:degree]
        for k in 0:degree
    )...)
    basis1d = Contract(Constant(coeffs), poly)

    # Formulate the tensor product
    factors = [InsertAxis(basis1d[:, k], [fill(1, k-1); fill(2, N-k)]) for k in 1:N]
    basisNd = .*(factors...)

    # Number the DoFs
    ndofs = [n*degree+1 for n in size(self.elemids)]
    dofgrid = reshape(1:prod(ndofs), ndofs...)
    elemwindows = [1+(k-1)*degree:1+k*degree for k in 1:maximum(ndofs)]
    dofmap = Array{Int,2}(undef, (degree+1)^N, length(self.elemids))
    for (i, windows) in enumerate(Iterators.product([elemwindows[1:n] for n in size(self.elemids)]...))
        dofmap[:,i] = dofgrid[windows...][:]
    end

    # Return inflation
    Inflate(reshape(basisNd, (degree+1)^N), (prod(ndofs),), Constant(dofmap)[:, elemindex])
end

end # module

max_stars_repo_name: TheBB/Jutils

max_issues_repo_name: TheBB/Jutils

max_forks_repo_name: TheBB/Jutils