text:
@doc """
[ReinforcementLearningBase.jl](https://github.com/JuliaReinforcementLearning/ReinforcementLearningBase.jl)
(**RLBase**) provides some common constants, traits, abstractions and interfaces
in developing reinforcement learning algorithms in Julia.

Basically, we defined the following two main concepts in reinforcement learning:

- [`AbstractPolicy`](@ref)
- [`AbstractEnv`](@ref)
""" RLBase

import Base: copy, copyto!, nameof
import Random: seed!, rand, AbstractRNG
import AbstractTrees: children, has_children
import Markdown

#####
# Policy
#####

"""
    (π::AbstractPolicy)(env) -> action

Policy is the most basic concept in reinforcement learning. Unlike the
definition in some other packages, here a policy is defined as a functional
object which takes in an environment and returns an action.

!!! note
    See discussions
    [here](https://github.com/JuliaReinforcementLearning/ReinforcementLearning.jl/issues/86)
    if you are wondering why we define the input as `AbstractEnv` instead of
    state.

!!! warning
    The policy `π` may change its internal state but it shouldn't change `env`.
    When it's really necessary, remember to make a copy of `env` to keep the
    original `env` untouched.
"""
@api abstract type AbstractPolicy end
@api (π::AbstractPolicy)(env)

"""
    update!(π::AbstractPolicy, experience)

Update the policy `π` with online/offline experience or parameters.
"""
@api update!(π::AbstractPolicy, experience)

"""
    prob(π::AbstractPolicy, env) -> Distribution

Get the probability distribution of actions based on policy `π` given an `env`.
"""
@api prob(π::AbstractPolicy, env)

"""
    prob(π::AbstractPolicy, env, action)

Only valid for environments with discrete actions.
"""
@api prob(π::AbstractPolicy, env, action)

"""
    priority(π::AbstractPolicy, experience)

Usually used in offline policies.
"""
@api priority(π::AbstractPolicy, experience)

#####
# Environment
#####

"""
    (env::AbstractEnv)(action, player=current_player(env))

Super type of all reinforcement learning environments.
"""
@api abstract type AbstractEnv end

abstract type AbstractEnvStyle end

#####
## Traits for Environment
## mostly borrowed from https://github.com/deepmind/open_spiel/blob/master/open_spiel/spiel.h
#####

#####
### NumAgentStyle
#####

abstract type AbstractNumAgentStyle <: AbstractEnvStyle end

@api struct SingleAgent <: AbstractNumAgentStyle end
@api const SINGLE_AGENT = SingleAgent()

@api struct MultiAgent{N} <: AbstractNumAgentStyle end

"""
    MultiAgent(n::Integer) -> MultiAgent{n}()

`n` must be ≥ 2.
"""
function MultiAgent(n::Integer)
    if n < 0
        throw(ArgumentError("number of agents must be > 1, get $n"))
    elseif n == 1
        throw(ArgumentError("do you mean `SINGLE_AGENT`?"))
    else
        MultiAgent{convert(Int, n)}()
    end
end

@api const TWO_AGENT = MultiAgent(2)

"""
    NumAgentStyle(env)

Number of agents involved in the `env`. Possible returns are:

- [`SINGLE_AGENT`](@ref). This is the default return.
- [`MultiAgent`][@ref].
"""
@env_api NumAgentStyle(env::T) where {T<:AbstractEnv} = NumAgentStyle(T)
NumAgentStyle(env::Type{<:AbstractEnv}) = SINGLE_AGENT

#####
### DynamicStyle
#####

abstract type AbstractDynamicStyle <: AbstractEnvStyle end

@api struct Sequential <: AbstractDynamicStyle end
@api struct Simultaneous <: AbstractDynamicStyle end

"Environment with the [`DynamicStyle`](@ref) of `SEQUENTIAL` must takes actions from different players one-by-one."
@api const SEQUENTIAL = Sequential()

"Environment with the [`DynamicStyle`](@ref) of `SIMULTANEOUS` must take in actions from some (or all) players at one time"
@api const SIMULTANEOUS = Simultaneous()

"""
    DynamicStyle(env::AbstractEnv) = SEQUENTIAL

Only valid in environments with a [`NumAgentStyle`](@ref) of
[`MultiAgent`](@ref). Determine whether the players can play simultaneously or
not. Possible returns are:

- [`SEQUENTIAL`](@ref). This is the default return.
- [`SIMULTANEOUS`](@ref).
"""
@env_api DynamicStyle(env::T) where {T<:AbstractEnv} = DynamicStyle(T)
DynamicStyle(::Type{<:AbstractEnv}) = SEQUENTIAL

#####
### InformationStyle
#####

abstract type AbstractInformationStyle <: AbstractEnvStyle end

@api struct PerfectInformation <: AbstractInformationStyle end
@api struct ImperfectInformation <: AbstractInformationStyle end

"All players observe the same state"
@api const PERFECT_INFORMATION = PerfectInformation()

"The inner state of some players' observations may be different"
@api const IMPERFECT_INFORMATION = ImperfectInformation()

"""
    InformationStyle(env) = IMPERFECT_INFORMATION

Distinguish environments between [`PERFECT_INFORMATION`](@ref) and
[`IMPERFECT_INFORMATION`](@ref). [`IMPERFECT_INFORMATION`](@ref) is returned by default.
"""
@env_api InformationStyle(env::T) where {T<:AbstractEnv} = InformationStyle(T)
InformationStyle(::Type{<:AbstractEnv}) = IMPERFECT_INFORMATION

#####
### ChanceStyle
#####

abstract type AbstractChanceStyle <: AbstractEnvStyle end
abstract type AbstractStochasticChanceStyle <: AbstractChanceStyle end

@api struct Deterministic <: AbstractChanceStyle end
@api struct Stochastic <: AbstractStochasticChanceStyle end
@api struct ExplicitStochastic <: AbstractStochasticChanceStyle end
@api struct SampledStochastic <: AbstractStochasticChanceStyle end

"No chance player in the environment. And the game is fully deterministic."
@api const DETERMINISTIC = Deterministic()

"""
No chance player in the environment. And the game is stochastic. To help
increase reproducibility, these environments should generally accept a
`AbstractRNG` as a keyword argument. For some third-party environments, at least
a `seed` is exposed in the constructor.
"""
@api const STOCHASTIC = Stochastic()

"""
Usually used to describe [extensive-form game](https://en.wikipedia.org/wiki/Extensive-form_game).
The environment contains a chance player and the corresponding probability is known.
Therefore, [`prob`](@ref)`(env, player=chance_player(env))` must be defined.
"""
@api const EXPLICIT_STOCHASTIC = ExplicitStochastic()

"""
Environment contains chance player and the probability is unknown. Usually only
a dummy action is allowed in this case.

!!! note
    The chance player ([`chance_player`](@ref)`(env)`) must appears in the result of
    [`players`](@ref)`(env)`.
    The result of `action_space(env, chance_player)` should only contains one
    dummy action.
"""
@api const SAMPLED_STOCHASTIC = SampledStochastic()

"""
    ChanceStyle(env) = DETERMINISTIC

Specify which role the chance plays in the `env`. Possible returns are:

- [`STOCHASTIC`](@ref). This is the default return.
- [`DETERMINISTIC`](@ref)
- [`EXPLICIT_STOCHASTIC`](@ref)
- [`SAMPLED_STOCHASTIC`](@ref)
"""
@env_api ChanceStyle(env::T) where {T<:AbstractEnv} = ChanceStyle(T)
ChanceStyle(::Type{<:AbstractEnv}) = STOCHASTIC

#####
### RewardStyle
#####

abstract type AbstractRewardStyle <: AbstractEnvStyle end

@api struct StepReward <: AbstractRewardStyle end
@api struct TerminalReward <: AbstractRewardStyle end

"We can get reward after each step"
@api const STEP_REWARD = StepReward()

"Only get reward at the end of environment"
@api const TERMINAL_REWARD = TerminalReward()

"""
Specify whether we can get reward after each step or only at the end of an game.
Possible values are [`STEP_REWARD`](@ref) (the default one) or
[`TERMINAL_REWARD`](@ref).

!!! note
    Environments of [`TERMINAL_REWARD`](@ref) style can be viewed as a subset of
    environments of [`STEP_REWARD`](@ref) style. For some algorithms, like MCTS,
    we may have some a more efficient implementation for environments of
    [`TERMINAL_REWARD`](@ref) style.
"""
@env_api RewardStyle(env::T) where {T<:AbstractEnv} = RewardStyle(T)
RewardStyle(::Type{<:AbstractEnv}) = STEP_REWARD

#####
### UtilityStyle
#####

abstract type AbstractUtilityStyle <: AbstractEnvStyle end

@api struct ZeroSum <: AbstractUtilityStyle end
@api struct ConstantSum <: AbstractUtilityStyle end
@api struct GeneralSum <: AbstractUtilityStyle end
@api struct IdenticalUtility <: AbstractUtilityStyle end

"Rewards of all players sum to 0. A special case of [`CONSTANT_SUM`]."
@api const ZERO_SUM = ZeroSum()

"Rewards of all players sum to a constant"
@api const CONSTANT_SUM = ConstantSum()

"Total rewards of all players may be different in each step"
@api const GENERAL_SUM = GeneralSum()

"Every player gets the same reward"
@api const IDENTICAL_UTILITY = IdenticalUtility()

"""
    UtilityStyle(env::AbstractEnv)

Specify the utility style in multi-agent environments. Possible values are:

- [GENERAL_SUM](@ref). The default return.
- [ZERO_SUM](@ref)
- [CONSTANT_SUM](@ref)
- [IDENTICAL_UTILITY](@ref)
"""
@env_api UtilityStyle(env::T) where {T<:AbstractEnv} = UtilityStyle(T)
UtilityStyle(::Type{<:AbstractEnv}) = GENERAL_SUM

#####
### ActionStyle
#####

abstract type AbstractActionStyle <: AbstractEnvStyle end
abstract type AbstractDiscreteActionStyle <: AbstractActionStyle end
@api struct FullActionSet <: AbstractDiscreteActionStyle end

"The action space of the environment may contains illegal actions"
@api const FULL_ACTION_SET = FullActionSet()

@api struct MinimalActionSet <: AbstractDiscreteActionStyle end

"All actions in the action space of the environment are legal"
@api const MINIMAL_ACTION_SET = MinimalActionSet()

"""
    ActionStyle(env::AbstractEnv)

For environments of discrete actions, specify whether the current state of `env`
contains a full action set or a minimal action set. By default the
[`MINIMAL_ACTION_SET`](@ref) is returned.
"""
@env_api ActionStyle(env::T) where {T<:AbstractEnv} = ActionStyle(T)
ActionStyle(::Type{<:AbstractEnv}) = MINIMAL_ACTION_SET

#####
# StateStyle
#####

abstract type AbstractStateStyle end

"See the definition of [information set](https://en.wikipedia.org/wiki/Information_set_(game_theory))"
@api struct InformationSet{T} <: AbstractStateStyle end
InformationSet() = InformationSet{Any}()

"Use it to represent the internal state."
@api struct InternalState{T} <: AbstractStateStyle end
InternalState() = InternalState{Any}()

"Use it to represent the [goal state](http://proceedings.mlr.press/v37/schaul15.pdf)"
@api struct GoalState{T} <: AbstractStateStyle end
GoalState() = GoalState{Any}()

"""
Sometimes people from different field talk about the same thing with a different
name. Here we set the `Observation{Any}()` as the default state style in this
package.

See discussions [here](https://ai.stackexchange.com/questions/5970/what-is-the-difference-between-an-observation-and-a-state-in-reinforcement-learn)
"""
@api struct Observation{T} <: AbstractStateStyle end
Observation() = Observation{Any}()

"""
    StateStyle(env::AbstractEnv)

Define the possible styles of `state(env)`. Possible values are:

- [`Observation{T}`](@ref). This is the default return.
- [`InternalState{T}`](@ref)
- [`Information{T}`](@ref)
- You can also define your customized state style when necessary.

Or a tuple contains several of the above ones.

This is useful for environments which provide more than one kind of state.
"""
@env_api StateStyle(env::AbstractEnv) = Observation{Any}()

"""
Specify the defalt state style when calling `state(env)`.
"""
@env_api DefaultStateStyle(env::AbstractEnv) = DefaultStateStyle(StateStyle(env))
DefaultStateStyle(ss::AbstractStateStyle) = ss
DefaultStateStyle(ss::Tuple{Vararg{<:AbstractStateStyle}}) = first(ss)

# EpisodeStyle
# Episodic
# NeverEnding

#####
# General
#####

@api struct DefaultPlayer end
@api const DEFAULT_PLAYER = DefaultPlayer()

@api struct ChancePlayer end
@api const CHANCE_PLAYER = ChancePlayer()

@api struct SimultaneousPlayer end
@api const SIMULTANEOUS_PLAYER = SimultaneousPlayer()

@api struct Spector end
@api const SPECTOR = Spector()

@api (env::AbstractEnv)(action, player = current_player(env))

"""
Make an independent copy of `env`,

!!! note
    rng (if `env` has) is also copied!
"""
@api copy(env::AbstractEnv) = deepcopy(env)
@api copyto!(dest::AbstractEnv, src::AbstractEnv)

# checking the state of all players in env is enough?
"""
    Base.:(==)(env1::T, env2::T) where T<:AbstractEnv
!!! warning
    Only check the state of all players in the env.
"""
function Base.:(==)(env1::T, env2::T) where T<:AbstractEnv
    len = length(players(env1))
    len == length(players(env2)) &&
    all(state(env1, player) == state(env2, player) for player in players(env1))
end
Base.hash(env::AbstractEnv, h::UInt) = hash([state(env, player) for player in players(env)], h)

@api nameof(env::AbstractEnv) = nameof(typeof(env))

"""
Get the action distribution of chance player.

!!! note
    Only valid for environments of [`EXPLICIT_STOCHASTIC`](@ref) style. The
    current player of `env` must be the chance player.
"""
@env_api prob(env::AbstractEnv, player = chance_player(env))

"""
    action_space(env, player=current_player(env))

Get all available actions from environment. See also:
[`legal_action_space`](@ref)
"""
@multi_agent_env_api action_space(env::AbstractEnv, player = current_player(env))

"""
    legal_action_space(env, player=current_player(env))

For environments of [`MINIMAL_ACTION_SET`](@ref), the result is the same with
[`action_space`](@ref).
"""
@multi_agent_env_api legal_action_space(env::AbstractEnv, player = current_player(env)) =
    legal_action_space(ActionStyle(env), env, player)

legal_action_space(::MinimalActionSet, env, player) = action_space(env)

"""
    legal_action_space_mask(env, player=current_player(env)) -> AbstractArray{Bool}

Required for environments of [`FULL_ACTION_SET`](@ref).
"""
@multi_agent_env_api legal_action_space_mask(env::AbstractEnv, player = current_player(env))

"""
    state(env, style=[DefaultStateStyle(env)], player=[current_player(env)])

The state can be of any type. However, most neural network based algorithms
assume an `AbstractArray` is returned. For environments with many different states
provided (inner state, information state, etc), users need to provide `style`
to declare which kind of state they want.

!!! warning
    The state **may** be reused and be mutated at each step. Always remember to make a copy
    if this is not what you expect.
"""
@multi_agent_env_api state(env::AbstractEnv) = state(env, DefaultStateStyle(env))
state(env::AbstractEnv, ss::AbstractStateStyle) = state(env, ss, current_player(env))
state(env::AbstractEnv, player) = state(env, DefaultStateStyle(env), player)

"""
    state_space(env, style=[DefaultStateStyle(env)], player=[current_player(env)])

Describe all possible states.
"""
@multi_agent_env_api state_space(env::AbstractEnv) =
    state_space(env, DefaultStateStyle(env))
state_space(env::AbstractEnv, ss::AbstractStateStyle) =
    state_space(env, ss, current_player(env))
state_space(env::AbstractEnv, player) = state_space(env, DefaultStateStyle(env), player)

"""
    current_player(env)

Return the next player to take action. For [Extensive Form
Games](https://en.wikipedia.org/wiki/Extensive-form_game), a *chance player* may
be returned. (See also [`chance_player`](@ref)) For [SIMULTANEOUS](@ref)
environments, a *simultaneous player* is always returned. (See also
[`simultaneous_player`](@ref)).
"""
@env_api current_player(env::AbstractEnv) = DEFAULT_PLAYER

"""
    chance_player(env)

Only valid for environments with a chance player.
"""
@env_api chance_player(env::AbstractEnv) = CHANCE_PLAYER

"""
    simultaneous_player(env)

Only valid for environments of [`SIMULTANEOUS`](@ref) style.
"""
@env_api simultaneous_player(env) = SIMULTANEOUS_PLAYER

"""
    spectator_player(env)

Used in imperfect multi-agent environments.
"""
@env_api spectator_player(env::AbstractEnv)

@env_api players(env::AbstractEnv) = (DEFAULT_PLAYER,)

"Reset the internal state of an environment"
@env_api reset!(env::AbstractEnv)

"Set the seed of internal rng"
@env_api seed!(env::AbstractEnv, seed)

"""
    is_terminated(env, player=current_player(env))
"""
@env_api is_terminated(env::AbstractEnv)

"""
    reward(env, player=current_player(env))
"""
@multi_agent_env_api reward(env::AbstractEnv, player = current_player(env))

"""
    child(env::AbstractEnv, action)

Treat the `env` as a game tree. Create an independent child after applying
`action`.
"""
@api function child(env::AbstractEnv, action)
    new_env = copy(env)
    new_env(action)
    new_env
end

@api has_children(env::AbstractEnv) = !is_terminated(env)

@api children(env::AbstractEnv) = (child(env, action) for action in legal_action_space(env))

"""
    walk(f, env::AbstractEnv)

Call `f` with `env` and its descendants. Only use it with small games.
"""
@api function walk(f, env::AbstractEnv)
    f(env)
    if has_children(env)
        for x in children(env)
            walk(f, x)
        end
    end
end

#####
# EnvironmentModel
#####

"""
TODO:

Describe how to model a reinforcement learning environment. TODO: need more
investigation Ref: https://bair.berkeley.edu/blog/2019/12/12/mbpo/
- Analytic gradient computation
- Sampling-based planning
- Model-based data generation
- Value-equivalence prediction [Model-based Reinforcement Learning: A
  Survey.](https://arxiv.org/pdf/2006.16712.pdf) [Tutorial on Model-Based
  Methods in Reinforcement
  Learning](https://sites.google.com/view/mbrl-tutorial)
"""
@api abstract type AbstractEnvironmentModel end


text:
abstract SegmentGeometry <: Gadfly.GeometryElement

# Geometry for vectors/arrows/segments
immutable SegmentGeom <: SegmentGeometry

     arrow::Bool
     filled::Bool
     tag::Symbol

    function SegmentGeom(;arrow=false, filled=false, tag=empty_tag)
         new(arrow, filled, tag)
     end
end

const segment = SegmentGeom

# Leave this as a function, pending extra arguments e.g. arrow attributes
function vector(;filled::Bool=false, tag::Symbol=empty_tag)
    return SegmentGeom(arrow=true, filled=filled, tag=tag)
end

function element_aesthetics(::SegmentGeom)
    return [:x, :y, :xend, :yend, :color]
end

function Gadfly.render(geom::SegmentGeom, theme::Gadfly.Theme, aes::Gadfly.Aesthetics,
                subplot_layer_aess::@compat(Union{(@compat Void), Vector{Gadfly.Aesthetics}}),
                subplot_layer_datas::@compat(Union{(@compat Void), Vector{Gadfly.Data}}),
    scales::Dict{Symbol, Gadfly.ScaleElement})
    render(geom, theme, aes, scales)
end

function render(geom::SegmentGeom, theme::Gadfly.Theme, aes::Gadfly.Aesthetics,
    scales::Dict{Symbol, Gadfly.ScaleElement})

    Gadfly.assert_aesthetics_defined("Geom.segment", aes, :x, :y, :xend, :yend)

    function arrow{T<:Real}(x::T, y::T, xmax::T, ymax::T, xyrange::Vector{T})
        dx = xmax-x
        dy = ymax-y
        vl = 0.225*hypot(dy/xyrange[2], dx/xyrange[1])
        θ =  atan2(dy/xyrange[2], dx/xyrange[1])
        ϕ = pi/15
        xr = -vl*xyrange[1]*[cos(θ+ϕ), cos(θ-ϕ)]
        yr = -vl*xyrange[2]*[sin(θ+ϕ), sin(θ-ϕ)]
        arr = [(xmax+xr[1],ymax+yr[1]), (xmax,ymax), (xmax+xr[2],ymax+yr[2]) ]
        return arr
    end

    n = length(aes.x)
    color = ColorTypes.RGBA{Float32}(theme.default_color)
    default_aes = Gadfly.Aesthetics()
    default_aes.color = DataArrays.DataArray(fill(color,n))

    aes = inherit(aes, default_aes)

#    line_style = Gadfly.get_stroke_vector(theme.line_style)
    line_style = theme.line_style
    if is(line_style, nothing)
        line_style = []
    end

# Geom.vector requires information about scales

if geom.arrow
    xscale = scales[:x]
    yscale = scales[:y]
    check = [xscale.minvalue, xscale.maxvalue, yscale.minvalue, yscale.maxvalue]
    if any( map(x -> is(x,nothing), check) )
        error("For Geom.vector, Scale minvalue and maxvalue must be manually provided for both axes")
    end
    fx = xscale.trans.f
    fy = yscale.trans.f
    xyrange = [fx(xscale.maxvalue)-fx(xscale.minvalue),
        fy(yscale.maxvalue)-fy(yscale.minvalue)]

     arrows = [ arrow(x, y, xend, yend, xyrange)
            for (x, y, xend, yend) in zip(aes.x, aes.y, aes.xend, aes.yend) ]

end

    segments = [ [(x,y), (xend,yend)]
        for (x, y, xend, yend) in zip(aes.x, aes.y, aes.xend, aes.yend) ]

    classes = [string("geometry ", svg_color_class_from_label( aes.color_label([c])[1] ))
        for c in aes.color ]

    ctx = context()

    compose!( ctx, Compose.line(segments, geom.tag), stroke(aes.color), linewidth(theme.line_width),
        strokedash(line_style), svgclass( classes )  )
    if geom.arrow
        if geom.filled
            compose!(ctx, (context(), Compose.polygon(arrows), fill(aes.color), strokedash([])) )
        else
            compose!(ctx, (context(), Compose.line(arrows), stroke(aes.color), linewidth(theme.line_width),
            strokedash([]))  )
        end
    end

    return ctx
end

ext: jl


text:
# SubgridScale parameters
const SubgridScale = CLIMAParameters.SubgridScale
SubgridScale.von_karman_const(ps::AbstractEarthParameterSet)        = 0.4


text:
# active information functions load correctly
@test Inform._active_info_found
@test Inform._active_info != C_NULL
@test isa(Inform._active_info, Ptr{Void})

@test Inform._active_info_ensemble_found
@test Inform._active_info_ensemble != C_NULL
@test isa(Inform._active_info_ensemble, Ptr{Void})

## Individual time series

# series too short
@test isnan(activeinfo([],2))
@test isnan(activeinfo([1],2))

# history too long
@test isnan(activeinfo([1,1],2))
@test !isnan(activeinfo([1,1,0],2))

# encoding error
let series = [2,1,0,0,1,0,0,1]
    @test !isnan(activeinfo(series,2,b=3))
    @test isnan(activeinfo(series,2,b=2))
end

@test_approx_eq_eps activeinfo([1,1,0,0,1,0,0,1],2) 0.918296 1e-6

@test_approx_eq_eps activeinfo([1,0,0,0,0,0,0,0,0], 2) 0.000000 1e-6
@test_approx_eq_eps activeinfo([0,0,1,1,1,1,0,0,0], 2) 0.305958 1e-6
@test_approx_eq_eps activeinfo([1,0,0,0,0,0,0,1,1], 2) 0.347458 1e-6
@test_approx_eq_eps activeinfo([0,0,0,0,0,1,1,0,0], 2) 0.399533 1e-6
@test_approx_eq_eps activeinfo([1,1,1,0,0,0,0,1,1], 2) 0.305958 1e-6
@test_approx_eq_eps activeinfo([0,0,0,1,1,1,1,0,0], 2) 0.305958 1e-6
@test_approx_eq_eps activeinfo([0,0,0,0,0,0,1,1,0], 2) 0.347458 1e-6

@test_approx_eq_eps activeinfo([3,3,3,2,1,0,0,0,1], 2, b=4) 0.635471 1e-6
@test_approx_eq_eps activeinfo([2,2,3,3,3,3,2,1,0], 2, b=4) 0.635471 1e-6
@test_approx_eq_eps activeinfo([2,2,2,2,2,2,1,1,1], 2, b=4) 0.234783 1e-6

## Ensemble time series
let series = [ 1 0; 1 0; 0 0; 0 1; 1 0; 0 0; 0 0; 1 1]
    @test_approx_eq_eps activeinfo(series, 2) 0.459148 1e-6
end
let series = [
    1 0 1 1 0 0 1 0 0;
    0 0 0 0 0 0 1 0 0;
    0 1 0 0 0 0 1 0 0;
    0 1 0 0 0 0 0 1 0;
    0 1 0 0 0 1 0 1 0;
    0 1 0 0 1 1 0 1 0;
    0 0 0 0 1 0 0 1 1;
    0 0 1 1 0 0 1 0 1;
    0 0 1 1 0 0 1 0 0]
    @test_approx_eq_eps activeinfo(series, 2) 0.308047 1e-6
end

let series = [
    3 3 3 2 1 0 0 0 1;
    2 2 3 3 3 3 2 1 0;
    0 0 0 0 1 1 0 0 0;
    1 1 0 0 0 1 1 2 2]'
    @test_approx_eq_eps activeinfo(series, 2) 0.6621455 1e-6
end


text:
using Jeff, Test, DelimitedFiles, ForwardDiff

curdir = pwd()
FWHM = 2.0 * sqrt(2.0 * log(2.0))

LAYERS = Array{Any}(undef, (3, 4))
LAYERS[1, 1] = 0.0
LAYERS[1, 2] = 0.0
LAYERS[1, 3] = 0.0
LAYERS[1, 4] = 0.0
LAYERS[2, 1] = 100.
LAYERS[2, 2] = 3.47
LAYERS[2, 3] = 0.
LAYERS[2, 4] = 2.
LAYERS[3, 1] = 0.0
LAYERS[3, 2] = 2.07
LAYERS[3, 3] = 0.0
LAYERS[3, 4] = 3.

@testset "abeles" begin
    data = readdlm(joinpath(curdir, "data", "theoretical.txt"))
    q_test = data[:, 1]
    r_test = data[:, 2]
    result = Jeff.abeles(q_test, LAYERS)
    @test all(isapprox.(result, r_test, atol=1e-9))

    local layers
    # some of the ORSO tests
    for i in [0, 1, 2, 3, 6]
        layers = readdlm(joinpath(curdir, "data", "test$i.layers"))
        data = readdlm(joinpath(curdir, "data", "test$i.dat"))
        q_test = data[:, 1]
        r_test = data[:, 2]
        result = Jeff.abeles(q_test, layers)
        # same tolerances as used in ORSO test for refnx
        @test all(isapprox.(result, r_test, rtol=8e-5))
    end

    # check that abeles is differentiable
    f = x -> sum(Jeff.abeles(q_test, x))
    ForwardDiff.gradient(f, layers)
end

@testset "constant_smearing" begin
    data = readdlm(joinpath(curdir, "data", "theoretical.txt"))
    q_test = data[:, 1]
    r_test = data[:, 2]

    result = Jeff.constant_smearing(q_test, LAYERS, 0.0, 1.0, 0.0)
    @test all(isapprox.(result, r_test, atol=1e-9))

    # ORSO tests
    data = readdlm(joinpath(curdir, "data", "test4.dat"))
    q_test = data[:, 1]
    r_test = data[:, 2]
    layers = readdlm(joinpath(curdir, "data", "test0.layers"))
    result = Jeff.constant_smearing(q_test, layers, 5.0, 1.0, 0.0)
    # same tolerance as used in ORSO test for refnx
    # should look into whether this can be tightened.
    @test all(isapprox.(result, r_test, rtol=0.03))
end

@testset "pointwise_smearing" begin
    # ORSO tests
    data = readdlm(joinpath(curdir, "data", "test4.dat"))
    q_test = data[:, 1]
    r_test = data[:, 2]
    dq_test = data[:, 4] * FWHM
    layers = readdlm(joinpath(curdir, "data", "test0.layers"))
    result = Jeff.pointwise_smearing(q_test, layers, dq_test)
    # same tolerance as used in ORSO test for refnx
    # should look into whether this can be tightened.
    @test all(isapprox.(result, r_test, rtol=0.03))
end


text:
# # [Fairness.jl](https://github.com/ashryaagr/Fairness.jl) - Fairness Toolkit in Julia
#
# Fairness.jl is a new bias audit and mitigation toolkit in Julia designed with the aim to solve practical problems faced by practitioners with existing fairness toolkits.
#
# This notebook shall present you with an introduction to **[Fairness.jl](https://github.com/ashryaagr/Fairness.jl)**, its power, and uniqueness by the means of a Real-Life example of [COMPAS Dataset](https://github.com/propublica/compas-analysis/blob/master/Compas%20Analysis.ipynb).
#
# But before visiting we begin the introduction, let us be clear why we need this package.
#
# ## Why should  I care about fairness, ethics, etc. ?
#
# Machine Learning is involved in a lot of crucial decision support tools. The use of these tools ranges from granting parole, shortlisting job applications to accepting credit applications. There have been a lot of political and policy developments during the past year that have pointed out the transparency issues and data bias in these Machine Learning tools. Thus it has become crucial for the Machine learning community to think about fairness and bias. Eliminating bias is not as easy as it might seem at first glance. This toolkit helps you in easily auditing and minimizing the bias through a collection of various fairness metrics and algorithms.
#
# # Example
#
# In this example, we will use the COMPAS Dataset to predict whether a criminal defendant will recidivate(re-offend). Neural Network Classifier is used for classification. It is wrapped with the Reweighing Algorithm to preprocess the data. This wrapped Model is then wrapped with Equalized Odds Algorithm for postprocessing of predictions.
#
# ## Downloading Required Packages
#
# To install the package, you have to install
#
# * Fairness.jl package
# * [MLJ : Machine Learning Toolkit](https://alan-turing-institute.github.io/MLJ.jl/dev/)
# * PrettyPrinting : To get pretty output.`pprint`is its only function used by us.

using Pkg
Pkg.activate("my_environment", shared=true)
Pkg.add("Fairness")
Pkg.add("MLJ") # Toolkit for Machine Learning
Pkg.add("PrettyPrinting") # For readibility of outputs

# Now we import the required packages. Note that this will take 5 minutes on the first run. This is the case with all Julia packages. Julia pre-compiles the code to make it more efficient. This is a one-time thing. 2nd run onwards, everything will be fast!!
#

using Fairness
using MLJ
using PrettyPrinting

#
# ## Load the [COMPAS Dataset](https://github.com/propublica/compas-analysis/blob/master/Compas%20Analysis.ipynb) using the macro by Fairness.jl
#
# This dataset has 8 features and 6907 rows. The protected attribute here is `race`. Using the 8 features, it predicts whether a criminal defendant will recidivate(re-offend).
#

data = @load_compas
X, y = data
data |> pprint

#
# ### Multi-valued Protected Attribute
#
# Notice in the output of the previous cell that the column `race` has 6 different possible values "Native American", "African-American", *"Caucasian"*, "Hispanic", "Asian" and "Other".
#
# We support Multi-valued Protected attributes in both fairness algorithms and metrics. *The fairness algorithms by Researchers have been improved to generalize for multiple values of a protected attribute.*
#
# ## Load Neural Network Classifier
#
# We will use [MLJFlux](https://github.com/alan-turing-institute/MLJFlux.jl) to load Neural Network classifier. MLJFlux is an interface of Flux with MLJ. You don't need to explicitly import MLJFlux. MLJ does all that for you!!
#
# We use the `@load` macro to load Neural Network Classifier into main. Then using the Neural Network, we use `@pipeline` to add Continuous Encoder to the Neural Network Encoder. Continuous Encoder converts categorical strings to continuous values that support a much wider range(\~50) of models!!
#
Pkg.add("MLJFlux")
@load NeuralNetworkClassifier
model = @pipeline ContinuousEncoder NeuralNetworkClassifier

#
# ## Fairness Algorithm Wrappers
#
# We first wrap the Neural Network Classifier with Reweighing Algorithm. This wrapped Model is again wrapped with LinProg Postprocessing Algorithm.
#
# *Notice how usage of wrappers allows provides us with composability and enables you to apply unlimited algorithms on a single classifier.*
#
wrappedModel = ReweighingSamplingWrapper(classifier=model, grp=:race)
wrappedModel2 = LinProgWrapper(classifier=wrappedModel, grp=:race,
                                measure=false_positive_rate)

# ## Automatic evaluation using MLJ.evaluate
#
# Using the evaluate function from MLJ, you only need to pass your model, data and concerned metrics. MLJ handles the rest of the work internally. Note that you need to wrap the metrics to specify the protected attribute.
#
evaluate(wrappedModel2,
  X, y,
  measures=[
    Disparity(false_positive_rate, refGrp="Caucasian", grp=:race),
    MetricWrapper(accuracy, grp=:race)]) |> pprint

# # Finer Control (Advanced)
#
# You can get greater control than what is provided by `evaluate` function.
#
# * First, we need to get the train and test indices. This will be provided by the partition function.
# * `machine` is used to package the dataset and the wrapped Model (reused from before)
# * The machine fitted on the training rows.
#

train, test = partition(eachindex(y), 0.7, shuffle=true)
mach = machine(wrappedModel2, X, y)
fit!(mach, rows=train)

#
# Now we use `predict` function on the `machine` on the rows specified by `test`
#

ŷ = predict(mach, rows=test)
ŷ |> pprint

#
# ## Auditing Bias
#
# We use the concept of Fairness Tensors to avoid redundant calculations. Refer <https://www.ashrya.in/Fairness.jl/dev/fairtensor/> to learn more about Fairness Tensors.
#
# We pass predictions, ground-truth and protected values to the `fair_tensor` function.
#

ft = fair_tensor(ŷ, y[test], X[test, :race])
#
# ### [Disparity Calculation](https://www.ashrya.in/Fairness.jl/dev/measures/#Fairness.disparity)
#
# Disparity can be calculated by passing the following to the disparity `function` :
#
# * An array of fairness metrics from the ones [listed in README](https://www.ashrya.in/Fairness.jl/dev/fairtensor/)
# * Fairness tensor that we calculated in the previous step
# * Reference Group
# * `func` : disparity value for a metric M, group A and reference group B is `func(M(A), M(B))` . The default value for func is division and hence is an optional argument
#

df_disparity = disparity(
  [accuracy, false_positive_rate],
  ft, refGrp="Caucasian",
  func=(x, y)->(x-y)/y
)

#
# The values above show that Asians and African-Americans have a higher percentage of False Positive Rate w.r.t. the reference group Caucasian. On the other hand, Native Americans have a lower percentage of False Positive Rate w.r.t Caucasian. But these disparity values are better than the case if the Neural Network Classifier was directly used.
#
# ## [Parity Calculation](https://www.ashrya.in/Fairness.jl/dev/measures/#Fairness.parity)
#
# To calculate parity, we need to pass following to the `parity` function :
#
# * DataFrame output from the disparity function in the previous step
# * Custom Function to calculate parity based on disparity values.
#
# Scroll the output to the right to see the column for parity values.
#

parity(df_disparity,
  func= (x) -> abs(x)<0.4
)

#
# The above parity outputs show that parity constraints for False Positive Rate are satisfied only by the groups: Other and Caucasian.
#
# # Visualizing improvement by Fairness Algorithm
#
# Now we will use [VegaLite](https://www.queryverse.org/VegaLite.jl/stable/) to visualize the improvement in fairness metrics due to the fairness algorithms added in the form of wrappers. We shall also visualize the drop in accuracy due to the trade-off between accuracy and fairness.
#
# Note that **`wrappedModel2`** is the ML model we previously wrapped with Reweighing algorithm and LinProg Algorithm.
#
# Summary of what following code does :
#
# * Evaluate metric values using MLJ.evaluate for both: The Wrapped Model and the original model
# * Collect metric values  from the result of evaluate function
# * Create a DataFrame using the collected values that will later be used with VegaLite to plot the graphs
#
Pkg.add("VegaLite")
Pkg.add("DataFrames")
using VegaLite
using DataFrames

result = evaluate(wrappedModel2,
    X, y,
    measures=[
    Disparity(false_positive_rate, refGrp="Caucasian", grp=:race),
    MetricWrapper(accuracy, grp=:race)])

result_1 = evaluate(model,
    X, y,
    measures=[
    Disparity(false_positive_rate, refGrp="Caucasian", grp=:race),
    MetricWrapper(accuracy, grp=:race)])

n_grps = length(levels(X[!, :race]))
dispVals = collect(values(result.measurement[1]))
dispVals_1 = collect(values(result_1.measurement[1]))
accVals = collect(values(result.measurement[2]))
accVals_1 = collect(values(result_1.measurement[2]))

df = DataFrame(
  disparity=vcat(dispVals, dispVals_1),
  accuracy=vcat(accVals, accVals_1),
  algo=vcat(repeat(["Wrapped Model"],n_grps+1), repeat(["ML Model"],n_grps+1)),
  grp=repeat(collect(keys(result.measurement[1])), 2));

#
# ## Improvement in False Positive Rate Disparity Values
#

df |> @vlplot(
  :bar,
  column="grp:o",
  y={"disparity:q",axis={title="False Positive Rate Disparity"}},
  x={"algo:o", axis={title=""}},
  color={"algo:o"},
  spacing=20,
  config={
  view={stroke=:transparent},
  axis={domainWidth=1}
  }
)

#
#
# The above plot shows that there was a high bias against the group "African-American" in the NeuralNetworkClassifier(ML Model). The False Positive Rate Disparity value is greater than `2.0` for this group while its nearly to `1.0` for others. This means that a person belonging to the group "African-American" is twice as likely as other groups to be falsely predicted as a criminal who would re-offend!!
#
# But in the case of the wrapped model, the False Positive Rate disparity has been reduced for "African-American" to about `1.3` which is the same as most other groups.
#
# ## Accuracy Comparison
#
df |> @vlplot(
  :bar,
  column="grp:o",
  y={"accuracy:q",axis={title="Accuracy"}},
  x={"algo:o", axis={title=""}},
  color={"algo:o"},
  spacing=20,
  config={
  view={stroke=:transparent},
  axis={domainWidth=1}
  }
)
#
#
#
# The above plot shows that there is a drop in accuracy on using the wrapped model. This is a direct consequence of the fairness-accuracy tradeoff. So, we obtain a model that is fairer at the cost of accuracy.
#
# ## Fairness vs Accuracy Comparison across Algorithms
#
Pkg.add("Plots")
Pkg.add("PyPlot")
using Plots

function algorithm_comparison(algorithms, algo_names, X, y;
  refGrp, grp::Symbol=:class)
	grps = X[!, grp]
	categories = levels(grps)
	train, test = partition(eachindex(y), 0.7, shuffle=true)
	plot(title="Fairness vs Accuracy Comparison", seriestype=:scatter,
        xlabel="accuracy",
        ylabel="False Positive Rate Disparity refGrp="*refGrp,
        legend=:topleft, framestyle=:zerolines)
	for i in 1:length(algorithms)
		mach = machine(algorithms[i], X, y)
		fit!(mach, rows=train)
		ŷ = predict(mach, rows=test)
		if typeof(ŷ) <: MLJ.UnivariateFiniteArray
			ŷ = mode.(ŷ)
		end
		ft = fair_tensor(ŷ, y[test], X[test, grp])
		plot!([accuracy(ft)], [fpr(ft)/fpr(ft, grp=refGrp)],
      seriestype=:scatter, label=algo_names[i])
	end
	display(plot!())
end

algorithm_comparison([model, wrappedModel, wrappedModel2],
  ["NeuralNetworkClassifier", "Reweighing(Model)",
    "LinProg+Reweighing(Model)"], X, y,
    refGrp="Caucasian", grp=:race)

#
#
#
# # Concluding Remarks
#
# This toolkit has been designed to solve the numerous problems faced by both Policy Makers, Researchers, etc while using Fairness toolkits. Various innovative features of this package have been explicitly listed at <https://github.com/ashryaagr/Fairness.jl#what-fairnessjl-offers-over-its-alternatives>
#
# We are open to contributions. Feel free to open an issue on Github in case you want to contribute or have any confusion regarding the package. We would love to help you in getting started with this package.
#
# Finally, this work would have been impossible without the immense support, novel ideas, and efforts made by [Jiahao Chen](https://jiahao.github.io/), [Sebastian Vollmer](https://www.turing.ac.uk/people/researchers/sebastian-vollmer), and [Anthony Blaom](https://ablaom.github.io/).
#
# [Ashrya Agrawal](https://github.com/ashryaagr/) (ashryaagr@gmail.com)
#
# Link to Github Repository for Fairness.jl: <https://github.com/ashryaagr/Fairness.jl>
#
# Documentation: <https://www.ashrya.in/Fairness.jl/dev/>
#
# [nextjournal#output#96347bce-c38f-43fe-8426-d220e19780ed#result]:
# <https://nextjournal.com/data/QmdNHusG4ixmzbpU7FJ3WxxaRQVVA41o7MP5h5ibX7Ahh8?content-type=image/svg%2Bxml>
#
# [nextjournal#output#b0ab56a1-e5eb-4df3-b974-7a4319ac2b81#result]:
# <https://nextjournal.com/data/QmNUZfL9Kgu1hnujPKzbG8dXUJjPAf28pHcZwbpMJWTSpp?content-type=image/svg%2Bxml>
#
# [nextjournal#output#f69e3749-dccf-415c-be94-6d75b5f65639#result]:
# <https://nextjournal.com/data/QmRSVxCFZn2Hg7Lck2c7vEkcBaRcTqZi6hdUURUgdUwqZe?content-type=image/svg%2Bxml>
#
# <details id="com.nextjournal.article">


text:
"""
CellGravity allows sorting on gravity while keeping records of the original cell coordinate
"""
mutable struct CellGravity{M,I}
    gravity::M
    index::I
end

"""
CellInterval allows ordering a list by the cumulative proportion of the total gravity,
and plotting based on the fraction of total gravity.
"""
struct CellInterval{P,I}
    cumprop::P
    index::I
end

#= Sorting
isless is used to iteratively sort lists and search in funcitons like
`searchsortedfirst` and `partialsort!`. We define `isless` on `CellGravity.gravity` in
order to sortby margnitudes  using `searchsortedfirst`, while tracking the original index.
=#

import Base: isless, +
isless(x::CellGravity, y::CellGravity) = isless(x.gravity, y.gravity)
isless(x::CellGravity, y) = isless(x.gravity, y)
isless(x, y::CellGravity) = isless(x, y.gravity)

isless(x::CellInterval, y::CellInterval) = isless(x.cumprop, y.cumprop)
isless(x::CellInterval, y) = isless(x.cumprop, y)
isless(x, y::CellInterval) = isless(x, y.cumprop)

# Adding methods for + allows us to use sum() on arrays of CellGravity
+(x::CellGravity, y::CellGravity) = +(x.gravity, y.gravity)
+(x, y::CellGravity) = +(x, y.gravity)
+(x::CellGravity, y) = +(x.gravity, y)

"""
    HumanDispersal{R,W}(mode, human_pop, cellsize, scale, aggregator, human_exponent,
                        dist_exponent, dispersalperpop, max_dispersers, nshortlisted,
                        dest_shortlists, proportion_covered, human_buffer, distances)
    HumanDispersal(; read=:_default_, write=read,
                   mode=BatchGroups(),
                   human_pop,
                   cellsize=1.0,
                   scale=4,
                   aggregator=mean,
                   human_exponent=1.0,
                   dist_exponent=1.0,
                   dispersalperpop=1e-3,
                   max_dispersers=100.0,
                   nshortlisted=100)

Human-driven dispersal patterns using population data.

Transport connections between grid cells are calculated using distance and human population,
modified with the `human_exponent` and `dist_exponent` parameters. A shortlist of the most
connected cells is selected for use in the simulation.

The time taken for precalulation will depend on the `scale` argument. Values above 1
will downsample the grid to improve precalulation time and runtime performance. A high
scale value is good for use in a live interface.

## Keyword Arguments

$(FIELDDOCTABLE)
"""
@flattenable @bounds @description struct HumanDispersal{R,W,M,HP,CS,S,AG,HE,DE,DP,MD,SL,PC,B,D} <: ManualRule{R,W}
    # Field                | Flatten | Limits          | Description
    mode::M                | false   | _               | "Dispersal mode"
    human_pop::HP          | false   | _               | "An array match the grid size containing human population data."
    cellsize::CS           | false   | _               | "The size of the cell width, assuming they are square"
    scale::S               | false   | _               | ""
    aggregator::AG         | false   | _               | "A function that aggregates scaled down cells"
    human_exponent::HE     | true    | (1.0, 3.0)      | "Human population exponent"
    dist_exponent::DE      | true    | (1.0, 3.0)      | "Distance exponent"
    dispersalperpop::DP    | true    | (0.0, 1e-8)     | "Scales the number of dispersing individuals by human activity (ie population^human_exponent)"
    max_dispersers::MD     | true    | (50.0, 10000.0) | "Maximum number of dispersers in a dispersal events"
    nshortlisted::SL       | false   | _               | "Length of dest shortlist"
    dest_shortlists::PC    | false   | _               | "Array of destination vectors for each cell. Automatically calculated"
    human_buffer::B        | false   | _               | "Buffer array used in precalculation"
    distances::D           | false   | _               | "Buffer array used in precalculation"
    function HumanDispersal{R,W,M,HP,CS,S,AG,HE,DE,DP,MD,SL,PC,B,D}(
        mode::M, human_pop::HP, cellsize::CS, scale::S, aggregator::AG,
        human_exponent::HE, dist_exponent::DE, dispersalperpop::DP,
        max_dispersers::MD, nshortlisted::SL, dest_shortlists::PC,
        human_buffer::B, distances::B
       ) where {R,W,M,HP,CS,S,AG,HE,DE,DP,MD,SL,PC,B,D}

        precalc_human_dispersal!(dest_shortlists, human_pop, cellsize, scale, aggregator,
                                 human_exponent, dist_exponent, nshortlisted, human_buffer, distances)
        new{R,W,M,HP,CS,S,AG,HE,DE,DP,MD,SL,PC,B,D}(mode, human_pop, cellsize, scale, aggregator, human_exponent,
                                 dist_exponent, dispersalperpop, max_dispersers, nshortlisted, dest_shortlists,
                                 human_buffer, distances)
    end
end
HumanDispersal(; read=:_default_, write=read, mode=BatchGroups(), human_pop, cellsize=1.0, scale=4,
               aggregator=mean, human_exponent=1.0, dist_exponent=1.0,
               dispersalperpop=1e-3, max_dispersers=100.0, nshortlisted=100) = begin
    # Allocate memory
    human_buffer = initdownsample(human_pop, scale)
    distances = initdownsample(human_pop, scale)
    dest_shortlists = alloc_dest_shortlist(nshortlisted, size(human_buffer))

    HumanDispersal{read,write}(mode, human_pop, cellsize, scale, aggregator, human_exponent,
                   dist_exponent, dispersalperpop, max_dispersers, nshortlisted,
                   dest_shortlists, human_buffer, distances)
end

# Minimal Array interface

size(rule::HumanDispersal) = size(parent(rule))
getindex(rule::HumanDispersal, I...) = getindex(parent(rule), I...)
parent(rule::HumanDispersal) = rule.dest_shortlists
mode(rule::HumanDispersal) = rule.mode

# Precalculation ###################################################

# Define types used in the precalculation
const Index = Tuple{Int16,Int16}
const Interval = CellInterval{Float32,Index}
const Precalc = Union{Vector{Interval},Missing}
const Gravity = CellGravity{Float32,Index}
const Prop = Union{Float32,Missing}

# Memory preallocation utilities
alloc_dest_shortlist(nshortlisted, (h, w)) =
    Precalc[Vector{Interval}(undef, nshortlisted) for i in 1:h, j in 1:w]

alloc_threads(human_buffer, nshortlisted) =
    [alloc_gravities(human_buffer, nshortlisted) for thread in 1:Threads.nthreads()]

alloc_gravities(human, nshortlisted) = begin
    gravities = Matrix{Gravity}(undef, size(human)...)
    for j in 1:size(gravities, 2), i in 1:size(gravities, 1)
        gravities[i, j] = Gravity(0.0f0, (i, j))
    end
    gravity_vector = deepcopy(vec(gravities))
    gravities, gravity_vector
end

# Precalculate a dispersal shortlist for each cell
precalc_human_dispersal!(dest_shortlists, human_pop, cellsize, scale, aggregator,
                        human_exponent, dist_exponent, nshortlisted, human_buffer, distances) = begin
    # Limit shortlist cells to the total available
    @assert nshortlisted <= length(human_buffer)

    # Copy downsampled human_pop matrix to human_buffer
    downsample!(human_buffer, human_pop, aggregator, scale)

    precalc_human_exponent!(human_buffer, human_exponent)
    precalc_distances!(distances, dist_exponent, cellsize, scale)

    # Calculate cell gravities by column in separate threads
    thread_alloc = alloc_threads(human_buffer, nshortlisted)
    Threads.@threads for j in 1:size(human_buffer, 2)
        precalc_col!(dest_shortlists, nshortlisted, human_buffer, distances, thread_alloc, j)
    end
end

# Raise the human population/activity matrix to some exponent
precalc_human_exponent!(human_pop::Matrix, human_exponent::Number) =
    human_pop .^= human_exponent

# Generate a matrix of distances
precalc_distances!(distances::Matrix, dist_exponent, cellsize, scale) = begin
    for j in 1:size(distances, 2), i in 1:size(distances, 1)
        distances[i, j] = (centercenter_distance(i, j) * cellsize * scale) ^ dist_exponent
    end
    # First cell uses the mean distance from cell centre
    distances[1, 1] = (cellsize * scale / 6 * (sqrt(2) + log(1 + sqrt(2)))) ^ dist_exponent
    return distances
end

centercenter_distance(a, b) = sqrt((a - 1)^2 + (b - 1)^2)

# Precalculate human dispersal shortlist for every cell in a column.
# Working on columns is a clean way to spread the load accross multiple processors.
function precalc_col!(dest_shortlists, nshortlisted, human, distances, thread_alloc, j)
    col_alloc = thread_alloc[Threads.threadid()]
    for i = 1:size(human, 1)
        precalc_cell!(dest_shortlists, nshortlisted, human, distances, col_alloc, i, j)
    end
end

function precalc_cell!(dest_shortlists, nshortlisted, human, distances, (gravities, gravity_vector), i, j)
    if ismissing(human[i, j])
        dest_shortlists[i, j] = missing
        return
    end
    # Calculate the gravity for all cells in the grid
    assign_gravities!(gravities, human, distances, i, j)
    gravity_vector .= vec(gravities)
    # Sort the top nshortlisted gravities in-place, highest first.
    partialsort!(gravity_vector, nshortlisted, rev=true)
    # Select the top nshortlisted gravities, low to high
    gravity_shortlist = view(gravity_vector, 1:nshortlisted)
    # Convert shortlies gravities to intervals
    gravity2inverval!(dest_shortlists[i, j], gravity_shortlist)
end

# Calculate the gravity for all cells in the grid
function assign_gravities!(gravities, human, distances, i, j)
    T = typeof(first(gravities).gravity)
    for jj = 1:size(human, 2), ii = 1:size(human, 1)
        gravities[ii, jj].gravity = if ismissing(human[ii, jj])
            # Missing values are assigned a zero gravity to miss the shortlist
            zero(T)
        else
            celldist = abs(i - ii) + 1, abs(j - jj) + 1
            (human[ii, jj] / distances[celldist...])
        end
    end
end

# Convert the gravity of a cell to an interval than can
# be randomly selected with `rand()` and `searchsortedfirst`
function gravity2inverval!(interval_shortlist, gravity_shortlist)
    # Sum gravities in the shortlist
    shortlist_sum = sum(gravity_shortlist)
    nshortlisted = length(gravity_shortlist)

    cumprop = 0.0f0
    # Create a list of intervals from the sorted list of gravities.
    # This will be used to randomly choose cells from the
    # distribution of gravities
    for (n, m) = enumerate(reverse(gravity_shortlist))
        # Calculaate proportion of current gravity in the complete shortlist
        gravityprop = m.gravity / shortlist_sum
        # Track cumulative proportion for use with `searchsortedfirst()`
        # In case of float innacuracy, finish the distribution at 1.0
        cumprop = n == nshortlisted ? oneunit(cumprop) : cumprop + gravityprop
        interval_shortlist[n] = CellInterval(cumprop, m.index)
    end
    interval_shortlist
end

# DynamicGrids Interface ###################################################

@inline applyrule!(data, rule::HumanDispersal{R,W}, population, index) where {R,W} = begin
    population == zero(population) && return
    dispersalprob = rule.human_pop[index...] * rule.dispersalperpop
    ismissing(dispersalprob) && return
    shortlist = rule.dest_shortlists[downsample_index(index, rule.scale)...]
    #ismissing(shortlist) && return

    #= This formulation introduces a bias where total_dispersers is
    close to max_disersers, for example, for chunks much less than max chunk
    chunks size will allways equal the max chunk size, so the amount wont be random.
    =#
    dispersed = disperse!(data[W], mode(rule), rule, shortlist, dispersalprob, population, index)
    data[W][index...] -= dispersed
    return
end

abstract type TransportMode end

@default_kw @bounds struct HeirarchicalGroups{S} <: TransportMode
    scalar::S | 1e-8 | (1e-6, 1e-9)
end

@inline disperse!(data::WritableGridData, mode::HeirarchicalGroups, rule::HumanDispersal,
                  shortlist, dispersalprob, population, index) = begin
    dispersed = zero(population)
    nevents = rand(Binomial(human_pop, dispersalperpop))
    for i in 1:nevents
        maybedispersing = rand(Poissonn(population * mode.scalar))
        # Just exit the loop if we exceed the existing populaiton
        (maybedispersing + dispersed > population) && break
        dispersed += disperse2dest!(data, rule, shortlist, maybedispersing)
    end
    dispersed
end

struct BatchGroups <: TransportMode end

@inline disperse!(data::WritableGridData, mode::BatchGroups, rule::HumanDispersal,
                  shortlist, dispersalprob, population, index) = begin
    # Find the expected number of dispersers given population and dispersal prob
    total_dispersers = trunc(Int, min(population * dispersalprob, population))

    # Maximum number of discrete dispersers in any single dispersal event
    # We never know this number, is this defensible
    max_dispersers = trunc(Int, rule.max_dispersers)

    # Simulate (possibly) multiple dispersal events from the cell during the timeframe
    dispersed = zero(population)
    while dispersed < total_dispersers
        # Select a subset of the remaining dispersers for a dispersal event
        maybedispersing = min(rand(1:max_dispersers), total_dispersers - dispersed)
        dispersed += disperse2dest!(data, rule, shortlist, maybedispersing)
        # Track how many have allready dispersed
    end
    dispersed
end

#= Maybe write the group to a randomised location in the destination array
unless it falls outside the grid or is masked, in which case we
say the event just never happened.
=#
disperse2dest!(data::DynamicGrids.WritableGridData, rule, shortlist, maybedispersing) = begin
    dest_id = min(length(shortlist), searchsortedfirst(shortlist, rand()))
    # Randomise cell destination within upsampled destination cells
    upsample = upsample_index(shortlist[dest_id].index, rule.scale)
    dest_index = upsample .+ (rand(0:rule.scale-1), rand(0:rule.scale-1))
    # Skip dispsal to upsampled dest cells that are masked or out of bounds, and try again
    if !DynamicGrids.ismasked(data, dest_index...) &&
        DynamicGrids.isinbounds(dest_index, data)
        # Disperse to the cell.
        data[dest_index...] += maybedispersing
        maybedispersing
    else
        zero(maybedispersing)
    end
end

# Examing the generated shorlists ######################################33

"""
    populate!(A::AbstractMatrix, rule::HumanDispersal, [I...])
    populate!(A::AbstractMatrix, cells::AbstractArray, [scale=1])

Populate a matrix with the precalculated destinations from a
[`HumanDispersal`](@ref) rule - either all of the or some subset
if passed the `I...` indexing arguments. This is useful for plotting
dispersal destinations, especially when used with GeoData.jl
"""
populate!(A::AbstractMatrix, rule::HumanDispersal) =
    populate!(A, rule.dest_shortlists, rule.scale)
populate!(A::AbstractMatrix, rule::HumanDispersal, I...) =
    populate!(A::AbstractMatrix, rule[I...], rule.scale)
# All shortlists for all cells
populate!(A::AbstractMatrix, shortlists::AbstractArray, scale::Int=1) = begin
    for I in CartesianIndices(shortlists)
        populate!(A, shortlists[I], scale)
    end
    return A
end
# Single shortlist for one cell
populate!(A::AbstractMatrix, cells::AbstractVector{<:Union{<:CellInterval,Missing}}, scale::Int=1) = begin
    lastcumprop = 0.0
    for cell in cells
        I = upsample_index(cell.index, scale)
        prop = cell.cumprop - lastcumprop
        if ismissing(A[I...])
            A[I...] = prop
        else
            A[I...] += prop
        end
        lastcumprop = cell.cumprop
    end
    return A
end
populate!(A::AbstractMatrix, cells::Missing, scale::Int=1) = missing

"""
    populate(rule::HumanDispersal, [I...])

Returns an array the size of human population matrix filled
with all destination locataion, or with destinations specific
to the passed-in indices `I`.
"""
populate(rule::HumanDispersal, args...) =
    populate!(zeros(Float32, size(rule.human_pop)), rule, args...)

"""
    populate(cells::AbstractVector, size::Tuple, [scale::Int=1])

Returns an array of size `size` populated from the vector
of positions in `cells` rescaled by `scale`.
"""
populate(cells::AbstractVector, size::Tuple, scale::Int=1) =
    populate!(zeros(Float32, size), cells, scale)


text:
using GeoData, RasterDataSources, Test, Dates, NCDatasets, ArchGDAL

# Too big to test
# @testset "load WorldClim Weather" begin
#     # Weather time-series
#     dates = (Date(2001), Date(2002))
#     ser = series(WorldClim{Weather}, (:prec,); date=dates, window=(Lat(600:1900), Lon(600:1900)))
#     ser[Date(2001, 1)][:prec]
#     # Select Australia, using regular lat/lon selectors
#     A = geoarray(WorldClim{Weather}, :prec; date=DateTime(2001, 05), mappedcrs=EPSG(4326))
# end

@testset "load WorldClim Climate" begin
    # Weather time-series
    ser = series(WorldClim{Climate}, (:prec,); res="10m", month=Jan:March)
    ser[Jan][:prec]
    # Select Australia, using regular lat/lon selectors
    A = geoarray(WorldClim{Climate}, :prec; month=1, mappedcrs=EPSG(4326))
    A[Lat(Between(-10, -45)), Lon(Between(110, 160))]
    st = stack(WorldClim{BioClim}, (1, 2))
    st[:BIO1]
    @test st isa DiskStack
    @test A isa GDALarray
end

@testset "load WorldClim BioClim" begin
    A = geoarray(WorldClim{BioClim}, 1; mappedcrs=EPSG(4326))
    A[Lat(Between(-10, -45)), Lon(Between(110, 160))]
    @test A isa GDALarray
end

@testset "load CHELSA BioClim" begin
    A = geoarray(CHELSA{BioClim}, 1; mappedcrs=EPSG(4326))
    A[Lat(Between(-10, -45)), Lon(Between(110, 160))]
    st = stack(CHELSA{BioClim}, (1, 2))
    @test A isa GDALarray
    @test st isa DiskStack
    @test st[:BIO2] isa GeoArray
end

@testset "load EarthEnv HabitatHeterogeneity" begin
    A = geoarray(EarthEnv{HabitatHeterogeneity}, :cv; mappedcrs=EPSG(4326))
    A[Lat(Between(-10, -45)), Lon(Between(110, 160))]
    st = stack(EarthEnv{HabitatHeterogeneity}, (:cv, :evenness))
    @test A isa GDALarray
    @test st isa DiskStack
    @test st[:evenness] isa GeoArray
end

@testset "load EarthEnv LandCover" begin
    A = geoarray(EarthEnv{LandCover}, 2; mappedcrs=EPSG(4326))
    A[Lat(Between(-10, -45)), Lon(Between(110, 160))]
    @test A isa GDALarray
end

@testset "load ALWB" begin
    A = geoarray(ALWB{Values,Day}, :ss_pct; date=DateTime(2019, 10, 19))
    st = stack(ALWB{Values,Day}, (:s0_pct, :ss_pct); date=DateTime(2019, 10, 19))
    dates = DateTime(2019, 10, 19):Day(1):DateTime(2019, 11, 20)
    s = series(ALWB{Values,Day}, (:s0_pct, :ss_pct); date=dates)
    @test A isa NCDarray
    @test st isa DiskStack
    @test s isa GeoSeries
end

if Sys.islinux()
    @testset "load AWAP" begin
        A = geoarray(AWAP, :rainfall; date=DateTime(2019, 10, 19))
        st = stack(AWAP, (:tmax, :rainfall); date=DateTime(2019, 10, 19))
        st[:tmax]
        dates = DateTime(2019, 09, 19), DateTime(2019, 11, 19)
        s = series(AWAP; date=dates)
        s[DateTime(2019, 10, 1)][:solar]
        @test A isa GDALarray
        @test st isa DiskStack
        @test s isa GeoSeries
    end
end


text:
# # Input PDF parametrisation and priors
#
# An important part of PDF fitting is defining a useful parametrisation for the
# PDF shapes, as well as meaningful prior distributions that encode our knowledge
# of the problem.
#
# In this notebook, we explore two different approaches:
# * *Full Dirichlet*
# * *Valence shape + Dirichlet*
#
# In the end, it seems like the latter option makes more sense for us and is
# therefore implemented elsewhere in the `PartonDensity` package.
# We demonstrate why below.

using Distributions, Plots, SpecialFunctions, Printf
const sf = SpecialFunctions;

# ## "Full Dirichlet" approach
#
# A clean way to ensure the momentum sum rule would be to sample different
# contributions of the momentrum density integral from a Dirichlet distribution,
# then use these weights to set the parameters on the individual Beta distributions.
# However, in practice this is non-trvial as we also want to fix the normalisation
# of the *number densities* of the valance contributions.

# 9 components of decreasing importance

dirichlet = Dirichlet([3., 2., 1, 0.5, 0.3, 0.2, 0.1, 0.1, 0.1])
data = rand(dirichlet, 1000);

# Have a look

plot()
for i in 1:9
    histogram!(data[i,:], bins=range(0, stop=1, length=20), alpha=0.7)
end
plot!(xlabel="I_i = A_i B_i")

# This would be great as the sum rule is automatically conserved

sum(data, dims=1)

# But, it is non-trival to define valence params from this

I = rand(dirichlet)

# Valance u component

λ_u = rand(Uniform(0, 1))
K_u = rand(Uniform(0, 10))

# Integral of number density must = 2

A_u = 2 / sf.beta(λ_u, K_u+1)

# integral of momentum density can be fixed by I[1]

I_1 = A_u * sf.beta(λ_u+1, K_u+1);

# Could use a root-finder to find K_u given I_1 and λ_u...
# Could be nasty to sample from though...

I_1 = 2 * (sf.beta(λ_u+1, K_u+1)/sf.beta(λ_u, K_u+1));

using Roots

function func_to_solve(K_u)
    return I_1 - 2 * (sf.beta(λ_u+1, K_u+1) / sf.beta(λ_u, K_u+1))
end

K_u ≈ find_zero(func_to_solve, (0, 10), Bisection())

# While this approach might be nice, there are two issues in practice:
# * It is difficult to set sensible priors on $\lambda_u$ that imply priors on $K_u$, and similarly for $\lambda_d$ and $K_d$
# * The problem is overconstrained and we hav to use a root finder. This is rather fragile, and could fail for certain parameter combinations, such as we might find in a fit.

# ## "Valence shape + Dirichlet" approach
#
# We can handle this more elegantly (maybe?) by specifying constraints on the
# valence params through the shape of their Beta distributions, then using a
# Dirichlet to specify the weights of the gluon and sea components. The problem
# here is it isn't clear how to specify that the d contribution must be less than
# the u contribution, but it is possible to do this indirectly through priors on
# the shape parameters. This will however require some further investigation.

x = range(0, stop=1, length=50)

# High-level priors
# Looks like we maybe want to change lambda and K priors to boost these components

λ_u = 0.7 #rand(Uniform(0, 1))
K_u = 4 #rand(Uniform(2, 10))
λ_d = 0.5 #rand(Uniform(0, 1))
K_d = 6 #rand(Uniform(2, 10))

u_V = Beta(λ_u, K_u+1)
A_u = 2 / sf.beta(λ_u, K_u+1)

d_V = Beta(λ_d, K_d+1)
A_d = 1 / sf.beta(λ_d, K_d+1)

# Integral contributions

I_u = A_u * sf.beta(λ_u+1, K_u+1)
I_d = A_d * sf.beta(λ_d+1, K_d+1)

plot(x, x .* A_u .* x.^λ_u .* (1 .- x).^K_u * 2, alpha=0.7, label="x u(x)", lw=3)
plot!(x, x .* A_d .* x.^λ_d .* (1 .- x).^K_d, alpha=0.7, label="x d(x)", lw=3)
plot!(xlabel="x", legend=:topright)

#

@printf("I_u = %.2f\n", I_u)
@printf("I_d = %.2f\n", I_d)

# The remaining 7 integrals can be dirichlet-sampled with decreasing importance

remaining = 1 - (I_u + I_d)
dirichlet = Dirichlet([3., 2., 1, 0.5, 0.3, 0.2, 0.1])
I = rand(dirichlet) * remaining;
sum(I) ≈ remaining

# Gluon contributions

λ_g1 = rand(Uniform(-1, 0))
λ_g2 = rand(Uniform(0, 1))
K_g = rand(Uniform(2, 10))
A_g2 = I[1] / sf.beta(λ_g2+1, K_g+1)
A_g1 = I[2] / sf.beta(λ_g1+1, 5+1);

# Sea quark contributions

λ_q = rand(Uniform(-1, 0))
A_ubar = I[3] / (2 * sf.beta(λ_q+1, 5+1))
A_dbar = I[4] / (2 * sf.beta(λ_q+1, 5+1))
A_s = I[5] / (2 * sf.beta(λ_q+1, 5+1))
A_c = I[6] / (2 * sf.beta(λ_q+1, 5+1))
A_b = I[7] / (2 * sf.beta(λ_q+1, 5+1));

total = A_u * sf.beta(λ_u+1, K_u+1) + A_d * sf.beta(λ_d+1, K_d+1)
total += A_g1 * sf.beta(λ_g1+1, 5+1) + A_g2 * sf.beta(λ_g2+1, K_g+1)
total += 2 * (A_ubar + A_dbar + A_s + A_c + A_b) * sf.beta(λ_q+1, 5+1)
total ≈ 1

#

x = 10 .^ range(-2, stop=0, length=500)

# How does it look?

xg2 = A_g2 * x.^λ_g2 .* (1 .- x).^K_g
xg1 = A_g1 * x.^λ_g1 .* (1 .-x).^5
plot(x, x .* A_u .* x.^λ_u .* (1 .- x).^K_u * 2, alpha=0.7, label="x u(x)", lw=3)
plot!(x, x .* A_d .* x.^λ_d .* (1 .- x).^K_d, alpha=0.7, label="x d(x)", lw=3)
plot!(x, xg1 + xg2, alpha=0.7, label="x g(x)", lw=3)
plot!(x, A_ubar * x.^λ_q .* (1.0 .- x).^5, alpha=0.7, label="x ubar(x)", lw=3)
plot!(x, A_dbar * x.^λ_q .* (1.0 .- x).^5, alpha=0.7, label="x dbar(x)", lw=3)
plot!(x, A_s * x.^λ_q .* (1.0 .- x).^5, alpha=0.7, label="x s(x)", lw=3)
plot!(x, A_c * x.^λ_q .* (1.0 .- x).^5, alpha=0.7, label="x c(x)", lw=3)
plot!(x, A_b * x.^λ_q .* (1.0 .-  x).^5, alpha=0.7, label="x b(x)", lw=3)
plot!(xlabel="x", legend=:bottomleft, xscale=:log, ylims=(1e-8, 10), yscale=:log)

# ### Prior predictive checks
#
# We can start to visualise the type of PDFs that are allowed by the
# combination of the choice of parametrisation and prior distributions
# with some simple prior predictive checks, as done below...

N = 100
alpha = 0.03
total = Array{Float64, 1}(undef, N)
first = true
leg = 0

plot()
for i in 1:N

    λ_u_i = rand(Uniform(0, 1))
    K_u_i = rand(Uniform(2, 10))
    λ_d_i = rand(Uniform(0, 1))
    K_d_i = rand(Uniform(2, 10))
    A_u_i = 2 / sf.beta(λ_u_i, K_u_i+1)
    A_d_i = 1 / sf.beta(λ_d_i, K_d_i+1)
    I_u_i = A_u * sf.beta(λ_u_i+1, K_u_i+1)
    I_d_i = A_d * sf.beta(λ_d_i+1, K_d_i+1)
    u_V_i = Beta(λ_u_i, K_u_i+1)
    d_V_i = Beta(λ_d_i, K_d_i+1)

    remaining_i = 1 - (I_u_i + I_d_i)
    dirichlet_i = Dirichlet([3., 2., 1, 0.5, 0.3, 0.2, 0.1])
    I_i = rand(dirichlet_i) * remaining_i

    λ_g1_i = rand(Uniform(-1, 0))
    λ_g2_i = rand(Uniform(0, 1))
    K_g_i = rand(Uniform(2, 10))
    A_g2_i = I_i[1] / sf.beta(λ_g2_i+1, K_g_i+1)
    A_g1_i = I_i[2] / sf.beta(λ_g1_i+1, 5+1)

    λ_q_i = rand(Uniform(-1, 0))
    A_ubar_i = I_i[3] / (2 * sf.beta(λ_q_i+1, 5+1))
    A_dbar_i = I_i[4] / (2 * sf.beta(λ_q_i+1, 5+1))
    A_s_i = I_i[5] / (2 * sf.beta(λ_q_i+1, 5+1))
    A_c_i = I_i[6] / (2 * sf.beta(λ_q_i+1, 5+1))
    A_b_i = I_i[7] / (2 * sf.beta(λ_q_i+1, 5+1))

    total[i] = A_u_i * sf.beta(λ_u_i+1, K_u_i+1) + A_d_i * sf.beta(λ_d_i+1, K_d_i+1)
    total[i] += A_g1_i * sf.beta(λ_g1_i+1, 5+1) + A_g2_i * sf.beta(λ_g2_i+1, K_g_i+1)
    total[i] += 2 * (A_ubar_i + A_dbar_i + A_s_i + A_c_i + A_b_i) * sf.beta(λ_q_i+1, 5+1)

    xg2_i = A_g2_i * x.^λ_g2_i .* (1 .- x).^K_g_i
    xg1_i = A_g1_i * x.^λ_g1_i .* (1 .- x).^5
    plot!(x, [x .* A_u_i .* x.^λ_u_i .* (1 .- x).^K_u_i * 2], alpha=alpha, color="blue", lw=3)
    plot!(x, x .* A_d_i .* x.^λ_d_i .* (1 .- x).^K_d_i, alpha=alpha, color="orange", lw=3)
    plot!(x, xg1_i + xg2_i, alpha=alpha, color="green", lw=3)
    plot!(x, A_ubar_i * x.^λ_q_i .* (1 .- x).^5, alpha=alpha, color="red", lw=3)
    plot!(x, A_dbar_i * x.^λ_q_i .* (1 .- x).^5, alpha=alpha, color="purple", lw=3)
    plot!(x, A_s_i * x.^λ_q_i .* (1 .- x).^5, alpha=alpha, color="brown", lw=3)
    plot!(x, A_c_i * x.^λ_q_i .* (1 .- x).^5, alpha=alpha, color="pink", lw=3)
    plot!(x, A_b_i * x.^λ_q_i .* (1 .- x).^5, alpha=alpha, color="grey", lw=3)
end

plot!(xlabel="x", ylabel="x f(x)", xscale=:log, legend=false,
      ylims=(1e-8, 10), yscale=:log)

# Looks like naive priors need some work...

# ## PDF Parametrisation interface
#
# `PartonDensity` provides a handy interface to the "Valence shape + Dirichlet" style parametrisation, as demonstrated here.

using PartonDensity

hyper_params = PDFParameters(λ_u=0.5, K_u=4.0, λ_d=0.6, K_d=6.0, λ_g1=-0.37, λ_g2=-0.7,
                             K_g=6.0, λ_q=0.5, seed=5, weights=[50., 0.5, 5., 5., 3., 2., 1.]);

plot_input_pdfs(hyper_params)

int_xtotx(hyper_params) ≈ 1


text:
# * Term symbols

"""
    struct Term

Represents a term symbol with specific parity in LS-coupling.

Only the ``L``, ``S`` and parity values of the symbol ``{}^{2S+1}L_{J}`` are specified. To
specify a _level_, the ``J`` value would have to be specified separately. The valid ``J``
values corresponding to given ``L`` and ``S`` fall in the following range:

```math
|L - S| \\leq J \\leq L+S
```

# Constructors

    Term(L::Real, S::Real, parity::Union{Parity,Integer})

Constructs a `Term` object with the given ``L`` and ``S`` quantum numbers and parity. `L`
and `S` both have to be convertible to `HalfInteger`s and `parity` must be of type
[`Parity`](@ref) or `±1`.

See also: [`@T_str`](@ref)

# Properties

To access the quantum number values, you can use the `.L`, `.S` and `.parity` properties to
access the ``L``, ``S`` and parity values (represented with [`Parity`](@ref)), respectively.
E.g.:

```jldoctest
julia> t = Term(2, 1//2, p"odd")
²Dᵒ

julia> t.L, t.S, t.parity
(2, 1/2, odd)
```
"""
struct Term
    L::HalfInt
    S::HalfInt
    parity::Parity
    function Term(L::HalfInteger, S::HalfInteger, parity::Parity)
        L >= 0 || throw(DomainError(L, "Term symbol can not have negative L"))
        S >= 0 || throw(DomainError(S, "Term symbol can not have negative S"))
        new(L, S, parity)
    end
end
Term(L::Real, S::Real, parity::Union{Parity,Integer}) =
    Term(convert(HalfInteger, L), convert(HalfInteger, S), convert(Parity, parity))

"""
    parse(::Type{Term}, ::AbstractString) -> Term

Parses a string into a [`Term`](@ref) object.

```jldoctest
julia> parse(Term, "4Po")
⁴Pᵒ

julia> parse(Term, "⁴Pᵒ")
⁴Pᵒ
```

See also: [`@T_str`](@ref)
"""
function Base.parse(::Type{Term}, s::AbstractString)
    m = match(r"([0-9]+)([A-Z]|\[[0-9/]+\])([oe ]{0,1})", s)
    m2 = match(r"([¹²³⁴⁵⁶⁷⁸⁹⁰]+)([A-Z]|\[[0-9/]+\])([ᵒᵉ]{0,1})", s)
    Sstr,Lstr,Pstr = if !isnothing(m)
        m[1],m[2],m[3]
    elseif !isnothing(m2)
        from_superscript(m2[1]),m2[2],m2[3]
    else
        throw(ArgumentError("Invalid term string $s"))
    end
    L = lowercase(Lstr)
    L = if L[1] == '['
        L = strip(L, ['[',']'])
        if occursin("/", L)
            Ls = split(L, "/")
            length(Ls) == 2 && length(Ls[1]) > 0 && length(Ls[2]) > 0 ||
                throw(ArgumentError("Invalid term string $(s)"))
            Rational(parse(Int, Ls[1]),parse(Int, Ls[2]))
        else
            parse(Int, L)
        end
    else
        findfirst(L, spectroscopic)[1]-1
    end
    denominator(L) ∈ [1,2] || throw(ArgumentError("L must be integer or half-integer"))
    S = (parse(Int, Sstr) - 1)//2
    Term(L, S, Pstr == "o" || Pstr == "ᵒ" ? p"odd" : p"even")
end

"""
    @T_str -> Term

Constructs a [`Term`](@ref) object out of its canonical string representation.

```jldoctest
julia> T"1S"
¹S

julia> T"4Po"
⁴Pᵒ

julia> T"⁴Pᵒ"
⁴Pᵒ

julia> T"2[3/2]o" # jK coupling, common in noble gases
²[3/2]ᵒ
```
"""
macro T_str(s::AbstractString)
    parse(Term, s)
end

Base.zero(::Type{Term}) = T"1S"

"""
    multiplicity(t::Term)

Returns the spin multiplicity of the [`Term`](@ref) `t`, i.e. the
number of possible values of ``J`` for a given value of ``L`` and
``S``.

# Examples

```jldoctest
julia> multiplicity(T"¹S")
1

julia> multiplicity(T"²S")
2

julia> multiplicity(T"³P")
3
```
"""
multiplicity(t::Term) = convert(Int, 2t.S + 1)
"""
    weight(t::Term)

Returns the statistical weight of the [`Term`](@ref) `t`, i.e. the
number of possible microstates: ``(2S+1)(2L+1)``.

# Examples

```jldoctest
julia> weight(T"¹S")
1

julia> weight(T"²S")
2

julia> weight(T"³P")
9
```
"""
weight(t::Term) = (2t.L + 1) * multiplicity(t)

import Base.==
==(t1::Term, t2::Term) = ((t1.L == t2.L) && (t1.S == t2.S) && (t1.parity == t2.parity))

import Base.<
<(t1::Term, t2::Term) = ((t1.S < t2.S) || (t1.S == t2.S) && (t1.L < t2.L)
                         || (t1.S == t2.S) && (t1.L == t2.L) && (t1.parity < t2.parity))
import Base.isless
isless(t1::Term, t2::Term) = (t1 < t2)

Base.hash(t::Term) = hash((t.L,t.S,t.parity))

include("xu2006.jl")

"""
    xu_terms(ℓ, w, p)

Return all term symbols for the orbital `ℓʷ` and parity `p`; the term
multiplicity is computed using [`AtomicLevels.Xu.X`](@ref).

# Examples

```jldoctest
julia> AtomicLevels.xu_terms(3, 3, parity(c"3d3"))
17-element Array{Term,1}:
 ²P
 ²D
 ²D
 ²F
 ²F
 ²G
 ²G
 ²H
 ²H
 ²I
 ²K
 ²L
 ⁴S
 ⁴D
 ⁴F
 ⁴G
 ⁴I
```
"""
function xu_terms(ℓ::Int, w::Int, p::Parity)
    ts = map(((w//2 - floor(Int, w//2)):w//2)) do S
        S′ = 2S |> Int
        map(L -> repeat([Term(L,S,p)], Xu.X(w, ℓ, S′, L)), 0:w*ℓ)
    end
    vcat(vcat(ts...)...)
end

"""
    terms(orb::Orbital, w::Int=one(Int))

Returns a list of valid LS term symbols for the orbital `orb` with `w`
occupancy.

# Examples

```jldoctest
julia> terms(o"3d", 3)
8-element Array{Term,1}:
 ²P
 ²D
 ²D
 ²F
 ²G
 ²H
 ⁴P
 ⁴F
```
"""
function terms(orb::Orbital, w::Int=one(Int))
    ℓ = orb.ℓ
    g = degeneracy(orb)
    w > g && throw(DomainError(w, "Invalid occupancy $w for $orb with degeneracy $g"))
    # For shells that are more than half-filled, we instead consider
    # the equivalent problem of g-w coupled holes.
    (w > g/2 && w != g) && (w = g - w)

    p = parity(orb)^w
    if w == 1
        [Term(ℓ,1//2,p)] # Single electron
    elseif ℓ == 0 && w == 2 || w == degeneracy(orb) || w == 0
        [Term(0,0,p"even")] # Filled ℓ shell
    else
        xu_terms(ℓ, w, p) # All other cases
    end
end

"""
    terms(config)

Generate all final ``LS`` terms for `config`.

# Examples

```jldoctest
julia> terms(c"1s")
1-element Array{Term,1}:
 ²S

julia> terms(c"1s 2p")
2-element Array{Term,1}:
 ¹Pᵒ
 ³Pᵒ

julia> terms(c"[Ne] 3d3")
7-element Array{Term,1}:
 ²P
 ²D
 ²F
 ²G
 ²H
 ⁴P
 ⁴F
```
"""
function terms(config::Configuration{O}) where {O<:AbstractOrbital}
    ts = map(config) do (orb,occ,state)
        terms(orb,occ)
    end

    final_terms(ts)
end

"""
    count_terms(orbital, occupation, term)

Count how many times `term` occurs among the valid terms of `orbital^occupation`.

```jldoctest
julia> count_terms(o"1s", 2, T"1S")
1

julia> count_terms(ro"6h", 4, 8)
4
```
"""
function count_terms end

function count_terms(orb::Orbital, occ::Int, term::Term)
    ℓ = orb.ℓ
    g = degeneracy(orb)
    occ > g && throw(ArgumentError("Invalid occupancy $occ for $orb with degeneracy $g"))
    (occ > g/2 && occ != g) && (occ = g - occ)

    p = parity(orb)^occ
    if occ == 1
        term == Term(ℓ,1//2,p) ? 1 : 0
    elseif ℓ == 0 && occ == 2 || occ == degeneracy(orb) || occ == 0
        term == Term(0,0,p"even") ? 1 : 0
    else
        S′ = convert(Int, 2term.S)
        Xu.X(occ, orb.ℓ, S′, convert(Int, term.L))
    end
end

function write_L(io::IO, term::Term)
    if isinteger(term.L)
        write(io, uppercase(spectroscopic_label(convert(Int, term.L))))
    else
        write(io, "[$(numerator(term.L))/$(denominator(term.L))]")
    end
end

function Base.show(io::IO, term::Term)
    write(io, to_superscript(multiplicity(term)))
    write_L(io, term)
    write(io, to_superscript(term.parity))
end

export Term, @T_str, multiplicity, weight, terms, count_terms


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

# Demand (call) model

# get demand mode for given time and priority
function getDemandMode!(demand::Demand, priority::Priority, startTime::Float)
	@assert(startTime != nullTime)
	@assert(priority != nullPriority)

	updateDemandToTime!(demand, startTime) # update demand.recentSetsStartTimesIndex for given start time
	demandSetIndex = demand.setsTimeOrder[demand.recentSetsStartTimesIndex]
	demandModeIndex = demand.modeLookup[demandSetIndex, Int(priority)]

	return demand.modes[demandModeIndex]
end

# update demand.recentSetsStartTimesIndex up to given demand start time
function updateDemandToTime!(demand::Demand, startTime::Float)
	# shorthand:
	setsStartTimes = demand.setsStartTimes
	i = demand.recentSetsStartTimesIndex
	n = length(setsStartTimes)

	@assert(setsStartTimes[i] <= startTime) # otherwise, have gone back in time?

	if i == n || startTime < setsStartTimes[i+1]
		return # do nothing
	end

	# use linear search to update position in setsStartTimes
	while i < n && setsStartTimes[i+1] <= startTime
		i += 1
	end

	demand.recentSetsStartTimesIndex = i
end

# Creates one set of points to represent the demand in demand.rasters.
# Each point will store one or more demand values, one from each raster.
# Where there are multiple points per raster cell, the demand is divided between the points.
# Points are not created in places where no rasters have demand.
# Requires all rasters to have the same dimensions (only z values can differ).
function createDemandPointsFromRasters(demand::Demand; numCellRows::Int = 1, numCellCols::Int = 1)
	@assert(demand.numRasters >= 1)
	@assert(numCellRows >= 1)
	@assert(numCellCols >= 1)

	# shorthand
	@unpack rasters, numRasters = demand
	@unpack nx, ny = rasters[1]

	# check that all rasters have same dimensions,
	# otherwise creating points that work for all rasters may be tricky
	for i = 2:numRasters
		@assert(rasters[i].nx == nx)
		@assert(rasters[i].ny == ny)
		@assert(all(j -> isapprox(rasters[i].x[j], rasters[1].x[j]; rtol = eps()), 1:nx))
		@assert(all(j -> isapprox(rasters[i].y[j], rasters[1].y[j]; rtol = eps()), 1:ny))
	end

	# create point(s) for each raster cell,
	# except where point is in cell where all rasters have z = 0
	# numCellRows is for y dimension, numCellCols for x
	numCellPoints = numCellRows * numCellCols # number of points per cell
	points = Vector{Point}()
	maxNumPoints = (nx * numCellCols) * (ny * numCellRows)
	rastersPointDemands = [Vector{Float}(undef, maxNumPoints) for i = 1:numRasters] # rastersPointDemands[i][j] gives demand at point j for raster i
	for i = 1:nx, j = 1:ny # raster cell indices
		if all(raster -> raster.z[i,j] == 0, rasters)
			continue
		end
		for cellRow = 1:numCellRows, cellCol = 1:numCellCols
			point = Point()
			point.index = length(points) + 1

			# set point location
			point.location.x = rasters[1].x[i] + rasters[1].dx * ((cellCol - 0.5)/numCellCols - 0.5)
			point.location.y = rasters[1].y[j] + rasters[1].dy * ((cellRow - 0.5)/numCellRows - 0.5)

			# set demand of point for each raster,
			# divide raster demand by number of points per cell
			for rasterIndex = 1:numRasters
				rastersPointDemands[rasterIndex][point.index] = rasters[rasterIndex].z[i,j] / numCellPoints
			end

			push!(points, point)
		end
	end

	# remove unneeded values from rastersPointDemands (rastersPointDemands[i][j] where j > length(points))
	numPoints = length(points) # shorthand
	for i = 1:numRasters
		rastersPointDemands[i] = rastersPointDemands[i][1:numPoints]
	end

	return points, rastersPointDemands
end

# Sets the nearest node data for the points.
# mutates: points (fields: nearestNodeIndex, nearestNodeDist)
function setPointsNearestNodes!(sim::Simulation, points::Vector{Point})
	# find and set nearest node index of each point
	for point in points
		@assert(point.location.x != nullX && point.location.y != nullY)
		(point.nearestNodeIndex, point.nearestNodeDist) = findNearestNode(sim.map, sim.grid, sim.net.fGraph.nodes, point.location)
	end
end

# Creates an instance of PointsCoverageMode for the given travelMode and coverTime.
function createPointsCoverageMode(sim::Simulation, travelMode::TravelMode, coverTime::Float)
	@assert(coverTime >= 0)
	@assert(!isempty(sim.demandCoverage.nodesPoints))

	pointsCoverageMode = PointsCoverageMode()
	pointsCoverageMode.coverTime = coverTime
	pointsCoverageMode.travelMode = travelMode
	pointsCoverageMode.points = sim.demandCoverage.points

	# shorthand:
	points = pointsCoverageMode.points
	numPoints = length(points)
	@unpack stations, numStations = sim
	numNodes = length(sim.net.fGraph.nodes)
	nodesPoints = sim.demandCoverage.nodesPoints

	# # determine point coverage
	# stationsCoverPoints = fill(false, numStations, numPoints)
	# for i = 1:numStations
		# station = stations[i] # shorthand

		# # get travel time from station to nearest node
		# (node1, dist1) = (station.nearestNodeIndex, station.nearestNodeDist)
		# time1 = offRoadTravelTime(travelMode, dist1) # time to reach nearest node

		# for j = 1:numPoints
			# # get shortest path travel time
			# point = points[j] # shorthand
			# (node2, dist2) = (point.nearestNodeIndex, point.nearestNodeDist)
			# pathTravelTime = shortestPathTravelTime(net, travelMode.index, node1, node2)
			# time2 = offRoadTravelTime(travelMode, dist2) # time to reach nearest node
			# if time1 + pathTravelTime + time2 <= coverTime
				# stationsCoverPoints[i,j] = true
			# end
		# end
	# end

	# determine point coverage
	stationsCoverPoints = fill(false, numStations, numPoints)
	for i = 1:numStations
		station = stations[i] # shorthand

		# get travel time from station to nearest node
		(node1, dist1) = (station.nearestNodeIndex, station.nearestNodeDist)
		time1 = offRoadTravelTime(travelMode, dist1) # time to reach nearest node from station

		for j = 1:numNodes
			pathTravelTime = shortestPathTravelTime(sim.net, travelMode.index, node1, j)
			if time1 + pathTravelTime <= coverTime
				for k in nodesPoints[j] # indices of points for which node j is the nearest node
					point = points[k] # shorthand
					@assert(point.nearestNodeIndex == j)
					time2 = offRoadTravelTime(travelMode, point.nearestNodeDist) # time to reach point from nearest node
					stationsCoverPoints[i,k] = (time1 + pathTravelTime + time2 <= coverTime)
				end
			end
		end
	end

	# group points with coverage by same set of stations
	stationSetsPoints = Dict{Vector{Bool},Vector{Int}}() # stationSetsPoints[stationSet] gives the points that belong to the stationSet
	for j = 1:numPoints
		stationSet = stationsCoverPoints[:,j] # stationSet[i] = if station i covers point j
		# will put points that are not covered into a set, though alternatively they could be ignored
		stationSets = get!(stationSetsPoints, stationSet, Int[])
		push!(stationSets, j)
	end
	pointSets = collect(values(stationSetsPoints)) # pointSets[i] has set of all point indices covered by the same unique set of stations
	stationSets = [findall(stationSet) for stationSet in keys(stationSetsPoints)] # stationSets[i] = unique set of stations for pointSets[i]
	# according to documentation, 'values' and 'keys' return elements in the same order, so pointSets[i] corresponds with stationSets[i]

	stationsCoverPointSets = Vector{Vector{Int}}(undef, numStations) # stationsCoverPointSets[i] = list of point set indices covered by station i
	for i = 1:numStations
		stationsCoverPointSets[i] = findall(stationSet -> in(i,stationSet), stationSets)
	end

	pointsCoverageMode.pointSets = pointSets
	pointsCoverageMode.stationSets = stationSets
	pointsCoverageMode.stationsCoverPointSets = stationsCoverPointSets

	return pointsCoverageMode # will still need to set pointsCoverageMode.index
end

function getPointsCoverageMode(demandCoverage::DemandCoverage, travelMode::TravelMode, coverTime::Float)
	dc = demandCoverage # shorthand
	@assert(travelMode.index <= length(dc.pointsCoverageModeLookup))
	@assert(haskey(dc.pointsCoverageModeLookup[travelMode.index], coverTime))
	pointsCoverageMode = dc.pointsCoverageModes[dc.pointsCoverageModeLookup[travelMode.index][coverTime]]
	return pointsCoverageMode
end

# mutates: sim.travel state
function getPointsCoverageMode!(sim::Simulation, demandPriority::Priority, currentTime::Float)
	@assert(demandPriority != nullPriority)
	@assert(currentTime != nullTime)
	travelPriority = sim.responseTravelPriorities[demandPriority]
	travelMode = getTravelMode!(sim.travel, travelPriority, currentTime)
	coverTime = sim.demandCoverage.coverTimes[demandPriority]
	pointsCoverageMode = getPointsCoverageMode(sim.demandCoverage, travelMode, coverTime)
	@assert(pointsCoverageMode != nothing)
	return pointsCoverageMode
end

# mutates: sim.travel and sim.demand states
function getPointSetsDemands!(sim::Simulation, demandPriority::Priority, currentTime::Float;
	pointsCoverageMode::Union{PointsCoverageMode,Nothing} = nothing)
	if pointsCoverageMode == nothing
		pointsCoverageMode = getPointsCoverageMode!(sim, demandPriority, currentTime)
	end
	demandMode = getDemandMode!(sim.demand, demandPriority, currentTime)
	pointSetsDemands = sim.demandCoverage.pointSetsDemands[pointsCoverageMode.index, demandMode.rasterIndex]
	return pointSetsDemands
end

"""
	function initDemand!(sim::Simulation, demand::Union{Demand,Nothing} = nothing;
		demandFilename::String = "")
Initialise demand data for `sim`.
Will initialise from the first item in this list: `demand`, `demandFilename` file, `sim.demand` (if set), `sim.inputFiles[\"demand\"].path` file.
Deletes old demand coverage data in `sim.demandCoverage`.

Mutates: `sim.demand`, `sim.demandCoverage`
"""
function initDemand!(sim::Simulation, demand::Union{Demand,Nothing} = nothing;
	demandFilename::String = "")
	if demand == nothing
		if demandFilename != ""
			demand = readDemandFile(demandFilename)
		elseif sim.demand.numSets == nullIndex
			if haskey(sim.inputFiles, "demand")
				demand = readDemandFile(sim.inputFiles["demand"].path)
			else
				error("no demand data given")
			end
		else
			demand = sim.demand
		end
	end
	demand.initialised = true
	sim.demand = demand

	# old demand coverage data may no longer be valid
	sim.demandCoverage = DemandCoverage(sim.demandCoverage) # keep some old params
end

# Initialises the points coverage modes data in sim.demandCoverage.
# mutates: sim.demandCoverage (fields: pointsCoverageModes, pointsCoverageModeLookup)
function initPointsCoverageModes!(sim::Simulation)
	# shorthand:
	dc = sim.demandCoverage
	travel = sim.travel

	# go through all travel sets and demand priorities, create points coverage mode for each call response travel priority and cover time
	# dc.pointsCoverageModes
	dc.pointsCoverageModeLookup = [Dict{Float,Int}() for i = 1:travel.numModes]
	for i = 1:travel.numSets, demandPriority in priorities
		# get travel mode and cover time
		travelPriority = sim.responseTravelPriorities[demandPriority]
		travelMode = travel.modes[travel.modeLookup[i,Int(travelPriority)]]
		coverTime = sim.demandCoverage.coverTimes[demandPriority]

		# if combination of travel mode and cover time is new, need to calculate points coverage
		if !haskey(dc.pointsCoverageModeLookup[travelMode.index], coverTime)
			pointsCoverageMode = createPointsCoverageMode(sim, travelMode, coverTime)
			pointsCoverageMode.index = length(dc.pointsCoverageModes) + 1
			push!(dc.pointsCoverageModes, pointsCoverageMode)
			dc.pointsCoverageModeLookup[travelMode.index][coverTime] = pointsCoverageMode.index
		end
	end
end

"""
	function initDemandCoverage!(sim::Simulation;
		coverTimes::Union{Dict{Priority,Float},Nothing} = nothing,
		rasterCellNumRows::Int = 1, rasterCellNumCols::Int = 1)
Initialises demand coverage data for `sim`.
If `coverTimes` is given, will use this (along with `rasterCellNumRows` and `rasterCellNumCols`), otherwise will use the previous parameter values set in `sim.demandCoverage`.

# Keyword arguments
- `coverTimes` is the target coverage time for each demand priority, e.g. `coverTimes = Dict([p => 8/60/24 for p in priorities])` sets a target time of 8 minutes (converted to days) for each priority. Can omit this if `sim.demandCoverage.coverTimes` is already set.
- `rasterCellNumRows` and `rasterCellNumCols` are used to set the number of demand points to represent demand per raster cell for `sim.demand.rasters`

Mutates: `sim.demandCoverage`
"""
function initDemandCoverage!(sim::Simulation;
	coverTimes::Union{Dict{Priority,Float},Nothing} = nothing,
	rasterCellNumRows::Int = 1, rasterCellNumCols::Int = 1)

	@assert(!sim.used)
	@assert(sim.travel.recentSetsStartTimesIndex == 1)
	sim.demand.initialised || initDemand!(sim)
	@assert(sim.demand.recentSetsStartTimesIndex == 1)

	# shorthand
	@unpack demand, travel = sim
	@unpack rasters, numRasters = demand
	numNodes = length(sim.net.fGraph.nodes)

	# set demand coverage params
	oldDemandCoverage = sim.demandCoverage
	if coverTimes != nothing
		sim.demandCoverage = DemandCoverage(coverTimes, rasterCellNumRows, rasterCellNumCols)
	else # use old params
		@assert(!isempty(oldDemandCoverage.coverTimes))
		sim.demandCoverage = DemandCoverage(oldDemandCoverage)
	end
	dc = sim.demandCoverage # shorthand
	demandPriorities = priorities
	@assert(all(p -> haskey(dc.coverTimes, p), demandPriorities), "coverTimes not set")
	@assert(dc.rasterCellNumRows >= 1)
	@assert(dc.rasterCellNumCols >= 1)

	# create demand points
	dc.points, dc.rastersPointDemands = createDemandPointsFromRasters(sim.demand; numCellRows = dc.rasterCellNumRows, numCellCols = dc.rasterCellNumCols)
	points = dc.points # shorthand
	setPointsNearestNodes!(sim, points)
	numPoints = length(points) # shorthand

	# set nodesPoints
	dc.nodesPoints = [Int[] for i = 1:numNodes]
	for point in points
		push!(dc.nodesPoints[point.nearestNodeIndex], point.index)
	end

	# set points coverage modes
	initPointsCoverageModes!(sim)
	numPointsCoverageModes = length(dc.pointsCoverageModes)

	# set pointSetsDemands
	# note that pointSetsDemands does not need to be set for all combinations of pointsCoverageMode.index and rasterIndex, as some combinations may not occur in simulation
	dc.pointSetsDemands = [Vector{Float}() for i = 1:numPointsCoverageModes, j = 1:demand.numRasters]
	times = vcat(travel.setsStartTimes, demand.setsStartTimes) |> unique |> sort
	for time in times, demandPriority in demandPriorities
		pointsCoverageMode = getPointsCoverageMode!(sim, demandPriority, time)
		demandMode = getDemandMode!(demand, demandPriority, time)
		rasterIndex = demandMode.rasterIndex # shorthand

		# if combination of pointsCoverageMode and rasterIndex is new, need to calculate point sets demands
		pointDemands = dc.rastersPointDemands[rasterIndex]
		if isempty(dc.pointSetsDemands[pointsCoverageMode.index, rasterIndex])
			pointSetsDemand = zeros(Float, length(pointsCoverageMode.pointSets))
			for (pointSetIndex, pointSet) in enumerate(pointsCoverageMode.pointSets)
				for pointIndex in pointSet
					pointSetsDemand[pointSetIndex] += pointDemands[pointIndex]
				end
			end
			dc.pointSetsDemands[pointsCoverageMode.index, rasterIndex] = pointSetsDemand
		end
	end
	# reset travel and demand state (previously changed by getTravelMode! and getDemandMode!)
	travel.recentSetsStartTimesIndex = 1
	demand.recentSetsStartTimesIndex = 1

	dc.initialised = true
end

# Calculate, for a given points coverage mode, and number of ambulances assigned to each
# station, how many ambulances cover each point set.
function calcPointSetsCoverCounts(pointsCoverageMode::PointsCoverageMode, stationsNumAmbs::Vector{Int})
	# stationsNumAmbs[i] gives the number of ambulances that can provide coverage and are assigned to station i

	@assert(all(n -> n >= 0, stationsNumAmbs))

	# count how many times each point set is covered
	numPointSets = length(pointsCoverageMode.pointSets) # shorthand
	pointSetsCoverCount = zeros(Int, numPointSets)
	for i = 1:numPointSets
		for j in pointsCoverageMode.stationSets[i]
			pointSetsCoverCount[i] += stationsNumAmbs[j]
		end
	end

	# # alternative code:
	# @assert(length(stationsNumAmbs) == length(pointsCoverageMode.stationsCoverPointSets))
	# pointSetsCoverCount = zeros(Int, length(pointsCoverageMode.pointSets))
	# for (i,stationNumAmbs) in enumerate(stationsNumAmbs)
		# for j in pointsCoverageMode.stationsCoverPointSets[i]
			# pointSetsCoverCount[j] += stationNumAmbs
		# end
	# end

	return pointSetsCoverCount
end

# Calculate, for each demand priority, how many ambulances cover each point set in
# pointsCoverageMode.pointSet, where pointsCoverageMode is for the demand priority and current time.
# mutates: sim.travel state
function calcPointSetsCoverCounts!(sim::Simulation, currentTime::Float, stationsNumAmbs::Vector{Int};
	demandPriorities::Vector{Priority} = priorities)
	# stationsNumAmbs[i] gives the number of ambulances that can provide coverage and are assigned to station i

	@assert(length(stationsNumAmbs) == sim.numStations)
	@assert(all(n -> n >= 0, stationsNumAmbs))

	demandsPointSetsCoverCounts = Dict{Priority,Vector{Int}}() # demandsPointSetsCoverCounts[p][i] gives the number of ambulances assigned to stations that cover point set i for demand priority p
	pointsCoverageModesUsed = Dict{PointsCoverageMode,Priority}() # keep track of which points coverage mode are used, and any demand priority that uses each
	for demandPriority in demandPriorities
		pointsCoverageMode = getPointsCoverageMode!(sim, demandPriority, currentTime)
		if haskey(pointsCoverageModesUsed, pointsCoverageMode)
			# have already calculated the point sets cover counts for this points coverage mode
			priority = pointsCoverageModesUsed[pointsCoverageMode]
			demandsPointSetsCoverCounts[demandPriority] = demandsPointSetsCoverCounts[priority] # should use copy(), so that vectors are independent?
		else
			demandsPointSetsCoverCounts[demandPriority] = calcPointSetsCoverCounts(pointsCoverageMode, stationsNumAmbs)
			pointsCoverageModesUsed[pointsCoverageMode] = demandPriority
		end
	end

	return demandsPointSetsCoverCounts
end

# Calculates, for each demand priority, how much of the demand is covered by 1,2,... ambulances
# mutates: sim.travel and sim.demand states
function calcDemandCoverCounts!(sim::Simulation, currentTime::Float, stationsNumAmbs::Vector{Int};
	demandPriorities::Vector{Priority} = priorities)
	# stationsNumAmbs[i] gives the number of ambulances that can provide coverage and are assigned to station i

	@assert(length(stationsNumAmbs) == sim.numStations)
	@assert(all(n -> n >= 0, stationsNumAmbs))

	# shorthand
	numAmbs = sum(stationsNumAmbs)

	demandsPointSetsCoverCounts = calcPointSetsCoverCounts!(sim, currentTime, stationsNumAmbs; demandPriorities = demandPriorities)

	demandCoverCounts = Dict{Priority,Vector{Float}}() # demandCoverCounts[p][i] gives the amount of demand covered by i ambulances for demand priority p
	coverCountsCalculated = [Dict{Vector{Int},Vector{Int}}() for i = 1:sim.demand.numRasters] # coverCountsCalculated[rasterIndex][pointSetsCoverCount] stores the coverCounts (before multiplying by demandMode.rasterMultiplier) for the combination of args
	for demandPriority in demandPriorities
		pointsCoverageMode = getPointsCoverageMode!(sim, demandPriority, currentTime)
		demandMode = getDemandMode!(sim.demand, demandPriority, currentTime)
		rasterIndex = demandMode.rasterIndex # shorthand

		pointSetsCoverCount = demandsPointSetsCoverCounts[demandPriority]
		pointSetsDemands = sim.demandCoverage.pointSetsDemands[pointsCoverageMode.index, rasterIndex]

		coverCounts = nothing # init
		if haskey(coverCountsCalculated[rasterIndex], pointSetsCoverCount)
			coverCounts = coverCountsCalculated[rasterIndex][pointSetsCoverCount]
		else
			coverCounts = zeros(Float, numAmbs) # coverCounts[i] is the demand covered by i ambulances
			for (i,pointSet) in enumerate(pointsCoverageMode.pointSets)
				if pointSetsCoverCount[i] > 0
					coverCounts[pointSetsCoverCount[i]] += pointSetsDemands[i]
				end
			end
			coverCountsCalculated[rasterIndex][pointSetsCoverCount] = coverCounts
		end
		demandCoverCounts[demandPriority] = coverCounts * demandMode.rasterMultiplier
	end

	return demandCoverCounts
end


text:
module PhyloModels

  using PhyloTrees,
        SubstitutionModels,
        GeneticBitArrays

  import SubstitutionModels._π
  import Base.rand

  # Re-export all of PhyloTrees, SubstitutionModels, and GeneticBitArrays
  for name in names(PhyloTrees)
    @eval export $(name)
  end

  for name in names(SubstitutionModels)
    @eval export $(name)
  end

  for name in names(GeneticBitArrays)
    @eval export $(name)
  end

  const NodeDNA = Dict{Int64, DNASeq}
  const NodeRNA = Dict{Int64, RNASeq}

  include("loglikelihood.jl")
  include("simulate.jl")

  export NodeDNA, NodeRNA, siimulate!, simulate, rand, loglikelihood

end # module


text:
@static if is_windows()
  const windows_time_freq = begin
    x = Array{Int64}()
    if ccall(:QueryPerformanceFrequency,Bool,(Ptr{Int64},),pointer(x))
      x[]
    else
      error("Unsupported operating system. Use Windows XP or later.")
    end
  end
end

const approx_timer_resolution = @static if is_windows()
  1/windows_time_freq
else
  1e-6
end

@static if is_windows()
  function precise_time()
    x = Array{Int64}()
    if !ccall(:QueryPerformanceCounter,Bool,(Ptr{Int64},),pointer(x))
      error("Unsupported operating system. Use Windows XP or later.")
    else
      x[] / windows_time_freq
    end
  end
else
  const precise_time = time
end


text:
#  Copyright 2016, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
#  This Source Code Form is subject to the terms of the Mozilla Public
#  License, v. 2.0. If a copy of the MPL was not distributed with this
#  file, You can obtain one at http://mozilla.org/MPL/2.0/.
#############################################################################
# JuMP
# An algebraic modeling language for Julia
# See http://github.com/JuliaOpt/JuMP.jl
#############################################################################
# src/sos.jl
# Defines SOSConstraint for "special ordered sets", and related methods
# SOS1: at most one of the variables in the set is positive, the rest are 0
# SOS2: at most two can be nonzero, if they a consecutive in the set
#############################################################################

type SOSConstraint <: AbstractConstraint
    terms::Vector{Variable}
    weights::Vector{Float64}
    sostype::Symbol
end
Base.copy(sos::SOSConstraint, new_model::Model) =
    SOSConstraint(Variable[copy(v, new_model) for v in sos.terms],
                    copy(sos.weights), sos.sostype)

# Given a vector of affine expressions, extract a vector of the single
# variable in each expression and a vector of their coefficients
function constructSOS(m::Model, coll::Vector{AffExpr})
    nvar = length(coll)
    vars = Array(Variable, nvar)
    weight = Array(Float64, nvar)
    for (i,aff) in enumerate(coll)
        if (length(aff.vars) != 1) || (aff.constant != 0)
            error("Must specify set in terms of single variables")
        end
        vars[i] = aff.vars[1]
        vars[i].m === m || error("Variable in constraint is not owned by the model")
        weight[i] = aff.coeffs[1]
    end
    return vars, weight
end

addSOS1(m::Model, coll) = addSOS1(m, convert(Vector{AffExpr}, coll))

function addSOS1(m::Model, coll::Vector{AffExpr})
    vars, weight = constructSOS(m,coll)
    push!(m.sosconstr, SOSConstraint(vars, weight, :SOS1))
    if m.internalModelLoaded
        idx = Int[v.col for v in vars]
        if applicable(MathProgBase.addsos1!, m.internalModel, idx, weight)
            MathProgBase.addsos1!(m.internalModel, idx, weight)
        else
            Base.warn_once("Solver does not appear to support adding constraints to an existing model. Hot-start is disabled.")
            m.internalModelLoaded = false
        end
    end
    return ConstraintRef{Model,SOSConstraint}(m,length(m.sosconstr))
end

addSOS2(m::Model, coll) = addSOS2(m, convert(Vector{AffExpr}, coll))

function addSOS2(m::Model, coll::Vector{AffExpr})
    vars, weight = constructSOS(m,coll)
    push!(m.sosconstr, SOSConstraint(vars, weight, :SOS2))
    if m.internalModelLoaded
        idx = Int[v.col for v in vars]
        if applicable(MathProgBase.addsos2!, m.internalModel, idx, weight)
            MathProgBase.addsos2!(m.internalModel, idx, weight)
        else
            Base.warn_once("Solver does not appear to support adding constraints to an existing model. Hot-start is disabled.")
            m.internalModelLoaded = false
        end
    end
    return ConstraintRef{Model,SOSConstraint}(m,length(m.sosconstr))
end


text:
@blueprint struct License <: MustacheBlueprint
    template_dir::String = default_template("licenses")
    name::String = "MIT"
end

function src_file(bp::License, ctx::Context)
    return joinpath(bp.template_dir, bp.name)
end

function dst_file(::License, ctx::Context)
    return joinpath(ctx.kwargs.path, "LICENSE")
end

function blueprint_view(::License, ctx::Context)
    return Dict(
        "AUTHORS" => join(ctx.kwargs.authors, ", "),
        "YEAR" => year(today()),
    )
end


text:
using CrystallographyBase
using Test
using Unitful: @u_str
using UnitfulAtomic

@testset "CrystallographyBase.jl" begin
    include("lattice.jl")
    include("miller.jl")
    include("metric.jl")
    include("transform.jl")
    include("reciprocal.jl")
    # include("wignerseitz.jl")
end