# This file is designed to test the visibility interpolation algorithm.
import PyPlot
import PyPlot.plt
using LaTeXStrings

using DiskJockey.constants
using DiskJockey.gridding
using DiskJockey.visibilities
using DiskJockey.image

# Because of the flipped nature of the sky (but not flipped nature of the UV plane)
# there are some tricky conventions about how to pack the array.

# Given two arrays of l and m coordinates, fill an array of the Gaussian image
# following the MATLAB convention.
# p0 is a vector of [mu_RA, mu_DEC, sigma_x, sigma_y] in units of arcseconds
# mu_RA and mu_DEC are the locations of the centroid emission relative to the
# image origin (RA=0, DEC=0).
function imageGauss(ll::AbstractVector{Float64}, mm::AbstractVector{Float64}, p::Vector{Float64}, k::Int)

    # Both ll and mm increase with array index
    nx = length(ll)
    ny = length(mm)

    img = Array(Float64, ny, nx)
    mu = p[1:2] * arcsec #ll and mm shifts
    Sigma = Diagonal((p[3:4] * arcsec).^2) #Convert from arcsec to radians
    pre = 1. / (2pi * sqrt(det(Sigma))) * k
    for j=1:ny
        for i=1:nx
            R = Float64[ll[i] , mm[j]] - mu
            img[j, i] = pre * exp(-0.5 * (R' * (Sigma\R))[1])
        end
    end
    return img
end

# Given u and v coordinates in [kλ], evaluate the analytic FT of the
# aforementioned Gaussian
# N.B. Here Sigma refers to the (same) covariance matrix in the *image* domain
# always return a complex value
function FTGauss(uu::Float64, vv::Float64, p::Vector{Float64}, k::Int)
    uu = uu .* 1e3 #[λ]
    vv = vv .* 1e3 #[λ]
    mu_RA, mu_DEC = p[1:2]
    mu = Float64[mu_RA, mu_DEC] * arcsec #ll and mm shifts
    R = Float64[uu, vv]
    Sigma = Diagonal((p[3:4] * arcsec).^2) #Convert from arcsec to radians
    phase_shift = exp(-2pi * 1.0im * (R' * mu)[1]) # Not actually in polar phase form
    return k * exp(-2 * (pi^2) * (R' * Sigma * R)[1]) * phase_shift
    # in this case, Sigma serves as the inverse
end

# Given two arrays of u and v coordinates in [kλ], fill an array with the
# analytic FT of aforementioned Gaussian evaluated at every pairwise (u,v) pair
function FTGauss(uu::AbstractVector{Float64}, vv::AbstractVector{Float64}, p::Vector{Float64}, k::Int)
    nu = length(uu)
    nv = length(vv)
    # Both uu and vv increase with array index
    img = Array(Complex128, nv, nu)
    for j=1:nv
        for i=1:nu
            img[j, i] = FTGauss(uu[i], vv[j], p, k)
        end
    end
    return img
end

# Test to see if we get the convolutional interpolation correct by using a 2D
# Elliptical Gaussian.
# Realistic Gaussian will have scale dimensions (fatter in x direction)
const mu_RA = 2.0 # [arcsec]
const mu_DEC = -0.5 # [arcsec]
const s_x = 1.2 # [arcsec]
const s_y = 1.0 # [arcsec]
const p0 = [mu_RA, mu_DEC, s_x, s_y] # [arcsec]
const p_center = [0.0, 0.0, s_x, s_y]

# If the synthesized beam is 0.3", to oversample this by a factor of 10 would require something like 512x512 pixels
nx = 128
ny = 128

# full span of the image
ra = fftspace(10., nx) # [arcsec]
dec = fftspace(10., ny) # [arcsec]

# convert ra and dec in [arcsec] to radians, and then take the sin to convert to ll, mm
ll = sin(ra * arcsec) # direction cosines
mm = sin(dec * arcsec)

# The natural, shifted Gaussian image
img = imageGauss(ll, mm, p0, 1)

# A centered Gaussian image which we will then shift via the phase theorem
img_plain1 = imageGauss(ll, mm, p_center, 1)
img_plain2 = imageGauss(ll, mm, p_center, 1)
img_plain3 = imageGauss(ll, mm, p_center, 1)

lam0 = lam0s["12CO2-1"] # [microns] 12CO 2-1

skim = SkyImage(img, ra, dec, lam0)

skim_plain1 = SkyImage(img_plain1, ra, dec, lam0)
skim_plain2 = SkyImage(img_plain2, ra, dec, lam0)
skim_plain3 = SkyImage(img_plain3, ra, dec, lam0)

# Apply a centered correction function on the offset image and Fourier transform the image
corrfun!(skim, 0., 0.)
shift_fft = transform(skim)

# Now, create images that we will use as test cases for doing the interpolation before or after
# sampling

# Apply a centered correction function on the centered image (to phase shift downsampled
# visibilities later)
corrfun!(skim_plain1, 0, 0)
vis_fft_center = transform(skim_plain1)

# Apply an offset correction function on the centered image (then  phase shift the full visibilities)
corrfun!(skim_plain2, mu_RA, mu_DEC)
vis_fft_shift = transform(skim_plain2)
phase_shift!(vis_fft_shift, mu_RA, mu_DEC)

# Apply an offset correction function on the centered image, but later we will phase shift the visibilities *after* they have been downsampled.
corrfun!(skim_plain3, mu_RA, mu_DEC)
vis_fft_offset = transform(skim_plain3)

# This right here is not actually a fair comparison, because we haven't yet *interpolated* the FFTed values using the gridding convolution functions.

# Complex subtraction
# println("Maximum FFT discrepancy: ", maximum(abs(plain_fft.VV - vis_analytic)))
# println("Maximum FFT discrepancy: ", maximum(abs(shift_fft.VV - vis_analytic)))

# Return a normalized instance that is symmetric about 0
function scale(data)
    s = maximum(abs(data))
    return norm = plt[:Normalize](vmin=-s, vmax=s, clip=false)
end

# Basic plot of the shifted image

# Because the sky convention is different than the way the SkyImage is stored,
# we need to flip the array for plotting
fig, ax = plt[:subplots](nrows=2, figsize=(5, 8))

ext = (skim.ra[end], skim.ra[1], skim.dec[1], skim.dec[end])
ax[1][:imshow](flipdim(skim.data[:,:,1], 2), interpolation="none", origin="lower", cmap=plt[:get_cmap]("Greys"), extent=ext)
ax[1][:contour](flipdim(skim.data[:,:,1], 2), origin="lower", extent=ext)
ax[1][:set_title]("Sky Projection")
ax[1][:set_xlabel](L"$\alpha$ [arcsec]")
ax[1][:set_ylabel](L"$\delta$ [arcsec]")

ext = (ll[1], ll[end], mm[1], mm[end])
ax[2][:imshow](skim.data[:,:,1], interpolation="none", origin="lower", cmap=plt[:get_cmap]("Greys"), extent=ext)
ax[2][:contour](skim.data[:,:,1], origin="lower", extent=ext)
ax[2][:set_title]("Raw Array")
ax[2][:set_xlabel](L"$ll$")
ax[2][:set_ylabel](L"$mm$")

fig[:subplots_adjust](left=0.15, right=0.85, hspace=0.25)
plt[:savefig]("gaussian_img_array.png")

# Basic plot of the *centered* Gaussian image.
fig, ax = plt[:subplots](nrows=1, figsize=(5, 5))
# Real, analytic Gaussian
ext = (skim.ra[end], skim.ra[1], skim.dec[1], skim.dec[end])
aximg = ax[:imshow](flipdim(skim_plain1.data[:,:,1], 2), interpolation="none", origin="lower", cmap=plt[:get_cmap]("Greys"), extent=ext) #, norm = scale(img))
ax[:set_title]("image")
ax[:set_xlabel](L"$\alpha$ [arcsec]")
ax[:set_ylabel](L"$\delta$ [arcsec]")
#[left, bottom, width, height]
cax = fig[:add_axes]([0.84, 0.25, 0.03, 0.45])
cb = fig[:colorbar](aximg, cax=cax)

fig[:subplots_adjust](left=0.15, right=0.85, hspace=0.25)
plt[:savefig]("gaussian_img_center.png")

# It's nice to see these plots and what they look like, but it doesn't really make sense
# to compare them, unless we are comparing *interpolated* visibilities to the analytic formula.
# Just comparing the image is WRONG.

# THIS (following) IS THE ONLY THING WE SHOULD CARE ABOUT.

# Next, choose some (u,v) points within the bounds and see how the
# interpolated values compare to what the FTGauss would return.

nu = length(shift_fft.uu)
zer = zeros(nu)

function plot_1d(analytic, approx, fname)
    fig, ax = plt[:subplots](nrows=2, figsize=(5, 8))
    ax[1][:plot](shift_fft.uu, zer, ".k", label="Grid spacing")
    ax[1][:plot](uu, real(approx), "ob", label="Interp")
    ax[1][:plot](uu, real(analytic), ".r", label="Analytic")
    # ax[1][:plot](vis_fft.uu, real(analytic_u), "or", label="Analytic")
    ax[1][:set_xlim](-100, 100)
    ax[1][:set_title]("Real")
    ax[1][:set_xlabel](L"u [k $\lambda$]")
    ax[1][:legend]()

    ax[2][:plot](shift_fft.uu, zer, ".k", label="Grid spacing")
    ax[2][:plot](uu, imag(approx), "ob", label="Interp")
    ax[2][:plot](uu, imag(analytic), ".r", label="Analytic")
    # ax[2][:plot](vis_fft.uu, imag(analytic_u), "or", label="Analytic")
    ax[2][:set_xlim](-100, 100)
    ax[2][:set_title]("Imag")
    ax[2][:set_xlabel](L"u [k $\lambda$]")

    plt[:savefig](fname)
end

n = 100
uu = linspace(-100, 100, n) # [kλ]
approx = Array(Complex128, n)
analytic = Array(Complex128, n)

# First, let's see how the interpolated points, with *no shift*, correspond to the analytic form
# with no shift.

for i=1:n
    u = uu[i]
    v = 0.0
    #println("Interpolating at $u, $v")
    approx[i] = interpolate_uv(u, v, vis_fft_center)
    analytic[i] = FTGauss(u, v, p_center, 1)
end

plot_1d(analytic, approx, "interpolation.png")

# Next, let's see how the image shifted, then interpolated points compare with the
# analytic form

for i=1:n
    u = uu[i]
    v = 0.0
    #println("Interpolating at $u, $v")
    approx[i] = interpolate_uv(u, v, shift_fft)
    analytic[i] = FTGauss(u, v, p0, 1)
end

plot_1d(analytic, approx, "interpolation_raw_image_offset.png")

# Next,  let's see how the FFT'ed, then interpolated, then shifted points compare with the
# analytic form
mu = Float64[mu_RA, mu_DEC] * arcsec # [radians]

for i=1:n
    u = uu[i]
    v = 0.0

    # Now shift the interpolated points according to the phase theorem
    R = Float64[u, v] * 1e3 #[λ]
    # Not actually in polar phase form
    shift = exp(-2pi * 1.0im * (R' * mu)[1])

    approx[i] = shift * interpolate_uv(u, v, vis_fft_center)
    analytic[i] = FTGauss(u, v, p0, 1)
end

plot_1d(analytic, approx, "interpolation_shift_visibilities.png")

# Finally, let's check the way we've been doing things, which is to phase shift the full visibilities first, then downsample.

for i=1:n
    u = uu[i]
    v = 0.0
    #println("Interpolating at $u, $v")
    approx[i] = interpolate_uv(u, v, vis_fft_shift)
    analytic[i] = FTGauss(u, v, p0, 1)
end

plot_1d(analytic, approx, "interpolation_shift_image.png")

# As a final, final check that the resampling is in fact linear, let's use the centered image, offset corrfun, downsample this, then apply the phase shift at the end. This is where all my confusion started. If we were just able to do the previous two tests OK, it seems like this *shouldn't work*, because if in fact everything is linear, then the only difference is that this has the corrfun applied w/ a shift.

for i=1:n
    u = uu[i]
    v = 0.0

    # Now shift the interpolated points according to the phase theorem
    R = Float64[u, v] * 1e3 #[λ]
    # Not actually in polar phase form
    shift = exp(-2pi * 1.0im * (R' * mu)[1])

    approx[i] = shift * interpolate_uv(u, v, vis_fft_offset)
    analytic[i] = FTGauss(u, v, p0, 1)
end

plot_1d(analytic, approx, "interpolation_linear_operations.png")

# Ok, this *does not work*. What is the reason? Is the resampling operation not linear apparently?
# I don't think it's that, rather, I think that in order for the gridding correction function/ gridding convolution function pair to work properly, and properly cancel each other out, then they need to correspond to the same phase center. So, either you do both fixed to 0,0, or you do both offset to some value. Once the interpolation is done, you can phase shift as much as you want, but don't mix order of operations.

# Now, we can check things on an actual 2D image if we so please.

# Create analytic function on a smaller grid
n = 256
uu = linspace(-150, 150, n)
vv = linspace(-150, 150, n)

function plot_2d(analytic::Matrix{Complex128}, approx::Matrix{Complex128}, fname::AbstractString)

    println("Analyzing $fname")

    # Plot the reals first

    fig, ax = plt[:subplots](nrows=3, figsize=(5, 11))

    ext = (uu[1], uu[end], vv[1], vv[end])

    axan = ax[1][:imshow](real(analytic), interpolation="none", origin="lower", cmap=plt[:get_cmap]("Greys"), extent=ext)
    ax[1][:set_title]("Analytic FT")
    ax[1][:set_xlabel](L"uu [k$\lambda$]")
    ax[1][:set_ylabel](L"vv [k$\lambda$]")

    cax = fig[:add_axes]([0.84, 0.70, 0.03, 0.25])
    cb = fig[:colorbar](axan, cax=cax)

    axfft = ax[2][:imshow](real(approx), interpolation="none", origin="lower", cmap=plt[:get_cmap]("Greys"), extent=ext)
    ax[2][:set_title]("Interpolated Visibilites from FFT")
    ax[2][:set_xlabel](L"uu [k$\lambda$]")
    ax[2][:set_ylabel](L"vv [k$\lambda$]")

    cax = fig[:add_axes]([0.84, 0.40, 0.03, 0.25])
    cb = fig[:colorbar](axfft, cax=cax)

    diff = real(analytic - approx)

    # Measure the max value in the analytic array, just as an opportunity to set scale.
    println("Maximum analytic value ", maxabs(analytic))

    # Then, measure the total Mod squared difference across the whole image
    println("Total sqrt(modsquared) difference across image ", sqrt(sumabs2(diff)))

    axdif = ax[3][:imshow](diff, interpolation="none", origin="lower", cmap=plt[:get_cmap]("bwr"), extent=ext, norm=scale(diff))
    ax[3][:set_title]("Difference")
    ax[3][:set_xlabel](L"uu [k$\lambda$]")
    ax[3][:set_ylabel](L"vv [k$\lambda$]")

    cax = fig[:add_axes]([0.84, 0.10, 0.03, 0.25])
    cb = fig[:colorbar](axdif, cax=cax)

    fig[:subplots_adjust](hspace=0.25, top=0.97, bottom=0.06)

    plt[:savefig](fname * "_real.png")

    fig, ax = plt[:subplots](nrows=3, figsize=(5, 11))

    axan = ax[1][:imshow](imag(analytic), interpolation="none", origin="lower", cmap=plt[:get_cmap]("Greys"), extent=ext)
    ax[1][:set_title]("Analytic FT")
    ax[1][:set_xlabel](L"uu [k$\lambda$]")
    ax[1][:set_ylabel](L"vv [k$\lambda$]")

    cax = fig[:add_axes]([0.84, 0.70, 0.03, 0.25])
    cb = fig[:colorbar](axan, cax=cax)

    axfft = ax[2][:imshow](imag(approx), interpolation="none", origin="lower", cmap=plt[:get_cmap]("Greys"), extent=ext)
    ax[2][:set_title]("Interpolated Visibilites from FFT")
    ax[2][:set_xlabel](L"uu [k$\lambda$]")
    ax[2][:set_ylabel](L"vv [k$\lambda$]")

    cax = fig[:add_axes]([0.84, 0.40, 0.03, 0.25])
    cb = fig[:colorbar](axfft, cax=cax)

    diff = imag(analytic - approx)

    axdif = ax[3][:imshow](diff, interpolation="none", origin="lower", cmap=plt[:get_cmap]("bwr"), extent=ext, norm=scale(diff))
    ax[3][:set_title]("Difference")
    ax[3][:set_xlabel](L"uu [k$\lambda$]")
    ax[3][:set_ylabel](L"vv [k$\lambda$]")

    cax = fig[:add_axes]([0.84, 0.10, 0.03, 0.25])
    cb = fig[:colorbar](axdif, cax=cax)

    fig[:subplots_adjust](hspace=0.25, top=0.97, bottom=0.06)

    plt[:savefig](fname * "_imag.png")

    println()

end

# First, let's see how the interpolated points, with *no shift*, correspond to the analytic form
# with no shift.
approx = Array(Complex128, n, n)
analytic = FTGauss(uu, vv, p_center, 1) # can take in full arrays

for i=1:n
    for j=1:n
        approx[j,i] = interpolate_uv(uu[i], vv[j], vis_fft_center)
    end
end

plot_2d(analytic, approx, "2D_interpolation")

# Next, let's see how the image shifted, then interpolated points compare with the
# analytic form

analytic = FTGauss(uu, vv, p0, 1) # can take in full arrays

for i=1:n
    for j=1:n
        approx[j,i] = interpolate_uv(uu[i], vv[j], shift_fft)
    end
end

plot_2d(analytic, approx, "2D_interpolation_raw_image_offset")

# Next,  let's see how the FFT'ed, then interpolated, then shifted points compare with the
# analytic form

for i=1:n
    for j=1:n
        R = Float64[uu[i], vv[j]] * 1e3 #[λ]
        # Not actually in polar phase form
        shift = exp(-2pi * 1.0im * (R' * mu)[1])

        approx[j,i] = shift * interpolate_uv(uu[i], vv[j], vis_fft_center)
    end
end

plot_2d(analytic, approx, "2D_interpolation_shift_visibilities")

# Finally, let's check the way we've been doing things, which is to phase shift the full visibilities first, then downsample.

for i=1:n
    for j=1:n
        approx[j,i] = interpolate_uv(uu[i], vv[j], vis_fft_shift)
    end
end

plot_2d(analytic, approx, "2D_interpolation_shift_image.png")

fig, ax = plt[:subplots](nrows=2, figsize=(5, 8))

axan = ax[1][:imshow](abs(analytic), interpolation="none", origin="lower", cmap=plt[:get_cmap]("Greys"), extent=ext)
ax[1][:set_title]("Amplitude [Analytic]")
ax[1][:set_xlabel](L"uu [k$\lambda$]")
ax[1][:set_ylabel](L"vv [k$\lambda$]")

cax = fig[:add_axes]([0.84, 0.70, 0.03, 0.25])
cb = fig[:colorbar](axan, cax=cax)

axfft = ax[2][:imshow](angle(analytic), interpolation="none", origin="lower", cmap=plt[:get_cmap]("Greys"), extent=ext)
ax[2][:set_title]("Phase [Analytic]")
ax[2][:set_xlabel](L"uu [k$\lambda$]")
ax[2][:set_ylabel](L"vv [k$\lambda$]")

cax = fig[:add_axes]([0.84, 0.20, 0.03, 0.25])
cb = fig[:colorbar](axfft, cax=cax)

fig[:subplots_adjust](hspace=0.25, top=0.95, bottom=0.1)

plt[:savefig]("interpolation_phase.png")


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} <: AbstractYield
    value::N
    compounding::T
end

# make rate a broadcastable type
Base.Broadcast.broadcastable(ic::T) where {T<:Rate} = Ref(ic)

"""
    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))`. If not given a second argument, `Rate(rate)` is equivalent to `Rate(rate,Periodic(1))`.

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}
    c = convert(a.compounding,b)
    return (a.compounding.frequency == c.compounding.frequency) && isapprox(rate(a), rate(c); 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

"""
    discount(rate, t)
    discount(rate, from, to)

Discount `rate` for a time `t` or for an interval `(from, to)`. If `rate` is not a `Rate`, it will be assumed to be a `Periodic` rate compounded once per period, i.e. `Periodic(rate,1)`.

# Examples

```julia-repl
julia> discount(0.03, 10)
0.7440939148967249

julia> discount(Periodic(0.03, 2), 10)
0.7424704182237725

julia> discount(Continuous(0.03), 10)
0.7408182206817179

julia> discount(0.03, 5, 10)
0.8626087843841639
```
"""

discount(rate::T, t) where {T<:Real} = discount(Rate(rate), t)
discount(rate::Rate{<:Real, <:Continuous}, t) = exp(-rate.value * t)
discount(rate::Rate{<:Real, <:Periodic}, t) = (1 + rate.value / rate.compounding.frequency)^(-rate.compounding.frequency * t)
discount(rate, from, to) = discount(rate, to - from)

"""
    accumulation(rate, t)
    accumulation(rate, from, to)

Accumulate `rate` for a time `t` or for an interval `(from, to)`. If `rate` is not a `Rate`, it will be assumed to be a `Periodic` rate compounded once per period, i.e. `Periodic(rate,1)`.

    # Examples

```julia-repl
julia> accumulation(0.03, 10)
1.3439163793441222

julia> accumulation(Periodic(0.03, 2), 10)
1.3468550065500535

julia> accumulation(Continuous(0.03), 10)
1.3498588075760032

julia> accumulation(0.03, 5, 10)
1.1592740743
```
"""

accumulation(rate, t) = accumulation(Rate(rate), t)
accumulation(rate::Rate{<:Real, <:Continuous}, t) = exp(rate.value * t)
accumulation(rate::Rate{<:Real, <:Periodic}, t) = (1 + rate.value / rate.compounding.frequency)^(rate.compounding.frequency * t)
accumulation(rate, from, to) = accumulation(rate, to - from)

Base.zero(rate::T,t) where {T<:Rate} = rate
forward(rate::T,to) where {T<:Rate} = rate
forward(rate::T,from,to) where {T<:Rate} = rate

"""
    +(Yields.Rate, T<:Real)
    +(T<:Real, Yields.Rate)
    +(Yields.Rate,Yields.Rate)

The addition of a rate with a number will inherit the type of the `Rate`, or the first argument's type if both are `Rate`s.

# Examples

```julia-repl
julia> Yields.Periodic(0.01,2) + Yields.Periodic(0.04,2)
Yields.Rate{Float64, Yields.Periodic}(0.05000000000000004, Yields.Periodic(2))

julia> Yields.Periodic(0.04,2) + 0.01
Yields.Rate{Float64, Yields.Periodic}(0.05, Yields.Periodic(2))
```
"""
function Base.:+(a::Rate{N,T}, b::Real) where {N<:Real,T<:Continuous}
    return Continuous(a.value + b)
end
function Base.:+(a::Real, b::Rate{N,T}) where {N<:Real,T<:Continuous}
    return Continuous(b.value + a)
end

function Base.:+(a::Rate{N,T}, b::Real) where {N<:Real, T<:Periodic}
    return Periodic(a.value + b,a.compounding.frequency)
end
function Base.:+(a::Real, b::Rate{N,T}) where {N<:Real, T<:Periodic}
    return Periodic(b.value + a,b.compounding.frequency)
end

function Base.:+(a::Rate{N,T},b::Rate{N,T}) where {N<:Real,T<:CompoundingFrequency}
    a_rate = rate(a)
    b_rate = rate(convert(a.compounding,b))
    r = Rate(a_rate + b_rate,a.compounding)
    return r
end

"""
    -(Yields.Rate, T<:Real)
    -(T<:Real, Yields.Rate)
    -(Yields.Rate, Yields.Rate)

The addition of a rate with a number will inherit the type of the `Rate`, or the first argument's type if both are `Rate`s.

# Examples

```julia-repl
julia> Yields.Periodic(0.04,2) - Yields.Periodic(0.01,2)
Yields.Rate{Float64, Yields.Periodic}(0.030000000000000214, Yields.Periodic(2))

julia> Yields.Periodic(0.04,2) - 0.01
Yields.Rate{Float64, Yields.Periodic}(0.03, Yields.Periodic(2))

```
"""
function Base.:-(a::Rate{N,T}, b::Real) where {N<:Real,T<:Continuous}
    return Continuous(a.value - b)
end
function Base.:-(a::Real, b::Rate{N,T}) where {N<:Real,T<:Continuous}
    return Continuous( a - b.value)
end

function Base.:-(a::Rate{N,T}, b::Real) where {N<:Real, T<:Periodic}
    return Periodic(a.value - b, a.compounding.frequency)
end
function Base.:-(a::Real, b::Rate{N,T}) where {N<:Real, T<:Periodic}
    return Periodic(a - b.value, b.compounding.frequency)
end
function Base.:-(a::Rate{N,T},b::Rate{N,T}) where {N<:Real,T<:CompoundingFrequency}
    a_rate = rate(a)
    b_rate = rate(convert(a.compounding,b))
    r = Rate(a_rate - b_rate,a.compounding)
    return r
end

"""
    *(Yields.Rate, T<:Real)
    *(T<:Real, Yields.Rate)

The multiplication of a Rate with a scalar will inherit the type of the `Rate`, or the first argument's type if both are `Rate`s.
"""
function Base.:*(a::Rate{N,T}, b::Real) where {N<:Real,T<:Continuous}
    return Continuous(a.value * b)
end
function Base.:*(a::Real, b::Rate{N,T}) where {N<:Real,T<:Continuous}
    return Continuous( a * b.value)
end

function Base.:*(a::Rate{N,T}, b::Real) where {N<:Real, T<:Periodic}
    return Periodic(a.value * b, a.compounding.frequency)
end
function Base.:*(a::Real, b::Rate{N,T}) where {N<:Real, T<:Periodic}
    return Periodic(a * b.value, b.compounding.frequency)
end

"""
    /(Yields.Rate, T<:Real)
    /(T<:Real, Yields.Rate)

The division of a Rate with a scalar will inherit the type of the `Rate`, or the first argument's type if both are `Rate`s.
"""
function Base.:/(a::Rate{N,T}, b::Real) where {N<:Real,T<:Continuous}
    return Continuous(a.value / b)
end

# unclear if dividing a scalar by a rate should be allowed
# function Base.:/(a::Real, b::Rate{N,T}) where {N<:Real,T<:Continuous}
#     return Continuous( a / b.value)
# end

function Base.:/(a::Rate{N,T}, b::Real) where {N<:Real, T<:Periodic}
    return Periodic(a.value / b, a.compounding.frequency)
end

# unclear if dividing a scalar by a rate should be allowed
# function Base.:/(a::Real, b::Rate{N,T}) where {N<:Real, T<:Periodic}
#     return Periodic(a / b.value, b.compounding.frequency)
# end

"""
    <(Rate,Rate)

Convert the second argument to the periodicity of the first and compare the scalar rate values to determine if the first argument has a lower force of interest than the second.

# Examples

```julia-repl
julia> Yields.Periodic(0.03,100) < Yields.Continuous(0.03)
true
```
"""
function Base.:<(a::Rate{N,T},b::Rate{N,U}) where {N<:Real,T<:CompoundingFrequency,U<:CompoundingFrequency}
    bc = convert(a.compounding,b)
    return rate(a) < rate(bc)
end

"""
    >(Rate,Rate)

Convert the second argument to the periodicity of the first and compare the scalar rate values to determine if the first argument has a greater force of interest than the second.

# Examples

```julia-repl
julia> Yields.Periodic(0.03,100) > Yields.Continuous(0.03)
false
```
"""
function Base.:>(a::Rate{N,T},b::Rate{N,U}) where {N<:Real,T<:CompoundingFrequency,U<:CompoundingFrequency}
    bc = convert(a.compounding,b)
    return rate(a) > rate(bc)
end


let
  # testing constructor and pdf
  d = BoolDistribution(0.3)
  @test pdf(d, true) == 0.3
  @test pdf(d, false) == 0.7

  # testing support
  @test support(d) == [true, false]

  # testing ==
  d2 = BoolDistribution(0.3)
  @test d == d2

  # testing hash
  @test hash(d) == hash(d.p)
end


function u_test1(x, y)
    cos(x) * cos(y)
end

function u_test2(x, y)
    cos(x) - cos(y)
end

function mfunc_one!(u, m)
    m .= 1
end

function hfunc_u!(u, ϕ, h)
    h .= u
end

function hfunc_ϕ!(u, ϕ, h)
    h .= ϕ
end

# function testFokkerPlanckOperator()
#     M  = 64
#     N  = 64
#
#     grid = Grid2d(M,N)
#     u    = get_field(grid)
#     divJ = get_field(grid)
#     ft   = FourierTransform(grid)
#     op   = FokkerPlanckOperator(grid, ft)
#
#     evaluate_function_on_grid(grid, u_test, u)
#
#     collision_operator!(op, u, divJ)
#
# end

function testNullspace(utest, mfunc, hfunc)
    M  = 64
    N  = 64

    grid = Grid2d(M,N)
    u    = get_field(grid)
    divJ = get_field(grid)
    ft   = FourierTransform(grid)
    op   = FokkerPlanckOperator(grid, ft; mfunc=mfunc, hfunc=hfunc)

    evaluate_function_on_grid(grid, utest, u)

    collision_operator!(op, u, divJ)

    # println(maximum(abs.(divJ)))
    @test maximum(abs.(divJ)) ≈ zero(eltype(divJ)) atol=1E-9
    # @test maximum(abs.(divJ)) ≈ zero(eltype(divJ)) atol=1E-10

end

# testFokkerPlanckOperator()
testNullspace(u_test1, mfunc_one!, hfunc_ϕ!)
testNullspace(u_test2, mfunc_one!, hfunc_ϕ!)
testNullspace(u_test,  mfunc_one!, hfunc_u!)


const CONE_ORDER = Dict(
    MOI.Zeros=>1,
    MOI.Nonnegatives=>2,
    MOI.SecondOrderCone=>3,
    MOI.PositiveSemidefiniteConeTriangle=>4,
    MOI.ExponentialCone=>5,
    MOI.DualExponentialCone=>6,
    MOI.PowerCone=>7
)

const SUPPORTED_SOLVERS = Dict(
    "ECOS" => ECOS.Optimizer,
    "Hypatia" => Hypatia.Optimizer,
    "SCS" => SCS.Optimizer
)

# TODO: specify for each solver
SOLVER_OPTION_SET = [:max_iters, :eps, :warm_start, :scale, :verbose]

function solve_and_diff(
    A::AbstractMatrix{S}, b::AbstractVector{S}, c::Vector{S}, cone_prod::Vector{T}; kwargs...
) where {T <: MOI.AbstractVectorSet, S <: AbstractFloat}
    kwargs = Dict(kwargs)
    m,n = size(A)

    # Arguments check
    (length(b) == m && length(c) == n) || throw(DimensionMismatch("Mismatch between A, b, c"))
    all([typeof(c) <: SUPPORTED_INPUT_SETS for c in cone_prod]) || throw(ArgumentError("Unsupported cones"))
    all([MOI.dimension(c) > 0 for c in cone_prod]) || throw(ArgumentError("Cones must have nonzero dimension"))
    sum([MOI.dimension(c) for c in cone_prod]) == m || throw(DimensionMismatch("Mismatch between dimension of c and cone"))
    if haskey(kwargs, :solver) && kwargs[:solver] in keys(SUPPORTED_SOLVERS)
        optimizer = SUPPORTED_SOLVERS[kwargs[:solver]]
        # TODO: check cone_prod for each solver
    elseif haskey(kwargs, :solver)
        throw(ArgumentError("Invalid solver"))
    else
        optimizer = SCS.Optimizer
    end

    warm_start = haskey(kwargs, :warm_start) ? kwargs[:warm_start] : nothing
    if haskey(kwargs, :method) && kwars[:method] == "dense"
        use_lsqr = false
    else
        use_lsqr = true
    end

    options = Dict()
    for option in SOLVER_OPTION_SET
        if haskey(kwargs, option)
            options[option] = kwargs[option]
        end
    end
    options[:verbose] = haskey(options, :verbose) ? options[:verbose] : false

    return _solve_and_diff(A, b, c, cone_prod, optimizer, use_lsqr; options...)
end

function _solve_and_diff(A, b, c, cone_prod, optimizer, use_lsqr; options...)
    m,n = size(A)
    typeof(A) <: SparseMatrixCSC && dropzeros!(A)
    x_star, y_star, s_star = solve_opt_problem(A, b, c, cone_prod, optimizer; options...)

    Q = spzeros(m+n+1,m+n+1)
    Q[1:n,n+1:n+m]      .= A'
    Q[n+1:n+m,1:n]      .= -A
    Q[n+1:n+m,m+n+1]    .= b
    Q[m+n+1,n+1:n+m]    .= -b
    Q[1:n,m+n+1]        .= c
    Q[m+n+1,1:n]        .= -c

    # precompute M
    u, v, w = (x_star, y_star - s_star, 1.0)
    M = (Q - I) * dpi_z(u, v, w, cone_prod) + I

    function pullback(dx)
        # Assume that dy = ds = 0
        dz = [
            dx;
            zeros(length(y_star));
            -x_star' * dx
        ]

        g = use_lsqr ? lsqr(-M', dz, atol=1e-10, btol=1e-10) : -M' \ dz

        # TODO: make more efficient by only pulling needed cols/rows
        dQ = g * pi_z(u, v, w, cone_prod)'
        dA = dQ[1:n,n+1:n+m]'  - dQ[n+1:n+m,1:n]    # dQ12' - dQ21
        db = dQ[n+1:n+m,m+n+1] - dQ[m+n+1,n+1:n+m]  # dQ23 - dQ32'
        dc = dQ[1:n,m+n+1]     - dQ[m+n+1,1:n]      # dQ13 - dQ31'
        return dA, db, dc
    end

    function pushforward(dA, db, dc)
        # TODO: Could we compute a matrix vector product instead of full matrix?
        DπKdual_v = d_project_onto_cone(v, [MOI.dual_set(c) for c in cone_prod])

        dQ = spzeros(n+m+1, n+m+1)
        dQ[1:n,n+1:n+m]      .= dA'
        dQ[n+1:n+m,1:n]      .= -dA
        dQ[n+1:n+m,m+n+1]    .= db
        dQ[m+n+1,n+1:n+m]    .= -db #dont use tranpose bc of broadcast
        dQ[1:n,m+n+1]        .= dc
        dQ[m+n+1,1:n]        .= -dc

        g = dQ*pi_z(u, v, w, cone_prod)

        dz = use_lsqr ? lsqr(-M, g, atol=1e-10, btol=1e-10) : -M \ g
        @views du, dv, dw = dz[1:n], dz[n+1:n+m], dz[n+m+1]
        dx = du .- dw*x_star
        dy = DπKdual_v*dv - dw*y_star
        ds = DπKdual_v*dv - dv - dw*s_star
        return dx, dy, ds
    end

    return x_star, y_star, s_star, pushforward, pullback
end

function reorder_opt_problem_scs(A, b, cones)
    # sort the cones so that they're in the correct order
    sort_ixs = sortperm(cones, by=cone->CONE_ORDER[typeof(cone)])
    sorted_cones = cones[sort_ixs]

    f = 0  # the number of zero cones
    l = 0  # the number of nonneg cones
    q = Int[]  # the array of SOC sizes
    s = Int[]  # the array of SDC sizes
    ep = 0  # the number of exponential cones
    ed = 0  # the number of dual exponential cones
    p = Float64[]  # the array of power cone parameters

    # create the parameters for SCS
    for cone in sorted_cones
        if typeof(cone) <: MOI.Zeros
            f += MOI.dimension(cone)
        elseif typeof(cone) <: MOI.Nonnegatives
            l += MOI.dimension(cone)
        elseif typeof(cone) <: MOI.SecondOrderCone
            push!(q, MOI.dimension(cone))
        elseif typeof(cone) <: MOI.PositiveSemidefiniteConeTriangle
            push!(s, cone.side_dimension)
        elseif typeof(cone) <: MOI.ExponentialCone
            ep += 1
        elseif typeof(cone) <: MOI.DualExponentialCone
            ep += 1
        elseif typeof(cone) <: MOI.PowerCone
            push!(p, MOI.dimension(a))
        elseif typeof(cone) <: MOI.DualPowerCone
            push!(p, -MOI.dimension(a))
        end
    end
    cone_dict = Dict(:f=>f, :l=>l, :q=>q, :s=>s, :ep=>ep, :ed=>ed, :p=>p)

    original_ixs = []
    start_ix = 1
    for cone in cones
        end_ix = start_ix + MOI.dimension(cone) - 1
        push!(original_ixs, start_ix:end_ix)
        start_ix = end_ix + 1
    end

    index_map = [original_ixs[i] for i in sort_ixs]
    index_map = collect(Iterators.flatten(index_map))
    reverse_map = sortperm(index_map)

    return A[index_map, :], b[index_map], cone_dict, reverse_map
end

function solve_opt_problem(A, b, c, cone_prod, optimizer_factory; options...)
    # TODO: potentially play with the scale depending on the cones?
    #       e.g. see https://github.com/jump-dev/Convex.jl/issues/104
    #       Essentially want to have the primal and dual res down at the same rate
    #           primal slower -> inc. scale
    #           dual slower -> dec. scale
    m,n = size(A)

    if optimizer_factory == SCS.Optimizer
        A, b = rewrite_sdps_to_row_major(A, b, cone_prod)
        A_reordered, b_reordered, cone_dict, reverse_map = reorder_opt_problem_scs(A, b, cone_prod)
        f, l, q, s = cone_dict[:f], cone_dict[:l], cone_dict[:q], cone_dict[:s]
        ep, ed, p = cone_dict[:ep], cone_dict[:ed], cone_dict[:p]

        # TODO: clean up argument passing
        # TODO: for some reason this fails when b is a sparse vector -- see SCS.jl
        scs_sol = SCS_solve(SCS.IndirectSolver, m, n,
                            A_reordered, b_reordered, c,
                            f, l, q, s, ep, ed, p;
                            options...)
        # TODO: additional solution checks
        #/* SCS returns one of the following integers: */
        #define SCS_INFEASIBLE_INACCURATE (-7)
        #define SCS_UNBOUNDED_INACCURATE (-6)
        #define SCS_SIGINT (-5)
        #define SCS_FAILED (-4)
        #define SCS_INDETERMINATE (-3)
        #define SCS_INFEASIBLE (-2) /* primal infeasible, dual unbounded   */
        #define SCS_UNBOUNDED (-1)  /* primal unbounded, dual infeasible   */
        #define SCS_UNFINISHED (0)  /* never returned, used as placeholder */
        #define SCS_SOLVED (1)
        #define SCS_SOLVED_INACCURATE (2)
        if !(scs_sol.ret_val in [1,2])
            error("Model not solved correctly.
                   SCS termination status: $(scs_sol.ret_val)\n
                   Info: $(scs_sol.info)")
        elseif scs_sol.ret_val == 2
            @warn "Maximum iterations reached. Set max_iters to increase accuracy"
        end

        return scs_sol.x, scs_sol.y[reverse_map], scs_sol.s[reverse_map]
    end

    #TODO: interface directly with other optimizer solvers

    # If another optimizer is used, construct the problem with JuMP
    # NOTE: this has a high computational overhead
    model = Model()
    set_optimizer(model, optimizer_factory)
    # TODO: allow for attributes in general optimizer, e.g.
    # optimizer_with_attributes(SCS.Optimizer, "eps" => 1e-6, etc...)
    @variable(model, x[1:n])
    @variable(model, s[1:m])
    @objective(model, Min, c'*x)
    con = @constraint(model, A*x + s .== b)
    curr = 1
    for cone in cone_prod
        @constraint(model, s[curr:curr+MOI.dimension(cone)-1] in cone)
        curr += MOI.dimension(cone)
    end
    optimize!(model)
    if !(termination_status(model) in [MOI.OPTIMAL, MOI.ALMOST_OPTIMAL])
        error("Model not solved correctly.
               Termination status: $(termination_status(model))")
    elseif termination_status(model) == MOI.ALMOST_OPTIMAL
        @warn "Maximum iterations reached. Set max_iters to increase accuracy"
    end

    # Ensure dual has the correct sign
    y = isapprox(A'* dual.(con) + c, zeros(length(c)), atol=1e-5) ? dual.(con) : -dual.(con)
    return value.(x), y, value.(s)
end

function pi_z(u, v, w, cone_prod; dual=false)
    cone = dual ? [MOI.dual_set(c) for c in cone_prod] : cone_prod
    return vcat(
        u,
        project_onto_cone(v, cone),
        max(w, 0.0)
    )
end

function dpi_z(u, v, w, cone_prod; dual=false)
    cone = dual ? [MOI.dual_set(c) for c in cone_prod] : cone_prod
    return blockdiag(
        sparse(Matrix{Float64}(I, length(u), length(u))),
        d_project_onto_cone(v, cone),
        sparse(Matrix{Float64}(I, 1,1))
    )
end


module ReactiveMPModelsHGFTest

using Test, InteractiveUtils
using Rocket, ReactiveMP, GraphPPL, Distributions
using BenchmarkTools, Random, Plots, Dates, LinearAlgebra, StableRNGs

## Model definition
## -------------------------------------------- ##
@model [ default_factorisation = MeanField() ] function hgf_online_model(real_k, real_w, z_variance, y_variance)

    # Priors from previous time step for `z`
    zt_min_mean = datavar(Float64)
    zt_min_var  = datavar(Float64)

    # Priors from previous time step for `x`
    xt_min_mean = datavar(Float64)
    xt_min_var  = datavar(Float64)

    zt_min ~ NormalMeanVariance(zt_min_mean, zt_min_var)
    xt_min ~ NormalMeanVariance(xt_min_mean, xt_min_var)

    meta = GCVMetadata(GaussHermiteCubature(9))

    # Higher layer is modelled as a random walk
    zt ~ NormalMeanVariance(zt_min, z_variance) where { q = q(zt, zt_min)q(z_variance), meta = meta }

    # Lower layer is modelled with `GCV` node
    gcv_node, xt ~ GCV(xt_min, zt, real_k, real_w) where { q = q(xt, xt_min)q(zt)q(κ)q(ω) }

    # Noisy observations
    y = datavar(Float64)
    y ~ NormalMeanVariance(xt, y_variance)

    return zt, xt, y, gcv_node, xt_min_mean, xt_min_var, zt_min_mean, zt_min_var
end
## -------------------------------------------- ##
## Inference definition
## -------------------------------------------- ##
function hgf_online_inference(data, vmp_iters, real_k, real_w, z_variance, y_variance)
    n = length(data)

    # We don't want to save all marginals from all VMP iterations
    # but only last one after all VMP iterations per time step
    # Rocket.jl exports PendingScheduler() object that postpones
    # any update unless manual `resolve!()` has been called
    ms_scheduler = PendingScheduler()

    mz = keep(Marginal)
    mx = keep(Marginal)
    fe = ScoreActor(Float64)

    model, (zt, xt, y, gcv_node, xt_min_mean, xt_min_var, zt_min_mean, zt_min_var) = hgf_online_model(real_k, real_w, z_variance, y_variance)

    # Initial priors
    current_zt_mean, current_zt_var = 0.0, 10.0
    current_xt_mean, current_xt_var = 0.0, 10.0

    s_mz = subscribe!(getmarginal(zt) |> schedule_on(ms_scheduler), mz)
    s_mx = subscribe!(getmarginal(xt) |> schedule_on(ms_scheduler), mx)
    s_fe = subscribe!(score(Float64, BetheFreeEnergy(), model), fe)

    # Initial marginals to start VMP procedire
    setmarginal!(gcv_node, :y_x, MvNormalMeanCovariance([ 0.0, 0.0 ], [ 5.0, 5.0 ]))
    setmarginal!(gcv_node, :z, NormalMeanVariance(0.0, 5.0))

    # For each observations we perofrm `vmp_iters` VMP iterations
    for i in 1:n

        for _ in 1:vmp_iters
            update!(y, data[i])
            update!(zt_min_mean, current_zt_mean)
            update!(zt_min_var, current_zt_var)
            update!(xt_min_mean, current_xt_mean)
            update!(xt_min_var, current_xt_var)
        end

        # After all VMP iterations we release! `PendingScheduler`
        # as well as release! `ScoreActor` to indicate new time step
        release!(ms_scheduler)
        release!(fe)

        current_zt_mean, current_zt_var = mean_var(last(mz))::Tuple{Float64, Float64}
        current_xt_mean, current_xt_var = mean_var(last(mx))::Tuple{Float64, Float64}
    end

    # It is important to unsubscribe at the end of the inference procedure
    unsubscribe!((s_mz, s_mx, s_fe))

    return map(getvalues, (mz, mx, fe))
end
## -------------------------------------------- ##

@testset "Hierarchical Gaussian Filter" begin

    @testset "Online inference" begin
        ## -------------------------------------------- ##
        ## Data creation
        ## -------------------------------------------- ##
        function generate_data(rng, k, w, zv, yv)
            z_prev = 0.0
            x_prev = 0.0

            z = Vector{Float64}(undef, n)
            v = Vector{Float64}(undef, n)
            x = Vector{Float64}(undef, n)
            y = Vector{Float64}(undef, n)

            for i in 1:n
                z[i] = rand(rng, Normal(z_prev, sqrt(zv)))
                v[i] = exp(k * z[i] + w)
                x[i] = rand(rng, Normal(x_prev, sqrt(v[i])))
                y[i] = rand(rng, Normal(x[i], sqrt(yv)))

                z_prev = z[i]
                x_prev = x[i]
            end

            return z, x, y
        end
        ## -------------------------------------------- ##
        rng = StableRNG(42)
        # Parameters of HGF process
        real_k = 1.0
        real_w = 0.0
        z_variance = abs2(0.2)
        y_variance = abs2(0.1)
        # Number of observations
        n = 2000
        z, x, y = generate_data(rng, real_k, real_w, z_variance, y_variance);
        ## -------------------------------------------- ##
        ## Inference execution
        vmp_iters = 10
        mz, mx, fe = hgf_online_inference(y, vmp_iters, real_k, real_w, z_variance, y_variance);
        ## -------------------------------------------- ##
        ## Test inference results
        @test length(mz) === n
        @test length(mx) === n
        @test length(fe) === vmp_iters
        @test abs(last(fe) - 2027.8628798126442) < 0.01
        @test all(filter(e -> abs(e) > 0.1, diff(fe)) .< 0)
        # Check if all estimates are within 6std interval
        @test all((mean.(mz) .- 6 .* std.(mz)) .< z .< (mean.(mz) .+ 6 .* std.(mz)))
        @test all((mean.(mx) .- 6 .* std.(mx)) .< x .< (mean.(mx) .+ 6 .* std.(mx)))
        # Check if more than 95% of estimates are within 3std interval
        @test (sum((mean.(mz) .- 3 .* std.(mz)) .< z .< (mean.(mz) .+ 3 .* std.(mz))) / n) > 0.95
        @test (sum((mean.(mx) .- 3 .* std.(mx)) .< x .< (mean.(mx) .+ 3 .* std.(mx))) / n) > 0.95
        @test all(var.(mx) .> 0.0)
        @test all(var.(mz) .> 0.0)
        ## -------------------------------------------- ##
        ## Form debug output
        base_output = joinpath(pwd(), "_output", "models")
        mkpath(base_output)
        timestamp        = Dates.format(now(), "dd-mm-yyyy-HH-MM")
        plot_output      = joinpath(base_output, "hgf_online_plot_$(timestamp)_v$(VERSION).png")
        benchmark_output = joinpath(base_output, "hgf_online_benchmark_$(timestamp)_v$(VERSION).txt")
        ## -------------------------------------------- ##
        ## Create output plots
        pz = plot(title = "Hidden States Z")
        px = plot(title = "Hidden States X")

        plot!(pz, 1:n, z, label = "z_i", color = :orange)
        plot!(pz, 1:n, mean.(mz), ribbon = std.(mz), label = "estimated z_i", color = :teal)

        plot!(px, 1:n, x, label = "x_i", color = :green)
        plot!(px, 1:n, mean.(mx), ribbon = std.(mx), label = "estimated x_i", color = :violet)

        pf = plot(fe, label = "Bethe Free Energy")

        p = plot(pz, px, pf, layout = @layout([ a; b; c ]))
        savefig(p, plot_output)
        ## -------------------------------------------- ##
        ## Create output benchmarks
        benchmark = @benchmark hgf_online_inference($y, $vmp_iters, $real_k, $real_w, $z_variance, $y_variance) seconds=15
        open(benchmark_output, "w") do io
            show(io, MIME("text/plain"), benchmark)
            versioninfo(io)
        end
        ## -------------------------------------------- ##
    end

end

end


# https://adventofcode.com/2021/day/2
using AdventOfCode

input = readlines("2021/data/day_2.txt")

function part_1(input)
    h,d=0,0
    for command in input
        X = parse(Int, match(r"[a-z]* (?<num>\d+)", command)[:num])
        if occursin(r"forward.*",command) h+=X end
        if occursin(r"down.*",command) d+=X end
        if occursin(r"up.*",command) d-=X end
    end
    h*d
end

function part_2(input)
    h,d,aim=0,0,0
    for command in input
        X = parse(Int, match(r"[a-z]* (?<num>\d+)", command)[:num])
        if occursin(r"forward.*",command) h+=X;d+=X*aim end
        if occursin(r"down.*",command) aim+=X end
        if occursin(r"up.*",command) aim-=X end
    end
    h*d
end


# # Concrete casting heat of hydration
#
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/examples/Heat3D.ipynb)
#

using ReachabilityAnalysis, StructuralDynamicsODESolvers, MAT
using LinearAlgebra, SparseArrays
using LazySets.Arrays
using FastExpm

(@isdefined TARGET_FOLDER) ? nothing : TARGET_FOLDER = ""

# load internal (non-exported) functions
using ReachabilityAnalysis: ReachSolution, discretize, normalize, solve, center, step_size

include("../krylov_concept.jl")

# reduce default tolerance for vertex enumeration
LazySets.set_ztol(Float64, 1e-12)

using Plots
using Plots.PlotMeasures
using LaTeXStrings

using LazySets: center

function load_heat3d_matrices(; path=joinpath(@__DIR__, "Heat3D.mat"))
    vars = matread(path)

    K = vars["K"]
    C = vars["C"]
    timeIncr = vars["timeIncr"]
    fext = vars["fext"]
    qextTamb = vars["qextTamb"][:] # vector
    QhG =  vars["QhG"][:] # vector
    udots=vars["udots"]
    xs=vars["xs"]

    return K, C, qextTamb, QhG, timeIncr
end

function load_heat3d_params(; path=joinpath(@__DIR__, "params.mat"))
    vars       = matread(path)
    rho        = vars["rho"]
    m          = vars["m"]
    QFH        = vars["QFH"]
    Tambmin    = vars["Tambmin"]
    Tambvar    = vars["Tambvar"]
    Tambvarmin = vars["Tambvarmin"]
    Tambvarmax = vars["Tambvarmax"]
    Tini       = vars["Tini"]
    QFHmin     = vars["QFHmin"]
    QFHmax     = vars["QFHmax"]

    return rho, QFHmin, QFHmax, m, Tini, Tambmin, Tambvarmin, Tambvarmax, Tambvar, QFH
end

function load_heat3d_octaveSols(; path=joinpath(@__DIR__, "solsOctave.mat"))
    vars = matread(path)
    solAOctave = vars["Ts3DA"]
    solBOctave = vars["Ts3DB"]
    return solAOctave, solBOctave
end

# compute the maximum temperature at the given nodes
# by default we compute the max temperature in all the nodes,
# but you can specify a vector of indices to search;
# for example pass nodes=1:100 to make the search only on the first 100 nodes
function maxtemp(sol::ReachSolution; nodes=nothing)
    ## initialize search
    Tmax = -Inf; kmax = -1; imax = -1

    if isnothing(nodes)
        nodes = 1:dim(sol)
    end

    ## for each reach-set we compute the max bounds
    ## in every direction
    for (k, R) in enumerate(sol)
        X = set(R)
        q = center(X) + radius_hyperrectangle(X)
        θi = argmax(view(q, nodes))
        θ = q[θi]
        if θ > Tmax
            Tmax = θ
            kmax = k
            imax = θi
        end
    end

    δ = step_size(sol.alg)
    dtmax = tspan(sol[kmax])
    return (Tmax=Tmax, node=imax, step=kmax, dt=dtmax)
end

# compute the minimum temperature on the given nodes
function mintemp(sol::ReachSolution; nodes=nothing)
    ## initialize search
    Tmin = Inf; kmin = -1; imin = -1

    if isnothing(nodes)
        nodes = 1:dim(sol)
    end

    ## for each reach-set we compute the max bounds
    ## in every direction
    for (k, R) in enumerate(sol)
        X = set(R)
        q = center(X) - radius_hyperrectangle(X)
        θi = argmax(view(q, nodes))
        θ = q[θi]
        if θ < Tmin
            Tmin = θ
            kmin = k
            imin = θi
        end
    end

    δ = step_size(sol.alg)
    dtmin = tspan(sol[kmin])
    return (Tmin=Tmin, node=imin, step=kmin, dt=dtmin)
end

## Control nodes
controlNodeA = 11^2*6+1
controlNodeB = 11^2*9+1
nNodes = 11^3

function AhomCalc(C, K, qextTamb, QhG, rho, m)

    invC = inv( Matrix(C) );
    A = - invC * K |> Matrix
    n = size(A, 1)

    b0   =  invC * qextTamb ;
    bext =  invC * QhG * rho * m ;
    bsin =  invC * qextTamb ;

    omega = pi/12.;

    Ahom = [ A                    b0  bext  bsin      zeros(n) ;
             transpose( zeros(n))  0   0    0         0 ;
             transpose( zeros(n))  0  -m    0         0 ;
             transpose( zeros(n))  0   0    0         1 ;
             transpose( zeros(n))  0   0    -omega^2  0 ] ;

    return sparse(Ahom) # cast to sparse matrix
end

function solve_heat3d_rfem(; model=StepIntersect(Forward(inv=false)), X0=nothing, case=1)

    ## load vectors and params
    K, C, qextTamb, QhG, timeIncr = load_heat3d_matrices()
    rho, QFHmin, QFHmax, m, Tini, Tambmin, Tambvarmin, Tambvarmax, Tambvar, QFH = load_heat3d_params()
    usA, usB = load_heat3d_octaveSols()

    δ = timeIncr
    n = size(K, 1)

    NSTEPS = size(usA,2) - 1

    ## compute "homogeneized" system matrix
    Ahom = AhomCalc(C, K, qextTamb, QhG, rho, m)

    QFH     = (QFHmax+QFHmin)*0.5
    Tambvar = (Tambvarmax+Tambvarmin)*0.5

    if case == 1
        ## opcion singleton
        X0 = Singleton( Tini*ones(n)) ×
            Singleton((Tambmin+Tambvar*0.5)*ones( 1 ) )× Singleton( QFH*ones(1) ) ×
            Singleton((-Tambvar*0.5)*ones( 1 ) )× Singleton( zeros(1) )
        X0 = Singleton(LazySets.center(box_approximation(X0)))

    else
        ## opcion distribuida
        X0 = Singleton( Tini*ones(n)) ×
             Interval(Tambvarmin * 0.5 + Tambmin, Tambvarmax * 0.5 + Tambmin ) ×
             Interval( QFHmin, QFHmax ) ×
             Interval(-Tambvarmax * 0.5, -Tambvarmin*0.5) ×
             Singleton( zeros(1) )
        X0 = box_approximation(X0)
    end

    # set propagation method
    method = :sf
    println("Solving with set prop method (Case $case) $method")
    if method == :box
        prob = @ivp(x' = Ahom*x, x(0) ∈ X0)
        alg = BOX(δ=δ, approx_model=model)
        @time sol = solve(prob, alg=alg, NSTEPS=NSTEPS)

    elseif method == :sf
        @time sol = _solve_heat3d_sf(Ahom, X0, δ, NSTEPS, nodeidx=[controlNodeA, controlNodeB])
    end

    if case == 1
        X0 = Singleton(vcat(Tini*ones(n),
                      (Tambmin+Tambvar*0.5)*ones(1) ,
                      QFHmax*ones(1) ,
                      (-Tambvar*0.5)*ones(1),
                      zeros(1)))

        probOrbitMax = @ivp(x' = Ahom*x, x(0) ∈ X0)

        println("Solving with set prop method (Case 1) ORBIT")
        @time solOrbit = solve( probOrbitMax, NSTEPS=NSTEPS, alg=ORBIT(δ=δ))
        return sol, solOrbit

    elseif case == 2

        ## Orbit min
        X0 = Singleton(
              vcat( Tini*ones(n),
                    (Tambmin+Tambvarmin*0.5)*ones(1) ,
                    QFHmin*ones(1) ,
                    (-Tambvarmin*0.5)*ones(1),
                    zeros(1) ) ) ;

        probOrbit   = @ivp(x' = Ahom*x, x(0) ∈ X0)
        println("Solving with set prop method (Case 2) ORBIT min")
        @time solOrbitmin = solve( probOrbit, NSTEPS=NSTEPS, alg=ORBIT(δ=δ));

        ## Orbit min
        X0 = Singleton(
              vcat( Tini*ones(n),
                    (Tambmin+Tambvarmax*0.5)*ones(1) ,
                    QFHmax*ones(1) ,
                    (-Tambvarmax*0.5)*ones(1),
                    zeros(1) ) ) ;

        probOrbit = @ivp(x' = Ahom*x, x(0) ∈ X0)
        println("Solving with set prop method (Case 2) ORBIT max")
        @time solOrbitmax = solve( probOrbit, NSTEPS=NSTEPS, alg=ORBIT(δ=δ))
        return sol, solOrbitmin, solOrbitmax
   end

end

maxtemp_heat3d(sol; nodes=1:dim(sol), kwargs...) = maxtemp(sol, nodes=nodes)
maxtemp_heat3d(; nodes=nothing, kwargs...) = maxtemp(solve_heat3d(kwargs...), nodes=nodes)
mintemp_heat3d(sol; nodes=1:dim(sol), kwargs...) = mintemp(sol, nodes=nodes)
mintemp_heat3d(; nodes=nothing, kwargs...) = mintemp(solve_heat3d(kwargs...), nodes=nodes)

function solve_heat3d_implicit_euler()

    ## load vectors and params
    K, C, qextTamb, QhG, timeIncr = load_heat3d_matrices()
    rho, QFHmin, QFHmax, m, Tini, Tambmin, Tambvarmin, Tambvarmax, Tambvar, QFH = load_heat3d_params()
    usA, usB = load_heat3d_octaveSols()

    δ = timeIncr;  n = size(K, 1)
    NSTEPS = ( size(usA,2)-1 ) ;

    omega = pi/12.;

    M    = zeros(n, n)
    Tambvar = (Tambvarmin + Tambvarmax ) * 0.5 ;
    QFH = (QFHmin+ QFHmax)*0.5;

    R    = [ QhG*QFH*rho*m*exp(-m*δ*(i-1))+qextTamb*(Tambmin+Tambvar*(0.5 + 0.5*sin( omega*δ*(i-1)-pi/2.0 ))) for i in 1:NSTEPS+1 ] ;

    B   = Diagonal(ones(n));

    sys  =  SecondOrderConstrainedLinearControlContinuousSystem(M, Matrix(C), Matrix(K), B, nothing, R)

    U₀ = Tini*ones(n) ;

    prob = InitialValueProblem(sys, (U₀, U₀))
    alg  = BackwardEuler(Δt=δ)

    println("Solving with Backward Euler (Case 1)")
    @time solBackwardEulerSDCaseA = solve(prob, alg, NSTEPS=NSTEPS)

    R    = [ QhG*QFHmax*rho*m*exp(-m*δ*(i-1))+qextTamb*(Tambmin+Tambvarmax*(0.5 + 0.5*sin( omega*δ*(i-1)-pi/2.0 ))) for i in 1:NSTEPS+1 ] ;

    sys  =  SecondOrderConstrainedLinearControlContinuousSystem(M, Matrix(C), Matrix(K), B, nothing, R)
    prob = InitialValueProblem(sys, (U₀, U₀))

    println("Solving with Backward Euler (Case 2)")
    @time solBackwardEulerSDCaseB = solve(prob, alg, NSTEPS=NSTEPS)

    return solBackwardEulerSDCaseA, solBackwardEulerSDCaseB
end

function _solve_heat3d_sf(A, X0, δ, NSTEPS; nodeidx, mk=70, tol=1e-6)
    num_nodes = isa(nodeidx, Int) ? 1 : length(nodeidx)

    Eplus = _build_Eplus_optim(A, X0, δ; m=mk, tol=tol, cutoff=eps(Float64))

    dirs = Vector{LazySets.Arrays.SingleEntryVector{Float64}}()
    k = size(A, 1)
    for j in 1:num_nodes
        node = nodeidx[j]
        push!(dirs, SingleEntryVector(node, k, 1.0))
        push!(dirs, SingleEntryVector(node, k, -1.0))
    end
    cdirs = CustomDirections(dirs)

    propagacion = :expm_fast
    Ω0, Φᵀ = _discretize(A, δ, propagacion, X0, Eplus, mk, tol)

    Mout = Matrix{Float64}(undef, 2*num_nodes, NSTEPS)

    Aᵀδ = copy(transpose(A .* δ))

    for j in 1:num_nodes
        node = nodeidx[j]
        d = zeros(k)
        d[node] = 1.0
        println("Propagando indice $node ...")
        idx = (2(j-1) + 1):2j

        if isnothing(Φᵀ)
            _reach_homog_krylov_LGG09_modif!(view(Mout, idx, :), Ω0, Aᵀδ, d, NSTEPS; m=mk, tol=tol)
        else
            _reach_homog_dir_LGG09_modif!(view(Mout, idx, :), Ω0, Φᵀ, d, NSTEPS)
        end
    end
    sol = [TemplateReachSet(cdirs, view(Mout, :, k), (0.0 .. δ) + (k-1) * δ) for k in 1:NSTEPS] |> Flowpipe
end

function _discretize(A, δ, propagacion, X0::AbstractSingleton, Eplus, mk, tol)
    ΦX0 = _expv(A, X0.element, 1, δ, m=mk, tol=tol) |> Singleton
    Ω0 = CH(X0, minkowski_sum(ΦX0, Eplus))
    return Ω0, nothing
end

function _discretize(A, δ, propagacion, X0::Hyperrectangle, Eplus, mk, tol)
    if propagacion == :expm_base
        Φ = exp(Matrix(transpose(A) .* δ))
        Φ = sparse(Φ)
        droptol!(Φ, eps(Float64))
    elseif propagacion == :expm_fast
        Φ = real.(fastExpm(A .* δ, threshold=eps(Float64), nonzero_tol=eps(Float64)) )
    end
    Ω0 = CH(X0, Φ*X0 ⊕ Eplus)
    return Ω0, transpose(Φ)
end

function _reach_homog_dir_LGG09_modif!(Mout, Ω₀, Φᵀ, d, NSTEPS)
    rᵢ = copy(d)
    rᵢ₊₁ = similar(rᵢ)

    @inbounds for i in 1:NSTEPS
        Mout[1, i] = ρ(rᵢ, Ω₀)
        Mout[2, i] = ρ(-rᵢ, Ω₀)

        # update cache for the next iteration
        mul!(rᵢ₊₁, Φᵀ, rᵢ)
        copy!(rᵢ, rᵢ₊₁)
    end
    return nothing
end

#-----------------------------------------------------------
# Solutions LSD
solSDCase1, solSDCase2 = solve_heat3d_implicit_euler()

# Numerical solutions
tdom                     = times(solSDCase1)
solnodeSDCase1(node)     = displacements(solSDCase1, node)

tdom                     = times(solSDCase2)
solnodeSDCase2(node)     = displacements(solSDCase2, node)

#-----------------------------------------------------------
# Solutions Octave
solAOctave, solBOctave = load_heat3d_octaveSols()

#-----------------------------------------------------------
# Solutions SetBased

## Case 1
solRFEMCase1, solOrbitCase1 = solve_heat3d_rfem(case=1)

## Case 2
solRFEMCase2, solOrbitCase2min, solOrbitCase2max = solve_heat3d_rfem(case=2);

# Check maximum node temp
# @time maxtemp_heat3d(solRFEMCase2, nodes=1:1331)

#-----------------------------------------------------------
# Plots Case 1

timeMark = 200

indTimeMark = 155

tdom[indTimeMark]
solnodeSDCase1(controlNodeA)[indTimeMark]

fig = plot(tdom, solnodeSDCase1(controlNodeA),
           xlab=L"\textrm{Time (h)}",
           lc=:red, ylab=L"\textrm{Temperature } (^\circ\!C)",
           lab="",
           linestyle=:solid,
           legendfontsize=15, legend=:bottomright,
           tickfont=font(18, "Times"), guidefontsize=18,
           xtick=([0, 50, 100, 150, 200], [L"0", L"50", L"100", L"150", L"200"]),
           ytick=([20, 30, 40, 50, 60, 70, 80], [L"20", L"30", L"40", L"50", L"60", L"70", L"80"]),
           xguidefont=font(18, "Times"),
           yguidefont=font(18, "Times"),
           bottom_margin=6mm, left_margin=8mm, right_margin=4mm, top_margin=3mm,
           size=(900, 600))

plot!( tdom, solnodeSDCase1(controlNodeB), lc=:blue, lab="" )

# markers
plot!( [tdom[indTimeMark]], [solnodeSDCase1(controlNodeA)[indTimeMark]], marker=:utriangle, lc=:red, c=:red, lab=L"\textrm{Backward Euler at A}" )
plot!( [tdom[indTimeMark]], [solnodeSDCase1(controlNodeB)[indTimeMark]], marker=:circle, lc=:blue, c=:blue, lab=L"\textrm{Backward Euler at B}" )

plot!( fig, solRFEMCase1, vars=(0, controlNodeA), alpha=.2, lc=:magenta, c=:magenta, lab=L"\textrm{Flowpipe at A}")
plot!( fig, solRFEMCase1, vars=(0, controlNodeB), alpha=.2, lc=:green, c=:green, lab=L"\textrm{Flowpipe at B}")
#plot!(fig, tdom, solAOctave[controlNodeA,1:end], lc=:blue, lab="Backward Euler point A")

# subida
lens!(fig, [9, 12], [39, 43], inset = (1, bbox(0.15, 0.69, 0.28, 0.24)),
           subplot=5,
           tickfont=font(18, "Times"),
           xticks=([9.0, 10.5, 12.0], [ L"9",L"10.5", L"12"]),
           yticks=([39, 41, 43], [ L"39", L"41", L"43"])
)

lens!(fig, [50, 52], [66.5, 67], inset = (1, bbox(0.2, 0.38, 0.18, 0.18)),
           subplot=6,
           tickfont=font(18, "Times"),
           xticks=([50, 51, 52], [L"50", L"51", L"52"]),
           yticks=([66.5, 66.75, 67], [ L"66.5", "", L"67"])
)

lens!(fig, [50, 52], [83.3, 84.5], inset = (1, bbox(0.7, 0.1, 0.2, 0.2)),
           subplot=7,
           tickfont=font(18, "Times"),
           xticks=([50, 51, 52], [L"50", L"51", L"52"]),
           yticks=([83.5, 84.0, 84.5], [ L"83.5", "", L"84.5"])
)

savefig(fig, joinpath(TARGET_FOLDER, "heat3d_Case1.pdf"))

#-----------------------------------------------------------
# Plots Case 2

fig = plot(legendfontsize=16, legend=:topright,
           tickfont=font(18, "Times"), guidefontsize=18,
           xtick=([0, 50, 100, 150, 200], [L"0", L"50", L"100", L"150", L"200"]),
           ytick=([20, 30, 40, 50, 60, 70, 80], [L"20", L"30", L"40", L"50", L"60", L"70", L"80"]),
           xguidefont=font(18, "Times"),
           yguidefont=font(18, "Times"),
           bottom_margin=6mm, left_margin=8mm, right_margin=4mm, top_margin=3mm, size=(900, 600))

plot!(fig, solRFEMCase2, vars=(0, controlNodeA), alpha=.2, lc=:red, c=:red,
    xlab= L"\textrm{Time (h)}",
    ylab = L"\textrm{Temperature } (^\circ\!C)",
    lab=L"\textrm{Flowpipe at A}")#

#plot!(tdom, solnodeSDCase2(controlNodeA), xlab="Time", lc=:red, ylab="Temperature", lab="Backward Euler")
#plot!(fig, tdom, solBOctave[controlNodeA,1:end], lc=:blue, lab="Backward Euler point A")
#plot!( fig, solOrbitCase2max, vars=(0, controlNodeA), alpha=.2, lc=:red, lab="OrbitMax", setiestype=:path)
#plot!( fig, solOrbitCase2min, vars=(0, controlNodeA), alpha=.2, lc=:red, lab="OrbitMin", setiestype=:path)
plot!(fig,  solRFEMCase2, vars=(0, controlNodeB), alpha=.2, lc=:blue, c=:blue, lab=L"\textrm{Flowpipe at B}")#
plot!(fig, tdom, solnodeSDCase1(controlNodeA), lc=:red, lab=L"\textrm{Backward Euler at A}" )
plot!(fig, tdom, solnodeSDCase1(controlNodeB), lc=:blue, lab=L"\textrm{Backward Euler at B}")

savefig(fig, joinpath(TARGET_FOLDER, "heat3d_Case2.pdf"))

#=

julia>     Ahom = AhomCalc(C, K, qextTamb, QhG, rho, m)
1335×1335 Matrix{Float64}:
 -7.28091  0.77986  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  1.55972  …  0.0  0.0  0.0  0.0  0.0  0.0  4.16147  0.0082213   4.16147    0.0
  ⋮                                ⋮                        ⋮             ⋱  ⋮                        ⋮

  0.0      0.0      0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0         0.0  0.0  0.0  0.0  0.0  0.0  0.0      0.0        -0.0685389  0.0

julia> sparse(Ahom)
1335×1335 SparseMatrixCSC{Float64, Int64} with 10787 stored entries:
⠻⣦⡀⠀⠀⢸
⠀⠈⠻⣦⡀⢸
⠀⠀⠀⠈⠻⣾

julia> 1335 * 1335
1782225

julia> Ahom * Ahom;

julia> Ahom * Ahom;

julia> @time Ahom * Ahom;
  0.100113 seconds (2 allocations: 13.597 MiB)

julia> @time Ahom * Ahom;
  0.109501 seconds (2 allocations: 13.597 MiB)

julia> Ahom_sp = sparse(Ahom);

julia> @time Ahom_sp * Ahom_sp;
  0.133171 seconds (126.53 k allocations: 8.742 MiB, 99.41% compilation time)

julia> @time Ahom_sp * Ahom_sp;
  0.000643 seconds (7 allocations: 1.460 MiB)

=#


#=
""" `construct_gp_prior()`
Returns a GP prior using a Matern 5/2 kernel and specified parameters
Optional Inputs:
- smooth_factor: Multiplies variance and length scale for GP kernel (1)
- σ²: Variance for GP kernel (0.5)
- l: length scale for GP kernel (5.8e-5)
"""
function construct_gp_prior(; smooth_factor::Real = 1, σ²::Real = 0.5, l::Real = 5.8e-5)
	σ² *= smooth_factor
	l *= smooth_factor
	k = σ² * transform(Matern52Kernel(), 1 / l)
	f_naive = GP(k)
	#f = to_sde(f_naive)   # if develop issues with StaticArrays could revert to this
	f = to_sde(f_naive, SArrayStorage(Float64))
end
=#

""" `construct_gp_prior()`
Returns a GP prior using a Matern 5/2 kernel and specified parameters
Optional Inputs:
- smooth_factor_σ²: Multiplies variance for GP kernel (1)
- smooth_factor_l: Multiplies length scale for GP kernel (1)
- σ²: Variance for GP kernel (0.5)
- l: length scale for GP kernel (5.8e-5)
"""
function construct_gp_prior(; smooth_factor_σ²::Real = 1, smooth_factor_l::Real = 1, smooth_factor::Real = 1, σ²::Real = 0.5, l::Real = 5.8e-5)
	σ² *= smooth_factor_σ² * smooth_factor
	l *= smooth_factor_l * smooth_factor
	k = σ² * transform(Matern52Kernel(), 1 / l)
	f_naive = GP(k)
	#f = to_sde(f_naive)   # if develop issues with StaticArrays could revert to this
	f = to_sde(f_naive, SArrayStorage(Float64))
end

#=
function construct_gp_prior_kernel_param_as_vector(gp_param = [0.5, 5.8e-5 ])
	σ², l = gp_param
	k = σ² * transform(Matern52Kernel(), 1 / l)
	f_naive = GP(k)
	#f = to_sde(f_naive)   # if develop issues with StaticArrays could revert to this
	f = to_sde(f_naive, SArrayStorage(Float64))
end
=#

""" `construct_gp_posterior(xobs, yobs; sigmasq_obs, use_logx, use_logy, smooth_factor, boost_factor )`
Inputs:
- xobs: x locations where data is provided
- yobs: y values of data to condition on
Optional Inputs:
- sigmasq_obs: variances for y values being conditioned on
- use_logx: If true, take log's of x values before fitting GP
- use_logy: If true, perform log transform on y's
- smooth_factor: scales GP hyperparameters so as to result in smoother GP posterior (1)
- boost_factor: scales xobs by 1/boost_factor (1)
Returns:
- Posterior GP at locations xpred given training data
"""
function construct_gp_posterior(xobs::AA1, yobs::AA2; sigmasq_obs::AA4 = 1e-16*ones(length(xobs)),
			 	use_logx::Bool = true, #= use_logy::Bool = true, =# smooth_factor::Real = 1, boost_factor::Real = 1 )   where {
				T1<:Real, AA1<:AbstractArray{T1,1}, T2<:Real, AA2<:AbstractArray{T2,1}, T3<:Real, AA3<:AbstractArray{T3,1}, T4<:Real, AA4<:AbstractArray{T4,1} }
	f = construct_gp_prior( smooth_factor=smooth_factor )
	if boost_factor != one(boost_factor)
		xobs_trans = use_logx ? log.(xobs./boost_factor) : xobs./boost_factor
	else
		xobs_trans = use_logx ? log.(xobs) : xobs
	end
	fx = f(xobs_trans, sigmasq_obs)
	f_posterior = posterior(fx, yobs )
    #yobs_trans = use_logy ? log.(yobs) : yobs
    #sigmasq_obs_trans = use_logy ? sigmasq_obs./yobs.^2 : sigmasq_obs
	#fx = f(xobs_trans, sigmasq_obs_trans)
	#f_posterior = posterior(fx, yobs_trans )
end

""" predict_gp_mean_var(gp, xpred ; use_logx, use_logy)
Inputs:
- gp:
- xpred: Locations to predict GP at
Optional inputs:
- use_logx: If true, apply log transform to xpred before evaluating GP
- use_logy: If true, apply exp transform after evaluating GP
Returns tuple with vector of means and variance of GP posterior at locations in xpred.
"""
function predict_gp_mean_var(gp::AGP,  xpred::AA3, use_logx::Bool = true, use_logy::Bool = false )   where {
		AGP<:AbstractGP, T3<:Real, AA3<:AbstractArray{T3,1} }
		gp_post = use_log_x ? gp(log.(xpred)) : gp(xpred)
		m = mean.(gp_post)
		v = var.(gp_post)
		if use_logy
			m = exp.(m)
			v .*= m.^2
		end
		return m, v
end

""" predict_gp_mean(gp, xpred ; use_logx, use_logy)
Inputs:
- gp:
- xpred: Locations to predict GP at
Optional inputs:
- use_logx: If true, apply log transform to xpred before evaluating GP
Returns vector of means of GP posterior at locations in xpred.
"""
function predict_mean(gp::AGP )   where { AGP<:Fx_PosteriorType, T3<:Real, AA3<:AbstractArray{T3,1} }
	gp_post = marginals(gp)
	m = mean.(gp_post)
	return m
end

function predict_deriv(gp::AGP, xpred::AA3; use_logx::Bool = true )   where { AGP<:Fx_PosteriorType,
				T3<:Real, AA3<:AbstractArray{T3,1}  }
  #kobs = make_kernel_data(xobs, kernel=kernel, sigmasq_obs=sigmasq_obs, sigmasq_cor=sigmasq_cor, rho=rho)
  #kobs_pred_deriv = make_kernel_obs_pred(xobs,xpred, kernel=dkerneldx, sigmasq_cor=sigmasq_cor, rho=rho)
  #alpha = kobs \ yobs
  #pred_deriv = kobs_pred_deriv' * alpha
  xpred_trans = use_logx ? log.(xpred) : xpred
  m = predict_mean(gp) #,xpred_trans)
  #if use_logy   m .= exp(m)   end
  return numerical_deriv(xpred_trans,m)
  #=
  dfluxdlnλ = zeros(size(m))
  dfluxdlnλ[1] = (m[2]-m[1])/(xpred_trans[2]-xpred_trans[1])
  dfluxdlnλ[2:end-1] .= (m[3:end].-m[1:end-2])./(xpred_trans[3:end].-xpred_trans[1:end-2]) # exp.(m[2:end-1]).*
  dfluxdlnλ[end] = (m[end]-m[end-1])/(xpred_trans[end]-xpred_trans[end-1])
  return dfluxdlnλ
	=#
end

function predict_deriv2(gp::AGP, xpred::AA3; use_logx::Bool = true )   where { AGP<:Fx_PosteriorType,
				T3<:Real, AA3<:AbstractArray{T3,1}  }
	#kobs = make_kernel_data(xobs, kernel=kernel, sigmasq_obs=sigmasq_obs, sigmasq_cor=sigmasq_cor, rho=rho)
	#kobs_pred_deriv2 = make_kernel_obs_pred(xobs,xpred, kernel=d2kerneldx2, sigmasq_cor=sigmasq_cor, rho=rho)
	#alpha = kobs \ yobs
	#pred_deriv = kobs_pred_deriv2' * alpha

	xpred_trans = use_logx ? log.(xpred) : xpred
	m = predict_mean(gp)
	#if use_logy   m .= exp(m)   end
	d2fluxdlnλ2 = zeros(size(m))
	d2fluxdlnλ2[2:end-1] .= (m[3:end].+m[1:end-2].-2.0.*m[2:end-1])./(xpred_trans[3:end].-xpred_trans[1:end-2]).^2
	d2fluxdlnλ2[1] = d2fluxdlnλ2[2]
	d2fluxdlnλ2[end] = d2fluxdlnλ2[end-1]

	# TODO add some sort of check based on second derivative and dx that the finite difference is good enough?
	return d2fluxdlnλ2
end

function predict_mean(xobs::AA1, yobs::AA2, xpred::AA3;	sigmasq_obs::AA4 = 1e-16*ones(length(xobs)),
			use_logx::Bool = true, use_logy::Bool = false, smooth_factor::Real = 1, boost_factor::Real = 1, verbose::Bool = false  )   where {
				T1<:Real, AA1<:AbstractArray{T1,1}, T2<:Real, AA2<:AbstractArray{T2,1}, T3<:Real, AA3<:AbstractArray{T3,1}, T4<:Real, AA4<:AbstractArray{T4,1} }
	# global ncalls += 1
	@assert size(xobs) == size(yobs) == size(sigmasq_obs)
  	#println("# predict_mean (TemporalGPs): size(xobs) = ",size(xobs), "  size(xpred) = ", size(xpred))
	if verbose    println("# Entering predict_mean for TemporalGPs")    end
	tstart = now()
	if verbose    println("# extrema(yobs) = ", extrema(yobs) )   end
	mean_y = mean(yobs)
	if verbose    println("# mean_y = ", mean_y )    end
	y_trans = yobs
	sigmasq_obs_trans = sigmasq_obs
	if use_logy
		#println("# mean_y = ", mean_y, " min_y = ", minimum(yobs), " max_y = ", maximum(yobs))
		y_trans = log.(yobs./mean_y)
		sigmasq_obs_trans = sigmasq_obs./yobs.^2
	else
		y_trans = yobs./mean_y .- 1
		sigmasq_obs_trans = sigmasq_obs./mean_y.^2
	end
	if verbose    println("# extrema(xobs) = ", extrema(xobs) )    end
	if verbose    println("# extrema(y_trans) = ", extrema(y_trans) )    end
	if verbose    println("# extrema(sigmasq_obs_trans) = ", extrema(sigmasq_obs_trans) )   end
	if verbose    println("# smooth_factor = ", smooth_factor, " boost_factor = ",boost_factor )    end
	f_posterior = construct_gp_posterior(xobs,y_trans,sigmasq_obs=sigmasq_obs_trans, use_logx=use_logx, smooth_factor=smooth_factor, boost_factor=boost_factor )
	#println("typeof(f_posterior) = ",typeof(f_posterior))
	#println("f_posterior <: AbstractGP = ",typeof(f_posterior) <: AbstractGP )
	#println("f_posterior <: AbstractMvNormal = ",typeof(f_posterior) <: AbstractMvNormal )
	xpred_trans = use_logx ? log.(xpred) : xpred
	if verbose    println("# extrema(xpred_trans) = ", extrema(xpred_trans) )    end
	fx_posterior = f_posterior(xpred_trans)
	#println("typeof(fx_posterior) = ",typeof(fx_posterior))
	#println("f_posterior <: AbstractGP = ",typeof(fx_posterior) <: AbstractGP )
	#println("f_posterior <: AbstractMvNormal = ",typeof(fx_posterior) <: AbstractMvNormal )
	#println("f_posterior <: Fx_PosteriorType = ",typeof(fx_posterior) <: Fx_PosteriorType )
	#output = predict_mean(f_posterior(xpred_trans), xpred_trans ) #, use_logx=use_logx,use_logy=use_logy)
	pred_mean = predict_mean(fx_posterior)
	if verbose    println("# extrema(pred_mean) = ", extrema(pred_mean) )    end
	if use_logy
		pred_mean .= exp.(pred_mean).*mean_y
	else
		pred_mean .= (pred_mean.+1).*mean_y
	end
	if verbose    println("# extrema(pred_mean) = ", extrema(pred_mean) )    end
	#println("# predict_mean (TemporalGPs) runtime: ", now()-tstart)
	return pred_mean
end

# NEED TO TEST
function predict_deriv(xobs::AA1, yobs::AA2, xpred::AA3; sigmasq_obs::AA4 = 1e-16*ones(length(xobs))
						, use_logx::Bool = true, use_logy::Bool = false, smooth_factor::Real = 1, verbose::Bool = false  )    where { T1<:Real, AA1<:AbstractArray{T1,1}, T2<:Real, AA2<:AbstractArray{T2,1}, T3<:Real, AA3<:AbstractArray{T3,1}, T4<:Real, AA4<:AbstractArray{T4,1}  }
  #kobs = make_kernel_data(xobs, kernel=kernel, sigmasq_obs=sigmasq_obs, sigmasq_cor=sigmasq_cor, rho=rho)
  #kobs_pred_deriv = make_kernel_obs_pred(xobs,xpred, kernel=dkerneldx, sigmasq_cor=sigmasq_cor, rho=rho)
  #alpha = kobs \ yobs
  #pred_deriv = kobs_pred_deriv' * alpha
  if verbose   println("# predict_deriv (TemporalGPs): size(xobs) = ",size(xobs), "  size(xpred) = ", size(xpred))   end
  pred_mean = predict_mean(xobs,yobs,xpred, sigmasq_obs=sigmasq_obs, use_logx=use_logx, use_logy=use_logy, smooth_factor=smooth_factor)
  pred_deriv = calc_dfluxdlnlambda(pred_mean,xpred)
  return pred_deriv
  end

# NEED TO TEST
function predict_deriv2(xobs::AA1, yobs::AA2, xpred::AA3;sigmasq_obs::AA4 = 1e-16*ones(length(xobs))
						, use_logx::Bool = true, use_logy::Bool = false, smooth_factor::Real = 1  )   where { T1<:Real, AA1<:AbstractArray{T1,1}, T2<:Real, AA2<:AbstractArray{T2,1}, T3<:Real, AA3<:AbstractArray{T3,1}, T4<:Real, AA4<:AbstractArray{T4,1}  }
  #kobs = make_kernel_data(xobs, kernel=kernel, sigmasq_obs=sigmasq_obs, sigmasq_cor=sigmasq_cor, rho=rho)
  #kobs_pred_deriv2 = make_kernel_obs_pred(xobs,xpred, kernel=d2kerneldx2, sigmasq_cor=sigmasq_cor, rho=rho)
  #alpha = kobs \ yobs
  #pred_deriv = kobs_pred_deriv2' * alpha
println("# predict_deriv2 (TemporalGPs) runtime: ", now()-tstart)
  pred_mean = predict_mean(xobs,yobs,xpred, sigmasq_obs=sigmasq_obs, use_logx=use_logx, use_logy=use_logy, smooth_factor=smooth_factor)
	pred_deriv2 = calc_d2fluxdlnlambda2(pred_mean,xpred)
	return pred_deriv2
end

function predict_mean_and_deriv(xobs::AA1, yobs::AA2, xpred::AA3;sigmasq_obs::AA4 = 1e-16*ones(length(xobs))
	 							, use_logx::Bool = true, use_logy::Bool = false, smooth_factor::Real = 1, verbose::Bool = false ) where { T1<:Real, AA1<:AbstractArray{T1,1}, T2<:Real, AA2<:AbstractArray{T2,1}, T3<:Real, AA3<:AbstractArray{T3,1}, T4<:Real, AA4<:AbstractArray{T4,1}  }
	#		kernel::Function = matern52_sparse_kernel, dkerneldx::Function = dkerneldx_matern52_sparse,
	#		sigmasq_cor::Real=1.0, rho::Real=1.0)
  #kobs = make_kernel_data(xobs, kernel=kernel, sigmasq_obs=sigmasq_obs, sigmasq_cor=sigmasq_cor, rho=rho)
  #alpha = kobs \ yobs
  #kobs_pred = make_kernel_obs_pred(xobs,xpred, kernel=kernel, sigmasq_cor=sigmasq_cor, rho=rho)
  #pred_mean = kobs_pred' * alpha
  #kobs_pred_deriv = make_kernel_obs_pred(xobs,xpred, kernel=dkerneldx, sigmasq_cor=sigmasq_cor, rho=rho)
  #pred_deriv = kobs_pred_deriv' * alpha
  if verbose   println("# predict_mean_and_deriv (TemporalGPs): size(xobs) = ",size(xobs), "  size(xpred) = ", size(xpred))   end
  xpred_trans = use_logx ? log.(xpred) : xpred
  pred_mean = predict_mean(xobs,yobs,xpred, sigmasq_obs=sigmasq_obs, use_logx=use_logx, use_logy=use_logy, smooth_factor=smooth_factor)
  pred_deriv = numerical_deriv(xpred_trans,pred_mean)
  return (mean=pred_mean, deriv=pred_deriv)
end

function predict_mean_and_derivs(xobs::AA1, yobs::AA2, xpred::AA3; sigmasq_obs::AA4 = 1e-16*ones(length(xobs))
								, use_logx::Bool = true, use_logy::Bool = false, smooth_factor::Real = 1  ) where { T1<:Real, AA1<:AbstractArray{T1,1}, T2<:Real, AA2<:AbstractArray{T2,1}, T3<:Real, AA3<:AbstractArray{T3,1}, T4<:Real, AA4<:AbstractArray{T4,1}  }
	#		kernel::Function = matern52_sparse_kernel, dkerneldx::Function = dkerneldx_matern52_sparse, d2kerneldx2::Function = d2kerneldx2_matern52_sparse,
	#		,	sigmasq_cor::Real=1.0, rho::Real=1.0)
  #=
  kobs = make_kernel_data(xobs, kernel=kernel, sigmasq_obs=sigmasq_obs, sigmasq_cor=sigmasq_cor, rho=rho)
  alpha = kobs \ yobs
  kobs_pred = make_kernel_obs_pred(xobs,xpred, kernel=kernel, sigmasq_cor=sigmasq_cor, rho=rho)
  pred_mean = kobs_pred' * alpha
  kobs_pred_deriv = make_kernel_obs_pred(xobs,xpred, kernel=dkerneldx, sigmasq_cor=sigmasq_cor, rho=rho)
  pred_deriv = kobs_pred_deriv' * alpha
  kobs_pred_deriv2 = make_kernel_obs_pred(xobs,xpred, kernel=d2kerneldx2, sigmasq_cor=sigmasq_cor, rho=rho)
  pred_deriv2 = kobs_pred_deriv2' * alpha
  =#
  #println("# predict_mean_and_derivs (TemporalGPs): size(xobs) = ",size(xobs), "  size(xpred) = ", size(xpred))
  pred_mean = predict_mean(xobs,yobs,xpred, sigmasq_obs=sigmasq_obs, use_logx=use_logx, use_logy=use_logy, smooth_factor=smooth_factor)
  pred_deriv = calc_dfluxdlnlambda(pred_mean,xpred)
  pred_deriv2 = calc_d2fluxdlnlambda2(pred_mean,xpred)
  return (mean=pred_mean, deriv=pred_deriv, deriv2=pred_deriv2)
end

function log_pdf_gp_posterior(xobs::AA, yobs::AA #, kernel::Function;
			; sigmasq_obs::AA = 1e-16*ones(length(xobs)),
			use_logx::Bool = true, use_logy::Bool = false, smooth_factor::Real = 1  )  where { T<:Real, AA<:AbstractArray{T,1} }
  	#kobs = make_kernel_data(xobs, kernel=kernel, sigmasq_obs=sigmasq_obs, sigmasq_cor=sigmasq_cor, rho=rho)
  	# -0.5*( invquad(kobs, yobs) + logdet(kobs) + length(xobs)*log(2pi) )
	mean_y = mean(yobs)
	xobs_trans = use_logx ? log.(xobs) : xobs
	if use_logy
		yobs_trans = log.(yobs./mean_y)
		sigmasq_obs_trans = sigmasq_obs./yobs.^2
	else
		yobs_trans = yobs./mean_y .- 1
		sigmasq_obs_trans = sigmasq_obs./mean_y.^2
	end
	f_posterior = construct_gp_posterior(xobs_trans,yobs_trans,xobs_trans,sigmasq_obs=sigmasq_obs_trans, use_logx=false, smooth_factor=smooth_factor )
	return -logpdf(f_posterior(xobs_trans, sigmasq_obs_trans, use_logx=false), yobs_trans)
end

function fit_blaze_order(blaze::AbstractArray{T1,1}, order::Integer; mask::AbstractArray{Bool,1} = trues(length(blaze)),
						smooth_factor_σ²::Real = 1, smooth_factor_l::Real = 1 ) where { T1<:Real }
   gp_pr = construct_gp_prior(smooth_factor_σ²=smooth_factor_σ², smooth_factor_l=smooth_factor_l)
   x_obs = (1:length(blaze))[mask]
   num_pix_in_mask = length(x_obs)
   sigmasq_obs = var(blaze[mask][floor(Int,num_pix_in_mask//4):ceil(Int,num_pix_in_mask*3//4)])
   println("σ²_obs = ", sigmasq_obs)
   fx = gp_pr(x_obs, sigmasq_obs)
   f_post = posterior(fx, blaze[x_obs])
   cost = logpdf(f_post, blaze[x_obs])
   println("cost = ", cost)
   gp_mean = marginals(f_post)
   mean.(gp_post)
end


# ## Data

using Distributions
using Random

Random.seed!(1)
y = rand(Normal(0,1), 100);

# ## Model

using Turing

@model m8_4(y) = begin
    ## Can't really set a Uniform[-Inf,Inf] on σ
    α₁ ~ Uniform(-3000, 1000)
    α₂ ~ Uniform(-1000, 3000)
    σ ~ truncated(Cauchy(0,1), 0, Inf)

    y ~ Normal(α₁ + α₂, σ)
end

chns = sample(m8_4(y), NUTS(), 2000)

# ## Output

# \defaultoutput{}

# ## Original output

"""
         mean      sd     5.5%   94.5% n_eff Rhat
 a1    -861.15 558.17 -1841.89  -31.04     7 1.43
 a2     861.26 558.17    31.31 1842.00     7 1.43
 sigma    0.97   0.07     0.89    1.09     9 1.17
""";


using Test

@testset "Zero shaves to zero" begin
    for T in [Float16,Float32,Float64]
        for k in -5:50
            A = zeros(T,2,3)
            Ar = shave(A,k)
            @test A == Ar
            @test zero(T) == round(zero(T),k)
        end
    end
end

@testset "one shaves to one" begin
    for T in [Float16,Float32,Float64]
        for k in 0:50
            A = ones(T,2,3)
            Ar = shave(A,k)
            @test A == Ar
            @test one(T) == round(one(T),k)
        end
    end
end

@testset "minus one shaves to minus one" begin
    for T in [Float16,Float32,Float64]
        for k in 0:50
            A = -ones(T,2,3)
            Ar = shave(A,k)
            @test A == Ar
            @test -one(T) == round(-one(T),k)
        end
    end
end

@testset "No (half)shaving/setting/grooming for keepbits=10,23,52" begin
    for (T,k) in zip([Float16,Float32,Float64],
                        [10,23,52])
        A = rand(T,200,300)
        Ar = shave(A,k)
        @test A == Ar

        Ar = halfshave(A,k)
        @test A == Ar

        Ar = set_one(A,k)
        @test A == Ar

        Ar = groom(A,k)
        @test A == Ar

        # and a single one
        r = rand(T)
        @test r == shave(r,k)
        @test r == halfshave(r,k)
        @test r == set_one(r,k)
    end
end

@testset "Approx equal for keepbits=5,10,25" begin
    for (T,k) in zip([Float16,Float32,Float64],
                        [6,11,26])
        A = rand(T,200,300)
        Ar = shave(A,k)
        @test A ≈ Ar

        Ar = halfshave(A,k)
        @test A ≈ Ar

        Ar = set_one(A,k)
        @test A ≈ Ar

        Ar = groom(A,k)
        @test A ≈ Ar

        # and a single one
        r = rand(T)
        @test r ≈ shave(r,k)
        @test r ≈ set_one(r,k)
        @test r ≈ halfshave(r,k)
    end
end

@testset "Idempotence" begin
    for T in [Float16,Float32,Float64]
        for k in 0:20
            A = rand(T,200,300)
            Ar = shave(A,k)
            Ar2 = shave(Ar,k)
            @test Ar == Ar2

            Ar = halfshave(A,k)
            Ar2 = halfshave(Ar,k)
            @test Ar == Ar2

            Ar = set_one(A,k)
            Ar2 = set_one(Ar,k)
            @test Ar == Ar2

            Ar = groom(A,k)
            Ar2 = groom(Ar,k)
            @test Ar == Ar2
        end
    end
end

@testset "Shave/set = round towards/away from zero?" begin
    N = 1000
    for _ in 1:N
        for (T,UIntT) in zip([Float16,Float32,Float64],
                                [UInt16,UInt32,UInt64])
            for k in 1:9
                x = randn(T)
                xr = shave(x,k)
                @test abs(xr) <= abs(x)

                xr = set_one(x,k)
                @test abs(xr) >= abs(x)
            end
        end
    end
end


using Lora

function ess(x::Matrix{Float64})
    (p, T) = size(x)

    ESS = zeros(p)

    for i = 1:p
        ESS[i] = T*var(x[i,:])/Lora.mcvar_imse(squeeze(x[i,:],1))
    end

    ESS
end

function ess(chain, varname)

    theta_chain = zeros(length(chain[1][1][varname]), length(chain))
    for i = 1:length(chain)
        theta = chain[i][1][varname]

        theta_chain[:,i] = theta[:]

    end

    ess(theta_chain)

end

function ref_autocor(chain, ref_mean, ref_var)
    (p, T) = size(chain)

    centered_chain = broadcast(-,chain, ref_mean)

    max_lag = min(T-1, ceil(Integer, 100log10(T)))
    acor = zeros(p,max_lag)
    for i = 1:p
        for lag = 1:max_lag
            A = centered_chain[i, 1:T-lag]*centered_chain[i, lag+1:T]'
            @assert length(A) == 1
            acor[i,lag] = A[1]/((T-lag)*ref_var[i])
        end
    end

    acor
end

function ref_autocov(chain, ref_mean, ref_var)
    (p, T) = size(chain)

    centered_chain = broadcast(-,chain, ref_mean)

    max_lag = min(T-1, ceil(Integer, 100log10(T)))
    acov = zeros(p,max_lag)
    for i = 1:p
        for lag = 1:max_lag
            A = centered_chain[i, 1:T-lag]*centered_chain[i, lag+1:T]'
            @assert length(A) == 1
            acov[i,lag] = A[1]/(T)
        end
    end

    ref_var + 2*sum(acov,2)

end

function ref_ess_f(chain, reference_chain, f)
    c = [ f(chain[:,i][:])::Float64 for j in 1:1, i in 1:size(chain,2)]
    r = [ f(reference_chain[:,i][:])::Float64 for j in 1:1, i in 1:size(reference_chain,2)]

    ref_ess(c,r)
end

function ref_ess(chain, reference_chain)
    (p,T) = size(chain)
    ref_mean = mean(reference_chain,2)
    ref_var = var(reference_chain,2)

    rho = ref_autocor(chain, ref_mean, ref_var)
    (p, M) = size(rho)
    threshhold_acor = [ minimum([M; find(rho[i,:] .< 0.05).-1]) for i = 1:p]

    ESS = zeros(p)

    one_sT = 1 - collect(1:M)/T

    for i = 1:p
        Mt = threshhold_acor[i]
        ESS[i] = T/(1+2*dot(one_sT[1:Mt],rho[i,1:Mt][:]))
    end

    centered_chain = broadcast(-, chain, ref_mean)
    moment_2 = centered_chain.^2

    ref_centered = broadcast(-, reference_chain, ref_mean)
    ref_moment_2 = ref_centered.^2

    ref_m2_mean = mean(ref_moment_2, 2)
    ref_m2_var = var(ref_moment_2, 2)

    rho = ref_autocor(moment_2, ref_m2_mean, ref_m2_var)
    (p, M) = size(rho)
    threshhold_acor = [ minimum([M; find(rho[i,:] .< 0.05).-1]) for i = 1:p]

    ESS2 = zeros(p)

    one_sT = 1 .- collect(1:M)/T
    for i = 1:p
        Mt = threshhold_acor[i]
        ESS2[i] = T/(1+2*dot(one_sT[1:Mt],rho[i,1:Mt][:]))
    end

    ESS, ESS2
end

function compute_bias2_var(chain, true_mean, true_var, g)
    (p,T) = size(chain)
    gs_arr = [g(chain[:,i]) for i in 1:T]
    gs = [gs_arr[j][i]::Float64 for i in 1:length(gs_arr[1]), j in 1:T]

    est_mean = sum(gs,2)/T
    bias2 = (est_mean .- true_mean).^2
    bias = sqrt(bias2)

#    println("bias: $bias")
#    println("gs[1]: $(gs[1])")

    #debiased_gs = gs .- bias

    #chain_var = [Lora.mcvar_imse(squeeze(debiased_gs[i,:],1))::Float64 for i = 1:p]
    chain_var = ref_autocov(gs, true_mean + bias, true_var)

    println("chain_var: $chain_var")

    return sum(bias2), sum(chain_var)
end


# ---
# title: 190. Reverse Bits
# id: problem190
# author: Qling
# date: 2021-03-29
# difficulty: Easy
# categories: Bit Manipulation
# link: <https://leetcode.com/problems/reverse-bits/description/>
# hidden: true
# ---
#
# Reverse bits of a given 32 bits unsigned integer.
#
# **Note:**
#
#   * Note that in some languages such as Java, there is no unsigned integer type. In this case, both input and output will be given as a signed integer type. They should not affect your implementation, as the integer's internal binary representation is the same, whether it is signed or unsigned.
#   * In Java, the compiler represents the signed integers using [2's complement notation](https://en.wikipedia.org/wiki/Two%27s_complement). Therefore, in **Example 2**  above, the input represents the signed integer `-3` and the output represents the signed integer `-1073741825`.
#
# **Follow up** :
#
# If this function is called many times, how would you optimize it?
#
#
#
# **Example 1:**
#
#
#
#     Input: n = 00000010100101000001111010011100
#     Output:    964176192 (00111001011110000010100101000000)
#     Explanation: The input binary string **00000010100101000001111010011100** represents the unsigned integer 43261596, so return 964176192 which its binary representation is **00111001011110000010100101000000**.
#
#
# **Example 2:**
#
#
#
#     Input: n = 11111111111111111111111111111101
#     Output:   3221225471 (10111111111111111111111111111111)
#     Explanation: The input binary string **11111111111111111111111111111101** represents the unsigned integer 4294967293, so return 3221225471 which its binary representation is **10111111111111111111111111111111**.
#
#
#
#
# **Constraints:**
#
#   * The input must be a **binary string** of length `32`
#
#
## @lc code=start
using LeetCode

function reverse_bit(n::UInt32)::UInt32
    ret, power = 0, 31
    while n != 0
        ret += (n & 1) << power
        power -= 1
        n = n >> 1
    end

    return ret
end
## @lc code=end


function multitime(callable; numruns=20, seconds=2)
    benchmark_f = () -> @elapsed callable()
    runtime = benchmark_f()
    if runtime < seconds
        runtime = mapreduce(j -> benchmark_f(), min, 2:numruns; init = runtime)
    end
    return runtime
end

function chi2(gaussian_estimate, actual_value)
    diff = gaussian_estimate.μ - actual_value
    chi2_vals = diff' * (gaussian_estimate.Σ \ diff)
    return chi2_vals
end

function make_savable(sol; appxsol=nothing)
    # sol = solve(args...; kwargs...)
    @assert DiffEqBase.has_analytic(sol.prob.f) || !isnothing(appxsol)

    y_ref(t) =
        DiffEqBase.has_analytic(sol.prob.f) ?
        sol.prob.f.analytic(sol.prob.u0, sol.prob.p, t) :
        appxsol(t)
    densetimes = collect(range(sol.t[1], stop=sol.t[end], length=1000))
    return (
        x=sol.x,
        t=sol.t,
        pu=sol.pu,
        u_ref=y_ref.(sol.t),
        t_dense=densetimes,
        pu_dense=sol.p(densetimes),
        u_ref_dense=y_ref.(densetimes),
    )
end

function compute_errors(savable_sol)
    # savable_sol should be an output of `make_savable`
    ssol = savable_sol
    diffs = ssol.pu.μ - ssol.u_ref
    dense_diffs = ssol.pu_dense.μ - ssol.u_ref_dense

    errors = Dict{Symbol, Float64}()

    errors[:final] = norm(mean(abs.(diffs[end])))
    errors[:final_est] = norm(mean(abs.(sqrt.(diag(ssol.pu.Σ[end])))))

    errors[:l2] = norm(sqrt.(mean.([float.(d) .^ 2 for d in diffs])))
    errors[:l2_est] = norm(sqrt.(mean.(diag.(ssol.pu.Σ))))
    errors[:l∞] = norm(maximum.([float.(d) .^ 2 for d in diffs]))
    errors[:l∞_est] = norm(maximum.(diag.(ssol.pu.Σ)))

    errors[:L2] = norm(sqrt.(mean.([float.(d) .^ 2 for d in dense_diffs])))
    errors[:L2_est] = norm(sqrt.(mean.(diag.(ssol.pu_dense.Σ))))
    errors[:L∞] = norm(maximum.([float.(d) .^ 2 for d in dense_diffs]))
    errors[:L∞_est] = norm(maximum.(diag.(ssol.pu_dense.Σ)))

    errors[:final_χ²] = chi2(ssol.pu[end], ssol.u_ref[end])
    errors[:χ²] = mean(remove_initial_nan(chi2.(ssol.pu, ssol.u_ref)))
    errors[:Χ²] = mean(remove_initial_nan(chi2.(ssol.pu_dense, ssol.u_ref_dense)))

    return errors
end

function remove_initial_nan(arr)
    @assert issorted(.!isnan.(arr))
    return arr[.!isnan.(arr)]
end


type Ash{F <: Function, R <: Range} <: AbstractAsh
    kernel::F
    rng::R                      # x values where you want density y = f(x)
    v::Vector{Int64}            # v[i] is the count at x[i]
    y::Vector{Float64}          # y[i] is the density at x[i]
    m::Int64                    # smoothing parameter
    nobs::Int64
end
function Ash(
        kernel::Function, rng::Range, v::Vector{Int64},
        y::Vector{Float64}, m::Int64, n::Int64
    )
    @assert kernel(0.1) > 0.0 && kernel(-0.1) > 0.0 "kernel must always be positive"
    @assert length(rng) > 1 "Need at least two bins"
    @assert length(v) == length(y) == length(rng)
    @assert m > 0 "Smoothing parameter must be positive"
    Ash{typeof(kernel), typeof(rng)}(kernel, rng, v, y, m, n)
end
function Base.show(io::IO, o::Ash)
    println(io, typeof(o))
    println(io, "  > m         : ", o.m)
    println(io, "  > kernel    : ", o.kernel)
    println(io, "  > edges     : ", o.rng)
    println(io, "  > nobs      : ", StatsBase.nobs(o))
    plt = UnicodePlots.lineplot(xy(o)...; grid = false)
    show(io, plt)
end

function ash(
        x::Vector, rng::Range = extendrange(x);
        kernel::Function = Kernels.biweight, m::Int64 = 5, warnout::Bool = true
    )
    d = length(rng)
    v = zeros(Int64, d)
    y = zeros(d)
    o = Ash(kernel, rng, v, y, m, 0)
    update!(o, x)
    ash!(o; warnout = warnout)
end

# Return the histogram (as density)
histdensity(o::Ash) = o.v / step(o.rng) / StatsBase.nobs(o)
# StatsBase.nobs(o::Ash) = o.nobs
# nout(o::Ash) = nobs(o) - sum(o.v)
xy(o::Ash) = collect(o.rng), o.y
function extendrange(y::Vector)
    σ = std(y)
    linspace(minimum(y) - .5 * σ, maximum(y) + .5 * σ, 150)
end
Base.mean(o::Ash) = mean(o.rng, StatsBase.WeightVec(o.y))
Base.var(o::Ash) = var(o.rng, StatsBase.WeightVec(o.y))
Base.std(o::Ash) = sqrt(var(o))
function Base.quantile(o::Ash, τ::Real)
    @assert 0 < τ < 1
    rng = o.rng
    cdf = cumsum(o.y) * step(rng)
    i = searchsortedlast(cdf, τ)
    if i == 0
        rng[1]
    else
        rng[i] + (rng[i+1] - rng[i]) * (τ - cdf[i]) / (cdf[i+1] - cdf[i])
    end
end
function Base.quantile{T<:Real}(o::Ash, τ::AbstractVector{T})
    [quantile(o, τi) for τi in τ]
end
function Distributions.pdf(o::Ash, x::Real)
    rng = o.rng
    y = o.y
    i = searchsortedlast(rng, x)
    if 1 <= i < length(rng)
        y[i] + (y[i+1] - y[i]) * (x - rng[i]) / (rng[i+1] - rng[i])
    else
        0.0
    end
end
function Distributions.cdf(o::Ash, x::Real)
    cdf = cumsum(o.y) * step(o.rng)
    i = searchsortedlast(o.rng, x)
    if i == 0
        return 0.0
    else
        return cdf[i]
    end
end

# Add data
function update!(o::Ash, x::Vector)
    rng = o.rng
    δinv = 1.0 / step(rng)
    a = first(rng)
    nbin = length(rng)
    o.nobs += length(x)
    for xi in x
        # This line below is different from the paper because the input to UnivariateASH
        # is the points where you want the estimate, NOT the bin edges.
        ki = floor(Int, (xi - a) * δinv + 1.5)
        if 1 <= ki <= nbin
            @inbounds o.v[ki] += 1
        end
    end
    o
end

# Compute the density
function ash!(o::Ash; m::Int = o.m, kernel::Function = o.kernel, warnout::Bool = true)
    o.m = m
    o.kernel = kernel
    δ = step(o.rng)
    nbin = length(o.rng)
    @inbounds for k = 1:nbin
        if o.v[k] != 0
            for i = max(1, k - m + 1):min(nbin, k + m - 1)
                o.y[i] += o.v[k] * kernel((i - k) / m)
            end
        end
    end
    # make y integrate to 1
    y = o.y
    denom = 1.0 / (sum(y) * δ)
    for i in eachindex(y)
        @inbounds y[i] *= denom
    end
    warnout && (o.y[1] != 0 || o.y[end] != 0) && warn("nonzero density outside of bounds")
    o
end

function StatsBase.fit!(o::Ash, x::AbstractVector; kw...)
    update!(o, x)
    ash!(o; kw...)
end

@recipe function f(o::Ash)
    label --> ["Histogram Density" "Ash Density"]
    seriestype --> [:sticks :line]
    linewidth --> [1 2]
    alpha --> [.7 1]
    o.rng, [histdensity(o) o.y]
end


#=
Masking operations
==================

Copyright (c) 2017 Ben J. Ward & Luis Yanes

This Source Code Form is subject to the terms of the Mozilla Public
License, v. 2.0. If a copy of the MPL was not distributed with this
file, You can obtain one at http://mozilla.org/MPL/2.0/.
=#

"""
    nibble_mask{T<:Unsigned}(value::T, x::T)

Create a mask for the nibbles (aligned 4 bit segments) in an unsigned integer
`x` that filter nibbles matching the corresponding nibble in `value`.
"""
@inline function nibble_mask(value::T, x::T) where {T<:Unsigned}
    # XOR with the desired values. So matching nibbles will be 0000.
    x = x ⊻ value
    # Horizontally OR the nibbles.
    x |= (x >>> 1)
    x |= (x >>> 2)
    # AND removes junk, we then widen x by multiplication and return
    # the inverse.
    x &= repeatpattern(T, 0x11)
    x *= 0x0F
    return ~x
end

"""
    mask{T<:Unsigned}(::Type{T}, n::Integer)

Creates a bit mask for given number of bits `n`.

The mask starts from the least significant bit, and end at bit `n`.

e.g:

```jldoctest
julia> Twiddle.mask(UInt64, 8)
0x00000000000000ff
```
"""
@inline mask(::Type{T}, n::Integer) where {T<:Unsigned} = (T(1) << n) - 0x1
@inline mask(n::Integer) = mask(UInt64, n)


# Example from A Clique Graph Based Merging Strategy for Decomposable SDPs
# by Michael Garstka, Mark Cannon, Paul Goulart
#       See Figures 1 and 3
#       Note: This does not include chordal extension (already chordal)

# Setup sparsity pattern
n = 9
nonzero_inds = vcat([
        (3,1),
        (3,2),
        (5,4),
        (6,1),
        (6,3),
        (7,3),
        (7,6),
        (8,3),
        (8,4),
        (8,5),
        (8,6),
        (8,7),
        (9,6),
        (9,7),
        (9,8)
    ],
    [(i,i) for i in 1:n]
)
sp = sparse(CD.unzip(nonzero_inds)..., [1.0i+j for (i, j) in nonzero_inds])
sp = sp + tril(sp, -1)'

@testset "Clique Graph Construction" begin
    cliques_true = Set([
        Set([6, 7, 8, 9]),
        Set([6, 7, 8, 3]),
        Set([8, 4, 5]),
        Set([3, 6, 1]),
        Set([3, 2])
    ])
    final_cliques_true = Set([
        Set([3, 6, 7, 8, 9]),
        Set([8, 4, 5]),
        Set([3, 6, 1]),
        Set([3, 2])
    ])

    sp_full = sparse(Matrix(sp))
    CD.preprocess!(sp_full)
    @test !(0.0 in nonzeros(sp_full)) && all(sp_full .== sp)

    cliques = Set([Set(c) for c in get_cliques(tril(sp))])
    @test issetequal(cliques, cliques_true)

    cg = generate_clique_graph(cliques, n)
    cliques = Set([Set(c) for c in get_cliques(cg)])
    @test issetequal(cliques, cliques_true)

    merge_cliques!(cg; verbose=false)
    cliques = Set([Set(c) for c in get_cliques(cg)])
    @test issetequal(cliques, final_cliques_true)

    # TODO: test selectors

end


include("util.jl")

lines = getlines("inputs/6.txt")

function maxnum()::Int64
    maxnum = 0
    for line in lines
        for num in split(line, ",")
            maxnum = max(maxnum,parse(Int64,num))
        end
    end
    return maxnum
end

const FULL = maxnum()*4
const HALF = floor(maxnum()*2)

function get_coords(lines)::Dict{Symbol,Tuple{Int64,Int64}}
    retval = Dict()
    for line in 1:length(lines)
        coords = split(lines[line], ",")
        x = parse(Int64,coords[1])+HALF
        y = parse(Int64,coords[2])+HALF
        retval[Symbol(line)] = (x,y)
    end
    return retval
end

function closest(x::Int64, y::Int64, coords::Dict{Symbol,Tuple{Int64,Int64}})::Symbol
    distances::Dict{Symbol,Int64} = Dict()
    for s in keys(coords)
        distance = abs(x-coords[s][1])+abs(y-coords[s][2])
        distances[s] = distance
    end
    shortest = sort(collect(distances), by=x->x[2])
    if shortest[1][2] == shortest[2][2]
        return Symbol(".")
    else
        return shortest[1][1]
    end
end

function get_grid(coords::Dict{Symbol,Tuple{Int64,Int64}})::Array{Symbol,2}
    grid::Array{Symbol,2} = map(x -> Symbol("."), zeros(FULL,FULL))
    for i in 1:FULL
        for j in 1:FULL
            grid[j,i] = closest(i,j,coords)
        end
    end
    return grid
end

function total_distance(x::Int64,y::Int64,coords::Dict{Symbol,Tuple{Int64,Int64}})::Int64
    distances = []
    for p in coords
        d = abs(x-p[2][1])+abs(y-p[2][2])
        push!(distances, d)
    end
    return sum(distances)
end

function part_2_region_size(coords::Dict{Symbol,Tuple{Int64,Int64}})
    grid = zeros(FULL,FULL)
    for i in 1:FULL
        for j in 1:FULL
            td = total_distance(i,j,coords)
            if td < 10000
                grid[i,j] = 1
            end
        end
    end
    return grid
end

function largest_area(grid::Array{Symbol,2})
    is::Set{Symbol} = Set([Symbol(".")])
    is = union(is, Set(grid[1,:]))
    is = union(is, Set(grid[FULL,:]))
    is = union(is, Set(grid[:,1]))
    is = union(is, Set(grid[:,FULL]))
    legal = setdiff(Set(grid),is)
    sizes = []
    for s in legal
        g = map(x -> begin
            if x == s
                1
            else
                0
            end
        end, grid)
        push!(sizes, sum(g))
    end
    return sizes
end

coords = get_coords(lines)
grid = get_grid(coords)
a = largest_area(grid)
maximum(a)
sum(part_2_region_size(coords))


using Gadfly, DataArrays

set_default_plot_size(6inch, 3inch)

spy(randn((10, 10)), Coord.cartesian(fixed=true))