text:
"""
# Onda Corporation
# Hydrophone Preamplifier Calibration Data
#
FILE_NAME AH-2020-20-100-1238_200929.txt
#
AMP_CAL_DATE 29-Sep-2020
AMP_MFG ONDA
AMP_MODEL AH-2020-100
AMP_SN 1238
AMP_GAIN_NOM_DB 20
AMP_BW_NOM_MHZ 100
AMP_OUTPUT_IMP_NOM_OHMS 50
AMP_INPUT_CAP_NOM_PF 7.1
AMP_INPUT_CAP_DATA_TYPE CALIBRATED
AMP_GAIN_UNCERT_DB 0.03
AMP_CAP_UNCERT_PCT 2.5
#
# Calibration Data
# Freq(MHz) Gain(dB) Phase(deg) Cap(pF)
#
DATA_FIELDS 4
DATA_FIELD FREQ_MHZ
DATA_FIELD GAIN_DB
DATA_FIELD PHASE_DEG
DATA_FIELD CAP_PF
#
# End of file header
HEADER_END 0
"""

"""
`Onda_AH2020_1238_0dB`
"""

@calibration :Onda_AH2020_1238_0dB [
    000.0500 -0.92 31.57 6.70
    000.0510 -0.90 30.95 6.70
    000.0520 -0.89 30.31 6.70
    000.0530 -0.88 29.76 6.70
    000.0540 -0.88 29.29 6.70
    000.0550 -0.87 28.75 6.70
    000.0560 -0.85 28.11 6.70
    000.0570 -0.83 27.64 6.70
    000.0580 -0.83 27.20 6.70
    000.0590 -0.82 26.75 6.70
    000.0600 -0.81 26.36 6.70
    000.0610 -0.81 25.92 6.70
    000.0620 -0.82 25.49 6.70
    000.0630 -0.81 25.09 6.70
    000.0640 -0.80 24.75 6.70
    000.0650 -0.79 24.40 6.70
    000.0660 -0.79 24.03 6.70
    000.0670 -0.80 23.69 6.70
    000.0680 -0.81 23.37 6.70
    000.0690 -0.79 23.06 6.70
    000.0700 -0.79 22.82 6.70
    000.0710 -0.78 22.58 6.70
    000.0720 -0.78 22.30 6.70
    000.0730 -0.77 22.04 6.70
    000.0740 -0.76 21.79 6.70
    000.0750 -0.77 21.57 6.70
    000.0760 -0.77 21.35 6.70
    000.0770 -0.76 21.08 6.70
    000.0780 -0.76 20.86 6.70
    000.0790 -0.75 20.71 6.70
    000.0800 -0.75 20.53 6.70
    000.0810 -0.74 20.33 6.70
    000.0820 -0.74 20.15 6.70
    000.0830 -0.73 19.99 6.70
    000.0840 -0.72 19.76 6.70
    000.0850 -0.72 19.51 6.70
    000.0860 -0.71 19.32 6.70
    000.0870 -0.70 19.14 6.70
    000.0880 -0.69 18.97 6.70
    000.0890 -0.69 18.81 6.70
    000.0900 -0.67 18.67 6.70
    000.0910 -0.66 18.53 6.70
    000.0920 -0.66 18.34 6.70
    000.0930 -0.65 18.13 6.70
    000.0940 -0.65 17.98 6.70
    000.0950 -0.64 17.85 6.70
    000.0960 -0.63 17.71 6.70
    000.0970 -0.62 17.54 6.70
    000.0980 -0.61 17.39 6.70
    000.0990 -0.61 17.25 6.70
    000.1000 -0.61 17.11 6.70
    000.1100 -0.54 15.75 6.70
    000.1200 -0.48 14.48 6.70
    000.1300 -0.46 13.42 6.70
    000.1400 -0.43 12.48 6.70
    000.1500 -0.41 11.64 6.70
    000.1600 -0.39 10.91 6.70
    000.1700 -0.37 10.34 6.70
    000.1800 -0.37 9.74 6.70
    000.1900 -0.36 9.23 6.70
    000.2000 -0.34 8.76 6.70
    000.2100 -0.34 8.37 6.70
    000.2200 -0.33 7.96 6.70
    000.2300 -0.32 7.67 6.70
    000.2400 -0.31 7.33 6.70
    000.2500 -0.31 7.01 6.70
    000.2600 -0.31 6.72 6.70
    000.2700 -0.29 6.44 6.70
    000.2800 -0.30 6.30 6.70
    000.2900 -0.29 6.09 6.70
    000.3000 -0.29 5.84 6.70
    000.3100 -0.28 5.67 6.70
    000.3200 -0.28 5.46 6.70
    000.3300 -0.28 5.31 6.70
    000.3400 -0.27 5.15 6.70
    000.3500 -0.27 4.98 6.70
    000.3600 -0.27 4.85 6.70
    000.3700 -0.26 4.70 6.70
    000.3800 -0.26 4.61 6.70
    000.3900 -0.26 4.42 6.70
    000.4000 -0.26 4.36 6.70
    000.4100 -0.25 4.19 6.70
    000.4200 -0.25 4.11 6.70
    000.4300 -0.26 4.03 6.70
    000.4400 -0.25 3.89 6.70
    000.4500 -0.24 3.77 6.70
    000.4600 -0.24 3.68 6.70
    000.4700 -0.24 3.59 6.70
    000.4800 -0.24 3.52 6.70
    000.4900 -0.23 3.45 6.70
    000.5000 -0.23 3.42 6.70
    000.5100 -0.23 3.29 6.70
    000.5200 -0.23 3.20 6.70
    000.5300 -0.23 3.12 6.70
    000.5400 -0.23 3.05 6.70
    000.5500 -0.23 2.97 6.70
    000.5600 -0.22 2.91 6.70
    000.5700 -0.23 2.84 6.70
    000.5800 -0.22 2.78 6.70
    000.5900 -0.21 2.71 6.70
    000.6000 -0.22 2.60 6.70
    000.6100 -0.22 2.55 6.70
    000.6200 -0.21 2.51 6.70
    000.6300 -0.22 2.43 6.70
    000.6400 -0.21 2.42 6.70
    000.6500 -0.21 2.37 6.70
    000.6600 -0.21 2.29 6.70
    000.6700 -0.20 2.27 6.70
    000.6800 -0.20 2.23 6.70
    000.6900 -0.20 2.16 6.70
    000.7000 -0.20 2.07 6.70
    000.7100 -0.19 1.99 6.70
    000.7200 -0.19 1.97 6.70
    000.7300 -0.19 1.94 6.70
    000.7400 -0.20 1.91 6.70
    000.7500 -0.19 1.86 6.70
    000.7600 -0.19 1.81 6.70
    000.7700 -0.18 1.75 6.70
    000.7800 -0.18 1.71 6.70
    000.7900 -0.19 1.70 6.70
    000.8000 -0.19 1.65 6.70
    000.8100 -0.19 1.57 6.70
    000.8200 -0.18 1.51 6.70
    000.8300 -0.17 1.47 6.70
    000.8400 -0.17 1.43 6.70
    000.8500 -0.17 1.39 6.70
    000.8600 -0.18 1.36 6.70
    000.8700 -0.18 1.34 6.70
    000.8800 -0.18 1.30 6.70
    000.8900 -0.18 1.26 6.70
    000.9000 -0.17 1.20 6.70
    000.9100 -0.17 1.14 6.70
    000.9200 -0.17 1.10 6.70
    000.9300 -0.17 1.07 6.70
    000.9400 -0.17 1.03 6.70
    000.9500 -0.17 0.99 6.70
    000.9600 -0.16 0.96 6.70
    000.9700 -0.16 0.93 6.70
    000.9800 -0.16 0.89 6.70
    000.9900 -0.16 0.86 6.70
    001.0000 -0.16 0.83 6.70
    001.1000 -0.15 0.54 6.70
    001.2000 -0.14 0.23 6.70
    001.3000 -0.13 -0.09 6.70
    001.4000 -0.13 -0.36 6.70
    001.5000 -0.13 -0.56 6.70
    001.6000 -0.11 -0.90 6.70
    001.7000 -0.10 -1.11 6.70
    001.8000 -0.10 -1.41 6.70
    001.9000 -0.10 -1.64 6.70
    002.0000 -0.09 -1.86 6.70
    002.1000 -0.10 -2.08 6.70
    002.2000 -0.10 -2.35 6.70
    002.3000 -0.09 -2.54 6.70
    002.4000 -0.10 -2.73 6.70
    002.5000 -0.09 -2.92 6.70
    002.6000 -0.08 -3.11 6.70
    002.7000 -0.09 -3.32 6.70
    002.8000 -0.09 -3.51 6.70
    002.9000 -0.10 -3.71 6.70
    003.0000 -0.09 -3.86 6.70
    003.1000 -0.09 -4.05 6.70
    003.2000 -0.09 -4.29 6.70
    003.3000 -0.09 -4.45 6.70
    003.4000 -0.09 -4.63 6.70
    003.5000 -0.09 -4.84 6.70
    003.6000 -0.09 -5.00 6.70
    003.7000 -0.09 -5.13 6.70
    003.8000 -0.10 -5.36 6.70
    003.9000 -0.09 -5.52 6.70
    004.0000 -0.10 -5.73 6.70
    004.1000 -0.10 -5.89 6.70
    004.2000 -0.09 -6.04 6.70
    004.3000 -0.09 -6.24 6.70
    004.4000 -0.09 -6.46 6.70
    004.5000 -0.09 -6.60 6.70
    004.6000 -0.09 -6.74 6.70
    004.7000 -0.09 -6.94 6.70
    004.8000 -0.09 -7.09 6.70
    004.9000 -0.09 -7.30 6.70
    005.0000 -0.08 -7.45 6.70
    005.1000 -0.08 -7.63 6.70
    005.2000 -0.09 -7.81 6.70
    005.3000 -0.09 -7.98 6.70
    005.4000 -0.08 -8.15 6.70
    005.5000 -0.09 -8.32 6.70
    005.6000 -0.09 -8.52 6.70
    005.7000 -0.08 -8.65 6.70
    005.8000 -0.08 -8.83 6.70
    005.9000 -0.09 -9.07 6.70
    006.0000 -0.09 -9.19 6.70
    006.1000 -0.08 -9.32 6.70
    006.2000 -0.08 -9.49 6.70
    006.3000 -0.08 -9.74 6.70
    006.4000 -0.08 -9.89 6.70
    006.5000 -0.08 -10.05 6.70
    006.6000 -0.08 -10.24 6.70
    006.7000 -0.08 -10.41 6.70
    006.8000 -0.08 -10.58 6.70
    006.9000 -0.08 -10.74 6.70
    007.0000 -0.08 -10.89 6.70
    007.1000 -0.08 -11.05 6.70
    007.2000 -0.08 -11.25 6.70
    007.3000 -0.09 -11.44 6.70
    007.4000 -0.09 -11.63 6.70
    007.5000 -0.09 -11.74 6.70
    007.6000 -0.09 -11.86 6.70
    007.7000 -0.09 -12.07 6.70
    007.8000 -0.09 -12.26 6.70
    007.9000 -0.09 -12.44 6.70
    008.0000 -0.09 -12.62 6.70
    008.1000 -0.09 -12.79 6.70
    008.2000 -0.09 -12.94 6.70
    008.3000 -0.09 -13.08 6.70
    008.4000 -0.09 -13.27 6.70
    008.5000 -0.09 -13.47 6.70
    008.6000 -0.09 -13.63 6.70
    008.7000 -0.09 -13.78 6.70
    008.8000 -0.09 -13.92 6.70
    008.9000 -0.10 -14.06 6.70
    009.0000 -0.10 -14.25 6.70
    009.1000 -0.11 -14.46 6.70
    009.2000 -0.11 -14.64 6.70
    009.3000 -0.10 -14.82 6.70
    009.4000 -0.10 -14.98 6.70
    009.5000 -0.10 -15.13 6.70
    009.6000 -0.10 -15.29 6.70
    009.7000 -0.10 -15.45 6.70
    009.8000 -0.11 -15.61 6.70
    009.9000 -0.10 -15.78 6.70
    010.0000 -0.10 -15.94 6.70
    011.0000 -0.12 -17.57 6.70
    012.0000 -0.14 -19.23 6.70
    013.0000 -0.16 -20.79 6.70
    014.0000 -0.17 -22.37 6.70
    015.0000 -0.18 -23.95 6.70
    016.0000 -0.19 -25.45 6.70
    017.0000 -0.19 -26.97 6.70
    018.0000 -0.20 -28.47 6.70
    019.0000 -0.21 -29.95 6.70
    020.0000 -0.20 -31.46 6.70
    021.0000 -0.20 -33.06 6.70
    022.0000 -0.19 -34.64 6.70
    023.0000 -0.18 -36.25 6.70
    024.0000 -0.18 -37.88 6.70
    025.0000 -0.17 -39.48 6.70
    026.0000 -0.17 -41.16 6.70
    027.0000 -0.17 -42.76 6.70
    028.0000 -0.18 -44.48 6.70
    029.0000 -0.18 -46.15 6.70
    030.0000 -0.20 -47.83 6.70
    031.0000 -0.20 -49.51 6.70
    032.0000 -0.21 -51.17 6.70
    033.0000 -0.22 -52.81 6.70
    034.0000 -0.24 -54.45 6.70
    035.0000 -0.26 -56.09 6.71
    036.0000 -0.28 -57.74 6.71
    037.0000 -0.30 -59.44 6.71
    038.0000 -0.30 -61.08 6.71
    039.0000 -0.32 -62.72 6.71
    040.0000 -0.33 -64.37 6.71
    041.0000 -0.36 -65.96 6.71
    042.0000 -0.39 -67.57 6.71
    043.0000 -0.40 -69.20 6.71
    044.0000 -0.43 -70.84 6.71
    045.0000 -0.45 -72.46 6.72
    046.0000 -0.47 -74.10 6.72
    047.0000 -0.49 -75.66 6.72
    048.0000 -0.51 -77.34 6.72
    049.0000 -0.53 -78.86 6.72
    050.0000 -0.55 -80.51 6.72
    051.0000 -0.58 -82.15 6.72
    052.0000 -0.60 -83.75 6.73
    053.0000 -0.63 -85.35 6.73
    054.0000 -0.66 -86.95 6.73
    055.0000 -0.67 -88.55 6.73
    056.0000 -0.70 -90.11 6.73
    057.0000 -0.72 -91.73 6.73
    058.0000 -0.74 -93.31 6.74
    059.0000 -0.77 -94.83 6.74
    060.0000 -0.80 -96.41 6.74
    061.0000 -0.83 -98.10 6.74
    062.0000 -0.85 -99.76 6.74
    063.0000 -0.87 -101.33 6.75
    064.0000 -0.90 -102.97 6.75
    065.0000 -0.93 -104.55 6.75
    066.0000 -0.94 -106.03 6.75
    067.0000 -0.97 -107.70 6.75
    068.0000 -0.99 -109.34 6.76
    069.0000 -1.02 -110.94 6.76
    070.0000 -1.05 -112.56 6.76
    071.0000 -1.07 -114.18 6.76
    072.0000 -1.10 -115.80 6.77
    073.0000 -1.13 -117.38 6.77
    074.0000 -1.16 -118.92 6.77
    075.0000 -1.19 -120.54 6.77
    076.0000 -1.22 -122.17 6.77
    077.0000 -1.25 -123.83 6.78
    078.0000 -1.29 -125.47 6.78
    079.0000 -1.32 -127.10 6.78
    080.0000 -1.36 -128.70 6.79
    081.0000 -1.39 -130.29 6.79
    082.0000 -1.42 -131.94 6.79
    083.0000 -1.44 -133.61 6.79
    084.0000 -1.48 -135.22 6.80
    085.0000 -1.52 -136.81 6.80
    086.0000 -1.56 -138.44 6.80
    087.0000 -1.59 -140.07 6.81
    088.0000 -1.63 -141.70 6.81
    089.0000 -1.67 -143.33 6.81
    090.0000 -1.71 -144.97 6.81
    091.0000 -1.75 -146.60 6.82
    092.0000 -1.79 -148.22 6.82
    093.0000 -1.84 -149.83 6.82
    094.0000 -1.87 -151.47 6.83
    095.0000 -1.91 -153.11 6.83
    096.0000 -1.95 -154.76 6.83
    097.0000 -2.00 -156.42 6.84
    098.0000 -2.04 -158.05 6.84
    099.0000 -2.08 -159.62 6.84
    100.0000 -2.12 -161.19 6.85
]


text:
"""
    propagate(l, g, aggr, [x, e]) -> x′, e′
    propagate(l, g, aggr) -> g′

Performs the message-passing for GNN layer `l` on graph `g` .
Returns updated node and edge features `x` and `e`.

In case no input and edge features are given as input,
extracts them from `g` and returns the same graph
with updated feautres.

The computational steps are the following:

```julia
m = compute_batch_message(l, g, x, e)  # calls `compute_message`
m̄ = aggregate_neighbors(g, aggr, m)
x′ = update_node(l, m̄, x)
e′ = update_edge(l, m, e)
```

Custom layers typically define their own [`update_node`](@ref)
and [`compute_message`](@ref) functions, then call
this method in the forward pass:

# Usage example

```julia
using GraphNeuralNetworks, Flux

struct GNNConv <: GNNLayer
    W
    b
    σ
end

Flux.@functor GNNConv

function GNNConv(ch::Pair{Int,Int}, σ=identity;
                 init=glorot_uniform, bias::Bool=true)
    in, out = ch
    W = init(out, in)
    b = Flux.create_bias(W, bias, out)
    GNNConv(W, b, σ, aggr)
end

compute_message(l::GNNConv, x_i, x_j, e_ij) = l.W * x_j
update_node(l::GNNConv, m̄, x) = l.σ.(m̄ .+ l.bias)

function (l::GNNConv)(g::GNNGraph, x::AbstractMatrix)
    x, _ = propagate(l, g, +, x)
    return x
end
```

See also [`compute_message`](@ref) and [`update_node`](@ref).
"""
function propagate end

function propagate(l, g::GNNGraph, aggr)
    x, e = propagate(l, g, aggr, node_features(g), edge_features(g))
    return GNNGraph(g, ndata=x, edata=e)
end

function propagate(l, g::GNNGraph, aggr, x, e=nothing)
    m = compute_batch_message(l, g, x, e)
    m̄ = aggregate_neighbors(g, aggr, m)
    x′ = update_node(l, m̄, x)
    e′ = update_edge(l, m, e)
    return x′, e′
end

## Step 1.

"""
    compute_message(l, x_i, x_j, [e_ij])

Message function for the message-passing scheme
started by [`propagate`](@ref).
Returns the message from node `j` to node `i` .
In the message-passing scheme, the incoming messages
from the neighborhood of `i` will later be aggregated
in order to update (see [`update_node`](@ref)) the features of node `i`.

The function operates on batches of edges, therefore
`x_i`, `x_j`, and `e_ij` are tensors whose last dimension
is the batch size, or can be tuple/namedtuples of
such tensors, according to the input to propagate.

By default, the function returns `x_j`.
Custom layer should specialize this method with the desired behavior.

# Arguments

- `l`: A gnn layer.
- `x_i`: Features of the central node `i`.
- `x_j`: Features of the neighbor `j` of node `i`.
- `e_ij`: Features of edge `(i,j)`.

See also [`update_node`](@ref) and [`propagate`](@ref).
"""
function compute_message end

@inline compute_message(l, x_i, x_j, e_ij) = compute_message(l, x_i, x_j)
@inline compute_message(l, x_i, x_j) = x_j

_gather(x::NamedTuple, i) = map(x -> _gather(x, i), x)
_gather(x::Tuple, i) = map(x -> _gather(x, i), x)
_gather(x::AbstractArray, i) = NNlib.gather(x, i)
_gather(x::Nothing, i) = nothing

function compute_batch_message(l, g::GNNGraph, x, e)
    s, t = edge_index(g)
    xi = _gather(x, t)
    xj = _gather(x, s)
    m = compute_message(l, xi, xj, e)
    return m
end

##  Step 2

_scatter(aggr, e::NamedTuple, t) = map(e -> _scatter(aggr, e, t), e)
_scatter(aggr, e::Tuple, t) = map(e -> _scatter(aggr, e, t), e)
_scatter(aggr, e::AbstractArray, t) = NNlib.scatter(aggr, e, t)
_scatter(aggr, e::Nothing, t) = nothing

function aggregate_neighbors(g::GNNGraph, aggr, e)
    s, t = edge_index(g)
    _scatter(aggr, e, t)
end

aggregate_neighbors(g::GNNGraph, aggr::Nothing, e) = nothing

## Step 3

"""
    update_node(l, m̄, x)

Node update function for the GNN layer `l`,
returning a new set of node features `x′` based on old
features `x` and the aggregated message `m̄` from the neighborhood.

The input `m̄` is an array, a tuple or a named tuple,
reflecting the output of [`compute_message`](@ref).

By default, the function returns `m̄`.
Custom layers should  specialize this method with the desired behavior.

See also [`compute_message`](@ref), [`update_edge`](@ref), and [`propagate`](@ref).
"""
function update_node end

@inline update_node(l, m̄, x) = m̄

## Step 4

"""
    update_edge(l, m, e)

Edge update function for the GNN layer `l`,
returning a new set of edge features `e′` based on old
features `e` and the newly computed messages `m`
from the [`compute_message`](@ref) function.

By default, the function returns `e`.
Custom layers should specialize this method with the desired behavior.

See also [`compute_message`](@ref), [`update_node`](@ref), and [`propagate`](@ref).
"""
function update_edge end

@inline update_edge(l, m, e) = e

### end steps ###

ext: jl


text:
using DiffEqBenchmarks
DiffEqBenchmarks.bench_footer(WEAVE_ARGS[:folder],WEAVE_ARGS[:file])


text:
# This script has a heavy dependency on PyCall and PyTorch,
# so we should re-distribute the model and make this unnecessary
using PyCall
@pyimport torchvision.models as models

const model_dict = Dict([
        101=>models.resnet101,
        50=>models.resnet50,
        152=>models.resnet152,
        18=>models.resnet18,
        34=>models.resnet34
])

function load_resnet!(model, depth=101)
    conv_mode!(model, 1)
    stride3!()
    ptdict = state_dict(models, depth)
    #println(keys(ptdict))
    parse_conv!(model.conv1, ptdict, "conv1")
    parse_bn!(model.bn1, ptdict, "bn1")
    # Load the layers
    for i = 1:length(model.layers)
        seq = model.layers[i]
        for j = 1:length(seq.layers)
            bneck = seq.layers[j]
            key = string("layer",i, ".", j-1)
            for k = 1:(isa(bneck, Bottleneck) ? 3 : 2) #convs and bnorms
                ckey = string(key, ".conv",k)
                bkey = string(key, ".bn", k)
                parse_conv!(
                    getfield(bneck, Symbol(string("conv", k))),
                    ptdict, ckey)
                parse_bn!(
                    getfield(bneck, Symbol(string("bn", k))),
                    ptdict, bkey)
            end
            if bneck.downsample !== nothing
                conv = bneck.downsample.layers[1]
                bn = bneck.downsample.layers[2]
                parse_conv!(conv, ptdict, string(
                    key, ".", "downsample.0"
                ))
                parse_bn!(bn, ptdict, string(
                    key, ".", "downsample.1"
                ))
            end
        end
    end
    if model.output !== nothing
        parse_linear!(model.output, ptdict, "fc")
    end
end

state_dict(models, depth) =
    model_dict[depth](pretrained=true)[:state_dict]()

function parse_linear!(lin, ptdict, key)
    setval!(lin.w, reshape(
        ptdict[string(key, ".weight")][:numpy](),
        size(aval(lin.w))))

    setval!(lin.b, reshape(
        ptdict[string(key, ".bias")][:numpy](),
        size(aval(lin.b))))
end

function parse_conv!(conv, ptdict, key)
    setval!(conv.w, permutedims(
        ptdict[string(key, ".weight")][:numpy](), (4, 3, 2, 1)))
end

function parse_bn!(bn, ptdict, key)
    setval!(bn.w, vcat(ptdict[string(key, ".weight")][:numpy]()[:],
                       ptdict[string(key, ".bias")][:numpy]()[:]))
    sz = (1, 1, div(length(aval(bn.w)), 2), 1)
    bn.moments.mean =
        reshape(ptdict[string(key, ".running_mean")][:numpy](), sz)
    bn.moments.var =
        reshape(ptdict[string(key, ".running_var")][:numpy](), sz)
    if bn.remove_initfns
        bn.moments.meaninit = nothing
        bn.moments.varinit  = nothing
    end
end


text:
@testset "testFunctions.jl" begin

    data = CSO.dataset("cso_raw")
    years = unique(Dates.year.(data.Date))
    train, test = train_test_year(data, [2019, 2020])
    train_years = setdiff!(years, [2019, 2020])

    #The test set contains only the chosen years
    @test unique(Dates.year.(test.Date)) == [2019, 2020]

    #The train set contains only the chosen years
    @test unique(Dates.year.(train.Date)) == train_years
end

max_stars_repo_name: AlexLsn/CSO

max_issues_repo_name: AlexLsn/CSO

max_forks_repo_name: AlexLsn/CSO


text:
module MathMLRepr

include("token.jl")
include("lexer.jl")
include("parser.jl")
include("serializer.jl")
include("io.jl")
include("julia_expr.jl")

export
    ### TYPES
    Token,
    TokenStream,
    CompoExpr,
    LATEX,

    ### FUNCTIONS
    # core functions
    scan_tex,
    parse!,
    expr_to_mml,
    expr_to_tex,
    exlist_to_mml,
    exlist_to_tex,
    mml_element,
    mml_element_attr,
    htlatex_to_htmml,
    # pretty-printing
    print_token_stream,
    print_expr,
    pretty_print_mml,
    # conversion from Julia expressions
    num_to_expr,
    num_to_tex,
    vector_to_expr,
    vector_to_tex,
    matrix_to_expr,
    matrix_to_tex,
    ## io
    *,
    tex_to_mml,
    tex,
    texcon,
    show,
    dshow,
    htprint,
    LATEX_to_HTML,
    ## MACROS
    @l_str,
    @rm_str

  function __init__()
#    println("hello world from latex2mathml")
end # __init__

end


text:
### A Pluto.jl notebook ###
# v0.17.5

using Markdown
using InteractiveUtils

# ╔═╡ f16ef748-7c86-11ec-10a5-830de8dd2a2e
md"""
# Package Managment AST tests

- `Random` or `Downloads` indicates the name should be parsed as package name
- Single letter names are for everything that shouldn't be parsed
- Random `Base` thrown in to make sure it is not parsed
"""

# ╔═╡ e99b72b2-8939-430c-b8b3-25155559e7dc
macro >(expr) end

# ╔═╡ e7f324d7-e4b0-498a-a8a2-f94f52a9453d
@> using Random

# ╔═╡ fb917588-1499-433b-aeba-e6833346c968
@> using Random: a as x

# ╔═╡ d4f1ec8b-7597-45bf-9cea-4770802d516b
@> using Random.A.B

# ╔═╡ a68c8859-d7dd-4e8d-9996-f81216b84355
@> using Random.B.C: X

# ╔═╡ 589ee4e7-8fec-424a-82f7-cf0a700bf20a
@> using .A.B.C: X

# ╔═╡ 5e9f4f30-39f5-4c1d-b94d-4d937d37c47f
@> using Random.B.C, Downloads.K

# ╔═╡ 55449e27-c8db-4593-bc3b-0011179c6b20
@> import Random

# ╔═╡ 91b12b1c-1590-421e-8904-6b728bbf004d
@> import Base

# ╔═╡ 9d42b339-f64d-4ce8-b765-0ebd060dec55
@> import Random: a, b

# ╔═╡ d555a75d-eed6-4cb8-82a0-53f84ceb982a
@> import Random: A as Z

# ╔═╡ db8041c4-645c-4af9-8d33-aa4dfd45ea5f
@> import Random as X, Downloads as Y, .A

# ╔═╡ 42a5bed8-c761-4e9d-907d-25d49157647b
@> import Random, Base, .A as X

# ╔═╡ 7c4496e7-5b21-4748-9877-bff4a7ae5752
@> import .A: X, Y

# ╔═╡ 04faa311-541d-44c7-a189-e390c5cbdef9
@> import .A

# ╔═╡ 3490d653-47da-4cea-939e-076ad38558d7
@> import .A as X, ..A as Y

# ╔═╡ 1fc087ca-70b3-43d2-8793-70131227329f
@> import ..A

# ╔═╡ 7150f6b7-aa4f-44bf-a751-9c34c4a1c981
md"---"

# ╔═╡ b371f04d-0a89-4aa9-a03a-7c903270fd20
md"---"

# ╔═╡ 3749489b-49d3-4187-a2f6-1dddc0e97f92
@> import Random as

# ╔═╡ ed56bd75-8d7d-45ec-a3f9-c440141f7633
@> import NotExisting

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
"""

# ╔═╡ 00000000-0000-0000-0000-000000000002
PLUTO_MANIFEST_TOML_CONTENTS = """
# This file is machine-generated - editing it directly is not advised

[[ArgTools]]
uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"

[[Artifacts]]
uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"

[[Downloads]]
deps = ["ArgTools", "LibCURL", "NetworkOptions"]
uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"

[[LibCURL]]
deps = ["LibCURL_jll", "MozillaCACerts_jll"]
uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"

[[LibCURL_jll]]
deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"]
uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0"

[[LibSSH2_jll]]
deps = ["Artifacts", "Libdl", "MbedTLS_jll"]
uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8"

[[Libdl]]
uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"

[[MbedTLS_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"

[[MozillaCACerts_jll]]
uuid = "14a3606d-f60d-562e-9121-12d972cd8159"

[[NetworkOptions]]
uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"

[[Random]]
deps = ["Serialization"]
uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

[[Serialization]]
uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"

[[Zlib_jll]]
deps = ["Libdl"]
uuid = "83775a58-1f1d-513f-b197-d71354ab007a"

[[nghttp2_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d"
"""

# ╔═╡ Cell order:
# ╟─f16ef748-7c86-11ec-10a5-830de8dd2a2e
# ╟─e99b72b2-8939-430c-b8b3-25155559e7dc
# ╠═55449e27-c8db-4593-bc3b-0011179c6b20
# ╠═9d42b339-f64d-4ce8-b765-0ebd060dec55
# ╠═d555a75d-eed6-4cb8-82a0-53f84ceb982a
# ╠═db8041c4-645c-4af9-8d33-aa4dfd45ea5f
# ╠═42a5bed8-c761-4e9d-907d-25d49157647b
# ╠═7c4496e7-5b21-4748-9877-bff4a7ae5752
# ╠═04faa311-541d-44c7-a189-e390c5cbdef9
# ╠═3490d653-47da-4cea-939e-076ad38558d7
# ╠═1fc087ca-70b3-43d2-8793-70131227329f
# ╟─7150f6b7-aa4f-44bf-a751-9c34c4a1c981
# ╠═e7f324d7-e4b0-498a-a8a2-f94f52a9453d
# ╠═fb917588-1499-433b-aeba-e6833346c968
# ╠═d4f1ec8b-7597-45bf-9cea-4770802d516b
# ╠═a68c8859-d7dd-4e8d-9996-f81216b84355
# ╠═589ee4e7-8fec-424a-82f7-cf0a700bf20a
# ╠═5e9f4f30-39f5-4c1d-b94d-4d937d37c47f
# ╟─b371f04d-0a89-4aa9-a03a-7c903270fd20
# ╠═3749489b-49d3-4187-a2f6-1dddc0e97f92
# ╠═ed56bd75-8d7d-45ec-a3f9-c440141f7633
# ╠═91b12b1c-1590-421e-8904-6b728bbf004d
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


text:
#
# Global heap
#

const GLOBAL_HEAP_SIGNATURE = reinterpret(UInt32, UInt8['G', 'C', 'O', 'L'])[1]

immutable GlobalHeapID
    heap_offset::RelOffset
    index::UInt32
end
define_packed(GlobalHeapID)

isatend(f::JLDFile, gh::GlobalHeap) =
    gh.offset != 0 && f.end_of_data == gh.offset + 8 + sizeof(Length) + gh.length

heap_object_length(data::AbstractArray) = length(data)
heap_object_length(::Any) = 1

function write_heap_object(f::JLDFile, odr, data, wsession::JLDWriteSession)
    psz = sizeof(odr)*heap_object_length(data)
    objsz = 8 + sizeof(Length) + psz
    objsz += 8 - mod1(objsz, 8)

    io = f.io

    # This is basically a memory allocation problem. Right now we do it
    # in a pretty naive way. We:

    # 1. Put the object in the last created global heap if it fits
    # 2. Extend the last global heap if it's at the end of the file
    # 3. Create a new global heap if we can't do 1 or 2

    # This is not a great approach if we're writing objects of
    # different sizes interspersed with new datasets. The torture case
    # would be a Vector{Any} of mutable objects, some of which contain
    # large (>4080 byte) strings and some of which contain small
    # strings. In that case, we'd be better off trying to put the small
    # strings into existing heaps, rather than writing new ones. This
    # should be revisited at a later date.

    if objsz + 8 + sizeof(Length) < f.global_heap.free
        # Fits in existing global heap
        gh = f.global_heap
    elseif isatend(f, f.global_heap)
        # Global heap is at end and can be extended
        gh = f.global_heap
        delta = objsz - gh.free + 8 + sizeof(Length)
        gh.free += delta
        gh.length += delta
        seek(io, gh.offset + 8)
        write(io, gh.length)
        f.end_of_data += delta
    else
        # Need to create a new global heap
        heapsz = max(objsz, 4096)
        offset = f.end_of_data + 8 - mod1(f.end_of_data, 8)
        seek(io, offset)
        write(io, GLOBAL_HEAP_SIGNATURE)
        write(io, UInt32(1))      # Version & Reserved
        write(io, Length(heapsz)) # Collection size
        f.end_of_data = position(io) + heapsz
        gh = f.global_heap = f.global_heaps[h5offset(f, offset)] =
            GlobalHeap(offset, heapsz, heapsz, Int64[])
    end

    # Write data
    index = length(gh.objects) + 1
    objoffset = gh.offset + 8 + sizeof(Length) + gh.length - gh.free
    seek(io, objoffset)
    write(io, UInt16(index))           # Heap object index
    write(io, UInt16(1))               # Reference count
    write(io, UInt32(0))               # Reserved
    write(io, Length(psz))             # Object size

    # Update global heap object
    gh.free -= objsz
    push!(gh.objects, objoffset)

    # Write free space object
    if gh.free >= 8 + sizeof(Length)
        seek(io, objoffset + objsz)
        write(io, UInt64(0))           # Object index, reference count, reserved
        write(io, Length(gh.free - 8 - sizeof(Length))) # Object size
    end

    # Write data
    seek(io, objoffset + 8+sizeof(Length))
    write_data(f, data, odr, wsession) # Object data

    GlobalHeapID(h5offset(f, gh.offset), index)
end

# Force specialization on DataType
write_heap_object(f::JLDFile, odr::Type{Union{}}, data, wsession::JLDWriteSession) =
    error("ODR is invalid")

function Base.read(io::IO, ::Type{GlobalHeap})
    offset = position(io)
    read(io, UInt32) == GLOBAL_HEAP_SIGNATURE || throw(InvalidDataException())
    read(io, UInt32) == 1 || throw(UnsupportedVersionException())
    heapsz = read(io, Length)
    index = 1
    objects = Int64[]
    startpos = position(io)
    free = heapsz
    while free > 8 + sizeof(Length)
        push!(objects, position(io))
        objidx = read(io, UInt16)
        objidx == 0 && break
        objidx == index || throw(UnsupportedFeatureException())
        skip(io, 6)                    # Reference count and reserved
        sz = read(io, Length)          # Length
        skip(io, sz + 8 - mod1(sz, 8)) # Payload
        free = heapsz - Length(position(io) - startpos)
        index += 1
    end
    GlobalHeap(offset, heapsz, free, objects)
end

# TODO make this take an ODR
function read_heap_object{T,RR}(f::JLDFile{MmapIO}, hid::GlobalHeapID, rr::ReadRepresentation{T,RR})
    io = f.io
    if haskey(f.global_heaps, hid.heap_offset)
        gh = f.global_heaps[hid.heap_offset]
    else
        seek(io, fileoffset(f, hid.heap_offset))
        f.global_heaps[hid.heap_offset] = gh = read(io, GlobalHeap)
    end
    seek(io, gh.objects[hid.index]+8)
    len = Int(read(io, Length))
    n = div(len, sizeof(RR))
    len == n * sizeof(RR) || throw(InvalidDataException())

    inptr = f.io.curptr
    v = Vector{T}(n)
    if isa(RR, DataType) && RR <: T && isbits(T)
        unsafe_copy!(pointer(v), convert(Ptr{T}, inptr), Int(n))
    else
        @simd for i = 1:n
            if !jlconvert_canbeuninitialized(rr) || jlconvert_isinitialized(rr, inptr)
                @inbounds v[i] = jlconvert(rr, f, inptr, NULL_REFERENCE)
            end
            inptr += sizeof(RR)
        end
    end
    v
end


text:
mutable struct SpectralCube <: AbstractSpectralCube
    data::Array{T, 3} where T <: Real
    value_unit::Unitful.FreeUnits
    header::FITSHeader
    wcs::WCSTransform
    xaxis::Array{T, 1} where T <: Real
    yaxis::Array{T, 1} where T <: Real
    zaxis::Array{T, 1} where T <: Real
    xy_unit::Unitful.FreeUnits
    z_unit::Unitful.FreeUnits
    xy_resolution::Real
    z_resolution::Real
end

# --------------------------------------------------
# Cube functions
function SpectralCube(data::Array{T, 3} where T <: Real, header::FITSHeader)
    wcs = wcs3d_from_header(header)

    idx = zeros((3, size(data, 3)))
    idx[3, :] = 1:size(data, 3)
    z_axis = pix_to_world(wcs, idx)[3, :];
    if ("CUNIT3" ∈ keys(header) && occursin("km", header["CUNIT3"]))
        z_unit = u"km / s";
    end
    z_unit = u"m / s"
#     z_axis = to_kps_val.(z_axis * z_unit);
#     z_unit = u"km / s";

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

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

    xy_unit = u"°";

    val_unit = u"K";

    xy_resolution = abs(header["CDELT1"])
    z_resolution = abs(z_axis[1] - z_axis[2])

    return SpectralCube(data, val_unit, header, wcs, x_axis, y_axis, z_axis, xy_unit, z_unit, xy_resolution, z_resolution);
end

function subcube(cube::SpectralCube, xyz_range::CubeRange)
    xmin_incube = closest_index(cube.xaxis, xyz_range.xmin)
    xmax_incube = closest_index(cube.xaxis, xyz_range.xmax)
    ymin_incube = closest_index(cube.yaxis, xyz_range.ymin)
    ymax_incube = closest_index(cube.yaxis, xyz_range.ymax)
    zmin_incube = closest_index(cube.zaxis, xyz_range.zmin)
    zmax_incube = closest_index(cube.zaxis, xyz_range.zmax)

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

    if (cube.data |> size |> length) == 4
        data = cube.data[xmin_incube:xmax_incube, ymin_incube:ymax_incube, zmin_incube: zmax_incube, 1]
    else
        data = cube.data[xmin_incube:xmax_incube, ymin_incube:ymax_incube, zmin_incube: zmax_incube]
    end

    header = deepcopy(cube.header)
    header["CRPIX1"] = header["CRPIX1"] - xmin_incube + 1
    header["CRPIX2"] = header["CRPIX2"] - ymin_incube + 1
    header["CRPIX3"] = header["CRPIX3"] - zmin_incube + 1
    header["NAXIS"] = 3
    header["NAXIS1"] = size(data, 1)
    header["NAXIS2"] = size(data, 2)
    header["NAXIS3"] = size(data, 3)

    return SpectralCube(data, header)
end

function map_to_new_zaxis!(cube1::AbstractSpectralCube, cube2::AbstractSpectralCube)
    # this method will modify cube2
    # from cube1 to cube2
    new_data = zeros(eltype(cube1.data), size(cube2.data))
    for i = 1:size(cube1.data, 3)
        v = cube1.zaxis[i]
        new_ch_idx = closest_index(cube2.zaxis, v)
        new_data[:, :, new_ch_idx] = cube1.data[:, :, i]
    end
    new_data[new_data .== 0] .= NaN;
    cube2.data = new_data;
    return cube2
end

function map_to_new_zaxis!(cube::AbstractSpectralCube, new_zaxis::AbstractArray)
    # this method will modify cube
    new_data = zeros(eltype(cube.data), tuple([size(cube.data)[1:2]...,  length(new_zaxis)]...) )
    for (i, v) in enumerate(new_zaxis)
        old_i = closest_index(cube.zaxis, v)
        new_data[:, :, i] = cube.data[:, :, old_i]
    end
    # new_data[new_data .== 0] .= NaN;
    cube.data = new_data;
    cube.zaxis = new_zaxis
    return cube
end

function voxel_num(maskcube::AbstractSpectralCube)
    maskcube.data .|> isnan .|> (!) |> sum
end

function pixel_num(maskcube::AbstractSpectralCube)
    maskcube.data .|> isnan .|> (!) |> (x->sum(x, dims=3) .!= 0) |> sum
end

function avg_channel_num(maskcube::AbstractSpectralCube)
    voxel_num(maskcube) / pixel_num(maskcube)
end

function valid_map(masked_datacube::SpectralCube, rmsmap::Map; n_sigma=3)
    cover_map = masked_datacube.data .|> isnan .|> (!) |> x->sum(x, dims=3)[:, :, 1];
    data = masked_datacube.data[:, :, :] # make a copy
    data[isnan.(data)] .= 0.;
    sum_map = sum(data, dims=3)[:, :, 1];
    vm = sum_map .> (n_sigma .* rmsmap.data .* sqrt.(cover_map))
    return vm
end

function three_channels_13(spec::Vector, rms::Real)
    Nσ = 2.0
    max_idx = argmax(spec)
    if max_idx >= 2 && max_idx <= (length(spec) - 1)
        chs = [max_idx - 1, max_idx, max_idx + 1]
    else
        for i=1:length(spec) - 2
            chs = [i, i + 1, i + 2]
            if all(spec[chs] .>= Nσ * rms)
                return true
            end
        end
        return false
    end
    return all(spec[chs] .>= Nσ * rms)
end

function three_channels_18(spec::Vector, rms::Real)
    Nσ = 1.5
    max_idx = argmax(spec)
    if max_idx >= 2 && max_idx <= (length(spec) - 1)
        chs = [max_idx - 1, max_idx, max_idx + 1]
    else
        for i=1:length(spec) - 2
            chs = [i, i + 1, i + 2]
            if all(spec[chs] .>= Nσ * rms)
                return true
            end
        end
        return false
    end
    return all(spec[chs] .>= Nσ * rms)
end

function Base.string(cube::SpectralCube)
    # TODO
    return "SpectralCube info"
end

function Base.copy(cube::SpectralCube)
    SpectralCube(cube.data, cube.header)
end


text:
"""
    ridge(A, b, λ = 0)

Compute the solution to min ||Ax - b||² + λ||x||².
"""
ridge(A, b, λ = 0) = (A'A + λ * I) \ A'b


text:
"""
    Ysize(ℓₘᵢₙ, ℓₘₐₓ)

Compute total size of array of mode weights

Parameters
----------
ℓₘᵢₙ : int
    Integer satisfying 0 <= ℓₘᵢₙ <= ℓₘₐₓ
ℓₘₐₓ : int
    Integer satisfying 0 <= ℓₘᵢₙ <= ℓₘₐₓ

Returns
-------
i : int
    Total size of array of mode weights arranged as described below

See Also
--------
Yrange : Array of (ℓ, m) indices corresponding to this array
Yindex : Index of a particular element of the mode weight array

Notes
-----
This assumes that the modes are arranged (with fixed s value) as

    [
        Y(s, ℓ, m)
        for ℓ in range(ℓₘᵢₙ, ℓₘₐₓ+1)
        for m in range(-ℓ, ℓ+1)
    ]

"""
function Ysize(ℓₘᵢₙ, ℓₘₐₓ)
    # from sympy import symbols, summation, horner
    # from sympy.printing.pycode import pycode
    # ℓ,m,ℓₘᵢₙ,ℓₘₐₓ = symbols('ℓ,m,ℓₘᵢₙ,ℓₘₐₓ', integer=True)
    #
    # def nice(expr)
    #     return horner(expr.expand().simplify(), (m′ₘₐₓ, ℓₘᵢₙ, ℓₘₐₓ))
    #
    # Ysize = horner(
    #     summation(summation(1, (m, -ℓ, ℓ)), (ℓ, ℓₘᵢₙ, ℓₘₐₓ)).expand().simplify(),
    #     (ℓₘᵢₙ, ℓₘₐₓ)
    # )
    return ℓₘₐₓ * (ℓₘₐₓ + 2) - ℓₘᵢₙ^2 + 1
end

"""
    deduce_limits(ysize, [ℓmin])

Deduce the value of `(ℓmin, ℓmax)` that produces ``Y`` arrays of the given size.

If `ℓmin` is not given, it is assumed to be 0.  If it is set to `nothing`, the
smallest possible value of `ℓmin` will be used.  However, note that this is not
a well-posed problem; multiple combinations of `(ℓmin, ℓmax)` can give rise to
``Y`` arrays of the same size.

See also [`Ysize`](@ref)
"""
function deduce_limits(ysize, ℓmin=0)
    if ℓmin === nothing
        maxℓmax = (ysize-1) ÷ 2
        ℓmin_range = 0:maxℓmax
    else
        ℓmin_range = [ℓmin]
    end
    for ℓmin in ℓmin_range
        ℓmax = round(Int, √(ysize + ℓmin^2) - 1, RoundDown)
        if ℓmax * (ℓmax + 2) - ℓmin^2 + 1 == ysize
            return (ℓmin, ℓmax)
        end
    end
    error("Input size $ysize does not correspond to a possible array of modes with ℓmin=$ℓmin")
end

"""
    Yrange(ℓₘᵢₙ, ℓₘₐₓ)

Create an array of (ℓ, m) indices as in Y array

Parameters
----------
ℓₘᵢₙ : int
    Integer satisfying 0 <= ℓₘᵢₙ <= ℓₘₐₓ
ℓₘₐₓ : int
    Integer satisfying 0 <= ℓₘᵢₙ <= ℓₘₐₓ

Returns
-------
i : int
    Total size of array of mode weights arranged as described below

See Also
--------
Ysize : Total size of array of mode weights
Yindex : Index of a particular element of the mode weight array

Notes
-----
This assumes that the modes are arranged (with fixed s value) as

    [
        Y(s, ℓ, m)
        for ℓ in range(ℓₘᵢₙ, ℓₘₐₓ+1)
        for m in range(-ℓ, ℓ+1)
    ]

"""
function Yrange(ℓₘᵢₙ, ℓₘₐₓ)
    r = zeros(typeof(ℓₘᵢₙ), (Ysize(ℓₘᵢₙ, ℓₘₐₓ), 2))
    i = 1
    for ℓ in ℓₘᵢₙ:ℓₘₐₓ
        for m in -ℓ:ℓ
            r[i, 1] = ℓ
            r[i, 2] = m
            i += 1
        end
    end
    return r
end

"""
    Yindex(ℓ, m, ℓₘᵢₙ=0)

Compute index into array of mode weights

Parameters
----------
ℓ : int
    Integer satisfying ℓₘᵢₙ <= ℓ <= ℓₘₐₓ
m : int
    Integer satisfying -ℓ <= m <= ℓ
ℓₘᵢₙ : int, optional
    Integer satisfying 0 <= ℓₘᵢₙ.  Defaults to 0.

Returns
-------
i : int
    Index of a particular element of the mode-weight array as described below

See Also
--------
Ysize : Total size of array of mode weights
Yrange : Array of (ℓ, m) indices corresponding to this array

Notes
-----
This assumes that the modes are arranged (with fixed s value) as

    [
        Y(s, ℓ, m)
        for ℓ in range(ℓₘᵢₙ, ℓₘₐₓ+1)
        for m in range(-ℓ, ℓ+1)
    ]

"""
@inline function Yindex(ℓ, m, ℓₘᵢₙ=0)
    # from sympy import symbols, summation, horner
    # from sympy.printing.pycode import pycode
    # ℓ,m,m′,ℓₘᵢₙ, = symbols('ℓ,m,m′,ℓₘᵢₙ', integer=True)
    #
    # def nice(expr)
    #     return horner(expr.expand().simplify(), (ℓₘᵢₙ, ℓ, m))
    #
    # Yindex = horner(
    #     (Ysize.subs(ℓₘₐₓ, ℓ-1) + summation(1, (m′, -ℓ, m)) - 1).expand().simplify(),
    #     (ℓₘₐₓ, ℓ, m)
    # )
    return ℓ*(ℓ + 1) - ℓₘᵢₙ^2 + m + 1
end

"""
    theta_phi(Nθ, Nϕ, [T=Float64])

Construct (theta, phi) grid in `spinsfast` order.

Note that this order is different from the one assumed by this package;
use [`phi_theta`](@ref) for the opposite ordering.

"""
function theta_phi(nθ, nϕ, T=Float64)
    [[θ, ϕ][i] for θ in range(0, T(π), length=nθ), ϕ in range(0, 2*T(π), length=nϕ+1)[begin:end-1], i in 1:2]
end

"""
    phi_theta(Nϕ, Nθ, [T=Float64])

Construct (phi, theta) grid in order expected by this package.

See also [`theta_phi`](@ref) for the opposite ordering.

"""
function phi_theta(nϕ, nθ, T=Float64)
    [[ϕ, θ][i] for ϕ in range(0, 2*T(π), length=nϕ+1)[begin:end-1], θ in range(0, T(π), length=nθ), i in 1:2]
end


text:
##  Miscellaneous User Settings
module user_settings

nz = 82; 	# The number of grid points in the vertical
nx = 801;  # The number of grid points in the horizontal
nt = 3600;  # Number of timesteps
dz = 250.0;	# Vertical grid spacing: m
dx = 250.0;  # Horizontal Grid Spacing: m
dt = 1;  # Time Step secs
c_sound = 50.0;  # Speed of sound m/s
asscoef = 0.1; # Asselin coefficient
cmixh = 0.003; # Horizontal diffusion Coefficient
cmixv = 0.003; # Vertical diffusion coefficient
# ub = 0;  % Mean Environmental wind
profile_method = 3; # 1 - Dry Neutral Sounding, 2 - Fovell Weisman-Klemp (1982) Sounding, 3 - Moist Stable

# Kessler Microphysics Settings

qc0 = .001; # Autoconversion Threshold
k1 = .001; # Autoconversion Rate (s^-1)
k2 = 2.2; # Accretion Rate (s^-1)

 # Initial Condition Bubble Settings
bubble_switch = 0; # 1 - Bubble On, 0 - Bubble Off
radz = 3000.0;  # Vertical Radius (m)
radx = 5000.0; # Horizontal Radius (m)
zcnt = 2000.0;  # Center of thermal (meters above ground)
xcnt = 82000; # Horizontal Center (m)
thermamp = -6.0; # Thermal Amplitude (K)

 # Rayleigh Damping Top Settings
 raydmpcoef = 10/100;  # Rayleigh Damping Coefficient
 raydmpz = 14000; # Damping layer extends down to xx (meters)

 # Lateral Sponge Boundary
sponge_dist = .15;  # 10 Percent of Domain Width
x_dist_in = sponge_dist*nx*dx;
latdmpcoef = 0.9;
#  Parameters to describe the model environmental sounding

 if profile_method == 2

# surface pressure, 96500 N/m^2
 psurf = 96500.0;
# potential temp. at the surface (K), default is 300.0 K
 tsurf = 300.0;
# mixing ratio at the surface (kg/kg), default is 0.0161
 qsurf = 0.0161;
# mixing ratio @ 4 km AGL (kg/kg), default is 0.001
 q4km = 0.001;
# height of the tropopause (m), default is 12000.0 m
 ztr = 12000.0;
# temperature at the tropopause (K), default is 213.0 K
 temptr = 213.0;
# potential temp. at the tropopause (K), default is 343.0 K
 ttr = 343.0;

elseif profile_method == 3
 # surface pressure, 96500 N/m^2
 psurf = 96500.0;
# potential temp. at the surface (K), default is 300.0 K
 tsurf = 281.0;
# mixing ratio at the surface (kg/kg), default is 0.0161
 qsurf = 0.006;
# mixing ratio @ 4 km AGL (kg/kg), default is 0.001
 q4km = 0.001;
# height of the tropopause (m), default is 12000.0 m
 ztr = 12000.0;
# temperature at the tropopause (K), default is 213.0 K
 temptr = 218;
# potential temp. at the tropopause (K), default is 343.0 K
 ttr = 345.0;

end
end

max_stars_repo_name: jeffreycunn/toy-atmospheric-model

max_issues_repo_name: jeffreycunn/toy-atmospheric-model

max_forks_repo_name: jeffreycunn/toy-atmospheric-model


text:
module IntegerIB

using Random
using DataFrames
import Base.merge!

"""
    Computes the KL divergence for two 1D probability distributions.
"""
function KL_single(P, Q)
    return sum(P.*log2.(P./Q))
end

"""
    Computes the KL divergence for a collection of distributions in matrix form.
Given P of shape MxN and Q of shape MxL, returns a matrix of shape NxL.
"""
function KL(P, Q)
    kl = zeros(size(P)[2], size(Q)[2])
    for i in 1:size(P)[2]
        for j in 1:size(Q)[2]
            kl[i,j] = KL_single(P[:,i],Q[:,j])
        end
    end
    return kl
end

"""
    Computes the entropy of a distribution. If distribution is of shape (MxN),
treats each column as an independant distribution and returns an 1xN entropy vector.
"""
function entropy(p)
    if any(size(p) .== 1)
        return -sum(p.*log2.(p))
    else
        return -sum(p.*log2.(p), dims = 1)
    end
end

"""
    get_y(x, algorithm = "nn")

Returns the corresponding context vector `y` to the data vector `x`.
Several types of context are possible, and are selected by the `type` argument.
type = "nn" is the next neighbor to every point of `x`.
type = "an" stands for "adjacent neighbor", or the previous and next element of every element of x.
"""
function get_y(x, type = "nn")
    if type == "nn"
        y = x[2:end]
    elseif type == "an"
        y = Any[]
        for idx in collect(2:length(x)-1)
            push!(y, (x[idx-1], x[idx+1]))
        end
    end
    return y
end

function findin(el, array)
    return findall(x -> x == el, array)[1]
end

"""
    get_pxy(x,y)
Returns joint probability of x and y from normalized 2D histogram.
Automatically replaces 0 by 10^-16 to prevent Inf and NaN values during KL computation.
"""
function get_pxy(x,y)
    X = x[1:end-1]
    if length(X) != length(y)
        X = X[2:end]
    end
    count = 0
    x_elts, y_elts = sort(unique(X)), sort(unique(y))
    pxy = zeros(size(x_elts)[1], size(y_elts)[1])
    for (idx, el) in enumerate(X)
        xpos, ypos = findin(el, x_elts), findin(y[idx], y_elts)
        pxy[xpos, ypos] += 1
        count += 1
    end
    pxy[pxy .== 0.0] .= 10^-16 #replacing zeros

    return pxy./count
end

"""
    returns the conditional probability of y given x.
"""
function get_py_x(x, y)
    pyx = get_pxy(x, y)
    inv_px = 1 ./ get_px(x, y)
    cond_pyx = inv_px.*pyx
    return cond_pyx' #transposition in order to get the y dimension in dims = 1.
end

function get_py_x(pxy)
    inv_px = 1 ./ get_px(pxy)
    cond_pyx = inv_px.*pxy
    return cond_pyx' #transposition in order to get the y dimension in dims = 1.
end

"""
    Computes marginal probability of 'x'.
"""
function get_px(x, y)
    return sum(get_pxy(x, y), dims = 2)
end

function get_px(pxy)
    return sum(pxy, dims = 2)
end

"""
    Computes marginal probability of 'y'.
"""
function get_py(x, y)
    return sum(get_pxy(x, y), dims = 1)
end

function get_py(pxy)
    return sum(pxy, dims = 1)
end

function get_qt_x_init(px, algorithm)
    if algorithm == "IB"
        qt_x_init = rand(size(px)[1], size(px)[1])  # randomly initialize qt_x_init
        qt_x_init = qt_x_init.*(1 ./ sum(qt_x_init, dims = 1)) #normalize qt_x_init
        return qt_x_init
    elseif algorithm == "DIB"
        index = shuffle(collect(1:size(px)[1]))
        qt_x_init = zeros(size(px)[1], size(px)[1])
        for (i,v) in enumerate(index)
            qt_x_init[v,i] = 1
        end
        return replace_zeros(qt_x_init)
    else
        throw("Algorithm must be \"IB\" pr \"DIB\"")
    end
end

"""
    init_values(x, y)

Initializes and return clustering porbability matrices `qt_x`, `q` and `qy_t`.
'qt_x' refers to the probability of a category in x to be mapped to an element of t.
'q' is the overall probability of an element of t to occur. 'qy_t' is the probability that an element of t
corresponds to a given element of `y`.
"""
function init_values(x, y, algorithm = "IB")
    pxy = get_pxy(x, y)
    px = get_px(x, y)
    qt_x_init = get_qt_x_init(px, algorithm)
    qt_init = qt_x_init*px
    qy_t_init = pxy'*qt_x_init'.*(1 ./ qt_init)'
    return (qt_x_init, qt_init, qy_t_init)
end

function init_values(pxy, algorithm = "IB")
    px = get_px(pxy)
    qt_x_init = get_qt_x_init(px, algorithm)
    qt_init = qt_x_init*px
    qy_t_init = pxy'*qt_x_init'.*(1 ./ qt_init)'
    return (qt_x_init, qt_init, qy_t_init)
end

"""
Interface for IB clustering computation.
Input data 'x' must be 1D. β controls the amount of clustering : the smaller the beta, the greater the clustering (and information loss).
Can be initalized directly with the data via IB(x), or if you already have the histograms/probability of co-occurences of 'x' and 'y' via IB(pxy).
To use the deterministic IB algorithm, set algorithm to "DIB" like IB(x, "DIB").
"""
mutable struct IB
    algorithm::String
    β::Real
    x
    px
    py
    pxy
    py_x
    qt_x
    qt
    qy_t
    IB(x, y, β = 100, algorithm = "IB") = new(algorithm, β, x, get_px(x,y), get_py(x,y), get_pxy(x,y), get_py_x(x,y), init_values(x, y, algorithm)...)
    IB(x::Array{Float64,1}, β = 100, algorithm = "IB") = new(algorithm, β, x, get_px(x, get_y(x)), get_py(x, get_y(x)), get_pxy(x, get_y(x)), get_py_x(x, get_y(x)), init_values(x,  get_y(x), algorithm)...)
    IB(x::Array{Int64,1}, β = 100, algorithm = "IB") = new(algorithm, β, x, get_px(x, get_y(x)), get_py(x, get_y(x)), get_pxy(x, get_y(x)), get_py_x(x, get_y(x)), init_values(x,  get_y(x), algorithm)...)
    IB(pxy::Array{Float64,2}, β = 100, algorithm = "IB") = new(algorithm, β, nothing, get_px(pxy), get_py(pxy), pxy, get_py_x(pxy), init_values(pxy, algorithm)...)
end

"""
    Re-initialies values of IB model
"""
function init!(m::IB)
    m.qt_x, m.qt, m.qy_t = init_values(m.pxy, m.algorithm)
end

"""
    Peforms q(t|x) update step.
"""
function qt_x_step!(m::IB)
    kl_part = log.(m.qt)' .- m.β*KL(m.py_x, m.qy_t)
    if m.algorithm == "DIB"
        m.qt_x = zeros(size(m.qt_x))
        for x_index in 1:size(kl_part)[1]
            m.qt_x[argmax(kl_part[x_index,:]), x_index] = 1
        end
        m.qt_x = replace_zeros(m.qt_x)
    elseif m.algorithm == "IB"
        q = replace_zeros(exp.(kl_part))
        m.qt_x = q'.*(1 ./ sum(q', dims = 1)) #transpose and normalize q.
    else
        throw(ArgumentError("'algorithm' unknown must be either \"IB\" or \"DIB\""))
    end
end

"""
    Peforms q(t) update step and drops unused clusters.
"""
function qt_step!(m::IB, clust_thres; report = true)
    m.qt = m.qt_x*m.px
    dropped = m.qt .< clust_thres
    if any(dropped) && size(m.qt)[1] > 1
        if report
            println("$(length(dropped[dropped])) clusters dropped, $(length(dropped[.~dropped])) ramaining.")
        end
        m.qt = m.qt[.~dropped]
        m.qt_x = drop_rows(m.qt_x, dropped)
        m.qt_x = m.qt_x.*(1 ./ sum(m.qt_x, dims = 1)) #normalize cluster distribution again
        m.qt = m.qt_x*m.px
    end
end

"""
    Remove rows of 2D array 'value_arr'. Boolean array 'drop_arr' indicates which rows to drop.
"""
function drop_rows(value_arr, drop_arr)
    new_size = length(drop_arr[.~drop_arr])
    dropped_arr = zeros(new_size, size(value_arr)[2])
    count = 1
    for (idx, drop) in enumerate(drop_arr)
        if ~drop
            dropped_arr[count,:] = value_arr[idx,:]
            count += 1
        end
    end
    return dropped_arr
end

"""
    Performs q(y|t) update step.
"""
function qy_t_step!(m::IB)
    qt_y = m.qt_x*m.pxy
    m.qy_t = qt_y'.*(1 ./ m.qt')
end

"""
    replace_zeros(arr::Array{Float64,2}, cut = 10^-305)

Replaces zeros (values below 'cut') by 'cut' to avoid exponential term causing NaN at high beta values.
"""
function replace_zeros(arr::Array{Float64,2}, cut = 10^-305)
    arr[arr .< cut] .= cut
    return arr
end

"""
    Performs one iteration of the IB algorithm.
"""
function IB_step!(m::IB, clust_thres; report = true)
    qt_x_step!(m)
    qt_step!(m, clust_thres; report = report)
    qy_t_step!(m)
end

"""
    Checks if absolute or relative convergence criterion on L function is satisfied.
"""
function check_conv(L_prev, L_current, abs_thres, rel_thres)
    abs_change = abs(L_prev - L_current)
    rel_change = abs_change/L_prev
    if abs_change <= abs_thres || rel_change <= rel_thres
        return true
    else
        return false
    end
end

function calc_metrics(m::IB)
    ht = entropy(m.qt)
    hy_t = entropy(m.qy_t)*m.qt
    iyt = entropy(m.py) - hy_t[1] # [1] cause of dumb julia rules focing dot product of vector to be 2 dimensional.
    ht_x = entropy(m.qt_x)*m.px
    ixt = entropy(m.qt) - ht_x[1] # same as iyt.
    L = ht - ht_x[1] - m.β*iyt    # same
    return (ht, ixt, iyt, L)
end

"""
    IB_optimize(model::IB; abs_thres = 10^-3, rel_thres = 0.01, ct_thres = 5, n_conv = 1000)

Optimizes IB model of the data.
Keywords arguments are :
    - abs_thres : threshold for convergence in absoulte change of the optimization function (lagrangian) L.
    - rel_thres : threshold for convergence in relative change of L.
    - ct_thres : number of steps where L has to be < to abs_thres or rel_thres to declare convergence.
    - n_conv : limit of the number of iterations for the optimization algorithm.
    - scan (bool): when using 'DIB' algorithm, checks and merges clusters if it leads to reduction of cost function.
    - report : wether or not to say when and how convergence was reached.
"""
function IB_optimize!(model::IB; merge_thres = 5*10^-2, abs_thres = 10^-10, rel_thres = 0, ct_thres = 25, n_conv = 1000, scan = true, report = true)
    consec_conv = 0
    total_steps = 0
    L = calc_metrics(model)[4]
    while consec_conv <= ct_thres && total_steps <= n_conv
        IB_step!(model, merge_thres; report = report)
        if model.algorithm == "DIB" && scan
            brute_optimize!(model, false)
        end
        L_current = calc_metrics(model)[4]
        if check_conv(L, L_current, abs_thres, rel_thres)
            consec_conv += 1
        else
            consec_conv = 0
        end
        L = L_current
        total_steps += 1
        if total_steps == n_conv && report
            println("Convergence not reached after $n_conv iterations, stopping optimization.")
        elseif consec_conv == ct_thres && report
            println("Convergence reached after $total_steps iterations, stopping optimization")
        end
    end
end

"""
    scans through the qt probabilities and merges clusters where values are too low.
"""
function merge!(m::IB, clust_thres = 10^-5)
    dropped = m.qt .< clust_thres
    if any(dropped) && size(m.qt)[1] > 1
        m.qt = m.qt[.~dropped]
        m.qt_x = drop_rows(m.qt_x, dropped)
        m.qt_x = m.qt_x.*(1 ./ sum(m.qt_x, dims = 1)) #normalize cluster distribution again
        m.qt = m.qt_x*m.px
        m.qy_t = drop_rows(m.qy_t', dropped)'
    end
end

"""
    brute_optimize!(m::IB, findbest = true)

Performs bute force optimization of IB model, looking at all possible cluster combinations
and choosing the merges that reduce the cost function.
"""
function brute_optimize!(m::IB, findbest = true)
    if m.algorithm == "IB"
        println("IB agorithm must be DIB.")
        return
    end
    if size(m.qt_x)[1] == 1
        return
    end
    anybetter = false
    best = deepcopy(m)
    for t1 in range(1, stop = size(m.qt_x)[1]-1)
        for t2 in range(2, stop = size(m.qt_x)[1])
            if ~anybetter || findbest
                tmp = deepcopy(best)
                tmp.qt_x[t1, findall(x -> x == 1, m.qt_x[t2,:])] .= 1
                tmp.qt_x[t2,:] .= 0
                IB_step!(tmp, 0)
                # check if cost function L reduced relative to best so far
                if calc_metrics(tmp)[4] < calc_metrics(best)[4]
                    best = deepcopy(tmp)
                    mergedt1 = t1
                    mergedt2 = t2
                    anybetter = true
                    print("merged clusters $mergedt1 and $mergedt2 \n")
                end
            end
        end
    end
    merge!(best)
    if anybetter
        m.qt_x = best.qt_x
        m.qt = best.qt
        m.qy_t = best.qy_t
        return true
    else
        return false
    end
end

"""
    Scans the IB plane with various vlaues of beta to get the optimal curve in the IB plane.
"""
function get_IB_curve(m::IB, start = 0.1, stop = 400, step = 0.05; glob = false)
    beta = collect(start:step:stop)
    ixt, iyt = [], []
    for (idx, b) in enumerate(beta)
        m.β = b
        init!(m)
        if glob
            search_optima!(m)
        else
            IB_optimize!(m; report = false)
        end
        x, y = calc_metrics(m)[2:3]
        if idx == 1
            push!(ixt, x)
            push!(iyt, y)
        elseif y >= iyt[end]
            push!(ixt, x)
            push!(iyt, y)
        end
    end
    return ixt, iyt
end

"""
    search_optima!(m::IB, n_iter = 10000)

The original IB algorithm can converge to a local optima. To find the global optima,
search_optima! performs 'n_iter' optimization of input model 'm', randomly re-initializing it at every iteration.
The best optimization is kept and 'm' is modified in-place.
"""
function search_optima!(m::IB, n_iter = 10000)
    if m.algorithm == "DIB"
        println("DIB algorithm is already giving optima.")
        return
    end
    best = deepcopy(m)
    for i in 1:n_iter
        init!(m)
        IB_optimize!(m; report = false)
        if calc_metrics(m)[4] < calc_metrics(best)[4]
            best.qt_x = m.qt_x
            best.qt = m.qt
            best.qy_t = m.qy_t
        end
    end
    m.qt_x = best.qt_x
    m.qt = best.qt
    m.qy_t = best.qy_t
end

"""
    Helper function that puts all similar clusters next to each other for ease of readability.
"""
function group_equivalent!(df)
  for i = 1:size(df,2)-1
    for j = i+1:size(df,2)
      if df[!,i] == df[!,j]
        idxs = collect(1:size(df,2))
        idxs[j], idxs[i+1] = idxs[i+1], idxs[j]
        permutecols!(df, idxs)
      end
    end
  end
end

"""
    print_results(m::IB, disp_thres = 0.1)

Displays the clustering of data for an optimized input model 'm'.

Since the IB algorithm is not deterministic, some probabilities in the clustering matrix 'm.qt_x' are non-zero.
For ease of readability, every probability under 'disp_thres' is displayed as 0, everything else as 1.
The result is a 2D matrix where the rows represent the clusters and the column the initial categories.
The 1 tell to which cluster which category belongs.
If several 1 are present for a single category, it means that the optimization wasn't able to assign
a single cluster to this category for the chosen β value. You might try a diffent β, or use the DIB variant of the algorithm.
"""
function print_results(m::IB, disp_thres = 0.1)
    if ~isnothing(m.x)
        df = DataFrame(model.qt_x .> disp_thres, Symbol.(sort(unique(m.x))))
        group_equivalent!(df)
        display(df)
    else
        @warn "Initial data is probability distribution, categories will be displayed as x1, x2 ..., xn."
        df = DataFrame(model.qt_x .> disp_thres)
        group_equivalent!(df)
        display(df)
    end
end

export IB, search_optima!, brute_optimize!, IB_optimize!, calc_metrics, get_IB_curve, get_y, print_results

end


text:
#module SSEquilibrium

#export ModelParams, utility, TauchenApprox, HHBellmanMap, SolveHHBellman, UpdateWage, SolveWage

using Parameters
using LinearAlgebra
using Distributions
using Optim
using QuantEcon
using SpecialFunctions
using SparseArrays
using ThreadTools

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# Initialize parameters & asset grids

@doc """
    ModelParams

A structure containing all model parameters, including asset and matching productivity shock grids.
"""
@with_kw mutable struct ModelParams
    A::Float64              = 1.0                           # output productivity factor
    r::Float64              = 0.06                          # annual interest rate
    r_mon::Float64          = (1.0 + r)^(1.0/12.0) - 1.0    # monthly interest rate
    worker_surplus::Float64 = 0.5                           # worker's share of surplus
    β::Float64              = 0.9948                        # discount factor
    β_mon::Float64          = 1.0/(1.0 + r_mon)             # monthly depreciation rate
    γ::Float64              = 1.0                           # curvature of utility function
    λ::Float64              = 0.02                          # exogenous separation rate
    b::Float64              = 0.4                           # UI benefits
    B::Float64              = 0.15                          # utility from leisure
    κ::Float64              = 0.522                         # vacancy posting cost with log utility
    θ::Float64              = 1.0                           # labor market tightness, v/u, normalized to 1
    α::Float64              = 0.5                           # power term in the matching function
    η::Float64              = 0.3133                        # scale parameter in the matching function
    p_θ::Float64            = η*θ^α                         # unemployed worker matching rate
    q_θ::Float64            = η*θ^(α-1)                     # vacancy matching probability
    amin::Float64           = -6.0                          # minimum asset
    amax::Float64           = 120.0                         # maximum asset
    agrid::Vector{Float64}  = LinRange(amin, amax, 50)      # asset grid w/ 50 nodes
    ρ_x::Float64            = 0.97                          # persistence of idiosyncratic productivity ln(x)
    σ_x::Float64            = 0.13                          # std. dev. of innovation to ln(x)
    π_x                     = TauchenApprox(9, ρ_x, σ_x).distribution           # probability weights on ln(x)
    xgrid                   = exp.(TauchenApprox(9, ρ_x, σ_x).state_values)     # match productivity grid
end;

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
#
@doc """
    utility(para::ModelParams, c::Real)

Household Bernoulli utility function.
"""
function utility(para::ModelParams, c::Real)

    γ = para.γ # Store risk-aversion parameter

    # Define the utility mapping
    # condition on the input `c`
    if c > 0
        if γ == 1
          return log(c) # if σ = 1, CRRA utility is log
        else
          return (c^(1-γ)-1)/(1-γ) # if σ != 1, regular specification
        end
    else
        return -Inf # infinite disutility if c∈(-∞,0]
    end
end;

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
#
@doc """
    utility(para::ModelParams, cvec)

Household Bernoulli utility function applied to a vector.
"""
function utility(para::ModelParams, cvec)

  u(c) = utility(para,c)

  return u.(cvec)
end;

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
#
@doc """
    TauchenApprox(N, ρ, σ, μ, n_std)

Tauchen's (1996) method for approximating AR(1) process with finite markov chain.

**Output:**
⋅ transition    -- Estimated transition matrix
⋅ distribution  -- Estimated stationary distribution (vector)
⋅ state_values  -- Etimated state values (vector)
"""
function TauchenApprox(N::Integer, ρ::T1, σ::T2, μ=zero(promote_type(T1, T2)), n_std::Integer=3) where {T1 <: Real, T2 <: Real}

    std_norm_cdf(x::T) where {T <: Real} = 0.5 * SpecialFunctions.erfc(-x/sqrt(2))
    std_norm_cdf(x::Array{T}) where {T <: Real} = 0.5 .* SpecialFunctions.erfc(-x./sqrt(2))

    # Get discretized space
    a_bar = n_std * sqrt(σ^2 / (1 - ρ^2))
    y = range(-a_bar, stop=a_bar, length=N)
    d = y[2] - y[1]

    # Get transition probabilities
    Π = zeros(promote_type(T1, T2), N, N)
    Threads.@threads for row = 1:N
        # Do end points first
        Π[row, 1] = std_norm_cdf((y[1] - ρ*y[row] + d/2) / σ)
        Π[row, N] = 1 - std_norm_cdf((y[N] - ρ*y[row] - d/2) / σ)

        # fill in the middle columns
        for col = 2:N-1
            Π[row, col] = (std_norm_cdf((y[col] - ρ*y[row] + d/2) / σ) -
                           std_norm_cdf((y[col] - ρ*y[row] - d/2) / σ))
        end
    end

    # state values
    ȳ = y .+ μ / (1 - ρ) # center process around its mean (wbar / (1 - rho)) in new variable

    # renormalize transition matrix
    Π = Π./sum(Π, dims = 2)

    # Generate stationary distribution from Π
    N = size(Π)[1]
    π0 = ones(1,N)/N
    diff = 1.
    while diff > 1e-10
        π1 = π0*Π
        diff = norm(π1-π0,Inf)
        π0 = π1
    end
    π0 = copy(π0')

    return (transition = Π, distribution = π0, state_values = ȳ)
end;

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
#
@doc """
    HHBellmanMap(para::ModelParams, wage, W_old, U_old)

Iterates on the bellman equation using continuation value function.

**Output:**
⋅ W_new         -- Updated value of employment (matrix)
⋅ U_new         -- Updated value of unemployment (vector)
⋅ emp_policy    -- Asset policy function in case of employment (matrix)
⋅ unemp_policy  -- Asset policy function in case of unemployment (vector)
"""
function HHBellmanMap(para::ModelParams, wage, W_old, U_old)

    @unpack agrid, xgrid, π_x, β, b, B, p_θ, r   = para

    u(c) = utility(para,c) #shorthand for utility

    W_new           = similar(W_old)
    U_new           = similar(U_old)
    emp_policy      = similar(W_new, Int)
    unemp_policy    = similar(U_new, Int)

    N_x = length(xgrid)
    N_a = length(agrid)

    Threads.@threads for x_i  in 1:N_x

        if x_i == 5
            for a_i in 1:N_a
                cvec_emp = (1+r)*agrid[a_i] + wage[a_i,x_i] .- agrid
                obj_emp     = u(cvec_emp) .+ β * sum( π_x[i]*max.(W_old[:,i], U_old) for i in 1:N_x)
                obj_emp     = vec(obj_emp)
                W_new[a_i,x_i], emp_policy[a_i,x_i] = findmax(obj_emp)

                cvec_unemp = (1+r)*agrid[a_i] + b .- agrid
                obj_unemp   = u(cvec_unemp) .+ B .+ β*(1-p_θ)*U_old .+ β*(p_θ)*W_old[:,5]
                obj_unemp   = vec(obj_unemp)
                U_new[a_i], unemp_policy[a_i] = findmax(obj_unemp)
            end
        else
            for a_i in 1:N_a
                cvec_emp = (1+r)*agrid[a_i] + wage[a_i,x_i] .- agrid
                obj_emp     = u(cvec_emp) .+ β * sum( π_x[i]*max.(W_old[:,i], U_old) for i in 1:N_x)
                obj_emp     = vec(obj_emp)
                W_new[a_i,x_i], emp_policy[a_i,x_i] = findmax(obj_emp)
            end
        end
    end

    return W_new, U_new, emp_policy, unemp_policy
end;

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
#
@doc """
    SolveHHBellman(para::ModelParams, wage, W0, U0, ϵ=1e-6)

Iterates on the Bellman map until convergence.

**Output:**
⋅ W             -- Approximated value of employment (matrix)
⋅ U             -- Approximated value of unemployment (vector)
⋅ emp_policy    -- Approximated asset policy function in case of employment (matrix)
⋅ unemp_policy  -- Approximated asset policy function in case of unemployment (vector)
"""
function SolveHHBellman(para::ModelParams, wage, W0, U0, ϵ=1e-5)

    W_old = W0
    U_old = U0
    emp_policy      = similar(W_old, Int)
    unemp_policy    = similar(U_old, Int)

    diff = 1.

    while diff > ϵ
        W_new, U_new = HHBellmanMap(para, wage, W_old, U_old)
        diff = norm((U_new-U_old)[:],Inf)
        W_old = W_new
        U_old = U_new
    end

    W, U, emp_policy, unemp_policy = HHBellmanMap(para, wage, W_old, U_old)

    return W, U, emp_policy, unemp_policy
end;

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
#
@doc """
    UpdateWage(para::ModelParams, W, U, emp_policy, wage_old)

Applies step #3(c) of the Bils, Chang, and Kim (2011) steady state equilibrium algorithm to update the wage, w¹(a,x;θ⁰), object.

**Output:**
⋅ wage_new  -- New iteration of the wage mapping (matrix)
⋅ J         -- New iteration of the value of a matched job (matrix)
"""
function UpdateWage(para::ModelParams, W, U, emp_policy, wage_old)

    @unpack agrid, xgrid, β, λ, α, r, π_x = para

    N_x = length(xgrid)
    N_a = length(agrid)

    J = zeros(N_a, N_x)
    wage_new = similar(wage_old)

    Threads.@threads for x_i in 1:N_x
        for a_i in 1:N_a
            c_e = (1+r)*agrid[a_i] + wage_old[a_i,x_i] - emp_policy[a_i,x_i]
            J[a_i,x_i] = ((1-α)/α)*(W[a_i,x_i] - U[a_i])*c_e
        end
    end

    Threads.@threads for x_i in 1:N_x
        for a_i in 1:N_a
            wage_new[a_i,x_i] = 1*xgrid[x_i] - J[a_i,x_i] + β*(1-λ)*sum( π_x[i]*max.( J[emp_policy[a_i,x_i],i] , 0.0 ) for i in 1:N_x)
        end
    end

    return wage_new, J
end;

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# Initialize value functions and wage
@doc """
    InitializeFunctions(para::ModelParams)

Initializes W(a,x), U(a), and wage(a,x).
"""
function InitializeFunctions(para::ModelParams)
    @unpack agrid, xgrid, β, r, b, B = para
    N_a = length(agrid)
    N_x = length(xgrid)
    W       = ones(N_a,N_x)
    U       = ones(N_a)
    wage    = ones(N_a,N_x)

    Threads.@threads for x_i  in 1:N_x

        if x_i == 5
            for a_i in 1:N_a
                wage[a_i,x_i] = xgrid[x_i]*2
                W[a_i,x_i] = log((1+r)*agrid[a_i] + wage[a_i,x_i] - agrid[a_i])/(1-β)
                U[a_i]     = log((1+r)*agrid[a_i] + b - agrid[a_i] + B)/(1-β)
            end
        else
            for a_i in 1:N_a
                wage[a_i,x_i] = xgrid[x_i]*2
                W[a_i,x_i] = log((1+r)*agrid[a_i] + wage[a_i,x_i] - agrid[a_i])/(1-β)
            end
        end
    end

    return W, U, wage
end

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# Solve wage
@doc """
    SolveWage(para::ModelParams, ϵ=1e-6)

Applies steps #1-3 of the Bils, Chang, and Kim (2011) steady state equilibrium algorithm to obtain value functions, policy functions, and the wage mapping.

**Output:**
⋅ W             -- Approximated value of employment (matrix)
⋅ U             -- Approximated value of unemployment (vector)
⋅ emp_policy    -- Approximated asset policy function in case of employment (matrix)
⋅ unemp_policy  -- Approximated asset policy function in case of unemployment (vector)
⋅ wage          -- Approximated wage mapping (matrix)
"""
function SolveWage(para::ModelParams, ϵ=1e-5)

    W_old, U_old, wage_old = InitializeFunctions(para)

    """
    W_old       = ones(length(para.agrid),length(para.xgrid))
    U_old       = ones(length(para.agrid))
    c           = 0.9
    wage_old    = ones(length(para.agrid),length(para.xgrid))*(c*θ)
    """

    ζ_x         = 0.01
    difff       = 10
    counter     = 0

    while difff > ϵ
        W_new, U_new, emp_policy = SolveHHBellman(para, wage_old, W_old, U_old)
        wage_new, J              = UpdateWage(para, W_new, U_new, emp_policy, wage_old)
        difff                    = norm(wage_new - wage_old)
        wage_old                 = ζ_x*wage_new + (1-ζ_x)*wage_old
        W_old                    = W_new
        U_old                    = U_new
        counter = counter + 1
        println("Wage Solver Iteration: $(counter). Norm: $(difff).")
    end

    W, U, emp_policy, unemp_policy              = SolveHHBellman(para, wage_old, W_old, U_old)
    wage, J                                     = UpdateWage(para, W, U, emp_policy, wage_old)

    return W, U, J, emp_policy, unemp_policy, wage
end;

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
# Reservation productivity, x(a;w)
@doc """
    ReservationProductivity(para::ModelParams, W, U)
"""
function ReservationProductivity(para::ModelParams, W, U)

        @unpack xgrid, agrid = para
        N_a = length(agrid)
        N_x = length(xgrid)
        x_comps = zeros(N_a,N_x)
        # If for any given (a,x), U(a) > W(a,x)
        # then store x_comps(a,x) = 1.0 (unemployment)
        # other store x_comps(a,x) = 0.0 (employment)
        for x_i in 1:N_x
            for a_i in 1:N_a
                x_comps[a_i,x_i] = -1*W[a_i,x_i] > -1*U[a_i]
            end
        end
        # x_grid = position in xgrid for which W surpasses U
        # indicating entry of xgrid that allows the switch
        # from unemployment to employment for a given `a`
        x_opt   = zeros(N_a)
        for a_i in 1:N_a
            x_index     = findfirst(x_comps[a_i,:] .== 0.0)
            if x_index > 1
                x_opt[a_i]  = (xgrid[x_index] - xgrid[x_index-1])/2
            else
                x_opt[a_i]  = 0.0
            end
        end

    return x_opt
end;

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
#
@doc """
    InitializeTransitionMatrices(para::ModelParams, emp_policy, unemp_policy, x_star)
"""
function InitializeTransitionMatrices(para::ModelParams, emp_policy, unemp_policy, x_star)

    @unpack agrid, xgrid, π_x = para
    N_x  = length(xgrid)
    N_a  = length(agrid)
    H_ee = zeros(N_a*N_x,N_a*N_x) # trans. from employmed to employmed
    H_eu = zeros(N_a*N_x,N_a*N_x) # trans. from employed to unemployed
    H_ue = zeros(N_a*N_x,N_a*N_x) # trans. from unemployed to employed
    H_uu = zeros(N_a*N_x,N_a*N_x) # trans. from unemployed to unemployed

    ### Fill out the employment matrix
    Threads.@threads for a_i in 1:N_a
        for x_i in 1:N_x
            j       = a_i+N_a*(x_i-1) # index for (a,x)
            a_emp   = emp_policy[a_i,x_i] # index for a′ given employment
            a_unemp = unemp_policy[a_i] # index for a′ given unemployment
            x′_star = x_star[a_emp] # value of x′_star under employment
            for x_i′ in 1:N_x
                j′ = a_emp+N_a*(x_i′-1) # index for (a′,x′) given employment
                if xgrid[x_i′] > x′_star # if x′ is above reservation x
                    H_ee[j,j′] = π_x[x_i′] # prob. of (a,x) → (a′,x′) given employment
                else
                    H_eu[j,j′] = π_x[x_i′] # prob. of (a,x) → (a′,x′) given employment
                end
                j′ = a_unemp+N_a*(x_i′-1) # index for (a′,x′) given unemployment
                H_uu[j,j′] = π_x[x_i′]  # prob. of (a,x) → (a′,x′) given unemployment
                if x_i′ == 5 # if x′ = x̄ (x̄ is the middle (5th) entry of xgrid)
                    H_ue[j,j′] = 1 # prob. of (a,x) → (a′,x′) given unemployment
                end
            end
        end
    end

    return H_ee, H_eu, H_ue, H_uu
end;

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
#
@doc """
    FindStationaryMeasures(para::ModelParams, emp_policy, unemp_policy)
"""
function FindStationaryMeasures(para::ModelParams,  emp_policy, unemp_policy, x_star)

    @unpack p_θ = para

    H_ee, H_eu, H_ue, H_uu = InitializeTransitionMatrices(para, emp_policy, unemp_policy, x_star)
    N_emp = size(H_ee)[1]
    π0_emp = ones(1,N_emp)/N_emp
    N_unemp = size(H_uu)[1]
    π0_unemp = ones(1,N_unemp)/N_unemp
    diff_emp = 1.
    diff_unemp = 1.
    while diff_emp > 1e-10 && diff_unemp > 1e-10
        π1_emp = π0_emp*H_ee + p_θ*π0_unemp*H_ue
        π1_unemp = π0_emp*H_eu + (1-p_θ)*π0_unemp*H_uu
        π1_emp = π1_emp/sum(π1_emp)
        π1_unemp = π1_unemp/sum(π1_unemp)
        diff_emp = norm(π1_emp-π0_emp,Inf)
        diff_unemp = norm(π1_unemp-π0_unemp,Inf)
        π0_emp = π1_emp
        π0_unemp = π1_unemp
    end

    π_emp   = π0_emp
    π_unemp = π0_unemp

    return π_emp, π_unemp
end;

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
#
@doc """
    AdjustTheta(para::ModelParams, J)
"""
function AdjustTheta(para::ModelParams, J, π_unemp)

    @unpack κ, β, η, α, agrid = para

    N_a = length(agrid)

    # integral of J over ψ is the dot product of
    # J(:,x̄) and π_unemp(N_a*x̄-1:N_a*x̄)
    int_J = dot(J[:,5],π_unemp[(N_a*(4)+1):N_a*(5)])
    θ_new = ( κ/( β*η*int_J ) )^(1/(α-1))

    return θ_new
end;

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
#
@doc """
    SolveEquilibrium(para::ModelParams)
"""
function SolveEquilibrium(para::ModelParams)

    para = para
    @unpack θ, agrid, xgrid = para

    N_a = length(agrid)
    N_x = length(xgrid)

    diff    = 1
    θ_old   = θ
    ζ_θ     = 0.05
    counterr = 0

    # Declare variables to be returned
    # at the end of the function
    W               = zeros(N_a,N_x)
    U               = zeros(N_a)
    J               = zeros(N_a,N_x)
    emp_policy      = zeros(N_a,N_x)
    unemp_policy    = zeros(N_a)
    wage            = zeros(N_a,N_x)
    x_star          = zeros(N_a)
    π_emp           = zeros(N_a*N_x)
    π_unemp         = zeros(N_a*N_x)

    while diff > 1e-4
    para = ModelParams(θ=θ_old)
    W, U, J, emp_policy, unemp_policy, wage = SolveWage(para)
    x_star = ReservationProductivity(para, W, U)
    π_emp, π_unemp = FindStationaryMeasures(para, emp_policy, unemp_policy, x_star)
    θ_new = AdjustTheta(para, J, π_unemp)
    diff = abs(θ_new - θ_old)
    θ_old = ζ_θ*θ_new + (1-ζ_θ)*θ_old
    counterr = counterr + 1
    println("Equilibrium Iteration: $(counterr). Norm: $(diff).")
    end

    θ = θ_old

    return para, θ, W, U, J, emp_policy, unemp_policy, wage, x_star, π_emp, π_unemp
end;

# ######################################################################
# ######################################################################
# ######################################################################
# ######################################################################
#end


text:
using OceanTurb

using OceanTurb.TKEMassFlux: TKEParameters

@use_pyplot_utils # add utilities for plotting OceanTurb Fields

function makeplot!(fig, axs, model)

    fig.suptitle("\$ t = \$ $(prettytime(model.clock.time))")

    ℓ = CellField(model.grid)
    P = CellField(model.grid)
    B = CellField(model.grid)
    ϵ = CellField(model.grid)

    markerkwargs = Dict(:marker=>"s", :markersize=>1)

    for i in eachindex(ℓ)
        ℓ[i] = TKEMassFlux.mixing_length(model, i)
        P[i] = TKEMassFlux.production(model, i)
        B[i] = TKEMassFlux.buoyancy_flux(model, i)
        ϵ[i] = TKEMassFlux.dissipation(model, i)
    end

    sca(axs[1])
    cla()
    plot(model.solution.T; linestyle="-", markerkwargs...)
    removespines("top", "right")
    xlabel("Temperature (\$ {}^\\circ \\mathrm{C} \$)")
    ylabel(L"z \, \mathrm{(m)}")

    sca(axs[2])
    cla()
    plot(model.solution.e; linestyle="-", markerkwargs...)
    removespines("top", "right", "left")
    xlabel(L"e")

    sca(axs[3])
    cla()
    plot(P; linestyle="-", label="production", markerkwargs...)
    plot(B; linestyle="--", label="buoyancy flux", markerkwargs...)
    plot(-ϵ; linestyle=":", label="dissipation", markerkwargs...)
    removespines("top", "right", "left")
    xlabel(L"\partial_t e")
    legend(loc="lower right")
    maxlim = maximum(abs, xlim())
    xlim(-maxlim, maxlim)

    sca(axs[4])
    cla()
    plot(ℓ; linestyle="-", markerkwargs...)
    removespines("top", "right", "left")
    xlabel(L"\ell")

    sca(axs[5])
    cla()
    plot(model.state.K; linestyle="-", markerkwargs...)
    removespines("top", "right", "left")
    xlabel(L"K")

    for i = 2:length(axs)
        axs[i].tick_params(left=false, labelleft=false)
    end

    ylim(-64, 1)
    pause(0.1)

    return ℓ
end

constants = Constants(f=1e-4)

 N = 128        # Model resolution
 H = 128        # Vertical extent of the model domain
Qᵇ = 1e-7       # Surface buoyancy flux (positive implies cooling)
N² = 1e-5       # Interior/initial temperature gradient
Δt = 1minute

dTdz = N² / (constants.α * constants.g)
Qᶿ = Qᵇ / (constants.α * constants.g)

# Build the model with a Backward Euler timestepper
model = TKEMassFlux.Model(                  grid = UniformGrid(N=N, H=H),
                                         stepper = :BackwardEuler,
                           convective_adjustment = TKEMassFlux.VariablePrandtlConvectiveAdjustment(),
                                       constants = constants)

# Set initial condition
T₀(z) = 20 + dTdz * z
model.solution.T = T₀

# Set boundary conditions
model.bcs.T.top = FluxBoundaryCondition(Qᶿ)
model.bcs.T.bottom = GradientBoundaryCondition(dTdz)

# Run the model
fig, axs = subplots(ncols=5, figsize=(16, 5), sharey=true)

for iplot = 1:12
    run_until!(model, Δt, iplot * 1hour)
    OceanTurb.update!(model)
    makeplot!(fig, axs, model)
end


text:
#Generic plot interface for all plots needing X-Y axis with a single series
function xy_plot(x::AbstractVector, y::AbstractVector;
			mark::String = "bar",
			stack::Union{Bool, AbstractVector, Nothing} = nothing, #No op, here since other functions dispatch here
			step::Union{String, Nothing} = nothing,
			horizontal::Bool = false,
			legend::Bool = false,
			scale::Bool = false,
			kwargs...)

	#Validate arrays are same length
	if size(x) != size(y)
		error("Arrays x and y need to have the same length.")
	end

    ec = newplot(kwargs, ec_charttype = "xy plot")

	#General visualization defaults
	ec.xAxis = [Axis(_type = "category", data = x)]
	ec.yAxis = [Axis(_type = "value", nameGap = 50, scale = scale)] #Add some padding for later if name added
	ec.series = [XYSeries(name = "Series 1", _type = mark, data = y, step = step)]

	#Make plot horizontal
	horizontal ? flip!(ec) : nothing

	#Add legend if requested
	legend ? legend!(ec) : nothing

	return ec

end

function xy_plot(x::AbstractVector, y::AbstractArray;
			mark::Union{String, AbstractVector} = "bar",
			stack::Union{Bool, AbstractVector, Nothing} = nothing,
			step::Union{String, Nothing} = nothing,
			horizontal::Bool = false,
			legend::Bool = true,
			scale::Bool = false,
			kwargs...)

	# Allow for convenience of using single string to represent same mark for all series values
	typeof(mark) <: AbstractVector ? nothing : mark = [mark for i in 1:length(x)]

	#If there are missing values, set equal to zero so the graph stacks correctly
	eltype(y) >: Missing ? y_ = collect(Missings.replace(y,0)) : y_ = deepcopy(y)

	# Call 1-D method to build base
	ec = xy_plot(x, y_[:,1]; mark = mark[1], scale = scale, kwargs...)

	# Append remaining Y data
	for i in 2:size(y_)[2]
		push!(ec.series, XYSeries(_type = mark[i], data = y_[:,i]))
	end

	#stack: this logic feels janky, but seems to work
	if stack != nothing
		if stack == true
			[x.stack = 1 for x in ec.series]
		elseif typeof(stack) <: AbstractVector
			[x.stack = stack[i] for (i,x) in enumerate(ec.series)]
		else
			nothing
		end

	end

	#step
	step != nothing ? [x.step = step for x in ec.series] : nothing

	#Make plot horizontal
	horizontal ? flip!(ec) : nothing

	# Add default names to series
	seriesnames!(ec)

	#Add legend if requested
	legend ? legend!(ec) : nothing

	return ec

end

#dataframe, single series
function xy_plot(df::AbstractDataFrame, x::Symbol, y::Symbol;
			mark::String = "bar",
			stack::Union{Bool, AbstractVector, Nothing} = nothing, #No op, here since other functions dispatch here
			step::Union{String, Nothing} = nothing,
			horizontal::Bool = false,
			legend::Bool = false,
			scale::Bool = false,
			kwargs...)

	#Intialize for single series
	ec = xy_plot(df.x, df.y, mark = mark, stack = stack, step = step, horizontal = horizontal, legend = legend, scale = scale, kwargs...)

	#Name axes since we know them
	xaxis!(ec, name = string(x))
	yaxis!(ec, name = string(y))

	#Add legend if requested
	legend ? legend!(ec) : nothing

	return ec

end

#dataframe, multiple series
function xy_plot(df::AbstractDataFrame, x::Symbol, y::Symbol, group::Symbol;
			mark::String = "bar",
			stack::Union{Bool, AbstractVector, Nothing} = nothing, #No op, here since other functions dispatch here
			step::Union{String, Nothing} = nothing,
			horizontal::Bool = false,
			legend::Bool = true,
			scale::Bool = false,
			kwargs...)

	#create grouped df
	#This appears to handle missings gracefully via DataFrames package
	pivotdf = unstack(df, x, group, y)

	#Get number of groups
	numgroups = size(pivotdf, 2)

	ec = xy_plot(pivotdf.x, convert(Matrix, pivotdf[:, 2:end]), mark = mark, stack = stack, step = step, horizontal = horizontal, legend = legend, scale = scale, kwargs...)

	#Name axes since we know them
	xaxis!(ec, name = string(x))
	yaxis!(ec, name = string(y))

	#Add legend if requested
	legend ? legend!(ec) : nothing

	#Add series names
	seriesnames!(ec, names(pivotdf)[2:end])

	return ec

end


text:
abstract type InverterComponent <: DynamicComponent end

mutable struct DynamicInverter{
    C <: Converter,
    O <: OuterControl,
    IC <: InnerControl,
    DC <: DCSource,
    P <: FrequencyEstimator,
    F <: Filter,
} <: DynamicInjection
    static_injector::Union{Nothing, StaticInjection}
    ω_ref::Float64
    MVABase::Float64
    converter::C
    outer_control::O
    inner_control::IC
    dc_source::DC
    freq_estimator::P
    filter::F #add MVAbase here
    n_states::Int64
    states::Vector{Symbol}
    ext::Dict{String, Any}
    internal::InfrastructureSystemsInternal
end

function DynamicInverter(
    static_injector::StaticInjection,
    ω_ref::Float64,
    MVABase::Float64,
    converter::C,
    outer_control::O,
    inner_control::IC,
    dc_source::DC,
    freq_estimator::P,
    filter::F,
    ext::Dict{String, Any} = Dict{String, Any}(),
) where {
    C <: Converter,
    O <: OuterControl,
    IC <: InnerControl,
    DC <: DCSource,
    P <: FrequencyEstimator,
    F <: Filter,
}

    n_states = (
        converter.n_states +
        outer_control.n_states +
        inner_control.n_states +
        dc_source.n_states +
        freq_estimator.n_states +
        filter.n_states
    )

    states = vcat(
        converter.states,
        outer_control.states,
        inner_control.states,
        dc_source.states,
        freq_estimator.states,
        filter.states,
    )

    return DynamicInverter{C, O, IC, DC, P, F}(
        static_injector,
        ω_ref,
        MVABase,
        converter,
        outer_control,
        inner_control,
        dc_source,
        freq_estimator,
        filter,
        n_states,
        states,
        ext,
        InfrastructureSystemsInternal(),
    )
end

function DynamicInverter(;
    static_injector::StaticInjection,
    ω_ref::Float64,
    MVABase::Float64,
    converter::C,
    outer_control::O,
    inner_control::IC,
    dc_source::DC,
    freq_estimator::P,
    filter::F,
    ext::Dict{String, Any} = Dict{String, Any}(),
) where {
    C <: Converter,
    O <: OuterControl,
    IC <: InnerControl,
    DC <: DCSource,
    P <: FrequencyEstimator,
    F <: Filter,
}
    DynamicInverter(
        static_injector,
        ω_ref,
        MVABase,
        converter,
        outer_control,
        inner_control,
        dc_source,
        freq_estimator,
        filter,
        ext,
    )
end

IS.get_name(device::DynamicInverter) = get_name(device.static_injector)
get_bus(device::DynamicInverter) = get_bus(device.static_injector)
get_inverter_Sbase(device::DynamicInverter) = device.converter.s_rated
get_inverter_Vref(device::DynamicInverter) = device.V_ref
get_ω_ref(device::DynamicInverter) = device.ω_ref
get_V_ref(device::DynamicInverter) = get_voltage(get_bus(device))
get_P_ref(device::DynamicInverter) = get_activepower(device.static_injector)
get_Q_ref(device::DynamicInverter) = get_reactivepower(device.static_injector)
get_basepower(device::DynamicInverter) = get_basepower(device.static_injector)
get_ext(device::DynamicInverter) = device.ext
get_states(device::DynamicInverter) = device.states
get_n_states(device::DynamicInverter) = device.n_states
get_converter(device::DynamicInverter) = device.converter
get_outer_control(device::DynamicInverter) = device.outer_control
get_inner_control(device::DynamicInverter) = device.inner_control
get_dc_source(device::DynamicInverter) = device.dc_source
get_freq_estimator(device::DynamicInverter) = device.freq_estimator
get_filter(device::DynamicInverter) = device.filter
get_static_injector(device::DynamicInverter) = device.static_injector


text:
mutable struct WallTimer
    starttime_ns::typeof(Base.time_ns())
    paused_elapsed_ns::typeof(Base.time_ns())
    WallTimer() = new(0,0)
end

function start!(timer::WallTimer)
    timer.starttime_ns = (Base.time_ns)()
    return nothing
end
started(timer::WallTimer) = (timer.starttime_ns ≠ 0)

""" Return nanoseconds since timer was started or 0 if not yet started. """
function elapsed(timer::WallTimer)
    local elapsedtime_ns = (Base.time_ns)() - timer.starttime_ns
    return started(timer) * elapsedtime_ns
end

function pause!(timer::WallTimer)
    timer.paused_elapsed_ns = (Base.time_ns)() - timer.starttime_ns
    return nothing
end

function unpause!(timer::WallTimer)
    timer.starttime_ns = (Base.time_ns)()
    timer.starttime_ns -= timer.paused_elapsed_ns;
    return nothing
end

abstract type Scheduled end

"""
The schedule type stores an array of `Sheduled` objects, and a `WallTimer` object
that is used to keep time. The scheduler object keeps its own timer, and does
not reuse the timer in the game main loop, since the game timer can be reset every
frame, while the scheduler timer needs to keep absolute time.
"""
struct Scheduler
    actions::Array{Scheduled}
    timer::WallTimer

    function Scheduler()
        w = WallTimer()
        start!(w)
        new(Vector{Scheduled}(), w)
    end
end

const scheduler = Ref{Scheduler}()

elapsed(s::Scheduler) = elapsed(s.timer)
Base.push!(s::Scheduler, x::Scheduled) = push!(s.actions, x)
Base.filter!(x::Scheduled, s::Scheduler) = filter!(a->x!=a, s.actions)
clear!(s::Scheduler) = deleteat!(s.actions, 1:length(s.actions))

"""
A scheduled action that is run once. The action is a zero-argument function that is wrapped as a
WeakRef object. This ensures that schedules do not inadvertantly store references
to game objects. This does however mean that anynymous functions should not be used
as scheduled actions.

The time stored in the schedule is in absolute nanoseconds since epoch
(Note the epoch for nanoseconds can be arbitrary).
User facing times should always be in seconds, or fractions of a second, and usually
specifed as interval from current time. Hence the time should be converted before
being stored in a `Sheduled` object.
"""
struct OnceScheduled <: Scheduled
    action::WeakRef
    time::Int64
end

"""
A scheduled action that is repeated indefinitely. The time of next invocation is
stored as `time`, while the interval between each invocation is stored in `interval`.
Both of these are stored in units of nanoseconds.
"""
struct RepeatScheduled <: Scheduled
    action::WeakRef
    time::Int64
    interval::Int64
end

"""
A contingent schdule. This type of schedule checks the return value of the action.
If the action method returns `nothing`, no further action is scheduled. Otherwise,
the return value is considered to be the interval, in seconds, to the next invocation of the action.

Since action methods are written by end users, the return value of those methods are
not considered significant for ease of use. Hence ContingentScheduled is a separate
type that should not be end user visible.
"""
struct ContingentScheduled <: Scheduled
    action::WeakRef
    time::Int64
end

"Run all actions in the global scheduler that are due"
function tick!(s::Scheduler)
    t = elapsed(s.timer)
    for x in s.actions
        tick(x, t, s)
    end
    return
end

"Run a single scheduled action if due"
function tick(x::OnceScheduled, elapsed, s=scheduler[])
    if x.time <= elapsed && x.action.value != nothing
        @debug "Running single scheduled function" x.action.value
        Base.invokelatest(x.action.value)
        filter!(x, s)
    end

end

"Run a repeated scheduled action if due. If run, this method will add a new scheduled action to the scheduler"
function tick(x::RepeatScheduled, elapsed, s=scheduler[])
    if x.time <= elapsed && x.action.value != nothing
        @debug "Running repeated scheduled function" x.action.value
        Base.invokelatest(x.action.value)
        filter!(x, s)
        push!(s, RepeatScheduled(x.action, x.interval+elapsed, x.interval))
    end

end

"Run a contingent schduled action if due. If run, and not stopped, add a new scheduled action to the scheduler"
function tick(x::ContingentScheduled, elapsed, s=scheduler[])
    if x.time <= elapsed && x.action.value != nothing
        @debug "Running contingent scheduled function" x.action.value
        r = Base.invokelatest(x.action.value)
        filter!(x, s)
        if r == nothing
            return
        else
            push!(s, ContingetScheduled(x.action, 1e9*r+elapsed))
        end
    end
end

"""
    schedule_once(f::Function, interval)

Takes a function and an interval in seconds, and sets the function to run after that interval.
"""
function schedule_once(f::Function, interval)
    t = elapsed(scheduler[])
    push!(scheduler[], OnceScheduled(WeakRef(f), t+interval*1e9) )
    @debug "Added Single Schedule" f
end

"""
    schedule_unique(f::Function, interval)

Takes a function and an interval in seconds, and sets the function to run after that interval only if the same function handle was not previously set.
"""
function schedule_unique(f::Function, interval)
    filter(WeakRef(f), scheduler)
    push!(scheduler[], OnceScheduled(WeakRef(f), elapsed(scheduler[])+interval*1e9) )
    @debug "Added Unique Schedule" f
end

"""
    schedule_interval(f::Function, interval, first_interval=interval)

Takes a function handle and time in seconds, and sets the function to run after that interval.
Optional third argument first_interval can be passed to wait for that amount of time before first execution.
"""
function schedule_interval(f::Function, interval, first_interval=interval)
     push!(scheduler[], RepeatScheduled(WeakRef(f), elapsed(scheduler[])+first_interval*1e9, interval*1e9 ))
     @debug "Added Repeated Schedule" f
end

"""
    unschedule(f::Function)

Removes a function from the schedule.
"""
unschedule(f::Function) = filter!(WeakRef(f), scheduler[])