"""
`annualtoquarter(v)`

Convert from annual to quarter frequency... by dividing by 4.
"""
function annualtoquarter(v)
    v / 4
end

"""
`quartertoannual(v)`

Convert from quarter to annual frequency... by multiplying by 4.
"""
function quartertoannual(v)
    4 * v
end

"""
`quartertoannualpercent(v)`

Convert from quarter to annual frequency in percent... by multiplying by 400.
"""
function quartertoannualpercent(v)
    400 * v
end

"""
`nominal_to_real(col, df; deflator_mnemonic = :GDPDEF)`

Converts nominal to real values using the specified deflator.

## Arguments
- `col`: Symbol indicating which column of `df` to transform
- `df`: DataFrame containining series for proper population measure and `col`

## Keyword arguments

- `deflator_mnemonic`: indicates which deflator to use to calculate real values. Default
  value is the FRED GDP Deflator mnemonic.
"""
function nominal_to_real(col::Symbol, df::DataFrame; deflator_mnemonic::Symbol = :GDPDEF)
    return df[col] ./ df[deflator_mnemonic]
end

"""
```
percapita(m, col, df)
percapita(col, df, population_mnemonic)
```

Converts data column `col` of DataFrame `df` to a per-capita value.

The first method checks `hpfilter_population(m)`. If true, then it divides by
the filtered population series. Otherwise it divides by the result of
`parse_population_mnemonic(m)[1]`.

## Arguments

- `col`: `Symbol` indicating which column of data to transform
- `df`: `DataFrame` containining series for proper population measure and `col`
- `population_mnemonic`: a mnemonic found in `df` for some population measure
"""
function percapita(m::AbstractModel, col::Symbol, df::DataFrame)
    if hpfilter_population(m)
        population_mnemonic = Nullable(:filtered_population)
    else
        population_mnemonic = parse_population_mnemonic(m)[1]
        if isnull(population_mnemonic)
            error("No population mnemonic provided")
        end
    end
    percapita(col, df, get(population_mnemonic))
end

function percapita(col::Symbol, df::DataFrame, population_mnemonic::Symbol)
    df[col] ./ df[population_mnemonic]
end

"""
```
series_lag_n = lag(series, n)
```
Returns a particular data series lagged by n periods
"""
function lag(series::AbstractArray, n::Int)
    return vcat(NaN*ones(n), series[1:end-n])
end

"""
```
yt, yf = hpfilter(y, λ)
```

Applies the Hodrick-Prescott filter (\"H-P filter\"). The smoothing parameter `λ` is applied
to the columns of `y`, returning the trend component `yt` and the cyclical component `yf`.
For quarterly data, one can use λ=1600.

Consecutive missing values at the beginning or end of the time series are excluded from the
filtering. If there are missing values within the series, the filtered values are all NaN.

See also:
```
Hodrick, Robert; Prescott, Edward C. (1997). \"Postwar U.S. Business Cycles: An Empirical
Investigation\". Journal of Money, Credit, and Banking 29 (1): 1–16.
```
"""
function hpfilter(y::AbstractVector, λ::Real)
    # Indices of consecutive NaN elements at beginning
    i = 1
    j = length(y)
    while isnan(y[i])
        i = i+1
    end
    while isnan(y[j])
        j = j-1
    end

    # Filter and adjust for NaNs
    yt_, yf_ = hpfilter_(y[i:j], λ)
    yt = [fill(NaN, i-1); yt_; fill(NaN, length(y)-j)]
    yf = [fill(NaN, i-1); yf_; fill(NaN, length(y)-j)]

    return yt, yf
end

function hpfilter_(y::AbstractVector, λ::Real)
    n = length(y)
    a = spzeros(n,n)
    for i = 3:n-2
        a[i,i]   = 6λ+1
        a[i,i+1] = -4λ
        a[i,i+2] = λ
        a[i,i-2] = λ
        a[i,i-1] = -4λ
    end

    a[2,2] = 1+5λ
    a[2,3] = -4λ
    a[2,4] = λ
    a[2,1] = -2λ
    a[1,1] = 1+λ
    a[1,2] = -2λ
    a[1,3] = λ

    a[n-1,n-1] = 1+5λ
    a[n-1,n-2] = -4λ
    a[n-1,n-3] = λ
    a[n-1,n]   = -2λ
    a[n,n]     = 1+λ
    a[n,n-1]   = -2λ
    a[n,n-2]   = λ

    yt = a\y
    yf = y-yt

    return yt, yf
end

"""
```
difflog(x::AbstractVector)
```
"""
function difflog(x::Vector)
    [NaN; log.(x[2:end]) - log.(x[1:end-1])]
end

"""
```
difflog(x::DataArray{AbstractFloat})
```
"""
function difflog(x::DataArray)
    DSGE.na2nan!(x)
    return difflog(convert(Vector, x))
end

"""
```
oneqtrpctchange(y)
```

Calculates the quarter-to-quarter percentage change of a series.
"""
function oneqtrpctchange(y::AbstractVector)
    100 * difflog(y)
end

## REVERSE TRANSFORMS

"""
```
loggrowthtopct(y)
```

Transform from annualized quarter-over-quarter log growth rates to annualized
quarter-over-quarter percent change.

### Note

This should only be used in Model 510, which has the core PCE inflation
observable in annualized log growth rates.
"""
function loggrowthtopct(y::AbstractArray)
    100. * (exp.(y/100.) - 1.)
end

"""
```
loggrowthtopct_percapita(y, pop_growth)
```

Transform from annualized quarter-over-quarter log per-capita growth rates to
annualized quarter-over-quarter aggregate percent change.

### Note

This should only be used in Model 510, which has the output growth observable in
annualized log per-capita growth rates.

### Inputs

- `y`: the data we wish to transform to annualized percent change from
  annualized log growth rates. `y` is either a vector of length `nperiods` or an
  `ndraws x `nperiods` matrix.

- `pop_growth::Vector`: the length `nperiods` vector of log population growth
  rates.
"""
function loggrowthtopct_percapita(y::AbstractArray, pop_growth::AbstractVector)
    # `y` is either a vector of length `nperiods` or an
    # `ndraws` x `nperiods` matrix
    if ndims(y) == 1
        nperiods = length(y)
    else
        nperiods = size(y, 2)

        # Transpose `pop_growth` to a 1 x `nperiods` row vector so it can be
        # broadcasted to match the dimensions of `y`
        pop_growth = pop_growth'
    end

    @assert length(pop_growth) == nperiods "Length of pop_growth ($(length(pop_growth))) must equal number of periods of y ($nperiods)"

    100. * ((exp.(y/100.) .* exp.(pop_growth).^4) - 1.)
end

"""
```
loggrowthtopct_annualized(y)
```

Transform from log growth rates to annualized quarter-over-quarter percent change.
"""
function loggrowthtopct_annualized(y::AbstractArray)
    100. * (exp.(y/100.).^4 - 1.)
end

"""
```
loggrowthtopct_annualized_percapita(y, pop_growth)
```

Transform from log per-capita growth rates to annualized aggregate (not
per-capita) quarter-over-quarter percent change.

### Note

This should only be used for output, consumption, investment
and GDP deflator (inflation).

### Inputs

- `y`: the data we wish to transform to annualized percent change from
  quarter-over-quarter log growth rates. `y` is either a vector of length
  `nperiods` or an `ndraws x `nperiods` matrix.

- `pop_growth::Vector`: the length `nperiods` vector of log population growth
  rates.
"""
function loggrowthtopct_annualized_percapita(y::AbstractArray, pop_growth::AbstractVector)
    # `y` is either a vector of length `nperiods` or an
    # `ndraws` x `nperiods` matrix
    if ndims(y) == 1
        nperiods = length(y)
    else
        nperiods = size(y, 2)

        # Transpose `pop_growth` to a 1 x `nperiods` row vector so it can be
        # broadcasted to match the dimensions of `y`
        pop_growth = pop_growth'
    end

    @assert length(pop_growth) == nperiods "Length of pop_growth ($(length(pop_growth))) must equal number of periods of y ($nperiods)"

    100. * (exp.(y/100. .+ pop_growth).^4 - 1.)
end

"""
```
logleveltopct_annualized(y, y0 = NaN)
```

Transform from log levels to annualized quarter-over-quarter percent change.

### Inputs

- `y`: the data we wish to transform to annualized quarter-over-quarter percent
  change from log levels. `y` is either a vector of length `nperiods` or an
  `ndraws x `nperiods` matrix.

- `y0`: the last data point in the history (of state or observable)
  corresponding to the `y` variable. This is required to compute a percent
  change for the first period.
"""
function logleveltopct_annualized(y::AbstractArray, y0::Real = NaN)
    # `y_t1` is an array of the same size as `y`, representing the previous
    # period observations for each draw
    if ndims(y) == 1
        y_t1 = vcat([y0], y[1:end-1])
    else
        ndraws = size(y, 1)
        y0s  = fill(y0, ndraws, 1)
        y_t1 = hcat(y0s, y[:, 1:end-1])
    end

    # Subtract log levels to get log growth rates, then take the exponential to
    # get growth rates
    100. * (exp.(y./100. - y_t1./100.).^4 .- 1.)
end

"""
```
logleveltopct_annualized_percapita(y, pop_growth, y0 = NaN)
```

Transform from per-capita log levels to annualized aggregate (not per-capita)
quarter-over-quarter percent change.

### Note

This is usually applied to labor supply (hours worked per hour), and
probably shouldn't be used for any other observables.

### Inputs

- `y`: the data we wish to transform to annualized aggregate
  quarter-over-quarter percent change from per-capita log levels. `y` is either
  a vector of length `nperiods` or an `ndraws x `nperiods` matrix.

- `pop_growth::Vector`: the length `nperiods` vector of log population growth
  rates.

- `y0`: The last data point in the history (of state or observable)
  corresponding to the `y` variable. This is required to compute a percent
  change for the first period.
"""
function logleveltopct_annualized_percapita(y::AbstractArray, pop_growth::AbstractVector, y0::Real = NaN)
    # `y_t1` is an array of the same size as `y`, representing the previous
    # period observations for each draw
    if ndims(y) == 1
        nperiods = length(y)
        y_t1 = vcat([y0], y[1:end-1])
    else
        (ndraws, nperiods) = size(y)
        y0s  = fill(y0, ndraws, 1)
        y_t1 = hcat(y0s, y[:, 1:end-1])

        # Transpose `pop_growth` to a 1 x `nperiods` row vector so it can be
        # broadcasted to match the dimensions of `y`
        pop_growth = pop_growth'
    end

    @assert length(pop_growth) == nperiods "Length of pop_growth ($(length(pop_growth))) must equal number of periods of y ($nperiods)"

    # Subtract log levels to get log growth rates, then take the exponential to
    # get growth rates
    100. * (exp.(y./100. - y_t1./100. .+ pop_growth).^4 .- 1.)
end

"""
```
get_nopop_transform(transform::Function)
```
Returns the corresponding transformation which doesn't add back population
growth. Used for shock decompositions, deterministic trends, and IRFs, which are
given in deviations.
"""
function get_nopop_transform(transform::Function)
    transform4q = if transform == loggrowthtopct_annualized_percapita
        loggrowthtopct_annualized
    elseif transform == logleveltopct_annualized_percapita
        logleveltopct_annualized
    else
        transform
    end
end

"""
```
get_irf_transform(transform::Function)
```
Returns the IRF-specific transformation, which doesn't add back population
growth (since IRFs are given in deviations).
"""
function get_irf_transform(transform::Function)
    transform4q = if transform == loggrowthtopct_annualized_percapita
        loggrowthtopct_annualized
    elseif transform == logleveltopct_annualized_percapita
        logleveltopct_annualized
    else
        transform
    end
end

"""
```
get_transform4q(transform::Function)
```
Returns the 4-quarter transformation associated with the annualizing transformation.
"""
function get_transform4q(transform::Function)
    if transform == loggrowthtopct_annualized_percapita
        loggrowthtopct_4q_percapita
    elseif transform == loggrowthtopct_annualized
        loggrowthtopct_4q
    elseif transform == logleveltopct_annualized_percapita
        logleveltopct_4q_percapita
    elseif transform == logleveltopct_annualized
        logleveltopct_4q
    elseif transform == quartertoannual
        quartertoannual
    elseif transform == identity
        identity
    else
        error("4q equivalent not implemented for $transform")
    end
end

"""
```
loggrowthtopct_4q(y, data = fill(NaN, 3))
```

Transform from log growth rates to 4-quarter percent change.

### Inputs

- `y`: the data we wish to transform to aggregate 4-quarter percent change from
  log per-capita growth rates. `y` is either a vector of length `nperiods` or an
  `ndraws x `nperiods` matrix.

- `data`: if `y = [y_t, y_{t+1}, ..., y_{t+nperiods-1}]`, then
  `data = [y_{t-3}, y_{t-2}, y_{t-1}]`. This is necessary to compute
  4-quarter percent changes for the first three periods.
"""
function loggrowthtopct_4q(y::AbstractArray, data::AbstractVector = fill(NaN, 3))
    @assert length(data) == 3 "Length of data ($(length(data))) must be 3"

    # Prepend previous three periods to `y`
    y = prepend_data(y, data)

    # `y` is either a vector of length `nperiods+3` or an
    # `ndraws` x `nperiods+3` matrix
    if ndims(y) == 1
        y_4q = y[1:end-3] + y[2:end-2] + y[3:end-1] + y[4:end]
    else
        y_4q = y[:,  1:end-3] + y[:, 2:end-2] + y[:, 3:end-1] + y[:, 4:end]
    end

    100. * (exp.(y_4q/100.) - 1.)
end

"""
```
loggrowthtopct_4q_percapita(y, pop_growth, data = fill(NaN, 3))
```
Transform from log per-capita growth rates to aggregate 4-quarter percent
change.

### Note

This should only be used for output, consumption, investment, and GDP deflator
(inflation).

### Inputs

- `y`: the data we wish to transform to aggregate 4-quarter percent change from
  log per-capita growth rates. `y` is either a vector of length `nperiods` or an
  `ndraws x `nperiods` matrix.

- `pop_growth::Vector`: the length `nperiods` vector of log population growth
  rates.

- `data`: if `y = [y_t, y_{t+1}, ..., y_{t+nperiods-1}]`, then
  `data = [y_{t-3}, y_{t-2}, y_{t-1}]`. This is necessary to compute
  4-quarter percent changes for the first three periods.
"""
function loggrowthtopct_4q_percapita(y::AbstractArray, pop_growth::AbstractVector,
                                     data::Vector = fill(NaN, 3))
    @assert length(data) == 3 "Length of data ($(length(data))) must be 3"

    # Four-quarter population growth
    pop_growth_4q = pop_growth[1:end-3] + pop_growth[2:end-2] + pop_growth[3:end-1] + pop_growth[4:end]

    # Prepend previous three periods to `y`
    y = prepend_data(y, data)

    # `y` is either a vector of length `nperiods+3` or an
    # `ndraws` x `nperiods+3` matrix
    if ndims(y) == 1
        y_4q = y[1:end-3] + y[2:end-2] + y[3:end-1] + y[4:end]
        nperiods = length(y_4q)
    else
        y_4q = y[:, 1:end-3] + y[:, 2:end-2] + y[:, 3:end-1] + y[:, 4:end]
        nperiods = size(y_4q, 2)

        # Transpose `pop_growth` to a 1 x `nperiods` row vector so it can be
        # broadcasted to match the dimensions of `y_4q`
        pop_growth_4q = pop_growth_4q'
    end

    @assert length(pop_growth_4q) == nperiods "Length of pop_growth ($(length(pop_growth))) must equal number of periods of y ($nperiods)"

    100. * (exp.(y_4q/100. .+ pop_growth_4q) - 1.)
end

"""
```
logleveltopct_4q(y, data = fill(NaN, 4))
```

Transform from log levels to 4-quarter percent change.

### Inputs

- `y`: the data we wish to transform to 4-quarter percent change from log
  levels. `y` is either a vector of length `nperiods` or an `ndraws x `nperiods`
  matrix.

- `data`: if `y = [y_t, y_{t+1}, ..., y_{t+nperiods-1}]`, then
  `data = [y_{t-4}, y_{t-3}, y_{t-2}, y_{t-1}]`. This is necessary to compute
  4-quarter percent changes for the first three periods.
"""
function logleveltopct_4q(y::AbstractArray, data::AbstractVector = fill(NaN, 4))
    @assert length(data) == 4 "Length of data ($(length(data))) must be 4"

    # `y_t4` is an array of the same size as `y`, representing the t-4
    # period observations for each t
    y_t4 = if ndims(y) == 1
        nperiods = length(y)
        prepend_data(y[1:nperiods-4], data)
    else
        nperiods = size(y, 2)
        prepend_data(y[:, 1:nperiods-4], data)
    end
    y_4q = y - y_t4

    # Subtract log levels to get log growth rates, then exponentiate to get
    # growth rates
    100. * (exp.(y_4q./100.) .- 1.)
end

"""
```
logleveltopct_4q_percapita(y, pop_growth, data = fill(NaN, 4))
```

Transform from per-capita log levels to 4-quarter aggregate percent change.

### Note

This is usually applied to labor supply (hours worked), and probably shouldn't
be used for any other observables.

### Inputs

- `y`: the data we wish to transform to 4-quarter aggregate percent change from
  per-capita log levels. `y` is either a vector of length `nperiods` or an
  `ndraws x `nperiods` matrix.

- `pop_growth::Vector`: the length `nperiods` vector of log population growth
  rates.

- `data`: if `y = [y_t, y_{t+1}, ..., y_{t+nperiods-1}]`, then
  `data = [y_{t-4}, y_{t-3}, y_{t-2}, y_{t-1}]`. This is necessary to compute
  4-quarter percent changes for the first three periods.
"""
function logleveltopct_4q_percapita(y::AbstractArray, pop_growth::AbstractVector,
                                    data::AbstractVector = fill(NaN, 4))
    @assert length(data) == 4 "Length of data ($(length(data))) must be 4"

    # Four-quarter population growth
    pop_growth_4q = pop_growth[1:end-3] + pop_growth[2:end-2] + pop_growth[3:end-1] + pop_growth[4:end]

    # `y_t4` is an array of the same size as `y`, representing the t-4
    # period observations for each t
    if ndims(y) == 1
        nperiods = length(y)
        y_t4 = prepend_data(y[1:nperiods-4], data)
    else
        # Transpose `pop_growth` to a 1 x `nperiods` row vector so it can be
        # broadcasted to match the dimensions of `y`
        pop_growth_4q = pop_growth_4q'

        nperiods = size(y, 2)
        y_t4 = prepend_data(y[:, 1:nperiods-4], data)
    end

    @assert length(pop_growth_4q) == nperiods "Length of pop_growth ($(length(pop_growth))) must equal number of periods of y ($nperiods)"

    y_4q = y - y_t4

    # Subtract log levels to get log growth rates, then exponentiate to get growth rates
    100. * (exp.(y_4q./100. .+ pop_growth_4q) .- 1.)
end

"""
```
prepend_data(y, data)
```

Prepends data necessary for running 4q transformations.

### Inputs:

- `y`: `ndraws x t` array representing a timeseries for variable `y`
- `data`: vector representing a timeseries to prepend to `y`
"""
function prepend_data(y::AbstractArray, data::AbstractVector)
    if ndims(y) == 1
        y_extended = vcat(data, y)
    else
        ndraws = size(y, 1)
        datas  = repmat(data', ndraws, 1)
        y_extended = hcat(datas, y)
    end

    return y_extended
end

"""
```
get_scenario_transform(transform::Function)
```
Given a transformation used for usual forecasting, return the transformation
used for *scenarios*, which are forecasted in deviations from baseline.

The 1Q deviation from baseline should really be calculated by 1Q transforming
the forecasts (in levels) under the baseline (call this `y_b`) and alternative
scenario (`y_s`), then subtracting baseline from alternative scenario (since
most of our 1Q transformations are nonlinear). Let `y_d = y_s - y_b`. Then, for
example, the most correct `loggrowthtopct_annualized` transformation is:

```
y_b_1q = 100*(exp(y_b/100)^4 - 1)
y_s_1q = 100*(exp(y_s/100)^4 - 1)
y_d_1q = y_b_1q - y_s_1q
```

Instead, we approximate this by transforming the deviation directly:

```
y_d_1q ≈ 4*(y_b - y_s)
```
"""
function get_scenario_transform(transform::Function)
    if transform in [loggrowthtopct_annualized_percapita, loggrowthtopct_annualized,
                     quartertoannual]
        quartertoannual
    elseif transform in [logleveltopct_annualized_percapita, logleveltopct_annualized]
        logleveltopct_annualized_approx
    elseif transform == identity
        identity
    else
        error("Scenario equivalent not implemented for $transform")
    end
end

function get_scenario_transform4q(transform::Function)
    if transform in [loggrowthtopct_annualized_percapita, loggrowthtopct_annualized]
        loggrowthtopct_4q_approx
    elseif transform in [logleveltopct_annualized_percapita, logleveltopct_annualized]
        logleveltopct_4q_approx
    elseif transform == quartertoannual
        quartertoannual
    elseif transform == identity
        identity
    else
        error("Scenario 4q equivalent not implemented for $transform")
    end
end

"""
```
logleveltopct_annualized_approx(y, y0 = NaN)
```

Transform from log levels to *approximate* annualized quarter-over-quarter
percent change.

**This method should only be used to transform scenarios forecasts, which are in
  deviations from baseline.**

### Inputs

- `y`: the data we wish to transform to annualized quarter-over-quarter percent
  change from log levels. `y` is either a vector of length `nperiods` or an
  `ndraws x `nperiods` matrix.

- `y0`: the last data point in the history (of state or observable)
  corresponding to the `y` variable. This is required to compute a percent
  change for the first period.
"""
function logleveltopct_annualized_approx(y::AbstractArray, y0::Real = NaN)
    # `y_t1` is an array of the same size as `y`, representing the previous
    # period observations for each draw
    if ndims(y) == 1
        y_t1 = vcat([y0], y[1:end-1])
    else
        ndraws = size(y, 1)
        y0s  = fill(y0, ndraws, 1)
        y_t1 = hcat(y0s, y[:, 1:end-1])
    end

    # Subtract log levels to get log growth rates, then multiply by 4 to
    # approximate annualizing
    4(y - y_t1)
end

"""
```
loggrowthtopct_4q_approx(y, data = fill(NaN, 3))
```

Transform from log growth rates to *approximate* 4-quarter percent change.

**This method should only be used to transform scenarios forecasts, which are in
  deviations from baseline.**

### Inputs

- `y`: the data we wish to transform to aggregate 4-quarter percent change from
  log per-capita growth rates. `y` is either a vector of length `nperiods` or an
  `ndraws x `nperiods` matrix.

- `data`: if `y = [y_t, y_{t+1}, ..., y_{t+nperiods-1}]`, then
  `data = [y_{t-3}, y_{t-2}, y_{t-1}]`. This is necessary to compute
  4-quarter percent changes for the first three periods.
"""
function loggrowthtopct_4q_approx(y::AbstractArray, data::AbstractVector = fill(NaN, 3))
    @assert length(data) == 3 "Length of data ($(length(data))) must be 3"

    # Prepend previous three periods to `y`
    y = prepend_data(y, data)

    # `y` is either a vector of length `nperiods+3` or an
    # `ndraws` x `nperiods+3` matrix
    if ndims(y) == 1
        y[1:end-3] + y[2:end-2] + y[3:end-1] + y[4:end]
    else
        y[:,  1:end-3] + y[:, 2:end-2] + y[:, 3:end-1] + y[:, 4:end]
    end
end

"""
```
logleveltopct_4q_approx(y, data = fill(NaN, 4))
```

Transform from log levels to *approximate* 4-quarter percent change.

**This method should only be used to transform scenarios forecasts, which are in
  deviations from baseline.**

### Inputs

- `y`: the data we wish to transform to 4-quarter percent change from log
  levels. `y` is either a vector of length `nperiods` or an `ndraws x `nperiods`
  matrix.

- `data`: if `y = [y_t, y_{t+1}, ..., y_{t+nperiods-1}]`, then
  `data = [y_{t-4}, y_{t-3}, y_{t-2}, y_{t-1}]`. This is necessary to compute
  4-quarter percent changes for the first three periods.
"""
function logleveltopct_4q_approx(y::AbstractArray, data::AbstractVector = fill(NaN, 4))
    @assert length(data) == 4 "Length of data ($(length(data))) must be 4"

    # `y_t4` is an array of the same size as `y`, representing the t-4
    # period observations for each t
    y_t4 = if ndims(y) == 1
        nperiods = length(y)
        prepend_data(y[1:nperiods-4], data)
    else
        nperiods = size(y, 2)
        prepend_data(y[:, 1:nperiods-4], data)
    end

    y - y_t4
end


"""
    convert2features(array)

Convert the given MNIST tensor to a feature matrix (or feature
vector in the case of a single image). The purpose of this
function is to drop the spatial dimensions such that traditional
ML algorithms can process the dataset.

```julia
julia> MNIST.convert2features(MNIST.traintensor()) # full training data
784×60000 Array{N0f8,2}:
[...]

julia> MNIST.convert2features(MNIST.traintensor(1)) # first observation
784-element Array{N0f8,1}:
[...]
```
"""
convert2features(array::AbstractMatrix{<:Number}) = vec(array)

function convert2features(array::AbstractArray{<:Number,3})
    nrows, ncols, nimages = size(array)
    reshape(array, (nrows * ncols, nimages))
end

"""
    convert2image(array) -> Array{Gray}

Convert the given MNIST horizontal-major tensor (or feature matrix)
to a vertical-major `Colorant` array. The values are also color
corrected according to the website's description, which means that
the digits are black on a white background.

```julia
julia> MNIST.convert2image(MNIST.traintensor()) # full training dataset
28×28×60000 Array{Gray{N0f8},3}:
[...]

julia> MNIST.convert2image(MNIST.traintensor(1)) # first training image
28×28 Array{Gray{N0f8},2}:
[...]
```
"""
function convert2image(array::AbstractVector{<:Number})
    @assert length(array) % 784 == 0
    if length(array) == 784
        convert2image(reshape(array, 28, 28))
    else
        n = Int(length(array) / 784)
        convert2image(reshape(array, 28, 28, n))
    end
end

function convert2image(array::AbstractMatrix{T}) where {T<:Number}
    if size(array) == (28, 28)
        # simple check to see if values are normalized to [0,1]
        if any(x->x > 50, array)
            _colorview(Gray, T(1) .- transpose(array) ./ T(255))
        else
            _colorview(Gray, T(1) .- transpose(array))
        end
    else # feature matrix
        @assert size(array, 1) == 784
        n = size(array, 2)
        convert2image(reshape(array, 28, 28, n))
    end
end

function convert2image(array::AbstractArray{T,3}) where {T<:Number}
    h, w, n = size(array)
    @assert h == 28 && w == 28
    # simple check to see if values are normalized to [0,1]
    if any(x->x > 50, array)
        _colorview(Gray, permutedims(T(1) .- array ./ T(255), [2,1,3]))
    else
        _colorview(Gray, permutedims(T(1) .- array, [2,1,3]))
    end
end


"""
calcentropy \\

Calculates the requested entropies using the x and y variables given. \\
Arguments:  x = random variable, usually input \\
            y = random variable, usually output/classification \\
            xdims = [number of variables in x, the maximum number of possible values across the dimensions in x] \\
            ydims = [number of variables in y, the maximum number of possible values across the dimensions in y] \\
            calc = The entropies to calculate ex: ["HX", "HXY"] \\
            method = sampling bias correction method ex: "pt" \\
            sampling = method for sampling ex: "naive" \\
Returns:    dictionary, entropies listed in calc => the entropy values \\
"""

function calcentropy(x::Array{Int}, y::Array{Int}, xdims::Array{Int, 1}, ydims::Array{Int, 1}, calc::Array{String, 1}, method::String, sampling::String; base::Int=2)
  xn, xmax = xdims
  yn, ymax = ydims
  x = atleast2d(x)
  y = atleast2d(y)
  checkinputs(x, y, xmax, ymax, xn, yn)
  probs = requireprobs(x, xn, xmax, y, yn, ymax, calc, sampling)
  return requireents(probs, calc, base=base)
end

function calcentropy(x::Array{Int}, xdims::Array{Int, 1}, calc::Array{String, 1}, method::String, sampling::String; base::Int=2)
  xn, xmax = xdims
  x = atleast2d(x)
  checkinputs(x, xmax, xn)
  x_dist_dim = xmax ^ xn
  probs = Dict{String, Array{Float64}}("PX" => prob(decimalise(x, xn, xmax), x_dist_dim, sampling))
  return requireents(probs, calc, base=base)
end


export AR1, AR1TV, ARH

function AR1(θ)
    μ = fill(θ[1],1)
    Σ0 = fill(θ[2].^2,1,1)
    F = fill(θ[3],1,1)
    G = zeros(1,1)
    Q = fill(θ[4].^2,1,1)
    H = [1.0]'
    R = fill(θ[5].^2,1,1)
    Γ = zeros(1,1)
    Θ = fill(1.0,1,1)
    LinearStateSpaceModel(F,Q,G,H,R,Γ,Θ,μ,Σ0)
end

function AR1TV(θ,N)
    μ = fill(θ[1],1)
    Σ0 = fill(θ[2].^2,1,1)
    F = fill(fill(θ[3],1,1),N)
    G = fill(zeros(1,1),N)
    Q = fill(fill(θ[4].^2,1,1),N)
    H = fill([1.0]',N)
    R = fill(fill(θ[5].^2,1,1),N)
    Γ = fill(zeros(1,1),N)
    Θ = fill(fill(1.0,1,1),N)
    TimeVariantLinearStateSpaceModel(F,Q,G,H,R,Γ,Θ,μ,Σ0)
end

"""
An autoregressive model with specified regression coefficients

    - θ = [δ;atanh(φ);σ]
    - β = [α1;α2;σε]

x_{t} = δ + φ*x_{t-1} + σ η_t
y_{t} = α1+ α2*x_t + σε ε_t

"""
function ARH(θ,β)
    δ = θ[1]
    φ = tanh(θ[2])
    σ = θ[3]
    σε = θ[4]
    α1= β[1]
    α2= β[2]

    μ = fill(δ/(1-φ),1)
    Σ0= fill(σ^2/(1-φ^2),1,1)
    F = fill(φ,1,1)
    Q = fill(σ^2,1,1)
    G = fill(δ,1,1)
    H = fill(α2,1,1)
    R = fill(σε^2,1,1)
    Γ = fill(α1,1,1)
    Θ = fill(1.0,1,1)

    LinearStateSpaceModel(F, Q, G, H, R, Γ, Θ, μ, Σ0)
end

function ARH(β)
    μ(θ) = fill(θ[1]/(1-tanh(θ[2])),1)
    Σ0(θ)= fill(θ[3]^2/(1-tanh(θ[2])^2),1,1)
    F(θ) = fill(tanh(θ[2]),1,1)
    Q(θ) = fill(θ[3]^2,1,1)
    G(θ) = fill(θ[1],1,1)
    H(θ) = fill(β[2],1,1)
    R(θ) = fill(θ[4]^2,1,1)
    Γ(θ) = fill(β[1],1,1)
    Θ(θ) = fill(1.0,1,1)

    LinearStateSpaceModel(F, Q, G, H, R, Γ, Θ, μ, Σ0)
end


println("Starting precompilation")
using OrdinaryDiffEq
using Flux, DiffEqFlux, GalacticOptim
using Test
using Distributions
using Distances
using Zygote
using DistributionsAD
using LinearAlgebra
using Random
println("Starting tests")

Random.seed!(1999)

##callback to be used by all tests
function cb(p,l)
    @show l
    false
end

###
#test for default base distribution and deterministic trace CNF
###

nn = Chain(Dense(1, 1, tanh))
data_train = Float32.(rand(Beta(7,7), 1, 100))
tspan = (0.0f0,1.0f0)
cnf_test = FFJORD(nn,tspan,Tsit5())

function loss_adjoint(θ)
    logpx = cnf_test(data_train,θ;monte_carlo=false)[1]
    loss = -mean(logpx)
end

optfunc = GalacticOptim.OptimizationFunction((x, p) -> loss_adjoint(x), GalacticOptim.AutoZygote())
optprob = GalacticOptim.OptimizationProblem(optfunc, cnf_test.p)
res = GalacticOptim.solve(optprob, ADAM(0.1), cb=cb, maxiters = 100)

θopt = res.minimizer
data_validate = Float32.(rand(Beta(7,7), 1, 100))
actual_pdf = [pdf(Beta(7,7),r) for r in data_validate]
#use direct trace calculation for predictions
learned_pdf = exp.(cnf_test(data_validate,θopt;monte_carlo=false)[1])

@test totalvariation(learned_pdf, actual_pdf)/100 < 0.35

###
#test for alternative base distribution and deterministic trace CNF
###

nn = Chain(Dense(1, 3, tanh), Dense(3, 1, tanh))
data_train = Float32.(rand(Normal(6.0,0.7), 1, 100))
tspan = (0.0f0,1.0f0)
cnf_test = FFJORD(nn,tspan,Tsit5();basedist=MvNormal([0.0f0],[2.0f0]))

function loss_adjoint(θ)
    logpx = cnf_test(data_train,θ;monte_carlo=false)[1]
    loss = -mean(logpx)
end

optfunc = GalacticOptim.OptimizationFunction((x, p) -> loss_adjoint(x), GalacticOptim.AutoZygote())
optprob = GalacticOptim.OptimizationProblem(optfunc, 0.01f0 .* cnf_test.p)
res = GalacticOptim.solve(optprob, ADAM(0.1), cb=cb, maxiters = 100)

θopt = res.minimizer
data_validate = Float32.(rand(Normal(6.0,0.7), 1, 100))
actual_pdf = [pdf(Normal(6.0,0.7),r) for r in data_validate]
learned_pdf = exp.(cnf_test(data_validate,θopt;monte_carlo=false)[1])

@test totalvariation(learned_pdf, actual_pdf)/100 < 0.25

###
#test for multivariate distribution and deterministic trace CNF
###

nn = Chain(Dense(2, 2, tanh))
μ = ones(2)
Σ = 7*I + zeros(2,2)
mv_normal = MvNormal(μ, Σ)
data_train = Float32.(rand(mv_normal, 100))
tspan = (0.0,1.0f0)
cnf_test = FFJORD(nn,tspan,Tsit5())

function loss_adjoint(θ)
    logpx = cnf_test(data_train,θ;monte_carlo=false)[1]
    loss = -mean(logpx)
end

optfunc = GalacticOptim.OptimizationFunction((x, p) -> loss_adjoint(x), GalacticOptim.AutoZygote())
optprob = GalacticOptim.OptimizationProblem(optfunc, 0.01f0 .* cnf_test.p)
res = GalacticOptim.solve(optprob, ADAM(0.1), cb=cb, maxiters = 300)

θopt = res.minimizer
data_validate = Float32.(rand(mv_normal, 100))
actual_pdf = [pdf(mv_normal,data_validate[:, i]) for i in 1:size(data_validate, 2)]
learned_pdf = exp.(cnf_test(data_validate,θopt;monte_carlo=false)[1])

@test totalvariation(learned_pdf, actual_pdf)/100 < 0.25

###
#test for default multivariate distribution and FFJORD with regularizers
###

nn = Chain(Dense(1, 1, tanh))
data_train = Float32.(rand(Beta(7,7), 1, 100))
tspan = (0.0f0,1.0f0)
ffjord_test = FFJORD(nn,tspan,Tsit5())

function loss_adjoint(θ)
    logpx, λ₁, λ₂ = ffjord_test(data_train,θ,true)
    return mean(@. -logpx + 0.1 * λ₁ + λ₂)
end

optfunc = GalacticOptim.OptimizationFunction((x, p) -> loss_adjoint(x), GalacticOptim.AutoZygote())
optprob = GalacticOptim.OptimizationProblem(optfunc, 0.01f0 .* ffjord_test.p)
res = GalacticOptim.solve(optprob,  ADAM(0.1), cb=cb, maxiters = 300)

θopt = res.minimizer
data_validate = Float32.(rand(Beta(7,7), 1, 100))
actual_pdf = [pdf(Beta(7,7),r) for r in data_validate]
learned_pdf = exp.(ffjord_test(data_validate,θopt;monte_carlo=false)[1])

@test totalvariation(learned_pdf, actual_pdf)/100 < 0.40


module BinningAnalysis

import Statistics: mean, var
using Reexport, Lazy, RecursiveArrayTools

# Generic functions
include("generic.jl")
export std_error, tau

# Logarithmic binning
include("log/binning.jl")
include("log/statistics.jl")
export LogBinner, capacity
export mean, var, varN, tau, std_error
export all_vars, all_varNs, all_taus, all_std_errors
export convergence, has_converged

# "Full" binning
include("full/binning.jl")
include("full/statistics.jl")
export FullBinner, std_error, all_binning_errors

# Jackknife resampling
include("Jackknife.jl")
@reexport using .Jackknife

include("ErrorPropagation/binning.jl")
include("ErrorPropagation/statistics.jl")
export ErrorPropagator
export mean, var, varN, tau, std_error
export means, vars, varNs, taus, std_errors
export all_means, all_vars, all_varNs, all_taus, all_std_errors
export covmat

include("incrementing/IncrementBinner.jl")
export IncrementBinner
export indices

end # module


# Example for 2D sparsity-promoting LSRTM with source subsampling
# Author: Philipp Witte, pwitte@eos.ubc.ca
# Date: January 2017
#

using PyCall
using PyPlot
using opesciSLIM.TimeModeling
using opesciSLIM.LeastSquaresMigration

## Set up model structure
n = (120,100)	# (x,y,z) or (x,z)
d = (10.,10.)
o = (0.,0.)

# Velocity [km/s]
v = ones(Float32,n) + 0.4f0
v0 = ones(Float32,n) + 0.4f0
v[:,40:70] = 3.0f0
v[:,71:end] = 4.0f0
rho = ones(Float32,n)

# Slowness squared [s^2/km^2]
m = (1./v).^2
m0 = (1./v0).^2
dm = vec(m - m0)

# Setup info and model structure
nsrc = 4
model = Model(n,d,o,m)	# to include density call Model(n,d,o,m,rho)
model0 = Model(n,d,o,m0)

## Set up receiver geometry
nxrec = 120
xrec = linspace(50.,1150.,nxrec)
yrec = 0.
zrec = linspace(50.,50.,nxrec)

# receiver sampling and recording time
timeR = 1000.	# receiver recording time [ms]
dtR = 4.	# receiver sampling interval

# Set up receiver structure
recGeometry = Geometry(xrec,yrec,zrec;dt=dtR,t=timeR,nsrc=nsrc)

## Set up source geometry (cell array with source locations for each shot)
xsrc = convertToCell(linspace(100,1000,nsrc))
ysrc = convertToCell(linspace(0,0,nsrc))
zsrc = convertToCell(linspace(20,20,nsrc))

# source sampling and number of time steps
timeS = 1000.	# source length in [ms]
dtS = 2.	# source sampling interval

# Set up source structure
srcGeometry = Geometry(xsrc,ysrc,zsrc;dt=dtS,t=timeS)

# setup wavelet
f0 = 0.01
wavelet = source(timeS,dtS,f0)

# Info structure for linear operators
ntComp = get_computational_nt(srcGeometry,recGeometry,model)	# no. of computational time steps
info = Info(prod(n),nsrc,ntComp)

###################################################################################################

# Setup operators
F = joModeling(info,model0,srcGeometry,recGeometry)
q = joData(srcGeometry,wavelet)
J = joJacobian(F,q)

# Observed data
dD = J*dm

# Least squares migration
image = linbreg_lsrtm(J, dD; iterations=5, shots_per_iteration=2)


using Gen

@gen (static) function foo()
    @param mu::Float64
    y = @trace(normal(mu, 1), :y)
    return y
end

load_generated_functions()

init_param!(foo, :mu, -1)

trace, = generate(foo, (), choicemap((:y, 3)))
step_size = 0.01
for iter=1:1000
    accumulate_param_gradients!(trace, 0.)
    grad_val = get_param_grad(foo, :mu)
    set_param!(foo, :mu, get_param(foo, :mu) + step_size * grad_val)
    zero_param_grad!(foo, :mu)
end

@assert abs(get_param(foo, :mu) - 3) < 1e-2


## Computing marginal posterior probability density
## More on defining prior_vector, iterations, likelihood_vector to be included
likelihood = []
X = []
w = []
Z = []
function marginal_pd(N, prior_vector, iterations, likelihood_vector_current)
    for i in 1:1:iterations
        likelihood[i] = minimum(likelihood_vector_current) ## re-check this step
        X[i] = exp(-i/N)
        w[i] = X[i - 1] - X[i]
        Z = Z + likelihood[i]*w[i]
        ## save point with least likelihood as sample point with weight wi.
        ## update point with least likelihood with some MCMC steps according to prior,
        ## accepting only steps that keep the likelihood about Li
        end
    end
    return Z
end


"""
```julia
BayesianSVM()
```

The [Bayesian SVM](https://arxiv.org/abs/1707.05532) is a Bayesian interpretation of the classical SVM.
```math
p(y|f) ∝ exp(2 max(1-yf,0))
````

---
For the analytic version of the likelihood, it is augmented via:
```math
p(y|f,ω) = 1/(sqrt(2πω) exp(-0.5((1+ω-yf)^2/ω))
```
where ``ω ∼ 𝟙[0,∞)`` has an improper prior (his posterior is however has a valid distribution, a Generalized Inverse Gaussian). For reference [see this paper](http://ecmlpkdd2017.ijs.si/papers/paperID502.pdf)
"""
struct BayesianSVM{T<:Real} <: ClassificationLikelihood{T}
    ω::AbstractVector{T}
    θ::AbstractVector{T}
    function BayesianSVM{T}() where {T<:Real}
        new{T}()
    end
    function BayesianSVM{T}(ω::AbstractVector{<:Real},θ::AbstractVector{<:Real}) where {T<:Real}
        new{T}(ω,θ)
    end
end

function BayesianSVM()
    BayesianSVM{Float64}()
end

function init_likelihood(likelihood::BayesianSVM{T},inference::Inference{T},nLatent::Integer,nSamplesUsed::Integer,nFeatures::Integer) where T
    BayesianSVM{T}(abs.(rand(T,nSamplesUsed)),zeros(T,nSamplesUsed))
end
function pdf(l::BayesianSVM,y::Real,f::Real)
    svmlikelihood(y*f)
end

function Base.show(io::IO,model::BayesianSVM{T}) where T
    print(io,"Bayesian SVM")
end

"""Return likelihood equivalent to SVM hinge loss"""
function svmlikelihood(f::Real)
    pos = svmpseudolikelihood(f)
    return pos./(pos.+svmpseudolikelihood(-f))
end

"""Return the pseudo likelihood of the SVM hinge loss"""
function svmpseudolikelihood(f::Real)
    return exp(-2.0*max.(1.0-f,0))
end

function compute_proba(l::BayesianSVM{T},μ::Vector{T},σ²::Vector{T}) where {T<:Real}
    N = length(μ)
    pred = zeros(T,N)
    sig_pred = zeros(T,N)
    for i in 1:N
        x = pred_nodes.*sqrt(max(σ²[i],zero(T))).+μ[i]
        pred[i] =  dot(pred_weights,svmlikelihood.(x))
        sig_pred[i] = max(dot(pred_weights,svmlikelihood.(x).^2)-pred[i]^2,zero(T))
    end
    return pred, sig_pred
end

###############################################################################

function local_updates!(l::BayesianSVM{T},y::AbstractVector,μ::AbstractVector,diag_cov::AbstractVector) where {T}
    l.ω .= abs2.(one(T) .- y.*μ) + diag_cov
    l.θ .= one(T)./sqrt.(l.ω)
end

@inline ∇E_μ(l::BayesianSVM{T},::AOptimizer,y::AbstractVector) where {T} = (y.*(l.θ.+one(T)),)
@inline ∇E_Σ(l::BayesianSVM{T},::AOptimizer,y::AbstractVector) where {T} = (0.5.*l.θ,)

function expec_log_likelihood(l::BayesianSVM{T},i::AnalyticVI,y::AbstractVector,μ::AbstractVector,diag_cov::AbstractVector) where {T}
    tot = -(0.5*length(y)*logtwo)
    tot += dot(μ,y)
    tot += -0.5*dot(l.θ,diag_cov)+dot(l.θ,abs2.(one(T).-y.*μ))
    return tot
end

AugmentedKL(l::BayesianSVM,::AbstractVector) = GIGEntropy(l)

function GIGEntropy(l::BayesianSVM)
    return 0.5*sum(log.(l.ω))+sum(log.(2.0*besselk.(0.5,sqrt.(l.ω))))-0.5*sum(sqrt.(l.ω))
end


@doc raw"""
    AlternatingGradientProblem <: Problem

An alternating gradient problem consists of
* a `ProductManifold M` ``=\mathcal M = \mathcal M_1 × ⋯ × M_n``
* a cost function ``F(x)``
* a gradient ``\operatorname{grad}F`` that is either
  * given as one function ``\operatorname{grad}F`` returning a tangent vector `X` on `M` or
  * an array of gradient functions ``\operatorname{grad}F_i``, `ì=1,…,n` s each returning a component of the gradient
  which might be allocating or mutating variants, but not a mix of both.

!!! note

    This Problem requires the `ProductManifold` from `Manifolds.jl`, so `Manifolds.jl` to be loaded.

!!! note

    The input of each of the (component) gradients is still the whole vector `x`,
    just that all other then the `i`th input component are assumed to be fixed and just
    the `i`th components gradient is computed / returned.

# Constructors
    AlternatingGradientProblem(M::ProductManifold, F, gradF::Function;
        evaluation=AllocatingEvaluation()
    )
    AlternatingGradientProblem(M::ProductManifold, F, gradF::AbstractVector{<:Function};
        evaluation=AllocatingEvaluation()
    )

Create a alternating gradient problem with an optional `cost` and the gradient either as one
function (returning an array) or a vector of functions.
"""
struct AlternatingGradientProblem{T,MT<:ProductManifold,TCost,TGradient} <:
       AbstractGradientProblem{T}
    M::MT
    cost::TCost
    gradient!!::TGradient
end
function AlternatingGradientProblem(
    M::TM, F::TCost, gradF!!::G; evaluation::AbstractEvaluationType=AllocatingEvaluation()
) where {TM<:ProductManifold,G,TCost}
    return AlternatingGradientProblem{typeof(evaluation),TM,TCost,G}(M, F, gradF!!)
end
function AlternatingGradientProblem(
    M::TM,
    F::TCost,
    gradF!!::AbstractVector{<:TG};
    evaluation::AbstractEvaluationType=AllocatingEvaluation(),
) where {TM<:ProductManifold,TCost,TG}
    return AlternatingGradientProblem{typeof(evaluation),TM,TCost,typeof(gradF!!)}(
        M, F, gradF!!
    )
end

@doc raw"""
    get_gradient(P::AlternatingGradientProblem, x)
    get_gradient!(P::AlternatingGradientProblem, Y, x)

Evaluate all summands gradients at a point `x` on the `ProductManifold M` (in place of `Y`)
"""
function get_gradient(
    p::AlternatingGradientProblem{AllocatingEvaluation,<:AbstractManifold,TC,<:Function}, x
) where {TC}
    return p.gradient!!(p.M, x)
end
function get_gradient(
    p::AlternatingGradientProblem{
        AllocatingEvaluation,<:AbstractManifold,TC,<:AbstractVector
    },
    x,
) where {TC}
    Y = ProductRepr([gi(p.M, x) for gi in p.gradient!!]...)
    return Y
end
function get_gradient!(
    p::AlternatingGradientProblem{AllocatingEvaluation,<:AbstractManifold,TC,<:Function},
    X,
    x,
) where {TC}
    copyto!(p.M, X, x, get_gradient(p, x))
    return X
end
function get_gradient!(
    p::AlternatingGradientProblem{
        AllocatingEvaluation,<:AbstractManifold,TC,<:AbstractVector
    },
    X,
    x,
) where {TC}
    copyto!(p.M, X, x, get_gradient(p, x))
    return X
end
function get_gradient(
    p::AlternatingGradientProblem{MutatingEvaluation,<:AbstractManifold,TC,<:Function}, x
) where {TC}
    Y = zero_vector(p.M, x)
    return p.gradient!!(p.M, Y, x)
end
function get_gradient(
    p::AlternatingGradientProblem{
        MutatingEvaluation,<:AbstractManifold,TC,<:AbstractVector
    },
    x,
) where {TC}
    Y = zero_vector(p.M, x)
    get_gradient!(p, Y, x)
    return Y
end
function get_gradient!(
    p::AlternatingGradientProblem{MutatingEvaluation,<:AbstractManifold,TC,<:Function}, X, x
) where {TC}
    return p.gradient!!(p.M, X, x)
end
function get_gradient!(
    p::AlternatingGradientProblem{
        MutatingEvaluation,<:AbstractManifold,TC,<:AbstractVector
    },
    X,
    x,
) where {TC}
    for (gi, Xi) in zip(p.gradient!!, submanifold_components(p.M, X))
        gi(p.M, Xi, x)
    end
    return X
end

@doc raw"""
    get_gradient(p::AlternatingGradientProblem, k, x)
    get_gradient!(p::AlternatingGradientProblem, Y, k, x)

Evaluate one of the component gradients ``\operatorname{grad}f_k``, ``k∈\{1,…,n\}``, at `x` (in place of `Y`).
"""
function get_gradient(
    p::AlternatingGradientProblem{AllocatingEvaluation,<:AbstractManifold,TC,<:Function},
    k,
    x,
) where {TC}
    return get_gradient(p, x)[p.M, k]
end
function get_gradient(
    p::AlternatingGradientProblem{
        AllocatingEvaluation,<:AbstractManifold,TC,<:AbstractVector
    },
    k,
    x,
) where {TC}
    return p.gradient!![k](p.M, x)
end
function get_gradient!(
    p::AlternatingGradientProblem{AllocatingEvaluation,<:AbstractManifold,TC,<:Function},
    X,
    k,
    x,
) where {TC}
    copyto!(p.M[k], X, p.gradient!!(p.M, x)[p.M, k])
    return X
end
function get_gradient!(
    p::AlternatingGradientProblem{
        AllocatingEvaluation,<:AbstractManifold,TC,<:AbstractVector
    },
    X,
    k,
    x,
) where {TC}
    copyto!(p.M[k], X, p.gradient!![k](p.M, x))
    return X
end
function get_gradient(
    p::AlternatingGradientProblem{MutatingEvaluation,<:AbstractManifold,TC}, k, x
) where {TC}
    X = zero_vector(p.M[k], x[p.M, k])
    get_gradient!(p, X, k, x)
    return X
end
function get_gradient!(
    p::AlternatingGradientProblem{MutatingEvaluation,<:AbstractManifold,TC,<:Function},
    X,
    k,
    x,
) where {TC}
    # this takes a lot more allocations than other methods, but the gradient can only be evaluated in full
    Xf = zero_vector(p.M, x)
    get_gradient!(p, Xf, x)
    copyto!(p.M[k], X, x[p.M, k], Xf[p.M, k])
    return X
end
function get_gradient!(
    p::AlternatingGradientProblem{
        MutatingEvaluation,<:AbstractManifold,TC,<:AbstractVector
    },
    X,
    k,
    x,
) where {TC}
    return p.gradient!![k](p.M, X, x)
end

"""
    AlternatingGradientDescentOptions <: AbstractGradientDescentOptions

Store the fields for an alternating gradient descent algorithm,
see also [`AlternatingGradientProblem`](@ref) and [`alternating_gradient_descent`](@ref).

# Fields
* `direction` (`AlternatingGradient(zero_vector(M, x))` a [`DirectionUpdateRule`](@ref)
* `evaluation_order` – (`:Linear`) – whether
* `inner_iterations`– (`5`) how many gradient steps to take in a component before alternating to the next
  to use a randomly permuted sequence (`:FixedRandom`), a per
  cycle newly permuted sequence (`:Random`) or the default `:Linear` evaluation order.
* `order` the current permutation
* `retraction_method` – (`default_retraction_method(M)`) a `retraction(M,x,ξ)` to use.
* `stepsize` ([`ConstantStepsize`](@ref)`(M)`) a [`Stepsize`](@ref)
* `stopping_criterion` ([`StopAfterIteration`](@ref)`(1000)`)– a [`StoppingCriterion`](@ref)
* `x` the current iterate
* `k`, ì` internal counters for the outer and inner iterations, respectively.

# Constructors

    AlternatingGradientDescentOptions(M, x; kwargs...)

Generate the options for point `x` and and where the keyword arguments
`inner_iterations`, `order_type`, `order`, `retraction_method`, `stopping_criterion`, and `stepsize``
are keyword arguments
"""
mutable struct AlternatingGradientDescentOptions{
    TX,
    TV,
    D<:DirectionUpdateRule,
    TStop<:StoppingCriterion,
    TStep<:Stepsize,
    RM<:AbstractRetractionMethod,
} <: AbstractGradientOptions
    x::TX
    gradient::TV
    direction::D
    stop::TStop
    stepsize::TStep
    order_type::Symbol
    order::Vector{<:Int}
    retraction_method::RM
    k::Int # current iterate
    i::Int # inner iterate
    inner_iterations::Int
end
@deprecate AlternatingGradientDescentOptions(
    x,
    X,
    direction::DirectionUpdateRule;
    retraction_method::AbstractRetractionMethod=ExponentialRetraction(),
    kwargs...,
) AlternatingGradientDescentOptions(
    DefaultManifold(2), x, X, direction; retraction_method=retraction_method, kwargs...
)
function AlternatingGradientDescentOptions(
    M::AbstractManifold,
    x;
    inner_iterations::Int=5,
    order_type::Symbol=:Linear,
    order::Vector{<:Int}=Int[],
    retraction_method::AbstractRetractionMethod=default_retraction_method(M),
    stopping_criterion::StoppingCriterion=StopAfterIteration(1000),
    stepsize::Stepsize=ConstantStepsize(M),
)
    X = zero_vector(M, x)
    return AlternatingGradientDescentOptions{
        typeof(x),
        typeof(X),
        AlternatingGradient,
        typeof(stopping_criterion),
        typeof(stepsize),
        typeof(retraction_method),
    }(
        x,
        X,
        AlternatingGradient(zero_vector(M, x)),
        stopping_criterion,
        stepsize,
        order_type,
        order,
        retraction_method,
        0,
        0,
        inner_iterations,
    )
end
"""
    AlternatingGradient <: DirectionUpdateRule

The default gradient processor, which just evaluates the (alternating) gradient on one of
the components
"""
struct AlternatingGradient{T} <: AbstractStochasticGradientProcessor
    dir::T
end

function (s::AlternatingGradient)(
    p::AlternatingGradientProblem, o::AlternatingGradientDescentOptions, iter
)
    # at begin of inner iterations reset internal vector to zero
    (o.i == 1) && zero_vector!(p.M, s.dir, o.x)
    # update order(k)th component inplace
    get_gradient!(p, s.dir[p.M, o.order[o.k]], o.order[o.k], o.x)
    return o.stepsize(p, o, iter), s.dir # return urrent full gradient
end

# update Armijo to work on the kth gradient only.
function (a::ArmijoLinesearch)(
    p::AlternatingGradientProblem, o::AlternatingGradientDescentOptions, ::Int
)
    X = zero_vector(p.M, o.x)
    X[p.M, o.order[o.k]] .= get_gradient(p, o.order[o.k], o.x)
    a.last_stepsize = linesearch_backtrack(
        p.M,
        x -> p.cost(p.M, x),
        o.x,
        X,
        a.last_stepsize,
        a.sufficient_decrease,
        a.contraction_factor,
        a.retraction_method,
    )
    return a.last_stepsize
end


"""
    $(FUNCTIONNAME)(model::MieModel{FDT}) where FDT<:PCW

Reference: Suniti Sanghavi 2014, https://doi.org/10.1016/j.jqsrt.2013.12.015

Compute the aerosol optical properties using the Domke-PCW method
Input: MieModel, holding all computation and aerosol properties
Output: AerosolOptics, holding all Greek coefficients and Cross-Sectional information
"""
function compute_aerosol_optical_properties(model::MieModel{FDT}, FT2::Type=Float64) where FDT<:PCW

    # Unpack the model
    @unpack computation_type, aerosol, λ, polarization_type, truncation_type, wigner_A, wigner_B = model

    # Extract variables from aerosol struct:
    @unpack size_distribution, nquad_radius, nᵣ, nᵢ,r_max =  aerosol

    # Get the refractive index's real part type
    FT = eltype(nᵣ);

    # Compute radii and weights
    r, wᵣ = gauleg(nquad_radius, 0.0, r_max ; norm=true)
    wₓ = compute_wₓ(size_distribution, wᵣ, r, r_max)

    # Find overall N_max from the maximum radius
    N_max = Scattering.get_n_max(2 * π * r_max/ λ)

    # Compute an, bn values
    an, bn = compute_anbn(aerosol::UnivariateAerosol, λ, r)

    # Compute average cross-sectional scattering and extinction
    k = 2π / λ
    avg_C_scatt, avg_C_ext = compute_avg_C_scatt_ext(k, an, bn, wₓ)

    # l range to go over
    ls = 1:(2 * N_max - 1)

    # Where to store the Greek Coefficients
    greek_coefs = zeros(FT, 6, size(ls, 1))

    # Pre-compute anbn averages
    # That is, pre-compute (an✶)am, (an✶)bm, (bn✶)am, (bn✶)bm
    # So that you can quickly compute (an✶ + bn✶) × (am + bm)
    ComplexFT = Complex{FT}
    mat_anam = LowerTriangular(zeros(ComplexFT, N_max, N_max));
    mat_anbm = LowerTriangular(zeros(ComplexFT, N_max, N_max));
    mat_bnam = LowerTriangular(zeros(ComplexFT, N_max, N_max));
    mat_bnbm = LowerTriangular(zeros(ComplexFT, N_max, N_max));
    ab_pairs = (mat_anam, mat_anbm, mat_bnam, mat_bnbm)

    # Get the N_max value corresponding to each radius
    N_max_ = Scattering.get_n_max.(2π * r/ λ)

    # Then pre-compute anbn averages
    Scattering.compute_avg_anbns!(an, bn, ab_pairs, wₓ, N_max, N_max_)

    # Pre-compute |an ± bn|² * w
    an_m_bn = transpose(abs2.(an-bn)) * wₓ
    an_p_bn = transpose(abs2.(an+bn)) * wₓ

    # For each l
    @showprogress 1 "Computing S functions ..." for l in ls

        # Compute all necessary Sl_νν values
        # Eq 22 in Sanghavi 2014

        Sl_00  = compute_Sl(l, 0, 0,  true,  k, N_max, ab_pairs, an_m_bn, an_p_bn, wigner_A, wigner_B)
        Sl_0m0 = compute_Sl(l, 0, 0,  false, k, N_max, ab_pairs, an_m_bn, an_p_bn, wigner_A, wigner_B)
        Sl_22  = compute_Sl(l, 2, 2,  true,  k, N_max, ab_pairs, an_m_bn, an_p_bn, wigner_A, wigner_B)
        Sl_2m2 = compute_Sl(l, 2, -2, false, k, N_max, ab_pairs, an_m_bn, an_p_bn, wigner_A, wigner_B)
        Sl_02  = compute_Sl(l, 0, 2,  true,  k, N_max, ab_pairs, an_m_bn, an_p_bn, wigner_A, wigner_B)

        # Use Sl_νν values to compute Greek coefficients
        # Eq 24 in Sanghavi 2014

        @inbounds greek_coefs[1,l] = (1/avg_C_scatt) * (Sl_00 + Sl_0m0)
        @inbounds greek_coefs[2,l] = (1/avg_C_scatt) * (Sl_00 - Sl_0m0)
        @inbounds greek_coefs[3,l] = (1/avg_C_scatt) * (Sl_22 + Sl_2m2)
        @inbounds greek_coefs[4,l] = (1/avg_C_scatt) * (Sl_22 - Sl_2m2)
        @inbounds greek_coefs[5,l] = (1/avg_C_scatt) * real(Sl_02)
        @inbounds greek_coefs[6,l] = (1/avg_C_scatt) * imag(Sl_02)
    end

    # Create GreekCoefs object with α, β, γ, δ, ϵ, ζ (in that order)
    greek_coefs = GreekCoefs(convert.(FT2, greek_coefs[3,:]),
                             convert.(FT2, greek_coefs[1,:]),
                             convert.(FT2, greek_coefs[5,:]),
                             convert.(FT2, greek_coefs[2,:]),
                             convert.(FT2, greek_coefs[6,:]),
                             convert.(FT2, greek_coefs[4,:]))

    # Return the packaged AerosolOptics object
    return AerosolOptics(greek_coefs=greek_coefs, ω̃=FT2(avg_C_scatt/avg_C_ext), k=FT2(avg_C_ext))
end

"""
    $(FUNCTIONNAME)(l, ν₁, ν₂, ν₂_positive_flag, k, N_max, ab_pairs, an_m_bn, an_p_bn, wigner_A, wigner_B)

Reference: Suniti Sanghavi 2014, https://doi.org/10.1016/j.jqsrt.2013.12.015

Compute Sl_νν, given by Eq 22
Input: l, ν₁, ν₂, ν₂_positive_flag (to differentiate b/w 0 and -0),
k, N_max, ab_pairs (precomputed), an_m_bn, an_p_bn, wigner_A, wigner_B,
Output: Complex{Float64}
"""
function compute_Sl(l::Integer, ν₁::Integer, ν₂::Integer, ν₂_positive_flag::Bool,
                    k, N_max, ab_pairs, an_m_bn, an_p_bn, wigner_A, wigner_B)

    @assert ((ν₁, ν₂) == (2, -2) || (ν₁, ν₂) == (2, 2) || (ν₁, ν₂) == (0, 0) || (ν₁, ν₂) == (0, 2)) "Not a valid input to calculate Sl"
    @assert ((ν₂ == 2 && ν₂_positive_flag == true) || (ν₂ == -2 && ν₂_positive_flag == false) || ν₂ == 0) "Sign flag doesn't match ν₂"

    # Extract conjugate tranpose pairs
    mat_anam, mat_anbm, mat_bnam, mat_bnbm = ab_pairs

    # Calculate the coefficient beforehand
    ll = l-1
    coef = (2ll+1)*π / k^2

    # Accumulate the terms of the Sl expression in these variables
    first_term = 0.0    # First summation over n
    second_term = 0.0   # Second summation over n

    # Use wigner_A or wigner_B, depending on value of ν
    wig_lnm = ν₁ == 0 ? wigner_A : wigner_B

    # Summation over n from 1 to N_max
    @inbounds for n = 1:N_max

        # Obtain appropriate m range to loop over
        m_star, m_max = (max(ll-n,n+1), min(ll+n,N_max))

        # Compute and accumulate the first/second terms if (ν₁, ν₂) == (0, 2)
        if (ν₁, ν₂) == (0, 2)

            # Loop over m
            @inbounds for m = m_star:m_max

                # First term
                # [(-1)^(l+m+n) * (an⋆ + bn⋆)(am - bm) + (am⋆ + bm⋆)(an - bn)] × A_lnm × B_lnm
                avg = (exp_m1(ll + n + m)) * (mat_anam[m,n] + mat_bnam[m,n] - mat_anbm[m,n] - mat_bnbm[m,n]) + (mat_anam[m,n]' - mat_bnam[m,n]' + mat_anbm[m,n]' - mat_bnbm[m,n]')
                first_term += avg * 2 * (2n+1) * (2m+1) * wigner_A[m, n, l] * wigner_B[m, n, l]
            end

            # Second term
            # (an⋆ + bn⋆)(an - bn) × A_lnm × B_lnm
            anan,anbn,bnan,bnbn = (mat_anam[n,n], mat_anbm[n,n],
                                   mat_bnam[n,n], mat_bnbm[n,n])
            avg = anan - anbn + bnan - bnbn
            second_term += avg * 2 * (2n+1) * (2n+1) * wigner_A[n, n, l] * wigner_B[n, n, l]

        # Compute and accumulate the first/second terms if (ν₁, ν₂) ≠ (0, 2)
        else

            # Small change in following equations, depending on ν₂ parity (See eq. 34)

            # Loop over m
            @inbounds for m = m_star:m_max

                # First term
                # [(-1)^(l+m+n)? * Re[(an⋆ ± bn⋆)(am ± bm)] * A/B_lnm]
                anam,bnbm,anbm,bnam = (mat_anam[m,n], mat_bnbm[m,n],
                                        mat_anbm[m,n], mat_bnam[m,n])
                real_avg = ν₂_positive_flag ? real(anam+anbm+bnam+bnbm) : real(anam-anbm-bnam+bnbm);
                coef_lnm2 = 2 * (2m + 1)  * (2n + 1) * (ν₂_positive_flag ? 1 : (exp_m1(ll + n + m)))
                first_term += (real_avg * (wig_lnm[m, n, l])^2 * coef_lnm2)
            end

            # Second term
            # (-1)^l * |an-bn|² * B_lnm

            coef_lnn2 = (2n + 1)^2 * (ν₂_positive_flag ? 1 : (exp_m1(ll)))
            if (ν₂_positive_flag)
                second_term += an_p_bn[n] * (wig_lnm[n, n, l])^2 * coef_lnn2
            else
                second_term += an_m_bn[n] * (wig_lnm[n, n, l])^2 * coef_lnn2
            end
        end
    end

    # Combine the result, add coefficients, and return
    return coef * (first_term + second_term)
end


""" [Nieuwstadt1993](@cite) """
Nieuwstadt_θ_liq_ice(::Type{FT}) where {FT} = z -> if z <= 1350.0
        FT(300)
    else
        FT(300) + 3 * (z - 1350) / 1000
    end

""" [Nieuwstadt1993](@cite) """
Nieuwstadt_u(::Type{FT}) where {FT} = z -> FT(0.01)

""" [Nieuwstadt1993](@cite) """
Nieuwstadt_tke(::Type{FT}) where {FT} = z -> if (z <= 1600.0)
        FT(0.1) * FT(1.46) * FT(1.46) * (1 - z / 1600)
    else
        FT(0)
    end


# Load the packages
using GraphNeuralNetworks, DiffEqFlux, DifferentialEquations
using Flux: onehotbatch, onecold
using Flux.Losses: logitcrossentropy
using Statistics: mean
using MLDatasets: Cora
using CUDA
# CUDA.allowscalar(false) # Some scalar indexing is still done by DiffEqFlux

# device = cpu # `gpu` not working yet
device = CUDA.functional() ? gpu : cpu

# LOAD DATA
data = Cora.dataset()
g = GNNGraph(data.adjacency_list) |> device
X = data.node_features |> device
y = onehotbatch(data.node_labels, 1:data.num_classes) |> device
train_ids = data.train_indices |> device
val_ids = data.val_indices |> device
test_ids = data.test_indices |> device
ytrain = y[:, train_ids]

# Model and Data Configuration
nin = size(X, 1)
nhidden = 16
nout = data.num_classes
epochs = 40

# Define the Neural GDE
diffeqsol_to_array(x) = reshape(device(x), size(x)[1:2])

# GCNConv(nhidden => nhidden, graph=g),

node_chain = GNNChain(GCNConv(nhidden => nhidden, relu),
                      GCNConv(nhidden => nhidden, relu)) |> device

node = NeuralODE(WithGraph(node_chain, g),
                (0.f0, 1.f0), Tsit5(), save_everystep = false,
                reltol = 1e-3, abstol = 1e-3, save_start = false) |> device

model = GNNChain(GCNConv(nin => nhidden, relu),
                 Dropout(0.5),
                 node,
                 diffeqsol_to_array,
                 Dense(nhidden, nout)) |> device

# Loss
loss(x, y) = logitcrossentropy(model(g, x), y)
accuracy(x, y) = mean(onecold(model(g, x)) .== onecold(y))

# # Training
# ## Model Parameters
ps = Flux.params(model);

# ## Optimizer
opt = ADAM(0.01)

# ## Training Loop
for epoch in 1:epochs
    gs = gradient(() -> loss(X, y), ps)
    Flux.Optimise.update!(opt, ps, gs)
    @show(accuracy(X, y))
end


module Divergences

using StatsFuns
using LoopVectorization
using Parameters
import LoopVectorization: vifelse
import VectorizationBase: andmask
import Distances: PreMetric

abstract type AbstractDivergence end
abstract type Divergence <: AbstractDivergence end
abstract type AbstractModifiedDivergence <: AbstractDivergence end

struct CressieRead{T} <: Divergence
    α::T
    function CressieRead(α::T) where T<:Real
        #@assert isempty(findall((in)([-1, 0]), α)) "CressieRead is defined for all α != {-1,0}"
        new{T}(α)
    end
end

struct ChiSquared  <: Divergence end
struct KullbackLeibler  <: Divergence end
struct ReverseKullbackLeibler <: Divergence end
struct Hellinger <: Divergence end

struct ModifiedDivergence{D, T} <: AbstractModifiedDivergence
    d::D
    m::NamedTuple{(:γ₀, :γ₁, :γ₂, :ρ), Tuple{T, T, T, T}}
end

struct FullyModifiedDivergence{D, T} <: AbstractModifiedDivergence
    d::D
    m::NamedTuple{(:γ₀, :γ₁, :γ₂, :ρ, :g₀, :g₁, :g₂, :φ), Tuple{T, T, T, T, T, T, T, T}}
end

function ModifiedDivergence(D::Divergence, ρ::Real)
    @assert ρ > 1 "A ModifiedDivergence requires ρ > 1"
    γ₀ = eval(D, [ρ])[1]
    γ₁ = gradient(D, [ρ])[1]
    γ₂ = hessian(D, [ρ])[1]
    ModifiedDivergence(D, (γ₀=γ₀, γ₁=γ₁, γ₂=γ₂, ρ=ρ))
end

function FullyModifiedDivergence(D::Divergence, φ::Real, ρ::Real)
    @assert ρ > 1 "A ModifiedDivergence requires ρ > 1"
    @assert φ < 1 && φ > 0 "A ModifiedDivergence requires  φ ∈ (0,1)"
    γ₀ = eval(D, [ρ])[1]
    γ₁ = gradient(D, [ρ])[1]
    γ₂ = hessian(D, [ρ])[1]
    g₀ = eval(D, [φ])[1]
    g₁ = gradient(D, [φ])[1]
    g₂ = hessian(D, [φ])[1]
    FullyModifiedDivergence(D, (γ₀=γ₀, γ₁=γ₁, γ₂=γ₂, ρ=ρ, g₀=g₀, g₁=g₁, g₂=g₂, φ=φ))
end

const 𝒦ℒ=KullbackLeibler
const ℬ𝓊𝓇ℊ=ReverseKullbackLeibler
const 𝒞ℛ=CressieRead
const ℋ𝒟=Hellinger
const χ²=ChiSquared
include("divs.jl")

export
    # KL
    KullbackLeibler,
    # RKL
    ReverseKullbackLeibler,
    # HD
    Hellinger,
    # CR
    CressieRead,
    #
    ChiSquared,
    # Modified
    ModifiedDivergence,
    # FullyModified
    FullyModifiedDivergence,
    # Abbr.
    𝒦ℒ,
    ℬ𝓊𝓇ℊ,
    𝒞ℛ,
    ℋ𝒟,
    χ²
end


# Create symbolic variables
@polyvar x y
p = 2*x^4 + 2*x^3*y - x^2*y^2 + 5*y^4

# We want constraint `p` to be a sum of squares
@constraint model p >= 0

optimize!(model)
# Solution status is `OPTIMAL` which means `p` is a sum of squares
@show termination_status(model)


"""
    get_potential(kout, kin, P, s::ShapeParams) -> sigma_mu

Given a shape `s` with `2N` discretization nodes, outer and inner wavenumbers
`kout`,`kin`, and the cylindrical harmonics parameter `P`, returns the potential
densities `sigma_mu`. Each column contains the response to a different harmonic,
where the first `2N` entries contain the single-layer potential density
(``\\sigma``), and the lower entries contain the double-layer density (``\\mu``).
"""
function get_potential(kout, kin, P, t, ft, dft)
    N = length(t) #N here is 2N elesewhere.

    A = SDNTpotentialsdiff(kout, kin, t, ft, dft)
    LU = lu(A)

    sigma_mu = Array{Complex{Float64}}(undef, 2*N, 2*P+1)

    #assuming the wave is sampled on the shape
    nz = sqrt.(sum(abs2, ft, dims=2))
    θ = atan.(ft[:,2], ft[:,1])
    ndz = sqrt.(sum(abs2, dft, dims=2))
    nzndz = nz.*ndz
    wro = dft[:,2].*ft[:,1] - dft[:,1].*ft[:,2]
	zz = dft[:,1].*ft[:,1] + dft[:,2].*ft[:,2]

    bessp = besselj.(-P-1, kout*nz)
    bess = similar(bessp)
    du = Array{Complex{Float64}}(undef, length(bessp))
    rhs = Array{Complex{Float64}}(undef, 2*length(bessp))
    for p = -P:P
        bess[:] = besselj.(p, kout*nz)
		du[:] = kout*bessp.*wro - (p*bess./nz).*(wro + 1im*zz)
        rhs[:] = -[bess.*exp.(1.0im*p*θ);
               (du./nzndz).*exp.(1.0im*p*θ)]
        sigma_mu[:,p + P + 1] = LU\rhs
        copyto!(bessp, bess)
    end
    return sigma_mu
end

"""
    get_potential(kout, kin, P, t, ft, dft) -> sigma_mu
Same, but with the `ShapeParams` supplied directly.
"""
get_potential(kout, kin, P, s::ShapeParams) =
    get_potential(kout, kin, P, s.t, s.ft, s.dft)

function SDNTpotentialsdiff(k1, k2, t, ft, dft)
    #now just returns system matrix, utilizes similarities between upper and lower triangles
    iseven(length(t)) ?
    (N = div(length(t),2)) : (error("length(t) must be even"))

    (Rvec, kmlogvec) = KM_weights(N)
    A = Array{Complex{Float64}}(undef, 4*N, 4*N)

    ndft = sqrt.(vec(sum(abs2,dft,dims=2)))
    rij = Array{Float64}(undef, 2)
    for i=1:2*N, j=1:i
        if i == j
            T1 = -(k1^2 - k2^2)
            T2 = k1^2*(π*1im - 2MathConstants.γ + 1 - 2*log(k1*ndft[i]/2)) -
                k2^2*(π*1im - 2MathConstants.γ + 1 - 2*log(k2*ndft[i]/2))

            A[i,j] = (-ndft[i]/2/N)*log(k1/k2) #dS (dM1=0)
            A[i,j+2*N] = 1 #dD (dL=0)
            A[i+2*N,j] = -1 #dN=0
            A[i+2*N,j+2*N] = (ndft[i]/8π)*(Rvec[1]*T1 + (π/N)*T2) #dT
            continue
        end
        rij[1] = ft[i,1]-ft[j,1]
        rij[2] = ft[i,2]-ft[j,2] #ridiculous but much faster than ft[i,:]-ft[j,:]
        r = sqrt(rij[1]^2 +  rij[2]^2)
        didj = dft[i,1]*dft[j,1] + dft[i,2]*dft[j,2]

        J01 = besselj(0, k1*r)
        J02 = besselj(0, k2*r)
        H01 = besselh(0, 1, k1*r)
        H02 = besselh(0, 1, k2*r)
        k2J0 = k1^2*J01 - k2^2*J02
        k1J1 = k1*besselj(1,k1*r) - k2*besselj(1,k2*r)
        k2H0 = k1^2*H01 - k2^2*H02
        k1H1 = k1*besselh(1,k1*r) - k2*besselh(1,k2*r)

        kmlog = kmlogvec[abs(i-j)+1]
        R = Rvec[abs(i-j)+1]

        N2 = 1im*pi*k1H1 - k1J1*kmlog

        P1 = (-didj/π)*k2J0
        P2 = 1im*didj*k2H0 - P1*kmlog

        Qtilde = (dft[i,1]*rij[1] + dft[i,2]*rij[2])*(dft[j,1]*rij[1] + dft[j,2]*rij[2])/r^2

        Q1 = (-Qtilde*k2J0 + (1/r)*k1J1*(2*Qtilde - didj))/π
        Q2 = 1im*Qtilde*k2H0 + (-1im/r)*k1H1*(2*Qtilde - didj) - Q1*kmlog

        M1 = (J02 - J01)/π
        L1 = k1J1/π
        M2 = 1im*(H01 - H02) - M1*kmlog
        L2 = 1im*k1H1 - L1*kmlog

        wro_ij =  (dft[j,2]*rij[1] - dft[j,1]*rij[2])/r
        wro_ji = -(dft[i,2]*rij[1] - dft[i,1]*rij[2])/r
        cross_ij = -wro_ji*ndft[j]/ndft[i]
        cross_ji = -wro_ij*ndft[i]/ndft[j]

        #edited to remove division by wro which might be 0
        A[i,j] = (0.25*ndft[j])*(R*M1 + (π/N)*M2) #dS
        A[j,i] = A[i,j]*(ndft[i]/ndft[j]) #dS
        A[i,j+2*N] = (0.25*wro_ij)*(R*L1 + (π/N)*L2) #dD
        A[j,i+2*N] = (0.25*wro_ji)*(R*L1 + (π/N)*L2) #dD
        A[i+2*N,j] = (-0.25*cross_ij/π)*(R*k1J1 + (π/N)*N2) #dN
        A[j+2*N,i] = (-0.25*cross_ji/π)*(R*k1J1 + (π/N)*N2) #dN
        A[i+2*N,j+2*N] = (R*(P1-Q1) + (π/N)*(P2-Q2))/(4*ndft[i]) #dT
        A[j+2*N,i+2*N] = A[i+2*N,j+2*N]*(ndft[i]/ndft[j]) #dT
    end
    any(isnan.(A)) && error("SDNTpotentialsdiff: encountered NaN, check data and division by ndft.")
    return A
end

function KM_weights(N)
    #computes the weights necessary for Kussmaul-Martensen quadrature (evenly
    #spaced).
    #Input: N (integer>=1)
    #Output: R,K (float vectors of length 2N)
    arg1 = Float64[cos(m*j*π/N)/m for m=1:N-1, j=0:2*N-1]
    R = vec((-2π/N)*sum(arg1,dims=1)) - (π/N^2)*Float64[cos(j*π) for j=0:2*N-1]

    K = Float64[2*log(2*sin(0.5π*j/N)) for j = 0:2*N-1]
    return (R,K)
end

"""
    get_potentialPW(kout, kin, s::ShapeParams, θ_i) -> sigma_mu

Given a shape `s` with `2N` discretization nodes, outer and inner wavenumbers
`kout`,`kin`, and an incident plane-wave angle, returns the potential
densities vector `sigma_mu`. The first `2N` entries contain the single-layer
potential density (``\\sigma``), and the lower entries contain the double-layer
density (``\\mu``).
"""
function get_potentialPW(kout, kin, s, θ_i)
    N = length(s.t) #N here is different...

    A = SDNTpotentialsdiff(kout, kin, s.t, s.ft, s.dft)
    LU = lu(A)

	ndft = sqrt.(sum(abs2, s.dft, dims=2))
    ui = exp.(1.0im*kout*(cos(θ_i)*s.ft[:,1] + sin(θ_i)*s.ft[:,2]))
    rhs = -[ui;
			(1.0im*kout*ui).*((cos(θ_i)*s.dft[:,2] - sin(θ_i)*s.dft[:,1])./ndft)]

    sigma_mu = LU\rhs
end

"""
    scatteredfield(sigma_mu, k, s::ShapeParams, p) -> u_s

Computes field scattered by the particle `s` with pre-computed potential
densities `sigma_mu` at points `p`. All points must either be inside `k = kin`
or outside `k = kout` the particle.
"""
scatteredfield(sigma_mu, k, s::ShapeParams, p) =
    scatteredfield(sigma_mu, k, s.t, s.ft, s.dft, p)

"""
    scatteredfield(sigma_mu, k, t, ft, dft, p) -> u_s

Same, but with the `ShapeParams` supplied directly. Useful for computing `u_s`
for rotated shapes.
"""
function scatteredfield(sigma_mu, k, t, ft, dft, p)
    #calculates the scattered field of a shape with parametrization ft(t),...,dft(t)
    #in space with wavenumber k at points p *off* the boundary. For field on the boundary,
    #SDpotentials function must be used.
    if size(p,2) == 1 #single point, rotate it
        p = transpose(p)
    end
    N = length(t)
    M = size(p,1)
    r = zeros(Float64,2)
    #loop is faster here:
    SDout = Array{Complex{Float64}}(undef, M, 2*N)
    for j = 1:N
        ndft = hypot(dft[j,1],dft[j,2])
        for i = 1:M
            r[:] = [p[i,1] - ft[j,1];p[i,2] - ft[j,2]]
            nr = hypot(r[1],r[2])
            if nr < eps()
                #TODO: use SDNTpotentialsdiff here
                @warn("Encountered singularity in scatteredfield.")
                SDout[i,j] = 0
                SDout[i,j+N] = 0
                continue
            end
            SDout[i,j] = (2*pi/N)*0.25im*besselh(0,1, k*nr)*ndft
            SDout[i,j+N] = (2*pi/N)*0.25im*k*besselh(1,1, k*nr)*(dft[j,2]*r[1] - dft[j,1]*r[2])/nr
        end
    end
    u_s = SDout*sigma_mu
end

function shapeMultipoleExpansion(k, t, ft, dft, P)
    #unlike others (so far), this does *not* assume t_j=pi*j/N
    N = div(length(t),2)
    nz = vec(sqrt.(sum(abs2, ft, dims=2)))
    θ = atan.(ft[:,2], ft[:,1])
    ndz = vec(sqrt.(sum(abs2, dft, dims=2)))
    AB = Array{Complex{Float64}}(undef, 2*P + 1, 4*N)
    bessp = besselj.(-P-1,k*nz)
    bess = similar(bessp)
    for l = -P:0
        bess[:] = besselj.(l,k*nz)
        for j = 1:2*N
            AB[l+P+1,j] = 0.25im*(π/N)*bess[j]*exp(-1.0im*l*θ[j])*ndz[j]
            l != 0 && (AB[-l+P+1,j] = 0.25im*((-1.0)^l*π/N)*bess[j]*exp(1.0im*l*θ[j])*ndz[j])
            wro = ft[j,1]*dft[j,2] - ft[j,2]*dft[j,1]
            zdz = -1.0im*(ft[j,1]*dft[j,1] + ft[j,2]*dft[j,2])
            b1 = (-l*bess[j]/nz[j])*(zdz + wro)
            b1_ = (-l*bess[j]/nz[j])*(zdz - wro)
            b2 = k*bessp[j]*wro
            AB[l+P+1,j+2*N] = 0.25im*(π/N)*(exp(-1.0im*l*θ[j])/nz[j])*(b1 + b2)
            l != 0 && (AB[-l+P+1,j+2*N] = 0.25im*((-1.0)^l*π/N)*(exp(1.0im*l*θ[j])/nz[j])*(-b1_ + b2))
        end
        copyto!(bessp,bess)
    end
    return AB
end

function solvePotential_forError(kin, kout, shape, ls_pos, ls_amp, θ_i)
    #plane wave outside, line sources inside
    N = length(shape.t) #N here is different...

    A = SDNTpotentialsdiff(kout, kin, shape.t, shape.ft, shape.dft)
    LU = lu(A)

	ndft = sqrt.(sum(abs2, shape.dft, dims=2))

    r = sqrt.((shape.ft[:,1] .- ls_pos[1,1]).^2 +
            (shape.ft[:,2] .- ls_pos[1,2]).^2)

    uls = (-ls_amp[1]*0.25im)*besselh.(0,kout*r)
    duls = (ls_amp[1]*0.25im*kout)*besselh.(1,kout*r).*((shape.ft[:,1].-ls_pos[1,1]).*shape.dft[:,2]-(shape.ft[:,2].-ls_pos[1,2]).*shape.dft[:,1])./r./ndft
    for i = 2:length(ls_amp)
        r = sqrt.((shape.ft[:,1] - ls_pos[i,1]).^2 + (shape.ft[:,2] - ls_pos[i,2]).^2)
        uls -= ls_amp[i]*0.25im*besselh.(0,kout*r)
        duls -= -ls_amp[i]*0.25im*kout*besselh.(1,kout*r).*((shape.ft[:,1]-ls_pos[i,1]).*shape.dft[:,2]-(shape.ft[:,2]-ls_pos[i,2]).*shape.dft[:,1])./r./ndft
    end

    #outer plane wave
    ui = exp.(1.0im*kin*(cos(θ_i)*shape.ft[:,1] + sin(θ_i)*shape.ft[:,2]))
    dui = (1.0im*kin)*(ui.*(cos(θ_i)*shape.dft[:,2] - sin(θ_i)*shape.dft[:,1])./ndft)

    rhs = -[ui+uls;dui+duls]
    sigma_mu = LU\rhs
    return sigma_mu
end


using BenchmarkTools, LinearAlgebra, Plots
gr()
import JuliaBLAS: addmul!, Block

function bench(nn=30)
    time2gflops(mnk, time) = 2 * mnk^3 / time

    BLAS.set_num_threads(1)
    mnks = 48*(1:2:nn)
    obtimes = zeros(length(mnks)); jbtimes = zeros(length(mnks))

    for (i,siz) in enumerate(mnks)
        A,B,C = (zeros(siz,siz) for i in 1:3)
        bj = @benchmarkable addmul!($C,$A,$B) samples=20 time_tolerance=0.1 seconds=0.8
        jbtimes[i] = run(bj) |> minimum |> time
        bo = @benchmarkable mul!($C,$A,$B) samples=20 time_tolerance=0.1 seconds=0.8
        obtimes[i] = run(bo) |> minimum |> time
    end

    jflops = time2gflops.(mnks, jbtimes)
    oflops = time2gflops.(mnks, obtimes)
    plot(mnks, jflops, lab="JuliaBLAS")
    plot!(mnks, oflops, lab="OpenBLAS", ylabel="GFLOPS", xlabel="M=N=K", legend=:bottomright, dpi=400, ylims=(0,60), yticks=0:5:50)
end
plt = bench()
savefig(plt, "bench1.png")


export rats, pumps, seeds, coin

coin = @model flips begin
    pHeads ~ Beta(1,1)
    flips ~ Bernoulli(pHeads) |> iid(20)
end

export lda
lda = @model (α, η, K, V, N) begin
    M = length(N)
    β ~ Dirichlet(repeat([η],V)) |> iid(K)
    θ ~ Dirichlet(repeat([α],K)) |> iid(M)
    z ~ For(1:M) do m
            Categorical(θ[m]) |> iid(N[m])
        end
    w ~ For(1:M) do m
            For(1:N[m]) do n
                Categorical(β[z[m][n]])
            end
        end
end

export hello
hello = @model μ begin
    σ ~ HalfCauchy()
    x ~ Normal(μ,σ)
end

export normalModel
normalModel = @model x begin
    μ ~ Normal(0,5)
    σ ~ HalfCauchy(3)
    x ~ Normal(μ,σ) |> iid(10)
end

# Is this right?
export tdist
tdist = @model ν begin
    w ~ InverseGamma(ν/2,ν/2)
    x ~ Normal(0, w^2)
    return x
end

export nested
nested = @model x begin
    μ ~ simpleModel(s = 2.0)
    σ ~ HalfCauchy(3)
    # x ~ Normal(μ.z,σ) |> iid
end

# export nested2
# nested2 = @model x begin
#     M ~ @model begin
#         z ~ Normal(0,2.0)
#     end
#     σ ~ HalfCauchy(3)
#     x ~ Normal(M.z, σ)
# end

export simpleModel
simpleModel = @model s begin
    z ~ Normal(0,s)
end

export mix
mix = @model (K,α) begin
    p ~ Dirichlet(repeat([α],K))
    μ ~ Cauchy(0,5) |> iid(K)
    σ ~ HalfCauchy(3) |> iid(K)
    components = Normal.(μ,σ)
    N ~ Poisson(100)
    x ~ MixtureModel(components, p) |> iid(N)
end

export linReg1D

linReg1D_local(ŷ, σ) = n -> Normal(ŷ[n])
linReg1D = @model (x, y) begin
        α ~ Cauchy(0, 10)
        β ~ Cauchy(0, 2.5)
        σ ~ HalfCauchy(3)
        ŷ = α .+ β .* x
        N = length(x)
        y ~ For(linReg1D_local(ŷ, σ), 1:N)
    end

# From OpenBUGS and section 6 of Gelfand et al.
rats = @model x begin
    μdist = Normal(0,1000)
    μα ~ μdist
    μβ ~ μdist

    σ2dist = InverseGamma(0.001,0.001)
    σ2α ~ σ2dist
    σ2β ~ σ2dist
    σ2c ~ σ2dist

    α ~ Normal(μα, sqrt(σ2α)) |> iid(30)
    β ~ Normal(μβ, sqrt(σ2β)) |> iid(30)

    x̄ = mean(x)
    y ~ For([(i,j) for i in 1:30, j in 1:5]) do (i,j)
            Normal(α[i] + β[i]*(x[j]-x̄), sqrt(σ2c))
        end
end

pumps = @model t begin
    n = length(t)
    α ~ Gamma(1,1)
    β ~ Gamma(0.1,1)
    θ ~ For(1:n) do i
            Gamma(α, 1/β)
        end
    y ~ For(1:n) do i
            Poisson(θ[i] * t[i])
        end
end

# dogs = @model
#     lo = typemin(Float64) |> nextfloat
#     negReals = Uniform(lo,-1e-5)
#     α ~ negReals
#     β ~ negReals
#     y ~ For()
#     Bernoulli
#     α
# end

seeds = @model (n,x) begin
    σ2 ~ InverseGamma(0.001,0.001)
    σ = sqrt(σ2)
    αDist = Normal(0,1000)
    α0 ~ αDist
    α1 ~ αDist
    α2 ~ αDist
    α12 ~ αDist
    b ~ Normal(0,σ) |> iid(21)
    y = α0 .+ α1 .* x[1,:] .+ α2 .* x[2,:] .+ α12 .* x[1,:] .* x[2,:] .+ b
    r ~ For(1:21) do i
            LogisticBinomial(n[i],y[i])
        end
end

stacks = @model begin
    σ2 ~ InverseGamma(0.001,0.001)
    βDist = Normal(0,1000)
    β0 ~ βDist
    β1 ~ βDist
    β2 ~ βDist
    β3 ~ βDist
    μ = β0 + β1*z[1,:] + β2*z[2,:] + β3*z[3,:]
    y ~ For(1:21) do i
            Laplace(μ[i],σ2)
        end
end

y = [28.,  8., -3.,  7., -1.,  1., 18., 12.]
σ = [15., 10., 16., 11.,  9., 11., 10., 18.]

export school8
school8 = @model (y, σ) begin
  μ ~ Flat()
  τ ~ HalfFlat()
  η ~ Normal() |> iid(8)
  y ~ For(1:8) do j
      Normal(μ + τ * η[j], σ[j])
  end
end

export funnel
funnel = @model begin
    y ~ Normal(0, 3);
    x ~ Normal(0, exp(y/2))
end


using DiffEqDevTools, Test

@test stability_region(constructDormandPrince6(), initial_guess=-3.5) ≈ -3.95413 rtol=1e-3
@test stability_region(constructTsitourasPapakostas6(), initial_guess=-3.5) ≈ -3.95413 rtol=1e-3
@test stability_region(constructRadauIIA5(), initial_guess=12.) ≈ 11.84 rtol=1e-2


using Distributions, SpecialFunctions, Roots, LinearAlgebra
"""
GNas12original: Modified algorithm of Nassar et al., Nat Neuro, 2012
            for the Gaussian task (section 2.3.1)
- Gaussian prior with sigma_0 >> sigma
- All update functions are same as the ones of GNas12
"""
struct GNas12original <: GLearner
    χ_0::Float64
    ν_0::Float64
    mu_0::Float64
    sigma_0::Float64 # std of prior
    sigma::Float64 # std of observations
    changeprobability::Float64
    m::Float64
    ρ::Float64
    χ_n::Array{Float64, 1}
    ν_n::Array{Float64, 1}
    r_n::Array{Float64, 1}
    var_n::Array{Float64, 1}
    mu_n::Array{Float64, 1} #parameter we are interested in
    Sgm_n::Array{Float64, 1} # Bayes Factor Surprise
    γ_n::Array{Float64, 1}
    Ssh_n::Array{Float64, 1} # Shannon Surprise
end
function GNas12original(; mu_0 = 0., sigma = 0.1, changeprobability = .001)
    sigma_0 = 1000.
    ρ = sigma^2 / sigma_0^2
    m = changeprobability/ (1. - changeprobability)
    ν_0 = calcν(sigma, sigma_0)
    χ_0 = calcχ(mu_0, ν_0, sigma)
    var_n_0 = sigma^2 / ν_0
    mu_n_0 = sigma * χ_0 / ν_0
    GNas12original(χ_0, ν_0, mu_0, sigma_0, sigma, changeprobability, m, ρ, [χ_0], [ν_0],
                [1.], [var_n_0], [mu_n_0], [1.], [m/(1. + m)], [1.])
end
export GNas12original


# Test simple evaluations:
options = Options(;
    default_params..., binary_operators=(+, *, /, -), unary_operators=(cos, sin)
)

# Here, we unittest the fast function evaluation scheme
# We need to trigger all possible fused functions, with all their logic.
# These are as follows:

## We fuse (and compile) the following:
##  - op(op2(x, y)), where x, y, z are constants or variables.
##  - op(op2(x)), where x is a constant or variable.
##  - op(x), for any x.
## We fuse (and compile) the following:
##  - op(x, y), where x, y are constants or variables.
##  - op(x, y), where x is a constant or variable but y is not.
##  - op(x, y), where y is a constant or variable but x is not.
##  - op(x, y), for any x or y
for fnc in [
    # deg2_l0_r0_eval
    (x1, x2, x3) -> x1 * x2,
    (x1, x2, x3) -> x1 * 3.0f0,
    (x1, x2, x3) -> 3.0f0 * x2,
    (((x1, x2, x3) -> 3.0f0 * 6.0f0), ((x1, x2, x3) -> Node(3.0f0) * 6.0f0)),
    # deg2_l0_eval
    (x1, x2, x3) -> x1 * sin(x2),
    (x1, x2, x3) -> 3.0f0 * sin(x2),

    # deg2_r0_eval
    (x1, x2, x3) -> sin(x1) * x2,
    (x1, x2, x3) -> sin(x1) * 3.0f0,

    # deg1_l2_ll0_lr0_eval
    (x1, x2, x3) -> cos(x1 * x2),
    (x1, x2, x3) -> cos(x1 * 3.0f0),
    (x1, x2, x3) -> cos(3.0f0 * x2),
    (((x1, x2, x3) -> cos(3.0f0 * -0.5f0)), ((x1, x2, x3) -> cos(Node(3.0f0) * -0.5f0))),

    # deg1_l1_ll0_eval
    (x1, x2, x3) -> cos(sin(x1)),
    (((x1, x2, x3) -> cos(sin(3.0f0))), ((x1, x2, x3) -> cos(sin(Node(3.0f0))))),

    # everything else:
    (x1, x2, x3) -> (sin(cos(sin(cos(x1) * x3) * 3.0f0) * -0.5f0) + 2.0f0) * 5.0f0,
]

    # check if fnc is tuple
    if typeof(fnc) <: Tuple
        realfnc = fnc[1]
        nodefnc = fnc[2]
    else
        realfnc = fnc
        nodefnc = fnc
    end

    global tree = nodefnc(Node("x1"), Node("x2"), Node("x3"))

    N = 100
    nfeatures = 3
    X = randn(MersenneTwister(0), Float32, nfeatures, N)

    test_y = eval_tree_array(tree, X, options)[1]
    true_y = realfnc.(X[1, :], X[2, :], X[3, :])

    zero_tolerance = 1e-6
    @test all(abs.(test_y .- true_y) / N .< zero_tolerance)
end


@polly function kernel_bicg(A,s,q,p,r)
    n,m = size(A)

    for i = 1:m
        s[i] = zero(eltype(s))
    end

    for i = 1:n
        q[i] = zero(eltype(q));
        for j = 1:m
            s[j] = s[j] + r[i] * A[i,j]
            q[i] = q[i] + A[i,j] * p[j]
        end
    end
end

let
    m = 1900
    n = 2100

    A = zeros(Float32, n, m)
    s = zeros(Float32, m)
    q = zeros(Float32, n)
    p = zeros(Float32, m)
    r = zeros(Float32, n)

    for i = 1:m
        p[i] = (i % m) / m
    end
    for i = 1:n
        r[i] = (i % n) / n
        for j = 1:m
            A[i,j] = (i*(j+1) % n)/n
        end
    end

    SUITE["bicg"] = @benchmarkable kernel_bicg(A, s, q, p, r) setup = (A = copy($A); s = copy($s); q = copy($q); p = copy($p); r = copy($r))
end


using Dates: Year, Date

toDate(dstr::String)::Date = Date(dstr, "mm/dd/yy") + Year(2000)

# Same dates for all test cases
# Test dates and values from https://www.isda.org/2008/12/22/30-360-day-count-conventions/ , https://www.isda.org/a/mIJEE/30-360-2006ISDADefs.xls
startDates = map(toDate, ["01/15/07", "01/15/07", "01/15/07", "09/30/07", "09/30/07", "09/30/07", "01/15/07", "01/31/07", "02/28/07", "08/31/06", "02/28/07", "02/14/07", "02/26/07", "02/29/08", "02/29/08", "02/29/08", "02/28/07", "10/31/07", "08/31/07", "02/29/08", "08/31/08", "02/28/09"])
endDates   = map(toDate, ["01/30/07", "02/15/07", "07/15/07", "03/31/08", "10/31/07", "09/30/08", "01/31/07", "02/28/07", "03/31/07", "02/28/07", "08/31/07", "02/28/07", "02/29/08", "02/28/09", "03/30/08", "03/31/08", "03/05/07", "11/28/07", "02/29/08", "08/31/08", "02/28/09", "08/31/09"])
terminationDate = toDate("02/28/09")

function test_helper(daysResults, daysFracResults, ::Type{T}) where T <: Learn.AbstractDayCountConvention
    for (sDate, eDate, daysR, yearFracR) in zip(startDates, endDates, daysResults, daysFracResults)
        @test days(sDate, eDate, T) == daysR
        @test yearfraction(sDate, eDate, T) ≈ yearFracR atol=0.0001
    end
end

function test_helper(daysResults, daysFracResults, tDate, ::Type{T}) where T <: Learn.AbstractDayCountConvention
    for (sDate, eDate, daysR, yearFracR) in zip(startDates, endDates, daysResults, daysFracResults)
        @test days(sDate, eDate, tDate, T) == daysR
        @test yearfraction(sDate, eDate, tDate, T) ≈ yearFracR atol=0.0001
    end
end

@testset "DayCounts" begin

    @testset "30/360 Bond Basis, 30A/360" begin
        # Test cases from https://www.isda.org/a/mIJEE/30-360-2006ISDADefs.xls
        testType = DayCount30A_360
        daysResults = [15,30,180,180,30,360,16,28,33,178,183,14,363,359,31,32,7,28,179,182,178,183]
        daysFracResults = [0.0417,0.0833,0.5000,0.5000,0.0833,1.0000,0.0444,0.0778,0.0917,0.4944,0.5083,0.0389,1.0083,0.9972,0.0861,0.0889,0.0194,0.0778,0.4972,0.5056,0.4944,0.5083]
        test_helper(daysResults, daysFracResults, testType)
    end

    @testset "30E/360, 30/360 ICMA, 30/360 ISMA, 30S/360, Eurobond basis (ISDA 2006), Special German" begin
        # Test cases from https://www.isda.org/a/mIJEE/30-360-2006ISDADefs.xls
        testType = DayCount30E_360
        daysResults = [15, 30, 180, 180, 30, 360, 15, 28, 32, 178, 182, 14, 363, 359, 31, 31, 7, 28, 179, 181, 178, 182]
        daysFracResults = [0.0417, 0.0833, 0.5000, 0.5000, 0.0833, 1.0000, 0.0417, 0.0778, 0.0889, 0.4944, 0.5056, 0.0389, 1.0083, 0.9972, 0.0861, 0.0861, 0.0194, 0.0778, 0.4972, 0.5028, 0.4944, 0.5056]
        test_helper(daysResults, daysFracResults, testType)
    end

    @testset "30E/360ISDA, Eurobond basis (ISDA 2000), German" begin
        # Test cases from https://www.isda.org/a/mIJEE/30-360-2006ISDADefs.xls
        testType = DayCount30E_360ISDA
        daysResults = [15, 30, 180, 180, 30, 360, 15, 30, 30, 180, 180, 16, 364, 358, 30, 30, 5, 28, 180, 180, 178, 180]
        daysFracResults = [0.0417, 0.0833, 0.5000, 0.5000, 0.0833, 1.0000, 0.0417, 0.0833, 0.0833, 0.5000, 0.5000, 0.0444, 1.0111, 0.9944, 0.0833, 0.0833, 0.0139, 0.0778, 0.5000, 0.5000, 0.4944, 0.5000]
        test_helper(daysResults, daysFracResults, terminationDate, testType)
    end

    @testset "Act/360, French" begin
        # Test cases from https://www.isda.org/a/mIJEE/30-360-2006ISDADefs.xls
        testType = DayCountAct_360
        daysResults = [15, 31, 181, 183, 31, 366, 16, 28, 31, 181, 184, 14, 368, 365, 30, 31, 5, 28, 182, 184, 181, 184]
        daysFracResults = [0.0417, 0.0861, 0.5028, 0.5083, 0.0861, 1.0167, 0.0444, 0.0778, 0.0861, 0.5028, 0.5111, 0.0389, 1.0222, 1.0139, 0.0833, 0.0861, 0.0139, 0.0778, 0.5056, 0.5111, 0.5028, 0.5111]
        test_helper(daysResults, daysFracResults, testType)
    end

    @testset "Act/364" begin
        testType = DayCountAct_364
        daysResults = [15, 31, 181, 183, 31, 366, 16, 28, 31, 181, 184, 14, 368, 365, 30, 31, 5, 28, 182, 184, 181, 184]
        daysFracResults = map(x-> x*360.0/364.0, [0.0417, 0.0861, 0.5028, 0.5083, 0.0861, 1.0167, 0.0444, 0.0778, 0.0861, 0.5028, 0.5111, 0.0389, 1.0222, 1.0139, 0.0833, 0.0861, 0.0139, 0.0778, 0.5056, 0.5111, 0.5028, 0.5111])
        test_helper(daysResults, daysFracResults, testType)
    end

    @testset "Act/365Fixed" begin
        testType = DayCountAct_365Fixed
        daysResults = [15, 31, 181, 183, 31, 366, 16, 28, 31, 181, 184, 14, 368, 365, 30, 31, 5, 28, 182, 184, 181, 184]
        daysFracResults = map(x-> x*360.0/365.0, [0.0417, 0.0861, 0.5028, 0.5083, 0.0861, 1.0167, 0.0444, 0.0778, 0.0861, 0.5028, 0.5111, 0.0389, 1.0222, 1.0139, 0.0833, 0.0861, 0.0139, 0.0778, 0.5056, 0.5111, 0.5028, 0.5111])
        test_helper(daysResults, daysFracResults, testType)
    end

    @testset "Act/365Act" begin
        testType = DayCountAct_365Actual
        daysResults = [15, 31, 181, 183, 31, 366, 16, 28, 31, 181, 184, 14, 368, 365, 30, 31, 5, 28, 182, 184, 181, 184]
        daysFracResults = map(x-> x*360.0/365.0, [0.0417, 0.0861, 0.5028, 0.5083, 0.0861, 1.0167, 0.0444, 0.0778, 0.0861, 0.5028, 0.5111, 0.0389, 1.0222, 1.0139, 0.0833, 0.0861, 0.0139, 0.0778, 0.5056, 0.5111, 0.5028, 0.5111])

        for i in [4, 6, 13, 19]
            daysFracResults[i] = daysFracResults[i]*365.0/366.0
        end
        test_helper(daysResults, daysFracResults, testType)
    end

    # @testset "Act/ActISDA" begin
    #     testType = DayCountAct_Act_ISDA
    #     daysResults = [15, 31, 181, 183, 31, 366, 16, 28, 31, 181, 184, 14, 368, 365, 30, 31, 5, 28, 182, 184, 181, 184]
    #     daysFracResults = map(x-> x*360.0/365.0, [0.0417, 0.0861, 0.5028, 0.5083, 0.0861, 1.0167, 0.0444, 0.0778, 0.0861, 0.5028, 0.5111, 0.0389, 1.0222, 1.0139, 0.0833, 0.0861, 0.0139, 0.0778, 0.5056, 0.5111, 0.5028, 0.5111])

    #     test_helper(daysResults, daysFracResults, testType)
    # end

end


using Distributions
export get_informed_insertion_dist

function get_informed_insertion_dist(
    posterior::Union{SisPosterior,SimPosterior},
    α::Float64
    )
    data = posterior.data
    V = posterior.V
    p = zeros(length(V))
    for S in data
        p .+= [i ∈ vcat(S...) for i in V]
    end
    p ./= sum(p)
    p .+= α
    p ./= sum(p)
    return Categorical(p)
end


function mhash(A::Array{Int64, 2})
  return reduce(*, map(string, A))
end

function nonempty(A::Array{Int64, 2})
  return prod(vec(sum(A, 1)) .+ vec(sum(A, 2))) > 0
end

function hashmat(x)
  V = round(Int64, sqrt(length(x)))
  A = reshape(x, (V, V))
  if nonempty(A)
    return mhash(A)
  end
end

function unhashmat(x)
  c = Int64[]
  for i in 1:length(x)
    push!(c, parse(Int64, x[i]))
  end
  V = round(Int64, sqrt(length(x)))
  return reshape(vec(c), (V, V))
end

function from_binarytree(h, s)
   tree = zeros(Int64, 2^(h+1)-2)
   for i in 1:length(tree)
      if iseven(i)
         tree[i] = 1
      end
   end
   terminal_nodes = (length(tree)-2^h+1):length(tree)
   paths = []
   for n in terminal_nodes
      i = h
      position = n
      sequence = zeros(Int64, h)
      while(i > 0)
         sequence[i] = tree[position]
         position = round(Int,floor((position - 1)/2))
         i = i-1
      end
      if sum(sequence) == s
         push!(paths, sequence)
      end
   end
   return paths
end


using Random, LinearAlgebra, Statistics, Optim, DataFrames, CSV, HTTP, GLM, ForwardDiff, FreqTables, Distributions

# read in data
url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2020/master/ProblemSets/PS4-mixture/nlsw88t.csv"
df = CSV.read(HTTP.get(url).body)
X = [df.age df.white df.collgrad]
Z = hcat(df.elnwage1, df.elnwage2, df.elnwage3, df.elnwage4,
         df.elnwage5, df.elnwage6, df.elnwage7, df.elnwage8)
y = df.occ_code

# estimate multionomial logit model
include("mlogit_with_Z.jl")
θ_start = [ .0403744; .2439942; -1.57132; .0433254; .1468556; -2.959103; .1020574; .7473086; -4.12005; .0375628; .6884899; -3.65577; .0204543; -.3584007; -4.376929; .1074636; -.5263738; -6.199197; .1168824; -.2870554; -5.322248; 1.307477]
td = TwiceDifferentiable(b -> mlogit_with_Z(b, X, Z, y), θ_start; autodiff = :forward);
θ̂_optim_ad = optimize(td, θ_start, LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50));
θ̂_mle_ad = θ̂_optim_ad.minimizer
H  = Optim.hessian!(td, θ̂_mle_ad)
θ̂_mle_ad_se = sqrt.(diag(inv(H)))

# get baseline predicted probabilities
function plogit(θ, X, Z, J)
    α = θ[1:end-1]
    γ = θ[end]
    K = size(X,2)
    bigα = [reshape(α,K,J-1) zeros(K)]
    P = exp.(X*bigα .+ Z*γ) ./ sum.(eachrow(exp.(X*bigα .+ Z*γ)))
    return P
end
println(θ̂_mle_ad)
P = plogit(θ̂_mle_ad, X, Z, length(unique(y)))

# compare with data
modelfitdf = DataFrame(data_freq = 100*convert(Array, prop(freqtable(df,:occ_code))), P = 100*vec(mean(P'; dims=2)))
display(modelfitdf)

# now run a counterfactual:
# suppose people don't care about wages when choosing occupation
# this means setting θ̂_mle_ad[end] = 0
θ̂_cfl = deepcopy(θ̂_mle_ad)
θ̂_cfl[end] = 0
P_cfl = plogit(θ̂_cfl, X, Z, length(unique(y)))
modelfitdf.P_cfl = 100*vec(mean(P_cfl'; dims=2))
modelfitdf.diff = modelfitdf.P_cfl .- modelfitdf.P
modelfitdf.ew   = vec(mean(Z'; dims=2))
display(modelfitdf)

# now do parametric bootstrap
Random.seed!(1234)
invHfix = UpperTriangular(inv(H)) .- Diagonal(inv(H)) .+ UpperTriangular(inv(H))' # this code makes inv(H) truly symmetric (not just numerically symmetric)
d = MvNormal(θ̂_mle_ad,invHfix)
B = 1000
cfl_bs = zeros(8,B)
for b = 1:B
    θ_draw = rand(d)
    Pb = plogit(θ_draw, X, Z, length(unique(y)))
    Pbs = 100*vec(mean(Pb'; dims=2))
    θ_draw[end] = 0
    P_cflb = plogit(θ_draw, X, Z, length(unique(y)))
    Pcfl_bs = 100*vec(mean(P_cflb'; dims=2))
    cfl_bs[:,b] = Pcfl_bs .- Pbs
end
bsdiffL = deepcopy(cfl_bs[:,1])
bsdiffU = deepcopy(cfl_bs[:,1])
for i=1:size(cfl_bs,1)
    bsdiffL[i] = quantile(cfl_bs[i,:],.025)
    bsdiffU[i] = quantile(cfl_bs[i,:],.975)
end
modelfitdf.bs_diffL = bsdiffL
modelfitdf.bs_diffU = bsdiffU
display(modelfitdf)