module Yields

import Interpolations
import ForwardDiff
using LinearAlgebra
using UnicodePlots

# don't export type, as the API of Yields.Zero is nicer and
# less polluting than Zero and less/equally verbose as ZeroYieldCurve or ZeroCurve
export rate, discount, accumulation, forward, Yield, Rate, rate, spot

abstract type CompoundingFrequency end
Base.Broadcast.broadcastable(x::T) where {T<:CompoundingFrequency} = Ref(x)

"""
    Continuous()

A type representing continuous interest compounding frequency.

# Examples

```julia-repl
julia> Rate(0.01,Continuous())
Rate(0.01, Continuous())
```

See also: [`Periodic`](@ref)
"""
struct Continuous <: CompoundingFrequency end

"""

```julia-repl
julia> Continuous(0.01)
Rate(0.01, Continuous())
```

See also: [`Periodic`](@ref)
"""
Continuous(rate) = Rate(rate, Continuous())

"""
    Periodic(frequency)

A type representing periodic interest compounding with the given frequency

# Examples

Creating a semi-annual bond equivalent yield:

```julia-repl
julia> Rate(0.01,Periodic(2))
Rate(0.01, Periodic(2))
```

See also: [`Continuous`](@ref)
"""
struct Periodic <: CompoundingFrequency
    frequency::Int
end

"""
    Periodic(rate,frequency)

A convenience constructor for Rate(rate,Periodic(frequency)).

# Examples

Creating a semi-annual bond equivalent yield:

```julia-repl
julia> Periodic(0.01,2)
Rate(0.01, Periodic(2))
```

See also: [`Continuous`](@ref)
"""
Periodic(x, frequency) = Rate(x, Periodic(frequency))

struct Rate{N<:Real,T<:CompoundingFrequency}
    value::N
    compounding::T
end

# Base.:==(r1::Rate,r2::Rate) = (r1.value == r2.value) && (r1.compounding == r2.compounding)

"""
    Rate(rate[,frequency=1])
    Rate(rate,frequency::CompoundingFrequency)

Rate is a type that encapsulates an interest `rate` along with its compounding `frequency`.

Periodic rates can be constructed via `Rate(rate,frequency)` or `Rate(rate,Periodic(frequency))`.

Continuous rates can be constructed via `Rate(rate, Inf)` or `Rate(rate,Continuous())`.

# Examples

```julia-repl
julia> Rate(0.01,Continuous())
Rate(0.01, Continuous())

julia> Rate(0.01,Periodic(2))
Rate(0.01, Periodic(2))

julia> Rate(0.01)
Rate(0.01, Periodic(1))

julia> Rate(0.01,2)
Rate(0.01, Periodic(2))

julia> Rate(0.01,Periodic(4))
Rate(0.01, Periodic(4))

julia> Rate(0.01,Inf)
Rate(0.01, Continuous())

julia> Rate(0.01,Continuous())
Rate(0.01, Continuous())
```
"""
Rate(rate) = Rate(rate, Periodic(1))
Rate(x, frequency::T) where {T<:Real} = isinf(frequency) ? Rate(x, Continuous()) : Rate(x, Periodic(frequency))

"""
    convert(T::CompoundingFrequency,r::Rate)

Returns a `Rate` with an equivalent discount but represented with a different compounding frequency.

# Examples

```julia-repl
julia> r = Rate(Periodic(12),0.01)
Rate(0.01, Periodic(12))

julia> convert(Periodic(1),r)
Rate(0.010045960887181016, Periodic(1))

julia> convert(Continuous(),r)
Rate(0.009995835646701251, Continuous())
```
"""
function Base.convert(T::CompoundingFrequency, r::Rate{<:Real,<:CompoundingFrequency})
    convert(T, r, r.compounding)
end
function Base.convert(to::Continuous, r, from::Continuous)
    return r
end

function Base.convert(to::Periodic, r, from::Continuous)
    return Rate(to.frequency * (exp(r.value / to.frequency) - 1), to)
end

function Base.convert(to::Continuous, r, from::Periodic)
    return Rate(from.frequency * log(1 + r.value / from.frequency), to)
end

function Base.convert(to::Periodic, r, from::Periodic)
    c = convert(Continuous(), r, from)
    return convert(to, c, Continuous())
end

function rate(r::Rate{<:Real,<:CompoundingFrequency})
    r.value
end

function Base.isapprox(a::Rate{N,T}, b::Rate{N,T}; atol::Real = 0, rtol::Real = atol > 0 ? 0 : √eps()) where {T<:Periodic,N<:Real}
    return (a.compounding.frequency == b.compounding.frequency) && isapprox(rate(a), rate(b); atol, rtol)
end

function Base.isapprox(a::Rate{N,T}, b::Rate{N,T}; atol::Real = 0, rtol::Real = atol > 0 ? 0 : √eps()) where {T<:Continuous,N<:Real}
    return isapprox(rate(a), rate(b); atol, rtol)
end

# the fallback for rates not of the same type
function Base.isapprox(a::T, b::N; atol::Real = 0, rtol::Real = atol > 0 ? 0 : √eps()) where {T<:Rate,N<:Rate}
    return isapprox(convert(b.compounding, a), b; atol, rtol)
end

"""
An AbstractYield is an object which can be used as an argument to:

- zero-coupon spot rates via [`zero`](@ref)
- discount factor via [`discount`](@ref)
- accumulation factor via [`accumulation`](@ref)

It can be be constructed via:

- zero rate curve with [`Zero`](@ref)
- forward rate curve with [`Forward`](@ref)
- par rate curve with [`Par`](@ref)
- typical OIS curve with [`OIS`](@ref)
- typical constant maturity treasury (CMT) curve with [`CMT`](@ref)
"""
abstract type AbstractYield end

# make interest curve broadcastable so that you can broadcast over multiple`time`s in `interest_rate`
Base.Broadcast.broadcastable(ic::T) where {T<:AbstractYield} = Ref(ic)

struct YieldCurve{T,U,V} <: AbstractYield
    rates::T
    maturities::U
    zero::V # function time -> continuous zero rate
end

# Forward curves

"""
    ForwardStarting(curve,forwardstart)

Rebase a `curve` so that `discount`/`accumulation`/etc. are re-based so that time zero from the new curves perspective is the given `forwardstart` time.

# Examples

```julia-repl
julia> zero = [5.0, 5.8, 6.4, 6.8] ./ 100
julia> maturity = [0.5, 1.0, 1.5, 2.0]
julia> curve = Yields.Zero(zero, maturity)
julia> fwd = Yields.ForwardStarting(curve, 1.0)

julia> discount(curve,1,2)
0.9275624570410582

julia> discount(fwd,1) # `curve` has effectively been reindexed to `1.0`
0.9275624570410582
```

# Extended Help

While `ForwardStarting` could be nested so that, e.g. the third period's curve is the one-period forward of the second period's curve, it will be more efficient to reuse the initial curve from a runtime and compiler perspective.

`ForwardStarting` is not used to construct a curve based on forward rates. See  [`Forward`](@ref) instead.
"""
struct ForwardStarting{T,U} <: AbstractYield
    curve::U
    forwardstart::T
end

function discount(c::ForwardStarting, to)
    discount(c.curve, c.forwardstart, to + c.forwardstart)
end

"""
    zero(curve,time)
    zero(curve,time,CompoundingFrequency)

Return the zero rate for the curve at the given time. If not specified, will use `Periodic(1)` compounding.
"""
Base.zero(c::YieldCurve, time) = convert(Periodic(1), Continuous(c.zero(time)))
function Base.zero(c::YieldCurve, time, cf::Periodic)
    z = Continuous(c.zero(time))
    i = convert(cf, z)
    return i
end

function Base.zero(c::YieldCurve, time, cf::Continuous)
    return Continuous(c.zero(time))
end

"""
    Constant(rate::Real, cf::CompoundingFrequency=Periodic(1))
    Constant(r::Rate)

Construct a yield object where the spot rate is constant for all maturities. If `rate` is not a `Rate` type, will assume `Periodic(1)` for the compounding frequency

# Examples

```julia-repl
julia> y = Yields.Constant(0.05)
julia> discount(y,2)
0.9070294784580498     # 1 / (1.05) ^ 2
```
"""
struct Constant{T} <: AbstractYield
    rate::T
end

function Constant(rate::T, cf::C = Periodic(1)) where {T<:Real,C<:CompoundingFrequency}
    return Constant(Rate(rate, cf))
end

Base.zero(c::Constant, time) = c.rate
Base.zero(c::Constant, time, cf::C) where {C<:CompoundingFrequency} = convert(cf, c.rate)
rate(c::Constant) = c.rate
rate(c::Constant, time) = c.rate
discount(c::T, time) where {T<:Real} = discount(Constant(c), time)
discount(r::Constant, time) = 1 / accumulation(r, time)

accumulation(r::Constant, time) = accumulation(r.rate.compounding, r, time)
accumulation(c::T, time) where {T>:Real} = accumulation(Constant(c), time)
accumulation(::Continuous, r::Constant, time) = exp(rate(r.rate) * time)
accumulation(::Periodic, r::Constant, time) = (1 + rate(r.rate) / r.rate.compounding.frequency)^(r.rate.compounding.frequency * time)

"""
    Step(rates,times)

Create a yield curve object where the applicable rate is the effective rate of interest applicable until corresponding time. If `rates` is not a `Vector{Rate}`, will assume `Periodic(1)` type.

# Examples

```julia-repl
julia>y = Yields.Step([0.02,0.05], [1,2])

julia>rate(y,0.5)
0.02

julia>rate(y,1.5)
0.05

julia>rate(y,2.5)
0.05
```
"""
struct Step <: AbstractYield
    rates
    times
end

Step(rates) = Step(rates, collect(1:length(rates)))

function rate(y::Step, time)
    i = findfirst(t -> time <= t, y.times)
    if isnothing(i)
        return y.rates[end]
    else
        return y.rates[i]
    end
end

function discount(y::Step, time)
    v = 1 / (1 + y.rates[1])^min(y.times[1], time)

    if y.times[1] >= time
        return v
    end

    for i = 2:length(y.times)

        if y.times[i] >= time
            # take partial discount and break
            v /= (1 + y.rates[i])^(time - y.times[i-1])
            break
        else
            # take full discount and continue
            v /= (1 + y.rates[i])^(y.times[i] - y.times[i-1])
        end

    end

    return v
end

"""
    Zero(rates,maturities)

Construct a yield curve with given zero-coupon spot `rates` at the given `maturities`. If `rates` is not a `Vector{Rate}`, will assume `Periodic(1)` type.
"""
function Zero(rates::Vector{<:Rate}, maturities)
    # bump to a constant yield if only given one rate
    length(rates) == 1 && return Constant(first(rates))

    continuous_zeros = rate.(convert.(Continuous(), rates))
    return YieldCurve(
        rates,
        maturities,
        linear_interp([0.0; maturities], [first(continuous_zeros); continuous_zeros])
    )
end

#fallback if `rates` aren't `Rate`s. Assume `Periodic(1)` per Zero docstring
function Zero(rates, maturities)
    Zero(Periodic.(rates, 1), maturities)
end

function Zero(rates)
    # bump to a constant yield if only given one rate
    maturities = collect(1:length(rates))
    return Zero(rates, maturities)
end

"""
    Par(rate, maturity)

Construct a curve given a set of bond equivalent yields and the corresponding maturities. Assumes that maturities <= 1 year do not pay coupons and that after one year, pays coupons with frequency equal to the CompoundingFrequency of the corresponding rate.

# Examples

```julia-repl

julia> par = [6.,8.,9.5,10.5,11.0,11.25,11.38,11.44,11.48,11.5] ./ 100
julia> maturities = [t for t in 1:10]
julia> curve = Par(par,maturities);
julia> zero(curve,1)
Rate(0.06000000000000005, Periodic(1))

```
"""
function Par(rates::Vector{<:Rate}, maturities)
    # bump to a constant yield if only given one rate
    if length(rates) == 1
        return Constant(rate[1])
    end
    return YieldCurve(
        rates,
        maturities,
        # assume that maturities less than or equal to 12 months are settled once, otherwise semi-annual
        # per Hull 4.7
        bootstrap(rates, maturities, [m <= 1 ? nothing : 1 / r.compounding.frequency for (r, m) in zip(rates, maturities)])
    )
end

function Par(rates::Vector{T}, maturities) where {T<:Real}
    return Par(Rate.(rates), maturities)
end

"""
    Forward(rate_vector,maturities)

Takes a vector of 1-period forward rates and constructs a discount curve.
"""
function Forward(rates, maturities)
    # convert to zeros and pass to Zero
    disc_v = similar(rates)
    v = 1.0

    for i = 1:length(rates)
        Δt = maturities[i] - (i == 1 ? 0 : maturities[i-1])
        v *= discount(Constant(rates[i]), Δt)
        disc_v[i] = v
    end

    z = (1.0 ./ disc_v) .^ (1 ./ maturities) .- 1 # convert disc_v to zero
    return Zero(z, maturities)
end

Forward(rates) = Forward(rates, collect(1:length(rates)))

"""
Takes CMT yields (bond equivalent), and assumes that instruments <= one year maturity pay no coupons and that the rest pay semi-annual.
"""
function CMT(rates::Vector{T}, maturities) where {T<:Real}
    rs = map(zip(rates, maturities)) do (r, m)
        if m <= 1
            Rate(r, Periodic(1 / m))
        else
            Rate(r, Periodic(2))
        end
    end

    CMT(rs, maturities)
end

function CMT(rates::Vector{<:Rate}, maturities)
    return YieldCurve(
        rates,
        maturities,
        # assume that maturities less than or equal to 12 months are settled once, otherwise semi-annual
        # per Hull 4.7
        bootstrap(rates, maturities, [m <= 1 ? nothing : 0.5 for m in maturities])
    )
end

"""
    OIS(rates,maturities)
Takes Overnight Index Swap rates, and assumes that instruments <= one year maturity are settled once and other agreements are settled quarterly with a corresponding CompoundingFrequency

"""
function OIS(rates::Vector{T}, maturities) where {T<:Real}
    rs = map(zip(rates, maturities)) do (r, m)
        if m <= 1
            Rate(r, Periodic(1 / m))
        else
            Rate(r, Periodic(4))
        end
    end

    return OIS(rs, maturities)
end
function OIS(rates::Vector{<:Rate}, maturities)
    return YieldCurve(
        rates,
        maturities,
        # assume that maturities less than or equal to 12 months are settled once, otherwise quarterly
        # per Hull 4.7
        bootstrap(rates, maturities, [m <= 1 ? nothing : 1 / 4 for m in maturities])
    )
end

abstract type ObservableQuote end

"""
    ZeroCouponQuote(price, maturity)

Quote for a set of zero coupon bonds with given `price` and `maturity`.

# Examples

```julia-repl
julia> prices = [1.3, 0.1, 4.5]
julia> maturities = [1.2, 2.5, 3.6]
julia> swq = Yields.ZeroCouponQuote.(prices, maturities)
```
"""
struct ZeroCouponQuote <: ObservableQuote
    price
    maturity
end

"""
    SwapQuote(yield, maturity, frequency)

Quote for a set of interest rate swaps with the given `yield` and `maturity` and a given payment `frequency`.

# Examples

```julia-repl
julia> maturities = [1.2, 2.5, 3.6]
julia> interests = [-0.02, 0.3, 0.04]
julia> prices = [1.3, 0.1, 4.5]
julia> frequencies = [2,1,2]
julia> swq = Yields.SwapQuote.(interests, maturities, frequencies)
```
"""
struct SwapQuote <: ObservableQuote
    yield
    maturity
    frequency
    function SwapQuote(yield, maturity, frequency)
        frequency <= 0 && throw(DomainError("Payment frequency must be positive"))
        return new(yield, maturity, frequency)
    end
end

"""
    BulletBondQuote(yield, price, maturity, frequency)

Quote for a set of fixed interest bullet bonds with given `yield`, `price`, `maturity` and a given payment frequency `frequency`.

Construct a vector of quotes for use with SmithWilson methods, e.g. by broadcasting over an array of inputs.

# Examples

```julia-repl
julia> maturities = [1.2, 2.5, 3.6]
julia> interests = [-0.02, 0.3, 0.04]
julia> prices = [1.3, 0.1, 4.5]
julia> frequencies = [2,1,2]
julia> bbq = Yields.BulletBondQuote.(interests, maturities, prices, frequencies)
```
"""
struct BulletBondQuote <: ObservableQuote
    yield
    price
    maturity
    frequency

    function BulletBondQuote(yield, maturity, price, frequency)
        frequency <= 0 && throw(DomainError("Payment frequency must be positive"))
        return new(yield, maturity, price, frequency)
    end
end

"""
    SmithWilson(zcq::Vector{ZeroCouponQuote}; ufr, α)
    SmithWilson(swq::Vector{SwapQuote}; ufr, α)
    SmithWilson(bbq::Vector{BulletBondQuote}; ufr, α)
    SmithWilson(times<:AbstractVector, cashflows<:AbstractMatrix, prices<:AbstractVector; ufr, α)
    SmithWilson(u, qb; ufr, α)

Create a yield curve object that implements the Smith-Wilson interpolation/extrapolation scheme.

Positional arguments to construct a curve:

- Quoted instrument as the first argument: either a `Vector` of `ZeroCouponQuote`s, `SwapQuote`s, or `BulletBondQuote`s, or
- A set of `times`, `cashflows`, and `prices`, or
- A curve can be with `u` is the timepoints coming from the calibration, and `qb` is the internal parameterization of the curve that ensures that the calibration is correct. Users may prefer the other constructors but this mathematical constructor is also available.

Required keyword arguments:

- `ufr` is the Ultimate Forward Rate, the forward interest rate to which the yield curve tends, in continuous compounding convention.
- `α` is the parameter that governs the speed of convergence towards the Ultimate Forward Rate. It can be typed with `\\alpha[TAB]`
"""
struct SmithWilson{TU<:AbstractVector,TQb<:AbstractVector} <: AbstractYield
    u::TU
    qb::TQb
    ufr
    α

    # Inner constructor ensures that vector lengths match
    function SmithWilson{TU,TQb}(u, qb; ufr, α) where {TU<:AbstractVector,TQb<:AbstractVector}
        if length(u) != length(qb)
            throw(DomainError("Vectors u and qb in SmithWilson must have equal length"))
        end
        return new(u, qb, ufr, α)
    end
end

SmithWilson(u::TU, qb::TQb; ufr, α) where {TU<:AbstractVector,TQb<:AbstractVector} = SmithWilson{TU,TQb}(u, qb; ufr = ufr, α = α)

"""
    H_ordered(α, t_min, t_max)

The Smith-Wilson H function with ordered arguments (for better performance than using min and max).
"""
function H_ordered(α, t_min, t_max)
    return α * t_min + exp(-α * t_max) * sinh(-α * t_min)
end

"""
    H(α, t1, t2)

The Smith-Wilson H function implemented in a faster way.
"""
function H(α, t1::T, t2::T) where {T}
    return t1 < t2 ? H_ordered(α, t1, t2) : H_ordered(α, t2, t1)
end

H(α, t1, t2) = H(α, promote(t1, t2)...)

H(α, t1vec::AbstractVector, t2) = [H(α, t1, t2) for t1 in t1vec]
H(α, t1vec::AbstractVector, t2vec::AbstractVector) = [H(α, t1, t2) for t1 in t1vec, t2 in t2vec]
# This can be optimized by going to H_ordered directly, but it might be a bit cumbersome
H(α, tvec::AbstractVector) = H(α, tvec, tvec)

discount(sw::SmithWilson, t) = exp(-sw.ufr * t) * (1.0 + H(sw.α, sw.u, t) ⋅ sw.qb)
Base.zero(sw::SmithWilson, t) = Continuous(sw.ufr - log(1.0 + H(sw.α, sw.u, t) ⋅ sw.qb) / t)
Base.zero(sw::SmithWilson, t, cf::CompoundingFrequency) = convert(cf, zero(sw, t))

function SmithWilson(times::AbstractVector, cashflows::AbstractMatrix, prices::AbstractVector; ufr, α)
    Q = Diagonal(exp.(-ufr * times)) * cashflows
    q = vec(sum(Q, dims = 1))  # We want q to be a column vector
    QHQ = Q' * H(α, times) * Q
    b = QHQ \ (prices - q)
    Qb = Q * b
    return SmithWilson(times, Qb; ufr = ufr, α = α)
end

"""
    timepoints(zcq::Vector{ZeroCouponQuote})
    timepoints(bbq::Vector{BulletBondQuote})

Return the times associated with the `cashflows` of the instruments.
"""
function timepoints(qs::Vector{Q}) where {Q<:ObservableQuote}
    frequency = maximum(q.frequency for q in qs)
    timestep = 1 / frequency
    maturity = maximum(q.maturity for q in qs)
    return [timestep:timestep:maturity...]
end

"""
    cashflows(interests, maturities, frequency)
    timepoints(zcq::Vector{ZeroCouponQuote})
    timepoints(bbq::Vector{BulletBondQuote})

Produce a cash flow matrix for a set of instruments with given `interests` and `maturities`
and a given payment frequency `frequency`. All instruments are assumed to have their first payment at time 1/`frequency`
and have their last payment at the largest multiple of 1/`frequency` less than or equal to the input maturity.
"""
function cashflows(interests, maturities, frequencies)
    frequency = lcm(frequencies)
    fq = inv.(frequencies)
    timestep = 1 / frequency
    floored_mats = floor.(maturities ./ timestep) .* timestep
    times = timestep:timestep:maximum(floored_mats)
    # we need to determine the coupons in relation to the payment date, not time zero
    time_adj = floored_mats .% fq

    cashflows = [
        # if on a coupon date and less than maturity, pay coupon
        ((((t + time_adj[instrument]) % fq[instrument] ≈ 0) && t <= floored_mats[instrument]) ? interests[instrument] / frequencies[instrument] : 0.0) +
        (t ≈ floored_mats[instrument] ? 1.0 : 0.0) # add maturity payment
        for t in times, instrument = 1:length(interests)
    ]

    return cashflows
end

function cashflows(qs::Vector{Q}) where {Q<:ObservableQuote}
    yield = [q.yield for q in qs]
    maturity = [q.maturity for q in qs]
    frequency = [q.frequency for q in qs]
    return cashflows(yield, maturity, frequency)
end

# Utility methods for calibrating Smith-Wilson directly from quotes
function SmithWilson(zcq::Vector{ZeroCouponQuote}; ufr, α)
    n = length(zcq)
    maturities = [q.maturity for q in zcq]
    prices = [q.price for q in zcq]
    return SmithWilson(maturities, Matrix{Float64}(I, n, n), prices; ufr = ufr, α = α)
end

function SmithWilson(swq::Vector{SwapQuote}; ufr, α)
    times = timepoints(swq)
    cfs = cashflows(swq)
    ones(length(swq))
    return SmithWilson(times, cfs, ones(length(swq)), ufr = ufr, α = α)
end

function SmithWilson(bbq::Vector{BulletBondQuote}; ufr, α)
    times = timepoints(bbq)
    cfs = cashflows(bbq)
    prices = [q.price for q in bbq]
    return SmithWilson(times, cfs, prices, ufr = ufr, α = α)
end

# https://github.com/dpsanders/hands_on_julia/blob/master/during_sessions/Fractale%20de%20Newton.ipynb
newton(f, f′, x) = x - f(x) / f′(x)
function solve(g, g′, x0, max_iterations = 100)
    x = x0

    tolerance = 2 * eps(x0)
    iteration = 0

    while (abs(g(x) - 0) > tolerance && iteration < max_iterations)
        x = newton(g, g′, x)
        iteration += 1
    end

    return x
end

function bootstrap(rates, maturities, settlement_frequency; interp_function = linear_interp)
    discount_vec = zeros(length(rates)) # construct a placeholder discount vector matching maturities
    # we have to take the first rate as the starting point
    discount_vec[1] = discount(Constant(rates[1]), maturities[1])

    for t = 2:length(maturities)
        if isnothing(settlement_frequency[t])
            # no settlement before maturity
            discount_vec[t] = discount(Constant(rates[t]), maturities[t])
        else
            # need to account for the interim cashflows settled
            times = settlement_frequency[t]:settlement_frequency[t]:maturities[t]
            cfs = [rate(rates[t]) * settlement_frequency[t] for s in times]
            cfs[end] += 1

            function pv(v_guess)
                v = interp_function([[0.0]; maturities[1:t]], vcat(1.0, discount_vec[1:t-1], v_guess...))
                return sum(v.(times) .* cfs)
            end
            target_pv = sum(map(t2 -> discount(Constant(rates[t]), t2), times) .* cfs)
            root_func(v_guess) = pv(v_guess) - target_pv
            root_func′(v_guess) = ForwardDiff.derivative(root_func, v_guess)
            discount_vec[t] = solve(root_func, root_func′, rate(rates[t]))
        end

    end
    zero_vec = -log.(discount_vec) ./ maturities
    return linear_interp([0.0; maturities], [first(zero_vec); zero_vec])
end

## Generic and Fallbacks
"""
    discount(rate,to)
    discount(rate,from,to)

The discount factor for the `rate` for times `from` through `to`. If rate is a `Real` number, will assume a `Constant` interest rate.
"""
discount(yc, time) = exp(-yc.zero(time) * time)
discount(rate::Rate{<:Real,<:CompoundingFrequency}, from, to) = discount(Constant(rate), from, to)
discount(rate::Rate{<:Real,<:CompoundingFrequency}, to) = discount(Constant(rate), to)

discount(yc, from, to) = discount(yc, to) / discount(yc, from)

"""
    forward(curve,from,to,CompoundingFrequency=Periodic(1))

The forward `Rate` implied by the curve between times `from` and `to`.
"""
function forward(yc, from, to)
    return forward(yc, from, to, Periodic(1))
end

function forward(yc, from, to, cf::T) where {T<:CompoundingFrequency}

    r = Periodic((accumulation(yc, to) / accumulation(yc, from))^(1 / (to - from)) - 1, 1)
    return convert(cf, r)
end

function forward(yc, from)
    to = from + 1
    return forward(yc, from, to)
end

"""
    accumulation(rate,from,to)

The accumulation factor for the `rate` for times `from` through `to`. If rate is a `Real` number, will assume a `Constant` interest rate.
"""
function accumulation(y::T, time) where {T<:AbstractYield}
    return 1 ./ discount(y, time)
end
accumulation(rate::Rate{<:Real,<:CompoundingFrequency}, to) = accumulation(Constant(rate), to)

function accumulation(y::T, from, to) where {T<:AbstractYield}
    return 1 ./ discount(y, from, to)
end
accumulation(rate::Rate{<:Real,<:CompoundingFrequency}, from, to) = accumulation(Constant(rate), from, to)

## Curve Manipulations
struct RateCombination{T,U,V} <: AbstractYield
    r1::T
    r2::U
    op::V
end

rate(rc::RateCombination, time) = rc.op(rate(rc.r1, time), rate(rc.r2, time))
function discount(rc::RateCombination, time)
    a1 = discount(rc.r1, time)^(-1 / time) - 1
    a2 = discount(rc.r2, time)^(-1 / time) - 1
    return 1 / (1 + rc.op(a1, a2))^time
end

"""
    Yields.AbstractYield + Yields.AbstractYield

The addition of two yields will create a `RateCombination`. For `rate`, `discount`, and `accumulation` purposes the spot rates of the two curves will be added together.
"""
function Base.:+(a::AbstractYield, b::AbstractYield)
    return RateCombination(a, b, +)
end

function Base.:+(a::Constant, b::Constant)
    a_kind = rate(a).compounding
    rate_new_basis = rate(convert(a_kind, rate(b)))
    return Constant(
        Rate(
            rate(a.rate) + rate_new_basis,
            a_kind
        )
    )
end

function Base.:+(a::T, b) where {T<:AbstractYield}
    return a + Constant(b)
end

function Base.:+(a, b::T) where {T<:AbstractYield}
    return Constant(a) + b
end

"""
    Yields.AbstractYield - Yields.AbstractYield

The subtraction of two yields will create a `RateCombination`. For `rate`, `discount`, and `accumulation` purposes the spot rates of the second curves will be subtracted from the first.
"""
function Base.:-(a::AbstractYield, b::AbstractYield)
    return RateCombination(a, b, -)
end

function Base.:-(a::Constant, b::Constant)
    a_kind = rate(a).compounding
    rate_new_basis = rate(convert(a_kind, rate(b)))
    return Constant(
        Rate(
            rate(a.rate) - rate_new_basis,
            a_kind
        )
    )
end

function Base.:-(a::T, b) where {T<:AbstractYield}
    return a - Constant(b)
end

function Base.:-(a, b::T) where {T<:AbstractYield}
    return Constant(a) - b
end

linear_interp(xs, ys) = Interpolations.extrapolate(
    Interpolations.interpolate((xs,), ys, Interpolations.Gridded(Interpolations.Linear())),
    Interpolations.Line()
)

# used to display simple type name in show method
# https://stackoverflow.com/questions/70043313/get-simple-name-of-type-in-julia?noredirect=1#comment123823820_70043313
name(::Type{T}) where {T} = (isempty(T.parameters) ? T : T.name.wrapper)

function Base.show(io::IO, curve::T) where {T<:AbstractYield}
    println() # blank line for padding
    r = zero(curve, 1)
    ylabel = isa(r.compounding, Continuous) ? "Continuous" : "Periodic($(r.compounding.frequency))"
    kind = name(typeof(curve))
    l = lineplot(
        t -> rate(zero(curve, t)),
        0.0, #from
        30.0,  # to
        xlabel = "time",
        ylabel = ylabel,
        compact = true,
        name = "Zero rates",
        width = 60,
        title = "Yield Curve ($kind)"
    )
    show(io, l)
end

end


# Tests for conditional entropy function

# Test with 1 values array conditioned on itself
@test get_conditional_entropy(arr1, arr1) ≈ 0 atol = 0.05
println("Conditional entropy with 1 array passed.")

# Test conditional entropy formula
@test get_conditional_entropy(arr1, arr2) ≈ get_entropy(arr1, arr2) - get_entropy(arr2) atol = 0.05
println("Conditional entropy formula passed.")


using CayleyDickson
using Base.Test: @test, @test_throws

@test begin
    a = CayleyDickson.randomBigFloat()
    isreal(SplitOctonion(a))
end

@test begin
    a = 1
    !isreal(SplitOctonion(a, a, a, a))
end

@test begin
    z = random(SplitOctonion{Int})
    z == +(z)
end

@test begin
    a = CayleyDickson.randomBigFloat()
    real(SplitOctonion(a)) == a
end

@test begin
    io = IOBuffer()
    show(io, SplitOctonion(1.0, 2, 3, 4, 5, 6, 7, 8))
    l = String(take!(io))
    r = "[1: 1.0, i: 2.0, j: 3.0, ij: 4.0, s: 5.0, is: 6.0, js: 7.0, (ij)s: 8.0]"
    l == r
end

@test begin
    io = IOBuffer()
    show(io, SplitOctonion(1.0, 2, 3, 4, 5, 6, 7))
    l = String(take!(io))
    r = "[1: 1.0, i: 2.0, j: 3.0, ij: 4.0, s: 5.0, is: 6.0, js: 7.0, (ij)s: 0.0]"
    l == r
end

@test begin
    io = IOBuffer()
    show(io, SplitOctonion(1.0, 2, 3, 4, 5, 6))
    l = String(take!(io))
    r = "[1: 1.0, i: 2.0, j: 3.0, ij: 4.0, s: 5.0, is: 6.0, js: 0.0, (ij)s: 0.0]"
    l == r
end

@test begin
    io = IOBuffer()
    show(io, SplitOctonion(1.0, 2, 3, 4, 5))
    l = String(take!(io))
    r = "[1: 1.0, i: 2.0, j: 3.0, ij: 4.0, s: 5.0, is: 0.0, js: 0.0, (ij)s: 0.0]"
    l == r
end

@test begin
    io = IOBuffer()
    show(io, SplitOctonion(1.0, 2, 3, 4))
    l = String(take!(io))
    r = "[1: 1.0, i: 2.0, j: 3.0, ij: 4.0, s: 0.0, is: 0.0, js: 0.0, (ij)s: 0.0]"
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    b = CayleyDickson.randomBigFloat()
    c = CayleyDickson.randomBigInt()
    d = CayleyDickson.randomBigFloat()
    l = SplitOctonion(Quaternion(a, b, c, d))
    r = SplitOctonion(a, b, c, d)
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    b = CayleyDickson.randomBigFloat()
    l = SplitOctonion(Binion(a, b))
    r = SplitOctonion(a, b)
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    b = CayleyDickson.randomBigFloat()
    c = CayleyDickson.randomBigInt()
    l = SplitOctonion(a, b, c)
    r = SplitOctonion(a, b, c, 0)
    l == r
end

@test begin
    x = random(SplitOctonion{BigFloat})
    y = random(SplitOctonion{BigFloat})
    x != y
end

@test begin
    a = CayleyDickson.randomBigInt()
    z = random(SplitOctonion{BigInt})
    l = a + z
    r = z + a
    l == r
end

@test begin
    a = random(Quaternion{BigInt})
    z = random(SplitOctonion{BigInt})
    l = a + z
    r = z + a
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    z = random(SplitOctonion{BigInt})
    l = a - z
    r = -(z - a)
    l == r
end

@test begin
    a = random(Quaternion{BigInt})
    z = random(SplitOctonion{BigInt})
    l = a - z
    r = -(z - a)
    l == r
end

@test begin
    x = random(SplitOctonion{BigInt})
    y = random(SplitOctonion{BigInt})
    l = x + y
    r = y + x
    l == r
end

@test begin
    x = random(SplitOctonion{BigInt})
    y = random(SplitOctonion{BigInt})
    l = x * y
    r = y * x
    l != r
end

@test begin
    x = random(SplitOctonion{BigInt})
    y = random(SplitOctonion{BigInt})
    l = x - y
    r = -(y - x)
    l == r
end

@test begin
    x = random(SplitOctonion{BigInt})
    y = random(SplitOctonion{BigInt})
    z = random(SplitOctonion{BigInt})
    l = (x + y) + z
    r = x + (y + z)
    l == r
end

@test begin
    x = random(SplitOctonion{BigInt})
    y = random(SplitOctonion{BigInt})
    z = random(SplitOctonion{BigInt})
    l = (x * y) * z
    r = x * (y * z)
    l != r
end

@test begin
    x = random(SplitOctonion{BigInt})
    y = random(SplitOctonion{BigInt})
    l = alternatorL(x, y)
    r = zero(SplitOctonion{BigInt})
    l == r
end

@test begin
    x = random(SplitOctonion{BigInt})
    y = random(SplitOctonion{BigInt})
    l = alternatorR(x, y)
    r = zero(SplitOctonion{BigInt})
    l == r
end

@test begin
    x = random(SplitOctonion{BigInt})
    conj(conj(x)) == x
end

@test begin
    x = random(SplitOctonion{Rational{BigInt}})
    inv(inv(x)) == x
end

@test begin
    x = random(SplitOctonion{BigInt})
    y = random(SplitOctonion{BigInt})
    l = conj(x * y)
    r = conj(y) * conj(x)
    l == r
end

@test begin
    x = random(SplitOctonion{BigInt})
    y = random(SplitOctonion{BigInt})
    l = conj(x + y)
    r = conj(x) + conj(y)
    l == r
end

@test begin
    x = random(SplitOctonion{BigInt})
    y = random(SplitOctonion{BigInt})
    l = conj(x - y)
    r = conj(x) - conj(y)
    l == r
end

@test begin
    x = random(SplitOctonion{Rational{BigInt}})
    y = random(SplitOctonion{Rational{BigInt}})
    l = inv(x * y)
    r = inv(y) * inv(x)
    l == r
end

@test begin
    x = random(SplitOctonion{Rational{BigInt}})
    y = random(SplitOctonion{Rational{BigInt}})
    l = inv(x / y)
    r = y / x
    l == r
end

@test begin
    x = random(SplitOctonion{Rational{BigInt}})
    y = random(SplitOctonion{Rational{BigInt}})
    l = inv(x \ y)
    r = y \ x
    l == r
end

@test begin
    x = random(SplitOctonion{BigInt})
    y = random(SplitOctonion{BigInt})
    z = random(SplitOctonion{BigInt})
    l = x * (y + z)
    r = (x * y) + (x * z)
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    x = random(SplitOctonion{BigInt})
    y = random(SplitOctonion{BigInt})
    l = a * (x + y)
    r = (a * x) + (a * y)
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    x = random(SplitOctonion{BigInt})
    l = a * x
    r = x * a
    l == r
end

@test begin
    a = random(Quaternion{BigInt})
    x = random(SplitOctonion{BigInt})
    l = commutator(a, x)
    r = zero(SplitOctonion{BigInt})
    l != r
end

@test begin
    x = random(SplitOctonion{BigInt})
    y = random(SplitOctonion{BigInt})
    z = random(SplitOctonion{BigInt})
    l = (x + y) * z
    r = (x * z) + (y * z)
    l == r
end

@test begin
    x = random(SplitOctonion{BigInt})
    y = random(SplitOctonion{BigInt})
    z = random(SplitOctonion{BigInt})
    l = x * (y - z)
    r = (x * y) - (x * z)
    l == r
end

@test begin
    x = random(SplitOctonion{BigInt})
    y = random(SplitOctonion{BigInt})
    z = random(SplitOctonion{BigInt})
    l = (x - y) * z
    r = (x * z) - (y * z)
    l == r
end

@test begin
    x = random(SplitOctonion{BigInt})
    quadrance(x) > big(0) || quadrance(x) < big(0) || quadrance(x) == big(0)
end

@test begin
    x = random(SplitOctonion{BigInt})
    y = random(SplitOctonion{BigInt})
    l = quadrance(x * y)
    r = quadrance(x) * quadrance(y)
    l == r
end

@test_throws ErrorException begin
    inv(SplitOctonion(0, 0.0))
end

@test_throws ErrorException begin
    random(SplitOctonion{Int}) / 0
end

@test_throws ErrorException begin
    0 \ random(SplitOctonion{Int})
end

@test begin
    a = CayleyDickson.randomBigRational()
    z = random(SplitOctonion{Rational{BigInt}})
    l = z / a
    r = a \ z
    l == r
end

@test begin
    a = CayleyDickson.randomBigRational()
    z = random(SplitOctonion{Rational{BigInt}})
    l = a / z
    r = z \ a
    l == r
end

@test begin
    z = random(SplitOctonion{BigInt})
    z == z + zero(SplitOctonion{BigInt})
end

@test begin
    z = random(SplitOctonion{BigInt})
    z == z * one(SplitOctonion{BigInt})
end

@test begin
    z = random(SplitOctonion{BigInt})
    z == z + zero(z)
end

@test begin
    z = random(SplitOctonion{BigInt})
    z == z * one(z)
end


type FractalDict
	data::Dict{AbstractString,Function}
end
FractalDict()=FractalDict(Dict{AbstractString,Function}())

const frac_dict=FractalDict()
function add_frac!(fracname::AbstractString,f::Function)
	frac_dict.data[fracname]=f
end

fractal_names() = keys(frac_dict.data)
function create_fractal(fractal_name::AbstractString)
	if !(fractal_name in keys(frac_dict.data))
		error("Unknown fractal name: $fractal_name")
	end
	return frac_dict.data[fractal_name]
end

@inline num_params(f::Function) = length(Base.uncompressed_ast(f.code).args[2][1])

"""
    @fractal name inner_function
    @fractal name zinit inner_function

Create a fractal variation and add it to the catalog.

The first macro invocation is used when the complex-valued pixel location is intended to be used
as the initial 'z' value for the fractal iteration (as in Julia sets). In this case,
`inner_function` should not take any parameters.

The second macro invocation is used when the complex-valued pixel location is intended to be used
as an extra parameter to the fractal iteration and an initial 'z' value is needed (as in Mandelbrot
sets).  In this case, `zinit` is the initial 'z' value, and `inner_function` should take a
complex-valued parameter representing the pixel location.
"""
macro fractal(fracname,zinit_func...)
	frac_name=string(fracname)
	closure_name=symbol(frac_name*"_closure")
	if length(zinit_func)==1
		func=zinit_func[1]
		@assert typeof(eval(func))==Function
		@assert num_params(eval(func))==0
		return quote
			@inline function $(esc(closure_name)){T}(param::T)
				@fastmath @inline f(z::Complex)=$func()
				return f,param
			end
			add_frac!($fracname,$(esc(closure_name)))
		end
	elseif length(zinit_func)==2
		zinit=zinit_func[1]
		func=zinit_func[2]
		@assert typeof(eval(func))==Function
		@assert num_params(eval(func))==1
		return quote
			@inline function $(esc(closure_name)){T}(param::T)
				@fastmath @inline f(z::Complex)=$func(param)
				return f,$zinit
			end
			add_frac!($fracname,$(esc(closure_name)))
		end
	end

end

const czero=zero(Complex)
const chalf=Complex(0.5,0.0)
const cone=Complex(1.0,0.0)
const phi=(1.0+sqrt(5.0))/2.0

@fractal "mandelbrot" czero (c::Complex)->z*z+c
@fractal "mandelbrot_minvc" czero (c::Complex)->z*z-inv(c)
@fractal "mandelbrot_minvc_p25" czero (c::Complex)->z*z-inv(c)+0.25
@fractal "mandelbrot_minvc_m25" czero (c::Complex)->z*z-inv(c)-0.25
@fractal "mandelbrot_minvc_m50" czero (c::Complex)->z*z-inv(c)-0.50
@fractal "mandelbrot_minvc_m75" czero (c::Complex)->z*z-inv(c)-0.75
@fractal "mandelbrot_minvc_mmyreberg" czero (c::Complex)->z*z-inv(c)-1.40115
@fractal "mandelbrot_minvc_m2" czero (c::Complex)->z*z-inv(c)-2.0

@fractal "mandelbrot_lambda" chalf (lambda::Complex)->z*(1.0-z)*lambda
@fractal "mandelbrot_invlambda" chalf (lambda::Complex)->z*(1.0-z)*inv(lambda)
@fractal "mandelbrot_invlambda_m1" chalf (lambda::Complex)->z*(1.0-z)*(inv(lambda)+cone)

@fractal "burningship" czero (c::Complex)->Complex(abs(real(z)),abs(imag(z)))^2-c

@fractal "julia_z2_1mphi" ()->z*z+Complex(1.0-phi,0.0)
@fractal "julia_z2_phim1_phim1" ()->z*z+Complex(phi-2.0,phi-1.0)
@fractal "julia_z2_285" ()->z*z+Complex(0.285,0.0)
@fractal "julia_z2_285_01" ()->z*z+Complex(0.285,0.01)
@fractal "julia_z2_n8_156" ()->z*z+Complex(-0.8,0.156)
@fractal "julia_z2_n7629_1889" ()->z*z+Complex(-0.7269,0.1889)

@fractal "julia_expz_n65" ()->exp(z)+Complex(-0.65,0.0)
@fractal "julia_expz3_n59" ()->exp(z^3)+Complex(-0.59,0.0)
@fractal "julia_expz3_n621" ()->exp(z^3)+Complex(-0.621,0.0)

@fractal "julia_sqrtsinhz2" ()->sqrt(sinh(z^2))
@fractal "julia_sqrtsinhz2_065_122" ()->sqrt(sinh(z^2))+Complex(0.065,0.122)

@fractal "tricorn" czero (c::Complex)->z'*z'+c


#=
Visualisation of a Galton Board, also known as a Bean Machine or quincunx (https://en.wikipedia.org/wiki/Bean_machine).

Invented by Francis Galton. the Galton Board is used to demonstrate how the central limit theorem works.
Specifically, it is used to show that for_ a sufficiently large sample size, the binomial distribution
will approximate the normal distribution.
=#

using Javis
using Colors
using Random

"""
    ground(args...)

Define the background and the initial "foreground" color.
"""
function ground(args...)
    background("black")
    sethue("white")
end

"""
    draw_separators(n, xpos, ypos, offset, height; color = "white")

Draws `n` vertical lines of length `height`, each spaced `offset` pixels apart. `xpos` represents the
x-coordinate of the leftmost line, and `ypos` represents the y-coordinate of the upper point used to
construct each line.
"""
function draw_separators(n, xpos, ypos, offset, height; color = "white")
    for i in 1:n
        x = xpos + offset * (i - 1)
        Object(JLine(Point(x, ypos), Point(x, ypos + height); color = color))
    end
end

"""
    draw_pins(n, xpos, ypos, offset; color = "white")

Draws `n` circles of radius 3, each spaced `offset` pixels apart. `xpos` represents the x-coordinate of
the leftmost circle, and `ypos` represents the y-coordinate of the centre of each circle.
"""
function draw_pins(n, xpos, ypos, offset; color = "white")
    for i in 1:n
        Object(
            JCircle(Point(xpos + offset * (i - 1), ypos), 3; color = color, action = :fill),
        )
    end
end

"""
    flip_coin(p)

Simulates flipping a coin, returning 1 with probability `p` and 1 with probability `(1 - p)`. Helper function
used in `move_ball`.
"""
function flip_coin(p = 0.5)
    return rand() < p ? 1 : -1
end

"""
    move_ball(ball, first, last, offset, p, fno)

Moves `ball` from the `first` level of pins to the `last` level of pins (i.e. top to bottom). `offset` should
be the same offset used when creating the bins using `draw_separators`. The value `p` is the probability that
a ball will go right at each pin, and the value `fno` is the frame number from which the ball should start moving.

Returns (fno, pos), where `fno` represents the last frame of the ball when it has reached it's destination,
and pos represents the "sum" of the direction movements. For example, a ball that moved 5 times left, and
7 times right would have a "sum" of -5 + 7 = 2.
"""
function move_ball(ball, first, last, offset, p, fno)
    pos = 0
    act!(ball, Action((fno + 1):(fno + 5), anim_translate(0, offset / 2)))
    fno += 5
    for i in first:last
        direction = flip_coin(p)
        pos += direction
        act!(ball, Action((fno + 1):(fno + 5), anim_translate(direction * offset / 2, 0)))
        fno += 5
        act!(ball, Action((fno + 1):(fno + 5), anim_translate(0, offset)))
        fno += 5
    end
    return (fno, pos)
end

"""
    galton(seed)

Create the Galton Board animation and saves it in the file `galton.gif`. `seed` is used to determine the seed
value for the pseudorandom number generation, and is used for reproducability.
"""
function galton(seed, p)
    nframes = 1200
    video = Video(250, 600)
    Background(1:nframes, ground)

    logocolors = Colors.JULIA_LOGO_COLORS # Use the Julia logo colors
    Random.seed!(seed)

    # Parameters
    n = 15
    xpos = -100
    ypos = -75
    width = 200
    gap = 15
    height = 350
    offset = width / (n - 1)
    radius = offset / 4

    # Draws the line separators or "bins"
    draw_separators(n, xpos, ypos + gap, offset, height)

    # Create "base" of the board
    Object(
        JRect(
            Point(xpos - offset, ypos + gap + height),
            width + 2 * offset,
            7;
            color = logocolors.purple,
            action = :fill,
        ),
    )

    # Draws the pins starting from the bottom most level (closest to base)
    for level in n:-1:1
        draw_pins(level, xpos, ypos, offset; color = logocolors.blue)
        xpos += offset / 2
        ypos -= offset
    end

    nballs = 200
    colors = range(logocolors.green, logocolors.red, length = nballs)
    # Frequency table (i.e. bins), used to determine how much to translate each ball.
    # Total of n + 1 bins.
    slots = zeros(Int, n + 1)
    for i in 1:nballs
        # Set start frame of each ball, interval of 5 frames between each ball
        fno = 5 * i
        # Start each ball at the topmost level
        ball = Object(JCircle(Point(0, ypos), radius; color = colors[i], action = :fill))
        # Move the ball to the top of a bin
        fno_end, pos = move_ball(ball, 1, n, offset, p, fno)

        # Map from "sum" of directions to bin number
        index = (pos + n) ÷ 2 + 1
        # Move each ball into it's "bin"
        act!(
            ball,
            Action(
                (fno_end + 1):(fno_end + 10),
                anim_translate(0, height - 2 * radius * slots[index]),
            ),
        )
        # Update frequency table
        slots[index] += 1
    end

    render(video; pathname = "galton.gif", framerate = 30)
    return slots
end

# Generate the "galton.gif" file.
galton(94, 0.5)


function embed(data::DataFrame;path::AbstractString="",xdim::Int64=20,ydim::Int64=20,perplexity=10,maxIter=10000,eta=200.0)
	# Kratochvíl M, Koladiya A and Vondrášek J. Generalized EmbedSOM on quadtree-structured self-organizing maps. F1000Research 2020, 8:2120

	imputed = Impute.srs(data)
	disallowmissing!(imputed)

	if isfile(path)
		@info """Using embedding saved at $path"""
		som = deserialize(path)

	else
		@info """Calculating embedding..."""
		som = initGigaSOM(imputed,xdim,ydim)
		som = trainGigaSOM(som,imputed)
		~isempty(path) && serialize(path,som)
	end

	######################## extract clusters and embedding
	try
		clusters = mapToGigaSOM(som,imputed)
		embedding = embedGigaSOM(som,imputed)

		rename!(clusters,"index"=>"cluster")
		return clusters, Matrix{Float64}(embedding')
	catch

		@info """Re-calculating embedding..."""
		som = initGigaSOM(imputed,xdim,ydim)
		som = trainGigaSOM(som,imputed)
		~isempty(path) && serialize(path,som)

		clusters = mapToGigaSOM(som,imputed)
		embedding = embedGigaSOM(som,imputed)

		rename!(clusters,"index"=>"cluster")
		return clusters, Matrix{Float64}(embedding')
	end
end


@doc raw"""
    dominoes(imgname, colr, constrnt)

Approximate a grayscale image with complete sets of double-nine dominoes.

# Resources
Watch Bob Bosch of Oberlin College talk at Google about domino mosaics, and his interpretation of this art form as an integer program.
https://www.youtube.com/watch?v=g3CiVrN-BnY

IP modeling details can be found in this article.
http://www.optimization-online.org/DB_FILE/2003/09/722.pdf

# Examples
```julia-repl
julia> include("ModuleDominoes.jl")
julia> using .Dominoes
julia> createMosaic("Frankenstein", "white", "constrained")
```

# Author
Kalyani Nagaraj
March 2021
"""
module Dominoes

using JuMP
using GLPK
using DelimitedFiles
using MAT
using ImageView, Images, FileIO

include("dominoes.jl")
include("createMosaic.jl")

export createMosaic

end


"""
    OlsSureFn(Y,X,NWQ=false,m=0)

LS of Y on X; where Y is Txn, and X is the same for all regressions

# Usage
(b,res,Yhat,Covb,R2) = OlsSureFn(Y,X,NWQ,m)

# Input
- `Y::Matrix`:     Txn, the n dependent variables
- `X::Matrix`:     Txk matrix of regressors (including deterministic ones)
- `NWQ:Bool`:      if true, then Newey-West's covariance matrix is used, otherwise Gauss-Markov
- `m::Int`:        scalar, bandwidth in Newey-West

# Output
- `b::Matrix`:     n*kx1, regression coefficients
- `u::Matrix`:     Txn, residuals Y - Yhat
- `Yhat::Matrix`:  Txn, fitted values X*b
- `V::Matrix`:     covariance matrix of vec(b)
- `R2::Vector`:    n-vector, R2 values

"""
function OlsSureFn(Y,X,NWQ=false,m=0)

    (T,n) = (size(Y,1),size(Y,2))
    k     = size(X,2)

    b     = X\Y
    Yhat  = X*b
    u     = Y - Yhat

    g     = zeros(T,n*k)
    for i = 1:n
      vv      = (1+(i-1)*k):(i*k)   #1:k,(1+k):2k,...
      g[:,vv] = X.*u[:,i]           #moment conditions for Y[:,i] regression
    end

    Sxx = X'X
    if NWQ
        S     = CovNWFn(g,m)            #Newey-West covariance matrix
        Sxx_1 = kron(I(n),inv(Sxx))
        V     = Sxx_1 * S * Sxx_1
    else
        V = kron(cov(u),inv(Sxx))      #traditional covariance matrix, Gauss-Markov
    end

    R2   = 1 .- var(u,dims=1)./var(Y,dims=1)

    return b, u, Yhat, V, R2

end
#------------------------------------------------------------------------------


# Arnoldi/Lanczos iteration algorithms (with custom IOP)

#######################################
# Output type/cache
"""
    KrylovSubspace{T}(n,[maxiter=30]) -> Ks

Constructs an uninitialized Krylov subspace, which can be filled by `arnoldi!`.

The dimension of the subspace, `Ks.m`, can be dynamically altered but should
be smaller than `maxiter`, the maximum allowed arnoldi iterations.

    getV(Ks) -> V
    getH(Ks) -> H

Access methods for the (extended) orthonormal basis `V` and the (extended)
Gram-Schmidt coefficients `H`. Both methods return a view into the storage
arrays and has the correct dimensions as indicated by `Ks.m`.

    resize!(Ks, maxiter) -> Ks

Resize `Ks` to a different `maxiter`, destroying its contents.

This is an expensive operation and should be used scarcely.
"""
mutable struct KrylovSubspace{B, T, U}
    m::Int        # subspace dimension
    maxiter::Int  # maximum allowed subspace size
    beta::B       # norm(b,2)
    V::Matrix{T}  # orthonormal bases
    H::Matrix{U}  # Gram-Schmidt coefficients (real for Hermitian matrices)
    KrylovSubspace{T,U}(n::Integer, maxiter::Integer=30) where {T,U} = new{real(T), T, U}(
        maxiter, maxiter, zero(real(T)), Matrix{T}(undef, n, maxiter + 1),
        fill(zero(U), maxiter + 1, maxiter))
    KrylovSubspace{T}(n::Integer, maxiter::Integer=30) where {T} = KrylovSubspace{T,T}(n, maxiter)
end
getH(Ks::KrylovSubspace) = @view(Ks.H[1:Ks.m + 1, 1:Ks.m])
getV(Ks::KrylovSubspace) = @view(Ks.V[:, 1:Ks.m + 1])
function Base.resize!(Ks::KrylovSubspace{B,T,U}, maxiter::Integer) where {B,T,U}
    V = Matrix{T}(undef, size(Ks.V, 1), maxiter + 1)
    H = fill(zero(U), maxiter + 1, maxiter)
    Ks.V = V; Ks.H = H
    Ks.m = Ks.maxiter = maxiter
    return Ks
end
function Base.show(io::IO, Ks::KrylovSubspace)
    println(io, "$(Ks.m)-dimensional Krylov subspace with fields")
    println(io, "beta: $(Ks.beta)")
    print(io, "V: ")
    println(IOContext(io, :limit => true), getV(Ks))
    print(io, "H: ")
    println(IOContext(io, :limit => true), getH(Ks))
end

# Helper functions that returns the real part if that is what is
# required (for Hermitian matrices), otherwise returns the value
# as-is.
coeff(::Type{T},α::T) where {T} = α
coeff(::Type{U},α::T) where {U<:Real,T<:Complex} = real(α)

#######################################
# Arnoldi/Lanczos with custom IOP
"""
    arnoldi(A,b[;m,tol,opnorm,iop]) -> Ks

Performs `m` anoldi iterations to obtain the Krylov subspace K_m(A,b).

The n x (m + 1) basis vectors `getV(Ks)` and the (m + 1) x m upper Hessenberg
matrix `getH(Ks)` are related by the recurrence formula

```
v_1=b,\\quad Av_j = \\sum_{i=1}^{j+1}h_{ij}v_i\\quad(j = 1,2,\\ldots,m)
```

`iop` determines the length of the incomplete orthogonalization procedure [^1].
The default value of 0 indicates full Arnoldi. For symmetric/Hermitian `A`,
`iop` will be ignored and the Lanczos algorithm will be used instead.

Refer to `KrylovSubspace` for more information regarding the output.

Happy-breakdown occurs whenver `norm(v_j) < tol * opnorm(A, Inf)`, in this case
the dimension of `Ks` is smaller than `m`.

[^1]: Koskela, A. (2015). Approximating the matrix exponential of an
advection-diffusion operator using the incomplete orthogonalization method. In
Numerical Mathematics and Advanced Applications-ENUMATH 2013 (pp. 345-353).
Springer, Cham.
"""
function arnoldi(A, b; m=min(30, size(A, 1)), tol=1e-7, opnorm=LinearAlgebra.opnorm, iop=0, cache=nothing )
    Ks = KrylovSubspace{eltype(b)}(length(b), m)
    arnoldi!(Ks, A, b; m=m, tol=tol, opnorm=opnorm, iop=iop)
end

"""
    arnoldi_step!(j, iop, n, A, V, H)

Take the `j`:th step of the Lanczos iteration.
"""
function arnoldi_step!(j::Integer, iop::Integer, A,
                       V::AbstractMatrix{T}, H::AbstractMatrix{U}) where {T,U}
    x,y = @view(V[:, j]),@view(V[:, j+1])
    mul!(y, A, x)
    @inbounds for i = max(1, j - iop + 1):j
        alpha = coeff(U, dot(@view(V[:, i]), y))
        H[i, j] = alpha
        axpy!(-alpha, @view(V[:, i]), y)
    end
    beta = norm(y)
    H[j+1, j] = beta
    @. y /= beta
    beta
end
"""
    arnoldi!(Ks,A,b[;tol,m,opnorm,iop]) -> Ks

Non-allocating version of `arnoldi`.
"""
function arnoldi!(Ks::KrylovSubspace{B, T, U}, A, b::AbstractVector{T};
                  tol::Real=1e-7, m::Int=min(Ks.maxiter, size(A, 1)),
                  opnorm=LinearAlgebra.opnorm, iop::Int=0, cache=nothing) where {B, T <: Number, U <: Number}
    if ishermitian(A)
        return lanczos!(Ks, A, b; tol=tol, m=m, opnorm=opnorm)
    end
    if m > Ks.maxiter
        resize!(Ks, m)
    else
        Ks.m = m # might change if happy-breakdown occurs
    end
    V, H = getV(Ks), getH(Ks)
    vtol = tol * opnorm(A, Inf)
    if iop == 0
        iop = m
    end
    # Safe checks
    n = size(V, 1)
    @assert length(b) == size(A,1) == size(A,2) == n "Dimension mismatch"
    # Arnoldi iterations (with IOP)
    fill!(H, zero(U))
    Ks.beta = norm(b)
    @. V[:, 1] = b / Ks.beta
    @inbounds for j = 1:m
        beta = arnoldi_step!(j, iop, A, V, H)
        if beta < vtol # happy-breakdown
            Ks.m = j
            break
        end
    end
    return Ks
end

"""
    lanczos_step!(j, m, n, A, V, H)

Take the `j`:th step of the Lanczos iteration.
"""
function lanczos_step!(j::Integer, A,
                       V::AbstractMatrix{T},
                       α::AbstractVector{U},
                       β::AbstractVector{B}) where {B,T,U}
    x,y = @view(V[:, j]),@view(V[:, j+1])
    mul!(y, A, x)
    α[j] = coeff(U, dot(x, y))
    axpy!(-α[j], x, y)
    j > 1 && axpy!(-β[j-1], @view(V[:, j-1]), y)
    β[j] = norm(y)
    @. y /= β[j]
    β[j]
end

macro diagview(A,d::Integer=0)
    s = d<=0 ? 1+abs(d) : :(m+$d)
    quote
        m = size($(esc(A)),1)
        @view($(esc(A))[($s):m+1:end])
    end
end

"""
    realview(R, V)

Returns a view of the real components of the complex vector `V`.
"""
realview(::Type{R}, V::AbstractVector{C}) where {R,C<:Complex} =
    @view(reinterpret(R, V)[1:2:end])
"""
    realview(R, V)

Returns a view of the real components of the real vector `V`.
"""
realview(::Type{R}, V::AbstractVector{R}) where {R} = V

"""
    lanczos!(Ks,A,b[;tol,m,opnorm]) -> Ks

A variation of `arnoldi!` that uses the Lanczos algorithm for
Hermitian matrices.
"""
function lanczos!(Ks::KrylovSubspace{B, T, U}, A, b::AbstractVector{T};
                  tol=1e-7, m=min(Ks.maxiter, size(A, 1)),
                  opnorm=LinearAlgebra.opnorm, cache=nothing) where {B, T <: Number, U <: Number}
    if m > Ks.maxiter
        resize!(Ks, m)
    else
        Ks.m = m # might change if happy-breakdown occurs
    end
    V, H = getV(Ks), getH(Ks)
    vtol = tol * opnorm(A, Inf)
    # Safe checks
    n = size(V, 1)
    @assert length(b) == size(A,1) == size(A,2) == n "Dimension mismatch"
    # Lanczos iterations
    fill!(H, zero(T))
    Ks.beta = norm(b)
    @. V[:, 1] = b / Ks.beta
    α = @diagview(H)
    # β is always real, even though α may (in general) be complex.
    β = realview(B, @diagview(H,-1))
    @inbounds for j = 1:m
        if vtol > lanczos_step!(j, A, V, α, β)
            # happy-breakdown
            Ks.m = j
            break
        end
    end
    copyto!(@diagview(H,1), β[1:end-1])
    return Ks
end


using LinearAlgebra
using FermiCG
using Printf
using Test

atoms = []
push!(atoms,Atom(1,"H",[0,0,0]))
push!(atoms,Atom(2,"H",[0,0,1]))
push!(atoms,Atom(3,"H",[0,0,2]))
push!(atoms,Atom(4,"H",[0,0,3]))
push!(atoms,Atom(5,"H",[0,0,4]))
push!(atoms,Atom(6,"H",[0,0,5]))
#basis = "6-31g"
basis = "sto-3g"

mol     = Molecule(0,1,atoms)
mf = FermiCG.pyscf_do_scf(mol,basis)
ints = FermiCG.pyscf_build_ints(mf.mol,mf.mo_coeff);
e_fci, d1_fci, d2_fci = FermiCG.pyscf_fci(ints,3,3)
# @printf(" FCI Energy: %12.8f\n", e_fci)

function compute_npairs(d2)
    npairs = 0
    for i in 1:size(d2)[1]
        for j in 1:size(d2)[1]
            npairs += d2[i,i,j,j]
        end
    end
    println(" NPairs:", round(npairs,digits=3))
end

compute_npairs(d2_fci)

FermiCG.pyscf_write_molden(mol,basis,mf.mo_coeff,filename="scf.molden")

# Cl = FermiCG.localize(mf.mo_coeff,"boys", mf)
# FermiCG.pyscf_write_molden(mol,"6-31g",Cl,filename="boys.molden")
C = mf.mo_coeff
rdm_mf = C[:,1:2] * C[:,1:2]'
Cl = FermiCG.localize(mf.mo_coeff,"lowdin",mf)
FermiCG.pyscf_write_molden(mol,basis,Cl,filename="lowdin.molden")
S = FermiCG.get_ovlp(mf)
U =  C' * S * Cl
println(" Build Integrals")
flush(stdout)
ints = FermiCG.orbital_rotation(ints,U)
println(" done.")
flush(stdout)
# ints = FermiCG.pyscf_build_ints(mf.mol,Cl);
# println(" done.")
# flush(stdout)

clusters    = [(1:2),(3:4),(5:6)]
#clusters    = [(1:4),(5:8),(9:12)]
init_fspace = [(1,1),(1,1),(1,1)]

clusters = [Cluster(i,collect(clusters[i])) for i = 1:length(clusters)]
display(clusters)

rdm1 = zeros(size(ints.h1))
rdm1a = rdm_mf*.5
rdm1b = rdm_mf*.5

#FermiCG.cmf_ci(ints, clusters, init_fspace, rdm1)
e_cmf, U = FermiCG.cmf_oo(ints, clusters, init_fspace, rdm1, verbose=0, gconv=1e-6)

@test isapprox(e_cmf, -3.205983033016, atol=1e-10)
C_cmf = Cl*U

FermiCG.pyscf_write_molden(mol,basis,C_cmf,filename="cmf.molden")


using CompScienceMeshes
using Test

fn = joinpath(dirname(@__FILE__),"mesa.msh")
m = CompScienceMeshes.read_gmsh_mesh(fn)

@test numcells(m) == 1132
@test numvertices(m) == 626
@test norm(vertices(m,200) - [0.454487,-0.0982151,-0.1]) < 1e-3


using Plots
gr(fmt=:png)
pyplot()

mcm = McCallModel()
V, U = solve_mccall_model(mcm)
U_vec = U .* ones(length(mcm.w_vec))

plot(mcm.w_vec,
     [V U_vec],
     lw=2,
     α=0.7,
     label=["V" "U"])


"""
The span of a dictionary is the set of all possible expansions in that
dictionary, with coefficient eltype determined uniquely by the Dictionary.

The span of a dictionary is a function space, mapping `S` to `T`. Here, `S` is
the domain type of the dictionary. The `T` is the codomain type.
"""
struct Span{S,T} <: FunctionSpace{S,T}
    dictionary  ::  Dictionary{S}
end

Span(dict::Dictionary{S}) where S = Span{S,span_codomaintype(dict)}(dict)

span_codomaintype(dict::Dictionary) =
    span_codomaintype(coefficienttype(dict), codomaintype(dict))
span_codomaintype(dict::Dictionary, coefficients) =
    span_codomaintype(eltype(coefficients), codomaintype(dict))
# - When the types are the same, that type is the result
span_codomaintype(::Type{T}, ::Type{T}) where {T <: Number} = T
# - the coefficient types are complex and the set itself is real
span_codomaintype(::Type{Complex{T}}, ::Type{T}) where {T <: Number} = Complex{T}
# - or vice-versa
span_codomaintype(::Type{T}, ::Type{Complex{T}}) where {T <: Number} = Complex{T}
# Default fallback
span_codomaintype(::Type{A}, ::Type{Z}) where {Z,A} = typeof(zero(A) * zero(Z))

coefficienttype(span::Span) = coefficienttype(dictionary(span))

# Convenient shorthand
coefficienttype = coefficienttype

dictionary(s::Span) = s.dictionary

similar(s::Span, ::Type{T}, dims) where {T} = Span(similar(dictionary(s), T, dims))

promote_domaintype(span::Span, ::Type{T}) where {T} = similar(span, T, size(span))

random_expansion(span::Span) = Expansion(dictionary(span), rand(dictionary(span)))

zero(span::Span) = Expansion(dictionary(span), zeros(dictionary(span)))

tensorproduct(s1::Span, s2::Span) = Span(tensorproduct(dictionary(s1), dictionary(s2)))

size(span::Span) = size(dictionary(span))
length(span::Span) = length(dictionary(span))

show(io::IO, mime::MIME"text/plain", s::Span) = composite_show(io, mime, s)
Display.displaystencil(s::Span) = ["Span(", dictionary(s), ")"]
show(io::IO, s::Span) = print(io, "Space : $(domaintype(s)) → $(codomaintype(s)) (span of $(repr(dictionary(s))))")


immutable SymTriangularDist <: ContinuousUnivariateDistribution
    μ::Float64
    s::Float64

    function SymTriangularDist(μ::Real, s::Real)
        s > zero(s) || error("SymTriangular: scale must be positive")
        @compat new(Float64(μ), Float64(s))
    end

    @compat SymTriangularDist(μ::Real) = new(Float64(μ), 1.0)
    SymTriangularDist() = new(0.0, 1.0)
end

@distr_support SymTriangularDist d.μ - d.s d.μ + d.s

#### Parameters

location(d::SymTriangularDist) = d.μ
scale(d::SymTriangularDist) = d.s

params(d::SymTriangularDist) = (d.μ, d.s)

#### Statistics

mean(d::SymTriangularDist) = d.μ
median(d::SymTriangularDist) = d.μ
mode(d::SymTriangularDist) = d.μ

var(d::SymTriangularDist) = d.s^2 / 6.0
skewness(d::SymTriangularDist) = 0.0
kurtosis(d::SymTriangularDist) = -0.6

entropy(d::SymTriangularDist) = 0.5 + log(d.s)

#### Evaluation

zval(d::SymTriangularDist, x::Float64) = (x - d.μ) / d.s
xval(d::SymTriangularDist, z::Float64) = d.μ + z * d.s

pdf(d::SymTriangularDist, x::Float64) = insupport(d, x) ? (1.0 - abs(zval(d, x))) / scale(d) : 0.0

logpdf(d::SymTriangularDist, x::Float64) = insupport(d, x) ? log((1.0 - abs(zval(d, x))) / scale(d)) : -Inf

function cdf(d::SymTriangularDist, x::Float64)
    (μ, s) = params(d)
    x <= μ - s ? 0.0 :
    x <= μ ? 0.5 * (1.0 + zval(d, x))^2 :
    x < μ + s ? 1.0 - 0.5 * (1.0 - zval(d, x))^2 : 1.0
end

function ccdf(d::SymTriangularDist, x::Float64)
    (μ, s) = params(d)
    x <= μ - s ? 1.0 :
    x <= μ ? 1.0 - 0.5 * (1.0 + zval(d, x))^2 :
    x < μ + s ? 0.5 * (1.0 - zval(d, x))^2 : 0.0
end

function logcdf(d::SymTriangularDist, x::Float64)
    (μ, s) = params(d)
    x <= μ - s ? -Inf :
    x <= μ ? loghalf + 2.0 * log1p(zval(d, x)) :
    x < μ + s ? log1p(-0.5 * (1.0 - zval(d, x))^2) : 0.0
end

function logccdf(d::SymTriangularDist, x::Float64)
    (μ, s) = params(d)
    x <= μ - s ? 0.0 :
    x <= μ ? log1p(-0.5 * (1.0 + zval(d, x))^2) :
    x < μ + s ? loghalf + 2.0 * log1p(-zval(d, x)) : -Inf
end

quantile(d::SymTriangularDist, p::Float64) = p < 0.5 ? xval(d, sqrt(2.0 * p) - 1.0) :
                                                       xval(d, 1.0 - sqrt(2.0 * (1.0 - p)))

cquantile(d::SymTriangularDist, p::Float64) = p > 0.5 ? xval(d, sqrt(2.0 * (1.0-p)) - 1.0) :
                                                        xval(d, 1.0 - sqrt(2.0 * p))

invlogcdf(d::SymTriangularDist, lp::Float64) = lp < loghalf ? xval(d, expm1(0.5*(lp - loghalf))) :
                                                              xval(d, 1.0 - sqrt(-2.0 * expm1(lp)))

invlogccdf(d::SymTriangularDist, lp::Float64) = lp > loghalf ? xval(d, sqrt(-2.0 * expm1(lp)) - 1.0) :
                                                               xval(d, -(expm1(0.5 * (lp - loghalf))))

function mgf(d::SymTriangularDist, t::Real)
    (μ, s) = params(d)
    a = s * t
    a == zero(a) && return one(a)
    4.0 * exp(μ * t) * (sinh(0.5 * a) / a)^2
end

function cf(d::SymTriangularDist, t::Real)
    (μ, s) = params(d)
    a = s * t
    a == zero(a) && return complex(one(a))
    4.0 * cis(μ * t) * (sin(0.5 * a) / a)^2
end

#### Sampling

rand(d::SymTriangularDist) = xval(d, rand() - rand())


using LinearAlgebra
using Plots
using Revise

includet("IterativeMethods.jl")
includet("InternalArray.jl")

# Definitions -----------------------------------------------------------------------------

# We take the reference system to be the position of the spring in rest.
# That is, for the one dimensional case: (l, 0).

"""
    SpringProblem

Contains the parameters associated with a typical spring problem.

# Fields
- `t0`: Initial time of the problem (defaults to 0)
- `p0`: Initial position
- `x0`: Initial speed

- `m`: Mass of the object
- `k`: Spring's constant
- `b`: Friction constant
- `A`: External force amplitude
- `ω`: External force angular speed
"""
struct SpringProblem
    t0
    p0
    v0

    m
    k
    b
    A
    ω

    SpringProblem(; t0 = 0, p0, v0, m, k, b = 0, A = 0, ω = 0) = new(t0, p0, v0, m, k, b, A, ω)
end

function eulermethod(prob::SpringProblem, step, nit)
    # We ought to transform the equation to a  differential equation of order one,
    # for which we write x = [position; speed].
    # That is, the variable x will represent the solution rather than the position.
    # And the solution contains the position and the speed.
    n = length(prob.p0)
    function f(t, x)
        p = x[1:n]
        v = x[n+1:2*n]
        return [v; (-prob.k*p -prob.b*v + [prob.A*sin(prob.ω*t)])/prob.m]
    end

    return eulermethod(f, prob.t0, [prob.p0; prob.v0], h = step, iter = nit)
end

function eulermethod_by_approximation(prob::SpringProblem, timelen; error = 1e-3)
    n = length(prob.p0)
    function f(t, x)
        p = x[1:n]
        v = x[n+1:2*n]
        return [v; (-prob.k*p -prob.b*v + [prob.A*sin(prob.ω*t)])/prob.m]
    end

    return eulermethod_by_approximation(f, prob.t0, timelen, [prob.p0; prob.v0]; error = error)
end


"""
# PureRandomSearch
## Constructor
```julia
PureRandomSearch(; draw, lb, ub)
```

Defines a pure random search method sampling from `draw`. It is up to the caller
to either supply a `draw` that takes zero input arguments and outputs an iterative
consistent with the objective function. If no `draw` method is supplied, the user
must provide lower bounds `lb` and upper bounds `ub` consistent with the objective
function.

## Example

"""

struct PureRandomSearch{T}
    draw::T
end

function PureRandomSearch(; draw=nothing, lb=nothing, ub=nothing)
    if isa(draw, Nothing)
        if isa(lb, Nothing) || isa(ub, Nothing)
            throw(ArgumentError("If you do not provide a draw, you need to provide lower and upper bounds. See ?PureRandomSearch for more details."))
        end
        width = ub .- lb
        T = eltype(width)
        N = length(width)
        _draw(xbest=nothing) = width .* rand(T, N) .+ lb
    else
        _draw = draw
    end
    PureRandomSearch(_draw)
end
function solve(prob, prs::PureRandomSearch, options)
    xbest = prs.draw()
    fbest = value(prob, xbest)
    for i = 1:options.maxiter
        xcandidate = prs.draw(xbest)
        fcandidate = value(prob, xcandidate)
        if fcandidate ≤ fbest
            fbest = fcandidate
            xbest = copy(xcandidate)
        end
    end
    return (fbest = fbest, xbest = xbest)
end


#!/usr/local/bin/julia

function binary_search(array::Array{Int64, 1}, key::Int64)
    bottom, top = 1, length(array)
    while bottom <= top
        middle = floor(Int64, (top + bottom)/2)
        if array[middle] < key
            bottom = middle + 1
        elseif array[middle] > key
            top = middle - 1
        else
            return true
        end
    end
    return false
end

# Sample run
# julia> @timev binary_search([1:10000000...], 1000011)
#   1.017336 seconds (10.00 M allocations: 310.125 MiB, 77.09% gc time)
# elapsed time (ns): 1017336317
# gc time (ns):      784260295
# bytes allocated:   325189632
# pool allocs:       9999500
# non-pool GC allocs:9
# malloc() calls:    2
# realloc() calls:   6
# GC pauses:         4
# full collections:  2
# true


"""
Example of a L2-regularized linearized regression
using the stochastic version of Frank-Wolfe.
"""

using FrankWolfe
using ProgressMeter
using Arpack
using Plots
using DoubleFloats
using ReverseDiff

using Random
using LinearAlgebra

using Test

# user-provided loss function and gradient

function simple_reg_loss(θ, data_point)
    (xi, yi) = data_point
    (a, b) = (θ[1:end-1], θ[end])
    pred = a ⋅ xi + b
    return (pred - yi)^2 / 2
end

function ∇simple_reg_loss(storage, θ, data_point)
    (xi, yi) = data_point
    (a, b) = (θ[1:end-1], θ[end])
    pred = a ⋅ xi + b
    @. storage[1:end-1] += xi * (pred - yi)
    storage[end] += pred - yi
    return storage
end

xs = [10 * randn(5) for i in 1:20000]
params = rand(6) .- 1 # start params in (-1,0)
bias = 4π
params_perfect = [1:5; bias]

data_perfect = [(x, x ⋅ (1:5) + bias) for x in xs]
f_stoch = FrankWolfe.StochasticObjective(simple_reg_loss, ∇simple_reg_loss, data_perfect, similar(params))

@test FrankWolfe.compute_value(f_stoch, params) > FrankWolfe.compute_value(f_stoch, params_perfect)

# Vanilla Stochastic Gradient Descent with reshuffling
storage = similar(params)
for idx in 1:1000
    for data_point in shuffle!(data_perfect)
        storage .= 0
        ∇simple_reg_loss(storage, params, data_point)
        params .-= 0.05 * storage / length(data_perfect)
    end
end

@test norm(params - params_perfect) <= 1e-6

# similar example with noisy data, Gaussian noise around the linear estimate
data_noisy = [(x, x ⋅ (1:5) + bias + 0.5 * randn()) for x in xs]
f_stoch_noisy = FrankWolfe.StochasticObjective(simple_reg_loss, ∇simple_reg_loss, data_noisy, storage)

params = rand(6) .- 1 # start params in (-1,0)

@testset "true parameters yield a good error" begin
    n1 = norm(FrankWolfe.compute_gradient(f_stoch_noisy, params_perfect))
    @test n1 <= length(data_noisy) * 0.05
end

# test that gradient at true parameters has lower norm than at randomly initialized ones
@test norm(FrankWolfe.compute_gradient(f_stoch_noisy, params_perfect)) <
      norm(FrankWolfe.compute_gradient(f_stoch_noisy, params))

# test that error at true parameters is lower than at randomly initialized ones
@test FrankWolfe.compute_value(f_stoch_noisy, params) >
      FrankWolfe.compute_value(f_stoch_noisy, params_perfect)

# Vanilla Stochastic Gradient Descent with reshuffling
for idx in 1:1000
    for data_point in shuffle!(data_perfect)
        storage .= 0
        params .-= 0.05 * ∇simple_reg_loss(storage, params, data_point) / length(data_perfect)
    end
end

# test that SGD converged towards true parameters
@test norm(params - params_perfect) <= 1e-6

#####
# Stochastic Frank Wolfe version
# We constrain the argument in the L2-norm ball with a large-enough radius

lmo = FrankWolfe.LpNormLMO{2}(1.05 * norm(params_perfect))

params0 = rand(6) .- 1 # start params in (-1,0)

k = 10000

@time x, v, primal, dual_gap, trajectoryS = FrankWolfe.stochastic_frank_wolfe(
    f_stoch_noisy,
    lmo,
    copy(params0),
    verbose=true,
    rng=Random.GLOBAL_RNG,
    line_search=FrankWolfe.Nonconvex(),
    max_iteration=k,
    print_iter=k / 10,
    batch_size=length(f_stoch_noisy.xs) ÷ 100 + 1,
    trajectory=true,
)

@time x, v, primal, dual_gap, trajectory09 = FrankWolfe.stochastic_frank_wolfe(
    f_stoch_noisy,
    lmo,
    copy(params0),
    momentum=0.9,
    verbose=true,
    rng=Random.GLOBAL_RNG,
    line_search=FrankWolfe.Nonconvex(),
    max_iteration=k,
    print_iter=k / 10,
    batch_size=length(f_stoch_noisy.xs) ÷ 100 + 1,
    trajectory=true,
)

@time x, v, primal, dual_gap, trajectory099 = FrankWolfe.stochastic_frank_wolfe(
    f_stoch_noisy,
    lmo,
    copy(params0),
    momentum=0.99,
    verbose=true,
    rng=Random.GLOBAL_RNG,
    line_search=FrankWolfe.Nonconvex(),
    max_iteration=k,
    print_iter=k / 10,
    batch_size=length(f_stoch_noisy.xs) ÷ 100 + 1,
    trajectory=true,
)

ff(x) = sum(simple_reg_loss(x, data_point) for data_point in data_noisy)
function gradf(storage, x)
    storage .= 0
    for dp in data_noisy
        ∇simple_reg_loss(storage, x, dp)
    end
end

@time x, v, primal, dual_gap, trajectory = FrankWolfe.frank_wolfe(
    ff,
    gradf,
    lmo,
    L=10,
    params,
    verbose=true,
    line_search=FrankWolfe.Adaptive(),
    max_iteration=k,
    print_iter=k / 10,
    trajectory=true,
)

data = [trajectory, trajectoryS, trajectory09, trajectory099]
label = ["exact", "stochastic", "stochM 0.9", "stochM 0.99"]

FrankWolfe.plot_trajectories(data, label)


function part1()
	lines = readlines("inputD10.txt")
	scores = Dict{Char, Int}(')' => 3, ']' => 57, '}' => 1197, '>' => 25137)
	opentoclose = Dict{Char, Char}('(' => ')', '[' => ']', '{' => '}', '<' => '>')

	score = 0

	for line in lines
		brackets = Char[]
		for char in line
			if char in keys(opentoclose)
				push!(brackets, char)
			else
				lastopen = pop!(brackets)
				if opentoclose[lastopen] != char
					score += scores[char]
					break
				end
			end
		end
	end

	return score
end

function part2()
	lines = readlines("inputD10.txt")
	opentoclose = Dict{Char, Char}('(' => ')', '[' => ']', '{' => '}', '<' => '>')

	closescores = Dict{Char, Int}(')' => 1, ']' => 2, '}' => 3, '>' => 4)
	linescores = Int[]
	for line in lines
		brackets = Char[]
		corrupted = false
		for char in line
			if char in keys(opentoclose)
				push!(brackets, char)
			else
				lastopen = pop!(brackets)
				if opentoclose[lastopen] != char
					corrupted = true
					break
				end
			end
		end

		if !corrupted
			linescore = 0
			for bracket in reverse(brackets)
				linescore *= 5
				linescore += closescores[opentoclose[bracket]]
			end
			push!(linescores, linescore)
		end
	end

	return sort(linescores)[div(length(linescores) + 1, 2)]
end

println(part1())
println(part2())


#---------------------------------------------------------
# # [NUFFT Overview](@id 1-nufft)
#---------------------------------------------------------

#=
This example illustrates how to use Nonuniform FFT (NUFFT)
for image reconstruction in MRI
using the Julia language.

This entire page was generated using a single Julia file:
[1-nufft.jl](@__REPO_ROOT_URL__/mri/1-nufft.jl).
=#
#md # In any such Julia documentation,
#md # you can access the source code
#md # using the "Edit on GitHub" link in the top right.

#md # The corresponding notebook can be viewed in
#md # [nbviewer](http://nbviewer.jupyter.org/) here:
#md # [`1-nufft.ipynb`](@__NBVIEWER_ROOT_URL__/mri/1-nufft.ipynb),
#md # and opened in [binder](https://mybinder.org/) here:
#md # [`1-nufft.ipynb`](@__BINDER_ROOT_URL__/mri/1-nufft.ipynb).

#=
Some MRI scans use non-Cartesian sampling patterns
(like radial and spiral k-space trajectories, among others),
often in the interest of acquisition speed.

Image reconstruction from fully sampled Cartesian k-space data
typically uses simple inverse FFT operations,
whereas non-Cartesian sampling requires more complicated methods.

The examples here illustrate both non-iterative (aka "gridding")
and iterative methods for non-Cartesian MRI reconstruction.
For simplicity the examples consider the case of single-coil data
and ignore the effects of B0 field inhomogeneity.

First we add the Julia packages that are need for these examples.
Change `false` to `true` in the following code block
if you are using any of the following packages for the first time.
=#

if false
    import Pkg
    Pkg.add([
        "ImagePhantoms"
        "Plots"
        "Unitful"
        "UnitfulRecipes"
        "LaTeXStrings"
        "MIRTjim"
        "MIRT"
        "InteractiveUtils"
    ])
end

# Now tell this Julia session to use the following packages for this example.
# Run `Pkg.add()` in the preceding code block first, if needed.

using ImagePhantoms: shepp_logan, SheppLoganEmis, spectrum, phantom #, Gauss2
using Plots; default(label="", markerstrokecolor=:auto)
using Unitful: mm # Allows use of physical units (mm here)
using UnitfulRecipes # Ensures that plots show appropriate units when applicable
using LaTeXStrings # for LaTeX in plot labels, e.g., L"\alpha_n"
using MIRTjim: jim, prompt # jiffy image display
using MIRT: Anufft, diffl_map, ncg
using InteractiveUtils: versioninfo

# The following line is helpful when running this jl-file as a script;
# this way it will prompt user to hit a key after each image is displayed.

isinteractive() && jim(:prompt, true);

# ## Radial k-space sampling

#=
We focus on radial sampling as a simple representative non-Cartesian case.
Consider imaging a 256mm × 256mm field of FOV
with the goal of reconstructing a 128 × 128 pixel image.
The following radial and angular k-space sampling is reasonable.
=#

N = 128
FOV = 256mm # physical units!
Δx = FOV / N # pixel size
kmax = 1 / 2Δx
kr = ((-N÷2):(N÷2)) / (N÷2) * kmax # radial sampling in k-space
Nϕ = 3N÷2 # theoretically should be about π/2 * N
kϕ = (0:Nϕ-1)/Nϕ * π # angular samples
νx = kr * cos.(kϕ)' # N × Nϕ k-space sampling in cycles/mm
νy = kr * sin.(kϕ)'
plot(νx, νy,
    xlabel=L"\nu_x", ylabel=L"\nu_y",
    aspect_ratio = 1,
    title = "Radial k-space sampling",
)

#
isinteractive() && prompt();

#=
For the NUFFT routines considered here,
the sampling arguments must be "Ω" values:
digital frequencies have pseudo-units of radians / pixel.
=#

Ωx = (2π * Δx) * νx # N × Nϕ grid of k-space sample locations
Ωy = (2π * Δx) * νy # in pseudo-units of radians / sample

scatter(Ωx, Ωy,
    xlabel=L"\Omega_x", ylabel=L"\Omega_y",
    xticks=((-1:1)*π, ["-π", "0", "π"]),
    yticks=((-1:1)*π, ["-π", "0", "π"]),
    xlims=(-π,π) .* 1.1,
    ylims=(-π,π) .* 1.1,
    aspect_ratio = 1, markersize = 0.5,
    title = "Radial k-space sampling",
)

#
isinteractive() && prompt();

# ## Radial k-space data for Shepp-Logan phantom

#=
Get the ellipse parameters for a MRI-suitable version of the Shepp-Logan phantom
and calculate (analytically) the radial k-space data.
Then display in polar coordinates.
=#

object = shepp_logan(SheppLoganEmis(); fovs=(FOV,FOV))
#object = [Gauss2(18mm, 0mm, 100mm, 70mm, 0, 1)] # useful for validating DCF
data = spectrum(object).(νx,νy)
data = data / oneunit(eltype(data)) # abandon units at this point
jim(kr, kϕ, abs.(data), title="k-space data magnitude",
    xlabel=L"k_r",
    ylabel=L"k_{\phi}",
    xticks = (-1:1) .* maximum(abs, kr),
    yticks = (0,π),
    ylims = (0,π),
    aspect_ratio = :none,
)

# ## Non-iterative gridding image reconstruction

#=
It would be impossible for a radiologist
to diagnose a patient
from the k-space data in polar coordinates,
so image reconstruction is needed.

The simplest approach
is to (nearest-neighbor) interpolate the k-space data
onto a Cartesian grid,
and then reconstruction with an inverse FFT.

One way to do that interpolation
is to use a `Histogram` method
in Julia's statistics package.
=#

using StatsBase: fit, Histogram, weights

#=
The following function is a work-around because `weights` in StatsBase is
[limited to Real data](https://github.com/JuliaStats/StatsBase.jl/issues/745),
so here we bin the real and imaginary k-space data separately,
and handle the units.
=#

function histogram(coord, vals::AbstractArray{<:Number}, edges)
    u = oneunit(eltype(vals))
    wr = weights(real(vec(vals / u)))
    wi = weights(imag(vec(vals / u)))
    tmp1 = fit(Histogram, coord, wr, edges)
    tmp2 = fit(Histogram, coord, wi, edges)
    return u * complex.(tmp1.weights, tmp2.weights)
end

kx = N * Δx * νx # N × Nϕ k-space sampling in cycles/mm
ky = N * Δx * νy
bin = (-(N÷2):(N÷2)) .- 0.5 # (N+1,) histogram bin edges
gridded1 = histogram((vec(kx), vec(ky)), data, (bin,bin))

using FFTW: ifft, fftshift
tmp = fftshift(ifft(fftshift(gridded1)))
x = (-(N÷2):(N÷2-1)) * Δx
y = (-(N÷2):(N÷2-1)) * Δx
jim(x, y, tmp, title="Elementary gridding reconstruction")

# That crummy gridding method does not work well.
# Here's what the phantom should look like:

ideal = phantom(x, y, object, 2)
clim = (0, 9)
p0 = jim(x, y, ideal, title="True Shepp-Logan phantom"; clim)

# ## NUFFT approach to gridding

#=
Basic nearest-neighbor gridding
does not provide acceptable image quality in MRI,
so now we turn to using the NUFFT.
For MRI purposes,
the NUFFT is a function that maps Cartesian spaced image data
into non-Cartesian k-space data.
The
[NFFT.jl](https://github.com/tknopp/NFFT.jl)
package has functions
for computing the NUFFT and its adjoint.
These are linear mappings (generalizations of matrices),
so instead of calling those functions directly,
here we use the NUFFT linear map object
defined in MIRT
that provides a non-Cartesian Fourier encoding "matrix".
=#

A = Anufft([vec(Ωx) vec(Ωy)], (N,N); n_shift = [N/2,N/2])

#=
This linear map is constructed to map a N × N image into `length(Ωx)` k-space samples.
So its
[adjoint](https://en.wikipedia.org/wiki/Adjoint)
goes the other direction.
However, an adjoint is *not* an inverse!
=#

gridded2 = A' * vec(data)
jim(x, y, gridded2, title="NUFFT gridding without DCF")

# ## Density compensation

#=
To get a decent image with NUFFT-based gridding of non-Cartesian data,
one must compensate for the k-space sampling density.
See
[this book chapter](https://web.eecs.umich.edu/~fessler/book/c-four.pdf)
for details.

Because this example uses radial sampling,
we can borrow ideas from tomography,
especially the
[ramp filter](https://en.wikipedia.org/wiki/Radon_transform#Radon_inversion_formula),
to define a reasonable density compensation function (DCF).

Here is a basic DCF version that uses the ramp filter in a simple way,
corresponding to the areas of annular segments
(Voronoi cells in polar coordinates).
The `dcf .* data` line uses Julia's
[broadcast](https://docs.julialang.org/en/v1/manual/arrays/#Broadcasting)
feature
to apply the 1D DCF to each radial spoke.
=#

dν = 1/FOV # k-space radial sample spacing
dcf = pi / Nϕ * dν * abs.(kr) # (N+1) weights along k-space polar coordinate
dcf = dcf / oneunit(eltype(dcf)) # kludge: units not working with LinearMap now
gridded3 = A' * vec(dcf .* data)
p3 = jim(x, y, gridded3, title="NUFFT gridding with simple ramp-filter DCF"; clim)

#=
The image is more reasonable than without any DCF,
but we can do better (quantitatively) using the correction of
[Lauzon&Rutt, 1996](http://doi.org/10.1002/mrm.1910360617)
and
[Joseph, 1998](http://doi.org/10.1002/mrm.1910400317).
=#

dcf = pi / Nϕ * dν * abs.(kr) # see lauzon:96:eop, joseph:98:sei
dcf[kr .== 0/mm] .= pi * (dν/2)^2 / Nϕ # area of center disk
dcf = dcf / oneunit(eltype(dcf)) # kludge: units not working with LinearMap now
gridded4 = A' * vec(dcf .* data)
p4 = jim(x, y, gridded4, title="NUFFT gridding with better ramp-filter DCF"; clim)

# A profile helps illustrate the improvement.

pp = plot(x, real(gridded4[:,N÷2]), label="Modified ramp DCF")
plot!(x, real(gridded3[:,N÷2]), label="Basic ramp DCF")
plot!(x, real(ideal[:,N÷2]), label="Ideal", xlabel=L"x", ylabel="middle profile")

#
isinteractive() && prompt();

#=
Finally we have made a NUFFT gridded image with DCF
that has the appropriate range of values,
but it still looks less than ideal.
So next we try an iterative approach.

## Iterative MR image reconstruction using NUFFT

As an initial iterative approach,
let's apply a few iterations of conjugate gradient (CG)
to seek the minimizer of the least-squares cost function:
```math
\arg \min_{x} \frac{1}{2} \| A x - y \|_2^2.
```
CG is well-suited to minimizing quadratic cost functions,
but we do not expect the image to be great quality
because radial sampling omits the "corners" of k-space
so the NUFFT operator ``A`` is badly conditioned.
=#

dx = FOV / N # pixel size
dx = dx / oneunit(dx) # abandon units for now
gradf = u -> u - vec(data/dx^2) # gradient of f(u) = 1/2 \| u - y \|^2
curvf = u -> 1 # curvature of f(u)
x0 = gridded4 # initial guess: best gridding reconstruction
xls, _ = ncg([A], [gradf], [curvf], x0; niter = 20)
p5 = jim(x, y, xls, "LS-CG reconstruction"; clim)

#=
To improve the results, we include regularization.
Here we would like to reconstruct an image
by finding the minimizer of a regularized LS cost function
such as the following:
```math
\arg \min_{x} \frac{1}{2} \| A x - y \|_2^2 + \beta R(x)
, \qquad
R(x) = 1' \psi.(T x).
```
We focus here on the case of edge-preserving regularization
where ``T`` is a 2D finite-differencing operator
and ``\psi`` is a potential function.
This operator maps a N×N image into a N×N×2 array
with the horizontal and vertical finite differences.
=#

T = diffl_map((N,N), [1,2] ; T = ComplexF32)

# Applying this operator to the ideal image illustrated its action:
jim(x, y, T * ideal; ncol=1, title="Horizontal and vertical finite differences")

#=
## Edge-preserving regularization

We use the Fair potential function:
a rounded corner version of absolute value,
an approximation to anisotropic total variation (TV).
=#

β = 2^13 # regularization parameter
δ = 0.05 # edge-preserving parameter
wpot = z -> 1 / (1 + abs(z)/δ) # weighting function

#=
## Nonlinear CG algorithm

We apply a (nonlinear) CG algorithm
to seek the minimizer of the cost function.
Nonlinear CG is well suited to convex problems
that are locally quadratic
like the regularized cost function considered here.
See
[this survey paper](http://doi.org/10.1109/MSP.2019.2943645)
for an overview of optimization methods for MRI.
=#

B = [A, T] # see MIRT.ncg
gradf = [u -> u - vec(data/dx^2), # data-term gradient, correct for pixel area
         u -> β * (u .* wpot.(u))] # regularizer gradient
curvf = [u -> 1, u -> β] # curvature of quadratic majorizers
x0 = gridded4 # initial guess is best gridding reconstruction
xhat, _ = ncg(B, gradf, curvf, x0; niter = 90)
p6 = jim(x, y, xhat, "Iterative reconstruction"; clim)

# Here is a comparison of the profiles.

plot!(pp, x, real(xls[:,N÷2]), label="LS-CG")
plot!(pp, x, real(xhat[:,N÷2]), label="Iterative edge-preserving", color=:black)

#
isinteractive() && prompt();

#=
In this case, iterative image reconstruction provides the best looking image.
One might argue this simulation was doomed to succeed,
because the phantom is piece-wise constant,
which is the best case for edge-preserving regularization.
On the other hand,
this was not an
[inverse crime](http://doi.org/10.1016/j.cam.2005.09.027)
([see also here](http://arxiv.org/abs/2109.08237))
because the k-space data came from the analytical spectrum of ellipses,
rather than from a discrete image.

### Caveats

* The phantom used here was real-valued, which is unrealistic
  (although the reconstruction methods did not "know" it was real).
* This simulation is for a single-coil scan
  whereas modern MRI scanners generally use multiple receive coils.
* There was no statistical noise in this simulation.
=#

# ## Reproducibility

# This page was generated with the following version of Julia:

io = IOBuffer(); versioninfo(io); split(String(take!(io)), '\n')

# And with the following package versions

import Pkg; Pkg.status()